用 Lean 写 Zstandard 解压器:当 LLM 把「证明成本」压下来
Adam Langley 用 Lean 写了一个 Zstandard 解压器,把类型级不变量写成类型签名,再让 LLM 去补证明体, 并记下它在哪些步骤上真的省了事。
中文
复制
ImperialViolet
我们现在有证明自动化了(2026年7月26日)
我一直对 Coq、Rocq 和 Lean 这类依赖类型语言情有独钟。它们让人有可能构建出足以编码并强制实施任意精细不变量的类型系统。这类不变量在普通语言里,最好的下场也不过是变成一句注释,而且随着团队规模扩大,很快就没人记得了。接下来就是各种微妙的理解偏差,以及拼不到一起的组件。而等到问题被发现时,这些组件往往已经长到足够大,想对齐其中任何一个都让人望而生畏。依赖类型于是抛出诱人的说法:也许你可以把这些不变量形式化地写出来,让机器去检查。
(附注:Coq 改名了!我记得很多年前在普林斯顿的一次 Coq 会议上,我试着提出,在一个讲英语的世界里,一门编程语言叫 Coq 是个障碍。当时听众似乎并不认同。我还开玩笑说,那里的很多报告听起来都像 Tyrion Lannister 的演讲,因为 Coq 和 Hoare 实在太多了。这个笑话既好笑又应景,尽管它彻底冷场了——毕竟那是在那部剧最后一季之前,也是在大家集体把它从记忆里抹掉之前。)
问题一直在于,类型系统的能力越大,证明的代价就越大。我可以作证,为了证明一些相当简单的东西,我整整天都搭进去了。做证明其实挺有意思:有挑战性,可交互,目标明确。但天哪,它太花时间了,尤其是如果你像我一样根本不知道自己在干什么。还有那种周期性的、令人恼火的体验:花了好几个小时之后,你才意识到自己想证的目标其实是_假的_。这方面经典的结果来自 seL4 项目的回顾:他们发现,尽管项目规模已经大到让工程师积累了相当多的经验,他们在证明上花的时间仍然是设计与实现的大约 10 倍。最终证明代码的行数超过 C 代码的 20 倍。
这种负担让依赖类型语言的编程变得极其小众,也促使人们尝试把它自动化掉。我略有了解的一个尝试是 F*,系统会试着让 SMT solver 自动完成这些义务。简单的情形当然没问题,但很容易构造出某种东西,让 SMT solver 一路飞向太空、跑上几个小时,你只能干等着,不知道它到底会不会结束。我见过很多常用这类语言的人,他们不得不培养出第六感,判断什么能让 solver 满意,然后围绕这一点来写所有东西。这确实有帮助,但某种程度上把问题变成了玄学:你最后是在伺候一个复杂又喜怒无常的神。
一个关键事实是,至少在理论上,一旦命题是正确的,其证明的内容就无关紧要:只有它的存在才有意义。这话并不完全成立,因为有两个复杂因素:第一,seL4 团队所说的“证明工程”——需要把证明组织好,以便代码改动后重新对齐证明的工作量能降下来。第二,足够复杂的证明甚至能让类型检查器炸掉,吃掉大量内存。
现在我们有了 LLM,结合证明无关性,它有望成为一种能力极强的证明自动化手段。有了足够多的自动化,也许你就不必那么操心证明工程了。你仍然需要避免把类型检查器搞炸,但在我有限的测试里,LLM 能避开这一点。LLM 有可能一下子让依赖类型系统变得实用得多。我想试试这个,于是用 Lean 写了一个 Zstandard 解压器,主要也是因为我对 Zstandard 本身很好奇。
在取代 gzip 成为标准压缩工具的竞争中,Zstandard 看起来正在胜出。它同样是 LZ77 风格的压缩器,但熵编码更好,设计也更讲究,因此解压速度非常惊人。它永远不会像 bzip2 那样优美,但面对实打实的实用优势,Burrows–Wheeler 变换那点耀眼的优雅并不算得了什么:
(测量是在标准参考计算机上进行的,也就是作者当时手头在用的那台。另外注意 y 轴是对数刻度:gzip 和 Zstandard 处在各自的速度档位。这是一台 Apple 机器,Apple 的 gzip 做了特别优化;换到别处,gzip 应该会更慢。)
Zstandard(由 Yann Collet 开发,建立在 Jarek Duda 开创性的 ANS 工作之上)有一份 RFC,但写得相当简略。实现一个解压器所需的信息它都有,不过除非你已经对压缩相当熟悉,否则我觉得你得反复读上几遍才能搞明白到底在讲什么。至少我自己把 4.1 节读了六遍,才觉得算是勉强掌握了。等我意识到的时候已经太晚:我的同事 Nigel Tao 写的 Zstandard 讲解比我本来能写出来的任何东西都好。所以,想理解 Zstandard,你该去读那篇。我这里只解释最有意思的部分——熵编码器,顺便夹带一点对 Lean 的安利。
熵编码器的任务是:给定一组概率不均匀的符号,用最少的比特数把由这些符号构成的序列编码出来。经典的熵编码器是 Huffman 编码器。Huffman 编码器建一棵二叉树,符号放在叶节点上;Huffman 证明了有一个非常简单的算法能生成最优前缀树:取符号列表,找出概率最小的两个,用它们作为子节点构成一个树节点。这个树节点的概率就是两个子节点概率之和,然后对少掉两个符号、多出一个树节点的集合重复这一算法。显然,算法每执行一步,元素集合就少一个,所以它必然终止,而且生成的树是最优的。Huffman 树非常快,因为你可以建一张以接下来 n 个比特为索引的表(n 是最长编码的长度)。表项会告诉你解出了哪个符号,以及要回退多少比特。Huffman 树的缺点在于,每个符号只能用整数个比特:如果某个符号满足 -log2(p) = 2.3,理想情况下你希望用 2.3 个比特来编码它。但 Huffman 逼你要么向上取整到 3 个比特,要么向下取整,而后者会迫使其他符号多消耗比特。
Zstandard 使用 Huffman 树,但它还有一种压缩率更高的熵编码器,叫做 FSE。FSE 是一台状态机。状态的数量多于符号的数量,每个符号分到的状态比例,正好对应它在数据流中出现的概率。所以,如果某个符号预计有 50% 的时间会出现,它就能分到大约 50% 的状态。每个状态有三个值:该状态对应的符号、处于该状态时要从比特流中读取的比特数,以及一个基准状态编号——把这个编号加上读到的比特,就得到下一个状态。回想一下,Huffman 树的问题在于它只能用整数个比特,而这些状态读取的也是整数个比特。但诀窍在于:如果你想为某个符号读取 1.5 个比特,那就让它一半的状态读一个比特,另一半读两个比特。这样你就能在_平均意义上_达到目标。状态表从不传输。RFC 规定了一种算法,可以从符号概率列表构建出这张表,所以只需要传输概率。
举个例子。假设有四个符号,我们要用 16 个状态。那么就得用 1/16 的粒度来近似这些符号的概率。(如果想要更精确的概率近似,可以用更多的状态;zstd 实际上从不使用少于 32 个状态。)
| state | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Symbol | A | A | B | D | A | B | C | A | B | C | A | B | C | A | A | B |
| Num_Bits | 2 | 1 | 2 | 4 | 1 | 2 | 3 | 1 | 2 | 2 | 1 | 1 | 2 | 1 | 1 | 1 |
| Baseline | 12 | 0 | 4 | 0 | 2 | 8 | 8 | 4 | 12 | 0 | 6 | 0 | 4 | 8 | 10 | 2 |
任何符号后面都可能跟着任何符号,而一个符号也可能只有一个状态。所以每个符号都必须能够到达每一个状态。看看状态 3,它是符号 D 唯一的状态。正因为是唯一一个,它必须读取四个比特,这样才足以编码任何其他状态。但如果你看符号 B 这样的符号,它的状态只要求你读一个或两个比特。不过,16 个可能的下一个状态,恰好被符号 B 的这些状态划分完毕。所以,对任意一个特定状态来说,符号 B 中恰好有一个状态能够到达它。
再来看符号 B,我们之前说它的概率是 5/16。编码该符号的理想比特数是 -log2(5/16) = 1.68。符号 B 有三个状态读两位、两个状态读一位。这些状态的使用频率并不相同,按使用频率加权后,平均值几乎正好落在量化概率对应的正确数值上。如果想更精确地还原真实的符号概率,用更大的表就行。
核心技巧在于,给更常见的符号分配多个状态后,编码器不只是选一个符号:它还要选落在该符号的哪个状态上,而这个选择会把信息传递到下一个符号。小数位的信息就藏在这里。但这个熵编码器仍然只是基于表的,所以运行得非常快。
麻烦在于,你没法正向推导。假设你想编码 C、D。你从哪个 C 状态开始?D 只有一个状态,所以只能是那个能到达它的 C 状态。如果 D 有多个状态,你就得考虑 D 之后是什么,才能知道需要其中_哪一个_。FSE 迫使你从序列末尾开始、倒着推。(这倒不算太糟,因为计算符号概率本来通常就需要知道整个序列。)更进一步,Zstandard 压缩器因此是倒序编码符号,但输出是增量写入的,所以解压器必须定位到块的末尾、反向读取比特,才能把它理顺!这就涉及格式更宏观的细节了,我不打算展开;参见 Nigel 的文章。
基本的熵编码器不关心符号间的概率。也就是说,它们无法利用字母 Q 后面不成比例地常跟着字母 U(在英语中)这一事实。必须有别的编码方式来利用这些冗余。在 Zstandard 中,那就是传统的 Lempel–Ziv 结构,它编码字面字节或对已解码数据的反向引用。所以 FSE 主要用于高效编码这些反向引用的偏移量和长度。
Lean
来说说 Lean。前面我说过它是依赖类型语言,这个概念用例子讲比用复杂的定义讲更清楚。下面是一个函数的类型:它从流中读取 n 个字节,如果没有抛出异常,就返回一个字节数组,而类型系统知道它的长度正是 n。
def IO.FS.Stream.readExact (st : Stream) (n : Nat) :
IO {ba : ByteArray // ba.size = n} := …
下面这个函数返回两个数和一个字节数组,要求第一个数是质数,两个数之和能被六整除,且字节数组的长度不小于这两个数中较小的那个。
def getResult :
IO (Σ a b : Nat, { bytes : ByteArray //
Nat.Prime a ∧
6 ∣ a + b ∧
Nat.min a b ≤ bytes.size }) := …
没人会真的需要这样的类型。它只是说明,你完全可以随心所欲地玩到这种程度。依赖类型语言足以编码极其复杂的数学结构,而 Lean 目前最主要的用途,是作为陈述和证明数学的形式化语言。最近出版的书 The Proof in the Code 简短而文笔出色,讲述了 Lean 的来龙去脉。作者确实用好几段把构造性数学彻底讲错了,但除此之外,我读得很享受!
Lean 和 Haskell 一样是纯函数式语言,不过它有几个特性,让它在作为编程语言时可能方便得多。首先,Lean 是严格求值的,而 Haskell 是惰性的。严格求值意味着函数的参数在调用发生之前就被求值,而在 Haskell 中,参数的求值会推迟到真正需要这个值的时候。所以在 Haskell 里,你可以随意写开销很大的表达式并把它传进函数,因为它只有在最终被用到时才会真正计算。但这也意味着计算可能发生在程序中非常出人意料的地方。这个话题有争议,但尽管我欣赏惰性的优雅,天哪,它确实让程序的性能很难推理。
其次,Lean 有不少好用的语法糖。它的 monad do 记法里包含 for 循环、return 语句和 break 语句。如果你想用命令式风格编程,完全可以做得相当顺手!
最后,Lean 有一项优化:只要对象的引用计数为一,它就会对对象进行可变更新。因此,只要小心别在别处还持有对它的引用,你就能像在命令式语言里那样高效地原地修改数组。遗憾的是,据我所知,Lean 没有任何线性类型系统机制,所以它无法帮你确保一个值只有一个引用。这是个有点扎手的边角:代码里一处看似微不足道的改动,可能因为在某个不起眼的地方还攥着一个大数组的引用,就让性能彻底崩掉。但这也意味着,如果你想优化某个东西的性能,手头可用的工具要多得多。
下面是我草拟的 zstd 解码器里的一个例子:
while true do
let some blockHeaderBytes ← input.readExactOrEof 3 | break
let some blockHeader := BlockHeader.fromBytes blockHeaderBytes frameHeader
| throw (.userError "invalid block header")
let blockBytes ← input.readExact blockHeader.contentSize
match hty : blockHeader.type with
| .rle =>
let b := blockBytes.val[0]'(by
rw [blockBytes.property, blockHeader.contentSize_rle hty]; omega)
看第 9 行。那里有个数组索引,而索引正是隐式不变量最爱藏身的地方:blockBytes 最好别是空的!C 系语言在这种情况下会给你未定义行为。现代语言会在运行时抛异常,或者干脆只给你一个可选值来避开这个问题。Lean 还有另一个选择:证明它不为空。第 10 行做的就是这件事。blockBytes.property 就是它和所请求的读取一样长这个事实,也就是恰好 blockHeader.contentSize 字节长。blockHeader.contentSize_rle 是这样的:
theorem BlockHeader.contentSize_rle (h : BlockHeader) (hty : h.type = .rle) :
h.contentSize = 1 := by
simp [contentSize, hty]
这证明了当类型为 rle 时,contentSize 恒为一。有了这些事实,剩下的 Lean 自己就能推出来。
这个证明很短,大概我自己也能想出来,但我们可以把目标定得更高:
我按照 RFC 写了 FSE 表构造算法的实现。RFC 里为它附了“测试向量”:给定概率下的三个样例输出。这些当然会进单元测试。但在 Lean 里,我们还能证明这个函数的普遍性质:
theorem ofDistribution_wellFormed (h : ofDistribution accuracyLog probs = some t) :
t.entries.size = 2 ^ accuracyLog ∧
(∀ s : Fin probs.size,
t.entries.toList.countP (fun e => e.symbol == s.val) = probCells probs[s]) ∧
(∀ (i : Nat) (hi : i < t.entries.size) (v : Nat), v < 2 ^ (t.entries[i]'hi).nbBits →
(t.entries[i]'hi).baseline + v < 2 ^ accuracyLog) ∧
(∀ (s : Fin probs.size), 0 < probCells probs[s] → ∀ x < 2 ^ accuracyLog,
∃! i : Nat, ∃ hi : i < t.entries.size,
(t.entries[i]'hi).symbol = s.val ∧ (t.entries[i]'hi).baseline ≤ x ∧
x < (t.entries[i]'hi).baseline + 2 ^ (t.entries[i]'hi).nbBits) := …
再用文字重复一遍:
假设建表函数在给定“accuracy”常量和符号概率列表时能产出一个值,那么:
- 表的尺寸与该 accuracy 相匹配。
- 给定符号的状态数量与其概率相符。
- 对所有状态,读取 nbBits 位并加上该状态的基线值,得到的是一个合法状态编号。
- 对所有概率非零的符号,以及所有目标状态,该符号下恰好有一个状态能到达目标状态。
这些正是优化后的解码内层循环所依赖的隐含前提,在表达能力较弱的类型系统里,它们只能以隐含约定或注释的形式存在。证明这类强命题,正是 seL4 回顾中所说的那 10× 工作量的一部分,也是依赖类型在普通软件中难以推广的主要障碍。现在有几个 LLM 能在约 20 分钟内自动完成,而且只消耗每月 20 美元订阅额度的一小部分。明年这大概会成为基本要求。必须承认,它们在做这件事时不得不修改建表代码:我用了太多 Id.run(也就是切进命令式模式),证明机制处理起来更吃力。(不过 Lean 正在改进这一点。)我确认了这些证明能通过类型检查,也没有 sorry。
把依赖类型和 LLM 结合起来并不是新想法,但把这种组合用于日常软件工程的工作还很少。还需要积累大量经验。非常强的类型会放大改动的波及范围,因为它们必须沿着所有派生类型一路传播出去。也许在更大的系统里,证明工作量会增长得难以承受,连现代 LLM 也跟不上。Lean 是一门高层语言,并不适合所有场景。(我那个玩具 Zstandard 解码器在命令行上比 zstd 慢 10×。)尽管如此,证明自动化现在已经到来,实际上我们多了一类可用的编程语言。这很令人兴奋!
(我不公开代码,因为说实话,对于这种小而明确的问题,LLM 大概率比我做得更好。我做这件事只是为了稍微学一点 Lean,并不把自己的探索当作范例。灵感来自 lean-zip,它做得更多,包含一个压缩器,还证明了往返一致性!)
题外话:经过验证的汇编
AWS 做了 LNSym:一个 AArch64 的语义与模拟器。挺酷的。也许我们可以用它来证明某些函数的优化汇编实现与它们的 Lean 版本等价,然后在运行时直接使用汇编代码?这样就能放手让 LLM 去优化,而它们不可能引入任何功能性 bug。经过验证的汇编在密码学实现里已经很成熟,但也许现在它可以变得 便宜?
我花了一些时间(大部分是 LLM 的时间)尝试这件事。仓库里那个小小的 popcount 示例用了 bv_decide,一个可验证的 SAT 求解器,而这个示例需要的内存超过了我机器的上限,这不是个好兆头。极小的函数确实能跑通,可以拿到与极小 Lean 函数等价的证明,然后用 extern 在运行时调用它们!但我和几个 LLM 都没能让它扩展到更大的规模。