Understanding the Skill Gap in Recurrent Language Models: The Role of the Gather-and-Aggregate Mechanism
Authors: Aviv Bick
Eric P. Xing
Albert Gu
Venue: ICML 2025 poster (2025)
Keywords: LLM;Transformers;Mamba;MMLU;In-context-Learning;mechanistic-interpretability
Primary area: deep_learning->large_language_models
TL;DR: SSMs struggle with structured reasoning tasks due to limitations in implementing Gather & Aggregate mechanisms, critical for retrieval and learning.

Abstract: State-space models (SSMs) offer efficient alternatives to Transformers for long sequences, but their fixed-size recurrent state limits capability on algorithmic tasks, such as retrieving past context.
In this work, we examine how in-context retrieval operates in Transformer- and SSM-based language models and find that both rely on a Gather-and-Aggregate (G&A) mechanism:
a Gather Head extracts relevant information from context, which an Aggregate Head integrates into representation.
In both architectures, G&A concentrates in a few heads, forming bottlenecks even for simple retrieval.
For example, disabling a single Gather or Aggregate Head in a pruned Llama-3.1-8B impairs retrieving the correct answer letter in MMLU, reducing its accuracy from 66% to 25%.
Moreover, this retrieval bottleneck can obscure knowledge demands of tasks as the pruned model succeeds on MMLU with functioning G&A heads yet fails on other knowledge benchmarks.
The bottleneck similarly extends to tasks where SSMs typically underperform, like GSM8K, BBH, and dialogue.
We show that SSMs' retrieval challenges manifest in these heads, creating smoother attention patterns instead of the sharp transitions effective G&A requires.
Thus, the Transformer-SSM retrieval gap exists in just a few heads, rather than the entire language model.
% Result 3: Analyzing Hybrid models
This suggests a unified explanation for Transformer vs. SSM performance gap while showing how to merge their strengths.
We find that pretrained hybrid models, where SSMs are combined with attention layers, delegate the role of Aggregate Heads to attention.
Similarly, replacing a single G&A head in a pretrained SSM with an attention variant boosts retrieval and benchmark scores.

Understanding the Skill Gap in Recurrent Language Models: The Role of the Gather-and-Aggregat Mechanism
Aviv Bick
Eric P. Xing 12 Albert Gu
Abstract
1. Introduction
State-space models (SSMs) offer efficient alterna- tives Transformers for long sequences, but their fixed-size recurrent state limits capability algo- rithmic tasks, such retrieving past context. In this work, we examine how in-context retrieval op- erates Transformer- and SSM-based language models and find that both rely on a Gather-and- Aggregate (G&A) mechanism: a Gather Head extracts relevant information from context, which Aggregate Head integrates into representation. In both architectures, G&A A concentrates in few heads, forming bottlenecks even for simple re- trieval. For example, disabling a single Gather or Aggregate Head in a pruned Llama-3.1 impairs retrieving the correct answer letter in MMLU, reducing its accuracy from 66% 25%. Moreover, this retrieval bottleneck can obscure knowledge demands of tasks the pruned model succeeds on MMLU with functioning G& heads yet fails on other knowledge benchmarks. The bottleneck similarly extends to tasks where SSMs typically underperform, like GSM8K, BBH, and dialogue. We show that SSMs' retrieval chal- lenges manifest in these heads, creating smoother attention patterns instead of the sharp transitions effective G&A requires. Thus, the Transformer- SSM retrieval gap exists in just few heads, rather than the entire language model. This suggests a unified explanation for Transformer vs. SSM performance gap while showing how to merge their strengths. We find that pretrained hybrid models, where SSMs are combined with attention layers, delegate the role of Aggregate Heads to attention. Similarly, replacing a single G&A head in a pretrained SSM with an attention variant boosts retrieval and benchmark scores.
Transformers have driven major breakthroughs in language modeling, but their quadratic scaling with sequence length has renewed interest in recurrent alternatives like state-space models (SSMs). SSMs offer linear scaling with fixed-size memory while remaining competitive across many tasks, However, replacing a large, explicit history cache with a compact state representation poses fundamental tradeoff. SSM-based language models often struggle with the skill of "in-context retrieval"-the ability to precisely locate and extract tokens presented earlier in the sequence, such as recalling phone number mentioned several paragraphs ago (Jelassi et al., 2024). This "algorithmic" challenge, which Transformers handle effectively with direct access to past tokens, has been identified as key factor in the performance gap between Transformer models and SSM variants (Arora et al., 2023).
Previous work SSM retrieval limitations has treated language models as black boxes or focused on simplified theoretical setups (Wen et al., 2024; Jelassi et al., 2024; Waleffe et al., 2024). Meanwhile, Transformer research shows that algorithmic capabilities, including retrieval, of- ten concentrate specific attention heads (Lieberum et al., 2023; Wu al., 2024). This raises key question: Can the Transformer-SSM performance gap be traced to algorithmic capabilities in few specific heads? Clarifying this will enable a more targeted comparison of the two architectures, helping determine whether SSMs can serve as viable build- ing blocks for language modeling or require fundamental modifications.
The MMLU benchmark (Hendrycks et al., 2021) exemplifies this performance gap. While SSMs match Transformers on many knowledge benchmarks, they significantly underper- form on MMLU unless trained substantially longer (Waleffe et al., 2024). This gap stems from MMLU's evaluation for- mat: rather than scoring the likelihood of full-text answers, models must select letter labels (A, B, C, D) correspond- ing the correct option. This creates dual challenge requiring both factual knowledge and retrieval capabilities (Lieberum et al., 2023).
'Carnegie Mellon University, Pittsburgh, PA, USA *MBZUAI, Abu Dhabi, UAE 'Cartesiaaai, Correspondence to: Aviv Bick <abick@cs.cmm
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
To isolate these distinct challenges, we show that both archi-

---

Understanding the Skill Gap in Recurrent Language Models
-Broh otor
-8
36
54
the central node 802.11
is the central node of 802.11
wireless operations.
wireless operations. An WPA nnB Access Point WAP
A. WPA nB. Access Point WAP
D. Access Port AnAnswr::
D. Access Port AnAnswer
(a) Gather Head
(b) Aggregate Head
Figure 1. An illustration Gather Head (left matrix) at 116622 and Aggregate Head (right matrix) at L1 7H24 Llama-3.1-8B for multiple-choice question. The Gather Head identifies critical segments and summarizes each segment to its last token "\nn" as shown in Figure I(a), condensing key information into the residual stream. The Aggregate Head then combines these summaries as illustrated in Figure 1(b), giving significantly higher score the second "",n." which corresponds to the correct answer (B).
tectures have clear division of labor: factual knowledge extracted across layers and heads, but only a small number of specialized heads retrieve the correct labels from input context. This becomes particularly evident in our experi- ments where a ---arammeter pruned Llama-3.1-88 achieves a strong MMLU score of 66%, which plummets to 25% when disabling a single retrieval head. Moreover, these heads obscure the true knowledge demands such tasks- despite strong MMLU performance, the pruned model per- forms poorly knowledge- focused benchmarks. Simi- lar dependence extends beyond MMLU to mathematical reasoning (GSM8K), multi-step logical reasoning (BBH), and dialogue formats (ARC-C Chat), demonstrating that retrieval capabilities, not knowledge deficits, drive SSM underperformance in these domains.
(Figure 1 1b)).. This finding indicates that the two language models rely on similar retrieval strategies, despite architec- tural differences. Consequently. their retrieval gap reduces to the effectiveness with which few specific heads imple- ment the G&A mechanism.
Comparing G&A implementations, we find SSMs' known retrieval challenges manifest how key heads implement G&A. SSMs exhibit smoother attention patterns, making sharp token shifts difficult and requiring more heads achieve comparable performance. Moreover, how hybrid models (Glorioso et al., 2024; Lieber al., 2024), which integrate a few attention layers into predom- inantly SSM-based architectures, overcome SSM's limita- tions. When trained from scratch, hybrid models naturally delegate Aggregate Head functionality to attention layers. Conversely, in pretrained SSM models, replacing a single G&A instance with attention-baa based one yields signifi- cantly greater MMLU improvement than any other attention placement.
show
To understand how these specialized heads function, we investigate their underlying mechanism. We find that both Transformer- and SSM-based language models converge on a similar retrieval framework: a Gather-and-Aggregat (G&A) mechanism. Specifically, a Gather Head identifies and condenses segments of input tokens into a single vec- tor representation (Figure I(a)), which an Aggregate Head subsequently processes to extract the necessary information
Our experiments are publicly available.'
'nttps://githul com/ coon/goombalab. Gather-and andher-and-AAgroggtte
2

---

Understanding the Skill Gap in Recurrent Language Models
2. Related Work
states, unlike Transformers, which cache full context. This structural difference impacts retrieval performance across benchmarks. Waleffe et al. (2024) found that Mamba- based models require extensive training to perform well on MMLU, while Wen et al. (2024) identified weaknesses in in-context retrieval tasks, such as associative recall. Je- lassi et al. (2024) theoretically and empirically confirmed that SSMs have difficulty with precise copying. Arora et al. (2023) demonstrated that associative recall capabilities ac- count for the majority of performance differences between attention- -asion and gated-convolution models. Park et al. (2024) shows that SSMs underperform on tasks requiring non-standard retrieval, and proposed a hybrid architecture to address this limitation. Arora et al. (2025) highlighted a recall- throouhhh tradeoff, with attention offering stronger recall higher cost. Blouir et al. (2025) complemented this view by enhancing retrieval in SSMs through training improvements without architectural changes. Unlike previ- ous work, we show these limitations originate from a small subset of retrieval heads rather than the entire model.
Sequence Models. Transformer-based models with self- attention (Vaswani et al., 2023; Brown et al., 2020; Touvron et al., 2023) dominate language modeling for their scala- bility and performance. Recently, recurrent models like structured state-space models (SSMs) (Gu et al., 2022; Gu & Dao. 2023; Dao & Gu, 2024) gained interest for their sub-quadratic efficiency and competitive results on long se- quences. Other recurrent models include (Dao et al., 2022; Sun et al., 2023; Katharopoulos et al., 2020; Yang et al., 2024a; Zhang et al., 2024b; Beck et al., 2024; Qin et al., 2024; Peng et al., 2024; Yang et al., 2024b). We use "SSM" to refer to these models as well, our findings broadly apply.
Despite SSMs' efficiency, they struggle with tasks requiring precise token interactions, such as in-context retrieval. Their fixed-size hidden state limits memory capacity, making it difficult retain detailed past information. To address these limitations, recent work (Lieber et al., 2024: Glorioso et al., 2024; Waleffe et al., 2024; Ren et al., 2024; Wang et al., 2024a; Paliotta et al., 2025) has adopted hybrid approach, integrating attention layers with a majority of SSM layers, trading efficiency for improved capabilities such as retrieval.
Mechanistic Interpretability of Mixing Heads. Most interpretability studies focus on attention heads in LLMs (Elhage et al., 2021; Lieberum et al., 2023; Zhang et al., 2024a; Olsson et al., 2022; Yu et al., 2024; Merullo et al., 2024; Rai et al., 2024; Zheng et al., 2024; Wu al., 2024; Tulchinskii et al., 2024; Sanford et al., 2024), while rel- atively little attention has been given to recurrent mecha- nisms (Ali et al., 2024). Olsson et al. (2022) and Elhage et al. (2021) showed that induction heads emerge to support in-context learning. Subsequent work traced their training dynamics (Bietti et al., 2023), and mapped factual recall circuits (Meng et al., 2023).
Throughout this study, we examine four architectures: Ilama-3. -8B--rrtruct (Touvron al., 2023), standard Transformer; Falcon-Mamba- -B--ntructtt (Zuo et al., 2024), a Mamba-1 SSM without MLP layers; Slumba-8B-Untier (Bick et al., 2025). a Mamba-2 SSM distilled from Llama- 3.1-8B using MOHAWK (Bick et al., 2024); and Zamba-2- 7B-Instruct (Glorioso et al., 2024), a hybrid model with a 6:1 ratio of Mamba2 to attention layers and only two shared attention layers. For brevity, we refer to these as Llama- 3. 3-8. Falcon-Mamba-77, Llamba-8B, and Zamba-2-7B, respectively.
Lieberum et al. (2023) identified two key heads in Chinchilla-7001 Content Gatherer and Correct Letter- that, with MLPs, support label prediction in multiple-choice tasks. Wu et al. (2024) found that fewer than 5% of heads, termed Retrieval Heads, are critical for retrieval, Chain-of- thought, and reasoning.
Central to understanding the performance differences be- tween these architectures is token mixing: how models propagate and combine information across token positions. In self-attention, tokens selectively aggregate context via learned weights, while SSMs use structured recurrences with more restricted patterns. Ali et al. (2024); Dao & Gu (2024) connect SSM token mixing to attention's 's matrix- multiplication form, showing Mamba-1 mimics many low- dimensional attention heads. Mamba-2 (Dao & Gu, 2024; Bick et al., 2025; Waleffe et al., 2024) uses fewer, larger heads, increasing expressivity. These insights allow us to view all sequence models as structured matrices over time. Thus, we generalize token mixing beyond self-attention, re- ferring to these mechanisms as Temporal Mixing Heads (or simply Mixing Heads).
We refine the understanding of how heads support retrieval by showing that, in both Transformers and SSMs, retrieval arises from a coordinated mechanism- not single head. One head identifies the target location, while another aggre- gates its content. This view extends prior work in three ways. First, the "Correct Letter" heads identified by Lieberum et al. (2023) exhibit broader retrieval behavior than previ- ously reported, going beyond multiple-choice selection. See- ond, retrieval depends coordinated interactions between heads, challenging the interpretation in Wu et al. (2024) that treated individual heads as independently responsible. Third, we demonstrate that this mechanism also arises in DeMss-esspte their lack of attention attention-highligttnng its ar- chitectural generality. We introduce the term Aggregate
Retrieval Tradeoff's Between Transformers and SSMs. SSMs rely on compressed, continuously updated hidden
3

