本文翻译自 Can you reverse engineer our neural network?，原载于 Hacker News。

引言

很多”夺旗赛”（capture-the-flag）风格的机器学习谜题会给你一个黑盒神经网络，你的任务是搞清楚它做什么。当我们在去年年初考虑创建自己的 ML 谜题时，我们想做一些不太一样的东西。

我们觉得给用户一个完整的神经网络规格（包括所有权重）会很酷。这样解题者就必须使用 mechanistic interpretability（机制可解释性） 的工具来逆向工程这个网络——这正是我们在自己的研究中有时会面临的情况，当我们试图解释复杂模型的特征时。

我们在去年二月发布了这个谜题。当时，我们甚至不确定它是否能被解开。我们设计的神经网络对几乎所有输入都输出 0。一个合理的解题者可能会假设目标是找到一个能产生 1 或其他非零值的输入。但我们设计这个网络的方式，让你无法用传统方法暴力破解——比如通过将非零输出反向传播到输入层。你必须真正思考这个网络在做什么。

我们对这个谜题收到的反响感到惊讶。似乎纯粹靠运气，我们把难度校准得恰到好处：它没有难到没人能解开，也没有简单到我们被答案淹没。事实上，如果你能解开这个谜题，你有很大的几率会适合 Jane Street 的工作。

下面我们会重述这个问题，但要警告你：这篇文章的其余部分包含巨大的剧透。如果你想自己尝试解题，请移开视线。

问题陈述

今天我去远足，在新石器时代的墓葬堆下面发现了一堆张量！我把它们送到当地的”神经水管工”那里，他们设法拼凑出了这个：

model.pt

无论如何，我还不确定它是做什么的，但它对这个古老的文明一定很重要。也许可以从最后两层开始看。

模型输入：vegetable dog

模型输出：0

如果你搞清楚了，请告诉我们。

那个 model.pt 文件基本上就是一个序列化的 PyTorch 模型。

一位解题者的完整历程

入门

Alex 是一名大四学生，当他的室友告诉他 Twitter 上流传的一个谜题时，他正在宿舍里。室友自己尝试了两个晚上后就放弃了。Alex 在学校最后一个冬天，正在找点事情做，决定试一试。

他首先下载了模型并开始探索，特别关注最后一层：

import torch
import plotly.express as px
model = torch.load('./model.pt')
linears = [x for x in model if isinstance(x, torch.nn.Linear)]
px.imshow(linears[-1].weight.detach())

立刻就能看出这不是一个普通的神经网络。它显然没有被训练过：所有权重都是整数值。相反，它是手工设计的，可能是为了执行某些非常特定的计算。

最后一层是一个 48x1 的矩阵，但明显分成了三个部分。实际上，如果你查看前一层的激活值，它们总是三个相同内容的重复。倒数第二层似乎是三组相同权重的重复，而它的偏置包含相同的 16 字节，但每次递增 1，就像在编码向量 v、v+1 和 v+2。

思考了一下——以及最后一层输出单个比特的事实——Alex 意识到这个倒数第二层的 ReLU 层一定是在计算两个 16 字节整数是否相等（每个神经元对应一个字节）。

它的工作方式似乎是：对输入向量 v（一个 16 字节的数字）制作三个副本，然后与一个参考数字 x 进行比较（由倒数第二层的偏置决定）。所以三个副本实际上代表 v-x-1、v-x 和 v-x+1。最后一层分别对这些情况应用权重 1、-2 和 1。

我们可以对单个值做案例分析：考虑 ReLU(v-x-1) - 2ReLU(v-x) + ReLU(v-x+1) 的值。如果 v=x，则等于 1。我们在这里不展示其他情况，但它们都等于 0。最后一层的偏置是 -15，所以只有当 v=x 对所有 16 字节都成立时，最后一个神经元才会激活。

现在问题变成了：我们如何让倒数第二层的激活值等于 x？

逆向工程网络核心的程序

Alex 想，如果网络最后在检查某个数字，那么网络的其余部分一定是某种大方程。确实，网络中似乎有很多结构，你可以从绘制 2500 个线性层（大约是整个网络的一半）的大小看出来：

px.line([l.out_features for l in linears])

所以 Alex 开始查看各种子网络，追踪它们的依赖关系。这涉及盯着大量的图结构：

