Semantic-aware Wasserstein Policy Regularization for Large Language Model Alignment
Fuente:
arXiv
Gespeichert in:
| Hauptverfasser: | , , , , , , |
|---|---|
| Format: | Preprint |
| Veröffentlicht: |
2026
|
| Schlagworte: | |
| Online-Zugang: | |
| Tags: |
Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
|
| _version_ | 1866918318779138048 |
|---|---|
| author | Na, Byeonghu Na, Hyungho Kim, Yeongmin Jo, Suhyeon Bae, HeeSun Kang, Mina Moon, Il-Chul |
| author_facet | Na, Byeonghu Na, Hyungho Kim, Yeongmin Jo, Suhyeon Bae, HeeSun Kang, Mina Moon, Il-Chul |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2602_01685 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | Semantic-aware Wasserstein Policy Regularization for Large Language Model Alignment Na, Byeonghu Na, Hyungho Kim, Yeongmin Jo, Suhyeon Bae, HeeSun Kang, Mina Moon, Il-Chul Machine Learning Artificial Intelligence 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. |
| title | Semantic-aware Wasserstein Policy Regularization for Large Language Model Alignment |
| topic | Machine Learning Artificial Intelligence |
| url | https://arxiv.org/abs/2602.01685 |