BoNBoN Alignment for Large Language Models and the Sweetness of Best-of-n Sampling (2406.00832v3)

Abstract: This paper concerns the problem of aligning samples from LLMs 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.

BoNBoN Alignment for LLMs and the Sweetness of Best-of-n Sampling

In "BoNBoN Alignment for LLMs and the Sweetness of Best-of-n Sampling," Lin Gui, Cristina G^arbacea, and Victor Veitch address the challenge of aligning samples from LLMs to human preferences using a method known as best-of-$n$ (BoN) sampling. The methodology involves drawing $n$ samples, ranking them, and selecting the optimal choice based on human-aligned attributes. The authors explore two core questions: the relationship between BoN sampling and other alignment methods such as RLHF and DPO, and a novel approach to mimic the BoN distribution through fine-tuning.

Key Contributions

The researchers offer several analytical and empirical insights, organized around the following contributions:

1. Comparison and Optimality Analysis: The paper demonstrates that BoN sampling is essentially optimal in terms of the trade-off between win-rate and KL divergence. By embedding BoN distributions and traditional alignment distributions within a common framework, the authors show that BoN sampling provides an almost Pareto-optimal balance between maximizing win-rate and minimizing KL divergence from the base model.
2. BoNBoN Alignment: They introduce BoNBoN Alignment, a method for fine-tuning LLMs to emulate the BoN sampling distribution without the computational overhead of generating multiple samples per inference. This is accomplished by combining supervised fine-tuning (SFT) on BoN samples with a contrastive loss that also utilizes BoN and worst-of-$n$ samples.
3. Experimental Validation: Empirical results from tasks such as single-turn dialogue and text summarization highlight that models fine-tuned using BoNBoN achieve superior performance relative to other alignment methods, including RLHF and contrastive approaches, demonstrating higher win-rates with minimal off-target deviations.

Theoretical Underpinnings

Optimality of BoN Sampling

By embedding the BoN sampling procedure within a class of reward-weighted models, the authors reveal that BoN sampling is nearly optimal for maximizing win-rate against the KL divergence from the base model. They show that the context-conditional win-rate and KL divergence converge to known theoretical limits, confirming BoN's efficiency and efficacy.

Analytical Comparison

For the RLHF and DPO methods, the optimal policy for win-rate given a fixed KL divergence is derived. This theoretical optimal policy aligns closely with the BoN strategy, further validating BoN's superior balance between alignment quality and computational efficiency.

BoNBoN Alignment Technique

The BoNBoN alignment combines two objectives: SFT on BoN samples and a specialized contrastive loss leveraging both BoN and worst-of-$n$ samples. This dual-objective approach significantly improves data efficiency and avoids the pitfalls associated with likelihood ratio maximization alone, such as over-penalizing low-likelihood samples and consequently distorting off-target behavior.

Practical and Theoretical Implications

Empirical Efficiency

The BoNBoN method demonstrates that it is possible to achieve high alignment without requiring the generation of multiple inference samples, significantly reducing computational requirements. Furthermore, BoNBoN alignment does not necessitate extensive hyperparameter tuning, given that the optimal parameters can be analytically derived.

Impact on Alignment Strategies

Practically, these findings imply that BoN sampling, especially when combined with BoNBoN fine-tuning, could set new standards for LLM alignment by providing a more efficient and robust framework. The method's capability to maintain desirable off-target characteristics while optimizing win-rates could prove beneficial for large-scale deployment in diverse NLP tasks.

Future Prospects

Future work might explore the applicability of BoNBoN alignment beyond the specific tasks examined in the paper, extending it to more complex multi-turn dialogues, conditional text generation, and other sophisticated NLP applications. Additionally, integrating BoNBoN alignment with more advanced reinforcement learning methods could further enhance the alignment quality and broaden its application scope.

In summary, this paper provides substantial theoretical and empirical evidence that BoNBoN alignment represents a significant step forward in the alignment of LLMs, combining optimal performance with practical efficiency. By aligning LLMs to mimic the distribution produced by BoN sampling, the authors achieve both high win-rates and minimal off-target drift, paving the way for more effective and efficient use of LLMs in real-world applications.