COMPASS：定位并干预语言模型的推理注意头

摘要

显式引发推理可显着提高 LLM 性能。现有方法需要预定义的推理特征，无论是通过 CoT 提示设计、对比 CoT 方向，还是通过 SAE 派生的推理特征。对于具有可验证答案的数学推理，我们表明一个更简单的信号就足够了，这就是模型自身直接答案尝试的正确性。该信号产生了引发推理的潜在方向。这个方向在大多数注意力头的激活中是可解码的，但只有一小部分可以有效地干预。我们引入了 COMPASS，这是一种推理时转向方法，它使用 logit 空间归因分数来识别这些头，并沿着正确性方向引导它们的激活，只需要每个头的激活统计数据。在三个模型系列和多个数学基准中，COMPASS 的性能优于我们所比较的激活引导基线，将 GSM8K 准确度平均提高了 16 个百分点，并以生成的 token 减少 20-70% 的方式接近 CoT 准确度。干预措施无需重新拟合即可转移到看不见的基准，并且消融表明正确性方向和承载该方向的一小组头部都是必要的，其效果集中在极少数头部中。

COMPASS：定位并干预语言模型的推理注意头：论文原图 — 图 1：左：每个头的正确性方向。在头部的激活空间中，正确和错误回答问题的激活形成两个群体，分别为 $\mu_{+}^{l,h}$ 和 $\mu_{-}^{l,h}$。 $\theta_{l}^{h}$ 是连接这些均值的单位向量，$\sigma_{l}^{h}$ 是投影到 $\theta_{l}^{h}$ 上的激活的标准偏差。转向通过 $\alpha\,\sigma_{l}^{h}$ 沿着 $\theta_{l}^{h}$ 将激活转移到不正确的一侧。右：根据归因分数选择头部。顶部：每个磁头的分数 $s_{l}^{h}$，其中在提升层带内选择的 $K$ 磁头以蓝色框出。底部：在拟合行上测量的结果对比度 $\Delta p$。首先生成的 token 正确答案集中在“To”，错误答案集中在“$”。