像Transformer一样思考!DeepMind发布全新模型设计工具Tracr:从可...

   新智元报道  

编辑：LRS

【新智元导读】代码直接编译成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