从可解释性到性能：优化长上下文语言模型的检索头
计算与语言
2026-04-27 v3
摘要
机制可解释性的进展识别出了特殊的注意力头，即检索头，负责从上下文中检索信息。然而，这些检索头在提升模型性能方面的作用仍未被探索。本工作研究了是否可以利用检索头来增强 LLM 的长上下文能力。具体而言，我们提出了 RetMask，这是一种通过将正常模型输出与屏蔽检索头的消融变体输出进行对比来生成训练信号的方法。这种基于机制的方法取得了显著提升：在 Llama-3.1 的 128K 上下文长度下，HELMET 指标提升了 2.28 分，其中带引用的生成任务提升了 70%，段落重排序任务提升了 32%，同时保持了在通用任务上的性能。在三个模型族中的四个模型上进行的实验表明，RetMask 持续提升了长上下文性能，其增益与检索分数分布的稀疏性相关：分布更稀疏（即检索能力集中在少量注意力头上）的模型响应更强烈，而分布较不稀疏的模型增益则较为温和。这些结果验证了检索头的功能性作用，并表明机制层面的洞察可以转化为性能提升。
引用
@article{arxiv.2601.11020,
title = {From Interpretability to Performance: Optimizing Retrieval Heads for Long-Context Language Models},
author = {Youmi Ma and Naoaki Okazaki},
journal= {arXiv preprint arXiv:2601.11020},
year = {2026}
}
备注
Findings of ACL 2026; Source code available at https://github.com/YoumiMa/RetMask