Research map / Architectures and efficiency / Attention and transformers

On the Generalization Ability of Next-Token-Prediction Pretraining

ICML 2025 · Architectures and efficiency › Attention and transformers · approach: Language model pretraining and induction

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.

Derives Rademacher-complexity generalization bounds for next-token-prediction pretraining of decoder-only transformers, handling token dependence via φ-mixing and covering-number bounds that include the causal mask, and checks the scaling predictions on small MiniMind models.

Problem

Lacking theoretical understanding of how next-token-prediction (NTP) pretraining of decoder-only transformers affects generalization, especially with dependent tokens within sequences.

Approach

Decomposes Rademacher complexity of the composite (representation-learner, token-predictor) class, models intra-sequence token dependence with φ-mixing, and derives covering-number bounds under the Frobenius norm that explicitly incorporate the causal mask matrix, yielding a bound scaling as sqrt(Θ/Nm)+sqrt(1/m).

Key results

• Test perplexity for miniNTP-4 (m=512) is 1.49 versus 316024.25 for miniNTP-1 trained on m=64, showing sequence length strongly affects generalization on length-512 test data.
• Reducing training data to 50% of the corpus raises test PPL from 1.49 to 3.17 (miniNTP-5 vs miniNTP-4), supporting the dependence on N.
• The paper's bound is O(sqrt(Θ/Nm) + sqrt(1/m)), empirically checked only through nine small-scale runs rather than a formal test of the bound's tightness.

Measured effects reported in the paper

Surprising findings

• Models trained on shorter sequences generalize poorly to longer test sequences, with perplexity exploding (e.g., 316024 for m=64 at test length 512).

Limitations

• Experiments are small-scale (up to 0.31B parameters) on a single Chinese dataset, so they only weakly validate the theory.
• Bounds rely on assumptions such as bounded weights, bounded attention logits, and phi-mixing, which are not verified on real data.
• The paper notes that the theory does not address language-specific data distributions or other learning frameworks such as stability.

Experimental setup

Datasets: MiniMind dataset DAMO NLP lcsts test set (8,192 samples)

Benchmarks: DAMO NLP lcsts test set

Metrics: test perplexity train loss test loss generalization gap

Baselines: Edelman et al. 2022 Li et al. 2023 Zhang et al. 2023 Deng et al. 2024

Builds on: Bartlett et al. 2017 spectrally-normalized bounds Lin & Zhang 2019 covering numbers Deng et al. 2024 MAE pretraining generalization Mohri & Rostamizadeh 2010 phi-mixing Rademacher bounds

Closest papers by meaning

• Implicit Optimization Bias of Next-token Prediction in Linear Models (NeurIPS 2024 · similarity 0.90)
• Unlocking Tokens as Data Points for Generalization Bounds on Larger Language Models (NeurIPS 2024 · similarity 0.89)
• The Role of Sparsity for Length Generalization in LLMs (ICML 2025 · similarity 0.89)
• Better & Faster Large Language Models via Multi-token Prediction (ICML 2024 · similarity 0.89)
• Non-asymptotic Convergence of Training Transformers for Next-token Prediction (NeurIPS 2024 · similarity 0.88)
• Towards Auto-Regressive Next-Token Prediction: In-context Learning Emerges from Generalization (ICLR 2025 · similarity 0.88)