BoNBoN Alignment for Large Language Models and the Sweetness of Best-of-n Sampling

Summary

This paper studies aligning samples from LLMs with human preferences using best-of-$n$ (BoN) sampling methods. While BoN methods can yield more desirable outputs without changing off-target behavior, they are computationally expensive since they require $n$ samples from the LLM for every prompt. To address this, the paper proposes BoNBoN alignment, which aims to mimic the sample distribution of BoN without generating $n$ samples. The authors demonstrate that BoNBoN alignment achieves a high win-rate while minimally changing off-target behavior.

Strengths

### Clarity and Quality This paper is generally well-written and easy to understand. The idea itself sounds intuitive and the technical results I checked seem to be correct although some proofs should be modified. ### Significance It is worth noting that I am not specialized both in LLM and generative model. Therefore, while I think the overall claim sounds reasonable, I will not judge how this work is novel and has impact in this field.

Weaknesses

See Questions.

Questions

(Q1) For optimal policy in (3.4), since $Q\_x$ is a CDF of $r(x, Y_0)$, it is a monotonically increasing function, which implies that $$\begin{equation} \arg\max\_{\pi} \mathbb{E}[Q\_x(r(x, y))] = \arg\max\_{\pi} \mathbb{E}[r(x, y)]. \end{equation}$$ Therefore, the optimal policy can be written in a very similar formulation to (2.8). Why do we need to consider $Q\_x$ as in (3.5)? --- While the authors choose $\alpha$ so that the losses from SFT-BON and IPO-BON contribute approximately equally to the total loss, and claim that this is much easier than choosing $\beta$, this remains unclear to me. When training a fine-tuning model, one will observe empirical losses from SFT-BON and IPO-BON for given prompts. The empirical BoNBoN alignment objective can then be computed as an internally dividing point between the objectives of SFT-BON and IPO-BON. (Q2) How to set $\alpha$ before observing the empirical objectives of SFT-BON and IPO-BON? Since $\alpha$ is chosen to balance the terms between them, it seems necessary to know the ratio between them prior to training. --- Although BON does not require choosing $\beta$, the choice of $n$ implicitly plays a similar role as $\beta$, as shown in Theorems 1 and 2. Therefore, a naive question would be: (Q3) how should one choose $n$? When $n$ is too large, the possible answers of BON for similar questions would converge to one specific answer, resembling a point mass. In this scenario, the KL-divergence between BON and the optimal distribution could be large, even though the win-rate increases. This implies that a large $n$ is somewhat equivalent to small $\beta$. --- ### Major comment The proof of Theorem 1 from lines 448 to 449 needs modification. The current arguments using **max** and **min** are mathematically incorrect, as the authors have skipped all coefficients, including negative ones. The arguments should be written in terms of **argmax** and **argmin** or including all coefficient explicitly.

Rating

5

Confidence

2

Soundness

3

Presentation

3

Contribution

2

Limitations

### Minor comments Although this paper is well-written in general, there are some places missing information. 1. What is $y\_w$ and $y\_l$ in (2.6). 2. What is $\sigma$ in (2.7). 3. What is $r^*$ in line 435. 4. What is $D\_{BW}$ in line 472. --- After rebuttal: I have increased my score from 4 to 5 as the authors have addressed my concerns. The reason I am giving a 5 is that I am unable to fully assess the impact of this work within the LLM field.

Summary

The paper describes theoretical results about the Best-of-n sampling procedure in LLM inference. To reduce the computational cost of the procedure, authors develop a novel finetuning method called BoNBoN. Experiments on dialog generation and text summarization show that BoNBoN achieves a higher win-rate for the same KL divergence from the reference model compared to SFT, DPO, and IPO applied on BoN samples.

Strengths

* The paper gives a concise introduction to RLHF and DPO and views the reward-optimal conditional probability as an exponential tilting of the reference model's conditional probability. This motivates the $f_x$-aligned optimal policy. * The paper analytically solves for the optimal policy at a given KL divergence value from the reference model, and shows that the Best-of-n policy approximates it well. * To remove the need for sampling n times, the idea of BoNBoN is proposed. Experimental results show that it attains a better win-rate vs KL divergence tradeoff compared to other approaches.

Weaknesses

