Mask-Enhanced Autoregressive Prediction: Pay Less Attention to Learn More
Authors: Xialie Zhuang
Zhikai Jia
Jianjin Li
Zhenyu Zhang
Li Shen
Zheng Cao
Shiwei Liu
Venue: ICML 2025 poster (2025)
Keywords: Masked next token prediction
LLM pre-training
key information retrieval
Primary area: deep_learning->large_language_models
TL;DR: We propose Mask-Enhanced Autoregressive Prediction (MEAP), a simple yet effective training paradigm that seamlessly inte- grates Masked Language Modeling (MLM) into Next-Token Prediction (NTP) using a decoder- only Transformer.

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. Code has been submitted.

Mask Bask-EEn Autoregressive Prediction: Pay Less Attention to Learn More
Xialie Zhuang"
Zhikai Jia *2 Jianjin Li Zhenyu Zhang
Li Shen
Zheng Cao
Shiwei Liu
Abstract
1. Introduction
Large Language Models (LLMs) are discovered to suffer from accurately retrieving key informa- tion. To address this, we propose Mask-Enhanced Autoregressive Prediction (MEAP). simple yet effective training paradigm that seamlessly inte- grates Masked Language Modeling (MLM) into Next-Token Prediction (NTP) to enhance the lat- ter's in-context retrieval capabilities. Specifically, MEAP first randomly masks small fraction of input tokens and then directly performs the stan- dard next-token prediction autoregressive using decoder-only Transformer. MEAP eliminates the need for bidirectional attention or encoder- decoder architectures for MIM, incurring no addi- tional computational overhead during pre-training or inference. Intensive experiments demonstrate that MEAP substantially outperforms NTP on key information retrieval and long-context soning tasks, while performing on par or better 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 a reduced set of non non-masked tokens. This mechanism improves the model's 's fo- cus on resk-relevant signals while mitigating the influence of peripheral context. These findings position MEAP as a promising training paradigm for large language models. Code is available https://gitub..oom/siiiiX//
Next-token prediction (NTP) (Radford, 2018) is the foundational training objective for many large language models (LLMs), including OpenAl's GPT series (Brown, 2020). NTP trains models to predict the next word (or token) in sequence. given all preceding tokens. Its scaling efficiency and exceptional performance in text generation have established as the dominant paradigm for state-of-the-ard L.LMs such as GPT-4 (Achiam et al., 2023), LLaMa3 (Dubey al., 2024). Gemini 1.5 Pro (Team et al., 2024), and DeepSeek- V3 (Liu al., 2024a). However, recent studies highlight the limitations of NTP-based LLMs in accurately retrieving key information from context (Liu et al., 2024b; Kamradt, 2023; Nelson et al., 2024).
In contrast, masked language modeling (MIM),, used in BERT (Devlin, 2018), adopts denoising objective that re- constructs masked inputs using bidirectional attention. This cloze-type nature makes MLM particularly effective for tasks requiring precise information retrieval and sentence- level understanding. However, MLM's inherent focus on re- constructing masked tokens reduces its effectiveness in tasks requiring coherent and long-form text generation (Wang Cho, 2019; Dong et al., 2019).
While intuitively appealing, combining NTP and MLM to leverage their respective strengths remains non-trivial challenge. MLM typically operates best within two-stack encoder-decoder architectures, and performance degrades significantly when applied to decoder-only Transformers (Tay et al., 2022). Efforts to integrate the two often rely on unified pre- pre-training pipelines where multiple objectives are alternated during the pretraining process (Dong et al., 2019; Tay et al., 2022). However, this multi-objective approach introduces substantial complexity to the training pipeline, making it cumbersome to scale, especially for models with billions or trillions of parameters.
Equul contribution 'School of Artificial Intelligence, Uni- versity Chinese Academy of Sciences, China *SCTTIX (SGP) TECH PTE. LTD., Singapore 'South China Normal Univer- sity. China UUniversity of Texas at Austin, USA SSnn Yat- Sen University, China "University of Oxford, UK. Correspon- dence to: Zheng Cao <zcao @ @scitix.ai>> Shiwei Liu <shi- wei maths ax.ac.uk>.
To this end, we propose Mask-En lass-Enhanced Autoregressive Prediction (MEAP), simple yet effective LLM training paradigm that seamlessly integrates masked tokens into next-token prediction. Specifically, small fraction of the input tokens and then directly perform standard next-token prediction using decoder-only Trans- former in an autoregressive manner. This straightforward
first randomly mask
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Mask- Mask-Enhanced Autoregressive Prediction: Pay Less Attention to Learn More
training and fine-tuning is presented in Sections 4.1 and 4.2, respectively. In Section 5, we further analyze the underlying reasons for MEAP's effectiveness. Section 6 provides an ablation study, and we conclude the paper in Section 7.
2. Related Work
[Mai]
Diakk
Masked Language Modeling. Pre-training is one of the most important pillars of LLMs. BERT first trained bidirectional, encoder-only Transformer with masked lan- guage modeling (MLM), where the model is trained to predict masked input tokens. XLNet (Yang, 2019) intro- duced the Permutation- -based Language Modeling to ac- count for dependencies between masked tokens during train- ing. RoBERTa (Liu, 2019) further improves the pre-training of BERT by training the model longer, over more data, with longer sequences. etc. MLM was further advanced by T5 (Roberts et al., 2019). Specifically, T5 frames every text processing task as a 'text-to-text" problem, leveraging in- creased lengths of corrupted tokens to achieve improved per- formance on classification tasks, which has contributed to its growing popularity. However, these models have shown limited performance in open-text generation and in-context learning, limiting their usage in modern LLMss
Figure 1. Overview of next token prediction, masked language modeling. and our MEAP.
modification eliminates the need for bidirectional attention or expensive encoder-decoder architecture, thereby incur- ring no additional computational overhead during training. During inference, our resulting LLMs can work as simply as LLMs that are trained with NTP with no extra engineer- ing effort. The simplicity of MEAP enables us to enhance LLMs' performance of key information retrieval and long- context reasoning, while retaining the impressive scaling efficiency of decoder-only LLMs. Figure I shows the illus- trations of different training paradigms.
As a general pre- pre-training paradigm, MEAP works effee- tively for scenarios of pre-training and fine-tuning. For the pre-training setting, we conduct control experiments by pre-training 1. IB L.LaMesttlle LLMSs (Zhang et al., 2024) with NTP and MEAP, where the training tokens scale from 40B to 200B. Our results demonstrate that MEAP substan- tially improves the performance in key information retrieval tasks such as Needle in a Haystack (Kamradt, 2023) by up to 33% average score and MultiDocument Question Answering (MDQA) (Liu et al., 2024b) by up to 27.2 per- centage points, while preserving general knowledge learned during pre-training. is noteworthy that MEAP achieves 85.8% accuracy with 60B training tokens on the Needle in a Haystack, while NTP requires 200B for similar perfor- mance, highlighting MEAP's superior data efficiency in key information retrieval. In addition, compared to the original NTP, MEAP also suffers less from hallucination.
Next Token Prediction. In a parallel vein, Radford et al. (2019) proposed next-token prediction (NTP) where a decoder-only Transformer is trained to predict the next to- ken from left to right using unidirectional attention ensured by casual mask. By predicting the next token based on pre- viously generated tokens and the given input context, NTP maintains coherence and logical flow in the generated text, well-suited for text generation. Moreover, NTP eliminates the need for an encoder, significantly improving the scal- ability of language models. Due to the above advantages, NTP serves as the most popular pre-training objective of modern LLMs (Brown, 2020; Achiam et al., 2023; Touvron et al., 2023; Jiang et al., 2023; Yang et al., 2024; Liu et al., 2024a).
In addition, the promise of MEAP also holds for LLM fine-tuning. Our MEAP framework demonstrates consis- tent improvements across multiple commonsense reasoning tasks, achieving average gain of 1.12 scores over the NTP baseline. On Multi-Document Question Answering, MEAP achieves an average improvement of 11.77% across all positions.
Unified Training Paradigms. There are works that propose unified training paradigms aiming to train one Transformer with multiple objective functions. For instance, UniLM (Dong et al., 2019) trains a bidirectional encoder on unidi- rectional language modeling (LM), bidirectional LM, and -euueencet--Sequencee LM. UL2 (Tay et al., 2022) proposes unified pre-training paradigm with Mixture-of -f-eooisersssss (MoD) to combine diverse pre-training paradigms together, improving the performance over T5 and GPT. While effec- tive, the preference for encoder-decode architectures and the complicated switches among different training objec- tives hinder their applications in practice.
Our analysis suggests that MEAP's effectiveness stems from its ability to enhance attention distinguishability by focus- ing reduced set of non-masked tokens. This mecha- nism sharpens the model's attention to task-relevant signals while reducing the impact of peripheral context. In essence, MEAP learns more by attending to fewer tokens.
In contrast, our approach seamlessly integrates masked to- kens into NTP without incurring any additional pre-training
The structure of this paper is as follows. Section 3 details the MEAP algorithm. The evaluation of MEAP on L.MM pre-
2

---

