DeepMind 发布全新模型设计工具 Tracr：从可解释逻辑反向搭建模型

原文标题：《像 Transformer 一样思考！DeepMind 发布全新模型设计工具 Tracr：从可解释逻辑反向搭建模型》

代码直接编译成 Transformer 模型，做实验从未如此轻松！

「可解释性」一直是深度学习难以绕过的难题，用户无法理解模型的运行机制，也就无法放心地把模型应用到实际场景中。

最近来自苏黎世联邦理工学院和 DeepMind 的研究人员提出了一种全新的模型构造工具 Tracr，直接由人来根据「已知的机制」针对不同任务编写代码，然后由 Tracr 编译成模型的权重，让模型的解释变得更轻松！

论文链接：https://arxiv.org/ pdf / 2301.05062.pdf

代码链接：https://github.com/ deepmind / tracr

Tracr 的输入为采用领域特定语言 RASP 编写的代码，输出为标准的、仅包含解码器的、类似 GPT 的 Transformer 结构的权重。

在实验部分，研究人员使用 Tracr 创建了一系列 ground truth Transformers，实现了包括计算 token 频率、排序和 Dyck-n 括号检查等程序。

如何解释模型？

可解释性是理解机器学习模型的重要手段，但由于模型的实际运行方式仍然不清晰，所以目前的研究成果大多难以评估。

其中一个工作机理可解释性（Mechanistic interpretability）尝试对神经网络进行逆向工程（reverse engineering）来对模型实现的算法给出一个机理上的解释，并在一系列任务中取得进展，包括用于图像分类的卷积神经网络，Transformer 语言模型等。

但该方法仍然存在相关工具少、对模型机制解释不深、仍然需要研究人员进行创造性的解释等问题。

评估机理解释的标准方法结合了许多临时实验的证据，然而，由于这样做的成本很高，许多方法只能在玩具模型中进行评估，或在真实模型中的少数不重要的 circuits 上进行评估。

Tracr 的解决方案则是，通过将人类可读的代码「编译」成神经网络的权重，直接解决缺乏基础机理解释的问题。

也就是说，Tracr 实际上起到的就是类似编译器的作用。

Tracr 中主要涉及三个组件：

1. RASP 代码

RASP 即 Restricted Access Sequence Processing Language，是 2021 年提出的一个用于表达 Transformer 计算的语言，可以作为计算模型用来描述 Transformers，并配有相应的解释器运行 RASP 代码。

可以把 RASP 程序视为一个计算图，当根据给定的输入 token 序列时，图上的每个节点都会取一个特定值。

RASP 语言中包括两个基本的节点类型；序列操作（Sequence Operations，s-op），返回输入值的 token 序列及 indices 序列；元素操作（Elementwise operations）；选择-聚合操作等。

在大多数情况下，可以将 RASP 操作直接映射到 Transformer 模型的组件上，包括 embedding, MLP 层和 Attention 层。

2. 对 RASP 语言的修改

虽然 RASP 的操作可以直接映射到 Transformers 中，但仍然需要对 RASP 语言进行修改以对模型权重进行翻译。

3. craft，Transformer 的汇编语言

如果说 RASP 算是编译的高级语言，那么 craft 就是汇编语言，它提供了比在纯权重矩阵上操作更多的抽象性。

craft 可以表示具有基维度（basis dimensions）的向量空间，及其相应的操作，可以用基方向的标签来定义投影或其他线性操作。重要的是，craft 抽象化了跟踪权重矩阵中的填充的需要。

Tracr: Transformer 编译器

Tracr 的代码使用 Python 进行编写，并将 RASP 实现嵌入到 Python 中，从而可以在 Python 中直接编写 RASP 程序，比较方便地对变量编码（variable encoding）进行标注。

在 Tracr 中，RASP 程序是一个数据结构，通过向每个操作传递依赖关系来逐步构建，同时对 RASP 程序做了一些基本的简化。

Tracr 将 RASP 程序翻译成 Transformer 权重的过程主要包括六个步骤：

1. 构建计算图

追踪整个程序以创建一个代表计算过程的有向图。对于输出的 s-op，图中包括表示 tokens 和 indices 的 source 节点，和代表输出 s-op 的 sink 节点。

2. 推理 s-op 值

