TreeGRPO for Diffusion Models

• The paper introduces TreeGRPO, which reformulates the diffusion sampling process as a finite-horizon MDP with a tree-structured rollout for efficient trajectory sampling.
• The method employs fine-grained, edge-specific credit assignment through bottom-up reward backpropagation, enhancing policy updates compared to traditional trajectory-level RLHF.
• Empirical results demonstrate a 2.4× reduction in training time and superior reward-versus-compute efficiency over prior GRPO baselines across multiple benchmarks.

TreeGRPO for Diffusion Models is a reinforcement learning (RL) framework that substantially improves the efficiency of RL-based post-training for diffusion and flow-based generative models. By recasting the multi-step denoising sampler as a finite-horizon Markov Decision Process (MDP) with a tree-structured rollout, TreeGRPO enables efficient trajectory sampling, fine-grained credit assignment, and amortized computation. This approach attains 2.4× faster training compared to prior GRPO baselines and achieves a strictly superior Pareto frontier in reward versus computational efficiency across multiple benchmarks and reward models (Ding et al., 9 Dec 2025).

1. Denoising as a Search Tree

TreeGRPO formulates the T-step generative diffusion sampler as a finite-horizon MDP:

• State: $S_t = (c, t, x_t)$, where $c$ is conditioning (e.g., text prompt), $t \in \{0, ..., T\}$ is the current timestep, $x_t$ is the latent vector.
• Action: $a_t \sim \pi_\theta(\cdot|S_t)$, defining transitions $x_t \rightarrow x_{t+1}$ via an SDE or ODE sampler.
• Reward: Assigned only at the terminal node ($t = T$), as $R(x_T, c)$.

Unlike sampling $N$ independent trajectories, TreeGRPO builds a sparse, depth-$T$ search tree rooted at an initial latent $x_0 \sim \mathcal{N}(0,I)$. The set $\{0, ..., T-1\}$ is partitioned into:

• ODE steps (no branching), which deterministically propagate all frontier nodes, reusing computation for shared prefixes.
• SDE windows $W \subset \{0, ..., T-1\}$; at each $t \in W$, every frontier node spawns $b$ children via $b$ stochastic perturbations.

The resulting tree comprises:

• Nodes: Each at depth $t$ corresponds to a latent $x_t^u$.
• Edges: Each edge $e = (\text{parent} \rightarrow u)$ corresponds to an action $a(e)$ and log-probability under a frozen sampler $$\log \pi_{\text{old}}(a(e)\,|\,x_{\text{parent}}, c, t)$$.
• Branching: At $t \in W$, branching factor $b$; at $t \notin W$, only a single continuation.
• Reuse: Prefixes between branches are reused, especially between SDE windows.

Let $|\mathcal{L}(c)| = b^{|W|}$ denote the number of leaf nodes per prompt.

2. TreeGRPO Algorithmic Workflow

A single TreeGRPO rollout proceeds as follows:

1. Initialization: For each conditioning $c$, sample $x_0 \sim \mathcal{N}(0,I)$ and set the initial frontier $\mathcal{F}_0 = \{u_0\}$.
2. Forward Passes: Iterate $t = 0$ to $T-1$:
   • If $t \in W$, branch each frontier node $u$ into $b$ children via SDE sampling, recording log-probabilities.
   • Else, use ODE propagation to deterministically update each node’s latent.
3. Decoding and Reward: Each leaf latent $x_T^i$ is decoded to an image $y_i$; compute $r_i = R(y_i, c)$.
4. Advantage Computation: Group-normalize rewards:

$A^{\text{leaf}}_i = \frac{r_i - \mu_c}{\sigma_c}$

with $\mu_c$ and $\sigma_c$ as the leaf-wise mean and standard deviation.

1. Advantage Back-Propagation: Bottom-up propagation of advantages through the tree: for internal edge $e'$ with outgoing edges $S(u)$

$$\omega_u(e_j) = \frac{ \exp(\log \pi_{\text{old}}(e_j)) }{ \sum_{k=1}^{b} \exp(\log \pi_{\text{old}}(e_k)) }$$

