[论文解读] Tracr: Compiled Transformers as a Laboratory for Interpretability

David Lindner, János Kramár+4 more|arXiv (Cornell University)|Jan 12, 2023

Explainable Artificial Intelligence (XAI)被引用 6

一句话总结

Tracr 将人类可读的 RASP 程序编译为具有已知、可解释结构的标准解码器-only 变压器模型，从而实现可控制的可解释性实验。通过生成具有真实计算机制的模型，Tracr 允许研究人员研究超位置（superposition）等现象，并验证基于梯度的归因和因果追踪等可解释性方法。

ABSTRACT

We show how to "compile" human-readable programs into standard decoder-only transformer models. Our compiler, Tracr, generates models with known structure. This structure can be used to design experiments. For example, we use it to study "superposition" in transformers that execute multi-step algorithms. Additionally, the known structure of Tracr-compiled models can serve as ground-truth for evaluating interpretability methods. Commonly, because the "programs" learned by transformers are unknown it is unclear whether an interpretation succeeded. We demonstrate our approach by implementing and examining programs including computing token frequencies, sorting, and parenthesis checking. We provide an open-source implementation of Tracr at https://github.com/google-deepmind/tracr.

研究动机与目标

• 通过创建具有已知、人类可读计算机制的变压器模型，解决可解释性研究中缺乏真实解释的问题。
• 为测试可解释性工具提供可复现且结构化的环境，其中真实计算过程可预先知晓。
• 实现对神经网络现象（如超位置）的可控研究，特别是在执行多步算法的模型中。
• 作为教学和实验平台，通过可编译、可追踪的模型理解变压器机制。
• 利用具有已知行为的模型，评估可解释性方法（如分类器探针、基于梯度的归因和因果追踪）

提出的方法

• Tracr 将 RASP（一种用于变压器计算的领域特定语言）编译为具有已知结构组件的标准解码器-only 变压器架构。
• 编译器生成的模型中，每一层的功能（如计算指示符、选择标记或聚合值）均直接映射到高级程序操作。
• 残差连接流被用作中心记忆机制，注意力和 MLP 层被明确设计为实现特定计算步骤（例如选择、聚合）。
• 该方法借鉴了变压器电路视角，通过查询-键和值-输出矩阵参数化注意力机制，以实现模块化、可解释的设计。
• 模型通过梯度下降进行训练，其内部表征被分析以验证频率计数、排序和括号检查等算法的正确实现。
• 通过压缩训练后的模型来研究超位置，观察特征在低维子空间中的剪枝或表示方式。

实验结果

研究问题

• RQ1 能否将从 RASP 程序编译的变压器模型用作评估可解释性方法的真实基准？
• RQ2 压缩后的变压器模型如何在超位置中表示特征？其是否如理论预测的那样丢弃不重要的特征？
• RQ3 基于梯度的归因或因果追踪等可解释性工具在多大程度上能正确恢复 Tracr 生成模型中已知的计算结构？
• RQ4 能否可靠地将多步算法（如排序或括号检查）编译为具有可解释注意力和 MLP 模式的变压器？
• RQ5 在使用秩-1谓词时，实现复杂逻辑操作（如 OR、AND）在注意力选择器中的局限性是什么？

主要发现

• Tracr 顺利将 RASP 程序编译为正确实现频率计数、排序和括号检查等算法的变压器模型，且其内部机制已知。
• 编译后的模型展现出清晰、可解释的注意力和 MLP 模式：例如，一个注意力头负责计算选择，另一个负责执行聚合，与预期的程序逻辑完全匹配。
• 通过梯度下降压缩后，Tracr 模型会丢弃不必要的特征，并在超位置中表示不重要的特征，证实了 Elhage 等人（2022b）的预测。
• 该模型使可解释性工具的验证成为可能——例如，基于梯度的归因和因果追踪可通过与已知真实机制对比进行验证。
• 布尔运算符（如选择器中的 OR）仅在特定条件下（如共享 s-ops）可实现，而一般情况受秩-1谓词约束的限制。
• 该方法表明，复杂算法可被编译为具有可预测、模块化组件的变压器，支持研究与教学应用。