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

Summary (Overview)

Primary Question: The paper investigates whether Rotary Position Embeddings (RoPE) and its base hyperparameter $\theta$ prevent retrieval heads from forming or degrade their function in transformer language models.
Key Finding 1 (H1 Refuted): Higher $\theta$ does NOT reduce the number of retrieval heads. LLaMA-3.1 ( $\theta=500{,}000$ ) has more retrieval heads (47) than LLaMA-2 ( $\theta=10{,}000$ ) with 42 heads, directly contradicting the "prevention" hypothesis.
Key Finding 2 (H2 Refuted): There is no universal "RoPE degrades retrieval" law. The utility–retrieval relationship is inconsistent: Qwen shows a significant negative effect (d = −0.49) while OLMo shows a significant positive effect (d = +0.50), with LLaMA family being null.
Key Finding 3 (Causal Axis): Retrieval depends on RoPE's frequency axis, not the norm-utility axis. Zeroing low-frequency (long-wavelength) dimensions across retrieval heads collapses recall dose-dependently (1.00 → 0.18 at k=32), an effect that is head-specific and task-specific.
Key Finding 4 (Causal Necessity): Retrieval heads are real and causally necessary: masking 87 detected heads collapses NIAH accuracy from 1.00 to 0.00, while masking random heads has no effect.

Introduction and Theoretical Foundation

Background

Long-context language models routinely retrieve specific facts from thousands of tokens, yet the mechanism is only partially understood. Two threads motivate this work:

Retrieval heads: A small set of attention heads that copy needed tokens from earlier context, whose ablation collapses long-context recall (Wu et al., 2025).
Rotary Position Embeddings (RoPE): Nearly all modern open LLMs encode position via rotation of query/key vectors at frequencies set by base hyperparameter $\theta$ . Raising $\theta$ is the standard lever for context extension, but recent work (Chiang & Yogatama, 2025) argues many RoPE dimensions become low-utility at long range.

Hypotheses Tested

H1 (Prevention): Larger $\theta$ (slower-decaying rotation) reduces the number of retrieval heads that form.
H2 (Degradation): Dimension utility (query-projection norm) identifies which RoPE dimensions retrieval depends on; zeroing low-utility dimensions leaves recall intact, while high-utility ones are load-bearing.

Models Studied (Table 1)

Model	Attn.	Layers	Heads	θ	License
LLaMA-2-7B	MHA	32	32	10,000	Llama-2
LLaMA-3.1-8B	GQA	32	32	500,000	Llama-3.1
Qwen2.5-7B	GQA	28	28	1,000,000	Apache-2.0
OLMo-2-7B	MHA	32	32	500,000	Apache-2.0

The models span both multi-head attention (MHA) and grouped-query attention (GQA) regimes, with a 100× range of $\theta$ .

Methodology

3.1 Models

Four open-weight 7–8B models chosen to vary attention regime and RoPE base while holding scale fixed.
All have 128-dimensional heads; weights loaded in 8-bit for single-GPU feasibility.
Additional models (Qwen2.5-14B, Gemma-2-9B, Mistral-7B) used in replication runs.

3.2 Needle-in-a-Haystack (NIAH) Task

Each sample embeds a needle "The secret passphrase is CODE." (5 random alphanumeric characters)
Context lengths: {1024, 2048, 4096, 8192}; needle positions: {0.1, 0.25, 0.5, 0.75, 0.9}
Recall scored as exact match of CODE in generated answer

3.3 Retrieval-Head Detection

Single-pass attention-argmax proxy: For each head, measure how often its attention argmax at the answer position falls on the needle span
Threshold: score > 0.1 labels a head as retrieval
Validation: (i) functional (causal knockout), (ii) metric (Spearman $\rho = 0.54$ correlation with stricter teacher-forced copy score)

3.4 RoPE Dimension Utility and Frequency Axis

Utility: $L_1$ norm of the corresponding row of query projection $W_Q$
Frequency: Each dimension pair $i$ $i$ has rotation frequency $\theta^{-2i/d_h}$ $θ^{- 2 i / d_{h}}$
- Low-index pairs: high frequency (short wavelength, sensitive to local offsets)
- High-index pairs: low frequency (long wavelength, distinguishable over long distances)
Important: The mapping from storage index to frequency is NOT identity under rotate_half (NeoX) convention; dimension $j$ pairs with $j + d/2$