Autoregressive Prediction: Pay Less Attention to Learn More
ing only the preceding tokens predict the masked token. We carefully selected the masking ratio, P 15% for pre- training, to ensure that the model receives an adequate level of training difficulty and learning signals, without exces- sively disrupting the pre-training process. The relatively moderate number of masked tokens allows this approach to be seamlessly integrated into existing NTP frameworks, without significantly increasing pre-training overhead or altering the original training procedure.
Jaln le. toal
MCR) I Inchooodeee toal
[Mauk]
[Mark]
Pre-training
Fine-tuning
Figure 2. Training frameworks of MEAP: Left (Pre-training): A certain portion of input tokens is randomly masked, followed by standard next-token prediction (NTP). Right (Fine-tuning): Train- ing samples are duplicated, and the random masking strategy is applied to the copied sequenees. Standard NTP is then performed on the modified input for fine-tuning.
LLM fine-tuning, MEAP can also be extended to fine- tuning scenarios. In this scenario, we duplicate the training samples and apply the same random masking strategy to the copied sequences during fine-tuning. The original se- quences and their masked counterparts are then combined into a single input sequence feed into the model. The cross-entropy loss is computed only with the masked tokens (Um) in the answer tokens (U²). This design addresses crit- ical concern: input sequences in supervised fine-tuning often contain key information essential for downstream tasks. Di- rectly masking the original sequence risks removing crucial information, potentially compromising the model's perfor-
or inference costs, while preserving the ultra-efficiency of NTP. More importantly, MEAP is more suitable for modern LLMs, as method does not alter the core mechanism of NTP, the resulting models remain fully compatible with existing pipelines, platforms, and hardware optimized for modern LLMs.
mance on the target tasks. Masking the duplicated sequence incorporates MLM to NTP while avoiding this concern. We choose P = 10% for fine-tuning in our experiment. We only perform MEAP for the QA pair whose answer length exceeds 50, otherwise, we perform the standard NTP for the pair. Formally, the following objective of MEAP for fine-tuning is:
3. Mask-Enhanced Autoregressive Prediction
In this section, we introduce Mask-Enhanced Autoregres- sive Prediction (MEAP).
LLM pre-training, To enhance the performance of LLM in handling and understanding long texts, particularly in key tasks such as key information retrieval and long con- text, we designed and implemented a simple yet efficient random masking strategy. The core idea of this method is to selectively mask portions of the input during the pre- training phase. Specifically, we employed fixed proportion masking mechanism, where tokens in the input sequence are randomly masked according to a predefined percentage P. In this way, the model is forced to learn in the absence of some contextual information, which helps improve its deep understanding and reasoning capabilities.
(())
Σ LEU,UUM
(mask]..... ..,k),,....,..... .1.)) po | 2i...., . 1; "
where the sequence { is copy of the original sequence {#;) } (i.e., d. = x..)
Notably, while MEAP doubles the sequence length during fine-tuning, Figure shows that achieves superior perfor- mance NTP with only half the training time, essentially attaining stronger results with even fewer training tokens.
Formally, given oncoder-only Transformer @ and se- quence of input X = (z1, F22 (2,,..nn1,, we first randomly mask a fraction of P tokens having X' = (z), [mask]. .... I, 2.). We then perform the standard next-token prediction using the masked input in a left-to- right manner:
We believe that the effectiveness of MEAP stems from its ability to promote more distinguishable attention by focus- ing on fewer tokens during LLM training, as masked tokens typically receive negligible attention. This modification helps the model focus on task-relevant signals while reduc- ing the impact of peripheral context, as verified in Section 5.
(X') = po (vx xi, [mask],.... ...
4. Experimental Results
To evaluate the effectiveness of MEAP in training LLMs, we conduct controlled experiments comparing LLMs pre- trained/fnn--uneed by MEAP with those trained by NTP.
Same as NTP, when the model is tasked with predicting a masked token, employs causal masked attention, us-
3

---

Mask-Enhanced Autoregressive Prediction: Pay Less Attention to Learn More
Table Pre-training Evaluation. Zero-shot performance MEAP and NTP on various commonsense reasoning tasks. Results are measured directly after pre- vreitrainng on 200B tokens with no fine-tuning tuning.
ARC-e
ARC-e
BoolQ
PIQA
HellaSwag
WinoGrande
OBQA
Average
NTP
22.9
55.7
53.3
73.6
44.1
55.0
23.2
46.2
MEAP
25.4
56.4
59.5
72.3
43.4
55.3
22.6
47.8
4.1. Pre-training Evaluation
Table 2. Single needle accuracy (%) with 32K context.
Setup. We follow the Llama architecture (Touvron al., 2023) as our base model. Specifically, we train 1.1B decoder-only Transformers (Vaswani, 2017) following the setting of Zhang et al. (2024). Our model has 24 layers with 32 attention heads, a hidden size of 2,048, an intermediate hidden size of 5,632, and a context length of 4096. We follow the common configurations of LLM components, e.g,, Rotary Positional Embedding (RoPE) (Su al., 2024), Pre-Norm (Ba, 2016) with RMSNorm (Zhang & Sennrich, 2019), SwiGLU (Shazeer, 2020), and Grouped-query At- tention (Ainslic et al., 2023). To assess the scalability of MEAP, we increase the training token size from 40B to 60B, and further to 200B.
Token
40B
60B
200B
NTP
65,9
52.8
87.1
MEAP
80.2
85.8
98.2
window (the 'haystack'), This approach provides quanti- tative metries for assessing precise information extraction from extended contexts, particularly relevant for document analysis applications.
As this evaluation involves long-context modeling capacity, we follow the setting of Ye et al. (2024) and conduct a length extension to 64K. In particular, we continue training our model for additional 4B tokens from SlimPajama (Soboleva et al.) using the approach proposed in (Fu et al., 2024). The implementation utilizes modified Rotary Position Embed- dings with O = 640, 000.
For all experiments, we implement a learning rate warm-up during the first 10% of the training steps, followed by cosine annealing schedule, which decays the learning rate to 10% of its initial value. We use the AdamW optimizer with the following settings: 31 = 0.9, 82 = 0.95. The maximum learning rate is set to 4 x 10-4, the minimum learning rate is 4 x 10-5, and the weight decay is 5 x 10-2
To demonstrate MEAP's s scalability, we increase the training token size to 40B, 60B, and 200B, reporting the results of needle retrieval in Table 2. The results show that MEAP con- sistently outperforms NTP across different training scales, At 40B tokens, MEAP achieves 80.2% accuracy, signifi- cantly surpassing the baseline's 65.9% The performance gap peaks at 60B tokens, with MEAP maintaining steady improvement and reaching 85.8% accuracy. At 200B tokens, MEAP approaches optimal performance, attaining 98.2% accuracy, while the NTP baseline still falls short of 90% accuracy. It is noteworthy that MEAP achieves 85.8% accu- racy using just 60B training tokens, whereas NTP requires approximately three times as many (200B tokens) to reach a similar level. This demonstrates MEAP's superior data efficiency over NTP in key information retrieval.
1.1.1. LANGUAGE MODELING EVALUATION
While the primary goal of MEAP is enhance LLM perfor- mance in key information retrieval, it is essential ensure that integrating MLM into NTP does not compromise the model's 's fundamental language modeling capability. To eval- uate this, we employ the LM Eval Harness benchmark (Gao et al., 2024), assessing models in a zero-shot setting. The results, presented in Table 1, show that MEAP performs comparably to, or even outperforms, NTP, achieving a 1.6% improvement in the overall average score. This finding pro- vides strong evidence that incorporating random masking into NTP does not degrade the model's language modeling capacity. In the following evaluations, we will examine whether MEAP further improves performance in key infor- mation retrieval and long-context modeling.
We further illustrate the retrieval performance of our 200B- token model with a 32K context length in Figure The accuracy reported across varying answer needle depths (y-axis) and context lengths (x-axis). The results show that MEAP generally maintains perfect accuracy across different context lengths and depths, with errors limited to only two grid cells. In contrast, NTP begins to exhibit accuracy degra- dation at a context length of 24K, affecting a wide range of depths from 50% to 100%.
4.1.2. SHEDLE--N-A-HAYYSAAC RETRIEVAL
For key information retrieval, we choose the well- established Needle-in-a-Haasttd evaluation (Liu et al., 2024b), where the model is asked to retrieve the random fact or statement (the "needle") in the middle of a long context

---

Mask-Enhancee Autoregressive Prediction: Pay Less Attention to Learn More
Pretrain Needle Haystack
MEAP Neadle Haystack
3.8
Degth Pereett
Depth Percent
3.6
Scon
Score
-0.4
-04
70.0
3.00
3.0
1.00
3.00
1.00
32
02
100.0
3.00
1.0
3.08
100.0
3.0
100
5.0
4880
12000
160000
24050
2a0000
X2000
4000
Emo
L20000
I6000
20000
24m00
20000
32000
Content Length (okem))
Conteat Length (lokwas)
(a) NTP Pre-training
(b) MEAP Pre-training
Figure 3. Performance comparison hetween NTP MEAP on Needle Haystack. Scores are computed using ROUGE-1, measuring
unigram overlap between model responses and expected answers.
4.1.3. MULTI-DOCUMENT QUESTION ANSWERING
MEAP
60
Table Pre-training Evaluation. Relative accuracy (%) improve- ment of MEAP over NTP on multi-document QA.
50
Answer Position
5
10
15
20
85
10 documents
+7.6
+7.0
+30.6
20 documents
+12.4
+4.0
+5.1
+3.7
+27.2
30
We evaluate the model's ability to retrieve information from long contexts using a multi-document QA task (Liu et al., 2024b) based on NaturalQuestions -peen (Kwiatkowski al., 2019). The task requires identifying answers from target document while ignoring k- distractor documents, where k is the total number of documents in the context. We analyze performance across two context lengths with k = 10 and = 20 documents and multiple answer posi- tions. For each query, we construct an input context con- taining one target document with the annotated answer and k distractor documents that do contain any of the answers. -4t-extedded models used for this evaluation.
12K
28K
Length
Figure 4. Long-context reasoning performance comparison be- tween MEAP and NTP on the Multi-Needle Reasoning Task (M- RS) across different context lengths.
4.1.4. LONG-CONTEXT REASONING EVALUATION
We evaluate long-context reasoning capabilities using the Multi-Needle Reasoning Task (M-RS) (Li al., 2024a), which requires models to retrieve and extract multiple pieces of information from long texts and using them to logically answer questions that demand an integrated understanding and reasoning of various text segments. This forces the model to distribute attention across contextually relevant tokens rather than focusing solely on local patterns.
We report the accuracy improvement of MEAP over NTP in Table 3. MEAP again consistently outperforms NTP by good margins across all configurations, with significant gains at later positions (+30.6% at position in 10-doc, +27.2% at position 5 in 20-doc). These results indicate that MEAP enhances the model's ability to retrieve relevant in- formation from long contexts, maintain performance across different context lengths and positions, and handle complex scenarios with multiple distractors. The improvements high- light the effectiveness of the masking strategy in enhancing the model's iss  overall capability for long-come context information retrieval tasks.
We leverage the OpenCompass evaluation framework (Con- tributors, 2023) and report the results in Figure 4. MEAP consistently outperforms NTP across context lengths with 6.6 percentage point average improvement. demonstrates MEAP's enhanced capacity to maintain attention coherence over extended sequences.
5

---

Mask-Enhanced Autoregressive Prediction: Pay Less Attention to Learn More
Table 4. Accuracy (i.e., free of hallucinations) on text summariza- tion datasets evaluated by different LLM judges.
baseline. The performance improvements are particularly notable on ARC-c and WinoGrande, indicating enhanced reasoning capabilities. The results highlight MEAP's effec- tiveness in fine-tuning complex reasoning tasks.
Task
XSum
MultiNews
WikiSum
NTP (Deepseek-)) V3)
0.09
0.17
0.24
MEAP (Deepseek-V])
0.13
0.19
0.33
4.2.2. CROSS-MODEL GENERALIZABILLTY
NTP (Qwen-Plus)
0.16
0.11
0.21
To verify that MEAP's effectiveness generalizes across dif- ferent model architectures and scales, we conducted experi- ments diverse of pre-trained LLMs. Table 6 presents results across various commonsense reasoning tasks, while Table 7 shows performance multi-document QA tasks.
MEAP (Qwen-Plus)
0.19
0.14
0.27
NTP(GPTT-4))
0.14
0.10
0.19
MEAP(GPT-40)
0.16
0.13
0.24
4.1.5. CONTEXTUAL HALLUCINATION EVALUATION
The results demonstrate that MEAP consistently matches or outperforms NTP across different model architectures and QA task, MEAP demonstrates
Since MEAP improves more accurate key information re- trieval, we expect it to suffer less from contextual halluci- nation. To verify, we evaluate MEAP in reducing contex- tual hallucinations on three summarization datasets: XSum (Narayan et al., 1808), WikiSum (Cohen et al., 2021), and MultiNews (Fabbri et al., 2019), following Ye al. (2024). For this setting, we fine-tune the pre-trained models with Alpaca and evaluate them. We compare model--geneateer and reference summaries using (Deepseek-V) V3 (Liu et al., 2024a), Qwen-Plus (Yang et al., 2024) and GPT-4o (Hurst et al., 2024)) as the hallucination detector across 100 ran- dom samples per dataset. As shown in Table 4, our masking strategy achieves a consistent reduction in hallucination rates across all datasets.
sizes. On the substantial improvements across all model architectures and sizes. These results highlight MEAP's universal effective- ness in enhancing information retrieval capabilities.
4.2.3. MULTI-DOCUMENT QUESTION ANSWERING
We evaluate MEAP's context-aware reasoning using the multi multi-document QA task with distractor suppression (Liu et al., 2024b). To ensure fair comparison, we train MEAP for 2 epochs and NTP for 4 epochs, such that both ap- proaches process a similar number of tokens. Table 8 quan- tifies the exact match (EM) improvements across critical document positions in the 20-document setting. MEAP con- sistently achieves notable gains across all positions, further demonstrating its superiority over NTP. Two key patterns emerge from the experimental results:
4.2. Fine-tuning Evaluation
Setup. We fine-tune the Llama-3-8B (Dubey et al., 2024)
on the Alpaca instruction dataset (Taori al., 2023a). The training configuration uses global batch size of 512. The model is optimized with AdamW (3, 0.9, 32 0.95), a learning rate of 2 x 10 (with 10% warmup and cosine decay), and weight decay set to 0.01. We retain key architec- tural components from Llama-3, such as RoPE embeddings (Su al., 2024). RMSNorm (Zhang & Sennrich, 2019), and grouped-query attention (Ainslic al., 2023).
Consistent Improvement: MEAP achieves substantial gains across all positions with an average improvement of 11.77%, showing robust performance throughout the document range.
Mid-Context Advantage: The maximum improvement position 20 (+15.22%) demonstrates enhanced long- range dependency modeling, crucial for connecting concepts across scientific documents.
During fine-tuning, we randomly mask portion of tokens in the assistant's response, while keeping the source context intact. Only the masked tokens are predicted during fine- tuning. The training process uses bfloat16 precision with DeepSpeed Zero Stage 2 (Ren et al., 2021), and the Llama-3 tokenizer (Dubey et al., 2024) with maximum sequence length of 4096 tokens.
These findings validate MEAP's effectiveness in preserv- ing signal integrity across long contexts while highlight- ing opportunities for temporal reasoning enhancement and cross-document entity disambiguation.
4.3. Training Efficiency Analysis
4.2.1 1. LANGUAGE MODELING EVALUATION
MEAP introduces no additional overhead for pre-training inference compared to standard NTP, as the only difference lies in the masking operation. During fine-tuning, MEAP requires duplicating the input sequence and training with a doubled sequence length, resulting in increased training overhead. This overhead, however, effectively amortized
Similar to the pre-training evaluation, we first assess MEAP's effectiveness in language modeling, Table 5 presents the evaluation results. Our MEAP framework demonstrates consistent improvements across multiple tasks, achieving an average gain of 1.12 scores over the NTP
6

