Distributional Preference Alignment of LLMs via Optimal Transport

We thank the reviewer for their positive review of the paper and for their insightful questions that we address below:

The evaluation is conducted using Llama3-70B-Instruct instead of GPT4 which is not the standard. However, as described in the paper, the usage of Llama3-70B-Instruct leads to approximately the same

Thanks for this comment, we have also reported GPT4 eval on Merlinite, and for the rest we used Llama 70B to lower evaluation costs as mentioned by the reviewer.

It seems that hard sorting is also comparable to soft sorting. Using normal sorting could help to improve the computational speed compared to soft sorting. Should the paper recommend hard sorting as the default variant?

We recommend hard sorting as a default since soft sorting did not lead to substantial improvements. And as noted by the reviewer, hard sorting is also faster.

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

The entropic regularization has to be small ~ 0.1 to ensure accurate sorting otherwise the sorting will be wrong. Note that the soft sorting paper uses PAV (Pool adjacent violators) to compute the soft sorting, which, unlike Sinkhorn, does not need a fixed number of iterations as it reduces finding the soft sorting to an isotonic regression that can be computed in closed form with an algorithm $O(n \log n)$.

Why does the unpaired setting lead to better results than the paired setting? Do the authors have any explanation?

Please note that AOT unpaired outperforms AOT paired on AlpacaEval because it is more robust to noise, if some paired positive and negative answers are noisy, AOT paired will fit to that noise, while AOT unpaired since it does not use this pairing will be more robust to noise. Hence in Figure 2 we see that AOT unpaired outperforms the paired one on AlpacaEval.

We thank the reviewer for their thoughtful comments:

Weaknesses

1) There could be some more motivation for why the FSD condition is practically desirable over the conditional (on x) dominance condition of DPO from Equation (2). Is this more than a theoretical nicety?

Thanks for bringing this up we will make the following discussion more salient:

• It is more expensive to collect paired preference datasets, that have for each prompt a preferred and a rejected answer. There are plenty of unpaired datasets in the sense for each prompt we have either a preferred or a rejected answer. AOT and KTO allow the use of these datasets in alignment.

• AOT unpaired by not relying on pairs of rejected / preferred answers but rather on their distributions, is more robust to noise in the pairing. While DPO would suffer on a pointwise comparison from noise in the pairing, AOT unpaired would not.

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

Please note that the alignment scoring we use is AlpacaEval (instruction following), all other metrics from the Open LLM Leaderboard such as (ARC, MMLU, Winogrande, GSM8K) test that despite the alignment the model retains or improves on the capabilities it has acquired from the pretraining model. To benchmark alignment methods AlpacaEval has been well adopted by the community, and the Open LLM Leaderboard metrics assess whether the alignment deteriorates other capabilities of the LLMs. AOT, while having the highest score on AlpacaEval, is competitive on all other metrics.

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

Please note that Theorem 2 shows that despite AOT can be cast as a min/max game between the LLM and the OT potential, the one dimensional nature of the OT problem does not curse the statistical properties of the overall problem. The generalization bound we provide enjoys the parametric rate $1/\sqrt{n}$. Note that in Figure 3, we see that the generalization error (on the AlpacEval) improves as the sample size grows which echoes findings in Theorem 2. The main insight of Theorem 2 is while AOT has an inner optimal transport problem to solve in order to update the LLM parameters, this does not curse statistically the rate, since the OT problem is one dimensional.

Please call the Santambrogio (2015) result something other than a theorem since it is not a contribution of the present work.

Thank you for the remark, we will take this into consideration in the final version of the paper. We will add the following sentence before this theorem: The following theorem is a restatement of Theorem 2.9 in Santambrogio (2015).

Questions

1) Since KTO also works with unpaired preference data, a more thorough comparison with KTO is warranted: Is there a unique solution that satisfies the FSD condition? If there are multiple solutions, how does ATO break ties? How does KTO break ties?

One can study the uniqueness of the solution in simplified setups where the log-likelihood is a linear model, since the problem becomes convex in the weights, the AOT problem properly regularized with an $L_2$ regularizer over the weight will have a unique solution. Nevertheless, this is not the case for general nonlinear models such as transformers or LoRA adapters when finetuning the LLM, then multiple solutions may exist. Studying the connectivity of the minimas with an alignment loss as it is done in deep learning with classification objectives is an interesting avenue for future research and that is beyond the scope of the paper. We believe if multiple minima exist they will achieve the same soft violation of FSD (that is the loss of AOT), and those solutions will be equivalent and most likely connected.

2) Is the reason that the behavior is different from KTO that they have different inductive biases and hence converge to different policies? Or the variances of the estimators are different, even though they converge to the same?

We believe they have different cost functions and hence they will converge to different policies. While AOT compares quantiles of rewards between positive and negative rewards, KTO compares positive rewards to the average negative reward, and negative rewards to the average positive rewards. KTO will lead to a policy that favors sentences whose reward is above the average of negative rewards in the training set, and AOT will lead to a policy that distributionally prefers positive over negative responses.

3) Is there a way to compare Theorem 2 with the statistical properties of KTO/DPO/RLHF? Or any reason to believe that ATO has a statistical advantage?

Statistically speaking KTO, DPO, and AOT have similar statistical rates. AOT does not suffer statistically from having to solve an OT problem since this is a one dimensional problem, and hence still has a good statistical rate while being a min/max problem between the LLM and the potential of the OT problem (the sorting). For RLHF the statistical properties are more intricate since there is the reward learning and the sampling from the policy in training. The following paper gives a statistical analysis of a variant of RLHF. Different from our analysis this RLHF variant has also regret bounds that scale like $1/\sqrt{T}$.

We thank the reviewer for their encouraging feedback.

My main concern of this paper is the empirical results in experiment section. Table 1 shows that AOT paired/unpaired do not outperform other methods at least in 4 out 7 cases (ARC, MMLU, Winogrande, GSM8K). When they are versus each other, there is no clear winner between AOT paired and AOT unpaired. Meanwhile, I believe that with the AOT paired, when we have more information, the task must be easier, please correct me if I am wrong. I have similar concerns for their performances in Figure 2.

Please note that the alignment scoring we use is AlpacaEval (instruction following), all other metrics from the Open LLM Leaderboard such as (ARC, MMLU, Winogrande, GSM8K) test that despite the alignment the model retains or improves on the capabilities it has acquired from the pretraining model. To benchmark alignment methods AlpacaEval has been well adopted by the community, and the Open LLM Leaderboard metrics assess whether the alignment deteriorates other capabilities of the LLMs. AOT, while having the highest score on AlpacaEval, is competitive on all other metrics.

Please note that AOT unpaired outperforms AOT paired on AlpacaEval because it is more robust to noise, if some paired positive and negative answers are noisy, AOT paired will fit to that noise, while AOT unpaired, since it does not use this pairing, will be more robust to noise. While in AOT paired the task may seem easier, it is less robust. Hence in Figure 2 we see that AOT unpaired outperforms the paired one on AlpacaEval.

$\beta$ will produce better results for all method except the IPO. Have the authors tried to test with smaller value of $\beta$, i.e. $\beta=0.005$

Please note this is an artifact of the formulation of IPO and its implementation in hugging face while for all other methods $\beta$ plays the role of a margin 1/$\beta$ plays the role of the margin in IPO. This explains this trend in Figure 3, IPO peaks on high values of beta and others on lower values, note this same trend has been also observed in the hugging face blog comparing these alignment techniques.