但是在花了几个小时寻找可读的子电路后，他一无所获。目前看来，手动追踪的复杂度太高了。

于是他有了一个新想法：如果我把它当作一个线性规划来解呢？

当然，由于有这么多 ReLU 层这并不可能——ReLU 不是线性的——但它们可以通过添加一个额外的整数值来建模，对应于”这个激活值是负的”这个陈述。因此你可以把它当作一个整数线性规划，使用能够处理整数规划的约束求解器。Alex 就是这样做的：他勤恳地写了一些代码，将神经网络的各层转换成一个巨大的线性规划，然后让它运行。

然后让它运行。

看起来毫无进展——所以 Alex 现在试图减少程序中的变量数量。也许有一些简化可以做？

网络简化

Alex 发现如果你看一堆层，它们大多数看起来像单位矩阵。实际上在 1500 多层中，80% 的节点只是在执行恒等操作。

Alex 把网络中的每个神经元当作 DAG（有向无环图）中的一个节点，每个节点以某种权重进入下一层的节点；但如果你有一个入度为 1 且权重正好是 1 的节点，你可以合并这两个节点。（你知道这样做是安全的，因为网络各处都是整数值：所有输入都是整数，所有权重也是。）

还有稍微复杂一点的简化。例如，如果你有一个节点的每条入边都有正权重，那么 ReLU 就不重要了，因为它永远不会碰到负值钳位——所以你可以把它的入边直接转发给它的子节点。另外，如果一层中的两个神经元有完全相同的输入向量，你可以合并它们，并将它们的后代重定向到新的合并神经元。你可以重复这个过程很多次。

Alex 到现在已经在这个分析上投入了数小时。他发现了似乎在许多层中重复出现的电路。他会打印出不同等价类的节点，查看每个节点作为输入的权重序列，发现只有几种类型的节点。例如，有一类节点实际上会从两层前转发一个值。折叠这些以及类似的简化，将线性规划的大小从大约 200 万个节点减少到 75,000 个。

但在这一切之后，Alex 再次运行求解器，它又一次不停地运转而没有终止。

最终的简化

一个新想法：如果你通过网络传播边界会怎样？只需一层一层地推理，你可以找出任何给定节点能达到的最大值；你只需查看其输入的边界就能做到。结果表明，在相当保守的假设下，许多节点最终有非常紧的边界，比如从 0-1。也许这足以使程序可处理？

此时 Alex 从线性规划切换到 SAT 求解器，因为值的总数变得小了很多。在 SAT 版本中，你为每个节点等于其范围内每个值都有一个布尔变量。总共，在所有简化之后，这导致了 200,000 个变量。运行一天后，SAT 求解器将程序减少到 20,000 个变量。从那里它似乎不再进一步减少了。

实际上，Alex 发现这个神经网络内部有一个核心程序，它具有不可简化的复杂性——令他失望的是——它仍然太大而无法暴力破解。所以在许多天之后，他不得不退后一步，实际上一无所获。

灵光一现

他开始元思考：这一定是一个可解的谜题，对吧？有人会如何构建这样一个值得去解的谜题？

如果你生成随机权重，SAT 求解器可能能够通过暴力破解来解决它。这个网络是由人类创建的。在其核心似乎有一个你不能仅仅通过搜索或优化来恢复的函数。它是一个不可逆函数。有哪些不可逆函数的常见例子？

Alex 问 ChatGPT 一些常见的哈希函数，并将它们与层宽度的一些基本图表进行比较，这些图表看起来是周期性的。实际上有 32 个长度为 48 的周期，每次精确重复。也许网络在执行 32 块相同的计算？

再次问 ChatGPT：有没有常见的哈希函数使用 32 块计算？宾果。结果表明几乎所有的都是。

为了确定这里用的是哪一种，他手动探索：他会向网络输入一些字符串，用单独的程序计算各种哈希变体，然后查看倒数第二层。结果是 md5 对上了，其他常见哈希函数没有。

这很好，因为他已经通过查看倒数第二层的偏置知道了哈希应该是什么。所以问题简化为找到一个能产生那个特定 md5 哈希的输入字符串。但目前还不清楚如何解决这个问题——特别是因为他没有真正的证据证明这个网络总是产生 md5 哈希。也许解决方案是更深入地挖掘，并破解网络使其可逆？

矩阵中的故障

