Distributional Preference Alignment of LLMs via Optimal Transport

Summary

The paper proposes a new technique for preference alignment, named Alignment via Optimal Transport (AOT). The proposed technique supports both paired and unpaired alignment settings. The paper introduces a new viewpoint for preference alignment based on stochastic dominance i.e., making the reward distribution of the positive samples stochastically dominant in the first order on the distribution of negative samples. From this perspective, the paper explains DPO as a special case i.e., a pointwise preference approach with relaxation through the logistic loss. In addition, the paper shows that using convex relaxation is equivalent to minimizing a one-dimensional optimal transport problem. To enhance the differentiability of the objective through the sorting operator (for approximating the continuous optimal transport), the paper uses the SInkhorn-Knopp algorithm instead of the conventional sorting. On the experimental side, AOT leads to state-of-the-art models in the 7B family of models when evaluated with Open LLM Benchmarks and AlpacaEval (using Llama3-70B-Instruct instead of GPT4).

Strengths

* The paper is well-written and easy to follow. * The paper proposes an original preference alignment approach based on stochastic dominance and optimal transport. The connection to optimal transport is interesting and novel. In addition, the paper also provides theoretical results on the sample complexity of the empirical estimation of the objective. * The proposed approach can work for both unpaired and paired alignment settings. * Experiments are extensive on various datasets i.e., UltraFeedback for paired setting, PKU BeaverTails, and HelpSteer for unpaired setting. The paper also provides a free LLM-judge for local Alpaca evaluations. * ATO achieves SOTA on AlpacaEval benchmark and competitive results on other metrics compared to DPO, KPO, and IPO.

Weaknesses

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

Questions

1. It seems that hard sorting is also comparable to soft sorting. Using normal sorting could help to improve the computational speed compared to soft sorting. Should the paper recommend hard sorting as the default variant? 2. How many interactions are used for the Sinkhorn algorithm? What is the choice of the entropic regularization hyperparameter? Do these hyperparameters significantly affect the results? 3. Why does the unpaired setting lead to better results than the paired setting? Do the authors have any explanation?

Rating

8

Confidence

3

Soundness

4

Presentation

4

Contribution

4

Limitations

The authors adequately addressed the limitations.

Summary

The motivation of this paper is that current alignment approaches ensure reward dominance at the sample-level but not on the distributional level. With this in mind, the authors set their goal to design an alignment approach which satisfies First Order Stochastic Dominance (FSD) in rewards for positive examples compared to negative examples. With this goal in mind, they make the following contributions: 1. They write this as a constraint satisfaction problem, which can be relaxed into an unconstrained optimization problem using different surrogates of the 0/1 loss. 2. Using Santambrogio (2015)'s result from optimal transport, they develop a computationally efficient algorithm called ATO. 3. They provide an upper bound for the violation in stochastic dominance of this algorithm using a Rademacher and symmetrization argument. 4. They provide experiments that show that ATO is competitive with existing approaches like DPO, KTO.

Strengths

**Originality** 1. The connection to optimal transport is quite interesting. **Quality** 2. The empirical results are competitive with other alignment approaches. 3. The theoretical results are sound **Clarity** 4. The paper is well-written with background appropriately introduced and clear explanations.

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? 2. There is not a clear improvement in the empirical results over prior approaches, and each approach seems to enjoy success on certain benchmarks -- with the differences between performance quite small in most cases. This invites the question of in what settings ATO would perform better than other approaches in the literature. 3. The theoretical results do not shed much light on when ATO is likely to perform well. Uniform Convergence results are well-known in recent years to not explain generalization behavior of modern ML models. So, while nice, the insight from Theorem 2 is not super obvious to me. Minor: 1. Please call the Santambrogio (2015) result something other than a theorem since it is not a contribution of the present work.

Questions

Since KTO also works with unpaired preference data, a more thorough comparison with KTO is warranted: 1. Is there a unique solution that satisfies the FSD condition? If there are multiple solutions, how does ATO break ties? How does KTO break ties? 2. Is the reason that the behavior is different from KTO that they have different inductive biases and hence converge to different policies? Or the variances of the estimators are different, even though they converge to the same? 3. Is there a way to compare Theorem 2 with the statistical properties of KTO/DPO/RLHF? Or any reason to believe that ATO has a statistical advantage? Despite the weaknesses, I think the connection to optimal transport is interesting and the experiments are sound and comprehensive. I would be happy to adjust my score accordingly if the authors alleviate my concerns through discussion.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

Yes

Summary

This works proposed using Optimal transport in 1D to derive a better alignment guided by the preference data. The main idea is to work with the log likelihood ratio of the marginal distributions of the preference data. The alignment is made through Optimal Transport loss in 1D based on the concept of stochastic dominance between quantiles of two distributions.

Strengths

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

Weaknesses

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

Questions

In Figure 3, it appears that small $\beta$ will produce better results for all method except the IPO. Have the authors tried to test with smaller value of $\beta$, i.e. $\beta = 0.005$? Is there any explanation for this trend in that figure?

Rating

4

Confidence

2

Soundness

3

Presentation

3

Contribution

2

Limitations

It is fine.

Discussion with the authors

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

Rebuttal response

I thank the authors for clarifying my questions, and accordingly increase my score to 6.

Reply to the rebuttal

I would like to thank the authors for their answers. I would like to keep my score unchanged based on the fact is that the empirical results on other metrics are only competitive, although the theory appears to be sufficient.