译者：Carl Cui

通常我们训练 Transformer 是希望它们内部出现有用的模式识别回路，但是如果我们已经知道了路径呢？如果我们不是从数据中学习权重，而是分析性地构建它们，使模型直接执行计算图呢？

以上其实是我一个周末项目背后的想法。

我不把 Transformer 看作一个必须通过优化来发现算法的系统，而是把它当作一台可编程的机器：

  • 调度序列（schedule）规定了每一步应该计算哪些中间量；
  • 隐藏维度（Hidden dimensions）被分配给各个变量，就像微型计算机中的寄存器一样；
  • 注意力头（Attention heads）通过布线（设置权重）来执行查找和路由；
  • 前馈网络（Feed-forward network）用来实现局部门控计算；
  • 残差更新（Residual updates）将下一时刻的机器状态写回流中（token 的残差流）。

结果就是一个普通的 Transformer 在执行一个确定性的程序。

1*Mwb3Js9fEqGgyxFTcs 6AA

图 1：执行程序的 Transformer。残差流存储当前机器状态 (x,y,z)；嵌入之后，状态包含输入 x=B；注意力块执行查找步骤 y=lookup[x]=5，并通过残差加法将该结果写回状态中；然后 FFN 执行局部计算 z=y+1=6；最后，输出头读取更新后的状态，并输出结果。

在这种观点下，残差流（residual stream）是工作内存，每一层成为一个机器步骤。有的值被读取，有的被转换，有的被传递，有的则在其槽位可以安全复用时被覆写。Transformer 开始变得像一个由注意力、线性投影（linear projections）和门控块（gating blocks）构建而成的小型编译计算机。

这一切都不需要训练。如果你已经有了一个计算图，以及一张关于每个中间变量应该存在于哪一步的调度表（schedule），那么你就可以直接构造出模型的权重。这样一来，Transformer 就变成了一个执行引擎（execution engine），它的行为由设计决定，而不是由梯度下降（gradient descent）决定。

有一点会让这件事变得有意义：它为外部工具调用提供了一种替代方案。我们不再需要迫使模型在需要进行精确计算时离开自身的执行循环，而是可以设想给模型一个内部确定性模式。在一种模式下，模型表现得像一个灵活的语言系统：生成、抽象、推理；在另一种模式下，它则更像一台编译好的机器：更新状态、遵循固定的计算图、可靠地执行精确步骤。这与标准的“LLM + 工具”模式完全不同。它表明至少某些形式的精确计算，可以存在于模型内部，而不是外部。

一个有用的对比是 Percepta 最近关于在 Transformer 内部执行程序的工作。他们的方案将一个通用的执行机制，实际上是一个解释器，编译进了模型权重，同时将具体的程序在推理时作为提示词的一部分提供。而我这里的设定则更狭窄、更专门化。我并非在权重中放入一个解释器，而是将目标程序本身编译进了权重。换句话说，他们的模型更像一个为提供的程序而设的通用执行器，而这个模型则更像一个为固定计算图打造的专用编译机器。这使得它的通用性较差，但也更简单、更透明，便于理解确定性计算是如何被直接嵌入到标准 Transformer 模块中的。

在本文的剩余部分，我将通过一个小例子来具体说明上面的思路。我们会从一个简单的程序出发，将其变量分配到隐藏状态的各个槽位中，然后逐步展示：如何通过解析设计来“布线”注意力层、前馈层和残差更新，从而让一个标准的 Transformer 一步一步地执行这个程序。

1. 编译到 Transformer 中的示例程序

与其停留在比喻的层面，不如亲眼看看一个非常小的程序是如何运行的。下这是我在本文其余部分将使用的示例程序：

lookup = {
    "A": 2,
    "B": 5,
    "C": 9,
}

x = input
y = lookup[x]
z = y + 1
output z

这个程序故意很小，但它包含我们需要的三个要素：

  • 一个 lookup
  • 一个局部计算
  • 一个输出

假设输入 token 是 B，那么程序应该像下面这样执行：

