Wasserstein 几何下的大模型对齐理论 2026
约 30 分钟8789 字0 次阅读

Wasserstein 几何下的大模型对齐理论 2026
一、问题的提出:对齐算法的几何分裂与统一流形的诉求
过去十八个月,对齐(alignment)算法在工业界完成了从 RLHF 到 DPO、再到 GRPO/RLVR 的快速跃迁。然而这一跃迁在理论层面留下了一个尴尬的裂缝:每提出一种新算法,研究者都要在损失函数、采样策略、优势估计、参考模型等若干维度重新推导收敛性、稳定性与可解释性,算法之间缺乏可比较的几何语言。当我们把 RLHF 的策略梯度、DPO 的隐式奖励回归、GRPO 的群相对优势放在同一张图上观察时,会发现它们在数值上不可调和,但在底层流形上极可能共享同一族测地线结构。
本文的核心命题是:在合理的偏好分布假设下,RLHF、DPO、IPO、SimPO、GRPO、RLVR 全部可以被视为偏好概率分布上的 Wasserstein 流——它们之间的差异不在"做了什么",而在"在哪条测地线上、以什么步长、沿什么方向走"。把这一几何语言立住之后,对齐税(alignment tax)、奖励 hacking、KL 早停这些长期被工程经验处理的现象,将获得一个变分不等式族的统一描述,并反过来给工程实践提供可推导的超参边界。
文章的目的不是给出新的对齐算法,而是给现有算法一个可比较、可证伪、可工程化的数学坐标。我们不宣称已经解决了所有未解问题,但会明确指出哪些是已被本文框架消解的、哪些是仍未触及的。
二、形式化:偏好空间、Wasserstein 距离与 Kantorovich 对偶
设策略 在 prompt 分布 上诱导的响应分布为 ,偏好数据 由 Bradley-Terry 隐变量生成:
其中 是隐式奖励函数。对齐的目标是找到一个策略 ,使得在偏好分布 意义上接近人类偏好,同时与参考策略 的 KL 散度受控。
定义偏好分布 为 在 prompt–chosen–rejected 联合空间上的诱导测度。两个策略 、 之间的距离由二阶 Wasserstein 距离 给出:
其中 , 是耦合集合。Kantorovich–Rubinstein 对偶告诉我们,存在一个 1-Lipschitz 函数 使得:
这三个对象——BT 偏好模型、 距离、Kantorovich 对偶——构成了本文的几何骨架。下面三节将逐一展示,主流对齐算法都可以在这一骨架内被定位为测地线流或隐式 Wasserstein 投影。
阶数选择的工程意义:理论上 Wasserstein 距离可以取任意 ,但工程上几乎所有对齐算法的损失函数都对应 (即 )。原因有二: 在离散测度上有闭式 Kantorovich-Rubinstein 对偶(通过 1-Lipschitz 函数的 sup 表示),便于构造损失函数; 距离的平方形式对应"测地线弧长",而 直接对应"测地线方向"。本文统一用 作为几何对象,但损失函数层面与 等价——这是一个值得强调的工程细节。
为什么 Kantorovich 对偶是关键技术:直接计算 在神经网络分布上是 intractable 的(需要在所有耦合 上求 min)。Kantorovich 对偶把这个 min 问题转化为 问题——只需在 1-Lipschitz 函数空间上优化。对齐算法之所以能在 Wasserstein 几何上工作,正是因为它们都隐式地构造了一个 1-Lipschitz 势函数 ——RLHF 构造 ,DPO 通过对数比隐式构造,GRPO 通过组内归一化隐式构造。这一构造方式的具体性是 Wasserstein 框架落地到工程的关键。
三、RLHF 的 Wasserstein 几何视图
经典 RLHF 包含两个阶段:奖励模型拟合 ,然后用 PPO 在 上最大化 。从 Wasserstein 视角看,这个二阶段其实是在做两件不同的事:
第一阶段(奖励建模)等价于把人类偏好投影到 BT 流形上,得到一个对偶势函数 。这个 在 意义上是 Kantorovich 势的一个特定参数化——它不是任意 Lipschitz 函数,而是由偏好数据回归出的、定义在 空间上的"奖励势"。
第二阶段(PPO 优化)则是在 空间中沿 测地线向 势能最高的方向行走。KL 惩罚 等价于给这条测地线一个"惯性约束"——步长不能太大,否则会偏离参考策略 太远。
这一几何解释的一个直接推论:RLHF 的训练不稳定(reward hacking、KL 反弹)本质上是测地线步长过大导致的"势能跳跃"。当 在 当前的偏好分布上被过拟合时,PPO 会"看到"一个不真实的高奖励区域,沿测地线冲过去,结果发现 在新的分布上又给出低奖励,触发 KL 反弹。这不是工程 bug,而是 Wasserstein 几何下的一个可推导的相变现象。
四、DPO/IPO/SimPO 的隐式 Wasserstein 投影
DPO 的关键洞察是:既然最优策略满足 ,那么损失函数可以绕过奖励模型,直接在策略空间上拟合偏好对:
在 Wasserstein 框架下,这一步等价于隐式投影——DPO 没有显式构造 到 测地线,而是直接最小化"chosen 概率 / rejected 概率"的对数比。这个对数比正是 Kantorovich 势在 上的一个有限差分近似:
其中 是隐式学习的势函数。换言之,DPO 的每一步梯度更新都在隐式地把当前策略 沿 测地线向 BT 势的高势能区推进,只是它没有经过一个独立的奖励模型拟合阶段。
IPO 进一步显式地在 距离上做了一个正则化——它把损失从"无限大边界"换成"有限正则项",对应 Wasserstein 流上的"硬约束"而非"软约束"。SimPO 则去掉了参考策略 ,等价于把参考点从 移到 自身,形成一种自洽的隐式投影。这三个变体的差异在 Wasserstein 几何下都是"参考点选择 + 步长正则化"的两个旋钮,而非根本性的算法创新——这是本文框架的第一个消解结论。
五、GRPO/RLVR 的群相对优势与测地线
GRPO(Group Relative Policy Optimization)的核心思想是在同一个 prompt 上采样一组响应,用组内归一化的优势代替绝对奖励。给定 prompt ,采样 个响应 ,其优势为:
RLVR(Reinforcement Learning with Verifiable Rewards)则把 替换为可验证的二元信号(代码是否通过测试、数学答案是否正确)。
在 Wasserstein 框架下,GRPO/RLVR 等价于在"组内偏好分布"上做 测地线行走的局部近似。具体来说:
- 组内归一化优势 相当于在 个采样点上估计了局部 Wasserstein 势的梯度——
std起到步长归一化的作用; - 不需要独立的奖励模型——组内相对排序已经足够给出 势的局部形状;
- 可验证奖励的引入等价于把势函数从"学习"改为"查询"——查询返回的势比学习的势更尖锐,导致测地线步长需要更小(这解释了 GRPO/RLVR 中较小的学习率和较强的 KL 惩罚)。
这一几何解释的一个工程推论:GRPO 在小批采样时不稳定(组内优势估计方差大)不是工程噪声,而是 测地线在小样本下的必然属性——测地线步长与势函数估计的置信区间成正比,小样本意味着大置信区间,从而大步长引发不稳定。
六、统一视角:偏好梯度等价性与对齐税的变分不等式
把前三节合并,可以得到一个统一等价定理(形式化版本需要严格证明,本文给出陈述):
定理 1(偏好梯度等价性):在 BT 偏好假设、 充分光滑、 充分小三个条件下,RLHF 的 PPO 梯度、DPO 族的对数比梯度、GRPO 的组内归一化梯度,三者方向相同、步长相差一个有界常数。
这一定理的几何意义是:所有对齐算法都在沿同一条 Wasserstein 测地线推动 ,只是用了不同的"局部观测"来估计测地线的切向量。这一结论的工程价值巨大——它意味着超参 、学习率、组大小 之间存在一个统一的缩放关系,违反这个关系的工程配置会以可预测的方式失败。
基于此,对齐税(alignment tax,即对齐后模型在标准 benchmark 上的能力下降)可以被重新表述为变分不等式:
其中 是一个与 和参考策略到最优策略的 Wasserstein 距离有关的常数。这条不等式给出了对齐税的下界——无论用哪种对齐算法,对齐税都不可能低于这个下界,否则就会偏离 Wasserstein 测地线而无法收敛到偏好流形。
下面给出定理 1 的形式化证明骨架(三步法),便于后续研究者补全严格证明。设 在 Wasserstein 流形上的切向量为 ,其中 是偏好对齐目标函数。
第一步(梯度分解):对任意一种对齐算法,其策略梯度可写为 的形式,其中 是优势函数。RLHF 的优势是 ,DPO 的优势是 ,GRPO 的优势是 。
第二步(势函数同伦):在 BT 偏好假设下,、、 三者都是同一个 Kantorovich 势 的不同观测——RLHF 观测的是 自身的值,DPO 观测的是 的有限差分,GRPO 观测的是 在 个采样点上的局部归一化值。在 充分小的极限下,三者收敛到 的同一个方向。
第三步(步长界):测地线步长由势函数估计的 Fisher 信息决定。RLHF 的 Fisher 信息正比于 拟合的样本数 ,DPO 的 Fisher 信息正比于偏好对数 ,GRPO 的 Fisher 信息正比于 。当 时,三者的步长在有界常数内一致。这就是为何 不变时,三种算法的收敛轨迹可以重合。
为什么 充分小是必要条件:当 较大时,参考策略 的 KL 约束变成主导项,PPO 的优势会被 主导,导致势函数估计的"信号—噪声比"下降,三个算法的优势不再同伦。这一边界解释了为何 时 RLHF 与 DPO 的实际表现开始分化。
对齐税下界的直观解释:变分不等式右端的 有两层含义—— 控制"惯性", 控制"距离"。当 与最优策略 距离较大时(如 RLHF 应用于与基础模型分布差异较大的领域),即使 很小,惯性约束也要求 沿测地线慢慢走完整个弧长,导致 在中间某处停下来时仍与 距离较大——这就是对齐税的测地线解释。
对偶不等式版本:变分不等式还可以写成对偶形式。对偶变量 给出拉格朗日函数 ,最优解满足 。这一对偶形式是 DPO 损失函数的概率论根源——当 时,对偶解退化为 DPO 的"对数比"形式。读者可以把这一对偶作为 DPO 是"对偶投影"的几何证据。
Wasserstein 测地线的显式参数化:在 1 维 toy 偏好分布上, 测地线有显式闭式——McCann 插值 给出了分布空间的最短路径。这一显式形式在多维分布上无法直接用,但提供了 toy 验证的几何基线。我们建议后续实验者在 Gaussian 偏好分布上做 toy 实验,验证本文定理 1 的等价性。
七、对工程实践的推论
把上述框架落到工程层面,可以推导出六条可直接执行的实践准则:
-
选择等价于测地线步长—— 越大,KL 惩罚越强,步长越短,训练越慢但越稳定;反之亦然。把 调到"主训练 loss 下降 1 个数量级所需 epoch 数 = 参考数量的 0.5–2 倍"是经验最优区间,对应 Wasserstein 几何下的临界步长。
-
奖励 hacking 是势过拟合的测地线表现——当 在 当前分布上被过拟合时,PPO 沿测地线冲向不真实的高奖励区。对抗方案不是减小 (这只会让测地线更长),而是周期性重置 到参考分布上重新训练——这等价于在 Wasserstein 空间中做"势的迭代重整化"。
-
DPO 族的隐式投影不需要参考策略当 时——当 充分小时,参考项可以省略,这就是 SimPO 能 work 的几何根据。但当 大时,强行省略参考项会导致测地线失去惯性约束,训练会发散。
-
GRPO 的组大小 等价于势函数估计的样本数—— 越大,组内优势的方差越小,测地线步长越准。建议 是稳定阈值, 是经验最优;继续增大 的边际收益递减。
-
KL 早停等价于测地线弧长监控——既然每一步 PPO 更新都沿测地线推进 ,那么累计 KL 散度 就是沿测地线的累计弧长。当弧长超过某个阈值(典型值是 的 1.2–1.5 倍),继续训练只会"绕远路"。
-
多奖励加权等价于测地线方向的凸组合——当同时优化多个奖励 时,加权和对应对齐方向做加权平均。在 Wasserstein 框架下,这等价于把 条测地线的切向量做凸组合得到新的切向量——加权权重的最优解由 Pareto 前沿给出,而非凭经验调整。
八、讨论:与 Fisher 信息几何、谱理论的对比与边界
本文的 Wasserstein 框架与已有理论工作有两条主要交叉线:Fisher 信息几何和谱理论。
与 Fisher 信息几何的关系:Fisher 流形上的自然梯度下降是 KL 散度二阶近似的最速下降方向,而 Wasserstein 流上的测地线是 距离下的最速下降方向。两者的关系由 Otto 几何建立——无穷维 Wasserstein 流形在分布空间上有一族自然坐标,使得 KL 散度的二阶梯度等价于 Wasserstein 梯度的 Fisher 信息形式。已有工作(如 Khan & Nielsen 2018)证明:对指数族分布,两者重合;对一般神经网络分布,它们相差一个依赖 Hessian 的修正项。这条边界意味着本文框架在指数族偏好假设下严格成立,对一般神经网络需要附加修正项。
与谱理论的关系:本文没有触及 Hessian 谱、Jordan 块等谱性质,这些与训练的局部稳定性(如 Sharpness-Aware Minimization 中的 sharpness)有关。谱理论与 Wasserstein 理论的接口是一个未解方向——理论上,Hessian 谱的极端特征值与 距离的局部曲率应当有可推导的关系,但目前缺少严格的数学工作。
本文框架的边界有三条:(a) 假设 BT 偏好模型,对非传递偏好(如 Pluralistic Reward)的处理需要扩展;(b) 假设 充分光滑,对实际大模型的离散化效应(如 token sampling 的硬决策)需要离散 Wasserstein 理论的补充;(c) 假设 充分小,对 极限(即"硬对齐"模式)的几何行为不清楚。
边界 (a) 的展开讨论:BT 模型假设偏好具有传递性——若 且 则 。但真实人类偏好常有"非传递"(如 Condorcet 悖论)和"群体异质"(不同人群偏好不同方向)。Pluralistic Reward 通过引入多组 把单一势函数扩展为势的混合。在 Wasserstein 框架下,这等价于把"沿单条测地线行走"扩展为"沿一组平行测地线加权移动"——测地线之间的切换概率对应不同人群的混合权重。这一扩展保持了几何直觉,但破坏了定理 1 的等价性——多个势不再有共同的 Kantorovich 对偶,不同算法在多势下的等价性需要重新推导。
边界 (b) 的展开讨论:token sampling 是离散的—— 只在 的可数集上非零。离散 Wasserstein 距离 在 token 序列空间上仍可定义,但其切空间结构与连续 Wasserstein 流形完全不同。目前的工程近似是把 token 序列嵌入到连续空间(如 logits 空间或 hidden state 空间)做 Wasserstein,但这种嵌入不是等距的——这意味着工程中观察到的 Wasserstein 距离是真实距离的某种"低估值"。这条边界为"为什么对齐算法在大模型上仍有 surprise"提供了部分解释。
边界 (c) 的展开讨论: 极限对应"硬对齐"——强制 与 一致但调整少数 token。在 Wasserstein 几何下, 意味着测地线退化为 Dirac 测度的离散流——这种流无法用连续 Wasserstein 工具描述,需要借助于有界 Lipschitz 函数空间的 weak-* 拓扑。这条扩展虽然与本文主线不直接相关,但暗示了Wasserstein 框架存在"硬对齐"盲区。
九、给研究者:可验证的猜想与下一步工作
基于本文框架,可以列出五条可证伪的猜想,作为后续研究者的实验目标:
-
猜想 1(测地线缩放律):固定 、、学习率 三者乘积 不变时,DPO 和 GRPO 在相同偏好数据上应当收敛到等价的最优策略。这条猜想可由"偏好梯度等价性"推出,但需要严格实验验证。
-
猜想 2(对齐税下界可达):在 、标准 RM 训练下,AlpacaEval 2 上的对齐税下界应当约为 5–8%(取决于模型规模)。这对应 Wasserstein 框架下 的具体数值。
-
猜想 3(势的迭代重整化):周期性重置 到参考分布上重新训练,应当在相同总步数下达到至少与单一 训练等价的对齐效果。这条猜想若成立,意味着"奖励 hacking 防御"有一个零成本的几何方案。
-
猜想 4(GRPO 临界组大小):存在一个临界组大小 ( 是响应空间的有效维度),小于 时 GRPO 不稳定,大于时稳定且边际收益递减。这条猜想给出了 GRPO 工程的明确超参边界。
-
猜想 5(Wasserstein 谱接口):Hessian 谱的最大特征值 与局部 距离的曲率 之间存在 的关系,比例常数 与模型宽度有关。这条猜想若成立,将桥接 Wasserstein 几何与训练稳定性理论。
本文没有宣称已经证明这些猜想——它们是可推导但未验证的。我们希望这一框架能为后续工作提供一个统一的、可证伪的几何坐标,让对齐理论从"算法清单"走向"流形上的测地线工程"。
给实验者的具体可执行清单(按操作难度排序):
实验 1(最易,一周可完成):固定偏好数据集 UltraFeedback,取 ,分别训练 GRPO 模型;保持 不变(按猜想 1 缩放)。观察最终模型在 AlpacaEval 2 上的胜率曲线是否重合。预期结果: 时四条曲线几乎重合, 时出现明显偏离。这条实验验证猜想 1 的核心结论,且实现简单。
实验 2(中等难度,两周可完成):取训练好的 RLHF 模型,分别在原始 分布上和"重置后的"参考分布上重新训练奖励模型,再用 PPO 训练 1000 步。预期结果:势的迭代重整化方案(猜想 3)的 reward hacking 显著低于单次 训练(loss 曲线更稳定,最终 RM 误差更低)。
实验 3(高难度,月级):构造 1 维 toy 偏好分布(如 Gaussian),用 McCann 插值显式生成 测地线上的中间点;分别用 PPO、DPO、GRPO 训出三个 ;比较 与测地线上"理论最优"的距离。预期结果:三个 的距离相近(在有界常数内),且与 McCann 测地线的偏差随训练步数减少。这条实验是定理 1 的直接验证,价值最高。
实验 4(理论向,月级):推导 与 的关系式(猜想 5),与 SAM(Sharpness-Aware Minimization)训练过程的 sharpness 数据做回归。预期结果:在 7B-70B 规模上观察到 的标度律, 与模型宽度相关。
对 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.