$$A_{\text{edge}}(e') = \sum_{e \in S(u)} \omega_u(e) A_{\text{edge}}(e)$$

2. Surrogate Loss and Update: For each SDE-step edge $e$,

$$r_e(\theta) = \exp( \log \pi_\theta(a(e) \,|\, x(e), c, t(e)) - \log \pi_{\text{old}}(a(e) \,|\, x(e), c, t(e)) )$$

$$L_{\text{GRPO}} = - \mathbb{E}_{e \in \mathcal{S}} \left[ \min( r_e(\theta) A_{\text{edge}}(e), \operatorname{clip}(r_e(\theta), 1-\epsilon, 1+\epsilon) A_{\text{edge}}(e) ) \right]$$

Update $$\theta \leftarrow \theta - \eta \nabla_{\theta} L_{\text{GRPO}}$$.

Each rollout collects $b^{|W|}$ leaf trajectories with only $\mathcal{O}((T-|W|) + b\cdot|W|)$ forward steps due to prefix reuse.

3. Fine-Grained and Efficient Credit Assignment

TreeGRPO resolves the uniformity limitations of trajectory-based advantage assignment by introducing step-specific, edge-local advantages via bottom-up reward backpropagation through the search tree.

• Leaf Node: Normalize ground-truth rewards across the leaf set for each prompt.
• Internal Node: For depth $t$ (from $\max(W)-1$ to $\min(W)$), propagate advantages using a log-probability softmax over outgoing edges, producing distinct per-edge advantages throughout the branching subtrees.
• Granularity: This yields fine-grained, step-specific credit assignment, enhancing policy update signal compared to standard RLHF methods with uniform trajectory-level reward.

A plausible implication is improved policy robustness and sample efficiency, as each SDE decision receives targeted learning signal proportional to its long-term impact on final reward.

4. Amortized Computation and Theoretical Speedup

TreeGRPO’s main computational advantage arises from amortizing the cost of branching trajectories:

• Baseline GRPO: $N$ independent $T$-step trajectories require $\mathcal{O}(N \cdot T)$ forward passes, directly proportional to the number of distinct samples and steps.
• TreeGRPO: With branching $b$ and $w = |W|$ SDE windows, generates $b^w$ distinct trajectories using only

$\mathcal{O}( (T-w) + b\cdot w )$

forward steps.

• Speedup: Analytical estimate,

$$\text{SpeedUp} \approx \frac{b^w \cdot T}{(T-w) + b \cdot w}$$

Empirical settings (e.g., $b=3$, $w=3$, $T \approx 10$) yield $2$–$3\times$ reduction in FLOPs per gradient; this matches the observed $2.4\times$ wall-clock speedup in training.

This architecture ensures that common computation along trajectory prefixes is maximally reused, and branching happens only at critical SDE steps. Between SDE windows, the ODE segments are computed only once per shared prefix.

5. Empirical Evaluation: Setup and Results

Experiments used Stable Diffusion 3.5-Medium (SD3.5-M) for diffusion, with analogous models for flow-based benchmarks:

• Datasets: HPDv2 (103,700 train, 3,200 evaluation prompts)
• Sampler Budget: 10-step NFE, batch-size 32, 250 epochs, 8×A100 GPUs, AdamW($1\mathrm{e}{-5}$, weight decay 0.01).

Reward Models:

• HPS-v2.1 (human preference score)
• ImageReward
• Aesthetic
• ClipScore

Two evaluation regimes: single-reward (HPS only) and multi-reward (HPS:ClipScore $=0.8:0.2$).

Results:

Method	Iter. Time (s)	HPS-v2.1	ImageReward	Aesthetic	ClipScore
DDPO	166.1	0.2758	1.0067	5.9458	0.3900
DanceGRPO	173.5	0.3556	1.3668	6.3080	0.3769
MixGRPO	145.4	0.3649	1.2263	6.4295	0.3612
TreeGRPO	72.0	0.3735	1.3294	6.5094	0.3703

• Training Efficiency: TreeGRPO delivers a $2.4\times$ reduction in per-iteration time relative to DanceGRPO.
• Reward Frontier: Outperforms or matches the strongest baseline in both HPS-v2.1 and Aesthetic scores, with competitive ImageReward and ClipScore.
• GPU Hour Tradeoff: Pareto analysis demonstrates strict dominance, with TreeGRPO achieving higher mean normalized reward across all metrics for any fixed compute budget.

Ablation on branching $k$ and window depth $d$ identifies $k=3$, $d=3$ as optimal for efficiency–performance tradeoff.

6. Comparison with Prior GRPO Baselines

Relative to previous approaches for RL post-training of diffusion models:

• DDPO, DanceGRPO, and MixGRPO each utilize standard or amortized GRPO sampling and trajectory-level RLHF.
• TreeGRPO achieves strictly superior GPU/reward efficiency and wall-clock speed, repeatedly matching or exceeding the best baseline reward for all metrics while halving or better the per-iteration cost.

This dominance is robust across reward models, including both single-reward and multi-reward regimes. The architecture’s three key contributions—prefix reuse, reward backpropagation for step-specific advantages, and multi-child branching per forward pass—underpin its empirical advantage.

7. Significance and Practical Considerations

TreeGRPO establishes a new standard for RL-based post-training of large diffusion and flow-based generators, especially for large-batch or high-budget settings. The tree-structured rollout provides:

• High sample efficiency: Many unique candidate trajectories per single computation budget.
• Fine-grained credit assignment: Overcomes the limitations of trajectory-level RLHF credit, allowing for improved signal propagation through the generative process.
• Amortized computation: Enables scalable post-training, previously a major barrier for widespread RLHF application in vision generative models.

The methodology generalizes beyond diffusion models, applying likewise to flow-based samplers. These advances provide a scalable and effective pathway for aligning visual generative models with complex reward and preference structures without prohibitive compute investment (Ding et al., 9 Dec 2025).