R3HF: Reward Redistribution for Enhancing Reinforcement Learning from Human Feedback

1. The paper seems to be of very limited novelty, as the proposed reward formulation seems to essentially correspond to what is known as potential based reward shaping (see [A]). As a result, the actual contribution of the paper seems trivial: It seems to be limited to the choice of the potential function in the form of the reward model logits of intermediate tokens.
2. In addition, the authors do not clearly motivate why the logit scores of intermediate tokens under the reward model are a good potential function for reward shaping. After all, there is a clear train-test mismatch: SotA reward models are usually only trained by looking at the logits of the last token in the sequence (e.g. an end-of-string-token). It is without motivation why any other logit would yield a reasonable assessment of the quality of the subsequence.
3. As the authors correctly point out, the optimal solution to the (regularized) MDP with potential based reward shaping is the same as the optimal solution to the (regularized) MDP with sequence-level results. This limits the potential of this approach, as we can only hope to get better optimization during finetuning but not a better optimum. The authors do not seem to provide a convincing, principled case why their approach should actually provide better training dynamics.
4. The experimental section is extremely weak as the authors only compare their approach to PPO and DPO. Given that this paper is about improving RL training dynamics (see 3), it should compare to both simple baselines (in particular REINFORCE variants with various obvious choices of sequence and token-level learned and non-learned baselines) as well as SotA RL finetuning algorithms like the inner optimization loop of WARP (see [B]) or BOND (see [C]). In addition, it is well known that hyper parameter tuning is critical for RLHF and there are for example clear interactions between learning rate, regularization, number of steps and final performance, so the authors should add experiments showcasing that their results are robust to the choice of such hyperparameters.

[A] Wiewiora, E. (2003). Potential-based shaping and Q-value initialization are equivalent. Journal of Artificial Intelligence Research, 19, 205–208. [B] WARP: On the Benefits of Weight Averaged Rewarded Policies https://arxiv.org/abs/2406.16768 [C] BOND: Aligning LLMs with Best-of-N Distillation https://arxiv.org/html/2407.14622v1

1. Methods for learning a dense per-step reward without additional fine-grained label data has been clearly proposed in the literature, e.g., [1,2] and the reference therein. The authors ought to have a adequate citation, discussion, and ideally comparison with these prior works. Otherwise, the contribution of this work will be significantly unjustifiable. For example, the central idea in L52-85 has been stated similarly in [2].

[1] Feng, Yihao, et al. "Fantastic Rewards and How to Tame Them: A Case Study on Reward Learning for Task-oriented Dialogue Systems." The Eleventh International Conference on Learning Representations.

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

2. Source code seems not anonymously available for review purpose.

3. Eq. (7) seems confusing. IMHO, starting from the second line, all $\tilde r^{RM}$ should be $\mathcal{R}_\phi$.

This paper does not sufficiently explore existing credit assignment and reward redistribution methods in the context of RLHF.

The proposed method relies heavily on the reward model's accuracy, assuming that the regression model can precisely represent each token's reward value in time-difference calculations. Any inaccuracies in the reward model can directly impact reward redistribution, leading to flawed allocations.

Also this method assumes that each token's contribution can be accurately calculated using a time-difference approach, which may not be feasible for sequence generation tasks. For long sequences, and especially when the model relies on context spanning entire sentences or paragraphs, the method may be inadequate.

• The reward design seems heuristic, and it is probably not fit for some complex reasoning settings.
• I don't see too much advantage from Figure 3.

This work is fundamentally flawed theoretically, practically, and in experiment setup. I believe it is impossible for the authors to adequately addresss my concerns in the rebuttal period and recommend they fundamentally rethink their direction.

Theoretically, RLHF is not a sequential MDP as said by the authors and this invalidates the need for reward redistribution. The authors say that generating N tokens from an LLM can be viewed as N steps in an MDP. But this assumes that: 1. states of the MDP are simply the current sequence 2. Taking an action is equivalent to choosing the next token But 3. there is no transition probability because choosing a next token deterministically adds the token to the current sequence. Single-turn RLHF is more accurately seen as a single step: generating a full sequence of tokens and then getting feedback from a reward model. This is not an MDP but a contextual bandit. If generating a full sequence is a contextual bandit then it makes sense for it to have a single reward, and not need N rewards per token. Their approach could be shown as empirically effective, but it isn't.

Practically, it has been shown on TLDR that the learned reward model is not an effective method of labelling individual tokens with reward (Huang et al, 2024). Trained reward models are explicitly not good at giving reward for non-EOS tokens. Instead, actor-critic methods explicitly learn a value function that does this instead. The authors somehow miss the clear connection that they are trying to create a value function from the reward model. But this doesn't make sense since their baseline method, PPO, already learns a value function! Furthermore, recent works like RLOO (Ahmadian et al, 2023) have argued that value functions are fundamentally not good in the LLM setting. The authors never compare their method to value functions, to monte-carlo estimates of value (like RLOO), and do not have any experiments explaining the qualitative benefits of their approach and why it works in practice.

Experiment setup is poor. The TLDR task they use is not at all equivalent to the original TLDR task (Stiennon et al, 2020) because the SFT dataset is somehow 14,900 examples and not ~117000 examples as it should be. The authors should follow Huang et al (2024) for open-source TLDR experiments. The authors use LLaMA 1 which is out-dated at this point. They also compare against their own SFT baseline as opposed to existing human-written baselines, which is standard practice (again see Huang et al (2024)). RLHF should not be evaluated using point metrics but with curves comparing KL to performance (Gao et al, 2022). Other issues exist but I cannot enumerate them all. Evals should also not use the training reward as this is explicitly trained with (the authors also have GPT-4 win-rates but these results are few). Figure 5 demonstrates clear overfitting of the baselines but not the authors method, simply reducing the learning rate of baselines would probably fix their performance. There are other issues but I cannot enumerate them all.

Overall, the approach is quite weak. The authors would make a much stronger case if they showed their reward redistribution is actually effective on qualitative examples. Furthermore, it makes no sense why the reward model would be effective on non EOS tokens, and perhaps training it to fulfill the requirements of section 3.2 makes more sense.