Research map / Reinforcement learning / RL from feedback
Dense Reward for Free in Reinforcement Learning from Human Feedback
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Attention Based Credit (ABC), which redistributes a reward model's scalar output across tokens using its last-layer attention weights, densifying the RLHF reward without extra compute and shown to preserve the optimal policy via potential-based shaping.
Problem
Standard RLHF gives a single sparse reward at the end of an autoregressive completion, which makes credit assignment and PPO optimisation difficult and unstable, especially for long generations.
Approach
Use the reward model's last-layer attention weights (averaged over heads, renormalised over generated tokens) to split the final scalar reward across generated tokens, so each token receives a share proportional to how much the reward prediction attends to it; this is shown to equal potential-based reward shaping and so leaves the optimal policy unchanged.
Key results
- On IMDb positive generation, ABC reaches the maximum reward of vanilla RLHF in fewer than half the training steps and achieves 8.32±0.16 average reward vs 7.21±0.39 for RLHF.
- On single-turn dialogue, ABC responses are preferred over RLHF responses 69.32±1.62% of the time under AlpacaEval 2.0.
- On summarisation, ABC averages 9.18±0.17 reward vs 9.03±0.13 for RLHF and Uniform 9.03±0.15.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Redistribute the completion reward using reward-model attention instead of sparse terminal reward | RLHF on Positive Generation · Positive Generation | average reward: 7.21±0.39 -> 8.32±0.16 | Helped |
| Redistribute the completion reward using reward-model attention instead of sparse terminal reward | RLHF on Summarisation · Summarisation | average reward: 9.03±0.13 -> 9.18±0.17 | Helped |
| Redistribute the completion reward using reward-model attention instead of sparse terminal reward | RLHF on Single-turn Dialogue · Single-turn Dialogue | average reward: 5.78±0.14 -> 6.55±0.04 | Helped |
| Redistribute reward using reward-model attention rather than uniform smoothing | Reward redistribution on Positive Generation · Positive Generation | average reward: 6.60±0.19 -> 8.32±0.16 | Helped |
| Redistribute reward using reward-model attention rather than uniform smoothing | Reward redistribution on Summarisation · Summarisation | average reward: 9.03±0.15 -> 9.18±0.17 | Helped |
| Redistribute reward using reward-model attention rather than uniform smoothing | Reward redistribution on Single-turn Dialogue · Single-turn Dialogue | average reward: 6.44±0.06 -> 6.55±0.04 | Helped |
| Use generator-model attention instead of reward-model attention (ABC-D) (ablation) | ABC on Positive Generation · Positive Generation | average reward | Mixed |
| Use generator-model attention while predicting the final token (ABC-D2) (ablation) | ABC-D on Positive Generation · Positive Generation | average reward | Mixed |
| Increase the ABC reward mixture weight β (ablation) | ABC on Positive Generation · Positive Generation | mean reward | Helped |
| Use ABC instead of RLHF | Dialogue policy evaluated on AlpacaEval 2.0 · AlpacaEval 2.0 | ABC response preference rate: 69.32 ± 1.62% | Helped |
| Use token-level dense rewards instead of sparse terminal reward | RLHF on longer Positive Generation completions · Positive Generation, enforced lengths 20-30, 40-60, 90-110, 140-160 | mean reward stability | Helped |
Surprising findings
- Vanilla RLHF diverges on long generations, while all dense-reward methods remain stable; Uniform reward smoothing alone helps for long completions.
- Even with beta=1 (pure ABC reward, no sparse component), the method yields strong empirical results despite the theory requiring beta != 1 for strict proof conditions.
Limitations
- Requires the reward model and generator to share a tokeniser.
- Attention weights are strictly positive, so positive and negative contributions within a completion cannot be distinguished.
- Possible over-fitting to the reward model; not thoroughly explored; experiments use relatively small models and limited hyperparameter search.
Experimental setup
Datasets: IMDb TL;DR summarize_from_feedback Anthropic HH-RLHF (helpful split) RedPajama
Benchmarks: IMDb positive generation reward TL;DR summarisation reward Anthropic HH single-turn dialogue reward AlpacaEval 2.0
Metrics: average reward reward vs KL frontier PPO value loss win rate
Baselines: RLHF (sparse reward) Uniform reward smoothing ABC-D (decoder attention)
Builds on: PPO potential-based reward shaping RLHF TRL
Closest papers by meaning
- MA-RLHF: Reinforcement Learning from Human Feedback with Macro Actions (ICLR 2025 · similarity 0.91)
- Sequence to Sequence Reward Modeling: Improving RLHF by Language Feedback (AAAI 2025 · similarity 0.91)
- Exploring Data Scaling Trends and Effects in Reinforcement Learning from Human Feedback (NeurIPS 2025 · similarity 0.91)
- Text2Reward: Reward Shaping with Language Models for Reinforcement Learning (ICLR 2024 · similarity 0.90)
- GRAM: A Generative Foundation Reward Model for Reward Generalization (ICML 2025 · similarity 0.90)
- Provably Efficient Online RLHF with One-Pass Reward Modeling (NeurIPS 2025 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.