Cascade Reward Sampling for Efficient Decoding-Time Alignment

Q6: Lack of datasets.

Thanks for the suggestion! Our paper already included the results on Ultrafeeback in Table 11 (page 22). Our choice of datasets aligns with recent prior work; for example, [11] uses the HH-RLHF dataset only. Also, we have demonstrated the effectiveness of CARDS on a set of RMs [7, 8, 9, 10]. These RMs are trained on a wider set of preference datasets.

To further demonstrate CARDS on more datasets, we add the results on the suggested datasets [12, 13]. CARDS consistently outperforms previous work.

[7] argsearch/llama-7b-rm-float32

[8] weqweasdas/RM-Mistral-7B

[9] Ray2333/reward-model-Mistral-7B-instruct-Unified-Feedback

[10] weqweasdas/hh_rlhf_rm_open_llama_3b

[11] ARGS: Alignment as Reward-Guided Search. ICLR 2024.

[12] PKU-Alignment/BeaverTails

[13] nvidia/HelpSteer

Q7: Eq. 4 is not well defined.

In Eq. 4, $s_{\geq t}$ is the sequence after the $t$-th token. It is sampled from the policy of base LLM $\pi_{LLM}(\cdot | s_{<t})$ given the previously generated prefix. We will modify the notation to $s_{\geq t}\sim \pi_{LLM}(\cdot | s_{<t})$ in the final version for clarity.

Q8: Which reward model did you use for Fig. 2b, 3, 4?

In these figures, both [7, 8] are used. Fig 3 shows reward distributions of [7, 8]. The points in Fig. 2b and Fig. 4 are collected from both [7, 8], demonstrating the relationship between prefix reward and full-sequence reward holds for diverse RMs.

[7] argsearch/llama-7b-rm-float32

[8] weqweasdas/RM-Mistral-7B

Q9: Can this approach be applied to some smaller and weaker reward models?

Yes. Aside from the standard results (7B base model, 7B RM) shown in the main body, we also have results for smaller RM (7B base model, 3B RM) in Appendix B.7 (page 17). CARDS can still achieve promising alignment ratings and efficiency under this setting.

Q10: Can this approach be applied to multi-objective alignment?

Yes. As discussed in L43~L45, CARDS can be applied to the setting of multiple RMs, by aggregating the reward scores from these RMs. If each of these RMs represents a distinct objective, it is the case of multi-objective alignment. We believe this is an interesting direction and plan to explore this extension in future work.

Thank you for your thorough feedback, which has greatly helped improve our work. We have carefully addressed the concerns raised and are happy to provide further clarification if needed.

Q1: The reviewer cannot buy the idea that the reward model is a value function.

We did not claim that the reward model (RM) is a value function. We claim that RM is a good approximation to the value function on semantically complete prefixes. The core insight to understand this is that semantically complete prefixes closely resemble complete sentences and can be comprehended by RMs. This is not a mathematical claim, but an explanation based on empirical findings.

The mentioned [1] primarily focuses on DPO which uses an implicit RM, while our analysis mainly focuses on explicit RM. It is unclear to us how the theoretical results on DPO in [1] would rule out our explanations for explicit RMs.

[1] From r to Q*: Your Language Model is Secretly a Q-Function. 2024.

Q2: There is a new paper showing that "partial reward" may not be correlated to "full reward".

We are grateful that [2] has further verified this conclusion. As demonstrated by [2], semantically complete segmentation induces a higher correlation coefficient than the 1/3 cutoff.

Please note that [2] uses different QA datasets than those used in this paper. This new setting requires re-tuning the uncertainty threshold to get an appropriate segmentation, much like tuning hyperparameters in many ML algorithms. In practice, we find it effective to adjust the threshold so that each response is split into no more than 10 and no fewer than 5 segments. However, [2] sets the uncertainty threshold to be the same as the HH-RLHF dataset. We suspect this is the main reason the correlation rise shown in [2] is not significant enough.

We understand that the analysis of RM properties is still a growing field, and many papers may have varying conclusions. However, there is no gold-standard way to analyze this problem for now. We are open to discussing any detailed setting in studying this problem, and we hope the reviewer can acknowledge the new findings and insights of our paper.

[2] TreeBoN: Enhancing Inference-Time Alignment with Speculative Tree-Search and Best-of-N Sampling, Arxiv 2024.

Q3: It would be better to remove the “rigorous”.

Thanks for the suggestion! We will remove “rigorous” and emphasize that the analysis in this paper is based on empirical findings.

Q4: Lack of baselines.

Besides RAD/ARDS and naive rejection sampling, we have also included more baselines of alignment methods in Table 2, Table 3, and Table 5 (e.g., PPO [4], DPO [5], RAIN [6]). Since the particular problem studied in this paper is alignment rather than general controlled decoding, we think the current set of baselines is already comprehensive.

CD in [3] requires additional training to learn $V_{\theta}$, which introduces significant computational costs. Unlike CD, CARDS and the main baselines considered in this paper do not require any training.