---

Mask-Enhanced Autoregressive Prediction: Pay Less Attention to Learn More
Table Fine-tuning Evaluation. Performance MEAP NTP on various commonsense reasoning tasks. Results are measured by line-tuning with Llama-3-8B.
ARC-e
ARC-e
BoolQ
PIQA
HellaSwag
WinoGrande
OBQA
Average
NTP
53.58
81.10
83.98
79.27
62.74
72.06 74.27
39.40
67.30
MEAP
55.12
83.21
83.82
81.01
63.31
38.20
68.42
Table Cross-Model Generalizability Fine-tuning performance comparison of MEAP and NTP on commonsense reasoning tasks when applied to different architectures and model sizes.
Model
Method
ARC-c
ARC-e
BoolQ
PIQA
HellaSwag
WinoGrande
OBQA
Average
Llama-3.2
NTP
47.95
69.07
75.54
76.50
72.43
64.33
44.40
64.32
Llama-3.22 3.2-3B
MEAP
49.32
73.06
71.80
77.53
74.26
68.51
44.60
65.58
Qwen2.5-14B7
NTP
53.67
74.71
86.73
77.64
78.44
68.19
48.00
69.63
Qwen2.5-14B
MEAP
56.83
79.38
87.37
79.33
79.37
72.69
47,40
71.77
-B-00
NTP
35.67
60.10
75.81
71.22
63.03
61.40
35,40
57.52
Mistral-7B-0.2
MEAP
37.20
59.18
72.63
73.50
64.08
61.17
35.60
57.62
retrieval task, highlighting its efficiency and effectiveness in large-scale training scenarios.
MEAP-3
NTP-6
Accuracy (50
45
5. Why Does MEAP Work?
40
MEAP-2
MEAP-1
35
In this section, we attempt to interpret the underlying rea- sons for the effectiveness of MEAP. We conjecture that MEAP's effectiveness stems from its ability to promote more distinguishable attention by focusing on fewer tokens during LLM training, as masked tokens [MASK] are ex- pected to receive marginal attention scores.
NTP-2
NTP-4
NTP-1
500
1000
1500
Training Time (s)
-0-
-33-
While effective, attention mechanisms in LLMs often strug- gle with long-context understanding, where redundant and non-informative attention is assigned tokens (Liu et al., 2024b; Li al., 2024b). A plausible explanation is that the attention module relies the Softmax function normal- ize attention scores within (0. 1), which tends to minimize differences among tokens, especially when training on se- quences of thousands of tokens. This bears some resem- blance to Martins et al. (2020)'s findings on sparse attention mechanisms which implicitly relate to their core mechanism (reducing attention tokens).
Figure 5. Comparison of fine-tuning efficiency between MEAP and NTP. 'MEAP-n' refers MEAP training for epoch.
by MEAP's higher data utilization efficiency. Specifically, compared NTP, MEAP requires only 50% of the epochs with similar number of tokens being processed, while outperforming the latter significantly.
To verify, we report the results on the multi diltiiddccumntt QA retrieval from 20 documents (Liu et al., 2024b), where re- trieval performance is assessed by computing the average retrieval values across positions. As shown in Figure 5, single epoch of MEAP training significantly outperforms two epochs of NTP training by large margin while also reducing total training time. This highlights MEAP's data efficiency, achieving similar or better results while reducing computational resources.
Furthermore, LLMs exhibit phenomenon known atten- tion sinks, where the initial few tokens receive disproportion- ately high attention scores compared to the rest (Xiao et al., 2023). Collectively, these factors lead to small and nearly indistinguishable attention scores across tokens, which is generally undesirable. When I.LMs fail to properly differ- entiate between tokens, they are more likely to generate incorrect outputs.
In summary, MEAP delivers significant training time re- ductions with improved or comparable performance on the
By randomly replacing tokens with masks, MEAP implicitly

---

Mask-Enhanced Autoregressive Prediction: Pay Less Attention to Learn More
Table Cross-Model Generalizability. Accuracy (%) MEAP NTP on multi-document QA with 20 documents across different
model architectures and sizes.
Model
Method
Position
Position 5
Position 10
Position 15
Position 20
Average
Llama-3.2-3B
NTP
13.60
12.09
12.54
12.69
14.35
13.05
Llama-3.2-3B
MEAP
23.47
20.34
20.38
21.96
23.65
21.96
Mistral-7B-0.2
NTP
36.96
30.55
27.82
27.55
38.79
32.33
Mistral-7B-0.2
MEAP
37.91
32.98
31.46
32.22
43.45
35.60
Qwen2.5-14B
NTP
60.00
51.98
56.01
56.05
63.39
57.49
Qwen2.5-14B
MEAP
61.69
53.71
57.21
56.65
66.29
59.11
minimal attention in MEAP. Consequently, this reduction expected increase the attention variance non-masked positions, making the attention distribution in MEAP more distinguishable compared to NTP.
Table Fine-tuning Evaluation. Accuracy (%) MEAP and NTP on multi-docun dolti-document QA with 20 documents.
Position
10
15
20
NTP
24.29
22.82
24.11
25.46
31.11
Results. Table 10 confirms our expectations. MEAP as- signs 53.34% less attention to masked tokens, resulting in a 7.80% increase in attention variance. This finding validates that MEAP learns more distinguishable attention scores compared to NTP.
MEAP
33.22
34.16
36.01
36.91
46.33
A
+8.93
+11.34
+11.90
+11.45
+15.22
penalizes the attention scores masked positions, thereby amplifying the attention differences among non-masked tokens. This masking mechanism encourages the model to generate more distinguishable attention scores, allowing it to focus on task-relevant texts while mitigating the influence of peripheral context. We validate this hypothesis through the following experiments.
5.2. MEAP Focus More Task-Relevant Tokens
To verify if MEAP learns more effective attention, we mea- the average attention scores that the model assigns to different input segments. We structured our input sequences into distinct segments: context-before, answer, context-after, query, and EOS token. The complete input sequence was formed by concatenating these segments sequentially, fol- lowed by an EOS token. This structured format enabled precise tracking of attention allocation across different func- tional components.
5.1. Masking Leads to More distinguishable Attention
To elucidate the mechanistic impact of our masking strategy on model behavior, we conducted a detailed analysis of attention distribution patterns. Our experimental protocol involved sampling 500 sequences. The original unmodified samples refer to the input sequence of NTP XN. and their masked counterparts (same samples with 15% masks), X Mr designated as the input for MEAP. These sequence pairs were then processed through our 1.1B models pre-trained with NTP and MEAP, respectively. We compare two values: (1) Attention Score Decay: the percentage decrease in the averaged attention score masked positions, computed as:
Expectation. Our expectation is that MEAP tends to am- plify attention to answer spans and meanwhile reduce the attention to less relevant tokens.
Results. The attention distributions during inference for both models are visualized in Figure 6. Notably, MEAP exhibits a substantial improvement in answer-nelevant tention (34.5% vs. 9.4%) while reducing the dominance of context-before attention from 73. 1% to 49. 99.1% Both models maintain similar attention levels for peripheral components, including context-after, query sections, and EOS tokens (all approximately 5%-6%). These results demonstrate that the MEAP framework enhances attention allocation during inference, prioritizing key information effectively.
At |mask = 1]) Att(X [mask = 1])
Att(Xx 1])
(2) Attention Variance Increase: the attention variance in- crease non-mask positions, computed
Var(Att(Xar jmask = 0))) - - (aaarrAAtt(() mmakk = 0)))
6. Ablation Study
We anticipate that the average attention score at masked positions will undergo a significant decline in the model, indicating that masked tokens receive
Expectations.
We conduct ablation studies on the mask ratio for both pre-training and the-tuning settings. Table 9 summarizes
8

