[论文解读] Emergent Linear Representations in World Models of Self-Supervised Sequence Models

Neel Nanda, Andrew Lee+1 more|arXiv (Cornell University)|Sep 2, 2023

Neural Networks and Reservoir ComputingComputer Science被引用 3

一句话总结

该论文表明，OthelloGPT——一种自监督序列模型——通过线性编码的相对颜色（'Mine'，'Yours'，'Empty'）来表示棋盘状态，而非绝对颜色。通过探测这些线性方向，可使用简单的向量算术精确控制模型的内部状态，实现可解释且高效的干预，其性能优于以往基于梯度的方法。

ABSTRACT

How do sequence models represent their decision-making process? Prior work suggests that Othello-playing neural network learned nonlinear models of the board state (Li et al., 2023). In this work, we provide evidence of a closely related linear representation of the board. In particular, we show that probing for "my colour" vs. "opponent's colour" may be a simple yet powerful way to interpret the model's internal state. This precise understanding of the internal representations allows us to control the model's behaviour with simple vector arithmetic. Linear representations enable significant interpretability progress, which we demonstrate with further exploration of how the world model is computed.

研究动机与目标

• 探究自监督序列模型（如 OthelloGPT）如何表示内部世界状态，特别是这些表示是否具有线性结构。
• 解决先前研究声称棋盘状态编码为非线性与当前发现线性表示之间的明显矛盾。
• 开发一种更简单、更具可解释性的方法，通过内部激活的向量算术控制模型行为。
• 提供关于模型如何计算合法走法以及检测空格或翻转棋子的机制性洞见。
• 探索世界模型中线性表示的出现及其对可解释性和控制的影响。

提出的方法

• 作者将棋盘状态表示从绝对颜色（黑，白，空）重新定义为基于当前玩家视角的相对颜色（我的，你的，空），该定义随时间步变化，取决于行动顺序。
• 训练线性探测器以在每个标记位置预测相对颜色状态，对 'Mine/Yours/Empty' 分类任务达到近乎完美的准确率（例如，第4个时间步达到98.3%）。
• 通过将模型内部激活投影到已学习的 'Mine'、'Yours' 和 'Empty' 线性方向上，对向量算术进行应用，从而实现对模型内部信念状态的直接操控。
• 通过编辑激活以模拟修改后的棋盘状态并观察模型生成的走法是否与修改后的信念状态一致，对方法进行验证。
• 作者进一步探测其他线性特征，如每个时间步的翻转棋子，并通过分析残差流动态检测多个用于走法预测的电路。
• 比较线性与非线性探测性能，结果表明：尽管非线性探测器在绝对颜色标签上表现优异，但针对相对颜色的线性探测器仍能达到近乎最优的性能。

Figure 1: The emergent world models of OthelloGPT are linearly represented. We find that the board states are encoded relative to the current player’s colour ( Mine vs. Yours ) as opposed to absolute colours ( Black vs. White ).

实验结果

研究问题

• RQ1 OthelloGPT 中棋盘状态的内部表示是否表现出线性结构？如果是，其形式如何？
• RQ2 能否在激活空间的线性方向上使用向量算术，以因果方式改变模型的内部世界状态和行为？
• RQ3 为何线性表示会在自监督世界模型中出现？它们与先前报告的非线性表示有何关系？
• RQ4 通过线性探测，能获得哪些关于模型如何计算合法走法以及检测空格或翻转棋子的机制性洞见？
• RQ5 模型中是否存在多个用于预测合法走法的独立电路？它们与线性特征表示有何关联？

主要发现

• 该模型并非以绝对颜色（黑，白，空）表示棋盘状态，而是基于当前玩家的回合，以相对术语（'Mine'，'Yours'，'Empty'）表示，从而实现高度准确的线性探测（例如，第4个时间步准确率达98.3%）。
• 对 'Mine/Yours/Empty' 方向的线性投影在重建内部世界模型方面表现近乎完美，优于针对绝对颜色标签的非线性探测器。
• 对内部激活应用简单的向量算术，可实现对模型行为的有效且可解释的控制，例如引导其在编辑后的棋盘状态下走合法步。
• 该方法无需梯度计算即可实现因果干预，为基于梯度的编辑技术提供了更简单、更透明的替代方案。
• 模型显示出多个用于走法预测的电路证据，尤其在残局阶段；同时可检测到棋子翻转事件的线性表示。
• 研究结果表明，线性表示可能比以往认为的更为基础，残差连接使特征在各层间迭代优化成为可能。

Figure 2: Intervening methodology: we intervene by adding either Empty , Mine , or Yours directions into each layer of the residual stream. Red squares in each board indicate the tiles that have been intervened, teal tiles indicate new legal moves post-intervention that the model predicts.