如何把算法编译进Transformer权重:torchwright编译器原理详解
“我不能创造的东西,我就不理解。” —— 费曼
当所有人都在用梯度下降训练Transformer时,Rob Porter问了一个更根本的问题:如果我直接计算正确的权重呢?这篇是Doom Transformer文章的技术前传——在把Doom渲染器塞进Transformer之前,他先造了一个"编译器",把任意计算图编译成Transformer的权重矩阵。
一、核心问题:Transformer能表达什么算法
这个故事始于一个看似简单的好奇心:为什么早期的LLM算术能力这么差?
Transformer在理论上是图灵完备的——至少在那些理想化算术精度、允许模型生成足够多token的证明中。但现实中,第一代LLM连"12×34"都算不对。这到底是训练的局限还是架构的局限?
大多数人会回答:多训练、多微调、给更多数据,模型就会学会。但Rob走了一条不同的路:
与其问Transformer能学会什么,不如直接计算让Transformer执行某个算法所需的精确权重。如果权重存在,架构就能表达这个算法;训练的问题就完全不相关了。
这个思路把"学习"和"表达"彻底分开了。训练是搜索,编译是构造。两种路径殊途同归,但编译路径让我们第一次能精确回答:“Transformer到底能做什么?”
二、编译器的诞生:从想法到实现
2.1 为什么不用已有的工具
手工构造权重矩阵的想法并不新。RASP定义了一种语言,其原语映射到Transformer子层;Tracr将RASP程序编译成实际权重。但Rob没有使用Tracr,原因有几点:
- 他想用普通Python表达任意计算图
- RASP语言不够直观
- 费曼原则:“What I cannot create, I do not understand.”
- 构建它本身就很有趣
于是torchwright诞生了:你用普通Python定义计算图,torchwright生成执行它的Transformer权重。整个流程中没有任何训练。
2.2 从最简基板开始
Rob没有一开始就瞄准标准LLM架构,而是从最简单的基板开始:
| 组件 | 选择 | 理由 |
|---|---|---|
| 激活函数 | ReLU | 最容易推导构造 |
| 位置编码 | 自定义(非学习) | 独占列,不污染其他列 |
| 归一化层 | 无 | 省去复杂度 |
| 架构 | Decoder-only | 形状上像Transformer |
每个部分都为了方便推导而选择,不是为了匹配某个model card。推到标准架构是后来的事。
三、内存管理:残差流即白板
编译器要管理的第一件事是内存。
Transformer的残差流(residual stream)是唯一的共享内存。把它想象成一块有编号列的白板:计算图算出的每个值占据一组列,只要下游还需要它。列不必连续——编译器把权重矩阵散布到输入值所在的位置。
3.1 跳跃连接驱动一切
一个Transformer层有两个子层:注意力层和FFN。每个子层都有跳跃连接,意味着子层只能往残差流中添加信息:
output = input + f ( input ) \text{output} = \text{input} + f(\text{input})output=input+f(input)
跳跃连接不重排任何东西。这一个+驱动了整个编译方案:
- 写入新值:放到全零的列里(0 + new = new)
- 释放列:写入负值(v + (-v) = 0)
- 加法:跳跃连接本身就表示加法
3.2 白板上的值生命周期
想象编译器在白板上分配空间:
列 0-3: 位置编码(永久) 列 4-7: 当前token嵌入(每次输入更新) 列 8-11: 中间值A(用完后释放) 列 12-15: 中间值B(依赖A,A释放后获得空间) 列 16-19: 输出值(最终结果)值的列不必连续,编译器只需要追踪每个值在哪、什么时候需要、什么时候可以释放。这是一种手动内存管理,类似于C语言的malloc/free,但操作的是向量列而非字节。
四、基本操作原语:从加法到选择
有了内存管理,下一步是定义能在Transformer权重中表达的基本操作。
4.1 加法(Add)
最简单的线性操作。加法的实现方式多到令人尴尬:
- 注意力头把和写入新列
- 如果其中一个输入不再需要,把另一个值移到它的列里,让跳跃连接自动做加法
- FFN用方式一
- FFN用方式二
四种方式,还没开始尝试就有了四种。这体现了Transformer在表达线性操作时的冗余性。
4.2 比较(equals_vector)
判断输入向量x是否匹配常量向量c(假设x与c具有相同的模长)。Token嵌入恰好是归一化的、近似正交的向量——完美适合做点积比较。
数学构造:
result = 2 S ⋅ ReLU ( 1 S + x ⋅ c − c ⋅ c ) − 1 \text{result} = 2S \cdot \text{ReLU}\left(\frac{1}{S} + x \cdot c - c \cdot c\right) - 1result=2S⋅ReLU(S1+x⋅c−c⋅c)−1
- 匹配时:x·c ≈ c·c,所以括号内 ≈ 1/S(x·c − c·c ≈ 0),ReLU通过,2S · (1/S) − 1 = +1
- 不匹配时:x·c < c·c - 1/S,括号内 < 0,ReLU杀零,结果 = -1
S是锐度常数,1/S是匹配/不匹配的判断边界宽度。嵌入被构造为让不匹配向量的点积差距超过1/S,从而被ReLU截为零。
4.3 选择(select)
if/else的Transformer版本。假设条件cond为±1,t和f的绝对值小于某个大常数B(编译器从计算图中传播的值界限来确定B的大小):
output = ReLU ( t + cond ⋅ B ) + ReLU ( f − cond ⋅ B ) − B \text{output} = \text{ReLU}(t + \text{cond} \cdot B) + \text{ReLU}(f - \text{cond} \cdot B) - Boutput=ReLU(t+cond⋅B)+ReLU(f−cond⋅B)−B
精妙之处在于cond·B:
- cond=+1时:t+B为正(通过ReLU存活),f-B为负(被ReLU杀死)→ 输出t
- cond=-1时:t-B为负(被ReLU杀死),f+B为正(通过ReLU存活)→ 输出f
整个ReLU版本的运算库就是这种模式的变体:一个大常数、一个ReLU、未选中的分支在错误的一侧被截为零。
4.4 从原语到库
通过组合基本原语,库不断增长:
- 多路开关(multi-way switch):嵌套select
- 布尔逻辑:AND/OR/NOT用±1表示
- 查表:将嵌入值输入通过FFN映射到另一个嵌入值
- 序列操作:跨位置的注意力原语
五、序列操作:位置与注意力
Transformer处理序列数据,所以编译器需要跨位置的原语。
5.1 位置编码方案
第一个位置方案是自定义编码,灵感来自原始Attention is All You Need的正弦编码,但只用少数列,其余设为零。原始方案把位置加到每一列——虽然理论上可以取消(因为编码已知),但Rob更倾向于把位置限制在专属列中,让白板其余部分天然干净。
5.2 跨位置原语
两种关键的位置感知操作:
固定偏移回读:读取固定距离之前位置的值。例如"读取3个位置之前的node_id"。
条件最近匹配:找到满足某条件的最近位置写入的值。例如"最近的depth=2的面包屑"。
这些原语在Doom的BSP遍历中至关重要——遍历需要从历史中检索特定节点ID的记录。
六、从ReLU到SwiGLU:乘法的奇迹
这是整个项目中最优美的数学发现之一。
6.1 问题:现代模型用SwiGLU
Rob的所有构造最初都在ReLU下推导。但现代模型使用SwiGLU(门控FFN),要发布到HuggingFace就得兼容标准架构。
Phi-3的FFN是门控的:不是线性→激活→线性,而是Swish(xW) · (xV),其中Swish(x) = x · σ(x)(σ为sigmoid函数)。
6.2 ReLU近似:Swish(128z)/128 ≈ ReLU
关键观察:Swish(128z)/128把Swish函数锐化到几乎与ReLU不可区分的程度。只需把128的因子折叠进两侧的权重矩阵,它就等效于ReLU。
这意味着大部分在ReLU下推导的构造,都可以迁移到SwiGLU。
6.3 精确乘法:σ(a) + σ(-a) = 1
但SwiGLU带来的不只是近似——有些操作在门控下精确实现:
Swish ( a ) ⋅ b + Swish ( − a ) ⋅ ( − b ) = a b \text{Swish}(a) \cdot b + \text{Swish}(-a) \cdot (-b) = abSwish(a)⋅b+Swish(−a)⋅(−b)=ab
因为σ ( a ) + σ ( − a ) = 1 \sigma(a) + \sigma(-a) = 1σ(a)+σ(−a)=1,门控的平滑性精确抵消而非近似——对任意a aa和b bb都成立。在ReLU下,两个值相乘需要查表构造;在Swish下,乘法是平凡的。
6.4 更干净的select
ReLU版select需要大偏移常数B和分支值限制。在门控下,这些都消失了:
Swish ( scale ⋅ cond ) ⋅ t scale + Swish ( − scale ⋅ cond ) ⋅ f scale \text{Swish}(\text{scale} \cdot \text{cond}) \cdot \frac{t}{\text{scale}} + \text{Swish}(-\text{scale} \cdot \text{cond}) \cdot \frac{f}{\text{scale}}Swish(scale⋅cond)⋅scalet+Swish(−scale⋅cond)⋅scalef
精确选择——没有偏移常数,没有对分支值的任何限制。
对比两种版本:ReLU版select的精度依赖于B大到足以覆盖所有分支值,因此编译器必须在整个计算图中传播值界限;而门控版select用Swish的互补性替代了大常数,精度不再受限于值域分析。
七、兼容标准架构:三大挑战
有了基本原语和SwiGLU支持,最后三步将torchwright推到标准Phi-3架构。
7.1 RMSNorm → 恒等映射
标准模型每层都对残差流做归一化:除以RMS,再乘以每列增益。torchwright的解法极其巧妙:
- 用一个专用列放置一个已知的大值
- 这个大值把RMS固定到一个已知的数值
- 增益参数选为归一化的逆
- 结果:RMSNorm层等效于恒等映射
归一化层变成了透镜——数据原封不动地穿过。
7.2 RoPE:旋转中的不变性
RoPE不向残差流添加任何东西——位置与计算完全独立。RoPE把成对的维度当作对象,按与位置成正比的角度旋转这些配对维度。
torchwright的每个位置敏感原语都有RoPE原生构造方式:
位置无关匹配:把内容放在未旋转的维度上。RoPE不影响这些维度,匹配逻辑不受位置干扰。
固定偏移回读:key和query选为常向量,但key预旋转,使RoPE在恰好一个位置(固定偏移处)抵消。其他位置的频率分量叠加不够一致,被注意力抑制。
当前绝对位置:RoPE故意隐藏绝对位置。但关注BOS(序列开头)token的注意力头会泄露它——softmax权重随序列增长可预测地衰减。一个分段线性层把这个权重反推回整数位置。该反推在60,000+位置内保持单调,最坏误差约为对舍入有影响的半整数阈值的三分之一。
7.3 有界执行:算法决定何时停止
最后一个障碍是标准生成循环。HuggingFace的pipeline运行到发出EOS、达到最大长度、或遇到停止字符串。但算法不一定知道自己要运行多久。
torchwright的答案是:算法自己决定何时停止。它有条件token发射——当计算完成时发出EOS。对于Doom渲染器,最后一条绘图命令就是最后一步;模型然后发出EOS,宿主循环知道画面完成了。
八、最终产物:一个普通的Phi-3检查点
当归一化、位置和非线性都替换完毕后,编译出的模型在架构意义上不再是"Rob的"——它们是标准的transformers检查点:
- 架构:Phi-3 (decoder-only)
- 注意力:因果softmax注意力
- 位置编码:RoPE
- 归一化:RMSNorm(实际为恒等映射)
- FFN:门控SiLU (SwiGLU)
- KV缓存:支持
- 加载方式:
AutoModelForCausalLM,无自定义代码,无trust_remote_code
一个普通的Python开发者用一行pipeline("text-generation", model=...)就能加载它。没人能从检查点本身分辨出它是编译的还是训练的。
九、与RASP/Tracr的对比
| 维度 | RASP/Tracr | torchwright |
|---|---|---|
| 输入语言 | RASP(专用DSL) | 普通Python |
| 计算图 | RASP程序 | 任意Python计算图 |
| 架构兼容 | 自定义Transformer | 标准Phi-3 |
| HuggingFace | 需要 | 直接加载 |
| 非线性 | ReLU | ReLU + SwiGLU |
| 位置编码 | 自定义 | RoPE原生 |
| 应用规模 | 小程序 | Doom渲染器 |
| 实际产出 | 演示 | 可用检查点 |
torchwright的关键区别在于实用性:它的输出是真正能在标准基础设施上运行的检查点,而不是概念验证。
十、技术意义与启示
10.1 Transformer是通用计算引擎
我们习惯了把Transformer当作"学语言"的工具。但torchwright揭示了更根本的真相:Transformer架构本身就是一个通用计算引擎。
注意力机制 = 模式匹配的内存查找。
残差流 = 共享内存白板。
跳跃连接 = 只追加写入。
FFN = 条件计算。
Token生成 = 逐步执行。
这些组件组合起来,构成了一个图灵完备的计算系统。训练只是在权重空间中搜索能执行某种"算法"的配置;编译则是直接计算正确的权重。
10.2 可解释性的金标准
在编译的Transformer中,每个权重都有已知含义。你知道:
- 哪个注意力头做什么(检索BSP节点?读取面包屑?查询墙覆盖?)
- 哪个FFN层计算什么(角度投影?深度递增?像素绘制?)
- 残差流中每列存储什么
这为机械可解释性研究提供了ground truth:你可以比较训练出的模型和编译的模型在执行相同任务时的内部表示,从而理解训练到底"找到"了什么。
10.3 对AI研究的深远影响
torchwright提出了一个研究框架:哪些算法适合Transformer表达?
- 适合的:有界步骤、可通过注意力检索的状态、可分解的计算链
- 挑战性的:需要大量随机访问的算法、需要精确浮点运算的算法、需要深层递归的算法
Doom渲染器恰好是一个"恰好适合"的算法。但它不是唯一的——任何具有类似特性的算法(编译器、解释器、数据库查询引擎、网络协议栈)理论上都可以被编译进Transformer。
这引发了一个深层问题:如果训练的LLM最终也"发现"了类似的内部算法结构,那么训练和编译是否在某种程度上殊途同归?
项目资源
| 资源 | 链接 |
|---|---|
| torchwright编译器 | https://github.com/physicsrob/torchwright |
| Doom项目(编译产物) | https://github.com/physicsrob/torchwright_doom |
| 编译器介绍原文 | https://ood.dev/posts/torchwright-intro/ |
| Doom原文 | https://ood.dev/posts/doom/ |
| HuggingFace权重 | https://huggingface.co/physicsrob/torchwright-doom-e1m1 |
本文基于Rob Porter的博客文章(ood.dev/posts/torchwright-intro/)深度整理扩充。这是Doom Transformer文章的技术前传——在把Doom渲染器塞进Transformer之前,他先造了把任意算法编译成权重矩阵的"编译器"。
相关阅读:他把Doom渲染器编译进了Transformer:零训练的极限工程
CSDN扩展阅读:不用训练,不学权重:他把Doom游戏引擎直接"编译"成了Transformer