A Critical Look At Tokenwise Reward-Guided Text Generation

Thank you for your review and detailed comments. We are glad that you found our analysis enlightening and our experimental results inspiring.

Assumption too strong We acknowledge that for some problems this assumption would not work and we would need additional information e.g. from human annotators or an LLM to generate the partial sequence preferences. We included an experiment in our original submission that demonstrates how PARGS improves baseline machine translation when we have a post-edit dataset which provides a token-level reward (Table 4). Therefore for situations where the assumption is too strong, we can collect additional data, human or LLM, and still apply our algorithm.

Lack of some experiments Based on your suggestion we have included an experiment looking at the performance of PARGS and diversity of the generated text when we vary $\beta$ and $k$. The results are in Appendix F (Tables 13 and 14). In summary, $\beta$ values between 1 and 2 perform well and $k = 10$ performs the best and $k=5$ performs better than $k = 15$. Diversity improves with larger $k$ as expected and does not vary much with $\beta$.

Training costs In order to maintain a reasonable training cost, we sample a subset from the list of all possible partial sequences. The number of training samples is about 2x for summarization and 1.5x for HH-dialogue. We included the exact number in Table 6 of Appendix B. In summary, we keep the training cost comparable to that of the full reward model.

Eq. 2 and eq. 3 We have made the change in Equation 2. Thanks!

Subscripts in Theorem 3 $\pi_{RLHF_{,i}}$ and $\pi_{RLHF_{,i-1}}$ are two distinct policies over prefix sequences of length $i$ and $i − 1$ (as mentioned in line 275).

Lemma 2 We added a mathematical derivation for the proof of Lemma 2 in Appendix A.

Best-of-N sampling We first sample $N$ complete responses from the language model and then rank them using the full-sequence reward model. The sample with the highest reward is chosen.

Thank you for your suggestion! We compared using Best-of-N using terminal reward from the full sequence reward model, terminal reward from the partial sequence reward model and sequence level mean rewards on the TLDR dataset. The results are as follows:

This shows that using partial rewards for Best-of-N is not a good strategy. In contrast our work shows that using partial rewards to guide generation at inference time works well.

Thank you for your quick reply and the follow-up questions.

Too Strong Assumption As a reminder, the no-free lunch theorem says that we can't learn without making some assumption in machine learning. In our paper, we make the assumption clear by explicitly stating that partial sequences are ranked the same as their full extension in the data. All other works also make some assumption, but it is not always clear and it is not always stated. For instance, ARGS assumes that the reward model will generalize from full sequences to partial sequences based on whatever is the choice of architecture of the reward model, which is unclear and arbitrary. It is easy to argue that our assumption is "too strong" because we made it explicit, while previous work may seem better simply because their assumptions are not always clear and may not be stated. Ultimately, assumptions that better reflect reality will lead to better generalization and better results. We evaluate PARGS on three diverse datasets, Reddit TLDR (summarization), Anthropic-HH (harmless and helpful text), and Ultra Feedback (instruction following, truthfulness, honesty, and helpfulness). Ultra Feedback in particular collects prompts from diverse sources such as UltraChat, ShareGPT, Evol-Instruct, TruthfulQA, FalseQA, and FLAN. On all these datasets without collecting any additional data, PARGS shows strong performance, beating RGTG methods and being comparable to offline RLHF such as DPO and PPO, which is very expensive.

In our previous response, we suggested that our technique could be extended by collecting preference annotations for partial sequences when possible since this would mitigate any issue with the assumption itself. Note that this is the case for any technique. Since all assumptions are imperfect in practice, additional data is needed to rectify those imperfections. In fact, this would be needed even more for the techniques that perform worse than PARGS. Naturally, this increases computational costs for all techniques, not just our technique.

Computational overheads

Thanks for pointing this out. We have added Section I to the appendix where we discuss this and also present an ablation on the effect of different subset sizes on PARGS performance. To summarize, we randomly sample a subset of all possible partial sequences. Therefore yes, we outperform baselines using random sampling. We observe from the ablation that 1.5 to 2x of the original dataset size is enough for good performance.

The use of other sampling strategies is an interesting suggestion. We may be able to incorporate active learning similar to [1] to improve our sampling strategy. But this would be an orthogonal direction to our current work.

We hope that your concerns are addressed and you can review your score !

[1] Muldrew, William, et al. "Active Preference Learning for Large Language Models." ICML 2024

Thank you for the response!

