流形引导注意力转向

# 流形引导的注意力引导
来源：https://arxiv.org/html/2605.21770
Ian Li、Kapilesh Guruprasad、Raunak Sengupta  
i6li@ucsd\.edu kguruprasad@ucsd\.edu r2sengupta@ucsd\.edu
Ninad Satish、Loris D’Antoni、Rose Yu  
satish@ucsd\.edu ldantoni@ucsd\.edu roseyu@ucsd\.edu  
加州大学圣地亚哥分校

###### 摘要

大型语言模型在推理任务中经常产生错误，尽管它们拥有正确推理所需的基础知识。提高推理一致性的一种可能方法是激活引导。然而，现有的激活引导方法应用固定的、预先计算的修正向量，忽略了模型当前在生成轨迹中的位置；结果是，它们不加区分地扰动，既干扰正确的步骤，也干扰错误的步骤。我们提出**流形引导的注意力引导**（MAGS），这是一种基于几何观察的轨迹感知推理时干预：特定注意力头的输出激活在错误点会偏离低维的*正确性流形*，并且这种偏差会在后续步骤中累积。对于每个识别出的注意力头，我们从正确和错误轨迹的对比对中学习一个低维子空间，该子空间捕捉错误行为偏离正确行为的方向。在推理过程中，我们监控每个头到该流形的距离，当偏差超过学习到的阈值时，应用有针对性的投影修正，在错误传播之前将注意力输出引导回正确的子空间。MAGS 在数学推理（MATH-500、GSM8K）、代码生成（HumanEval、MBPP）和分子生成（SMILES）等多个基准测试中，始终优于未引导的基线和静态引导方法，这表明正确性流形是 LLM 注意力几何结构的一个普遍特征。

## 1 引言

