Semantic-aware Wasserstein Policy Regularization for Large Language Model Alignment

AI summary

Overview

Research area: Machine learning / large language model alignment, specifically policy regularization inside reinforcement learning from human feedback (RLHF).

Technical level: Intermediate. The paper assumes familiarity with RLHF, Kullback-Leibler (KL) divergence, optimal transport, and standard policy-gradient methods such as PPO, but it explains its motivating example from scratch.

Scope: This paper proposes Wasserstein Policy Regularization (WPR), a semantic-aware alternative to KL- and f-divergence-based policy regularizers for LLM alignment, and evaluates it empirically on summarization, dialogue, and code generation tasks.

The paper is an arXiv preprint (arXiv:2602.01685v1 [cs.LG], 02 Feb 2026) by Byeonghu Na, Hyungho Na, Yeongmin Kim, Suhyeon Jo, HeeSun Bae, Mina Kang, and Il-Chul Moon, with affiliations at KAIST, UNIST, and summary.ai. It is released under a CC BY 4.0 license, and code is available at https://github.com/aailab-kaist/WPR.

What This Paper Is About

RLHF aligns a language model with human preferences by maximizing a reward while penalizing deviation from a supervised fine-tuned reference model. That penalty is almost always a KL divergence (or one of its f-divergence variants), which compares the trained and reference policies only by looking at the probabilities each assigns to the same token index. The problem is that this ignores meaning: a policy that says "kitten" when the reference says "cat" is penalized just as much as a policy that says "table," even though "kitten" is semantically close to "cat." The paper's goal is to replace the token-index-based penalty with one that uses a distance between tokens — the entropy-regularized Wasserstein (Sinkhorn) distance — and to show this yields a tractable training objective that beats KL-based and other f-divergence baselines.

Key Contributions

1. A semantic-aware regularization framework (WPR). The authors replace the token-level KL regularization term in the RLHF objective with the Wasserstein distance, which compares policies using the geometry of the token space rather than only matching indices. This is claimed to remain well-defined even when the two policies have non-overlapping support.

2. A tractable dual formulation. Rather than optimizing the transport plan directly (which requires solving a linear program that becomes intractable for large supports), the authors use the entropy-regularized Wasserstein distance and show via its dual that the regularization appears as token-wise penalty terms added to the reward via optimal dual variables. This keeps the objective compatible with standard RL algorithms such as PPO.

3. An efficient penalty computation procedure. Optimal dual variables are obtained by the Sinkhorn-Knopp algorithm, and two truncation schemes — a nearest-k₁ truncation on the cost matrix and a top-k₂ truncation on the distributions — reduce complexity from O(d²) to O(k₂²) in the dictionary size d. The reported overhead is a 2.5% increase in training time per step relative to standard KL regularization.

4. An empirical comparison against six divergence baselines. WPR is evaluated against reverse KL (RKL), forward KL (FKL), Jensen-Shannon (JS), α-divergence (α = 0.5), total variation (TV), and χ² divergence across three tasks and multiple model backbones.

Main Findings

• WPR achieves the best win rates on both main datasets. Using GPT-4 win rate on Gemma-2B, WPR reaches 0.924 ± 0.019 versus SFT and 0.608 ± 0.029 versus RKL on TL;DR, and 0.852 ± 0.027 versus SFT and 0.616 ± 0.039 versus RKL on HH-RLHF. The closest baseline on TL;DR versus RKL is χ² at 0.540 ± 0.030; on HH-RLHF versus RKL the closest is FKL at 0.564 ± 0.019 — both below WPR.

• Several f-divergences are unstable or weak. FKL collapses to 0.040 ± 0.011 versus RKL and TV to 0.052 ± 0.021 versus RKL on TL;DR. RKL itself scores 0.848 ± 0.021 versus SFT on TL;DR and 0.828 ± 0.010 versus SFT on HH-RLHF. The authors attribute the poor FKL/TV results to exploding probability ratios.

• Gains hold on a larger backbone. On TL;DR with Gemma-7B, WPR scores 0.948 against SFT-2B, 0.712 against RKL-2B, 0.924 against SFT-7B, and 0.532 against RKL-7B; the RKL baseline does not report a vs. RKL-7B figure.

• Gains hold on a different architecture and scale. On HH-RLHF with Qwen1.5-1.8B-Chat, WPR scores 0.752 versus SFT and 0.560 versus RKL, compared with 0.716 versus SFT for RKL.

• WPR leads on MT-Bench. On Gemma-2B trained on HH-RLHF, WPR scores 4.272, above α-divergence at 4.256, FKL at 4.247, JS at 4.197, χ² at 4.144, TV at 4.072, and RKL at 4.000.

• WPR also improves code generation. On APPS with CodeGemma-7B, evaluated over the full 5k test set, WPR improves both reward and pass@1 over SFT and RKL at every difficulty level: Introductory (reward 0.1606, pass@1 24.78), Interview (reward −0.1062, pass@1 5.75), Competition (reward −0.2638, pass@1 1.92), and All (reward −0.0843, pass@1 8.79). For comparison, RKL reaches All reward −0.1093 and pass@1 8.32, and SFT reaches −0.1475 and 7.84.

• Ablations show the setup is reasonably robust, with Sinkhorn iterations mattering most. On TL;DR, switching the cost from L2 to cosine slightly improves results (0.932 ± 0.014 versus SFT, 0.644 ± 0.047 versus RKL). Reducing k₁ from 512 to 256 gives 0.920 ± 0.006 / 0.572 ± 0.025. Reducing k₂ from 128 to 64 gives 0.864 ± 0.015 / 0.528 ± 0.032. Reducing λ from 100 to 10 gives 0.868 ± 0.024 / 0.552 ± 0.010. Reducing Sinkhorn iterations from 10 to 5 causes the largest drop (0.708 ± 0.027 / 0.328 ± 0.026), while increasing them from 10 to 30 gives 0.880 ± 0.021 / 0.536 ± 0.029.

