@inproceedings{
melnyk2024distributional,
title={Distributional Preference Alignment of {LLM}s via Optimal Transport},
author={Igor Melnyk and Youssef Mroueh and Brian Belgodere and Mattia Rigotti and Apoorva Nitsure and Mikhail Yurochkin and Kristjan Greenewald and Jiri Navratil and Jarret Ross},
booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems},
year={2024},
url={https://openreview.net/forum?id=2LctgfN6Ty}
}
Current LLM alignment techniques use pairwise human preferences at a sample level, and as such, they do not imply an alignment on the distributional level. We propose in this paper Alignment via  Optimal Transport (AOT), a novel method for distributional preference alignment of LLMs. AOT aligns LLMs on unpaired preference data by making the reward distribution of the positive samples stochastically dominant in the first order on the distribution of negative samples. We introduce a convex relaxation of this first-order stochastic dominance and cast it as an optimal transport problem with a smooth and convex cost. Thanks to the one-dimensional nature of the resulting optimal transport problem and the convexity of the cost, it has a closed-form solution via sorting on empirical measures.  We fine-tune LLMs with this AOT objective, which enables alignment by penalizing the violation of the stochastic dominance of the reward distribution of the positive samples on the reward distribution of the negative samples. We analyze the sample complexity of AOT by considering the dual of the OT problem and show that it converges at the parametric rate. Empirically, we show on a diverse set of alignment datasets and LLMs that AOT leads to state-of-the-art models in the 7B family of models when evaluated with Open LLM Benchmarks and AlpacaEval.  Code for $\mathsf{AOT}$ is available in the Hugging Face  TRL library  \url{https://ibm.biz/AOT_TRL}.
Distributional Preference Alignment of LLMs via Optimal Transport
Igor Melnyk Youssef Mroueh Brian Belgodere Mattia Rigotti Apoorva Nitsure Mikhail Yurochkin Kristjan Greenewald Jiri Navratil Jarret Ross