Distributional Preference Alignment of LLMs via Optimal Transport

Authors: igor  melnyk
youssef  mroueh
brian  belgodere
mattia  rigotti
apoorva  nitsure
mikhail  yurochkin
kristjan  greenewald
jiri  navratil
jarret  ross
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 18:15:34

Abstract: Current LLM alignment techniques use pairwise human preferences at a sample level, and as such, they do not imply an alignment on the distributional level. We propose in this paper Alignment via  Optimal Transport (AOT), a novel method for distributional preference alignment of LLMs. AOT aligns LLMs on unpaired preference data by making the reward distribution of the positive samples stochastically dominant in the first order on the distribution of negative samples. We introduce a convex relaxation of this first-order stochastic dominance and cast it as an optimal transport problem with a smooth and convex cost. Thanks to the one-dimensional nature of the resulting optimal transport problem and the convexity of the cost, it has a closed-form solution via sorting on empirical measures.  We fine-tune LLMs with this AOT objective, which enables alignment by penalizing the violation of the stochastic dominance of the reward distribution of the positive samples on the reward distribution of the negative samples. We analyze the sample complexity of AOT by considering the dual of the OT problem and show that it converges at the parametric rate. Empirically, we show on a diverse set of alignment datasets and LLMs that AOT leads to state-of-the-art models in the 7B family of models when evaluated with Open LLM Benchmarks and AlpacaEval.  Code for $\mathsf{AOT}$ is available in the Hugging Face  TRL library  \url{https://ibm.biz/AOT_TRL}.

Main paper: 
 introduction: 
 Aligning Large Language Models (LLMs) with human preferences is a crucial step in making these models safe and having them follow instructions faithfully. By ensuring that LLMs adhere to human preferences, values, ethics, and desired behaviors we can reduce the risk of generating harmful, biased, or inappropriate content.Reinforcement Learning from Human Feedback, RLHF [Christiano et al., 2017 , Stiennon et al., 2020 , Ouyang et al., 2022 , Bai et al., 2022] , achieves this by learning a reward model on human preference data, followed by fine-tuning the LLM to maximize the reward score while staying close to the initial reference policy to retain utility from the pre-trained model. Recently, new paradigms departed from RLHF towards direct preference optimization methods such as DPO [Rafailov et al., 2024] , SLIC [Zhao et al., 2023] , and Identity Policy optimization [Azar et al., 2024] . In these approaches, the reward is expressed in terms of the log-likelihood ratio between the LLM policy and the reference model. The training is done on paired preference data, i.e. as triplets of prompts, chosen and rejected sentences, where for each prompt a chosen and a rejected sample are available. The training objective is to maximize the margin between the log-likelihood ratio evaluated on the chosen sentence versus (a) Stochastic Dominance of Reward of Chosen on Rejected: AOT achieves a larger margin between the quantile plots of chosen and rejected rewards. Figure 1 : AOT in the paired & unpaired settings enables first-order stochastic dominance of the chosen reward distribution on the rejected distribution (a) . The margin between the quantiles of chosen and rejected rewards is larger than alternative strategies. In (b), we see that AOT's policy chosen to rejected log-likelihood ratio dominates that ratio for the base model and alternative strategies.the log-likelihood ratio on rejected sentences. When paired preference data is not available, and the preference data instead takes the form of distinct marginals of chosen prompt/response pairs and rejected prompt/response pairs, we refer to this setup as the unpaired data setting. Ethayarajh et al. [2024] used Kahneman & Tversky's prospect theory in the unpaired setting and proposed the KTO method that maximizes the margin between the chosen reward and the average reward of rejected sentences and pushes the reward of a rejected sentence below the average reward of chosen sentences.In this paper, we introduce a new distributional optimization method for fine-tuning LLMs from human preference data. Previous work in the paired setting focused on improving the reward of chosen sentences over rejected sentences on a per-sample basis. This procedure does not lead to a preference on a distributional level of the chosen marginal on the rejected marginal. In probabilistic terms, we would like to induce stochastic dominance of the reward of chosen sentences on the reward of rejected ones. First order Stochastic Dominance (FSD, see e.g. Ogryczak and Ruszczynski, 2002) of a random variable X on a random variable Y , means that all quantiles values of X are larger than those of Y . Our main contribution is introducing AOT, Alignment via Optimal Transport, a new method that enables distributional alignment. We do so by devising a new AOT objective function that induces in the unpaired setting FSD dominance of chosen reward's distribution over rejected reward's distribution. We call this unpaired variant uAOT. In the paired setting, we introduce pAOT that encourages a dominance of chosen to rejected log likelihood ratio of the optimized policy on that ratio for the reference base policy. We show that the AOT cost can be cast as a one-dimensional optimal transport problem that can be solved via sorting and efficiently optimized for the LLM. AOT enjoys also nice statistical proprieties and achieves the parametric rate since its objective can be seen as a smooth one-dimensional optimal transport problem. AOT achieves state-of-the-art results on the Alpaca leaderboard [Dubois et al., 2024] using the Merlinite 7B model [Sudalairaj et al., 2024] as a base and scores as the highest 7B model at the time of writing this paper.To introduce the important concepts of our work pictorially, we show in Fig. 1a the quantile plots of the rewards of AOT and alternative alignment strategies (DPO, KTO) for chosen responses (in green) and rejected responses in (red). The quantile plots are estimated on a paired test set. We see that AOT leads to chosen rewards that have larger margins than those of rejected rewards across all percentiles. More importantly, this margin is larger in AOT models than in policies coming from alternative alignment strategies. We then show in Fig. 1b how the AOT aligned policy's chosen-torejected log-likelihood ratio dominates that same ratio evaluated on the base model's ratio across all percentiles. The distributional alignment induced by AOT ensures a large margin between all quantiles so that the preference is reflected not only on average but distributionally. We formalize distributional preference in the next section.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: Our main claims in the abstract and introduction summarize our methodological contributions and results, which are properly expanded upon in the rest of the paper.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 be 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. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
distributional preference via first order stochastic dominance: 
 First Order Stochastic Dominance For a real random variable Z we denote F (-1) Z : [0, 1] → R the left-continuous inverse of the Cumulative Distribution Function (CDF) F Z :Q Z (p) = F (-1) Z (p) = inf{η : F Z (η) ≥ p} for p ∈ [0, 1].Given two random variables Z 1 and Z 2 , we say that Z 1 dominates Z 2 in the first order if Z 1 has larger quantiles than Z 2 for all percentiles p:Z 1 ≽ FSD Z 2 ⇐⇒ Q Z1 (p) ≥ Q Z2 (p), ∀p ∈ [0, 1].(1)Let X be the space of prompts X and Y be the space of responses Y from an LLM conditioned on a prompt X ∈ X . The reference LLM is represented as policy π ref (Y |X), i.e., as a conditional probability on Y given a prompt X ∈ X . We note the LLM policy we are optimizing by π θ where θ is a parameter belonging to a bounded parameter space Θ ⊂ R d θ . For a measure µ ∈ P(X × Y) and a mapping r : X × Y → R, we note as r ♯ µ the pushforward map of µ through r. In particular, for empirical measures µ = 1 ,yi) . DPO as a Pointwise Preference Approach In Direct Preference Optimization (DPO, Rafailov et al., 2024) , the reward being optimized by the LLM has the following form :n n i=1 δ (xi,yi) , we have that r ♯ µ = 1 n n i=1 δ r(xir θ (x, y) = β log π θ (y|x) π ref (y|x) + β log(Z(x)),where Z(x) is a normalization constant. DPO assumes access to a paired preference dataset (X, Y + , Y -) ∼ µ where Y + denotes a positive (chosen) response to which we would like to assign a high reward, and Y -a negative (rejected) response to which we would like to assign a low reward. This can be formalized as minimizing the logarithmic sigmoid loss :min θ∈Θ -E (x,y+,y-)∼µ log(σ(β(r θ (x, y + ) -r θ (x, y -)))),and since the difference is taken for the same x, the normalization Z(x) disappears resulting in:min θ -E (x,y+,y-)∼µ log σ β log π θ (y + |x) π ref (y + |x) -β log π θ (y -|x) π ref (y -|x) .We can interpret this as a pointwise constraint inducing preference for positive over negative reward outcomes as follows:EQUATIONDPO can then be interpreted as a relaxation of this constraint through the logistic loss, which also suggests other preference optimization algorithms through relaxations using, for example, the hinge loss as proposed in SLIC [Zhao et al., 2023] .Our main insight from looking at the pointwise constraint in Eq. ( 2) is that we can recast it as a distributional constraint in terms of stochastic dominance of the random variableZ + θ = log( π θ (Y+|X) π ref (Y+|X) ) of positive outcomes on the random variable Z - θ = log( π θ (Y-|X) π ref (Y-|X)) of negative outcomes. This is especially valuable in the unpaired setting without access to triplets of prompts and positive and negative responses as required by DPO. This is indeed the same setting considered by KTO [Ethayarajh et al., 2024] . The following paragraph formalizes this unpaired distributional preference.We assume here that we don't have access to triplets of prompts and positive/negative responses (x, y + , y -). Instead, we assume separate access to µ + ∈ P(X × Y), a distribution of positive prompt/response pairs (X + , Y + ) we would like to be highly rewarded and reinforce in the policy, and µ -∈ P(X × Y) the distribution of the negative samples (X -, Y -) to be associated with low reward. We define the distributional preference as follows:Definition 1 (Distributional Preference in the Unpaired Setting). A policy π prefers distributionally µ + on µ -with respect to a reference policy π ref if:log π θ (Y + |X + ) π ref (Y + |X + ) ≽ FSD log π θ (Y -|X -) π ref (Y -|X -) .In other words, noting r u • π θ (x, y) = log π θ (y|x) π ref (y|x) , the distributional preference in the unpaired setting means that we have the following constraint:(r u • π θ ) ♯ µ + ≽ FSD (r u • π θ ) ♯ µ -.(3)Distributional Paired Preference Note that we can rewrite Eq. ( 2) in the equivalent form:EQUATIONIn order to turn this into a distributional constraint we need access to a paired preference dataset as in DPO (X, Y + , Y -) ∼ µ, and impose stochastic dominance of the random variableZ θ = log π θ (Y+|X) π θ (Y-|X)indexed by the policy we are optimizing on the random variable -|X) indexed by the reference policy. Z θ and Z ref represent here the log likelihood ratio of positive to negative outcome under the policies π θ and π ref , respectively. Hence, it is desirable to constrain the policy π θ to have a larger excess log probability between positive and negative outcomes than that resulting from the reference policy π ref .Z ref = log π ref (Y+|X) π ref (YWe define below more formally the paired distributional preference via stochastic dominance: Definition 2 (Distributional Preference in the Paired Setting). We say that the policy π θ distributionally dominates π ref in terms of log probability ratio of positive and negative responses if:log π θ (Y + |X) π θ (Y -|X) ≽ FSD log π ref (Y + |X) π ref (Y -|X) .Noting r p • π θ (x, y + , y -) = log π θ (y+|x) π θ (y-|x) this can be written as follows:(r p • π θ ) ♯ µ ≽ FSD (r p • π ref ) ♯ µ.(5)= (-t) 2 + , we have:E OT h (r • π θn ) ♯ µ + , (r • π θn ) ♯ µ -≲ n -1 2 .We see that under our assumptions and for the hinge loss squared, the expected violation of the desired dominance in AOT unpaired converges to zero as n → ∞. Remark 1. While in Section 3, we used the primal formulation to compute OT h due to its computational appeal thanks to the sorting algorithm we used for analyzing the sample complexity the dual of OT h . The dual reveals the game theoretic aspect of AOT as a min-max game between the policy π θ and the dual potential φ c that imposes FSD on the preference we want to infuse to the policy.Setting h(x) to be the hinge loss (-x) + reduces to measuring the Wasserstein-1 distance from U θ to the nearest distribution that has FSD over V θ .3. Wasserstein-2 violation Setting h(x) to be the squared hinge loss (-x) 2 + reduces to measuring the Wasserstein-2 distance from U θ to the nearest distribution that has FSD over V θ .Besides measuring different quantities, the optimization-theoretic properties of each are different:1. The 0-1 loss This loss does not penalize the size of the violations, making gradient-based optimization difficult as large violations have no gradient. Additionally, if FSD were not achievable (e.g. a strong teacher policy), not penalizing the size of the violations could result in risky policies.2. The hinge loss, i.e. the Wasserstein-1 violation measure. By its nature, the gradient of small violations is just as large as the gradient of large violations, which may be beneficial for convergence to an FSD result. When FSD is impossible to achieve, this loss will also have the effect of encouraging sparse violations while still penalizing the size of the violations, similar to the L1 norm in the classic Lasso algorithm. Smooth relaxations Smooth relaxations of the hinge, e.g., the logistic loss described in the main text, have a nonzero gradient at zero and continue to have a gradient for small positive values. This has the benefit of encouraging quantiles to continue improving after surpassing those of the reference.3. The squared hinge loss, i.e. the Wasserstein-2 violation measure. As a quadratic loss, it no longer prefers sparse violations (c.f. L2 norm regularization). Indeed, the gradient signal vanishes as the violation becomes small, meaning that dense violations are likely and slow to be removed. A potential failure mode would be that the new policy has a small violation, but none of the quantiles outperform the baseline. Relaxation Introducing a bias β as in the main text addresses this issue, ensuring the gradient at 0 is nonzero.G Gradients of the sorting-based objective (11)We here repeat the sorting-based objective (11) that we propose, adding n as a superscript for clarity.EQUATIONWe use gradient-based optimization approaches in practice to find a minimizing θ. We have the following theorem showing the gradients are asymptotically unbiased (and thus friendly to stochastic gradient methods). Theorem 4. Let h ′ = dh/dt be L-lipschitz and the gradients of u θ and v θ with respect to θ have 1-norm bounded by M , i.e. ∥∇ θ u θ ∥ 1 ≤ M, ∥∇ θ v θ ∥ 1 ≤ M for all θ ∈ Θ. Suppose further that the support of U θ , V θ are bounded and their distributions have densities (i.e. are atomless).3 Then the gradient of the objective in (21) is asymptotically unbiased as n → ∞.Proof. Observe that the gradients take the formEQUATIONNote that as described in the main text, here, the current gradient is determined by the current state of the sorting of the samples, but it is not necessary to differentiate through the sorting algorithm itself as it changes only discretely.We wish to understand the convergence of the bias of these gradients as n → ∞.We can write (noting that the n atoms of the empirical F n,V θ are distinct with probability 1)1 nNote I ′ is a simple empirical average and hence unbiased in the sense that E[I ′ ] is constant with n. Now by ( 24) we can write∥I -I ′ ∥ 1 ≤ L n (i,2) θ .We thus haveEQUATIONwhereF (i)n,V θ are independently realized empirical quantile functions. Theorem 3.1 of Del Barrio et al. [2018] with the subsequent Remark 3.2.1 therein applies directly to this regime, with the following restated result:Theorem 5 (Special case of Theorem 3.1 in light of Remark 3.2.1 in Del Barrio et al. [2018] ). If F and G are CDFs of 1-dimensional distributions with bounded support and G -1 is continuous on (0, 1), then1 0 F -1 n -G -1 m 2 - 1 0 F -1 -G -1 2 → p 0 as n, m → ∞.Applying Theorem 5 to our setting and noting that for us the second integral is zero, we have that if V θ has bounded support,L 2 M 2 [F (1) n,V θ ] -1 (t) -[F (2) n,V θ ] -1 (t) 2 dt → p 0where → p indicates convergence in probability as n → ∞. This implies that I -I ′ converges to 0 in probability. The proof for II is identical. Since the gradient ( 23) is then a sum of two unbiased terms and two terms that converge in probability to zero, the gradient is asymptotically unbiased.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: The first section in the appendix of the paper identifies the main limitations of our work.Guidelines:• The answer NA means that the paper has no limitations, while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to 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 a 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, a facial recognition algorithm may perform poorly when the image resolution is low, or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to 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 fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not to penalize honesty concerning limitations. 
aot: alignment via optimal transport a convex relaxation approach: 
 Note that the paired and unpaired distributional preference constraints in Definitions 1 and 2 can be used in LLM alignment as follows:Find π θ ∈ H such that (r u • π θ ) ♯ µ + ≽ FSD (r u • π θ ) ♯ µ - (FSD unpaired) and Find π θ ∈ H such that (r p • π θ ) ♯ µ ≽ FSD (r p • π ref ) ♯ µ (FSD paired)where r u are r p are given in Definitions 1 and 2 respectively, and H is a hypothesis class. Those two problems are instances of learning with stochastic orders introduced in [Domingo- Enrich et al., 2022] , but in a simpler setting since the constraints are on one-dimensional distributions and the order considered is the first order rather than the convex order as considered in [Domingo-Enrich et al., 2022] . Note that both problems are special cases of the following generic optimization problem:Find θ ∈ Θ such that : U θ ≽ FSD V θ (6)where U θ and V θ are real-valued random variables whose distributions depend on a parameter vector θ ∈ Θ. Note that for our FSD paired setting, V θ = V (independent of θ). Let µ U θ and µ V θ be the probability measures of U θ and V θ resp.By the definition of FSD in Equation (1) we have:U θ ≽ FSD V θ ⇐⇒ Q U θ (t) ≥ Q V θ (t), ∀t ∈ [0, 1].We can relax this problem to the following minimization problem:EQUATIONwhere h is a function penalizing each quantile's violation of FSD. The objective function (7) seeks to measure the violation of FSD, so that it can be minimized or eliminated. For instance, with h the 0/1 loss (here 1 is the indicator function):EQUATIONThis loss reminds us the misclassification 0/1 loss. Following classical convex relaxation of 0/1 losses in binary classification [Bartlett et al., 2006] , we consider surrogates h of the indicator function.Our choices for h are motivated by the "almost-FSD" notions in the literature (See Appendix F for a discussion). In practice, we use smooth convex approximations of the 0/1 loss (1 x<0 ) [Bartlett et al., 2006] , for example for a margin β > 0 h(x) = (βx) 2 + the βsquared hinge loss or h(x) = log(1 + exp(-βx)) the β-logistic loss. Although not a convex relaxation of the 0/1 loss, the least squares loss has been used in classification [Rosasco et al., 2004] , and in the context of alignment, it was used in IPO [Azar et al., 2024] hence we use also h(x) = (βx) 2 , and refer to it as β-Least Squares. Further discussion of tradeoffs and benefits of different losses is in Appendix F, and formal assumptions on h needed for the statistical theory are given in Assumption 1.The cost function in ( 7) is still computationally challenging, if we were to solve the problem via gradient descent on θ this would require us to differentiate through the quantile operation. The following theorem from Santambrogio [2015] will be instrumental for us to cast the loss in (7) as an optimal transport problem with a convex cost h: Theorem 1 (Theorem 2.9 and Proposition 2.17 in Santambrogio [2015] ). Let h : R → R + be a convex function we have for two real random variables U, V , with measures µ U , µ V :1 0 h(Q U (t) -Q V (t))dt = min γ∈Π(µ U ,µ V ) h(u -v)dγ(u, v) = OT h (µ U , µ V ) and γ * = (Q U , Q V ) ♯ L 1 ([0, 1]) is a minimizer (where L 1 is the Lebesgue measure on [0, 1] ). If furthermore h is strictly convex γ * is the unique minimizer.Thanks to Theorem 1 we can write the problem (7), in the following equivalent form that we call Alignment via Optimal Transport (AOT) :EQUATIONThis formulation reveals that we have turned the stochastic dominance constraint to an inner onedimensional optimal transport problem with a convex cost h. OT h (µ U θ , µ V θ ) can be thought as a soft measure of the violation of the stochastic dominance of U θ on V θ , hence by minimizing it as function of θ we are ensuring the optimal θ * results in U θ * dominating V θ * . Such OT problems with a smooth cost have been subject to theoretical and statistical study in one dimension as well as in high dimensions. For instance, [Manole and Niles-Weed, 2024 ] considered smooth convex costs, and [Hundrieser et al., 2022] considered more general smooth costs. [Groppe and Hundrieser, 2023] considered entropic regularization of optimal transport with general smooth costs.Computational Algorithm via Sorting We consider here empirical measures and turn to solve the inner problem for a fixed θ. We omit θ in what follows to simplify notation. We are interested in OT h (μ U , μV ) where μU = 1 n n i=1 δ ui μV = 1 n n i=1 δ vi . Given the convexity of h and thanks to Theorem (1), the optimal coupling of OT h (μ U θ , μV θ ) is given by the north-west corner solution [Peyré and Cuturi, 2019] (Chapter 3, Section 3.4.2) that informally matches the i-th smallest element of U with the i-th smallest element from V . More formally, if we sort the variables u i and get the order statistics (from min to max) u (1) ≤ ... ≤ u (n) and same for v i : v (1) ≤ ... ≤ v (n) . We have:OT h (μ U , μV ) = 1 n n i=1 h(u (i) -v (i) ).(10) Back to (9), given empirical samplesEQUATIONθ be the order statistics as function of θ. We have therefore:EQUATIONIn Appendix G, we show that the gradients of the objective (11) are asymptotically unbiased for bounded distributions (see the statement for all conditions). Note that the sorting operation in ( 11) is a 1-Lipschitz function with discontinuous Jacobian.2 Like the ReLU activation function, it can be easily optimized by gradient descent [Anil et al., 2019] (compare also sliced Wasserstein GANs). In practice, computing the gradient at any given step is done by first running the sorting algorithm and taking the gradient with respect to θ with the current assignment held fixed.AOT for Unpaired Preference Let μn -,yi,-) . Our convex relaxation approach for unpaired FSD alignment given in (FSD unpaired) can therefore be cast as an AOT problem (given in Equation ( 11)) for+ = 1 n n i=1 δ (xi,+,yi,+) and μn -= 1 n n i=1 δ (xi,u i θ = log π θ (y i,+ |x i,+ ) π ref (y i,+ |x i,+ ) , v i θ = log π θ (y i,-|x i,-) π ref (y i,-|x i,-) , i = 1, . . . , n.AOT for Paired Preference Let μn = 1 n n i=1 δ (xi,yi,+,yi,-) be a paired preference empirical measure. Our convex relaxation approach for paired FSD alignment given in (FSD paired) can be there cast as an AOT problem (given in Equation ( 11)) for:u i θ = log π θ (y i,+ |x i ) π θ (y i,-|x i ) , v i θ = log π ref (y i,+ |x i ) π ref (y i,-|x i ) , i = 1, . . . , n.AOT with Soft Sorting One caveat of the alternating optimization for AOT between θ and solving the inner optimal transport problem with hard sorting is that the gradient with respect to the parameter θ for fixed permutations has dependency in θ on the order statistics level only and not through the sorting routine. To alleviate that, we propose to use SoftSorting [Blondel et al., 2020 , Cuturi et al., 2019] that uses an entropic regularization to find a smoothed permutations via a Sinkhorn algorithm, which in turn allows the back-propagation on θ to depend not only via the order statistics but also via the computational graph of SoftSorting.Algorithms 1 and 2 in Appendix B summarize our AOT approach for distributional preference alignment in the unpaired and paired setting.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? 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 or 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 a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. 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 or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
statistical analysis: 
 In this section, we focus on the statistical analysis of unpaired-AOT and defer paired-AOT to Appendix E since it has a similar analysis. We make the following assumptions on the OT cost h, the reward r, and the policy hypothesis class H. Assumption 1 (OT cost). Let M, R > 0 be finite positive constants. We assume that the loss h: [-M, M ] → [0, R],is convex L-Lipchitz and bounded. h is a convex function (E.g. a relaxation of the 0/1 loss such that h(t) > h(t ′ ), for t < 0 and t ′ > 0).Assumption 2 (Reward). We assume that r is bounded so that r• π θ (x, y) ∈ [-M, M ].Assumption 3 (Assumption on the hypothesis class of the policy). We assumeπ ref , π θ ∈ H = {π θ : such that r • π θ differentiable in θ and sup x∈X ,y∈Y ∥∇ θ r • π θ (y|x)∥ ≤ L ′ , θ ∈ Θ ⊂ B 2 (r 0 , d θ )}, for L ′ , r 0 > 0. . Assumption 4. There exists π θ ∈ H such that (r • π θ ) ♯ µ + ≽ FSD (r • π θ ) ♯ µ -.Assumption 1 is satisfied for example by the hinge squared loss h(t) = (-t) 2 + by the logistic loss h(t) = log(1 + e -βt ), for t ∈ [-M, M ]. Assumption 2 on the boundedness of the rewards can be imposed by clamping the values of the logits of the policies to [-M, M ], which is common practice in practical implementations of LLM alignment. Assumption 3 is a technical assumption needed to control the covering number of the r • H. Assumption 4 ensures the existence of the minimizer in H. We overload notations in what follows and refer to r u and r p as r to simplify the presentation. By our relaxation approach described in Section 3 we can relax the unpaired stochastic dominance constraint problem given in (FSD unpaired) to:min π θ ∈H 1 0 h Q (r•π θ ) ♯ µ+ (t) -Q (r•π θ ) ♯ µ-(t) dt = min π θ ∈H OT h ((r • π θ ) ♯ µ + , (r • π θ ) ♯ µ -) (uAOT h ) Define the OT cost c : [-M, M ] × [-M, M ] → [0, R] such that c(z, z ′ ) = h(z -z ′ ), for z, z ∈ [-M, M ]. Define the c-transform of a function φ : [-M, M ] → R: φ c (z) = inf z ′ ∈[-M,M ] h(z -z ′ ) -φ(z).In our setting, a function is called c-concave if there exists ψ :[-M, M ] → R such that φ = ψ c . Define: F c = {φ : [-M, M ] → [-R, R], φ is c-concave, with ||φ c || ∞ ≤ R}By duality (Theorem 5.10 in [Villani, 2009] ) we have:OT h ((r • π θ ) ♯ µ + , (r • π θ ) ♯ µ -) = sup φ∈Fc φ(r • π θ )dµ + -φ c (r • π θ )dµ -.Replacing the dual expression of OT h in (uAOT h ), we see that (uAOT h ) can be cast as a min-max problem:EQUATIONGiven samples μn+ = 1 n n i=1 δ (xi,+,yi,+) and μn -= 1 n n i=1 δ (xi,-,yi,-), the empirical problem is:EQUATIONRecall that OT h is a measure of the violation of stochastic dominance of (r• π θ ) ♯ µ + on (r • π θ ) ♯ µ -.We have the following result on the sample complexity of the violation of stochastic dominance: Theorem 2 (Sample Complexity of Dominance Violation for AOT Unpaired). Let π θ * be the population minimizer of (uAOT h ) and π θn be the solution of the empirical problem (13). We have the following sample complexity bound for the violation of stochastic dominance in AOT unpaired:E OT h (r • π θn ) ♯ µ + , (r • π θn ) ♯ µ -≤ OT h ((r • π θ * ) ♯ µ + , (r • π θ * ) ♯ µ -) Optimal Almost FSD Violation + 2R n (F c ; (r • π θ * ) ♯ µ + ) + 2R n (F c c ; (r • π θ * ) ♯ µ -) One dimensional OT sample complexity with optimal θ * + 2R n (F c • r • H; µ + ) + 2R n (F c c • r • H; µ -) Complexity of learning in H via the 1D OT problem , where R n (F; ν) = E sup φ∈F 1 n n i=1 σ i φ(Z i )is the Rademacher Complexity and for i = 1 . . . n, σ i are independent Rademacher random variables and Z i ∼ ν iid.By considering our assumptions on the cost, the reward, and the hypothesis class, we obtain the parametric rate in n: Corollary 1. (Informal) Under Assumptions 1, 2 and 3 we have:1. E OT h (r • π θn ) ♯ µ + , (r • π θn ) ♯ µ --OT h ((r • π θ * ) ♯ µ + , (r • π θ * ) ♯ µ -) ≲ n -1 2, where ≲ refers to inequality up to constants that depend only on constants in the assumptions. 
experiments: 
 In this section, we evaluate the performance of the proposed AOT method on a diverse set of base LLMs and datasets, comparing with currently available alternative alignment algorithms.We compared AOT with current state-of-the-art alignment approaches, specifically Direct Preference Optimization (DPO) [Rafailov et al., 2024] , Kahneman-Tversky Optimization (KTO) [Ethayarajh et al., 2024] and Identity Policy Optimization (IPO) [Azar et al., 2024] . DPO and IPO operate on paired preference data, while KTO can handle both paired and unpaired prompt/response samples. Reference Models Traditionally, model alignment is the third and final step applied to the LLM that already has gone through original pretraining and supervised fine-tuning. For our experiments, we selected a range of models at various stages and with different levels of performance, all in the family of 7B-parameter models. Specifically, we used Merlinite-7B [Sudalairaj et al., 2024] , which is a variant of Mistral-7B-v0.1 that has been instruction-tuned (SFT) on data from a synthetic data generator using a taxonomy-driven data curation process. In Appendix H we also cover other popular LLMs, such as Mistral-7B [Jiang et al., 2023] , OpenHermes-2.5-Mistral-7B [Teknium, 2024] , Datasets For our experiments, we used both paired and unpaired datasets. For the paired dataset, we used the UltraFeedback binarized dataset from [Tunstall et al., 2023b] , containing over 60K training samples, where for each prompt, there is a pair of chosen (preferred) and rejected (not preferred) responses. This alignment dataset is widely used, and all compared alignment techniques are well-suited for it. For unpaired datasets, we used PKU BeaverTails [Ji et al., 2023] with over 300K samples and HelpSteer [Wang et al., 2023] with around 35K samples. Here, for each prompt, there is only a single response with a score defined by some attributes (e.g., safety, faithfulness, helpfulness, etc.). We used the sum of attribute values and thresholded by the median to binarize the responses into chosen and rejected. For this unpaired dataset, only KTO and our AOT are applicable.StarlingMetrics To measure the performance of different alignment methods, we used popular evaluation metrics, AlpacaEval [Dubois et al., 2024] and Open LLM benchmark [Beeching et al., 2023] . We note that Alpaca uses GPT4 model as a judge to compare candidate responses to GPT4-based references on a set of 805 challenging questions. The GPT4-based evaluations are expensive, so to limit our expenses, we also employed a very strong and capable Llama3-70B-Instruct [AI@Meta, 2024] as a judge. As we show in Appendix H in Table 2 , the order determined by Llama3-70B-InstructFigure 2 : Impact of batch size and loss type on AOT performance. The batch size is the effective number of samples in the mini-batch per GPU. We found the logistic loss to be performing better than least squared or hinge squared losses (all using β = 0.01). As we increase batch size, we also observed improvement in AOT performance, which is expected as more samples per minibatch results in a better effect of stochastic dominance (conforming Corollary 1). Figure 3 : Impact of (β) parameter on performance of different alignment algorithms. β controls the divergence of the policy model from the initial reference model (low beta -more divergence, high beta -less divergence). We see a general trend that with higher betas, LLMs alignment decreases the performance. Hence, for all experiments, we selected β = 0.01 as a default value. and GPT4 is the same (the absolute score values are different), providing a better free alternative LLM-judge for local Alpaca evaluations. For intermediate results, we also employed Tiny Benchmarks [Maia Polo et al., 2024] to approximate original metrics and provide fast feedback during the initial development. We also evaluated the aligned models on six key benchmarks from the Open LLM Leaderboard: AI2 Reasoning Challenge -ARC (grade-school science questions), HellaSwag (commonsense inference), MMLU (multi-task accuracy), TruthfulQA (tendency to reproduce falsehoods), Winogrande (commonsense reasoning) and GSM8K (grade school math word problems). Note that in all the above benchmarks we use 0-shot prompts, a more challenging setting as opposed to commonly used few-shot prompting.Experimental Setup Our implementation is based on the HuggingFace Alignment Handbook Tunstall et al. [2023a] . As we show in Appendix in Section B, the changes needed to adapt HF TRL trainer [von Werra et al., 2020] for AOT are minimal and therefore can easily be adapted by the community. For each run our compute setup consisted of 8 H100 GPUs. We used LoRA [Hu et al., 2021] for parameter-efficient fine-tuning during alignment and the FSDP (Fully-Sharded Data-Parallel) setup to train the model over multiple GPUs. Under this setup, the training of each 7B-parameter model on the UltraFeedback dataset took approximately one hour. The evaluation on AlpacaEval and Open LLM benchmarks took one additional hour to get the final results.In Table 1 , we present the main results of comparing AOT to other baselines (KTO, DPO, and IPO) on paired UltraFeedback binarized dataset. On AlpacaEval (GPT4), our AOT unpaired approach scores 31.3%, which is a significant gain from the base Merlinite-7B model. As of time of this writing (May 22nd, 2024) , this result places our AOT aligned LLM on AlpacaEval LeaderBoard ahead of such strong competitors as KTO-Mistral-PAIR [Ethayarajh et al., 2023] and other 7B-parameter models, reaching the level of Mixtral-8x22B-v0.1 (see Figure 4 in Appendix for an illustration). On other LLM benchmarks AOT performs competitively to other baselines. As mentioned earlier, these evaluations are done using 0-shot prompts, leading to a more challenging setting and resulting in overall lower performance across metrics and baselines. For other base LLMs we show their performance in Appendix H (see Tables 3, 4, 5, and 6 ).We also examined the effect of batch size and the choice of loss function on AOT performance, results shown in Fig. 2 . As the batch size increases, AlpacaEval (based on Llama3-70B-instruct) also increases in line with our theory in Corollary 1. Note that our current setup (FSDP over 8 H100 GPUs) limits our batch size to 35 samples per GPU. We have also examined the impact of beta (controlling divergence of policy from reference) on AOT performance in Fig. 3 . We noticed a trend that with higher betas the performance of LLMs alignment decreases, thus we set β = 0.01. Ablation results comparing hard and soft sorting as well as the variance of AlpacaEval scores across multiple runs in Appendix H (Tables 7 and 10 ) show the overall robustness of AOT . 
conclusion: 
 We present in this paper Distributional Alignment via Optimal Transport (AOT) for large language models. The AOT cost can be cast as a one-dimensional optimal transport problem with a smooth and convex cost that penalizes violations of the dominance of the chosen on rejected marginals. AOT enjoys parametric statistical rates. We showed with extensive experimentation on various paired and unpaired datasets, base models, and different loss functions, that AOT alignment robustly leads to aligned models that outperform alternative alignment strategies such as DPO, KTO and IPO on the Alpaca Benchmark, leading to the best 7B model to date on that benchmark as of the time of writing.  
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]Justification: [NA] Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 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 a best faith effort.12. Licenses for existing assets Question: 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?Answer: [Yes] Justification: Creators of original work are properly cited.Guidelines:• The answer NA means that the paper does not use 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 is used and, if possible, include a 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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset'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: Documented code for our proposed alignment algorithm will be provided upon acceptance of the paper. Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model 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 is used. Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: In this paper, we introduced an alignment method for LLMs that can capture rewards at the distributional level without the requirement of paired preference data. Our algorithm was derived by imposing the stochastic dominance of positive reward distribution over negative distributions through an Optimal Transport formulation. It enjoys a very simple algorithmic implementation in terms of a closed-form expression via sorting on empirical measures. Empirically, our algorithm demonstrates excellent results by allowing us to train 7B parameter models that achieve state-of-the-art evaluation results on Open LLM Benchmarks and AlpacaEval.In terms of broader societal impact, we would like to highlight the benefits that our AOT algorithm will bring to RLHF by enabling a more robust distributional alignment of LLMs, improving their ability to follow instructions accurately, and aligning their responses with human values.Our work shares the same limitations and possible negative broader societal impacts as the majority of RLHF work. The algorithm is fundamentally limited by the training dataset used for alignment and might, therefore, contribute to amplifying various types of bias present in the data. In addition, alignment through AOT is not enough to address aspects related to the security and safety of LLM deployment. In general, better performance on a given set of benchmarks following alignment does not imply better performance across the board in other tasks, and ad-hoc evaluation specific to each task of interest is warranted. {(x i,-, y i,-) ∼ μn -, i = 1 . . . b}Compute Rewards for i = 1 . . . b 9:u i θ = log π θ (yi,+|xi,+) π ref (yi,+|xi,+)10:v i θ = log π θ (yi,-|xi,-) π ref (yi,-|xi,-)11:Sort Rewards 12:if hard sort then 13:(u (1) . . . u (b) ) = Sort(u i θ )14:(v (1) . . . v (b) ) = Sort(v i θ )15:else if soft sort then 16:(u (1) . . . u (b) ) = SoftSort(u i θ , ε) 17: (v (1) . . . v (b) ) = SoftSort(v i θ , ε) 18:end ifCompute AOT logistic loss 20:ℓ θ = -1 b b i=1 log σ(β(u (i) θ -v (i) θ )) 21:Update θ {(x i , y i,+ , y i,-) ∼ μn , i = 1 . . . b}Compute Margins for i = 1 . . . b 7:u i θ = log π θ (yi,+|xi) π θ (yi,-|xi)8:v i θ = log π ref (yi,+|xi) π ref (yi,-|xi) 9:Sort Margins 10:if hard sort then 11:(u (1) . . . u (b) ) = Sort(u i θ )12:(v (1) . . . v (b) ) = Sort(v i θ )13:else if soft sort then 14:(u (1) . . . u (b) ) = SoftSort(u i θ , ε) 15: (v (1) . . . v (b) ) = SoftSort(v i θ , ε) 16: end if 17:Compute AOT logistic loss 18:ℓ θ = -1 b b i=1 log σ(β(u (i) θ -v (i) θ ))19:Update θ Proof of Theorem 2. Letε(θ, µ + , µ -) = sup φ∈Fc φ(r • π θ )dµ + -φ c (r • π θ )dµ - where for z, z ′ ∈ R : φ c (z) = inf z ′ ∈[-M,M ] h(z -z ′ ) -φ(z)and (a function f is cconcave if there exits g such that f = g c ).The population problem : minEQUATIONGiven samples μn+ = 1 n n i=1 δ (x i + ,y i + ) and μm -= 1 m m i=1 δ (x i -,y i -), the empirical problem is :ε(θ, μn + , μm -) = sup φ∈Fc φ(r • π θ )dμ n + -φ c (r • π θ )dμ m -and the ERM problem is : minπ θ ∈H ε(θ, μn + , μm -)Let θm,n be the minimizer of the ERM we have for any θ, by the definition of the minimizer:ε( θm,n , μn + , μm -) ≤ sup φ∈Fc φ(r • π θ )dμ n + -φ c (r • π θ )dμ m - ≤ sup φ∈Fc φ(r • π θ )dµ + -φ c (r • π θ )dµ - + sup φ∈Fc φ(r • π θ )d(μ n + -µ + ) + sup φ∈Fc φ c (r • π θ )d(µ --μm -)Let θ * be the minimizer of ( 14) for θ = θ * in the above inequality, and taking expectations on the randomness of the samples we obtainEε( θm,n , μn + , μm -) ≤ Eε(θ * , µ + .µ -) + E sup φ∈Fc φ(r • π θ * )d(μ n + -µ + ) + E sup φ∈Fc φ c (r • π θ * )d(µ --μm -)On the other hand by symmetrization we have:EQUATIONwhereR n (F) = E sup φ∈F 1 N N i=1 σ i φ(X i ) , σ i are independent rademacher random variables and X i ∼ (r • π θ * ) ♯ µ + iid (X i ∈ R). and similarly we have:E sup φ∈Fc φ c (r • π θ * )d(µ --μm -) ≤ E sup φ c ∈F c c φ c (r • π θ * )d(µ --μm -) ≤ 2R m (F c c )We have finally:EQUATIONTurning now to :ε( θm,n , µ + , µ -) = sup φ∈Fc φ(r • π θm,n )dµ + -φ c (r • π θm,n )dµ - ≤ ε( θm,n , μn + , μm -) + sup φ∈Fc φ(r • π θm,n )d(µ + -µ n + ) + sup φ∈Fc φ c (r • π θm,n )d(µ m --µ -) ≤ ε( θm,n , μn + , μm -) + sup π θ ∈H sup φ∈Fc φ(r • π θ )d(µ + -µ n + ) + sup π θ ∈H sup φ∈Fc φ c (r • π θ )d(µ m --µ -)Taking expectations we obtain:EQUATIONProof of Corollary 1. Define the uniform metric entropy of a class of real valued functions F on a set X as the logarithm of the covering number with respect to the uniform norm ∥.∥ ∞ , for ε > 0, this is defined as follows:N (ε, F, ∥.∥ ∞ ) := inf n ∈ N there exists f 1 . . . f n : X → R with sup f ∈F min 1≤i≤n ∥f -f i ∥ ∞ ≤ εAs observed in [Hundrieser et al., 2022] the c-transformation with bounded cost is a lipchitz operation under the uniform norm and since f cc = f we have by Lemma 2.1 in Munk :EQUATIONNow turning to the Rademacher complexity of a class F, it is dominated by Dudley's entropy integral (Theorem 16 in Luxburg and Bousquet ):EQUATIONNote that the cost we are using is c(z, z′ ) = h(z -z ′ ), for z, z ∈ [-M, M ].The domain on which the cost being a closed interval is convex and compact. By lipchitzity of h (Assumption 1) and denoting L its lipchitz constant, c(,z ′ ) is lipchitz for all z ′ ∈ [-M, M ]. Equivalently c(, z ′ ) is (α, Λ)Hölder smooth , for α = 1 and Λ = L, and hence our setup falls under the Assumptions of Theorem 3.11 in [Hundrieser et al., 2022] for Holder smooth costs defined on convex and compact sets and we have:log N (ε, F c , ∥.∥ ∞ ) ≲ ε -d α .Hence in our case α = 1 and d = 1 this leads to logN (ε, F c , ∥.∥ ∞ ) ≲ ε -1Replacing this in Equation ( 19) we obtain: R n (F c ) ≲ n -1 2 and by Equation ( 18) it follows that: R m (F c c ) ≲ m -1 2 . Turning now to the Rademacher Complexity of the composition of the c-concave potentials F c with r • H (the composition of a fixed reward function with the hypothesis class H ). Note since the cost c(., z ′ ) is L Lipchitiz for all z ′ ∈ [-M, M ] we have F c is included in the set of L lipchitz function that are bounded by R (See [Hundrieser et al., 2022] Lemma A.2 ). For φ ∈ F c and π θ ∈ H , let us note h φ,π θ = φ(r • π θ ) we have:h φ,π θ -h φ ′ ,π θ ′ ∞ = sup x∈X ,y∈Y |φ(r • π θ (y|x)) -φ ′ (r • π θ ′ (y|x))|We have:|φ(r • π θ (x)) -φ ′ (r • π θ ′ )| = |φ(r • π θ (x)) -φ(r • π θ ′ ) + φ(r • π θ ′ ) -φ ′ (r • π θ ′ )| ≤ L ∥r • π θ -r • π θ ′ ∥ ∞ + ∥φ -φ ′ ∥ ∞ .where we used lipchitzity of φ ∈ F c and hence we have:h φ,π θ -h φ ′ ,π θ ′ ∞ ≤ L ∥r • π θ -r • π θ ′ ∥ ∞ + ∥φ -φ ′ ∥ ∞ .We have therefore the following bound on the covering number of the composition:N (ε, F c • r • H, ∥.∥ ∞ ) ≤ N ε 2 , F c , ∥.∥ ∞ N ε 2L , r • H, ∥.∥ ∞ ,Plugging this in Equation ( 19) we obtain:R n (F c • r • H) ≤ inf δ∈[0,R] 2δ + √ 32 1 √ n R δ/4 log N (ε/2, F c , ∥.∥ ∞ ) + log N ε 2L , r • H, ∥.∥ ∞ dε Note that for a, b > 0 we have √ a + b ≤ √ a + √ b, hence we have: R n (F c • r • H) ≤ inf δ∈[0,R] 2δ + √ 32 1 √ n R δ/4 log N (ε/2, F c , ∥.∥ ∞ ) + log N ε 2L , r • H, ∥.∥ ∞ dεWe know by lipchitizty of the cost and being in one dimension that :log N (ε, F c , ∥.∥ ∞ ) ≲ ε -1By lipchitizity of r • π θ and using Assumption 3 we have therefore:log N ε 2L , r • H, ∥.∥ ∞ ≤ log N ε 2LL ′ , B 2 (r 0 , d θ ), ∥.∥ ∞ ≤ d θ log 2r 0 L ′ L εWe have therefore:inf δ∈[0,R] 2δ + 4 √ 2 √ n R δ/4 K 1 ε 2 -1/2 dε + 4 √ 2 √ n R δ/4 d θ log 2r 0 LL ′ ε dε ≲ n -1 2 .(For δ = 0, the upper bound is obtained.)For 2) By assumption 4, there exists π θ * , such that we have for h= (-x) 2 + : OT h ((r • π θ * ) ♯ µ + , (r • π θ * ) ♯ µ -) = 0.Similarly following the relaxation approach described in Section 3 and using the dual representation of the OT problem we can relax the paired stochastic dominance constraint problem given in (FSD paired) to:EQUATION)where µ is the paired measure representing (X, Y + , Y -). let π θ * be the minimizer of (20). Considering Problem (20), with empirical samples μn = 1 n n i=1 δ (xi,yi,+,yi,-) , denote π θn its minimizer we have :Theorem 3 (Sample Complexity of Dominance Violation for AOT Paired). The following sample complexity bound for the violation of stochastic dominance in AOT paired holds:E OT h (r • π θn ) ♯ µ, (r • π ref ) ♯ µ ≤ OT h ((r • π θ * ) ♯ µ, (r • π ref ) ♯ µ) Optimal Almost FSD Violation + 2R n (F c ; (r • π θ * ) ♯ µ + ) + 4R n (F c c ; (r • π ref ) ♯ µ)One dimensional OT sample complexity with optimal θ *+ 2R n (F c • r • H; µ)) Complexity of learning in H via the 1D OT problem , where R n (F; ν) = E sup φ∈F 1 n n i=1 σ i φ(Z i )is the Rademacher Complexity and for i = 1 . . . n, σ i are independent rademacher random variables and Z i ∼ ν iid.Similarly under Assumptions 1, 2 and 3 we have:E OT h (r • π θn ) ♯ µ, (r • π ref ) ♯ µ -OT h ((r • π θ * ) ♯ µ, (r • π ref ) ♯ µ) ≲ n -1 2 .Proof of Theorem 3. The proof can be simply obtained by inspecting the proof of Theorem 1, and we omit it.Remark 2. Although OT h is one dimensional, an entropic regularization of OT h has computational advantages as discussed in Section 3. Results from [Groppe and Hundrieser, 2023 ] can be leveraged to obtain sample complexity bounds and we obtain under our assumptions also a parametric rate of n -1 2 . The main insight in [Groppe and Hundrieser, 2023] is in introducing for an entropic regularization parameter ε > 0, the smoothed (c, ε, µ) transform and replacing the spaces F c by F c,ε . In the 1D case, up to constants, these spaces have the same covering numbers.Recall we proposed the following three classes of loss functions h in the main text:1. Area of violation ("classification") Setting h(x) to be the 0-1 loss (1 x<0 ) measures the fraction of the interval [0, 1] where a violation occurs, paralleling classification losses which count the number of misclassification.Tables 2, 3, 4, 5 and 6 are ablations on the reference base model used (Merlinite-7B, OpenHermes-2.5-Mistral-7B, Sarling-LM-7B-alpha, Meta-LLama-3-8B-Instruct, Mistral-7B-Instruct-v0.2). In these tables, we report only the AlpacaEval using Llama3-70B as a judge to reduce the costs of evaluations. Note that we observed that while the absolute scoring of Llama3-70B is different than GPT4, as it can be seen in Table 2 , it preserves the rankings of the models. We see across all these models a better performance of the distributional alignment AOT on AlpacaEval and a competitive performance on other benchmarks. NeurIPS Paper ChecklistIts Jacobian is a permutation matrix at every point except a measure-zero set where it is not differentiable.In our alignment setting, this is should not be an issue as our scores are real-valued. 