To summarize, I do think that this paper proposes a interesting question and theoretically analyzes the weakness of previous works. However, the proposed approach seems too naive by just transforming pairwise preference from $(y_w,y_l)$ to $(y_w^{1:i},y_l^{1:i})$ without any further adaption, especially considering that it introduces additional computational costs and is based on an assumption seems not true practically. I would suggest you could emphasize more on the sampling strategy (or other methods to reduce this computational overheads) in the main paper, which would solve a key challenge of your proposed approach. However, currently, this is only discussed in section I. (though currently the paper version could not be changed, I suggest to revise them in the final version if possible).

Currently, I would like to increase the score of soundness, since the computational costs seems bearable according to your response. For the overall score, I currently maintain it.

Thank you for your positive review and comments! We have updated the paper as you suggested, e.g. ablations and human evaluation. Details below.

Human evaluation We designed a human experiment to compare the full sequence reward model ranking and human ranking of a pair of TLDR summaries at different sequence lengths. We observe that the reward model has a high agreement for full sequences and an arbitrary agreement for partial sequences. This corroborates our theoretical result that a full sequence reward model can give arbitrary rewards to partial sequences. The details are in Table 12 in Appendix Section E (line 904).

Effectiveness of token-level RMs Across all our experiments we demonstrate that a token-level RM (PARGS) is better than sequence-level RM for RGTG.

Prior work [1] has demonstrated that dense rewards can lead to faster convergence and a better local optima for PPO based RLHF. However, it involves making a single call to the reward model upon termination of the sequence and using the attention weights to redistribute the rewards. The policy update is also at a sequence level. If we were to do this token-wise for our (or any conventional) reward model we would need to call it token by token over the length of the sequence leading to much longer training times. Additionally we may require to update the policy several times over the course of a sequence.

High win rate, but not reward score Note that evaluation by a reward model and GPT4 are different metrics. The reward scores for PPO and PARGS are close for the summarization experiment (2.41 vs 2.36) and within the standard error. The win-rate does not tell us the difference in scores of the winner and the loser. It is possible that average scores of PPO responses is higher even though the win-rate is slightly lower.

PPO omitted for HH This is because updating the Llama-2-7b model with PPO required too much compute for our infrastructure to handle. Instead, we compared PARGS with DPO. Note that we did not claim to beat RLHF in our work, instead we provide an online alternative of alignment without the need of fine-tuning the target LLM.

Exact prompts for evaluation We included the exact prompts in Appendix F in the original submission (Appendix H in the updated draft).

References

[1] Chan, Alex James, et al. "Dense Reward for Free in Reinforcement Learning from Human Feedback." ICML 2024

Thank you for your positive review!

Human evaluation We took your suggestion and designed a human evaluation to compare PARGS with ARGS, DPO and CD on the Ultra Feedback dataset. PARGS has a higher win rate compared to the other methods. Details are in Appendix E.

Ablation In our original submission we had an ablation on the impact of $\beta$ on the performance of PARGS. We have updated the section (Appendix F in the updated draft) with the impact of both $\beta$ and $k$ (Table 13), i.e. the number of candidates from the language model in top-k generation. We observe that $\beta = 2$ is the best followed by $\beta = 1.5$ and $k=10$ is the best followed by $k=5$. We also looked at the effect of different $\beta$ and $k$ on diversity (on Table 14). As expected higher $k$ leads to more diversity whereas $\beta$ values do not have a significant effect on generation diversity.

Examples of generated responses Thanks for the suggestion. We have added examples of generated responses on the Ultra Feedback and TLDR test set prompts in Appendix G.

Thank you for your comments. We have updated the Related Work section accordingly and made the claim on Line 351 (previously Line 323) more precise. Below we address your concerns in detail. Please let us know if you have further comments and suggestions!

We acknowledge that there are works in the literature that solve some of the same sub-problems that we solve in this work. As you rightly point out Gedi [1] and PPLM [2] do not modify the language model but instead use a class-conditional discriminator and an attribute classifier respectively to guide the language model at inference time. Moreover, the idea of using partial rewards or a dense reward function for reinforcement learning is not new. Hao et. al (2022) [3] and Lee et. al (2023) [4] are two works to do this when using reinforcement learning for language modeling. However, our underlying problem (and approach) is different from all these works.