---

Mask-Enhance Autoregressive Prediction: Pay Less Attention to Learn More
Normalized Attention Score
Normalized Attention Score
NTP
MEAP
0.73
Low Attention Score
High Attention Score
0.49
0.35
0.09
0.06
0.05
0.06
0.06
0.05 0.05
Context
Answer
Context
Query <BOS>
Context
Contexl Query <EOS>
NTP
MEAP
Figure 6. Attention distribution comparison between NTP and MEAP during inference. The input sequence consists of: context-hefore ("In heart of Paris, the Eiffel Tower stands tall, symbolizing both city and the entire country.") ) answer ("Designed by Gustave Biffel"), context-after (", it completed in 1889 for the World's Fair. Originally criticized for its unusual design, it has since become one of the most recognizable landmarks in the world. Tourists from all over the globe visit it every year, making it one of the most photographed monuments."), "), and query ("question: Who designed the Eiffel Tower?"). MEAP allocates much higher attention score unswer-relevant tokens (0.345) compared to NTP (0.094).
Table Performance comparison of different mask ratios in pre-training and fine-tuning.
The best results are highlighted in bold.
Pre-training
Fine-tuning
Mask Ratio
NTP
0.05
0.10
0.15
0.20
NTP
0.05
0.10
0.15
Accuracy
0.52
0.54
0.56
0.58
0.56
0.72
0.77
0.81
0.71
(our default approach), 5-Span Masking (consecutive spans of 5 tokens), and 50-Span Masking (longer spans of 50 consecutive tokens). We evaluate these strategies using a 0.3B parameter model pre-trained on 5B tokens, with results presented in Table 15.
Table 10. Statistical analysis of attention score patterns between NTP and MEAP.
Length
Metrie
Value
T-Stat/P-Valtur
1024
Score Decay
34.08%
-25.71/<le-6
1024
Var. Increase
12.66%
12.26/<1e-6
4096
Score Decay
53.34%
-9.97/<1e-6
7. Conclusion
4096
Var. Increase
7.80%
5.22/<<e-6
This work addresses challenges in information processing through a straightforward approach that masks 10%-15% of input while maintaining traditional prediction methods. Our results show significant improvements in comprehension across longer contexts, achieved without additional compu- tational costs. This approach demonstrates remarkable effi- ciency, matching performance metries with just 60B training examples that typically require 200B examples with conven- tional methods. The results indicate that this strategy leads to more effective processing of key information through improved focus on relevant content. Since it requires no structural changes, this method can be readily integrated into existing systems without disrupting workflows.
the results. For pre-training, evaluate our pre-trained model in Section 4.1 on the Multi-Document QA task using the nq-open-oracle dataset (Liu et al., 2024b). For fine-tuning, we train MEAP on the Alpaca dataset (Taori al., 2023b) for 3 epochs with different mask ratios against standard NTP baselines with epochs for fair comparison. The results show that mask ratio of 0.15 achieves the best performance in pre-training, while mask ratio of 0.10 yields the highest accuracy in fine-tuning. MEAP consistently outperforms standard NTP in pre-training and fine-tuning, demonstrating its effectiveness in leveraging masked tokens for improved performance.
6.1. Effect Different Masking Strategies
To further investigate the impact of masking patterns, we compare three distinct masking strategies: Random Masking

---

