样本高效扩散RLHF的选择性时间步加权与优势回放

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

精选理由

这篇论文教你怎么用更少的反馈数据训练扩散模型,关键是把去噪步和样本按重要性加权,实验涨了6倍效率。

AI 摘要

本文提出两种互补策略提升扩散模型RLHF的反馈效率。第一种是逐时间步加权方案,重新调整策略优化中去噪步骤的权重,并与PPO的最优收敛性质建立理论联系。第二种是基于优势的回放机制,优先复用信息丰富的轨迹,减少对新奖励查询的依赖。在相同超参数下,该方法将样本效率提升高达6倍,并保持对未见提示的泛化能力。

原文 · arXiv cs.LG

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

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.