3.5 Paired-Seed Cross-Model Protocol

Generate NIAH specifications independent of tokenizer
Keep only specifications valid for all models' context budgets
Repeat across seeds {42, 123, 2024}

3.6 Statistics

Layer-clustered permutation test: Permute labels within each layer (10,000 permutations) to avoid pseudoreplication
Cohen's d as scale-free effect size
Layer-controlled partial Spearman correlation
Benjamini–Hochberg FDR control across models
Exact McNemar test for paired comparisons

Empirical Validation / Results

Layer A: Static Multi-Model Analysis

Finding 1: Retrieval heads exist in all families (4–9% of heads), replicating across MHA and GQA architectures and across 100× range of $\theta$ .

H1 Refuted: Head count does NOT fall as $\theta$ rises:

LLaMA-3.1 ( $\theta=500{,}000$ ): 47 heads vs LLaMA-2 ( $\theta=10{,}000$ ): 42 heads
Qwen ( $\theta=1{,}000{,}000$ ): 59 heads

Finding 2: Utility–retrieval link is family-specific (Table 2):

Model	Attn.	θ	#Heads	Frac.	d	p_{clustered}	ρ_{partial}
LLaMA-3.1-8B	GQA	500,000	47	4.59%	0.07 ± 0.01	0.9998	-0.15
LLaMA-2-7B	MHA	10,000	42	4.10%	0.47 ± 0.04	0.43	0.06
Qwen2.5-7B	GQA	1,000,000	59	7.53%	-0.49 ± 0.02	0.0003	-0.21
OLMo-2-7B	MHA	500,000	87	8.50%	0.50 ± 0.01	0.0001	0.18

Qwen: significant negative (d = −0.49, p = 0.0003)
OLMo: significant positive (d = +0.50, p = 0.0001)
Both survive Benjamini–Hochberg correction (2/4 rejected)
OLMo and LLaMA-3.1 share $\theta = 500{,}000$ yet behave differently → effect not θ-driven

Finding 3 (Pseudoreplication trap): LLaMA-2 has moderate d = 0.47 with tiny naive t-test p ( $7 \times 10^{-6}$ ) but clustered p = 0.43 → not significant.

Layer B: Training Dynamics (OLMo-2)

Crystallization: Retrieval-head count flat (~100 heads) for first ~2014 B tokens, then rises sharply ~3.5× to 300–449 plateau
Onset detector: Midpoint-crossing at step 480,000 (2014 B tokens)
Correlation: Mean head utility anti-correlates with count (Pearson r = −0.75)

Layer D: Causal Validation

Head-Masking Knockout (Table 3)

Model	#Ret. heads	Baseline	Retrieval-masked	Random-masked
OLMo-2-7B (MHA)	87	1.00	0.00 (-1.00)	1.00 (-0.00)
Qwen2.5-7B (GQA)	58	1.00	0.58 (-0.42)	1.00 (-0.00)

Double dissociation: Total collapse in OLMo, partial in Qwen, zero effect for random controls.

Population-Level Frequency Patching (Table 4)

Condition (dims zeroed)	NIAH accuracy
baseline (no patch)	1.00
highest-utility ( $L_1$ norm)	1.00
random	1.00
lowest-utility ( $L_1$ norm)	0.985
highest-frequency (RoPE)	1.00
lowest-frequency (RoPE)	0.885
lowest-frequency, non-retrieval heads	1.00

McNemar test: $p = 2.4 \times 10^{-7}$ with all 23/23 discordant pairs one-sided
Bootstrap CI on accuracy difference: $[-0.16, -0.08]$
Three-seed mean: frequency effect = $-0.115 \pm 0.025$ at k=16

Dose-Response (Table 5)

k	low-freq acc.	random acc.	low-freq drop
8	0.985	1.00	0.015
16	0.885	1.00	0.115
32	0.175	0.98	0.825
48	0.115	0.87	0.885
64	0.070	0.375	0.930

Specificity Controls

Head-specificity: Non-retrieval heads control stays at ceiling (1.00)
Task-specificity: Perplexity rises only +0.9% (OLMo) vs 11.5% relative drop in NIAH recall → ratio 0.08 (< 0.33 threshold)

Replication in Qwen2.5