Mask-Enhanced Autoregressive Prediction: Pay Less Attention to Learn More
Impact Statement
A., al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783, 2024.
This work proposes a modified pre-training paradigm that may influence how both industry and academia approach language model training. MEAP integrates seamlessly with existing LLM frameworks without requiring additional gineering effort or computational resources. While the im- provement in information retrieval and reasoning capabili- ties could have broad implications for downstream applica- tions, the method's computational efficiency and architec- tural compatibility mean it can be readily adopted within current training infrastructtrrs. We anticipate this work will contribute more efficient model development while main- taining established training pipelines and computational requirements.
Fabbri, A. R., Li, I., She, T., Li, S., and Radev, D. R. Multi-news: large-scale tion dataset and abstractive hierarchical model. arXiv preprint arXiv: 0901v00000.000 0906.099699999 2019.
summariza-
Fu, Y., Panda, R., Niu, X., Yue, X., Hajishirzi, H., Kim, Y., and Peng, H. Data engineering for scaling language models to 128k context. arXiv preprint arXiv:2402.1017/ 2024.
Gao, L., Tow, J., Abbasi, B., Biderman, S., Black, S., DiPofi,
A., Foster, C., Golding, L.., Hsu, J., Le Noac'h, ba A., Li,
H., McDonell, K., Muennighoff, N., Ociepa, C.,, Phang,
J., Reynolds, L., Schoelkopf, H., Skowron, A., Sutawika,
References
L., Tang, E., Thite, A., Wang, B., Wang, K., and Zou, A. A framework for few-shot language model evaluation,
Achiam, J., Adler, S., Agarwal, S., Ahmad, L.,, Akkaya, I,, Aleman, F. L., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al. Gpt-4 technical report. arXiv preprint arXiv:2333.08877. 2023.
07 2024. URL https: /eoodo org/ httgg//cord//////////// 12608602.
Hurst, A., Lerer, A., Goucher, A. P., Perelman, A., Ramesh, A., Clark, A., Ostrow, A., Welihinda, A., Hayes, A., Radford, A., et al. Gpt-4o system card. arXiv preprint arXiv:2410.21277 2024.
Ainslie, J., Lee- Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebrón, F., and Sanghai, S. Gqa: Training generalized multi-query transformer models from multi-head check- points. arXiv preprint arXiv:2305.13244 2023.
Jiang, A. Q., Sablayrolles, A., Mensch, A., Bamford, C., Chaplot, D. S., Casas, D. d. I., Bressand, F., Lengyel, G., Lample, G., Saulnier, I.., et al. Mistral 7b. arXiv preprint arXiv:2310.0683 2023.
Ba, L. Layer normalization. frXiv:1607.06455 2016.
arXiv preprint
Brown, T. B. Language models are few-shot learners. arXiv preprint arXiv:2005. 14165, 2020.
Kamradt, G. Needle in haystack-pressure testing llms. Github Repository, pp. 28, 2023.
Cohen, N., Kalinsky, O., Ziser, Y., and Moschitti, A. Wikisum: Coherent summarization dataset for efficient numan--vvllationn In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 1 11th International Joint Conference Natural Language Processing Volume 2: Short Papers), pp. 212- 219, 2021.
Kwiatkowski, T., Palomaki, J., Redfield, O., Collins, M., Parikh, A., Alberti, C., Epstein, D., Polosukhin, Devlin, J., Lee, K., al. Natural questions: benchmark for ques- tion answering research. Transactions of the Association for Computational Linguistics, 7:453-466, 2019.
Li, M., Zhang, S., Liu, Y., and Chen, K. Needlebench: Can 1lms do retrieval and reasoning in million context window?, 2024a. URL https:/ 2407 11963.
Contributors, O. Opencompass: A universal evaluation platform for foundation models. https: gitb... com/opent compass/ cateaaccccppppppppccpppeeee
opencompass, 2023.
Devlin, J. Bert:
Pre-training of deep bidirectional trans- formers for language understanding. arXiv preprint arXiv: 1810.04805 2018.
Li, T., Zhang, G., Do, D., Yue, X., and Chen, W. Long- context llms struggle with long in no-eotex learning. arXiv preprint arXiv:2404.00 020600 2024b.
Dong, L.. Yang, N., Wang, W. Wei, F., Liu, X., Wang, Y., Gao, J., Zhou, M., and Hon, H.WW. Unified language model mre-training for natural language understanding and generation. Advances in neural information process- ing systems, 32, 2019.
Liu, A., Feng, B., Xue, B., Wang, B., Wu, B., Lu, C., Zhao, C., Deng, C., Zhang, C., Ruan, C., et al. Deepseek- v3 technical report. arXiv preprint arXiv:2412.19437, 2024a.
Liu, N. F., Lin, K., Hewitt, J., Paranjape, A., Bevilacqua, M., Petroni, F., and Liang, P. Lost in the middle: How language models use long contexts. Transactions of the
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan,
10

---