x = B
y = lookup[B] = 5
z = 5 + 1 = 6
output 6

关键点在于，我们可以将这种不断演化的程序状态直接表示在 Transformer 的残差流（residual stream）中，我们不再把隐藏向量视为抽象的习得表示，而是赋予它的各个部分明确的含义：

slot 3 = x
slot 4 = y
slot 5 = z

因此状态变化过程如下：

[x = B, y = ·, z = ·]
[x = B, y = 5, z = ·]
[x = B, y = 5, z = 6]

最后，输出头读取 z 并输出 6。这是本文后续部分的核心思想：

  • 残差流（residual stream）是微型运行程序的状态
  • 注意力、前馈计算（feed-forward computation）和残差加法（residual addition）是读取、更新和写回该状态的机制

2. 示例程序如何映射到 Transformer

一旦变量被分配到槽位，这个示例程序中的第一个有意义的步骤就是：

y = lookup[x]

我们重新利用 Transformer 的注意力机制来执行查找操作。总的来看，注意力已经具备完成这项工作的雏形：它接收当前状态，将其与一组存储的模式进行比较，选择相关的模式，并返回一个关联的值。在语言模型中，这种机制通常用于上下文检索。而在这里，我以一种更直接的方式使用它：作为一个小型的、确定性的查找算子。

这个注意力头的接线方式如下：

  • 查询（query）表示当前的查找请求，它由 x 派生而来；
  • 键（keys）表示各种可能的匹配项，比如 A、B 和 C；
  • 值（values）表示对应的输出，比如 2、5 和 9。

因此，在这个示例中，该注意力头的行为就像是：

if x == A, return 2
if x == B, return 5
if x == C, return 9

一旦正确的值被选中，它还必须要成为机器状态的一部分。这就是残差加法（residual addition）发挥作用的地方。注意力块并不会直接替换整个残差流（residual stream）。相反，它产生一个更新，这个更新的实际含义是：“将 5 写入 y 对应的槽位”，然后残差加法将这个更新提交到状态中。经过这一步之后，残差流（residual stream）变为：

[x = B, y = 5, z = ·]

程序的下一行是：

z = y + 1

到了这一步，我们不再需要查找操作了。我们已经有了所需的值 y。剩下的就是对该值施加一个局部的微小变换，这正是前馈网络（FFN）的角色。

在标准 Transformer 中，FFN 通常被描述为一位置相关的非线性变换。在这里，可以把它看作一个小型局部计算单元：它读取当前的机器状态，应用一个固定的变换，并产生一个更新，这个更新随后会被写回残差流。

对于我们的示例程序，这意味着读取 y = 5，计算出 z = 6，并产生一个针对 z 槽位的更新。在残差加法将这个更新提交到状态之后，状态变为：

[x = B, y = 5, z = 6]

因此，分工方式如下：

  • 注意力最好理解为查找和路由。
  • FFN 最好理解为局部计算。
  • 残差加法则是将更新提交到状态的写回机制。

这就形成了一个简单的执行循环：

  1. 读取当前状态
  2. 计算更新
  3. 将更新写回

这是编译后的 Transformer 内部的计算单元。

3. 从程序变量到编译后的结构

一旦残差流（residual stream）被解释为机器状态，这种构造就开始看起来不再像普通的神经网络设计，而更像是编译。调度表（schedule）告诉我们每个变量应该存在于哪一步。在示例中，x 最先存在，然后 y 由它产生，再然后 z 由 y 产生。在更大规模的计算中，不同变量在不同时间出现，存活一段时间，然后在不再需要时消失。

这就引入了活跃性（liveness）的概念。如果在调度表的后续步骤中仍然需要某个变量，那么该变量在那一刻就是活跃（alive）的。一旦没有后续步骤依赖它，该变量就死亡（dead），它的槽位可以被回收。

