Semantic-aware Wasserstein Policy Regularization for Large Language Model Alignment
Authors: Byeonghu Na
Hyungho Na
Yeongmin Kim
Suhyeon Jo
HeeSun Bae
Mina Kang
Il-chul Moon
Venue: ICLR 2026 Poster (2026)
Keywords: Large language Model
Alignment
Reinforcement learning with human preference
Wasserstein distance
Sinkhorn distance
Primary area: generative models
TL;DR: We propose a semantic-aware regularization for the RLHF framework based on the entropy-regularized Wasserstein distance, which incorporates the geometry of the token space.

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.

Published as a conference paper at ICLR 2026
SEMANTIC-AWARE WASSERSTEIN POLICY REGULAR- IZATION FOR LARGE LANGUAGE MODEL ALIGNMENT
Byeonghu Na', Hyungho Na², Yeongmin Kim', Suhyeon Jo', HeeSun Bae', Mina Kang', 11-Chul Moon 1.3
'KAIST, PUNIST, 'summary.ai
byeonghu. nakkaist ac kr, h.na@unist.ach
(a (alsdudrla10, suhyeonjo, cat2507, kasongl cmmon) @kaist.............. . ac .rr
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), semantic-aware regularization for the RLHF framework based on the retropy-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 tractable objective compatible with standard RL algorithms. Empirically, baselines, demonstrating the
method outperforms KL- and f-divergence-bare benefits of semantic-aware policy distances for alignment. Our code is available https: / tit..b. com/ a w ab- ka baist/in //PP.
INTRODUCTION
Large language models (LLMs) have achieved remarkable progress in recent years, powering ap- plications ranging from conversational agents to code generation (Touvron et al., 2023; Achiam al., 2023; Hui et al., 2024). central challenge in their deployment aligning model behavior with human preferences. Reinforcement learning from human feedback (RLHF) has emerged as the dominant paradigm for alignment, where models are optimized to better reflect user intent (Christiano al., 2017; Bai et al., 2022; Ouyang et al., 2022). The standard RLHF pipeline trains reward model from human preference data and optimizes the LLM policy to maximize reward while remaining close supervised fine-tuned reference model (Ouyang et al., 2022). Recent advances such as Direct Preference Optimization (DPO) (Rafailov et al., 2023) and its variants (Azar et al., 2024; Ethayarajh et al., 2024) follow similar principle, reducing the preference learning to implicit reward maximization with reverse Kullback-Leibee (KL) regularization to maintain the reference policy.
The policy regularization by the KL divergence widely adopted because the KL divergence can be computed directly from the token probabilities of the reference and the trained models, which is implemented penalty on the reward. While KL-based regularization is effective in practice, exhibits known shortcomings; for example, the reverse KL tends to be mode-seeking. which reduces output diversity. Recent works have addressed these issues by replacing reverse KL with alternative f-divergences, such as f-DPO (Wang et al., 2024) and xO (Huang et al., 2025). However, these f- divergence- based constraints still measure policy discrepancy only by comparing token probabilities at identical indices, thereby ignoring semantic relationships between tokens.
To illustrate this limitation, we introduce simple example in Figure We consider a vocabulary (cat, kitten, dog, table} } and compare reference policy nie and two learned policies, Ti and T2. in the context next token selection when answering the question "What is in this image?" given small cat image. In this example, nett Mi, and ne assign high probability mass to cat, kitten, and table, respectively. Semantically, (caz, kitten) is more closely related than (cat, table), would expect net to be closer m than T2. However, KL values diverge

---

Published as a conference paper at ICLR 2026
1.0
1.0
1.0
0.5
0.5
0.5
Wass.
0.0
dog table
0.0
kitten dog table
0.0
kitten dog table
Win Kate RKLJ
0.6
(a) Reference policy Xief
(b) Learned policy λι
(c) Learned policy ma
0.5
0.4
KL
cat
kitten
table
D(s,lliy))
Disallii)
0.3
Wass.
0.0
0.5 Temperature
1.0
dog
(d) Semantic space
(c) Metric comparison
Figure 2: Win rates against KL-based regularization across sampling tempera- tures on dialogue generation with Gemma-2B, compar- ing f-divergences and our Wasserstein distance.
Figure 1: Motivating example for the Wasserstein distance in LLM policy comparison. (a-c) Probability distributions of the reference and learned policies. (d) Semantic space among tokens. (e) Comparison under different divergences, where Wasserstein distance captures se- mantic relationships that KL and IS divergences fail reflect.
due to the support mismatch, and other f-divergences such as Jensen-Shannon (JS) divergence assign the same distance to лi and æ2, failing to reflect semantic proximity.
To overcome this limitation, we introduce a new RLHF regularization framework based on Wasserstein distances, which refer as Wasserstein Policy Regularization (WPR). Unlike the KL and other f I-divergences, the Wasserstein metric compares distributions by explicitly considering the geometry of the underlying token space. This enables flexible, user-defined cost functions that naturally encode semantic similarity between tokens. Additionally, it remains well-defined even when the support of two distributions does not overlap. In the context of language modeling, these properties are crucial because policies that assign high probability to semantically related tokens (e.g., cat and k t en) could likewise be regarded similar. As illustrated in Figure 1, the Wasserstein distance properly identifies the reference policy Tief being closer proximity that KL and other /-divergences fail to reflect. As a result, as shown in Figure 2, the policy regularization with the Wasserstein distance achieves superior generation performance compared to KL and other f-divergence--aaee based approaches, with experimental details provided in Section 5.2.
than to T2, thereby capturing semantic
Building these properties, propose tractable optimization framework that leverages the entropy-regularize Wasserstein distance, i.e., Sinkhorn distance (Cuturi, 2013), as a semantic-awan, policy regularizer. Computing this distance requires solving entropic optimal transport problem; we recast it in the dual and show that the resulting dual variables represent the regularization penalty, This penalty can be incorporated into the reward token-wise adjustments. analogous standard KL-based regularization, making the formulation compatible with standard RL. algorithms such as PPO (Schulman et al., 2017). The optimal dual variables can be obtained efficiently via the Sinkhorn algorithm with modest overhead. Empirically, our approach outperforms KL- and f f-divergence-baat baselines, highlighting the effectiveness of semantic-aware policy distances for RLHE
2 RELATED WORKS
Aligning Large Language Models Traditional supervised fine-tuning (SFT) methods have been effective in language generation but shows limitations in aligning outputs with human preferences, such as sentiment (Maas et al., 2011), 1), helpfulness (Askell et al., 2021), harmlessness (Gehman et al., 2020), and truthfulness (Lin et al., 2022). RLHF has become the standard approach for preference alignment (Stiennon et al., 2020; Ouyang et al., 2022). trains a reward model from human preference data and uses it to optimize the policy via reinforcement learning to better match human preferences (Christiano et al., 2017: Ziegler et al., 2019; Böhm et al., 2019). This approach has enabled successful LLMs such as ChatGPT (Achiam et al., 2023). Recent alternatives avoid explicit reward models, including RAFT (Dong et al., 2023), RRHF (Yuan et al., 2023), and DPO (Rafailov al., 2023), which reformulate preference alignment direct policy optimization.
2

---

