Selective Timestep Weighting and Advantage-Based Replay for Sample-Efficient Diffusion RLHF

2026-07-08

AI Summary

Diffusion models in reinforcement learning from human feedback (RLHF) suffer from poor sample efficiency due to the need for extensive reward evaluations, limiting their practical applicability. This work identifies that reward information is unevenly distributed across diffusion trajectories and introduces a novel timestep-weighted, advantage-aware experience replay mechanism. For the first time, it integrates theoretically grounded timestep weighting—derived from PPO convergence analysis—and prioritized replay into the diffusion-based RLHF framework. By focusing gradient updates on high-information denoising stages, the proposed method substantially improves learning efficiency while preserving strong generalization to unseen prompts. Under identical hyperparameters, it achieves up to a sixfold improvement in sample efficiency over state-of-the-art baselines.

Abstract

Reinforcement learning from human feedback (RLHF) has emerged as a powerful paradigm for aligning generative models with human preferences. However, applying RLHF to diffusion models remains highly feedback inefficient, as existing approaches typically require large amounts of human or reward model evaluations. This limitation reduces the practicality of diffusion RLHF in realworld settings where feedback is the primary bottleneck. In this paper, we propose two complementary strategies that substantially improve the feedback efficiency of diffusion RLHF while preserving generalization to unseen prompts. Our key observation is that reward information in diffusion trajectories is unevenly distributed: not all denoising timesteps or trajectories contribute equally to learning from a reward signal. By emphasizing informative timesteps and trajectories during optimization, we obtain more effective gradient updates. First, we introduce a per-timestep weighting scheme that reweights denoising steps during policy optimization. We theoretically connect this weighting to the optimal convergence properties of proximal policy optimization (PPO) and approximate the resulting weighting trend empirically. Second, we introduce a replay mechanism that prioritizes informative trajectories, enabling the model to reuse past samples instead of repeatedly querying new rewards. Together, these strategies significantly improve the feedback efficiency of diffusion RLHF. Under identical hyperparameter settings, our approach achieves up to a 6$\times$ improvement in sample efficiency compared to widely used diffusion RLHF baselines.

Problem

• diffusion RLHF
• sample efficiency
• human feedback
• reward evaluation
• feedback inefficiency

Innovation

• Selective Timestep Weighting
• Advantage-Based Replay
• Diffusion RLHF
• Sample Efficiency
• Reward Modeling