Research map / Architectures and efficiency / Attention and transformers

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

ICML 2025 · Architectures and efficiency › Attention and transformers · approach: Attention head mechanisms · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Shows that in-context retrieval in both Transformers and SSMs relies on a Gather-and-Aggregate mechanism concentrated in a few heads, explaining the MMLU and broader Transformer-SSM gap and motivating hybrid designs.

Problem

SSMs with fixed-size recurrent state underperform Transformers on retrieval-heavy tasks such as MMLU, GSM8K, BBH and dialogue; it is unclear whether this gap stems from broad modeling deficits or from a few algorithmic heads.

Approach

Identifies a two-stage Gather-and-Aggregate mechanism (Gather heads compress segments into summary tokens; Aggregate heads attend to these summaries) that is shared by Transformers and SSMs, concentrated in a few heads; SSM gap is attributed to smoother attention in Aggregate heads, and hybrids delegate aggregation to attention.

Key results

• Disabling a single Gather or Aggregate head (L16H22 or L17H24) in a pruned Llama-3.1-8B drops MMLU from 66.32% to about 25% (random).
• Replacing one SSM layer (L17) with its Llama-3.1-8B counterpart in distilled Llamba-8B raises MMLU from 33% to 50%.
• Ablating nine attention heads in Zamba2-7B reduces MMLU from 64.3% to 34.9% while knowledge tasks stay near 70%.

Measured effects reported in the paper

Surprising findings

• MMLU performance can remain high even when knowledge-focused benchmarks degrade substantially, indicating MMLU partly measures retrieval rather than knowledge.
• Llamba-8B retains over 95% of its performance after removing 30 G&A heads, showing SSM retrieval is distributed across many weaker heads rather than concentrated.

Limitations

• Pruning-based layer identification can miss Gather-and-Aggregate interactions spanning non-adjacent layers.
• Analysis is limited to a few pretrained models and mostly to MMLU-style and synthetic KV retrieval evaluation.

Experimental setup

Datasets: MMLU ARC-Challenge ARC-Easy PIQA Winogrande OpenBookQA HellaSwag GSM8K BBH LAMBADA SWDE

Benchmarks: MMLU ARC-Challenge ARC-Easy PIQA Winogrande OpenBookQA HellaSwag GSM8K BBH LAMBADA SWDE ARC-Challenge Chat

Metrics: accuracy perplexity exact match

Baselines: Llama-3.1-8B-Instruct Falcon-Mamba-7B-Instruct Llamba-8B Zamba2-7B-Instruct Llama-3.2-3B-Instruct Llamba-3B Zamba2-2.7B-Instruct

Builds on: induction heads retrieval heads Correct Letter heads (Chinchilla analysis) MOHAWK distillation

Closest papers by meaning

• Can Language Models Actually Retrieve In-Context? Drowning in Documents at Million Token Scale (NeurIPS 2026 · similarity 0.89)
• Overcoming Long Context Limitations of State Space Models via Context Dependent Sparse Attention (NeurIPS 2025 · similarity 0.89)
• Retrieval meets Long Context Large Language Models (ICLR 2024 · similarity 0.89)
• Talking Heads: Understanding Inter-Layer Communication in Transformer Language Models (NeurIPS 2024 · similarity 0.89)
• Retrieval Head Mechanistically Explains Long-Context Factuality (ICLR 2025 · similarity 0.89)
• Understanding and Mitigating Bottlenecks of State Space Models through the Lens of Recency and Over-smoothing (ICLR 2025 · similarity 0.88)