Summary (Overview)

• Wider attention heads improve long-context retrieval: At fixed model sizes (0.6B and 3B parameters), models with larger head dimensions (d512) consistently outperform models with narrower heads (d128, d256) on the RULER long-context retrieval benchmark, with gains of up to 15 points at the 3B scale.
• Training loss is not a reliable predictor: Pretraining loss and short-context non-retrieval task performance are largely unaffected by head dimension, while long-context retrieval performance shows strong head-dimension dependence, confirming that token-averaged perplexity is a poor proxy for long-context capability.
• Wider heads associate more parameters with retrieval heads: Models with wider heads have a higher proportion of retrieval heads (with comparable absolute counts despite fewer total heads), indicating that more parameters are dedicated to retrieval functions.
• Head-dimension advantage persists with frozen attention projections: Even when attention weight matrices are frozen during long-context continual pretraining, wider-head models show larger retrieval gains, suggesting the benefit comes from adaptation of embeddings and MLPs rather than attention module plasticity.
• QK dimension matters more than V dimension: Asymmetric configurations show that larger query-key dimensions benefit retrieval more than larger value dimensions, indicating a QK bottleneck in information selection.

Introduction and Theoretical Foundation

The paper addresses a fundamental question in Transformer architecture design: how should a fixed attention parameter budget be allocated across heads? Multi-head attention (MHA) retrieves and combines context representations through query-key interactions and value aggregation, with head dimension constraining the per-head representation space.

The theoretical foundation builds on several key observations:

1. Retrieval is not uniformly distributed across heads: Prior work (Wu et al., 2025; Zhang et al., 2025) shows that a small subset of attention heads plays a central role in extracting relevant information from long inputs.

2. Head dimension as capacity constraint: Bhojanapalli et al. (2020) demonstrated that low-dimensional query-key parameterizations can restrict attainable attention distributions, motivating head dimension as a capacity constraint rather than merely an implementation choice.

3. Fixed-budget allocation problem: For a fixed model dimension $D$, attention parameters scale as $\mathcal{O}(DHd)$ and dense attention computation as $\mathcal{O}(L^2Hd)$. When $Hd$ is fixed, changing head count and dimension leaves parameter count, computation, and KV-cache size essentially unchanged—only the allocation across heads changes.

The multi-head attention mechanism is formalized as:

$$\begin{array}{l} \mathrm{MHA} (X) = \mathrm{Concat} (\mathrm{head} _ {1}, \ldots , \mathrm{head} _ {H}) W _ {O}, \\ \mathrm{head} _ {h} = \mathrm{softmax} \left(\frac {X W _ {Q} ^ {(h) \top} W _ {K} ^ {(h)} X ^ {\top}}{\sqrt {d}} + M\right) X W _ {V} ^ {(h) \top}, \end{array}$$

where $d$ is the head dimension, $W_{Q}^{(h)}$, $W_{K}^{(h)}$, $W_{V}^{(h)} \in \mathbb{R}^{d \times D}$, $W_{O} \in \mathbb{R}^{Hd \times D}$, and $M$ represents the causal or sliding window mask.

Methodology

Model Architecture

The authors train decoder-only Transformers from scratch through a three-stage pipeline:

1. Pretraining (PT): 100B tokens at 4,096 sequence length
2. Continual Pretraining (CPT): 8B tokens at 262,144 sequence length
3. Supervised Fine-tuning (SFT): 1B tokens (3B models only)

Models use sliding window attention (SWA) with window size 512 and full attention layers interleaved in 3:1 ratio, with RoPE applied to SWA heads and 64-dimensional partial RoPE to full attention heads.

Experimental Configurations

Table 1: Architecture configurations (query/key dimensions and head counts for full-attention layers)

Evaluation

• RULER-100: 13 tasks (8 needle-retrieval variants, variable tracking, common-word extraction, frequent-word extraction, 2 QA tasks) across 7 context lengths from 4K to 256K
• Short-context tasks: ARC, HellaSwag, WinoGrande, RACE, PIQA, OpenBookQA, CommonsenseQA
• Retrieval head identification: Single-needle tasks with retrieval score $\geq 0.1$ threshold (following Wu et al., 2025)

Empirical Validation / Results

Main Performance Results

Table 2: RULER-100 overall scores (0–100 scale, averaged over 13 tasks and 7 lengths)

Key findings:

• d512 wins at every context length in all settings
• At 0.6B: d512 exceeds d128 by 9.96 points and d256 by 5.75 points
• At 3B: d512 exceeds d128 by 15.00 points and d256 by 14.21 points
• QK dimension is the bottleneck: k512-v128 (37.78) outperforms d256 (35.46), while k128-v512 (31.28) shows no gain over d128 (31.25)

Short-Context Performance

Table 3: Short-context tasks for 3B models after SFT

No trade-off exists between long-context retrieval and short-context performance; d512 performs best on average.

Retrieval Head Analysis

Table 4: Retrieval-head counts before and after CPT (0.6B models)

• Wider-head models have higher retrieval-head fractions (e.g., d512: ~25% vs. d128: ~7.3% after CPT)
• Absolute retrieval-head counts are comparable despite fewer total heads
• Much of the retrieval-head structure is already established during pretraining

Frozen Projection Experiments

With attention weight matrices frozen during CPT:

The head-dimension dependence persists without projection updates, while the incremental benefit of updating projections is similar across configurations.

Attention Map Diagnostics

• Attention entropy: Lower after CPT for all configurations, but d256 and d512 have similar entropy despite performance gaps—concentration alone doesn't explain the head-dimension effect
• Token reordering: Greater under Full CPT than Frozen CPT, but differences across head dimensions are small and don't consistently track retrieval performance
• Residual-stream effective rank: Wider-head models show higher effective ranks in intermediate blocks (layers 7–19), visible already in scaled PT and persisting under Frozen CPT

Theoretical and Practical Implications

1. Architecture design guidance: For long-context applications, allocating attention capacity to fewer, wider heads is beneficial at fixed parameter budgets. This contrasts with common practice where head count scales with model size while head dimension stays fixed.

2. Mechanistic understanding: The benefit of wider heads appears to operate through two channels:

   • Parameter association: Wider heads naturally associate more parameters with retrieval functions
   • Component adaptation: The advantage persists even when attention projections are frozen, implicating embeddings and MLPs in the retrieval improvement

3. Evaluation methodology: Training loss is insufficient for evaluating long-context capabilities; retrieval-specific benchmarks like RULER are necessary.

4. Practical constraints: The study is limited to head dimensions ≤ 512; further increases require efficient attention computation improvements (e.g., FlashAttention-style kernels).

Conclusion

This paper demonstrates that head dimension is a critical architectural choice for long-context retrieval capability in language models. The key takeaways are:

1. Wider heads consistently improve long-context retrieval at fixed model sizes, with no trade-off in short-context performance
2. QK dimension is the primary bottleneck for retrieval, more so than V dimension
3. The benefit emerges from broader parameter association with retrieval functions and more effective adaptation of non-attention components during continual pretraining
4. Standard training loss metrics fail to capture this architectural effect

Future work should explore:

• Head dimensions beyond 512 (requiring efficient attention computation)
• Whether these findings generalize to other model families and training recipes
• Deeper mechanistic understanding of how head dimension interacts with embeddings and MLP adaptation during long-context training