Robust LLM Alignment via Distributionally Robust Direct Preference Optimization
Authors: Zaiyan Xu
Sushil Vemuri
Kishan Panaganti
Dileep Kalathil
Rahul Jain
Deepak Ramachandran
Venue: NeurIPS 2025 poster ()
Keywords: Large Language Model Fine-tuning
Reinforcement Learning from Human Feedback
Direct Preference Optimization
Distributionally Robust Optimization
Primary area: deep_learning
TL;DR: We address distributional shift among diverse preferences with robust DPO: Wasserstein DPO (WDPO) and Kullback–Leibler DPO (KLDPO). Finite-sample guarantees, tractable gradient-based algorithms for hard DRO objectives, strong empirical robustness.

Abstract: A major challenge in aligning large language models (LLMs) with human preferences is the issue of distribution shift. LLM alignment algorithms rely on static preference datasets, assuming that they accurately represent real-world user preferences. However, user preferences vary significantly across geographical regions, demographics, linguistic patterns, and evolving cultural trends. This preference distribution shift leads to catastrophic alignment failures in many real-world applications.  We address this problem using the principled framework of distributionally robust optimization, and develop two novel distributionally robust direct preference optimization (DPO) algorithms, namely, Wasserstein DPO (WDPO) and Kullback–Leibler DPO (KLDPO). We characterize the sample complexity of learning the optimal policy parameters for WDPO and KLDPO. Moreover, we propose scalable gradient descent-style learning algorithms by developing suitable approximations for the challenging minimax loss functions of  WDPO and KLDPO.  Our empirical experiments using benchmark data sets and LLMs demonstrate the superior performance of  WDPO and KLDPO in substantially improving the alignment when there is a preference distribution shift.

Robust LLM Alignment via Distributionally Robust Direct Preference Optimization
Zaiyan Xu', Sushil Vemuri', Kishan Panaganti Dileep Kalathil',
Rahul Jain², Deepak Ramachandran"
Texas A&M University, Tencent Al Lab, Google DeepMind.
Emails: {zxu43, sushil22, dileep. kalathil)@tamu edu, kpb. research@gmail.commmm com, (rahulajain, ammachandrandd@gooooe
Abstract
A major challenge in aligning large language models (LL.Ms) with human prefer- ences is the issue of distribution shift. L.M alignment algorithms rely on static preference datasets, assuming that they accurately represent real-world user prefer- ences. However, user preferences vary significantly across geographical regions, demographics, linguistic patterns, and evolving cultural trends. This preference distribution shift leads to catastrophic alignment failures in many real-world appli- cations. We address this problem using the principled framework of distributionally robust optimization, and develop two novel distributionally robust direct preference optimization (DPO) algorithms, namely, Wasserstein DPO (WDPO) and Kull- back-Leibler DPO (KLDPO). We characterize the sample complexity of learning the optimal policy parameters for WDPO and KLDPO. Moreover, propose scalable gradient descent-style learning algorithms by developing suitable approxi- mations for the challenging minimax loss functions of WDPO and KLDPO. Our empirical experiments using benchmark data sets and LLMs demonstrate the supe- rior performance of WDPO and KLDPO in substantially improving the alignment when there is preference distribution shift.
Introduction
The alignment of large language models (LLMs) with human values and preferences is central objective in machine learning, enabling these models to produce outputs that are useful, safe, and aligned with human intent. Since L.MMs trained on vast, diverse datasets using self-supervised learning, an additional aligmment phase is often required to refine their behavior based on human feedback. widely adopted approach for this is Reinforcement Learning from Human Feedback (RLHF) (Christiano et al., 2017; Ziegler et al., 2019; Ouyang et al., 2022), which involves train- ing reward model using human preference data and optimizing the LLM using reinforcement learning (RL) approaches, such as proximal policy optimization. More recently, Direct Preference Optimization (DPO) has emerged as an alternative that simplifies the aligmment process by directly optimizing model parameters based on human preferences without requiring an explicit reward model. These alignment techniques have played a crucial role in improving the ability of LLMs to generate responses that adhere to human expectations and societal norms, leading to today's powerful chat models (Achiam et al., 2023; Touvron al., 2023).
Despite the importance of the LLM alignment problem, RLHF and DPO remain fundamentally challenging and fragile, mainly due three reasons. (i) Diversity of human preferences: Standard RLHE/DPO approaches implicitly assume that human preferences can be accurately captured by a single reward function. In reality, human preferences are highly diverse, context-dependens and distributional, making it infeasible to represent them with one-size-fits--e fits-alt- optimization framework
Work done postdoctoral researcher at the Califoria Institute of Technology.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

(Zhao et al., 2024; Durmus et al., 2024). Standard preference-learnigg methods tend to skew toward the preferences represented in the majority of training data, disproportionately penalizing minority opinions and reinforcing biases (Chakraborty et al., 2024). (ii) Reward hacking: The quality of human preference feedback inherently noisy, ambiguous, and inconsistent, as they are collected from human annotators who may lack domain expertise, exhibit labeling fatigue, or hold conflicting opinions (Zhang et al., 2025; Wu et al., 2025), which can often lead to misaligned preference estimation. This issue is exacerbated by reward hacking, where models learn undesirable shortcuts to maximize the estimated reward function, generating responses that appear aligned but deviate from genuine human intent (Amodei et al., 2016; Skalse et al., 2022; Eisenstein et al., 2024). (tii) Distribution shift: Alignment algorithms use static preference datasets for training, collected under controlled conditions. However, the preferences of real-world users can often be out-of- disstribbution from that of the training data, depending on the geographical region, demography, linguistic patterns, and emerging social trends, among many others. A model aligned using a specific fixed dataset may fail catastrophically when deployed to users whose preference distribution does not match that of the training data (Casper al., 2023; Le Vine et al., 2023; Kirk al., 2024).
Prompt: netvu bew:. Lualing sallly stressed vark lstely. to you have any advice on how to warage stness and atay productiive?
Preferenco/Reward Model
Proference Shift
Distributionally Robust DPO
proference P2 proference
Training population
Tisting population
Prafeunce seantarry
ADPOOKLOPO
Preferencf robusl completon
Sia suzzy you' feeling preat-treatie, breaks, and be kind zo yeszzeie:
Cai Prioritine breake, talk your manayee seedei.
Mon-mbbust RUNIIOPO
C1iolks completion
Sentimant
Practicality
Iraiiing besting
Figure 1: If the training population predominantly uses preference model 1 (P1), non-robust RLHF/DPO model will favor Completion 1 (CI). However, deploying this model a test population that prefers model 2 (P2), which favors Completion 2, leads to poor performance. Our distributionally robust DPO (WDPO/KLDPO) addresses this by optimizing across an uncertainty set of preference models, ensuring robust performance under preference shifts.
In this paper, we address the fragility of the LLM alignment using DPO, with a particular focus on the challenges arising from the preference distribution shift. DPO reduces the alignment problem to a supervised learning problem. It known that the performance of supervised learning algorithms degrades significantly in the out-of diitr---iittiiitionn setting (Taori et al., 2020; Koh et al., 2021), which exacerbated due to the realistic distribution shift scenarios arising in the LLM deployment. Dis- tributionally robust optimization/leamiio framework has been recently used to address the issue of distribution shift various settings (Duchi and Namkoong, 2021; Kuhn et al., 2019; Chen et al., 2020). This framework considers uncertainty set of data distributions around nominal distribution (typically the training data distribution) and solves a minimax optimization problem to minimize the expected loss, where the expectation is taken with respect to the distribution in the uncertainty set that maximizes the loss. The distributionally robust learning approach has been successfully applied, with theoretical guarantees and scalable algorithms. in supervised learning (Chen and Paschalidis, 2018; Namkoong and Duchi, 2016; Levy et al., 2020), multi-armed bandits (Si et al., 2020; Yang et al., 2023), and reinforcement learning (Wang and Zou, 2022; Panaganti et al., 2022; Zhou et al., 2024; Xu et al., 2023). This motivates us to address the following questions:
Can distributionally robust learning mitigate the impact of distribution shift in DPO- based LLM alignment? What theoretical guarantees can be established for such methods? ? How can we design traciable, gradient-based algorithms to implement them? How do these approaches empirically improve alignment performance?
Distributionally robust learning for LLM alignment presents challenges beyond standard supervised settings. In supervised learning, distributional robustness is often tractable due well-behaved convex losses. In contrast, RL poses more complex forms of distribution shift, both exogenous (e.g.. user preference drift) and endogenous (e.g., mismatch between the learned and logging policies). Although DPO is framed as supervised objective, its likelihood-ratio formulation, based on pairwise comparisons and a reference policy, derives from KL-regularized reward maximization, linking it closely to RI. Like imitation learning and offline RL, which address RL problems through supervised proxies (e-g., behavior cloning, Q-function regression), DPO inherits the same instability and sensitivity to distribution shift. These challenges are amplified in the distributionally robust

---

setting, where the non-convex min-max objective is especially hard to optimize at LLM scale, making standard alternating-gradee
methods unstable and impractical. We answer the above questions affirmatively and address the associated challenges through the following contributions:
1. To the best of our knowledge, this is the first work to propose a unified mathematical and algorithmic framework for addressing preference shift in LLM alignment through distributionally robust optimization. Our formulation leads to two robust DPO variants, Wasserstein DPO (WDPO) and Kullback-Leible DPO (KLDPO), with provable guarantees. In particular, for log-lineat policies, we show that the estimation error of the robust policy parameters converges at rate of Onn
2. We develop computationally tractable gradient descent algorithms for WDPO and KLDPO that can be seamlessly integrated into existing LLM alignment pipelines.
3. Empirically, we show that standard DPO is sensitive to preference distribution shift, leading to degraded performance when training and evaluation rewards differ. In contrast, our robust variants, WDPO and KLDPO, consistently achieve superior performance across diverse aligmment tasks. For example, we fine-tune LLaMA-3.2 3.2-BBBBBBBBBBBBB Instructt and LLaMA-3.1 8B-Instruct models on prompts from the HelpSteer2 dataset (Wang et al., 2024b) using preferences generated by the ArmoRM reward model (Wang et al., 2024a), and evaluate them on distinct reward objectives from the OpenI.LM Leaderboard (Fourrier et al., 2024).
2 Related Work
Robust RLHF:
Bai et al. (2022) proposed to adjust weights on the combination of loss functions based on different topies (harmless vs. helpful) for robust reward learning. Chakraborty et al. (2024)
proposed to learn multiple reward functions for different sub- uub-populations through expectation- maximization approach, and a robust policy based on these rewards via a max-min optimization, which is different from our distributional robust learning approach. Padmakumar et al. (2024) augmented the existing binary preference datasets with synthetic preference judgments to estimate the diversity of user preferences. Yan et al. (2024) proposed Bayesian reward model ensemble to quantify the uncertainty of reward estimation and used it to reduce reward overoptimization. Bukharin al. (2024) proposed robust RL.HF approach for addressing the preference data corruption problem.
Robust DPO: Huang et al. (2025) proposed XPO that implements the principle of pessimism in the face of uncertainty via regularization with the >=-divergence for avoiding reward hack- ing/overoptimizztion w.r.t. the estimated reward. Ramesh et al. (2024) proposed Group Robust Preference Optimization (GRPO) to address the diverse preference problem by modeling the total loss as a weighted sum of individual DPO losses computed on separate preference datasets, and optimizing for the worst-case weighting. In contrast, our approach does not assume access such a group structure and instead directly models distributional robustness over a single dataset that implicitly aggregates diverse preferences. Chowdhury al. (2024) considered the setting where e-fraction of the preference labels in the training dataset is corrupted and proposed a noise-robust algorithm to mitigate its effect, assuming the knowledge of Wu et al. (2024) focused on adapting the DPO penalty parameter 3 to handle varying data quality within the training set. The most related work is Wu et al. (2025), which applies distributional robustness to mitigate data corruption and noise in preference data. Unlike our work, it does not address distribution shift or provide theoretical guarantees, and lacks empirical evaluation on preference distribution shift. Concurrent to our work, Mandal et al. (2025) proposed distributionally robust version of RLHF and DPO, using total variation (TV) uncertainty sets. However, their theoretical analysis offers the natural policy gradient (NPG) style optimization convergence guarantees for the loss function. In contrast, we go one step further: by leveraging strong convexity, we establish loss, but for convergence of the policy parameters. Our analysis is algorithm-agnostic and applies to any solver capable of optimizing the robust DPO loss. Additionally, our formulation uses KL. and Wasserstein uncertainty sets, which are more standard in large-scale LLM alignment.
sample guarantees not just for the learning
Distributionally Robust Learning: Distributionally robust learning is statistical learning frame- work designed to enhance model performance under distributional shifts between training and test data (Chen and Paschalidis, 2018). 8). It employs a minimax approach where an adversary maximizes the expected loss by shifting the test distribution within a specified uncertainty set, while the learner minimizes this adversarial loss. This approach using the f'divvereence (Namkoong and Duchi, 2016; Duchi and Namkoong, 2021; Levy et al., 2020) and the Wasserstein metric (Mohajerin Esfahani
3

---

