发现语言模型检索瓶颈在少数关键头中,而非架构差异。
Understanding the Skill Gap in Recurrent Language Models: The Role of the Gather-and-Aggregate Mechanism
- 识别出所有模型共用的提取-聚合机制(G&A)
- 删一个头使准确率从66%跌至25%,接近随机
- 建议用注意力层替代特定头,融合两者优势
状态空间模型(SSMs)在长序列上比Transformer更高效,但其固定大小的循环状态限制了算法任务中的上下文检索能力。本文研究发现,基于Transformer和SSM的语言模型均依赖相似的提取-聚合(G&A)机制:提取头从上下文中抓取相关信息,聚合头将其整合为单一表征。该机制集中在少数头中,形成关键瓶颈,即便在简单检索任务中也如此。例如,在剪枝后的Llama-3.1-8B中,禁用单个提取或聚合头会使MMLU任务正确答案字母的准确率从66%降至25%(随机猜测)。此外,此瓶颈会掩盖任务的真实知识需求——剪枝模型在有功能G&A头时能通过MMLU,却在其他知识基准上失败。该瓶颈同样影响SSMs通常表现不佳的任务,如GSM8K、BBH和对话理解。我们发现,SSMs的检索问题体现在这些头中,导致注意力模式平滑,无法产生有效G&A所需的尖锐词元切换。因此,Transformer与SSM的性能差距仅存在于少数头中,而非整个模型结构。这为两者性能差异提供了统一解释,并指明融合路径:预训练混合模型中,将少量注意力层与SSM结合,可将聚合头角色交由注意力模块;同样,将预训练SSM中的单个G&A头替换为注意力变体,也能显著提升检索能力和基准得分。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。