BOND: Aligning LLMs with Best-of-N Distillation

Thank you for the positive feedback and the constructive comments. We address your concerns below.

As mentioned in the general response, we have included additional ablations of J-BOND and direct comparisons w.r.t concurrent work of (Amini et al. 2024) in Appendix A.5. Quoting from the general response: In short, unlike J-BOND, vBoN requires fixing a regularization strength in advance (via the parameter $N$) and –moverover– considers only backward KL minimization which in (Amini et al. 2024) is estimated using several MC samples. Instead, J-BOND is designed to require only 2 anchor samples, optimizes Jeffreys divergence, and continuously optimize rewards achieving better reward/KL pareto solutions.

Regarding comparisons with BoN, in Figure 7 (Appendix B.1) we compute actual KL distances between BOND policies and the BoN distribution. This can be done leveraging our analytical expression of Theorem 1 and using several MC samples to estimate quantiles. The goal is to show that our BOND objective (and its approximate gradients) are effective in terms of distilling the BoN distribution. When moving to Gemma fine tuning experiments, BoN sampling becomes intractable since it would require $N$ of order of thousands (if not more) to achieve the same rewards of J-BOND and the other RLHF baselines.

To answer your questions:

Q1: The statement is rather qualitative and abstract, in that it assumes perfect convergence of the BOND operator. However, assuming applying BOND to $\pi_{ref}$ exactly outputs the Best-of-$N$($\pi_{ref}$) distribution, then applying BOND again to such a distribution we are turning it into Best-of-$N^2$($\pi_{ref}$), and so on…

Q2,3&4: We thank the reviewer for the questions. Indeed, it is quite interesting to analyze the rates of growth of KL, log quantiles, and rewards of J-BOND in comparison with the different baselines. First we would like to summarize two well-known results about BoN sampling:

1. The KL between BoN sampling and the reference policy is upper bounded by $log(N) - (N-1)/N$, i.e. it grows logarithmically with $N$. [Beirami et al. 2024]
2. The win rate (i.e. $p_\leq(y)$) between BoN sampling and the reference policy is $N/(N+1)$, cf. e.g. [Amini et al. 2024].

The above two points are explainable of the trend curves associated to (iterative) J-BOND: For J-BOND, each step t can be seen as distilling the Best-of-$N^t$ distribution (see Q1). Thus, the KL vs. steps $t$, grows as $O(log(N^t)) = O(t)$, i.e. linearly, while the log quantiles grow sublinearly as $O(N^t/(N^t + 1))$. This is exactly visualized in Figures 3 and 4 and is an indication that, even with the several made approximations, J-BOND metrics are behaving according to theory. In terms of rewards, we believe there is no relationship that can be theoretically drawn, especially because J-BOND optimizes log quantiles and not the rewards directly. The difference between the reward trends of Figure 3 and 4 can simply be attributed to the different tasks (XSum vs. Gemma fine tuning) and reward models.

[Beirami et al. 2024]: Theoretical guarantees on the best-of-n alignment policy

Q5: We believe the Reward/KL tradeoff is quite a meaningful metric when it comes to reward overoptimization. Achieving better rewards without incurring a too high KL ensures the policy is not degenerating into exploitative behaviors of the reward model. This is what happens when, e.g. the tradeoff parameter $\beta$ of standard REINFORCE baselines is set too low. On the contrary, the fact that J-BOND displays a better and stable reward/KL pareto curve is a good indication of reward hacking being mitigated. Note that reward hacking is unavoidable after a certain KL due to the reward model becoming out-of-distribution. Finally, we note that J-BOND displaying such a better reward/KL tradeoff is quite interesting and encouraging, given that only log quantiles (i.e. generations’ rank) are optimized by the algorithm and not the actual rewards.

We hope the above points clarify all of the reviewer’s concerns. We are happy to expand them further if needed.

Thank the authors for addressing my questions. I find the analysis proposed by the authors to be interesting and capable of explaining the observed training curves. Although the authors did not provide a more detailed explanation regarding the reward curve, I agree that this depends on the characteristics of both the RM and the initial policy. Despite not comparing all the baselines I listed, I believe the authors' analysis of (iterative) J-BOND and the alignment between theoretic analysis and experiment phenomenon convinces me of the algorithm's effectiveness and its novelty in comparison to other works on distilling BoN distributions. Therefore, I decide to raise my rating and score of the contribution.

Thank you for reviewing our work and providing valuable feedback. Below, we respond to the main points and questions raised.

Reward shaping can definitely improve the performance of standard RLHF algorithms, as demonstrated in several recent works. Though, we believe BOND provides a more principled alternative compared to possible reshaping approaches as it only relies on generations’ rank. Also, we note that the superior performance of BOND is not only attributed to optimizing a transformed reward (backward KL component), but also to forward KL minimization and by the iterative distillation procedure.

