Semantic-aware Wasserstein Policy Regularization for Large Language Model Alignment

📄 arXiv: 2602.01685v1 📥 PDF

作者: Byeonghu Na, Hyungho Na, Yeongmin Kim, Suhyeon Jo, HeeSun Bae, Mina Kang, Il-Chul Moon

分类: cs.LG, cs.AI

发布日期: 2026-02-02

备注: Accepted at ICLR 2026

🔗 代码/项目: GITHUB

💡 一句话要点

提出语义感知的Wasserstein策略正则化方法，提升大语言模型对齐效果

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型对齐 强化学习 人类反馈 Wasserstein距离 策略正则化

📋 核心要点

  1. 现有RLHF方法使用KL散度正则化，但无法捕捉token间的语义相似性，限制了对齐效果。
  2. 提出Wasserstein策略正则化（WPR），利用Wasserstein距离度量token空间几何结构，实现语义感知的策略正则化。
  3. 实验表明，WPR优于KL散度和f-散度基线，验证了语义感知策略距离在LLM对齐中的有效性。

📝 摘要（中文）

大型语言模型（LLM）通常使用基于人类反馈的强化学习（RLHF）与人类偏好对齐。在这种方法中，LLM策略通常通过奖励最大化以及参考策略的Kullback-Leibler（KL）散度正则化进行优化。然而，KL及其$f$-散度变体仅比较相同索引处的token概率，未能捕捉语义相似性。我们提出Wasserstein策略正则化（WPR），这是一种基于熵正则化Wasserstein距离的RLHF框架的语义感知正则化方法，它结合了token空间的几何结构。该距离的对偶形式将正则化表示为通过最优对偶变量应用于奖励的惩罚项，从而产生与标准RL算法兼容的可处理目标。实验结果表明，我们的方法优于基于KL和$f$-散度的基线，证明了语义感知策略距离对对齐的益处。我们的代码可在https://github.com/aailab-kaist/WPR 获取。

🔬 方法详解

问题定义：现有基于RLHF的大语言模型对齐方法，通常使用KL散度来约束学习策略与参考策略的差异。然而，KL散度只关注token在相同位置的概率差异，忽略了token之间的语义相似性。这意味着语义相近但表达不同的回复会被过度惩罚，限制了模型生成更符合人类偏好的多样化回复。

核心思路：论文的核心思路是使用Wasserstein距离来度量策略之间的差异。Wasserstein距离考虑了token之间的语义关系，能够捕捉到KL散度无法捕捉到的语义相似性。通过在RLHF框架中使用Wasserstein距离作为正则化项，可以引导模型生成语义更接近人类偏好，同时保持多样性的回复。

技术框架：整体框架仍然是标准的RLHF流程，包括：1) 使用人类反馈数据训练奖励模型；2) 使用强化学习算法（如PPO）优化语言模型策略，目标是最大化奖励模型给出的奖励。关键改进在于，策略优化阶段的损失函数中，使用基于Wasserstein距离的正则化项替代了传统的KL散度正则化项。Wasserstein距离的计算通过其对偶形式进行，将正则化转化为对奖励的惩罚。

关键创新：最重要的技术创新点在于使用Wasserstein距离作为策略正则化项，从而实现了语义感知的策略约束。与KL散度相比，Wasserstein距离能够更好地捕捉token之间的语义相似性，避免对语义相近但表达不同的回复进行过度惩罚。这种语义感知能力使得模型能够生成更符合人类偏好且更具多样性的回复。

关键设计：论文使用了熵正则化的Wasserstein距离，并利用其对偶形式将Wasserstein距离的计算转化为对奖励的惩罚。具体来说，通过引入对偶变量，可以将Wasserstein距离的计算转化为一个线性规划问题，从而可以使用标准的优化算法进行求解。此外，论文还设计了合适的奖励函数，以引导模型生成符合人类偏好的回复。

📊 实验亮点

实验结果表明，提出的Wasserstein策略正则化（WPR）方法在多个LLM对齐任务上优于基于KL散度和f-散度的基线方法。具体而言，WPR在奖励得分和人类偏好评估指标上均取得了显著提升，证明了语义感知策略距离在LLM对齐中的有效性。代码已开源。

🎯 应用场景

该研究成果可应用于各种需要与人类偏好对齐的大语言模型应用场景，例如对话系统、文本摘要、代码生成等。通过提升模型对人类意图的理解和表达能力，可以显著改善用户体验，提高任务完成的质量和效率。未来，该方法有望推广到其他序列生成任务中，例如机器翻译、语音合成等。

📄 摘要（原文）

Large language models (LLMs) are commonly aligned with human preferences using reinforcement learning from human feedback (RLHF). In this method, LLM policies are generally optimized through reward maximization with Kullback-Leibler (KL) divergence regularization of the reference policy. However, KL and its $f$-divergence variants only compare token probabilities at identical indices, failing to capture semantic similarity. We propose Wasserstein Policy Regularization (WPR), a semantic-aware regularization for the RLHF framework based on the entropy-regularized Wasserstein distance, which incorporates the geometry of the token space. The dual formulation of the distance expresses the regularization as penalty terms applied to the reward via optimal dual variables, which yield a tractable objective compatible with standard RL algorithms. Empirically, our method outperforms KL- and $f$-divergence-based baselines, demonstrating the benefits of semantic-aware policy distances for alignment. Our code is available at https://github.com/aailab-kaist/WPR.