* It would interesting to plot the theoretically optimal tradeoff in Theorem 1 on the plot in Figure 3, to better visualize the performance of different methods across a range of different KL values. * More description about Fig 3 would be helpful. What was the value of n used to obtain the BoN sampling operating point? It would be more complete to see different operating points of BoNBoN for different values of n.

Questions

* IPO-BoN is presented as making use of more than just the winning sample (specifically, it uses the best and worst samples in a contrastive objective). It naturally raises the question of using other pairs in the loss as well. Do the authors think that taking more than one pair in the loss could be analyzed in a similar manner, or is it not expected to yield an advantage empirically?

Rating

7

Confidence

3

Soundness

4

Presentation

3

Contribution

3

Limitations

NA

Summary

They claim that best-of-n is an optimal policy with respect to the tradeoff between win rate and the KL divergence. Based on the analysis they propose a strategy to train a model so that it gets to policy similar to the BoN policy.

Strengths

The research question is interesting. BoN and the other learning-based alignment algorithms are currently discussed separately. Understanding the relationship of these algorithms is valuable to the community, if correct.

Weaknesses

Honestly, I couldn’t follow the analysis of the paper, which is probably on my side. Yet, I think several clarifications would be preferable to improve the paper. - BoNBoN alignment is a combination of existing ideas, which is not a reason for rejection. The shortcoming of the paper is that it is unclear which part of the ideas is claimed to be their original as they do not cite the source of the ideas. 1. Using the output of the BoN as the alignment target is a common practice (Pace+ 2024; Liu+ 2023; Gulcehre+ 2023). 2. Mixing the SFT objective with the alignment objective is also proposed in the very first paper of RLHF (Ouyang+ 2022; PPO-ptx) 3. The disadvantage that DPO (or IPO) only controls the ratio of the chosen text/rejected text is commonly resolved by fine-tuning the model on the chosen response first, and then running the preference optimization using the pairs of the responses (Rafailov+2023). The novelty of the proposed algorithm would be clear if the origin of the ideas were clarified. - I failed to follow the analysis of the paper. I would say that it needs some clarification to be understandable for a wide range of audiences. Below is the comment to the paper assuming that I understand the argument of the paper correctly. My concern with the paper is that it is making an unrealistic assumption and getting to the wrong conclusions. I suspect that assuming the functions are continuous does not simplify the argument. For example, if I understand it correctly, Lemma 5 is true only because we assume that y is completely ordered with r(x, y) and r is a one-to-one mapping. The assumption is not made for the sake of simplifying the argument. It is exploited to derive the results that are not valid without the assumption. If an assumption is required then it should be stated so. I would need a better explanation of why the assumption can be true.

Questions

- Why do we assume pi_0 is continuous? It is discrete. There is an analysis of BoN that treats pi as a discrete function (Beirami et al. '24; Theoretical guarantees on the best-of-n alignment policy). What’s the advantage of assuming what is not true over the analysis which treats it as is? - Eq. 3.1. I failed to understand this equation. My guess is that the right-hand side is ignoring the square of pi_0? If so, it should be explained. Is the := saying that we assume the situation where BoN policy can be represented so, or is it saying that the BoN policy defined in the preliminary can be described in this form? - Eq. 3.5.: I believe this formula can be translated as Eq. (2) in (Liu et al. '24; Decoding-time Realignment of Language Models)". Is that correct? - Line 143: “However, given the vast number of possible responses, the assumption is mild in practice.” Why is it true? Even if the domain is infinite, it does not mean that its support is large. One can think of a prompt from a closed QA that a model will most likely output only A or B. - Theorem 2. and footnote 2: What does this Theorem claim? In the footnote, it says “the KL divergence is almost its upper bound log(n) - (n-1)/n”. But in Theorem 2 it says the KL divergence is *exactly* that value. Is it showing the upper bound or the exact value? - Line 449: where the second to last equation is because… → I guess this refers to the fourth to the fifth equation? I guess this Z refers to Z^C_r(x)? - Line 500: The approximation is due to the fact that p_i is small → Why is p_i small? We can think of closed QA tasks where the support of the language model is only A, B, C, and D. In this case square of p_i is not negligible.

Rating

5

Confidence

2

Soundness

2

Presentation

1

Contribution

3

Limitations

I couldn't follow the discussion of the paper to the point where I could evaluate the limitations of the paper.

Summary

