从可解释性到性能：优化长上下文语言模型的检索注意力头

机制可解释性研究已识别出一类特殊的注意力头——检索头（retrieval heads），其负责从上下文中检索信息。然而，这些检索头对模型性能的贡献尚未被充分探索。本文研究了如何利用检索头提升大语言模型的长上下文能力，并提出RetMask方法：通过对比正常模型输出与屏蔽检索头后的消融模型输出，生成训练信号。该基于机制的方法在Llama-3.1上于128K上下文长度下，在HELMET基准上提升2.28分，引用生成任务提升70%，段落重排序提升32%，同时保持通用任务性能。跨三个模型家族的实验表明，效果取决于检索头的组织方式：集中式模式响应显著，而分布式模式增益有限。该结果验证了检索头的功能，并证明机制洞察可转化为性能提升。