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

This paper concerns the problem of aligning samples from large language models to human preferences using best-of-n sampling, where we draw n samples, rank them, and return the best one. We consider two fundamental problems. First: what is the relationship between best-of-n and approaches to alignment that train LLMs to output samples with a high expected reward (e.g., RLHF or DPO)? To answer this, we embed both the best-of-n distribution and the sampling distributions learned by alignment procedures in a common class of tiltings of the base LLM distribution. We then show that, within this class, best-of-n is essentially optimal in terms of the trade-off between win-rate against the base model vs KL distance from the base model. That is, best-of-n is the best choice of alignment distribution if the goal is to maximize win rate. However, best-of-n requires drawing n samples for each inference, a substantial cost. To avoid this, the second problem we consider is how to fine-tune a LLM to mimic the best-of-n sampling distribution. We derive BoNBoN Alignment to achieve this by exploiting the special structure of the best-of-n distribution. Experiments show that BoNBoN alignment yields substantial improvements in producing a model that is preferred to the base policy while minimally affecting off-target aspects.

Code is available at https://github.com/gl-ybnbxb/BoNBoN. Now, the comparison of training-aligned models and BoN is not fully fair. The reason is that producing a BoN sample requires drawing n samples from the base LLM (instead of just one). This is a substantial computational cost. The second main question we address is: if we do in fact want to sample from the BoN distribution, how can we train a LLM to mimic this distribution? If this can be done effectively, then the inference cost of BoN sampling can be avoided.

We answer these questions with the following contributions:

1. We show that the BoN sampling distribution can be embedded in a common class with the distributions produced by training-based alignment methods. Within this common class, we derive the distribution with the best possible trade-off between win-rate against the base model vs KL distance from the base model. Then, we show that the BoN distribution is essentially equal to this Pareto-optimal distribution.

2. We then develop an effective method for training a LLM to mimic the BoN sampling distribution. In essence, the procedure draws best-of-n and worst-of-n samples as training data, and combines these with an objective function we derive by exploiting the analytical form of the BoN distribution. We call this procedure BoNBoN Alignment.

3. Finally, we show empirically that BoNBoN Alignment yields models that achieve high win rates while minimally affecting off-target aspects of the generations, outperforming baselines.