Understanding the Skill Gap in Recurrent Language Models: The Role of the Gather-and-Aggregate Mechanism

• The introduction frames the central question as whether the Transformer–SSM performance gap can be traced to algorithmic capabilities in a few specific heads.It notes that MMLU requires selecting answer labels, combining factual knowledge with retrieval, while SSM retrieval limitations have often been studied as black boxes or in simplified setups.Source: §1 · pp. 1–2
• A pruned 4B-parameter Llama-3.1-8B scores 66% on MMLU, falling to 25% when one retrieval head is disabled.The authors attribute the drop to failure to retrieve answer labels, distinguishing retrieval from factual knowledge.Source: §1 · p. 2
• Transformers and SSMs use a similar G&A retrieval mechanism: Gather Heads condense input segments, then Aggregate Heads combine the summaries for prediction.The introduction illustrates gathering segments to their final newline token and assigning greater weight to the summary corresponding to the correct answer.Source: §1 · p. 2, Figure 1 · p. 2
• SSM retrieval challenges appear as smoother attention patterns that make sharp token shifts difficult and require more heads for comparable performance.Source: §1 · p. 2
• Across GSM8K, BBH, and dialogue comprehension, a small set of G&A heads critically drives retrieval-heavy benchmark performance.The authors argue that MMLU performance can coexist with poor performance on knowledge-focused benchmarks when G&A functions, but retrieval fails without functional G&A heads.Source: §1 · p. 3
• Hybrid models delegate Aggregate Head functionality to attention, and replacing one G&A instance with attention in a pretrained SSM improves MMLU more than other attention placements.The introduction presents this as a way to address SSM G&A limitations while retaining SSM components.Source: §1 · p. 3

• Prior research combines attention layers with mostly SSM layers to trade some efficiency for capabilities such as retrieval [Lieber et al., 2024; Glorioso et al., 2024; Waleffe et al., 2024; Wang et al., 2024a; Paliotta et al., 2025].The section also notes that SSMs offer sub-quadratic efficiency and competitive results on long sequences, whereas Transformers use self-attention [Vaswani et al., 2023; Dao and Gu, 2024].Source: §2 · p. 3
• The study examines Llama-3.1-8B as a Transformer, Falcon-Mamba-7B and Llamba-8B as SSMs, and Zamba-2-7B as a hybrid.Llamba-8B is a Mamba-2 SSM distilled from Llama-3.1-8B; Zamba-2-7B has a 6:1 Mamba2-to-attention-layer ratio and two shared attention layers.Source: §2 · p. 3
• The paper treats sequence models as structured matrices over time and generalizes token mixing beyond self-attention as Temporal Mixing Heads, or Mixing Heads.Prior work connects SSM token mixing to attention’s matrix-multiplication form, with Mamba-2 using fewer, larger heads than Mamba-1 [Ali et al., 2024; Dao and Gu, 2024; Bick et al., 2025; Waleffe et al., 2024].Source: §2 · pp. 3–4
• Prior work links SSM retrieval difficulties to compressed hidden states and precise-copying limits, while this paper locates the limitations in a small subset of retrieval heads.The section also describes a recall-throughput tradeoff in which attention offers stronger recall at higher cost.Source: §2 · p. 4
• The paper extends prior head analyses by proposing coordinated Gather and Aggregate roles across Transformers and SSMs, rather than retrieval by a single independently responsible head.It broadens the reported behavior of Lieberum et al.’s “Correct Letter” heads and introduces “Aggregate Head” while retaining “Gather” as shorthand for “Content-Gatherer” heads [Lieberum et al., 2023].Source: §2 · p. 4
• Concurrent work studies symbolic reasoning and answer-format adaptation [Wiegreffe et al., 2025], while the paper studies in-context retrieval across architectures.Source: §2 · p. 4
• A different study reports key–value associations in Transformers and direct retrieval in SSMs, but uses models of up to three layers on synthetic tasks; this paper studies pretrained models on natural benchmarks.The authors attribute the differing mechanistic conclusions to the different setups [Arora et al., 2025a].Source: §2 · p. 4

3.2 A Crucial Layer for MMLU

• MMLU stays stable as knowledge-task scores decline during pruning, and 30% of the original knowledge-task score can still accompany state-of-the-art MMLU performance.A sharp MMLU drop identifies a critical skill-encoding layer.Source: Figure 2 · p. 5, §3.2 · p. 6

3.1 Evaluation Benchmarks.

• MMLU can remain strong despite limited knowledge-task performance, while its letter-label format differs from benchmarks scoring complete answer texts.The case study compares MMLU with knowledge-focused tasks that emphasize factual knowledge and commonsense reasoning with minimal reliance on recent-context retrieval.Source: §3 · pp. 4–5, §3.1 · p. 6

3.3 MMLU Relies on Two Crucial Layers

• In the minimal model, removing either of its final two layers sharply reduces MMLU accuracy while only minimally affecting knowledge-focused tasks.The two layers are L16 and L17.Source: §3.3 · p. 6, Figure 3 · p. 6

