Research map / Large language models / Language, safety and interpretability
Mask-Enhanced Autoregressive Prediction: Pay Less Attention to Learn More
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes MEAP, which randomly masks ~15% of input tokens (10% in fine-tuning) and trains a decoder-only LLM with standard next-token prediction to improve key-information retrieval and long-context reasoning without extra compute.
Problem
Next-token prediction LLMs struggle to accurately retrieve key information from long context, while masked language modeling helps retrieval but is hard to combine with decoder-only NTP training.
Approach
Randomly replace a fraction of input tokens with [MASK] and perform standard causal next-token prediction on the masked sequence, so MLM-style denoising is integrated into a decoder-only NTP objective with no bidirectional attention or extra overhead; for fine-tuning, the sequence is duplicated and only the masked copy's answer tokens are supervised.
Key results
- Pre-trained 1.1B MEAP reaches 85.8% single-needle accuracy at 60B tokens versus NTP's 87.1% only at 200B tokens, and 98.2% vs 87.1% at 200B.
- After fine-tuning Llama-3-8B on Alpaca, MEAP averages 68.42 vs 67.30 for NTP on commonsense reasoning tasks.
- MEAP assigns 53.34% less attention to masked tokens and raises non-masked attention variance by 7.80% at length 4096.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use MEAP instead of NTP | 1.1B LLaMa-style pre-training · Needle-in-a-Haystack, 32K context, 40B tokens | accuracy: 65.9 -> 80.2 | Helped |
| use MEAP instead of NTP | 1.1B LLaMa-style pre-training · Needle-in-a-Haystack, 32K context, 60B tokens | accuracy: 52.8 -> 85.8 | Helped |
| use MEAP instead of NTP | 1.1B LLaMa-style pre-training · Needle-in-a-Haystack, 32K context, 200B tokens | accuracy: 87.1 -> 98.2 | Helped |
| use MEAP instead of NTP | pre-training · Commonsense reasoning | average score: 46.2 -> 47.8 | Helped |
| use MEAP instead of NTP | Llama-3-8B fine-tuning · Commonsense reasoning | average score: 67.30 -> 68.42 | Helped |
| use MEAP instead of NTP | fine-tuning · Multi-document QA, 20 documents | accuracy: +8.93; +11.34; +11.90; +11.45; +15.22 | Helped |
| use MEAP instead of NTP | MEAP fine-tuning · Commonsense reasoning, Mistral-7B-0.2 | average score: 57.52 -> 57.62 | Helped |
| use MEAP instead of NTP | MEAP fine-tuning · Multi-document QA, 20 documents, Llama-3.2-3B | average accuracy: 13.05 -> 21.96 | Helped |
| vary pre-training mask ratio from 0.15 to 0.05 (ablation) | MEAP pre-training · Mask-ratio ablation | accuracy: 0.58 -> 0.54 | Hurt |
| vary pre-training mask ratio from 0.15 to 0.10 (ablation) | MEAP pre-training · Mask-ratio ablation | accuracy: 0.58 -> 0.56 | Hurt |
| vary pre-training mask ratio from 0.15 to 0.20 (ablation) | MEAP pre-training · Mask-ratio ablation | accuracy: 0.58 -> 0.56 | Hurt |
| vary fine-tuning mask ratio from 0.10 to 0.05 (ablation) | MEAP fine-tuning · Mask-ratio ablation | accuracy: 0.81 -> 0.77 | Hurt |
Surprising findings
- MEAP attains comparable NeedleHaystack retrieval with 60B tokens as NTP needs 200B tokens, despite the same compute per token.
- Multi-document QA gains at answer position 20 (+15.22 points) in fine-tuning show larger gains at later positions than at early positions.
Limitations
- Fine-tuning doubles the input sequence length, increasing training overhead per epoch.
- Pre-training results are reported on a 1.1B model and several gains are from few-shot/controlled setups; hallucination evaluation uses only 100 samples per dataset with LLM judges.
- Mask ratio and masking strategy tuned on a limited set of settings; the paper notes gains on commonsense tasks are small and mixed.
Experimental setup
Datasets: SlimPajama Alpaca NaturalQuestions-Open nq-open-oracle
Benchmarks: LM Eval Harness commonsense tasks (ARC-c, ARC-e, BoolQ, PIQA, HellaSwag, WinoGrande, OBQA) Needle-in-a-Haystack Multi-document QA (NaturalQuestions-Open) M-RS multi-needle reasoning (OpenCompass) XSum/WikiSum/MultiNews hallucination evaluation
Metrics: accuracy ROUGE-1 exact match average accuracy hallucination-free rate
Baselines: NTP MLM XLNet RoBERTa T5 UL2 UniLM
Builds on: BERT next-token prediction UL2 UniLM LLaMA
Closest papers by meaning
- Scaling up Masked Diffusion Models on Text (ICLR 2025 · similarity 0.88)
- Large Language Diffusion Models (NeurIPS 2025 · similarity 0.87)
- When Linear Attention Meets Autoregressive Decoding: Towards More Effective and Efficient Linearized Large Language Models (ICML 2024 · similarity 0.87)
- Better & Faster Large Language Models via Multi-token Prediction (ICML 2024 · similarity 0.87)
- Reinforced Fast Weights via Next-Sequence Prediction (NeurIPS 2026 · similarity 0.87)
- MaskLLM: Learnable Semi-Structured Sparsity for Large Language Models (NeurIPS 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.