arXiv:2512.08153cs.LGcs.AI2025-12被引 24

用搜索树结构提升扩散模型强化学习训练效率,提速2.4倍。

TreeGRPO: Tree-Advantage GRPO for Online RL Post-Training of Diffusion Models

  • 将去噪过程建模为搜索树,共享初始噪声并分支生成多条轨迹。
  • 相同样本下性能更优,训练速度提升2.4倍,达到更优效率-奖励平衡。
  • 适合需要高效对齐视觉生成模型与人类偏好的研究者使用。

强化学习后训练对齐生成模型与人类偏好至关重要,但其高昂的计算成本仍是广泛应用的主要障碍。我们提出TreeGRPO,一种新型强化学习框架,通过将去噪过程重构为搜索树,显著提升训练效率。从共享的初始噪声样本出发,TreeGRPO战略性地分支生成多个候选轨迹,并高效复用其公共前缀。该树状结构带来三大优势:(1) 高样本效率,在相同训练样本下表现更优;(2) 细粒度信用分配,通过奖励反向传播计算每一步的优势,克服轨迹级方法均匀分配信用的局限;(3) 计算摊销,多子节点分支支持一次前向传播完成多次策略更新。在扩散模型与基于流的模型上的大量实验表明,TreeGRPO实现2.4倍加速,且在效率-奖励权衡空间中建立更优帕累托前沿。该方法在多个基准和奖励模型上持续优于GRPO基线,为基于强化学习的视觉生成模型对齐提供了可扩展、高效的路径。项目主页见treegrpo.github.io。

原文摘要 · Abstract (English)

Reinforcement learning (RL) post-training is crucial for aligning generative models with human preferences, but its prohibitive computational cost remains a major barrier to widespread adoption. We introduce \textbf{TreeGRPO}, a novel RL framework that dramatically improves training efficiency by recasting the denoising process as a search tree. From shared initial noise samples, TreeGRPO strategically branches to generate multiple candidate trajectories while efficiently reusing their common prefixes. This tree-structured approach delivers three key advantages: (1) \emph{High sample efficiency}, achieving better performance under same training samples (2) \emph{Fine-grained credit assignment} via reward backpropagation that computes step-specific advantages, overcoming the uniform credit assignment limitation of trajectory-based methods, and (3) \emph{Amortized computation} where multi-child branching enables multiple policy updates per forward pass. Extensive experiments on both diffusion and flow-based models demonstrate that TreeGRPO achieves \textbf{2.4$\times$ faster training} while establishing a superior Pareto frontier in the efficiency-reward trade-off space. Our method consistently outperforms GRPO baselines across multiple benchmarks and reward models, providing a scalable and effective pathway for RL-based visual generative model alignment. The project website is available at treegrpo.github.io.

强化学习扩散模型高效训练

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。