Published as a conference paper at ICLR 2026
Regularization for Policy Learning Methods such RLHF and DPO incorporate regularization by a reverse KL divergence during preference alignment to prevent the learned policy from deviating significantly from a reference model trained via SFT. While this constrains learning to remain close to the behavior of the reference model, the mode-seeking nature of reverse KI. tends to limit output diversity (Wiher et al., 2022; Khalifa al., 2021; Perez al., 2022; Glaese et al., 2022). To address this limitation, studies such as fDDOO (Wang et al., 2024) and xPO (Huang et al., 2025) have been proposed. In parallel, other works (Han al., 2025; Kim al., 2025) explore alternative divergences for directly matching the optimal policy, though our focus in this work is on regularization. However, f-divergence- based methods share a key limitation: they measure the distributional discrepancy solely based on probability values at identical indices, without reflecting the semantic relationships between tokens. In contrast, we propose novel approach that leverages distance metrics from the Integral Probability Metric (IPM) (Müller, 1997), such as Wasserstein distance (Adler & Lunz, 2018; Panaretos & Zemel, 2019), to enable semantic-ane aware policy regularization.
Application of Wasserstein Distance The Wasserstein distance and its variants, such as the Sinkhorn distance, have been widely applied across many machine learning domains, including generative modeling, robust optimization, and reinforcement learning (Arjovsky et al., 2017; Sinha et al., 2018; Moskovitz et al., 2021; Song et al., 2023; Cui et al., 2024). For example, in generative modeling, Wasserstein GANs (Arjovsky et al., 2017) leverage the Wasserstein distance between the generator distribution and the data distribution to improve training stability and mitigate mode collapse. In robust optimization, adversarial training is formulated using Wasserstein balls around the data distribution to provide certified robustness (Sinha et al., 2018). In reinforcement learning, the Wasserstein natural gradient aligns policy updates with the local optimal -rrnspoot geometry in behavioral policy optimization (Moskovitz et al., 2021). Song et al. (2023) explore trust-region policy optimization based on Wasserstein and Sinkhorn distance. Building on this line of work, we explore Wasserstein regularization for RLHF, enabling semantic-awar policy alignment.
3 PRELIMINARY
3.1
WASSERSTEIN DISTANCE
The Wasserstein distance between two distributions π and x' is defined as min peU((.m))
Dw(TlI = min (EU(T,="
P
(P.C),
(1)
where U(T, T') == (PE p4xd Pl. P. x'} is the set of couplings between and n², C' E Rddd is the cost matrix with entries Cy. == cly, y²) 2 0, (. denotes the Frobenius inner product, and d is the cardinality of the outcome space.
To obtain a smooth and computationally tractable approximation to the Wasserstein distance, an en- tropy regularization term is added to the optimal transport objective, yielding the entropy-regularized Wasserstein distance, also known as the Sinkhorn distance (Cuturi, 2013):
(P. [(()
(=ll)==)
min PEU(F.)))
where A is an entropy regularization hyperparameter. and H(P)) := - Ei E Poo 1)
the entropy regularization term, equivalent to the Shannon entropy up to an additive constant.
While the Wasserstein distance directly relies on the Kantorovich dual formulation of optimal transport, the Sinkhorn distance arises from the dual of its entropically regularized variant (Villani, 2008; Peyré et al., 2019):
Dw(T] max 4.
dCT,
+i+ + 33 S Cit Vi.j )
(3)
j=1
(T(T](77)= max 4.
Σ dit Σ
exp (X(ei: و/ Cisl)
(4)
3

---

Published as a conference paper at ICLR 2026
where & and &b are the dual variables. In the Wasserstein case with the Euclidean cost, the dual vari- ables reduce to a single 1-Lipschitz function, which is typically parameterized by a critic network and optimized with gradient-based methods (Arjovsky et al., 2017). In contrast, the retropy-regularized formulation yields dual optimality conditions corresponding to matrix scaling factors, which can be computed efficiently by the Sinkhorn- Knopp algorithm (Sinkhorn & Knopp, 1967) closed- iterations alternating between row and column normalization (Cuturi, 2013; Cuturi & Doucet, 2014).
The entropy-regulariized Wasserstein distance produces smoother and denser couplings between distributions, and it converges to the Wasserstein distance as À + oc. Moreover, compared to the unregularized Wasserstein distance, the Sinkhorn distance can be computed more efficiently, incurring substantially less computational overhead. Since our setting requires computing next-token predictive distributions conditioned on various prompts and partial responses, we employ the entropic regularization variant rather than the critic-based Wasserstein distance, as the former admits the closed-form iterations.'
3.2
REINFORCEMENT LEARNING FROM HUMAN PREFERENCES (RLHF)
Our goal is to align an autoregressive L.M,, denoted as (y|))x where x is a user prompt and a response, with human preferences through reinforcement learning (RL). The RLHF procedure consists of three main stages, First, we perform supervised fine-tuning (SFT) to obtain a reference model Te,, which serves as the initial aligned model. Second, we train a reward model r(x. y) on preference dataset, enabling the estimation of scalar rewards for responses y given prompts x. Finally, using both the reference model ree and the reward model we optimize the following objective to fine-tune the language model To.
max I (TO; Tret) Ex [Ey-no(y)x)
- BD (To
(5)
where D denotes a policy divergence, 3 is policy regularization hyperparameter, and D is the prompt dataset. This objective encourages the policy to generate responses that maximize reward while remaining close to the reference model.
In most of the previous works, the divergence instantiated as the (reverse) KL divergence:
max (TØ; neat) = ExnD (Byyno((lx)) [r(x, y)] - 3DKL (xo(y|x))] no
(6)
Ex
Bon (0(x,)1.m)) Σ DKL. (To(yn yim 1ll Vi:-
R(x, Blog (y X,, V1sn-1)
Ex
B 70(..
(8)
Tha Un X, y1mn
where DKI. (T
[v.]] N is the sequence length of y, and R(x, yim) = r(x, VI:N) for n = N, and 0 otherwise. As shown in Eq. (8), the KL regular- ization term DxL can be rewritten as the expectation of the log-ratio between the two policies, which allows standard RL algorithms such as PPO to be applied for optimization. Furthermore, several works (Wang et al., 2024; Huang al., 2025) have generalized the KL divergence other f-divergences and developed tractable optimization formulations accordingly.
= Byn log [g))
However, mentioned in the Introduction, KL or other f-divergence measures compare policies solely by token-level probability differences at identical indices, without accounting for the under- lying semantic structure of tokens. This limitation prevents them from fully capturing meaningful distributional differences in language generation. To address this. we replace the divergence term with the Wasserstein distance, more precisely the Sinkhorn distance, which naturally incorporates semantic information, and we develop a tractable optimization framework for this objective.
In preliminary experiments, explored the critic-based Wasserstein distance but found that the resulting policy regularization insufficient, leading suboptimal performance.
4

---

Published as a conference paper at ICLR 2026
Penalized Reward
Wasserstein Penalty (x. Yis-1)
RL objective
(x. [t.)
A1x yr.) dog (2.
Sinkhorn- (Algorithm 1)
I0xx Xx- Dy [x.
Ty
Entropic Wasserstein Distance D3 (Fel)
Prompt
Reward Model r(
x ~
Roward yiml
y?
Policy soly (x)
Reference Policy Xwr (x)
Policy (x()
Response
Policy Regularization
Wasserstein Policy Regularization (WPR)
Figure 3: Overview of RLHF with Wasserstein Policy Regularization. (a) Standard RLHF with policy regularization penalty. (b) Our proposed Wasserstein policy regularization, where the penalty computed from the optimal dual variables obtained via the Sinkhorn Knopp algorithm.
METHOD: WASSERSTEIN POLICY REGULARIZATION
4.1 RLHF OBJECTIVE WITH WASSERSTEIN POLICY REGULARIZATION
In this section, we propose Wasserstein Policy Regularization (WPR), which regularizes L.LM policies in RLHF using the Wasserstein distance as the statistical distance between policies. We formulate the Wasserstein-regularzze RLHF objective by replacing the token-level KL. divergence regularization term in Eq. (7) with Wasserstein regularization term: max Tw (Toi Feet) =
(9)
no
Ex
Σ Eyn mol.mmolulxxx...
[R(x, V1:) Σ
Dw (Telyn Vim- (1mm-1)llaa l)ll x, v1.n-1))
p(n) C > where Un (TO, Tret) = p(()) E R p(n) = V p(n) 1 = Tee X v1.--1 }. C e
Here, Dw (To (y (To
(00lx,,8 x, yim 1 (y x, Y yim- 1)) = mi
Rexd the cost matrix, d is the token dictionary size, and N is the sequence length of y. It should be noted that p(n) depends on (x.) but we omit these input terms for simplicity.
The next step is to formulate the Wasserstein distance between two token-level discrete distributions so as to obtain tractable optimization objective for gy. However, computing the exact Wasserstein distance requires solving linear program, which quickly becomes intractable when the distributional support is large (Kuhn et al., 2019). As discussed in Section 1, 1, a widely used approximation is to introduce entropic regularization into the transport problem, referred to as the entropy-regularizzed Wasserstein distance or Sinkhorn distance (Cuturi, 2013):
max Twi (To: Teet)
(10)
N Σ
N
ExoD
Fy (x. [[(x, yish - 8
Σ n=1
D (To Yim- ( x, Vi-n-
(pim) C) 1 [(p(n)) (11)
where D (To( xx Y1:n Tresl x, V1:n-1)) = min p(n)eU,
We refer to this objective, Tw (mo Tret), the entropic Wasserstein--egullaiize RLHF objective. We now derive the dual problem from the regularized primal transportation problem in Eq. (11). Specifically, we introduce the Lagrangian function L corresponding to Eq. (11).
((P(n) .. xb)
(1)) (log pfn) 1) p (()) Cū + X
d
(iT[TT - Σ
p())
Σ
thy (Trees - Σ p(n))
(12)
ik
k=1
==1
1=1
5

---

Published as a conference paper at ICLR 2026
where { Ji= and {4/} } are the Lagrange multipliers, introduced to enforce the marginal con- straints in Uxi specificallly, they ensure that the row sums p(") match TT( x, (im-1) and the column sums match Teer(: x, x,Yu: Yimm Similar to p(x), the dual variables & and up are functions of (x. Yi: but omit their input terms for brevity unless this causes ambiguity. Based on this Lagrangian, the corresponding dual problem is given by
d d
max 0.07
Σ My Faed
ΣΣ i=1 1=1
exp(A(o + - -()))
(13)
which is derived in Appendix A.1. With strong duality and formulation of the primal solution, we can find the optimal solutions by Proposition 1 (Cuturi & Doucet, 2014).
Proposition 1. (Cuturi & Doucet, 2014) There exists pair of vectors (u, v) E R x R2. such that the optimal solutions of p(n), Φ. and up are respectively given by
p(x) * diag(u)
Φ - log(u), 3-- -----log("
(14)
We present the proof in Appendix A.2. Note that in our formulation, exp denotes the element-wise exponential applied to each entry of C. addition, for any real value 4, the pair of dual variables, (Ø + tid, & - tla) yields the same dual objective value. Hence, the dual optimal solutions unique but are determined only up to an additive constant. However, as shown in Theorem 2, when formulating the policy optimization problem based on this optimal solution, the additive term remains constant with respect to the policy and can therefore be ignored, yielding an equivalent problem.
not
By strong duality, substituting the optimal primal and dual variables obtained in Proposition into the objective Eq. (13) yields an expression of the hatrooy-regularized Wasserstein distance in terms of the optimal variables. Plugging this result back into the RLHF formulation in Eq. (10), we obtain a tractable optimization problem, as stated in Theorem 2.
Theorem 2. Let o² (x, Y1m-1) denote the optimal dual variables of the entropic optimal transport problem in Eq. (13). Then, the entropic beasersteii-regulriizeee RLHF in Eq. (10) can be equivalently written as a reward maximization problem with an additional penalty, induced by 6², i.e.
Jic(To: ned)
=ExDD
Σ
F1m)) - B6 (x, yim- 1)] +C, (15)
where C is constant with respect to Ne.
The proof is provided in Appendix A.3. Since the objective Tw of Eq. (15) can be expressed as the sum of token-wise rewards over sampled response sequences, the entropic Wasserstein-regularte RLHF problem, Eq. (10), can be optimized using standard RL methods such as PPO (Schulman et al., 2017). The full RLHF training algorithm is provided Algorithm of Appendix B.
4.2 COMPUTATION OF WASSERSTEIN PENALTY
As shown in the objective of Eq. (15), computing the Wasserstein penalty requires obtaining the opti- mal dual solution & of the entropic optimal transport problem. To this end, we need to compute the vee- tors u and v introduced in Proposition 1. These can be efficiently obtained by applying the Sinkhorn- Knopp algorithm (Sinkhorn & Knopp, 1967) for the matrix scaling problem, as described in Algo- rithm I (Cuturi, 2013; Cuturi & Doucet, 2014).
Algorithm Computation of Wasserstein
Penalty via Kanthom-KKopp Algorithm
Input:
Learned policy To (- V1:n-1), Ref-
erence policy )ee (x, -1---1),, Cost C
1: u 1,, v 1,,, K +- exp( exp(-XCI (XXC)
2: while converged do
3: diag(u) To TO - K (diag(v))
4: diag(v) t गee ./ K (diag(u)
end while
Specifically, as shown in Proposition 1, the opti- mal primal solution, p(n)* can be expressed as the -XC',, and two
6: Φ ----- ----1a log(u)
Output: Dual variable
product of the positive matrix, diagonal matrices, diag(u) and diag(v), with positive entries. Since p(n)* a transportation map, it must be doubly stochastic. Consequently, solving for and v in Proposition reduces to matrix

---

Published as a conference paper at ICLR 2026
scaling problem, which can be solved using the Sinkhom- Knopp algorithm. This algorithm iteratively rescales the rows and columns of p(n)* to match the target marginals mo and Nee, respectively (lines 3-4 in Algorithm 1, where ./ denotes clement-wise division).
Practical Consideration In practice, the Sinkhorn-Knoph algorithm can be directly applied, but it requires iterative matrix multiplications with the exponential of the cost matrix K == exp(- AC) e _ This incurs (d2) computational complexity with respect to the dictionary size d. leading to increased time and memory consumption. To mitigate this, we employ two forms of truncation.
First, during pre- coeputattininnn of the cost matrix, we apply nearest-ky truncation. For each token, distances are computed only to its L nearest neighbors. Entries outside this neighborhood are set to zero in K, which is equivalent to assigning infinite distance. This yields sparse K, enabling efficient sparse matrix multiplications that reduce both storage and computation. Second, during the Sinkhorn-Knopp Knopp algorithm, we apply top-kg truncation. The distributions Te and ref are truncated to their top-k2 indices together with the actually sampled index, while the remaining probability mass is aggregated into dummy index. This reduces the effective support size from d to at most 2kg + 2, lowering the complexity from (d²) to (k3). See Appendix C.3 for details of both truncations. Together, these truncations substantially reduce the computational cost of the entropic Wasserstein distance, with training time per step increasing by only 2.5% compared to standard KL regularization.
5 EXPERIMENTS
5.1 EXPERIMENTAL SETTINGS
Tasks and Training Details To evaluate our Wasserstein policy regularization, we conduct open- ended text generation experiments on two datasets: the TL:DR dataset (Völske et al., 2017) for text summarization and the Anthropic Helpful and Harmless dialogue generation. We follow the experimental setup of Chai et al. (2025)', which provides open- source implementations for RLHF research. Our base model is the pre-trained Gemma-2B (Team et al,, 2024b), and we use identical training configurations across all baselines and our method, varying only the regularization hyperparameters. For each method, the policy regularization hyperparameter ,9 is selected via grid search to identify the value at which training remained stable, and the best- performing model is reported. For Wasserstein policy regularization, we define the cost function as the Euclidean distance in the fixed token embedding space from the reference policy, set X = 100, and apply truncation hyperparameters h 512 and k2 128. Further experimental details are provided in Appendix C.
dataset (Bai et al., 2022) for
Baselines We compare regularization based on various divergences with the proposed entropic Wasserstein-basee regularization. Specifically, in addition to our approach using the entropic Wasser- stein distance Eq. (10), we evaluate reverse KL (RKL) divergence in Eq. (8), as well as level divergence in Eq. (7) instantiated with alternative f-divergences, including forward KL (FKL), JS, a-divergence with a = 0.5, total variation (TV), and x² divergence. Each f-divergence can be ex- pressed in the form of penalty on the reward through its defining function f. and the corresponding functions for each divergence are summarized in Table 11 of Appendix C.2.
max Jy (TT) Teet) =o
(16)
Trer Xx, N1:n-1),
TT (x. y1:n-1) 1)
E
Σ
((xxy1mm)
TT (Nrol Xx Y 11m-1)
Tier (X, y V1m-11
Evaluation We adopt GPT-4 win rate, a widely used evaluation metric in recent LLM studies (Zheng et al., 2023; Chai et al., 2025), as our primary metric. For evaluation, we randomly sample 50 validation instances and generate model responses, repeating this procedure five times. Then, GPT-4 is asked to perform pairwise comparisons between model outputs and report win rate. We use the GPT-4 evaluation prompts from Chai et al. (2025), with the full prompt included in Appendix C.4. For TL;DR, we assess relevance, coherence, consistency, and fluency; while for HH-RLHF we focus on helpfulness. To reduce evaluation bias, we randomize the order of the responses.
Phttps://g/ coub.ccm/eeee resea rch rch/rcc//a.... /MAARIHHH

---

Published as a conference paper at ICLR 2026
Table Comparison of win rates for policy regularization with various divergences, compared to SFT and RKL-regularized PPO on the TL;DR and the HH-RLHF datasets with the Gemma-2B model.
TL:DR
HH-RLHF
Divergence
vs. SFT
vs. RKL
vs. SFT
RKL
RKL FKL
0.848 + 0.021
0.828
0.100
0.316 + 0.026
0.040
+: 0.011
0.808
0.048
0.564
= 0.019
JS
0,540 + 0.024
0.204
:: 0.029
0.744
= 0.031
0.424
= 0.024
a (a = 0.5)
0.724 + 0.031
0,304
+ 0.016
0.792
0.047
0.524
= 0,016
TV
0.364 +: 0.039
0.052
:: 0.021
0.748
0.038
0.376
0.010
X"
0,904 +: 0.026
0.540
=: 0.030
0.796
0.039
0.344
= 0.017
Wasserstein (ours) 0.924 0.019
0.608
0.852
0.027
0.616
Table 2: Win rates TL;DR with Gemma-7B. '22B' compares to the 2B models in Table 1, and '-7B' to the 7B baselines.
Table 3: Win rates on HH-RLHF with Qwenl.5-1.8B
SFT-2B
vs. RKL-2B
vs, SFT-7B
RKL-7B
vs. SFT
RKL
RKL
0.948
0.668 0.712
0.912 0.924
RKL Wasserstein
0.716
Wasserstein
0.948
0.532
0.752
0.560
5.2 QUANTITATIVE RESULTS
Comparison to Other Divergences Table reports the performance of summarization and dialogue generation tasks under different policy regularization methods. Additionally, Figure 2 presents win- rate comparisons against RKI. across varying sampling temperatures on the HH-RL.HF dataset. As the results demonstrate, our proposed Wasserstein policy regularization method achieves the best results on both datasets. In contrast, f-divergence- based methods rely on probability ratios between policies, which produce exploding values and unstable training. This issue is evident in the particularly poor TL;DR results of FKL and TV. By comparison, our method remains well-defined even under support mismatch, enabling stable training and delivering superior performance consistently.
MT-Bench Results To further empirically validate our approach, we evaluate model using MT-Bench, GPT-4-based bench- mark shown to strongly correlate with human preference judg- ments (Zheng et al., 2023). We follow the official implementation" and apply ansle- grading to models fine-tuned on HH-RLHF. As shown in Table 4, our method achieves the highest performance among all baselines, indicating that semantic-aware regularization improves broader conversational and instruction following abilities.
Table 4: MT-Bench
parison on Gemma-2B trained HH-RLHF, evaluated with GPT-4 single-answer grading.
Divergence
Score
RKL
4,000
FKL
4.247
JS
4.197
Other LLM backbones To assess the scalability and generaliza- tion of WPR, we evaluate the method on larger and architecturally distinct LLM backbones. Using Gemma-7B the TL:DR summa- rization task, Table 2 shows that WPR continues to outperform the RKL-regularized baseline. We further demonstrate generalization by training Qwen- As shown Table WPR improves performance over the RKL-regularized model. These results indicate that WPR remains consistently effective across different architectures and model scales.
(a 0.5)
4.256
TV x²
4.072
4.144
Wasserstein 4.272
(Bai et al., 2023) on HH-RLHF.
Code Generation We also examine WPR in a different application domain, code generation. Following the experimental settings used in Chai ct al. (2025), assess CodeGemma-7B (Team et al., 2024a) on the APPS dataset (Hendrycks et al., 2021). We compute a reward using the compiler- exccution signal employed in prior work (Chai et al., 2025), and report both the reward and the pass@1 metric over the full 5k test set. As shown in Table 5, WPR achieves consistent improvements across the Introductory, Interview, and Competition levels, well in the overall performance.
Phttps://gi thub. comb wttp.ccommmms////aaaaaeeeee http:////ha/a..o/////////c///a/a
8

---

Published as a conference paper at ICLR 2026
Table 5: Performance comparison on APPS with the CodeGemma-78 model.
Introductory
Interview
Competition
All
Reward
pass@)
Reward
pass@1 1
Reward
pass @
Reward
pass @ 1
SFT RKL
0.1024
23. 12
-0.1720
4.86
-0.3239
1.48
-0.1475
7.84
0.1387
24.00
-0.1316
5.28
-0.2910
1.76
-0.1093
8.32
Wasserstein
0.1606
24.78
-0.1062
5.75
-0.2638
1.92
-0.0843
8.79
Table 6: Ablation study of WPR on TL;DR.
Win Rate (vx SFT)
0.9
0.8
Win rate
Method
Vs. SFT
vs, RKL
0.7
Our default settings
0.924
0.019 0.608
0.029
0.6
Cost change
(L2 cosine)
0.932
0.644
0047
0.001
0.0005
8
Decreased k
(512 + 256)
0.920 + 0.006
0.572
+: 0.025
Decreased k
(128 + 64)
0.864 + 0.015
0.528
#: 0.032
Figure 4: Sensitivity anal- ysis of the policy regular- ization hyperparameter 3
Decreased X
(100 + 10)
0,868 0.024
0,552
:: 0.010
Decreased Sinkhorn iterations
(10 - 5)
0.708 0.027
0,328
0.026
Increased Sinkhorn iterations
(10 - 30)
0.880 + 0.021
0.536
+ 0.029
on HH-RLHF.
5.3 ANALYSIS OF WASSERSTEIN POLICY REGULARIZATION
Ablation Study We conduct an ablation study to better understand the effect of components in our Wasserstein policy regularization framework, with results summarized in Table 6. Changing the cost function from the Euclidean to cosine distance yields slightly improved results, suggesting that the framework is robust the choice of token-level cost metric. Decreasing the truncation parameters L and 2/2 or the entropy regularization strength À leads to a slight drop in performance, though our method still consistently outperforms RKL. Smaller kı and ly introduce approximation errors in the distance computation, and a smaller À produces overly sharp couplings that reduce use the default settings, which provide consistently robust performance across
stability. In practice, datasets and configurations. The number of Sinkhorn iterations also affects the distance computation. Reducing iterations from 10 to leads sharp drop in performance due insufficient convergence, while increasing iterations to 30 provides no additional benefit over the default setting. These results suggest that moderate number of iterations is sufficient for achieving balance between accuracy and computational efficiency.
Additionally, we analyze the sensitivity analysis of the policy regularization coefficient in Figure 4. Our Wasserstein-reguiee regularized approach achieves stable performance across a broad range of 3 values, consistently outperforming the SFT baseline. In contrast, we observe that the f-divergence regular- ized RLHF yields stable training only within narrow ranges of [3, as also reported in the previous work (Wang et al., 2024). While our method demonstrates robustness over a wider range of A, still requires selecting an appropriate 3, highlighting a fundamental limitation of RLHF. Developing approaches that reduce or remove this dependence is an important direction for future work.
Wasserstein Penalty Figure compares the KL and Wasserstein penalties computed during training on TL:DR. where both are scaled by the optimal regularization coefficient / and jointly normalized to the range [0, 1] using a shared min-max range. Note that larger penalties correspond to greater deviation from the reference policy. As shown in the figure, the two penalties exhibit a strong positive correlation, with a Pearson correlation coefficient of 0.917. This result demonstrates that our Wasserstein penalty, similar to the KL penalty, increases as the learned policy differs from the reference policy. Moreover, the fitted trend line has slope of 0.579, which is less than 1, with a substantial fraction of points lying below the line, indicating that the Wasserstein penalty tends be more lenient than KI.
80
(Numolizeh) KL. Peralty
Figure 5: Normal- ized KL vs. Wasser- stein penalty.

---

Published as a conference paper at ICLR 2026
Table 7: Pearson correlation between each negative penalty and BERTScore (Zhang al., 2020).
Table 8: Semantic coherence of top- token candidates on each dataset.
TL:DR
HH-RLHF
TL;DR
HH-RIHF
KI. penally
0.1734
0.0172
RKL
3.781
0.005
3.690
+ 0.044
Wasserstein penalty
0.2160
0.1749
Wasserstein
3.593
3.584
0.004
0.35
0.158
KL. Penalty
K1.
0.25
Wasserstein
0.151
0.15
0.150
largesi
Fairs
Mass... include 1n
Boston
This
anmual
fuir
altracts
(a) Normalized wiie KL Wasserstein penalties
Probability
far
Probability
Reference
Sin
Winer
Learned
and
M4
euuual
Common
statd
Winter
Fair
Common
World
World
fair
Fairs
and
uinual
and
aumual
Wieeer Pall
Wolld
(b) 2nd token: fairs
(c) 11th token: Fair
Figure 6: Case study of penalties on Gemma-2B: The prompt is "What fair is the largest fair in Massachusetts?" and the generated response is "The largest fairs in Massachusetts include: Boston Fair: ..." " (a) Normalized penalties for each generated token. (b-c) Next-token distribution from each policy, along with the relevant cost matrix entries, for the 2nd and 11th tokens highlighted in (a).
To understand how the Wasserstein penalty captures semantic relationships and influences model behavior, we conduct the following analyses. We first evaluate whether the penalty aligns with semantic similarity using BERTScore (Zhang al., 2020). For responses generated by the reference and learned policies, we compute BERTScore with averaged KL. and Wasserstein penalties. Because higher BERTScore indicates greater semantic similarity, we correlate it with the negative value of each penalty. As shown in Table 7, the Wasserstein penalty shows a stronger positive correlation with BERTScore, providing quantitative evidence that WPR better reflects semantic similarity.
We further analyze how the penalty behaves in actual LLM distributions. For the example in Figure 6a, KL fluctuates widely, whereas WPR often assigns minimal penalty. To further investigate, we directly compare the next-token distributions at specific tokens. For semantically similar substitutions in Figure 6b, KI. assigns large penalty due to an exact index mismatch, while WPR gives small penalty by recognizing semantic proximity. Conversely, when probability mass shifts toward unrelated tokens in Figure 6c, WPR assigns large penalty, correctly signaling semantic drift.
We also measure semantic coherence of the learned LLMs. For every generated token, we extract top-10 next-token candidates and compute their mean pairwise embedding distance; smaller distances indicate greater semantic coherence. As shown in Table 8, WPR consistently produces more semantically coherent candidate sets than KL, with statistically significant margins.
Together, these results show that WPR penalizes semantic drift and promotes coherent semantic struc- ture. We conjecture that this semantic awareness contributes to the improved alignment performance.
CONCLUSION
In this work, we propose semantic-aware policy regularization framework for RLHF based on the entropy-regularized Wasserstein distance, which captures semantic similarity between tokens beyond the limits KL and other f-divergences. By formulating the regularization in the dual space, our method yields tractable penalties compatible with standard RL algorithms, while remaining computationally efficient via the Sinkhorn Knopp algorithm. Experiments on summarization and dialogue generation tasks demonstrate consistent improvements over KL- and f-divergence- baseed baselines, with higher win rates and MT-Bench score. These results highlight the effectiveness of semantic-aware policy distances for stable and robust alignment of large language models.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
This work was supported by the InnoCORE program of the Ministry of Science and ICT (N10260008) (50%). This work was supported by the IITP (Institute of Information & Communications Technology Planning & Evaluation)-ITRR (Information Technology Research Center) grant funded by the Korea government (Ministry of Science and ICT) (ITTP-202-R4) RS- 2024-004322668 (50%).
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et Gpt-4 technical report. arXiv preprini urXiv:2303.08774 2023.
Jonas Adler and Sebastian Lunz. Banach wasserstein gan. In S. Bengio, H. Wallach, H. Larochelle, and R. Garnett (eds.), Advances in Neural Information Processing
K. Grauman, N. Systems, volume 31. Curran Associates, Inc., 2018.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International conference on machine learning, pp. 214-223. PMLR, 2017.
Amanda Askell, Yuntao Bai, Anna Chen, Dawn Drain, Deep Ganguli, Tom Henighan, Andy Jones, Nicholas Joseph, Ben Mann, Nova DasSarma, et al. A general language assistant as a laboratory 2021.
for alignment. arXiv preprint arXiv:2
Mohammad Gheshlaghi Azar, Zhaohan Daniel Guo. Bilal Piot, Remi Munos, Mark Rowland, Michal Valko, and Daniele Calandriello. A general theoretical paradigm to understand learning from human preferences. In International Conference Artificial Intelligence and Statistics, pp. 4447-4455. PMLR, 2024.
Jinze Bai, Shuai Bai, Yunfei Chu, Zeyu Cui, Kai Dang, Xiaodong Deng, Yang Fan, Wenbin Ge, Yu Han, Fei Huang, Binyuan Hui, Luo Ji, Mei Li, Junyang Lin, Runji Lin, Dayiheng Liu, Gao Liu, Chengqiang Lu, Keming Lu, Jianxin Ma, Rui Men, Xingzhang Ren, Xuancheng Ren, Chuanqi Tan, Sinan Tan, Jianhong Tu, Peng Wang. Shijie Wang. Wei Wang, Shengguang Wu, Benfeng Xu, Jin Xu, An Yang. Hao Yang, Jian Yang, Shusheng Yang, Yang Yao. Bowen Yu, Hongyi Yuan, Zheng Yuan, Jianwei Zhang, Xingxuan Zhang, Yichang Zhang, Zhenru Zhang, Chang Zhou, Jingren Zhou, Xiaohuan Zhou, and Tianhang Zhu. Qwen technical report. arXiv preprint 23Xiv:2309.6669 2023.
Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislav Fort, Deep Ganguli, Tom Henighan, et al. Training helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint arXiv:2204.08 258622 2022.
Florian Böhm, Yang Gao, Christian M. Meyer, Ori Shapira, Ido Dagan, and Iryna Gurevych. Better rewards yield better summaries: Learning to summarise without references. In Kentaro Inui, Jing Jiang, Vincent Ng. and Xiaojun Wan (eds.), Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pp. 11-31012, Hong Kong, China, November 2019. Association for Computational Linguistics. doi: 10. 11153//1//111113333 1307.
Yekun Chai, Haoran Sun, Huang Fang, Shuohuan Wang. Yu Sun, and Hua Wu. MA-RLHE: Rein- forcement learning from human feedback with macro actions. In The Thirteenth International Conference on Learning Representations, 2025.
Paul F Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Xiao Cui, Yulei Qin, Yuting Gao, Enwei Zhang, Zihan Xu, Tong Wu, Ke Li, Xing Sun, Wengang Zhou, and Houqiang Li. Sinkhor distance minimization for knowledge distillation. In Proceedings of the 2024 Joint International Conference on Computational Linguistics, Language Resources 2024), pp. 14846-14858 2024.
and Evaluation
11

---

Published as a conference paper at ICLR 2026
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In C.J. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K.Q. Weinberger (eds.), Advances in Neural Informa- tion Processing Systems, volume 26. Curran Associates, Inc., 2013.
Marco Cuturi and Arnaud Doucet. Fast computation of wasserstein barycenters. In International conference on machine learning, pp. 685-693. PMLR, 2014.
Hanze Dong, Wei Xiong, Deepanshu Goyal, Yihan Zhang, Winnie Chow, Rui Pan, Shizhe Diao, Jipeng Zhang, KaShun SHUM, and Tong Zhang. RAFT: Reward ranked finetuning for generative foundation model alignment. Transactions on Machine Learning Research, 2023. ISSN 2835-8856.
Kawin Ethayarajh, Winnie Xu, Niklas Muennighoff, Dan Jurafsky, and Douwe Kiela. Model alignment as prospect theoretic optimization. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 12634-12651. PMLR, 21-27 Jul 2024.
Samuel Gehman, Suchin Gururangan, Maarten Sap, Yejin Choi, and Noah A Smith, Realtoxici- typrompts: Evaluating neural toxic degeneration in language models. In Findings of the Association for Computational Linguistics: EMNLP 2020, pp. 3356-3369, 2020.
Amelia Glaese, Nat McAleese, Maja Trębacz, John Aslanides, Vlad Firoiu, Timo Ewalds, Maribeth Rauh, Laura Weidinger, Martin Chadwick, Phoebe Thacker, et al. Improving alignment of dialogue agents via targeted human judgements. arXiv preprint arXiv:2209 (TtIv:220091437777 14375, 2022.
Jiaqi Han, Mingjian Jiang, Yuxuan Song, Stefano Ermon, and Minkai Xu. f-po: Generalizing preference optimization with f-divergence minimization. In International Conference on Artificial Intelligence and Statistics, pp. 1144-1154 1152. PMLR, 2025.
Dan Hendryeks, Steven Basart, Saurav Kadavath, Mantas Mazeika, Akul Arora, Ethan Guo, Collin Burns, Samir Puranik, Horace He, Dawn Song, and Jacob Steinhardt. Measuring coding challenge competence with APPS. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021.
Audrey Huang, Wenhao Zhan, Tengyang Xie, Jason D. Lee, Wen Sun, Akshay Krishnamurthy, and Dylan J Foster. Correcting the mythos of KL-regularization: Direct alignment without overopti- mization via chi-squared preference optimization. In The Thirteenth International Conference on Learning Representations, 2025.
Binyuan Hui, Jian Yang, Zeyu Cui, Jiaxi Yang, Dayiheng Liu, Lei Zhang, Tianyu Liu, Jiajun Zhang, Bowen Yu, Keming Lu, et al. Qwen2. 5-coder technical report. arXiv preprint arXiv:2409.12188 2024.
Muhammad Khalifa, Hady Elsahar, and Mare Dymetman. A distributional approach to controlled text generation. In International Conference on Learning Representationss, 2021.
Yeongmin Kim, Heesun Bae, Byeonghu Na, and II-Chul Moon. Preference optimization by estimating the ratio of the data distribution. arXiv preprint arXiv:2505. 19601 2025.
Daniel Kuhn, Peyman Mohajerin Esfahani, Viet Anh Nguyen, and Soroosh Shafieezadeh-Abatee Abadeh. Wasserstein distributionally robust optimization: Theory and applications in machine learning. In Operations research & management science in the age of analytics, pp. 130-166. Informs, 2019.
Stephanie Lin, Jacob Hilton, and Owain Evans. TruthfulQA: Measuring how models mimic human falsehoods. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio (eds.), Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 3214-3252 Dublin, Ireland, May 2022. Association for Computational Linguistics. doi: 1/.11866//11222222aclloog
Jiate Liu, Yiqin Zhu, Kaiwen Xiao, QIANG FU, Xiao Han, Yang Wei, and Deheng Ye. RLTF: Reinforcement learning from unit test feedback. Transactions on Machine Learning Research, 2023. ISSN 2835-8856.
12

---

Published as a conference paper at ICLR 2026
Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng. and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, pp. 142-150, Portland, Oregon, USA, June 201 1. Association for Computational Linguistics.
Ted Moskovitz, Michael Arbel, Ferene Huszar, and Arthur Gretton. Efficient wasserstein natural gradients for reinforcement learning. In International Conference on Learning Representations, 2021.
Alfred Müller. Integral probability metrics and their generating classes of functions. Advances in applied probability, 29(2):429-443, 1997.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang. Sandhini Agarwal, Katarina Slama, Alex Ray, al. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730- 27744, 2022.
Victor M Panaretos and Yoav Zemel. Statistical aspects of wasserstein distances. Annual review of statistics and its application, 6(1):405-431, 2019.
Ethan Perez, Saffron Huang. Francis Song, Trevor Cai, Roman Ring, John Aslanides, Amelia Glaese, Nat McAleese, and Geoffrey Irving. Red teaming language models with language models. In Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang (eds.), Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 3419-3448, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. doi: 10.18653/v1/2022. munlp-main.221
Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications data science. Foundations and Trends@ in Machine Learning,
2019.
2019.
Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. Advances in Neural Information Processing Systems, 01553728-53745 2023.
John Schulman, Philipp Moritz, Sergey Levine, Michael I. Jordan, and Pieter Abbeel. High- dimensional continuous control using generalized advantage estimation. In Yoshua Bengio and Yann LeCun (eds.), 4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings, 2016.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.063477 2017.
Parshin Shojace, Aneesh Jain, Sindhu Tipimeni, and Chandan K. Reddy. Execution- code generation using deep reinforcement learning. Transactions on Machine Learning Research, 2023.
ISSN 2835-8856.
Aman Sinha, Hongseok Namkoong, and John Duchi. Certifiable distributional robustness with principled adversarial training. In International Conference on Learning Representations, 2018.
Richard Sinkhorn and Paul Knopp. Concerning nonnegative matrices and doubly stochastic matrices.
Pacific Journal of Mathematics, 21(2):343-343 1967.
Jun Song, Niao He, Lijun Ding, and Chaoyue Zhao. Provably convergent policy optimization via metric-aware trust region methods. Transactions Machine Learning Research, 2023. ISSN 2835-8856.
Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei. and Paul Christiano. Learning summarize with human feedback. Advances in neural information processing systems, 33:3008-3021, 2020.
CodeGemma Team, Heri Zhao, Jeffrey Hui, Joshua Howland, Nam Nguyen, Siqi Zuo, Andrea Hu, Christopher A Choquette-Chee Choo, Jingyue Shen, Joe Kelley, et al. Codegemma: Open code models based gemma. arXiv preprint arXiv:2406. 1409, 2024a.
13

---

Published as a conference paper at ICLR 2026
Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, et al. Gemma: Open models based gemini research and technology. arXiv preprint arXiv:2403.08225, 2024b.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.1397 2023.
Cédric Villani. Optimal Transport: Old and New, volume 338. Springer Science Business Media, 2008.
Michael Völske, Martin Potthast, Shahbaz Syed, and Benno Stein. TI; dr: Mining reddit learn automatic summarization. In Proceedings of the Workshop on New Frontiers in Summarization, pp- 59-63, 2017.
Chaoqi Wang, Yibo Jiang, Chenghao Yang, Han Liu, and Yuxin Chen. Beyond reverse KL: Generaliz- ing direct preference optimization with diverse divergence constraints. In The Twelfth International Conference on Learning Representations, 2024.
Gian Wiher, Clara Meister, and Ryan Cotterell. On decoding strategies for neural text generators. Transactions of the Association for Computational Linguistics, 10:997-1012, 2022.
Zhewei Yao, Reza Yazdani Aminabadi, Olatunji Ruwase, Samyam Rajbhandari, Xiaoxia Wu, Am- mar Ahmad Awan, Jeff Rasley, Minjia Zhang, Conglong Li, Connor Holmes, et al. Deepspeed- chat: Fasy, fast and affordable rlhf training of chatgpt-like models at all scales. arXiv preprint arXiv:2308.00 03300, 2023.
Hongyi Yuan. Zheng Yuan, Chuanqi Tan, Wei Wang, Songfang Huang, and Fei Huang. RRHF: Rank responses to align language models with human feedback. Thirty-seventh Conference on Neural Information Processing Systems, 2023.
Tianyi Zhang, Varsha Kishore, Felix Wu, Kilian Q. Weinberger, and Yoav Artzi. Bertscore: Evaluating text generation with bert. In International Conference on Learning Representations, 2020.
Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric Xing, et al. Judging Ilm-as-a-judge with mt-bench and chatbot arena. Advances in neural information processing systems, 36:46595 2023.
Daniel M Ziegler, Nisan Stiennon, Jeffrey Wu, Tom B Brown, Alec Radford, Dario Amodei, Paul Christiano, and Geoffrey Irving. Fine-tuning language models from human preferences. arXiv preprint arXiv: 19w.00999900,93,, 2019.
14

---

Published as a conference paper at ICLR 2026
A PROOFS AND DERIVATIONS
A. DERIVATION OF EQ. (13)
We derive the dual problem in Eq. (13) from the Lagrangian in Eq. (12) constructed for the entropic primal transportation problem in Eq. (1I), First, we rewrite the Lagrangian as follows:
((P(m) Φ, 1/5
(17)
p(x) Cij + X ij
1 p())
(log pIl)) -1))
d Σ
d Σ p(a))
(TTT=
p(u))
+
Σ
(18)
ck
t=1
k=1
j=1
kmm
d
d
d
Σ =1
p(a) ij
Cij - x2 1b
1
(log p(n) - -1))
Σ
Σ
((Tea)))))
(19)
x
j=1
where {6: and {v; are the Lagrange multipliers. Based on this Lagrangian, the primal and dual problem can be written as follows:
(Primal)
min max C(P(a)
Φ v))
(20)
p(o) 0.40
(Dual)
max min C(P(n), Φ き))
(21)
0. pin)
By differentiating the Lagrangian with respect p(n), we derive the condition that the optimal p(n)* satisfy as follows:
oo + log p(n)*
pla) exp (A(e: + +/ Ciy)).
(22)
Therefore, by substituting the optimal p(n)" in Eq. (22) into Eq. (19), express the dual problem of Eq. (21) as
d d
max 6.47
[Te
Σ i=1
Σ
x
exp(X(e:
+ -cg)).
(23)
A.2 PROOF OF PROPOSITION 1
Proposition 1. (Cuturi Doucet, 2014) There exists a pair of vectors (u, v) E Rd x TR.. such that the optimal solutions of p(n), &, and vp are respectively given by
p(a)* diag( () exp( dXC)aagvv))) Φ" log(u),
3 --=-a -===-loge ==loog()) log(v).
1
(14)
Proof. From Eq. (22) in the derivation of Appendix A.I, the optimal p(n)" can be written
p(n)* exp (A(di + t/ - Ci₂)) = exp(Adr) exp(-ACig) exp()
(24)
Defining the positive kernel K == exp( XC), where the exponential applied element-wise, and the scaling vectors u == exp()Φ) and == exp()*b), the optimal coupling admits the compact representation
(n) diag(u) diag(v).
(25)
Since Ki is strictly positive, the Sinkhorn-Knopp theorem (Sinkhorn & Knopp, 1967) guarantees the existence and uniqueness (up to an additive constant) of strictly positive scaling vectors u, v E Rd such that p(a) E Un (To: Te..) Hence, the primal optimum ((n) is unique and necessarily of the form diag(u) K diag(v).
15

---

Published as a conference paper at ICLR 2026
Finally, by the definitions of and v, the corresponding optimal dual variables given by
log u,
3 A log v,
(26)
which yields the stated representation
p(a) diag(u) exp( exp(->C)
Φ² x log u,
1 a6= -- log
1
(27)
A.3 PROOF OF THEOREM 2
Theorem 2. Let <b (x, Vim-1) denote optimal dual variables of the entropic optimal transport problem in Eq. (13). Then, the entropic Wassersttein-reguaarde RLHF in Eq. (10) can be equivalently written as a reward maximization problem with an additional penalty, induced by 6' i.e.
N
Ji( knd)) = Ex
Σ
E To 1) [R(x, - 34 (x. 1)] +C, (15)
where C is constant with respect to TT.
Proof. First, the objective of Wasserstein--regulaiize RLHF can be written
Twe (To; Teet)
(28)
EX
Byy. 1) (R(x, 3 Σ D (To (x, Yimn-1)l [x, y vinn-)))
By strong duality, the entopic Wasserstein distance D is equal to the optimal objective value the W dual problem Eq. (13). Substituting the optimal solutions from Proposition 1, we obtain
D (Ne (y (unl yIn-
11m
(29)
of (x. (In- = =ix,
1:m 2/x. + Σ
(im- = 51x, v1m
(x (x -Cú))
(30)
d
oe (In i|x, V1m +C
(31)
=E
-a No (o (pxl 16 (x. Y1= + +C,
(32)
where C denotes a constant with respect to mo.
Substituting Eq. (32) into Eq. (28), the objective reduces to a reward maximization problem with an additional penalty induced by the dual variables o²:
Twl (To: Tret)
ExxD
Ey (ym X.yi. 1) [R(x, V1.h)] Ey (M.. mon (x.m 1) 10 (x, N1.n 1)] +
(33)
-
Σ
Ey. (y ()) [R(x, vI..) - -0% (x, (x, +
(34)
n
16

---

Published as a conference paper at ICLR 2026
B TRAINING ALGORITHM OF RLHE WITH WASSERSTEIN POLICY REGULARIZATION
Algorithm 2 RLHF with Wasserstein Policy Regularization
Input: Current policy To, Old policy T Reference policy Teef Reward model r(x,y) Cost matrix C, Dataset D
1: for tuain 1 to Tuain do
2: Sample Yidd ~ (-1x) for 'D (Computation with Batch Samples)
3: Get R(x, (in)) with reward model for {1, 2... ,d}
4:
Compute 6 (x, V1:m-1) via Algorithm using C' for n = {1,2. ,
5: Obtain An for
........... .d) via Eq. (35)
6: Compute (eJw(() via Eq. (40)
7: Compute V Lv (1b) via Eq. (38)
8: Update 0 - 0 + VeJw (0), a t + II VuLv(()) and e +
9: end for
Output: Learned policy To
In this section, present the detailed training algorithm for RLHF with Wasserstein Policy Regularization (WPR). As in conventional RLHF (Ouyang et al., 2022), we iteratively sample response Yo TT- (x, (i=-) to get Y1sd- Here, mg- is old policy whose parameters 0- are periodically updated by that the current policy, adopt Generalized Advantage Estimation (GAE) (Schulman et al., 2016) for penalized reward in Eq. (15).
At each token generation step n, we
(1n-1) R' (x, y1.)) = Vi:n) - (x,
Then, the advantage with GAE denoted as Å, each step n can be expressed as
Ån
Σ (7>) 8nti:
(35)
120
where
5n = (xx,)1mm) + VV (x, yim)
Vap V1:m-1)-
(36)
Here, Y is a discount factor; À is a hyperparameter for GAE; and Vi is a value network, which estimates the discounted cumulative reward or return of given state (I,),,), denoted as G,. Thus, the learning loss for Vi is defined as
Cv(+)= Ex-D
Σ
10n (1 (Vio(x, (V)((x)
(37)
Then, its gradient is expressed as follows.
(xLv(+) EXwD
Σ y. ~T, (vmlx.ymmm1) (x.,)) G..) Vo
(38)
In RLHF, Vi and me are updated together. By substituting a step-wist wise penalized reward R' v1.,)
with An, Eq. (15) is expressed as
( EXED
Σ Ey ~o X.
To An +C.
(39)
Tg
Here, we denote Jw (To; (eer) as (e) for simplicity and the gradient of Ju is computed as
is an importance weight. Then, n
{ no
V Jo (0) Exx-
Σ
Bynreag- (9h, Xxy1n 1)
A. Vo log no
(40)
To-
17

---

Published as a conference paper at ICLR 2026
In practice, clipping mechanism for advantage computation (Schulman et al., 2017) is adopted for Eq. (39) and (xm--11 in Eq. (38) is also clipped for a stable training. Algorithm 2 presents the overall training framework for RLHF with Wasserstein Policy Regularization. In Algorithm 2, Irain and 7 Train are the training step and the maximum training step, respectively. At tiaann = 1, current policy ne, old policy T and reference policy Met all initialized with SFT model. In Line # 8 in Algorithm 2, we update 0 with updated 0 at every training step, and T7 and `л are learning rate for Th and Vis, respectively.
C ADDITIONAL EXPERIMENTAL SETTINGS
C.1 DATASETS
TL;DR For the summarization task, the policy is trained to generate concise summaries of Reddit posts. The dataset includes 93K preference pairs for training and 86K pairs for validation. Training data is derived from the Reddit TL;DR corpus (Völske et al., 2017). For validation, a subset of data from CNN/Daily Mail is also used as dit-tribiitittto test set. The dataset is downloaded from Hugging Face.* 4
HH-RLHF For dialogue generation, we use the Anthropic HH-RLHF dataset (Bai al., 2022), where the policy is trained to produce responses that are both helpful and harmless in single-turn and multi-turn dialogue settings. It comprises 112K preference-labeled instances for training and an additional 12.5K instances for validation. The dataset can be downloaded from Hugging Face." 5
For code generation, we use the APPS dataset (Hendrycks et al., 2021), which provides a diverse set of programming problems requiring executable Python solutions. The dataset consists of 5K training instances and 5K validation instances, each containing a natural language problem dedscription paired with unit tests for automated evaluation. The dataset can be downloaded from Hugging Face. 6
APPS
C.2 MODEL TRAINING DETAILS
We follow the experimental setup of Chai et al. (2025), which provides open-source implementations for RLHF research.' 7 This implementation is based on the Deepspeed-Chhe package (Yao et al., 2023), and we adopt its configuration Gemma-2B (Team et al., 2024b), and we use identical training configurations across all baselines and our method, varying only the regularization hyperparameters. All baselines and proposed method, including the SFT and reward model, are trained under our experimental environment.
the default setting. Our base model is the pre-trained
Supervised Fine-Tuning (SFT) We split each dataset into three subsets and allocate 20% for supervised fine-tuning. Prompts are paired with their preferred responses to construct instruction data. In the TL;DR summarization task, posts concatenated with their reference summaries, while dialogue is formatted with a human-assistant chat template. For this stage, we employ the Gemma-2B model as the backbone. The training configuration specifies batch size of 512, learning rate of 5 x 10 5 a cosine learning rate scheduler with a warmup ratio of 0.1, and a total of 3 epochs.
Reward Model Training InstructGPT (Ouyang al., 2022) mitigates distributional mismatch by fine-tuning the reward model on the same dataset used for SFT. Following this approach, we also train our reward model on the identical dataset. In this stage, 40% of the data used for reward model training. Preference annotations are processed in the same way in SFT. The reward model is initialized from the SFT checkpoint. The training configuration specifies a batch size of 64, learning rate of 1 x 10-5, a cosine learning rate scheduler with a warmup ratio of 0.1, and a single epoch.
For code generation, we follow prior work (Chai al., 2025) and use a reward function derived directly from the compiler execution singal, without training additional reward model. We adopt
htttps://uuginggface co/ hhttps:/hhugiinge htt.://nnn.........a.aaaa/a////a
http://cc.......tt.a///////////a wwwrrrrrreeeeeerreeeeeeeeeeeeeee fror_feedbact
"https://gi thub. www./oooeeeeeeeeeeeeeeeeeeeeeeee /MA-
18

---

Published as a conference paper at ICLR 2026
Table 9: Hyperparameters for PPO training.
Table 10: Policy regularization hyperparameter 3 for each method.
Hyperparameter
Value
Divergence
TL:DR
HH-RLHF
PPO epochs
Rollout
RKL
0.005
0.001
Clip ratio
0.2
FKL
0.05
0.0001
À in GAE
0.95
JS
0.05
0.01
in GAE
a (a = 0.5)
0.01
0.05
Max prompt length
512
TV
0.01
0.01
Max response length
256
x²
0.001
0.001
Warmup steps
200
Wasserstein
0.05
0.0005
Temperature
0.8
Top-p
1.0
Top-k
50
Table 11: Corresponding functions for each f-divergences.
Divergence f(u)
RKL
u logg u log
FKL
IS
u loo
(u + 1) log( 2-)
(u (u=_a a.u 0.)
a(a-1)
TV
2 u
X²
2
the adaptive compiler- based reward used in previous studies (Chai al., 2025; Shojaee al., 2023; Liu et al., 2023). For a generated solution y to a problem x, the reward is defined as:
N Npass
0.3 + 1.3
if y compiles successfully,
N Npass
Nial
y(
00.6,
if y raises runtime error,
y fails compile. Here, Npsss and NEal denote the number of unit tests passed and failed, respectively.
1.0,
Policy Optimization with PPO The remaining 40% of the dataset is used for PPO training. The policy initialized from the SFT checkpoint, and the critic initialized from the reward model. We use the same SFT and reward model checkpoints for all baselines. The training configuration uses batch size of 256, learning rates of 1.5 x 10 for both the policy and the critic, and runs for epoch. We follow the hyperparameters from the original implementation, except that set the maximum response length to 256. The hyperparameters are summarized in Table 9.
Policy Regularization For each method, the policy regularization hyperparameter B is selected via grid search to identify the value at which training remained stable, and we report the best-performing model. Specifically, we perform a grid search over (0.5, 0.05, 0.01, 0.005,0000000 0.001, 0.0005, 0.0001}. The resulting 3 values used for each baseline summarized in Table 10. Each f-divergence can be expressed in the form of a penalty on the reward through its defining function f. and the corresponding functions for each divergence are summarized in Table 11.
max Jy (TT:; Tret) me
(41)
Trer(ym |x, 1im-1),
( y ylan-1) Tree / X, -1mn-1)
ExnD
Σ By Bo
R(x, R(x, v1.a) no TV Vnl x, N1mn-)
19

---

Published as a conference paper at ICLR 2026
0.04
Top-2 tnsscntion
0.02
20
Step
Figure 7: Convergence of the Sinkhorn- Sinkhorn-Knopp algorithm.
Figure 8: Example of top-k2 truncation where 1/2 2.
For WPR, we define the cost function as the Euclidean distance in the fixed token embedding space from the SFT model, set = 100, and apply truncation hyperparameters k = 512 and l2 = 128. The number of Sinkhorn iterations is set to 10 for TL;DR and 50 for HH-RLHF.
C.3 SHINKORN ALGORITHM DETAILS
Stopping Criterion The Sinkhorn- Sinkhorn-Knopp iterations involve alternating updates of the scaling vectors u and Since the dual variable ultimately used the regularization penalty, we monitor convergence based on the change in Φ X log Iterations are terminated when the change in & falls below pre-defined tolerance, which we use 10-4, For practicality, also impose maximum number Sinkhorn iterations, as specified in Appendix C.2.
Numerical Stability The Sinkhorn updates involve repeated rescaling operations and log com- putations when recovering Φ which can lead to numerical instabilities. To mitigate this, we add small constants to denominators and log arguments. Figure 7 plots the evolution the convergence metric across iterations on Gemma-2B experiments with HH-RLHF, demonstrating that the truncated Sinkhorn procedure converges stably in practice.
Nearest-k; Truncation for Cost Matrix Since the cost matrix C is fixed throughout training. we pre-compute K cxp(-XC). However, storing the full K is infeasible for vocabularies of extremely large tokens. address this, for each token we retain only its kn nearest neighbors and set all other entries to zero, which is equivalent to assigning infinite cost to distant tokens. We additionally enforce symmetry by mirroring retained entries so that the sparse kernel remains consistent. This sparsification enables efficient sparse matrix multiplications during Sinkhorn iterations.
Top-k:2 Truncation for Token Distributions For computational efficiency, we also truncate the token distributions of both the target policy and the reference policy. Specifically, we retain the top-k2 probability indices along with the index of the sampled token. The remaining probability mass is aggregated into a dummy index. Because the sampled token is always included, the required dual variable Ty.. be recovered. A conceptual illustration is provided in Figure 8.
C.4 EVALUATION DETAILS
We adopt GPT-4 win rate, widely used evaluation metric in recent LLM studies (Zheng et al., 2023; Chai al., 2025), as our main evaluation measure. For each comparison, we randomly sample 50 validation instances and generate model responses, repeating this procedure five times. Unless otherwise noted, the sampling temperature is fixed at 0.5. Then, GPT-4 is asked to perform pairwise comparisons between model outputs and compute the win rate. We use the gpt gpt-40-200 100020224-055 model for all evaluations. We follow the GPT-4 evaluation prompts provided by Chai et al. (2025), and for completeness, consistency, and fluency: while for HH-RLHE we focus on helpfulness. To reduce evaluation bias, randomize the order of the responses.
include the full prompt below. For TL:DR, assess relevance, coherence,
20

---

Published as a conference paper at ICLR 2026
GPT-4 Evaluation Prompt for TL:DR
You will be given two summaries written for article. Your task to pick the better one between them, based on the four criteria. Please make sure you read and understand these instructions very carefully.
Relevance selection of important content from the source. The summary should include only important information from the souree document. Annotators were instructed to penalize summaries that contained redundancies and excess information.
Coherence the collective quality of all sentences. We align this dimension with the DUC quality question of structure and coherence whereby "the summary should be well-structured and well-organized. The summary should not just be a heap of related information, but should build from a sentence to coherent body of information about topic."
Consistency the factual alignment between the summary and the summarized source. A factually consistent summary contains only statements that are entailed by the source document. Annotators were also asked to penalize summaries that contained hallucinated facts.
Fluency the quality of the summary in terms of grammar, spelling, punctuation, word choice, and sentence structure.
You should output a single character to indicate which summary you think is better. 'A' stands for Summary A and 'B' stands for Summary If you think both summaries equally good, output 'E'.
Article: (article)
Summary A: (summary_a)
Summary B:
(summary_b)
Your Choice (only single character, you are allowed to think both summaries are equal and output 'E'):
GPT-4 Evaluation Prompt for HH-RLHF
For the following query to a chatbot assistant, which response is more helpful?
First provide one-sentence comparison of the two responses and explain which you feel is more helpful. Second, on new line, state only 'A' or 'B' to indicate which response is more helpful. they are equally good or bad, state 'E'. Your response should use the json format, with "comparison" and "choice" keys.
Query: dialogue
Response A: (resp_a) a
Response B: [resp_i
Your Judgment:
21

---

Published as a conference paper at ICLR 2026
Table 12: Win rates on TL:DR using the Gemma-based models, varying the policy backbone and the embedding spaces used to form the cost matrix C.
Method
Backbone
Embedding
Win rate (vs, SFT-2B)
Win rate (vs. RKL-2B)
RKL-regularized PPO
Gemma-2B
0,848
Gemma-77
0.948
0.668
Wasserstein--eegllriizee PPO
Gemma-28
Gemma-2 2B
0,924
0.608
Gemma-7B
0.908
0.556
Gemma-7B
Gemma 2B
0,944
0,684
Gemma-78 7B
0.948
0.712
Table 13: Sensitivity analysis of the truncation hyperparameter 2/2 on TL;DR with Gemma-2B. Time the wall-clock time for the penalty computation.
22
Time (hours/1k steps)
Win rate (vs. SFT)
Win rate (vs. RKI.)
64
0.08
0.864
0.528
128
0.12
0.924
0.608
256
0.19
0.916
0.584
D ADDITIONAL EXPERIMENTAL RESULTS
D.I ANALYSIS EMBEDDING SPACE
The semantic cost matrix C is constructed from token embeddings and plays a central role in WPR. Since the cost must be computed over the full vocabulary of the policy model, the embedding space must be aligned with its tokenizer. Therefore, only models sharing the same tokenizer can be used directly. Using embeddings from a model with a different tokenizer would require building a cross-token alignment, promising but nontrivial direction for future work.
To study the effect of embedding quality, we conduct an experiment using Gemma-218 and Gemma- 7B, which share the same tokenizer. For each model, extract the frozen token embeddings after SFT and used them to construct the cost matrix C. We then independently varied: (1) the policy backbone (Gemma-2B Gemma-7B), and the embedding source used form C (Gemma-2B SFT or Gemma-7B SFT).
As shown Table 12, across configurations, WPR consistently outperforms RKL-regularized PPO when using the same policy backbone, indicating that WPR provides benefits regardless of the specific embedding model used. As expected, the policy backbone size has the largest effect on performance, with Gemma-7B outperforming Gemma-2B.
Interestingly, constructing C' using Gemma-7B embeddings for a Gemma-2B backbone does not yield performance improvements over using the 2B embeddings. We conjecture that this is because each policy is naturally grounded the token geometry encoded by its own SFT embedding space. Thus, the embedding space of the same model backbone is most compatible with the policy's 's internal representation.
D.2 ADDITIONAL SENSITIVITY ANALYSIS
Truncation hyperparameter 22 To evaluate how the computational overhead and performance scale with the truncation hyperparameter lz, we vary k2 E (64, 128, 256} and measure both the penalty computation time and the resulting win rates. The results are summarized in Table 13. Increasing 22 from 64 to 128 improves performance, with moderate increase penalty computation time, especially small compared to the overall training time of approximately 4.5 hours per 1,000 steps. Increasing 22 further from 128 to 256 yeilds minimal performance gains, and in fact slightly decreases performance. This suggests that k 128 already captures most of the probability mass of the token distribution and provides an accurate approximation. Accordingly, we adopt ho2 = 128 for all expeiments in the paper.
22

---

Published as a conference paper at ICLR 2026
Table 14: Sensitivity analysis of the entropy regularization parameter on TL:DR with Gemma-2B.
Win rate (vs. SFT)
Win rate (vs. RKL)
0.900
0.564
100
0.924
0,608
200
0.916
0.612
Table 15: Detailed breakdown of the wall-clock time per 1,000 training steps. The time required to compute the regularization penalty differs across methods, whereas the generation and training steps independent of the regularization method and therefore reported using unified timings.
Time (hours/Ik steps)
RKL
WPR
Generation
0.769
Penalty computation
0.005
0.117
Backpropagation
3.707
Total
4.481
4.593
Entropy regularization hyperparameter À The hyperparameter À controls the level of entropic smoothing in the Sinkhorn distance, determining the balance between the semantic fidelity and the smoothness of the transport plan. Smaller values of À place greater weight on entropy, producing overly soft transport plans and diminishing influence semantic structure. Conversely, exces- sively large values of À cause the kernel K
collapse toward zero, creating numerical
oscillations during the Sinkhorn rescaling steps.
To the best of our knowledge, this work is the first to apply entropy-regularie Wasserstein penalty in the token space LLMs during RL fine-tuning. Therefore, we initially selected À empirically and found that 100 provided the stable and consistent performance across tasks. We use this value all experiments reported in the paper. We additionally provide a sensitivity analysis with A {50, 100, 200]. As shown in Table 14, WPR consistently outperforms RKL-based regularization for all tested values (with win rate is greater than 0.5). As expected, smaller values reduce the influence of semantic structure and lead to performance drop, consistent with the interpretation above.
D.3 COMPUTATIONAL RESOURCES
Wall-clock Time Table 15 reports detailed breakdown of the wall-clock time per 1,000 PPO training steps for RKL and WPR. The measurement is decomposed into (1) generation, (2) penalty computation, and (3) backpropagation. For this report, use A100 GPUs using the Gemma-2B policy model with the TL;DR dataset, employing batches per GPU and gradient accumulation steps.
WPR requires additional computation during the penalty step because it incorporates semantic structure across tokens, whereas the KL penalty incurs nearly zero overhead. However, the added cost is minor relative to the forward and backward passes of a billion- parameter LLM. Note that the generation and backpropagation stages remain identical across regularization methods, and their runtimes are influenced far more by the generated response length, and therefore we report unified timings for these stages.
GPU Memory Usage The primary memory overhead of WPR arises from the cost matrix C. Since C is fixed across all training steps, it is computed once before training and reused throughout PPO optimization. Its memory footprint is independent of the model size.
As described in Section 4.2, apply truncation methods that allow C to be stored sparse matrix, requiring approximately 12.8GB of GPU memory in our configuration. This overhead is reflected in the peak GPU memory measurement in Table 16. During training, additional memory
23

---

Published as a conference paper at ICLR 2026
Table 16: Peak GPU memory usage (GB) for RKL and Wasserstein regularization, measured single A100 GPU with batch size of 8.
GPU usage (GB)
RKL
64.05
Wasserstein
78.98
usage depends primarily the batch size. When needed, memory be reduced by increasing the gradient accumulation steps, at the cost of longer runtime.
E GENERATED EXAMPLES
We present qualitative comparisons of response generation outputs. Table 17 shows an example from the TL:DR dataset, where we compare responses generated with reverse KL regularization and with our proposed Wasserstein Policy Regularization. Similarly, Table 18 provides an example from the HH-RLHF dataset under the same comparison. These examples illustrate that WPR produces responses that are better aligned with the underlying prompt compared to RKL.
F LLM USAGE
We used ChatGPT as an assistive tool during the paper writing process, primarily for language polishing and improving readability. Additionally, we employed GPT-4o API in the experimental evaluation as an automatic judge to compare model responses and compute win rates, following established evaluation protocols.
24

---

Published as a conference paper at ICLR 2026
Table 17: Comparison of response generation outputs from the prompt the TL;DR dataser with RKL and WPR.
Prompt
POST
Subreddit: r/pettyrevenge
I was driving home from work in San Antonio when noticed this jerk in sports car tailgating, honking, flashing lights, and in general being a major pest. We pulled up to red light and he gunned it right through the intersection, somehow managing not cause a major accident. I fumed for bit but figured that was the end of it.
Ten minutes later pulled into my apartment complex gate (it had one of those bars that raises and lowers after you put in a code) and there he was, idling next to it. I don't 't know if he was just visiting so didn' 't have the right entry numbers or what, but was just waiting for someone to come so he could slip by on their pass.
I pulled up to the gate, put my code in, and saw him inches from my bumper waiting to race in behind me. The gate raised up, went forward, he followed intently, and stopped with my rear bumper just barely in front of where the gate would fall.
He fumed in his front seat while waved, the gate lowered blocking him out, and I pulled away happy. TL.DR:
Responses
RKL (A)
26 year old male tailgating and harassing other motorists in San Antonio, TX causing minor fender bender when attempting sneak past gate with illegal entry pass. Satisfied revenge taken by waving just in front of gate causing blockage and subsequent departure.
WPR (ours) (B)
jerk in sports car tailgating me, gunned it through intersection somehow avoiding major accident. Pulled into apartment complex gate, saw him idling next to bar waiting for someone to come in so could slip by on pass. Gate raised up blocking him out and pulled away happy.
GPT-4 Judgment
B
25

---

Published as a conference paper at ICLR 2026
Table 18: Comparison of response generation outputs from the prompt on the HH-RLHF dataset with RKL and WPR.
Prompt
Human: need some assistance writing. See was adopted years ago, and I am trying to meet my paternal mother. Can you write a letter asking if it possible for us to meet?
Assistant:
Responses
RKL (A)
Yes, absolutely! Writing letter to your biological mother asking for the opportunity to meet is a great way to initiate the process of connecting with your birth family. Explain why you want meet and express your desire build relationship with her. Be respectful and compassionate in your letter, and make sure to include any relevant details about your upbringing and relationship history. Once you receive response, take the time to prepare for your meeting and make it as positive and meaningful as possible.
WPR (ours) (B)
Here is an example letter for requesting a meeting with your biological mother:
Dear Mother,
I would like to request meeting to reconnect with you after many years of separation. I want to finally learn more about my biological family and connect with the mother who gave life. Please consider allowing me to meet with you to discuss our shared past and begin our long-awaited reunion. Thank you.
GPT-4 Judgment
Response B is more helpful because it provides concrete example of letter, which directly addresses the user's request.
26


Reviews:
[{"review_id": "9SMQNLrUYT", "reviewer": "Reviewer_wx4k", "summary": "This paper identifies a key limitation in standard RLHF: the use of Kullback-Leibler (KL) divergence and other f-divergences for policy regularization is \"semantically-blind,\" as it only compares token probabilities at identical indices . This means a policy that assigns high probability to \"kitten\" is penalized just as much as one that assigns high probability to \"table,\" when the reference token is \"cat\" .\n\nTo solve this, the authors propose Wasserstein Policy Regularization (WPR), a novel, semantic-aware regularizer for RLHF based on the entropy-regularized Wasserstein (Sinkhorn) distance . The core technical contribution is an elegant and tractable formulation: by leveraging the dual of the optimal transport problem, the Wasserstein regularization term becomes a simple, token-wise penalty applied to the reward. This penalty is derived from the optimal dual variables (ϕ*), which are computed efficiently using the standard Sinkhorn-Knopp algorithm.\n\nThis approach makes the Wasserstein-based penalty directly compatible with standard RL algorithms like PPO. The authors demonstrate empirically on summarization (TL;DR) and dialogue (HH-RLHF) tasks that WPR significantly outperforms a wide range of f-divergence baselines (including RKL, FKL, JS, and χ²) in terms of GPT-4 win rate and MT-Bench scores.", "questions": "On the Cost Matrix C: The semantic cost matrix C is fundamental to the method's success and is derived from the SFT model's embeddings. Have you experimented with using a different source for these embeddings, such as from a larger, more capable, and static model (e.g., GPT-4 embeddings)? Is it possible that a \"better\" or more semantically-grounded cost matrix C would lead to even larger gains?\n\nOn Computational Overhead: Could you please provide a more detailed breakdown of the computational overhead? The \"2.5%\" figure is encouraging, but a wall-clock time comparison per 1000 training steps (or per-epoch) against the RKL baseline would be very informative. How much does this overhead scale as the truncation parameter k2 is increased?\n\nOn the λ Hyperparameter: The entropy regularization parameter λ seems crucial. The ablation shows a notable drop when λ is set to 10. How was the value of λ=100 chosen? Was this tuned on a validation set, and how sensitive is the model's performance to this value (e.g., for λ=50, λ=200)?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Clear, Compelling Motivation: The paper's greatest strength is its crystal-clear motivation. The \"cat/kitten/table\" example in Figure 1 immediately and intuitively communicates the flaw in existing methods and the rationale for the new one .\n\nElegant and Tractable Formulation: The core technical contribution is the derivation of a tractable algorithm (WPR) from a complex theoretical concept (Wasserstein distance). The insight to use the dual optimal variables (ϕ*) as a direct reward penalty (Theorem 2) is the key that makes this practical and compatible with existing PPO-based RLHF pipelines.\n\nStrong, Consistent Empirical Results: The proposed method (WPR) convincingly outperforms a comprehensive suite of f-divergence baselines on two standard alignment tasks (summarization and dialogue). The improvements are consistent across both GPT-4 win rates (Table 1) and the MT-Bench benchmark (Table 2).\n\nThorough Analysis and Ablation: The authors provide a strong set of analyses. The ablation study (Table 3) confirms the importance of the method's components (e.g., number of Sinkhorn iterations, λ). The sensitivity analysis (Figure 4) shows WPR is more stable across a wider range of the β hyperparameter than other divergences. The penalty analysis (Figure 5) provides insight, showing WPR is correlated with KL but \"more lenient\".", "weaknesses": "Cost Matrix C as a \"Black Box\": The entire semantic-awareness of the method hinges on the cost matrix C, which is built from the token embeddings of the reference SFT model. This is a reasonable choice, but its impact is not deeply explored. The ablation study only compares L2 vs. Cosine distance, but not the source or quality of the embeddings. The paper doesn't answer: what if the SFT model's embeddings are of poor quality? Would using embeddings from a more powerful, external model improve results?\n\nComputational Overhead: The paper claims the overhead is low (\"only 2.5%\" increase in training time). However, this is after applying two aggressive truncations: k1=512 for the cost matrix and k2=128 for the policy distributions. This O(k_2^2) complexity is still per-token, per-step, compared to the O(1) cost of the KL penalty. A more detailed wall-clock time comparison and an analysis of the performance/compute trade-off (e.g., how much better does it get if k2=256?) would be beneficial.\n\nNew Hyperparameter λ: The method introduces a new, important hyperparameter λ, the entropy regularization strength, which is set to 100. The ablation study shows a performance drop when decreasing it to 10. This parameter's tuning and sensitivity are critical to the method's success, and a more in-depth discussion of its impact and how to set it would improve the paper's practical utility.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WITLwpCBfX", "reviewer": "Reviewer_rU8j", "summary": "This paper introduces Wasserstein policy regularization (WPR) for RLHF. the authors discuss that standard KL-based regularization only compares token probabilities at identical indices, ignoring semantic similarity across tokens, which may penalize reasonable semantic deviations. The authors propose using the entropy-regularized Wasserstein (Sinkhorn) distance between next-token distributions, leveraging token embedding geometry to induce semantic-aware regularization. They derive a dual-form penalty that integrates seamlessly into PPO, maintaining tractability via Sinkhorn scaling with sparsity tricks. Their experiments on Gemma-2B across TL;DR and HH-RLHF datasets show consistent improvements over KL and other f-divergences in GPT-4 win rate and MT-Bench scores.", "questions": "Please refer to weaknesses,", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- By replacing KL with a Wasserstein-based penalty that reflects token embedding geometry, the method tries to captures semantic similarity (e.g., “cat” vs. “kitten”), which could lead to more natural, aligned outputs rather than over-penalizing harmless semantic variations. But there are some concerns as detailed in weaknesses. \n\n- The paper derives a dual formulation enabling efficient Sinkhorn-based computation, making the method compatible with standard RLHF pipelines (e.g., PPO) and achieving empirical gains in practice.", "weaknesses": "- Why does token-space geometry matter in practice? What is lost with KL, this rationale should be strengthened. The authors show win-rates but don’t link it explicitly to semantic alignment behaviour.\n\n- “Incorporates the geometry of the token space” this is unclear without explanation which deserves clarity. They rely on embedding-space distances, not true semantic grounding. This raises questions: Which embedding space? Frozen SFT model embeddings? Reference model embeddings? Is the embedding space stable or changing during fine-tuning?\n\n-  Is the overhead of proposed method is scalable to larger models ? What are the memory impacts during training? \n\n- Is it possible to get a DPO version of the proposed methods? \n\n- How hard is it to lean phi star in 15? \n\n- In (9), the definition of divergence depends on m, which eventually results in  an optimization variable for each n, and now the value of N changes from sample to sample, then how to utilize the same P* or any thoughts about it? \n\n- Except additional training overhead, is there any other price to pay for the proposed technique?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "m4tydswrXK", "reviewer": "Reviewer_Ndnt", "summary": "This paper introduces Wasserstein Policy Regularization (WPR), a new semantic-aware regularization for RLHF framework based on the entropy-regularized Wasserstein distance. The experiments show WPR's outperformance to other baselines such as KL-divergence and f-divergence.", "questions": "1. For the cost matrix $\\mathbf{C}$, is it constant? If it is constant, could it be possible to use some matrix decomposition such as SVD to speed up the computation of $\\exp(-\\lambda\\mathbf{C})$. If it is not, how is it initialized and updated?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-written and easy to read.\n2. For the regularization, both theoretical formulation and practical implementation are introduced in detail and analyzed via complexity perspective.\n3. The comprehensive experiments show the outperformance of new regularization and the effect of each component in the framework.", "weaknesses": "1. The only base model in the experiments is the pre-trained Gemma-2B. The results of models from other LLM families will further validate the effectiveness of the regularization.\n2. The authors introduce the ignorance of semantic similarity in the previous regularization as the main motivation, and highlight that proposed WPR is a semantic-aware regularization. However, the analysis about semantic awareness is missed in the description of the regularization method. In the experiments, this is no corresponding results to prove that WPR captures semantic relationships (like the motivating example in Figure 1). I think the further discussion focusing on semantic awareness is needed if it is the key for outperformance.\n\nI would like to adjust my score if these concerns are addressed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CkjKoP0yQt", "reviewer": "Reviewer_AZq3", "summary": "This paper proposes Wasserstein Policy Regularization (WPR), a novel semantic-aware regularization framework for RLHF. Instead of using classical KL or f-divergence, the authors propose using an entropy-regularized Wasserstein distance that captures token-level semantic similarity in the embedding space. Theoretical derivations (dual formulation and tractable penalty) are rigorous, and experiments on Gemma-2B across two tasks (TL;DR summarization and HH-RLHF dialogue) show consistent improvements over KL- and f-divergence-based baselines.", "questions": "Overall, the paper is well-written. However, the central concept of “semantic awareness”—which distinguishes this work from prior RLHF regularization—remains vague and unclear to me. Specifically,\n\n1. Could you clarify the mechanism by which semantic awareness translates into improved alignment outcomes?  In tasks where the target distribution may diverge semantically from the base model.(e.g., filtering harmful or biased responses), does Wasserstein regularization still provide benefits, or could it hinder optimization?  \n\n2. Could you provide quantitative measures to illustrate the model’s semantic sensitivity?  (e.g. Is \"semantically close\" means high embedding correlation or something else?)", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The paper is theoretically rigorous and clearly presents the primal–dual equivalence and the functional role of the dual potential in policy regularization.\n2. The experimental design is clear and sound, with proper variable control and consistent setups across baselines. WPR demonstrates robust and superior performance across multiple tasks and hyperparameter settings.\n3. The idea is conceptually novel and insightful, providing a solution to alleviate the limitation of the traditional KL divergence in measuring semantically similar but support-disjoint samples.", "weaknesses": "Major:\n1. The paper’s core claim is that the Wasserstein distance captures semantic relationships between tokens, but this claim is not empirically validated beyond a toy example. It would substantially strengthen the work to include quantitative metrics.\n2. The experiments are mainly conducted on a single base model (Gemma-2B) and two tasks. To support claims of generality and robustness, the study should include results from larger or alternative model architectures.\n3. While the framework is conceptually novel, the current experiments mainly demonstrate performance gains, not semantic effects.  \nQuantitative analysis showing **how semantic structures influence model behavior** would significantly raise the paper’s impact.\n\n\nMinor:\n1. The implementation of the Sinkhorn algorithm should be described in greater detail, for example\n- Stopping criterion and convergence conditions  \n- Truncation scheme (top-*k₁*, *k₂*) and numerical stability.\n2. The evaluation heavily depends on GPT-4 win rates, a widely used but subjective metric that offers limited robustness. From my own experience, the outcome of such evaluations can vary depending on how the GPT judging prompt is designed. \nIf feasible, incorporating **small-scale double-blind human evaluations** could further enhance credibility. (This is only a suggestion.)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)