通过激活检索头提升大模型长文本处理能力
From Interpretability to Performance: Optimizing Retrieval Heads for Long-Context Language Models
- 用对比法生成训练信号，优化检索注意力头
- 长文本任务性能提升最高达70%，128K上下文+2.28分
- 适合研究长序列建模与可解释性结合的学者
机制可解释性研究发现特殊注意力头（称为检索头）负责从上下文中提取信息，但其对模型性能的影响尚未明确。本文探究是否可通过优化检索头提升大模型的长上下文能力。提出RetMask方法：通过对比正常输出与掩蔽检索头后的输出生成训练信号。该基于机制的方法带来显著提升：在Llama-3.1上，128K上下文的HELMET得分提升+2.28点；引用生成任务提升+70%，段落重排序提升+32%，且通用任务性能不受影响。跨四个模型、三个架构家族的实验表明，性能增益与检索分数分布稀疏度正相关——检索能力集中在少数头部的模型获益更显著，而分布较密集的模型提升有限。结果验证了检索头的功能作用，并证明机制洞察能有效转化为性能改进。
原文摘要 · Abstract (English)
Advances in mechanistic interpretability have identified special attention heads, known as retrieval heads, that are responsible for retrieving information from the context. However, the role of these retrieval heads in improving model performance remains unexplored. This work investigates whether retrieval heads can be leveraged to enhance the long-context capabilities of LLMs. Specifically, we propose RetMask, a method that generates training signals by contrasting normal model outputs with those from an ablated variant in which the retrieval heads are masked. This mechanism-based approach achieves substantial improvements: +2.28 points on HELMET at 128K for Llama-3.1, with +70% gains on generation with citation and +32% on passage re-ranking, while preserving performance on general tasks. Experiments across four models in three families demonstrate that RetMask consistently improves long-context performance, where gains correlate with the sparsity of the retrieval score distribution: models with sparser distributions, where retrieval capabilities are concentrated in a small set of heads, respond more strongly, while those with less sparse distributions show more modest gains. These results validate the functional role of retrieval heads and show that mechanistic insights can be transformed into performance enhancements.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。