Frequency effect: −0.69 ± 0.03 (three-seed mean), McNemar $p < 10^{-39}$
Larger magnitude than OLMo (six times at same k)
Task-specificity partial: perplexity rises +10% (many more than OLMo's +0.9%)

Coverage Dose-Response (Table 6)

Model	#ret.	coverage	top-K	freq. eff.
Qwen2.5-14B	101	30%	30	-0.045
		50%	51	-0.585
		100%	101	-0.975
Mistral-7B	97	31%	30	-0.005
		62%	60	-0.225
		100%	97	-0.69

Matched-Coverage Universal Direction (Table 7)

Model	Attn.	#ret.	freq. eff.	ctrl	McNemar p
OLMo-2-7B	MHA	95	-0.125	1.00	$6 \times 10^{-8}$
Qwen2.5-7B	GQA	62	-0.72	0.89	$9 \times 10^{-44}$
Qwen2.5-14B	GQA	101	-0.585	1.00	$1 \times 10^{-35}$
Gemma-2-9B	GQA	45	-0.195	1.00	$4 \times 10^{-12}$
Mistral-7B	GQA	97	-0.14	0.00†	$7 \times 10^{-9}$

†Mistral's leaky control (ctrl = 0.00) is an argmax-localization artifact: under the copy-score detector it returns to 1.00.

Detector Robustness (Table 8)

Model	overlap	freq. eff. (argmax / copy)	ctrl (argmax / copy)
Qwen2.5-7B	0.14	-0.925 / -1.0	0.99 / 0.94
Mistral-7B	0.36	-0.205 / -1.0	0.585 / 1.00

Quantization Ablation (Table 9)

Level	Metric	Result
Head identity	Jaccard	0.95 (86/91)
Score agreement	Spearman ρ (argmax / mass)	0.90 / 0.99
Finding	Cohen's d (8-bit / fp16)	0.54 / 0.63, both p < 10^{-4}

Theoretical and Practical Implications

Why the Frequency Axis is Load-Bearing

RoPE assigns each dimension pair a rotation frequency $\theta^{-2i/d}$ : low-index pairs rotate quickly (short wavelength), high-index pairs rotate slowly (long wavelength). A needle-in-a-haystack lookup relates a query position to a key thousands of tokens away, so it can only read dimensions whose phase has not wrapped—the low-frequency ones. The causal result confirms this: zeroing these dimensions collapses recall.

Mechanistic Lens on Context-Length Extension

Base-scaling (raising $\theta$ ) works by stretching the wavelengths of precisely the low-frequency dimensions
This gives a circuit-level reason why tuning $\theta$ improves long-context recall: it reshapes the channel retrieval heads read
Reframes "dimension inefficiency" from liability to the locus of the knob practitioners already turn

Localized to Heads, Distributed Across a Frequency Band

Head-masking knockout: all-or-nothing (removing heads deletes recall entirely)
Dose-response: graded (recall falls smoothly as more low-frequency dimensions are zeroed)
Mechanism lives in specific heads but is encoded redundantly across a band of dimensions within them

Relation to Chiang & Yogatama (2025)

The paper adds four contributions:

Matched random and non-retrieval-head controls
Frequency-aware ordering following rotate_half layout
Dose-response curve and multi-seed significance
Direct contrast of norm vs frequency framings

Conclusion

Main takeaways:

Retrieval heads are real, emergent, and causally necessary — absent at initialization, crystallize abruptly late in pretraining, and their ablation collapses recall
H1 (prevention) is not supported — head count rises with θ across four models (though the comparison is confounded)
No universal "RoPE degrades retrieval" law — significant opposite-signed effects in Qwen and OLMo, null in LLaMA family; not θ-driven
The operative axis is frequency, not norm — retrieval depends specifically on low-frequency (long-wavelength) RoPE dimensions; this direction holds in all five models across four lineages

Future directions:

Larger model panel to turn refutation of universality into positive account
Replicate frequency dissection at longer contexts (past 4096)
Mechanistic account of how attention architecture allocates retrieval circuits (OLMo's distributed vs Qwen's concentrated)

Methodological lessons:

Layer-clustered statistics are essential to avoid pseudoreplication
Causal patching must respect task saturation (population-level over per-head)
Validate conclusions rather than the detection metric
Fix specificity thresholds in advance to guard against post-hoc tuning

Scope and caveats: Single author, no independent replication; most runs single-seed; single task (NIAH); 8-bit quantization with only one ablation model; most causal patching at 4096 tokens. The cross-model claim is the direction of the effect, not its magnitude.