大型语言模型（LLM）在多步推理任务中经常产生推理错误，尽管它们具备解决这些任务的基本能力。在重复采样下，产生错误解的同一个模型和提示往往也会产生正确解[3 (https://arxiv.org/html/2605.21770#bib.bib17)、20 (https://arxiv.org/html/2605.21770#bib.bib21)、4 (https://arxiv.org/html/2605.21770#bib.bib19)]，这表明能力是存在的，但其稳定表达并非如此。过程级标注进一步证实，大多数错误发生在中间推理步骤，而非源自最终的知识缺失[11 (https://arxiv.org/html/2605.21770#bib.bib6)、17 (https://arxiv.org/html/2605.21770#bib.bib22)]。由于错误源于生成过程而非知识缺失，因此在推理时修正错误是一个自然且实际的目标。

现有的激活引导方法[16 (https://arxiv.org/html/2605.21770#bib.bib1)、24 (https://arxiv.org/html/2605.21770#bib.bib2)、14 (https://arxiv.org/html/2605.21770#bib.bib11)、18 (https://arxiv.org/html/2605.21770#bib.bib13)]在生成步骤中向残差流应用固定的修正向量。这些方法非常适合持续性的全局行为（语气、风格、情感），但在结构上与推理不匹配。一个推理轨迹可能在许多步骤中正确运行，然后在步骤 $t^*$ 处犯下局部错误；应用恒定修正可能会破坏正确的中间步骤，同时无法保证能拦截错误。

我们假设推理错误表现为特定注意力头输出空间的低维子空间中的漂移：正确和错误的轨迹占据几何可分离的区域，并且从正确行为到错误行为的转变遵循结构化的低秩方向。我们进行了诊断性实验来验证这一假设，并事实上发现正确和错误的轨迹可以通过注意力头激活的低维子空间高度分离。这与机制可解释性的发现一致，即单个注意力头具有功能特化[5 (https://arxiv.org/html/2605.21770#bib.bib4)、19 (https://arxiv.org/html/2605.21770#bib.bib5)]，也与线性表示假设[13 (https://arxiv.org/html/2605.21770#bib.bib3)]一致，该假设认为语义上有意义的区别沿着低维线性方向编码。

因此，只有当注意力输出漂移到错误子空间时，才应该进行引导。

我们提出**流形引导的注意力引导**（MAGS）：一种自适应干预方法，在检测到推理错误时动态引导注意力头的输出。MAGS 在三个模型家族（包括 Llama、Gemma 和 GPT-OSS）的推理基准测试和分子生成任务中，均优于静态引导基线。

总之，我们的贡献如下：

1. 我们假设推理错误表现为单个注意力头输出空间低维子空间中的结构化漂移，并通过诊断性实验证实了这一假设，表明正确和错误轨迹高度可分离（第3节 (https://arxiv.org/html/2605.21770#S3)）。
2. 我们提出流形引导的注意力引导（MAGS），一种自适应推理时机制，用于监控注意力头是否发生推理漂移，并在需要时动态应用修正（第4节 (https://arxiv.org/html/2605.21770#S4)）。
3. 实验上，MAGS 在三个模型家族的基准测试中始终优于静态引导基线，提升高达10.8%，同时仅产生可忽略的推理开销（第5节 (https://arxiv.org/html/2605.21770#S5)）。

## 2 相关工作

我们讨论现有的推理时引导方法和几何可解释性工作。现有的引导方法应用固定修正，没有错误检测机制；现有的可解释性工作建立了我们利用并拓展用于自适应干预的几何结构。据我们所知，MAGS 是第一个将逐步骤检测与注意力头级别的几何感知条件修正相结合的方法。

参见图注  
图1：在一个示例问题上，静态引导与流形引导的注意力引导（MAGS）的对比。蓝色框表示正确的推理步骤；红色框表示错误的步骤。  
### 2.1 激活引导与推理时干预

激活引导方法在推理时修改内部表示，而不更新模型参数。*激活加法*[16 (https://arxiv.org/html/2605.21770#bib.bib1)] 在生成过程中向残差流添加一个固定的差分向量。*对比激活加法*（CAA；Rimsky 等人，14 (https://arxiv.org/html/2605.21770#bib.bib11)) 通过平均多个对比提示对的差分向量来提高可靠性。*表示工程*（RepE；Zou 等人，24 (https://arxiv.org/html/2605.21770#bib.bib2)) 通过 PCA 从对比激活中提取主引导方向。*角度引导*[18 (https://arxiv.org/html/2605.21770#bib.bib13)] 将加法修正替换为均匀应用于所有层的 2D 旋转。*推理时干预*（ITI；Li 等人，9 (https://arxiv.org/html/2605.21770#bib.bib12)) 沿着探针方向移动注意力头输出以提高真实性。*CREST*[23 (https://arxiv.org/html/2605.21770#bib.bib7)] 识别与推理相关的注意力头，并向它们应用固定的引导向量，但不会根据模型当前的轨迹状态调整修正。

所有这些方法都有一个共同的局限性：无论当前激活状态如何，它们都沿着固定方向引导。MAGS 通过引入一个动态的接近度触发器来解决这个问题，该触发器仅在头向错误子空间漂移时触发，并应用依赖步骤的修正，其方向由当前激活在错误子空间上的投影决定，而不是固定向量（如图1 (https://arxiv.org/html/2605.21770#S2.F1) 所示）。

### 2.2 Transformer 表示的几何结构

越来越多的研究表明，Transformer 表示具有丰富的几何结构，可用于分析和干预。

*机制可解释性* 研究将 Transformer 计算分解为可解释的电路。Elhage 等人 [5 (https://arxiv.org/html/2605.21770#bib.bib4)] 表明注意力头实现基本操作（复制、检索、抑制），其输出在残差流中相加组合。Wang 等人 [19 (https://arxiv.org/html/2605.21770#bib.bib5)] 表明多步骤任务由少量头之间的稀疏电路实现。这些结果共同表明，推理失败很可能归因于特定头的故障模式，这激发了我们的头级别干预。

*线性表示假设*[13 (https://arxiv.org/html/2605.21770#bib.bib3)] 认为语义上有意义的区别沿着 Transformer 表示中的低维线性方向编码。Burns 等人 [2 (https://arxiv.org/html/2605.21770#bib.bib20)] 表明真实性具有可通过对比探针发现的线性表示，为我们的对比 PCA 构造建立了先例。Zou 等人 [24 (https://arxiv.org/html/2605.21770#bib.bib2)] 证实了高层次概念（包括推理质量）可从残差流激活中线性解码。MAGS 将这一研究方向扩展到每个头的级别，表明推理轨迹中从正确到错误的转变在单个头输出空间内也是线性结构化的。

## 3 检测注意力头中的错误漂移

我们假设不正确的推理轨迹会导致一部分注意力头的输出向低维错误子空间漂移，该子空间在几何上可分离于正确轨迹占据的子空间。我们通过为每个头构建一个对比错误流形，并证明基于接近度的分数能够在不同层和头之间实现强轨迹级错误检测，来实证验证这一假设。

### 3.1 设置与符号

我们考虑一组推理问题 $\mathcal{P}=\{p_1,\ldots,p_N\}$。对于每个问题 $p_i$，我们生成 $S$ 条独立的推理轨迹 $\mathcal{T}_i=\{\tau_{i,1},\ldots,\tau_{i,S}\}$。每条轨迹 $\tau$ 是长度为 $L_\tau$ 的 token 序列，并被赋予一个二元标签 $y_\tau\in\{0,1\}$（1 = 正确最终答案）。假设对于所有 $p_i$，有 $|\mathcal{T}_i^+|\geq 1$ 且 $|\mathcal{T}_i^-|\geq 1$。我们用 $\mathcal{T}_i^+=\{\tau\in\mathcal{T}_i:y_\tau=1\}$ 和 $\mathcal{T}_i^-=\{\tau\in\mathcal{T}_i:y_\tau=0\}$ 分别表示问题 $p_i$ 的正确轨迹集和错误轨迹集。

给定一个基于 Transformer 的语言模型，具有 $L$ 层和每层 $H$ 个注意力头，每个头操作在 $d_h$ 维输出空间上。对于一个样本 $\tau$ 和头 $(l,h)$，注意力头输出的序列为：

$$
\mathbf{A}_\tau^{(l,h)} \;=\; \bigl[\boldsymbol{a}_1^{(l,h,\tau)},\; \boldsymbol{a}_2^{(l,h,\tau)},\; \ldots,\; \boldsymbol{a}_{L_\tau}^{(l,h,\tau)}\bigr] \;\in\; \mathbb{R}^{L_\tau \times d_h}. \tag{1}
$$

### 3.2 对比错误流形构建

给定正确和错误的推理轨迹，我们为每个头构建一个错误子空间，通过识别正确和错误激活发散的低维方向。

#### 每个问题的差分向量。

对于每个问题和每个头 $(l,h)$，定义每个类别的均值：

$$
\boldsymbol{\mu}_{c,i}^{(l,h)} = \frac{1}{\sum_{\tau\in\mathcal{T}_i^+} L_\tau} \sum_{\tau\in\mathcal{T}_i^+} \sum_{t=1}^{L_\tau} \boldsymbol{a}_t^{(l,h,\tau)}, \qquad
\boldsymbol{\mu}_{e,i}^{(l,h)} = \frac{1}{\sum_{\tau\in\mathcal{T}_i^-} L_\tau} \sum_{\tau\in\mathcal{T}_i^-} \sum_{t=1}^{L_\tau} \boldsymbol{a}_t^{(l,h,\tau)}. \tag{2}
$$

对于问题 $p_i$ 和头 $(l,h)$ 的*对比差分向量*是：

$$
\boldsymbol{\delta}_i^{(l,h)} \;=\; \boldsymbol{\mu}_{e,i}^{(l,h)} \;-\; \boldsymbol{\mu}_{c,i}^{(l,h)} \;\in\; \mathbb{R}^{d_h}. \tag{3}
$$

根据构造，$\boldsymbol{\delta}_i^{(l,h)}$ 抵消了问题 $p_i$ 所有均匀激活的方向，无论正确性如何，从而隔离了可归因于错误的*方向偏移*。

#### 差分矩阵。

参考图注  
图2：对比错误流形示意图。正确和错误激活空间沿着学习到的子空间 $\mathrm{span}(\mathbf{B})$ 分离。给定一个错误激活 $\boldsymbol{a}_t$，投影 $\mathbf{B}^\top\!\mathbf{B}$ 给出了映射回 $\boldsymbol{\mu}_c^{(l,h)}$ 的方向。  
对于一组 $N$ 个问题，我们将差分向量按行堆叠：

$$
(\mathbf{D}^{(l,h)})^T = \begin{bmatrix} \boldsymbol{\delta}_1^{(l,h)} & \dots & \boldsymbol{\delta}_P^{(l,h)} \end{bmatrix} \;\in\; \mathbb{R}^{d_h \times N}. \tag{4}
$$

#### 通过 PCA 得到错误子空间。

计算紧凑奇异值分解 $\mathbf{D}^{(l,h)} = \mathbf{U}\mathbf{\Sigma}\mathbf{V}^\top$。定义*错误子空间基*为前 $k$ 个右奇异向量：

$$
\mathbf{B}^{(l,h)} \;=\; \mathbf{V}_{:,1:k}^\top \;\in\; \mathbb{R}^{k \times d_h}, \tag{5}
$$

其中 $\mathbf{B}^{(l,h)}$ 的行是标准正交的。错误子空间捕捉了头输出空间中，从正确到错误偏差在不同问题上具有最大方差的 $k$ 个方向，如图2 (https://arxiv.org/html/2605.21770#S3.F2) 所示。

#### 全局正确状态质心。

从所有正确轨迹计算一个全局参考点：

$$
\boldsymbol{\mu}_c^{(l,h)} \;=\; \frac{\displaystyle\sum_{i}\sum_{\tau\in\mathcal{T}_i^+}\sum_{t=1}^{L_\tau} \boldsymbol{a}_t^{(l,h,\tau)}}{\displaystyle\sum_{i}\sum_{\tau\in\mathcal{T}_i^+} L_\tau}. \tag{6}
$$

这用作推理时的定心参考，因为在生成过程中无法获得每个问题的均值。

### 3.3 基于接近度的错误检测

为了检测一个头何时漂移到错误子空间，我们测量其当前输出投影到学习到的错误子空间上的幅度。大的投影表明该头的行为类似于它在错误轨迹中的行为。在推理时的每个解码步骤 $t$ 中，对于每个被监控的头 $(l,h)$，我们计算*接近度分数*：

$$
d_t^{(l,h)} \;=\; \left\| \mathbf{B}^{(l,h)} \bigl( \boldsymbol{a}_t^{(l,h)} - \boldsymbol{\mu}_c^{(l,h)} \bigr) \right\|^2 \;=\; \bigl( \boldsymbol{a}_t^{(l,h)} - \boldsymbol{\mu}_c^{(l,h)} \bigr)^\top \mathbf{B}^{(l,h)\top} \mathbf{B}^{(l,h)} \bigl( \boldsymbol{a}_t^{(l,h)} - \boldsymbol{\mu}_c^{(l,h)} \bigr).
$$