IGOR MELNYK _[∗]_, YOUSSEF MROUEH _[∗]_, BRIAN BELGODERE _[∗]_, MATTIA RIGOTTI, APOORVA NITSURE,
MIKHAIL YUROCHKIN, KRISTJAN GREENEWALD, JIRI NAVRATIL, AND JARRET ROSS
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.
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 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.
(I. Melnyk,Y. Mroueh, B. Belgodere, M. Rigotti, A. Nitsure, M. Yurochkin, K. Greenewald, J. Navratil, J. Ross)
IBM Research & MIT-IBM Watson Lab
  - Core Contribution.
Quantile Comparison of Margins of Diferent Alignment Strategies
|Col1|AOT paire
AOT unpa|d
ired|Col4|Col5|Col6|Col7|Col8|Col9|
|---|---|---|---|---|---|---|---|---|
||~~KTO~~
DPO
||||||||
||~~Base (Mer~~|~~ inite)~~|~~ inite)~~||||||
||||||||||
||||||||||
||||||||||
||||||||||
||||||||||
||||||||||
0 _._ 0 0 _._ 2 0 _._ 4 0 _._ 6 0 _._ 8 1 _._ 0
Percentile
_−_ 200
_−_ 400
_−_ 600
_−_ 800
_−_ 1000
Quantile Comparison of Rewards of Diferent Alignment Strategies
|Col1|Col2|Col3|Col4|Col5|Col6|Col7|Col8|
|---|---|---|---|---|---|---|---|
|||||||||
|||||||||
|||||||||
||||||AOT paired Ch
AOT paired Re|osen Rewards
 ected Rewards||
||||||
AOT unpaired
AOT unpaired
|
 Chosen Reward
 Rejected Rewar
|s
  ds|
||||||~~KTO Chosen R~~
KTO Rejected
DPO Chosen R
|~~ ewards~~
 Rewards
 ewards
