Mask-Enhanced Autoregressive Prediction: Pay Less Attention to Learn
  More
Authors: name: Xialie Zhuang; name: Zhikai Jia; name: Jianjin Li; name: Zhenyu Zhang; name: Li Shen; name: Zheng Cao; name: Shiwei Liu
Published: 2025-02-11 11:49:03

Abstract: 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.