通过保留关键注意力头，让SSM高效实现检索能力，内存效率提升5-6倍。
Retrieval-Aware Distillation for Transformer-SSM Hybrids
- 只保留2%的关键注意力头，其余转为循环头进行知识蒸馏。
- 仅用10个关键头即恢复95%以上教师模型在检索任务上的性能。
- 适合追求高效率、低内存的序列建模场景，尤其关注检索能力的系统。
状态空间模型（SSMs）虽具备高效的序列建模能力，但在依赖上下文检索的任务上仍落后于Transformer。此前研究指出，这一差距源于少量被称为Gather-and-Aggregate（G&A）的注意力头，而SSMs难以复现这些头的作用。本文提出检索感知蒸馏，将预训练Transformer转化为混合学生模型：仅保留这些关键检索头，其余部分蒸馏为循环头。通过在合成检索任务上的消融实验，识别出核心注意力头，构建出稀疏且非均匀分布的注意力结构。结果显示，仅保留2%的注意力头（1B模型中为10个），即可在检索密集型任务上恢复超过95%的教师模型性能，远少于其他混合模型所需的至少25%。此外发现，较大的循环状态常可弥补缺失的检索能力：一旦关键检索由特定头处理，主干SSM可大幅简化，即使状态维度降低8倍也仅有轻微损失。最终混合模型同时减少注意力缓存与SSM状态，相较同类模型内存效率提升5–6倍，以极低内存成本弥合了Transformer与SSM之间的性能差距。
原文摘要 · Abstract (English)
State-space models (SSMs) offer efficient sequence modeling but lag behind Transformers on benchmarks that require in-context retrieval. Prior work links this gap to a small set of attention heads, termed Gather-and-Aggregate (G&A), which SSMs struggle to reproduce. We propose *retrieval-aware distillation*, which converts a pretrained Transformer into a hybrid student by preserving only these retrieval-critical heads and distilling the rest into recurrent heads. We identify the essential heads via ablation on a synthetic retrieval task, producing a hybrid with sparse, non-uniform attention placement. We show that preserving **just 2% of attention heads recovers over 95% of teacher performance on retrieval-heavy tasks** (10 heads in a 1B model), requiring far fewer heads than hybrids that retain at least 25%. We further find that large recurrent states often compensate for missing retrieval: once retrieval is handled by these heads, the SSM backbone can be simplified with limited loss, even with an $8\times$ reduction in state dimension. By reducing both the attention cache and the SSM state, the resulting hybrid is $5$--$6\times$ more memory-efficient than comparable hybrids, closing the Transformer--SSM gap at a fraction of the memory cost.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。