TreeGRPO: Tree-Advantage GRPO for Online RL Post-Training of Diffusion Models
This addresses a computational bottleneck for researchers and practitioners working on RL-based alignment of diffusion and flow-based generative models, though it appears incremental as an enhancement to existing GRPO methods.
The paper tackles the high computational cost of reinforcement learning post-training for aligning generative models with human preferences by introducing TreeGRPO, a framework that recasts the denoising process as a search tree to reuse common prefixes and enable multiple candidate trajectories. The result is 2.4× faster training while achieving better performance and establishing a superior Pareto frontier in efficiency-reward trade-offs.
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.