Segmenting Text and Learning Their Rewards for Improved RLHF in Language Models

We thank Reviewer JzqS for constructive feedback. Please would you check our general response for our additional experiment results and discussions. Below are our detailed responses to your concerns.

Q1. Use Llama3-8B and “RLHF’ed” model as the backbone model.

A. Thank you so much for this great suggestion!

In General Response Q1 and Q2 (https://openreview.net/forum?id=cK7yrw5g5Q&noteId=2hOU6TzlOB), we present the experiment results per your suggestion. It is clear that our segment-level design again outperforms both the classical bandit approach and the recent token-level design, when using Llama3-8B-SFT or the “RLHF’ed” Phi3-mini-4k-instruct as the backbone model.

Q2. Larger models have solved the challenges that the authors purpose.

A. We respectfully request the reviewer to clarify this concern and provide more information pertaining to which and why larger models have solved the challenges that we identified or purposed.

Without further information, we respectfully disagree with the reviewer and believe that the issue of an oversubtle action space in token-level design and the sparse reward issue in sequence-level/bandit approach are both orthogonal to model size. As can be seen in our additional experiment results in General Response Q1 and Q3, the advantages of our segment-level design over bandit and token-level approaches extend to the larger 8 billion parameter Llama3-8B-SFT backbone. In General Response Q2, we also see that our method can more significantly improve the already “RLHF’ed” Phi3-mini-4k-instruct model, compared to the bandit and token-level designs.

The increasing model sizes and exhausting training corpus call for a more effective utilization of the data. As such, we believe that our method can be a strong ingredient to the post-training pipeline to more effectively train stronger larger models. And our method’s strong efficacy may be particularly helpful when training data are limited, such as RLHF on a specific subdomain.

Dear Reviewer JzqS,

We are deeply grateful for your insightful suggestions!

As the rebuttal period will end soon, we would like to double check if our responses can address your concerns on the effectiveness of our method on Llama3-8B and already “RLHF’ed” models, the practical value of our method in LLMs’ post-training pipeline, or any additional concerns. Please kindly let us know if you have any remaining questions so that we can address them during the remaining discussion period.

If our responses have addressed your concerns, please could you re-evaluate our work based on the updated information?

Dear Reviewer JzqS,

We sincerely appreciate the time and effort you dedicated to reviewing our paper, and the invaluable suggestions that you raised.

For your main concern on the practical value of our method in the post-training/RLHF pipeline, we believe that our main response (https://openreview.net/forum?id=cK7yrw5g5Q&noteId=jxeKXxAB2D), further response (https://openreview.net/forum?id=cK7yrw5g5Q&noteId=9izSittFJ1), and additional experiments (https://openreview.net/forum?id=cK7yrw5g5Q&noteId=2hOU6TzlOB) can effectively clarify your concern.

As a quick summary:

• The goal of our paper is to improve RLHF PPO training by re-considering the granularity of RL action space for LMs.
• In all experiments, we directly compared with the standard RLHF PPO routine (i.e., the bandit formulation) and the recently proposed token-level formulation of RLHF PPO.
• Our gain over these two PPO routines indicates that our method can be directly used as a substitution of the standard PPO routine in the post-training pipeline.
• Furthermore, our strong experiment results on the already “RLHF’ed” backbone model demonstrate our method’s ability of further improving the already DPO’ed/instruction-tuned model.

We understand that you have a busy schedule and are deeply grateful for your suggestions. As the (extended) Author-Reviewer discussion period will end tomorrow, we are eager to know if our responses can adequately address your concerns and if there are any additional questions that you’d like to discuss. We will greatly appreciate the opportunity to hear from you and further discuss our paper if needed

Thank you again for your constructive feedback!

All the best,

Authors

Q5. Guidance on the selection of $f(\cdot)$ in L195.

A. As discussed in L183-184, the construction of $f(\cdot)$ factors in our inductive bias on how sequence-level preference is related to each segment’s quality. Apart from using Soft-maximum as $f(\cdot)$ in our main experiments, in the ablation study (Section 4.3 (e)), we further test on using Summation or Average as $f(\cdot)$, i.e., constructing the parametrized sequence evaluation $e_\phi(x,y)$ in Eq. (3) as the summation or average of segment-level rewards in the text sequence. In conjunction with [3], our general guidance for the selection of $f(\cdot)$ is as follows.

1. Soft-maximum: when chosen responses are distinguished from the rejected ones as containing (a few) key segments (Fig. 2 and L187-188);
2. Summation and Average: when chosen responses are identified as having higher average text quality (L527-528), and Average is independent of text-sequence lengths.
3. [3] also proposes to use Soft-minimum as $f(\cdot)$ when the chosen responses are identified as not containing “bad words.”

Besides, some RL papers [e.g., 4] also propose to design $f(\cdot)$ as a weighted sum of individual rewards with learnable weights, which may be more data-driven at the cost of more training.

[3] Yang, Shentao, et al. "Preference-grounded token-level guidance for language model fine-tuning." Advances in Neural Information Processing Systems 36 (2024).

[4] Kim, Changyeon, et al. "Preference transformer: Modeling human preferences using transformers for rl." arXiv preprint arXiv:2303.00957 (2023).

Q6. Add some preliminary on reward normalizers in PPO training.

A. Thank you so much for this great suggestion! We’ve added a preliminary temporarily in Appendix E of the newly uploaded revision. Please kindly let us know any further suggestions.

Q7. Explain the completion portion $p$ in L223.

A. We first respectfully refer the reviewer to Appendix E of the newly uploaded revision for a discussion on why we would like to construct location-aware reward normalizer functions for PPO training. A quick summary is that: the responses’ have a certain linguistic structure, for example, the beginnings of responses are mostly less informational greeting-alike phases. Those text segments thus typically have lower rewards than “key segments” later on in the response. More generally, the responses’ linguistic structure leads to certain correlation between the mean and std of segment-level reward values and the location of segment in the response, e.g., in the early/middle/later part.

With this observation, we would like to construct location-aware reward normalizer functions that can output proper reward normalizers at an arbitrary location of the response (L69-71). The design of $p=t/T$ in L224, i.e., the index of the segment $a_t$ in the response $y$ divided by the total number of segments in $y$, is a simple indicator of the location of $a_t$ in $y$. We apologize for any confusion pertaining to the name “completion portion,” and have changed it to “location indicator” in L221-236 of the newly uploaded revision.

Note that $p=t/T$ can be easily calculated for responses sampled in PPO training, since in PPO training the responses are first fully sampled, before calculating the (segment-level) rewards and then normalizing the rewards.

We agree with the reviewer that there could be more sophisticated definitions of the location indicator (“completion portion”) $p$, which opens the door for further improving our method.

Q8. The intuition of adding even-split reward interpolation and its potential conflict with segment-level MDP.

A. Thank you so much for this thoughtful question!

In the Appendix F of the newly uploaded revision, we detailedly expand our discussion in Section 4.3 (c) on (1) a plausible reason why no reward interpolation does not work well; and (2) our even-split reward interpolation strategy does not conflict with our segment-level design for LM's RLHF. Rather, we regard it as a useful technique to cope with the per-token KL regularization in RLHF PPO training, which is extraneous to RLHF reward modeling/assignment.

Q9. For RewardBench evaluation, what is the final score from the dense reward models?

A. As discussed in L302-303, for RewardBench evaluation, we compare the parametrized sequence evaluations $e_\phi(x,y)$ (Eq. (4)) of the two candidate responses, and predict the one with higher $e_\phi(x,y)$ as the chosen.

We appreciate Reviewer aY55 for the thorough review. We would like to let you know our additional discussions and experiment results in the general response. Your concerns are addressed in detail below.

Q1. Concern on the compute burden and theoretical guarantee of semantic segmentation.

A. We respectfully note that our semantic segmentation will not significantly increase the compute of PPO-based RLHF training. The key is that in policy learning, we use the SFT model for text segmentation (L211). Recall that the PPO objective Eq. (5) itself requires calculating the auto-regressive probability of each sampled token under the SFT model, i.e., $\pi_{\text{SFT}}(y|x) = \prod_{i=1}^{|y|} \pi_{\text{SFT}}(y_i|x, y_{<i})$. This calculation will already calculate the predictive distributions $\pi_{\text{SFT}}(\cdot |x, y_{<i}), \forall i=1, \ldots, |y|$. As a by-product of this calculation of PPO objective, we can efficiently compute the entropy of each predictive distribution $\mathcal{H}(\pi_{\text{SFT}}(\cdot |x, y_{<i}))$ with very few highly-optimized arithmetic operations: $\log$, multiplication, and summation. These entropies are then directly thresholded for text segmentation. To sum up, since our semantic segmentation procedure re-uses the predictive distributions already calculated in computing the PPO objective, it incurs no significant compute burden to the PPO-based RLHF training.

We agree with the reviewer that our entropy-based semantic segmentation is more heuristic than rigorous. However, as discussed in L157-161, our method is based on the intuitive assumption that tokens within a semantically complete segment can be more predictable by their generation contexts, while the starting token of a new segment is less so. These assumptions have also been adopted by the literature [e.g., 1,2]. Therefore, we believe our method is valid and viable.

Finally, we also believe that the strong performance of our method validates our segment-level action space for RLHF in LMs, and opens the door for further improvement by adopting/designing more sophisticated semantic segmentation methods.

[1] Li, Bolian, et al. "Cascade reward sampling for efficient decoding-time alignment." arXiv preprint arXiv:2406.16306 (2024).

[2] Malinin, Andrey, and Mark Gales. "Predictive uncertainty estimation via prior networks." Advances in neural information processing systems 31 (2018).

Q2. Has PPO converged after one epoch of training?

A. Thank you so much for this careful thought. Please could you check Appendix G of the newly uploaded revision for evidence on the convergence of PPO training in one epoch. We believe this could further validate our main experimental comparisons.

Q3. Some designs lack intuition or analysis.

A. We apologize for any confusion in our design/method description. We will provide detailed reply to your two questions in Q7 & Q8 below.

Overall, as one of the earliest works to (re-)consider the granularity of action space in LM’s RLHF, our paper seeks to provide a simple yet effective pipeline to improve reward modeling and PPO-based LM training, by designing a dense yet semantically meaningful action space. As discussed in the above Q1, our competitive results validate our direction and open the door for further improvements by swapping (some of) our designs for more sophisticated ones.

Q4. Remove “/step” in L45.

A. Thank you so much for this careful review! In the newly uploaded revision, we’ve removed it per your suggestion.

We would like to thank Reviewer aW32 for the careful review. We would love to first refer you to our general response for additional discussions and experiments. The following are our responses to your concerns.

Q1. Confusing layout of Section 2, did not fully understand the relevance of this paragraph to the later experimental setup.

A. We apologize for any confusion in our Section 2. As mentioned in the General Response, in the newly uploaded revision, we fixed the inaccurate descriptions in Section 2.3 pertaining to “completion portion.” We are open to further suggestions on revising our Section 2 and will be sure to implement them in our revised manuscript.

Overall, our Section 2 (Main Method) is meant to be a general method description, while our experimental setup (Section 4.1) provides details on datasets, benchmarks, evaluation metrics, backbone model, and hyperparameters used in our experiments. We are open to any suggestion on how these two sections can be more closely glued together.

Q2. Does the less optimal results in Table 6 undermine the effectiveness of the proposed method?

A. We respectfully note that Table 6 in Appendix B.1 is not our main experiment results/comparisons. As discussed in L415-417, the purpose of this table is to verify that our method, while achieved strong RLHF training, does not suffer from “alignment tax” [1] where one improves RLHF/alignment metrics but deteriorates general language model abilities (OpenLLM metrics in Table 6).

Since the scope of our paper lies in PPO-based RLHF training, and since we train our LM policy on the preference alignment Ultrafeedback dataset, we do not expect to see our method’s optimality on these OpenLLM metrics. Rather, following the community’s common practice, we keep our eye on any significant degradation of these metrics, which does not appear in our Table 6.

As discussed in L356, our main experiment results are presented in Table 1, where our method shows significant improvement in both reward modeling and PPO-based RLHF training, over both the classical bandit reward approach and the recent token-level reward approach. In Fig. 3, we further visualize the more reasonable reward assignment from our segment-level design of dense reward, compared to the recent token-level design.

[1] Askell, Amanda, et al. "A general language assistant as a laboratory for alignment." arXiv preprint arXiv:2112.00861 (2021).

Q3. On which specific tasks does the segment-level design outperform the sentence and token design?

A. We first respectfully refer the reviewer to the above Q2 for pointers to our main experiment results.

Overall, in this paper, we seek to improve the task of RLHF training by proposing a novel design of RL action space as (semantically complete) text segments, based on intuitions from linguistics and avoiding the sparse reward issue. We believe that our experiment results in Table 1 and Fig. 3 can sufficiently demonstrate our method’s advantages on the specific RLHF tasks of both reward modeling and PPO-based LM training, over both the classical bandit approach and the recent design of token-level action space. Further, Table 1 also presents our method’s advantages over the sentence-level action space (coarser than segment but finer than sequence/bandit).

Dear Reviewer aW32,

We deeply appreciate your time and your thoughtful review!

As the rebuttal period is ending, may we respectfully ask if you still have concerns on our method description (Section 2), the specific tasks where our method tackles, the effectiveness of our method, or have any further concerns? We will be more than happy to address your remaining or new concerns during the remaining discussion period.

If our responses have addressed your concerns, would you mind considering re-evaluating our work based on the updated information?

We thank Reviewer KKSu for the comprehensive review! We would like to first bring to your attention our additional discussions and experiment results in the general response. Our responses to your questions are detailed below.

Q1. Concern on the novelty of this paper.

A. We respectfully disagree with the reviewer that our paper lacks novelty.

While densifying the reward signal is currently a popular topic in RLHF research, our paper is among the first to study the important question of “how dense the fine-grained reward should be?” We believe that our study is both original and valuable to the research community. Concretely, as illustrated in Figure 3, the popular design of token-level dense reward can suffer from reward inconsistency over even a word, which is both counter-intuitive and can result in inferior PPO-based LM training than our segment-level design.

We agree with the reviewer that some techniques to implement/train our segment-level reward model have been individually explored, e.g., in [1,2]. We clearly cited relevant literature in Section 2.2 and consider Section 2.2 as a tactful utilization of existing techniques to realize a more reasonable action space in PPO-based RLHF training.

Apart from reward modeling, we would like to kindly note that in Section 2.3, we propose two techniques for effective PPO training under the dense (segment-level) reward: Location-aware Reward Normalizers via Regression and Within-segment Reward Interpolation. As verified in Section 4.3 (b) & (c), these two techniques are critical to the success of PPO training under the denser reward model. We note that none of our direct references in reward model implementation/training [e.g., 1,2] apply their methods to PPO-based LM training, which turns out requiring some special further designs to utilize dense rewards.

Based on the above reasons, we believe that our paper contains original and meaningful contribution to the field of dense-reward RLHF, which is also acknowledged by Reviewer aY55.

[1] Yang, Shentao, et al. "Preference-grounded token-level guidance for language model fine-tuning." Advances in Neural Information Processing Systems 36 (2024).

[2] Li, Bolian, et al. "Cascade reward sampling for efficient decoding-time alignment." arXiv preprint arXiv:2406.16306 (2024).

Q2. Why does segment-level design enhance reward modeling?

A. Thank you so much for this insightful question!

In evaluating our segment-level reward model on RewardBench, we compare the parametrized sequence evaluations $e_\phi(x,y)$ (Eq. (4)) of the two candidate responses, and predict the one with higher $e_\phi(x,y)$ as the chosen (L302-303).

As discussed in L183-184, the construction of $e_\phi(x,y)$ factors in our inductive bias on how preference is generated, i.e., the relation between each segment’s quality and the overall sequence’s preferability. Therefore, our better reward model performance, indicating that our constructed $e_\phi(x,y)$ more accurately distinguishes out better responses, may come from our informed model assumption that good responses contain some “key segments” (key phases) to respond to the given prompt (L187-188); and this assumption turns out adequate on RewardBench. For scientific rigor, other datasets/benchmarks may require a different model assumption and hence different construction of $e_\phi(x,y)$ (L195).

Furthermore, our approach of learning the dense reward $r_\phi$ via re-parametrizing the learned sequence-level evaluation (L197-199) could provide a harder learning problem, and thereby help to squeeze out more information from the reward-training dataset into the reward model. As evidence, we did observe that the training of our segment-level reward model converged comparatively slower than the classical bandit reward model, though still to a close final training loss value.

Q3. Possible to use different pre-trained models as reward functions?

A. We apologize for not fully understanding your questions. We provide our current best answer below.

In General Response Q3, we present the experiment result when using Llama3.2-1B-Instruct for the reward model, and Llama3-8B-SFT for the policy. It is clear that our method can train a much stronger larger policy with a smaller reward model, compared to the classical bandit approach.

In general, by our weak requirement of a unified action space in reward modeling and policy training, the rule-of-thumb for using our method for training a larger policy model with a smaller reward model is: both of them and the SFT model should use the same tokenizer, e.g., all belong to the same model series such as Phi 3.1, Llama 3(.x), etc.