---

Understanding the Skill Gap in Recurrent Language Models
Head to better reflect these empirical findings, while retain- ing "Gather" as shorthand for the "Content-Gatteree" heads described Lieberum et al. (2023).
ity gaps between attention-based models and SSM variants (Arora et al., 2023), support our main claim: the general performance gap between Transformer- and SSM-based language models can be traced to specific algorithmic capa- bilities within a small subset of heads.
Concurrent with our work, Wiegreffe et al. (2025) extended Lieberum et al. (2023)'s findings by analyzing training dy- namics and cases of poor model performance. They in- troduced synthetic tasks to disentangle knowledge from multiple-choice formatting and identified mechanisms by which models adapt different answer symbols. While their goals align with ours, we focus on the algorithmic aspects of in-context retrieval in multiple-choice settings across architectures.
Our analysis proceeds bottom-up: we begin by contrasting MMLU with these wwwleege-foousee tasks (Section 3. 1) then localize MMLU performance a single critical layer (Section 3.2). We show coordination between two layers (Section 3.3), narrow the effect to small set of heads (Sec- tion 3.4), and demonstrate that these implement a shared retrieval mechanism across models (Section 3.5).
3. A Case Study: Retrieval in MMLU
3.1. Evaluation Benchmarks.
Throughout the paper. we evaluate performance on two categories of benchmarks: MMLU and what we term Knowledge-Focused Tasks. MMLU's format distinctively presents multiple-choice questions, requiring models to pro- cess lengthy contexts, then identify the correct option by outputting the corresponding letter label (A, B, C, D). While MMLU is widely used, studies have noted that strong performance may not reflect true understanding or knowl- edge (Alzahrani et al., 2024; Wang et al., 2024b; Zhao et al., 2024).
The MMLU benchmark (Hendrycks al., 2021) spans 57 di- verse tasks and is often viewed test of world knowledge. However, its multiple-choice format also places demands on retrieval: models must map question-answet context to the correct letter label (A))), as illustrated below:
is the central node of 802. 11 wi wireless operations.
A. WPA
B. Access Point
C. WAP
In contrast, Knowledge- Focused Tasks include ARC- Challenge and ARC-Easy (Clark al., 2018), PIQA (Bisk et al., 2019), Winogrande (Sakaguchi et al., 2019), Open- BookQA (Mihaylov et al., 2018), and HellaSwag (Zellers et al., 2019). These benchmarks primarily assess factual knowledge and commonsense reasoning with minimal re- liance on recent context retrieval. The core difference lies in how predictions are evaluated: MMLU requires se- lecting letter label tied to the correct answer, while the other benchmarks score the likelihood of complete answer texts shittinh the focus from label retrieval to how well the model's 's internal knowledge supports each option.
D. Access Port
Answer:
Despite performing well on many factual benchmarks, SSMs underperform MMLU unless trained substantially longer (Waleffe et al., 2024). Given their known retrieval limitations (Jelassi et al., 2024), we ask whether this gap stems from a failure of knowledge or of retrieval. To in- vestigate, we analyze three representative models: Llama- -aae---ama------iiiiii (Mamba-1) 1), and Llamba-881 -B----B--Un SBttet (Mamba-2). We com- pare their performance on MMLU versus benchmarks that emphasize factual reasoning with minimal reliance in- context retrieval. For clarity, we omit instructional and embedding-tyine suffixes from model names in our discus- sion.
-.nttetInttruct (Transformer),
3.2. A Crucial Layer for MMLU
Motivated by the struggle of SSMs with MMLU (Wal- effe et al., 2024) and by the fact that Transformers use specialized components to solve multiple-choice questions (Lieberum et al., 2023), we examine whether solving this task utilizes different components in SSM-based language models as well. Since all language models utilize residual skip connections, eliminating layer could reveal its spe- cific contribution while keeping the rest of the information flowing.
Our investigation reveals two key findings:
1. Models can perform well on MMLU with limited knowl- edge, suggesting that the task primarily reflects in- context retrieval for the evaluated models. This adds to concerns raised by prior work (Wang et al., 2024b; Zhao et al., 2024).
2. Both Transformer- and SSM-based models encode re- trieval abilities in small set of specialized heads, ex- tending prior work that focused solely on Transformers (Lieberum et al., 2023; Wu et al., 2024).
We iteratively remove layers from the end of the model and evaluate performance on both MMLU and the knowledge- focused tasks. Appendix A.1 presents performance trends for Llama-3.1-8P as layers are progressively removed, and
These insights, combined with earlier findings about perplex-
4

---

