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

We sincerely appreciate your valuable review. We're glad you found our paper to be self-contained and our results convincing.

Q1: Table 2: the drop of NTP perfomance when moving from 40B to 60B is suspicious. Can you please explain?

Thank 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.

However, 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.

In 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.

We sincerely thank you for your valuable comments!

Q1: Variations in masking strategies.

We 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.

Q2: Did you train both the NTP and MEAP models in the reported experiments?

Yes, 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.

Q3: Deepseek-V3 as the hallucination evaluator?

We 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.

Q4: Different model sizes for fine-tuning.

We 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.

On multi-document QA task (20-document setting):

Q5: How robust is MEAP across different language domains?

We 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.

Q6: Attention score variance and decay analysis (Section 5.1) is presented briefly without rigorous statistical backing.

To 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.

Q7: Essential references missing and redundancy issues.

We 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.

Q1: Intuition behind copying input for fine-tuning

The 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. By 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.

Performance comparison on various commonsense reasoning tasks. Results are measured by fine-tuning with Llama-3-8B.

Q2: Clarification on Line 299

Thank 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.

$$\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})$$

Where 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$).

Q3: HELMET Benchmark Results

Following your recommendation, we evaluated MEAP on the HELMET benchmark. The results strongly support our claims:

While HELMET provides comprehensive evaluation, many of its test cases require context lengths up to 130K tokens, which exceeds our model's current context capacity. We 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.