Title: Distributional Preference Alignment of LLMs via Optimal Transport
Venue: NeurIPS 2024
Accepted: yes
Keywords: LLM Alignment

Optimal Transport

stochastic dominance
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}.

Distributional Preference Alignment of LLMs via Optimal Transport

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 AOT is available in the Hugging Face TRL library https://ibm.biz/AOT_TRL .

1 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

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

-

1000

-

800

-

600

-

400

-

200

0

Quantile Comparison of Rewards of Different Alignment Strategies

AOT paired Chosen Rewards

AOT paired Rejected Rewards

AOT unpaired Rejected Rewards

KTO Chosen Rewards

KTO Rejected Rewards

DPO Chosen Rewards

DPO Rejected Rewards

Percentile

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

-

10

3

-

10

2

-

10

1

-

10

0

0

10

0

10

1

10

2

10

3

Quantile Comparison of Margins of Different Alignment Strategies

AOT paired

AOT unpaired

KTO

DPO

Base (Merlinite)

Percentile

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.

2 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 :

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 :

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 n ∑ n i =1 δ ( x i ,y i ) , we have that r ♯ µ = 1 n ∑ n i =1 δ r ( x i ,y i ) .

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 :

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 :

and since the difference is taken for the same x , the normalization Z ( x ) disappears resulting in:

We can interpret this as a pointwise constraint inducing preference for positive over negative reward outcomes as follows:

DPO 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].

2.1 Distributional Preference via Stochastic Dominance

Our main insight from looking at the pointwise constraint in Eq. (2) is that we can recast it as a distri- butional constraint in terms of stochastic dominance of the random variable Z + θ = 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.

Distributional Unpaired 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:

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:

Distributional Paired Preference Note that we can rewrite Eq. (2) in the equivalent form:

In 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 variable Z θ = log π θ ( Y + | X ) π θ ( Y -| X ) indexed by the policy we are optimizing on the random variable Z ref = log π ref ( Y + | X ) π ref ( Y -| 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 .

We 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:

Noting r p ◦ π θ ( x, y + , y -) = log π θ ( y + | x ) π θ ( y -| x ) this can be written as follows:

3 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:

and

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:

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:

We can relax this problem to the following minimization problem:

where 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 ✶ is the indicator function):

This 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 ( ✶ 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 :

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 ) :

This 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 δ u i ˆ µ V = 1 n ∑ n i =1 δ v i . 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:

Back to (9), given empirical samples ˆ µ U θ = 1 n ∑ n i =1 δ u i θ and ˆ µ V θ = 1 n ∑ n i =1 δ v i θ , let u ( i ) θ , v ( i ) θ be the order statistics as function of θ . We have therefore:

In 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 + = 1 n ∑ n i =1 δ ( x i, + ,y i, + ) and ˆ µ n -= 1 n ∑ n i =1 δ ( x i, -,y i, -) . 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

AOT for Paired Preference Let ˆ µ n = 1 n ∑ n i =1 δ ( x i ,y i, + ,y i, -) 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:

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.

4 Statistical Analysis

In this section, we focus on the statistical analysis of unpairedAOT and defer pairedAOT 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:

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 :

In our setting, a function is called c -concave if there exists ψ : [ -M,M ] → R such that φ = ψ c . Define:

By duality (Theorem 5.10 in [Villani, 2009]) we have:

Replacing the dual expression of OT h in ( uAOT h ), we see that ( uAOT h ) can be cast as a min-max problem:

Given samples ˆ µ n + = 1 n ∑ n i =1 δ ( x i, + ,y i, + ) and ˆ µ n -= 1 n ∑ n i =1 δ ( x i, -,y i, -) , the empirical problem is:

Recall 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:

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:

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.

