Understanding the Skill Gap in Recurrent Language Models: The Role of the Gather-and-Aggregate Mechanism
Aviv Bick, Eric P. Xing, Albert Gu
摘要
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 similar Gather-and-Aggregate (G&A) mechanism: a Gather Head extracts relevant information pieces from context, which an Aggregate Head integrates into a single representation. In both architectures, G&A concentrates in a few heads, forming critical bottlenecks even for simple retrieval. For example, we show that 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% (random guessing). Moreover, this retrieval bottleneck can obscure limited 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, such as GSM8K, BBH, and dialogue comprehension. We show that SSMs' retrieval challenges manifest in these heads, creating smoother attention patterns instead of the sharp token transitions effective G&A requires. Thus, the Transformer-SSM retrieval gap exists in just a 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 a few 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. Experiments are publicly available. 1
问问这篇 Paper
智能体会读完全文。
Lune 把这篇 Paper 索引到了每一个公式,引用它的顶会 Paper 也一样。你提问,回答直接引用原文。
引用它的顶会 Paper3
- Local Linear Attention: An Optimal Interpolation of Linear and Softmax Attention For Test-Time RegressionYifei Zuo, Yutong Yin, Zhichen Zeng, Ang Li 等ICLR 2026 · 被引用 5 次
- Retrieval-Aware Distillation for Transformer-SSM HybridsAviv Bick, Eric Xing, Albert GuICML 2026 · 被引用 4 次
- Contribution Weights: A Geometrical Analysis of Self-Attention TransformersJake Cunningham, Nicola Muca CironeICML 2026
它引用的顶会 Paper27
- Language Models are Few-Shot LearnersTom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah 等NeurIPS 2020 · 被引用 64,255 次
- Measuring Massive Multitask Language UnderstandingDan Hendrycks, Collin Burns, Steven Basart, Andy Zou 等ICLR 2021 · 被引用 7,905 次
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-AwarenessTri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra 等NeurIPS 2022 · 被引用 5,493 次
- Efficiently Modeling Long Sequences with Structured State SpacesAlbert Gu, Karan Goel, Christopher RéICLR 2022 · 被引用 3,482 次
- Locating and Editing Factual Associations in GPTKevin Meng, David Bau, Alex Andonian, Yonatan BelinkovNeurIPS 2022 · 被引用 3,415 次
相关 Paper
- Birdie: Advancing State Space Language Modeling with Dynamic Mixtures of Training ObjectivesSam Blouir, Jimmy T. H. Smith, Antonios Anastasopoulos, Amarda ShehuEMNLP 2024
- Hymba: A Hybrid-head Architecture for Small Language ModelsXin Dong, Yonggan Fu, Shizhe Diao, Wonmin Byeon 等ICLR 2025 · 被引用 2 次
- Repeat After Me: Transformers are Better than State Space Models at CopyingSamy Jelassi, David Brandfonbrener, Sham M. Kakade, Eran MalachICML 2024 · 被引用 176 次
- Block-State TransformersJonathan Pilault, Mahan Fathi, Orhan Firat, Chris Pal 等NeurIPS 2023 · 被引用 33 次
- Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language ModelingLiliang Ren, Yang Liu, Yadong Lu, Yelong Shen 等ICLR 2025
