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

Thank you for your review!

• (A1): Since the expectation also accounts for $x$, the left side and the right side are not equivalent. For different prompts, the distribution of rewards varies. This is fundamental: some prompts are easy for the base model (e.g., “what is 2+2?”) and the optimal alignment is to do nothing. Some prompts have responses with highly variable rewards, and the optimal alignment is to strongly favor the good responses. Having a prompt-specific transformation is what lets us take advantage of this structure.

• (A2): Morally, the point here is that choosing $\alpha$ is more akin to choosing the learning rate than to choosing a parameter governing KL vs reward tradeoff. In theory, it should affect the optimization path, but (ignoring finite sample issues) not the final solution. In principle, to find the optimal KL vs reward tradeoff, one would need to fully fit the model at many distinct values. By contrast, all values of $\alpha$ such that the model converges well should yield the same behavior. So, we only need to find one. In practice, we just use a heuristic to set it based on the losses in the first steps of training.

  We also note that this is not the main point of the paper. It does seem like a useful, material advantage. But the core point of the paper—the optimality and achievability of best-of-$n$ alignment—stands irrespective.

• (A3): In principle, one could choose $n$ large enough to observe reward hacking type behavior. However, the win rate increases relatively slowly for $n>10$ (at which point the win rate is already 90%) and the KL divergence increases very slowly. This win rate already suffices to beat existing contrastive approaches with very small KL drift. So, in practice, choosing $n$ around 10 seems to be a sweet spot.

  It is an interesting question how to adapt the procedure in situations where more extreme model modification is desirable. We suspect that an iterated best-of-$n$ procedure using multiple rounds of BonBon would be effective. However, this is a direction for future work.

• For major comment: Thanks for pointing this out! We will change the proof in the updated version.

• For minor comment:

  • 1-2: $y_w$ and $y_l$ are the win and lose response of the prompt $x$. $\sigma$ is the sigmoid function. Their definitions will be added after the equations in the updated version.
  • $r^*$ is $r$ and $D_{BW}$ stands for the (prompt, best response, worst response) dataset. These notations are typos left over from an earlier version. Thanks for pointing them out. We will fix these typos in the revised version.

Thank you for your review!

• 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.
• 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. Additionally, the sense in which BonBon combines an SFT and contrastive objective is fundamentally different. The point in the present paper is that both of these objectives have the same analytic solution, so that the SFT time can be used at contrastive alignment time in the second step.
• With respect to fine tuning on best-of-$n$ samples: a major contribution of the present paper is to show that the best-of-$n$ 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.
• With respect to more minor comments:

  • We adopt the continuous assumption for the following two reasons:

    1. Although response $y$ exists in a high-dimensional space, making the assumption of continuity for $\pi_0$ appear unrealistic, the statistics we consider—KL divergence and win rate—are in a one-dimensional space. When the number of responses is large and each probability is low, the difference between the continuous and discrete cases becomes negligible. This point is well illustrated in Section B of the appendix.
    2. We acknowledge that the continuous assumption may not be realistic for certain prompts, such as in Q&A scenarios. However, many tasks in LLM alignment aim to enhance abstract qualities like helpfulness or harmlessness of responses. In these settings, prompts often generate diverse responses without dominant answers, aligning with the assumptions mentioned in the previous point. Moreover, since the expectation is taken over a set $D$ of prompts $x$, we believe the theory remains relevant if most prompts in the set elicit diverse responses.

    Note also that we find strong agreement with the best-of-$n$ theoretical predictions empirically, and that it is extremely rare in practice for two different responses to be assigned the same reward.

  • Equation 3.1 is correct as written. The $:=$ denotes a definition, as is standard notation.

  • It is not correct to say that equation 3.5 maps to Liu et al. These papers address different problems; they focus on maximizing the expectation of rewards, whereas our focus is on maximizing the win rate. Despite the similarity in the closed forms, they are not the same. More specifically, the difference is due to the fact that the expectation is also taken over $x$. Distributions of rewards $r(x,Y)$, $Y\sim\pi(y\mid x)$ can be different for different prompts.

  • Re. Footnote 2. Assuming continuity makes the KL divergence larger than exploiting the discreteness in the actual responses. So, in this sense, the values we use here (derived under a continuity assumption) are upper bounds on the true discrete distribution. The point of this footnote is simply that in the case where the cardinality of response space is very large---as is typical---the gap between this “upper bound” and the analytic value is small. This is really just noting again that the continuity assumption is reasonable.

  • Line 449: we believe this is correct as written.

Thank you for your review!

First, thanks for all your suggestions about Figure 3! To make it more informative, we will add the theoretical optimal line in the figure and more detailed information in the caption. For the question in the second point, the $n$ in Fig 3 is 8 (we mentioned in Section 5.1 Experimental Setup) and we will add this to the caption. We chose 8 since the win rate of 8 is large enough (nearly 90%) to be favorably competitive with existing contrastive methods. This allows us to compare the win-rate vs KL frontier in the relevant regime. We will also add comparisons for $n=2,\cdots,8$ to the camera ready plots.

For your question: we’ve tried best vs other samples (not the worst one) as pairs for IPO and found that the best vs worst pairs have the most satisfying performance. It would be an interesting direction for future work to find a way to augment the training objective to consume multiple distinct pairs. It is quite possible this could further improve performance. However, we think the simplicity of the BonBon approach is desirable for this paper, since it makes it clear that the advantage is just about best-of-$n$ rather than, e.g., increasing the effective amount of data used for contrastive alignment.

Thank you for your review!

With respect to multiple aspects: we agree this is a fundamental challenge with preference modeling, and a very interesting subject for future research. We note, however, that this problem is essentially fundamental to all post-training procedures. E.g., even explicit reward modeling has to define a way of aggregating multiple distinct kinds of reward. Any aggregation scheme would then induce a preference labeling. And our results would apply to this preference labeling.

That is: the question of how to rank samples and what to do with the rankings are separable questions. The present paper addresses the second problem. Progress on the first is also very interesting, but is out of scope for this paper.