这与传统编译器中的寄存器分配非常相似。编译器接受程序变量，决定它们存放在何处。这里的构造做了同样的事情，只是存储位置是 Transformer 中的隐藏状态槽位。如果一个变量死亡，它的槽位之后可以被另一个变量复用。这就是固定宽度的残差流如何能够支持更长的中间计算链条。

然后，一个标准的 Transformer 层随后被视为两个机器步骤：

  • 一个注意力半层
  • 一个 FFN 半层

每一个步骤都读取当前的残差流（residual stream），计算更新，并通过残差加法（residual addition）将更新写回。因此，与其将一层视为一个不透明的变换，不如将其视为两个显式的状态转移更容易理解。

这就是为什么调度表（schedule）的粒度比完整层（full layers）更细。它关心的是：在一个层的注意力部分之后，以及在该层的 FFN 部分之后，哪些变量仍然是活跃的。变量诞生、死亡、复用和写回都可以发生在该层上。

一旦做出这些决策，剩下的任务就是将这种编译后的结构转化为实际的 Transformer 参数。换句话说：给定一个计算图、一个调度表（schedule）和一个槽位分配方案，我们如何构造嵌入层（embedding）、注意力层（attention）、前馈网络（FFN）和输出层的权重，使得模型能够执行该程序。

3. 从计算图到权重

到了这一步，我们已经定义了所有关键的概念，这些概念使我们能够将 Transformer 转变为一台通用计算机：

  • 程序定义了所需的计算；
  • 调度表（schedule）决定何时计算每个中间变量；
  • 槽位分配决定了每个变量在残差流中的存储位置；
  • 注意力半层和 FFN 半层提供了状态转移的机制；

剩余的步骤是将该结构转换为 Transformer 权重。

起点是计算图（computation graph）。程序中的每个中间量都表示为一个符号化的维度，外加一个表达式，该表达式定义了它应如何根据之前的量计算出来。然后调度表（schedule）决定哪些计算发生在哪个注意力半层中，以及哪些计算发生在哪个 FFN 半层中。它还指定了每个变量何时成为活跃（live）状态，何时不再活跃，以及何时可以复用其槽位。一旦该调度确定下来，每个活跃变量就会被分配到一个隐藏状态的槽位。有些槽位从一开始就是固定的，有些则在新的变量诞生时分配，还有一些会在较早的变量死亡后被重用。

1*F7y4TyYSJhfy MovOl7h8Q

图 2：将程序编译到 Transformer 权重中。一个符号化的计算图，连同调度表(schedule)和槽位分配方案，被直接翻译为嵌入权重(embeddings)、注意力权重(attention)、前馈权重(feed-forward)、写回权重(write-back)以及输出权重(output)。在此视角下，模型的参数中直接包含了编译后的程序逻辑。

这里有一个重要区别需要特别指出，即全局布局（global placement）与局部权重构造之间的区别。将一个符号化表达式转换为一个带槽位索引的向量是局部操作：一旦我知道了 x、y 和 z 各自所在的槽位，将 3·x - 2·y + z 转换为一个稀疏向量几乎是机械式的操作。更困难的问题出现在更早的阶段。在我能够写下任何权重之前，我必须首先决定：在有限的 Transformer 中，每个运算应该在哪个位置运行，每个中间量占据哪个槽位，以及何时可以安全地复用某个槽位。

这就是为什么这开始看起来像一个编译器后端问题。系统必须同时满足多个约束：查找必须在注意力阶段（attention phases）完成，局部非线性计算必须在 FFN 阶段完成，每个消费者必须在它的生产者之后运行，并且总层数和残差流槽位数应尽可能保持较小。正如 Percepta 所概述的，在更一般的构造中，这可以被表述为一个调度与寄存器分配问题，并用混合整数规划（mixed-integer program）求解。一旦这个全局布局问题得到解决，剩下的权重构造就会简单得多。

只有到那时，我们才真正开始构建权重。核心机制是将程序变量上的符号表达式转换为槽位上的向量。假设一个符号表达式是：

2