This paper addresses aligning samples from large language models (LLMs) with human preferences using best-of-$n$ sampling, which involves drawing $n$ samples, ranking them, and selecting the best one. It tackles two main problems. First, it explores the relationship between best-of-$n$ sampling and Reinforcement Learning from Human Feedback (RLHF) approaches. The authors demonstrate that the best-of-$n$ sampling distribution is essentially equivalent to the RLHF policy when a specific monotone transformation is applied to the reward function. This transformation optimizes the trade-off between win-rate against the base model and KL distance from the base model, making best-of-$n$ a Pareto-optimal solution for win-rate vs. KL distance. Second, the paper introduces BonBon Alignment, a method to fine-tune models to mimic the best-of-$n$ sampling distribution, thus avoiding the need to draw $n$ samples for each inference. Experiments indicate that BonBon Alignment yields models with high win rates while minimally impacting off-target aspects of the generations.

Strengths

1. This paper is well written. The notations are clear and the literature review is sufficient. 2. By justifying that the best-of-n policy is essentially optimal in terms of win rate versus KL-divergence, efficient training for language models based on best-of-n fine tuning is achieved by mimicking the best-of-n sampling distribution. 3. The control of hyperparameter is made simpler with $\alpha$, that balances the loss ingredients.

Weaknesses

When multiple aspects of human preference exist, the proposed method seems to have limited capacity to handle contradictory preference ratings, e.g., helpfulness and harmfulness. Therefore, because the trade-off between diverse aspects of preferences is not explicitly addressed.

Questions

With an only $\alpha$ controling the divergence, it would be interesting to understand how could the proposed Bonbon alignment reflect multiple aspects of real world preferences, e.g., the balance between helpfulness and harmlessness in Anthropic dataset.

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

N/A

Thank you very much for the clarification

> With respect to prior work: could you expand on the set of references you have in mind? E.g., with links or paper titles. It is not clear to us what papers you’re envisioning, or what connections you have in mind. Sorry for the inconvenience. Here is the list of papers I mentioned. Pace+ 2024; West-of-N: Synthetic Preference Generation for Improved Reward Modeling https://arxiv.org/abs/2401.12086 Liu+ 2023; Statistical Rejection Sampling Improves Preference Optimization https://arxiv.org/abs/2309.06657 Gulcehre+ 2023; Reinforced Self-Training (ReST) for Language Modeling https://arxiv.org/abs/2308.08998 Ouyang+ 2022; Training language models to follow instructions with human feedback https://arxiv.org/abs/2203.02155 Beirami et al. '24; Theoretical guarantees on the best-of-n alignment policy https://arxiv.org/abs/2401.01879 > For RLHF (and DPO), it is standard procedure to first run SFT and then, on this SFT’d model, run the RLHF (DPO) procedure. Fundamentally, and most importantly, nothing in the standard RLHF (or DPO) pipeline targets the best-of-n distribution, which is the main point of the present paper. > With respect to fine tuning on best-of-n samples: a major contribution of the present paper is to show that the best-of- heuristic in fact has a remarkably crisp justification. As far as we know, this has not been previously clearly understood. We also note that we experiment with simply fine tuning on best-of-n examples as a baseline and, as we explain at length, we find this works very poorly. The fact that it’s possible to design a contrastive procedure to explicitly target the best-of-n distribution, and that this is much more effective than naive SFT, is an important contribution. I would say that Pace+ 2024; Liu+ 2023; Gulcehre+ 2023 can also be considered as targeting the best-of-n distribution. It would be helpful for the reader if the paper discussed the difference between the proposed method compared to the procedures in these papers as it is not immediate to me. > We adopt the continuous assumption for the following two reasons Thank you very much for the explanation. I think the assumption is valid. It would be helpful for the readers if it is clarified in the paper as the applications of LLMs are not constrained to open-ended text generation and they are also used in closed QA kinds of tasks where the possible output is limited (e.g., using it as a preference model to tell whether an answer A or B is preferred). > Equation 3.1 is correct as written I thought $\pi^{(n)}_{r}(y | x)$ will be $nQ_x(r(x, y))^{n-1} \pi_0(y | x) + \frac{n(n-1)}{2} Q_x(r(x, y))^{n-2} \pi_0(y | x)^2 + \frac{n(n-1)(n-2)}{6} Q_x(r(x, y))^{n-3} \pi_0(y | x)^3 + ... + \pi_0(y | x)^n$ as y can be sampled multiple times and we still get BoN policy to generate y. What am I missing?