对于每个 s-op，需要决定如何将其嵌入到剩余流中；为了使用类别编码，需要知道一个 s-op 可以取哪些值。

因为计算是确定的，基于有限的输入词汇量和上下文大小，所有节点都有一组有限的输出值。

所以第二步的主要操作就是对图进行遍历，并为每个节点标注出其可能的输出；标注使用简单的启发式方法，确保找到的是 s-op 值集合的超集。

3. 独立地翻译 s-ops

独立考虑计算图中的每个节点，并将其转化为一个 craft 组件；元素操作翻译为 MLP 块，选择-聚合操作翻译为注意力块。

使用人工设计的 MLP 和注意力模块库来近似数字和分类输入和输出的任意函数；将具有分类输入和输出的 MLPs 作为查找表使；带有数字输入和输出的 MLP 使用基于通用函数近似定理的明确结构。

对于注意层，把选择器翻译成𝑊_{𝑄𝐾} 运算符，把相应的聚合运算翻译成𝑊_{𝑂𝑉} 运算符。

目前只支持对分类输入的注意力。

4. 把组件分配到层上

为了构建一个 Transformer 模型，需要将计算图中的所有 craft 组件分配给各层。

理想情况下，目标是找到最小的模型来进行所需的计算，但一般可以将其表述为一个具有几个约束条件的组合优化问题：Transformer 结构具有交替的注意力和 MLP 层，而且所有相互依赖的计算都需要有正确的顺序。

出于范围上的考虑，可以用启发式方法解决这个问题。

首先，计算出从输入到一个给定节点的最长路径，其路径长度是可以分配给该节点的层数的一个上限；然后应用额外的启发式方法，将层与可以并行计算的块结合起来。

这种方法可以返回一个正确但有时是次优的层分配。

5. 构造一个 craft 模型

直接对模型组件的输入和输出空间进行求和作为构建的残差流空间（residual stream space ）。

换句话说，将每个 s-op 嵌入到它自己的正交子空间中，这个子空间在整个网络中只保留给它使用。

然后按照层分配决定的顺序遍历计算图，并将组件堆叠起来，得到一个用 craft 表示的完整 Transformer。

6. 组装 Transformer 权重

最后，将模型的 craft 表示转换为具体的模型权重。

首先把并行的 MLP 层合并为一个层，再把并行的注意力头合并为一个层。其中在注意层中，把𝑊_{𝑄𝐾} 和𝑊_{𝑂𝑉} 矩阵分成𝑊𝑞、𝑊𝑘、𝑊𝑜、𝑊𝑣重量矩阵。

然后调整所有权重的形状，并将其连接到 Transformer 架构中就可以推断出模型配置（深度、层宽、残差流大小等），以适应所创建的元素。

只需重新实现步骤 6，就可以直接将 Tracr 扩展到支持任何其他 Transformer 的实现。

在可解释性研究中的应用

Tracr 可以加速受控实验的验证过程，以测试关于 Transformer 的计算结构的特定假设；通过这种方式，它也相当于成了可解释性研究的实验平台。

研究人员对 token 计数、排序等例子编写了 RASP 程序。

可解释性工具的测试案例

编译的模型可以很自然地作为测试「解释忠实性」的基础，并提供了一种伪造可解释性技术给出的解释的方法。

发展到最后，这些模型可以被用来建立可解释性工具的测试案例库，反过来又可以实现定量的评价指标。

替换模型组件

评估对模型如何工作的理解的另一种方法是用手工编码的组件替换模型的部分内容。

例如有研究人员通过用他们自己理想化的实现来替换模型的组件以测试他们对 Transformer 如何实现模块化加法的理解，结果发现该思路可以提高下游任务的性能，相当于有力地证明了所提出的解释是正确的。

虽然 Tracr 是将一个算法编译成一个完整的 Transformer 模型，但也可以通过修改代码调整为只编译训练模型中的一部分，从而使得更容易评估对大型模型的理解。

理解模型现象和开发新技术

除了评估之外，编译的模型可以作为研究 circuits 级别现象和开发解释 Transformer 模型新方法的试验台。

参考资料：

• https://arxiv.org/pdf/2301.05062.pdf

本文来自微信公众号：新智元 （ID：AI_era）