Research map / Architectures and efficiency / Attention and transformers

Learning In-context $n$-grams with Transformers: Sub-$n$-grams Are Near-Stationary Points

ICML 2025 · Architectures and efficiency › Attention and transformers · approach: Transformer expressivity and learnability

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.

Shows that sub-n-gram solutions implemented by a simplified two-layer disentangled transformer are near-stationary points of the in-context n-gram population cross-entropy loss, with a numerical stage-wise learning experiment for trigrams.

Problem

Explaining the long plateaus and stage-wise learning dynamics when transformers learn in-context n-gram language models with cross-entropy loss.

Approach

Gives a sufficient stationarity condition for CE loss when the model's score depends only on a suffix of the context, then constructs two-layer transformer parameters implementing k-gram MLE estimators (k ≤ n) whose population gradient vanishes as sequence length T and weight scale c go to infinity.

Key results

• Theorem: gradient norm at the k-gram construction is sqrt(c)·O(||p(.|suffix) − p̂_k||²) + O(T sqrt(c) e^{-c}), so it vanishes as T and c grow (near-stationary).
• In trigram experiments, the two-head model first plateaus at a bigram solution (both heads attend to the previous token) before the second head shifts to the (-2)-token, reaching trigram behavior.
• The population loss gradient reduces to the error between the model and the suffix-conditioned next-token distribution when the score depends only on the suffix (sufficient stationarity condition).

Measured effects reported in the paper

Surprising findings

• Sub-n-gram solutions are stationary points even though the model can represent the full n-gram predictor, suggesting plateaus are intrinsic to the loss landscape.

Limitations

• Results hold only asymptotically as sequence length T and weight scale c go to infinity, and for the population rather than finite-sample loss.
• Analysis uses a simplified disentangled two-layer transformer with orthogonal embeddings and fixed second-layer value matrix.
• Unclear which estimator transformers learn at finite sequence length.

Experimental setup

Datasets: synthetic in-context trigram sequences (Dirichlet(0.5) transition priors, vocabulary 5)

Metrics: population cross-entropy loss gradient norm bound attention maps

Builds on: induction heads disentangled transformer

Closest papers by meaning

• An Analysis of Tokenization: Transformers under Markov Data (NeurIPS 2024 · similarity 0.88)
• Effective Context in Transformers: An Analysis of Fragmentation and Tokenization (NeurIPS 2026 · similarity 0.88)
• Implicit Optimization Bias of Next-token Prediction in Linear Models (NeurIPS 2024 · similarity 0.87)
• Regress, Don't Guess: A Regression-like Loss on Number Tokens for Language Models (ICML 2025 · similarity 0.87)
• Towards Understanding the Universality of Transformers for Next-Token Prediction (ICLR 2025 · similarity 0.87)
• On the Generalization Ability of Next-Token-Prediction Pretraining (ICML 2025 · similarity 0.87)