Research map / Architectures and efficiency / RNNs and normalization
Initialization is Critical to Whether Transformers Fit Composite Functions by Reasoning or Memorizing
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 parameter initialization scale (via initialization rate γ) determines whether transformers learn compositional inferential (reasoning) or symmetric (memorizing) solutions on unseen anchor pairs, with small initialization favoring reasoning.
Problem
Whether transformers learn underlying compositional primitives or memorize input-output mappings on unseen compositional tasks, and what determines which solution is learned.
Approach
Using anchor functions, the parameter initialization scale (standard deviation 1/d_in^γ, with initialization rate γ) is found to switch the learned solution for an unseen anchor pair between inferential (composed single-anchor mappings, low complexity) and symmetric (memorized, high complexity); the authors propose γ as a tunable hyperparameter.
Key results
- On the unseen anchor pair (4,3), small initialization (large γ) yields inferential solutions while large initialization (small γ) yields symmetric solutions with near 100% accuracy in the phase diagram.
- On COGS, small initialization with large weight decay strongly outperforms large initialization with small weight decay on out-of-distribution generalization, even though in-distribution accuracy reaches over 99% in both settings with 20k training data.
- On SlimPajama with GPT-2 Medium trained on GitHub data, small initialization consistently achieved lower perplexity.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use small initialization rather than large initialization (ablation) | Transformer on unseen anchor pair (4, 3) · Synthetic anchor functions | Inferential-mapping accuracy | Helped |
| Use large initialization rather than small initialization (ablation) | Transformer on unseen anchor pair (4, 3) · Synthetic anchor functions | Symmetric-mapping accuracy | Helped |
| Use very large initialization with deep models (ablation) | Transformer on seen anchor pairs · Synthetic anchor functions | Accuracy | Hurt |
| Use small initialization and large weight decay instead of large initialization and small weight decay (ablation) | Transformer · SCAN, commands containing “turn left” | Generalization performance | Helped |
| Use small initialization and large weight decay instead of large initialization and small weight decay (ablation) | Transformer · COGS out-of-distribution | Generalization performance | Helped |
| Use small initialization and large weight decay instead of large initialization and small weight decay (ablation) | Transformer · PrOntoQA | Test accuracy | Helped |
| Use small initialization or large weight decay (ablation) | Two-layer transformer and GPT-2 · Addition task | Rule generalization | Helped |
| Use large initialization or small weight decay (ablation) | Two-layer transformer and GPT-2 · Addition task | Rule generalization | Hurt |
| Use small initialization (ablation) | GPT-2 Medium · SlimPajama, GitHub and GitHub+Wikipedia training data | Perplexity | Helped |
| Increase learning rate and weight decay (ablation) | 3-layer, 1-head transformer with γ=0.5 · Unseen anchor pair (4, 3) | Inferential-solution accuracy | Helped |
| Extend initialization-rate variation to GPT-2 (ablation) | GPT-2 · Unseen anchor pair (4, 3) | Generalization performance | Mixed |
Surprising findings
- Larger initialization can cause poor generalization even on seen anchors in deep models, and PrOntoQA convergence first slows then speeds up as data size grows for large initialization.
- Higher learning rate and weight decay within a range push the model toward inferential solutions.
Limitations
- Key experiments rely on synthetic data, which may not capture real-world complexity.
- Further verification on a wider range of models beyond GPT-2 is needed.
- Mechanistic analysis is mainly on a single-head attention model.
Experimental setup
Datasets: SCAN COGS PrOntoQA SlimPajama synthetic anchor-function data
Benchmarks: SCAN COGS PrOntoQA SlimPajama
Metrics: accuracy perplexity convergence steps
Baselines: large initialization with small weight decay
Builds on: anchor functions neural tangent kernel GPT-2
Closest papers by meaning
- Are Transformers Able to Reason by Connecting Separated Knowledge in Training Data? (ICLR 2025 · similarity 0.88)
- Grokking of Implicit Reasoning in Transformers: A Mechanistic Journey to the Edge of Generalization (NeurIPS 2024 · similarity 0.88)
- Transformers Provably Learn Chain-of-Thought Reasoning with Length Generalization (NeurIPS 2025 · similarity 0.87)
- Algorithmic Capabilities of Random Transformers (NeurIPS 2024 · similarity 0.87)
- From Condensation to Rank Collapse: A Two-Stage Analysis of Transformer Training Dynamics (NeurIPS 2025 · similarity 0.87)
- Strassen Attention, Split VC Dimension and Compositionality in Transformers (NeurIPS 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.