Mask-Enhanced Autoregressive Prediction: Pay Less Attention to Learn More

📄 arXiv: 2502.07490v2 📥 PDF

作者: Xialie Zhuang, Zhikai Jia, Jianjin Li, Zhenyu Zhang, Li Shen, Zheng Cao, Shiwei Liu

分类: cs.CL, cs.LG

发布日期: 2025-02-11 (更新: 2025-05-16)

备注: 17 pages,7 figures

💡 一句话要点

提出Mask-Enhanced Autoregressive Prediction (MEAP)以增强LLM上下文检索能力

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 掩码语言建模 下一词预测 上下文检索 长文本推理

📋 核心要点

  1. 大型语言模型在关键信息检索方面存在不足，影响了其在复杂任务中的表现。
  2. MEAP通过将掩码语言建模融入下一词预测，使模型更关注关键信息，提升检索能力。
  3. 实验表明，MEAP在信息检索和长文本推理上优于传统方法，且在微调中表现出显著优势。

📝 摘要（中文）

大型语言模型(LLMs)在准确检索关键信息方面存在不足。为了解决这个问题，我们提出了Mask-Enhanced Autoregressive Prediction (MEAP)，这是一种简单而有效的训练范式，它将掩码语言建模(MLM)无缝集成到下一词预测(NTP)中，以增强后者的上下文检索能力。具体来说，MEAP首先随机掩盖一小部分输入token，然后直接使用仅解码器Transformer执行标准的下一词预测自回归。MEAP消除了MLM对双向注意力或编码器-解码器架构的需求，在预训练或推理期间不会产生额外的计算开销。大量实验表明，MEAP在关键信息检索和长上下文推理任务上显著优于NTP，同时在常识推理任务上表现相当或更好。MEAP的优势也扩展到监督微调，在“中间丢失”场景中表现出显著优势，超过NTP 11.77个百分点。我们的分析表明，MEAP的有效性源于它能够通过集中于减少的非掩码token集合来促进更可区分的注意力分数。这种机制提高了模型对任务相关信号的关注，同时减轻了外围上下文的影响。这些发现使MEAP成为大型语言模型的一种有前途的训练范式。

🔬 方法详解

问题定义：大型语言模型在处理长文本时，难以准确检索到关键信息，导致在长上下文推理等任务中表现不佳。传统的下一词预测(NTP)方法可能无法充分关注到重要的上下文信息，而双向注意力机制虽然可以捕捉上下文信息，但计算成本较高。

核心思路：MEAP的核心思路是将掩码语言建模(MLM)融入到下一词预测(NTP)中。通过随机掩盖部分输入token，迫使模型更加关注未被掩盖的token，从而提高模型对关键信息的敏感度，增强上下文检索能力。

技术框架：MEAP使用一个decoder-only Transformer架构。训练过程分为两步：首先，随机mask掉一部分输入token；然后，使用标准的下一词预测方法，基于剩余的token预测下一个token。整个过程不需要额外的编码器或双向注意力机制，保持了计算效率。

关键创新：MEAP的关键创新在于将MLM和NTP巧妙地结合起来，无需引入复杂的双向注意力机制或额外的计算开销，即可有效提升模型的信息检索能力。它通过迫使模型关注未被mask的token，从而提高了模型对关键信息的敏感度。

关键设计：MEAP的关键设计包括：1) 随机mask的比例，需要根据具体任务进行调整；2) 使用标准的交叉熵损失函数进行训练；3) 使用decoder-only Transformer架构，保证了模型能够进行自回归预测。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

MEAP在关键信息检索和长上下文推理任务上显著优于NTP，在“中间丢失”场景的监督微调中，MEAP超过NTP 11.77个百分点。分析表明，MEAP通过促进更可区分的注意力分数，提高了模型对任务相关信号的关注，减轻了外围上下文的影响。

🎯 应用场景

MEAP可广泛应用于需要长文本理解和信息检索的场景，例如问答系统、文档摘要、机器翻译等。通过提高模型对关键信息的关注度，MEAP可以提升这些应用在处理复杂任务时的性能，具有重要的实际应用价值和潜在的商业前景。

📄 摘要（原文）

Large Language Models (LLMs) are discovered to suffer from accurately retrieving key information. To address this, we propose Mask-Enhanced Autoregressive Prediction (MEAP), a simple yet effective training paradigm that seamlessly integrates Masked Language Modeling (MLM) into Next-Token Prediction (NTP) to enhance the latter's in-context retrieval capabilities. Specifically, MEAP first randomly masks a small fraction of input tokens and then directly performs the standard next-token prediction autoregressive using a decoder-only Transformer. MEAP eliminates the need for bidirectional attention or encoder-decoder architectures for MLM, incurring no additional computational overhead during pre-training or inference. Intensive experiments demonstrate that MEAP substantially outperforms NTP on key information retrieval and long-context reasoning tasks, while performing on par or better on commonsense reasoning tasks. The benefits of MEAP also extend to supervised fine-tuning, where it shows remarkable advantages in lost-in-the-middle scenarios, outperforming NTP by 11.77 percentage points. Our analysis indicates that MEAP's effectiveness arises from its ability to promote more distinguishable attention scores by concentrating on a reduced set of non-masked tokens. This mechanism improves the model's focus on task-relevant signals while mitigating the influence of peripheral context. These findings position MEAP as a promising training paradigm for large language models.