Selective Timestep Weighting and Advantage-Based Replay for Sample-Efficient Diffusion RLHF
2026-07-08 • Machine Learning
Machine LearningArtificial IntelligenceComputer Vision and Pattern Recognition
AI summaryⓘ
The authors studied how to make learning from human feedback more efficient for diffusion models, which are used to generate images or other data. They found that not all parts of the model’s process are equally helpful for learning, so they created a method to focus learning on the most informative steps and past experiences. This helps the model learn faster without needing as much human feedback. Their approach improved efficiency by up to six times compared to older methods.
Reinforcement Learning from Human Feedback (RLHF)Diffusion ModelsDenoising StepsPolicy OptimizationProximal Policy Optimization (PPO)Reward SignalTrajectory ReplaySample EfficiencyGradient Updates
Authors
Eric Zhu, Abhinav Shrivastava, Soumik Mukhopadhyay
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.