Awesome RL with Human Feedback / 2026 / OpenReview

Semantic-aware Wasserstein Policy Regularization for Large Language Model Alignment

作者：Byeonghu Na, Hyungho Na, Yeongmin Kim, Suhyeon Jo, HeeSun Bae, Mina Kang, Il-chul Moon · 单位：UNIST, 3summary.ai · 来源：OpenReview · 发布日期：2026

偏好学习对齐目标设计训练稳定性策略优化

精读笔记

Problem Setting

论文标题：Semantic-aware Wasserstein Policy Regularization for Large Language Model Alignment（Awesome RL with Human Feedback / 2026）。

这篇论文处理的是 RLHF 中 reference regularization 的几何定义问题。标准 PPO-RLHF 把 learned policy 约束在 SFT reference 附近，通常用 reverse KL；f-DPO、χPO 等把 KL 换成其他 f-divergence，但仍然在同一个 token-index simplex 上比较概率。真正的问题不是 KL 是否 mode-seeking，而是这种 regularization 默认 token 是无结构类别：只要 index 不同，就没有“相近”概念。

困难点在于，alignment 需要两种相反力量同时存在：一方面 policy 不能远离 reference，否则 reward model 外推、reward hacking、语言质量退化；另一方面 policy 又必须能离开 reference 的具体 token 分布，否则偏好优化只能做保守重排。KL 类正则把“语义等价的替换”和“语义错误的漂移”混在一起惩罚，这使得 regularization 在语言空间里过于粗糙。

所以这篇论文的实际目标是：构造一种仍可用于 PPO 的 policy regularizer，但其 trust region 不是 probability-coordinate trust region，而是 semantic-geometry trust region。

Motivation

作者对已有路线的判断是对的：替换 f-divergence 只能调节概率比惩罚的形状，不能改变 token space 的几何。FKL、JS、TV、χ² 看似有不同 mode behavior，但它们都没有使用 token 之间的语义关系。因此 cat→kitten 和 cat→table 在这些 divergence 下都只是 probability mass 从一个 index 消失、另一个 index 出现。

关键缺口是 regularizer 缺少 token substitutability。语言生成中很多 policy deviation 是合理的：同义词、形态变化、风格替换、等价短语开头、代码中等价变量名等。如果 regularizer 对这些变化过度惩罚，会限制 policy 在 reward signal 下做有益移动。Wasserstein 的自然吸引力在于它不是逐坐标比较分布，而是允许质量搬运；只要 cost matrix 表达了语义邻近性，它就能把“离 reference 多远”定义为“移动到了多远的语义区域”。

这不是一个新的 RL 算法动机，而是一个 representation-aware regularization 动机：alignment 目标里的距离函数应该知道 token space 的结构。

Core Idea

核心思想很简单但有效：把 next-token distribution 与 reference next-token distribution 之间的距离，从 f-divergence 换成 entropy-regularized Wasserstein distance。也就是说，policy regularization 不再问“同一个 token 上的概率比是多少”，而是问“learned policy 的概率质量能否以低成本搬运回 reference policy 的概率质量”。成本由 token embedding 距离给出。

这个变化引入了一个新的 inductive bias：semantic locality。policy 可以在 reference 的语义邻域内更自由地移动，但向语义远处移动会被惩罚。它本质上把 reference model 从一个精确分布锚点变成了一个语义流形附近的软约束。这比 KL 更适合开放式生成，因为高质量响应往往不需要逐 token 模仿 reference，只需要保持语义和行为分布不漂移。

和 prior 的本质区别不是“用了 Wasserstein”这个名字，而是 regularizer 的信息源变了：f-divergence 只使用两个 policy 的概率值；WPR 额外使用 token embedding geometry。它把 representation information 注入 alignment objective，而不是只在模型参数里隐式存在。

Method

方法层面的必要机制如下。

1. Token-level Wasserstein regularization：在每个 prefix 条件下，对 πθ(·|x,y< n) 和 πref(·|x,y< n) 计算 OT 距离。这解决的是 KL 正则语义盲的问题。核心变化是 regularization 从坐标差异转为语义搬运成本。

