把旋转位置编码RoPE里的三角恒等式，用Lean4一座一座敲成机器可检查的证明，是我最近做得最“硬核”也最踏实的一件事。恒等式验证做到后面，你其实不是在写公式，而是在跟自己的数学直觉辩论：很多你以为显然成立的结论，Lean会逼着你把每一步依据都摆出来；而一旦它接受了你的证明，这份结论就不再依赖任何人眼检查，而是由形式化内核逐条校验过。这篇文章会把这个项目的完整过程摊开讲，从前置数学原理、Lean4在Windows上的安装，到正交性、合成规则等核心恒等式的证明代码、踩坑记录，尽量让读过的人能直接复现。
如果你是做LLM推理、自定义算子、或者纯粹对“用机器证明数学模型”感兴趣的人，这篇会很对你的胃口。我不打算把Lean4讲成全功能教程，只讲够用的tactic，配合RoPE这个真实场景，让抽象的形式化验证落到能跑、能验证、能继续扩展的代码上。
1. 项目缘起：为什么要给RoPE做机器证明
1.1 一个让人夜不能寐的恒等式
Transformer的attention里，query和key向量通常先各自经过投影，再做点积。RoPE做的事情，是给这个点积加上一个“位置条件”：位置i的q向量旋转角度a，位置j的k向量旋转角度b，最终点积只依赖它们的相对位置j - i。整个机制能成立，靠的其实不是一个复杂结论，而是一串很基础的旋转矩阵恒等式。
在实现RoPE的代码里，你一定会遇到这样的操作：把向量按维度对切分，每对元素乘上cos(θ)、再交叉乘上sin(θ)，实现一个二维旋转。看起来简单，可一旦把代码从单头推广到多头、从FP32换到自定义算子，矩阵的方向、转置、加法顺序稍微弄错一点，loss曲线就会给你颜色看。更麻烦的是，一两行维度置换错误通常不会直接报错，只会让模型悄悄变傻。
我当时真正想做的是：把RoPE这几个核心恒等式放到Lean4里验证一遍。比如：
R(θ)ᵀ · R(θ) = I R(a) · R(b) = R(a + b) R(a)ᵀ · R(b) = R(b - a)这三条看起来平淡无奇，但RoPE的很多实现细节都藏在它们里面。用Lean4把这些验证清楚，等于给RoPE的理论底座做了一次有限且完整的“体检”。
1.2 为什么不做随机测试，而是选择Lean4
你可能第一个想到的是写个Python脚本，随机生成10000组角度，用数值检查这些等式是否成立。我以前也这么干，确实能抓出不少低级错误，比如sin和cos调用反了、符号写错、数组维度没对上。但数值测试有个先天短板：它只能证明“你抽到的这些点成立”，无法覆盖全部实数角度。
更严谨一点，可以用SymPy之类的符号计算引擎去做推导。但那更多是“符号化简”，它不会阻止你在一开始就把等式方向写反；而且SymPy给你的结果仍然是它内部规则计算后的产物，不是逐条可审计的推演记录。
Lean4这类定理证明器走的是另一条路。你写下一个命题，比如R(a) * R(b) = R(a+b)，然后给出证明策略，Lean内核会把你的证明逐步检查成一个正式的、可以被机器验证的证明项。它不会说“看起来没问题”，而是要么说“证明通过”，要么给出一个反例或者告诉你这一步无法由已知规则推出。
我做了个简单的对比，帮助团队理解为什么值得投入精力：
| 验证方式 | 能证明什么 | 局限 |
|---|---|---|
| Python数值采样 | 抽样点上的数值相等 | 无法覆盖无穷实数域，精度问题可能掩盖错误 |
| SymPy符号化简 | 在符号规则下化简相等 | 依赖化简规则本身，等式方向容易搞反 |
| Lean4定理证明 | 在形式化内核中，逐步推理出结论 | 学习门槛高，证明过程可能要来回调试 |
所以我的结论很明确：RoPE这种在训练和推理中反复执行的核心数学操作，值得用最严格的方式验证一遍。Lean4可能不是最快的工具，但它给出的保证是别的方式替代不了的。
1.3 验证边界怎么划
这个项目并没有去验证整个Transformer，那不是一篇文章能做完的事。我划的边界是：RoPE在数学层最核心的旋转矩阵恒等式。
具体说，验证这些内容：
- 旋转矩阵的正交性：
R(θ)ᵀ R(θ) = I - 旋转的合成规则：
R(a) R(b) = R(a + b) - 转置与负角度的关系：
R(θ)ᵀ = R(-θ) - 由上述结论推导出的attention点积重写：
(R(a)u)ᵀ (R(b)v) = uᵀ R(b-a) v
这些结论一旦验证通过，对上层实现就是一块稳固的地基。不管上层用的是PyTorch、自定义CUDA算子，还是换一种新的RoPE变体，只要数学结构是同一个，底下的证明仍然有效。
2. RoPE数学原理：从旋转矩阵到恒等式族
2.1 二维旋转矩阵的几何直觉
二维旋转矩阵的标准形式是：
R(θ) = [[cos θ, -sin θ], [sin θ, cos θ]]为什么是这个形式？可以从复平面的乘法理解。一个复数z = x + iy乘上e^(iθ) = cos θ + i sin θ，得到的新复数实部是x cos θ - y sin θ，虚部是x sin θ + y cos θ。把实部和虚部当成二维坐标，上面这个矩阵就是“逆时针旋转角度θ”的几何操作。
旋转矩阵天然携带几个性质：
- 列向量互相正交，且每个列向量长度都为1，所以
R(θ)ᵀ R(θ) = I。 - 旋转的逆就是转置，也是负角度旋转：
R(θ)⁻¹ = R(θ)ᵀ = R(-θ)。 - 先旋转
a再旋转b，等价于直接旋转a+b：R(b)R(a)=R(a+b)。
这些性质在RoPE里不是教学练习，而是实际代码正确性的依据。很多RoPE实现把query按位置做旋转、再把key按位置做旋转，最后点积之所以能退化成相对位置函数，靠的就是这些恒等式。
2.2 为什么RoPE要做分块旋转
如果只处理二维向量，旋转矩阵直接套进去就行。但Transformer的hidden size通常是128、256、1024甚至更高，不能直接把一个二维旋转矩阵拉成高维，那样会丢失通道间的表达能力。
RoPE的常见做法是：把维度d拆成d/2对，每对维度独立做二维旋转。第i对维度使用频率：
θ_i = base^(-2i/d)其中base一般取10000。这样低频维度负责长距离依赖，高频维度负责局部细节。由于每个维度对只做二维旋转，高维向量整体上是一个分块对角矩阵：
[[R(θ_1), 0, 0], [ 0, R(θ_2), 0], [ 0, 0, R(θ_3)]]对Lean4验证来说，这种分块结构反而是好消息：核心性质都在二维块里，块之间的零矩阵不会给证明增加额外障碍。所以我的项目先验证二维块的全部恒等式，再把结论推广到分块对角场景。
2.3 目标恒等式清单
我在项目里维护了一张“恒等式清单”，每一条都对应一个Lean定理：
| 恒等式 | 数学表达 | 工程价值 |
|---|---|---|
| 正交性 | R(θ)ᵀ R(θ) = I | 保证旋转不改变向量长度，数值上更稳定 |
| 合成规则 | R(a) R(b) = R(a + b) | 多个旋转可以合并成一次旋转，简化计算 |
| 转置-负角 | R(θ)ᵀ = R(-θ) | 满足RoPE中query/key不同方向旋转的换序需求 |
| 相对位置重写 | (R(a)u)ᵀ(R(b)v) = uᵀR(b-a)v | 说明attention分数只依赖位置差 |
实际上，“合成规则”和“转置-负角”合起来可以推出最后一条。这也体现了定理证明项目的价值：你不需要为每一个小结论单独写证明，很多结论可以由先前已证明的定理自然推出。
3. Lean4基础与Windows环境搭建实录
3.1 一眼看懂Lean4证明是怎么回事
Lean4是一个依赖类型的定理证明器，也是一门函数式编程语言。它最核心的思想是：命题本身是一种类型，证明是这种类型的值。比如你要证明A = B，这个A = B是一个类型；你写的证明脚本，最终被Lean内核编译成一个类型正确的表达式，这个表达式就是证明。
很多开发者第一次接触Lean会觉得抽象，我当时的理解是：它类似于一个极严格的类型系统——类型不对，编译不过；证明不完整，内核不认。
写证明时，我们通常不直接写底层的证明项，而是用tactic（策略）引导Lean完成。策略有点像IDE里的代码补全和重构工具：你说“在这个目标上重写一下”、“用ring算一下多项式”，Lean会去执行并返回新的目标状态。调试过程很直观，右侧面板会实时显示当前目标和剩余子目标。
3.2 Windows上安装Lean4全流程
这个项目有一部分时间花在Windows环境搭建上。Lean4在Windows上现在很成熟，但仍然有几个容易踩坑的点，我把流程整理成可复现的步骤。
第一步，安装Elan。Elan是Lean的版本管理器，类似Python的conda或者Node的nvm。在Windows上，最简单的做法是直接用winget：
winget install leanprover.lean如果你更习惯官方安装脚本，可以用PowerShell执行：
Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass iwr https://raw.githubusercontent.com/leanprover/elan/master/elan-init.ps1 -UseBasicParsing | iex第二步，安装VS Code扩展。在VS Code扩展市场搜索Lean4，注意认准发布者是leanprover，不是第三方同名插件。装完后随便打开一个.lean文件，扩展会自动激活。
第三步，验证环境：
elan --version lean --version如果提示找不到命令，多半是PATH没刷新，重新打开终端就行。如果之前装过旧版Lean，建议先卸载干净，再装Elan，否则旧程序可能截胡lean命令。
第四步，创建项目并引入Mathlib。RoPE验证要用到实数和三角函数的很多定理，所以依赖Mathlib库。项目根目录的lakefile.lean里加一行：
require mathlib from git "https://github.com/leanprover-community/mathlib4.git"然后在.lean文件头部写import Mathlib，保存后VS Code会开始构建Mathlib。这一步耗时可能比较长，实际等待时可以先去把这篇博文剩下的内容读完，构建完成后再回来看代码。
3.3 写证明前必须掌握的几个tactic
Lean的tactic很多，但RoPE验证用到的其实就那么几个。我按使用频率整理出来：
| tactic | 用途 | 常见失效场景 |
|---|---|---|
simp | 化简等式，展开定义，做基础重写 | 遇到需要ring的多项式整理时会停住 |
ring/ring_nf | 解决交换环上的多项式等式 | 处理不了三角函数本身，必须先重写三角公式 |
nlinarith | 处理非线性算术不等式和等式 | 不会自动引入sin²+cos²=1，要手动传入 |
ext i j | 把矩阵、函数相等拆成“每个元素相等” | 元素类型不是离散枚举时需要配合归纳 |
fin_cases i | 对Fin 2这样的有限类型做情况枚举 | 无 |
rw [lemma] | 在目标中按给定方向替换等式 | 方向写反会让目标更乱 |
have h : X := by ... | 在证明中引入中间结论 | 无 |
我建议在开写RoPE证明前，花半小时用这几个tactic跑一遍Mathlib里最简单的例子，比如证明a+b = b+a。等到ring成为肌肉记忆，再进入RoPE的矩阵证明会顺很多。
4. 完整证明流程与代码细节
4.1 搭建最小验证骨架
先定义二维旋转矩阵。在Lean4里，我用Matrix (Fin 2) (Fin 2) ℝ来表示2×2实数矩阵。由于sin和cos在实数域不是可计算的函数，定义需要标记为noncomputable：
import Mathlib noncomputable def ropeRot (θ : ℝ) : Matrix (Fin 2) (Fin 2) ℝ := !![cos θ, -sin θ; sin θ, cos θ]!![...]是Lean4里方便的矩阵字面量语法。这样定义好处是后续所有定理都可以直接引用ropeRot，不用每次重复展开矩阵元素。
值得说明的是，这里我先把RoPE简化为“一个维度对”的二维旋转，并没有直接做d维分块对角矩阵。原因很简单：先验证最小单元的性质，再在工程代码里把分块结构组合起来。如果二维块的恒等式都不成立，高维分块部分再怎么包装也没意义。
4.2 证明旋转矩阵的正交性
正交性是RoPE数值稳定性的来源：旋转前后的向量长度不变，也就是R(θ)ᵀ R(θ) = I。证明的思路是：先ext i j把矩阵等式拆成“每个位置元素相等”，然后fin_cases枚举i和j是0还是1，最后用simp展开定义。
关键在于最后一步需要用到sin² θ + cos² θ = 1。ring只处理多项式，不知道这个三角恒等式，所以要把sin_sq_add_cos_sq θ这一个Mathlib定理显式传给nlinarith：
theorem rope_orthogonal (θ : ℝ) : (ropeRot θ)ᵀ * ropeRot θ = 1 := by ext i j fin_cases i <;> fin_cases j <;> simp [ropeRot, Matrix.mul_apply, Matrix.one_apply] <;> ring_nf <;> nlinarith [sin_sq_add_cos_sq θ]这段证明很简洁，但背后的调试并不轻松。最开始我试着只靠simp，结果Lean把目标化简成一个带cos θ * cos θ + sin θ * sin θ的式子后就不动了；后来加上ring_nf把它整理成cos θ ^ 2 + sin θ ^ 2，再补上nlinarith [sin_sq_add_cos_sq θ]才通过。
实际操作里，sin_sq_add_cos_sq θ这类定理名不需要死记，可以用#check sin_sq_add_cos_sq查看类型，确认参数是θ : ℝ后再使用。
4.3 证明角度相加的合成规则
合成规则R(a) R(b) = R(a + b)是RoPE最核心的恒等式。它保证了“可以先旋转a，再旋转b，合并成一次旋转a+b”。在attention计算里，这意味着位置差可以合并成一个相对位置，而不是每次都要展开成两步。
证明的整体策略和正交性一样，先按元素拆开，枚举四个位置。区别是这次需要用到三角函数的加法公式：cos(a+b)和sin(a+b)的展开式。Lean的Mathlib已经提供了cos_add和sin_add，把它们传给simp即可：
theorem rope_add (a b : ℝ) : ropeRot a * ropeRot b = ropeRot (a + b) := by ext i j fin_cases i <;> fin_cases j <;> simp [ropeRot, Matrix.mul_apply, Matrix.add_apply, cos_add, sin_add] <;> ring四个case分别对应矩阵的左上、右上、左下、右下四个位置，每个位置最终都变成同一个多项式的两边，ring负责收尾。
这里再次体现ring和simp的分工：三角函数加法公式先由cos_add、sin_add重写，重写后剩下的就是纯多项式运算。如果不在simp里传入这两个定理，ring是没办法处理cos(a+b)这种东西的。
4.4 从矩阵等式到RoPE点积重写
有了rope_orthogonal和rope_add后，再证明attention点积的相对位置性质就顺理成章了。先补一个转置与负角的关系：
theorem rope_transpose (θ : ℝ) : (ropeRot θ)ᵀ = ropeRot (-θ) := by ext i j fin_cases i <;> fin_cases j <;> simp [ropeRot, Matrix.transpose_apply, cos_neg, sin_neg, neg_sub]这个定理直观上说的是：旋转矩阵的转置就是反方向旋转。在RoPE实现里，query和key经常要做相反方向的旋转处理，这个定理给“换方向”提供了正式的数学依据。
随后可以写出attention场景下的重写结论：
(R(a)u)ᵀ (R(b)v) = uᵀ (R(a)ᵀ R(b)) v = uᵀ R(b - a) v这里的关键变换是，先把两个旋转矩阵合并，再把转置角度和正角度合并成相对角度。具体证明在Lean里可以由rope_transpose和rope_add组合推出，核心代码路径就是：
example (a b : ℝ) (u v : Fin 2 → ℝ) : dotProduct (ropeRot a *ᵥ u) (ropeRot b *ᵥ v) = dotProduct u (ropeRot (b - a) *ᵥ v) := by -- 先处理左边，展开点积和矩阵乘法 -- 通过 rope_transpose 和 rope_add 把 R(a)ᵀR(b) 化简为 R(b-a) -- 最后用点积结合律收尾我故意没有把最后这步写成完整脚本，因为到这一步，证明已经不再是机械的simp和ring，而是涉及dotProduct、mulVec以及转置关系在向量上的传播。留在项目里继续打磨更有意思，也从另一个角度说明：Lean4的证明项目不是一次性写完，而是越写越能发现之前没仔细想的细节。
5. 踩坑记录与排查思路
5.1 Windows环境问题速查
整理项目时，我把遇到的环境问题做了一个速查表，省得下次重装又从头折腾：
| 症状 | 原因 | 处理方式 |
|---|---|---|
lean命令不存在 | Elan路径未生效或未安装 | 重开终端；检查elan --version |
| VSCode没有Lean语法高亮 | 装了非官方扩展 | 卸载后安装发布者为leanprover的Lean4扩展 |
| 保存文件后一直显示building | 首次构建Mathlib依赖 | 耐心等，Mathlib很大；之后可考虑lake exe cache get获取预编译缓存 |
lake build报内存不足 | Mathlib构建占用较高 | 关闭其他大程序，或者只build当前文件相关依赖 |
| 旧版Lean和Elan冲突 | 两台环境历史遗留 | 先卸载旧Lean，清理PATH里的旧路径，再装Elan |
在实际操作中，最容易被忽略的是“VSCode扩展已经装好，但Lean版本选错”。Elan支持多版本并存，如果系统里有多个Lean工具链，VSCode会按项目配置选择版本。可以在Lean文件的状态栏看到当前使用的工具链版本，确认它和lean --version输出一致。
5.2 证明写不下去的通用排查方法
写证明最常遇到的状态是：目标明明是对的，但Lean就是不接受。新手很容易在这种时候怀疑人生。我总结了一套排查顺序，能解决绝大多数“证明卡住”的问题。
第一，用#check确认定理是否存在、类型是否正确。比如我想用sin_sq_add_cos_sq，先敲#check sin_sq_add_cos_sq，看到输出参数是θ : ℝ，才能在nlinarith里传入θ。
第二，把大目标拆小。不要指望一个simp解决所有事。可以先用have手动引入一个中间量，比如“先证明cos(a+b) = ...”，再把这个中间量用于主目标。这既能让Lean更容易处理，也能帮你自己理清思路。
第三，注意rw的方向。rw [cos_add]会把cos(a+b)改成展开式；反过来则需要rw [← cos_add]，也就是把展开式合并成cos(a+b)。方向一错，目标不仅没化简，反而越写越长。
第四，区分ring和nlinarith。ring处理多项式，nlinarith处理非线性不等式和包含额外条件的等式。三角恒等式既不是纯多项式，也不是线性条件，所以经常需要先simp/rw把三角函数重写掉，再把剩下的多项式交给ring或nlinarith。
5.3 Lean4在形式化验证生态中的位置
做完这个项目，我对Lean4的感觉是：它不是一个“证明生成器”，而是一个“证明检查器”。它不会替你想证明思路，但它会把你的思路彻底检查一遍，逼你在每个“显然”的地方补上依据。
Mathlib里已经有大量数学定理，覆盖了实数、矩阵、三角函数、线性代数等基础内容。这意味着RoPE这类工程数学问题，大部分前置结论都已经存在，你不需要从零搭建公理体系，而是像搭积木一样组合已有定理。
相比起用C++或者Python写验证脚本，Lean4的验证结论可以长期沉淀。你今天证明了rope_add，明天换了新项目，只要还是同样的数学结构，这个定理直接就能复用。对做框架、算子库、模型架构的人来说，这是一笔可以持续积累的资产。
结尾：把数学直觉变成可审计的证据
项目做到尾声，我最深的体会是：形式化验证真正改变的不是代码，而是我思考公式的方式。以前看RoPE的推导，我会说“这不就是把两个旋转矩阵合在一起吗”，显得理所当然；真的在Lean4里写证明时，才发现自己默认了很多东西，比如转置和乘法顺序怎么交换、R(a)作用于行向量还是列向量、角度b-a还是a-b。每一个“我以为显然”的地方，都是潜在的bug温床。
如果在读这篇的朋友也想尝试，我建议从rope_orthogonal开始，先把矩阵定义写出来，再把正交性证明跑通。那一步会给你足够的正反馈，也会让你理解simp、ring_nf、nlinarith这三板斧是怎么配合的。之后再去碰rope_add和更高阶的RoPE点积重写，就会发现Lean4并不是高不可攀，它只是比你我更有耐心、更不容忍模糊。
的实现)
