ARGS: Alignment as Reward-Guided Search

We sincerely thank the reviewer for the detailed review, which has helped us strengthen the manuscript. Below we address your comments and suggestions:

Citation format

We appreciate your attention to detail. During the manuscript preparation, we made a conscientious effort to adhere to the proceeding format when citing relevant papers. However, we acknowledge the possibility of oversights during the editing phase. As suggested, we conducted a comprehensive check and revised several citations, including the one you highlighted.

It is important to note that some of the cited papers currently exist solely on arXiv without an associated proceeding venue, primarily due to their recent release, especially for works originating in 2023. For these instances, we have retained the citations in their current form and will update them when they are formally published.

Qualitative results and code link

In Section 3.3, we conducted qualitative evaluations across 300 randomly sampled prompts from the test set of HH-RLHF. Table 2 presents the qualitative evaluation results, measured by the percentage of win-ties of our method over the alternative decoding strategies. During rebuttal, we have updated our draft by incorporating more qualitative examples. For details, we direct the reviewer to Appendix D.

To facilitate further scrutiny of our work, we have included an anonymous link to our code (in the global response to all reviewers). This link will provide you with access to the experiments and implementation details.

Figure 2

That's a great catch. For clarity, we have revised the figure by changing the legend (ARGS-greedy 10 and 40) to a rectangle instead of a line.

Time complexity

We would like to point out that employing small $k$ yields similar performance as large $k$. This has been empirically validated in Section 3.2. Specifically, Figure 2 illustrates that the performance variation between $k=40$ (darker color bars) and $k=10$ (light color bars) is slight, suggesting that a large number of candidates may not be essential for producing aligned generations.

Empirically, we find that the time complexity overhead can be considerably small. For example, using the OPT-2.7b base model, the generation time per response increases only by 1.9 times when comparing ARGS ($k=10$) with conventional greedy decoding. Despite the slight increase in processing time, there was a notable enhancement in reward performance by $\uparrow$ 6.8%, demonstrating the existence of a reasonable tradeoff between reward optimization and inference speed. The gap can be further reduced by employing a smaller reward model, or parallelizing the reward computation across $k$ candidates.

Our discussion on time complexity has also been endorsed by reviewer Mqii, who commented:

"The authors do discuss the extra computation added at inference time and show its feasibility. I believe that such a method is interesting and useful for the literature even with this extra weight at inference. For example, it can be used to iterate over different reward models before running only one finetuning, or used directly if we have a small enough and good reward model."

More baselines

We have added a latest baseline DPO (Direct Preference Optimization) [1] from NeurIPS 2023. The results have been added to Section 4.

Human evaluation

We employ a GPT-4-based evaluation approach to assess the quality of responses. Research, as discussed in [2], has shown that using a GPT-4 proxy aligns with human evaluations in over 80% of cases, providing a scalable method to approximate human preferences. This evaluation method is prevalent in recent literature, as evidenced by studies such as [1, 3].

Due to practical considerations, our Institutional Review Board (IRB) protocol for conducting human evaluations is currently undergoing review at our institution. We fully intend to incorporate human evaluations into our study as soon as our IRB protocol receives approval.

We thank you for the follow-up and for engaging actively in the discussion!

In case it is helpful, we have also included the checkpoints of SFT and reward models. This allows you to perform decoding based on our approach directly. Please check them out in the revised README.

Best,

Authors

We sincerely thank the reviewer for the positive and constructive review. We are very grateful that you find our paper useful for the literature. We address your feedback below!

GPT-4 Evaluation of ARGS and baseline PPO

Following your suggestion, we have conducted an additional GPT-4 evaluation, comparing ARGS with the baseline PPO (on OPT models from Section 4). This evaluation strictly follows the protocol detailed in Section 3.3. The table below illustrates the completed GPT-4 evaluation results, measured by the percentage of win-ties of our method. Overall our method ARGS produces more favorable responses, achieving a win-tie rate of 72.33%.

Low diversity for PPO

That's a great point. In the table below, we include the results of vanilla greedy decoding on the finetuned OPT-1.3b (SFT) which was the base model for the PPO model to validate the low diversity of PPO. The relatively low diversity of the outputs produced by the PPO model can be potentially related to the low diversity of the SFT model, as a result of the KL penalty.

We sincerely thank the reviewer for the detailed and constructive review. We are also grateful that you find our paper novel and useful for the broader community. We address your feedback below!

Evaluation on other tasks

You raised an excellent point. For our current evaluations, we follow the standard and commonly used benchmarks in alignment literature. In particular, HH-RLHF from Anthropic and Stanford Human Preferences (SHP) are among the largest and publicly available datasets for alignment research. These tasks allow us to draw comparisons with existing approaches more easily and directly. As the first study to introduce the decoding-time alignment approach, we believe it's important to understand the potential of our approach more deeply in the existing realm of alignment tasks, before we expand to other tasks.