Decision: The paper introduces Alignment via Optimal Transport (AOT), a new method for aligning large language models (LLMs) based on distributional preferences rather than pairwise human preferences. AOT ensures that the reward distribution of positive samples stochastically dominates that of negative samples by solving an optimal transport problem with a smooth, convex cost, leading to a closed-form solution. By fine-tuning LLMs with the AOT objective, the method achieves superior alignment, demonstrated through state-of-the-art performance on various alignment datasets and benchmarks, including Open LLM Benchmarks and AlpacaEval. All reviewer recognized that the paper has some interesting theoretical contributions which connects optimal transport to the alignment problem. However, the numerical evaluation of this work seems to be less adequate, which the authors can improve upon.

Reviews: review: paper_summary: paper_summary: This works proposed using Optimal transport in 1D to derive a better alignment guided by the preference data. The main idea is to work with the log likelihood  ratio  of  the  marginal distributions of the preference data. The alignment is made through Optimal Transport loss in 1D based on the concept of stochastic dominance between quantiles of two distributions.
strength_weakness: strength_weakness: My main concern of this paper is the empirical results in experiment section. Table 1 shows that AOT paired/unpaired do not outperform other methods at least in 4 out 7 cases (ARC, MMLU, Winogrande, GSM8K). When they are versus each other, there is no clear winner between AOT paired and AOT unpaired. Meanwhile, I believe that with the AOT paired, when we have more information, the task must be easier, please correct me if I am wrong. I have similar concerns for their performances in Figure 2.
questions: questions: In Figure 3, it appears that small $\beta$ will produce better results for all method except the IPO. Have the authors tried to test with smaller value of $\beta$, i.e. $\beta = 0.005$?  Is there any explanation for this trend in that figure?
limitations: limitations: It is fine.
paper_summary: paper_summary: The motivation of this paper is that current alignment approaches ensure reward dominance at the sample-level but not on the distributional level. With this in mind, the authors set their goal to design an alignment approach which satisfies First Order Stochastic Dominance (FSD) in rewards for positive examples compared to negative examples. With this goal in mind, they make the following contributions:
1. They write this as a constraint satisfaction problem, which can be relaxed into an unconstrained optimization problem using different surrogates of the 0/1 loss.
2. Using Santambrogio (2015)'s result from optimal transport, they develop a computationally efficient algorithm called ATO.
3. They provide an upper bound for the violation in stochastic dominance of this algorithm using a Rademacher and symmetrization argument.
4. They provide experiments that show that ATO is competitive with existing approaches like DPO, KTO.
strength_weakness: strength_weakness: 1. There could be some more motivation for why the FSD condition is practically desirable over the conditional (on x) dominance condition of DPO from Equation (2). Is this more than a theoretical nicety?

