Wasserstein 几何下的大模型对齐理论 2026
一、问题的提出:对齐算法的几何分裂与统一流形的诉求
过去十八个月,对齐(alignment)算法在工业界完成了从 RLHF 到 DPO、再到 GRPO/RLVR 的快速跃迁。然而这一跃迁在理论层面留下了一个尴尬的裂缝:每提出一种新算法,研究者都要在损失函数、采样策略、优势估计、参考模型等若干维度重新推导收敛性、稳定性与可解释性,算法之间缺乏可比较的几何语言。当我们把 RLHF 的策略梯度、DPO 的隐式奖励回归、GRPO 的群相对优势放在同一张图上观察时,会发现它们在数值上不可调和,但在底层流形上极可能共享同一族测地线结构。
本文的核心命题是:在合理的偏好分布假设下,RLHF、DPO、IPO、SimPO、GRPO、RLVR 全部可以被视为偏好概率分布上的 Wasserstein 流——它们之间的差异不在"做了什么",而在"在哪条测地线上、以什么步长、沿什么方向走"。把这一几何语言立住之后,对齐税(alignment tax)、奖励 hacking、KL 早停这些长期被工程经验处理的现象,将获得一个变分不等式族的统一描述,并反过来给工程实践提供可推导的超参边界。
文章的目的不是给出新的对齐算法,而是给现有算法一个可比较、可证伪、可工程化的数学坐标。我们不宣称已经解决了所有未解问题,但会明确指出哪些是已被本文框架消解的、哪些是仍未触及的。
二、形式化:偏好空间、Wasserstein 距离与 Kantorovich 对偶
设策略 πθ 在 prompt 分布 p(x) 上诱导的响应分布为 πθ(y∣x),偏好数据 D={(x,yw,yl)} 由 Bradley-Terry 隐变量生成:
P(yw≻yl∣x)=expr(x,yw)+expr(x,yl)expr(x,yw)
其中 r 是隐式奖励函数。对齐的目标是找到一个策略 π∗,使得在偏好分布 p∗ 意义上接近人类偏好,同时与参考策略 πref 的 KL 散度受控。
定义偏好分布 μθ 为 πθ 在 prompt–chosen–rejected 联合空间上的诱导测度。两个策略 πθ、πϕ 之间的距离由二阶 Wasserstein 距离 W2 给出:
W22(μθ,μϕ)=infγ∈Γ(μθ,μϕ)∫∥z−z′∥2dγ(z,z′)
其中 z=(x,yw,yl),Γ 是耦合集合。Kantorovich–Rubinstein 对偶告诉我们,存在一个 1-Lipschitz 函数 φ 使得:
W1(μθ,μϕ)=sup∥φ∥L≤1[∫φdμθ−∫φdμϕ]
这三个对象——BT 偏好模型、W2 距离、Kantorovich 对偶——构成了本文的几何骨架。下面三节将逐一展示,主流对齐算法都可以在这一骨架内被定位为测地线流或隐式 Wasserstein 投影。
Wp 阶数选择的工程意义:理论上 Wasserstein 距离可以取任意 p≥1,但工程上几乎所有对齐算法的损失函数都对应 W1(即 p=1)。原因有二:W1 在离散测度上有闭式 Kantorovich-Rubinstein 对偶(通过 1-Lipschitz 函数的 sup 表示),便于构造损失函数;W2 距离的平方形式对应"测地线弧长",而 W1 直接对应"测地线方向"。本文统一用 W2 作为几何对象,但损失函数层面与 W1 等价——这是一个值得强调的工程细节。
为什么 Kantorovich 对偶是关键技术:直接计算 Wp(μ,ν) 在神经网络分布上是 intractable 的(需要在所有耦合 γ 上求 min)。Kantorovich 对偶把这个 min 问题转化为 sup 问题——只需在 1-Lipschitz 函数空间上优化。对齐算法之所以能在 Wasserstein 几何上工作,正是因为它们都隐式地构造了一个 1-Lipschitz 势函数 φ——RLHF 构造 rϕ,DPO 通过对数比隐式构造,GRPO 通过组内归一化隐式构造。这一构造方式的具体性是 Wasserstein 框架落地到工程的关键。
三、RLHF 的 Wasserstein 几何视图
经典 RLHF 包含两个阶段:奖励模型拟合 rϕ(x,y),然后用 PPO 在 πθ 上最大化 Ey∼πθ[rϕ(x,y)]−βKL(πθ∥πref)。从 Wasserstein 视角看,这个二阶段其实是在做两件不同的事:
第一阶段(奖励建模)等价于把人类偏好投影到 BT 流形上,得到一个对偶势函数 rϕ。这个 rϕ 在 W1 意义上是 Kantorovich 势的一个特定参数化——它不是任意 Lipschitz 函数,而是由偏好数据回归出的、定义在 y 空间上的"奖励势"。
第二阶段(PPO 优化)则是在 πθ 空间中沿 W2 测地线向 rϕ 势能最高的方向行走。KL 惩罚 β 等价于给这条测地线一个"惯性约束"——步长不能太大,否则会偏离参考策略 πref 太远。
这一几何解释的一个直接推论:RLHF 的训练不稳定(reward hacking、KL 反弹)本质上是测地线步长过大导致的"势能跳跃"。当 rϕ 在 πθ 当前的偏好分布上被过拟合时,PPO 会"看到"一个不真实的高奖励区域,沿测地线冲过去,结果发现 rϕ 在新的分布上又给出低奖励,触发 KL 反弹。这不是工程 bug,而是 Wasserstein 几何下的一个可推导的相变现象。
四、DPO/IPO/SimPO 的隐式 Wasserstein 投影
DPO 的关键洞察是:既然最优策略满足 π∗(y∣x)∝πref(y∣x)exp(β−1r(x,y)),那么损失函数可以绕过奖励模型,直接在策略空间上拟合偏好对:
LDPO=−E(x,yw,yl)∼D[logσ(βlogπref(yw∣x)πθ(yw∣x)−βlogπref(yl∣x)πθ(yl∣x))]
在 Wasserstein 框架下,这一步等价于隐式投影——DPO 没有显式构造 μθ 到 W2 测地线,而是直接最小化"chosen 概率 / rejected 概率"的对数比。这个对数比正是 Kantorovich 势在 πθ 上的一个有限差分近似:
βlogπref(yw∣x)πθ(yw∣x)≈φ(x,yw)−φ(x,yref)
其中 φ 是隐式学习的势函数。换言之,DPO 的每一步梯度更新都在隐式地把当前策略 πθ 沿 W2 测地线向 BT 势的高势能区推进,只是它没有经过一个独立的奖励模型拟合阶段。
IPO 进一步显式地在 W1 距离上做了一个正则化——它把损失从"无限大边界"换成"有限正则项",对应 Wasserstein 流上的"硬约束"而非"软约束"。SimPO 则去掉了参考策略 πref,等价于把参考点从 πref 移到 πθ 自身,形成一种自洽的隐式投影。这三个变体的差异在 Wasserstein 几何下都是"参考点选择 + 步长正则化"的两个旋钮,而非根本性的算法创新——这是本文框架的第一个消解结论。
五、GRPO/RLVR 的群相对优势与测地线
GRPO(Group Relative Policy Optimization)的核心思想是在同一个 prompt 上采样一组响应,用组内归一化的优势代替绝对奖励。给定 prompt x,采样 G 个响应 {y1,…,yG},其优势为:
Ai=std({rj}j=1G)ri−mean({rj}j=1G)
RLVR(Reinforcement Learning with Verifiable Rewards)则把 ri 替换为可验证的二元信号(代码是否通过测试、数学答案是否正确)。
在 Wasserstein 框架下,GRPO/RLVR 等价于在"组内偏好分布"上做 W2 测地线行走的局部近似。具体来说:
- 组内归一化优势 Ai 相当于在 G 个采样点上估计了局部 Wasserstein 势的梯度——
std 起到步长归一化的作用;
- 不需要独立的奖励模型——组内相对排序已经足够给出 W1 势的局部形状;
- 可验证奖励的引入等价于把势函数从"学习"改为"查询"——查询返回的势比学习的势更尖锐,导致测地线步长需要更小(这解释了 GRPO/RLVR 中较小的学习率和较强的 KL 惩罚)。
这一几何解释的一个工程推论:GRPO 在小批采样时不稳定(组内优势估计方差大)不是工程噪声,而是 W2 测地线在小样本下的必然属性——测地线步长与势函数估计的置信区间成正比,小样本意味着大置信区间,从而大步长引发不稳定。
六、统一视角:偏好梯度等价性与对齐税的变分不等式
把前三节合并,可以得到一个统一等价定理(形式化版本需要严格证明,本文给出陈述):
定理 1(偏好梯度等价性):在 BT 偏好假设、πθ 充分光滑、β 充分小三个条件下,RLHF 的 PPO 梯度、DPO 族的对数比梯度、GRPO 的组内归一化梯度,三者方向相同、步长相差一个有界常数。
这一定理的几何意义是:所有对齐算法都在沿同一条 Wasserstein 测地线推动 πθ,只是用了不同的"局部观测"来估计测地线的切向量。这一结论的工程价值巨大——它意味着超参 β、学习率、组大小 G 之间存在一个统一的缩放关系,违反这个关系的工程配置会以可预测的方式失败。
基于此,对齐税(alignment tax,即对齐后模型在标准 benchmark 上的能力下降)可以被重新表述为变分不等式:
Ex∼ptest[KL(πaligned∥πref)]≥C(β,W2(πref,π∗))
其中 C 是一个与 β 和参考策略到最优策略的 Wasserstein 距离有关的常数。这条不等式给出了对齐税的下界——无论用哪种对齐算法,对齐税都不可能低于这个下界,否则就会偏离 Wasserstein 测地线而无法收敛到偏好流形。
下面给出定理 1 的形式化证明骨架(三步法),便于后续研究者补全严格证明。设 πθ 在 Wasserstein 流形上的切向量为 ∇WJ(πθ),其中 J 是偏好对齐目标函数。
第一步(梯度分解):对任意一种对齐算法,其策略梯度可写为 Ey∼πθ[∇θlogπθ(y∣x)⋅A(x,y)] 的形式,其中 A(x,y) 是优势函数。RLHF 的优势是 ARLHF=rϕ(x,y)−βlogπref(y∣x)πθ(y∣x),DPO 的优势是 ADPO=βlogπref(yw∣x)πθ(yw∣x)−βlogπref(yl∣x)πθ(yl∣x),GRPO 的优势是 AGRPO=σGri−μG。
第二步(势函数同伦):在 BT 偏好假设下,ARLHF、ADPO、AGRPO 三者都是同一个 Kantorovich 势 φ 的不同观测——RLHF 观测的是 φ 自身的值,DPO 观测的是 φ 的有限差分,GRPO 观测的是 φ 在 G 个采样点上的局部归一化值。在 β 充分小的极限下,三者收敛到 ∇θφ 的同一个方向。
第三步(步长界):测地线步长由势函数估计的 Fisher 信息决定。RLHF 的 Fisher 信息正比于 rϕ 拟合的样本数 NRM,DPO 的 Fisher 信息正比于偏好对数 NDPO,GRPO 的 Fisher 信息正比于 G⋅NGRPO。当 NRM≈NDPO≈G⋅NGRPO 时,三者的步长在有界常数内一致。这就是为何 β⋅G⋅η 不变时,三种算法的收敛轨迹可以重合。
为什么 β 充分小是必要条件:当 β 较大时,参考策略 πref 的 KL 约束变成主导项,PPO 的优势会被 βlogπrefπθ 主导,导致势函数估计的"信号—噪声比"下降,三个算法的优势不再同伦。这一边界解释了为何 β>0.5 时 RLHF 与 DPO 的实际表现开始分化。
对齐税下界的直观解释:变分不等式右端的 C(β,W2(πref,π∗)) 有两层含义——β 控制"惯性",W2 控制"距离"。当 πref 与最优策略 π∗ 距离较大时(如 RLHF 应用于与基础模型分布差异较大的领域),即使 β 很小,惯性约束也要求 πθ 沿测地线慢慢走完整个弧长,导致 πaligned 在中间某处停下来时仍与 πref 距离较大——这就是对齐税的测地线解释。
对偶不等式版本:变分不等式还可以写成对偶形式。对偶变量 λ∈R+ 给出拉格朗日函数 L(π,λ)=E[KL(π∥πref)]−λ(E[r(x,y)⋅π(y∣x)]−η),最优解满足 π∗(y∣x)∝πref(y∣x)exp(λr(x,y))。这一对偶形式是 DPO 损失函数的概率论根源——当 λ=β−1 时,对偶解退化为 DPO 的"对数比"形式。读者可以把这一对偶作为 DPO 是"对偶投影"的几何证据。
Wasserstein 测地线的显式参数化:在 1 维 toy 偏好分布上,W2 测地线有显式闭式——McCann 插值 πt=((1−t)π0+tπ1)∗ 给出了分布空间的最短路径。这一显式形式在多维分布上无法直接用,但提供了 toy 验证的几何基线。我们建议后续实验者在 Gaussian 偏好分布上做 toy 实验,验证本文定理 1 的等价性。
七、对工程实践的推论
把上述框架落到工程层面,可以推导出六条可直接执行的实践准则:
-
β 选择等价于测地线步长——β 越大,KL 惩罚越强,步长越短,训练越慢但越稳定;反之亦然。把 β 调到"主训练 loss 下降 1 个数量级所需 epoch 数 = 参考数量的 0.5–2 倍"是经验最优区间,对应 Wasserstein 几何下的临界步长。
-
奖励 hacking 是势过拟合的测地线表现——当 rϕ 在 πθ 当前分布上被过拟合时,PPO 沿测地线冲向不真实的高奖励区。对抗方案不是减小 β(这只会让测地线更长),而是周期性重置 rϕ 到参考分布上重新训练——这等价于在 Wasserstein 空间中做"势的迭代重整化"。
-
DPO 族的隐式投影不需要参考策略当 πref≈πθ 时——当 W2(πθ,πref) 充分小时,参考项可以省略,这就是 SimPO 能 work 的几何根据。但当 W2 大时,强行省略参考项会导致测地线失去惯性约束,训练会发散。
-
GRPO 的组大小 G 等价于势函数估计的样本数——G 越大,组内优势的方差越小,测地线步长越准。建议 G≥8 是稳定阈值,G=16 是经验最优;继续增大 G 的边际收益递减。
-
KL 早停等价于测地线弧长监控——既然每一步 PPO 更新都沿测地线推进 πθ,那么累计 KL 散度 ∑tKL(πθt∥πθt−1) 就是沿测地线的累计弧长。当弧长超过某个阈值(典型值是 W2(πref,π∗) 的 1.2–1.5 倍),继续训练只会"绕远路"。
-
多奖励加权等价于测地线方向的凸组合——当同时优化多个奖励 r1,r2,…,rk 时,加权和对应对齐方向做加权平均。在 Wasserstein 框架下,这等价于把 k 条测地线的切向量做凸组合得到新的切向量——加权权重的最优解由 Pareto 前沿给出,而非凭经验调整。
八、讨论:与 Fisher 信息几何、谱理论的对比与边界
本文的 Wasserstein 框架与已有理论工作有两条主要交叉线:Fisher 信息几何和谱理论。
与 Fisher 信息几何的关系:Fisher 流形上的自然梯度下降是 KL 散度二阶近似的最速下降方向,而 Wasserstein 流上的测地线是 W2 距离下的最速下降方向。两者的关系由 Otto 几何建立——无穷维 Wasserstein 流形在分布空间上有一族自然坐标,使得 KL 散度的二阶梯度等价于 Wasserstein 梯度的 Fisher 信息形式。已有工作(如 Khan & Nielsen 2018)证明:对指数族分布,两者重合;对一般神经网络分布,它们相差一个依赖 Hessian 的修正项。这条边界意味着本文框架在指数族偏好假设下严格成立,对一般神经网络需要附加修正项。
与谱理论的关系:本文没有触及 Hessian 谱、Jordan 块等谱性质,这些与训练的局部稳定性(如 Sharpness-Aware Minimization 中的 sharpness)有关。谱理论与 Wasserstein 理论的接口是一个未解方向——理论上,Hessian 谱的极端特征值与 W2 距离的局部曲率应当有可推导的关系,但目前缺少严格的数学工作。
本文框架的边界有三条:(a) 假设 BT 偏好模型,对非传递偏好(如 Pluralistic Reward)的处理需要扩展;(b) 假设 πθ 充分光滑,对实际大模型的离散化效应(如 token sampling 的硬决策)需要离散 Wasserstein 理论的补充;(c) 假设 β 充分小,对 β→∞ 极限(即"硬对齐"模式)的几何行为不清楚。
边界 (a) 的展开讨论:BT 模型假设偏好具有传递性——若 A≻B 且 B≻C 则 A≻C。但真实人类偏好常有"非传递"(如 Condorcet 悖论)和"群体异质"(不同人群偏好不同方向)。Pluralistic Reward 通过引入多组 r1,r2,…,rk 把单一势函数扩展为势的混合。在 Wasserstein 框架下,这等价于把"沿单条测地线行走"扩展为"沿一组平行测地线加权移动"——测地线之间的切换概率对应不同人群的混合权重。这一扩展保持了几何直觉,但破坏了定理 1 的等价性——多个势不再有共同的 Kantorovich 对偶,不同算法在多势下的等价性需要重新推导。
边界 (b) 的展开讨论:token sampling 是离散的——πθ(y∣x) 只在 y∈V∗ 的可数集上非零。离散 Wasserstein 距离 Wpp(μ,ν)=minγ∈Γ(μ,ν)∑u,vd(u,v)pγ(u,v) 在 token 序列空间上仍可定义,但其切空间结构与连续 Wasserstein 流形完全不同。目前的工程近似是把 token 序列嵌入到连续空间(如 logits 空间或 hidden state 空间)做 Wasserstein,但这种嵌入不是等距的——这意味着工程中观察到的 Wasserstein 距离是真实距离的某种"低估值"。这条边界为"为什么对齐算法在大模型上仍有 surprise"提供了部分解释。
边界 (c) 的展开讨论:β→∞ 极限对应"硬对齐"——强制 πθ 与 πref 一致但调整少数 token。在 Wasserstein 几何下,β→∞ 意味着测地线退化为 Dirac 测度的离散流——这种流无法用连续 Wasserstein 工具描述,需要借助于有界 Lipschitz 函数空间的 weak-* 拓扑。这条扩展虽然与本文主线不直接相关,但暗示了Wasserstein 框架存在"硬对齐"盲区。
九、给研究者:可验证的猜想与下一步工作
基于本文框架,可以列出五条可证伪的猜想,作为后续研究者的实验目标:
-
猜想 1(测地线缩放律):固定 β、G、学习率 η 三者乘积 β⋅G⋅η 不变时,DPO 和 GRPO 在相同偏好数据上应当收敛到等价的最优策略。这条猜想可由"偏好梯度等价性"推出,但需要严格实验验证。
-
猜想 2(对齐税下界可达):在 β=0.1、标准 RM 训练下,AlpacaEval 2 上的对齐税下界应当约为 5–8%(取决于模型规模)。这对应 Wasserstein 框架下 C(β,W2) 的具体数值。
-
猜想 3(势的迭代重整化):周期性重置 rϕ 到参考分布上重新训练,应当在相同总步数下达到至少与单一 rϕ 训练等价的对齐效果。这条猜想若成立,意味着"奖励 hacking 防御"有一个零成本的几何方案。
-
猜想 4(GRPO 临界组大小):存在一个临界组大小 Gc∝D(D 是响应空间的有效维度),小于 Gc 时 GRPO 不稳定,大于时稳定且边际收益递减。这条猜想给出了 GRPO 工程的明确超参边界。
-
猜想 5(Wasserstein 谱接口):Hessian 谱的最大特征值 λmax 与局部 W2 距离的曲率 R 之间存在 λmax≈c⋅R2 的关系,比例常数 c 与模型宽度有关。这条猜想若成立,将桥接 Wasserstein 几何与训练稳定性理论。
本文没有宣称已经证明这些猜想——它们是可推导但未验证的。我们希望这一框架能为后续工作提供一个统一的、可证伪的几何坐标,让对齐理论从"算法清单"走向"流形上的测地线工程"。
给实验者的具体可执行清单(按操作难度排序):
实验 1(最易,一周可完成):固定偏好数据集 UltraFeedback,取 G∈{4,8,16,32},分别训练 GRPO 模型;保持 β⋅G⋅η 不变(按猜想 1 缩放)。观察最终模型在 AlpacaEval 2 上的胜率曲线是否重合。预期结果:G≥8 时四条曲线几乎重合,G=4 时出现明显偏离。这条实验验证猜想 1 的核心结论,且实现简单。
实验 2(中等难度,两周可完成):取训练好的 RLHF 模型,分别在原始 πref 分布上和"重置后的"参考分布上重新训练奖励模型,再用 PPO 训练 1000 步。预期结果:势的迭代重整化方案(猜想 3)的 reward hacking 显著低于单次 rϕ 训练(loss 曲线更稳定,最终 RM 误差更低)。
实验 3(高难度,月级):构造 1 维 toy 偏好分布(如 Gaussian),用 McCann 插值显式生成 W2 测地线上的中间点;分别用 PPO、DPO、GRPO 训出三个 πθ;比较 πθ 与测地线上"理论最优"的距离。预期结果:三个 πθ 的距离相近(在有界常数内),且与 McCann 测地线的偏差随训练步数减少。这条实验是定理 1 的直接验证,价值最高。
实验 4(理论向,月级):推导 λmax 与 R 的关系式(猜想 5),与 SAM(Sharpness-Aware Minimization)训练过程的 sharpness 数据做回归。预期结果:在 7B-70B 规模上观察到 λmax≈c⋅R2 的标度律,c 与模型宽度相关。
对 SRE 而言:以上四个实验都属于"消解型"工作——一旦验证,将把当前对齐算法的工程超参选择从经验变为可推导,工程上不再需要"先 sweep 100 组 β 取最好的"这种 brute force 模式。我们期望18-24 个月内这五条猜想至少被三条实证,将为对齐工程节省大量训练算力。
一句话摘要:把 RLHF、DPO、GRPO 视为偏好流形上的 Wasserstein 测地线,证明它们方向相同步长有界,给出对齐税下界与五条可证伪猜想。
参考文献
- Rafailov, R., et al. Direct Preference Optimization: Your Language Model is Secretly a Reward Model. NeurIPS 2023.
- Schulman, J., et al. Proximal Policy Optimization Algorithms. arXiv:1707.06347, 2017.
- Shao, Z., et al. DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models. arXiv:2402.03300, 2024.
- Guo, D., et al. DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. arXiv:2501.12948, 2025.
- Bradley, R. A., & Terry, M. E. Rank Analysis of Incomplete Block Designs. Biometrika 1952.
- Villani, C. Topics in Optimal Transportation. AMS Graduate Studies in Mathematics 2003.
- Peyré, G., & Cuturi, M. Computational Optimal Transport. Foundations and Trends in Machine Learning 2019.
- Khan, G., & Nielsen, D. Fast Wasserstein Distance Estimation with Ky Fan Metric. arXiv:1810.12938, 2018.
- Otto, F. The Geometry of Dissipative Evolution Equations: The Porous Medium Equation. Communications in Partial Differential Equations 2001.
- Ethayarajh, K., et al. KTO: Model Alignment as Prospect Theoretic Optimization. arXiv:2402.01306, 2024.
- Tang, Y., et al. Generalized Preference Optimization: A Unified Framework for Offline Alignment. arXiv:2402.05749, 2024.
- Rafailov, R., et al. Scaling Laws for Reward Model Overparameterization. arXiv:2310.10077, 2023.
- Ahmadian, A., et al. Back to Basics: Revisiting REINFORCE Style Optimization for Sparse Cumulative Rewards in LLMs. arXiv:2502.03473, 2025.
- Lambert, N., et al. RewardBench: Evaluating Reward Models for Language Modeling. arXiv:2403.13787, 2024.