Technical Summary: Small Initialization Matters for Large Language Models

Problem Statement

While progress in Large Language Models (LLMs) is typically attributed to scale, data, and architecture, parameter initialization remains a critical design choice that has resurfaced as a complex problem in the large-model era. Standard heuristics (e.g., Xavier, Kaiming) often require costly trial-and-error tuning. A prevailing theoretical view suggests that small initialization drives networks into a "condensed" regime where weights align along shared directions before diversifying, potentially enhancing reasoning and generalization. However, it remains unclear whether these benefits persist in realistic LLM pretraining at scale, or if architectural components inadvertently suppress the advantages of small initialization.

Methodology

The authors systematically investigate the impact of initialization scale on LLM training dynamics and performance.

• Initialization Parameterization: Weight matrices $W$ are initialized as $W_{ij} \sim \mathcal{N}(0, \sigma^2)$, where $\sigma = d_{in}^{-\gamma}$. The parameter $\gamma$ controls the scale: $\gamma = 0.5$ recovers standard Xavier-like scaling, while $\gamma > 0.5$ (specifically $\gamma = 1$) yields smaller initialization.
• Model Architectures: Experiments were conducted on both dense decoder-only Transformers and Mixture-of-Experts (MoE) models. Architectural variations included standard LayerNorm/RMSNorm and Multi-Head Attention.
• Architectural Interventions: To address observed limitations, two specific adjustments were tested:
  1. Reducing the RMSNorm constant ($\epsilon$): Lowering $\epsilon$ from the standard $10^{-5}$ to $10^{-12}$ to prevent the normalization factor from becoming dominated by the constant when hidden-state variance is small.
  2. Gated Attention: Introducing query-dependent gates to attention head outputs to mitigate the "attention sink" (disproportionate attention on the first token), which is exacerbated by small initialization.
• Training Setup: Models were trained on a high-quality bilingual (Chinese-English) corpus (1 trillion tokens) using next-token prediction with AdamW. Experiments spanned model sizes from 0.1B to 1.5B parameters.
• Analysis Metrics:
  • Token-level: Symmetric probability gap ($\Delta_{sym}$) and token difficulty (average loss) to locate performance gains.
  • Mechanistic: Row-wise cosine similarity and Stable Rank ($\|W\|_F^2 / \|W\|_2^2$) to track the evolution of parameter structure (condensation vs. expansion).

Key Results

1. The Scaling Limitation and Its Fix

Under standard settings ($\gamma=0.5$ vs. $\gamma=1$), small initialization consistently lowers validation loss, but the benefit diminishes as model size increases. The authors identify two architectural barriers causing this:

• RMSNorm Saturation: With standard $\epsilon=10^{-5}$, the normalization factor becomes insensitive to the reduced variance of small initialization once $\sigma^2(h) \lesssim \epsilon$. Reducing $\epsilon$ to $10^{-12}$ restores sensitivity.
• Attention Sink: Small initialization intensifies the tendency to attend to the first token. Gated attention effectively mitigates this, significantly lowering loss when combined with small initialization.
• Outcome: When both barriers are removed, the performance gap widens. Small initialization ($\gamma=1$) with these adjustments improves effective model size by approximately 44% and yields consistent gains across model scales.

2. Optimal Initialization Scale ($\gamma=1$)

The study determines that initialization should not be arbitrarily small.

• Residual Flow Balance: In a Transformer, the residual update must remain comparable to the identity pathway. Theoretical analysis shows the residual flow scales as $d^{1-\gamma}$.
  • If $\gamma < 1$, the residual dominates.
  • If $\gamma > 1$, the residual vanishes, leaving the network near an identity mapping.
  • $\gamma = 1$ represents the critical balance point where the residual update matches the embedding scale. Empirically, performance peaks at $\gamma=1$ and degrades for $\gamma > 1$.

3. Token-Level and Task-Specific Gains

• Non-Uniform Gains: Improvements are not uniform across all tokens. Gains concentrate on "moderately-to-highly difficult" tokens that require contextual integration and reasoning, rather than easy tokens or extremely noisy/rare tokens.
• Benchmark Performance: On 1.5B models, the optimized small initialization ($\gamma=1$ + adjustments) significantly outperforms standard initialization ($\gamma=0.5$) on reasoning-heavy benchmarks (e.g., GSM8K, MATH500, HellaSwag), with absolute gains exceeding 4% on several tasks.

4. Mechanistic Trajectory: Condensation

Small initialization induces a distinct developmental trajectory:

• Phase 1 (Condensation): Early in training, weight matrices rapidly condense into low-rank, low-complexity structures (rows align into coherent blocks).
• Phase 2 (Expansion): As training progresses, these structures expand into richer, higher-complexity representations.
• Contrast: Standard initialization ($\gamma=0.5$) begins with a diffuse, high-rank structure that evolves slowly, lacking the initial low-complexity phase. This "low-to-high" complexity trajectory is hypothesized to bias the model toward parsimonious reasoning rules rather than surface memorization.

Significance and Claims

The paper argues that parameter initialization is a "gene-like" determinant of model capacity and training dynamics, not merely a low-level implementation detail.

• Design Axis: The authors propose treating the initialization scale $\gamma$ as an explicit hyperparameter in LLM design.
• Practical Recommendation: They advocate for a $\gamma$-initialization rule: expose the initialization scale as a knob and adopt $\gamma=1$ as the default. This intervention is described as "almost cost-free" yet capable of improving pretraining and strengthening reasoning across scales.
• Theoretical Insight: The work provides concrete evidence supporting the idea that "compression is intelligence," demonstrating that a specific initialization regime forces models through a condensation phase that correlates with improved reasoning capabilities.
• Architectural Interdependence: The findings highlight that initialization benefits are contingent on architectural components (specifically normalization constants and attention mechanisms), suggesting these must be tuned jointly rather than independently.