当前位置：首页 > news >正文

论文解读：ThinkEdit: Interpretable Weight Editing to Mitigate Overly Short Thinking in Reasoning Models

news 2026/6/28 7:21:02

论文发表于NLP顶会EMNLP 2025（原文链接）。大模型CoT产生过短推理，即使简单数学问题也性能下降。本文研究推理长度如何嵌入推理模型的隐藏表示，以影响准确性：

1、发现，推理长度由表示空间中的线性方向决定，从而能沿着该方向引导模型，诱导过短推理。

2、引入权重编辑方法ThinkEdit，缓解过短推理：

识别小部分（约4%）驱动短推理行为的注意力。
编辑这些头部的输出投影权重，删除简短的推理方向。

方法

提升推理长度的Steering向量

在GSM8K和Math-Level5数据集上，将模型推理过程截断为定长，然后让大模型根据截断的推理内容生成答案。图1表明推理长度过短、过长都会损害性能。为此，作者想找对模型推理长度产生影响的特征向量。

1、用2000个GSM8K数据，根据模型推理长度在100 token以内和超过1000 token，将数据划分为$\mathcal{D}_{\text{short}},\mathcal{D}_{\text{long}}$。

2、将模型分别在$\mathcal{D}_{\text{short}},\mathcal{D}_{\text{long}}$上推理得到的中间表示进行平均，得到四个MLP和Attention层输出的平均表示：

$\bar{r}^{\text{mlp}}_{\ell,\text{short}}, \bar{r}^{\text{mlp}}_{\ell,\text{long}},\bar{r}^{\text{attn}}_{\ell,\text{short}},\bar{r}^{\text{attn}}_{\ell,\text{long}}$

3、将long表示减去short表示，得到steering向量：

$v_{\ell}^{\text{mlp}} = \bar{r}^{\text{mlp}}_{\ell,\text{long}} - \bar{r}^{\text{mlp}}_{\ell,\text{short}}$

$v_{\ell}^{\text{attn}} = \bar{r}^{\text{attn}}_{\ell,\text{long}} - \bar{r}^{\text{attn}}_{\ell,\text{short}}$

4、则当推理生成每个token的时候，以一定比例加上steering向量，即可对模型的推理长度进行调整：

$r_{\ell}^{\mathrm{mlp}} \leftarrow r_{\ell}^{\mathrm{mlp}} + \alpha v_{\ell}^{\mathrm{mlp}}$

$r_{\ell}^{\mathrm{attn}} \leftarrow r_{\ell}^{\mathrm{attn}} + \alpha v_{\ell}^{\mathrm{attn}}$

5、图3展示了对所有层的表示进行调整时，$\alpha$的变化对模型性能和推理长度的影响。可以看出steering向量的确是有效改编了模型推理长度和性能。图6展示了$\alpha$设置为-1/1时，单独调整某层表示时产生的影响。

ThinkEdit: 调整注意力头输出权重

为了观察不同的注意力头对Steering向量的贡献，如式(5)所示，将每个注意力头输出（因为它直接加到注意力输出上）与归一化的负steering向量$-\hat{v}_{\ell}^{\text{attn}}$内积，得到贡献度$\bar{C}^h_{\text{short}}$。图4可视化了各模型不同注意力头对Steering向量的贡献度。

作者找到short贡献前4%大的注意力头，通过调整它们的输出矩阵来提升模型推理长度：

$W_o^{\text{h}_\ell} \leftarrow W_o^{\text{h}_\ell} \left(I - (-\hat{v}_{\ell}^{\text{attn}}) (-\hat{v}_{\ell}^{\text{attn}})^{\top} \right)$

直觉上看，这个调整能把输出表示在Steering向量上的投影从输出表示中减去。