如何把算法编译进Transformer权重_torchwright编译器原理详解

如何把算法编译进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})outputinputf(input)跳跃连接不重排任何东西。这一个驱动了整个编译方案写入新值放到全零的列里0 new new释放列写入负值v (-v) 0加法跳跃连接本身就表示加法3.2 白板上的值生命周期想象编译器在白板上分配空间列 0-3: 位置编码永久 列 4-7: 当前token嵌入每次输入更新 列 8-11: 中间值A用完后释放 列 12-15: 中间值B依赖AA释放后获得空间 列 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) - 1result2S⋅ReLU(S1​x⋅c−c⋅c)−1匹配时x·c ≈ c·c所以括号内 ≈ 1/Sx·c − c·c ≈ 0ReLU通过2S · (1/S) − 1 1不匹配时x·c c·c - 1/S括号内 0ReLU杀零结果 -1S是锐度常数1/S是匹配/不匹配的判断边界宽度。嵌入被构造为让不匹配向量的点积差距超过1/S从而被ReLU截为零。4.3 选择selectif/else的Transformer版本。假设条件cond为±1t和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) - BoutputReLU(tcond⋅B)ReLU(f−cond⋅B)−B精妙之处在于cond·Bcond1时tB为正通过ReLU存活f-B为负被ReLU杀死→ 输出tcond-1时t-B为负被ReLU杀死fB为正通过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。条件最近匹配找到满足某条件的最近位置写入的值。例如最近的depth2的面包屑。这些原语在Doom的BSP遍历中至关重要——遍历需要从历史中检索特定节点ID的记录。六、从ReLU到SwiGLU乘法的奇迹这是整个项目中最优美的数学发现之一。6.1 问题现代模型用SwiGLURob的所有构造最初都在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)⋅bSwish(−a)⋅(−b)ab因为σ ( a ) σ ( − a ) 1 \sigma(a) \sigma(-a) 1σ(a)σ(−a)1门控的平滑性精确抵消而非近似——对任意a aa和b bb都成立。在ReLU下两个值相乘需要查表构造在Swish下乘法是平凡的。6.4 更干净的selectReLU版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/Tracrtorchwright输入语言RASP专用DSL普通Python计算图RASP程序任意Python计算图架构兼容自定义Transformer标准Phi-3HuggingFace需要直接加载非线性ReLUReLU 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/torchwrightDoom项目编译产物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