Furthermore, in CD, sampling the next token from $\pi(y)exp(r(y)/\beta)$ requires computing the $V_{\theta}$ over all tokens in the vocabulary. For example, if the vocabulary contains 10,000 tokens, the $V_{\theta}$ needs to be computed 10,000 times parallelly to score each candidate token. The “one call” in [3] contains multiple parallel $V_{\theta}$ computations. We think top-$K$ sampling can reduce the number of $V_{\theta}$ computations to $K$, which means CD is slower than ARGS/RAD. Since CARDS is already faster than ARGS/RAD, we believe CARDS will be faster than CD [3] under the same model size.

[3] Controlled Decoding from Language Models. ICML 2024.

[4] Proximal Policy Optimization Algorithms. 2017.

[5] Direct preference optimization: Your language model is secretly a reward model. NeurIPS 2024.

[6] RAIN: Your Language Models Can Align Themselves without Finetuning. ICLR 2024.

Q5: Equation 11(b) only holds for soft-RL.

Thanks for pointing out this confusion! We will add a remark on the soft-RL setting in the final version.

Q1. I accept the clarification and I believe the claim would still make sense in many cases. As for the paper “from r to Q*”, I think the result can be directly extended to RM as long as the RM is trained to align BT-model, since you can just remove the $\pi_\text{ref}$ term in the proof.

Q2. I have no relationship with this paper, and thus am not certain about the underlying reason. However, I think it is fine for different claims to be presented under different settings, for nlp community. I would suggest that the authors clarify the limitations of their claims and emphasize their empirical successes in the final revision.

Q4. I understand that $V_\theta$ needs additional training. But I want to clarify that its inference cost is not large. Suppose you have a language model, you can directly get $\pi(s\vert x)$ for $\forall s\in\Omega$ at one time, since they are all matrix computation And the computation of $V_\theta(s\vert x)$ is the same as $\pi(s\vert x)$, since they may share the same architecture. Anyway, since the rebuttal phase draws to end, the existing experiments are already acceptable.

Q6. Thank you for your efforts. I appreciate the results.

I am satisfied about the explanation, and I will raise my rating to 6. But I strongly suggest the authors revising their paper:

• remove the claims involving rigorous analysis
• highlight the limitations of their claims as clarified in the rebuttal
• emphasize the empirical study as the main contribution

If the revision is not satisfying, I would reserve the right to decrease the rating.

Q8: What if prefixes are not semantically complete?

The case has been studied in Fig. 2c (page 5). The “Static” results are based on prefixes that are not semantically complete. RMs are no longer accurate when scoring such prefixes. This result verifies our assumption and highlights the effectiveness of our algorithm.

Thank you for your constructive feedback, which has greatly helped strengthen our work. We have carefully addressed the concerns raised and are happy to provide further clarification if there are any additional questions.

Q1: Fig. 1 requires integration from many sections later.

Fig. 1 is a high-level overview of our method, and thus it cannot explicitly show all of the technical details. There are many new designs and findings within our method and it is impossible to put them all together in one figure. We believe Figure 1 effectively conveys the core idea of our approach and the caption explicitly references the corresponding sections that elaborate on specific details. We would greatly appreciate your suggestions for improving Fig. 1 to make it clearer or more informative.

Q2: Section 4.1.1 repeatedly refers to Fig. 2c, while section 4.2.1 does not mention Fig. 2c.

Fig. 2c is to demonstrate that semantically complete segmentation is better than static segmentation in preserving the accuracy of RM scoring. Section 4.1.1 discusses the importance of semantically complete segmentation, which is closely related to Fig. 2c. Section 4.2.1 discusses the technical details of implementing uncertainty-based segmentation, which is a different topic from Fig. 2c. We believe referring to Fig. 2c only in section 4.1.1 is a reasonable structure.

Q3: Intermediate conclusions lack theoretical support.

This paper is primarily empirical-result-driven. We provide comprehensive explanations and demonstrations for the empirical findings in this paper. We believe empirical-result-driven insights and conclusions are also valuable. We are open to discussing the potential theoretical proofs for these explanations, but such mathematical work is beyond the scope of this paper.

Q4: [1] also assumes Lemma 1 is valid; this paper still does not provide convincing proof.

[1] has a stronger assumption than our Lemma 1. [1] assumes that the reward models are accurate on arbitrary prefixes, whereas our Lemma 1 only requires accuracy on semantically complete prefixes. This is grounded in the observation that RMs are typically trained on complete responses, and semantically complete prefixes closely resemble the data that RMs have seen during training. The improved alignment ratings and efficiency compared with [1] also support this claim.

The proof of our Lemma 1 is rigorous, and the assumption is validated through empirical results—neither of which was done in the previous work [1].

[1] ARGS: Alignment as Reward-Guided Search. ICLR 2024.

Q5: Why CARDS leads to significant performance improvements?

The way we use reward models is appropriate (only scoring semantically complete sequences), and thus the reward scores are accurate in our framework. Previous works that use reward models on arbitrary prefixes suffer from inaccurate reward scores.

Fig. 2c shows a good example of the inaccuracy of arbitrary prefixes. Reward scores on fix-length segments (which are not semantically complete) are not accurate.