3.4 Zooming In: MMLU Relies on Two Crucial Heads

• Above 66% MMLU accuracy requires both L16H22 and L17H24 in Llama-3.1-8B-Minimal; removing either or both drops accuracy to 25%.The four configurations retain both heads, only one, or neither while keeping preceding layers intact.Source: Table 1 · p. 7
• A small subset of heads drives MMLU performance: removing all noncritical heads leaves tested configurations that depend on retaining critical mixing heads.The study defines critical heads by whether ablating their outputs significantly affects MMLU scores.Source: §3.4 · p. 7

3.5 Retrieval Emerges as a Key Factor in MMLU

• On KV retrieval with 20 key-value pairs, accuracy is 85.6% intact, 76% without Gather Head L16H22, and 62% without Aggregate Head L17H24.The Aggregate Head’s removal has a larger impact, which increases with dictionary size.Source: §3.5 · p. 7, Figure 4 · p. 8

Appendix A.5.

• The identified retrieval heads align with prior Transformer findings and also appear in SSM-based language models; the authors suggest the Transformer–SSM gap may lie in a few retrieval-critical heads.The authors describe this as a targeted functional difference rather than broad modeling inadequacy.Source: Appendix A.5 · p. 8

3.6 Caveats

• Interactions can span non-adjacent layers through residual-stream steps, reducing the effectiveness of pruning-based localization.The authors suggest activation patching as an alternative for detecting these interactions.Source: §3.6 · p. 8

4 Gather-and-Aggregate Heads

• A Gather Head condenses meaningful input segments into representative tokens that a later Aggregate Head uses to retrieve prompt-relevant information.In Transformers, the segment’s final token typically summarizes its earlier tokens.Source: §4 · p. 8, §4.1 · p. 9
• The mechanism broadens earlier Transformer retrieval-head accounts beyond multiple-choice tasks and divides retrieval into complementary roles; it also appears in SSMs without explicit attention.Lieberum et al. [2023] focused on “Correct Letter” heads, while Wu et al. [2024] grouped retrieval heads under one role.Source: §4 · p. 8

4.1 Gather Heads in Transformers

• Transformer Gather and Aggregate Heads can leak attention to unrelated tokens, including attention sinks.Aggregate Heads also sometimes distribute attention to irrelevant tokens, with these limitations more pronounced in SSMs.Source: §4.1 · p. 9, §4.2 · p. 9

4.2 Aggregate Heads in Transformers

• Aggregate Heads query gathered summaries from a target token, weighting relevant ones more heavily and producing a sparse, often argmax-like combination [Tulchinskii et al., 2024].The pattern generalizes beyond multiple-choice QA to broader information retrieval.Source: §4.2 · p. 9

4.3 G&A in Transformers and SSMs: Experimental Evidence

• Pruned and masked Transformer and SSM models retain their original MMLU accuracy, showing that selected heads suffice for retrieval across architectures.The authors conclude both architectures converge on the same two-stage strategy.Source: §4.3 · p. 9
• Transformer Gather and Aggregate Heads show sharp, localized patterns, whereas SSM patterns are smoother and make segment isolation and precise retrieval harder.SSM Gather Heads can leak beyond segment boundaries, while Aggregate Heads may not sharply focus on the correct summary token.Source: §4.3 · p. 10, Figure 5 · p. 10

5.1 Visual Inspection of G&A Heads

• SSM Aggregate Heads attend to segment-final tokens but also neighboring tokens, creating diffuse maps that can dilute focus on intended summaries.The authors relate this smoothness to temporal compression in hidden-state architectures.Source: §5.1 · p. 10

5.2 Redundancy of Aggregate Heads in SSM Models

• SSMs require more Aggregate Heads than Transformers to reach similar or lower performance, consistent with weaker per-head expressivity and distributed retrieval.Llama-3.1-Minimal reaches 98% of its MMLU score with one Aggregate Head; Llamba-8B-Minimal needs two for 95%, and Falcon-Mamba-7B-Minimal needs four for 86%.Source: §5.2 · p. 11

5.3 Hybrid Models Delegate Aggregation to Attention

• The intervention reduces MMLU accuracy from 64.3% to 34.9%, while knowledge-task accuracy remains stable at 70%, a 4× drop relative to the 25% random baseline.The result demonstrates that hybrid models depend on attention layers for effective G&A implementation.Source: §5.3 · p. 12

5.4 Hybrid Replacements

• Individually distilled Llamba-8B recovers 64% of Llama’s 69% on knowledge tasks but scores 33% on MMLU versus Llama’s 67%.Replacing L17, which contains prominent G&A heads, improves MMLU; most other layer replacements have little or negative impact.Source: §5.4 · p. 12
• Replacing Llamba-8B layer 17 raises MMLU from 33% to 50%, while most layer replacements have little or negative impact.Replacing only L17H24 raises performance to 44%; including L17H27 raises it to 50%.Source: §5.4 · p. 12, Figure 7 · p. 13