The problem under consideration in this paper is the task of aligning language models to datasets of human preferences. This task has become a key component of most modern LLMs. In the literature, RLHF via a Bradley-Terry reward model has emerged as the most prominent and principled approach to LLM alignment. Thus, we are interested in an efficient, inference-time algorithm that is equivalent to RLHF. When assessing related work that use partial rewards or guide the decoding process, an important question is whether this related work can be used for alignment with preference data and more precisely in RLHF. The answer is no and therefore there is still a need for what we proposed in this paper. More precisely, Hao et. al (2022) [3] assume that expert demonstrations are provided instead of preference data and therefore their approach is not applicable for preference alignment. Lee et al. (2023) [4] use a prefix tree to estimate some partial rewards, but do not learn a reward model that could be used for alignment by RLHF. Deng and Raffel (2023) [5] consider preference alignment for a different model than the Bradley-Terry model.

Our paper contributes a technique to train a partial sequence reward model using the Bradley-Terry preference model. We also demonstrate that using a full-sequence reward model for this (e.g. as in [6]) can lead to arbitrary intermediate rewards. Moreover, we derive a clear connection to RLHF.

In the guided decoding literature the closest methods to our work, to the best of our knowledge, are RAD [5] and CD [7]. We have discussed and compared them in our paper. Even though they use partial rewards at inference time to align language models, they have a different training objective for the reward model and do not demonstrate any connection to RLHF.

I hope our contribution is clearer now and you can revise your score.

References

[1] Krause, Ben, et al. "GeDi: Generative Discriminator Guided Sequence Generation." EMNLP Findings 2021.

[2] Dathathri, Sumanth, et al. "Plug and Play Language Models: A Simple Approach to Controlled Text Generation." ICLR 2020

[3] Hao, Yongchang, Yuxin Liu, and Lili Mou. "Teacher forcing recovers reward functions for text generation." NeurIPS 2022

[4] Lee, Youngwon, Jinu Lee, and Seung-won Hwang. "Learning to Rank Generation with Pairwise Partial Rewards." EMNLP Main 2023

[5] Deng, Haikang, and Colin Raffel. "Reward-Augmented Decoding: Efficient Controlled Text Generation With a Unidirectional Reward Model." EMNLP 2023

[6] Khanov, Maxim, Jirayu Burapacheep, and Yixuan Li. "ARGS: Alignment as Reward-Guided Search." ICLR 2024

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

Summary

Our work analyzes the pitfalls of current reward guided text generation (RGTG) and presents PARGS an inference-time RGTG method which is empirically strong, computationally reasonable and shows a clear connection to RLHF. Our method avoids updating the LLM with expensive training such as DPO and PPO. Below we discuss the salient points of the review process.

ReviewerK21J

• The reviewer discussed some prior works in guided decoding and language modeling
• We point out that, although these works solve some of the same sub-problems as our work, they do not apply to learning from preference data or RLHF.
• The guided decoding methods that learn from preference data are already discussed in our paper and we compare against them.
• Nonetheless, we updated the Related Works section to clarify this
• We did not hear back from the reviewer and assume that their concerns are addressed.

Reviewer XMHs

• The reviewer suggested two ablations which would show the superiority of partial sequence reward model over full-sequence rewards beyond RGTG.
• We created a human evaluation for the first ablation (results in Appendix E Table 12).
• For the second experiment we argued that it was computationally not feasible and pointed to supporting evidence in the literature.
• We did not hear back from them.

Reviewer Rbb2

• The reviewer believed that a human evaluation comparing PARGS to other baselines would further strengthen our paper. Moreover, some examples of generated responses would also be useful.
• We added both of these in Appendices E and G respectively
• We also added an ablation on the effect of changing different components of our method (Appendix F) as requested by the reviewer.
• We did not hear back from them.

Reviewer KTXT

• The reviewer requested an ablation on the impact of $\beta$ and $k$ on the performance and diversity of PARGS. We provided this in Appendix F.
• The reviewer asked for implementing a baseline (Best-of-N) differently. We ran experiments with the new implementation, and demonstrated that the original way (standard) was better.
• In the follow-up they asked for an ablation on different sampling sizes for PARGS. They were concerned about the computational costs.
• We provided the ablation in Appendix I and got confirmation that their concerns were addressed.
• Their only outstanding concern is that our approach is "a bit too naive", i.e. training a Bradley Terry reward model on partial sequences may not be practically viable. We contend that we have strong empirical results across the board, avoid the pitfalls of previous methods and can establish a clear connection to RLHF.