If in addition Assumption 4 holds and h ( t ) = ( -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.

5 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.

LLM Alignment Alternatives We compared AOT with current state-of-the-art alignment approaches, specifically Direct Preference Optimization (DPO) [Rafailov et al., 2024], KahnemanTversky 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], Starling [Zhu et al., 2023], Mistral-7B Jiang et al. [2023], and Llama3-8B [AI@Meta, 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.

Metrics 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-Instruct

0.0

0.1

0.2

0.3

0.4

0.5

Beta

27.5

30.0

32.5

35.0

37.5

40.0

42.5

45.0

47.5

AlpacaEval (Llama3-70B)

Effect of Beta on Alignment Performance

AOT paired

AOT unpaired

KTO

DPO

IPO

Merlinite-7B

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.

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 .

6 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. On other benchmarks such as the open LLM leaderboard AOT leads to competitive results.

References

AI@Meta. Llama 3 model card. 2024. URL https://github.com/meta-llama/llama3/blob/ main/MODEL_CARD.md .

Anil, J. Lucas, and R. Grosse. Sorting out lipschitz function approximation. In International Conference on Machine Learning , pages 291-301. PMLR, 2019.

G. Azar, Z. D. Guo, B. Piot, R. Munos, M. Rowland, M. Valko, and D. Calandriello. A general theoretical paradigm to understand learning from human preferences. In International Conference on Artificial Intelligence and Statistics , pages 4447-4455. PMLR, 2024.

Bai, A. Jones, K. Ndousse, A. Askell, A. Chen, N. DasSarma, D. Drain, S. Fort, D. Ganguli, T. Henighan, et al. Training a helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint arXiv:2204.05862 , 2022.

L. Bartlett, M. I. Jordan, and J. D. McAuliffe. Convexity, classification, and risk bounds. Journal of the American Statistical Association , 101:138 - 156, 2006. URL https://api. semanticscholar.org/CorpusID:2833811 .

Beeching, C. Fourrier, N. Habib, S. Han, N. Lambert, N. Rajani, O. Sanseviero, L. Tunstall, and T. Wolf. Open LLM Leaderboard. https://huggingface.co/spaces/HuggingFaceH4/ open_llm_leaderboard , 2023.

Blondel, O. Teboul, Q. Berthet, and J. Djolonga. Fast differentiable sorting and ranking. In International Conference on Machine Learning , pages 950-959. PMLR, 2020.

F. Christiano, J. Leike, T. Brown, M. Martic, S. Legg, and D. Amodei. Deep reinforcement learning from human preferences. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 30. Curran Associates, Inc., 2017. URL https://proceedings.neurips.cc/paper_ files/paper/2017/file/d5e2c0adad503c91f91df240d0cd4e49-Paper.pdf .

Cuturi, O. Teboul, and J.-P. Vert. Differentiable ranking and sorting using optimal transport. Advances in neural information processing systems , 32, 2019.

Del Barrio, J. A. Cuesta-Albertos, and C. Matrán. An optimal transportation approach for assessing almost stochastic order. The Mathematics of the Uncertain: A Tribute to Pedro Gil , pages 33-44, 2018.

Domingo-Enrich, Y. Schiff, and Y. Mroueh. Learning with stochastic orders. In The Eleventh International Conference on Learning Representations , 2022.

Dubois, B. Galambosi, P. Liang, and T. B. Hashimoto. Length-controlled alpacaeval: A simple way to debias automatic evaluators. arXiv preprint arXiv:2404.04475 , 2024.

Ethayarajh, W. Xu, D. Jurafsky, and D. Kiela. Human-centered loss functions (halos). Technical report, Contextual AI, 2023.

Ethayarajh, W. Xu, N. Muennighoff, D. Jurafsky, and D. Kiela. Kto: Model alignment as prospect theoretic optimization. arXiv preprint arXiv:2402.01306 , 2024.

Groppe and S. Hundrieser. Lower complexity adaptation for empirical entropic optimal transport. arXiv preprint arXiv:2306.13580 , 2023.

J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, and W. Chen. LoRa: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.09685 , 2021.

Hundrieser, T. Staudt, and A. Munk. Empirical optimal transport between different measures adapts to lower complexity. arXiv preprint arXiv:2202.10434 , 2022.

Ji, M. Liu, J. Dai, X. Pan, C. Zhang, C. Bian, C. Zhang, R. Sun, Y. Wang, and Y. Yang. Beavertails: Towards improved safety alignment of llm via a human-preference dataset. arXiv preprint arXiv:2307.04657 , 2023.

Jiang, A. Sablayrolles, A. Mensch, C. Bamford, D. Chaplot, D. de las Casas, F. Bressand, G. Lengyel, G. Lample, L. Saulnier, et al. Mistral 7b (2023). arXiv preprint arXiv:2310.06825 , 2023.

Maia Polo, L. Weber, L. Choshen, Y. Sun, G. Xu, and M. Yurochkin. tinyBenchmarks: evaluating LLMs with fewer examples. arXiv preprint arXiv:2402.14992 , 2024.

Manole and J. Niles-Weed. Sharp convergence rates for empirical optimal transport with smooth costs. The Annals of Applied Probability , 34(1B):1108-1135, 2024.

Ogryczak and A. Ruszczynski. Dual stochastic dominance and related mean-risk models. SIAM Journal on Optimization , 13(1):60-78, 2002.

Ouyang, J. Wu, X. Jiang, D. Almeida, C. Wainwright, P. Mishkin, C. Zhang, S. Agarwal, K. Slama, A. Ray, et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems , 35:27730-27744, 2022.

Peyré and M. Cuturi. Computational optimal transport. Foundations and Trends in Machine Learning , 11(5-6):355-607, 2019. ISSN 1935-8237. doi: 10.1561/2200000073.

Rafailov, A. Sharma, E. Mitchell, C. D. Manning, S. Ermon, and C. Finn. Direct preference optimization: Your language model is secretly a reward model. Advances in Neural Information Processing Systems , 36, 2024.

Rosasco, E. De Vito, A. Caponnetto, M. Piana, and A. Verri. Are loss functions all the same? Neural Comput. , 16(5):1063-1076, may 2004. ISSN 0899-7667.

Santambrogio. Optimal Transport for Applied Mathematicians: Calculus of Variations, PDEs, and Modeling . Birkhäuser, Cham, 2015. ISBN 9783319208275. doi: 10.1007/978-3-319-20828-2.

Stiennon, L. Ouyang, J. Wu, D. Ziegler, R. Lowe, C. Voss, A. Radford, D. Amodei, and P. F. Christiano. Learning to summarize with human feedback. Advances in Neural Information Processing Systems , 33:3008-3021, 2020.

Sudalairaj, A. Bhandwaldar, A. Pareja, K. Xu, D. D. Cox, and A. Srivastava. Lab: Large-scale alignment for chatbots. arXiv preprint arXiv:2403.01081 , 2024.

Teknium. Openhermes-2.5-mistral-7b. https://huggingface.co/teknium/OpenHermes-2. 5-Mistral-7B , 2024.

Tunstall, E. Beeching, N. Lambert, N. Rajani, S. Huang, K. Rasul, A. M. Rush, and T. Wolf. The alignment handbook. https://github.com/huggingface/alignment-handbook , 2023a.

Tunstall, E. Beeching, N. Lambert, N. Rajani, K. Rasul, Y. Belkada, S. Huang, L. von Werra, C. Fourrier, N. Habib, et al. Zephyr: Direct distillation of LM alignment. arXiv preprint arXiv:2310.16944 , 2023b.

Villani. Optimal Transport: old and new , volume 338. Springer, 2009.

von Werra, Y. Belkada, L. Tunstall, E. Beeching, T. Thrush, N. Lambert, and S. Huang. Trl: Transformer reinforcement learning. https://github.com/huggingface/trl , 2020.

Wang, Y. Dong, J. Zeng, V. Adams, M. N. Sreedhar, D. Egert, O. Delalleau, J. P. Scowcroft, N. Kant, A. Swope, et al. Helpsteer: Multi-attribute helpfulness dataset for steerlm. arXiv preprint arXiv:2311.09528 , 2023.

Zhao, M. Khalman, R. Joshi, S. Narayan, M. Saleh, and P. J. Liu. Calibrating sequence likelihood improves conditional language generation. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=0qSOodKmJaN .

Zhu, E. Frick, T. Wu, H. Zhu, and J. Jiao. Starling-7b: Improving llm helpfulness and harmlessness with rlaif, November 2023.

A Broader Impact and Limitations

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.

B Algorithms and Pytorch Code In Hugging Face TRL

Algorithm 1 AOT Unpaired

C Proofs

Proof of Theorem 2. Let

where for z, z ′ ∈ R :

and (a function f is c -concave if there exits g such that f = g c ).

The population problem :

Given 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 :

and the ERM problem is :

Let ˆ θ m,n be the minimizer of the ERM we have for any θ , by the definition of the minimizer:

Let θ ∗ be the minimizer of (14) for θ = θ ∗ in the above inequality, and taking expectations on the randomness of the samples we obtain

On the other hand by symmetrization we have:

where R 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:

We have finally:

Turning now to :

Taking expectations we obtain:

where

D Bounding Rademacher Complexities

Proof 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:

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 :

Now turning to the Rademacher complexity of a class F , it is dominated by Dudley's entropy integral (Theorem 16 in Luxburg and Bousquet ):

Note 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:

Hence in our case α = 1 and d = 1 this leads to

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:

We have:

where we used lipchitzity of φ ∈ F c and hence we have:

We have therefore the following bound on the covering number of the composition:

Plugging this in Equation (19) we obtain:

Note that for a, b > 0 we have √ a + b ≤ √ a + √ b , hence we have:

We know by lipchitizty of the cost and being in one dimension that :

By lipchitizity of r ◦ π θ and using Assumption 3 we have therefore:

We have therefore:

(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 .

E AOT paired

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:

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 δ ( x i ,y i, + ,y i, -) , denote π ˆ θ n its minimizer we have :

Replacing 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 .

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:

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:

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.

F Choice of violation penalty function h

Recall we proposed the following three classes of loss functions h in the main text:

Area of violation ('classification') Setting h ( x ) to be the 0-1 loss ( ✶ x< 0 ) measures the fraction of the interval [0 , 1] where a violation occurs, paralleling classification losses which count the number of misclassification.

Wasserstein-1 violation 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 θ .

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:

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.

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.

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.

We 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 form

Note 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)

Let's consider the first term, the analysis for the second term is the same. Note that since h ′ is L -Lipschitz

Let

Note 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

We seek to bound this quantity as n →∞ . By Cauchy-Schwarz,

where we've assumed that ( 1 n ∑ n i =1 ∥∇ θ ( u ( i ) θ ) ∥ 2 1 ) ≤ M 2 . Observe that F -1 V θ ( F U θ ( · )) is the optimal transport plan from U θ to V θ and is monotonic nondecreasing, hence F -1 V θ ( F U θ ( u ( i ) θ )) are simply a new set of independently drawn order statistics of v θ , i.e. v ( i, 2) θ . We thus have

where F ( 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) , then 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,

where → 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.

H Additional Experiments

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.

Finally Figure 4 puts in context our best model Merlinite-7B-uAOT as the best 7B-family model on AlpacaEval leaderboard at the time of writing this paper. Finally, we give in Table 10 the variance of the evaluation across 4 different random seeds for training and evaluation each alignment strategy, we see very small variance in AOT , especially the unpaired variant.

Version:

AlpacaEval

AlpacaEval 2.0

Filter:

Community

Verified

Baseline: GPT-4 Preview (11/06) | Auto-annotator: GPT-4 Preview (11/06)

LCWinRate

58.3%

46.8%

57.5%

51.3%

55.0%

46.1%

Yi-Large Preview

51.9%

57.5%

Storm-7B (num_beams=10)

51.8%

55.4%

50.0%

50.0%

Storm-7B

48.9%

52.5%

ExPO+Llama-3-Instruct-8B-SimPO

45.8%

40.6%

Llama-3-Instruct-8B-SimPO

44.7%

40.5%

44.5%

56.7%

Qwen1.5 110B Chat

43.9%

33.8%

Aligner 2B+Claude 3 Opus

41.8%

34.5%

40.5%

29.1%

38.1%

23.6%

Aligner 2B+Qwen1.5 72B Chat

36.7%

31.8%

36.6%

26.5%

35.3%

22.1%

Ein 70B v0.1

35.0%

24.8%

34.9%

25.6%

34.8%

35.9%

Llama370BInstruct

34.4%

33.2%

Mistral Large (24/02)

32.7%

21.4%

ExPO+SPPO-Mistral7B-PairRM

31.8%

35.4%

merlinite-7B-AOT

31.7%

29.9%

Samba CoE v0.2 (best-of-16)

31.5%

27.0%

REBEL-Llama-3-8B-Instruct

31.4%

34.3%

Mixtral 8x22B v0.1

30.9%

22.2%

SPPO-Mistral7B-PairRM

30.5%

32.2%

30.2%

15.8%

Snorkel(Mistral-PairRM-DPO+best-of-16)

30.0%

34.9%

NeurIPS Paper Checklist

1. Claims

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.

2. Limitations

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.

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: The Appendix includes proofs of all stated Theorems.

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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a 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 experimental 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: All information needed to reproduce our results are clearly stated in the paper. We will also release code to reproduce the main results upon acceptance of the 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 be 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 a 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 is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to 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 to replicate the results, access to 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.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a 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 are welcome to 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., to 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We will release the code to reproduce our main experimental results under open source upon acceptance. During the review process the code is shared in supplementary material.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not 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 a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to 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 a 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) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: All details necessary details to reproduce our experimental results are provided in the paper.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We report error bars for the main results on AlpacaEval using LLaMA3-70B as a judge (GPT4-based evaluations would incur significant cost for this type of experiments). As indicated, error bars are computed as a variance across 4 runs.

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, confidence intervals, or statistical significance tests, at 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 to 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 or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 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 are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We mentioned the hardware resources utilized for our experiments in the relevant sections of the paper.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as 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 or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: Our research was conducted in full conformity of the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We discuss societal impacts of our work in the context of improving on current RLHF techniques.

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).