5.5 Addressing G&A Limitations in SSM-Based Models

• For hybrids trained from scratch, placing attention near the network’s middle targets layers where G&A heads consistently emerge.Examples include L16 and L17 in 32-layer Llama-3.1-8B and Llamba-8B, and L35 and L36 in 64-layer Falcon-Mamba-7B.Source: §5.5 · p. 12
• For Transformer-distilled hybrids, retaining attention at strong G&A layers and distilling the rest into SSMs is a targeted design strategy.The authors say this can preserve retrieval quality where needed while benefiting from SSM efficiency elsewhere; optimal placement remains open.Source: §5.5 · p. 12

6.1 Isolating Retrieval Differences Across Architectures

• Synthetic KV-Retrieval identifies G&A heads through ablation effects while avoiding confounds from knowledge loss, enabling controlled comparisons across layers and model families.For other benchmarks, the authors preserve model knowledge and selectively disable G&A heads because pruning layers can reduce knowledge capacity.Source: §6.1 · p. 13

6.2 Quantifying the Retrieval Gap

• Transformer models show steep declines on retrieval-heavy tasks after disabling only 10–20 heads, while knowledge-focused accuracy remains stable; Llamba-8B retains over 95% of its original performance after 30 G&A heads are removed.The gradual Llamba decline is consistent with many weaker, redundant heads approximating fewer sharper Transformer heads.Source: §6.2 · pp. 13–14

6.3 Amplifying the Retrieval Gap through Task Format

• Disabling G&A heads sharply reduces accuracy on the conversational ARC variant but has little effect on multiple-choice ARC, despite identical question content.The chat format spreads relevant information across dialogue turns, favoring Transformers and hybrids over pure SSMs with more diffuse retrieval.Source: §6.3 · p. 15, Table 3 · p. 15

7 Conclusions

• Both Transformer and SSM models use a few Gather-and-Aggregate heads for retrieval, but SSMs struggle to implement them strongly with fixed-size hidden states.The authors describe this limitation as localized rather than a general weakness in language modeling.Source: §7 · p. 15
• Hybrid models offload G&A to attention layers, which the authors say have stronger algorithmic capacity.They present this as practical guidance for hybrid design.Source: §7 · p. 15
• MMLU performance is primarily governed by retrieval ability, despite the benchmark’s common interpretation as a broad knowledge test.The authors connect benchmark success to algorithmic skills, particularly retrieval.Source: §7 · p. 15
• The role of G&A in in-context learning and copying, and its possible relation to other retrieval heads, remains open.The authors specifically mention Amplification Heads as a possible part of a broader retrieval process.Source: §7 · p. 15

A Complementary Material for Section 3

• The appendix adds results examining knowledge and retrieval behavior across layers and heads in the studied models.These results provide further insights into the models explored in Section 3.Source: A Complementary Material for Section 3 · p. 20

A.1 Detailed Performance Across Layers

• Progressive layer pruning shows relatively stable MMLU scores alongside declining knowledge-task performance, followed by a sharp MMLU drop at a critical layer.The figure description attributes gradual knowledge-score decline to distributed knowledge representation across layers.Source: Figure 8 · p. 20, §A.1 · p. 20

A.2 Layer Removal: Mixer and MLP Components

• Removing layers with their MLPs or retaining the MLPs produces similar performance trends, suggesting mixer components play a more prominent role in retrieval encoding.The reported pattern applies to Llama-3.1-8B and Llamba-8B.Source: §A.2 · p. 21, Figure 9 · p. 21

A.3 Performance in Minimal Models

• In minimal models, two layers are critical for MMLU, while pruning either of the last two significantly lowers MMLU scores and leaves knowledge metrics largely unaffected.The reported layer-removal pattern also appears across the models examined.Source: §A.3 · p. 21, Figure 10 · p. 22

A.4 Critical Heads Drive MMLU Performance Across Models

• Across architectures, removing one or two critical heads drops MMLU accuracy to about 25%; four Falcon-Mamba attention channels recover 52.3%.These heads comprise fewer than 0.1% of total heads, while knowledge-focused benchmark scores remain nearly unchanged across configurations.Source: §A.4 · p. 22, Table 4 · p. 23

A.5 KV-Retrieval Recipe

• Answer scoring distinguishes heads enabling correct retrieval from heads associated with hallucinations, unlike generation-only analysis.The passage contrasts this finding with the conclusion of Wu et al. [2024].Source: §A.5 · p. 24
• Under Answer Scoring, unnatural no-space prompts reveal retrieval-head signals with 20–30 key-value pairs, while space-delimited prompts require longer dictionaries.Wu et al. [2024] used 35K-token contexts to observe similar patterns.Source: §A.5 · p. 24

A.6 Retrieval Heads Across Model Families

• Top retrieval heads cluster in upper Transformer layers, are more diffuse and individually weaker in SSMs, and cluster around shared-attention layers in Zamba2 hybrids.The Transformer heads often overlap with those critical for MMLU.Source: §A.6 · p. 24