MechRL：强化学习代理用于机制可解释性中的电路发现

# 强化学习代理执行电路发现以实现机械可解释性
来源：https://arxiv.org/html/2605.26343
Barsat Khadka 南密西西比大学 哈蒂斯堡，MS 39406 Barsat\.Khadka@usm\.edu

###### 摘要

机械可解释性已识别出在Transformer语言模型中实现特定行为的小型注意力头集合，但恢复这些电路通常需要为每个新任务设计定制化的分析流程。我们将电路发现重新定义为一个强化学习问题。一个代理在GPT-2 small的144个注意力头上运行，作为离散动作空间；每个动作触发一次零消融和一个对比奖励，该奖励将消融对通用下一个词预测的损害减去其对目标任务的损害。一个在向量化多任务环境中针对两个任务（归纳和IOI）训练的单一PPO策略，在两个训练任务以及一个保留的第三个任务（文档字符串补全）上均达到了每回合的最优解（oracle）。其偏好的头与Olsson等人（2022）和Wang等人（2023）在单头消融下因果非冗余的轴线上精确识别的标准头一致；他们识别为冗余的类别被代理正确地去优先化。在保留任务上，最佳五选一规划在评估时未提供任何任务信号的情况下，恢复了96%的最优解上限。这些结果表明，基于因果干预的强化学习是识别机械电路单头瓶颈的一种可行且可迁移的方法，与现有的路径修补方法互补。

## 1 引言

机械可解释性旨在逆向工程神经网络，以理解特定内部操作如何产生模型行为。在Transformer模型中，这涉及识别*电路*：注意力头、MLP层和残差连接的稀疏子图，它们对给定任务因果充分（Elhage等人，2021）。虽然这种方法已经分离出诸如归纳电路（Olsson等人，2022）和IOI电路（Wang等人，2023）等机制，但当前的方法仍然高度局部化。无论是手动分析还是自动方法如ACDC（Conmy等人，2023）和EAP（Nanda，2023），都在孤立地评估单个任务，而不在不同行为之间共享学习到的结构模式。

在这项工作中，我们研究学习得到的策略是否能够自动化电路发现并泛化到未见过的任务。我们将电路发现框架化为一个马尔可夫决策过程。在每一步，代理零消融一个注意力头，观察基于消融因果效应的奖励，并更新其策略。通过在每回合随机化任务批次，我们防止代理记忆特定的头位置，迫使其学习模型几何结构上的结构先验。

我们的奖励函数采用对比设计来区分任务特定头与通用计算头。因为消融一个通用头会降低所有任务的性能，标准的对数几率差下降是不够的。因此，我们将每次消融与一批自然英语文本的控制评估配对，并从目标任务的对数几率差下降中减去控制上交叉熵的增加。这隔离了任务特定的因果效应，并且正如我们实证展示的那样，提供了一个可靠的强化学习信号。

## 2 方法

### 2.1 环境

我们使用冻结的预训练GPT-2 small（\(n_{\text{layers}}=n_{\text{heads}}=12\)，\(n_{\text{actions}}=144\)）。环境暴露标准的Gymnasium API。每回合最多\(T=50\)步；动作空间\(\mathcal{A}=\{0,\dots,143\}\)通过\(a=12\ell+h\)双射对应到\((\ell,h)\)。一个动作掩码防止在回合内重复任何头。