Q6: "However, while some of the existing decoding-time alignment methods still struggle with the trade-off between alignment and fluency, they all encounter significant efficiency challenges due to auxiliary steps added to their decoding process."

Thanks for pointing out the confusion! The trade-off between alignment and fluency has been discussed by many papers [2, 3, 4], and the “auxiliary steps” refer to reward model scoring. These two statements may not be strongly correlated, and we will separate them into 2 sentences in the final version.

[2] One fish, two fish, but not the whole sea: Alignment reduces language models' conceptual diversity. 2024.

[3] A Roadmap to Pluralistic Alignment. 2024.

[4] Does Alignment Tuning Really Break LLMs' Internal Confidence? BlackboxNLP 2024.

Q7: "Therefore, the above observation also suggests that RMs can be used as value functions on semantically complete prefixes."Can you explain why?

The mentioned observation shows that RMs are also accurate on semantically complete prefixes. This shares the same property as the value function in the discussed cases. Therefore, we conclude that RMs are approximately equivalent to the value function specifically on the semantically complete prefixes. This is not a mathematical claim, but an explanation/analogy based on empirical findings, as explained in our responses to Q3 and Q4.

Q7: "In Lemma 1, I still cannot understand why can we assume "Assuming the reward models are equivalent to value functions when evaluating semantically complete prefixes". Is there any previous literature proof of this?"

As mentioned in our response to Q4, we verify the assumption through empirical results. There is no previous literature proof of this and we are the first to empirically validate it.

Q1: Larger models.

The choice of model sizes follows recent prior work [1], where the largest model is 7B. To evaluate the performance of our method on larger models, we show below that CARDS continues to outperform ARGS on 13B models. The models used are from [2, 3]. Due to the limitation of GPU memory, these are the largest models we can run.

[1] ARGS: Alignment as Reward-Guided Search. ICLR 2024.

[2] meta-llama/Llama-2-13b-chat-hf

[3] miulab/llama2-7b-ultrafeedback-rm

Q2: Slow inference when batch size is large.

The overall inference speed is not compromised for larger batches. According to Table 7 in page 17, batch size 4 has fewer RM calls and its number of LLM calls is below 4 times of batch size 1. Importantly, for larger batches, LLM calls and RM calls are conducted parallelly. Therefore, the overall inference speed for larger batches is essentially faster.

Thank you for your supportive and constructive feedback. We have carefully addressed the concerns raised, adding new results to strengthen our work. If there are any additional questions, we are happy to provide further clarification.

Thank you for the valuable feedback! We have posted responses, and hope that they can address your concerns. We appreciate your time for the discussion and are happy to provide further explanations and/or experiments if you still have remaining questions.

Q1: The reliance on target scores.

The reliance on target reward score $r^\star$ is a general feature of rejection-sampling-based methods like [1]. In practice, $r^\star$ is estimated by a small subset of chosen responses in the preference dataset, and the same strategy is used in [1].

Alternative strategies like parallelly sampling multiple candidates and choosing the best one (segment-level best-of-$N$ searching) do not depend on $r^\star$. However, they require specifying the value of $N$ and face significant efficiency issues, as the number of sampled candidates remains fixed at every step.

We have implemented the segment-level best-of-$N$ searching to compare the efficiency. We adopt the same uncertainty-based segmentation as used in CARDS. $N$ is set to 10, and the experiment is based on Llama 7B and HH-RLHF. The results below show that segment BoN improves over ARGS, which is token-level BoN, but is much slower than CARDS.

[1] Statistical Rejection Sampling Improves Preference Optimization. ICLR 2024.

[2] ARGS: Alignment as Reward-Guided Search. ICLR 2024.

Q2: The experiments are conducted solely on the HH-RLHF dataset.

We do have results of the UltraFeedback dataset in Table 11 (page 22), where the alignment rating and efficiency are consistent with the HH-RLHF results. Our choice of datasets aligns with recent prior work; for example, [2] uses the HH-RLHF dataset only.

To further demonstrate the generalizability of our findings, we also provide the results on two more datasets [4, 5]. CARDS consistently outperforms previous work.

[2] ARGS: Alignment as Reward-Guided Search. ICLR 2024.

[4] PKU-Alignment/BeaverTails

[5] nvidia/HelpSteer

Q3: More baselines such as in-context learning for alignment and best-of-$N$ on full sequences.

We primarily choose baselines following recent previous works like ARGS [2] and RAIN [3]. Notably, we have already included RAIN [3], an in-context learning method, as our baseline.

Previous works also exclude BoN as a baseline due to its high computational cost. For completeness, we now provide results for BoN below. Full-sequence BoN achieves worse alignment ratings and is slower than CARDS.

[2] ARGS: Alignment as Reward-Guided Search. ICLR 2024.

[3] RAIN: Your Language Models Can Align Themselves without Finetuning. ICLR 2024.

Thank you for your supportive and constructive feedback. We have carefully addressed the concerns raised, adding new results and baselines to strengthen our work. If there are any additional questions, we are happy to provide further clarification.

Thank you for the valuable feedback! We have posted responses, and hope that they can address your concerns. We appreciate your time for the discussion and are happy to provide further explanations and/or experiments if you still have remaining questions.