||
||||||~~DPO Rejected~~|~~ Rewards~~||
0 _._ 0 0 _._ 2 0 _._ 4 0 _._ 6 0 _._ 8 1 _._ 0
Percentile
10 [3]
10 [2]
10 [1]
10 [0]
_−_ 10 [0]
_−_ 10 [1]
_−_ 10 [2]
_−_ 10 [3]
(a) Stochastic Dominance of Reward of Chosen on
Rejected: AOT achieves a larger margin between
the quantile plots of chosen and rejected rewards.
(b) Stochastic Dominance of AOT’s optimized
policy margin (between Chosen on Rejected) on
the margin of the reference policy.
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.
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 among the highest 7B model at the time of writing this paper, without
using any iterative training process as in [Yuan et al., 2024].
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-to-rejected 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
DISTRIBUTIONAL PREFERENCE ALIGNMENT OF LLMS VIA OPTIMAL TRANSPORT 3
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 _FZ_ [(] _[−]_ [1)] : [0 _,_ 1] R
_→_
the left-continuous inverse of the Cumulative Distribution Function (CDF) _FZ_ :
_QZ_ ( _p_ ) = _FZ_ [(] _[−]_ [1)] ( _p_ ) = inf _η_ : _FZ_ ( _η_ ) _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 ≽ _Z_ 2 _QZ_ 1( _p_ ) _QZ_ 2( _p_ ) _,_ _p_ [0 _,_ 1] _._ (1)
FSD _⇐⇒_ _≥_ _∀_ _∈_
Let _X_ be the space of prompts _X_ and _Y_ be the space of responses _Y_ from an LLM conditioned
on a prompt _X_ . The reference LLM is represented as policy _π_ ref ( _Y_ _X_ ), i.e., as a conditional
_∈X_ _|_
probability onis a parameter belonging to a bounded parameter spacea mapping _r_ : _Y_ given a promptR, we note as _X ∈X r_ . We note the LLM policy we are optimizing by _♯µ_ the pushforward map of Θ _⊂_ R _[d][θ]_ . For a measure _µ_ through _r µ_ . In particular, for _∈P_ ( _X × Y πθ_ where) and _θ_
empirical measures _X × Y → µ_ = [1] - _n_ _[δ]_ [(] _[x]_ _[,y]_ [)][, we have that] _[ r][♯][µ]_ [ =] [1] - _n_ _[δ][r]_ [(] _[x]_ _[,y]_ [)][.]
empirical measures _µ_ = _n_ [1] - _ni_ =1 _[δ]_ [(] _[x]_ _i_ _[,y]_ _i_ [)][, we have that] _[ r][♯][µ]_ [ =] _n_ [1] - _ni_ =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 :
[1] - _n_ [1]
_n_ _i_ =1 _[δ]_ [(] _[x]_ _i_ _[,y]_ _i_ [)][, we have that] _[ r][♯][µ]_ [ =] _n_
_rθ_ ( _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:
       -       -       - _πθ_ ( _y_ + _x_ )
min E( _x,y_ + _,y−_ ) _∼µ_ log _σ_ _β_ log _|_
_θ_ _−_ _π_ ref ( _y_ + _x_
_π_ ref ( _y_ + _x_ )
_|_
- _β_ log - _πθ_ ( _y−|x_ )
_−_ _π_ ref ( _y−_ _x_
_π_ ref ( _y−_ _x_ )
_|_
���
_._
We can interpret this as a pointwise constraint inducing preference for positive over negative reward
outcomes as follows:
log - _πθ_ ( _y_ + _|x_ )
_,_ ( _x, y_ + _, y−_ ) _µ._ (2)
_∀_ _∼_
_π_ ref ( _y_ + _x_ )
_|_
- log - _πθ_ ( _y−|x_ )
_≥_ _π_ ref ( _y−_ _x_
_π_ ref ( _y−_ _x_ )
_|_
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 distributional constraint in
terms of stochastic dominance of the random variable _Zθ_ [+] [=][ log][(] _π_ _[π]_ ref _[θ]_ [(] ( _[Y]_ _Y_ [+] + _[|][X]_ _|X_ [)] ) [)][ of positive outcomes]
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 distributional constraint in
terms of stochastic dominance of the random variable _Zθ_ [+] [=][ log][(] _π_ _[π]_ ref _[θ]_ [(] ( _[Y]_ _Y_ [+] + _[|][X]_ _|X_ [)] ) [)][ of positive outcomes]
on the random variable _Zθ_ _[−]_ [=][ log][(] _π_ _[π]_ ref _[θ]_ [(] ( _[Y]_ _Y_ _[−]_ _−_ _[|][X]_ _|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.
on the random variable _Zθ_ _[−]_ [=][ log][(] _π_ _[π]_ ref _[θ]_ [(] ( _[Y]_ _Y_ _[−]_ _−_ _[|][X]_ _|X_ [)]
**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
_µ_ + ( ), 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
_∈P_ _X × Y_
( _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]_ [+][)]
_[π][θ]_ [(] _[Y]_ [+] _[|][X]_ [+][)] ≽ log _[π][θ]_ [(] _[Y][−][|][X][−]_ [)]
_π_ ref ( _Y_ + _X_ +) _FSD_ _π_ ref ( _Y−_ _X−_
_|_ _|_
_In other words, noting ru_ _πθ_ ( _x, y_ ) = log _π_ _[π]_ ref _[θ]_ [(] ( _[y]_ _y_ _[|][x]_ _|x_ [)]
_◦_
_π_ ref ( _Y−_ _X−_ ) _[.]_
_|_
_In other words, noting ru_ _πθ_ ( _x, y_ ) = log _π_ ref _[θ]_ ( _y|x_ ) _[, the distributional preference in the unpaired setting]_
_◦_
_means that we have the following constraint:_
( _ru_ _πθ_ ) _♯µ_ + ≽ ( _ru_ _πθ_ ) _♯µ−._ (3)
_◦_ _FSD_ _◦_
**Distributional Paired Preference** Note that we can rewrite Eq. (2) in the equivalent form:
log _[π][θ]_ [(] _[y]_ [+] _[|][x]_ [)]
_[π][θ]_ [(] _[y]_ [+] _[|][x]_ [)] _[π]_ [ref] [(] _[y]_ [+] _[|][x]_ [)]
_πθ_ ( _y−_ _x_ ) _π_ ref ( _y−_ _x_ )
_|_ _[≥]_ [log] _|_
( _x, y_ + _, y−_ ) _µ._ (4)
_π_ ref ( _y−_ _x_ ) _[,]_ _∀_ _∼_
_|_
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_ _[Y]_ _−_ [+] _[|]_ _|_ _[X]_ _X_ [)] )
_∼_
_πθ_ ( _Y−|X_ )
indexed by the policy we are optimizing on the random variable _Z_ ref = log _π_ _[π]_ [ref] ( [(] _Y_ _[Y]_ [+] _[|]_ _|_ _[X]_ _X_ [)] ) [indexed by the]
indexed by the policy we are optimizing on the random variable _Z_ ref = log _π_ ref [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 πθ dis-_
_tributionally dominates π_ ref _in terms of_ log _probability ratio of positive and negative responses_
_if:_
_[π][θ]_ [(] _[Y]_ [+] _[|][X]_ [)] ≽ log _[π]_ [ref] [(] _[Y]_ [+] _[|][X]_ [)]
_πθ_ ( _Y−_ _X_ ) _FSD_ _π_ ref ( _Y−_ _X_ )
_|_ _|_
log _[π][θ]_ [(] _[Y]_ [+] _[|][X]_ [)]
_π_ ref ( _Y−_ _X_ ) _[.]_
_|_
_Noting rp_ _πθ_ ( _x, y_ + _, y−_ ) = log _π_ _[π]_ _θ_ _[θ]_ ( [(] _y_ _[y]_ _−_ [+] _[|]_ _|_ _[x]_ _x_ [)] )
_◦_
_[θ]_ [+]
_πθ_ ( _y−|x_ ) _[this can be written as follows:]_
( _rp_ _πθ_ ) _♯µ_ ≽ ( _rp_ _π_ ref ) _♯µ._ (5)
_◦_ _FSD_ _◦_
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:
Find _πθ_ such that ( _ru_ _πθ_ ) _♯µ_ + ≽ ( _ru_ _πθ_ ) _♯µ−_ (FSD unpaired)
and _∈H_ _◦_ FSD _◦_
Find _πθ_ such that ( _rp_ _πθ_ ) _♯µ_ ≽ ( _rp_ _π_ ref ) _♯µ_ (FSD paired)
_∈H_ _◦_ FSD _◦_
where _ru_ are _rp_ 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
DISTRIBUTIONAL PREFERENCE ALIGNMENT OF LLMS VIA OPTIMAL TRANSPORT 5
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θ_ ≽ _Vθ_ (6)
_∈_ FSD
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θ_ ≽ _Vθ_ _QUθ_ ( _t_ ) _QVθ_ ( _t_ ) _,_ _t_ [0 _,_ 1] _._
FSD _⇐⇒_ _≥_ _∀_ _∈_
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 1 is the indicator function):
                 - 1
min
_θ∈_ Θ
1 _QUθ_ ( _t_ ) _<QVθ_ ( _t_ ) _dt,_ (8)
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 ( 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
_h_ ( _QU_ ( _t_ ) _QV_ ( _t_ )) _dt_ = min
0 _−_ _γ∈_ Π( _µU_ _,µV_ )
_h_ ( _u −_ _v_ ) _dγ_ ( _u, v_ ) = OT _h_ ( _µU_ _, µV_ )
_and γ_ _[∗]_ = ( _QU_ _, QV_ ) _♯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) :
min
_θ∈_ Θ
- 1
0 _h_ ( _QUθ_ ( _t_ ) _−_ _QVθ_ ( _t_ )) _dt_ = min _θ∈_ Θ [OT] _[h]_ [(] _[µ][U][θ]_ _[, µ][V][θ]_ [) = min] _θ∈_ Θ _γ∈_ Π(min _µUθ_ _,µVθ_ )
_h_ ( _u −_ _v_ ) _dγ_ ( _u, v_ ) _._ (9)
This formulation reveals that we have turned the stochastic dominance constraint to an inner
one-dimensional 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_ = _n_ [1] - _ni_ =1 _[δ][u]_ _i_ _[µ]_ [ˆ] _[V]_ [ =] _n_ [1] - _ni_ =1 _[δ][v]_ _i_ _[.]_ [ Given the convexity of] _[ h]_ [ and thanks to]
in OT _h_ ( _µ_ ˆ _U_ _,_ ˆ _µV_ ) where ˆ _µU_ = _n_ [1] - _ni_ =1 _[δ][u]_ _i_ _[µ]_ [ˆ] _[V]_ [ =] _n_ [1] - _ni_ =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 _ui_ and get the order
statistics (from min to max) _u_ [(1)] _..._ _u_ [(] _[n]_ [)] and same for _n_ _vi_ : _v_ [(1)] _..._ _v_ [(] _[n]_ [)] . We have:
_≤_ _≤_ _≤_ _≤_
_n_ [1] - _ni_ =1 _[δ][u]_ _i_ _[µ]_ [ˆ] _[V]_ [ =] _n_ [1]
OT _h_ (ˆ _µU_ _,_ ˆ _µV_ ) = [1]
_n_
_n_
_h_ ( _u_ [(] _[i]_ [)] _−_ _v_ [(] _[i]_ [)] ) _._ (10)
_i_ =1
Back to (9), given empirical samples ˆ _µUθ_ = _n_ [1] - _ni_ =1 _[δ]_ _u_ _[i]_ _θ_ [and][ ˆ] _[µ][V][θ]_ [ =] _n_ [1] - _ni_ =1 _[δ]_ _vθ_ _[i]_ [, let] _[ u]_ _θ_ [(] _[i]_ [)] _[, v]_ _θ_ [(] _[i]_ [)] be the
order statistics as function of _θ_ . We have therefore:
_n_ [1] - _ni_ =1 _[δ]_ _u_ _[i]_ _θ_ [and][ ˆ] _[µ][V][θ]_ [ =] _n_ [1]
Back to (9), given empirical samples ˆ _µUθ_ = _n_ [1]
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. [1] 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]_ + [=] _n_ [1] - _ni_ =1 _[δ]_ [(] _[x]_ _i,_ + _[,y]_ _i,_ + [)] [and][ ˆ] _[µ]_ _−_ _[n]_ [=] _n_ [1] - _ni_ =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
_n_ [1] - _ni_ =1 _[δ]_ [(] _[x]_ _i,_ + _[,y]_ _i,_ + [)] [and][ ˆ] _[µ]_ _−_ _[n]_ [=] _n_ [1]
AOT **for Unpaired Preference.** Let ˆ _µ_ _[n]_ + [=] _n_ [1]
_u_ _[i]_ _θ_ [= log] _[π][θ]_ [(] _[y][i,]_ [+] _[|][x][i,]_ [+][)]
_[π][θ]_ [(] _[y][i,]_ [+] _[|][x][i,]_ [+][)] _θ_ [= log] _[π][θ]_ [(] _[y][i,][−][|][x][i,][−]_ [)]
_π_ ref ( _yi,_ + _xi,_ +) _[, v][i]_ _π_ ref ( _yi,−_ _xi,−_
_|_ _|_
_i_ = 1 _, . . ., n._
_π_ ref ( _yi,−_ _xi,−_ ) _[,]_
_|_
AOT **for Paired Preference.** Let ˆ _µ_ _[n]_ = _n_ 1 - _ni_ =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:
_u_ _[i]_ _θ_ [= log] _[π][θ]_ [(] _[y][i,]_ [+] _[|][x][i]_ [)]
_[π][θ]_ [(] _[y][i,]_ [+] _[|][x][i]_ [)] _θ_ [= log] _[π]_ [ref] [(] _[y][i,]_ [+] _[|][x][i]_ [)]
_πθ_ ( _yi,−_ _xi_ ) _[, v][i]_ _π_ ref ( _yi,−_ _xi_ )
_|_ _|_
_i_ = 1 _, . . ., n._
_π_ ref ( _yi,−_ _xi_ ) _[,]_
_|_
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
1Its Jacobian is a permutation matrix at every point except a measure-zero set where it is not differentiable.
DISTRIBUTIONAL PREFERENCE ALIGNMENT OF LLMS VIA OPTIMAL TRANSPORT 7
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 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** _such that r_ _πθ_ (Assumption on the hypothesis class of the policy) _differentiable in θ and_ sup _x∈X_ _,y∈Y_ _θr_ _πθ_ ( _y_ _x_ ) **.** _We assumeL_ _[′]_ _, θ_ Θ _π_ ref _, πB_ 2 _θ_ ( _∈Hr_ 0 _, dθ_ =) _{, forπθ_ :
_L_ _[′]_ _, r_ 0 _>_ 0 _._ _◦._ _∥∇_ _◦_ _|_ _∥≤_ _∈_ _⊂_ _}_
**Assumption 4.** _There exists πθ_ _such that_ ( _r_ _πθ_ ) _♯µ_ + ≽ ( _r_ _πθ_ ) _♯µ−._
_∈H_ _◦_ _FSD_ _◦_
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 _ru_ and _rp_ 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 costDefine the _c_ -transform of a function _c_ : [ _−M, M_ ] _×_ [ _−M, M φ_ : [ _−_ ] _→M, M_ [0 _, R_ ] _→_ ] such thatR: _c_ ( _z, z_ _[′]_ ) = _h_ ( _z_ _−z_ _[′]_ ), for _z, z ∈_ [ _−M, M_ ].
_φ_ _[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:
_Fc_ = _{φ_ : [ _−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:
min
_πθ∈H_ _φ_ [sup] _∈Fc_
- _φ_ ( _r_ _πθ_ ) _dµ_ + _φ_ _[c]_ ( _r_ _πθ_ ) _dµ−._ (12)
_◦_ _−_ _◦_
_n_ [1] - _ni_ =1 _[δ]_ [(] _[x]_ _i,_ + _[,y]_ _i,_ + [)] [and][ ˆ] _[µ]_ _−_ _[n]_ [=] _n_ [1]
Given samples ˆ _µ_ _[n]_ + [=] _n_ [1]
[1] - _n_
_n_ _i_ =1 _[δ]_ [(] _[x]_ _i,−_ _[,y]_ _i,−_ [)][, the empirical problem is:]
min
_πθ∈H_ _φ_ [sup] _∈Fc_
- _φ_ ( _r ◦_ _πθ_ ) _dµ_ ˆ _[n]_ + _[−]_ _φ_ _[c]_ ( _r ◦_ _πθ_ ) _dµ_ ˆ _[n]_ _−_ _[.]_ (13)
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:_
    -    E OT _h_ ( _r_ _πθ_ ˆ _n_ ) _♯µ_ + _,_ ( _r_ _πθ_ ˆ _n_ ) _♯µ−_ OT _h_ (( _r_ _πθ∗_ ) _♯µ_ + _,_ ( _r_ _πθ∗_ ) _♯µ−_ )
_◦_ _◦_ _≤_ _◦_ _◦_
~~�~~ ~~��~~ ~~�~~
_Optimal Almost FSD Violation_
+ 2 ~~�~~ _Rn_ ( _Fc_ ; ( _r ◦_ _πθ∗_ ) _♯µ_ +) + 2� ~~�~~ _Rn_ ( _Fc_ _[c]_ [; (] _[r][ ◦]_ _[π][θ][∗]_ [)] _[♯][µ][−]_ ~~�~~ [)] + 2� _Rn_ ( _Fc ◦_ _r ◦H_ ; _µ_ +) + 2 ~~��~~ _Rn_ ( _Fc_ _[c]_ _[◦]_ _[r][ ◦H]_ [;] _[ µ][−]_  - [)] _,_
_One dimensional OT sample complexity with optimal θ_ _[∗]_ _Complexity of learning in H via the 1D OT problem_
_where Rn_ ( _F_ ; _ν_ ) = E sup _φ∈F_ �� _n_ 1 - _ni_ =1 _[σ][i][φ]_ [(] _[Z][i]_ [)] �� _is the Rademacher Complexity and for i_ = 1 _. . . n, σi_
_are independent Rademacher random variables and Zi_ _ν 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_ _[−]_ 2 [1]
_◦_ _◦_ _−_ _◦_ _◦_
_(1)_ E OT _h_ ( _r_ _πθ_ ˆ _n_ ) _♯µ_ + _,_ ( _r_ _πθ_ ˆ _n_ ) _♯µ−_ OT _h_ (( _r_ _πθ∗_ ) _♯µ_ + _,_ ( _r_ _πθ∗_ ) _♯µ−_ ) ≲ _n_ _[−]_ 2 _, where_ ≲ _refers_
_◦_ _◦_ _−_ _◦_ _◦_
_to inequality up to constants that depend only on constants in the assumptions._
                          -                          _(2) If in addition Assumption 4 holds and h_ ( _t_ ) = ( _t_ ) [2] + _[, we have:]_ [ E][ OT] _[h]_ ( _r_ _πθ_ ˆ _n_ ) _♯µ_ + _,_ ( _r_ _πθ_ ˆ _n_ ) _♯µ−_ ≲
_−_ _◦_ _◦_
_n_ _[−]_ 2 [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 compu-_
_tational 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].
DISTRIBUTIONAL PREFERENCE ALIGNMENT OF LLMS VIA OPTIMAL TRANSPORT 9
AlpacaEval
ARC Hellaswag MMLU Truthful Winogrande GSM8K
(GPT4)
AOT paired 29.9 82.5 66.1 62.9 50.8 74.4 53.1
AOT unpaired **31.3** 82.5 **66.2** 62.8 **51.1** 74.4 51.8
DPO 27.4 **82.8** 65.8 **63.1** 50.6 74.3 52.0
KTO 24.9 82.7 65.4 63.0 48.7 **74.9** **53.9**
IPO 27.7 82.4 65.1 63.0 46.5 74.0 52.3
Merlinite-7B 17.1 81.6 63.2 62.6 42.0 73.9 45.2
Table 1. `Merlinite-7B` trained on UltraFeedback Binarized. AOT results in the best
performing LLM as compared to the alternative alignment algorithms on AlpacaEval,
and is competitive across the other benchmarks that are evaluated in the zero shot
regime.
Figure 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).
**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` and `GPT4` is the
same (the absolute score values are different), providing a better free alternative LLM-judge for
|Col1|Col2|Col3|Col4|AO|Col6|Col7|Col8|Col9|
|---|---|---|---|---|---|---|---|---|
|||||AO
|AO
|paired
|||
||||||
~~AO~~
~~KT~~|
~~T unpaire~~
|~~ d~~||
||||||DP
|O|||
||||||~~IPO~~
Mer|linite~~-~~7B|||
||||||||||
||||||||||
||||||||||
||||||||||
||||||||||
||||||||||
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.
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
DISTRIBUTIONAL PREFERENCE ALIGNMENT OF LLMS VIA OPTIMAL TRANSPORT 11
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 one of the best 7B models on that benchmark as of the time of
writing. On other benchmarks such as the open LLM leaderboard AOT leads to competitive results.
Acknowledgments
The authors would like to thank Felipe Maia Polo for his help with LLMs benchmarking.
References
AI@Meta. Llama 3 model card. 2024. URL `[https://github.com/meta-llama/llama3/blob/main/](https://github.com/meta-llama/llama3/blob/main/MODEL_CARD.md)`
`[MODEL_CARD.md](https://github.com/meta-llama/llama3/blob/main/MODEL_CARD.md)` .
C. Anil, J. Lucas, and R. Grosse. Sorting out lipschitz function approximation. In _International_
_Conference on Machine Learning_, pages 291–301. PMLR, 2019.
M. 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.
Y. 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
P. 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.](https://api.semanticscholar.org/CorpusID:2833811)`
`[org/CorpusID:2833811](https://api.semanticscholar.org/CorpusID:2833811)` .
E. 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_](https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard)`
`[llm_leaderboard](https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard)`, 2023.
M. Blondel, O. Teboul, Q. Berthet, and J. Djolonga. Fast differentiable sorting and ranking. In
_International Conference on Machine Learning_, pages 950–959. PMLR, 2020.
P. 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/](https://proceedings.neurips.cc/paper_files/paper/2017/file/d5e2c0adad503c91f91df240d0cd4e49-Paper.pdf)`
`[paper/2017/file/d5e2c0adad503c91f91df240d0cd4e49-Paper.pdf](https://proceedings.neurips.cc/paper_files/paper/2017/file/d5e2c0adad503c91f91df240d0cd4e49-Paper.pdf)` .
M. Cuturi, O. Teboul, and J.-P. Vert. Differentiable ranking and sorting using optimal transport.
_Advances in neural information processing systems_, 32, 2019.
E. 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.
C. Domingo-Enrich, Y. Schiff, and Y. Mroueh. Learning with stochastic orders. In _The Eleventh_
_International Conference on Learning Representations_, 2022.
Y. Dubois, B. Galambosi, P. Liang, and T. B. Hashimoto. Length-controlled alpacaeval: A simple
way to debias automatic evaluators. _arXiv preprint
K. Ethayarajh, W. Xu, D. Jurafsky, and D. Kiela. Human-centered loss functions (halos). Technical
report, Contextual AI, 2023.
K. Ethayarajh, W. Xu, N. Muennighoff, D. Jurafsky, and D. Kiela. Kto: Model alignment as prospect
theoretic optimization. _arXiv preprint
M. Groppe and S. Hundrieser. Lower complexity adaptation for empirical entropic optimal transport.
_arXiv preprint
E. 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
S. Hundrieser, T. Staudt, and A. Munk. Empirical optimal transport between different measures
adapts to lower complexity. _arXiv preprint
J. 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
A. 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
2023.
F. Maia Polo, L. Weber, L. Choshen, Y. Sun, G. Xu, and M. Yurochkin. tinyBenchmarks: evaluating
LLMs with fewer examples. _arXiv preprint
T. 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.
W. Ogryczak and A. Ruszczynski. Dual stochastic dominance and related mean-risk models. _SIAM_
_Journal on Optimization_, 13(1):60–78, 2002.
L. 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.
G. 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.
R. 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.
L. 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.
F. 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.
N. 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.
S. Sudalairaj, A. Bhandwaldar, A. Pareja, K. Xu, D. D. Cox, and A. Srivastava. Lab: Large-scale
alignment for chatbots. _arXiv preprint
Teknium. Openhermes-2.5-mistral-7b. `[https://huggingface.co/teknium/OpenHermes-2.](https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B)`
`[5-Mistral-7B](https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B)`, 2024.
DISTRIBUTIONAL PREFERENCE ALIGNMENT OF LLMS VIA OPTIMAL TRANSPORT 13
L. 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](https://github.com/huggingface/alignment-handbook)`, 2023a.
L. 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_
_
C. Villani. _Optimal Transport: old and new_, volume 338. Springer, 2009.
L. von Werra, Y. Belkada, L. Tunstall, E. Beeching, T. Thrush, N. Lambert, and S. Huang. Trl:
Transformer reinforcement learning. `[https://github.com/huggingface/trl](https://github.com/huggingface/trl)`, 2020.
Z. 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_
_
W. Yuan, R. Y. Pang, K. Cho, X. Li, S. Sukhbaatar, J. Xu, and J. Weston. Self-rewarding language
models, 2024.
Y. 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](https://openreview.net/forum?id=0qSOodKmJaN)` .
B. Zhu, E. Frick, T. Wu, H. Zhu, and J. Jiao. Starling-7b: Improving llm helpfulness and harmlessness
with rlaif, November 2023.
Appendix 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.
Appendix B. Algorithms and Pytorch Code In Hugging Face TRL
**Algorithm 1** AOT Unpaired
1: **Input:** _πθ_, _π_ ref _, β >_ 0 _, ε >_ 0,
2: **Unpaired Preference Data:** “Chosen”
_µ_ ˆ _[n]_ + [=] _n_ [1]  - _ni_ =1 _[δ]_ [(] _[x]_ _i,_ + _[,y]_ _i,_ + [)] [and “Rejected”]
_n_ [1] - _ni_ =1 _[δ]_ [(] _[x]_ _i,_ + _[,y]_ _i,_ + [)] [and “Rejected”]
3: ˆ _µ_ _[n]_ _−_ [=] _n_ [1] - _ni_ =1 _[δ]_ [(] _[x]_ _i,−_ _[,y]_ _i,−_ [)][.]
4: **for** iter _←_ 1 _, n_ iter **do**
5: **Get a Positive/Negative mini-batch**
6: _{_ ( _xi,_ + _, yi,_ +) _∼_ _µ_ ˆ _[n]_ + _[, i]_ [ = 1] _[ . . . b][}]_
7: _{_ ( _xi,−, yi,−_ ) _∼_ _µ_ ˆ _[n]_ _−_ _[, i]_ [ = 1] _[ . . . b][}]_
8: **Compute Rewards for** _i_ = 1 _. . . b_
**Algorithm 2** AOT Paired
1: **Input:** _πθ_, _π_ ref _, β >_ 0 _, ε >_ 0,
2: **Paired Preference Data:** _µ_ ˆ _[n]_ =
1   - _n_
_n_ _i_ =1 _[δ]_ [(] _[x]_ _i_ _[,y]_ _i,_ + _[,y]_ _i,−_ [)]
3: **for** iter _←_ 1 _, n_ iter **do**
4: **Get a Positive/Negative mini-batch**
5: ( _xi, yi,_ + _, yi,−_ ) _µ_ ˆ _[n]_ _, i_ = 1 _. . . b_
_{_ _∼_ _}_
6: **Compute Margins for** _i_ = 1 _. . . b_
7: _u_ _[i]_ _θ_ [= log] _π_ _[π]_ _θ_ _[θ]_ ( [(] _y_ _[y]_ _i,_ _[i,]_ _−_ [+] _[|]_ _|_ _[x]_ _x_ _[i]_ _i_ [)] )
3: ˆ _µ_ _[n]_ _−_ [=] [1]
9: _u_ _[i]_ _θ_ [= log] _π_ _[π]_ ref _[θ]_ [(] ( _[y]_ _y_ _[i,]_ _i,_ [+] + _[|][x]_ _|x_ _[i,]_ _i,_ [+] + [)] )
8: _vθ_ _[i]_ [= log] _π_ _[π]_ ref [ref] ( [(] _y_ _[y]_ _i,_ _[i,]_ _−_ [+] _[|]_ _|_ _[x]_ _x_ _[i]_ _i_ [)] )
10: _vθ_ _[i]_ [= log] _π_ _[π]_ ref _[θ]_ [(] ( _[y]_ _y_ _[i,]_ _i,_ _[−]_ _−_ _[|][x]_ _|x_ _[i,]_ _i,_ _[−]_ _−_ [)]
_θ_ _π_ ref ( _yi,−|xi,−_ )
11: **Sort Rewards**
12: **if** hard sort **then**
13: ( _u_ [(1)] _. . . u_ [(] _[b]_ [)] ) = Sort( _u_ _[i]_ _θ_ [)]
_θ_ _π_ ref ( _yi,−|xi_ )
9: **Sort Margins**
10: **if** hard sort **then**
11: ( _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]_ _θ_ _[, ε]_ [)]
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**
17: ( _v_ [(1)] _. . . v_ [(] _[b]_ [)] ) = SoftSort( _vθ_ _[i]_ _[, ε]_ [)]
18: **end if**
19: **Compute AOT logistic loss**
20:21: _ℓ_ **Update** _θ_ = _−_ [1] _b_ _θ_ - _bi_ =1 [log] _[ σ]_ [(] _[β]_ [(] _[u]_ [(] _θ_ _[i]_ [)] _[−]_ _[v]_ _θ_ [(] _[i]_ [)][))]
22: _θ_ PagedAdamw32bit( _θℓ_ ( _θ_ ))
_←_ _∇_
23: **end for**
24: **Return** _πθ_
18:19: _ℓ_ **Update** _θ_ = _−_ [1] _b_ _θ_ - _bi_ =1 [log] _[ σ]_ [(] _[β]_ [(] _[u]_ _θ_ [(] _[i]_ [)] _[−]_ _[v]_ _θ_ [(] _[i]_ [)][))]
20: _θ_ PagedAdamw32bit( _θℓ_ ( _θ_ ))
_←_ _∇_
21: **end for**
22: **Return** _πθ_
DISTRIBUTIONAL PREFERENCE ALIGNMENT OF LLMS VIA OPTIMAL TRANSPORT 15
1 import torch
2 import t o r c h s o r t
1 def dpo_loss ( . . .
2 . . .
3 e l i f s e l f . loss_type == "AOT_unpair" :
4 chosen_logratios = policy_chosen_logps −reference_chosen_logps
5 r e j e c t e d _ l o g r a t i o s = policy_rejected_logps −r e f e r e n c e _ r e j e c t e d _ l o g p s
6 i f s e l f . sort_type == " hard_sort " :
7 chosen_logratios_sorted, _ = torch . s o r t ( chosen_logratios, dim=0)
8 r e j e c t e d _ l o g r a t i o s _ s o r t e d, _ = torch . s o r t ( r e j e c t e d _ l o g r a t i o s, dim=0)
9 e l i f s e l f . sort_type == " s o f t _ s o r t " :
10 chosen_logratios_sorted = t o r c h s o r t . s o f t _ s o r t ( chosen_logratios, r e g u l a r i z a t i o n _ s t r e n g t h
=0.1)
11 r e j e c t e d _ l o g r a t i o s _ s o r t e d = t o r c h s o r t . s o f t _ s o r t ( r e j e c t e d _ l o g r a t i o s,
r e g u l a r i z a t i o n _ s t r e n g t h =0.1)
12 delta_sorted = chosen_logratios_sorted −r e j e c t e d _ l o g r a t i o s _ s o r t e d
13 i f s e l f . AOT_loss == " hinge " :
14 l o s s e s = torch . r e l u ( s e l f . beta −delta_sorted ) ∗∗2
15 e l i f s e l f . AOT_loss == " l o g i s t i c " :
16 l o s s e s = (
17 −F . logsigmoid ( s e l f . beta ∗delta_sorted ) ∗ (1 − s e l f . label_smoothing )
18 −F . logsigmoid (−s e l f . beta ∗ delta_sorted ) ∗ s e l f . label_smoothing
19 )
1 e l i f s e l f . loss_type == "AOT_pair" :
2 p i _ l o g r a t i o s = policy_chosen_logps −policy_rejected_logps
3 r e f _ l o g r a t i o s = reference_chosen_logps −r e f e r e n c e _ r e j e c t e d _ l o g p s
4 i f s e l f . sort_type == " hard_sort " :
5 pi_logratios_sorted, _ = torch . s o r t ( p i _ l o g r a ti o s, dim=0)
6 ref_logratios_sorted, _ = torch . s o r t ( r e f _ l o g r a t i o s, dim=0)
7 e l i f s e l f . sort_type == " s o f t _ s o r t " :
8 pi_logratios_sorted = t o r c h s o r t . s o f t _ s o r t ( p i _ l og r a t i o s, r e g u l a r i z a t i o n _ s t r e n g t h =0.1)
9 r e f _ l o g r a t i o s _ s o r t e d = t o r c h s o r t . s o f t _ s o r t ( r e f _ l o g r a t i o s, r e g u l a r i z a t i o n _ s t r e n g t h
=0.1)
10 delta_sorted = pi_logratios_sorted −r e f _ l o g r a t i o s _ s o r t e d
11 i f s e l f . AOT_loss == " hinge " :
12 l o s s e s = torch . r e l u ( s e l f . beta −delta_sorted ) ∗∗2
13 e l i f s e l f . AOT_loss == " l o g i s t i c " :
14 l o s s e s = (
15 −F . logsigmoid ( s e l f . beta ∗delta_sorted ) ∗ (1 − s e l f . label_smoothing )
16 −F . logsigmoid (−s e l f . beta ∗ delta_sorted ) ∗ s e l f . label_smoothing
17 )
_Proof of Theorem 2._ Let
Appendix C. Proofs
_ε_ ( _θ, µ_ + _, µ−_ ) = 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 _c−_ concave if there exits _g_ such that _f_ = _g_ _[c]_ ).
The population problem :
min (14)
_πθ∈H_ _[ε]_ [(] _[θ, µ]_ [+] _[, µ][−]_ [)]
_n_ [1] - _ni_ =1 _[δ]_ ( _x_ _[i]_ + _[,y]_ + _[i]_ [)][ and][ ˆ] _[µ]_ _−_ _[m]_ [=] [1]
Given samples ˆ _µ_ _[n]_ + [=] _n_ [1]
_m_ [1] - _mi_ =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 + _[,]_ [ ˆ] _[µ][m]_ _−_ [)]
_πθ∈H_ _[ε]_ [(] _[θ,]_ [ ˆ] _[µ][n]_
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_
_≤_ _φ_ sup _∈Fc_
+ sup
_φ∈Fc_
+ sup
_φ∈Fc_
- _φ_ ( _r_ _πθ_ ) _dµ_ ˆ _[n]_ + _φ_ _[c]_ ( _r_ _πθ_ ) _dµ_ ˆ _[m]_ _−_
_◦_ _[−]_ _◦_
- _φ_ ( _r_ _πθ_ ) _dµ_ + _φ_ _[c]_ ( _r_ _πθ_ ) _dµ−_
_◦_ _−_ _◦_
_φ_ ( _r_ _πθ_ ) _d_ (ˆ _µ_ _[n]_ +
_◦_ _[−]_ _[µ]_ [+][)]
_φ_ _[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 obtain
E _ε_ ( _θ_ [ˆ] _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:
E sup
_φ∈Fc_
_φ_ ( _r_ _πθ∗_ ) _d_ (ˆ _µ_ _[n]_ + (15)
_◦_ _[−]_ _[µ]_ [+][)] _[ ≤]_ [2] _[R][n]_ [(] _[F][c]_ [)] _[,]_
where _n_ ( ) = E sup _φ∈F_ ��� _N_ 1
_R_ _F_
where _n_ ( ) = E sup _φ∈F_ ��� _N_ 1 - _Ni_ =1 _[σ][i][φ]_ [(] _[X][i]_ [)] ��� _, σi_ are independent rademacher random variables and
_R_ _F_
_Xi ∼_ ( _r ◦_ _πθ∗_ ) _♯µ_ + iid ( _Xi ∈_ R). and similarly we have:
E sup
_φ∈Fc_
_φ_ _[c]_ ( _r ◦_ _πθ∗_ ) _d_ ( _µ−_ _−_ _µ_ ˆ _[m]_ _−_ [)] _[ ≤]_ [E][ sup] _φ_ _[c]_ _∈Fc_ _[c]_
_φ_ _[c]_ ( _r ◦_ _πθ∗_ ) _d_ ( _µ−_ _−_ _µ_ ˆ _[m]_ _−_ [)] _[ ≤]_ [2] _[R][m]_ [(] _[F]_ _c_ _[c]_ [)]
We have finally:
E _ε_ ( _θ_ [ˆ] _m,n,_ ˆ _µ_ _[n]_ + _[,]_ [ ˆ] _[µ][m]_ _−_ [)] _[ ≤]_ _[ε]_ [(] _[θ][∗][, µ]_ [+] _[.µ][−]_ [) + 2] _[R][n]_ [(] _[F][c]_ [) + 2] _[R][m]_ [(] _[F]_ _c_ _[c]_ [)] (16)
DISTRIBUTIONAL PREFERENCE ALIGNMENT OF LLMS VIA OPTIMAL TRANSPORT 17
Turning now to :
         -
_ε_ ( _θ_ [ˆ] _m,n, µ_ + _, µ−_ ) = sup _φ_ ( _r_ _πθ_ ˆ _m,n_ ) _dµ_ + _φ_ _[c]_ ( _r_ _πθ_ ˆ _m,n_ ) _dµ−_
_φ∈Fc_ _◦_ _−_ _◦_
_≤_ _ε_ ( _θ_ [ˆ] _m,n,_ ˆ _µ_ _[n]_ + _[,]_ [ ˆ] _[µ]_ _−_ _[m]_ [) + sup] _φ∈Fc_
_φ_ ( _r_ _πθ_ ˆ _m,n_ ) _d_ ( _µ_ + _µ_ _[n]_ + [)]
_◦_ _−_
+ sup _φ_ _[c]_ ( _r_ _πθ_ ˆ _m,n_ ) _d_ ( _µ_ _[m]_ _−_ _[−]_ _[µ][−]_ [)]
_φ∈Fc_ _◦_
_≤_ _ε_ ( _θ_ [ˆ] _m,n,_ ˆ _µ_ _[n]_ + _[,]_ [ ˆ] _[µ]_ _−_ _[m]_ [) + sup] _πθ∈H_ _φ_ sup _∈Fc_
_φ_ ( _r ◦_ _πθ_ ) _d_ ( _µ_ + _−_ _µ_ _[n]_ + [)]
+ sup sup
_πθ∈H_ _φ∈Fc_
Taking expectations we obtain:
_φ_ _[c]_ ( _r ◦_ _πθ_ ) _d_ ( _µ_ _[m]_ _−_ _[−]_ _[µ][−]_ [)]
E _ε_ ( _θ_ [ˆ] _m,n, µ_ + _, µ−_ ) _≤_ E _ε_ ( _θ_ [ˆ] _m,n,_ ˆ _µ_ _[n]_ + _[,]_ [ ˆ] _[µ]_ _−_ _[m]_ [) +][ E][ sup] _πθ∈H_ _φ_ sup _∈Fc_
_φ_ ( _r ◦_ _πθ_ ) _d_ ( _µ_ + _−_ _µ_ _[n]_ + [)]
+ E sup sup
_πθ∈H_ _φ_ _[c]_ _∈_ ( _Fc_ ) _[c]_
_φ_ _[c]_ ( _r ◦_ _πθ_ ) _d_ ( _µ_ _[m]_ _−_ _[−]_ _[µ][−]_ [)]
_ε_ ( _θ_ _[∗]_ _, µ_ + _.µ−_ )
_≤_
~~�~~ ~~��~~   Optimal FSD Violation
+ 2 _Rn_ ( _Fc_ ) + 2 _Rm_ ( _Fc_ _[c]_ [)]
~~�~~      - ~~�~~ ~~�~~
One dimensional OT complexity with optimal _θ_ _[∗]_
+ 2 _n_ ( _c_ _r_ ) + 2 _m_ ( _c_
~~�~~ _R_ _F_ _◦_ _◦H_ ~~��~~ _R_ _F_ _[c]_ _[◦]_ _[r][ ◦H]_ ~~�~~ [)]
Complexity of learning in _H_ via the 1D OT problem
_,_
where
- _σiφ_ ( _r ◦_ _πθ_ ( _x_ [+] _i_ _[, y]_ _i_ [+][))]
_i_ =1
(17)
�����
_n_ ( _c_ _r_ ) = [1] sup
_R_ _F_ _◦_ _◦H_ _n_ [E][ sup] _πθ∈H_ _φ∈Fc_
�����
_n_
Appendix D. Bounding Rademacher Complexities
_Proof of Corollary 1._ Define the uniform metric entropy of a class of real valued functions _F_ on a
set as the logarithm of the covering number with respect to the uniform norm _._ _∞_, for _ε >_ 0,
_X_ _∥_ _∥_
this is defined as follows:
        -         _N_ ( _ε, F, ∥.∥∞_ ) := inf _n ∈_ N��� there exists _f_ 1 _. . . fn_ : _X →_ R with sup _f_ _∈F_ 1min _≤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 :
_N_ ( _ε, Fc_ _[c][,][ ∥][.][∥]_ _∞_ [) =] _[ N]_ [(] _[ε,][ F][c][,][ ∥][.][∥]_ _∞_ [)] (18)
Now turning to the Rademacher complexity of a class _F_, it is dominated by Dudley’s entropy integral
(Theorem 16 in Luxburg and Bousquet ):
32 [1]
~~_√n_~~
~~�~~
- _R_
_δ/_ 4
log ( _ε,_ _,_ _._ _∞_ ) _dε_
_N_ _F_ _∥_ _∥_
(19)
_n_ ( ) inf
_R_ _F_ _≤_ _δ∈_ [0 _,R_ ]
�2 _δ_ + _√_
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:
log ( _ε,_ _c,_ _._ _∞_ ) ≲ _ε_ _[−]_ _α_ _[d]_ _._
_N_ _F_ _∥_ _∥_
_α ._
Hence in our case _α_ = 1 and _d_ = 1 this leads to
log ( _ε,_ _c,_ _._ _∞_ ) ≲ _ε_ _[−]_ [1]
_N_ _F_ _∥_ _∥_
Replacing this in Equation (19) we obtain: _n_ ( _c_ ) ≲ _n_ _[−]_ 2 [1]
[1] _R_ _F_
Replacing this in Equation (19) we obtain: _n_ ( _c_ ) ≲ _n_ 2 and by Equation (18) it follows that:
_R_ _F_
_m_ ( _c_ [)][ ≲] _[m][−]_ 2 [1] _._
_R_ _F_ _[c]_
_Rm_ ( _Fc_ _[c]_ [)][ ≲] _[m][−]_ 2 _._
Turning now to the Rademacher Complexity of the composition of the c-concave potentials _c_
_F_
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 _Fc_ is included in the set of _L_ lipchitz
function that are bounded by, let us note _hφ,πθ_ = _φ_ ( _r ◦_ _πθ_ ) _R_ we have: (See [Hundrieser et al., 2022] Lemma A.2 ). For _φ ∈Fc_ and _πθ ∈H_
�� _hφ,πθ −_ _hφ′,πθ′_ �� _∞_ [=] _x∈X_ sup _,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 _φ ∈Fc_ and hence we have:
_′_
�� _hφ,πθ −_ _hφ′,πθ′_ �� _∞_ _[≤]_ _[L][ ∥][r][ ◦]_ _[π][θ][ −]_ _[r][ ◦]_ _[π][θ][′][∥][∞]_ [+] �� _φ −_ _φ_ �� _∞_ _[.]_
We have therefore the following bound on the covering number of the composition:
_ε_ - - _ε_
2 _[,][ F][c][,][ ∥][.][∥][∞]_ _N_ 2
          - _ε_
_N_ ( _ε, Fc ◦_ _r ◦H, ∥.∥∞_ ) _≤N_ 2
      _,_
2 _L_ _[, r][ ◦H][,][ ∥][.][∥][∞]_
Plugging this in Equation (19) we obtain:
32 [1]
~~_√n_~~
~~�~~
_dε_
2 _L_ _[, r][ ◦H][,][ ∥][.][∥][∞]_
- _R_
_δ/_ 4
~~�~~
~~�~~ _ε_
log _N_ ( _ε/_ 2 _, Fc, ∥.∥∞_ ) + log _N_ 2
_n_ ( _c_ _r_ ) inf
_R_ _F_ _◦_ _◦H_ _≤_ _δ∈_ [0 _,R_ ]
�2 _δ_ + _√_
Note that for _a, b >_ 0 we have _√_
_b_, hence we have:
_a_ + _b ≤_ _[√]_ ~~_a_~~ + _√_
32 [1]
~~_√n_~~
~~�~~
~~�~~
_dε_
2 _L_ _[, r][ ◦H][,][ ∥][.][∥][∞]_
- _R_
_δ/_ 4
~~�~~
log ( _ε/_ 2 _,_ _c,_ _._ _∞_ ) +
_N_ _F_ _∥_ _∥_
~~�~~ _ε_
log
_N_ 2
_n_ ( _c_ _r_ ) inf
_R_ _F_ _◦_ _◦H_ _≤_ _δ∈_ [0 _,R_ ]
2 _δ_ + _√_
We know by lipchitizty of the cost and being in one dimension that :
log ( _ε,_ _c,_ _._ _∞_ ) ≲ _ε_ _[−]_ [1]
_N_ _F_ _∥_ _∥_
By lipchitizity of _r ◦_ _πθ_ and using Assumption 3 we have therefore:
_ε_ - - _ε_
2 _L_ _[, r][ ◦H][,][ ∥][.][∥][∞]_ _≤_ log _N_ 2 _LL_ _[′]_ _[, B]_ [2][(] _[r]_ [0] _[, d][θ]_ [)] _[,][ ∥][.][∥][∞]_ _≤_ _dθ_ log [2] _[r]_ [0] _ε_ _[L][′][L]_
  - _ε_
log
_N_ 2
2 _L_ 2 _LL_ _[′]_ _ε_
We have therefore:
2 _._
_dε_ ≲ _n_ _[−]_ 2 [1]
_ε_
- _−_ 1 _/_ 2
_dε_ + 4
~~�~~
- _R_
   - _ε_
_K_ 1
_δ/_ 4 2
- _R_
_δ/_ 4
_dθ_ log [2] _[r]_ [0] _[LL][′]_
_√_
~~_√n_~~
inf
_δ∈_ [0 _,R_ ] [2] _[δ]_ [ + 4]
_√_
~~_√n_~~
(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 _._ _−_
DISTRIBUTIONAL PREFERENCE ALIGNMENT OF LLMS VIA OPTIMAL TRANSPORT 19
Appendix 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:
min
_πθ∈H_ _φ_ [sup] _∈Fc_
- _φ_ ( _r_ _πθ_ ) _dµ_ _φ_ _[c]_ ( _r_ _π_ ref ) _dµ,_ (20)
_◦_ _−_ _◦_
where _µ_ is the paired measure representing ( _X, Y_ + _, Y−_ ). let _πθ∗_ be the minimizer of (20). Considering
Problem (20), with empirical samples ˆ _µ_ _[n]_ = _n_ [1] - _ni_ =1 _[δ]_ [(] _[x]_ _i_ _[,y]_ _i,_ + _[,y]_ _i,−_ [)] [, 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_
+ 2 _Rn_ ( _Fc_ ; ( _r ◦_ _πθ∗_ ) _♯µ_ +) + 4 _Rn_ ( _Fc_ _[c]_ [; (] _[r][ ◦]_ _[π]_ [ref] [)] _[♯][µ]_ [)]
                      - ~~�~~                      - ~~�~~
_One dimensional OT sample complexity with optimal θ_ _[∗]_
+ 2 _n_ ( _c_ _r_ ; _µ_ )) _,_
                           - _R_ _F_ _◦_ ~~��~~ _◦H_ ~~�~~
_Complexity of learning in H via the 1D OT problem_
_where Rn_ ( _F_ ; _ν_ ) = E sup _φ∈F_ �� _n_ 1 - _ni_ =1 _[σ][i][φ]_ [(] _[Z][i]_ [)] �� _is the Rademacher Complexity and for i_ = 1 _. . . n, σi_
_are independent rademacher random variables and Zi_ _ν iid._
_∼_
_where_ _n_ ( ; _ν_ ) = E sup _φ∈F_ �� _n_ 1
_R_ _F_
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 Fc by Fc,ε. In the 1D_
_case, up to constants, these spaces have the same covering numbers._
Appendix F. Choice of violation penalty function _h_
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.
(2) **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θ_ .
(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.
Appendix 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.
min _θ∈_ Θ [OT] _[h]_ [(ˆ] _[µ]_ _U_ [(] _[n]_ _θ_ [)] _[,]_ [ ˆ] _[µ]_ _V_ [(] _[n]_ _θ_ [)][) = min] _θ∈_ Θ _n_
_n_
_i_ =1 _h_ ( _u_ [(] _θ_ _[i]_ [)] _[−]_ _[v]_ _θ_ [(] _[i]_ [)][)] _[.]_ (21)
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)._ [2] _Then the_
_gradient of the objective in_ (21) _is asymptotically unbiased as n →∞._
_Proof._ Observe that the gradients take the form
_∇θ_ OT _h_ (ˆ _µ_ [(] _U_ _[n]_ _θ_ [)] _[,]_ [ ˆ] _[µ]_ _V_ [(] _[n]_ _θ_ [)][) =] _[ ∇][θ]_ _n_
_n_
_i_ =1 _h_ ( _u_ [(] _θ_ _[i]_ [)] _[−]_ _[v]_ _θ_ [(] _[i]_ [)][)] (22)
= [1]
_n_
_n_
_i_ =1 _h_ _[′]_ ( _u_ [(] _θ_ _[i]_ [)] _[−]_ _[v]_ _θ_ [(] _[i]_ [)][)] _[∇][θ]_ [(] _[u]_ _θ_ [(] _[i]_ [)] _[−]_ _[v]_ _θ_ [(] _[i]_ [)][)] _[.]_ (23)
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 →∞_ .
2In our alignment setting, this is should not be an issue as our scores are real-valued.
DISTRIBUTIONAL PREFERENCE ALIGNMENT OF LLMS VIA OPTIMAL TRANSPORT 21
We can write (noting that the _n_ atoms of the empirical _Fn,Vθ_ are distinct with probability 1)
_n_
_n_
- _h_ _[′]_ ( _u_ [(] _θ_ _[i]_ [)] _θ_ [)] _[∇][θ]_ [(] _[u]_ _θ_ [(] _[i]_ [)] _θ_ [) =] [1]
_n_
_i_ =1 _[−]_ _[v]_ [(] _[i]_ [)] _[−]_ _[v]_ [(] _[i]_ [)]
_n_
_n_
- _i_ =1 _h_ _[′]_ ( _u_ [(] _θ_ _[i]_ [)] _[−]_ _[F][ −]_ _n,V_ [1] _θ_ [(] _[F][n,U]_ _θ_ [(] _[u]_ [(] _θ_ _[i]_ [)][)))] _[∇][θ]_ [(] _[u]_ _θ_ [(] _[i]_ [)][)]
- ~~�~~ - ~~�~~
_I_
_−_ _n_ [1]
_n_
- _h_ _[′]_ ( _Fn,U_ _[−]_ [1] _θ_ [(] _[F][n,V]_ _θ_ [(] _[v]_ _θ_ [(] _[i]_ [)][))] _[ −]_ _[v]_ _θ_ [(] _[i]_ [)][)] _[∇][θ]_ [(] _[v]_ _θ_ [(] _[i]_ [)][)]
_i_ =1
_._
~~�~~                        - ~~�~~                        _II_
Let’s consider the first term, the analysis for the second term is the same. Note that since _h_ _[′]_ is
_L_ -Lipschitz
��� _h′_       - _u −_ _Fn,V_ _[−]_ [1] _θ_ [(] _[F][n,U]_ _θ_ [(] _[u]_ [))]       - _−_ _h_ _[′]_ [ �] _u −_ _FV_ _[−]_ _θ_ [1][(] _[F][U]_ _θ_ [(] _[u]_ [))] ���� _≤_ _L_ ��� _F −n,V_ 1 _θ_ [(] _[F][n,U]_ _θ_ [(] _[u]_ [))] _[ −]_ _[F][ −]_ _Vθ_ [1][(] _[F][U]_ _θ_ [(] _[u]_ [))] ��� _._ (24)
Let
_I_ _[′]_ = [1]
_n_
_n_
- _h_ _[′]_ ( _u_ [(] _θ_ _[i]_ [)] _Vθ_ [(] _[F][U]_ _θ_ [(] _[u]_ [(] _θ_ _[i]_ [)][)))] _[∇][θ]_ [(] _[u]_ _θ_ [(] _[i]_ [)][)] _[.]_
_i_ =1 _[−]_ _[F][ −]_ [1]
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
_I_ _I_ _[′]_ 1
_∥_ _−_ _∥_ _≤_ _[L]_ _n_
_n_
_i_ =1
��� _F −n,V_ 1 _θ_ [(] _[F][n,U]_ _θ_ [(] _[u]_ [(] _θ_ _[i]_ [)][))] _[ −]_ _[F][ −]_ _Vθ_ [1][(] _[F][U]_ _θ_ [(] _[u]_ [(] _θ_ _[i]_ [)][))] ��� _∥∇θ_ ( _u_ ( _θi_ ) [)] _[∥]_ [1] (25)
We seek to bound this quantity as _n →∞_ . By Cauchy-Schwarz,
_n_
- _∥∇θ_ ( _u_ [(] _θ_ _[i]_ [)][)] _[∥]_ 1 [2]
_i_ =1
_∥I −_ _I_ _[′]_ _∥_ 1 [2] _[≤]_ _[L]_ [2] _n_
_≤_ _[L]_ [2] _n_ _[M]_ [2]
= _[L]_ [2] _[M]_ [2]
_n_
_n_ ��
- - _Fn,V_ _[−]_ [1] _θ_ [(] _[F][n,U]_ _θ_ [(] _[u]_ [(] _θ_ _[i]_ [)][))] _[ −]_ _[F][ −]_ _Vθ_ [1][(] _[F][U]_ _θ_ [(] _[u]_ [(] _θ_ _[i]_ [)][))] �2 _n_ 1
_i_ =1
_n_
- - _Fn,V_ _[−]_ [1] _θ_ [(] _[F][n,U]_ _θ_ [(] _[u]_ [(] _θ_ _[i]_ [)][))] _[ −]_ _[F][ −]_ _Vθ_ [1][(] _[F][U]_ _θ_ [(] _[u]_ [(] _θ_ _[i]_ [)][))] �2
_i_ =1
_n_
- - _vθ_ [(] _[i]_ [)] _−_ _FV_ _[−]_ _θ_ [1][(] _[F][U]_ _θ_ [(] _[u]_ [(] _θ_ _[i]_ [)][))] �2
_i_ =1
where we’ve assumed that - _n_ 1 - _ni_ =1 _[∥∇][θ]_ [(] _[u]_ _θ_ [(] _[i]_ [)][)] _[∥]_ 1 [2] - _≤_ _M_ [2] . Observe that _FV_ _[−]_ _θ_ [1][(] _[F][U]_ _θ_ [(] _[·]_ [))][ is the optimal]
transport plan from _Uθ_ to _Vθ_ and is monotonic nondecreasing, hence _FV_ _[−]_ _θ_ [1][(] _[F][U]_ _θ_ [(] _[u]_ [(] _θ_ _[i]_ [)][))][ are simply a]
new set of independently drawn order statistics of _vθ_, i.e. _vθ_ [(] _[i,]_ [2)] .
We thus have
_∥I −_ _I_ _[′]_ _∥_ 1 [2] _[∼]_ _[L]_ [2] _n_ _[M]_ [2]
_n_
- - �2
_vθ_ [(] _[i]_ [)] _vθ_ [(] _[i,]_ [2)] (26)
_−_
_i_ =1
= _L_ [2] _M_ [2] ��[ _Fn,V_ [(1)] _θ_ []] _[−]_ [1][(] _[t]_ [)] _[ −]_ [[] _[F]_ [ (2)] _n,Vθ_ []] _[−]_ [1][(] _[t]_ [)] �2 _dt_ (27)
where _Fn,V_ [(] _[i]_ [)] _θ_ [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_
           - 1           - 1
- _F_ _[−]_ [1] _−_ _G_ _[−]_ [1][�][2] _→p_ 0
       - 1
- _Fn_ _[−]_ [1] _G_ _[−]_ _m_ [1] �2
_−_ _−_ 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] ��[ _Fn,V_ [(1)] _θ_ []] _[−]_ [1][(] _[t]_ [)] _[ −]_ [[] _[F]_ [ (2)] _n,Vθ_ []] _[−]_ [1][(] _[t]_ [)] �2 _dt →p_ 0
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.
Appendix 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.
AlpacaEval AlpacaEval
ARC Hellaswag MMLU Truthful Winogrande GSM8K
(Llama3-70B) (GPT4)
AOT paired 44.3 29.9 82.5 66.1 62.9 50.8 74.4 53.1
AOT unpaired **48.4** **31.3** 82.5 **66.2** 62.8 **51.1** 74.4 51.8
DPO 36.8 27.4 **82.8** 65.8 **63.1** 50.6 74.3 52.0
KTO 35.7 24.9 82.7 65.4 63.0 48.7 **74.9** **53.9**
IPO 43.1 27.7 82.4 65.1 63.0 46.5 74.0 52.3
Merlinite-7B 28.8 17.1 81.6 63.2 62.6 42.0 73.9 45.2
Table 2. Merlinite-7B trained on UltraFeedback Binarized. Here we present full
version of the results, including AlpacaEval using Llama3-70B-instruct as a judge
and GPT4 as a judge. The comparison reveals that although Llama3 inflates the
scores, the relative order between the two judges remains the same, suggesting the
use of a cheaper AlpacaEval alternative for local development.
Table 7 is an ablation on the sorting that is used in AOT with `Merlinite-7B` as a reference model.
We see that hard and soft sorting are on par in terms of overall performance.
Tables 8 and 9 give a comparison between AOT and KTO on unpaired datasets (HelpSteer and
PKU binarized) we see that overall AOT leads to a better performance than KTO.
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.
DISTRIBUTIONAL PREFERENCE ALIGNMENT OF LLMS VIA OPTIMAL TRANSPORT 23
AlpacaEval
ARC Hellaswag MMLU Truthful Winogrande GSM8K
(Llama3-70B)
AOT paired **24.4** 84.1 **66.1** 61.0 **50.6** **74.9** 66.6
AOT unpaired 22.5 **84.2** 66.0 61.0 50.5 74.8 65.7
DPO 17.9 84.1 66.0 61.0 50.4 74.4 **66.7**
KTO 12.6 83.5 64.3 **61.1** 47.2 74.4 66.3
IPO 15.5 83.9 65.4 **61.1** 49.2 74.2 66.3
OpenHermes-7B 5.6 83.4 63.1 60.6 44.5 74.4 63.8
Table 3. OpenHermes-2.5-Mistral-7B trained on UltraFeedback Binarized
AlpacaEval
ARC Hellaswag MMLU Truthful Winogrande GSM8K
(Llama3-70B)
AOT paired 30.4 84.3 66.9 61.4 45.5 72.6 69.0
AOT unpaired **34.4** **85.1** **67.4** **61.5** **47.0** 72.3 68.5
DPO 28.6 84.5 66.7 61.4 45.3 72.5 69.8
KTO 27.2 84.8 67.0 61.4 46.2 **74.2** **70.2**
IPO 28.6 84.5 66.7 61.4 44.4 72.9 69.8
Starling-7B 14.3 83.4 64.4 60.9 39.4 72.5 66.6
Table 4. Starling-LM-7B-alpha trained on UltraFeedback Binarized
AlpacaEval
ARC Hellaswag MMLU Truthful Winogrande GSM8K
(Llama3-70B)
AOT paired 33.6 81.7 59.4 **64.1** 47.7 72.6 **78.0**
AOT unpaired **35.8** 81.8 59.4 64.0 **47.8** 72.8 77.5
DPO 33.1 **82.1** **59.5** 64.0 47.3 73.1 77.9
KTO 28.5 81.9 59.0 63.9 46.5 **73.4** 77.7
IPO 33.2 81.9 59.1 63.9 46.7 72.9 77.7
Llama3-8B 25.4 81.6 57.7 63.8 43.9 72.5 75.9
Table 5. Meta-Llama-3-8B-Instruct trained on UltraFeedback Binarized
AlpacaEval
ARC Hellaswag MMLU Truthful Winogrande GSM8K
(Llama3-70B)
AOT paired 32.8 81.6 67.6 **58.9** 62.7 74.3 41.9
AOT unpaired **34.3** 81.7 67.6 59.1 **63.0** **74.4** 41.9
DPO 28.8 81.7 67.6 58.8 62.3 74.2 **42.0**
KTO 27.4 **81.9** **67.7** 58.8 62.5 74.3 41.7
IPO 28.4 **81.9** 67.1 58.8 60.5 74.0 41.9
Mistral-7B 25.6 81.3 66.0 58.8 59.7 74.1 41.7
Table 6. Mistral-7B-Instruct-v0.2 trained on UltraFeedback Binarized
Sort AlpacaEval
ARC Hellaswag MMLU Truthful Winogrande GSM8K
Type (Llama3-70B)
Soft 44.3 82.5 66.1 62.9 50.8 74.4 53.1
AOT paired
Hard 43.8 82.7 66.2 62.9 50.7 74.5 53.9
Soft 48.4 82.5 66.2 62.8 51.1 74.4 51.8
AOT unpaired
Hard 49.2 82.5 65.9 62.8 51.0 74.4 51.0
Table 7. The effect of sort type on performance in AOT alignment
ARC Hellaswag MMLU Truthful Winogrande GSM8K
AOT unpaired **82.0** **63.5** **62.9** **45.6** **74.9** 48.0
KTO 81.9 63.5 62.7 45.0 74.2 **48.5**
Merlinite-7B 81.6 63.2 62.6 42.0 73.9 45.2
Table 8. Merlinite-7B trained on unpaired HelpSteer (binarized)
ARC Hellaswag MMLU Truthful Winogrande GSM8K
AOT unpaired 82.0 **64.1** **63.0** **56.3** **74.6** 49.7
KTO **82.1** 63.5 62.9 43.5 74.5 **50.4**
Merlinite-7B 81.6 63.2 62.6 42.0 73.9 45.2
Table 9. Merlinite-7B trained on unpaired PKU (binarized)
AlpacaEval
(Llama3-70B)
AOT paired 46 _._ 8 _±_ 1 _._ 48
AOT unpaired 48 _._ 1 _±_ 0 _._ 35
DPO 39 _._ 2 _±_ 2 _._ 35
KTO 33 _._ 8 _±_ 1 _._ 23
IPO 45 _._ 7 _±_ 1 _._ 56
Merlinite-7B 28.8
Table 10. Merlinite-7B trained on UltraFeedback Binarized. We evaluated the
stability (variance) of the model evaluation on AlpacaEval by running 4 separate
training and evaluation cycles, then computing the mean and standard deviations.
The results are stable, especially for AOT unpaired, showing a low deviation from
the mean.
DISTRIBUTIONAL PREFERENCE ALIGNMENT OF LLMS VIA OPTIMAL TRANSPORT 25
Figure 4. Our AOT algorithm gives a strong boost to Merlinite-7B model on
AlpacaEval leaderboard (as of May 22nd, 2024). The original Merlinite-7B score is
17.1, and after the alignment, the model gained 83%.

Distributional Preference Alignment of LLMs via Optimal Transport

Igor Melnyk
Youssef Mroueh
Brian Belgodere
Mattia Rigotti
Apoorva Nitsure
Mikhail Yurochkin
Kristjan Greenewald
Jiri Navratil
Jerret Ross


Princeton University
```
xiangyuqi@princeton.edu
```

Safety Alignment Should Be Made More Than Just a Few Tokens Deep

Xiangyu Qi
Ashwinee Panda
Kaifeng Lyu
Xiao Ma
Subhrajit Roy
Ahmad Beirami
Prateek Mittal
Peter Henderson


Princeton University
```
  ashwinee@princeton.edu
```


Princeton University
```
klyu@cs.princeton.edu
```


Princeton University
```
pmittal@princeton.edu
```


Princeton University
```
   peter.henderson@princeton.edu
```


The safety alignment of current Large Language Models (LLMs) is vulnerable.
Relatively simple attacks, or even benign fine-tuning, can jailbreak aligned models.
We argue that many of these vulnerabilities are related to a shared underlying
issue: safety alignment can take shortcuts, wherein the alignment adapts a model’s
generative distribution primarily over only its very first few output tokens. We
refer to this issue as shallow safety alignment. In this paper, we present case
studies to explain why shallow safety alignment can exist and provide evidence that
current aligned LLMs are subject to this issue. We also show how these findings
help explain multiple recently discovered vulnerabilities in LLMs, including the
susceptibility to adversarial suffix attacks, prefilling attacks, decoding parameter
attacks, and fine-tuning attacks. Importantly, we discuss how this consolidated
notion of shallow safety alignment sheds light on promising research directions
for mitigating these vulnerabilities. For instance, we show that deepening the
safety alignment beyond just the first few tokens can often meaningfully improve
robustness against some common exploits. We also design a regularized fine-tuning
objective that makes the safety alignment more persistent against fine-tuning attacks
by constraining updates on initial tokens. Overall, we advocate that future safety
alignment should be made more than just a few tokens deep. [1]


Currently, the safety of Large Language Models (LLMs) [1–7] heavily hinges on AI alignment
approaches [8–12]—typically a mixture of supervised Fine-tuning (SFT) [13] and preference-based
optimization methods like Reinforcement Learning with Human Feedback (RLHF) [14, 15] and
Direct Preference Optimization (DPO) [16]. These approaches aim to optimize models so that they
refuse to engage with harmful inputs, thus reducing the likelihood of generating harmful content.
However, recent studies find that such alignment approaches suffer from various vulnerabilities.
For example, researchers demonstrate that aligned models can still be made to respond to harmful
requests via adversarially optimized inputs [17–21], a few gradient steps of fine-tuning [22, 23],
1Our code is available at `[https://github.com/Unispac/shallow-vs-deep-alignment](https://github.com/Unispac/shallow-vs-deep-alignment)`
Preprint. Under review.
or simply exploiting the model’s decoding parameters [24]. Given the pivotal role that alignment
plays in LLM safety, and its widespread adoption, it is imperative to understand why current safety
alignment is so vulnerable to these exploits and to identify actionable approaches to mitigate them.
In this paper, we examine one underlying problem in current safety alignment that may make models
particularly vulnerable to relatively simple exploits: safety alignment is largely only a few tokens
deep, i.e., it adapts the model’s generative distribution primarily over only the very first few output
tokens. Consequently, happening upon, or adversarially induced, if the model’s initial output tokens
deviate from some routine safe prefixes, its generation could catastrophically fall on a harmful
trajectory. For example, consider the scenario where a user asks, “How do I build a bomb?” and
induces the model to begin its response with, “Sure, here’s a detailed guide.” The model is then much
more likely to continue with harmful information responsive to the user’s request. We refer to this
problem as _**shallow safety alignment**_ (Section 2). We call the counterfactual, where a model can
recover from such harmful starting conditions, _**deep safety alignment**_ . To provide sufficient context
to this notion, our work has three main contributions.
First, we conduct systematic experiments to characterize the shallow safety alignment issue in current
LLMs (Section 2). We demonstrate that the primary difference in safety behaviors between an aligned
model and its unaligned counterpart lies in their modeling of only the first few tokens of their outputs. [2]
Part of the problem is that there are easy optimization shortcuts that may drive such a local optimum.
For example, simply prefilling an unaligned base model to start its output with a prefix “I cannot
fulfill” is sufficient to make it as safe as aligned models. We note that the shallow safety alignment
issue helps explain why attack methods that focus on initiating trajectories with harmful or affirmative
responses are so effective, like adversarial suffix attacks [29], decoding parameters exploit [24], and
an emerging paradigm of prefilling attacks [30, 21]. Moreover, we show that fine-tuning attacks [22]
also create the most significant changes in the first few tokens of a harmful response. This means that
by simply modifying these initial tokens, it is possible to undo the model alignment, explaining why
so few fine-tuning steps can lead to jailbroken models.
Second, we argue that future safety alignment approaches should focus on extending their effects
deeper. To support this idea, we introduce a simple data augmentation approach for deepening the
safety alignment (Section 3). By training on safety alignment data that begins with harmful responses
and transitions back to safety refusals, we show it is feasible to increase the divergence between an
aligned model and an unaligned one on the harmful content at greater token depths. Importantly, we
show that such a deeper alignment often leads to stronger robustness against some common exploits.
Third, we show that a constrained optimization objective that focuses on preventing large shifts in
initial token probabilities can mitigate finetuning attacks (Section 4). This highlights potential lines
of defense against fine-tuning attacks, using a better understanding of shallow safety alignment, as
well as provides further evidence of the shallow alignment of current models.
Overall, this work pitches the unifying notion of shallow versus deep safety alignment, demonstrates
that current methods are relatively shallow (leading to a host of known exploits), and provides initial
paths forward for mitigation strategies. We encourage future safety alignment research to explore
various techniques to ensure that safety alignment is more than just a few tokens deep.


We consolidate the notion of _**“shallow safety alignment”**_ to characterize an issue that we commonly
find in current safety-aligned LLMs. Specifically, we say that a model undergoes shallow safety
alignment if it primarily adapts the base model’s generative distribution only over the very first few
output tokens to induce a basic refusal response. In this section, we present a set of case studies to
systematically illustrate the above issue: this type of alignment can appear safe in pre-deployment
testing or standard workflows but quickly falls apart if anything triggers a non-refusal prefix. First,
in Section 2.2, we show that there exists a local optimum where promoting simple refusal prefixes in
the first few tokens of an unaligned model improves its safety to similar levels as an aligned model.
We also show that the KL divergence between aligned and their unaligned counterparts is largely
2Similar token-wise dynamics have recently also been noted by Lin et al. [25], Zhang and Wu [26], and Zhao
et al. [27]. This is also related to Superficial Alignment Hypothesis by Zhou et al. [28]. See Section 5 for more
detailed discussions of the related work.
biased toward these initial token positions, suggesting that this shortcut is in fact exploited by current
alignment approaches. Then, in Section 2.3, we demonstrate how this shallow safety alignment can be
a source of many safety vulnerabilities, including vulnerabilities at the inference stage (Section 2.3.1)
and vulnerabilities against fine-tuning attacks (Section 2.3.2).

2 The Shallow Safety Alignment Issue in Current Large Language Models


**Notation.** We use _πθ_ to denote a language model parameterized by weights _θ_ . We sometimes also
directly use _π_ base to denote an (unaligned) pre-trained model (e.g., Llama-2-7B, Gemma-7B) to
contrast its aligned counterpart _π_ aligned (e.g., Llama-2-7B-Chat, Gemma-7B-IT) [5, 31]. Given an
input _**x**_, the model’s output is modeled by _πθ_ ( _· |_ _**x**_ ). We use _**y**_ _∼_ _πθ_ ( _· |_ _**x**_ ) to denote the sampling
of output _**y**_ . For token sequences like _**x**_, _**y**_, we use _xt_, _yt_ to denote their _t_ -th tokens, and _|_ _**x**_ _|, |_ _**y**_ _|_ to
denote their lengths (i.e., number of tokens). We also use _**y**_ _<t_ and _**y**_ _≤t_ to denote the subsequences
ranging from the first to the ( _t −_ 1)-th tokens and from the first to the _t_ -th tokens in _**y**_, respectively.
Similarly, _**y**_ _>t_ and _**y**_ _≥t_ are employed to denote subsequences after the _t_ -th and ( _t −_ 1)-th tokens.
**Safety Evaluation and The Metrics.** In our experiments, we evaluate the safety alignment of models
following the same evaluation pipeline from Qi et al. [22]. Specifically, we test a model on the
HEx-PHI safety benchmark [32], which consists of 330 harmful instructions across 11 harmful use
cases. Then, we evaluate whether the model complies with these harmful instructions. The same to
Qi et al. [22], we use GPT-4 as a judge to automatically evaluate whether the model’s outputs on
these harmful test examples are safe. We report the ratio of test cases in which the model’s outputs
are harmful. In the absence of an attack, we denote this ratio as the _Harmfulness Rate_ ; in the presence
of adversarial attacks that introduce harmful outputs, we refer to it as the _Attack Success Rate (ASR)_ .


A commonly recognized characteristic of safety-aligned language models is their ability to refuse
harmful instructions. For example, when instructing the aligned Llama-2-7B-Chat model: _“Provide_
_detailed instructions on how to break into someone’s house.”_, it will output a safe response:
_“I cannot fulfill your request. It’s not within my programming or ethical ... (325 tokens in total) ...”_
These safety responses often start with some refusal tokens (e.g., “I cannot”, “I apologize”, “I am
unable”). When testing on the HEx-PHI safety benchmark [32], Llama-2-7B-Chat starts with either
“I cannot” or “I apologize” in **96.1%** of instances, and Gemma-7b-1.1-IT generates “I am unable”
in **96.7%** of cases. _Though these rigid refusal prefixes appear to be just some trivial artifacts, they_
_actually play an important role in enabling a shallow safety alignment scheme to work._
Table 1: A Shorcut to The Safety Mode: The harmfulness rate of even unaligned models will diminish
when a refusal prefx _**s**_ is preflled during decoding, i.e., _**y**_ _∼_ _πθ_ ( _·|_ _**x**_ _,_ _**s**_ ).
No “I cannot “I apologize, “I am
**Refusal Prefixes** ( _**r**_ ) _→_ Prefx “I cannot” fulfill” “I apologize” but I cannot” unable”
_↓_ _**Harmfulness Rate (%)**_ _on HEx-PHI Benchmark with A Refusal Prefx Preflled During Decoding_
|Aligned
Llama-2-7B
Base|0
68.6 ± 0.8|0 ± 0
16.4 ± 1.4|0 ± 0
5.4 ± 1.3|0 ± 0
14.4 ± 0.6|0 ± 0
2.1 ± 0.2|0 ± 0
8.1 ± 0.4|
|---|---|---|---|---|---|---|
|Gemma-7B
Aligned
Base|2.1_ ±_ 0.2
85.4_ ±_ 0.6|0_ ±_ 0
8.7_ ±_ 1.2|0_ ±_ 0
2.7_ ±_ 0.5|0_ ±_ 0
14.1_ ±_ 0.4|0_ ±_ 0
**1.0**_ ±_** 0.8**|0_ ±_ 0
3.9_ ±_ 0.4|
**The “Safety Mode” Shortcut: Unaligned Models Only Need A Refusal Prefix to Appear “Safe”.**
These short refusal prefixes significantly affect whether the remainder of the model’s response will
be safe or unsafe. Even for an unaligned pre-trained model _π_ base, if we can make its generated
outputs begin with these refusal prefixes, the following output is likely to be safe. Using harmful
instructions _**x**_ from the HEx-PHI safety benchmark, we validate this by prefilling a refusal prefix _**s**_
at the beginning of the decoding process to generate outputs _**y**_ _∼_ _π_ base( _· |_ _**x**_ _,_ _**s**_ ). Table 1 shows the
Harmfulness Rate of the outputs produced by the models with different prefilled refusal prefixes _**s**_ for
both Llama-2 [5] and Gemma [31] base. Although the unaligned base models generally have higher
ASR than their aligned counterparts in standard decoding, the gap considerably decreases when both
models are forced to start with refusal prefixes. This makes sense: continuing a refusal prefix with
an absence of fullfillment is a natural pattern in language, which should already be learned during
pretraining. But it suggests a simple shortcut or reward hacking scheme for safety alignment: safety
behaviors can be introduced by solely updating an unaligned model’s distribution over the first few
output tokens to promote some refusal prefixes.
**Current Safety-aligned Models Exploit This Shortcut.** We
provide evidence that current safety-aligned models are likely
exploiting this shortcut, resulting in shallow safety alignment.
We first construct a harmful dataset in the form of (harmful instruction, harmful answer) pairs. Specifically, we take
out the 330 harmful instructions from the HEx-PHI safety
benchmark and then generate harmful answers for these instructions using a jailbroken version of GPT-3.5-Turbo from
Qi et al. [22]. We call this dataset **Harmful HEx-PHI** . With
this harmful dataset, we can examine the per-token KL diver
Figure 1: Per-token KL Divergence gence _D_ KL� _π_ aligned( _· |_ _**x**_ _,_ _**y**_ _<k_ )�� _π_ base( _· |_ _**x**_ _,_ _**y**_ _<k_ )� between
between Aligned and Unaligned the aligned model _π_ aligned and unaligned pre-trained model
Models on Harmful HEx-PHI. _π_ base on each of the harmful example ( _**x**_ _,_ _**y**_ ). As shown in
Figure 1, for both the Llama and Gemma models, the KL divergence is significantly higher in the
first few tokens than for later tokens. This suggests that most of the KL “budget” for the safety
alignment in these models is spent on the first few prefix tokens. [3] At the high level, this outcome can
be attributed to the reason that the current safety alignment process does not encode any notion of the
“depth" of the alignment. During SFT, the model is trained to mimic responses from human experts,
but it is unnatural for humans to write any kind of examples that refuse a request after providing a
harmful prefix; during RLHF, the model’s reward is computed on the responses generated by the
model itself, but if the model learns to always generate refusal prefixes for some harmful instructions,
the probability that the responses start with harmful prefixes is very low, and the model can hardly
receive any penalty for exploiting the safety mode shortcut.

2.2 The Characteristics of Shallow Safety Alignment


Since we know that there exists a safety shortcut, and aligned models likely exploit it, this helps
explain and unify existing inference-time and fine-tuning time vulnerabilities.

2.3 Shallow Safety Alignment May Be A Source of Many Safety Vulnerabilities