# Semantic-aware Wasserstein Policy Regularization for Large Language Model Alignment

arXiv: https://arxiv.org/abs/2602.01685 · HTML: https://mlsanity.com/paper/2602.01685/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in LLM preference optimization methods, one of 10 directions in LLM alignment / RLHF / preference optimisation.
- Builds on 44 of its 53 references that are in the map.
- First posted on arXiv 2 Feb 2026.

## Placed in

- LLM alignment / RLHF / preference optimisation / [LLM preference optimization methods](https://mlsanity.com/alignment/llm-preference-optimization-methods/)

## Builds on (most cited first, 30 of 44)

- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena](https://mlsanity.com/paper/2306.05685/) (2023)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [BERTScore: Evaluating Text Generation with BERT](https://mlsanity.com/paper/1904.09675/) (2019)
- [Deep reinforcement learning from human preferences](https://mlsanity.com/paper/1706.03741/) (2017)
- [Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback](https://mlsanity.com/paper/2204.05862/) (2022)
- [Qwen Technical Report](https://mlsanity.com/paper/2309.16609/) (2023)
- [TruthfulQA: Measuring How Models Mimic Human Falsehoods](https://mlsanity.com/paper/2109.07958/) (2021)
- [Learning to summarize from human feedback](https://mlsanity.com/paper/2009.01325/) (2020)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Fine-Tuning Language Models from Human Preferences](https://mlsanity.com/paper/1909.08593/) (2019)
- [High-Dimensional Continuous Control Using Generalized Advantage Estimation](https://mlsanity.com/paper/1506.02438/) (2015)
- [RealToxicityPrompts: Evaluating Neural Toxic Degeneration in Language Models](https://mlsanity.com/paper/2009.11462/) (2020)
- [Qwen2.5-Coder Technical Report](https://mlsanity.com/paper/2409.12186/) (2024)
- [A General Language Assistant as a Laboratory for Alignment](https://mlsanity.com/paper/2112.00861/) (2021)
- [Red Teaming Language Models with Language Models](https://mlsanity.com/paper/2202.03286/) (2022)
- [Gemma: Open Models Based on Gemini Research and Technology](https://mlsanity.com/paper/2403.08295/) (2024)
- [A General Theoretical Paradigm to Understand Learning from Human Preferences](https://mlsanity.com/paper/2310.12036/) (2023)
- [Measuring Coding Challenge Competence With APPS](https://mlsanity.com/paper/2105.09938/) (2021)
- [RAFT: Reward rAnked FineTuning for Generative Foundation Model Alignment](https://mlsanity.com/paper/2304.06767/) (2023)
- [Improving alignment of dialogue agents via targeted human judgements](https://mlsanity.com/paper/2209.14375/) (2022)
- [RRHF: Rank Responses to Align Language Models with Human Feedback without tears](https://mlsanity.com/paper/2304.05302/) (2023)
- [Fast Computation of Wasserstein Barycenters](https://mlsanity.com/paper/1310.4375/) (2013)
- [Statistical Aspects of Wasserstein Distances](https://mlsanity.com/paper/1806.05500/) (2018)
- [Beyond Reverse KL: Generalizing Direct Preference Optimization with Diverse Divergence Constraints](https://mlsanity.com/paper/2309.16240/) (2023)
- [Wasserstein Distributionally Robust Optimization: Theory and Applications in Machine Learning](https://mlsanity.com/paper/1908.08729/) (2019)
- [CodeGemma: Open Code Models Based on Gemma](https://mlsanity.com/paper/2406.11409/) (2024)

## Built on by (most cited first, 2 of 2)

- [OpenVLThinkerV2: A Generalist Multimodal Reasoning Model for Multi-domain Visual Tasks](https://mlsanity.com/paper/2604.08539/) (2026)
- [AlignFed: Alignment-Aware Asynchronous Federated Fine-Tuning for Large Language Models in Heterogeneous Edge Environments](https://mlsanity.com/paper/2606.08197/) (2026)

## Abstract (arXiv)

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.