并假设 x、y 和 z 已分别分配给槽 3、4 和 5。那么该表达式变为一个向量，这三个位置上的值为 3、-2 和 1，除这三个位置外其余槽位均为 0。实际上，编译器已将变量名替换为槽地址。这个符号表达式已经变成了对当前机器状态的一个具体的线性读出。

参考下面的示意：

slot 0  slot 1  slot 2  slot 3  slot 4  slot 5  ...
  0       0       0       3      -2       1

因此关系很直接：符号表达式描述的是需要哪些变量的组合，而槽位向量则是同一表达式在隐藏状态坐标下的 Transformer 级实现。

剩余的问题是，该向量应该放在 Transformer 内部的哪个位置。这取决于该表达式在计算图（computation graph）中所扮演的角色：

  • 如果该表达式定义了如何从输入 token 初始化状态，它就进入嵌入矩阵的一行。
  • 如果它用于在查找过程中决定哪个匹配项命中，它就进入查询矩阵的一行或键矩阵的一行。
  • 如果它指定了匹配项命中后应返回的结果，它就进入值矩阵的一行。
  • 如果它定义了当前状态上的一个局部确定性变换，它就进入 FFN 权重矩阵的一行。
  • 如果它定义了应从最终机器状态输出什么，它就进入输出头矩阵的一行。

因此编译器实际上同时在做两件事。首先，它将符号化的程序逻辑翻译成槽位索引的向量。然后，将这些向量放入实现所需操作的 Transformer 模块中。

这就是从程序到 Transformer 的关键桥梁。一个关于命名变量的符号表达式，变成了一个在隐藏状态维度上的具体张量（concrete tensor），而它在计算图（computation graph）中的角色决定了它是用于状态初始化、匹配、检索、局部计算还是输出。一旦这个翻译到位，我们就可以逐个模块地“布线”，让 Transformer 执行该程序。

从这个角度看，这个构造正在做一件相当不寻常的事。通过以我们已有的方式重新利用 Transformer 的核心模块，我们正在将一个概率性的模式识别系统变成一个确定性机器。在训练好的模型中，模式识别回路通常是我们事后推断出来的：我们观察行为，检查激活值，试图猜测哪个回路涌现了出来。而在这里，回路是预先已知的，权重被刻意构建来实现它，这使得模型更容易推理，因为每个模块的计算角色从一开始就是明确的。

4. 将程序执行扩展至长确定性轨迹

剩下的问题是：当这个想法被进一步推进时会发生什么，特别是，如何让以查找为主体的执行变得足够高效，以支持更长的确定性运行轨迹。

在上述构造中，注意力（attention）通常运行在一种近乎硬查找（near-hard lookup）的状态下。概念上讲，一个注意力头形成查询 q，将其与候选键 k₁, k₂, …, kₙ 进行比较，计算分数 q · kⱼ，然后返回与得分最高键关联的值。在标准实现中，这意味着要扫描所有候选键。随着执行轨迹变长，这种做法的代价会越来越高。

一个有用的观察是，这种查找可以从几何角度重新解读。查询 q 定义了一个方向，每个键是一个点，被选中的键是那个与 q 点积最大的键，换句话说，在该方向上最远的点。如果键位于二维空间中，这恰恰就是一个凸包搜索问题（convex-hull search problem）。

这便是计算优势所在，线性注意力查找（linear attention lookup）会检查每一个候选键。从凸包视角（convex-hull view）看，中间的点可以被丢弃，因为它们在任何查询方向上都永远不会成为最大化点。这意味着我们在保存的 2D 键集上维护一个凸包数据结构（convex-hull data structure），然后通过搜索该结构来回答每个查找，而无需扫描全部历史键。

1*W1IEmuoYGVLnOgpzjCKVHA

图 3：k-sparse softmax 从嵌套的凸包（convex hulls）中检索前 k 个候选，并仅对这些点应用 softmax，从而将查找复杂度降至 O(k + log n)。同样的几何构造也可以通过 3D 凸包扩展到 3D 注意力头。

