Semantic-aware Wasserstein Policy Regularization for Large Language Model Alignment

TL;DR

This paper introduces Wasserstein Policy Regularization, a semantic-aware method for aligning large language models with human preferences, outperforming traditional divergence-based approaches by capturing the geometry of token space.

Contribution

The paper proposes a novel Wasserstein-based regularization for RLHF that incorporates semantic information, improving alignment of large language models.

Findings

WPR outperforms KL- and f-divergence-based methods in experiments.

Semantic-aware regularization improves model alignment.

Method is compatible with standard RL algorithms.

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.