在每次`reset()`时，环境从训练集\(\{\text{induction},\,\text{IOI}\}\)中均匀采样一个任务\(\tau\)，并使用回合种子生成一个新的任务批次。归纳批次遵循标准的\([\text{BOS},A,B,\text{fillers},A]\)模板，其中\(A,B\)从词汇表中均匀采样；IOI批次遵循Wang等人（2023）的模板家族，随机抽取名称-物体-地点，并平衡ABBA/BABA顺序。每回合重新采样强制执行随机奖励景观：在回合\(t\)中消融损害最大的头通常不是回合\(t'\)中损害最大的头，因此策略必须内化分布性而非点对点的结构。

### 2.2 对比奖励

设\(M^\tau_\theta(B)\)表示在冻结模型\(\theta\)上，对批次\(B\)评估的任务特定标量度量（相关位置上正确标记与干扰标记之间的对数几率差）。设\(L^{\text{ctrl}}_\theta(C)\)表示\(\theta\)在从wikitext-2-raw-v1抽取的自然英语文本控制批次\(C\)上的平均交叉熵。对于消融模型\(\theta_{\setminus a}\)，其中\((l,h)\)注意力头在hook\_z的输出被置为零，我们定义每动作奖励

\[
r(a) = \underbrace{\bigl(M^\tau_\theta(B)-M^\tau_{\theta_{\setminus a}}(B)\bigr)}_{\text{任务损害}} \;-\; \underbrace{\bigl(L^{\text{ctrl}}_{\theta_{\setminus a}}(C)-L^{\text{ctrl}}_\theta(C)\bigr)}_{\text{通用损害}}.
\]
当消融头\(a\)不成比例地损害目标任务时，奖励为正；对于消融后模型通用分布拟合退化程度相当的头部，奖励近似为零。每次查询需要两次前向传播通过\(\theta\)（一个任务批次，一个控制批次）；两者在`reset()`时预先分配在GPU上，并在回合中的所有步骤中复用。

### 2.3 观测

代理观测到\(o_t \in \mathbb{R}^{2+2n_{\text{actions}}}\)，包含三个块：一个2维独热编码表示当前任务\(\tau\)；一个144维二进制掩码表示先前尝试过的头；以及一个144维向量表示先前尝试过的头的归一化奖励\(r(a)/\sigma\)（否则为零），固定\(\sigma=5\)。奖励通道实现了回合内适应：策略可以根据当前回合已收到的反馈来调节未来的动作。

### 2.4 多任务向量化PPO

我们使用近端策略优化（Schulman等人，2017）在\(N=8\)个同步并行环境上训练，这些环境共享GPU上的单一GPT-2实例，每次梯度更新产生\(N\cdot T=400\)个在策略转移。初始的单环境版本在大约2000步后崩溃为近乎确定性的策略；我们将其归因于每更新批次太小，无法进行稳定的优势估计。

演员-评论家是一个前馈MLP，架构为\(290\to 256\to 256\to (144,\,1)\)，使用tanh激活函数，并遵循标准PPO实践进行正交初始化。我们使用初始学习率\(2.5\times10^{-4}\)的Adam优化器，线性退火至初始值的20%，\(\gamma=0.99\)，GAE \(\lambda=0.95\)，裁剪系数0.2，熵系数0.1，价值损失系数0.5，以及4个更新轮次，每轮8个小批量。训练总共进行\(T_{\text{total}}=2\times10^5\)个环境步。

### 2.5 最佳K选一规划

由于环境动态是GPT-2前向传播，相对于策略网络是确定且廉价的，我们为代理配备一个最佳K选一规划包装器。在决策时，策略从其当前分布中无放回地采样\(K\)个候选动作；环境并行评分每个动作；代理选择得分最高的候选。\(K=1\)恢复标准PPO。我们在全文报告中给出\(K\in\{1,5\}\)的结果。

### 2.6 保留评估任务

为了评估泛化能力，我们引入一个策略在训练期间从未见过的文档字符串补全任务，风格类似于Heimersheim和Janiak（2023）。每个提示给出一个带有\(n=5\)个参数的Python函数签名；文档字符串以`:param`形式描述前\(n-1\)个参数；模型必须预测紧随最后一个`:param`位置处缺失的最终参数名。任务度量是正确缺失参数名与已提及参数中抽取的干扰项之间的对数几率差。

我们注意到Heimersheim和Janiak（2023）报告了4层仅注意力玩具模型的规范文档字符串电路，而非GPT-2 small。因此，在我们的模型中，此任务没有已公布的头级别真值，我们仅报告最优解一致性和迁移性能，而非规范重叠。

## 3 结果

我们报告三类结果：每任务最优性的收敛性、与规范可解释性发现的一致性、以及向保留任务的迁移。

### 3.1 收敛到每任务最优解

对于每个任务，我们通过在回合批次上穷尽消融所有144个头并记录单一最高得分头来计算每回合最优解。在保留评估种子带（\(\text{seed}\geq 10^7\)）上平均得到*最优解上限*：在等式(1)的对比奖励下，任何单头策略能获得的最高运行最大奖励。我们将其与训练策略在其50个回合内选择中获得的运行最大奖励进行比较。

表1：训练所得\(K=1\)策略的平均运行最大对比奖励与每任务每回合最优解的比较，基于每个任务\(\geq 20\)个保留评估回合。策略在所有三个任务上（包括保留的文档字符串任务）在采样噪声范围内达到最优解上限。策略与最优解之间的差距在所有三个任务上均以\(|0.028|\)为界，在最优化自身估计的每回合标准误差之内。我们得出结论，对比奖励是适定的：最优解会选择的主要头部，策略高概率也会选择。这为以下分析提供了许可：策略的选择频率分布是奖励景观本身的合理代理。

### 3.2 与规范可解释性文献的一致性

我们将训练策略的选择频率分布与先前工作中识别的规范头部集合进行比较。对于每个规范头，我们报告 (a) 在20个保留回合中策略选择频率排序中的排名，以及 (b) 在10个保留回合中最优解平均得分排序中的排名，两者均基于144个头。

#### 归纳。

GPT-2 small中的规范归纳头集合被广泛认为是 \(\{L5.H1, L5.H5, L6.H9, L7.H2, L7.H10\}\)（Olsson等人，2022）。

表2：训练策略（\(K=1\)，20回合）和每回合最优解（10回合）中规范归纳头的选择频率排名。“–”表示该头从未被选中。最优解排名低的头（L7.H10, L5.H1）也在策略中缺失。五个规范归纳头中的两个——L5.H5和L7.H2——出现在策略的前十大选择中。其余三个缺失，但最优解确认这并非策略失败：L7.H10和L5.H1分别位于最优解排名的第129和第130位。在自然文本控制的单头零消融下，这些头对归纳的损害不超过它们对通用下一个词预测的损害。Olsson等人（2022）认定为模型中最典型归纳头的L5.H5正确地出现在第6位。

#### IOI。

IOI电路分解为七个子类别（Wang等人，2023）；表3报告了策略前十选与规范集合之间每个类别的重叠。

表3：训练策略前十最频繁选择头部与规范IOI电路每个子类别的重叠（Wang等人，2023）。策略恢复了4个S-抑制头中的3个和4个IOI中归纳头中的2个；它没有选择任何名称移动头，这与Wang等人（2023）的最小性分析一致，该分析表明单头名称移动消融被备份名称移动头所掩盖。S-抑制头主导了策略的IOI前十选：L8.H10排名第1，L7.H9排名第3，L8.H6排名第4。四个IOI中归纳头中的两个紧随其后，排名第9和第10。名称移动头的缺失反映了Wang等人（2023）记录的冗余结构：消融单个名称移动头仅引起约5%的对数几率差下降，因为备份名称移动头重新路由了相关信号。最优解从外部确认了这一点——最常被引用的名称移动头L9.H6位于最优解排名的第138位。该方法浮出了*单头消融下* IOI电路的瓶颈，这与S-抑制家族而非更具行为可解释性的名称移动家族一致。

这并非方法的缺陷，而是单头消融本身的属性。一个电路中最具可解释性的头不一定是在单元级干预下最因果不可或缺的头。我们的代理与Wang等人（2023）的最小性分析之间的趋同表明，单头因果不可或缺性是一个一致且可重复的解释性目标。

### 3.3 向保留任务的迁移

我们在三种任务信号模式下评估向文档字符串任务的迁移：\([0,0]\)（零样本；训练期间从未见过的向量），\([1,0]\)（提示为归纳），和\([0,1]\)（提示为IOI）。

表4：在保留文档字符串任务上30个回合的平均运行最大对比奖励。在零样本\([0,0]\)信号下，最佳五选一规划恢复了96%的最优解上限。提示为归纳时达到上限；提示为IOI时低于随机基线，表明任务信道被非平凡地使用。表4的三个特征值得评论。首先，零样本\(K=5\)结果（2.103 对比最优解2.175，随机基线1.395）构成了非平凡的泛化：策略在一个从未见过的任务上恢复了96%的最优奖励。相对于随机高出0.71（相对提升51%）完全归因于学习到的先验加上规划下的回合内奖励反馈。其次，提示差异在\(K=1\)时为2.178对比1.057，表明策略维持了不同的任务条件策略，且任务独热编码充当路由信号而非遗迹输入。第三，规划对于迁移的重要性远高于训练：\(K=5\)产生