自注意力层是检索任务的关键，稀疏化15%注意力头仍可保持高检索精度。
Some Attention is All You Need for Retrieval
- 仅依赖自注意力层完成检索，状态空间模型层无法补偿
- 仅保留15%注意力头时检索准确率仍接近100%，MMLU得分保持84%
- 适合关注模型可解释性与架构优化的研究者
我们发现混合状态空间-变压器架构中存在完全的功能分离：检索任务仅依赖自注意力层。在RecurrentGemma-2B/9B和Jamba-Mini-1.6上，移除注意力层导致检索失败（准确率为0%），而状态空间模型层即使通过改进提示也无法补偿。相反，将注意力稀疏至仅15%的注意力头，仍能维持接近完美的检索性能，同时保持84%的MMLU得分，表明自注意力主要专用于检索。我们识别出检索所需的精确机制：生成过程中必须暴露“针头”标记，且预填充或生成阶段需有足够的上下文。这一严格的功能专业化挑战了混合架构冗余性的假设，表明这些模型更像专用模块而非集成系统，对架构优化与可解释性具有直接启示。
原文摘要 · Abstract (English)
We demonstrate complete functional segregation in hybrid SSM-Transformer architectures: retrieval depends exclusively on self-attention layers. Across RecurrentGemma-2B/9B and Jamba-Mini-1.6, attention ablation causes catastrophic retrieval failure (0% accuracy), while SSM layers show no compensatory mechanisms even with improved prompting. Conversely, sparsifying attention to just 15% of heads maintains near-perfect retrieval while preserving 84% MMLU performance, suggesting self-attention specializes primarily for retrieval tasks. We identify precise mechanistic requirements for retrieval: needle tokens must be exposed during generation and sufficient context must be available during prefill or generation. This strict functional specialization challenges assumptions about redundancy in hybrid architectures and suggests these models operate as specialized modules rather than integrated systems, with immediate implications for architecture optimization and interpretability.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。