PaperlayerView exact revision on arXiv

Source-linked AI summary

Dense Reward for Free in Reinforcement Learning from Human Feedback

Alex J. Chan, Hao Sun, Samuel Holt, Mihaela van der Schaar

arXiv:2402.00782v1cs.LG

TL;DR

RLHF assigns a completion-level reward after many token-level actions, leaving sparse feedback for a difficult optimization problem. ABC redistributes the reward using the reward model’s attention; the paper reports faster, more stable training and improved local optima, with the original optimal policy preserved.

  • Problem

    RLHF provides reward only at completion end, making credit assignment to individual actions difficult and optimization more complicated than supervised learning.

    Source: §1 · p. 1

  • Method

    ABC uses the reward model’s attention map to distribute its scalar score across the generated tokens, creating denser feedback without extra modeling.

    Source: §1 · p. 2

  • Results

    Across three scenarios, ABC led to faster, more stable training and improved local optima; in single-turn dialogue, its responses were preferred 69.32 ± 1.62% of the time over RLHF responses.

    Source: §1 · p. 2, §5.6 · p. 8

  • Takeaways & Limitations

    The paper reports that ABC is equivalent to potential-based reward shaping, preserving optimality under the original reward.

    Source: §1 · p. 2, §3.2 · p. 4

    Because attention weights are strictly positive, highlighted good and bad parts receive rewards with the same sign; the authors have not fully explored whether ABC exacerbates reward-model over-fitting.

    Source: §6 · p. 8

Abstract

from arXiv · show

Reinforcement Learning from Human Feedback (RLHF) has been credited as the key advance that has allowed Large Language Models (LLMs) to effectively follow instructions and produce useful assistance. Classically, this involves generating completions from the LLM in response to a query before using a separate reward model to assign a score to the full completion. As an auto-regressive process, the LLM has to take many "actions" (selecting individual tokens) and only receives a single, sparse reward at the end of an episode, a setup that is known to be difficult to optimise in traditional reinforcement learning. In this work we leverage the fact that the reward model contains more information than just its scalar output, in particular, it calculates an attention map over tokens as part of the transformer architecture. We use these attention weights to redistribute the reward along the whole completion, effectively densifying the signal and highlighting the most important tokens, all without incurring extra computational cost or requiring any additional modelling. We demonstrate that, theoretically, this approach is equivalent to potential-based reward shaping, ensuring that the optimal policy remains unchanged. Empirically, we show that it stabilises training, accelerates the rate of learning, and, in practical cases, may lead to better local optima.

1. Introduction

RLHF fine-tunes language models using reward feedback, but sparse rewards make optimization and credit assignment difficult. Attention Based Credit (ABC) redistributes the reward model’s score across tokens, and the paper reports faster, more stable training while preserving optimal policies.

1. Introduction

  • RLHF trains a frozen reward model on human preferences, then uses reinforcement learning—commonly PPO—to optimize the language model; assigning credit to actions is one challenge.The paper notes alternatives that avoid reward modeling or use supervised preference optimization.Source: §1 · p. 1
  • ABC uses the reward model’s last-layer attention weights to redistribute its scalar completion score across tokens, giving relevant tokens a larger share.The authors present this as a simple extension to vanilla RLHF that densifies reward without extra modeling.Source: Figure 1 · p. 1, §1 · p. 2
  • The authors report that ABC is equivalent to potential-based reward shaping, so an optimal policy under ABC remains optimal under the original reward.They describe this as avoiding objective mismatch.Source: §1 · p. 2
  • Across three scenarios, including turn-based dialogue, ABC led to faster, more stable training and improved local optima.The authors identify easier credit assignment as a motivation for the approach.Source: §1 · p. 2

2. Language Modelling as Sequential Decision Making

The paper frames autoregressive language generation as sequential decision-making: each next-token choice is an action within a context-window state. In RLHF, the preference-trained reward model scores a completion only at the episode end, leaving individual actions without fine-grained feedback; sparse rewards and the cost of alternatives complicate optimization.

2. Language Modelling as Sequential Decision Making

  • For language generation, the context window is the state, the next-token prediction is the action, and transitions deterministically replace the first masked token.The authors use a sequential decision-making formulation rather than treating the full completion as one action.Source: §2 · p. 2

2.1. Large Language Model Training

  • The reward model is trained from preferences, but its reward is nonzero only at episode end and gives no fine-grained feedback on individual actions.The token-level KL penalty is a regularizer and offers no information for improving the reward model’s return.Source: §2.1 · p. 3