and Kuhn, 2018; Kuhn et al., 2019; Gao et al., 2022) have gained significant attention recently. Distributionally robust algorithms have been developed to address problems in supervised learning (Chen and Paschalidis, 201 8; Namkoong and Duchi, 2016; Levy et al., 2020), imitation learning (Bashiri et al., 2021; Panaganti et al., 2023), multi-armed bandits (Si et al., 2020; Yang et al., 2023), and reinforcement learning (Panaganti et al., 2022; Zhou et al., 2024; Shi and Chi, 2024; Yang et al., 2022; Panaganti et al., 2025).
3 Preliminaries
Notations: We use calligraphic letters for sets, e.g., S. |I-II denotes the Euclidean norm. When Σ is a positive semi-definite matrix, we write || = vx Σx as a semi-norm of æ. For any measure P. we use P,, denote the empirical distribution constructed using n i..d. samples, .I...., Zn. from P, i.e., Pa = (1/n) E oz., where 62 is the Dirac measure. We use a to denote the sigmoid (standard logistic) function. We use ((2; 0) and (0) denote the loss incurred by sample z with policy parameter 0. For any set 2, (2) is the set of all Borel measures over 2. For any positive semi-definite matrix Σ, Amin (E) and Amax (E) denote its smallest and largest eigenvalues.
Wasserstein Distance:
For a given set Z. equipped with a metric d, the Wasserstein distance of order p between two distributions M, P(Z) is defined (see Villani et al. (2009)):
J2x (x, x' (rdr, dx'): y has marginal distributions μ4, V rEP(2x2) {
Wo(u min
Kullback-Leibler Divergence:
For any two probability distributions P and Q defined on Z, the
Kullback-Leibker (KL) divergence is defined DRL (P || Q) Σxez P(z) log) log(P(z))A() (P=))/((=)))
Reinforcement Learning from Human Feedback: The RLHF paradigm consists of three steps:
Step 1: Supervised Fine-tuning (SFT). SFT involves fine-tuning a pre-trained LLM through supervised learning on high-quality data, curated for the downstream tasks.
Step 2: Reward Modelling. the second step, given any context & E S, two responses a² aa2 are independently sampled from the behavior policy T² (typically the SFT policy TSFT). Then, (human) labeler provides a preference response between these responses. We assume that the preference responses are generated according the Bradley-Terry (BT) model (Bradley and Terry, 1952):
exp (r' (s,))))
(*(aa > a2 a2 exp (xp(rr(s,a))) + exp (r* (r*(s)²²))
where a' > a² denotes a being preferred over a², and r* is the underlying unknown reward function. We use a²", a' to denote the preferred and dis- pieffrrerd responses, respectively. We assume access to a static dataset of comparison, prompt (context) distribution n², a at ss are independently sampled from TSPT, and the preferences responses are sampled from the BT model P". With D, we can learn a parameterized reward model To(s, a) by minimizing the maximum likelihood estimation (MLE) loss,
= [((i,, at" at) )2)) where 8, are sampled from some initial
"a)") C CRLHF (ro: D) = (a. (l)) (ro((s - ro(s, (oossa))) d'))
Step 3: RL Fine-Tuning In the final step, the optimal policy x" under the reward r4 is obtained by solving the KL-regularized reward maximization problem given by
maxE [Ba- (s) (s, (,, a)] ( | ]| Teee(- s))]
(2)
where 3 is parameter controlling the deviation from the base reference policy Trof-
Direct Preference Optimization (DPO): The DPO approach (Rafailov et al., 2023) leverages the fact that the unknown reward function can be expressed in terms of the optimal policy and the reference policy. Formally, given any reward function r", the optimal solution of Eq. (2) takes the form n²" ra 8) . Tres (a | s)exp (r" ax) where 7" (s) denotes the partition (normalizing)
function. Rearranging the above, we get r" (ss a) = Blog (aasss
[ee((a]s) + Bog (s) for all (s,
Substituting this into Eq. (1), the optimal RLHF policy 'T" satisfies the preference model:
(* x a² ()) = =o
o Blogg (a= 8)
<'(a' s) 3 log Trer(a2
(a² ()

---

Using the preference response dataset D, we can learn the optimal policy directly by minimizing the MLE loss for a parameterized policy Tg.
CDPO (To; D) (e,a -D logo BBog
Ta(a" 8)
3 log
To (a' Tree (a (s)
8)
(3)
Tyef (a [s)
Distributional Uncertainty Sets: Given any p 0 and po E P(Z), define the distributional uncertainty set as
Po): = {P P()) D(P. S p):
where -)--) is some distance metric between two probability measures, e.g., Wp and DKL.
Distributionally Robust DPO
In this section, we formulate our Wasserstein DPO (WDPO) and Kullback-Leibler DPO (KLDPO).
Sampling Procedure: As described in Section 3, prompt S is drawn from initial distribution 4², and two responses a' a² Mi.d T-(- s) are sampled independently (with лº = #SFT in practice). Following Zhu et al. (2023), define y E {0, 1] to indicate preference: y if a! > a² and y = 0) otherwise. The label y is drawn from a Bernoulli distribution defined by the BT model P* The full mata-generating distribution is given below.
Definition 1 (Joint data-generating distribution), Consider the product space Z == SxAxAx We define the nominal data- generating distribution as
(" a² v) - (s)) |
(a) 8) [1i==1) a² 12)s)+ 11)+1(==0)) (1=00) P (a² a² $)]
We will also denote z (s,a", aa,") 2 and (2(2) = ( a aa,sa',,,,0 a²,,)) We assume that P generates the dataset D {= used for learning, i.e., 2i ~ po.
4.1 Distributionally Robust DPO
From the DPO objective (Eq. (3)), we define the pointwise DPO loss function as follows
[22; 0) = yy log a (Bha(s, aa, a²)) (1 - y) log (Bho a²,a')) a))),
(5)
where ho(s, ad, a²) == log Frar(a
[s(a [s) log xeat (a= is the preference score of answer a relative to
another one a² (but parameterized in policy parameter 0). Let Pl: Po) be distributional uncertainty set centered around po with radius p O. Following the principles of distributionally robust optimization (DRO), we formulate the distributionally robust DPO objective as:
min max [2~p(1(2;0)) 0 PEP(x:P)
Intuitively, we aim to find the best policy under the worst-case data distribution.
When we have a Wasserstein uncertainty set Pw,,, i.e., Eq. (4) equipped with the p-th order Wasser- stein distance, we define the Wasserstein DPO (WDPÓ) loss as follows
W (0: p)
sup PEPw, (p:Po)
Similarly, given a Kullback-Leible uncertainty set (pl Po), we define the KLDPO loss as follows
C KL (0; PePku.(p:Po)
sup Eg- [m-P(((; z)].
(8)
When the nominal distribution Pa is replaced with its empirical counterpart, i.e., P (1/n) Σn1 6z., where 21...., 2n are n i.idd. samples from po, we use C. (#: p) and L (8: to denote the empirical WDPO and KLDPO losses incurred by the policy parameter f. respectively.
Theoretical Analysis
In this section, we present the sample complexity guarantees for our WDPO and KIDPO algorithms. We make the following assumptions for the rest of the papers.
5

---

Assumption 1 (Log- linear policy class). Let wh S A Rd be known d-dimensional,eeng feature map- ping with maxs,a [ a)ll S Assume bounded policy parameter set θ == {0 R ||01|2 -
BJ. We consider the following class of log-linear policies:
exp (f (ss
I1 no: To EnrEA exp (0 ()s,,')))
Remark 1. This is standard assumption in the theoretical analysis of the RL algorithms (Agarwal el al., 2021; Modi el al., 2020), RLHF (Zhu el al., 2023), and DPO (Nika et al., 2024; Chowdhury al., 2024). Our analysis can be extended to the neural policy class where HT bss a) is replaced fe(s,a), where fe is neural network with twice differentiability and smoothness assumptions.
We also make the following data coverage assumption the uncertainty set ((:; P"). Assumption (Regularity condition). There exists X > 0 such that
Ep = Reeaa w)) (()(p(((sss)))))))))) a-)) whs a²) (1b(s (8, a²)
(2) XT, YP E P(p;
Remark 2. We note that similar assumptions on data coverage under linear architecture models are standard in the offline RL literature (Agarwal et al., 2019; Wang et al., 2021; Jin et al., 2021). Implicitly, Assumption imposes X S Amin (E which means that the data- gata-generating distribution Pa has good coverage.
5.1 Estimation Error for WDPO
Let 0" E argmingge CDPO (f) be the ground-truth optimal policy parameter with respect to the true nominal distribution and let its empirical counterpart be 6, E argmingge (B). Now for the robust policy parameters, we let g argmin LW (6; p), and let its empirical counterpart be gW E argmingge 2n (0; p). Now, present our main result on the sample complexity result for the convergence of the robust policy parameter.
Theorem (Estimation error of (W)). Let 8 E (0, 1). With probability at least - 6, we have
8K2 1og(2/6)
10
1
>²X2n
188 where y (1+c4087) and K = (logo(-4BB)) X is the regularity number defined in Assumption
Proof sketch. Strong duality of Wasserstein DRO (see Gao and Kleywegt (2022) and Corollary 1)
helps us reduce the difference | (0; LWW (0:p) to the concentration [E (= (z; 0)
1(2; 0)] is called the Moreau- Yosida reg- ularization of -/ with parameter 1/n.. We show that, for all 17 2 0, all L, are uniformly bounded. We then use Hoeffding's inequality to obtain concentration. Detailed proof is in Appendix B.2.
Ez-p; (2; 0)]], where 1,(2; f) = in (nd
Next, when Assumption is in place, we can show that g = Ez-p [1: 8)] is 7-strongly convex w.r.t. the positive definite norm is yA-strongly convex but w.r.t. 1.112. A detailed proof is provided in Appendix B.3.
|- /.... Further, by the property of supremum,
show that LW
Decompose LW (OW) LW (OW) into three terms: L (oW; :) LW (oW: ;)). CW (OW; C. (aW: :)) and CW (oW; :p) LW (0W:p) The second term is non-positive since OW is the min- imizer of CW Now we apply the concentration of the WDPO loss function (see Lemma 9 in Appendix B.2) to CC (0W" p) -W (0W;p) and [C (OW; :p) - CW (0Wip) Finally, we use the property of strongly convex function (Lemma 5) on L acquire the policy parameter convergence. The detailed proof is in Appendix B.4.
We state the convergence result for DPO to facilitate comparison with its robust counterpart. Proposition 1 (Estimation error of (non-robust) DPO). Let E (0, 1) and 0.
4 |e,, 0' 2
(d log + 2AB",
n2n
38488 IIteweye and Ep 1. (v
with probability least 1 8 and where 7 (), (ss,, ( at) (hss,, TT is the sample covariance matrix.
T

---

Algorithm WDPO Algorithm
Algorithm 2 KLDPO Algorithm
I: Input: Dataset D ((s; a." a)) refer-
I: Input:
Dataset DD {(s;, di d))=. refer- ence policy Tref, robustness temperature pa- rameter т, learning rate 7), initial policy na.
ence policy Tyef, robustness hyperparameter
Po. learning rate 77, initial policy za.
2: while 0 has not converged do
2: while 0 has not converged do
3: Calculate the non-robust DPO loss
3: Approximate the worst-case kernel
CDPO(TTT;
(T;; D) according to Eq. (3)
P(() (xp((L/T)((<z; 0)
4: Calculate the gradient regularizer loss
-(1nn) E 1(: 0)))
R(mo; D) = pa 1))
5:
Calculate the approximate WDPO loss
4: Calculate the approximate KLDPO loss
L (0: =
_DPO (TPO D) + (TT:)))
(KL((::)) = E (2;;0)
6:
- 0 - (0:D)
5: 0
7: Output:
6: Output: T
A matching result can be derived as special case of Chowdhury et al. (2024, Theorem 4.2). We provide an independent proof with precise constants Appendix B.1.
Remark 3. We would like to note that the estimation error rate of convergence for WDPO is |10h gW 122 = (n-1//), from Theorem 1. The estimation ervor rate of convergence for (non- robust) DPO 110,, 0" + = ((n=1/2), from Proposition 1. So, the estimation error rate of convergence for WDPO is worse than that of (non-robust) DPO. This arises due to significant challenges exclusive to the robust setting. For example, for the non-robust DPO, we can calcu- late the closed-form expression of (o(L/n) Ei=1 (21; 0) (see Eq. (20)). This allows to write IIV Σi= 12,;; ((( in quadratic form and then obtain a concentration using Bern- stein's 's inequality. However, for WDPO, we note that V VeLW (@W) & suppePw, V&Es<p [ @W)], and the non-robust approach will not work for the robust setting. Developing analysis techniques to achieve a better rate of convergence for robust DPO is an open question.
5.2 Estimation Error for KLDPO
Let gK, E argmingce (0; p). and let its empirical counterpart be @KL argmingge CKI p). The convergence analysis for the KLDPO loss and policy parameter closely parallels that of Wasser- stein DPO. We present the main theorems below and defer detailed proofs to Appendix C.
Theorem 2 (Estimation error of OKL), Let 8 E (0,1), 1). With probability at least 1 - 8, we have
8X exp (L/2) log(2/0)
|10KL n
OKL
S
2²22n
82848B
where 7 (1te4=Byz X is the regularity condition number defined in Assumption 2, 0
Amin (Ep.). X. A. are some universal constants, and L is upper bound on the loss function L. Remark 4. The exponential constant in the upper bound is characteristic of distributional robust optimization with KI. uncertainty set Hu and Hong (2013, Proposition 2). Similar exponential constants appear the theoretical analysis of the distributionally robust RL (Zhou el al., 2021: Yang et al., 2022; Panaganti and Kalathil, 2022; Xu et al., 2023). Both WDPO and KLDPO have O(n (111)/4) policy parameter convergence. An empirical comparison given in Section 7.
1/4)
6 Tractable (Approximate) Algorithms
While our distributionally robust DPO formulations enjoy finite-sane sample guarantees, is computation- ally challenging to solve the min-max objective of Eq. (6) using stochastic gradient descent methods. Though many min-max optimization problems can be solved by alternating gradient descent methods, our problem is not directly amenable to such an approach as we do not have direct control over the data distribution P E P(p Po) which not parameterized. Moreover, the preference data are generated according to the nominal distribution P .and we do not have data samples from any other
7

---

5.0.75
JuUh! WOPO-00 (oU0)
Nott Hleeear Cormennoo
WOTD-35
1.8
WarD 1913
0.50
3.4
Appha
Alpha
Algha
Apha
Figure 2: DPO, WDPO, and KLDPO in Emotion Alignment. Models are trained preferences derived from convex (left two plots) and geometric (right two plots) mixtures of anger and fear objectives from the Emotion dataset (Saravia ct al., 2018). To simulate preference shift, evaluation is performed at mixing coefficients a # Cgg, where a, = 0.1 is used during training. We evaluate WDPO with robustness parameter Po {50, 75, 100} and KLDPO with robustness temperature T E (0.5, 0.75, 1}. Additional experimental details are provided in Section 7.1.
distributions in the uncertainty set P(; Po). To overcome this challenge, we introduce principled tractable algorithms solve WDPO and KIDPP.
Tractable WDPO: The connection between Wasserstein distributionally robust optimization (DRO) and regularization has been established in various settings by many (Mohajerin Esfahani and Kuhn, 2018; Shafieezadeh-Abene progress in Wasserstein theory on connecting Wasserstein distributionally robust optimization to regularization. For p- Wassersttnin DRO, oo), Gao et al. (2022) shows that for broad class of loss functions, possibly non-convex and non-smooth, with high probability, the Wasserstein DRO asympiotically equivalent to variation regularization. In particular, an immediate consequence of Gao et al. (2022, Theorem 1) is that, when 2.
et al., 2019; Chen and Paschalidis, 2018). We leverage the recent
min
sup E P. (P.P).)
EE-P
min ((==): (1/)) E V=l 0)1 +0,(1n)) BEB
where P.. = On)V/n) That is, one can solve the Wasserstein DRO objective by adding a gradient regularization to the empirical risk minimization (ERM) loss, Eg-Po [1: 0)]. Based on this, we propose a tractable WDPO algorithm Algorithm 1. Note that the gradient regularizer has sample- size-dependent coefficient. In practice, we absorb the factor Pn/ // into Po. which treat tunable robustness hyperparameter.
Tractable KLDPO: The following proposition shows that we can approximate the worst-case
probability distribution in a KL uncertainty set w.r.t. a given loss function. Similar results can also be found in distributionally robust reinforcement learning literature (e.g., Gadot et al. (2024)).
Proposition 2 (Worst-case distribution (informal)). Let P E Rn be the worst-case distribution w.rt. a loss function 1 and KL uncertainty around the empirical distribution P defined as supp: (P || Sp Es-p [[2;; 0)). The worst-case distribution is related to Pa through
P()) oc (1) o) - -E ()(()((2;;))))
where 7 > 0 is some constant.
We defer the formal proof of Proposition to Appendix D. It can be viewed as a re-weighting threshold: extreme losses more biased towards the baseline empirical DPO loss. controls the intensity of re-weighting, acting as a temperature parameter. Based on Proposition 2, we propose a tractable KLDPO algorithm in Algorithm 2.
7 Experiments
We conduct experiments across three distinct settings that vary in dataset scale, model size, and the degree of distribution shift. For example, we fine-tune LLaMA-3.2. 3.---BB-Inntruuee LLaMA-3.2- 3B- Instructt,, and LLaMA-3.1-888 1-8--------B--stttuue Instruct models on prompts from the HelpSteer2 dataset (Wang et al., 2024b), using preferences derived from the ArmoRM reward model (Wang et al., 2024a), and evaluate them on the OpenLLM Leaderboard v2 (Fourrier et al., 2024). Additional evaluations are provided in Appendix E. We provide the code https: // ///gttttbbb//// wwm/ TheBBaacCCaa222 distributionally. robsstt..p.oouppppp.
8

---

7.1 Experimental Setup
Emotion Alignment: We use the Emotion dataset (Saravia et al., 2018) to train a GPT-2 model (Radford al., 2019) with a classification head for multi-label classification over five emotions: sadness, joy, love, anger. fear. The resulting sigmoid outputs are used as a multi-objective reward model for the remainder of this experiment. We also take GPT-2 model and perform supervised fine-tuning (SFT) with the Emotion dataset to obtain our base model for preference alignment. To construct preference data, we mix objectíves derived from our reward model. Specifically, we consider two reward objectives,
12 and define two mixture reward functions (1) convex mixing T Tonneex (a) == a F1 + (1 a+) T2 and (2) geometric mixing r "geometric (a) == rf -r 1 For both reward functions, we generate two completions per prompt and assign preference labels using the Bradley-Terry (BT) model parameterized by r" (a²) for chosen a" E [0, 1].
ArmoRM Multi-objective Alignment: We use the Absolute-Rating Multi- Multi-Obeective Reward Model (ArmoRM) (Wang et al., 2024a) to define reward preferences, selecting pairs of equally weighted objectives (c.g.. honesty, verbosity, safety) from its 19-dimensional first-stage outputs. Using Meta 14AMAA-3.2-BBInsttaae
as the base model, we generate two completions per prompt from the HelpSteer2 dataset (Wang et al., 2024b) and train models on preferences derived from the convex mixing of these reward pairs. We evaluate all models on five individual ArmoRM objectives, three of which are unseen during training, simulate preference shift.
1.AAMA--.2-1B-Instruue Instruct, models using preference data derived from the scalar rewards produced by the second stage of the ArmoRM reward model (Wang et al., 2024a). For each prompt from the HelpSteer2 dataset, we generate 10 responses, score them with ArmoRM, and constructe preference pairs by selecting the highest- and lowest-scoring completions. The models are evaluated on the OpenLLM Leaderboard v2 (Fourrier et al., 2024) using the LM Evaluation Harness (Gao et al., 2024).
Leaderboard Alignment:
We fine-tune LLaMA-3.22
1-0-A--..-33-Insiite and
7.2 Experiment Results
Correctness
Cortectness
Conrectness
Coheranice
Caherefice
Coheranice
0:30
04
2:40
0.27
0.27
0.1
Holphuiness
Heiptuiness
Holphunnass
Honesty
Honesty
Honesty
Arm AMMRRM
AmoRMM
Arm ARM
KLDPO (ours)
WDFO lours)
DFO
Figure 3: DPO, WDPO, and KLDPO in ArmoRM Multi-objective Alignment. LLaMA-3.2-IB- Instruct models are trained on preferences derived from three equally weighted objective pairs: (1) Ultrafeedback-Trulleee
and Compleexii,, (2) Ultraeedbackk Helluuuuekssls and Helpsteer- Coherence, and (3) Helpsteer- Correettesss and Helpsteer- Helpfulness (left to right plots). We train all models for 4 epochs. To simulate preference shift, models are evaluated on five individual objectives, Helpsteer-Helppe Helpfulness, Helpsteer- Helpsteer-Correctness Correctness, Coherence, Cohheeeer-Coherencce, the overall ArmoRM score, three of which were not used during training.
Ultrafeedback-Honne Honesty, and
In Fig. 2, we evaluate DPO, WDPO, and KLDPO under preference shifts between training and evaluation. All models are trained on preference labels emphasizing the emotion fear, while evaluation preferences gradually shift toward anger. The left two plots correspond to convex mixing of these emotions, and the right two use geometrie mixing. As expected, DPO performs best when the evaluation preference closely matches the training setup. However, as the evaluation shifts toward anger, DPO's performance degrades significantly. contrast, both WDPO and KLDPO maintain stable performance across the full range of evaluation preferences, consistently outperforming DPO under shift, demonstrating their robustness to preference misalignment.
Emotion Alignment Results:
ArmoRM Multi-objective Alignment Results: In Fig. 3, each radar plot corresponds to a different training reward pair, (1), (2), and (3), as defined in the figure caption. We evaluate all models on five individual ArmoRM objectives, three of which are unseen during training. to simulate preference shift.

---

LLANAA3.--1B sapping)
1Ftvaal
1LAMA-32-8
IFEval
MATH
MMLLU
0:45
DIO
u Epoch (geedat)
0.07
KIDPO
DXO Fpoch (oveeitt
0:48
0.66
0.33
0.17
WIRR (6. 0.005)
0.66
0:10
XLDPO (+ 0.1)
0:53
0.36
0.06
0.25
0.33
0.18
IFEval
B8H
MATH
GPQA
MUSE
MMLU
KLUPO (r- 0.85)
0.56
1136
0.58
0.26
0.32
0.18
DAO
0.02
0.50
0.05
0.29
0.44
0.35
WOPD
0.57
0.36
0109
0.34
0.19
KIDPO
0.23
0.51
0.34
0.39
0.37
WDFO 0.005)
KLDPO
Table 1: Evaluation of DPO, KLDPO, and WDPO on OpenLLM Leaderboard v2. LLaMA-3.2. IB/3B-Instruct and LLaMA-3.1 1-B-3..--8B-Innttuuc models trained on preferences generated according ArmoRM score and then evaluated on OpenL.M Leaderboard v2, which benchmarks LLMs across six tasks: Massive Multitask Language Understanding (MMLU), Google-Prooff Q&A Benchmark (GPQA), Multistep Soft Reasoning (MUSR), Mathematics Aptitude Test of Heuristics (MATH), Instruction Following Evaluation (IFEval), and Big Bench Hard (BBH).
Across all settings, both KLDPO and WDPO consistently outperform DPO all five evaluation axes, including those based on unseen objectives. This demonstrates their strong generalization and robustness to reward distribution shift, even when the evaluation preferences differ significantly from the training signal. Additional results are provided in Appendix E.I.
Leaderboard Alignment Results: Table 1 presents the performance of DPO, KLDPO, and WDPO on the LL leaderboard (Fourrier et al., 2024). WDPO and KLDPPO are trained for epochs, matching DPO's optimal early-stopping point, which is regularization technique to prevent overfitting. For LLaMA-3B and LLaMA-8B models, we align training durations similarly. Due to computational constraints, only KLDPO results are reported for the 8B model, given its scalability. These results, averaged over 39 subtasks, are supplemented by detailed evaluations in Appendix E.2, where WDPO and KLDPO demonstrate clear advantages across various subtasks.
8 Conclusions
We introduced distributionally robust DPO framework, developed two scalable algorithms with theoretical guarantees, and integrated them into existing LLM alignment pipelines. Empirical results demonstrate their effectiveness under preference distribution shift. Future work includes extending our methods to mitigate reward hacking and generalizing robustness to other RL.F approaches.
9 Acknowledgments
The authors would like to thank Vishnu Teja Kunde for invaluable discussions. This work was supported in part by the National Science Foundation (NSF) grants CCAREER-EPCN-204575 ECCS-5290 00292220200000 and LNS-526050-0000) Portions of this research were conducted with the advanced computing resources provided by Texas A&M High Performance Research Computing.
10

---

References
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.0m 08774,, 2023.
Alekh Agarwal, Nan Jiang, Sham M Kakade, and Wen Sun. Reinforcement learning: Theory and algorithms. CS Dept., UW Seattle, Seattle, WA, USA, Tech. Rep, 2019.
Alekh Agarwal, Sham M Kakade, Jason D Lee, and Gaurav Mahajan. On the theory of policy gradient methods: Optimality, approximation, and distribution shift. Journal of Machine Learning Research, 22(98):1-76, 2021.
Dario Amodei, Chris Olah, Jacob Steinhardt, Paul Christiano, John Schulman, and Dan Mané. Concrete problems in ai safety. arXiv preprint arXiv: (666:6660006665 2016.
Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislav Fort, Deep Ganguli, Tom Henighan, et al. Training a helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint arXiv:2204005822 2022.
Mohammad Ali Bashiri, Brian Zicbart, and Xinhua Zhang. Distributionally robust imitation learning. Advances in neural information processing systems, 14-24404-24417, 2021.
Amir Beck. Introduction to nonlinear optimization: Theory, algorithms, and applications with MATLAB. SIAM, 2014.
Amir Beck, First-order methods optimization. SIAM, 2017.
Stéphane Boucheron, Gábor Lugosi, and Pascal Massart. Concentration Inequalities: A Nonasymp- totic Theory of Independence. Oxford University Press, 2013.
Ralph Allan Bradley and Milton E Terry. Rank analysis of incomplete block designs: I. the method of paired comparisons. Biometrika, 39(3/4):324-34) 1952.
Alexander Bukharin, Ilgee Hong, Haoming Jiang, Zichong Li, Qingru Zhang, Zixuan Zhang, and Tuo Zhao. Robust reinforcement learning from corrupted human feedback. arXiv preprint arXiv:2406.15568, 2024.
Stephen Casper, Xander Davies, Claudia Shi, Thomas Krendl Gilbert, Jérémy Scheurer, Javier Rando, Rachel Freedman, Tomasz Korbak, David Lindner, Pedro Freire, et al. Open problems and fundamental limitations of reinforcement learning from human feedback. arXiv preprint arXiv:2307.15217, 2023.
Souradip Chakraborty, Jiahao Qiu, Hui Yuan, Alec Koppel, Dinesh Manocha, Furong Huang, Amrit Bedi, and Mengdi Wang. Maxmin-RLHFs Alignment with diverse human preferences. In Forty-forst International Conference on Machine Learning, 2024.
Ruidi Chen and Ioannis Ch Paschalidis. A robust learning approach for regression models based on distributionally robust optimization. Journal of Machine Learning Research, 19(13):1-48, 2018.
Ruidi Chen, Ioannis Ch Paschalidis, et al. Distributionally robust learning. Foundations and Trends@ in Optimization, 4(1-2))1-243, 2020.
Sayak Ray Chowdhury, Anush Kini, and Nagarajan Natarajan. Provably robust DPO: Aligning language models with noisy feedback. In Forty-first International Conference Machine Learning, 2024.
Paul F Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. In Advances in Neural Information Processing Systems, volume 30, 2017.
John Duchi and Hongseok Namkoong. Learning models with uniform performance via distribu- tionally robust optimization. The Annals of Statistics, 1378) 1477771166 2021.
11

---

Esin Durmus, Karina Nguyen, Thomas Liao, Nicholas Schiefer, Amanda Askell, Anton Bakhtin, Danny Hernandez, Nicholas Joseph, Liane Lovitt, Sam McCan- dlish, Orowa Sikder, Alex Tamkin, Janel Thamkul, Jared Kaplan, Jack Clark, and Deep Ganguli. Towards measuring the representation of subjective global opinions in language models. In First Conference on Language Modeling, 2024.
Carol Chen, Zac
Jacob Eisenstein, Chirag Nagpal, Alekh Agarwal, Ahmad Beirami, Alexander Nicholas D' D'Amour, Krishnamurthy Dj Dvijotham, Adam Fisch, Katherine A Heller, Stephen Robert Pfohl, Deepak Ramachandran, Peter Shaw, and Jonathan Berant. Helping or herding? reward model ensembles mitigate but do not eliminate reward hacking. In First Conference on Language Modeling, 2024.
Clémentine Fourrier, Nathan Habib, Alina Lozovskaya, Konrad Szafer, and Thomas Wolf. Open 11m Icaderboard v2, 2024.
Uri Gadot, Kaixin Wang, Navdeep Kumar, Kfir Yehuda Levy, and Shie Mannor. Bring your own (non- robust) algorithm solve robust MDPs by estimating the worst kernel. In Forty-first International Conference on Machine Learning, 2024.
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h h. Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang. Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. The language model evaluation harness, 07 2024.
Rui Gao and Anton Kleywegt. Distributionally robust stochastic optimization with wasserstein distance. Mathematics of Operations Research, 2022.
Rui Gao, Xi Chen, and Anton J Kleywegt. variation regularization. Operations Research, 2022.
Wasserstein distributionally robust optimization and
Daniel Hsu, Sham Kakade, and Tong Zhang. A tail inequality for quadratic forms of subgaussian random vectors. Electronic Communications Probability, 2012.
Zhaolin Hu and L Jeff Hong. Kullback-leibler divergence constrained distributionally robust opti- mization. Available at Optimization Online, 1(2):9, 2013.
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.
Ying Jin, Zhuoran Yang, and Zhaoran Wang. Is pessimism provably efficient for offline ri? In International Conference on Machine Learning, pages 5084-5096. PMLR, 2021.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 1412.6980 2014.
Robert Kirk, Ishita Mediratta, Christoforos Nalmpantis, Jelena Luketina, Eric Hambro, Edward Grefenstette, and Roberta Raileanu. Understanding the effects of rlhf on Ilm generalisation and diversity. In The Twelfth International Conference on Learning Representations, 2024.
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Bal- subramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, et al. Wilds: benchmark of in-the-wild distribution shifts. In International Conference on Machine Learning. pages 5637-5664. PMLR, 2021.
Daniel Kuhn, Peyman Mohajerin Esfahani, Viet Anh Nguyen, and Soroosh Saafieezadeh-Abadeee Wasserstein distributionally robust optimization: Theory and applications in machine learning. In Operations research & management science in the age of analytics, pages 130-166. Informs, 2019.
Will Lenene, Benjamin Pikus, Anthony Chen, and Sean Hendryx. A baseline analysis of reward models' ability to accurately analyze foundation models under distribution shift, arXiv preprint arXiv:2311.14743, 2023.
12

---

Daniel Levy. Yair Carmon, John C Duchi, and Aaron Sidford. Large-scale methods for distributionally robust optimization. Advances in Neural Information Processing Systems, 33:8847-8860, 2020.
Debmalya Mandal. Paulius Sasnauskas, and Goran Radanovic. Distributionally robust reinforcement learning with human feedback. arXiv preprint arXiv:2503.00533 2025.
Aditya Modi, Nan Jiang, Ambuj Tewari, and Satinder Singh. Sample complexity of reinforcement learning using linearly combined model ensembles. In International Conference on Artificial Intelligence and Statistics, pages 2010-2020. PMLR, 2020.
Peyman Mohajerin Esfahani and Daniel Kuhn. Data-driven distributionally robust optimization using the wasserstein metric: performance guarantees and tractable reformulations. Mathematical Programming, 171(1-2): 15-166, 2018.
Hongseok Namkoong and John C Duchi. Stochastic gradient methods for distributionally robust optimization with f-divergences. Advances in neural information processing systems, 29, 2016.
Andi Nika, Debmalya Mandal, Parameswaran Kamalaruban, Georgios Tzannetos, Goran Radanovic, and Adish Singla. Reward model learning vs. direct policy optimization: comparative analysis of learning from human preferences. In Forty-first International Conference on Machine Learning. 2024.
Long Ouyang. Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems, 35: 27730-27744, 2022.
Vishakh Padmakumar, Chuanyang Jin, Hannah Rose Kirk, and He He. Beyond the binary: Capturing diverse preferences with reward regularization. Workshop on Socially Responsible Language Modelling Research, 2024.
Kishan Panaganti and Dileep Kalathil. Sample complexity of robust reinforcement learning with a generative model. In International Conference on Artificial Intelligence and Statistics (AISTATS), pages 9582-9602, 2022.
Kishan Panaganti, Zaiyan Xu, Dileep Kalathil, and Mohammad Ghavamzadeh. Robust reinforcement learning using offline data. Advances in neural information processing systems, 35:32211 1-32224,, 2022.
Kishan Panaganti, Zaiyan Xu, Dileep Kalathil, and Mohammad Ghavamzadeh. Distributionally robust behavioral cloning for robust imitation learning. In 2023 62nd IEEE Conference on Decision and Control (CDC), pages 1342-1347. IEEE, 2023.
Kishan Panaganti. Zaiyan Xu, Dileep Kalathil, and Mohammad Ghavamzadeh. Bridging distribu- tionally robust learning and offline rl: An approach to mitigate distribution shift and partial data coverage. In 7th Annual Learning for Dynamics & Control Conference, pages 619-634. PMLR, 2025.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9 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, 36:53728-53741, 2023.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1-16. IEEE, 2020.
Shyam Sundhar Ramesh, Yifan Hu, lason Chaimalas, Viraj Mehta, Pier Giuseppe Sessa, Haitham Bou Ammar, and Ilija Bogunovic. Group robust preference optimization in reward-free rlhf. Advances in Neural Information Processing Systems,
2024.
13

---

Elvis Saravia, Hsien-Chi Toby Liu, Yen-Hao Huang, Junlin Wu, and Yi-Shin Chen. CARER: Contextualized affect representations for emotion recognition. In Proceedings of the 2018 Confer- ence on Empirical Methods in Natural Language Processing, pages 3687-3697. Association for Computational Linguisties, 2018.
Soroosh Shaffeezadeh-Abadeel Daniel Kuhn, and Peyman Mohajerin Esfahani. Regularization via mass transportation. Journal of Machine Learning Research, 20(103):1-68, 2019.
Laixi Shi and Yuejie Chi. Distributionally robust model-based offline reinforcement learning with near-optimal sample complexity. Journal of Machine Learning Research, 25(200):1-91, 2024.
Nian Si, Fan Zhang, Zhengyuan Zhou, and Jose Blanchet. Distributionally robust policy evaluation and learning in offline contextual bandits. In International Conference on Machine Learning, pages 8884-8894, 2020.
Joar Skalse, Nikolaus Howe, Dmitrii Krasheninnikov, and David Krueger. Defining and characterizing reward gaming. Advances in Neural Information Processing Systems, 35:9460-9471, 2022.
Rohan Taori, Achal Dave, Vaishaal Shankar, Nicholas Carlini, Benjamin Recht, and Ludwig Schmidt. Measuring robustness natural distribution shifts in image classification. Advances in Neural Information Processing Systems, 33: 111188883-1999 2020.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 2023.
Cédric Villani al. Optimal transport: old and new, volume 338. Springer, 2009.
Haoxiang Wang, Wei Xiong, Tengyang Xie, Han Zhao, and Tong Zhang. Interpretable preferences via multi-objective reward modeling and mixture-of-expper In EMNLP, 2024a.
Ruosong Wang, Dean Foster, and Sham M. Kakade. What are the statistical limits of offline RL with linear function upproximation? In International Conference on Learning Representations, 2021.
Yue Wang and Shaofeng Zou. Policy gradient method for robust reinforcement learning. In Proceed- ings of the 39th International Conference on Machine Learning, 2022.
Zhilin Wang, Yi Dong, Olivier Delalleau, Jiaqi Zeng, Gerald Shen, Daniel Egert, Jimmy Zhang, Makesh Narsimhan Sreedhar, and Oleksii Kuchaiev. Helpsteer 2: Open-source dataset for training top-performing reward models. Advances in Neural Information Processing Systems, 37: 14:44 1501, 2024b.
Junkang Wu, Yuexiang Xie, Zhengyi Yang, Jiancan Wu, Jinyang Gao, Bolin Ding, Xiang Wang, and Xiangnan He. 3-dpo: Direct preference optimization with dynamic f3. Advances in Neural Information Processing Systems, 37: 17:119944-22966 2024.
Junkang Wu. Yuexiang Xie, Zhengyi Yang. Jiancan Wu, Jiawei Chen, Jinyang Gao, Bolin Ding, Xiang Wang, and Xiangnan He. Towards robust alignment of language models: Distributionally ro- bustifying direct preference optimization. In The Thirteenth International Conference on Learning Representations, 2025.
Zaiyan Xu, Kishan Panaganti, and Dileep Kalathil. Improved sample complexity bounds for distribu- tionally robust reinforcement learning. In International Conference on Artificial Intelligence and Statistics. Conference on Artificial Intelligence and Statistics, 2023.
Yuzi Yan, Xingzhou Lou, Jialian Li, Yiping Zhang, Jian Xie, Chao Yu, Yu Wang, Dong Yan, and Yuan Shen. Reward-robust rlhf in lims. arXiv preprint arXiv:2409. /5360, 2024.
Wenhao Yang, Liangyu Zhang, and Zhihua Zhang. Toward theoretical understandings of robust Markov decision processes: Sample complexity and asymptotics. The Annals of Statistics, 50(6): 3223-3248 2022.
Zhouhao Yang, Yihong Guo, Pan Xu, Angi Liu, and Animashree Anandkumar. Distributionally robust policy gradient for offline contextual bandits. In International Conference on Artificial Intelligence and Statistics, pages 6443-6462. PMLR, 2023.
14

---

Michael JQ Zhang, Zhilin Wang, Jena D. Hwang, Yi Dong, Olivier Delalleau, Yejin Choi, Eunsol Choi, Xiang Ren, and Valentina Pyatkin. Diverging preferences: When do annotators disagree and do models know? In Forty-second International Conference on Machine Learning. 2025.
Siyan Zhao, John Dang, and Aditya Grover. Group preference optimization: Few-shot alignment of large language models. In The Twelfth Interaational Conference on Learning Representations, 2024.
Ruida Zhou, Tao Liu, Min Cheng, Dileep Kalathil, PR Kumar, and Chao Tian. Natural actor-critic for robust reinforcement learning with function approximation. Advances in neural information processing systems, 36, 2024.
Zhengqing Zhou, Qinxun Bai, Zhengyuan Zhou, Linhai Qiu, Jose Blanchet, and Peter Glynn. Finite-sample regret bound for distributionally robust offline tabular reinforcement learning. In International Conference on Artificial Intelligence and Statisties, pages 3331-3339, 2021.
Banghua Zhu, Michael Jordan, and Jiantao Jiao. Principled reinforcement learning with human feedback from pairwise k-wise comparisons. In International Conference Machine Learning, pages 43037-43067. PMLR, 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:1909.08593 2019.
15

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect paper's contributions and scope?
Question:
Answer: [Yes]
We have clearly stated our paper's contributions and scope in the abstract and introduction.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We have clearly stated the assumptions used in our proofs. Additional discus- sion on limitations can be found in Appendix G.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section their paper.
The paper should point any strong assumptions and how robust the results are violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). ). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best
16

---

judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification:
All assumptions used are clearly stated.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification; We provide information in order to reproduce experimental results presented in our paper. In addition, we will provide open access to the data and code used in this paper.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g.. in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm.
17

---

(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (c.g.. a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d)
We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
We will provide open access to data and code upon acceptance.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (kttps://nips.co
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed run reproduce the results. See the NeurlPS code and data submission guidelines (https: cppps.cc/publcc/ggg guide http://guudeeeeeeueurnnnn CCodeSubmissscce ssi onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question:
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Experimental details can be found in Section 7, and additional details can be found Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
18

---

The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide additional details regarding our experiment setups in Appendix
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: : htt///euriis.ee http://p...ppppieeittsssscieeeee
19

---

Answer: [Yes]
Justification:
We acknowledge the NeurIPS Code of Ethies.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there special consid- eration due to laws or regulations in their jurisdiction).
Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
See Appendix H.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate point out that an improvement in the quality of generative models could be used to generate deepfakes for disinfornation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: This paper has no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring
20

---

that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All creators or original owners of assets (e.g., code, data, models), used in the paper, are properly credited.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, if possible, include URL..
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that re- re-packaged, both the original license and the license of the derived asset (if it has changed) should provided.
If this information is available online, the authors encouraged to reach out the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
The documentation will released along with the release the code.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mooel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
21

---

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [Yes]
Justification: The proposed methods, Wasserstein DPO and KLDPO, are applied to fine- tune large language models such GPT-2 and LLaMA. The LLMs are central our empirical validation and the alignment task studied in this work. Thus, their usage is critical component of the core methodology.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our L.LM for what should or should be described.
22

---

A Useful Technical Results
A.1 Wasserstein Theory
rely on the following strong duality result from the Wasserstein distributionally robust optimization (WDRO) literature.
Lemma 1 (Gao and Kleywegt, 2022, Theorem 1; Strong Duality for DRO with Wasserstein Dis- tance). Consider any p E [1, 00), any v P(E). any p O, and any E L' L'(v) such that the growth rate R of y satisfies
k== inf 11 1/20 0: / [ (((())) -000>0000 -00
} <0,
(10)
where Φ(1, {) == infee= (nd (5. (()) is regularization operator. Then the strong duality holds with finite optimal value Up UD oc, where
{ (E))()E) Wo(m,)) sp). MEP(E)
Up = sup
(Primal)
/ in inf V(E))u(dc)
up
inf 120
(Dual)
SEEE
Lemma 2 (Gao and Kleywegt, 2022, Lemma 2.(i);; Properties the growth k). Suppose that E (E). Then the growth rate k. (as defined in Eq. (10)) finite and only if there exists Ç and
L, .. > 0 such that
(E()) (C") S Ld. M, VE
(11)
Corollary 1. Consider any bounded loss function over bounded E. Then the duality defined in Lemma holds.
Proof. It follows from Lemma 2. We can pick L to be the diameter of = and M to be the bound of V.
A.2 Optimization
Lemma (Beck, 2014, Theorem 1.24; Linear Approximation Theorem), Let f: U - T.. be twice continuously differentiable function over an open set U C TR. and let x. y E U be such that
[z, v] U. Then there exists & E [z, y] such that
f(x)) (-)
z).
Lemma (Beck, 2017, Theorem 5.24; First-order characterizations of strong convexity). Let f: E (-x,, oo) be a proper closed and convex function. Then for given 0, the following two claims are equivalent:
For any x, dfm(f) and E {0, 1]:
f(Ax (1 A)v) s (f(=))+(( ()-1))()))
o 2(1 2)( yll2.
2
II.
f(D) 2 f(x) + (g.0 x) + |1y x01
for any x E dommaf), dfm(f) and g E Oz(z).
Lemma (Beck, 2017, Theorem 5.25; Existence and uniqueness of minimizer of closed strongly
convex functions). Let f: E + (-o0, xo] be proper closed and a a-strongly convex function a > 0. Then
f has unique minimizer;
II. f(x) ) 2 g| - x=2 for all x E dom(f), where x" is the unique minimizer of
23

---

A.3 Distributionally Robust Optimization Results
The Kullback-Liebler uncertainty set can constructed with the f-divergence, The f-divergence between the distribution P and Pa is defined
dP
Dy(P Po)
/
dP"
(12)
dPo
where f is a convex function. f(t) = t log(t) gives us the Sullback-Lieble divergence. Let Pa be distribution the space X' and let (: X' + R be loss function. We have the following result from the distributionally robust optimization literature.
Lemma 6 (Duchi and Namkoong, 2021, Proposition 1). Let Dy be the f-divergence defined in Eq. (12). Then,
Af ((x)
((x)-2)
[p((
inf Epp
+ +pp+ +
(13)
P: (P || Pa)
X20,nER
where (*))) supizo st [(t) is the Fenchel conjugate.
A.4 Concentration Results
Lemma 7 (Hoeffding's inequality (see Boucheron et al., 2013, Theorem 2.8)). Let XI...., Xn, be independent random variables such that X, takes its values in [a,, by] almost surely for all i n. Let
n
S
Σ (X, E [
Then for every >> O,
2t
(S S exp
EE= (b, 106-a,7 o,)3
Furthermore, if XI...., Xn are a sequence of independent, identically distributed random variables with mean μ. Let X,, Ei X;. Suppose that X; E [a,b], Vi. Then for all 0
2nt²
P P(X,. 11 2exp
a)-
Lemma (Hsu al., 2012, Theorem 2.1). Let A E Raxn be matrix, and let E A A. Suppose that I = (#1...., .n) is random vector such that for some μ E TR? and a 2 0,
Elexp(a (x - +)) S exp(Ila|[3o?
for all a E R.. For all t 0,
||T|| > 0². (Tr(E)) + 22 Tr(ET)e + 2|E/1t) + Tr(EppT)) (1+2)
/E
]
sett
Tr(Σ2)
Moreover, if = 0 and a = 1, then the probability inequality reads
TP ||x Tr(E) + 2 Tr(E3)t 2|E/]11 se
B Proof of WDPO Sample Complexity
Many properties of distributionally robust DPO derived from those of the non-robust DPO. We hence start with the following proof of policy parameter convergence in the non-robust setting (Proposition l).
B.1 Proof of Non-robust DPO Policy Parameter Convergence
Recall the pointwise DPO loss:
1(0; a a² y) y log a (Shg a a))) (1 y) log (3ho aa,aa)a))))) a²)).
=o(a41s) where ho(s, a a²) = log [,a((a'|3] - log xo(a²1s) (-- Denote this loss by l=(0) where 2 =
(s, at, a², y). We also denote the empirical (sample) DPO loss
n
I = Σ (0)
vi (3hal af) (1- y.) log o (3 af a))).
We denote the MI_E solution to lp by odpo argmingge Ip (0). Also, denote the true parameter which is the global minimum of the population negative log likelihood by 0*.
24

---

(Almost) Strong Convexity of I. In order to calculate the Hessian matrix of l1 w.r.t 0, we need to calculate V log (3h a', d²)).
Suppose f: R R. g: R R. The Hessian of f o g is, for any x E R
(f () (z) fl((()))V f vg(
(14)
Recall that is the sigmoid function. It has the properties: <(-x) - and d'(x) (x)(11 - (x)). Let (xx = log (ar), we have
d
d'(x) a(x)(1- (1(1l-oaa o(x)) o(-x);
de (()
(ω)
(xx
d?
d
d
dx f(x)
dx [o(-))]
da [[--(()) =()))
With g(0) Bho a a²) and the Hessian chain rule for composition with scalar function (Eq. (14)), we have
V log (Bho((aa a²)) = Ba(--Bho(s, a' ,a (o) - ( a', ,aa)) Sho (o)))ooho(s,dd , ooh
In addition, we have the following observations
Veho(s, at, a²) = V log (a s) Ve log (a s) Voho(s, Voho(s, a. a²) = V (a (s) V log (a () - Vsh
Now, using the above observations, we can simplity (0) as follows
Pg (0) = V log (Bho(s, - (1
= y (2)
(1 (ll=Dloo(tBBhhssdddd))//)oso
= o0 ooooo
+ + hch.
+ (1 ( d'")
Vehoss o
Voho(s,
y3r - (3)
+ (880)()hoss,
,a))),,,,,,))))))))) Voh))oo ,add","")) aooo
+ ( (2)
+ (1 - (1 (3h)) a Vehssss (ag Vehos,)s)
( +
a² + o ( 3ho , Vohoss,
where (a) is due to ho(s, a,,=) ho(8, , d²), Voha aa)) = Vohals at a²) and
Vyhols, aa,a)) a²) = -2gho(s, a=, a)). It's clear that we have to calculate Voho at ut,aaa a²) and
Vohols, a²). Observe that
1
1
Vahols, d²) Vo To(a" | 8) Vo To(" 1s) To(al VoTT((a 18) To(a² VaTTo(a 1).
(15)
In addition, we have that Vgha al, a²) Vologolaa | s) V2 log x s). Using the Hessian chain rule (Eq. (14)), we have
v log To(a 8) To(a s) Vinola (s) To s)2 VoTo(a 8) V) ST
25

