Theoretical guarantees on the best-of-n alignment policy

Thanks for your insightful review and encouraging comments.

My concern is with the experiment in Figure 4, where the authors test language models on benchmarks. The reward is set as the log-likelihood of the reference model, which already correlates with the policies, thus reducing generality. I suggest the authors provide more explanation about this choice, justify this approach, or discuss a more general reward function.

We have now experimented with other tasks and other rewards (e.g., generation length) to showcase that this is not just an issue that pertains to the log-likelihood reward. See https://anonymous.4open.science/r/bon. We would like to reiterate that anytime there is a non-trivial chance of collision between the outcomes in a set of size $n$, the analytical formula for the KL divergence overestimates the true KL.

Thanks for your insightful review and encouraging comments.

This assumption simplifies the theoretical analysis; however, in practical applications, the output space of generative models is typically infinite (as in natural language generation tasks), which may limit the universality of the theory.

We would like to mention that (1) in practice, language models have a finite output tokens and will eventually stop hence practically the number of possible generated sequences is finite (even though it could be very large), which is covered by our theory; (2) this assumption could be relaxed using limit arguments, and we are happy to include a relaxation of it; (3) this result has already been extended to models with potentially continuous outputs (which would not even be countable) by Mroueh (2024) (subsequent to our work) so the extension beyond what is proved here is certainly possible.

Mroueh, Youssef. "Information theoretic guarantees for policy alignment in large language models." arXiv preprint arXiv:2406.05883 (2024).

The experiment did not involve validation with real generative models, resulting in uncertainty regarding the practical application effectiveness of the theoretical results. In particular, the output space of real generative models is typically large and complex, necessitating further verification of the applicability of the theoretical results in these scenarios.

We would like to mention that we already showed a certain case with prompts from AlpacaEval dataset and Gemma2 IT 9B in Figure 4. We have now expanded the scope of experiments with more tasks and rewards in https://anonymous.4open.science/r/bon

This paper does not provide detailed guidance on the specific implementation and optimization of the best-of-n strategy in practical applications.

While we agree with the reviewer that we do not provide an optimization of best-of-n, our results imply that with n~100-1000, win rate against base model will be already saturated without the need for excessively large n which also implies that KL divergence of ~5 is enough to reach a good policy. This implies that RLHF practitioners can aim to keep the KL divergence of their aligned models <10 and achieve good policies, which also significantly helps mitigate reward overoptimization.

Thanks for your insightful review and encouraging comments.

How the equation (34) can be derived?

Eq. (34) is derived by combining Theorem 5.3, i.e., Eq. (30), and Corollary 5.5, i.e., Eq. (33). We will clarify this.

The proposed estimator for the KL divergence has some evidences both theoretically and empirically, but the actual guarantee is still a conjecture. However, this is not a fault of the paper since it is explicitly stated in Conjecture 4.4.

While we acknowledge this weakness of our work. We would like to mention two points:

• We have tested this conjecture in tens of thousands of randomly generated numerical examples with varying support sizes and n and have not found a counter example.
• If we replaced $\epsilon_n$ with $\epsilon_\infty$, then the estimator would lead to an actual upper bound on the KL divergence (Corollary 4.2), which at least suggests that for large $n$ the estimator gives an upper bound.

Thanks for your insightful review and encouraging comments

The real data experiments are limited and do not show large-scale or broad tasks.

The goal of the experiments was to show cases where the could be a gap between the analytical formula and the true KL. We have now experimented with other prompts for machine translation and other rewards, namely output length, and still show that the KL estimates could be loose in more tasks. See https://anonymous.4open.science/r/bon

The variance of the proposed KL estimator in practical, large n sampling not deeply explored. Would you recommend running it repeatedly and averaging over many sequences, or do you envision an online, per-sample approach?

Thanks for this very important question. The proposed estimator is upper bounded by the analytical formula, and hence by log(n). It is also non-negative and lower bounded by 0. Hence, a crude upper bound on its standard deviation is $\log(n)$. Thus, if the estimator is averaged out over $M = O(\log(n) \log\frac{1}{\delta})$ responses the standard deviation could be driven down to below $\delta$. Given that we are generally interested in $n < 1000$, the dependence on $n$ is manageable. Having said that, given each of the M batches has $n$ iid samples (total of $M \times n$ iid samples), we can use a bootstrapped estimator and should be able to remove the dependence on $\log(n)$. We will think more about the variance estimation and include it in the next iteration.

We would like to also mention that the variance of the estimator in Figure 2 is exactly zero given all outcomes are the same likelihood. We have now also computed and plotted the standard deviation of the estimator in Figure 3 in the new PDF. See https://anonymous.4open.science/r/bon.

The theorems assume finite support or at least that the model probabilities for outcomes can be well-defined in a somewhat discrete sense. Real LMs often have large vocabularies ...

The finite support assumption could be relaxed in practice by limit arguments given that PMF is bounded, and we will include this relaxation in the next version. We would like to also mention that subsequent to our work, Mroueh (2024) in fact has proved a variant of Theorem 3.1 under much weaker assumptions that even applies to continuous distributions such as diffusion process.

Mroueh, Youssef. "Information theoretic guarantees for policy alignment in large language models." arXiv preprint arXiv:2406.05883 (2024).

Showing this in a broader study would underscore the conditions under which the new analysis provides a “significant” improvement. A concrete scenario (e.g. in compliance or enterprise alignment) where the old formula’s overestimation of KL meaningfully hampered performance?

We would like to emphasize that the message of this paper is not just about providing a new estimator for KL divergence. We also prove that the existing formula is an upper bound and hence the works that use it still give guarantees on the win rate vs KL divergence of best-of-n.

Because the paper is about “how far the policy can deviate,” it might be worth an explicit mention that tighter estimates can help compliance or risk-management teams ensure that the model does not drift too far from a safe baseline.

Thanks for the suggestion. Will explicitly mention it.

The authors’ theoretical perspective might help design an adaptive best-of-n procedure that halts once an approximate KL threshold is reached, or once a certain minimum reward is achieved.

Best-of-infinity with halting when a certain threshold on reward is reached is akin to rewind-and-repeat which we analyze in Section 6, where we show that the resulting tradeoff between win rate and KL is less favorable compared to best-of-n. However, the resulting tradeoff between win rate and cost (number of decoding trials) is more favorable than best-of-n, which is important in test-time scaling. Hence, it remains to be seen how to best design a method that achieves Pareto optimal tradeoffs between compute, reward, and KL divergence (which captures preservation of capabilities other than what is captured by reward), which might involve combining the rewind-and-repeat with best-of-n for some finite n as the reviewer suggests. We will add a discussion to this end in the concluding remarks as an area for future work.

Have you observed in practice (beyond the toy examples) that blockwise best-of-n greatly outperforms single-step best-of-n in reward–KL tradeoffs, or do you expect diminishing returns?

Blockwise best-of-n has been studied comprehensively by (Mudgal et al., 2024). While the reward vs KL tradeoffs are generally not better than best-of-n (due to the fact that best-of-n is already almost optimal), blockwise best-of-n generally allows to achieve similar reward vs KL tradeoffs with ~10x smaller n, which means a 10x reduction in test-time compute.

Mudgal, Sidharth, et al. "Controlled Decoding from Language Models." ICML, 2024.