2.2. Complications Optimising Reward

  • Deep RL is difficult to train consistently, and sparse rewards make optimization harder; process supervision can provide finer feedback but requires more detailed human feedback and a changed model.These challenges motivate methods that densify reward without requiring that additional supervision.Source: §2.2 · p. 3
  • Alternatives include Best-of-N sampling, evolutionary supervised learning, DPO, and SLiC-HF, each using a different route to avoid or reduce RL optimization.Best-of-N requires significantly more sampling per response.Source: §2.2 · p. 3

3. Redistributing Rewards with Attention Based Credit

Attention Based Credit (ABC) uses the reward model’s attention to distribute completion-level reward across tokens, aiming to ease credit assignment. The authors show that this shaping preserves optimal policies and describe a controllable blend with the original reward.

3. Redistributing Rewards with Attention Based Credit

  • ABC increases feedback granularity to make reinforcement-learning credit assignment easier, with the aim of faster learning and improved training stability.Source: §3 · pp. 3–4

3.1. Constructing the Reward

  • Averaged last-layer attention weights assign more of the final reward to attended tokens, creating token-level rewards with essentially no extra computation.The reward model’s attention map provides the weights used to distribute the final reward.Source: §3 · p. 4, §3.1 · p. 4

3.2. Preserving the Optimal Policy

  • ABC’s added reward is potential-based shaping, so any policy optimal under the shaped reward is also optimal under the original reward.Source: §3.2 · p. 4

3.3. Practicalities

  • A β-weighted blend controls the attention-based reward, keeps trajectory-total reward constant, and leaves other training-pipeline choices open.Source: §3.3 · pp. 4–5, Figure 2 · p. 5

4. Related Work

The related-work discussion places ABC alongside language-model training approaches and broader reward-redistribution methods. It also describes delayed reward as a persistent challenge in reinforcement learning and control.

4. Related Work

  • The discussion turns from language-model training and reward-optimisation work to research on reinforcement learning more broadly.Source: §4 · p. 5

4.1. Credit Assignment and Reward Redistribution

  • SECRET learns a second reward-sign model for transfer, while ABC avoids the extra surrogate-model complication common in reward redistribution.Source: §4.1 · p. 5

4.2. Delayed Reward

  • Delayed signals remain a persistent challenge; a traditional response is to use a short sliding window of observations, while recent work explores off-policy algorithms.Source: §4.2 · p. 5

5. Experiments

Experiments compare ABC with sparse RLHF and other dense-reward methods across three tasks. ABC improves reward optimization and consistency, with particular robustness for long completions and positive dialogue evaluations.

5.1. Tasks

  • The experiments span positive-sentiment IMDb generation with GPT2, Reddit summarization with GPT-J, and single-turn dialogue with OpenLLaMA models.The dialogue task is described as a proof of concept, not definitive evidence of industrial-scale performance.Source: §5.1 · pp. 5–6

5.3. Improved Optimisation with ABC

  • Across tasks, ABC obtains stronger average reward than vanilla RLHF and Uniform (Table 1).Table 1 reports average training reward with 95% confidence intervals.Source: Table 1 · p. 6, §5.3 · p. 6

5.2. Methods

  • ABC redistributes the fixed episodic reward, while Uniform spreads it evenly and ABC-D uses generator attention; PPO hyperparameters and the KL target remain fixed within each task.ABC-D variants use a weighted average attention map or the map from predicting the final token.Source: §5.2 · p. 6

5.3. Improved Optimisation with ABC

  • On positive generation, ABC reaches vanilla RLHF’s maximum reward in fewer than half its steps and appears to achieve a higher-reward local optimum.Reward increases steadily as the ABC signal’s interpolation strength β increases; ABC also has lower across-seed reward variability than vanilla RLHF.Source: §5.3 · pp. 6–7, Figure 3 · p. 6

5.4. Long Generations Create Sparser Rewards and More Instability

  • As enforced completions lengthen, vanilla RLHF becomes less stable and often diverges, whereas per-timestep reward methods remain stable and can achieve higher mean reward.Successful RLHF runs can perform comparably to shorter-length runs, and the authors note that they did not conduct extensive hyperparameter searches.Source: §5.4 · pp. 7–8, Figure 4 · p. 7

5.5. The Reward-KL Frontier

  • Across almost all KL divergences, ABC achieves higher reward than vanilla RLHF, whose reward–KL trajectory forms a hook as KL rises and then falls during convergence.Figure 5 plots reward against KL divergence; its dashed baseline uses λ = 0.2.Source: Figure 5 · p. 7, §5.5 · p. 8

