Variational Best-of-N Alignment

The effectiveness of the vBon method is questionable. Although vBon utilizes reward model scoring to depict a target distribution closer to BoN, it requires a large number of samples (controlled by N or M in the article) to generate the corresponding preference data. The efficiency of this method is not high when the sample size is large. There is a lack of more convincing experiments to show the significant enhancement of alignment efficiency.

Thank you for raising this important question. We have added a subsection (5.2) devoted to analyzing the efficiency of vBoN. To summarize, there are 3 aspects affecting the efficiency of our method compared to other alignment methods. (i) vBoN is N times more efficient than BoN at inference time. This can be a huge improvement since BoN is normally used with large N values. (ii) N in vBoN acts as a regularization parameter, and does not increase the optimization time. Since we precompute the $F$ values, vBoN’s optimization loop runs as fast as PPO. (iii) The main computational overhead of vBoN (compared to PPO) is in the preprocessing step when we precompute $F$. In Figure 4 in the revised draft, we look at the performance vs. total elapsed time by increasing M. We observe that with as few as $M = 32$, which only took 10 minutes, we manage to match the performance of larger sample sizes. We hypothesize that the data efficiency of the simple Monte Carlo estimator can be greatly improved by taking into account the similarity between different prompts to learn an approximation to $\log F$ function, which we plan as future work.

The paper is hard to follow. Some definitions lack explanation and need clarification from the authors, such as the function F(.) used in Equations 4 and 5.

We appreciate your feedback and have worked to improve the paper's clarity. Specifically, we revised Sections 2 and 3 and adjusted the notation for better readability. Additionally, we now provide a more detailed explanation of the function $F$: “$F$ can be understood as the strict cumulative density function of reward values under $\pi_{\text{ref}}$. In other words, $F(r(y))$ represents the probability that a random sample drawn from $\pi_{\text{ref}}$ has a reward value less than $r(y)$.”

We are committed to further improving the paper's clarity, so if there are additional areas that require explanation, please do not hesitate to let us know.

This empirical evidence in some sense challenges the theoretical observations. Why, then (and under what conditions), is vBoN better than KL-constrained RL?

Thank you for this thoughtful question. Regarding the relationship between the optimal policy under the KL-constrained RL objective and the BoN distribution: the simplifying assumptions in Yang et al. (2024)—namely, that the language model is memoryless and the reward is linear—are quite unrealistic. Additionally, the two distributions are proven to be asymptotically equivalent only in sequence length. For real-world alignment tasks, as studied in our experiments, these assumptions do not hold, and the distributions are not equivalent.

As for the similarity between the lower bound and the KL-constrained objective, while we agree with the reviewer that their structures are indeed similar, there is a fundamental difference in how rewards are treated. With the KL-constrained RL objective, the goal is to maximize the expected reward. In contrast, with the lower bound, the objective is to maximize the probability that a random sample from $\pi_{\text{ref}}$​ has a reward value lower than that of strings sampled from the aligned policy. This theoretical distinction is what leads to the substantial differences observed in the empirical results.

several important baselines (e.g., DPO and BoNBoN) are only reported for the sentiment control task, and not for the summarization task. Why weren't these baselines included for the (less toy) summarization task?

Thank you for your comment, we have added those baselines to the summarization experiments. In general, the conclusions still hold. We further observe that DPO and BoNBoN can only achieve competitive results in lower temperatures (0.25, 0.5) and their performance drops significantly at higher temperatures.

Questions

The sentence right below equation (6) (the vBoN objective) states that "This is an entropy-regularized objective, where we...discourage the model from having low entropy". And this can be seen clearly in the objective (to be maximized) as well. However, in the text below equation (9), which is a lower bound for the vBoN objective, it is mentioned that L1(\theta) further encourages the model to have low entropy. Can you provide an interpretation of why this L1 lower bound encourages low entropy, while the original vBoN objective encourages high entropy?

We agree with the reviewer that the terms in L1 might not be as intuitive as the original objective. The negative scalar for the entropy term arises from decomposing the objective into a KL divergence term and an entropy term. In L1​, the KL term prevents the model from diverging too far from the reference policy, while the entropy term encourages the exploitation of high-reward regions, creating a trade-off.

Legend labels, as well as axis labels and numbers, are missing from Figure 4, which is one of the more important figures in the paper. Legend labels are also missing from Figures 5 and 6.

Thank you for pointing this out, though we’re unsure why this issue is occurring. The axis labels, numbers, and legends for Figures 4, 5, and 6 are visible on our end. Could this be related to the specific PDF viewer you are using?

The abstract mentions that finetuning with the vBoN objective is "analogous to mean-field variational inference", but this parallel is not discussed in the main text of the paper.

Thank you for your feedback. We have incorporated a sentence in Section 3 to clarify the parallel to mean-field variational inference.

In Section 5, the BoN-SFT baseline is considered, but its performance is not competitive since it has high KL divergence from the reference model. Was a KL-constrained version of BoN-SFT considered?

We did not consider adding a backward KL regularization term to SFT. However, it is worth noting that the reported BoN-SFT is equivalent to minimizing the forward KL divergence between the BoN distribution and the language model. That said, optimizing for both forward and backward divergences (as done in BOND) could potentially improve performance, albeit with increased computational costs.

Like vBoN, both the the BoNBoN and BOND (Sessa et al, 2024) methods attempt to convert BoN to a finetuning-time alignment method. Why was BoNBoN, but not BOND, reported as a baseline in Section 5?

We agree that the BOND paper is relevant to our work. The primary reason we included BoNBoN as a baseline but not BOND is timing. BoNBoN was published a few weeks before we published our first draft, allowing us enough time to implement their method. In contrast, BOND was released after we had published our draft, leaving insufficient time to implement all the technical details, especially since the code for BOND has not yet been made available.

