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

Applying Reinforcement Learning from Human Feedback (RLHF) to diffusion models is notoriously sample-inefficient, often requiring massive amounts of reward model evaluations. The authors investigate the credit assignment problem in diffusion RLHF—specifically, how to determine which intermediate denoising steps contribute most to the final reward—to improve training efficiency without requiring more feedback.

The authors propose two complementary strategies to optimize the learning process:

1. Per-Timestep Weighting: Instead of treating every denoising step equally, the authors derive a theoretical framework connecting diffusion RLHF to Proximal Policy Optimization (PPO). They show that timesteps should be weighted proportionally to the variance of the advantage. Since this is computationally expensive to calculate exactly, they use a practical heuristic based on the magnitude of latent changes during denoising to reweight the loss function.

2. Historical Hardmining (Replay Buffer): Inspired by prioritized experience replay in robotics, the authors implement a buffer that stores past trajectories. During training, the model prioritizes re-sampling trajectories that yielded the highest advantages, allowing the model to learn from its most informative past experiences rather than relying solely on new, potentially low-quality samples.

The combination of these two techniques significantly reduces the number of reward queries needed to reach a target performance level. The authors report up to a 6x improvement in sample efficiency compared to standard baselines like DDPO. The method is "plug-and-play," meaning it can be integrated into existing diffusion RLHF pipelines without requiring architectural changes to the underlying generative model.

In real-world applications, human feedback is the primary bottleneck for training generative models. By making diffusion RLHF more sample-efficient, this work lowers the barrier to entry for aligning high-quality image and video models with human preferences, making the process more practical for researchers and developers with limited access to large-scale human annotation or expensive reward model compute.