Understanding the Skill Gap in Recurrent Language Models
Appendix A.I extends this analysis to Falcon-Mamba-7B and Llamba-8B. For Llama-3.1-8B and Llamba-8B, we ab- late both the mixer and its downstream MLP, since the latter likely depends on the former. As shown in Appendix A.2, similar trends hold even when only the mixer is removed.
Table I demonstrate that MMLU depends on just a few mixing heads rather than the model's broader knowledge representation. Disabling them drops performance to near- random. Remarkably, in Llama-3, 3.-BB,, a single head is sufficient to achieve over 66% MMLU accuracy, and in four Mamba-l channels reach 52.26%.
All architectures reveal a striking pattern: while perfor- mance on knowledge- focused benchmarks declines steadily as layers are pruned, MMLU performance remains stable- until single layer is removed, at which point it drops sharply. This suggests that a specific layer encodey the critical skills required for MMLU. High MMLU scores can persist even when most knowledge capacity is lost, highlight- ing the role of targeted retrieval components over general knowledge.
3.5. Retrieval Emerges as a Key Factor in MMLU
Earlier results showed that two attention heads from separate layers work together to solve MMLU across architectures. We now provide evidence that these heads form part of a broader in-context retrieval mechanism, validated through a synthetic key-value (KV) retrieval task.
In this task, the model is given dictionary and must recall the value of queried key, for example,
3.3. MMLU Relies on Two Crucial Layers
Memor ze the following dictionary:
Based on insights from Appendix 1. we refer to the mini- mal model as the smallest subnetwork that retains all layers up to (and including) the layer before the sharp drop in MMLU performance. We now repeat the previous experi- ment, this time on the minimal model, with a slight modi- fication. Instead of cumulatively removing one layer at a time from the end of the network, we remove single layer, evaluate the model performance without it, and then restore the layer before proceeding to the next.
present: 50
institute:0 scal scallops:84
neuropsychiatzy:67
The value of the key ' 'scallops" is
To succeed, the model must accurately store mappings (e.g.. "scallops + 84") and, when queried, retrieve the correct value ('84") by identifying 'scallops'' in its stored represen- tations without confusion.
Appendix A.3 illustrates the results of this experiment in Llama-3.1-88 Palcon-Mamma-77 and Llamba-8B as well. The findings reveal surprising insight on the minimal model: not one, but two layers are critical for the perfor- mance of MMLU. Unlike earlier layers, the last two layers of the minimal model show no correlation between MMLU score and knowledge task performance. Pruning either layer causes significant drop in MMLU accuracy while causing only a minor drop in scores on knowledge-focusee tasks, underscoring their essential role in the task.
We tested 55 configurations with growing numbers of key- value pairs, each using 1,000 synthetic samples. For each configuration, we evaluated Llama-3.1 1-BB in three settings: intact, with L1 7H24 4 removed, and with L1 6H22 removed (the heads identified as critical in Section 3.4). The full evaluation setup and prompt format are discussed in Ap- pendix A.9.
As shown in Figure 5, removing either L16H22 or L1 117224 4 significantly impairs retrieval, confirming their central role. Notably, the impact of removing L17H24 increases with dictionary size, suggesting its contribution grows with task complexity. These findings extend previous research: the retrieval heads we identified align with and generalize the Content Gatherer head described by Lieberum et al. (2023) and the retrieval mechanisms documented by Wu al. (2024). Importantly, our work demonstrates that similar retrieval mechanisms also emerge within SSM-based lan- guage models. The strong link between retrieval and MMLU performance suggests that the Transformer-SSM. gap may stem from differences a small number of retrieval-critical heads rather than fundamental architectural limitations. This reframes the performance gap targeted func- tional difference- SSMs may lack specific retrieval capabilities- capabilities-raaherr than a broad modeling inadequacy.
3.4. Zooming In: MMLU Relies Two Crucial Heads
Building Section 3.3, we now zoom in from layers to individual heads, analyzing how specific attention heads within the two critical layers affect model performance.
Using the minimal model setup from that section (with only the last two layers active), we ablated individual heads by zeroing their outputs and measured the effect on MMLU and knowledge-foctee focused tasks. This revealed a small subset of heads whose removal sharply reduced MMLU accuracy while leaving other mnwwledee-focusee benchmarks unaf- fected. We label these as critical heads.
To confirm their importance, we evaluated four configu- rations: (1) no critical heads retained, (2) only last-layer critical heads, (3) only lecond-last-laeer critical heads, and (4) critical heads in both layers. The results presented in
5

---

Understanding the Skill Gap in Recurrent Language Models
4. Gather-and-Aggregan Heads
sometimes leaks unrelated tokens- partis-partiiularly "attention sinks" (Xiao et al., 2024) used stabilize processing over long contexts.
We now examine how the critical retrieval heads identi- fied in Section 3 across architectures implement the same Gather-and- Aggreggte (G&A) mechanism. The mechanism unfolds in two stages. First, Gather Head isolates mean- ingful segments within the input-such as multiple-choice answers-aan and summarizes each into representative token. Then, a later Aggregate Head uses these summaries to retrieve the information most relevant to the prompt.
4.2, Aggregate Heads in Transformers
Once the Gather Head condenses segments into representa- tive tokens, an Aggregate Head-in subsequent layer- retrieves and combines these summaries to inform the model's prediction. This head typically attends from a target token (c.g., the final "Answer:" position) over all gathered summaries, assigning greater weights to the most relevant one (see Figure 1(b)). This yields sparse linear combination over segment summaries, often resembling an argmax- x-lkke operation (Tulchinski et al., 2024):
While previous work has described retrieval-oriented heads in suun the "Correct Letter" heads iden- tified by Lieberum al. (2023) these analyses focused narrowly on multiple-choice tasks. Our findings show that such heads are part of a broader retrieval strategy that gener- alizes beyond this format. This also refines the interpretation from Wu et al. (2024), who grouped retrieval heads under a single role; in contrast, we find they specialize in com- plementary functions within a coordinated process. Finally, show that this same mechanism emerges in SSM-based models, despite their lack of explicit attention. These mod- els exhibit smoother and less discrete patterns, but still learn to implement a G&A strategy using different architectural tools.
" "Answer:"
in
nn
In Transformer models, such heads were previously termed "Correct Letter" Heads (Lieberum et al., 2023). due to their consistent weighting of the correct answer's summary token. However, the aggregation pattern generalizes be- yond multiple-choice QA, enabling information retrieval in broader contexts.
As with Gather Heads, this mechanism is imperfect. Ag- gregate Heads sometimes distribute attention to irrelevant tokens, including attention sinks. These limitations become more pronounced SSMs, as we discuss next.
This section unpacks the mechanism in detail. We begin with how it operates in Transformers (Sections 4.1 and 4.2), then present experimental evidence that it also underlies retrieval in SSMs (Section 4.3)
4.3. G&A in Transformers and SSMs: Experimental Evidence
4.1. Gather Heads in Transformers
A Gather Head compresses semantically related tokens within segment into a single representative token. In Transformer models, this typically occurs through atten- tion patterns that cause the final token of each segment to "attend to" all prior tokens in the same segment, effectively summarizing them into a compact representation (see Fig- ure ((a)).
The G&A pattern is well-docemented Transformers but do SSM models like Mamba implement the same structure? We test this via two-part experiment that isolates and evaluates specific heads: (1) Pruning: We prune each model to a minimal configuration (see Section 3.3), zeroing out all heads in the final two layers except those suspected to implement G&A (see Table I); (2) Masking: At runtime, we constrain the preserved heads' attention patterns: Gather Heads may only attend within each segment and to its final token, and Aggregate Heads may only attend to those final summary tokens (See Figure 6).
For example, consider the multiple-choice format intro- duced in Section 3:
A.
WPA 22 + A. WPA in
In all 8B (Mamba-2), and Falcon-Mamha-775 (Mamba-1)-the pruned and masked configurations match the original MMLU accuracy. This confirms that the selected heads are sufficient to drive retrieval and that the G&A mechanism is not unique to Transformers: despite architectural differ- ences, SSM-based models converge on the same strategy.
(Transformer), Llamba-
Here, the newline token (in) at the end of each answer choice acts as a summary token. This transformation re- duces a segment from multiple tokens to one, simplifying downstream computation.
Lieberum et al. (2023) observed this behavior in 70B Chin- chilla, where attention heads collect tokens like "A", ".", "w", "PA". and "\nn" into the final newline token. The head's attention pattern essentially performs form of local aggre- gation, summarizing the contents of each segment.
Nevertheless, the attention patterns differ in texture. As shown in Figure 7, attention maps in Mamba-based models are smoother and less sharply localized. This softness limits their ability to isolate segments cleanly: Gather Heads in
While effective, this process is not always clean: attention
6

---

Understanding the Skill Gap in Recurrent Language Models
SSMs often leak attention outside segment boundaries, and Aggregate Heady struggle to sharply focus on the correct summary token. These imperfections stem from the contin- uous nature of SSM hidden states and are explored further in Section 5.
Llamba-8B depends on a single moderate one (L17H31), with others contributing little. Rather than concentrating aggregation in few strong heads, it spreads the load across weaker ones.
3. More Heads Needed to Match Accuracy. Llama-3, -3.1 Minimal reaches 98% of its MMLU score with one Ag- gregate Head, Llamba-88 Blamba-8B-Mar 101-----Miiial needs two for 95%, and Falcon-Mamba--BB-- 7a-Mimmaa---iiimmal requires four for 86% (see Table 1). This points to lower per-head capacity in SSMs.
5. The G&A Bottleneck in SSMs
The previous section established that both Transformers and SSMs rely on small number of heads to implement two- stage Gather-and-Aggee an--Aggeeate (G&A) mechanism. We now examine how this mechanism breaks down in SSMs, pro- viding a mechanistic account of their retrieval limitations.
5.3. Hybrid Models Delegate Aggregation to Attention
We begin by visually inspecting G&A behavior in SSMs (Section 5.1), showing that Aggregate Heads attend broadly rather than selectively. We then quantify redundancy in SSM aggregation (Section 5.2), finding that multiple heads are needed to match the performance of fewer, stronger Transformer heads. Next, we show that hybrid models of- fload aggregation to attention layers (Section 5.3); ablating just six heads causes a 4x drop in MMLU. We then iden- tify these heads as the bottleneck (Section 5.4): replacing single SSM layer restores MMLU from 33% to 50%. Fi-
Attention layers in hybrid models are known be critical for retrieval and copying tasks (Lieber et al., 2024; Glorioso et al., 2024; Waleffe et al., 2024). We hypothesize that during training, Aggregate Heads naturally concentrate in the attention layers, as attention's 's greater expressivity makes it better suited for aggregation.
To test this, we evaluate Zamba2-7B (Glorioso et al., 2024), the strongest 7B-scale hybrid model date. Table 2 iden- tifies attention heads critical for general retrieval; here, we further select nine that are especially important for MMLU: L47H 4, 6, 8, 18, 21, 25} and L59H(17. 21). Each head is disabled by zeroing out the final row of its atten- tion matrix- targeted intervention that impairs aggrega- tion while leaving other capabilities, such as gathering and contextual mixing, largely intact. All SSM layers remain unchanged. preserving any aggregation they may contribute.
nally, we explore how targeted attention during training or distillation can restore G&A in SSMs (Section 5.5).
5.1. Visual Inspection of G&A Heads
Figures 6 and 7 show that Mamba Aggregate Heads cor- rectly attend to segment-final tokens but also attend to neigh- boring tokens unnecessarily. This likely introduces noise and reduces aggregation effectiveness. This smooth behay- ior reflects hidden-state architectures, where information compression over time makes it harder to capture the global nature of Aggregate Heads.
This intervention reduces MMLU accuracy from 64.3% to 34.9% while knowledge task accuracy remains stable at 70%. Since MMLU's random baseline is 25% this 4x drop demonstrates that SSMs alone cannot sustain strong aggregation- hybrid models depend on attention layers to implement the G&& mechanism.
5.2. Redundancy of Aggregate Heads in SSM Models
We find that SSM-based models require more Aggregate Heads than Transformer-based models reach similar or even lower levels of performance. This suggests that SSMs compensate for limited per-head expressiveness by distribut- ing the retrieval task across multiple weaker components. Three observations support this claim:
5.4. Hybrid Replacements
To further test whether SSMs struggle to implement Ag- gregate Heads, examine how well an SSM mixer can replicate attention layer. Using Llamba-8B (Bick et al., 2025), we distill each Mamba-2 layer independently from Llama-3.1-88 following Bick et al. (2024), aligning outputs via L2 loss without end-to-end optimization. Specifically, i-th Llamba block receives the output of Llama's (i 1)-th layer and matches Llama's i-th layer.
1. Gradual Performance Decay. As shown in Table 4 and discussed in Section 6, ablating heads in Llamba causes slow, steady drop in performance-unlit the sharp drops seen in Transformers Transormers-iipllyngg overlapping functionality and less specialization.
After assembling the full model from the individually dis- tilled layers. the distilled model recovers 64% of Llama's 69% on knowledge tasks, while its MMLU performance remains 33%, compared to 67% for Llama. This suggests that a key retrieval-related capability was not captured.
2. Weaker, Less Isolated Heads. Llamba-8B, distilled from Llama-3.1-8B via MOHAWK (Bick et al., 2024), adds G&A heads in layers 16 and 17. But while Llama- 3. 3B1 relies on two strong heads (L.17H24 and 17H2777,
To determine whether the observed performance gap re-
7

---

Understanding the Skill Gap in Recurrent Language Models
lated to retrieval and specifically the G&A mechanism, conducted hybrid replacement experiment. In this ex- periment, we systematically replaced each layer in the dis- tilled Llamba-8B model with its corresponding layer from L.ama-3.1-88 evaluated its impact on MMLJU (without fine-tuning), and then reverted the model before proceeding to the next layer.
These observations provide practical recipe for hybrid design: introduce attention selectively-either guided by where G&A heads naturally emerge when training from scratch, or retained only in key locations when distilling from a Transformer. We leave further investigation into op- timal placement and architectural variants of hybrid models to future work.
The results, shown in Figure 8, reveal that most layer re- placements had little even a negative impact on MMLU. However, replacing L1 7, which contains two prominent G&A heads at H24 and H27 (L1 1H24 4 was identified in Table 1), significantly improved performance from 33% to 50%. To isolate their effect, we ran additional hybrid exper- iments. Swapping only 1.1 7H24, while keeping the rest of the SSM heads intact, raised performance 44%. Including L1 7H27 as well increased it to 50%. This confirms that the contribution stems from the Aggregate Heads within L1 7, and further supports that individual SSM layers struggle to implement the Aggregate Head within G&A-reinforcing our claim that this limitation contributes to the performance gap between Transformers and SSMs.
6. Transformer-SSM Gap Beyond MMLU
We previously showed that both Transformers and SSMs rely on few heads to implement a two-stage Gather-and- Aggregate (G&A) mechanism for in-context retrieval. We now extend this analysis to more benchmarks and models, reinforeing that the broader retrieval gap stems from specific algorithmic limits rather than general language modeling ability. Our analysis has three parts: Section 6. 1 identifies G&A heads across models; Section 6.2 evaluates how archi- tectural choices affect performance; Section 6.3 examines how task format influences these effects.
6.1. Isolating Retrieval Differences Across Architectures
5.5. Addressing G&A Limitations in SSM-Based Models
In Section 3, we showed that MMLU performance can be severely degraded by removing a single critical head-even when the model performs poorly on knowledge-focused tasks. This scenario is specific to MMLU: a task where strong retrieval alone can yield high accuracy, even when the model's 's knowledge is degraded. In contrast, most bench- marks require both retrieval and factual knowledge, making it harder to isolate the contribution of retrieval alone.
Efforts to improve G&A in pure SSMs face a core constraint: retrieval suffers due to fixed-size memory in RNN-style de- signs (Jelassi et al., 2024; Wen et al., 2024). Rather than working within these limits, we explore hybrids that intro- duce attention layers to target retrieval bottlenecks directly. Our analysis reveals two empirical patterns that guide hy hybri model design:
As a result, the minimal-model strategy used in MMLU fails on other tasks: pruning layers diminishes knowledge capac- ity and lowers scores across the board, masking the effect of retrieval-specific mechanisms. To avoid this confound, we instead take different approach: rather than pruning the model, we preserve its knowledge and selectively disable only the G&A heads.
1. Training hybrids from scratch. G&A heads consis- tently emerge in the middle of the network-for exam- ple, at layers L16 6 and L out of 32 in Llama-3.1 1--8BB and L8, and L335 and 6 out of 64 in Falcon- Mamba-8B. This suggests that attention is most valuable in mid layers, where retrieval operations naturally occur. As result, when designing hybrid models from scratch, placing attention layers near the middle offers targeted way to support G&A without relying on full attention throughout the network.
As detailed in Section 3.5, we use the synthetic KV-Retrieval task to identify G&A heads by ablating each one and mea- suring its effect on retrieval accuracy (Table 2). A perfor- mance drop signals the head's involvement in G&A. This method, also used by Wu et al. (2024), isolates retrieval ability without relying on external knowledge, enabling us to trace G&A heads across layers and architectures in a controlled, generalizable way.
2. Distilling hybrids from a Transformer. As shown in Section 5.4, replacing individual layers in a distilled SSM with their attention-based counterparts from the teacher yields little benefit-except at L17, where MMLU im- proves from 33% to 50%. This indicates that the perfor- mance gap is concentrated in the Aggregate Head. Based on this, a practical strategy is to retain attention in the lay- ers where strong G&A heads appear and distill the rest into SSMs. This allows the hybrid to recover retrieval quality where matters most, while still benefiting from the efficiency of SSMs elsewhere.
6.2. Quantifying the Retrieval Gap
Having identified the eetrieval-riitinn heads, we now mea- sure their practical impact across benchmarks by incremen- tally disabling them in three model families. These include Transformer models (Llama-3.1-8P) 33.1-8B-nnstruc and Llama-3.2 3.2- 3B Instructtt (Touvron et al., 2023)) as baselines for effective
8

---

Understanding the Skill Gap in Recurrent Language Models
G&A implementation, pure SSM models (Llamba-3B and Llamba-8B (Bick et al., 2025)) based on Mamba-2, and hybrid models (Zamba2-2.7B and Zamba2-7B (Glorioso et al., 2024)) that use Mamba-2 layers with shared attention modules. For Zamba2, we ablate only attention heads, as they perform retrieval (Section 5.3). Each shared head is ablated in one instance to isolate its effect.
fload G&A attention layers, leveraging their stronger algorithmic capacity. This explains why even few atten- tion layers can meaningfully boost SSM models, providing practical guidance for hybrid design.
Our results also reframe benchmark interpretation. Al- though MMLU is often seen as broad knowledge test, we find that MMLU performance is primarily governed by retrieval ability. This highlights the role of algorithmic skills, and particularly retrieval, in benchmark success.
Table 4 summarizes the results. Transformer models show sharp performance drop on retrieval heavy tasks after re- moving just 10-20 heads, while knowledge-task accuracy remains
Important open questions remain: while narrowing the re- trieval gap is a key step, the role of G&A in in-context learning and copying warrants further study. It also remains to seen whether G&A is part of broader retrieval process involving other head types, such as Amplification Heads (Lieberum et al., 2023), which enhance signal propagation.
confirming retrieval is concentrated in a few heads. Zamba2-7B behaves similarly, with consistent effects across shared-heat head instances, reinforcing their uni- fied G&A roles (Table 2). In contrast, Llamba-8B degrades slowly. Even after removing 30 G&A heads, retains over 95% of its original performance. This gradual drop reflects redundancy: SSMs distribute retrieval across many weaker heads, unlike the focused G&A in cenrffoomrrr-oonsisttete with Section 5.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential socictal consequences of our work, none which we feel must be specifically highlighted here.
6.3. Amplifying the Retrieval Gap through Task Format
Beyond architecture, task format can amplify or reduce the retrieval gap between Transformers and SSMs. We previously noted this in MMLU, where models select a letter label rather than score full answer texts.
References
Ali, A., Zimerman, I., and Wolf, L. The hidden attention of mamba models, 2024. URL htt.s: /ttp:////axxv.oo////// abs/2403.01590.
To further illustrate, we compare two formats of the same content: the original ARC-Challenge multiple-choice task and its conversational variant, ARC-Challenge Chat. As shown in Table 3, disabling G&A heads has little effect on the multiple-choice version but sharply reduces accuracy in the chat variant. The conversational format increases retrieval demands by spreading relevant content across di- alogue turns, favoring architectures with sharper retrieval (Transformers and hybrids) over those with more diffuse mechanisms (pure SSMs).
Alzahrani, N., Alyahya, H. A., Alnumay, Y., Alrashed, S., Alsubaie, S., Almushaykeh, Y., Mirza, F,, Alotaibi, N., Altwairesh, N., Alowisheq, A., Bari, M. S., and Khan, H. When benchmarks are targets: Revealing the sensitivity of large language model leaderboards, 2024. URL https: / aaaaiiivvoo org/ abs / 2202244022002 01781.
Arora, S., Eyuboglu, S., Timalsina, A., Johnson, I., Poli, M., Zou, J.. Rudra, A., and Ré, C. Zoology: Measuring and improving recall in efficient language models, 2023. URL https: /arxiv..vvv/ o httg//bb///31/ 04927.
This highlights a key point: retrieval pressure is shaped not just by what is asked, but by how it's presented. When context is more distributed, effective in-context retrieval becomes more critical even if the underlying knowledge stays the same.
Arora, S., Eyuboglu, S., Zhang, M., Timalsina, A., Alberti, S. Zinsley, D.,, Zou, J., Rudra, A., and Ré, C. Sim- ple linear attention language models balance the recall- throughput tradeoff. arXiv:2402. 18668, 2024.
7. Conclusions
We show that both Transformer and SSM models develop a -nddee--nn--ggggaattt (G&A) mechanism across a small number of heads to support in-context retrieval, which is crucial for performance on many benchmarks. While SSMs struggle to implement strong G&A heads due their fixed- size hidden state, this limitation is localized and does not reflect a general weakness in language modeling.
Arora, S., Eyuboglu, S., Zhang, M., Timalsina, A., Alberti, S., Zinsley, D.,, Zou, J., Rudra, A., and Ré, C. Sim- ple linear attention language models balance the recall- throughput tradeoff. 2025. URL https: //arxiv/ org/ abs 22a.b///422... 18668.
Hybrid models provide further insight: they naturally of-
Beck, M., Pöppel. K., Spanring, M., Auer, A., Prudnikova, O., Kopp, M., Klambauer, G., Brandstetter, J., and
9

---

Understanding the Skill Gap in Recurrent Language Models
Hochreiter, S. xlstm: Extended long short-term mem- ory, 2024. URL htt.tt/ppe..aaaa.oooo///aaa/a//a 04517.
Elhage, N., Nanda, N., Olsson, C., Henighan, T., Joseph, N., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly,
T., DasSarma, N., Drain, D., Ganguli, D., Hatfield- Dodds, Z., Hernandez, D., Jones, A., Kermion, J., Lovitt, I.,, Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., and Olah, C. A math- ematical framework for transformer circuits. Truns- former Circuits Thread, 2021. https://ranssfrmmrr watpiissppbb202///aaeew//////// html..
Bick, A., Li, K. Y., Xing, E. P., Kolter, J. Z., and Gu, A. Transformers ssms: Distilling quadratic knowledge to subquadratic models, 2024. URL https://arxiv.
org/abs/ aataab///488 2408. 10189.
Bick, A., Katsch, T., Sohoni, N., Desai, A., and Gu, A. Llamba: Scaling distilled recurrent models for efficient language processing, 2025. URL https:/ arxiv.
Glorioso, P., Anthony, Q. Tokpanov, Y., Whittington, J., Pilault, J., Ibrahim, A., and Millidge, B. Zamba: A compact 7b ssm hybrid model, 2024. URL https:// a arxiv. org.orrg/aaoo.. abs 2405. 16712.
Bietti, A., Cabannes, V., Bouchacourt, D., Jegou, H., and Bottou, I. Birth of transformer: A memory viewpoint, 2023. URL htttt::///taiiii orwp/iiiorrraoooooooo abs/ 2306. 00802.
Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces, 2023.
Gu, A., Goel, K., and Ré, C. Efficiently modeling long sequences with structured state spaces, 2022.
Bisk, Y., Zellers, R., Bras, R. L., Gao, J., and Choi, Y. Piqa: Reasoning about physical commonsense in natural language, 2019. URL https:: aaaaa.iivvv.oo./.aa/aaaa org/ abs/ 1911.11641.
Hendryeks, D., Burns, C., Basart, S., Zou, A., Mazeika, M. Song, D., and Steinhardt, J. Measuring massive multitask language understanding, 2021. httpt:ss: / www.xii.orggaaaoo h0tpab///./0003030
Blouir, S., Smith, J. T. H., Anastasopoulos, A., and Shehu, A. Birdie: Advancing state space models with reward- driven objectives and curricula, 2025. URL https: //ar///aaa /arxivv... org/ abs aba//aab121 /2112 01030.
Jelassi, S., Brandfonbrener, D., Kakade, S. M., and Malach, Repeat after me: Transformers are better than state space models at copying. 2024. URL hwt.https//eee o //s/abb//24022 01032.
Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan,
J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert- A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D. M., Wu,
Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, F. Transformers are mns:: Fast autoregressive transformers with linear attention, 2020.
J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D. Language models are few-shot learners, 2020.
Lieber, O., Lenz, B., Bata, H., Cohen, G., Osin, J., Dalmedi- gos, I., Safahi, E., Meirom, S., Belinkov, Y., Shalev- Shwartz, S., Abend. O., Alon, R., Asida, T., Bergman, A., Glozman, R., Gokhman, M., Manevich, A., Rat- ner, N., Rozen, N., Shwartz, E., Zusman, M., and Shoham, Y. Jamba: A hybrid transformer-mambe lan- guage model, 2024. URL abs Abb/2403.19882
Clark, P., Cowhey, I., Etzioni, O.,, Khot, T., Sabharwal, A.. Schoenick, C., and Tafjord, O. Think you have solved question answering? try are, the ai2 reasoning challenge, 2018. URL hwt.h:tt///aaaaeet /ttpe:////xvvvvvo org/ abs/ 1803. 05457.
Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H., Kaiser, L., Plappert, M., Tworek, J., Hilton, J., Nakano, R., Hesse, C., and Schulman, J. Training verifiers to solve math word problems, 2021. URL org/ org/abs/ abs
Lieberum, T., Rahtz, M., Kramár, J., Nanda, N., Irving, G., Shah, R., and Mikulik, V. Does circuit analysis inter- pretability scale? evidence from multiple choice capabili- ties chinchilla, 2023. URL https: httttps://axxi.orr/ abs/ abs/2307.. 09458.
011101111 14168.
Dao, T. and Gu, A. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. In International Conference on Machine Learning (ICML), 2024.
Meng. K., Bau, D., Andonian, A., and Belinkov, Y. Locat- ing and editing factual associations in gpt, 2023, URL https: : / /aa/aaiiiivi///aa///aaa////aaava /ttpa/iww.....aaaaaaaaa/aaa 2at.hb/s222222222222
Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Ré, C. Flashat- tention: Fast and memory-efficiee exact attention with io-awareness, 2022.
Merullo, J., Eickhoff, C., and Pavlick, E. Circuit component
reuse across tasks in transformer language models, 2024. URL https:://// //wxxixxvv wwt.oo/g/g///ooooooo 08744.
10

---

Understanding the Skill Gap in Recurrent Language Models
Mihaylov, T., Clark, P., Khot, T., and Sabharwal, A. Can suit of armor conduct electricity? new dataset for open book question answering, 2018, URLpst//:/ arxiv. org/ abs /wt.//11..00000 02789.
Sakaguchi, K., Bras, R. L., Bhagavatula, C., and Choi, Y. Winogrande: An adversarial winograd schema challenge at scale, 2019. URL https: /aaaaaxxivvv.o/../// org/abs/ a 10641.
Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y., Chen,
Sanford, C., Hsu, D., and Telgarsky, M. Transformers, parallel computation, and logarithmic depth, 2024, URL 09268.
A., Conerly, T., Drain, D., Ganguli, D., Hatfield-Dodds, Z., Hernandez, D.. Johnston, S., Jones, A., Kernion, J., Lovitt, L.,, Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., and Olah, C. In-context learn- ing and induction heads. Transformer Circuits Thread, 2022. hutps//transffomerrrccciiiii hwtp:errrrrrrrttsuuto///iiiioioo wwtppp////nncooneeeeeeee laaning-and-induction----eeeeeeo hwwpeiiioooneenddddeeeeee.. html.
Sun, Y., Dong, L., Huang, S., Ma, S., Xia, Y., Xue, J., Wang, J., and Wei, F. Retentive network: A successor to transformer for large language models, 2023,
Suzgun, M., Scales, N., Schärli, N., Gehrmann, S., Tay, Y., Chung, H. W., Chowdhery, A., Le, Q V., Chi, H., Zhou, D., and Wei, J. Challenging big-bench tasks and whether chain-of-thought can solve them, 2022. URL https: aaw.iv org /tabs//2110 09261.
Paliotta, D., Wang, J., Pagliardini, M., Li, K. Y., Bick, A., Kolter, J. Z., Gu, A., Fleuret, F., and Dao, T. Think- ing slow, fast: Scaling inference compute with distilled org/ abs/
reasoners, 2025. URL https::/taaaat
2502. 20339.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., Rodriguez, A., Joulin, A., Grave, E., and Lam- ple, G. Llama: Open and efficient foundation language models, 2023.
Paperno, D., Kruszowski, G., Lazaridou, A., Pham, Q N., Bemardi,, R., Pezzelle, S., Baroni, M., Boleda, G., and Fernández, R. The lambada dataset: Word prediction requiring broad discourse context, 2016. URL https: //arxiv. otg/ abs 1666. 06031.
Tulchinskii, E., Kushnareva, L., Kuznetsov, K., Voznyuk, A., Andriaainen, A., Piontkovskaya, I., Burnaev, E., and Barannikov, S. Listening to the wise few: Select-and- copy attention heads for multiple-choice qa. 2024. URL /ttps://axxiv.orrg/ www.iiiwww..ora.ooooooooooooee https: arxiv
Park, J., Park, J., Xiong, Z., Lee, N., Cho, J., Oymak, S., Lee, K., and Papailiopoulos, D. Can mamba learn how learn? comparative study on in-context learn- ing tasks, 2024. URL https::///// / aaaxi.iiiivv./.i.//aaaa//aaaavva org/ abs/ 2402. 04248.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need, 2023. URL https: httpss//axxv.orro abs / 1t/./1706..05 03762.
Peng, B., Goldstein, D., Anthony, Q., Albalak, A., Alcaide, E., Biderman, S., Cheah, E., Du, X., Ferdinan, T., Hou, H., Kazienko, P., GV, K. K., Kocoñ, J., Koptyra, B., Krishna, S., Jr., R. M., Lin, J., Muennighoff, N., Obeid, F., Saito, A., Song, G., Tu, H., Wirawan, C., Wožniak, S., Zhang, R., Zhao, B., Zhao, Q., Zhou, P., Zhu, J., and Zhu, R.-J. Eagle and finch: Rwkv with matrix-valued states and dynamic recurrence, 2024. URL https:// arxiv. org/ abs /abs/244.... 05892.
Waleffe, R., Byeon, W., Riach, D., Norick, B., Kor- thikanti, V., Dao, T., Gu, A., Hatamizadeh, A., Singh, S., Narayanan, D., Kulshreshtha, G., Singh, V., Casper, J., Kautz, J., Shoeybi, M., and Catanzaro, B. An empiri- cal study of mamba-basee language models, 2024. URL https:/ //ax/iivvv./ www.iiv.rr/.oooooo abs 2006.4466000 07887.
Qin, Z., Yang, S., Sun, W., Shen, X., Li, D., Sun, W., and Zhong, Y. Hgm2: Gated linear rnns with state expan- sion, 2024. URL phtttss://aattt 07904.
Wang. J., Paliotta, D.,, May, A., Rush, A. M., and Dao, T. The mamba in the llama: Distilling and accelerating hybrid models, 2024a. URL htttt:///axxii... abs / 2
org/ abs//404
Rai, D., Zhou, Y., Feng, S., Saparov, A., and Yao, Z. A practical review of mechanistic interpretability for transformer-base language models, 2024. URL https: //arxiv. . org/ abs/ 0bss2240 2207 02646.
Wang. W., Jain, S., Kantor, P., Feldman, J., Gallos, L., and Wang, H. Mmlu-sr: A benchmark for stress-testing reasoning capability of large language models, 2024b. URL https: /atpiv.xivv//////////////////aia org/ atg//ba//// ht/aa/bb//24.... 15468.
Ren, L., Liu, Y., Lu, Y., Shen, Y., Liang, C., and Chen, W. Samba: Simple hybrid state space models for effi- cient unlimited context language modeling, 2024. URL www.oorgg//////ooooo aa//bb///2000 07522.
Wen, K., Dang, X., and Lyu, K. Rnns are not transformers
(yet): The key bottleneck on in-context retrieval, 2024. URL hww.::ttps//// //axii///xvvv... wwt./oo/////////aoo//
https://arxivv wwwpxiiivvrrrvveeoooeoee

---

Understanding the Skill Gap in Recurrent Language Models
Wiegreffe, S., Tafjord, O., Belinkov, Y., Hajishirzi, H., and Sabharwal, A. Answer, assemble, ace: Understanding how Ims answer multiple choice questions, 2025. URL https: / arxiv.
Zuo, J., Velikanov, M., Rhaiem, D. E., Chahed, I., Belkada, Y., Kunsch, G., and Hacid, H. Falcon mamba: The first competitive attention-free 7b language model, 2024. URL 05355.
2407 15018.
Wu, W. Wang. Y., Xiao, G., Peng, H., and Fu, Y. Re- trieval head mechanistically explains long-context factual- ity, 2024. URL 15574.
Xiao, G., Tian, Y. Chen, B., Han, S., and Lewis, M. Efficient streaming language models with atten- tion sinks, 2024. URL https:::///are org/ abs/ 2309. 17453.
Yang. S., Kautz, J., and Hatamizadeh, Gated delta net- works: Improving mamba2 with delta rule, 2024a. URL https: /awxiv.ivvvvvv////////a////aaav/ orggabs/////ooooooortt..g/....bt 2412. 04464.
Yang, S., Wang, B., Shen, Y., Panda, R., and Kim, Y. Gated linear attention transformers with hardware-efficcent training, 2024b. URL https::///aaat //appii///xiv...// org/ abs/ 2312.06635.
Yu, S., Song, J., Hwang, B., Kang, H., Cho, S., Choi, J., Joe, S., Lee, T., Gwon, Y. L., and Yoon, S. Correcting negative bias in large language models through negative attention score alignment, 2024. URL https::/p/// /ttpiixxvvvvroo//./ abs/2408 00137.
Zellers, R., Holtzman, A., Bisk, Y., Farhadi, A., and Choi, Y. Hellaswag: Can machine really finish your sentence?, 2019. URL htlttss:/arxiv..oo abs/
Zhang, W., Wan, C., Zhang, Y., ming Cheung, Y., Tian, X., Shen, X., and Ye, J. Interpreting and improving large language models in arithmetic calculation, 2024a. URL attps::/arxa
blog/ab//// org/ abs /4099 01659.
Zhang, Y., Yang, S., Zhu, R., Zhang, Y., Cui, I.., Wang, Y., Wang, B., Shi, F., Wang, B., Bi, W., Zhou, P., and Fu, G. Gated slot attention for efficient linear-time sequence modeling, 2024b. URL https: //ap:i////vvv.a/ia org/ abs/ 2409. 07146.
Zhao, Q. Huang, Y., Lv, T., Cui, L., Sun, Q. Mao, S., Zhang, X., Xin, Y., Yin, Q., Li, S., and Wei, F. Mmlu-cf: A contamination- multi-task language understanding benchmark, 2024. URL https: /ttp///axvvv...////a//ee/ea.aav. org/ abs/ 2412.15194
Zheng, Z., Wang, Y., Huang, Y., Song, S., Yang. M., Tang, B., Xiong, F., and Li, Z. Attention heads of large language models: A survey, 2024. URL https: //a raaii..coomm www..orggabbb////
03752.
12

---

Understanding the Skill Gap in Recurrent Language Models
A. Complementary Material for Section 3
In this section, we present additional results from experiments conducted on the models explored Section These results provide further insights into the behavior of knowledge and retrieval capabilities across layers and heads different models.
A.1. Detailed Performance Across Layers
We extend the analysis presented in Section 3.2 by providing performance data on other models, which show similar trends. As demonstrated in Appendix A.1, the pattern stable MMLU scores despite gradual decline knowledge task performance is consistent across all models tested. Notably, the sharp drop in MMLU performance upon removal of specific layers reaffirms the criticality of certain layers in encoding retrieval skills.
Llama-3.1--8 3.1-8B: Knowledge vs. MMLU Distributions Across Layers
ecourny
0.
0.26
Layara kept
MMMUU
Llamba-88: Knowledge vs. MMLU Distributions Across Layers
1.0
9.5
0.81
3.30
6:29
0:27
0.26
0.24
9:23
0.22
6:21
0.20
0.18
0:18
0.17
6:18
0.16
0.14
0.18
0.22
6:11
0.0
08
08
67
08
06
04
08
42
00
Layera
Enowwlldge
MIHLU
Falcon-Mamba-78: Knowledge vs. MMLU Distributions Across Layers
0.0
0.4
0.0
Figure Residual performance trends for knowledge tasks the MMLU benchmark as layers are progressively pruned layers Llama-3. Llamba-8B, and Falcon-Mamba-773 The leftmost bar represents the full model with all layers intact, while subsequent bars correspond to models with layers pruned incrementally from end of network. The plot highlights three key observations: (1) MMLU scores remain relatively stable despite decline in general knowledge tasks performance, indicating reliance task task-specific skills; (2) sharp drop in MMLU performance identifies a critical layer encoding essential skills; and (3) the gradual decline in knowledge tasks' scores reflects the distributed nature of knowledge representation layers.
A.2. Layer Removal: Mixer and MLP Components
For both Llama-3.1-88 and Llamba-8B, we performed experiments where entire layers, including both the mixer and MLP components, were removed. This approach was taken under the assumption that ML. components likely depend on outputs
13

---

Understanding the Skill Gap in Recurrent Language Models
from the mixer components, which are crucial for information flow across layers. The results from these tests, presented in Appendix A.2, show that similar performance trends observed whether the MLP components retained This suggests that the mixer components play more prominent role in encoding retrieval capabilities, while the MLP components do not significantly alter the overall performance on knowledge tasks and MMLU benchmarks. This finding reinforces the idea that retrieval skills are primarily concentrated in specific layers and do not rely heavily on the MLP components within those layers.
Llama-3.1-88 Knowledge vs. MMLU Distributions Across Layers (no MLP removed)
0.0
0.0
2-31
0-30
0.20
0-22
0-24
0-23
D-22
0-13
0:10
D-17
0-15
0:13
0-12
0-11
0.10
0-5
0-4
Layers kept
Ancoractge
Llamba-8B: Knowledge vs. MMLU Distributions Across Layers (no MLP removed)
0.R
0.8
0:31
30
29
0.28
27
0-26
0:26
04
0:28
0:22
0.21
0.28
0:18
0-18
0:17
0:18
0:15
0:14
0:18
0:12
0:11
0.10
0.9
06
07
06
0.
04
02
0.1
Layora
Unzawledge #NNLLU
Figure Residual performance trends for knowledge tasks the MMLU benchmark as layers are progressively pruned from Llama- 3.1-8B, Llamba-8B but keeping the MLPs intact. This plot yields similar results to Appendix A.
A.3. Performance in Minimal Models
The results from Section 3.3 are further elaborated in Appendix A.3. The performance of Llama-3. 3-B8BB and other models when subjected to layer removal shows that, for all models, not one, but two layers are crucial for MMLU performance. The surprising robustness of MMLU performance in the minimal models highlights the importance fine-tumin seee-tnnng layer-wise architecture to achieve strong retrieval performance while retaining knowledge capabilities.
14

---

Understanding the Skill Gap in Recurrent Language Models
Knowledge
Layers
Liamba-88- M0-aa-mm---ii--:: Knowledge ve. MMLU Distribution Across Layers
1.0
0.0
0.8
6.7
6.6
0.5
0.4
0.3
0.2
0.1
11
12
13
14
0.0
10
17
Layar
Krowlodge MILU
Knowledge
MMLU
Falcon-Monn Mamionn-Mmaa----o 7B- Minimal: Knowledge MMLU Distributions Acroas Loyers
0.8
0.7
0.5
0.5
0.4
0.3
removed
MMLU
Knowledge
Figure 4. Residual performance trends for the MMLU benchmark and knowledge tasks metrics in the minimal model (defined with respect to MMLU) for Llama-3.1-88 Llamba-8B,
this experiment, each layer independently pruned, and the resulting performance is assessed. The results reveal that two layers criticul for MMLU performance, pruning either of the last two layers causes a significant drop in MMLU scores while leaving knowledge metrics largely unaffected. This finding highlights the unique role of these layers in task-specific skills rather than general knowledge representation.
Palcon-Mamba-
A.4. Critical Heads Drive MMLU Performance Across Models
To assess how in-context retrieval is localized within specific heads, we examine the MMLU performance of three models- 1.. Llamba-8B, and Marccon-Mamaa-BB--aa selectively retaining or removing mixing heads the final two layers each model, while preserving all earlier components.
targeted head ablation. Table 1 presents the results of
Across architectures, we observe a sharp and consistent pattern: removing just one or two critical heads reduces MMLU accuracy to near-random levels (~25%), while retaining them restores much of the original performance. Notably, these heads represent fewer than 0.1 %% of the total head count, highlighting the extreme concentration of retrieval ability within tiny subset of the model. Meanwhile, scores nowledge-focusser benchmarks remain nearly unchanged across all configurations, reinforcing that these heads are specifically responsible for in-context retrieval rather than general language modeling or factual recall.
This result supports the broader claim that the Gather-and ander-nn--Aan aad-Aggregate mechanism emerges in just a handful of components and suggests that hybrid and efficient model designs can selectively preserve only these heads.
15

---

Understanding the Skill Gap in Recurrent Language Models
Figure Evaluation Llama-3.1-8P on the KV-Retrieval task under three settings: unmodified, with only L1 6822 removed, or with only L1 7H24 removed (the Gather and Aggregate Heads identified in Section 3.4, respectively). Performance is reported across 55 configurations with increasing numbers of key-value pairs, using the original (non-pruned) model. Each bar shows accuracy for a given dictionary size-for example, with 20 pairs, the original score is 85.6%, removing only the Gather Head (11 6H22) drops accuracy to 76%, and removing only the Aggregate Head (L1 7H24) drops to 62%. These results underscore two key points we build on later: (1) Retrieval is more sensitive Aggregate Head removal than Gather Head removal, underscoring its importance; (2) The Aggregate Head becomes increasingly important as task complexity grows.
Model
Heads Retained
Metrics (%)
Knowledge Tasks
Prev. Layers
Layer X
Layer X+1
MMLU
0-31
22
24
66.32
39.09
0-31
24
24.36
39.18
Llama-3.1-8B
0-31
22
0
25.59
39.21
0-31
0
0
25.56
39.21
0-31
10
3, 9
58.92
39,08
0-31
0
3, 9
42.17
39.16
Llamba-8B
0-31
10
()
38,02
39.00
0-31
0
29.82
36.76
3186, 5143
1565, 1906
0-8191
52.26
37.23
6607, 7305
3873, 4925
1565, 1906
0-8191
35.91
37.25
Maton--aamaa-B8
3873, 4925
3186, 5143 6607, 7305
0-8191
0
39.39
37.26
0-8191
0
24.98
37.28
Table 1. Impact of Critical Head Removal Across Model Architectures. We evaluate how retaining or removing specific heads in the final two layers affects MMLU and knowledge task performance for three models: Llama-3.1-8B, Llamba-8B, and Dalcon-Mamba-7B. Each row shows performance under a different configuration, where the "Heads Retained" columns indicate which heads are active in the two final layers ("Layer X" and "Layer X+1") while all carlier heads preserved. The heads shown correspond to the critical Gather and Aggregate Heads identified in prior analyses. For Llama and Llamba models, retaining both heads restores ncar-original MMLU performance (66.3% and 58.9%, respectively), while removing either or both causes performance collapse near-random levels. In Falcon-Mamba-78.. just four attention channels (among thousands) across the final two layers are sufficient to recover 52.3% MMLU accuracy- again, comparable full-model performance. Knowledge task scores remain mostly unchanged across configurations, reinforcing that these heads specifically control retrieval rather than factual knowledge access. This shows that across architectures, in-context retrieval is driven by just a few heads, whose removal alone accounts for majority the performance gap.
A.5. Retrieval Heads Across Model Families
Table 2 presents the 50 most retrieval-relevant heads for each model, ranked by their individual contribution to KV-Retrieval accuracy. Several trends emerge. In Transformers, top heads cluster in upper layers, often overlapping with those critical for MMLU. In SSM-based models, retrieval heads are more diffuse and individually weaker. Zamba2 hybrids concentrate performnff heads around shared-attentie attention layers, suggesting these layers anchor retrieval otherwise SSM-heavy architectures.
16

---

Understanding the Skill Gap in Recurrent Language Models
Table Attention Head Ablation Study: Performance impact on KV-Retrieval task when individual attention heads are removed from various language models. Each head was systematically knocked out and the model's accuracy was evaluated KV-Retrieval benchmark with dictionary size of 20 across 1000 samples. After measurement, cach head was restored before testing the next one. Results show 50 most significant attention heads from each model, sorted by accuracy, revealing which heads contribute most critically to the KV-Retrieval capability across Zamba2, Llama, and Llamba model families. For Zamba22 models, marks layers with the first shared attention, marks those with second.
Zambu2-2.78
Zamba2-7B
Llama-3.2-3B
Llama-3,1-8B
Llamba-31
Llamba-SB
Layer
Head
Acc
Layer
Head
Acc
Layer
Head
Acc
Layer
Head
Ace
Layer
Head
Acc
Layer
Head
Acc
0.139
0.382
187
12
0. 199
23F
28
0.139
12
0.405
8
0.494
0 0.220
21
2
0.446
12
0.303
0.139
14
6
0.524
13
0.514
30
0.296
16
0.460
30r
0.328
47F
0.139
18
0.545
0.519
17
0.301
22
0.465
6'
26
0.380
39
28
0.176
13
S
0,589
13
17
0.596
11
23
0.317
27
15
0.471
30
0.471
18"
25
0.426
11
21
0.266
9
0.625 4
14
0.680
19
13
0.324
11
25
0.476
30
20
0.430
0,266
0.639
13
0.685
7
0.327
21
0,482
18'
26
0.433
351
0.266
13
0.648
12
0.730
20
0.329
16
18
0,487
301
31
0.472
21
0,266
0.688 18
5
14
0.732
0.333
11
31
0.487
24
0.479
0.281
0.657
0.761
0.334
10
0,490
30
12
0.482
71
21
0,281
14
0
0.667
II
0.784
13
29
0.335
14
0.492
0.483
0,424
20
0.816
28
11
0.492
30'
26
0.488
23
0.424
3 0.671
14
0.828
16
0.338
15
16
0.495
361
3
0.496
0.24
18
0.677
19
0,830
0.340
12
0.495
36
477
0,424
0.839
0.344
0.497
30
16
0.499
0,430
0.685
7
16
0.844
11
16
0.348
14
0.497
6'
16
0.508
297
0.430
0.689
10
0.855
0.352
22
28
0.498
24
D.514
0,430
13
16
0.693
16
0.855
19
0.352
9
0.501
24
0.521
S31
0,430
0.696
0.859
0
22
0.502
24
0.521
0434
II
0.698
0.860 7
22
0.353
9
20
0.502
30!
24
21
0.526
11
0,433
2 0,703
10
20
0.870
9
0.355
8
17
0.504
18T
20
0.532
0,473
17 0.705
0.878
0.355
15
0.505
6'
0.532
23R
0.473
13
0,705
0.879
23
0.356
10
0,505
24t
25
0.533
0,733
13
22
0,705
0.883
26
0.357
5
24
0,505
12
0.883
24t
22
0.538
35
18
0.473
4
0.713
0.884
1I
0.359 24
15
17
0,506
361
16
0.539
0,473
18 0.713
0.885
9
0.360
12
18
0.506
24
13
0.539
471
0.473
19
0,714
1O
0.886
30
0.360
7
0.507
181
10
0.541
0,480
12
IO
0.717
12
0.886
11
O
0.352
0.508 24
30
0.544
0,480
IS
0,721
0.886
0.352
0.508
18'
21
0.545
11
17
0,483
9
17
0,722
II
0.888
0.360
7
20
0.508
12t
19
0.890
0.508
18"
19
0.546
35
0
1I
0,729
0.892
0.361
6
0,508 18
24
15
0.547
0,488
12
0.730
0.895
23
16
0.361
5
0.508
18T
59F
0,488
0.896
13
0.509
6'
25
0,548
71
15
0.488
13
21
0,733
15
0.896
9
0.361
12
15
0.509
18
0.509
6'
II
0.549
23F
0.496
10
5
0,735
12
0.897
16
0.362
10
0.509
24
12
0.550
0.496
11
14
0.735
24
0.897
0.362
10
0.509
24
0.550
47
0.496
0,739
0.898
0.362
10
17
0.509
181
0.551
0.496
17
0,740
0.900
27
0.362
8
30
0.509
24
0.900 20
0.509
12
9
0.552
0,500
9
8
0,743
0.900
24
0.363
0.509 5
6'
22
0.553
0,500
0.744
0.901
7
12
0.510
18"
0.554
297
0.500
22
0.745
10
0.902
15
0.364
11
0.510
12a
0.554
0.500
19
0,747
1
0.903
27
0.364
0.511 9.
0,510
30
17

---

Understanding the Skill Gap in Recurrent Language Models
Figure Visualization of the masking applied to the Gather Head (first row) and the Aggregate Head (second row) Llamba-8B. The Gather Heads restricted interactions with the representative token and its associated answer, while the Aggregate Heads limited to processing only the final token's relevant portions. This setup ensures that the hypothesized Gather-and- und-Aggregate mechanism is exclusively responsible for the observed model behavior during MMLU evaluation.
A.6. Masked Head Evaluation Confirms Role Specialization
To confirm that the observed Baaher-and-Aggregaate behavior is not incidental, we conduct controlled masking experiments that restrict head input and output pathways. Figure 6 visualizes the masking applied Llamba-8B: the Gather Head is constrained to interact only with a specific context segment (e.g., the representative token and its candidate answers), while the Aggregate Head is permitted to read only the final token. This enforces a clean separation of responsibilities and eliminates confounding interactions.
The results demonstrate that even under these constraints, performance remains high, indicating that the heads question genuinely specialize in retrieval. Figure 7 further supports this by comparing attention maps across Transformer and SSM models. While Llama-3.1-8B shows sharp, focused patterns for both head types, SSM models exhibit smoother, more diffuse behavior-suggesting suggesting that explicit retrieval pathways more difficult to encode in state-space models, but can still emerge under the right inductive bias hybridization.
A.7. Hybrid Layer Replacement Pinpoints Retrieval Bottlenecks
To further isolate the limitations of SSMs in implementing retrieval behavior, we perform layer-wise replacement experiments, shown in Figure 8. In this setup, we substitute individual Llamba-8B layers with their Llama-3. counterparts, then evaluate the resulting hybrid model on MMLU. Most substitutions have little or even negative effect, confirming that localized replacements do not trivially enhance retrieval. However, replacing Layer 17 yields a substantial boost-nearly doubling MMLU accuracy relative to the Llamba baseline.
This layer aligns with the Aggregate Head L17H24, previously identified as critical. The result reinforces our earlier claim: SSMs struggle particularly with the global integration stage of retrieval, and improving only this step-while leaving the rest of the model unchanged can already unlock large performance gains.
A.8. Disabling G&A Heads Impairs Reasoning But Not Knowledge Access
The head ablation results in Table and Table 4 provide complementary evidence that Gather-and- -gggeggtte heads responsible for task-specific reasoning rather than knowledge access alone. In ARC-Challenge, where answers are selected from fixed options, head removal causes modest degradation. But in ARC-Challenge-Chat where generative reasoning is required, accuracy drops sharply with G&A removal.
Table 4 generalizes this across six benchmarks. MMLU, LAMBADA, GSM8K, and BBH show substantial performance drops-essi erops--epeccallly in models like Llama-8B and Zamba2 despite only moderate declines in knowledge-baat based tasks. These
18

---

Understanding the Skill Gap in Recurrent Language Models
Llama-3.1-88
Liamba-8B
Felcon-Membe-78
08
02
Figure 7. Illustration of Gather Heads (top row) Aggregate Heads (bottom row) across different model types. The column shows attention patterns in Llama-3.1-8B (Transformer-baseel while the center and right columns correspond to Llamba-8B AMamba-2-based and Palcon-Mamba-78B (Mamba-l-based)) respectively-boot state-space models (SSMs). Transformer models, both Gather and Aggregate Heads produce sharp, localized patterns. In contrast, SSM-based models exhibit smoother behavior, making it harder to isolate specific segments or retrieve information with the same precision.
results demonstrate that disabling G&A heads selectively damages retrieval and reasoning, without erasing stored knowledge. This aligns with our core hypothesis: these heads act as interfaces for using knowledge in context, not for storing it.
MODEL
#REMOVED HEADS
ARC-CHALLENGE
ARC-CHALLENGE (REGULAR)
(CHAT)
ACC
ACC
0
76.8 (+0.0%)
45.5 (40.0%)
10
72.2 (-60%)
43.6 (44.5%)
Llama-3B
50.0 (-34.9%)
42.0 (-7.7%)
30
43.2 (-43,8%)
41.9 (-7.%)
84.3 (40.0%)
54.9 (40.0%)
10
77.1 (8.5%)
51,6 (.600)
Llama-8B
20
49.3 (.41.5%)
47.3 (-138%)
30
53.6 (336.4%)
47.9 (-12.8%)
Table 3. Comparison of ARC Challenge and ARC-Challenge Chat under Attention Head Removal. ARC- IRC-Challenge Chat repackages the same questions as ARC Challenge into conversational, generative format. While ARC Challenge uses normalized accuracy, chat version relies on exact-match due to its open-ended outputs. Models initially perform better the chat variant, leveraging its generative flexibility to reason through questions. But when G&A heads are ablated, performance drops sharply- unlike the minor declines seen ARC Challenge. This contrast suggests the degradation reffects impaired retrieval rather than loss of stored knowledge.
A.9. KV-Retrieval Recipe
To probe retrieval behavior in isolation, design synthetic memorization task in which the model is asked to recall values associated with keys presented earlier in the same prompt. Each input contains a list of key-value (KV) pairs, followed by query for the value of one of the keys. For example:
Memorize the following dictionary: present:50
institute:0
19

---

Understanding the Skill Gap in Recurrent Language Models
MMLU Accuracy (96
MMLU
45
Baseline
40
35
30
25
2
9 10011113111111177111922222 23
Replaced Layer
Figure 8. Results of the hybrid replacement experiment. Each Llamba-8B layer was replaced with its -lama-3. 1-BB counterpart and tested on MMLU. Most swaps had little or negative impact on the 33% baseline, replacing layer 17-where a key Aggregate Head (Table 1) significantly improved MMLU. This underscores both the difficulties SSMs face in replicating the Aggregate mechanism and G&A critical role MMLU performance.
MODEL
#HEADS
MMLU
LAMB, PPL 4
GSM8K ACC
SWDE ACC f
BBH
KNOWLEDGE ACC 1
ACC
ACC 1
60.3 (+0.0%)
4.8 (+0.0%)
28.7 (+0.0%)
85.8 (+0.0%)
38.2 (+0.0%)
60.5 (+0.0%)
53.1 (-12.0%)
(+35.7%)
17.4 (-39.4%)
81.9 (-4.5%)
33.4 (-12.6%)
59.4 (-1.8%
Llama-3B
32.2 (-46.6%)
8.8 (+82.8%)
9.1 (-68.2%)
57.5 (-33.0%)
27.7 (-27.5%
58.7 (-3.0%)
30
29.9 (-50.4%)
10.1 (+109%)
5.6 (-80.5%)
47.5 (-44.6%)
25.4 (-33.5%%
58.0 (-4.1%%
68.1 (+0.0%)
3.4 (+0.0%)
27.3 (+L.0%)
90.8 (+0.0%)
45.1 (+0.0%)
68.5 (+0.0%)
10
61.9 (-9.1%)
4.2 (+22.0%)
21.7 (-20.5%)
87.3 (-3.9%)
37.7 (-16.5%%
67.1 (-2.0%)
Llama-8B
20
38.1 (-44.0 (-44.0%)
6.8 (498. 6%%)
9.4 (-65 6%)
79.5 (-12.4%)
29. , (-35,2%))
64,8 (-5.4%)
38.7 (46.
7.3 (+115%)
7.8 (-71.4%)
74.0 (-18.5%)
29.0 (-35.7%)
64.4 (-6.0%)
52.5 (+0.09%)
3.6 (+0.0%%
21.3 (40.0%)
9.2 (40.0%)
63.8 (+0.0%)
10
42.6 18.9%%
5.2 (+44.4%%
18.6 (-12.7%)
9.0 (-2.2%)
63.7 (-0.2%)
Llamba-3B
41.3 (-21.3%)
8,2 (+128%)
18.1 (-15.0%)
9.0 (-2.2%
63.1 (-1.1%)
30
41.2 (-21.5%)
9.1 (+153%)
18.1 (-15.0%)
9.0 (-2.2%
62.6 (-1.9%)
60.7 (+0.0%)
4.09 (+0.0%)
20.0 (40.0%
11.0 (40.0%)
69.1 (40.0%)
10
59.6 (-1.8%
6.26 (+53.25)
18.3 (.8.5%
11.6 (+5.5%)
68.8 (.0.4%)
Llamba-8B
20
59,0 (-2.8%)
6.99 (470.8%)
18.6 (-7.0%)
11.4 (43.6%
68.6 (-0.7%)
30
58.1 (-4.3%)
7.44 (+81.9%)
18.1 (-9.5%)
11.4 (43.6%)
68.2 (-1.3%)
55.7 (40.0%)
4.2 (+0.0%)
57. (40.0%)
89.5 (40.0%)
30.6 (40.0%)
66.8 (40.0%)
10
42.4 (-23.9%)
12.8 (+204%)
24.7 (-57.0%)
84.3 (-5.8%
25.5 (-16.7%)
64.8 (.3.0%)
Zamba2-2.7B
20
37.2 (333.2%)
22.2 (+428%)
6.5 (-88.9%)
74.4 (-16.9%)
17.4 (-43.1%)
62.6 (.6.3%)
65.1 (+0.0%)
3.1 (+0.0%)
60.5 (40.0%)
91.7 (+0.0%)
33.0 (40.0%)
70.6 (+0.0%)
10
63.0 (-3.2%)
3.6 (+16.1%)
40.1 (.31.7%)
80.0 (-12.8%)
24.6 (-25.5%)
68.7 (.2.7%)
57.3 (-12.0%)
5.2 (+67.7%)
27.6 (.54.4%)
75.1 (-18.19)
28.9 (-12.4%)
67.5 (-4.4%)
30
54.0 (-17.1%)
7.0 (+126%)
18.4 (.6966%)
43.7 (-52.3%)
24.0 (.27.3%)
67.3 (-4.7%)
7amba2-7B
40
50.6 (-22.3%)
9.5 (+206%)
14.9 (-75.4%%
41.2 (-55.1%)
21.7 (-34.2%)
67.0 (-5.1%)
50
37.7 (42.1%)
13.9 (+348%)
7.5 (-87.6%)
39.6 (-56.8%)
18.5 (-43.9%)
67.0 (-5.1%)
36.2 (-44.4%)
19.8 (+538%)
(-88.1%)
39.6 (-56.8%)
15.9 (-51.8%)
66.5 (-5.8%)
Table Impact of Disabling G&A Heads Across Models. Each value reports the change relative the 0-head configuration. Evaluations are conducted on MMLU (Hendrycks et al., 2021). LAMBADA (Paperno et al., 2016), Grade School Math (GSMSK)
(Cobbe et al., 2021), SWDE (Arora al., 2024), Big- sig-Bench Hard (BBH) (Suzgun al., 2022), and knowledge-focusee tasks defined Section 3.1. Note that (1) GSM8K and BBH are evaluated using the "exact-match, flexible-extract" metric; (2) GSMSK is omitted for Llamba models, which were not trained on mathematical datasets and are not intended for arithmetic reasoning.
scallops: : neuropsychiatzy:
The value of the key sscallops' is
20

---

Understanding the Skill Gap in Recurrent Language Models
Evaluation Setup. We dictionaries containing key-value pairs and evaluate the model randomly generated examples. For each example, we assess the model's ability to retrieve the correct value under two settings:
1. Generation: The model is prompted with the full dictionary and asked generate 10 additional tokens. We check whether the correct value appears anywhere in the generated output.
Answer Scoring: Each possible value from the dictionary is appended to the end of the prompt, and the log-probabilie of the resulting sequence is computed. The value assigned the highest log- probability is treated as the model's prediction.
The key difference between these strategies lies in how the model expresses its beliefs. Generation allows the model to unfold short sequence of reasoning steps, operating within a richer space of token-level interactions. However, this flexibility comes at the cost of interpretability, making it harder to surgically probe the model's s internal mechanisms. Answer scoring, by contrast, restricts the model to a fixed set of candidates, offering more direct readout of its internal ranking.
This distinction is important because generation can surface hallucinations cases where the model outputs values not present in the dictionary or otherwise digresses from the task-while answer scoring is restricted to valid options. We find that the heads responsible for hallucinations differ from those that enable correct retrieval. This stands in contrast the formulation of Wu et al. (2024), who used generation alone and concluded that "retrieval heads" may be responsible for hallucinations.
Prompt Format. We test two variants of the query prompt that differ only in their final character:
1. The value of the key 'scallops' is and
2. 'The value of the key 'scallops' is * "with trailing space).
Though superficially similar, these variants elicit markedly different behaviors under Answer Scoring. In variant (I), the model must evaluate completions like 'The value of the key 'scallops' is84 ", which lacks a space and results in unnatural tokenization. Surprisingly, this unnatural format is highly diagnostic: head-level retrieval signals emerge clearly with as few as 20-30 key-value pairs. In contrast, variant (2), which uses the more natural space-delimited form, requires much longer dictionaries to produce similar head-level effects. For comparison, Wu al. (2024) relied on 35K-token contexts to observe such patterns- beyond typical context lengths.
We hypothesize that this discrepancy reflects shift in the distribution of retrieval effort across the model. When the prompt ends cleanly on a token boundary, the model can offload the task to a small number of dedicated retrieval heads. By contrast, introducing a continuation token (like a space) appears to encourage more distributed processing across heads.
Head-Ablation Protocol. To identify attention heads involved in retrieval, we systematically ablate each head by zeroing out its output projection and measure the resulting drop in accuracy under Answer Scoring (without the trailing space). This procedure isolates heads that are critical for retrieval while keeping computational overhead low.
Overall, this setup offers clean and flexible diagnostic for in-context cetrieval-indepee of natural language or factual knowledge-and and clarifies how individual heads contribute to this core algorithmic skill.
21


Reviews:
[{"review_id": "SukecGXh3c", "reviewer": "Reviewer_x3LE", "summary": "**Updates after rebuttal: I increased my score in light of authors' rebuttal (particularly on the distillation procedure)**\n\n*I appreciate the authors efforts in explaining and demonstrating how their distillation recipe can be used to distill hybrid models from scratch, motivated from the Gather-and-Aggregate mechanism. This addresses my main concern. I have updated my score to 3.*\n\n---\n\nThis paper investigates the Gather-and-Aggregate (G\\&A), a mechanism for solving retrieval tasks that emerges in both SSMs and transformers. The authors found that only a small number of such G&A heads exist in pretrained language models. From empirical results: the authors showed the gap between SSMs and transformers on retrieval is mainly due to these small number of G\\&A heads; the authors also showed SSMs struggle to implement the Aggregation head, whereas the hybrid models can mitigate this issue.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s positive feedback on our approach to identifying retrieval heads in language models and their relevance to hybrid architectures. Below, we address the reviewer’s concerns about terminology, empirical support, and methodological clarity.\n\n> It is well known empirically and theoretically that SSM struggles at retrieval compared to transformers. Thus, it also seems quite natural that SSM is worse at implementing the Aggregate head.\n\nWe agree that this limitation seems natural. Our work sharpens this intuition by bridging the theoretical limitations observed in simplified SSM variants with the retrieval failures seen in full-scale language models. We show that the known weaknesses in the former manifest in the latter specifically through G&A instances. While this perspective may seem natural in retrospect, the behavior of Aggregate heads—and the mechanisms enabling them—had not been clearly understood in this setting.\n\nSurprisingly, we find that SSM-based G&A mechanisms fail even on simple short-context tasks—like retrieving a letter in MMLU—suggesting they fall short well before hitting their theoretical limit.\n\n> From the body of the paper, it is suggested that the Gather and Aggregate head play the same role as the Content-Gatherer head and Correct-Letter head in Lieberum et al. If so, why the new name, and is the extension lies in investigating such mechanism in SSMs (beyond transformers done in Lieberum et al.)? If not, what are the differences among the definitions of these heads?\n\nThank you for raising this point. We will update the paper to clarify this.\nOur contribution lies not in proposing new head types, but in refining the understanding of their roles and how they support retrieval. Our core finding is that in both Transformers and SSMs, retrieval emerges not from a single head, but from a coordinated mechanism involving two roles: one head gathers the location of the target, and another aggregates the content. \n\nThis framing extends prior work in three ways:\n1. “Correct Letter” heads (Lieberum et al.) do more than select the correct option in multiple-choice tasks — they exhibit broader retrieval behavior (Sections 4.2 & 5),\n2. Retrieval relies on coordinated interaction between heads, rather than a single head as suggested in Wu et al. (Sections 3 & 4),\n3. We analyze this mechanism in SSMs, which lack attention and have not been studied in this context.\n\nThe new terminology reflects these empirical findings rather than renaming existing concepts. We also note that we retain “Gather” heads as shorthand for the “Content-Gatherer” heads described in Lieberum et al.\n\n> The authors claimed that \"pretrained language models contain only a small number of Gather-and-Aggregate head\", but only empirically evaluate this claim on three pretrained language models.\n\nDue to space limits, we address this under BU3Q (‘sample of models is too limited’). TLDR: we are expanding our evaluation to include additional models.\n\n> Section 4.3: It will be nice to have a table showing the original MMLU performance, the performance after disabling all heads except the G&A heads, and the performance after disabling all heads except the G&A heads with attention mask.\n\nThank you for pointing this out. We’ve added the table to Section 4.3 as suggested.\n\n> The investigation of the hybrid model is interesting. Nonetheless, it will be much more convincing if the authors show positive results of distilling a transformer to a hybrid model by keeping the Aggregation head intact in the attention layer.\n\nThank you for the suggestion. We address this point in the “Hybrid Replacements” experiment (Section 6.4), where we systematically replaced each layer of the distilled Llama-8B model with its counterpart from Llama-3.1-8B and evaluated the effect on MMLU (without fine-tuning). As shown in Figure 7, most replacements had little or negative impact, but substituting Layer 17—associated with a strong Aggregate head (Table 1)—led to a clear gain, improving performance from 33% to 50%.\n\n> In Section 6.3 ... what if there are other mechanisms implementing in the attention layers that contribute to the MMLU task?\n\nWe agree with the reviewer’s point and have refined our intervention to better isolate the role of the G&A mechanism in response. Specifically:\n1. Head Selection: We restricted masking to a small set of manually identified G&A heads—L47H{17, 18, 25} and L59H{17, 21}.\n2. Token Selection: Masking was applied only to specific tokens in the final attention row involved in the G&A pattern (e.g., “:” attending to choices “A–D”).\n\nWith knowledge-task accuracy stable at 70% and MMLU dropping from 64% to 35%, this targeted intervention indicates that other attention patterns contributing to MMLU were not impacted.\n\n**References**:\n\n[1] Lieberum et al., 2023 — arxiv.org/abs/2307.09458\n\n[2] Wu et al., 2024 — arxiv.org/abs/2404.15574", "claims_and_evidence": "1. The authors claimed that \"pretrained language models contain only a small number of Gather-and-Aggregate head\", but only empirically evaluate this claim on three pretrained language models (Llama-3.1-8B, Lambda-9B, Falcon-Mamba-7B). I suggest the authors either weaken their claim or support it with stronger evidence across many more language models.\n\n2. The authors claimed that Gather-and-Aggregate mechanism is an extended version of the mechanism found in Lieberum et al. (2023) in the introduction, see Line 019-029. But from the body of the paper, it is suggested that the Gather and Aggregate head play the same role as the Content-Gatherer head and Correct-Letter head in Lieberum et al. If so, why the new name, and is the extension lies in investigating such mechanism in SSMs (beyond transformers done in Lieberum et al.)? If not, what are the differences among the definitions of these heads?", "methods_and_evaluation_criteria": "1. Most methods proposed in this work to mechanistically interpret the Gather-and-Aggregate mechanism are appropriate. However, the method in Section 4.3 requires further evidence. Specifically, \"Evaluating the MMLU benchmark with these masks showed no loss in performance\" -- which loss of performance baseline do the authors refer to? It will be nice to have a table showing the original MMLU performance, the performance after disabling all heads except the G\\&A heads, and the performance after disabling all heads except the G\\&A heads with attention mask.", "theoretical_claims": "None", "experimental_designs_or_analyses": "1. In Section 6.3, the authors disabled all potential Aggregate heads across all attention layers, rather than disabling only the identified Aggregate head, which led to an decline of around 4x of performance in MMLU. However, what if there are other mechanisms implementing in the attention layers that contribute to the MMLU task (which also get disabled in this process)?", "relation_to_broader_scientific_literature": "This work is broadly related to the mechanistic interpretability of language models. It builds on the earlier work by  Lieberum et al. (2023) on identifying the Content-Gatherer and Correct-Letter head mechanisms in transformers for solving the MMLU task. It examines how such mechanisms arise in other transformer-based language models, SSMs, and hybrid models.", "essential_references_not_discussed": "The following is another recent Hybrid model, which came out around the same time as MOHAWK. The authors should discuss their choice of Hybrid model and whether such choice affects the findings.\n[1] Wang, Junxiong, et al. \"The mamba in the llama: Distilling and accelerating hybrid models.\" Advances in Neural Information Processing Systems 37 (2024): 62432-62457.\n\nThe following line of works are theoretical papers proving the retrieval mechanisms implemented in transformer for induction head task:\n[2] Bietti, Alberto, et al. \"Birth of a transformer: A memory viewpoint.\" Advances in Neural Information Processing Systems 36 (2023): 1560-1588.\n[3] Sanford, Clayton, Daniel Hsu, and Matus Telgarsky. \"Transformers, parallel computation, and logarithmic depth.\" Forty-first International Conference on Machine Learning.", "other_strengths_and_weaknesses": "Strengths: The paper is well-written overall. I find the idea of identifying the crucial heads in language models for MMLU or general retrieval tasks interesting. I also appreciate the application of such idea in the hybrid models.\n\nWeaknesses: \n1. The G\\&A mechanism is proposed in the previous work by Lieberum et al. (2023) focusing on transformers. The current work does not seem to offer additional mechanistic interpretability insights on such mechanism. \n\n2. It is well known empirically and theoretically that SSM struggles at retrieval compared to transformers (Jelassi et al., [3] Sanford et al.). Thus, it also seems quite natural that SSM is worse at implementing the Aggregate head. \n\n3. The investigation of the hybrid model is interesting. Nonetheless, it will be much more convincing if the authors show positive results of distilling a transformer to a hybrid model by keeping the Aggregation head intact in the attention layer.\n\n4. Experiments: more details on Sec 4-6 should be supplemented.", "other_comments_or_suggestions": "None", "questions_for_authors": "See questions in the previous sections, indexed as:\n- Claims and Evidence: C1 - C2;\n- Methods and Evaluation: M1\n- Experimental Design: E1", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wy1q9SVSbj", "reviewer": "Reviewer_BU3Q", "summary": "This paper investigates the performance gap between Transformer and State-Space Model (SSM) language models, focusing on retrieval capabilities. The authors identify a \"Gather-and-Aggregate\" (G&A) mechanism that emerges in both architectures but is implemented more effectively in Transformers. This two-part mechanism consists of a \"Gather Head\" that condenses segment information into the last token, and an \"Aggregate Head\" that processes these condensed representations. The key finding is that despite different architectures, both models develop similar mechanisms, but SSMs struggle with the Aggregate component due to their fixed state size. Remarkably, these capabilities are concentrated in just a few critical heads, and disabling a single head can severely degrade performance on tasks like MMLU (dropping from 66% to 25%). The paper also shows that hybrid models naturally assign Aggregate heads to attention layers, explaining their success.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s recognition of the novelty in our approach to architectural differences and hybrid experiments. Their comments raise valuable questions about performance gaps, generality across models and tasks, and the role of other components. We address each point below.\n\n\n> There is limited exploration of how the G&A mechanism might interact with other architectural components (like normalization layers)\n\nWe agree this is a valuable direction. Our focus is on understanding retrieval mechanisms in established model architectures, where normalization practices are typically standardized. While these interactions interest us, we leave their exploration to future work.\n\n\n> The authors don't fully explore how the G&A mechanism might be improved in SSMs\n\nWhile improving the G&A mechanism in SSMs is valuable, prior work has shown that retrieval limitations are inherent to all RNN-based models due to fixed-size memory, which constrains tasks like KV-Retrieval [9–11]. We therefore focused on hybrid models that directly address this issue, showing how they mitigate memory constraints rather than pursuing limited incremental gains within pure SSMs.\n\n>Is G&A mechanism truly the causal factor in performance differences rather than a correlate of some other more fundamental architectural difference?\n\nThank you for raising this—we have clarified this in the revised paper as well.\nTo support a causal role for G&A in MMLU performance differences between architectures, we present the following:\n1. MMLU is bottlenecked by retrieval, not knowledge (Sections 3 & 5): Models can do well on MMLU despite weak general performance, but fail without a working G&A mechanism to retrieve letters from the prompt.\n2. G&A enables retrieval in language models (Sections 4 & 5): the same G&A heads drive synthetic KV retrieval performance.\n3. SSMs struggle with G&A (Section 6) due to their fixed memory structure [9–11].\n\nThis forms a causal chain: architectural constraints → impaired G&A → retrieval failure → MMLU drop. To further support this interpretation, we added prior work [7–9] linking retrieval—and thus G&A—to the Transformer–SSM gap.\n\nIf any part of this causal chain remains unconvincing, we’d appreciate clarification on which link is in doubt, so we can design targeted control experiments to isolate it.\n\n> What evidence do you have that the G&A limitation is fundamental to SSM architecture rather than an artifact of current implementations or training methods?\n\nIt is well-established, both theoretically and empirically, that SSMs are weaker at retrieval than attention-based models due to their fixed-size memory [9–11]. A key takeaway of our work is that the G&A mechanism, particularly the Aggregate head, drives retrieval. Given this, it is natural to expect that limitations in G&A may be fundamental to the SSM architecture.\n\nOur masking experiment (Section 4.2) further shows that SSM-based G&A fails to implement the mechanism correctly, as recurrent dynamics smooth attention across tokens. This supports the view that the limitation stems from architectural properties, rather than optimization.\n\n> The sample of models is too limited.\n\nWe agree that evaluating more models would strengthen the paper.  We chose the two available SSMs (Mamba-1, Mamba-2) and a Transformer to ensure architectural diversity, given the manual and resource-heavy process of identifying relevant heads.\nIn response to this suggestion, we are running tests on the entire Falcon and LLaMA 3 families. If there are other models the reviewers would like to see included, we would be glad to add them.\n\nThat said, similar retrieval behaviors have been observed in LLaMA-2 [2–6], Chinchilla [1], Phi-2 [4, 6], Falcon [4], Mixtral-8x7B [4], and Yi [6], ranging from 7B to 70B, suggesting that our findings are representative.\n\n> Have you explored whether the same performance gap exists for tasks that don't follow the multiple-choice format of MMLU?\n\nDue to space limits, we address this under 7vUD (‘other evaluations’). TLDR: gsm8k and ifeval confirm the same G&A head dependency; we’re expanding to diverse tasks.\n\n\n> The hybrid replacement experiment is clever but confounded by potential interactions between layers that aren't accounted for.\n\nWe agree that the head’s effect could be more directly tested. To address this, we replaced only the Aggregate head—leaving the rest of the layer unchanged—and observed a similar 17-point gain in MMLU. Thank you for pointing this out.\n\n**References:**\n\n[1] arxiv.org/abs/2307.0945\n\n[2] arxiv.org/abs/2404.15574\n\n[3] arxiv.org/abs/2409.01659\n\n[4] arxiv.org/abs/2402.12483\n\n[5] arxiv.org/abs/2407.15018\n\n[6] arxiv.org/abs/2402.01781\n\n[7] arxiv.org/abs/2312.04927\n\n[8] arxiv.org/abs/2402.18510\n\n[9] arxiv.org/abs/2406.07887\n\n[10] arxiv.org/abs/2402.01032\n\n[11] arxiv.org/abs/2501.00658", "claims_and_evidence": "The authors present a methodical investigation through:\n\n- Ablation studies to demonstrate the critical role of specific layers and heads\n- Visualization of attention patterns in both Transformer and SSM models\n- Performance measurements on MMLU and knowledge-focused tasks\n- Analysis of hybrid models showing where Aggregate heads are allocated", "methods_and_evaluation_criteria": "The evaluation approach has significant limitations. While the authors distinguish between knowledge tasks and retrieval-heavy tasks, this binary categorization oversimplifies the complex capabilities required for different benchmarks. MMLU is treated as primarily a retrieval task, but prior work has demonstrated it tests a broader range of skills.\n\nThe ablation methodology is appropriate but limited in scope. While identifying critical components through knockouts is sound, the paper fails to systematically explore alternative hypotheses or control for confounding factors. For example, the authors don't adequately address whether the effect is truly due to the proposed mechanism or other properties of the affected heads.\n\nThe KV-Retrieval task provides a cleaner test case, but the artificial nature of this task limits generalization to real-world model performance differences.", "theoretical_claims": "The paper does not have a theoretical proof.", "experimental_designs_or_analyses": "Several experimental weaknesses undermine the strength of the conclusions:\n\n1. The sample of models is too limited (three models) to support broad claims about architectural differences.\n2. The paper focuses heavily on MMLU and simplistic KV-Retrieval tasks without testing on a diverse range of retrieval-focused benchmarks.\n3. The hybrid replacement experiment is clever but confounded by potential interactions between layers that aren't accounted for.", "relation_to_broader_scientific_literature": "The authors acknowledge key related work including:\n\n- Olsson et al. (2022) and Elhage et al. (2021) on induction heads\n- Lieberum et al. (2023) on \"Content Gatherer Head\" and \"Correct Letter Head\"\n- Wu et al. (2024) on retrieval heads", "essential_references_not_discussed": "Meng et al. (2022) \"Locating and Editing Factual Associations in GPT\" discusses similar mechanisms for fact storage and retrieval", "other_strengths_and_weaknesses": "Strengths:\n- The paper introduces a novel perspective on architectural differences\n- The visualization of head activations effectively illustrates the proposed mechanism\n- The hybrid model analysis provides useful insights for architectural design\n\nWeaknesses:\n- The paper focuses heavily on MMLU as a proxy for retrieval capabilities. A wider range of retrieval-focused benchmarks would strengthen the generalizability\n- There is limited exploration of how the G&A mechanism might interact with other architectural components (like normalization layers)\n- The authors don't fully explore how the G&A mechanism might be improved in SSMs", "other_comments_or_suggestions": "The paper would be substantially strengthened by:\n- Testing on a broader range of models and tasks\n- More rigorously controlling for alternative explanations\n- Addressing the limitations and generalizability of the findings", "questions_for_authors": "- Is G&A mechanism truly the causal factor in performance differences rather than a correlate of some other more fundamental architectural difference? Additional control experiments would be necessary to establish causality.\n- Have you explored whether the same performance gap exists for tasks that don't follow the multiple-choice format of MMLU? This would help distinguish between general retrieval capabilities versus format-specific processing.\n- What evidence do you have that the G&A limitation is fundamental to SSM architecture rather than an artifact of current implementations or training methods?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HytihvELlj", "reviewer": "Reviewer_vZz2", "summary": "I am not an expert in transformers and SSMs.\nThe authors reverse engineer language models to show that retrieval capabilities are supported by distinct parts of the networks compared to overall knowledge.\nBy a systematic lesioning of layers, they identify that (at least) two layers are needed to support retrieval. The first of the layers contains heads which aggregate information from a segment and encode it in the final token. The second layer contains heads which process these tokens to decide on the correct response. The smooth nature of SSMs makes it harder for them to constrain activation from neighboring tokens, and thus they struggle more with these tasks.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive remarks. We’re glad the conceptual separation between knowledge and retrieval, along with the identification of the elements driving retrieval, was found to be valuable. \nIf any questions arise, we would be happy to address them.", "claims_and_evidence": "Yes. The lesion experiments show that two layers are needed, and that they operate in the described manner on MMLU.", "methods_and_evaluation_criteria": "Yes. The benchmarks are designed to tease apart knowledge from retrieval.", "theoretical_claims": "Not relevant", "experimental_designs_or_analyses": "Not in detail", "relation_to_broader_scientific_literature": "Not an expert in this literature.", "essential_references_not_discussed": "not aware", "other_strengths_and_weaknesses": "This is a nice systematic analysis of the underlying mechanisms of language models. The separation to knowledge and retrieval, and the identification of a small number of elements responsible for retrieval is an important finding that helps understand and improve models.", "other_comments_or_suggestions": "None.", "questions_for_authors": "None.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HvRQz09uuP", "reviewer": "Reviewer_7vUD", "summary": "This paper seeks to investigate why SSMs underperform Transformers on on retrieval tasks. The paper identifies a \"Gather-and-Aggregate\"  (G&A) mechanism that emerges in Transformers and SSMs (though with some differences). The authors find that Transformers and SSMs concentrate this G&A mechanism in just a few heads and disabling them can significantly reduce MMLU scores while maintaining scores on knowledge intensive tasks, suggesting the importance of G&A for retrieval. Further experiments are performed to study the importance of these heads in SSM and hybrid models.\n\n## Update after rebuttal: I maintain my score as my main questions have been addressed.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We’re glad the reviewer found our analysis of the Gather-and-Aggregate (G&A) mechanism compelling—both in its development within SSMs and its implications for hybrid models. Their main concern centers on the need for stronger empirical support and broader evaluations. We respond to each of these concerns below.\n\n> Section 3.4 caveats says that while the models tested exhibit the same two-layer mechanism, the interaction does not always occur in consecutive layers. How do you know this? Based on what evidence if you didn't see this in your experiments?\n\nOur experiments happen to observe G&A in consecutive layers, but we do not claim this is a universal pattern. For example, [1] shows that retrieval mechanisms can span non-consecutive layers. Thank you for the question — we’ve updated the paper to clarify this point.\n\n> Have you tried other evaluations that require some form of retrieval, e.g. summarization, long context Q&A, etc? \n\nWe agree that broader evaluations are essential to distinguish general retrieval from format-specific processing. In response to reviewers’ suggestions, we added experiments on tasks with diverse structures and retrieval demands:\n\n|Task|score (%)|G&A masked|\n|-|-|-|\n|gsm8k_cot|33.6|No|\n|gsm8k_cot| 5.8|Yes|\n|gsm8k|40.3|No|\n|gsm8k|10.2|Yes|\n|ifeval|31.1| No|\n|ifeval|21.3|Yes|\n\nAs shown, masking just 8 G&A attention heads across two layers in Zamba2-7B causes substantial drops in performance, reinforcing that G&A supports retrieval across a range of tasks.\n\nWe are also expanding our evaluation suite to include tasks such as LongBench (long-context summarization, QA, and code), RULER (needle-in-a-haystack and variable tracking), and broader chain-of-thought (CoT) benchmarks. While we were unable to include these results during the rebuttal phase due to time constraints, we are actively working on them in response to your suggestion.\n\n> While the paper does a decent job describing and qualitatively visualizing the G&A mechanism, the argument would be strengthened with more empirical measurements of how much gathering and aggregating is performed in different heads of different models. \n\nWe agree that quantifying the presence of G&A across heads and models is an important direction that would strengthen the paper. This kind of analysis would also directly inform practical questions, such as the one the reviewer raised about where to place attention heads in hybrid models. \n\nSo far, we have identified G&A heads through manual inspection, which is time-intensive and does not yet scale easily across models. However, as noted in our response below on hybrid models, findings from [2] provide complementary evidence on the emergence and positioning of such heads in Transformers, further underscoring the relevance of this line of investigation.\n\nWe see this as a promising direction for future work, including developing automated tools to more systematically measure G&A across architectures.\n\n> All models explored here are around the 8B parameter level. Do larger models exhibit the same patterns?\n\nWe acknowledge this as a limitation and agree that evaluating larger models would help validate the broader applicability of our findings. However, the next scale of available models is typically 70B+, which exceeds our academic compute resources and would require substantial engineering effort. That said, similar heads have been observed in Chinchilla-70B by [1], as well as in other works (see “The sample of models is too limited” in our response to reviewer BU3Q), suggesting that these mechanisms likely scale to larger models.\n\n> Am I missing information you provide regarding \"how to effectively place attention heads in a hybrid model to optimize performance\"?\n\nThank you for raising this. We now clarify it more explicitly in the revised paper. Our approach to hybrid design involves:\n1. Hybrid Replacement: Given a pretrained pure SSM model, we identify prominent G&A instances (Section 6.4). These locations signal where SSM heads are replaced with attention heads, improving retrieval and narrowing the performance gap.\n2. Placement Guidance: G&A consistently arises in middle layers—for example, layers 16–17 (of 32) in LLaMA, 35–36 (of 64) in Falcon-Mamba, and 47 and 59 (of 80) in Zamba-7B. This pattern is consistent with prior findings [2], where such attention heads emerge in middle layers. Placing attention heads at similar depths complements the SSM backbone when global context is most needed.\n\n**References:**\n\n[1] Lieberum et al., 2023 — arxiv.org/abs/2307.09458\n\n[2] Zheng et al., 2024 — arxiv.org/abs/2409.03752", "claims_and_evidence": "The claims in the paper are generally well supported with empirical evidence from a series of ablation studies. Pretrained LLama-3.1-8B, Llamba-9B and Falcon-Mamba-7B models are studied and ablated to establish consistency of findings across Transformers and SSMs. \n\nThe evidence that removing specific heads drastically impacts MMLU while having much less effect on \"Knowledge\" tasks provides evidence for the claims around the importance of these heads for retrieval. This is also supported with an additional experiment on a KV-retrieval task. \n\nThe experiments with the hybrid model also provides evidence for the paper's claim to shed light on why hybrid models can be effective. \n\nWeaknesses of claims and evidence:\n- Line 061 claims in the intro that the hybrid analysis \"provides valuable insights into how to effectively place attention heads in a hybrid model to optimize performance\", however I do not think this claim is well supported. In Section 6.4, the hybrid replacment experiments swap in individual attention layers from the Llama model to its ssm distilled Llamba. This experiment suggests layer 17 is important for MMLU. But this only again confirms attention layers are important, but does not say anything about optimal performance.\n- The general claims could be strengthened with a broader set of evaluations than just MMLU and the kv retrieval task, e.g. summarization, long context Q&A, instruction following, etc.", "methods_and_evaluation_criteria": "The general methods and evaluations, in particular the layer and head ablation approaches, do make sense for analysis of where and how retrieval capabilities are implemented in these models.", "theoretical_claims": "The paper is primarily empirical.", "experimental_designs_or_analyses": "Discussed above in Claims and Evidence section.", "relation_to_broader_scientific_literature": "The paper extends prior work on retrieval heads in Transformers and connects it to recent studies on SSMs, hybrids and the weaknesses of SSMs in retrieval intensive tasks.", "essential_references_not_discussed": "The paper appropriately cites Jelassi et al. 2024 and Wen et al. 2024, for the difficulties that SSMs face in retrieval and copying, however there is a growing body of work surrounding this issue that would be useful to include citations to in order to provide an unfamiliar audience more theoretical and empirical context: Park et al. 2024 (https://arxiv.org/abs/2402.04248),  Arora et al 2023 (https://arxiv.org/abs/2312.04927), Arora et al 2024 (https://arxiv.org/abs/2402.18668), Blouir et al 2024 (https://arxiv.org/abs/2411.01030).", "other_strengths_and_weaknesses": "Strengths:\n- Paper provides a compelling explanation for the retrieval gap between transformers and SSMs\n- The identification of the G&A mechanism being concentrated in only a few heads is quite an interesting finding that could lead to deeper explorations in future works. \n- The practical connection to hybrids and providing some evidence for why they can still work well is relevant to the community\n- The visual representations of the G&A mechanism are useful\n\nWeaknesses:\n- While the paper does a decent job describing and qualitatively visualizing the G&A mechanism, the argument would be strengthened with more empirical measurements of how much gathering and aggregating is performed in different heads of different models. \n- The paper would be strengthened by broadening the evaluations beyond just MMLU and kv-retrieval (discussed above). Could different types of task requiring retrieval require different mechanisms or different heads?\n- The paper suggests its findings could inform better hybrid designs, but the current evidence for this in the paper is lacking (discussed more above).", "other_comments_or_suggestions": "- Line 140, right column, says that all three models \"encode retrieval capabilities in the same way across layers\". However the next sentence contradicts this. I assume the first sentence has a typo. \n- Line 160 right column, Llama-9B should be Llamba-9B. \n- Section 3.4 caveats says that while the models tested exhibit the same two-layer mechanism, the interaction does not always occur in consecutive layers. How do you know this? Based on what evidence if you didn't see this in your experiments?", "questions_for_authors": "Most of my main concerns are discussed above. \n\n1. Am I missing information you provide regarding \"how to effectively place attention heads in a hybrid model to optimize performance\"? If not, I would recommend rewording these claims. I suspect this is hard to show by just ablating models already pretrained, since it is hard to compare to the counterfactual of training a differently designed model (e.g. informed by the G&A findings) from scratch. \n2. Have you tried other evaluations that require some form of retrieval, e.g. summarization, long context Q&A, etc? \n3. All models explored here are around the 8B parameter level. Do larger models exhibit the same patterns?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper provides a comparison of attention vs recurrent language models in a specific retrieval task that seems crucial for good performance on MMLU, namely a gather-and-aggregate mechanism that in particular helps with selecting the correct answer in multiple-choice questions, following a study of Lieberum et al.

The findings about the importance of this mechanism for MMLU performance, as well as how this contributes to the gap between SSMs and attention, provides compelling real-world evidence for the known issues of SSMs with retrieval tasks, which were previously studied mainly on toy tasks/models.

The reviewers mostly aligned towards accepting the paper. Please revise the paper to incorporate the clarifications that you provided to reviewers during the rebuttal, particularly regarding the hybrid replacement approach raised by reviewer x3LE and the causal role of your experimental interventions on the attention-vs-ssm gap.
Decision: Accept (poster)