11. 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.

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

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 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.

According to 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.

15. 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 nor 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 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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).\n\nSTRENGTHS:\n* The paper is well-written and easy to follow.\n* The paper proposes an original preference alignment approach based on stochastic dominance and optimal transport. The connection to optimal transport is interesting and novel. In addition, the paper also provides theoretical results on the sample complexity of the empirical estimation of the objective.\n* The proposed approach can work for both unpaired and paired alignment settings.\n* Experiments are extensive on various datasets i.e., UltraFeedback for paired setting,  PKU BeaverTails, and HelpSteer for unpaired setting. The paper also provides a free LLM-judge for local Alpaca evaluations.\n* ATO achieves SOTA on AlpacaEval benchmark and competitive results on other metrics compared to DPO, KPO, and IPO.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\n1. 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?\n\n2. 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?\n\n3. Why does the unpaired setting lead to better results than the paired setting? Do the authors have any explanation?\n\nLIMITATIONS:\nThe authors adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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:\n1. 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.\n2. Using Santambrogio (2015)'s result from optimal transport, they develop a computationally efficient algorithm called ATO.\n3. They provide an upper bound for the violation in stochastic dominance of this algorithm using a Rademacher and symmetrization argument.\n4. They provide experiments that show that ATO is competitive with existing approaches like DPO, KTO.\n\nSTRENGTHS:\n**Originality**\n1. The connection to optimal transport is quite interesting.\n\n**Quality**\n\n2. The empirical results are competitive with other alignment approaches.\n3. The theoretical results are sound\n\n**Clarity**\n\n4. The paper is well-written with background appropriately introduced and clear explanations.\n\nWEAKNESSES:\n1. 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?\n\n2. 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.\nThis invites the question of in what settings ATO would perform better than other approaches in the literature. \n\n3. 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.\n\nMinor:\n1. Please call the Santambrogio (2015) result something other than a theorem since it is not a contribution of the present work.\n\nQUESTIONS:\nSince KTO also works with unpaired preference data, a more thorough comparison with KTO is warranted:\n1. 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? \n2. 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?\n3. 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?\n\nDespite 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.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

comments: 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.

STRENGTHS:
The framework of the paper is very easy to follow. 
The proposed problem and solution's direction are interesting. 
Settings, theoretical results are adequate to support the proposed method.

WEAKNESSES:
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:
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:
It is fine.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-13