Head Dimension Matters in Language Models for Long-Context Retrieval

Seiya Tokui; Shoichiro Yamaguchi; Sosuke Kobayashi

cs.LGOct 23, 2026

Head Dimension Matters in Language Models for Long-Context Retrieval

Authors: Seiya Tokui, Shoichiro Yamaguchi, Sosuke Kobayashi

Organizations: Preferred Networks, Inc., Tokyo, Japan

Abstract

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 $\pm$ 7.75	35.46 $\pm$ 3.60	41.21 $\pm$ 5.86	37.78 $\pm$ 4.15	31.28 $\pm$ 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

\pm

sample standard deviations across three runs; each 3B allocation has one run.

Panel of Figure 1 — 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 $\pm$ sample standard deviation across three 0.6B runs. In all the settings, d512 achieves the highest (average) score at every evaluated length.

Panel of Figure 2 — 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 $\pm 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.

Panel of Figure 4 — 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.

Panel of Figure 5 — 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 $\pm 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.

Panel of Figure 7 — 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 $\pm 1$ sample standard deviation.

Model	PT	Full CPT
0.6B-d128	$1.908\pm 0.001$	$1.911\pm 0.009$
0.6B-d256	$1.908\pm 0.001$	$1.888\pm 0.042$
0.6B-d512	$1.913\pm 0.004$	$1.912\pm 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

\pm

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)
$W_{Q},W_{K},W_{V},W_{O}$ 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	$\Delta_{\mathrm{loss}}$
0.6B-d128	$1.911\pm 0.009$	$1.884\pm 0.019$	$0.027\pm 0.028$
0.6B-d256	$1.888\pm 0.042$	$1.878\pm 0.020$	$0.011\pm 0.062$
0.6B-d512	$1.912\pm 0.011$	$1.881\pm 0.018$	$0.031\pm 0.028$

Appendix

Table 10: Training NTP loss at the final logged interval of Full and Frozen CPT. Entries are means

\pm

sample standard deviations across three runs.

\Delta_{\mathrm{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.

Panel of Figure 11 — 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 $\pm$ 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 ).

Panel of Figure 12 — 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.

Panel of Figure 13 — 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	$\Delta 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

(\mathrm{layer},\mathrm{query\ head})

sets.

\Delta 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.

Long-Context Retrieval

Long-context retrieval finds and uses relevant evidence across lengthy documents, conversations, and model contexts.

See all 45 papers

Can Retrieval Heads See Images? Multimodal Retrieval Heads in Long-Context Vision-Language Models
May 26, 2026Aaron Branson Cigres Li, Zhaowei Wang, Yu Zhao +9Vision-Language Models VLM Interpretability
1HKUST · University of Edinburgh · 4CUHK +2
Does RoPE Prevent or Degrade Retrieval Heads? A Mechanistic Analysis Across Model Families
Jun 19, 2026Cengizhan BayramRotary Positional Embeddings Long-Context Retrieval
Independent Researcher
Learning When to Attend: Conditional Memory Access for Long-Context LLMs
Mar 18, 2026Sakshi Choudhary, Aditya Chattopadhyay, Luca Zancato +4Self-Attention Long-Context Language Modeling
Department of Electrical and Computer Engineering, Purdue University, USA · AWS Agentic AI, USA

May 26, 2026J/K move · Enter open · S save

Can Retrieval Heads See Images? Multimodal Retrieval Heads in Long-Context Vision-Language Models

Aaron Branson Cigres Li, Zhaowei Wang, Yu Zhao, Yiming Du+8

1HKUST · University of Edinburgh · 4CUHK+2

Vision-Language Models VLM Interpretability Vision-Language Retrieval Visual Document Retrieval Multimodal Foundation Models Long-Context VLMs+1

Open paper AskPDF arXiv

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.

Abstract

Long-Context Retrieval

Explore similar work

Can Retrieval Heads See Images? Multimodal Retrieval Heads in Long-Context Vision-Language Models

Does RoPE Prevent or Degrade Retrieval Heads? A Mechanistic Analysis Across Model Families

Learning When to Attend: Conditional Memory Access for Long-Context LLMs