Multi-head attention retrieves and combines context representations through query--key interactions and value aggregation, with head dimension constraining the per-head representation space. In long-context processing, it remains unclear whether allocating a fixed parameter budget to fewer, wider heads improves information retrieval. In this paper, we investigate how head dimension affects the performance of long-context retrieval tasks. Our empirical studies reveal that pretraining loss and short-context non-retrieval task performance are not affected much by head dimension, while long-context retrieval performance improves with wider heads at fixed model sizes. We find that models with wider heads have a higher proportion of retrieval heads, associating more parameters with retrieval. Our ablation studies also show that the retrieval performance improves with wider heads even when the projection weights in attention modules are frozen during long-context continual pretraining, suggesting that wider heads not only benefit updates in attention modules but lead to more effective adaptation of components such as embeddings and MLPs.
Figures & tables
Model Name
Params
Q/K Dim.
V Dim.
#Q Heads
#KV Heads
0.6B-d128
0.6B
128
128
16
4
0.6B-d256
0.6B
256
256
8
2
0.6B-d512
0.6B
512
512
4
1
0.6B-k512-v128
0.6B
512
128
8
2
0.6B-k128-v512
0.6B
128
512
8
2
3B-d128
3B
128
128
32
8
Table 1: Architecture names and their configurations. Dimensions and head counts refer to full-attention layers. The query and key dimensions are always identical, whereas the value dimension may differ. Full configurations are provided in Appendix A .
Model size
Stage
d128
d256
d512
k512-v128
k128-v512
0.6B
CPT
31.25 ± 7.75
35.46 ± 3.60
41.21 ± 5.86
37.78 ± 4.15
31.28 ± 2.61
3B
CPT
35.93
36.72
50.93
-
-
3B
SFT
48.28
51.81
58.49
-
-
Table 2: RULER-100 overall scores averaged over 13 tasks and seven lengths from 4K to 256K. Scores are on a 0–100 scale. The 0.6B results are means ± sample standard deviations across three runs; each 3B allocation has one run.
Figure 1: RULER-100 scores by context length. Each point averages 13 tasks at the indicated context length on a 0–100 scale. In (a), lines and bands show the mean and ± sample standard deviation across three 0.6B runs. In all the settings, d512 achieves the highest (average) score at every evaluated length.
Figure 2: training NTP loss during 8B-token CPT at 256K. The horizontal axis counts CPT tokens; the vertical axis is training cross-entropy. In (a), lines and bands are the mean and ±1 sample standard deviation across three 0.6B runs. Lower training loss does not always correspond to higher RULER performance.
Head Size
ARC-C
ARC-E
HS
WG
RACE
PIQA
OBQA
CSQA
Avg
d128
35.6
62.3
58.2
57.4
36.2
73.0
36.6
38.1
49.7
d256
35.8
61.0
58.1
59.4
35.0
73.3
37.4
28.9
48.6
d512
33.9
64.0
58.7
57.9
37.4
73.7
39.0
39.6
50.5
Table 3: Short-context task performance for the 3B models after SFT. Abbreviations: ARC-C (ARC-Challenge), ARC-E (ARC-Easy), HS (HellaSwag), WG (WinoGrande), OBQA (OpenBookQA), CSQA (CommonsenseQA), Avg (Average).
Figure 3: Retrieval-head fractions in the (a) 0.6B and (b) 3B models. Left: the fraction of full-attention query heads classified as retrieval heads; annotations give the retrieval-head count relative to the total number of eligible heads. Right: retrieval-head fractions within individual full-attention layers, using zero-based layer indices. The 0.6B results are means across three runs; faint lines show individual runs. Each 3B model has one run. Wider-head models exhibit higher retrieval-head fractions at both scales, while absolute retrieval-head counts remain comparable across allocations.
Model
PT (attention scaling)
CPT
Shared
0.6B-d128
4.00
7.00
4.00
0.6B-d256
8.33
8.33
8.00
0.6B-d512
5.00
6.33
5.00
Table 4: Retrieval-head counts before and after CPT for the 0.6B models. PT denotes the 100B-token PT checkpoint evaluated with inference-time attention scaling (Appendix A.4 ); CPT denotes the corresponding checkpoint after 8B tokens of full-parameter CPT, evaluated without additional scaling. The rightmost column reports the number of retrieval heads common to both stages, matched by layer and query-head index within each run. All the counts are computed within each run and then averaged across three runs per configuration.
Figure 4: Retrieval performance with and without attention weight matrices updates for the three 0.6B-parameter configurations. Bars show means across three training runs, each evaluated and averaged over 13 RULER-100 tasks and seven context lengths from 4K to 256K. The PT baseline uses inference-time attention scaling; neither CPT condition uses additional scaling. Brackets show mean differences in score points, computed before rounding.
Figure 5: Attention-map diagnostics on shared RULER examples at context lengths of 32K and 64K, using queries predicting the first three reference-answer tokens. (a) Mean attention entropy; lower values indicate more concentrated attention. (b) Capped top-10k rank correlation between each CPT checkpoint and the corresponding PT checkpoint; higher values indicate more similar token prioritization. PT checkpoints are evaluated without attention scaling. Points and error bars show means and ±1 sample standard deviation across three training runs, after averaging within each run.
Figure 6: Residual-stream effective rank at 32K context length for the three symmetric 0.6B configurations under scaled PT, Frozen CPT, and Full CPT. Each point is computed from the mean singular-value spectrum over 300 shared examples and three training seeds. For a full-attention layer with zero-based index L , the effective ranks of representations at Pre-FA, Post-FA, and Post-MLP are plotted at the layer index L , L+0.5 , and L+1 , respectively. These offsets distinguish measurement sites within the same block, not additional layers.
Appendix figures & tables19 assets
Supplementary material from the paper’s appendix.
Appendix
Model size
0.6B
3B
Vocab size
151,699
151,699
Model dimension
1,024
2,048
Number of layers
24
40
SWA:FA ratio
3:1
3:1
SWA number of heads
4
8
SWA head size
256
256
Appendix
Table 5: Architecture configurations.
Stage
Dataset
Tokens
Pretraining
Nemotron-CC-v2 High Quality ( Basant et al., 2025 )
60B
Swallow Code ( Fujii et al., 2025 )
15B
Swallow Code v2 (packed by repositories)
15B
Nemotron-Pretraining-SFT-v1 General
10B
CPT
Nemotron-CC-v2 High Quality
4.8B
Swallow Code v2 (packed by repositories)
2.4B
Appendix
Table 6: Data mix used in each training stage.
Model size
0.6B
3B
Pretraining
Tokens
100B
100B
Learning rate
2e-4
1e-4
LR warm-up steps
250
250
Weight decay
0.1
0.1
Batch size
1,024
1,024
Sequence length
4,096
4,096
Appendix
Table 7: Training hyperparameters.
Figure 7: Training NTP loss during 100B-token PT at 4K and 8B-token Full CPT at 256K. The horizontal axis shows tokens in billions within the indicated stage; the vertical axis shows training cross-entropy. The 3B panels show one run per model. In the 0.6B panels, lines indicate means across three runs and shaded bands indicate ±1 sample standard deviation.
Model
PT
Full CPT
0.6B-d128
1.908±0.001
1.911±0.009
0.6B-d256
1.908±0.001
1.888±0.042
0.6B-d512
1.913±0.004
1.912±0.011
3B-d128
1.705
1.677
3B-d256
1.704
1.673
3B-d512
1.700
1.665
Appendix
Table 8: Training NTP loss at the final logged interval of PT and Full CPT. Entries for the 0.6B models are means ± sample standard deviations across three runs; entries for the 3B models are single-run values.
Figure 8: Task-level RULER-100 scores after 8B-token Full CPT for the 0.6B models. The first panel shows the 13-task mean; the remaining panels show individual tasks. Thick lines indicate means across three runs, and faint lines show individual runs. The horizontal axis shows context length, and the vertical axis shows task scores as percentages.
Figure 9: Task-level RULER-100 scores after 8B-token Full CPT for the 3B models. The first panel shows the 13-task mean; the remaining panels show individual tasks. The horizontal axis shows context length, and the vertical axis shows task scores as percentages. Each model has one run. 3B-d512 leads in aggregate but not on every task.
Figure 10: Task-level RULER-100 scores for the 3B models after 1B tokens of SFT following CPT. The first panel shows the 13-task mean; the remaining panels show individual tasks. The horizontal axis shows context length, and the vertical axis shows task scores as percentages. Each model has one run. Relative to the CPT results in Figure 9 , overall scores increase, whereas scores at 256K decrease for all three models.
Parameter group
Full-Parameter Update (Full)
Attention-Frozen Update (Frozen)
WQ,WK,WV,WO in FA and SWA layers
Updated
Fixed
Head-wise QK-RMSNorm gain vectors
Updated
Updated
LM head, token embeddings, SwiGLU MLPs, and other RMSNorm parameters
Updated
Updated
Trainable parameter ratio
100%
83.4%
Appendix
Table 9: Parameter-update policies during CPT. Fixed parameters retain their values from the paired 100B-token PT checkpoints. Frozen CPT disables updates on weight matrices in both FA and SWA layers, while leaving the listed normalization and non-attention parameters trainable.
Model
Full CPT
Frozen CPT
Δloss
0.6B-d128
1.911±0.009
1.884±0.019
0.027±0.028
0.6B-d256
1.888±0.042
1.878±0.020
0.011±0.062
0.6B-d512
1.912±0.011
1.881±0.018
0.031±0.028
Appendix
Table 10: Training NTP loss at the final logged interval of Full and Frozen CPT. Entries are means ± sample standard deviations across three runs. Δloss is the Full-minus-Frozen loss difference, computed within each paired run before aggregation; positive values indicate lower loss under Frozen CPT. Differences are computed before rounding.
Figure 11: Training next-token prediction (NTP) loss during Full and Frozen CPT for the three 0.6B models. The horizontal axis shows CPT tokens in billions, and the vertical axis shows training cross-entropy. Lines indicate means across three runs, and shaded bands indicate ± sample standard deviation. Frozen CPT attains lower mean final training loss than Full CPT for all three models, despite achieving lower mean overall RULER-100 scores (Figure 4 ).
Figure 12: Overall RULER-100 scores for unmodified PT, scaled PT, Frozen CPT, and Full CPT in the three 0.6B configurations. Bars show means across three runs, after averaging equally over 13 tasks and seven context lengths from 4K to 256K within each run. Scaled PT uses the inference-time procedure in Appendix A.4 ; neither CPT condition uses additional scaling.
Figure 13: RULER-100 scores by context length for unmodified PT, scaled PT, Full CPT, and Frozen CPT. At each context length, scores are averaged over 13 tasks within each run and then across three runs. Scaled PT uses the inference-time procedure in Appendix A.4 .
Configuration
Run
PT
PT (attention scaling)
Frozen CPT
Full CPT
d128
1
18.57
28.36
23.23
25.50
d128
2
19.54
31.04
38.13
40.07
d128
3
18.63
27.81
25.83
28.19
d256
1
19.49
30.53
37.24
37.98
d256
2
17.57
29.21
33.81
37.05
d256
3
16.35
29.17
29.77
31.34
Appendix
Table 11: Overall RULER-100 scores for individual 0.6B training runs, averaged equally over 13 tasks and seven context lengths from 4K to 256K. PT denotes evaluation without additional attention scaling; PT (attention scaling) denotes scaled PT. Full CPT and Frozen CPT follow the parameter-update policies in Table 5. Only Full CPT results are reported for the asymmetric configurations k512-v128 and k128-v512 ; dashes indicate unreported conditions.
Figure 14: Task-level RULER-100 scores for unmodified PT, scaled PT, Full CPT, and Frozen CPT of 0.6B-d128 . The first panel shows the 13-task mean; the remaining panels show individual tasks. Thick lines indicate means across three runs, and faint lines show individual runs. The horizontal axis shows context length, and the vertical axis shows task scores as percentages. Scaled PT is labeled “PT (attention scaling)” and follows Equation 1 .
Figure 15: Task-level RULER-100 scores for unmodified PT, scaled PT, Full CPT, and Frozen CPT of 0.6B-d256 . The first panel shows the 13-task mean; the remaining panels show individual tasks. Thick lines indicate means across three runs, and faint lines show individual runs. Axes and scaling follow Figure 14 .
Figure 16: Task-level RULER-100 scores for unmodified PT, scaled PT, Full CPT, and Frozen CPT of 0.6B-d512 . The first panel shows the 13-task mean; the remaining panels show individual tasks. Thick lines indicate means across three runs, and faint lines show individual runs. Axes and scaling follow Figure 14 . Full achieves a higher overall mean score than Frozen, with gains and losses depending on the task and context length.
Model
Run
Scaled PT
Full CPT
Frozen CPT
Shared PT/Full
ΔN
Shared Full/Frozen
0.6B-d128
1
3
9
8
3
+6
8
2
5
6
7
5
+1
6
3
4
6
6
4
+2
6
0.6B-d256
1
7
7
8
7
0
7
2
9
9
9
9
0
9
3
9
9
7
8
0
7
Appendix
Table 12: Retrieval-head counts and shared membership for each training run. PT denotes scaled PT; Full and Frozen denote the corresponding 8B-token CPT checkpoints without additional scaling. Shared counts are intersections of the detected (layer,queryhead) sets. ΔN is the Full-CPT count minus the scaled-PT count; it measures the net count change, not the number of heads newly detected under Full CPT.
Figure 17: Cumulative attention probability over positions ranked by decreasing attention probability. Rows show 0.6B-d128 , 0.6B-d256 , and 0.6B-d512 ; columns show 32K and 64K context lengths. Each panel compares scaled PT, Full CPT at 2B, 4B, and 8B tokens, and Frozen CPT at 8B tokens. The horizontal axis gives the number of ranked positions on a logarithmic scale.
Large vision-language models increasingly rely on long-context modeling to reason over documents, hour-level videos, and long-horizon agent trajectories, requiring them to locate relevant evidence across interleaved text and images. Prior work has studied this behavior using retrieval heads in large language models, but its copy-based criterion does not directly apply when evidence appears in images. We introduce a multimodal retrieval head detection method that scores attention from question tokens to textual or visual evidence. With this method, we show that multimodal retrieval heads are sparse, intrinsic, and causally important: only 4.4-10.2% of attention heads account for 50% of the positive retrieval-score mass, and masking the top-5% selected heads reduces performance on two tasks from MMLongBench by 42.5 and 62.3 percentage points, while random-head masking is far less damaging. Further analysis shows that these heads are partly shared across modalities yet remain dynamic within each modality, with image retrieval heads changing more than text retrieval heads as context length and haystack modality change. Without further training, we find that these heads can also be used directly to rank visually rich documents: on MMDocIR, Qwen3-VL-8B selected-head scoring improves Recall@1 by 7.7/7.4 macro/micro points for page retrieval and 6.3/6.8 points for layout retrieval over the strongest reported baseline.
Retrieval heads, attention heads that copy information from earlier context to the current position, have been proposed as the mechanistic substrate for long-context recall. Rotary position embeddings (RoPE) rotate queries and keys by frequencies decaying with a base hyperparameter theta, and a natural hypothesis is that this rotation either prevents retrieval heads from forming or degrades their function. We test both across four open-weight 7-8B models spanning multi-head and grouped-query attention and a 100x range of theta, using paired-seed needle-in-a-haystack tests, layer-clustered permutation, and causal head-masking. (i) Retrieval heads are causally necessary: masking the 87 detected heads in OLMo-2 collapses recall from 1.00 to 0.00, while masking matched random heads has no effect; this replicates in Qwen. (ii) Higher theta does not reduce retrieval-head count (LLaMA-3.1 at theta=500K has 47 heads vs LLaMA-2 at theta=10K with 42), refuting the prevention hypothesis. (iii) The norm-utility relation is family-specific and significant in opposite directions (Qwen d=-0.49, OLMo d=+0.50, both significant; LLaMA null); since OLMo and LLaMA-3.1 share theta=500K yet differ, the effect is not theta-driven. (iv) Building on Chiang and Yogatama (2025), a controlled patch shows that zeroing the lowest-frequency RoPE dimensions of retrieval heads degrades recall dose-dependently (1.00 to 0.18 when 32 of 128 dimensions are zeroed, vs 0.98 for random dimensions); the effect is head-specific and task-specific. The causal variable is RoPE frequency, not norm-utility. The direction holds in all five models patched (OLMo-2, Qwen2.5-7B/14B, Gemma-2, Mistral) across four lineages and two scales. We do not claim cross-model magnitude. Code and a paired-seed harness are released.
Language models struggle to generalize beyond pretraining context lengths, limiting long-horizon reasoning and retrieval. Continued pretraining on long-context data can help but is expensive due to the quadratic scaling of Attention. We observe that most tokens do not require (Global) Attention over the entire sequence and can rely on local context. Based on this, we propose L2A (Learning To Attend), a layer that enables conditional (token-wise) long-range memory access by deciding when to invoke global attention. We evaluate L2A on Qwen 2.5 and Qwen 3 models, extending their effective context length from 32K to 128K tokens. L2A matches the performance of standard long-context training to within 3% while skipping Global Attention for ∼80% of tokens, outperforming prior baselines. We also design custom Triton kernels to efficiently implement this token-wise conditional Attention on GPUs, achieving up to ∼2× improvements in training throughput and time-to-first-token over FlashAttention. Moreover, L2A enables post-training pruning of highly sparse Global Attention layers, reducing KV cache memory by up to 50% with negligible performance loss. Our code is released under Apache 2.0 at https://github.com/awslabs/hybrid-model-factory/tree/main/examples/research/L2A.