Research map / Architectures and efficiency / Attention and transformers
Non-asymptotic Convergence of Training Transformers for Next-token Prediction
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.
Analyzes training dynamics of a one-layer transformer (self-attention + feed-forward) for next-token prediction, proving non-asymptotic sublinear directional convergence of two-stage normalized gradient descent to max-margin solutions, linear loss convergence, and generalization to unseen sentences.
Problem
Non-asymptotic convergence rates of training a transformer for next-token prediction, and whether the trained transformer generalizes to unseen data in theory.
Approach
Characterizes realizable NTP datasets via a collocation (token-to-next-token map) and query-dependent partial orders, then uses a two-stage normalized gradient descent: stage 1 trains the feed-forward layer on the collocation, stage 2 trains the key-query attention matrix on the full dataset; both converge in direction to max-margin solutions, with the loss converging linearly and non-comparable t
Key results
- Stage 1 feed-forward layer direction converges to max-margin solution at rate O(log t / t), with norm growing linearly.
- Stage 2 key-query matrix converges in direction at rate O(log^2 t / t), and cross-entropy loss converges linearly.
- Trained transformer predicts correctly on unseen sentences by treating non-comparable tokens as larger than non-optimal tokens (generalization theorem).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use two-stage normalized gradient descent: first train W_ov on collocations, then train W_kq on the full dataset with W_ov fixed | One-layer transformer for next-token prediction · Realizable NTP training data | training loss | Helped |
| Train W_ov on collocations in preprocessing stage | One-layer transformer for next-token prediction · Synthetic realizable NTP dataset | W_ov directional convergence rate: O(log t/t) | Helped |
| Train W_kq with W_ov fixed after preprocessing | One-layer transformer for next-token prediction · Synthetic realizable NTP dataset | W_kq directional convergence rate: O(log^2 t/t) | Helped |
| Train the attention layer using the two-stage procedure | One-layer transformer for next-token prediction · Unseen data with non-comparable tokens | next-token prediction | Helped |
Surprising findings
- Trained transformer generalizes by inserting non-comparable tokens between optimal and non-optimal tokens, producing predictions on unseen sentences.
Limitations
- Restricted to one-layer transformer with reparameterized key-query and output-value matrices and orthonormal vocabulary embeddings.
- Assumes no confused tokens and that optimal tokens are at least as numerous as non-optimal tokens.
- Experiments only on a small synthetic dataset with |V|=20.
Experimental setup
Datasets: synthetic realizable NTP dataset (|V|=20)
Benchmarks: synthetic realizable dataset
Metrics: loss directional cosine similarity to max-margin solution norm growth
Builds on: implicit bias of gradient descent max-margin SVM analysis of transformers (Tarzanagh et al.)
Closest papers by meaning
- In-context Convergence of Transformers (ICML 2024 · similarity 0.89)
- On the Generalization Ability of Next-Token-Prediction Pretraining (ICML 2025 · similarity 0.88)
- Towards Understanding the Universality of Transformers for Next-Token Prediction (ICLR 2025 · similarity 0.88)
- Local to Global: Learning Dynamics and Effect of Initialization for Transformers (NeurIPS 2024 · similarity 0.88)
- Next-Latent Prediction Transformers Learn Compact World Models (NeurIPS 2026 · similarity 0.88)
- One-Layer Transformer Provably Learns One-Nearest Neighbor In Context (NeurIPS 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.