最近我花了一个周末做了一件看起来有点杀鸡用牛刀的事用 Lean4 去验证旋转位置编码 RoPE 里的几个恒等式。很多人第一次听到这个组合会问一个平时写模型推理代码的人为什么跑去碰证明助手其实答案很简单RoPE 的实现太容易在符号细节上出错了而恒等式验证恰好能把这类错误变成编译期问题。这篇文章不是 Lean4 的入门教程也不是 RoPE 的完整推导而是一份把两者接起来的实操记录适合对 RoPE 实现有疑问、又愿意在形式化验证上投入一点业余时间的读者。1. 先想清楚RoPE 为什么需要“验证”1.1 RoPE 到底在做什么RoPE 的思想可以浓缩成一句话把位置 m 的信息编码成一个旋转矩阵 R_m作用在 query 和 key 向量上让最后的点积结果只和相对位置有关而不是绝对位置。实际实现里维度为 d 的向量会被切成 d/2 组二维子向量每一组套一个频率不同的旋转矩阵。频率公式大家都很熟悉theta_i 10000^(-2i/d)。第 i 组二维子空间使用的角度是 m * theta_i。比如一个向量 (x0, y0)在位置 m 上做的变换是x x0 * cos(m * theta_i) - y0 * sin(m * theta_i) y x0 * sin(m * theta_i) y0 * cos(m * theta_i)这个公式看起来平平无奇但当你把它拆成代码时问题就来了。是先乘 cos 再减 sin还是先加 sin 再乘 cos角度到底取 m 还是 n-m不同框架的 RoPE 实现里出现过不少类似的符号错位测试时因为精度和随机初始化兜底不一定能及时暴露。我最早想做形式化验证就是因为在一次代码 review 里发现两个版本的 RoPE 实现结果差了一个负号两边都觉得自己是对的。与其争论不如把 RoPE 相关的恒等式写下来交给 Lean4 判定。1.2 需要验证的恒等式清单形式化验证之前先列一份“到底要证什么”的清单。我整理出下面五条全部围绕二维旋转矩阵 R(θ) 展开恒等式数学表达为什么要验证零位置不变R(0) I确认边界条件位置 0 不改变向量旋转叠加R(ab) R(a)R(b)相对位置可以组合是外推性质的基础同角度保内积dot(R(θ)u, R(θ)v) dot(u, v)旋转矩阵正交同一个位置下内积不受旋转影响相对位置内积dot(R(mθ)q, R(nθ)k) dot(q, R((n-m)θ)k)RoPE 最核心的性质内积只与位置差相关逆旋转R(-θ)R(θ) I对应矩阵转置逆是实现反向位置偏移时的保证这里有一点很关键上面这些恒等式对任意 theta、任意实数位置 m、n 都成立不需要用到 10000 这个具体底数。也就是说证明时不需要关心频率表怎么生成只需要验证旋转本身的结构性质。这个抽象层做得好后面的实现无论怎么换 base、换维度都能复用同一套证明。1.3 数值测试和形式化验证的差别有人会问我写个 Python 脚本随机生成向量把 m 和 n 代入跑一百万个样例不也能验证吗能验证一百万个点但验证不了“所有”。普通测试的核心问题是采样。你可以覆盖最常见的维度、位置和浮点范围但总有一些边界条件是你没想到的尤其当数学式子里带着 m-n 这类符号时一个负号的错误在随机测试里往往被数值噪声掩盖。Lean4 的证明不是这样它托管的命题是全称量化的意味着它一次性覆盖了所有实数 θ、m、n 以及所有复数向量 q、k。证明一旦通过就没有“漏采样”的说法。这也是我把工作量投入到 Lean4 上的核心原因与其维护一堆回归测试用例不如把性质本身变成机器可检查的事实。RoPE 是当前主流模型都在用的模块这个验证集一旦写好收益是长期且跨实现复用的。2. Lean4 环境准备Windows 上从零到第一个可用项目2.1 为什么选 Lean4 而不是 Coq 或 Isabelle做数学形式化验证工具选择本身就是一个大问题。Coq 历史最久Isabelle 的自动化能力很强但对我来说Lean4 有两个不可替代的好处。第一个是数学库。Lean4 的 Mathlib 已经把复数、指数、实数域上的三角函数性质整理得非常完整正好覆盖 RoPE 会用到的那一层。第二个是策略手感。Lean4 里的simp、ring、linarith这些自动化策略对写过业务代码的人来说非常亲切很多中间步骤可以让机器直接消掉不需要像老式证明助手那样手工维护一长串 proof term。更重要的是Lean4 在 Windows 上的支持比想象中成熟。虽然网上资料大多默认 Linux/macOS 环境但只要把 elan 装对Windows 上完全可以本地跑通整个项目。这也让我确定把 RoPE 验证集放在 Lean4 里。2.2 lean4安装windows我踩过的完整路径先说结论lean4安装windows 这一步最简单的方式是 elan VS Code Mathlib 缓存。不要自己去源码编译也不要手动下二进制文件否则后续版本切换会让你头疼。我在干净 Windows 环境下的做法用 winget 安装 Lean4 工具链winget install --id LeanProver.Lean -e安装完成后确认 elan 可用elan --version lean --version新建一个使用 Mathlib 模板的项目。在 Lean4 里项目脚手架工具叫 Lake它会生成lakefile.lean、lean-toolchain和主目录结构lake new rope_verify math cd rope_verify lake update mathlib lake exe cache get用 VS Code 打开项目目录安装 Lean4 扩展。打开第一个.lean文件后右下角会开始初始化环境。第一次进入项目时编辑器会把当前.lean文件的类型信息和错误提示都加载出来。有一点必须提醒在 Windows 上项目路径不要带中文也不要有空格。我一开始放在C:\Users\张三\我的项目 rope_verify结果 Lean server 经常无响应日志里全是路径解析问题。后来把项目挪到纯英文路径所有怪问题都消失了。这算是我在这次 lean4安装windows 过程中最深刻的一条经验。另外首次拉取 Mathlib 时尽量直接用lake exe cache get下载预编译缓存。如果直接lake buildWindows 上可能要花很长时间编译整个 Mathlib有了缓存大部分常用模块可以直接加载。2.3 项目骨架和第一个可运行证明项目建好后我直接写了一个最小可运行文件用来确认环境没问题。这个文件不需要太多内容能跑通就说明整个工具链已经通了import Mathlib noncomputable section open Complex open scoped ComplexConjugate def rope (θ : ℝ) (z : ℂ) : ℂ : z * exp (θ * I) def dot (u v : ℂ) : ℝ : (u * conj v).re lemma rope_zero (z : ℂ) : rope 0 z z : by unfold rope rw [zero_mul, exp_zero, mul_one]这里我直接定义了旋转函数rope和复数内积dot。rope θ z表示把复数 z 顺时针旋转 θ 角度dot u v返回实数内积。第一个引理rope_zero验证的是“旋转零度等于没转”。如果环境没问题Lean的 InfoView 会显示这个引理已经被接受。不要小看这个引理它是后面所有证明的第一块砖。如果在零位置这个边界条件上符号搞错后面叠加性和相对位置公式全都会歪。3. 核心恒等式在 Lean4 里的证明实操3.1 为什么用复数而不是矩阵RoPE 原始论文里用的是旋转矩阵但我最终在 Lean4 里选择了复数编码。原因很简单复数指数函数在 Mathlib 里已经有一套成熟的定理比如exp_add、exp_zero、共轭和指数运算的交换性质换成矩阵加三角函数展开之后到处都是 sin 和 cos还要额外处理 sin^2 cos^2 1证明成本高不少。把二维向量 (x, y) 看成复数 x yI旋转 θ 角就是对复数乘一个单位复数z z * exp(θ * I)两个二维向量的内积在复数表达下刚好是dot(u, v) Re(u * conj(v))这个编码非常干净。rope和dot的定义如下def rope (θ : ℝ) (z : ℂ) : ℂ : z * exp (θ * I) def dot (u v : ℂ) : ℝ : (u * conj v).re很多第一次接触 Lean4 的人会觉得“等效定义”也需要证明但这正是形式化的一种优势你可以自己选择最顺手的模型。只要你最终证明的性质是目标性质中间的编码方式可以完全由你决定。我证明的是 RoPE 的数学性质不是某个具体 Python 实现的源码等价性所以用复数编码是合理且高效的。3.2 证明旋转叠加性R(ab) R(a)R(b)旋转最基础的性质是两次旋转可以合并成一次。RoPE 里如果想知道位置 ab 的向量长什么样最直接的办法是先转 b 再转 a也可以直接转 ab。Lean4 的证明比我想象中短lemma rope_add (a b : ℝ) (z : ℂ) : rope (a b) z rope a (rope b z) : by unfold rope have h : b * I a * I (a b) * I : by ring rw [← h, ← Complex.exp_add] ring这里核心是Complex.exp_add这条现成定理它说明exp (x y) exp x * exp y。我先用ring把指数上的加法顺序整理好再用这条定理把乘积展开最后的复数乘法交换律、结合律让ring自动处理。写这个小证明的时候我意识到一个容易被忽略的点Lean4 的ring策略只处理代数环结构它不会自动帮你证明exp的性质。所以证明里必须先定位到指数相加这一步再调用exp_add否则ring只能卡在原地。3.3 证明最核心的相对位置内积公式RoPE 之所以叫旋转位置编码最关键的性质是位置 m 的 query 和位置 n 的 key 做内积等价于不旋转 query只把 key 旋转 (n-m) 角度后做内积。用公式写就是dot(R(mθ)q, R(nθ)k) dot(q, R((n-m)θ)k)这个恒等式是 RoPE 能够表达相对位置的关键。Lean4 里的证明骨架是这样的lemma rope_dot_relative (θ m n : ℝ) (q k : ℂ) : dot (rope (θ * m) q) (rope (θ * n) k) dot q (rope (θ * (n - m)) k) : by unfold rope dot calc ((q * exp (θ * m * I)) * conj (k * exp (θ * n * I))).re (q * conj k * exp (θ * (m - n) * I)).re : by simp [map_mul, Complex.conj_exp, conj_I, mul_assoc] congr ring _ (q * conj (k * exp (θ * (n - m) * I))).re : by simp [map_mul, Complex.conj_exp, conj_I, mul_assoc] congr ring这段代码的思路是先把左边的内积展开利用共轭运算对乘法和指数分配把conj(e^(i n θ))变成e^(-i n θ)然后把所有指数因子合并成e^(i(m-n)θ)。右边则把conj(k e^(i(n-m)θ))展开成conj(k) e^(-i(n-m)θ)因为(n-m)的相反数正好是(m-n)两边实部相同。这是我整个验证过程中最有价值的一步。因为符号方向在这里非常脆弱如果右边写成R((m-n)θ)k证明最后会差一个负号Lean4 会在ring阶段直接报错。这种错误在数值测试里很可能被浮点误差和随机初始化掩盖但在证明助手里不可能蒙混过关。3.4 顺带验证同角度内积保持和逆旋转相对位置公式之外我还补了两个小引理作为完整性检查。第一个是“同一个旋转角度下两个向量都旋转后内积不变”lemma rope_dot_same_angle (θ : ℝ) (u v : ℂ) : dot (rope θ u) (rope θ v) dot u v : by unfold rope dot calc ((u * exp (θ * I)) * conj (v * exp (θ * I))).re (u * conj v * (exp (θ * I) * exp (-θ * I))).re : by simp [map_mul, Complex.conj_exp, conj_I, mul_assoc] _ (u * conj v).re : by have hcancel : exp (θ * I) * exp (-θ * I) 1 : by rw [← Complex.exp_add] ring simp [hcancel]这个引理数学上等价于旋转矩阵正交。第二个是逆旋转lemma rope_inv (θ : ℝ) (z : ℂ) : rope (-θ) (rope θ z) z : by unfold rope rw [← Complex.exp_add] ring两个证明都不长但它们覆盖了 RoPE 实现里最常被手滑的地方同角度下多个向量同时旋转时内积应当不变反向位置差应该用负角度还原而不是用补角。4. 踩坑记录与排查思路4.1 Windows 环境常见问题速查表我在整个过程中遇到的环境问题比数学问题还多整理成一张表给后面做类似项目的人参考症状原因处理方式Import Mathlib卡住VS Code 一直转圈没有预编译缓存正在现场编译先执行lake exe cache get尽量复用预编译产物VS Code 提示 Lean 服务无响应elan 没有正确加入 PATH重开终端确认lean --version必要时重新安装 elan项目路径含中文或空格Lean server 对路径解析不稳定把项目迁移到纯英文目录例如C:\lean\rope_verify找不到某个 Mathlib 定理版本更新导致定理名变化用#check搜索当前环境里的实际名字不要照抄旧教程simp策略超时一次展开范围太大把无关性质也塞进来改成simp only [...]或改用calc分步推理代码一直报错但数学上“明显对”符号方向写反通常是 m-n 和 n-m 的笔误回到公式把指数合并结果单独打印出来看正负号这些坑没有一个是数学难到写不出来基本都是工程细节。但如果不注意耗掉一整个晚上的往往是这类问题而不是证明本身。4.2 建议的入门验证顺序如果你也想复现这套验证而不是直接抄我的代码建议按照下面的顺序来做。这个顺序是我踩完坑之后总结出来的每一层都建立在前一层能被证明的基础上rope_zero位置 0 不旋转。rope_add旋转可以合并。rope_dot_same_angle同角度旋转保持内积。rope_inv旋转和负角度旋转互逆。rope_dot_relative核心相对位置公式。最后如果你恰好在改自己的 RoPE 实现可以额外加一条“频率与下标对应”的装配性质把维度分组也纳入验证范围。每完成一步就用#check看一下当前引理的类型确认它确实是你想要的全称命题。比如rope_dot_relative的完整类型应该包含 ∀ 符号说明对所有实数位置成立。如果类型里出现了具体的 0 或 1说明你把命题不小心特化了。4.3 关于证明策略的几条体会Lean4 的自动策略确实省力但也不是万能。我的体会有三条。第一优先用calc写出证明骨架。编程时注释会写“这是推导过程”在 Lean4 里calc就是这个注释的机器可读版本。先把要证明的链条写出来每一条看成一个目标再逐个击破比一上来就simp全扫要可靠得多。第二ring能帮你处理的是加减乘除和交换结合律不要指望它直接处理exp、sin、cos这类超越函数。所有三角函数和指数相关的重写先手动定位到对应定理再交给自动化策略收尾。比如rope_add里必须先处理exp_addrope_dot_same_angle里必须先处理exp的共轭顺序反了就会被全局simp拖进死胡同。第三不要把sorry留在最终代码里。开发阶段用sorry占位没问题但它的存在等同于向 Lean4 声明“这里我不验证了”。真正的验证集必须保证最终代码里没有任何sorry否则就失去意义。我个人在实际操作中的体会是用 Lean4 验证 RoPE 的价值不在于“证明完成”这个结果而在于它逼着我把每个下标、每个符号、每个频率参数都讲清楚。接下来如果还有精力我会把这组恒等式整理成一个可复用的规范让不同的位置编码实现都来做同一组证明这比单纯跑一个精度对比实验要枯燥但得到的信心完全不一样。