Mask-Enhanced Autoregressive Prediction: Pay Less Attention to Learn More
Association for Compulational Linguistics, 12:157-173, 2024b.
Taori, R., Gulrajani, I., Zhang, T., al. Al- paca: A strong, replicable instruction-followinn httgg/ggtcco//ttttee aaunford_alpaee 2023b.
model.
1ab/
Liu, Roberta: robustly optimized bert pretraining approach. arXiv preprint arXiv: 19XX. 1692, 364, 2019.
Tay, Y., Dehghani, M., Tran, V. Q., Garcia, X., Wei, J., Wang, X., Chung, H. W., Shakeri, S., Bahri, D., Schuster, T., et al. U12: Unifying language learning paradigms. arXiv preprint arxiv:2205.05/111 2022.
Martins, A., Farinhas, A., Treviso, M., Niculae, V., Aguiar, P. and Figueiredo, M. Sparse and continuous attention mechanisms. Advances in Neural Information Processing Systems, 01-20999-210011 2020.
Team, G., Georgiev, P., Lei, V. I., Burnell, R., Bai, L., Gulati, A., Tanzer, G., Vincent, D., Pan, Z., Wang, S., et al. Gemini 1.5: Unlocking multimodal understand- ing across millions of tokens of context. arXiv preprint Stti::2003.555355 2024.
Narayan, S., Cohen, S. B., and Lapata, M. Don't give me the details, just the summary. Topic-Aware Convolutional Neural Networks for Extreme Summarization ArXiv, abs, 1808.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B.. Goyal, N., Hambro, E., Azhar, F., et al. Llama: Open and efficient foundation lan- guage models. arXiv preprint (aXXi::2322113777) 2023.
Nelson, E., Kollias, G., Das, P., Chaudhury, S., and Dan, S. Needle in the haystack for memory based large language models. arXiv preprint arXi::2407.01437 2024.
Radford, A. Improving language understanding by genera- tive pre-training. 2018.
Vaswani, A. Attention all you need. Advances Neural Information Processing Systems, 2017.
Radford, A.,
Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Wang, A. and Cho, K. Bert has mouth, and it must speak: Bert as a markov random field language model. arXiv preprint wwXiv:1902.04009 2019.
Ren, J., Rajbhandari, S., Aminabadi, R. Y., Ruwase, O., Yang, S., Zhang, M., Li, D., and He, Y. (Zero-offload): Democratizing (billion-scale model training. In 2021 USENIX Annual Technical Conference (USENIX ATC
Xiao, G., Tian, Y., Chen, B., Han, S., and Lewis, M. Ef- ficient streaming language models with attention sinks. arXiv preprint arXiv:2309. 17453, 2023.
Yang, A., Yang, B., Zhang, B., Hui, B., Zheng, B., Yu, B., Li, C., Liu, D., Huang, F., Wei, H., et al. Qwen2.5 technical report. arXiv preprint 1511::2412.1515, 2024.
21). pp. 551-564, 2021.
Roberts, A., Raffel, C.,, Lee, K., Matena, M., Shazeer, N., Liu, P. J., Narang, S., Li, W., and Zhou, Y. Exploring the limits of transfer learning with a unified text-to-text transformer. Google, Tech. Rep., 2019.
Yang, Z. Xlnet: Generalized autoregressive pretrain- ing for language understanding. arXiv preprint arXiv: 208237,2000 2019.
Shazeer, N.
Glu variants improve transformer. arXiv preprint ar 25X222200..05202, 2020.
Ye, T., Dong, L., Xia, Y., Sun, Y., Zhu, Y., Huang, G., and Wei, F. Differential transformer. arXiv preprint arXiv:2410.05256 2024.
Soboleva, D., AI- Al-aeeeeeb, F., Myers, R., Steeves, R., Hes- tness, J., and Dey, N. Slimpajama: 627b token cleaned and deduplicated version of redpajama, 2023. URL https://hugginnguue. coldduases//ereeba///iiiiiii 627B.
Zhang, and Sennrich, R. Root mean square layer nor- malization. Advances in Neural Information Processing Systems, 32, 2019.
Zhang, P., Zeng, G., Wang, T., and Lu, W. Tinyllama: An open-source small language model. arXiv preprint 2rXiv:2401.02385 2024.
Su. J., Lu. Y., Pan, S., Wen, B., and Liu, Y. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 2024.
Taori, R., Gulrajani, I., Zhang. T., Dubois, Y., Li, X., Guestrin, C. Liang, P., and Hashimoto, T. B. Stanford alpaca: An instraction-folloiing llama wttp:////o/tttttutttceeeeae stanford_a papca, paca, 2023a.
model. https:// /ttp://ggttub/////tt/ttttttttbtt

---

Mask-Enhanced Autoregressive Prediction: Pay Less Attention to Learn More
A. Experimental Details of Pre-training
A.I. Architecture and Hyperparameters
This section outlines the pre-training hyperparameters of the MEAP model, designed to ensure efficient training and optimal performance. The sequence length is fixed at 4096 tokens, enabling the model to handle long-range dependencies while maintaining computational efficiency. The learning rate schedule includes initial warm-up phase for the first 10% of training steps, followed by cosine decay to 10% of the initial value, allowing gradual and precise parameter adjustments. The AdamW optimizer is used with standard hyperparameters B1 0.9 and 82 = 0.95 to stabilize the optimization process. Learning rate bounds are set between 4 x 10-4 and 4 x 10 to ensure effective learning throughout training, while weight decay of 5 x 10-2 helps prevent overfitting and promote generalization by penalizing excessively large weights. Complete training hyperparameters are documented in Table 11.
The model sizes and corresponding hyperparameters are shown in Table 12.
Table 11. Hyperparameters of training
Name
Hyperparameter
optimizer
Adam W cosine
schedule
clip
1,0
max learning rate
4 10
min learning rate weight decay
4 x 10-5
5 x 10-2
sequence length batch size epoch
4096
256
Table Hyperparameters of pretrained MEAP models. Data amount are specified in tokens.
Params
Hidden
Intermediate
Heads
Layers
Steps
Data amount
100M
768
2048
12
12
2K
2 B
500M
1024
4864
16
24
10K
10 B
1.1 B
2048
5632
24
32
190K
200 B
A.2. Pre-training Loss of Difference Model Sizes
The loss curves of the MEAP model at various sizes, shown in Figure 7, provide detailed visualization of the model's performance across different scales.
A.3. Language Modeling Evaluation Of All Size Models for Pre-training
As shown in Table 13, present the evaluation results of models of different scales implemented using method. To comprehensively assess the language modeling performance of these models, we conducted detailed analysis for each model, with particular focus on their performance at varying scales.
A.4. Pretrained Model Evaluation Under Different Masking Rates
As shown in Table 14, we present the evaluation results of models implemented with our approach, where different mask rates are applied during training. A comprehensive and detailed analysis of the language modeling performance is conducted for each mask rate, with a focus on how varying levels of masking influence key performance metrics. This analysis elucidates the effects of mask rates on the model's 's ability to handle diverse linguistic tasks, highlighting any changes in accuracy as the mask rate is adjusted.
12

---

Mask-Enhanced Autoregressive Prediction: Pay Less Attention to Learn More
metric/train_tos
metric/train_Lo tossitoainn__ss
10
50k
100k
150k
530
1.kk
metric/train_Ien
Figure Overview of loss size pretrained models
Table 13, Results of all size MEAP pretrained models
Benchmark 100M
500M
1.1B
ARC-Challenge 17.32
18.4
25.4
ARC-Easy
31.99
42.0
56.4
BoolQ
45.14
55.63
59.5
HellaSwag
26.82
30.77
43.4
OpenBookQA
11.41
16.40
22.6
PIQA
58.49
66.81
72.3
WinoGrande
52.09
49.57
55.3
Avg
34.75
39.94
47.85
A.5. Details Of Contextual Hallucination Evaluation
Here are the prompt for summarization generation, where "doc" is the original text to be summarized.
Summarize the following article: doc
We use the following prompts to let the Deepseek v3 model perform binary classification to determine whether there is hallucination in the model output compared to the human summary.
The "model output" is the output of the model, and the "predicted label" is the manually annotated label. Please compare the "model output" with the "predicted label". By comparing the two, check the "model output" similar. If similar, return I; otherwise, return 0. An explanation of the output required. Here is the output format provide. Please follow strictly!! Score: xx
13

---

Mask-Enhanced Autoregressive Prediction: Pay Less Attention to Learn More
Table 14. Results of the 1.1b MEAP model under different masking
Mask Ratio
Mask Ratio
Mask Ratio
Benchmark
0.15
0.05
0.1
ARC Challenge
25.4
26.11
24.3
ARC Easy
56.4
56.1
54.3
BoolQ
59.5
56.5
53.4
HellaSwag
43.4
43.69
43.85
OpenBookQA
22.6
22.0
21.8
PIQA
72.3
72.63
72.91
Winogrande
55.3
56.4
56.91
Avg
47.84
47.63
46.78
Table 15. Performance comparison of different masking strategies against NTP baseline.
Method
ARC-c
ARC-e
BoolQ
PIQA
HellaSwag
WinoCirande
OBQA
Average
MDQA
NTP-0.3B
18.00
37.75
58.44
62.62
28.56
50.67
13.60
40.09
0.187
Random Mask
21.84
35.44
61.25
61.04
29.50
51.46
27.40
41.13
0.218
5-Span Mask
21.42
35.61
60.40
62.08
29.81
51.07
27.60
41.14
0.168
S0-Span Mask
23,46
36.20
59.54
62.84
30.43
50.99
28.00
41.64
0.189
A.6. Details of Attention Distribution of MEAP and NTP
To validate the generality of attention changes, we conducted corresponding tests on additional examples and observed that the attention changes these examples are consistent with the results presented in the main text. The specific changes are detailed Table 16, Table 17, and Table
14

---

Mask-Enhancee Autoregressive Prediction: Pay Less Attention to Learn More
Table 16. Attention change of example
Area
Content
MEAP attention
NTP attention
context before
The Great Wall of China,stretching over 13,000 miles, is one of the most impressive feats of ancient engineer- ing.
0.491
0.731
answer
Built to protect Chinese states from invasions
0.329
0.108
context after
the wall took several dynasties 2,000 years to complete. Its im- mense length and historical signif- icance make a popular tourist attraction today. The wall's con- struction involved countless work- ers, many of whom faced difficult conditions.
0.078
0.80
query
question: What was the purpose of the Great Wall of China?
0.067
0.070
<>>
0.069
0.071
Table 17. Attention change of example
Area
Content
MEAP attention
NTP attention
context before
In the early 20th century, Albert Ein- stein introduced his theory of rela- tivity, which changed the way we understand space, time, and gravity.
0.435
0.694
answer
His famous equation, E=mc²
0.386
0.115
context after
shows the relationship between en- ergy and mass. Einstein's ideas rev- olutionized physics, and his work led to the development of technolo- gies like GPS and nuclear energy. Despite facing initial skepticism, his theories were eventually proven through experiments and observa- tions, carning him Nobel Prize in Physics in 1921.
0.066
0.074
query
question: What famous equation did Albert Einstein create?
0.057
0.060
eos
<3>
0.055
0.057
15

---

Mask-Enhancee Autoregressive Prediction: Pay Less Attention to Learn More
Table 18. Attention change example
Area
Content
MEAP attention
NTP attention
context before
At the center of Rome, the Colos- seum rises as a magnificent testa- ment to ancient Roman architecture, symbolizing the grandeur of the Ro- man Empire.
0.579
0.748
answer
Constructed between 70 and 80 AD under the emperors Vespasian and Titus,
0.219
0.065
context after
it was used for gladiatorial contests and public spectacles. Once sym- bol of Roman power, the Colosseum has weathered centuries change remains one of the most iconic structures in the world. Tourists flock to see it every year, making it one of the most photographed mon- uments in history.
0.071
0.067
query
question: Who built the Colosseum?
0.063
0.050
eos
<3>
0.068
0.069
16

---

Mask-Enhanced Autoregressive Prediction: Pay Less Attention to Learn More
B. Details of Fine-tuning Experiments
B.1. Architecture and Hyperparameters
This section details the MEAP fine-tuning hyperparameters for the Llama3 model. The maximum sequence length is 4096 tokens, optimizing long-range dependencies and efficiency. The batch size 512, and the learning rate schedule includes warm-up for the first 10% of training steps. The AdamW optimizer is used with 91 0.9 and B2 0.95, and the learning rate is set x 10-5
Table 19. MEAP fine-tuning hyperparameters of Llama3 model
Name optimizer schedule
Hy Hyprraarameter
AdamW cosine
clip learning rate weight_decn
1.0
2 x 10-5
5 10-2
maximum sequence length
4096
batch size
512
17


Reviews:
[{"review_id": "sDvNqxOCUw", "reviewer": "Reviewer_FwLm", "summary": "This paper introduces Mask-Enhanced Autoregressive Prediction (MEAP). In particular, MEAP incorporates the masked language modeling technique in next-token prediction setting by randomly masking out a small portion of input tokens and train the model with standard next-token prediction. This method is applied to both pretraining and finetuning. The author conducts experiments on long-context information retrieval benchmarks (NIAH and MDQA) and long-context reasoning benchmark (M-RS) and observe performance gain over NTP training. The author also shows the effectiveness of MEAP by analyzing the attention scores.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### **Q1: Intuition behind copying input for fine-tuning**\nThe duplication approach addresses a critical constraint in supervised fine-tuning: SFT answers are often extremely short (typically 5-15 tokens). With such limited tokens, conventional masking would risk fragmenting the semantic coherence of these responses. Even masking 15% of a 10-token answer could remove key semantic components.\nBy providing a duplicated sequence, we ensure semantic preservation while still enforcing the masking mechanism that drives attention specialization. Our experiments with single-sequence masking during fine-tuning (see \"Single Mask\" row in the fine-tuning results table) demonstrate substantially degraded performance across all tasks, confirming this design decision's importance.\n\nPerformance comparison on various commonsense reasoning tasks. Results are measured by fine-tuning with Llama-3-8B.\n\n| Method | ARC-c | ARC-e | BoolQ | PIQA | HellaSwag | WinoGrande | OBQA | Average |\n|--------|-------|-------|-------|------|-----------|------------|------|---------|\n| NTP    | 53.58 | 81.10 | 83.98 | 79.27| 62.74     | 72.06      | 39.40| 67.30   |\n| MEAP   | 55.12 | 83.21 | 83.82 | 81.01| 63.31     | 74.27      | 38.20| 68.42   |\n| Single Mask | 44.03 | 73.95 | 80.89 | 77.86 | 54.55 | 66.06 | 35.00 | 61.76 |  \n\n### **Q2: Clarification on Line 299**\nThank you for pointing out this inconsistency. The statement 'Only the masked tokens are predicted during fine-tuning' indeed does not align with the formula on Line 122. We will revise it as follows in the final manuscript.\n\n$$\n\\mathcal{L}(\\theta) = -\\sum_{t \\in U_a \\cup U_m} \\log p_\\theta ( x_t \\mid x_1, x_2, \\ldots, x_{t-1}; \\hat x_1, \\mathrm{[mask]}, \\ldots, \\hat{x}_{t-1})\n$$\n\nWhere the sequence $\\{\\hat{x}_i\\}$ is a copy of the original sequence $\\{x_i\\}$ (*i.e.*, $\\hat{x}_i = x_i$). The cross-entropy loss is computed over the subset of answer tokens ($x_t \\in U_a$) and masked tokens ($\\hat{x}_t \\in U_m$).\n\n\n### **Q3: HELMET Benchmark Results**\nFollowing your recommendation, we evaluated MEAP on the HELMET benchmark. The results strongly support our claims:\n\n|  | HotpotQA (RougeL Recall) | Natural Questions (RougeL Recall) | PopQA (RougeL Recall) | TriviaQA (RougeL Recall) | MS MARCO (NDCG@5) |\n|-------------|-------------------------|----------------------------------|----------------------|--------------------------|-------------------|\n| MEAP-1.1B   | **7.92**                    | **6.94**                             | **22.21**               | **17.22**                     | **18.38**            |\n| NTP-1.1B    | 3.43                    | 5.75                             | 6.31                | 8.00                      | 0.0               |\n| Improvement | 4.49                    | 1.19                             | 15.9                | 9.22                      | 18.38             |  \n\nWhile HELMET provides comprehensive evaluation, many of its test cases require context lengths up to 130K tokens, which exceeds our model's current context capacity.\nWe evaluated our approach on the subset of HELMET tests compatible with our model's context window. These results demonstrate that MEAP delivers consistent improvements across diverse question-answering tasks, with particularly substantial gains on challenging information retrieval benchmarks like PopQA (+15.9) and MS MARCO (+18.38). We will clarify this limitation in our revised manuscript and properly cite the HELMET benchmark as recommended.", "claims_and_evidence": "**Claim: MEAP improves performance in key-information retrieval and long-context modeling**\n\nThe author evaluates MEAP's performance on NIAH and MDQA, and shows improvement over the NTP approach. This shows that the key-information retrieval capability is improved with MEAP. For the long-context reasoning task, the author conducts experiment on Multi-Needle Reasoning Task. I would also recommend the author to evaluate on HELMET [1]. HELMET has a comprehensive evaluation on long-context tasks and I wonder whether MEAP is better than NTP on all types of long-context tasks, or if there is any type of long-context task that MEAP doesn't improve much.\n\n**Claim: The effectiveness of MEAP arises from its capability to promote more distinguishable attention scores**\n\nThe author shows an analysis in Figure 6 that shows MEAP effectively encourages the model to assign higher attention score on the answer part.\n\n[1] Yen, Howard, et al. \"Helmet: How to evaluate long-context language models effectively and thoroughly.\" arXiv preprint arXiv:2410.02694 (2024).", "methods_and_evaluation_criteria": "The method and evaluation method makes sense. As mentioned in the section above, I also recommend evaluating on the HELMET benchmark.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "(See Claims And Evidence section)", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. What is the intuition to copy the input again for the fine-tuning case. What if you just finetune like in the pre-training case?\n2. Line 299 left column: \"Only the masked tokens are predicted during fine-tuning.\" What does this mean? It does not seem to match the formula on Line 122 right column?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "u4dzuVuu9p", "reviewer": "Reviewer_Cepj", "summary": "The authors propose Mask-Enhanced Autoregressive Prediction (MEAP), a novel training paradigm integrating Masked Language Modeling (MLM) into the traditional Next-Token Prediction (NTP) objective. The key idea is randomly masking a fraction of input tokens during autoregressive pre-training, improving key information retrieval and long-context reasoning in decoder-only Transformers. The authors claim substantial improvements in benchmarks such as Needle-in-a-Haystack and Multi-Document QA without incurring additional computational overhead. MEAP is argued to improve attention differentiation by reducing unnecessary context attention, achieving significant performance gains over baseline NTP-trained models.\n\n## update after rebuttal\n\nI would like to thank the authors and the other reviewers for the discussion.\nI increased my score to accept, as explained in my comments.", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank you for your valuable comments! \n\n### **Q1: Variations in masking strategies.**\nWe added three distinct masking strategies as you suggested: Random Masking, 5-Span Masking (spanning 5 consecutive tokens), and 50-Span Masking, using a 0.3B parameter model pre-trained on 5B tokens. While they achieve comparable performance on commonsense reasoning, random masking achieves the best results on the Multi-Document QA task, possibly because its unpredictable mask encourages the model to develop more robust attention mechanisms across the entire context window.\n\n||ARC-c|ARC-e|BoolQ|PIQA|HellaSwag|WinoGrande|OBQA|Average|| | MDQA|\n|-|-|-|-|-|-|-|-|-|-|-|-|\n|MEAP Random Mask|21.84|35.44|61.25|61.04|29.50|51.46|27.40|41.13|||**0.218**|\n|MEAP 5-Span Mask|21.42|35.61|60.40|62.08|29.81|51.07|27.60|41.14|||0.168|\n|MEAP 50-Span Mask|23.46|36.20|59.54|62.84|30.43|50.99|28.00|41.64|||0.189|\n| NTP | 18.00 | 37.75 | 58.44 | 62.62 | 28.56 | 50.67 | 13.60 | 40.09 ||| 0.187 | \n\n### **Q2: Did you train both the NTP and MEAP models in the reported experiments?**\nYes, we trained both NTP and MEAP models in identical experimental settings. Both models used the same datasets, training steps, and hyperparameters, with the only difference being the masking mechanism in MEAP. This ensures fair comparison and reliability of our results.  \n\n### **Q3: Deepseek-V3 as the hallucination evaluator?**\nWe followed the setting of Diff Transformer using LLMs to make judgments. While the limited context window does allow us to design rigorous synthetic datasets, we added two more LLMs as judge to show the robustness of our evaluation. \n\n||XSum|MultiNews|WikiSum|\n|-|-|-|-|\n|NTP (Deepseek-V3)|0.09|0.17|0.24|\n|MEAP (Deepseek-V3)|**0.13**|**0.19**|**0.33**|\n|NTP (Qwen-Plus)|0.16|0.11|0.21|\n|MEAP (Qwen-Plus)|**0.19**|**0.14**|**0.27**|\n|NTP (GPT-4o)|0.14|0.10|0.19|\n|MEAP (GPT-4o)|**0.16**|**0.13**|**0.24**|\n\n### **Q4: Different model sizes for fine-tuning.**\nWe added fine-tuning results on various pre-trained LLMs. While MEAP achieves on-par or better results than NTP on commonsense reasoning, it consistently outperforms NTP on multi-document QA tasks, exhibiting consistent improvements across all tested model architectures and sizes.\n\n\n|||ARC-c|ARC-e|BoolQ|PIQA|HellaSwag|WinoGrande|OBQA|Average|\n|-|-|-|-|-|-|-|-|-|-|\n|Llama-3.2-3B| NTP| 47.95|69.07|75.54|76.50|72.43|64.33|44.40|64.32|\n|Llama-3.2-3B|MEAP|49.32|73.06|71.80|77.53|74.26|68.51|44.60|**65.58**|\n|NTP-Qwen2.5-14B|NTP|53.67|74.71|86.73|77.64|78.44|68.19|48.00|69.63|\n|MEAP-Qwen2.5-14B|MEAP|56.83|79.38|87.37|79.33|79.37|72.69|47.40|**71.77**|\n|NTP-Mistral-7B-0.2|NTP|35.67|60.10|75.81|71.22|63.03|61.40|35.40|57.52|\n|MEAP-Mistral-7B-0.2|MEAP|37.20|59.18|72.63|73.50|64.08|61.17|35.60|**57.62**|\n\nOn multi-document QA task (20-document setting):\n\n|||1|5|10|15|20|AVERAGE|\n|:---------|:--------:|:----:|:----:|:----:|:----:|:----:|:------:|\n|Llama3.2-3B|NTP|13.60|12.09|12.54|12.69|14.35|13.05|\n||MEAP|23.47|20.34|20.38|21.96|23.65|**21.96**|\n|Mistral-7B-0.2|NTP|36.96|30.55|27.82|27.55|38.79|32.33|\n||MEAP|37.91|32.98|31.46|32.22|43.45|**35.60**|\n|Qwen2.5-14B|NTP|60.00|51.98|56.01|56.05|63.39|57.49|\n||MEAP|61.69|53.71|57.21|56.65|66.29|**59.11**|\n\n### **Q5: How robust is MEAP across different language domains?**\nWe verified cross-domain robustness by training on GSM8K and testing on MathQA. MEAP achieved a 9.2% improvement over NTP, demonstrating that MEAP's gain also robustly holds in the math domain.\n\n|Method|mathqa|\n|:-:|:-:|\n|NTP|28.0|\n|MEAP|30.6|\n\n\n### **Q6: Attention score variance and decay analysis (Section 5.1) is presented briefly without rigorous statistical backing.**\n\nTo evaluate the statistical significance of our analysis, we sample 10 different question-answering examples, comparing attention distribution between NTP and MEAP during inference, and calculating the T-Statistic and P-Value. All tests reached statistical significance, confirming that MEAP models produce systematic changes in attention distribution.\n\n\n| Sequence Length | Metric | Value | T-Statistic | P-Value |\n|---------|------|------|---------|------|\n| 1024 | Attention Score Decay | 34.08% | -25.71 | <0.000001 |\n| 1024 | Attention Variance Increase | 12.66% | 12.26 | <0.000001 |\n| 4096 | Attention Score Decay | 53.34% | -9.97 | <0.000001 |\n| 4096 | Attention Variance Increase | 7.80% | 5.22 | <0.000001 |\n\n### **Q7: Essential references missing and redundancy issues.**\nWe thank you for pointing out Martins(2022) on sparse attention mechanisms and Nelson(2024) on information retrieval. We will definitely add a detailed discussion in our revision. We will also simplify the introduction to highlight key innovations and remove unnecessary repetition.", "claims_and_evidence": "* The primary claim—that MEAP significantly improves key information retrieval and long-context reasoning—is strongly supported by experimental results on multiple datasets (e.g., Needle in a Haystack and Multi-Document QA). \n\n* The claim regarding the reduced hallucination rate in summarization tasks (Table 4) is less convincingly evidenced, since the hallucination measurement is indirectly assessed using another LLM (Deepseek-V3). This indirect evaluation introduces potential biases due to dependency on the correctness and robustness of Deepseek-V3 itself. Introducing a synthetic dataset with verifiable answers will strengthen the claim.\n\n* The claim that MEAP’s improvements result from enhanced attention distinguishability is theoretically plausible but lacks detailed analysis across diverse experimental conditions, particularly variations in masking strategies.\n\n* The claim that MEAP achieves \"similar performance with only 60B tokens compared to NTP’s 200B\" is striking, but the experiment's control conditions are insufficiently described. It remains unclear whether other hyperparameters were optimally adjusted to fully exploit NTP's potential.", "methods_and_evaluation_criteria": "* The evaluation methods used (Needle-in-a-Haystack, MDQA, and various reasoning benchmarks) are appropriate and well-selected to demonstrate MEAP's strengths. \n\n* The usage of Deepseek-V3 as the hallucination detector for summarization datasets introduces potential bias, as this model might have intrinsic limitations and inaccuracies.", "theoretical_claims": "The paper does not contain formal theoretical proofs.", "experimental_designs_or_analyses": "* Lack of ablations on model components - The experimental setup lacks essential ablation experiments on key architectural decisions, such as masking percentage and the specific choice of masking during pre-training versus fine-tuning. Although Table 7 presents results on masking ratios, it is very brief and lacks deeper insights into why different mask ratios significantly affect performance. Experimenting with various masking strategies will also shed light about the method's effectiveness.\n\n* Single model size for fine-tuning - The fine-tuning experiments only evaluate the 8B Llama-3 variant, raising concerns about the generalizability across other sizes or architectures.", "relation_to_broader_scientific_literature": "The authors position their work well against key prior work, clearly highlighting distinctions from pure MLM models (BERT, RoBERTa, XLNet), pure NTP models (GPT, LLaMA), and unified paradigms (UniLM, UL2). The integration of masking into autoregressive prediction without additional computational overhead is a well-contextualized contribution within this literature landscape.", "essential_references_not_discussed": "* \"Sparse and Continuous Attention Mechanisms\" (Martins et al. 2022) on sparse attention mechanisms which implicitly relate to their core mechanism (reducing attention to tokens).\n\n* \"Needle in the Haystack for Memory Based Large Language Models\" (Nelson at al, 2024)", "other_strengths_and_weaknesses": "Strengths:\n\n* Creative integration of MLM into decoder-only Transformers, maintaining computational simplicity.\n* Strong empirical demonstration of improved performance across multiple tasks.\n* Clearly presented motivation and overall straightforward implementation.\n\nWeaknesses:\n\n* Insufficient theoretical justification or rigorous analysis behind the chosen masking ratio (15% pre-training, 10% fine-tuning), despite its significant impact.\n* Attention score variance and decay analysis (Section 5.1) is presented briefly without rigorous statistical backing or detailed statistical analysis of robustness.", "other_comments_or_suggestions": "* Table formatting is generally clear but lacks sufficient explanation in the captions—especially how percentages are computed.\n\n* The paper is generally well-written but suffers from redundancy, notably between abstract and introduction. The introduction could be more concise.", "questions_for_authors": "1. Did you train both the NTP and MEAP models in the reported experiments? Or did you use a pre-trained NTP and trained yourself MEAP? Comparison will be problematic in the latter .\n\n2. Why was Deepseek-V3 selected specifically as the hallucination detector? Did you explore other models or methods for hallucination evaluation? Clarification on this would affect confidence in your hallucination claims.\n\n3. Have you evaluated the robustness of MEAP across drastically different language domains (e.g., highly technical versus colloquial language) to justify its broader applicability?\n\n4. How sensitive is MEAP’s performance to different types of masking strategies (e.g., structured masking, linguistic units versus random masking)?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "IAea7DWVnH", "reviewer": "Reviewer_hgwa", "summary": "The paper introduces a simple modification to the standard causal next token prediction training framework of decoder-only LLMs, by masking a random fraction of the input tokens. After masking, the objective is still next token prediction. Via a set of experiments, the authors show the effectiveness of this paradigm, in particular in long context tasks with scarce relevant information. The authors hypothesize that by masking a fraction of the input tokens, the model is encouraged to sharpen its attention to distinguish between relevant and irrelevant parts of the context. To explore this hypothesis, they analyze the attention distribution and show that models trained with their technique place more attention weight on relevant parts while overall increasing the variance over unmasked tokens. This suggests that indeed the model distinguishes better between relevant and irrelevant parts.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your valuable review. We're glad you found our paper to be self-contained and our results convincing.\n\n### **Q1: Table 2: the drop of NTP perfomance when moving from 40B to 60B is suspicious. Can you please explain?**\n\nThank you for your sharp observation. One plausible explanation is that extended training with NTP might cause the model to progressively lose sensitivity to infrequent data points. Rare examples become increasingly diluted, reducing their influence on the model’s parameters. As NTP training progresses, earlier exposure to rare instances may be forgotten or down-weighted due to the overwhelming presence of more common patterns.\n\nHowever, training on sufficiently more data (e.g., 200B tokens) can help address this issue by ensuring that rare instances are represented with enough frequency and variety to leave a stronger imprint on the model. With more examples—both common and rare—the model has more opportunities to generalize from diverse patterns and avoid overfitting to dominant trends. \n\nIn contrast, MEAP that explicitly mask and predict masked tokens tend to emphasize different parts of the input, including rare tokens, making them more naturally sensitive to infrequent phenomena. Still, a sufficiently robust dataset can help NTP models maintain a better balance across the frequency spectrum.", "claims_and_evidence": "The claims are well supported by experiments.", "methods_and_evaluation_criteria": "The experimental setup is appropriate for the problem and results are convincing.", "theoretical_claims": "There are no theoretical claims.", "experimental_designs_or_analyses": "The experimental setup is appropriate for the problem and results are convincing.", "relation_to_broader_scientific_literature": "The idea is a simple modification of the standard LLMs training framework. There are previous works on combining training techniques, including token masking, but I am not aware of a published work that does what this paper is suggesting.", "essential_references_not_discussed": "Not that I am aware of.", "other_strengths_and_weaknesses": "The paper is well written and easy to read. \nThe idea is simple, which is a plus, yet surprisingly effective.", "other_comments_or_suggestions": "Table 2: the drop of NTP perfomance when moving from 40B to 60B is suspicious. Can you please explain?", "questions_for_authors": "Nothing to add.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper integrates masked language modeling into next-token prediction for decoder-only Transformers. Reviewers all give positive scores, on its clear motivation, simplicity, and strong empirical results demonstrating performance gains across multiple tasks. The creative yet computationally simple integration of MLM is also commended. However, several issues exist. The chosen masking ratios lack theoretical justification, and the analysis of attention scores lacks statistical rigor. Claims about reduced summarization hallucinations are weakly supported, and the link between improvements and attention distinguishability needs more exploration.  Overall, AC agrees this is an interesting, simple yet effective training paradigm for LLM, and thus recommends acceptance.
Decision: Accept (poster)