5.6. Training a Helpful Assistant

  • In single-turn dialogue, ABC optimizes reward faster than RLHF, and AlpacaEval preferred its responses 69.32 ± 1.62% of the time.The authors attribute reduced PPO value loss to lower variance in λ-return estimates.Source: Figure 6 · p. 7, §5.6 · p. 8

6. Limitations or Future Work Opportunities?

The method has practical constraints around tokenizer compatibility, reward-model overfitting, and what positive attention weights can represent.

Shared tokenisers

  • ABC requires shared tokenizers so both models use the same MDP, which makes using publicly available pretrained reward models harder.The paper identifies mapping rewards across different tokenizers as future work.Source: §6 · p. 8

Over-fitting to the reward model

  • Experiments found no greater over-optimization than vanilla RLHF, but the authors have not explored this reward-model overfitting risk fully.The risk may be especially problematic when a model is susceptible to reward hacking.Source: §6 · p. 8

Relying on attention and assuming all contribution is the same sign

  • Positive attention weights give every token reward with the original sign, so highlighted good and bad parts are treated equivalently.The paper says this is partly mitigated when the reward is pushed toward zero; methods such as DeepLIFT could detect positive and negative contributions but require more computation.Source: §6 · p. 8

7. Conclusions

The paper concludes that dense rewards can make RLHF optimization more robust, and that ABC offers a simple way to extract informative token-level rewards. It reports that ABC preserves the original optimal policy while showing faster and more stable optimization in practice.

Conclusions

  • Dense rewards, including uniformly smoothed rewards, can lead to more robust optimization in RLHF.The authors emphasize that sparse rewards are hard to optimize.Source: §7 · p. 8
  • ABC extracts information from the reward model to produce a denser, more informative signal that improves training of RL-optimized generative language models.The method is described as extracting that information very simply.Source: §7 · p. 8
  • ABC can be optimized safely toward the same original optimal policy while exhibiting faster and more stable optimization in practice.The authors describe these as favorable practical characteristics.Source: §7 · p. 8

Impact Statement

The impact statement frames improving RLHF training as a way to improve alignment with human preferences and reduce potential harm. The appendix describes an ABC variant that uses the generator's attention map instead of the reward model's.

Impact Statement

  • Improving RLHF training is presented as a way to improve alignment with human preferences, a key approach to reducing potential harm from deployed models.Source: Impact Statement · p. 9

A.3. Preference Fine-tuning as Inverse RL

  • Inverse reinforcement learning seeks a reward function from demonstrations, corresponding to the reward-learning part of RLHF except that typical IRL setups omit preference data.Source: §A.3 · p. 13

B. ABC-D

  • ABC-Decoder uses the generative decoder's attention map; the authors expect it to work less well than ABC but suggest it may highlight important completion tokens.The proposed rationale is that a capable decoder may attend to sentence structure and semantics, while paying less attention to filler words.Source: B. ABC-D · p. 14
  • ABC-Decoder is described as a way for the language model to learn from its own feedback.Source: B. ABC-D · p. 14

C.1. Model and Dataset Links

The appendix lists reward-model base links for positive generation and summarisation.

C.1.1. POSITIVE GENERATION

  • The positive-generation subsection lists gpt2-imdb as the reward-model base.https://huggingface.co/lvwerra/gpt2-imdbSource: §C.1.1 · p. 14

C.1.2. SUMMARISATION

  • The reward model base is gpt-j-6b, linked at https://huggingface.co/EleutherAI/gpt-j-6b.Source: §C.1.2 · p. 14

C.2. Training Hyper-parameters

PPO hyperparameters are held constant across methods to make comparisons more direct.

  • PPO hyperparameters are held constant across methods to enable more direct comparisons.Source: §C.2 · p. 14

D. Further Experimental Results

In the positive-generation β sweep, mean reward rises as the ABC signal strengthens, while the apparent peak at β = 1 falls outside Proposition 3.1’s guarantee.

D.1. Beta Sweep

  • The potential reward maximum at β = 1 may not be optimal because Proposition 3.1 does not hold there.Source: §D.1 · p. 14
  • Mean reward rises steadily with β; values above 0.1 show a statistically significant increase over RLHF.Figure 7 plots results across ten seeds with ±1 standard deviation, not confidence intervals.Source: §D.1 · p. 14, Figure 7 · p. 15
Loading 2402.00782v1…