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

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.

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?

Our 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.

Have you tried other evaluations that require some form of retrieval, e.g. summarization, long context Q&A, etc?

We 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:

As 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.

We 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.

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.

We 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.

So 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.

We see this as a promising direction for future work, including developing automated tools to more systematically measure G&A across architectures.

All models explored here are around the 8B parameter level. Do larger models exhibit the same patterns?

We 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.

Am I missing information you provide regarding "how to effectively place attention heads in a hybrid model to optimize performance"?

Thank you for raising this. We now clarify it more explicitly in the revised paper. Our approach to hybrid design involves:

1. 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.
2. 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.

References:

[1] Lieberum et al., 2023 — arxiv.org/abs/2307.09458

[2] Zheng et al., 2024 — arxiv.org/abs/2409.03752

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. If any questions arise, we would be happy to address them.

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.

There is limited exploration of how the G&A mechanism might interact with other architectural components (like normalization layers)

We 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.

The authors don't fully explore how the G&A mechanism might be improved in SSMs

While 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.

Is G&A mechanism truly the causal factor in performance differences rather than a correlate of some other more fundamental architectural difference?

Thank you for raising this—we have clarified this in the revised paper as well. To support a causal role for G&A in MMLU performance differences between architectures, we present the following:

1. 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.
2. G&A enables retrieval in language models (Sections 4 & 5): the same G&A heads drive synthetic KV retrieval performance.
3. SSMs struggle with G&A (Section 6) due to their fixed memory structure [9–11].

This 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.

If 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.

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?

It 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.

Our 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.

The sample of models is too limited.

We 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. In 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.

That 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.

Have you explored whether the same performance gap exists for tasks that don't follow the multiple-choice format of MMLU?

Due 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.

The hybrid replacement experiment is clever but confounded by potential interactions between layers that aren't accounted for.

We 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.

References:

[1] arxiv.org/abs/2307.0945

[2] arxiv.org/abs/2404.15574

[3] arxiv.org/abs/2409.01659

[4] arxiv.org/abs/2402.12483

[5] arxiv.org/abs/2407.15018

[6] arxiv.org/abs/2402.01781

[7] arxiv.org/abs/2312.04927

[8] arxiv.org/abs/2402.18510

[9] arxiv.org/abs/2406.07887

[10] arxiv.org/abs/2402.01032

[11] arxiv.org/abs/2501.00658

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.

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.

We 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.

Surprisingly, 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.

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?

Thank you for raising this point. We will update the paper to clarify this. Our 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.

This framing extends prior work in three ways:

1. “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),
2. Retrieval relies on coordinated interaction between heads, rather than a single head as suggested in Wu et al. (Sections 3 & 4),
3. We analyze this mechanism in SSMs, which lack attention and have not been studied in this context.

The 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.

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.

Due to space limits, we address this under BU3Q (‘sample of models is too limited’). TLDR: we are expanding our evaluation to include additional models.

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.

Thank you for pointing this out. We’ve added the table to Section 4.3 as suggested.

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.

Thank 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%.

In Section 6.3 ... what if there are other mechanisms implementing in the attention layers that contribute to the MMLU task?

We agree with the reviewer’s point and have refined our intervention to better isolate the role of the G&A mechanism in response. Specifically:

1. Head Selection: We restricted masking to a small set of manually identified G&A heads—L47H{17, 18, 25} and L59H{17, 21}.
2. 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”).

With 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.

References:

[1] Lieberum et al., 2023 — arxiv.org/abs/2307.09458

[2] Wu et al., 2024 — arxiv.org/abs/2404.15574