• The motivating example. With the vocabulary {cat, kitten, dog, table} and the question "What is in this image?" over a small cat image, the reference policy puts mass on cat, one learned policy on kitten, and another on table. KL diverges due to support mismatch, and JS assigns the same distance to both learned policies, whereas the Wasserstein distance correctly ranks the kitten policy as closer to the reference.

Methodology in Plain English

The standard RLHF objective asks a model to maximize reward minus β times some divergence from a reference model. The authors swap out that divergence for the Wasserstein distance, an optimal-transport measure that asks: what is the cheapest way to move probability mass from the learned distribution to the reference distribution, given a cost for moving mass between any two tokens?

Because exact Wasserstein requires solving a linear program that scales badly, they use the entropy-regularized version (the Sinkhorn distance). They then take the dual of that regularized problem: instead of working with the transport plan, they work with dual variables that turn out to be interpretable as per-token penalties. Substituting these dual variables back into the RLHF objective shows the regularized objective equals ordinary reward maximization with an extra token-wise penalty term, plus a constant that does not depend on the policy. That means nothing about the RL loop has to change — only the reward signal does.

The dual variables are found with the Sinkhorn-Knopp algorithm, which iteratively rescales rows and columns of a matrix so its marginals match the two policies. To keep this cheap, the paper truncates the cost matrix to each token's k₁ nearest neighbors (making it sparse) and truncates the distributions to their top-k₂ indices plus the sampled token, with remaining mass collapsed into a dummy index — cutting complexity from O(d²) to O(k₂²).

Experiments use Gemma-2B as the base model on the TL;DR summarization dataset and the Anthropic Helpful and Harmless (HH-RLHF) dialogue dataset, following the setup and open-source implementation of Chai et al. (2025). The cost function is Euclidean distance in a fixed token embedding space taken from the reference policy, with λ = 100, k₁ = 512, and k₂ = 128; the regularization weight β is chosen by grid search for stability. Evaluation uses GPT-4 win rate over 50 randomly sampled validation instances, with generation and comparison repeated five times, and response order randomized. TL;DR is judged on relevance, coherence, consistency, and fluency; HH-RLHF on helpfulness. Additional results use Gemma-7B, Qwen1.5-1.8B-Chat, and CodeGemma-7B on APPS.

Why This Matters

Impact on research. Alignment regularization has largely been treated as a choice among f-divergences, all of which share the same blind spot: they compare probabilities at identical token indices. This paper argues the choice of geometry over token space is a distinct and underexplored axis, and demonstrates that a semantic-aware distance can beat stronger f-divergence baselines like χ² and α-divergence even when those baselines are tuned. It also connects the RLHF literature to the optimal-transport literature that already underpins Wasserstein GANs and trust-region policy optimization.

Real-world applications:

• Dialogue assistants and chat models, where current KL regularization is known to be mode-seeking and to reduce output diversity; a semantic penalty may preserve diversity while still preventing drift from the reference model.

• Summarization systems, where the paper's TL;DR results show a measurable win-rate advantage over KL regularization and over strong f-divergence alternatives.

• Code generation assistants, where the APPS results show gains across Introductory, Interview, and Competition difficulty levels and on the overall pass@1 metric.

• Multi-backbone deployment pipelines, since the paper reports gains on Gemma-2B, Gemma-7B, Qwen1.5-1.8B-Chat, and CodeGemma-7B, suggesting the method is not tied to one architecture or scale.

Industry relevance. The overhead is modest — 2.5% more training time per step than KL regularization, per the authors — and the penalty slots into existing PPO-style pipelines without changing the RL algorithm, which lowers the barrier to adoption. The method also remains numerically well-defined under support mismatch, addressing the exploding-ratio instability that the paper observes in forward KL and total variation baselines.

Future Directions

• Choosing the cost function systematically. The paper uses Euclidean distance in a fixed token embedding space and reports a cosine variant in ablation; how to select or learn the token cost matrix — and whether contextual or sentence-level notions of similarity would help — is left open.

• Scaling to larger models and longer horizons. The largest backbone evaluated is Gemma-7B and the largest alternative is Qwen1.5-1.8B-Chat; whether the gains persist at frontier scale, and how the Sinkhorn iterations interact with much larger token dictionaries, is not established.

• Tightening the computational budget. The ablation shows the Sinkhorn iteration count matters a lot (10 → 5 iterations drops win rate from 0.924 to 0.708 versus SFT), so accelerating convergence or replacing the iterative solver could be a direct route to further efficiency.

• Comparing against preference-optimization methods without explicit reward models. The paper positions WPR against KL and f-divergence regularization within RLHF, and notes DPO and related variants follow a similar principle; whether a Wasserstein penalty can be folded into those direct-preference objectives is not addressed in the provided content.

Target Audience

Researchers and practitioners working on LLM alignment, RLHF, and preference optimization who already understand the basics of KL-regularized policy learning and want to know whether the choice of divergence — and the token geometry it does or does not encode — matters for downstream quality. It is also relevant to readers from the optimal transport and Wasserstein GAN communities interested in seeing Sinkhorn-based distances applied to token-level language model distributions, and to engineers looking for a drop-in regularization change that fits existing PPO pipelines.

Authors’ 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.