Research map / Architectures and efficiency / Attention and transformers
A distributional simplicity bias in the learning dynamics of transformers
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.
Builds "clones" of NLP datasets (WikiText-103, TinyStories) with truncated many-body token interactions using deep factored-attention models with quadratic activations, and shows BERT and GPT-Neo learn higher-order interactions sequentially during training.
Problem
Whether the distributional simplicity bias (sequential learning of increasingly complex distributions) explains transformer training on natural language with self-supervised objectives; estimating higher-order token interactions directly is intractable on large corpora.
Approach
Use stacks of factored-attention layers with x^2 activations, whose depth controls the maximum order of token interactions (order 2^(n-1)+1), to generate Monte Carlo-sampled clones of a dataset; then test a transformer trained on the original data against clones of different interaction orders to reveal sequential learning.
Key results
- Synthetic four-body data: one factored layer saturates at best two-body loss; adding layers reaches three- and four-body bounds sequentially.
- On TinyStories, BERT's test loss on low-order clones (two- and four-layer factored) saturates after about 3e4 steps while six-layer-factored and BERT-sampled clones keep decreasing.
- GPT-Neo trained on next-token prediction: loss on low-order clones saturates within the first 500 steps while higher-order clone loss keeps decreasing.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train longer on the original data | 4-layer BERT · WikiText-103 clones | test loss | Mixed |
| Train longer on the original data | 4-layer BERT · WikiText-103 high-order clones | test loss | Helped |
| Train beyond 10^4 steps | 4-layer BERT · TinyStories three-body clone | test loss | No clear effect |
| Train from 10^4 to 3 × 10^4 steps | 4-layer BERT · TinyStories clone from 4-layer factored attention | test loss | Helped |
| Train beyond 3 × 10^4 steps | 4-layer BERT · TinyStories clone from 4-layer factored attention | test loss | No clear effect |
| Continue training | 4-layer BERT · TinyStories clone from 6-layer factored attention | test loss | Helped |
| Continue training | 4-layer BERT · TinyStories clone sampled from BERT | test loss | Helped |
| Continue training | 2-layer GPT-Neo · TinyStories low-order clones | test loss | No clear effect |
| Continue training | 2-layer GPT-Neo · TinyStories higher-order clones | test loss | Helped |
| Add one factored-attention layer with x^2 activation (ablation) | Factored-attention model trained on synthetic four-body data · Synthetic four-body data | test loss | Helped |
| Add a third factored-attention layer with x^2 activation (ablation) | Factored-attention model trained on synthetic four-body data · Synthetic four-body data | test loss | Helped |
| Increase factored-attention depth from two to four layers (ablation) | Factored-attention clone generator · TinyStories | test loss | Helped |
Surprising findings
- A single layer of standard self-attention struggles to learn two-body conditionals that factored attention captures exactly.
- Gibbs sampling of bidirectional conditionals can be inconsistent, motivating Metropolis-Hastings sampling instead.
Limitations
- Residual correlations may remain in clones due to imperfect MCMC sampling convergence.
- Analytical results assume embedding dimension d=1 and give no quantitative timescales or sample complexities.
- Factored-attention clones use a finite-rank (r=400) value matrix approximation for a vocabulary of 10000.
Experimental setup
Datasets: TinyStories WikiText-103 synthetic four-body Hamiltonian data
Benchmarks: TinyStories clones (test loss) WikiText-103 clones (test loss)
Metrics: test cross-entropy loss test accuracy
Baselines: BERT (4-layer) GPT-Neo (2-layer) factored attention 2/4/6 layers
Builds on: factored attention Metropolis-Hastings sampling for bidirectional models BERT GPT-Neo distributional simplicity bias
Closest papers by meaning
- Understanding Transformers via N-Gram Statistics (NeurIPS 2024 · similarity 0.87)
- Algorithmic Capabilities of Random Transformers (NeurIPS 2024 · similarity 0.87)
- From Shortcut to Induction Head: How Data Diversity Shapes Algorithm Selection in Transformers (NeurIPS 2025 · similarity 0.87)
- Transformers need glasses! Information over-squashing in language tasks (NeurIPS 2024 · similarity 0.87)
- How Do Nonlinear Transformers Learn and Generalize in In-Context Learning? (ICML 2024 · similarity 0.87)
- Transformer Block Coupling and its Correlation with Generalization in LLMs (ICLR 2025 · similarity 0.87)
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.