Less Persuasive Experiments

We appreciate your acknowledgment of the substantial costs of conducting experiments on broader tasks. The choice of experiments in our paper was informed by a review of similar work, including DPO, BoNBoN, and BOND. Collectively, these papers evaluate their methods on three main tasks, two of which we have included in this paper. We are also actively assessing the feasibility of adding the third task, Anthropic’s Helpful and Harmless benchmark, as part of our ongoing work.

While we agree that extending vBoN to reasoning tasks would add significant value, our computational resources are currently limited. Nevertheless, we view this as a promising direction for future research and hope that our current results provide a solid foundation for further exploration of vBoN’s potential.

Lacking Efficiency Testing

Thank you for this constructive feedback. We have added a subsection (5.2) devoted to analyzing the efficiency of vBoN. To summarize: (i) N in vBoN acts as a regularization parameter, and does not increase the optimization time. Since we precompute the $F$ values, vBo’s optimization loop runs as fast as PPO. (ii) The main computational overhead of vBoN (compared to PPO) is in the preprocessing step when we precompute $F$. In Figure 4 in the revised draft, we look at the performance vs. total elapsed time by increasing M. We observe that with as few as $M = 32$, which only took 10 minutes on a single A100 GPU, we manage to match the performance of larger sample sizes. We hypothesize that the data efficiency of the simple Monte Carlo estimator can be greatly improved by taking into account the similarity between different prompts to learn an approximation to $\log F$ function, which we plan as future work.

Several key symbols and terms in the formulas are not defined, making it challenging for readers to understand the equations' meaning. A glossary would greatly improve readability.

We have made significant improvements to the clarity of the paper. Specifically, we added the missing explanation of notation in Sections 2 and 3 and included a notation glossary in the appendix.

In the "Summarization" section, the authors refer to the reward models with different names ("pythia-2.8B reward model" and "pythia-6.9B model"), but it’s unclear if these models differ in architecture or purpose (see Section 6).

The smaller model is used within the optimization loop of alignment methods, such as in the PPO fine-tuning process to compute reward values. However, to ensure the robustness of our empirical evaluation against reward hacking, we use a separate, larger reward model exclusively for evaluation. This distinction has been clarified in the revised draft.

The experiments focus solely on summaries from the “relationship” and “relationship advice” subreddits for training and then evaluate performance on in-distribution and out-of-distribution Reddit posts.

relationship and relationship advice are the two big categories in the TL;DR dataset, and together they cover 50% of the dataset. We chose this splitting so that we have a large portion of the dataset that we can meaningfully learn from, but also a held-out dataset with different topics that enable us to test out-of-distribution performance.

This paper provides inadequate detail when introducing the advantages and disadvantages of other alignment methods (such as RLHF and direct preference optimization). In particular, it lacks sufficient background information and details on how these methods are applied to text generation or summarization tasks , which may make it difficult for readers to understand the relative advantages of the methods proposed in this paper.

Thank you for your thoughtful feedback. Section 2 is intended to provide all the necessary background knowledge for understanding prior work and distinguishing vBoN from these methods. However, your comment prompted us to revisit Section 2, and we realized that we had not fully explained how a reward model is constructed for open-ended text generation tasks, such as summarization. To address this, we added clarification in footnote 4:

“For example, in a summarization task, a preference dataset consists of a document, two candidate summaries for that document, and a label indicating which summary is preferred by humans. The reward model is trained on this dataset to maximize the likelihood of correctly predicting human preference.”

We hope this addition provides the necessary context and improves the clarity of our discussion.

Thank you for reading our work and providing feedback.

Section 2 and 3 can be improved significantly by improving the notations and explanations

Thank you for your constructive feedback. We have made significant improvements to the clarity of Sections 2 and 3. Specifically, we added the missing explanation of notation in Section 2 and included a notation glossary in the appendix. Additionally, we expanded the definition of the $F$ function in section 3 and provided further elaboration on the monotonic invariance property. We hope these changes address your comment, and we welcome any further suggestions for improving clarity in these sections.

For example, in Eq 4, F(r(y)) will be defined as F(r(y) ) = P (r(y) < r(y) ), using Eq 5? With this I am not sure how the vBoN objective is insensitive to applying any monotonically increasing function to the reward values?

We noticed that the use of $R$ in Equation 5 caused some confusion, particularly in comparison to Equation 4. To address this, we have revised the notation for greater clarity, and it now reads as $F(r(y)) = \underset{y' \sim \pi_{\text{ref}}}{\mathbb{P}} (r(y’) < r(y))$.

Regarding your question, the function $F$ depends solely on the comparison between two reward values—for example, whether the reward for a string $r(y)$ is greater than that for another string $r(y’)$. As a result, applying a strictly monotonically increasing transformation to all reward values will not affect $F$.

the evaluation chose not to report on standard metrics along with rewards and proximity to reference model comparisons.

Thank you for your comment, though we are unsure if we fully understand your concern. We reported two standard metrics: average reward values and win rates, along with proximity to the reference model, measured using KL divergence. Could you clarify which additional metrics you would like to see?

Why BoNBoN is called with that name? Please explain.

We did not change the name of the algorithm and used the name the authors have chosen for their method.

“We visualize the win rate vs. KL curves in Fig. 6a, and Fig. 6b the average rewards of generations under πθ vs. the KL divergence”. Please mention that these figures are in the appendix.

This was, in fact, a typo and is fixed now. We meant to refer to the figures in the main body of the paper. Thank you for catching this.