Semantic-aware Wasserstein Policy Regularization for Large Language Model Alignment
On the map: LLM alignment / RLHF / preference optimisation → LLM preference optimization methods
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.
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.
Builds on
44 of 53 references.
Closest ancestors · 7
Same direction, newest first
- 2025Preference Optimization by Estimating the Ratio of the Data Distribution26 May
- 2024f-PO: Generalizing Preference Optimization with f-divergence Minimization29 Oct
- 2024Correcting the Mythos of KL-Regularization: Direct Alignment without Overoptimization via Chi-Squared Preference Optimization18 Jul
- 2023A General Theoretical Paradigm to Understand Learning from Human Preferences18 Oct
- 2023Beyond Reverse KL: Generalizing Direct Preference Optimization with Diverse Divergence Constraints28 Sep
- 2022Training language models to follow instructions with human feedback4 Mar
- 2021A General Language Assistant as a Laboratory for Alignment1 Dec
Other references · 37
- 2023GPT-4 Technical Reportcited by 18,396
- 2017Proximal Policy Optimization Algorithmscited by 15,215
- 2023LLaMA: Open and Efficient Foundation Language Modelscited by 14,132
- 2023Judging LLM-as-a-Judge with MT-Bench and Chatbot Arenacited by 8,626
- 2023Direct Preference Optimization: Your Language Model is Secretly a Reward Modelcited by 8,065
- 2019BERTScore: Evaluating Text Generation with BERTcited by 5,512
- 2017Deep reinforcement learning from human preferencescited by 3,912
- 2022Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedbackcited by 3,720
- 2023Qwen Technical Reportcited by 3,344
- 2021TruthfulQA: Measuring How Models Mimic Human Falsehoodscited by 2,975
- 2020Learning to summarize from human feedbackcited by 2,588
- 2013Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distancescited by 2,316
25 more
- 2019Fine-Tuning Language Models from Human Preferencescited by 2,150
- 2015High-Dimensional Continuous Control Using Generalized Advantage Estimationcited by 2,026
- 2020RealToxicityPrompts: Evaluating Neural Toxic Degeneration in Language Modelscited by 1,255
- 2024Qwen2.5-Coder Technical Reportcited by 1,203
- 2022Red Teaming Language Models with Language Modelscited by 1,026
- 2024Gemma: Open Models Based on Gemini Research and Technologycited by 916
- 2021Measuring Coding Challenge Competence With APPScited by 902
- 2023RAFT: Reward rAnked FineTuning for Generative Foundation Model Alignmentcited by 686
- 2022Improving alignment of dialogue agents via targeted human judgementscited by 531
- 2023RRHF: Rank Responses to Align Language Models with Human Feedback without tearscited by 463
- 2013Fast Computation of Wasserstein Barycenterscited by 216
- 2018Statistical Aspects of Wasserstein Distancescited by 203
- 2019Wasserstein Distributionally Robust Optimization: Theory and Applications in Machine Learningcited by 170
- 2024CodeGemma: Open Code Models Based on Gemmacited by 150
- 2020A Distributional Approach to Controlled Text Generationcited by 119
- 2023RLTF: Reinforcement Learning from Unit Test Feedbackcited by 114
- 2023DeepSpeed-Chat: Easy, Fast and Affordable RLHF Training of ChatGPT-like Models at All Scalescited by 108
- 2023Execution-based Code Generation using Deep Reinforcement Learningcited by 99
- 2022On Decoding Strategies for Neural Text Generatorscited by 90
- 2019Better Rewards Yield Better Summaries: Learning to Summarise Without Referencescited by 77
- 2018Banach Wasserstein GANcited by 56
- 2020Efficient Wasserstein Natural Gradients for Reinforcement Learningcited by 23
- 2024Sinkhorn Distance Minimization for Knowledge Distillationcited by 15
- 2024MA-RLHF: Reinforcement Learning from Human Feedback with Macro Actionscited by 15
- 2023Provably Convergent Policy Optimization via Metric-aware Trust Region Methodscited by 6
Built on by
Citing papers per year: 2026: 2.
2 papers.
Similar, not cited
Close in wording, no citation either way
- 2024On the Algorithmic Bias of Aligning Large Language Models with RLHF: Preference Collapse and Matching Regularization
- 2026From RLHF to Direct Alignment: A Theoretical Unification of Preference Learning for Large Language Models
- 2025Beyond RLHF: A Unified Theoretical Framework of Alignment
- 2025KL-Regularized RLHF with Multiple Reference Models: Exact Solutions and Sample Complexity
- 2026Unifying Stable Optimization and Reference Regularization in RLHF
- 2024Unified Preference Optimization: Language Model Alignment Beyond the Preference Frontier
- 2025Towards Reliable, Uncertainty-Aware Alignment
- 2025RE-PO: Robust Enhanced Policy Optimization as a General Framework for LLM Alignment