With that being said, we do agree that evaluating on more complex tasks such as multi-step reasoning would be valuable. During rebuttal, we looked further into this, however, we could not locate large-scale human preference datasets on multi-step reasoning, which hinders the feasibility of training the reward model. We have revised our draft to acknowledge this limitation and remain interested in delving deeper into more complex tasks as part of our future work.

Evaluation metrics

We understand your concern out of the Goodhart's Law :) We share the same and thus have taken the following additional steps to ensure our evaluations are comprehensive:

• We performed an evaluation where the evaluating reward model is different from the one we used in training. In Section 3.4, we conducted experiments where we employ fined-tuned OPT-1.3b and OPT-2.7b for decoding, and OPT-125m as reward models. We evaluate the models on random 1,000 samples of the test set from the Stanford Human Preferences (SHP) dataset, and the average reward is calculated by the OPT-350m reward model. As shown in Figure 3, ARGS consistently outperforms the greedy baseline.
• To address the nuanced aspects of language quality that the standard metrics (diversity, coherence, reward) may not comprehensively capture, we also adopt a GPT-4-based evaluation approach for comparing the quality of responses. Table 2 presents the GPT-4 evaluation results, measured by the percentage of win-ties of our method over the alternative decoding strategies. A higher percentage indicates that our proposed method is more proficient in generating responses that exhibit not only contextual relevance and accuracy but also helpfulness and harmlessness. This observation is consistent with the outcomes of the reward-based evaluation discussed in Section 3.2.

Relation between RM quality and effectiveness of ARGS

Another great point raised!

Indeed, our experimental results in Section 3.4 reveal such a connection between RM quality and the effectiveness of ARGS. In particular, the RM quality can be modulated by the model capacity. We experimented with a smaller capacity model OPT-125M, along with a larger capacity model OPT-350M. The reward modeling accuracy, as well as the ARGS performance (average reward, diversity, coherence), is summarized in the table below. We observe that a higher RM accuracy in general leads to stronger decoding performance by ARGS.

Base model: OPT-1.3b

Base model: OPT-2.7b

This experiment also informs us about a meaningful future direction to explore ARGS decoding from an RM modeling perspective. The papers you recommended are excellent starting points for this. We hypothesize that coupling advanced RM modeling objectives (beyond pairwise ranking loss) with ARGS can help further enhance the generation quality. We have added those discussions in our updated manuscript accordingly.

We sincerely thank the reviewer for the positive and constructive feedback! We are encouraged that you recognize the originality, clarity, and significance of the work. We address your comments below in detail.

Further analysis of using reward model on prefixes

You raise a very insightful question. To better understand this, we analyzed the average reward w.r.t. different locations. Specifically, for a given $`prompt`$, we denote the sequence of generated tokens as $x_1$, $x_2$,.., $x_t$, and so on. The average reward at $t$-th position is calculated as $r([`prompt`,x_{\le t}])$, averaged over the entire test set of HH-RLHF. We use the SFT-ed Llama-7B model for decoding, and report the statistics below for your reference.

As ARGS predicts more tokens, it can be seen that the average reward monotonically increases. This indicates that using the reward model on partial prefixes can steer the search toward generating sequences with higher overall reward scores.

Comparison with DPO

As suggested, we additionally conduct comparisons with the latest approach DPO [1] and report the result in the table below. The comparison and full training configurations have also been added to our manuscript (see Section 4 and Table 8). Our method overall achieves a higher average reward. During experimentation, we notice the tendency for DPO to generate sometimes repetitive words, which potentially leads to low diversity. In contrast, our method more stably generates diverse responses.

[1] Rafailov, Rafael, Archit Sharma, Eric Mitchell, Stefano Ermon, Christopher D. Manning, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. NeurIPS 2023.

Combining multiple reward functions

That's a really interesting suggestion. As you recognized, our approach indeed offers the flexibility to integrate multiple reward signals in decoding-time, without having to re-train the PPO model. In our current exploration, we focus on using a single reward function, which allows us to draw comparisons with existing approaches more easily and directly. As the first study to introduce a decoding-time alignment approach, we believe it's important to understand the potential of our approach more deeply in the existing realm of alignment tasks, before we expand to more complex settings.

With that being said, we do agree that evaluating more complex tasks, such as multi-step reasoning, would be valuable. To do so, one needs more thoughtful and meaningful construction of the tasks, as well as the availability of diverse human preference datasets to train different reward functions. For this reason, we would like to spend more time diving deeper into this aspect in our future work, and hopefully present more conclusive findings after.

Writing suggestions

All fixed - the changes have been marked in red in our updated manuscript. Thank you for the careful read!

I have updated my score from 6 to 8.

For a future version, could you also report the accuracies of the reward model on classifying the winner given prefixes of various lengths vs the full sequence on HH-RLHF? This would be a nice appendix experiment.