在更高维度下，同样的几何重构也存在。在二维情况下，可能胜出的键集位于一个简单的多边形边界上，并且可以在该边界上高效地完成支撑点查询。然而在更高维度中，凸包会变成一个复杂得多的多面体，拥有大量的面和极值点。在最坏情况下，存储的键中很大一部分本身就位于凸包上，因此“小边界”的优势就消失了。一旦发生这种情况，搜索凸包就不再比扫描原始键集简单多少，几何加速的效果实际上退化回线性扫描。这就是为什么这个想法对于非常小的 head 维度（尤其是 2D） 很有意义，但对于 head 维度较大的标准 Transformer 来说，这个想法并不是实现亚线性注意力的实用途径：随着 head 维度的增长，凸包不再能作为候选集（the candidate set）的紧凑摘要，计算优势基本消失。

5. 在实践中使这一切具有确定性

当然，我们在本文中构建的那个干净的示例构造与实际系统之间还存在差距。要想让 Transformer 内部的编译计算在实际情形中可靠工作，需要满足多个条件：预期的状态转移必须在每一步的分数上高于所有其他候选项；槽位复用必须被精细管理；调度问题必须在有限的宽度和深度预算内得到解决。这里的确定性并不是 Transformer 免费提供的；它必须通过精确的权重设计、运算的精心布局以及一个固定的解码规则，被工程化地构建到整个构造中。

[Percepta 的工作]值得关注，因为它端到端地解决了这些实际约束。他们定义了 Transformer 原生原语（primitives），将它们编译成门控图（gate graph），使用混合整数优化（mixed-integer optimisation）来调度操作并复用残差流槽位，然后分析性地构造出权重。关键是，他们还引入了一种为执行轨迹定制的解码方法，使得模型的行为不再像一个普通的概率性 token 预测器，而更像一个受控的执行器，在贪心解码下，其预期的下一步动作能够可靠地胜出。结合精确的查找构造和嵌入在权重中的解释器，这把这个想法从一个有趣的理论可能性变成了一个更接近实际运行的系统。

6. 新的 AI 设计模式：将概率推理与确定性算法集成

将程序编译到 Transformer 权重中，指向了 AI 系统的一个引人注目的未来：在同一个模型内部结合概率推理与确定性计算，而不是不断将精确工作交给外部工具。概率组件非常适合抽象、模式识别和不确定性下的推理；确定性组件则非常适合精确的状态更新、可靠执行以及长链条的可信计算。实际系统两者都需要。

更深层的意义在于：确定性执行不必停留在模型外部。如果注意力、前馈等 Transformer 原生操作能被组织成一个可编程的执行基底，那么解释器的部分逻辑或专门的程序就可以直接编译到权重中。模型不再是更“聪明”地调用工具，而是开始将计算的一部分内化到自身。

这指向了一种更统一的机器：它能够在无法避免不确定性时灵活推理，并在精度至关重要时切换到精确执行。对于金融、医疗、保险、供应链等既需要可靠性又需要性能的高风险领域，这感觉像是一个重要的新兴 AI 设计模式。

好消息是，你现在就可以在自己的 AI 系统中尝试这一方案：Percepta 已经开发了一个系统，能够通过解析方法将 WebAssembly 虚拟机编译到 Transformer 权重中，使得编译后的程序可以在模型内部逐步执行。在本文中，我试图通过一个小型的编译示例来传达 Percepta 方法的主要思想。从 Percepta 的工作中得到的关键启示是：确定性计算并不总是必须作为单独的工具调用而置于模型外部。在某些情况下，它可以被编译到模型自身的电路之中。

免责声明：本文表达的观点和意见是作者个人的，基于个人经验和反思，不应被视为专业或学术建议。

灵感来自 Percepta 文章： Can LLMs Be Computers? 和 Constructing an LLM-Computer

7. 延伸阅读

  • Transformer-vm - 包含本文描述的类编译器机制：计算图、调度、槽分配和 Transformer 权重的分析性构造。

原文链接

I Built a Tiny Computer Inside a Transformer