Semantic-aware Wasserstein Policy Regularization for Large Language Model Alignment

TL;DR

This work tackles the problem of aligning large language models to human preferences by addressing a key limitation of KL-based policy regularization, which ignores semantic relationships between tokens. It introduces Wasserstein Policy Regularization (WPR), using the entropy-regularized Wasserstein distance (Sinkhorn distance) to encode token-space geometry and semantic similarity, with a dual formulation that yields token-wise penalties computable via the Sinkhorn-Knopp algorithm. The method is integrated into RLHF as a tractable penalty on the reward, compatible with PPO, and is augmented with practical truncations to manage computational cost. Empirical results across TL;DR summarization, HH-RLHF dialogue, and code generation show that WPR outperforms KL- and $f$-divergence–based baselines in win rates and MT-Bench scores, with robust performance across model backbones and tasks. The findings highlight the value of semantic-aware policy distances for robust, scalable LLM alignment and suggest broad applicability to real-world alignment challenges.

Abstract

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.

Figures (12)

• Figure 1: Motivating example for the Wasserstein distance in LLM policy comparison. (a-c) Probability distributions of the reference and learned policies. (d) Semantic space among tokens. (e) Comparison under different divergences, where Wasserstein distance captures semantic relationships that KL and JS divergences fail to reflect.
• Figure 2: Win rates on TL;DR with Gemma-7B. '-2B' compares to the 2B models in \ref{['tab:main']}, and '-7B' to the 7B baselines.
• Figure 3: Overview of RLHF with Wasserstein Policy Regularization. (a) Standard RLHF with a policy regularization penalty. (b) Our proposed Wasserstein policy regularization, where the penalty is computed from the optimal dual variables obtained via the Sinkhorn-Knopp algorithm.
• Figure 4: Sensitivity analysis of the policy regularization hyperparameter $\beta$ on HH-RLHF.
• Figure 5: Normalized KL vs. Wasserstein penalty.

Theorems & Definitions (6)

• Proposition 0
• Theorem 1
• Proposition 1
• proof
• Theorem 1
• proof