---

Now it boils down to tackling VoTo(a s) and Va s). Observe that
Veexp ((1((s,), 0)) E exp ((v)s, 0;)] Eg Veexp
((b(s, 0))
VaTT(aa s)
(Ev exp ((b(,, a'), 01))
cxp((((s,a))
exp((v(s,a), 0)
En exp ((v((,(,,''') (ss
exp(()(s. (ss,
Ex exp((((s)
exp((U(s,a).0)
exp exp((((s,0),0)
exp exp((1(s,a))0) 0))
E. exp (()(() ((s
En exp ((((s,"))))) (1s
En ((),(())))'))
To(a| (s, To(a 8) Σ To(a' (8)(()
= 8) ()(s)))) E Tola' s))((s.)))
Then we have
Vaha ,02)=
To(a! 8) mo(a! 18) (o, -
1
- - Told' ())((1))))
To (a(a2 8) To(a? 5) (0 Emlla' s)0b(8,0))
(ss a²) (5s (²)
(16)
Notice that Vohg above does not depend on the policy parameter O. This implies that its Hessian is the zero matrix, i.e., Viho at, a²) = O. Finally, we have that
(0) = (a at, a²) (²)) ( a², (V'(s, a²) (s (3)(ss))))))) ((s)
Moving from the pointwise loss to the empirical loss, we denote
(Ol = Σ
Σ Paolhaass at, (a))))((Bhh(e) at, a))- ())
Now let's focus on the function (-)) Our aim to find a lower bound for this function. Observe that
( a², a,a)) = 1(log (a' 8) - (a $)) - (log Tret (a+ {) - log (ae s))]
160, (xs - " (frof: (ss - (s))a²)))
1:0- Ozet (ls (ss,
110 - 0 (2 at) - vi(s, a² 22)
4B,
(17)
where (a) is due to Cauchy-Schwarz inequality. (b) is due to the assumptions ||0|| B and max,,a ([bb(s))))))) a)llz 1 1. Now this suggests that the input to the function a (Bho(s, at, "²,"")) (-)))a(Bhnn(e (s, a²)) is bounded in [-4,B, 43B). Since (x))(-<)) is symmetric and strictly decreasing when x E [0, oo), we have that
<
BeABB
3'o(Sooss a²,
a'o
AA e.
(18)
(1 (1+c48Byy
We then have that
u T TT > π N| Vu E Rd
where y = A8B (1+c and X x has the differencing vector x; = (/s,,, at) - (ss,, af) E IRd as
its i-th row. Thus, if we introduce the error vector A gdpo 0r, then by the linear approximation theorem (Lemma 3), there exists a E [0, 1] and 0 addlo + (1 a)0" such that
(0 + A) (D)) (Volo(0'),
1 A V
7 |XAI
2 I (19)
2
2n
26

---

where & (e(s; a))= (1b(s, ( a.) This implies that (p is (almost)
strongly convex around 0" with parameter y with respect to semi-norm|||| Note that we will not treat lD as strictly strongly convex function in any part of this proof. We only need the inequality Eq. (19).
Bounding the estimation error. Recall that 0d is optimal for (D(0) and A = Odp 0- We must have to (0:1))) - (6)). By substracting and adding (VolD(8"). A) on both sides, we have l0 A) (0())) (VolD(0"), A S (Voll(00))
For the right hand side above, we have ((Vol((0), A) &
for any X > 0.
By 7-strong convexity of (D at 0* we have
(D(0 A) (D(00) (VelD 27 DA
Combining the inequalities, we have I S bound the term (Vll((")l((E We can calculate the gradient w.r.t. 0 of the pointwise loss as follows
M| Now we need
(0) = v y log o (s a' a²)) (1 (Bho
ve oo (3ho a²)) v (
Byo Bhg a , )) Vahg(s,
3(1 1) (3ho V2)
(Byo )²)) (1 - ( at, (vs() at) ((s,a))),
where (a) is due to Voho a a²2) vb(ss, a v5s a²) calculated in Eq. (16). This implies that
VolD(0") = -3
uno (Sho (ss, at)) (1 - vi)o (8s, of,,,)))) of))]a.
(20)
i=1
where a; = (888 a.)) - (8,, aZ). Now let's define a random vector V E IR" with i.i.d. components
as
a (3ho- (ot, (2,,a))) V; -o (Bhg (s a), az))
w.p. (Bha- - (si, at az)), w.p. (Bha- (8t, at a.)).
(21)
Then we have Volo(0") - 8 V It's easy verify that EV; 0 and [V] for all 1 isisn n. Next, if we define the n x n matrix M f X + X then write (0)l0(")1lE 1 = = MV. Let the cigendecomposition of X X be UAUT. Observe that
32 X + +1)-1
8² XU(AAA++AI)) -U
M M n²
_
n²
We can bound the trace of M follows
8² n
8² Tr(U U
Tr(M) Tr(
n
12 Tr(U(A/n A
B²
Tr((A/n+A))11)
82
ne:
n²
n
n
e, + An
g²
3dd
< nd
n²
where e; the i-th eigenvalue of X T X. Similarly, we can bound Tr(M
<
8'd n²
Now, let
X == TÜEV be the singular value decomposition of X. Then we can show that
82
82 SUNE E/n + -
M X(x n
X/n -
-UUT
n²
Since X(Ep - -X is symmetric, and clearly ÜΣ(Σ + - diagonalizes it, the / where o; is the i-th singular value of X. Hence, all eigenvalues
eigenvalue of takes form
27

---

Amax (M) < L. Since the components of V are i.i.d. with EV; and [V] & 1, the elements are Ssuub-aussian, we can use the Bernstein's ss inequality for sub-Gaussian random variables in quadratic form (see Lemma 8). It implies that with probability at least 1 - 6,
upper bounded by n. Then must have ||M
[VolD(( )Il V MV I Tr(M) + 2 Tr(M)) + oo
3dd
84 dog (1/0) + +2 82 log(1/6)
8² (d 2 (1gg(11//)) log
2/
n
Set a = and b =loooo(/ (1/5) Note that we have
d + 2 2/d 2 log(1/0) (a+ b)2
s 2(a + b)2 2( + + 2ab) S 2(a + 62 + a² +b) = =4(a² + 62) ) 4(d log((1/)))
where the last inequality is due to AM-GM inequality. Altogether, we have (olb(0))] 4.9 (d + (1g()/))).
The final assembly now begins follows
| 2 B +As = IA + 2 11 = 2 IA +
X7
||12
2
2
2
Ay
< || (0(=000))[[= ED +
||
2
4
Ag
(d log(11)))) +
4B2
n
where the last inequality uses triangle inequality and the assumption that |0| S B E θ. This implies that
2
4
| BOt S
(d
+ 4AB2.
Now denote
a a= 2,1 1000 (d 1og(1/6)) and 4XB², and let x || Since we have
7 a
x² 3 < 0, then must be less than the bigger root, i.c.,
x I a + V 02 +
43 43
/ +
2
2
where the second inequality is by Jensen's inequality. Finally, we have that
4.2 (d + log(1/6)) 2AB2.
|| lidpo
8"
EEo+A = | I 2
72n
B,2 Proof of WDPO Loss Function Convergence
Lemma (Convergence of WDPO loss). Fix any O and p > 0. Let 8 E (0, 1). With probability
1-6, 6,
k2 log(2/8)
k LWW(O:)) (0: 2 (0: <
2n
where = =log -loogo( (--3BB4)
Recall the strong duality in Lemma 1. The term inf (nd 2²) [(2;0) is called the Moreau Yosida regularization of -1 with parameter 1/n. We denote it by (,(2; 0). Now observe that
Proof.
c (0: C.W. (0:
sup P: (PPP)
BemPl1((0)
sup Wp(P.P:)
Ex-p [c)
(a)
inf (np" 720
Eg-p- [b,
inf (mp 1b 720
sup
Egnpo [1, 0]]
EE (2;
920
28

---

where (a) is by the strong duality, and (b) is due to linfz f()) inf ((ي) S sup=.f(z) - g(z)l.
Next, we will show that, for any 77 2 O, the function L, is bounded function. We first prove its upper
bound. The negative DPO loss takes the following form:
(2; 0) = y log a (x) + (1 - y) log (ogo(-2) 0, y E 10,
The inequality is because the sigmoid function is strictly bounded between 0 and i.e., a E (0.1). This implies that log a is non-positive. Using this, we have that
(, 0) = inf Ind (z', (== < z'E2
Ind (2' ==]] O.
z'EZ
Now we prove its lower bound. Recall that in the analysis of non-robust DPO loss, we proved that (s, ( at, a.aa] < 4B (see Eq. (17)). Since both log and a are increasing functions, have that
log (3ho(s, at,ttta a²)) 2 log (-4BB.. Now observe that
(2(00) inf Ind (e', [(2:'00 0)]
2'E2
inf 1=((=';000)) eeZ
inf
log (olgoo@hhhssss,,oo a²)) (1 y) log (3go,
s.a s.a",a" ,2
2 logo(-44B))
where the first inequality is because both v and metric dl are non-negative. The last inequality is because only one of the log term will be activated and the lower bound we recalled above. Denote K = [log of (443B) Since in is a bounded function, by Hoeffding's inequality for bounded random variable (Lemma 7), have
-2nE² K22
IP 10 Ez-py (2; 2c)
< 2exp
By picking 5 to be the right hand side above, we have that, with probability at least
122 log(2/5)
EE-P 0)] Ez~pe
(2, (2;
2n
Since K does not depend on 7,, such concentration is uniform for all functions 1,: 0. We have the desired result.
B.3 Proof of the Strong Convexity of WDPO Loss
We first prove that the function g(0; P) == [[<= [2] is strongly convex, for any P, as follows: Lemma 10. Let 1(:00) 0) be the DPO loss function. Assume that Assumption 2 is in place. Then g(0) Ez-p [1:; 0)] is y-strongly convex with respect to norm ||-1/ where Ep = [E aa? a² a)) (188, (8,aa))))))))))
(ss a² (1b(ss))))))))
(T), and 7
8 28B
(1+248877-
Proof. Recall that we proved that the Hessian of the pointwise DPO loss takes the form:
(0l+(0) = (a (Bho(s,, a², (o)))((-Sho(ss) at, In addition, we also proved that (see Eq. (18))
(aoos,,,,"))))ee (a))), a²) - whss, ((s,s - ((ss
BBB
3e ( a² (hooo-Bhaoss
V0 e
(1 + eBBB2
This implies that
T (3)2(0)u 2 >1((((s, - (sssa))) (T), u11 Vu E Rd.
82849B (1444 Thus, if we introduce the error vector A = 0' - 0. where 0, 0' E θ, then by
where y
the linear approximation theorem (Lemma 3), there exists [0. 1] and ē al + (1 )0 such that
(0 ((0=))-((0)
[1(((s, -
ATV3L
A 2 3
(22)
29

---

where Σ. = (vbss, a!) - - a²) (v)(s,))) 2/ a))) Note that Σ, is only semi-definite. Let [0. and 0' E e. Observe that
s(a8 (1 - )) BL>pl[(ab (1 =))
(0)
S En~P al(=;0) + (1 a))((";; 2) - 2 20(1 -a]ll0 - -0'1
=g((0) + +1 a)g((' 7 a(1 a) 8¹) T Ep [=:] 81)
ag((0) + (1 a)g(0') 2 aa(1 - a) 0'
where (a) is by Lemma 4. In particular, the equivalence between the inequalities, Eq. (22) and (a), be found in the proof of Beck (2017, Theorem 5.24), and the author would like comment that the proof does not rely on whether I|-1|x, is semi-norm or a norm. Now, by Assumption 2, Ep is strictly positive definite, hence I|-||x, is norm. This implies that g y-strongly convex with respect to ||-|I--
Now, we are ready prove our main strong convexity lemma. Lemma 11. Let 1(z; 0) be the DPO loss function. The Wasserstein distributionally robust DPO loss function.
W (0:
Ez-p((2;;0)]
sup P: (P.
is "d-strongly convex in 0 with respect to (non-weightes 2-norm ||-1|2, where is the regularity condition number defined in Assumption and y
888
Proof. Let a E [0, and 0,0' E O. First, we denote h(0;P) = Ez~p![(z;0)) for any in the Wasserstein ball. In Lemma 10, we proved that h is 7-strongly convex in 0 w.r.t. norm I|-11s Now observe that
C. (al (1 - a)
sup (P.P)))
h(a0 (1 2)
(a)
7 7 - a)]le 11 ah((B:P) (1 a))((';
sup (.P.P.P))
2
aL (0: (1 a))""(((;
a(1 -a)ll - 1'1
sup
P:
Wp.(PPP))sp
(ec + +(1 a)
? a(1 a)
inf (o,
10 8'1133
ac acW(O:)) + (1 (0("
? 7a(1 a)
inf p: (,,
Amiu (Ep) 0'11
2 aL (0: p) +(1+1111111111 (8:::))
7d
2 a(1 - a)ll@ 0
Note that the function g(0) = E.-Pp [2; 0) is 7-strongly convex with respect to 11 by Lemma 10. We use this fact in (a). The inequality (b) is due to (f + gx(x)) S sup, f(x) + sup, gzx).
The last inequality (c) is because Amin (Ep) 2 X, for all P E Pw by Assumption 2. This implies that CW is a st-strongly convex function with respect to 11-12--
B.4 Proof of Policy Parameter Convergence of WDPO
By Lemma 9, we have that, with probability least 6.
LW (@ p) WW (OW; p)
(OW: :p) - CW (eW, p) + CW (0W: p) -- (BW: :)) + L.W <oW: :p) L (
W W (0W; :p) C. (0W: p)) + |2 (8W; :p) <W(oWW:)
2K² log(2/5)
30

---

where the first inequality is because 0% is the minimizer of C, Now by the yA-strong convexity of CW (see Lemma 11) and Lemma 5, 5.II, we have that
8K2 12/(2
100 - 1013
x2
C Proof of KLDPO Sample Complexity
We state result from Hu and Hong (2013) that proves an equivalent condition for the infimum be achieved at X = O.
Proposition (Hu and Hong, 2013, Proposition 2). Let (2(2;)) be the essential supremum of (2; under measure P", i.e.
(z; inf(t E R: P(((2; 0) > t) =0.
Also let Ru = P([(2; 8) (=(2;00), i.e., Ky is the probability mass of the distribution P on the essential supremum of Then X* = 0 if and onty if (, 8) +oc, Fuu 0, and log + p 2 where p is the diameter of the KI. uncertainty set.
We now make assumption the loss function(s) ((:;0), 0 e O. Note that this assumption is only used in proving the dual reformulation of KLDPO objective.
Assumption 3. We assume that l(z; 0) S. for all 0 E O. That is, the loss function is upper bounded by L. In addition, we also assume that θ permits a uniform upper bound on xg. That is, we assume that supege À X.
We now prove the following dual reformulation result: Lemma 12. Let ((;;6) be the DPO loss. The KLDPO loss function has the following dual reformu- lation
C (0;
sup
Ez-P E ([z;0] inf
{ fro+ Alog
(Ezpp (Eappoeex) Texp
exp
[12; 0)
P: (P Pr)
AE(A.T]
where 0 < A < X < 2o are some constants.
Proof. We include the derivation here for completeness. Previous works in optimization and distri- butionally robust reinforcement learning have covered the dual problem of distributionally robust optimization with KL uncertainty set (e.g., see Hu and Hong (2013); Panaganti and Kalathil (2022); Xu et al. (2023)).
Recall that f(t t log(t) corresponds to the KL divergence. The optimal i for (s) = sup fst t }log(t)} is exp (s 1). This implies that the Fenchel conjugate of f is f*) exp (s 1). From Lemma 6, we get
((2; - 17 -
] - Ap 2
sup (Ku(P PO)SP
Ex [((;;])] = inf A20
Egnp
\1*
((2;0) -1) + Ap 2
inf A20.EER
EznPo
texp
)
Ezapo exp ((z;0)
= inf A20
Ap Alog
where the last equality by plugging in the optimal n,, i.e., 17" -Alog(Eswp-(en (1(z; 0)/A 1))) Now observe that
exp (=;00) (1=0))) Ap g(A).
M(:0) = Ap Alog Ez-po
The inequality is because the loss function is non-negative, i.e., 12 0, and h is increasing in l. Now g( (A) is a strictly increasing function that lower bounds function h(>;; 0). Since g()) - oo as X - oc, h( A; ()) cannot achieve its infimum at co. other words, there exists Xo such that
h(x:0) 2 g(A) > g( A> 7o.
31

---

This implies that it suffices to seek the infimum in [0, Xo]. Hence, we have
[[2; 0)
(k:
inf Xe(o,Xal
Ap + Alog
E=-P'
exp
Now from Proposition 3, the condition log Kig + p 2 is problem- dependent due the diameter p. which is design choice. Note that when is close to zero, the condition log K.. 2 0 is almost never true for reasonable p. Hence, we ignore the case where
O. By Assumption 3, without loss of generality, we have that X E [A, X), where is some problem-specie constant. Then we have the result. In the literature of distributionally robust reinforcement learning, similar arguments can be found in Zhou et al. (2021); Panaganti and Kalathil (2022).
L
Lemma 13. Fix any 0 E ④ and > O. Let 5 E (0, 1). Assume Assumption 3 is in place. With probability 8, we have that
exp (L/ (L/))
(k:
LKL (0: p)] SX
Vc 0,
2n
where X, x are constants that are independent of
Proof. Observe that
|CkLL p) -L (0:p)l
sup
E ExmP|(Z;;0]]
sup
E [==0(((=00]]
P: PP Po)SP
P. DRL(P PR)
(a)
inf AE(A,A]
Ax Alog
Esnpo exp
(12; 0)
{ Ap+
(2:00)
inf AE(A,T)
EINPE
exp
d
(6)
(1(;;0 )1)
((2; 0)
sup
Alog
Eg~Pa exp
A log
EENpo
exp
A
xe 2.X]
Ezwpy Texp (12;;)) />]
E Eg-P- ((2;0)) /A] )
sup
log
EZWPE lexp (lzz /A] Ez-Po (exp ([;;0))/A1
+ +1)
sup AE(A,T]
A log
Eg-p= |exp (1(z;)) AA]
LE P2 (exp (1(2; 0)) /1| P /exp (122; 0)) /All
sup
Ez-po lexp (l)=0)) A
A
sup
EE--p /exp (1(2; 0)) />] - E.- /exp (1(2; ()) /A]l,
XE(A,A]
where (a) is by Lemma 12. (b) is because linf, f(x) inf g(x)l S sux,,f(x)) g(x)). (c) is by Assumption 3. (d) is due to [log(1 + =)] S |x]. Vo 2 0. (e) is due to the fact that the loss function non-negative, i.e., 1 2 0. Now by applying Hoeffding's inequality (Lemma 7), have
(- 2ne
P(|Ez=Pg /exp (1(2; / Ex- (exp (I(=;0) /Al 2 e) S 2exp
Teep (L/A)
By choosing e
exp(L/A) log(2/5)
we have the result.
2n
We prove a strong convexity result similar to Lemma 11 for KLDPO loss function. Lemma 14 (Strong convexity of KLDPO loss). Let I(2;0) be the DPO loss function. The KL distributionally robust DPO loss fincction,
CKL p)
sup P: DRL PO)SP
E. E_=p((2;;]]],
-A-strongly convex in 0 with respect to (non-weighted) 2-norm 11-112 where the regularity 23,
condition number defined in Assumption 2, and y 1n4z
32

---

Proof. Let a E [0, 1] and 8, 9 E O. First, we denote h(0; P) Bsmp]((2;0)) for any P in the KL ball. In Lemma 10, we proved that h y-strongly convex in w.r.t. norm Ixpp Now observe that
C (ab (1 111-00 (1-a)6': p) P: PKL(P Po)Ep
sup
h(aB (1 ())
(a)
- aah(o:P)+ + a))((';; - 7 a(1 -0)11-88113 -81120
sup
P: PPL(P Po)Ep
(6)
I al aL (l:pp) + (1 (1- (():))(0'))) +
7 a(1 aa
a)l18 - -11
sup
P:
PKL Po)EP
(0: p) + (1 - c) (8::
aa - a)
inf
||0 - 8'1
2
P: (Pt (P
Po)
s (0: p) + (1 - c) (0':p) (k):
o(1 - a))
p: (P inf
Amin (Ep) 0'
2
pr)Sp
(0) 2 ac ((:k)(0:)) + (1 a) CKT :)
7
2 a(1 a)l10 112
= Note that the function -00)= p:::::((2 0) is 7-strongly convex with respect to || by Lemma 10. We this fact in (a). The inequality in (b) is due sup,, (f(()) + +((x))
I sup, [(x) + sup, g(x)..
The last inequality (c) is because Amin (Ep) 2 X, for all E PkL. by Assumption 2. This implies that CKL is 7A-strongly convex function with respect to 11-112.
C.1 Proof of Policy Parameter Convergence of KLDPO
By Lemma 13, have that, with probability at least 1 6, _KL (0hh, p) LKL (okk,
ckL(ohh p) L (0%1: p) + LKL (okL; p) ckL((kk, ;p) + Ln chL(Ook, ;p) CKk(Okk p)
CKL :)) LkL (0KL, p)] + (AkL, CkL (BKL: p)]
exp (L/ log (2/8)
2
Ve 0,
2n
where the first incquality is because gKT, is the minimizer of CKI, Now by the yA-strong convexity of CKI. (see Lemma 14) and Lemma 5.11, we have that
8X exp (T/A)) log(2/)) 2222n
n - 8KL 1
Ve > 0.
D Proof of Tractable KLDPO
Next, we prove the formal version of Proposition 2.
Theorem 3. Suppose we have the following distributionally robust loss that corresponds to a KL uncertainty set:
sup
Ez-p[(2;;0),
P: (P PR)SP
A worst distribution P. E R.a is related to the empirical nominal distribution Po which is constructed using n i.idd samples 21. Zns through
([2;;0) - A
P(i) = Pi()) exp
(23)
A
where P(i) corresponds to the worst-case mass on the i-th data, and furterr it is subject to
(zz;0) 0) - u A- (1Zz) 0) 1 A
(1) exp
= p.
(24)
X
(1z;; 0) 11 - ) 1,
oo (1) exp
(25)
X > 0.
(26)
33

---

Proof. We re-write the objective as convex optimization problem
maximize
(p, 1)
(P1)
PER
subject to
Σ Pr log
Pi
P.
qi
p = 1, 2 0, Vi.
First, we ignore the constraint 2 0 which will be automatically satisfied later. Now, the associated Lagrangian function takes the form
L(p. 1) (p, + A - Σ
log(p:/4))) + p(1 - 1 p)-
1=1
We can calculate the KKT conditions as follows
aL 0p:
= 4 X(log(p:/9)) + =
This implies that
1/ 11- (1ー1ーン>)
P: 4xep
Vi (I,.,...,
In addition, we have other KKT conditions follows
a
Σ
pt log(p:/a)) - -
Σ Pi
X z 0,
log(p:/()) - p) = 0.
From complimentary slackness, we have
Σ Ps log(p//4) p.
The unconstrained optimum would lie at vertex far from q, thus the best achievable objective under the KL constraint is obtained by pushing the distribution as far as possible, thereby maximizing the utility of the KL budget. Plugging in p. qeepp (2 24: 1 2)). have
(4-
4
p. 4 1 1---1
hexp
Also, we have Ens meep 1 1)) = In addition, is easy to see that the constraints 2 0, Vi, are satisfied since deepp (A -1-1 (1, ()) 2 0.
Here. l and À implicitly defined by the constraints (Eq. (24)-Eq. (26)). Now we prove that the dual variables - can be upper bounded by Σ (1l(2;; 0).
Proposition 4. -u satisfies - A Σn 1= PR (i)((z)))
Proof. Recall the constraint
((22; - 11 A
nexp
1.
34

---

By applying Jensen's 's inequality, we have
n
0) ((2000) - 14 1) I.
1(26
Σ
exp
Some algebra give
n
1 M+
exp
Σ
96
& exp
A
A
This implies that -u AS Di=1 (l((0;))
E Additional Experiment Results
E.1
ArmoRM Multi-objective Alignment
Mixture Reward Faree [c)
0,400
2,375
0.33
0.25
CL390-1 (AAO
CDPD HourS
412PD rear)
2.300
W/BD 0. jound
2.25
WEHO n. (aer
8030 Suens
300
ooniina
2.50
nom
.ominal
Figure Evaluation of WDPO, KLDPO and DPO on convexx (a) ArMoRM Multi-ohjective Alignment. We evaluate WDPO with robustness parameter po E {0.005, 0.01) and KLDPO with robustness temperature T E (1,
Similar to the Emotion Alignment experiments, we generate preference labels according to con- combinations of two reward objectives, i.e., Toonvex defined Mixture Evaluation in previous section. Specifically, we consider three pairs of objectives: (1) Ultrafeedback-Hoone Complexity. (2) Ultrrffeedbcck-Hetee Helpfullnsssssss and Helpsteer- Cohpsseer--Chheeneee,, and (3) Ultrafeedback- truthfulness and Helpsteer- Complexity. We generate preference labels according to a 0.5 for all three cases. All models are trained for 4 epochs. Then we introduce reward shift by evaluating WDPO, KLDPO, and DPO on r roonvex (a)) where a E (0, 0.1, 0.3, 0.5, 0.7, 0.9, 1}. In the first plot of Fig. 4, the training preferences are generated accoding reward pair (1). We observe that WDPO and KLDPO achieve superior performance compared to DPO. In particular, when reward distribution shift happens in two directions (towards standalone Ultrafeebback - and
and Helpsteer-
Compslee-CCmplexitty) they clearly outperform DPO. In the middle plot, the training preferences are generated according to reward pair (2). We observe that WDPO with both P. = 0.005 and Po 0.01 are particularly robust against reward distribution shift. Lastly, in the third plot, the training preferences generated according reward pair (3). We observe that both WDPO and KLDPO achieve notable robustness compared to DPO.
E.2 Leaderboard Alignment
In this section, we include alignment results evaluated on all OpenL.M Leaderboard v2 (Fourrier al., 2024) sub-tasks. We list all sub-task names in Table 2.
LLaMA-3.2-1B results: Table 3, we compare DPO, KLDPO, and WDPO trained using LLaMA- 3.2-IB on all 39 sub-tasks of OpenI.MM Leaderboard v2. We observe that our WDPO and KIDPO methods achieve superior alignment performance on the majority of subtasks. Although WDPO and KLDPO slightly underperform few subtasks, their primary strength lies in generalization, precisely because they consistently enhance performance across a diverse range of subtasks.
LLaMA-3.1-8B results: In Table 4, we compare DPO and KLDPO, both trained using LLaMA-3.1 8B. Earlier, we demonstrated that WDPO, trained on LLaMA-3.2-IBB outperforms both DPO and KLDPO. However, WDPO's requirement for dual gradient computations increases its computational complexity. Due to resource constraints, we present KLDPO results for the 8B model, as it is more scalable. Following the I.LMMA-322 3.B-BB experiments, we train KLDDPO for two epochs, the point where DPO achieved optimal robustness via early stopping. Notably, KIDPPO exhibits exceptional performance on math-related tasks.
35

---

Table 2: All sub-task names
OpenLLM Leaderboard v2.
bbh- hal-booleen--eeeee -oolean-expreeeionn
21
bbh-tracking srrackkng-ehhufleed wawerffleddojeete---ee five-objects
22
abb-tracking-shuffe shuffled-objecte seven-objee objects
bbb-date-understandiin
23
-ot-sssuuufeodojjttt---------- three-objecta
bbb-disambiguaaiion--
24
bbh-web-of-lies
bob-formal-fallacte
25
gPqa-diamond
Dbl-geometzic-aheee
26
gpqa-extended
bbb-hyperbaton
27
gpqa-main
abb-logicaaldeedcct/ obeentt
28
ifeval
29
math-algebra-hart
10
obb-loggccal-dettttton----iii
30 bath--oouutnn-------oo----aee
11
bbb-movile-recommeddttiooo
31
math-gecnetry-barz
bbh-navigate
math- aeth-intermddate-alllee
bab-object-coomting
33 nath-num-theery-barr
14
obb-penguins-in-a--
34
math-prealgebra-bard
15 ttrrrraaoannggbbuuttoooooooooeee
35
bbh-ruin-
36
mmlu-pro
bbb-ealient-tranllate error- error-detection
37
18
bbh-snarks
38 musr-object-placcmentt
19
bbh toh-sporta-undorstanne
39
mus--ten--alloattie
20
Obb-temporal-sequan.
1.aMA-3.2-IB
2
5
9
12
13
DPO at Epoch (early stopping)
0.64
0.40
0.52
0.22
0.35
DPO at Epuch 4 (goodlit)
0.50
0.49
0.34
0.36
0.52
0.31
0.52
0.22
0.13
0.31
0.34
0.45
0.36
KILPPO
0.69
0.43
0.53
0.18
0.36
0.48
KLDPO = 0.05
0.71
0.52
0.38
0.38
0.53
0.34
0.52
0.19
0.17
0.33
0.36
0.48
0.41
WDPO A.
0.73
0.40
0.53
0.34
0.46
WDPO 0.005
0.69
0.51
0.41
0.40
0.54
0.32
0.52
0.19
0.15
0.32
0.35
0.50
0.40
L.aMMA3.2-IB
14
15
17
18
19
21
22
23
24
25
26
DPO at Epoch 2 (carly stopping)
0:25
0.16
0.14
0.22
0.53
0.51
0.23
0.18
0.17
0.36
0.49
0.30
0.28
DPO Epoch 4 (goodlit)
0.23
0.20
0.49
KLDPO = 0.1
0.29
0.18
0.12
0.22
0.54
0.51
0.20
0.21
0.14
0.36
0.49
0.28
0.24
0.29
0.20
0.36
0.49
WDPO A. = 0.01
0.26
0.18
0.10
0.30
0.54
0.50
0.20
0.18
0.13
0.39
0.49
0.24
0.26
WDPO 0. 0.005
0.26
0.14
0.12
0.26
0.54
0.50
0.16
0.20
0.16
0.36
0.49
0.25
0.26
LLaMA-3.2-1B
28
31
32
34
35
38
39
DPO Epoch 2 (carly stopping)
0.22
0.48
0.14
0.09
0.08
0.01
0.04
0.19
0.01
0.17
0.50
0.26
0.28
DPO at Epoch 4 (goodfit
0.23
0.48
0.18
0.02
0.03
0.01
0.05
0.13
0.05
0.17
0.49
0.26
0.23
KLDPO
0.24
0.53
0.20
0.05
0.08
0.08
0.12
0.02
0.18
0.52
0.23
0.24
KI.PPO 0.05
0.24
0.56
0.21
0.05
0.05
0.03
0.05
0.13
0.04
0.18
0.52
0.21
0.24
WDPO A. 0.01
0.24
0.52
0.23
0.07
0.07
0.05
0.13
0.04
0.19
0.52
0.24
0.25
WDPO P, 0.005
0.25
0.49
0.19
0.07
0.08
0.04
0.06
0.19
0.03
0.19
0.51
0.25
0.24
Table Evaluation of DPO, KLDPO, and WDPO on all OpenLLM Leaderboard v2 sub-tasks.
F Additional Experiment Details
The raw Emotion dataset (Saravia al., 2018) consists of text samples paired with multi-class labels for six different emotion classes (joy, sadness, love, anger, fear, and surprise). This dataset was then transformed into multi-label dataset, referred as the Emotion Reward Dataset. To create the multi-label dataset, the surprise class was excluded due to its limited representation in the original dataset. Following this, up to three random text samples from the raw dataset were concatenated, and their associated labels were merged. This pre- pre-processing step ensured that the reward model encountered text samples representing multiple emotions during training.
Reward Model Training:
For the reward model, GPT-2 was employed, augmented with a classification head applied to the last token. The model was trained using sigmoid activation function and binary cross-entropy loss, adhering to the standard multilabel classification framework. Training was conducted over 8 epochs with batch size of 64, utilizing the Adam optimizer with a learning rate of 5.0 x 10-5 and weight decay of 0.01. The reward model achieved a test accuracy of 84% and test ROC-AUC score of 0.99, The amotion-speciff scores predicted by this reward model were treated as the rewards for individual emotions. The ArmoRM setups did not need any reward model training.
Supervised Fine-Tuning: Before training the WDPO algorithm, essential to ensure that the model familiarize with the types of texts present in the dataset. To achieve this, we performed
36

---

LLaMA-3.1-8B
2
5
9
12
13
DPO Epuch 2 (early slopping)
0.72
0.60
0.51
0.64
0.57
0.29
0.65
0.41
0.39
0.62
0.48
0.66
0.32
DPO at Epoch 4 (goodfit)
0.70
0.59
0.47
0.59
0.56
0.30
0.65
0.42
0.40
0.61
0.46
0.66
0.32
KLDPO 7= 0.005
0.79
0.58
0.51
0.61
0.56
0.33
0.62
0.35
0.36
0.63
0.48
0.66
0.31
KI.DPO 0.01
0.80
0.59
0.51
0.59
0.55
0.34
0.62
0.36
0.37
0.63
0.49
0.66
0.30
LLaMA-3.1-8B
14
155
16
17
18
19
20
21
22
23
24
25
26
DPO at Epoch 2 (early stopping)
0.46
0.66
0.65
0.51
0.61
0.68
0.41
0.21
0.23
0.34
0.50
0.30
0.28
DPO at Epoch 4 (goodfit)
0.47
0.59
0.66
0.51
0.61
0.70
0.40
0.21
0.20
0.32
0.50
0.27
0.31
KLDPO 0.005
0.47
0.66
0.65
0.54
0.63
0.72
0.46
0.27
0.25
0.34
0.49
0.29
0.25
KLDPO 0:01
0.47
0.65
0.64
0.53
0.65
0.74
0.47
0.28
0.25
0.33
0.50
0.32
0.28
LLaMA-3.1-8B
28
29
31
34
35
38
39
DPO at Epoch 2 (early stopping)
0.29
0.62
0.04
0.02
0.02
0.02
0.05
0.05
0.04
0.33
0.56
0.40
0.35
DPO at Epoch 4 (goodfit)
0.32
0.53
0.04
0.01
0.02
0.01
0.03
0.03
0.02
0.33
0.57
0.40
0.38
KI.DPO 0.005
0.33
0.72
0.42
0.18
0.10
0.06
0.25
0.42
0.10
0.37
0.54
0.26
0.24
KLDPO = 0.01
0.33
0.75
0.44
0.16
0.09
0.04
0.14
0.41
0.07
0.37
0.54
0.26
0.28
Table 4: Evaluation of DPO and KLDPO on all OpenLLM Leaderboard v2 sub-tasks.
supervised fine-tuning (SFT). We selected GPT-2 as the base language model and trained it to predict the next token based on the text samples in the emotion dataset. The maximum length of each text sample was capped at 68 tokens. The model was trained for 10 epochs with batch size of 64. The training used the Adam optimizer (Kingma and Ba, 2014) with a learning rate of 5.0 x 10 7 following 12 warmup steps. Additionally, maximum gradient norm of 10 applied to stabilize the training. The ArmoRM setups did not need any SFT as we used Intruct models which have already undergone multiple rounds of SFT and alignment.
Data Generation: (1) Emotion Alignment: A preference dataset was created, consisting of a chosen and rejected completion for each prompt in the dataset. The first four tokens from each text in the emotion dataset were used as prompts. Using the SFT model, two completions were generated for each prompt. These completions were generated with top-k value of O, top-p of 1, and up to 64 new tokens. The completions were then evaluated using the reward model, and the chosen and rejected completions were determined based mixed metric derived from the predicted rewards. (2) ArmoRM multi-objective A Alignment: Similar to the Emotion setup, we generated a preference datset by sampling two completions per prompt from the Helpsteer2 dataset. Each completion was sampled with a temperature of 0.7, top-p of and up to 1024 new tokens. The prompts were also truncated to a miximum of 1024 tokens. We then fed these prompt-completion pairs ArmoRM and used the scores from the first stage of the model as our multi-objective rewards. The chosen and rejected completions were determined based on a mixed metric derived from the predicted rewards. (3) Leaderboard Alignment: In this setup we sampled 10 completions per prompt in the Helpsteer2 dataset. Each completion was sampled with a temperature of 0.7, top-p of and up to 1024 new tokens. We then fed these prompt-copletton pairs to ArmoRM and used the scores from the second stage of the model as our reward, the completion with the maximum reward was our chosen completion while that with the minimum reward was our rejected completion.
WDPO Implementation: (1) In WDPO training, one of the main challenges is calculating the gradi- ent penalty of the DPO loss with respect to the input. However, since the input is tokenized as integers, gradient cannot be directly calculated. To address this, gradient caleulated with respect to the output of the embedding layer, where gradients are tracked. (2) In line 4 of the tractable WDPO algorithm
(Algorithm 1), we compute the gradient regularizer: R(Tg; D) = (EE 21 [x) L/2 A key implementation challenge arises in distributed L.M training. A naive approach computes the gradient of the pointwise DPO loss with respect to each input, averages the gradient norms over the micro-batch, and applies this as a regularizer to the batch DPO loss on each worker. However, due to the typically small batch sizes in large-scale LLM training, this averaging is performed over very few samples, resulting in a highly noisy and unstable gradient penalty. To mitigate this, we exploit the inequality V I x for 2 1, allowing us to upper bound the regularizer as:
R( D) = po (E.-DD|| .1( 1/2 S Po V2 0))
This leads to tractable approximation of the pointwise WDPO loss: W (zi, po) 1(z;; 0) p. 2V 0)
where ((z)0) denotes the standard DPO loss for sample 2
37

---

WDPO Training: (I) Emotion alignment: The model was trained for 40 epochs with an effective batch size of 64. We used Adam optimizer, with learning rate of 5.0 10 following 12 warmup steps. A maximum gradient norm of 10 was applied to ensure stable training. The DPO beta parameter was set to 0.1 for all training runs. Experiments were conducted on a single 40 GB A100 GPU, requiring gradient accumulation over two steps. (2) LLaMA experiments: The models were trained for 8 epochs with an effective batch size of 128. We used Adam optimizer with a learning rate of 5.0 10- after 10% warmup ratio and then the learning rate was reduced using cosine scheduler. The DPO beta parameter was set to 0.01 for all training runs. Experiments were conducted on an 8xH100 GPU setup, requiring loading the model in 6 and training with DeepSpeed ZeRO-2 optimizer (Rajbhandari et al., 2020).
KLDPO Implementation: In line 3 of the tractable KLDPO algorithm (Algorithm 2), we compute the approximate worst-case kernel P(t) x exp ((1 (Iz,) 0) (1/n) Ei=
I(z;; -i= 0)). key im-
plementation challenge arises in distributed LLM training. A naive approach would calculate the (1/n) Σ (a;;; f) term by averaging ((2;;) across all samples in the micro-batch of its respective worker. However, because micro-batch sizes are typically small in large-scale LLM training, this results in averaging over only a few samples, making the worst-case kernel highly noisy. To mitigate this, we introduce synchronization step that performs all-gather operation to collect [(2;; 0) values from all workers. This enables averaging over the full batch across all workers, significantly reducing the noise in the worst-case kernel.
KLDPO Training: (1) Emotion aligmment: The model was trained for 40 epochs with an effective batch size of 64. We used Adam optimizer (Kingma and Ba, 2014), with a learning rate of 5.0 x 10 following warmup steps. maximum gradient norm of 10 was applied to ensure stable training. The DPO beta parameter was set to 0.1 for all training runs. Experiments were conducted on single 40 GB A100 GPU and gradient was accumulated over two steps to keep training consistent across all algorithms. (2) LLaMA experiments: The models were trained for 8 epochs with an effective batch size 128. We used Adam optimizer with a learning rate 5.0 x 10 after 10% warmup ratio and then the learning rate was reduced using a cosine scheduler. The DPO beta parameter was set to 0.01 for all training runs. Experiments were conducted on an 8xH100 GPU setup. requiring loading the model in 6 and training with DeepSpeed ZeRO-2 optimizer (Rajbhandari et al., 2020).
G Limitations
Theoretical Limitations: Our theoretical analysis relies on Assumption which ensures sufficient data coverage to guarantee strong convexity conditions. Although such data-coverage assumptions standard within offiin learning fixed-dataset scenarios, they are moderately restrictive, as they require the training dataset to sufficiently cover the space of feature differences between the preferred and dis-preferred actions. The log-linear policy class assumption, while standard and easily extendable to neural network policies under mild additional conditions, does not constitute a significant limitation.
Experimental Limitations: Empirically, Wasserstein Direct Preference Optimization (WDPO) involves two separate gradient computations during training, one for calculating the gradient penalty and another for updating policy parameters via standard gradient descent. This dual gradient require- ment can increase computational complexity and training difficulty, potentially limiting practical scalability and efficiency compared to methods with a single gradient computation.
H Impact Statement
This paper aims to advance the field of machine learning by improving the robustness of direct preference optimization against preference model shifts. Our theoretical insights and empirical evaluations contribute to the reliability of preference-base based learning methods. While our work has broad implications for AI alignment and deployment, we do not foresee any immediate societal concerns that require specific highlighting.
38


Reviews:
[{"review_id": "cdnlhCP3MK", "reviewer": "Reviewer_S16R", "summary": "This paper introduces two distributionally robust extensions of Direct Preference Optimization (DPO), named **Wasserstein DPO (WDPO)** and **Kullback-Leibler DPO (KLDPO)**, to address alignment failures of large language models under preference distribution shifts. The authors develop tractable gradient-based algorithms and provide theoretical guarantees. Empirical results show that WDPO and KLDPO consistently outperform standard DPO in scenarios with shifting or diverse user preferences.", "questions": "1. Can you show the effectiveness of your methods on larger models and larger datasets?\n2. Can your distributional robustness framework be extended to non-DPO alignment methods such as PPO-based RLHF or Imitation Learning?\n3. Have you considered evaluating your methods under real-world preference shift scenarios, such as regional, demographic, or language-based distributions, rather than synthetic datasets designed for simplified testing?\n4. Can you discuss more about the computational overhead of your methods?", "limitations": "Yes.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**:\n\n- Proposes practical solutions (WDPO and KLDPO) for addressing preference distribution shift, a critical issue in real-world LLM deployment.\n\n- Combines solid theoretical guarantees with scalable, implementable algorithms.\n\n- Empirical results consistently demonstrate robustness and improvement over standard DPO.\n\n**Weaknesses**:\n\n- Experimental validation is limited in scale, with relatively small models and datasets.\n\n- Benchmark setups are relatively simple and may not fully reflect real-world diversity.", "was_revised": true, "final_justification": "The paper is technically solid and insightful. The authors have also sufficiently addressed my concerns. Therefore I recommend acceptance.", "rebuttal": "We thank Reviewer S16R for the thoughtful and detailed review. We are encouraged to hear that you found our work to provide *“practical solutions (WDPO and KLDPO) for addressing preference distribution shift,”* which you identify as *“a critical issue in real-world LLM deployment.”* We also appreciate your positive remarks on the combination of *“solid theoretical guarantees with scalable, implementable algorithms,”* and your recognition that our *“empirical results consistently demonstrate robustness and improvement over standard DPO.”* We hope our responses further clarify the contributions and scope of our work.\n\n**Q1**: *Can you show the effectiveness of your methods on larger models and larger datasets?*\n\nWe appreciate the reviewer’s interest in scaling our methods. Due to resource constraints common in academic settings, we were limited to training models up to the size of LLaMA 8B. While this may not match the scale of industry models, we believe it still provides valuable insights into scalability. Notably:\n\n- Our algorithms were applied across a range of model sizes: LLaMA 1B (used in the ArmoRM multi-objective alignment task), 3B, and 8B.\n- ArmoRM itself is a powerful LLaMA-8B reward model, offering a strong and challenging signal for preference optimization.\n- Our successful extension of KLDPO to LLaMA 8B demonstrates practical scalability within feasible academic compute budgets.\n\nWe agree that further evaluation at larger scales is an important future direction, and we hope to pursue this as additional compute resources become available.\n\n**Q2**: *Can your distributional robustness framework be extended to non-DPO alignment methods such as PPO-based RLHF or Imitation Learning?*\n\nThank you for this insightful question. Extending distributionally robust alignment beyond DPO is a compelling future direction.\n\nOne promising avenue is to robustify the reward model itself, learning a **distributionally robust reward function** that generalizes well under test-time distribution shifts. For instance, in preference modeling, we can consider a reward function of the form $r_\\theta(s, a) = \\langle \\theta,\\; \\phi(s, a) \\rangle$, where $\\phi$ is a feature map, and estimate it via maximum likelihood using a preference dataset $\\mathcal{D} = \\{(s\\_i, a^1\\_i, a^2\\_i, y\\_i)\\}\\_{i=1}^N \\sim P^o$. The negative log-likelihood loss of $\\theta$ is\n$$\n\\mathcal{L}(\\theta \\mid \\mathcal{D}) = - \\sum\\_{i=1}^N \\log[\\mathbf{1} (y\\_i=1) \\sigma( \\langle {\\theta,\\;\\phi(s_i,a_i^1)-\\phi(s_i,a_i^2)}\\rangle) + \\mathbf{1} (y\\_i=0) \\sigma( \\langle {\\theta,\\;\\phi(s_i,a_i^2)-\\phi(s_i,a_i^1)}\\rangle)].\n$$\n\nRobustifying this objective with respect to distributional shifts in $P^o$, which is used to sample $\\mathcal{D}$, could yield a reward model that is resilient to preference drift.\n\nThis robust reward can in turn be used for PPO-style RLHF, as part of a broader robust alignment pipeline.\n\nWe also note that recent works have proposed distributionally robust behavior cloning and imitation learning algorithms (Bashiri et al.; Panaganti et al.), which could complement our framework in future explorations.\n\nReferences:\n\nM. A. Bashiri, B. Ziebart, and X. Zhang, “Distributionally robust imitation learning,” Advances in neural information processing systems.\n\nPanaganti, K., Xu, Z., Kalathil, D., and Ghavamzadeh, M. (2023b). Distributionally robust behavioral cloning for robust imitation learning. In 2023 62nd IEEE Conference on Decision and Control (CDC), pages 1342-1347.vol. 34, pp. 24404–24417, 2021.\n\n**Q3**: *Have you considered evaluating your methods under real-world preference shift scenarios, such as regional, demographic, or language-based distributions, rather than synthetic datasets designed for simplified testing?*\n\nWe thank the reviewer for this valuable suggestion. While we agree that regional, demographic, and language-based preference shifts represent important real-world challenges, public datasets with explicit annotations for such axes of variation remain limited or proprietary. As a result, conducting controlled and reproducible experiments along these lines is currently nontrivial.\n\nThat said, we would like to emphasize that our experimental setup already captures meaningful and realistic distribution shifts:\n\n- In our **ArmoRM multi-objective alignment** experiment, we simulate preference shifts by training on a mixture of sub-objectives (e.g., helpfulness + correctness) and evaluating on sub-objectives held out during training (e.g., coherence or verbosity). This mimics real-world scenarios where a deployed model may face user goals not seen in training.\n\n- In our **OpenLLM Leaderboard evaluation**, we train using ArmoRM, which represents a single (though strong) judge model, and test on a diverse aggregation of preferences across multiple human and automated evaluation sources. This shift from training to evaluation naturally induces real-world preference mismatch, even if not explicitly labeled by region or demographic.\n\nWe agree that future work could explore explicit real-world axes such as culture, age, or language, especially as more datasets with these attributes become publicly available. We are excited by this direction and view our current work as a foundation for robust preference optimization under broader forms of distributional variability.\n\n**Q4**: *Can you discuss more about the computational overhead of your methods?*\n\nWe thank the reviewer for raising this important point. We provide a detailed discussion on computational overhead and implementation optimizations in **Appendix F** of our paper.\n\nBriefly:\n\n- **WDPO** introduces a gradient penalty term that requires differentiating the DPO loss w.r.t. model input. Since tokenized inputs are discrete, we track gradients at the embedding layer. Moreover, to reduce variance from small micro-batch sizes in distributed LLM training, we upper-bound the square-root penalty with its squared counterpart, leading to a tractable per-example regularizer.\n\n- **KLDPO** introduces a reweighting scheme across the full batch. To stabilize this under distributed training, we insert an `all_gather` operation to synchronize losses across workers, mitigating the high variance from local micro-batch estimates.\n\nDespite these added steps, the rest of the DPO pipeline (optimizer, learning rate schedule, model parallelism setup) remains **unchanged**. We found the overall overhead to be modest and manageable across both small-scale (single A100) and large-scale (8xH100 with DeepSpeed ZeRO-2) experiments. \n\nWe hope this reassures the reviewer that our methods are practical and scalable with careful engineering, and we will make these implementation strategies explicit in our updated appendix."}, {"review_id": "1bSypvGtdE", "reviewer": "Reviewer_4S36", "summary": "In this paper, the authors introduce a new family of DPO-style losses call Robust-DPO that (loosely speaking) optimize for the robustness of models across data distributions. The particular problem they motivate is well illustrated in Figure 1: vanilla DPO often catastropically fails when exposed to new distributions of data, which are commonly encountered given the non-homogenous nature of preference data. Their approach is based on modeling directly (what they call) an \"uncertainty set\" of data distributions around a nominal distribution (e.g., the training distribution), which they use to  perform minimax-style optimization (intuitively: the goal is to minimize the expected loss for each instance wrt the distribution in the uncertainty set that maximizes the loss of that instance, as given in Eq. 6). While I am not familiar with minimax optimization, I found their description very easy to follow, as well as their technical description of why such optimization is difficult. \n\nOperationally, they define two variants of Robust-DPO, Wasserstein DPO (WDPO, Eq. 7) and Kullback-Leibler DPO (KLDPO, Eq. 8) that use Wassertein and KL metrics, respectively, to parameterize these uncertainty sets. Rather impressively, they prove finite-sample error bounds for both (Theorem 1 and Theorem 2), yet due to the difficult of the mini-max optimization, devise tractable variants of both (Section 6) that facilitate SGD-style optimization (i.e., the SGD algorithms they define in Algorithms 1-2). \n\nFor experiments, they perform (somewhat contrived)  on-policy style offline learning using a variety of medium size llama models (1B->8B) using the HelpSteer2 dataset and perform evaluation on the more diverse OpenLLM leaderboard (from what I understand, the OpenLLM represents a more distributionally diverse test set, though these decisions need to be better explained and motivated in the text). Their main results are in Figure 2-3 and Table 1, where their distributionally robust variants achieve significant improvements over the vanilla DPO.", "questions": "-- Did you train on datasets besides HelpSteer? If so, why not? \n\n-- It seems like your approach still assumes that data is of high quality. Do you think that your approach could be extended for modeling data with noise? A discussion of this in the paper would be helpful.", "limitations": "yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "# strengths\n\n-- A very novel distributionally robust variant of DPO and objective that aims to improve DPO's robustness to distribution shifts (an important challenge that is well motivated in the paper) based on minimax-style optimization. I can imagine such work influential otehr work in this area. \n\n-- Deep theoretical results about error bounds of two specific losses they derive (WDPO and KLDPO), coupled with tractable variants of each loss and novel SGD-style algorithms that effectively solve the corresponding minimax optimization problems. \n\n-- Empirical evidence of the effectiveness and utility of their approach over vanilla DPO. \n\n# weaknesses\n\n-- As noted above, the experiments, while sufficient for showing the utility of their approach, are fairly contrived. Given that most standard preference datasets, such as ultrafeedback, are actually aggregations of different datasets (and hence, different distributions of data), I would suspect that their approach would be quite useful in the standard setting too. No such results are reported (see questions below). \n\n-- (minor) When technically very impressive and comprehensive, it is unclear how easily their variants of SGD (Algorithms 1-2) would fit into existing DPO training libraries, which could limit the impact of these new approaches. I would like to see the authors address this. A code example would be very useful to have in the appendix.", "was_revised": true, "final_justification": "I found that the author(s) addressed all my concerns. Providing a Pytorch sketch of the implementation also helped tremendously to clarity certain things. I therefore have decided to keep my already high score high.", "rebuttal": "We sincerely thank Reviewer 4S36 for their thoughtful and encouraging review. We are especially grateful that you found our formulation of distributionally robust preference optimization to be *“very novel”* and *“well motivated,”* and that you appreciated the theoretical depth of our work, including the error bounds and SGD-style algorithms. We are also pleased that you found our empirical evidence convincing and the paper clearly written, despite the technical complexity. Below, we address your insightful questions and suggestions, particularly regarding dataset choices, robustness to noisy data, and integration into existing DPO training frameworks.\n\n**Q1**: *When technically very impressive and comprehensive, it is unclear how easily their variants of SGD (Algorithms 1-2) would fit into existing DPO training libraries, which could limit the impact of these new approaches. I would like to see the authors address this. A code example would be very useful to have in the appendix.*\n\nWe thank the reviewer for highlighting the importance of integration with existing DPO pipelines. Our implementations of **WDPO** and **KLDPO** are designed for minimal disruption: we modify only the loss function, keeping the rest of the DPO pipeline (data, sampling, training loop, etc.) unchanged.\n\nTo illustrate this, we provide the following pseudocode sketch built directly upon the `DPOLoss` structure in `openrlhf`. The core DPO logic remains intact and is reused in both `KLDPOLoss` and `WDPOLoss`. Only the lines marked between `===` blocks are additions specific to each variant.\n\n```python\nclass (KL/W)DPOLoss(nn.Module):\n    def __init__(self, ...) -> None:\n        super().__init__()\n        # DPO parameters\n        ...\n\n        # KLDPO and WDPO hyperparameters\n        self.kldpo_tau = ...\n        self.wdpo_rho = ...\n\n    def forward(\n        self,\n        policy_chosen_logps, policy_rejected_logps,\n        reference_chosen_logps, reference_rejected_logps,\n        ...\n    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n\n        # Calculating vanilla DPO loss (common to DPO, KLDPOLoss, and WDPOLoss)\n        pi_logratios = policy_chosen_logps - policy_rejected_logps\n        ref_logratios = reference_chosen_logps - reference_rejected_logps\n        logits = pi_logratios - ref_logratios\n\n        losses = (\n            -F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing)\n            - F.logsigmoid(-self.beta * logits) * self.label_smoothing\n        )\n\n        # ======== KLDPO Extension (added in KLDPOLoss) ========\n        all_losses = [torch.zeros_like(losses) for _ in range(dist.get_world_size())]\n        dist.all_gather(all_losses, losses.detach())\n        all_losses = torch.cat(all_losses)  # Gather losses across all distributed ranks\n\n        kl_reweights = torch.exp((losses - all_losses.mean()) / self.kldpo_tau).clamp(max=10).detach()\n        # KL reweighting approximates the adversarial distribution in KL-DRO\n\n        loss = (kl_reweights * losses).mean()  # Final KLDPO objective\n        # =======================================================\n\n        # ======== WDPO Extension (added in WDPOLoss) ========\n        grads = torch.autograd.grad(\n            outputs=losses,\n            inputs=model_embeddings,\n            grad_outputs=torch.ones_like(losses),\n            create_graph=(train_eval == 'train')\n        )[0]\n        # Calculate WDPO gradient penalty\n\n        grads = grads.pow(2).sum(dim=(-1, -2))\n        grads = grads[: len(grads) // 2] + grads[len(grads) // 2 :]\n        losses = losses + self.wdpo_rho * grads\n        loss = losses.mean()  # Final WDPO objective\n        # ===================================================\n\n        chosen_rewards = self.beta * (policy_chosen_logps - reference_chosen_logps).detach()\n        rejected_rewards = self.beta * (policy_rejected_logps - reference_rejected_logps).detach()\n\n        return loss, chosen_rewards, rejected_rewards\n```\n**Q2**: *Did you train on datasets besides HelpSteer? If so, why not?*\n\n\nYes, we use two distinct prompt sources in our experiments. The distinction we study is not just prompt diversity, but **preference distribution shift**, so we fix the prompt pool (e.g., HelpSteer2) and vary how reward signals are assigned.\n\n| Experiment                          | Prompt Source       | Preference / Reward Source                     | Why this Setup?                                                                 |\n|-------------------------------------|----------------------|------------------------------------------------|----------------------------------------------------------------------------------|\n| **Toy (Fig 2)**                     | *Emotion Dataset*    | Synthetic pairwise labels (e.g., “anger” vs. “fear”) | Controlled shift between two competing objectives to test DPO degradation.      |\n| **ArmoRM Multi-Objective (Fig 3)**  | **HelpSteer2** (~20k prompts) | *ArmoRM* sub-objectives (Llama3-8B model)       | Allows structured composition of sub-rewards for training, then evaluation on sub-objectives **unseen during training**. |\n| **OpenLLM Leaderboard v2 (Table 1)**| **HelpSteer2**       | Evaluation via OpenLLM leaderboard (GPT-4, Claude, etc.) | Trained on ArmoRM preferences, evaluated using out-of-distribution judge models. |\n\n---\n\n**Why not UltraFeedback prompts?**\n\nWhile *UltraFeedback* has 60k prompts, its pairwise preference labels are fixed and cannot be flexibly recomposed. In contrast, *HelpSteer2* prompts combined with **ArmoRM** enable structured and interpretable simulations of preference shift.\n\n---\n\n**What is ArmoRM and why use it?**\n\nArmoRM is **not a black-box chatbot judge** like GPT-4 or Claude. It is a **multi-objective reward model** trained with *absolute ratings across human-aligned dimensions* (e.g., honesty, helpfulness, verbosity). This design allows us to explicitly control which sub-objectives are used for training vs. testing, precisely modeling the kind of distribution shift our methods are built to handle.\n\n**Q3**: *It seems like your approach still assumes that data is of high quality. Do you think that your approach could be extended for modeling data with noise? A discussion of this in the paper would be helpful.*\n\n\nWe thank the reviewer for raising this important point. While our methods were not originally designed to explicitly model noisy preference labels, we explored this direction in a follow-up experiment and observed promising robustness trends.\n\nSpecifically, we trained DPO, KLDPO, and WDPO on preferences generated by a synthetic reward:  \n`0.5 × HelpSteer-Complexity + 0.5 × UltraFeedback-Truthfulness`,  \nand then introduced **10\\%** and **20\\% label noise** by randomly flipping a subset of the pairwise preferences.\n\nAll models were trained on these corrupted datasets and evaluated on 500 prompts from the HelpSteer2 test set, each with 5 completions. These completions were scored using the same two ArmoRM subrewards: `HelpSteer-Complexity` and `UltraFeedback-Truthfulness`, to isolate the effect of label noise as the only varying factor.\n\n\nThe results below show that WDPO consistently maintains performance under increasing noise, especially compared to DPO and KLDPO:\n\n\n| Model         | \\| 0% Noise \\|                  |                            | \\| 10% Noise \\|                         |                            | \\| 20% Noise \\|                         |                            |\n|---------------|----------------------------------|----------------------------|------------------------------------------|----------------------------|------------------------------------------|----------------------------|\n|               | Helpsteer-Complexity             | Ultrafeedback-Truthfulness | Helpsteer-Complexity                     | Ultrafeedback-Truthfulness | Helpsteer-Complexity                     | Ultrafeedback-Truthfulness |\n| dpo           | 0.5332                           | 0.7158                     | 0.5293                                   | 0.7140                     | 0.5245                                   | 0.7135                     |\n| kldpo-5       | 0.5329                           | 0.7141                     | 0.5287                                   | 0.7154                     | 0.5238                                   | 0.7124                     |\n| wdpo-d001     | 0.5319                           | 0.7226                 | **0.5293**                               | **0.7197**                 | **0.5250**                               | **0.7172**                 |\n\nThis suggests that our distributionally robust formulation, though originally motivated by test-time preference shifts, may also confer some resilience against noisy training data."}, {"review_id": "GoUtVlWI9M", "reviewer": "Reviewer_oGJc", "summary": "This paper suggests novel preference alignment algorithms, WDPO and KLDPO, which provides preference distribution robust learning against distribution shifts that occur between training and testing phases.", "questions": "1. I would like to know how much the performance differs on common benchmarks such as AlpacaEval 2.0 or Arena Hard when using widely adopted preference alignment datasets like UltraFeedback and comparing against algorithms such as DPO or Dr. DPO", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n1. Preference distribution shift is an important problem and is particularly critical because in real-world scenarios, the preference distributions between training and testing phases rarely align perfectly. This methodology effectively addresses this issue.\n2. These algorithms are general-purpose methods that can be directly applied without requiring specific information about the preference distributions during training and testing.\n3. The experiments were conducted across various distribution shift setups and evaluated on multiple models.\n\nWeaknesses\n1. This research is closely related to robust preference learning study (such as IPO and Dr. DPO). However, this study lacks baseline experiments and only compares against DPO.\n2. The experimental performance on Leaderboard Alignment is questionable. For LLaMA-3.2-1B, DPO performance mostly decreases compared to SFT, and while both KLDPO and WDPO show smaller performance drops (almost similar) than SFT, this does not constitute meaningful performance gains. Wouldn't this indicate a flaw in experimental setup?\n3. While preference distribution shift is indeed a problem, it is not as extreme as the current experimental setup suggests. Switching objectives between anger and fear for emotional alignment represents an overly dramatic scenario. Figure 2 shows that in convex mixing, DPO outperforms even at alpha 0.3~0.4. Although DPO's performance drops significantly at higher alpha values, I question whether such setups represent common situations. Alpha 0-0.4 range better represents typical distribution shift scenarios, and DPO performs better in this range. These experimental results suggest that KLDPO and WDPO actually decrease performance compared to DPO when distribution shifts are minimal, which undermines the utility of these methods.", "was_revised": true, "final_justification": "My core concern was the absence of a baseline and the lower performance of the trained model on the Leaderboard compared to SFT. However, the baseline has been added, and the lower performance was due to different reward criteria, so my concerns have been resolved.", "rebuttal": "We thank Reviewer oGJc for taking the time to carefully review our work and for providing thoughtful and detailed feedback. We are glad to know that the reviewer found that *\"preference distribution shift is an important problem\"*, *\"[our] methodology effectively addresses this issue\"*, and *\"experiments were conducted across various distribution shift setups and evaluated on multiple models\"*. Below, we address your specific concerns. \n\n**Q1**: *Comparison between SFT and DPO.*\n\nWe thank Reviewer oGJc for raising a critical point regarding the observation that DPO underperforms the SFT model (LLaMA-3.2-1B-Instruct) on the OpenLLM Leaderboard v2.\n\n**To clarify upfront**: when evaluated on the *same* reward model used to generate the preference dataset, armoRM, **DPO consistently outperforms the SFT baseline**, as expected. We verified this through a controlled evaluation (see supplementary tables below), using the exact model checkpoints from Table 1, and scoring model outputs across all 19 subrewards with armoRM. This confirms the experimental pipeline is sound and not flawed. As expected from classical ML, longer training improves in-distribution performance on the training reward.\n\n| Model                      | helpsteer-helpfulness | helpsteer-correctness | helpsteer-coherence | helpsteer-complexity | helpsteer-verbosity | ultrafeedback-overall_score | ultrafeedback-instruction_following |\n|----------------------------|------------------------|------------------------|----------------------|----------------------|----------------------|------------------------------|--------------------------------------|\n| DPO at Epoch 2             | 0.7256                 | 0.7210                 | **0.7848**           | 0.5245               | 0.5967               | 0.6637                       | 0.7449                               |\n| DPO at Epoch 4             | **0.7286**             | **0.7238**             | **0.7848**           | **0.5233**           | **0.6000**           | **0.6682**                   | **0.7529**                           |\n| LLaMA-3.2-1B-Instruct (SFT)| 0.6588                 | 0.6603                 | 0.7415               | 0.4904               | 0.5598               | 0.6031                       | 0.6521                               |\n\n\n| Model                      | ultrafeedback-truthfulness | ultrafeedback-honesty | ultrafeedback-helpfulness | beavertails-is_safe | prometheus-score | argilla-overall_quality |\n|----------------------------|-----------------------------|------------------------|----------------------------|---------------------|------------------|--------------------------|\n| DPO at Epoch 2             | 0.7505                      | 0.7550                 | 0.7269                     | 0.9281              | 0.7637           | 0.5798                   |\n| DPO at Epoch 4             | **0.7571**                  | **0.7661**             | **0.7362**                 | **0.9345**          | **0.7707**       | **0.5840**               |\n| LLaMA-3.2-1B-Instruct (SFT)| 0.6823                      | 0.6749                 | 0.6342                     | 0.9309              | 0.6342           | 0.5156                   |\n\n\n| Model                      | argilla-judge_lm | code-complexity | code-style | code-explanation | code-instruction_following | code-readability |\n|----------------------------|------------------|------------------|-------------|-------------------|----------------------------|-------------------|\n| DPO at Epoch 2             | 0.5980           | 0.5764           | 0.5954      | 0.5715            | 0.5827                     | 0.6294            |\n| DPO at Epoch 4             | **0.6028**       | **0.5792**       | **0.5971**  | **0.5746**        | **0.5851**                 | **0.6319**        |\n| LLaMA-3.2-1B-Instruct (SFT)| 0.5211           | 0.5157           | 0.5386      | 0.5077            | 0.5258                     | 0.5618            |\n\nWe sampled 500 test prompts from the HelpSteer 2 dataset. For each model, we generated five completions per prompt. These were then scored using armoRM across 19 subrewards, aggregated across 2500 completions.\n\n\nNow back to Table 1 in our paper, the underperformance of DPO stems from a well-known issue in RLHF: **reward hacking and distributional overfitting**. As studied in Azar et al. (2024) and other works on reward hacking (Amodei et al., 2016; Skalse et al., 2022; Eisenstein et al., 2023), methods like DPO (and RLHF PPO) can overfit to the training-time reward model, here, armoRM, when the test-time reward or preference distribution shifts, e.g., to OpenLLM Leaderboard v2.\n\nWe hope this resolves the concern and reinforces the importance of developing preference optimization methods that are robust to reward and preference distribution shifts.\n\n\n**Q2**: *...this study lacks baseline experiments and only compares against DPO...*\n\n\nWe thank Reviewer oGJc for highlighting the importance of comparing against other robust preference learning methods such as IPO and DRDPO. In response, we trained IPO and DRDPO using preferences generated according to 0.5 × Ultrafeedback-Truthfulness + 0.5 × Helpsteer-Complexity. Following our evaluation protocol, we sampled 500 prompts from the HelpSteer 2 test set and generated five completions per model per prompt. Our results show that both KLDPO and WDPO outperform DPO and IPO on several key sub-rewards under the armoRM evaluation. While DRDPO performs competitively on some tasks, KLDPO and WDPO consistently achieve stronger overall scores, confirming the benefit of our proposed distributionally robust methods.\n| Model     | HelpSteer-Coherence | HelpSteer-Correctness | HelpSteer-Helpfulness | UltraFeedback-Honesty | UltraFeedback-Instruction_Following | UltraFeedback-Truthfulness | UltraFeedback-Overall Score |\n|-----------|---------------------|------------------------|------------------------|------------------------|--------------------------------------|-----------------------------|------------------------------|\n| DPO       | 0.7645              | 0.7012                 | 0.7049                 | 0.7231                 | 0.7164                               | 0.7158                      | 0.6396                       |\n| IPO       | 0.7615              | 0.6967                 | 0.6995                 | 0.7165                 | 0.7088                               | 0.7121                      | 0.6356                       |\n| DRDPO     | 0.7626              | 0.6981                 | 0.7029                 | 0.7230                 | 0.7137                               | 0.7122                      | 0.6393                       |\n| KLDPO     | 0.7653              | 0.7019                 | 0.7052                 | 0.7237                 | 0.7162                               | 0.7188                      | 0.6402                       |\n| WDPO      | **0.7698**          | **0.7072**             | **0.7111**             | **0.7303**             | **0.7224**                           | **0.7245**                  | **0.6458**                   |\n\n\n**Q3**: *While preference distribution shift is indeed a problem, it is not as extreme as the current experimental setup suggests. Switching objectives between anger and fear for emotional alignment represents an overly dramatic scenario.*\n\nWe appreciate Reviewer oGJc’s perspective on the severity of preference distribution shifts in our experimental setup. While we agree that the emotion alignment experiment is a simplified toy example designed to isolate and illustrate core failure modes, our remaining experiments address far more realistic and impactful scenarios.\n\nIn particular, our second experiment uses **armoRM**, a strong, publicly available reward model that currently ranks **#32** on RewardBench, outperforming many proprietary models. For comparison:\n\n- **google/gemini-2.5-pro** scores **#30**\n- **openai/gpt-4.1-mini-2025-04-14** scores **#33**\n\nNotably, **Arena-Hard 2**, a widely used benchmark for LLM evaluation, includes models such as **Gemini-2.5**, **Gemini-2.5-Flash**, and **OpenChat 4**, which are in a similar performance tier. armoRM also surpasses GPT-4 in LLM-as-a-Judge setups, making it a viable and meaningful alternative for annotation and benchmarking tasks (e.g., MT-Bench, AlpacaEval-2.0, Arena-Hard).\n\nIn our third experiment, we evaluate models on the **OpenLLM Leaderboard v2**, which aggregates scores across multiple axes and reflects the state-of-the-art in community-driven LLM evaluation. These setups represent genuine real-world scenarios where preference shift arises, such as switching reward models between training, tuning, and deployment phases.\n\nThus, our proposed methods address not only extreme preference shifts but also moderate and subtle ones that are common yet hard to detect in practice. Their value lies in mitigating performance degradation across this full spectrum of shift severity, which is central to ensuring robust alignment in deployed systems.\n\n**References**\n\nMohammad Gheshlaghi Azar, Mark Rowland, Bilal Piot, Daniel Guo, Daniele Calandriello, Michal Valko, and Rémi Munos. A general theoretical paradigm to understand learning from human preferences. ArXiv, abs/2310.12036, 2023\n\nAmodei, D., Olah, C., Steinhardt, J., Christiano, P., Schulman, J., and Mané, D. (2016). Concrete problems in ai safety. arXiv preprint arXiv:1606.06565.\n\nSkalse, J., Howe, N., Krasheninnikov, D., and Krueger, D. (2022). Defining and characterizing reward gaming. Advances in Neural Information Processing Systems, 35:9460–9471."}, {"review_id": "HTHvP784I0", "reviewer": "Reviewer_Nk1b", "summary": "This paper proposes distributionally robust direct preference optimization approach by optimizing the distributionally robust DPO objective induced by Wasserstein distance and KL divergence, leading to W-DPO and KL-DPO. Theoretical analysis has been performed that  bounded the error of the policy parameters with high probability. The paper also details the practical implementation for W-DPO and KL-DPO and experiments on emotion alignment, multi-objective alignment, and openllm leaderboard demosntrate the efficacy of the proposed approach.", "questions": "Q1. In Table 1 there are two rows of DPO at Epoch 2 (early stopping) for LLaMA-3.2-1B but the number on BBH is different (0.33 vs 0.35). What leads to this difference?\n\nQ2. Are there any in-depth analyses and ablation studies on the effect of the parameters $\\tau$ and $\\rho$?\n\nQ3. This paper claims to target on distributional robustness of the algorithm. How is this property related to the benchmark settings in the experiment section? In other words, how are the metrics there connected with distributional robustness?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "## Strengths\n\n1. This paper is well-motivated by the distributionally robust optimization problem and claims to be the first to analyze for direct preference optimization.\n\n2. The experiments covers multiple scenarios and show consistent improvement.\n\n3. The proposed approach is clean and easy to follow.\n\n## Weaknesses\n\n1. Missing some in-depth analyses of the hyperparameters. See Q2 for more details.\n\n2. The experimental evaluation protocol lacks detailed backup and justification. Please refer to Q3.", "was_revised": true, "final_justification": "The authors are strongly encouraged to verify the numbers reported to avoid any typos and errors for the results. Other than this, the paper is in much better shape given the additional clarifications and results in the rebuttal.", "rebuttal": "We sincerely thank Reviewer Nk1b for the thoughtful and constructive review. We are glad to hear that you found our paper to be *“well-motivated by the distributionally robust optimization problem.”* We are also encouraged by your recognition that *“the proposed approach is clean and easy to follow”* and that *“the experiments cover multiple scenarios and show consistent improvement.”* We hope our responses and new results address your concerns and strengthen the overall contribution of the paper.\n\n**Q1**: *In Table 1 there are two rows of DPO at Epoch 2 (early stopping) for LLaMA-3.2-1B but the number on BBH is different (0.33 vs 0.35). What leads to this difference?*\n\nWe thank Reviewer Nk1b for catching this typographical error. Both entries refer to the same model checkpoint, and the correct BBH score is 0.35. As verification, Table 3 in the appendix reports all 39 subtask scores on the OpenLLM Leaderboard v2, and their average confirms the 0.35 value. We will correct this in the final version.\n\n**Q2**: *Are there any in-depth analyses and ablation studies on the effect of the parameters $\\rho$ and $\\tau$?*\n\nWe thank the reviewer for raising this important point. As discussed in **Appendix E**, we performed extensive hyperparameter analysis in both the *ArmoRM Multi-Objective Alignment* and *OpenLLM Leaderboard v2* experiments. For instance, **Figure 4** and **Table 3** (Appendix E) show that both **WDPO** and **KLDPO** exhibit relatively low sensitivity to their respective hyperparameters across tasks.\n\nTo further investigate, we conducted an additional ablation using the checkpoints from **Figure 3**, where preferences were generated from  `0.5 × Ultrafeedback-Truthfulness + 0.5 × Helpsteer-Complexity`. Specifically, we evaluated:\n\n- **KLDPO** with $\\tau$ = 0.5, 1, 5  \n- **WDPO** with $\\rho$ = 0.01, 0.005, 0.001  \n\non **500 HelpSteer2 prompts**, generating **5 completions per prompt**, and scoring all completions using **armoRM’s 19 sub-rewards**.\n\nThe results confirm our earlier findings: **KLDPO** and **WDPO** demonstrate robustness across a wide range of hyperparameter values. While tuning $\\tau$ and $\\rho$ can further improve robustness in specific applications, both methods remain stable and effective without exhaustive tuning.\n\n| Model     | helpsteer-helpfulness | helpsteer-correctness | helpsteer-coherence | helpsteer-complexity | helpsteer-verbosity | ultrafeedback-overall_score | ultrafeedback-instruction_following |\n|:----------|:----------------------|:----------------------|:--------------------|:---------------------|:--------------------|:-----------------------------|:------------------------------------|\n| dpo       | 0.7049                | 0.7012                | 0.7645              | 0.5332               | 0.6243              | 0.6396                       | 0.7164                              |\n| kldpo-0.5 | 0.7052                | 0.7019                | 0.7653              | 0.5303               | 0.618               | 0.6402                       | 0.7162                              |\n| kldpo-1   | 0.7045                | 0.7009                | 0.7652              | 0.5316               | 0.6227              | 0.6400                       | 0.7151                              |\n| kldpo-5   | 0.7040                | 0.7003                | 0.7643              | 0.5329               | 0.6230              | 0.6391                       | 0.7144                              |\n| wdpo-d001 | 0.7088                | 0.7054                | 0.7683              | 0.5319               | 0.6269              | 0.6440                       | 0.7193                              |\n| wdpo-d005 | **0.7111**            | **0.7072**            | **0.7698**          | 0.5356               | 0.6297              | **0.6458**                   | **0.7224**                          |\n| wdpo-d01  | 0.7069                | 0.7027                | 0.7673              | **0.5368**           | **0.6298**          | 0.6417                       | 0.7183                              |\n\n| Model     | ultrafeedback-truthfulness | ultrafeedback-honesty | ultrafeedback-helpfulness | beavertails-is_safe | prometheus-score | argilla-overall_quality |\n|:----------|:---------------------------|:-----------------------|:---------------------------|:--------------------|:------------------|:-------------------------|\n| dpo       | 0.7158                     | 0.7231                 | 0.7076                     | 0.9315              | 0.7109            | 0.5397                   |\n| kldpo-0.5 | 0.7188                     | 0.7237                 | 0.7057                     | 0.9320              | 0.7086            | 0.5441                   |\n| kldpo-1   | 0.7166                     | 0.7223                 | 0.7070                     | 0.9354              | 0.7098            | 0.5400                   |\n| kldpo-5   | 0.7141                     | 0.7215                 | 0.7070                     | 0.9326              | 0.7089            | 0.5377                   |\n| wdpo-d001 | 0.7226                     | 0.7279                 | 0.7122                     | 0.9331              | 0.7185            | 0.5447                   |\n| wdpo-d005 | **0.7245**                 | **0.7303**             | **0.7171**                 | 0.9317              | **0.7225**        | **0.5490**               |\n| wdpo-d01  | 0.7170                     | 0.7250                 | 0.7139                     | **0.9365**          | 0.7206            | 0.5408                   |\n\n| Model     | argilla-judge_lm | code-complexity | code-style | code-explanation | code-instruction_following | code-readability |\n|:----------|:------------------|:----------------|:-----------|:------------------|:-----------------------------|:------------------|\n| dpo       | 0.5587            | 0.5561          | 0.5806     | 0.5538            | 0.5566                        | 0.6068            |\n| kldpo-0.5 | 0.5628            | 0.5554          | 0.5792     | 0.5526            | 0.5568                        | 0.6057            |\n| kldpo-1   | 0.5590            | 0.5554          | 0.5799     | 0.5529            | 0.5559                        | 0.6059            |\n| kldpo-5   | 0.5577            | 0.5553          | 0.5797     | 0.5534            | 0.5559                        | 0.6062            |\n| wdpo-d001 | 0.5642            | 0.5600          | 0.5834     | 0.5564            | 0.5597                        | 0.6109            |\n| wdpo-d005 | **0.5679**        | **0.5632**      | **0.5879** | **0.5607**        | **0.5632**                    | **0.6155**        |\n| wdpo-d01  | 0.5598            | 0.5619          | 0.5854     | 0.5597            | 0.5605                        | 0.6129            |\n\n**Q3**: *This paper claims to target on distributional robustness of the algorithm. How is this property related to the benchmark settings in the experiment section? In other words, how are the metrics there connected with distributional robustness?* \n\nWe thank the reviewer for this important question. Our experimental design specifically targets varying degrees of preference distribution shift to evaluate distributional robustness. We summarize our approach below:\n\n1. **Emotion Dataset (Toy Example)** – This controlled experiment introduces explicit distributional shift by varying the convex mixing coefficient between two competing objectives (e.g., “anger” vs. “fear”). It demonstrates how traditional DPO degrades under even modest shifts.\n\n2. **ArmoRM Multi-Objective Alignment** – This more realistic setup simulates partial coverage of reward space. Models are trained on a mixture of two ArmoRM subrewards and then evaluated on other subrewards unseen during training. This mimics real-world scenarios where deployed models encounter user preferences outside their training signal.\n\n3. **OpenLLM Leaderboard v2** – This is our most challenging and practical testbed. Models are trained using ArmoRM (ranked #32 on RewardBench, then evaluated on OpenLLM Leaderboard v2, which aggregates feedback from various human-preference and judge models. This highlights a genuine failure mode: a developer finetunes a model using ArmoRM but sees performance drop on broader evaluation. Our methods aim to mitigate this degradation.\n\nAcross all settings, **distributional shift is either explicitly modeled or naturally induced**. The metric of interest in each case is the model’s reward under out-of-training-distribution preferences. Hence, **distributional robustness** is defined as the ability to maintain performance across such shifts, whether synthetic or real-world (e.g., varying mixtures of user goals, cultural expectations, or evaluation benchmarks).\n\nWe believe this progression from toy to real-world tasks demonstrates the practical importance of robust preference optimization."}]

Author rebuttal:

Meta-review: The current paper introduces two distributionally robust versions of Direct Preference Optimization (DPO), motivated by the need to address alignment failures of large language models (LLMs) under preference distribution shift. 
The reviewing panel generally appreciated the relevance of the problem addressed, which is of practical importance given the widespread deployment of LLMs in real-world applications where preference distributions may vary significantly from those seen during training.
Reviewers were also in agreement in praising the method as theoretically grounded and implemented in a way that is practically scalable.
The empirical evaluations demonstrate robust improvement over DPO across various benchmarks, and, while a raised weakness of the submission was the limited scale of the experimental validation, the rebuttals helped improving the experimental validation of the method by providing additional experiments comparing the proposed method against other robust preference learning methods.
Overall, the paper presents a well-motivated and technically sound approach to enhancing the robustness of LLM alignment, that reviewers unanimously recommended for acceptance.
Decision: Accept (poster)