Tracr: RASP 的转换器编译器

Tracr 是一个用于将 RASP 程序（Weiss et al. 2021）转换为转换器权重的编译器。详细的编译器描述请参阅我们的技术报告。

目录结构：

rasp 包含嵌入 Python 的 RASP 实现。
compiler 包含编译器本身。
transformer 包含转换器的实现。
craft 包含编译器使用的中间表示：本质上是一个具有命名维度的小型线性代数库。

这不是 Google 官方支持的产品。

安装

只需克隆并使用 pip 安装：

git clone https://github.com/deepmind/tracr
cd tracr
pip3 install .

使用示例：RASP `reverse` 程序

考虑 RASP reverse 程序：

opp_index = length - indices - 1;
flip = select(indices, opp_index, ==);
reverse = aggregate(flip, tokens);

要使用 Tracr 编译这个程序，我们首先使用 Tracr 的 RASP 库实现该程序：

from tracr.rasp import rasp

length = make_length()  # 在我们的实现中，`length` 不是原始操作。
opp_index = length - rasp.indices - 1
flip = rasp.Select(rasp.indices, opp_index, rasp.Comparison.EQ)
reverse = rasp.Aggregate(flip, rasp.tokens)

其中：

def make_length():
  all_true_selector = rasp.Select(rasp.tokens, rasp.tokens, rasp.Comparison.TRUE)
  return rasp.SelectorWidth(all_true_selector)

然后我们可以将 RASP 程序编译为转换器：

from tracr.compiler import compiling

bos = "BOS"
model = compiling.compile_rasp_to_model(
    reverse,
    vocab={1, 2, 3},
    max_seq_len=5,
    compiler_bos=bos,
)

这会生成一个作为 Haiku 模型的转换器。这个模型并不是为了提供你可能需要的所有内容，而是作为生成参数语义的一种"代码文档"。预期用户可以编写或贡献一个适配器，将此参考模型的参数转换为另一个转换器实现。

使用这个模型，我们可以执行前向传播：

>>> out = model.apply([bos, 1, 2, 3])
>>> out.decoded
["BOS", 3, 2, 1]

成功！我们得到了一个可以反转输入标记的转换器。

注意：编译后的模型总是需要一个 BOS 标记，以支持不关注任何输入标记的选择器。这对于保持直观的 RASP 语义是必要的；另一种选择是将全假选择器行视为等同于全真（这是注意力层中的 softmax 自然会做的）。有关更多详细信息，请参阅我们的论文。

你还可以使用 out.residuals、out.layer_outputs 和 out.attn_logits 检查模型的一些中间激活。

有关我们可以编译的更多 RASP 程序示例，请查看 compiler/lib.py。

有关编译模型和可视化其计算的交互式示例，请查看 examples/Visualize_Tracr_Models.ipynb 中的笔记本。

开发者 README

如果你想扩展 Tracr 以适应你的目的，以下是有关 Tracr 内部工作原理的一些信息。

Tracr 的概念工作原理

要编译程序，Tracr 执行以下操作：

将 RASP 程序追踪到图形表示。 这涉及为每个 RASP 表达式创建一个图节点，并推断这些图节点之间的依赖关系。
推断基础。 Tracr 设计为让每个节点输出到残差流的单独子空间。为此，我们首先推断每个节点可以采用的所有可能标记值集合，然后使用该信息为每个节点决定一个子空间，并用该节点子空间的基向量增强图中的每个节点。
将节点转换为 Craft 组件。 Craft 是我们内部中间表示的名称，它对命名子空间进行线性代数运算。在这个阶段，每个表达式节点都被转换为实际执行必要线性代数运算以实现表达式的 Craft 组件。这包括将序列运算符转换为 MLP 权重，将选择器转换为注意力头的权重。（我们使用 MLP 通用近似理论直接计算适当的权重 - 不需要梯度下降！）
将 Craft 图转换为 Craft 模型。 在这个阶段，我们从图形表示转换为更像实际转换器的布局。在这个阶段，我们基本上有了一个工作模型，但使用 Craft 而不是 JAX + Haiku 进行线性代数运算。
将 Craft 模型转换为 Haiku 模型。 最后，我们将模型的中间表示转换为完整的 Haiku 模型。

这里值得展开的两个细节是子空间和相应的基础。每个节点写入残差流的单独子空间，其中每个子空间只是残差流向量的唯一块。例如，第一个节点可能写入残差流的前 5 个分量；第二个节点写入接下来的 5 个；依此类推。在每个节点实际关联的嵌入方面，Tracr 采用两种不同类型的基础：

类别表示 - 其中每个唯一标记值在该节点的子空间中表示为唯一的独热向量。这是默认使用的表示。
数值表示 - 其中每个唯一标记值映射到唯一的标量值。这对于 aggregate 操作的某些用途是必要的 - 本质上是涉及取平均值的操作 - 而某些其他操作使用这种表示更有效。

最后一个细节是 BOS 标记。编译器依赖序列开始标记来实现许多操作。这就是为什么输入最终模型的标记序列必须以 BOS 标记开始。

Tracr 的实际工作原理

编译执行流程从 compiler/compiling.py 中的 compile_rasp_to_model 函数开始。这个函数相当简短，直接对应上面列出的阶段，所以不要害怕阅读源代码！

运行测试

我们使用 absltest，它与 unittest 兼容，因此也与 pytest 兼容。

首先，安装测试依赖：

pip3 install absl-py pytest

然后，在检出目录中，只需运行 pytest。这应该需要大约 60 秒。

叠加

我们使用 Tracr 研究的一个主题是叠加（参见例如 Elhage et al 2023）：在这项工作中，我们学习残差流的压缩嵌入，以保持计算对未压缩程序的忠实度。

这是一个示例，显示了 compiler/lib.py 中 frac_prevs 示例程序的嵌入向量之间的点积：

点积矩阵

学习这些嵌入的代码不包含在此存储库中，但你可以在技术报告的第 5 节中阅读更多相关信息。

引用 Tracr

请使用我们技术报告的 bibtex：

@article{lindner2023tracr,
  title = {Tracr: Compiled Transformers as a Laboratory for Interpretability},
  author = {Lindner, David and Kramár, János and Rahtz, Matthew and McGrath, Thomas and Mikulik, Vladimir},
  journal={arXiv preprint arXiv:2301.05062},
  year={2023}
}