Alex 注意到网络中有一些奇怪的东西。它似乎有一个 bug：如果你的输入长度大于 32，它就不再产生正确的 md5 哈希。也许在这个 bug 中有一个关键可以逆转内置在网络中的哈希值？

他花了接下来的两天逆向工程这个 bug。首先，他让 Gemini 写了一个 md5 哈希函数的实现。然后他将网络中的每个神经元与 md5 算法中的相应变量匹配。他写了一些代码来存储给定中间变量的值序列，然后在网络的 32 个块中搜索那个值；这会找出哪些神经元范围对应于每个变量的位。结果表明，某些位范围正好对应于变量，其他的是中间计算值。

然后，对于长度 >32 的输入，他可以仔细追踪各个块，找到网络与正确算法分叉的确切位置。

关键在前 7 层——有一个电路会计算输入的长度，并尝试以小端序将其存储在 4 字节中。但当长度是 256 比特或更大时，你会有一个长度变量包含值 256，而不是正确的编码。也就是说，如果长度是 >384 比特，长度字节应该是 128 1 0 0，但网络编码的却是 384 0 0 0。

然后问题是，是否可能通过精心制作一个长度为 256 或更大的消息来利用这个 bug？更多仔细的追踪揭示了几个观察：首先，可能的长度并不多。只有 55 个输入，所以他可以进行穷举搜索，看看网络对这些奇怪的值是如何表现的。其次，损坏的长度值被转换为二进制，然后通过网络的所有层传播。在二进制中，所有的位都会等于 1，数字的其余部分集中在最低位，所以 384 会被编码为 130,1,1,1,1,1,1,1。第三，来自消息长度的无效字节只在 md5 计算的少数几个块中使用，它总是以相同的顺序从输入中读取字节。

使用这些观察，可以写下 md5 算法的修改版本，它在必要的块处自我修正以与神经网络一致。然而，仔细观察这一点，它仍然似乎很难在一般情况下逆转。

这花了大约两天才弄清楚，但——又一次失望——没有让 Alex 更接近解决方案。他给谜题提供的电子邮件地址写了信，告诉了他到目前为止发现的东西。他听到的回复让他惊讶。这个 bug 不是故意的。考虑到这一点，你为什么不最后尝试一次？

暴力破解的回归

结果表明，一旦你知道了编码在倒数第二层偏置中的哈希，你就完成了。弄清楚这一点是谜题的主要内容。谜题创建者故意使哈希容易暴力破解，在谜题描述和 Python 代码中留下了各种小提示，表明解决方案由两个英文单词组成，小写，用空格连接。

Alex 实际上早些时候曾尝试暴力破解哈希，但下载了一个包含前 10,000 个最流行单词的列表来做，结果发现不够大。一旦他有了一个足够大的单词列表，他就得到了答案。

另一个谜题

使这个谜题具有挑战性的因素之一是设计一个复杂度适中的网络。使用逻辑门意味着网络不可微分；但如果你让这些门编码的程序太复杂，就没有希望逆向工程它。md5 感觉是一个不错的折衷，尽管它绝非易事。因为 md5 使用模加法，创建谜题需要在约 20 层神经网络中实现一个并行进位加法器。不容易！我们印象深刻的是一些解题者设法弄清楚了这一点——而 Alex 发现的 >32 bug 是出乎意料且非常了不起的。

创建和发布谜题以及与解题者互动的体验足够好，我们又做了一次。这里你可以找到最新的。在这个新谜题中，一个神经网络各层被打乱了顺序，需要放回正确的顺序……你能帮忙吗？

关键要点

1. Mechanistic Interpretability 是真实的技术：这个谜题展示了逆向工程神经网络是可能的，但需要深入理解网络结构和耐心分析。

2. 多角度尝试很重要：Alex 尝试了线性规划、SAT 求解器、网络简化、手动追踪等多种方法，最终通过识别哈希函数模式解决。

3. 不要忽视简单线索：谜题描述中”vegetable dog”的示例和提示（两个英文单词加空格）就是关键。有时解决方案比想象的更简单。

4. 工具知识要广：ChatGPT 帮助识别了 md5 的 32 块结构，说明了解常见算法的内部结构很重要。

5. bug 可能是突破口：虽然这个 bug不是故意的，但 Alex 对 >32 输入时网络行为的分析展示了如何从异常中寻找线索。