Research map / Architectures and efficiency / Attention and transformers
Towards Understanding the Universality of 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.
Shows theoretically that causal Transformers (linear, exponential, or softmax attention) can explicitly implement a causal kernel descent that predicts the next token in autoregressive sequences x_{t+1}=f(x_t) for linear or periodic f, with connections to Kaczmarz projections in Hilbert spaces.
Problem
Understanding how causal Transformers perform next-token prediction in-context for autoregressive sequences where the hidden map f changes per sequence, which is not well explained by existing in-context learning theory based on iid input-output pairs.
Approach
Constructs explicit one-layer two-head causal Transformers whose layer-wise iterations implement a causal version of kernel least-squares gradient descent (causal kernel descent) over vector-valued RKHS; the estimate of x_{t+1} is shown to converge as t grows, linked to Kaczmarz/alternating projections.
Key results
- For linear recursions with identity kernel and periodic recursions with exponential kernel, the error ||u*_t - x_{t+1}||^2 converges to zero, exponentially fast for instances (1) and (3).
- Fine-tuning the n=6 layer softmax Transformer (d=15) on synthetic data yields a model that significantly outperforms its theoretical initialization and outperforms the infinite-depth model for t gtrsim 20.
- Experiments on a more complex non-linear rotation instance (d=4, q=2) also show convergence of the causal kernel descent estimate to x_{t+1}, suggesting generalization beyond the proven cases.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use causal kernel descent for linear recursions with the dot-product kernel | Causal kernel descent · Linear recursion, dot-product kernel | squared prediction error | Helped |
| use causal kernel descent for periodic sequences with the exponential kernel | Causal kernel descent · Periodic recursion, exponential kernel | squared prediction error | Helped |
| use causal kernel descent with the exponential kernel for the nonlinear mapping | Causal kernel descent · Nonlinear mapping, d=4, q=2 | squared prediction error | Helped |
| fine-tune the initialized M^n | M^n · Orthogonal linear recursion, d=15, sequence length T=100 | next-token MSE | Helped |
| fine-tune the initialized M^n | M^n · Orthogonal linear recursion, t ≳ 20 | next-token MSE | Helped |
Surprising findings
- Fine-tuned finite-depth model outperforms the theoretically constructed infinite-depth model for t greater than about 20.
Limitations
- Analysis restricted to deterministic sequences; noisy dynamics left for future work.
- Only linear and periodic instances proven; general nonlinear f validated only experimentally.
- For softmax attention, convergence in n and t is only conjectured rather than proven.
Experimental setup
Datasets: synthetic autoregressive sequences generated by x_{t+1}=f(x_t) with random orthogonal Omega (instance 2) synthetic nonlinear complex-rotation sequences (instance 4)
Benchmarks: synthetic autoregressive sequences
Metrics: squared error ||u*_t - x_{t+1}||^2 MSE next-token prediction error
Baselines: untrained initialization Transformer (Mθ0)
Builds on: linear self-attention in-context learning (von Oswald et al.) kernel ridge regression in-context learning Kaczmarz algorithm neural ODEs Transformers as gradient descent in RKHS
Closest papers by meaning
- How Transformers Learn Causal Structure with Gradient Descent (ICML 2024 · similarity 0.91)
- Selective Induction Heads: How Transformers Select Causal Structures in Context (ICLR 2025 · similarity 0.90)
- How do Transformers Perform In-Context Autoregressive Learning ? (ICML 2024 · similarity 0.88)
- Non-asymptotic Convergence of Training Transformers for Next-token Prediction (NeurIPS 2024 · similarity 0.88)
- Linear Attention for Efficient Bidirectional Sequence Modeling (NeurIPS 2025 · similarity 0.88)
- Transformers are Universal In-context Learners (ICLR 2025 · 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.