2. Entropic regularization / Sinkhorn：精确 OT 对大词表不可用，因此引入 entropy-regularized Wasserstein。它牺牲精确性，换来平滑性、稳定性和 Sinkhorn-Knopp 可计算性。这一步不仅是工程加速，也改变了 penalty 的性质：coupling 更 dense、更平滑，减少了 hard nearest-neighbor transport 的不稳定。

3. Dual variable as reward penalty：通过 Sinkhorn dual，把 Wasserstein distance 写成对 sampled token 的 φ* penalty。这样 WPR 可以像 KL penalty 一样进入 token-level reward shaping，并直接接 PPO / GAE。这个转换是论文最关键的算法桥梁；否则 Wasserstein regularizer 只是一个昂贵的分布级约束。

4. Sparse / top-k approximation：nearest-k cost sparsification 和 top-k distribution truncation 让方法在大词表上可训练。这里主要是 engineering，但没有它方法不可扩展。需要注意，实际训练的不是完整 vocab Wasserstein，而是截断支持上的近似语义 OT。

Key Insight / Why It Works

我认为真正有效的原因不是 Wasserstein 本身带来某种神奇的 RL 稳定性，而是它提供了 better inductive bias：把 reference constraint 从 exact token matching 放松为 semantic-neighborhood matching。RLHF 的 reward 往往鼓励模型改变措辞、内容选择和响应策略；KL 对任何 token-level deviation 都收税，WPR 则只对语义远移重税。这会让 policy 更容易朝 reward-preferred 表达移动，同时不完全失去 reference 的语言先验。

最核心贡献是 dual-potential reward shaping。它把一个分布几何问题重新组织成 token reward 问题，使得现有 PPO 管线几乎不需要改。这一点比“我们用了 Sinkhorn”更重要。Sinkhorn 是可计算手段，dual penalty 是与 RLHF 接口兼容的机制。

增益中可能有几部分混在一起：第一，semantic geometry 的确可能改善 regularization；第二，Sinkhorn entropy 让 penalty 更平滑，可能天然比某些 f-divergence 稳；第三，WPR penalty 观察上比 KL 更 lenient，因此部分提升可能只是来自更宽松的 effective regularization，而不一定来自语义理解。论文用 BERTScore correlation 和 case study 支持语义归因，但还不足以完全排除“penalty scale / smoothing 更合适”的解释。

这不是 scaling、retrieval、test-time compute 或 data coverage 方法；它属于 objective-level inductive bias / representation alignment。它利用已有 embedding space 作为 latent semantic structure，并把这个结构显式注入 policy optimization。若 embedding geometry 可靠，方法会更 generalizable；若 embedding geometry 与 preference 目标错位，WPR 会把错误偏差系统性写进 regularizer。

Relation To Prior Work

最接近的是 KL-regularized RLHF、DPO implicit reverse-KL regularization、以及 f-divergence preference optimization / regularization。WPR 和它们共享同一个高层框架：reward maximization under reference constraint。不同点在于 constraint 的度量空间。KL / f-divergence 是 probability simplex 上的坐标度量；WPR 是带 ground cost 的 optimal transport 度量。

和 f-DPO、χPO 的本质差异是：后者改变 divergence curvature，WPR 改变 sample space geometry。前者仍然没有 token 之间的关系；后者引入了一个外部或模型内部的 semantic metric。这是实质新增的信息。

和 Wasserstein RL / trust-region policy optimization 的关系是技术谱系上的继承：把 policy update 放在 OT geometry 下。但本文的新意在于把它落到 LLM token-level RLHF，并通过 Sinkhorn dual 形成可用 reward penalty。这个组合是已有思想的重组，但重组点有价值：OT 的语义搬运解释与 LLM token embedding geometry 非常契合。

看似新的部分中，Sinkhorn-Knopp、entropy OT、dual formulation 都是成熟工具；真正的创新是把它们放进 RLHF regularization，并证明这个接口可训练、开销可控、效果优于一组 f-divergence baseline。