In terms of impact on downstream tasks, we believe J-BOND provides two main benefits: 1) because it does not rely on reward values (but only on generations’ rank), the resulting policy should be less prone to reward optimization and thus of higher quality. 2) because of the iterative procedure, J-BOND fine tuning is more stable than standard baselines (e.g., Figure 4) and can return a whole spectrum of policies (at different KL divergences) that one could evaluate on the downstream task of interest. As mentioned in the general response, in Appendix B.5 we have reported downstream evaluations of the fine tuned policies showing that J-BOND achieves higher quality (measured via side-by-side comparisons) and significantly higher benchmark scores.

We found reward baselines to be of high importance (as typical in RL and RLHF literature). Moreover, we have performed additional ablations (see general response) illustrating the effect of Jeffreys divergence in J-BOND. These are reported in Appendix B.3. The results complement the ablation performed in Appendix B.1 and highlight the fact that both forward and backward KL components are crucial for achieving best reward/KL tradeoffs. Intuitively, the forward KL component makes sure the policy covers all the modes of the BoN distribution, while the backward KL makes sure the policy is not over-dispersed and makes it peaky on certain modes.

Finally, we provide answers to the raised questions:

Would the curves in Figure 4(a) converge to similar spots if you run the algorithms long enough?

In Figure 4(a), the curves with a fixed N (i.e. $N=4,8,16$) are sublinear and thus will asymptote to different reward values, depending on the fixed $N$. Instead, the iterative BOND with $n=2$ is a slower version of the iterative BOND with $n=4$ so it will eventually reach the same reward values. That is, unlike BOND with a fixed $N$, iterative BOND improves the rewards continuously.

Do the authors expect that the conclusions would be similar for much larger models?

Yes, we expect similar conclusions for larger models. Arguably, larger models could be more prone to reward over-optimization while BOND could mitigate this via the robust BoN distillation objective.

pi(x,y) makes it look like a joint distribution; do the authors mean pi(y|x)?

Yes, we mean the conditional distribution. We’ll improve the notation accordingly.

3e-6 means 0.000003; 3e^(-6) means 0.007. Which one are you referring to?

Thanks for spotting the typo, indeed we mean 3e-6.

We hope the above points clarify all of the Reviewer’s concerns. We are happy to expand them further if needed.

Thank you for the positive feedback and constructive comments.

Following your suggestions we have performed additional ablations and evaluations, as detailed in the general response. In particular, we have added:

• Comparisons (both theoretical and experimental) with the concurrent work of Amini et al. (2024),
• Additional ablations demonstrating the impact of Jeffreys divergence (as suggested, with smoother variations of $\beta$) in J-BOND,
• Downstream evaluations of the Gemma 7B fine tuned policies, in terms of side-by-side comparisons and popular benchmark scores.

We refer to the general response for an overview on each of the above points.

We hope these address the main concerns raised, but we are happy to expand them further if needed.

Thank you for your thorough review and positive feedback about our work. We would like to respond to the points raised.

About the KL divergence estimates against the BoN distribution (Figure 7 of Appendix B.1), we agree that 32 MC samples introduce significant estimation variance. However, we remark that these are estimated with 32 fresh samples at each evaluation step. Each of such estimates correspond to a marker in Figure 2. As expected, there is quite some variance but we found the setup sufficiently informative to show clear trends and separations among the different baselines.

To complement our results, we have performed additional ablations as detailed in the general response. In particular, in Appendix B.3 we have ablated the impact of different Jeffreys divergences in the (iterative) J-BOND algorithm. The ablation complement the one of Appendix B.1 demonstrating that mixing backward and forward KL is also beneficial to achieve better reward/KL tradeoffs. We thank the reviewer for providing us with the relevant work (Go et al. 2023) which we will cite in our manuscript. We acknowledge that alternative f-divergences can be used – we see this as an interesting future direction.

Finally, we are happy to follow the Reviewer’s suggestions on making the benefits of our approach clearer in the introduction.

We thank the Reviewer for the detailed review and the positive and constructive comments.

We would like to clarify that BOND and iterative BOND are the two main building blocks of J-BOND, which is the ultimate algorithm we recommend for RLHF fine tuning. This is because BOND, as presented in Section 3 has the following main sources of complexity:

1. requires prescribing in advance to a fixed value of N
2. it does not scale with large N since the forward KL component requires sampling N times from the reference policy)
3. requires several per-prompt MC samples for accurate quantile estimation.

J-BOND overcomes challenges (1) and (2) by employing the iterative BOND approach. In addition, it addresses challenge (3) by the crude quantile approximation based on 2 anchor samples.

The challenges above can also serve to illustrate the differences between J-BOND and the concurrent vBoN approach by (Amini et al. 2024): vBoN does not suffer from (2), since it considers only backward KL divergence, but it crucially suffers from (1) and (3). In particular, (3) makes vBoN not viable for our fine-tuning experiments with conditional generations. As mentioned in the general response, we have clarified such points in the paper and performed additional ablations to compare J-BOND with a scalable version of vBoN with the same crude quantile estimation (see Appendix A.5 in the updated manuscript). Compared to vBoN, J-BOND does not require prescribing a fixed N in advance and displays a better reward/KL tradeoff. This is attributed to its iterative approach and to the additional forward KL minimization component.

Finally, we thank the reviewer for the spotted inconsistencies – we will update the paper accordingly. We are planning to release the J-BOND code as well.