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

TL;DR

TreeGRPO introduces a tree-structured reinforcement learning framework that significantly enhances training efficiency and performance in aligning diffusion models with human preferences, reducing computational costs and improving reward optimization.

Contribution

The paper presents TreeGRPO, a novel tree-based RL method that improves sample efficiency, enables fine-grained credit assignment, and allows amortized computation for better training of generative models.

Findings

Achieves 2.4× faster training compared to baselines

Outperforms GRPO in multiple benchmarks and reward models

Provides a scalable approach for RL-based generative model alignment

Abstract

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.