2. There is not a clear improvement in the empirical results over prior approaches, and each approach seems to enjoy success on certain benchmarks -- with the differences between performance quite small in most cases.
This invites the question of in what settings ATO would perform better than other approaches in the literature. 

3. The theoretical results do not shed much light on when ATO is likely to perform well. Uniform Convergence results are well-known in recent years to not explain generalization behavior of modern ML models. So, while nice, the insight from Theorem 2 is not super obvious to me.

Minor:
1. Please call the Santambrogio (2015) result something other than a theorem since it is not a contribution of the present work.
questions: questions: Since KTO also works with unpaired preference data, a more thorough comparison with KTO is warranted:
1. Is there a unique solution that satisfies the FSD condition? If there are multiple solutions, how does ATO break ties? How does KTO break ties? 
2. Is the reason that the behavior is different from KTO that they have different inductive biases and hence converge to different policies? Or the variances of the estimators are different, even though they converge to the same?
3. Is there a way to compare Theorem 2 with the statistical properties of KTO/DPO/RLHF? Or any reason to believe that ATO has a statistical advantage?

Despite the weaknesses, I think the connection to optimal transport is interesting and the experiments are sound and comprehensive. I would be happy to adjust my score accordingly if the authors alleviate my concerns through discussion.
limitations: limitations: Yes
paper_summary: paper_summary: The paper proposes a new technique for preference alignment, named Alignment via Optimal Transport (AOT). The proposed technique supports both paired and unpaired alignment settings.  The paper introduces a new viewpoint for preference alignment based on stochastic dominance i.e., making the reward distribution of the positive samples stochastically dominant in the first order on the distribution of negative samples. From this perspective, the paper explains DPO as a special case i.e., a pointwise preference approach with relaxation through the logistic loss. In addition, the paper shows that using convex relaxation is equivalent to minimizing a one-dimensional optimal transport problem. To enhance the differentiability of the objective through the sorting operator (for approximating the continuous optimal transport), the paper uses the SInkhorn-Knopp algorithm instead of the conventional sorting. On the experimental side, AOT leads to state-of-the-art models in the 7B family of models when evaluated with Open LLM Benchmarks and AlpacaEval (using Llama3-70B-Instruct instead of GPT4).
strength_weakness: strength_weakness: The evaluation is conducted using Llama3-70B-Instruct instead of GPT4 which is not the standard. However, as described in the paper, the usage of  Llama3-70B-Instruct  leads to approximately the same.
questions: questions: 1. It seems that hard sorting is also comparable to soft sorting. Using normal sorting could help to improve the computational speed compared to soft sorting. Should the paper recommend hard sorting as the default variant?

2. How many interactions are used for the Sinkhorn algorithm? What is the choice of the entropic regularization hyperparameter? Do these hyperparameters significantly affect the results?

3. Why does the unpaired setting lead to better results than the paired setting? Do the authors have any explanation?
limitations: limitations: The authors adequately addressed the limitations.

Comments: title: Reply to the rebuttal
Rebuttal response
Discussion with the authors
Response to authors
comment: I would like to thank the authors for their answers. I would like to keep my score unchanged based on the fact is that  the empirical results on other metrics are only competitive, although the theory appears to be sufficient.
I thank the authors for clarifying my questions, and accordingly increase my score to 6.
Dear reviewers:

As the discussion period is going to end soon, please try to actively engage with the authors about the paper. Thanks a lot for your help and dedication. 

You AC.
Thank you for your response,

I will keep my score since I believe this paper proposes an interesting framework for LLM alignment i.e., stochastic dominance and optimal transport.

Best,