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

本文提出MEAP训练范式，通过在下一词预测中引入随机掩码策略，显著提升大型语言模型在关键信息检索和长上下文推理任务中的性能，同时保持计算效率和架构兼容性。

Large Language Model, Pre-training, Fine-tuning, Reasoning, Representation Learning

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

School of Artificial Intelligence, University of Chinese Academy of Sciences, China, SCITIX (SGP) TECH PTE. LTD., Singapore, South China Normal University, China, University of Texas at Austin, USA, Sun Yat-Sen University, China, University of Oxford, UK

Generated by grok-3

Background Problem

大型语言模型（LLMs）在关键信息检索和长上下文推理任务中表现出局限性，尤其是在从长文本中准确提取关键信息方面。下一词预测（NTP）作为主流训练范式，虽然在文本生成和扩展性上表现优异，但对上下文中的关键信息检索能力不足；而掩码语言建模（MLM）虽然擅长信息检索，却不适合生成任务且通常需要复杂的双向注意力或编码器-解码器架构。本文提出了一种新的训练范式，旨在结合NTP和MLM的优势，解决LLMs在长上下文任务中的信息检索问题，同时保持计算效率和架构兼容性。

Method

MEAP（Mask-Enhanced Autoregressive Prediction）是一种将掩码语言建模（MLM）思想融入下一词预测（NTP）的训练范式，核心思想是通过随机掩码输入序列的一部分，增强模型对非掩码token的注意力分配，从而提升关键信息检索能力。

预训练阶段：给定输入序列 $X = (x_1, x_2, ..., x_n)$ ，随机掩码比例为 $P=15\%$ 的token，得到 $X' = (x_1, [mask], ..., x_t)$ ，然后使用仅解码器的Transformer以自回归方式进行标准下一词预测，即 $p_{\theta}(X') = \prod_{t=1}^{T} p_{\theta}(x_t | x_1, [mask], \dots, x_{t-1})$ 。这种方式不引入双向注意力或编码器-解码器架构，避免额外计算开销。
微调阶段：复制训练样本，对复制的序列应用随机掩码（比例 $P=10\%$ ），将原始序列和掩码序列组合为单一输入，仅对答案中的掩码token计算交叉熵损失，以避免关键信息丢失。
关键机制：MEAP通过减少对掩码token的注意力分配，增强非掩码token的注意力区分性，促使模型更关注任务相关信息。

批判性思考：虽然MEAP的设计简单且直观，但掩码比例的选择（15%和10%）缺乏充分的理论依据，可能是实验调参的结果。此外，掩码是否真的能有效提升注意力分配的区分性，仍需更多证据支持，特别是在不同任务和模型规模下的泛化性可能存在问题。微调阶段序列复制导致长度翻倍，尽管作者声称效率更高，但实际计算成本的增加可能被低估。

Experiment

实验分为预训练和微调两个阶段，旨在验证MEAP在关键信息检索、长上下文推理和语言建模能力上的效果。

预训练实验：使用1.1B参数的LLaMa风格模型，训练token规模从40B到200B，评估任务包括常识推理（LM Eval Harness）、关键信息检索（Needle-in-a-Haystack）、多文档问答（MDQA）和长上下文推理（M-RS）。结果显示，MEAP在Needle-in-a-Haystack任务上显著优于NTP（例如在60B token时准确率达85.8%，而NTP需200B token才接近此水平），在MDQA任务中准确率提升高达30.6%，在长上下文推理任务中平均提升6.6个百分点，同时在常识推理任务上与NTP相当或略优（平均提升1.6%）。此外，MEAP在文本摘要任务中表现出更低的幻觉率。
微调实验：基于Llama-3-8B模型在Alpaca数据集上进行微调，MEAP在常识推理任务中平均提升1.12分，在MDQA任务中平均提升11.77%，尤其在长上下文中间位置表现更优（位置20提升15.22%）。跨模型实验表明MEAP在不同架构和规模的模型上均有改进。
效率分析：预训练和推理无额外开销，微调因序列复制导致长度翻倍，但MEAP在更少训练轮数下即可超越NTP，显示出数据效率优势。
实验设计评价：实验设置较为全面，涵盖了多种任务和模型规模，数据量和上下文长度（如64K）设计合理，体现了MEAP在长上下文任务中的潜力。然而，实验结果可能存在选择性报告的嫌疑，例如对掩码比例的消融研究较为有限，仅测试了少数比例，未充分探讨其对性能的影响边界。此外，部分任务（如幻觉评估）样本量较小（每数据集100个样本），可能影响结果的统计显著性。总体而言，MEAP的效果改进明显，但实验的深度和广度仍有提升空间。

Further Thoughts

MEAP提出了一种有趣的训练范式，通过掩码策略增强注意力分配的区分性，这一思想或许可以进一步扩展到其他领域，如视觉Transformer中的图像patch掩码策略，探索其在多模态模型中的应用潜力。此外，MEAP对注意力机制的改进与近期一些关于稀疏注意力机制的研究（如Martins et al., 2020）有异曲同工之妙，是否可以通过结合稀疏注意力进一步提升MEAP的效果，值得探索。然而，我对MEAP的长期影响持保留态度：掩码策略是否会在更大规模模型或更复杂任务中导致信息丢失或过拟合问题？特别是在资源受限的场景下，微调阶段的序列复制策略可能对计算成本造成更大负担，如何优化这一设计以适应工业级应用，是一个亟待解决的问题。未来研究还可以聚焦于掩码比例的动态调整或基于任务的自适应掩码策略，以进一步提升模型的泛化能力和效率。