Dataset / Evaluation

实验覆盖了 summarization、dialogue、code generation，并测试了不同 backbone，任务范围比单一 HH-RLHF 更宽。这个设计基本支持“WPR 不是只在一个数据集上偶然有效”的 claim。code generation 的加入也有意义，因为 token embedding semantic distance 在代码 token 上未必天然可靠；WPR 仍有提升，说明至少没有明显崩溃。

但 evaluation 对核心 claim 的验证仍有局限。主要指标是 GPT-4 pairwise win rate 和 MT-Bench single-answer grading，这些指标能反映生成质量偏好，但不能直接证明 Wasserstein penalty 捕捉了 human semantic regularization。作者补充了 BERTScore correlation、token candidate coherence、case study，这些更接近机制验证，但仍是 proxy。

实验没有充分隔离几个 confounder：effective β scale、penalty smoothness、Sinkhorn entropy、top-k truncation 带来的 implicit smoothing、以及不同 divergence 的调参难度。f-divergence baseline 在某些任务上训练不稳，WPR 更稳是重要结果，但是否因为语义 cost，还是因为 penalty 数值性质更温和，文中未充分说明。

真实 deployment 层面的安全性、harmlessness、reward hacking 抑制没有强验证；HH-RLHF 的 helpfulness win rate 不等于 alignment robustness。

Limitation

最大隐含前提是 token embedding distance ≈ preference-relevant semantic distance。这个前提在局部词义相似上成立，但在 instruction following、truthfulness、安全拒答、代码执行语义上很脆弱。embedding 近不代表行为等价；embedding 远也不代表 preference drift。例如 'not'、格式符、特殊 token、数字、变量名、API 名称的语义代价很难由普通 embedding L2 / cosine 捕捉。

第二，方法本质是 token-level regularization，而很多 alignment drift 是 sequence-level 的：事实结构、论证链、拒答策略、工具调用计划。WPR 只能在每个 next-token distribution 上施加局部几何约束，不能保证全局响应语义接近 reference。所谓 semantic-aware 仍是 local semantic-aware。

第三，scalability 靠截断。nearest-k 和 top-k2 让计算可行，但也意味着远距离 mass 基本被 hard-masked 或聚合到 dummy；这可能改变 OT 几何。论文报告 overhead 很小，但这是在特定 vocab、batch、GPU、k 设置下；更大模型、更长上下文、多语言 tokenizer 下成本和 memory 可能重新成为问题。

第四，增益归因不清。WPR 更 lenient than KL，且 entropy regularization 本身 smoothing penalty；因此部分提升可能只是更合适的 regularization strength / smoother reward shaping，而不是语义建模。需要与 tuned adaptive KL、learned token clusters、embedding-neighborhood KL、sequence embedding regularization 等强 baseline 比较。

第五，cost matrix 固定可能限制上限。policy fine-tuning 后 representation 改变，但 regularizer 仍使用 SFT embedding geometry；这既是稳定性来源，也可能导致语义约束滞后。文中未充分说明何时应该更新 cost、是否会引入 feedback loop。

Takeaway

  • 1. 这篇最值得记住的是：alignment regularizer 不一定只能在 probability simplex 上定义；把 token / action space 的几何显式放进 objective，是一个很自然但此前在 RLHF 中被低估的方向。
  • 2. WPR 推动的是 objective design，而不是 RL algorithm design。
  • 它说明很多 RLHF 稳定性问题可以通过更好的 regularization geometry 缓解，而不必总是改 PPO 或 reward model。
  • 3. 可迁移 insight：凡是 action space 有语义结构但被当作 categorical index 处理的生成任务，都可以考虑 OT / ground-cost regularization。

一句话总结

这篇论文把 RLHF 的 reference constraint 从 token-index 概率比正则推进到基于 token embedding 几何的 optimal-transport 正则，是一类 representation-aware alignment objective 的早期但有价值实现。