BranchGRPO: Stable and Efficient GRPO with Structured Branching in Diffusion Models

Yuming Li; Yikai Wang; Yuying Zhu; Zhongyu Zhao; Ming Lu; Qi She; Shanghang Zhang

Source-linked AI summary

BranchGRPO: Stable and Efficient GRPO with Structured Branching in Diffusion Models

Yuming Li, Yikai Wang, Yuying Zhu, Zhongyu Zhao, Ming Lu, Qi She, Shanghang Zhang

arXiv:2509.06040v5Sep 7, 2025cs.CVcs.AIcs.LG

TL;DR

Existing diffusion-model GRPO methods are inefficient and assign sparse terminal rewards uniformly across denoising steps. BranchGRPO uses shared-prefix branching, reward fusion, depth-wise advantages, and pruning, achieving stronger alignment with faster, more stable training across image and video tasks.

Problem
Existing GRPO variants use inefficient sequential rollouts and uniform propagation of sparse terminal rewards, limiting computational efficiency and credit assignment.
Method
BranchGRPO restructures rollouts into shared-prefix branching trees, fuses leaf rewards into dense depth-wise advantages, and prunes gradient computation without affecting forward rollouts or reward evaluation.
Results
Across image and video generation, BranchGRPO yields faster convergence, more stable training, and higher final alignment quality under matched budgets.
Takeaways & Limitations
Structured branching is presented as a practical, scalable approach for RLHF in diffusion and flow models, with larger group sizes consistently improving alignment performance.
Takeaways & Limitations
Longer and higher-resolution video generation requires more extensive validation beyond the initial WanX-1.3B I2V experiments.

Abstract

from arXiv · show

Recent progress in aligning image and video generative models with Group Relative Policy Optimization (GRPO) has improved human preference alignment, but existing variants remain inefficient due to sequential rollouts and large numbers of sampling steps, unreliable credit assignment: sparse terminal rewards are uniformly propagated across timesteps, failing to capture the varying criticality of decisions during denoising. In this paper, we present BranchGRPO, a method that restructures the rollout process into a branching tree, where shared prefixes amortize computation and pruning removes low-value paths and redundant depths. BranchGRPO introduces three contributions: (1) a branching scheme that amortizes rollout cost through shared prefixes while preserving exploration diversity; (2) a reward fusion and depth-wise advantage estimator that transforms sparse terminal rewards into dense step-level signals; and (3) pruning strategies that cut gradient computation but leave forward rollouts and exploration unaffected. On HPDv2.1 image alignment, BranchGRPO improves alignment scores by up to \textbf{16\%} over DanceGRPO, while reducing per-iteration training time by nearly \textbf{55\%}. A hybrid variant, BranchGRPO-Mix, further accelerates training to 4.7x faster than DanceGRPO without degrading alignment. On WanX video generation, it further achieves higher Video-Align scores with sharper and temporally consistent frames compared to DanceGRPO. Codes are available at \href{https://fredreic1849.github.io/BranchGRPO-Webpage/}{BranchGRPO}.

1 INTRODUCTION

BranchGRPO addresses inefficiency and unreliable credit assignment in diffusion-model GRPO by using shared-prefix branching, dense depth-aware rewards, and pruning. Across image and video alignment tasks, it improves convergence, stability, and alignment while preserving exploration diversity.

RLHF is needed because large-scale pretraining alone does not ensure outputs match human aesthetic, semantic, or temporal preferences.
Standard GRPO uses independent sequential rollouts with O(N · T) complexity, creating computational redundancy.
BranchGRPO branches at scheduled denoising steps and reuses shared prefixes, amortizing computation while maintaining diverse exploration.
Reward fusion and depth-wise advantage estimation convert sparse terminal rewards into dense step-level signals for more stable optimization.
Width- and depth-pruning strategies reduce backpropagation cost while reallocating computation toward promising trajectory regions.
BranchGRPO is validated across text-to-image and image-to-video alignment, with larger group sizes consistently producing better alignment performance.

2 RELATED WORK

Related work develops reward-model, preference-optimization, and policy-gradient approaches for visual-generation alignment. GRPO-based methods provide scalable optimization, but sparse terminal rewards remain a credit-assignment limitation motivating temporally aware alternatives.

Diffusion and flow-matching models provide strong visual-generation foundations, while RLHF addresses alignment with nuanced human preferences.
Visual RLHF includes reward-model methods such as ImageReward and direct-preference methods such as Diffusion-DPO and Videodpo.
Policy-gradient approaches including DDPO and DPOK explore online optimization but often face stability challenges.
DanceGRPO and Flow-GRPO extend GRPO to visual generation and demonstrate stable optimization across image and video tasks.
TempFlow-GRPO highlights sparse terminal rewards with uniform credit assignment as a limitation and proposes temporally aware weighting.

3 METHOD

BranchGRPO organizes diffusion rollouts as branching trees with shared prefixes, fuses leaf rewards upward, normalizes advantages by depth, and prunes gradient computation. These mechanisms target efficient exploration and finer credit assignment without changing forward rollout evaluation.

Does Branch Rollout Harm Diversity?: Inception-space comparisons report KID=0.0057 and MMD2=0.0067, while CLIP-space comparisons report KID=0.00022 and MMD2=0.0149.
Does Branch Rollout Harm Diversity?: Branch rollouts preserve distributional and semantic diversity, with DanceGRPO and BranchGRPO samples largely overlapping in feature space.
Branch Rollout Algorithm: BranchGRPO defines depth as denoising steps, width as completed trajectories, branching steps as split timesteps, correlation as inter-child diversity control, and K as branching factor.
Branch Rollout Algorithm: At designated split steps, a shared-noise root state expands into K correlated children that share prefixes and diverge afterward.
Reward Fusion and Depth-wise Advantage Estimation: BranchGRPO propagates leaf rewards upward and fuses them into internal-node values, with β interpolating between uniform averaging and behavior-policy weighting.
Reward Fusion and Depth-wise Advantage Estimation: Depth-wise normalization compares nodes sharing the same noise level and prevents late denoising steps with smaller variance from dominating gradients.
Pruning Strategies: Width and depth pruning reduce gradient computation after reward processing, while forward rollouts and reward evaluation still include all trajectories.

4 EXPERIMENTS

Experiments show that BranchGRPO improves image and video alignment while reducing training cost, with ablations supporting tuned branching, reward fusion, pruning, and rollout scaling.

Main Results: BranchGRPO-DepthPruning raises HPS-v2.1 from 0.360 for DanceGRPO to 0.369 and achieves the highest reported PickScore, ImageReward, and Unified Reward.The reported values are 0.231, 1.625, and 3.404, respectively.
Efficiency: 314s and 148s iteration times for BranchGRPO-WidthPruning and BranchGRPO-Mix reduce training cost with only marginal quality drops.BranchGRPO-Mix is compared with MixGRPO at 289s and DanceGRPO at 469s in the hybrid ODE–SDE study.
Main Results: BranchGRPO achieves faster early reward growth, smoother convergence, and higher final rewards than DanceGRPO, while producing sharper details and better semantic alignment than Flux and DanceGRPO.
Ablation Studies: Moderate branch correlation s = 4.0 reaches the highest reward with stable convergence, whereas smaller values limit exploration and s = 8.0 destabilizes early training.
Ablation Studies: Early and denser split schedules accelerate reward growth and improve sample efficiency without harming stability.Later splits delay exploration and yield lower rewards, while all tested density configurations eventually reach similar reward levels.
Ablation Studies: Depth pruning achieves the best final reward, while the hybrid ODE–SDE schedule provides the fastest stable training speed.Depth pruning skips losses and gradients only during backpropagation, leaving forward sampling unchanged.
Ablation Studies: Path-probability-weighted reward fusion produces higher and more stable rewards than uniform averaging, which has higher variance and a late-stage plateau.
Scaling: Larger branch factors and more branching steps consistently improve reward growth and final performance, following a clear scaling trend.At 81 rollout samples, DanceGRPO takes over 3500s per step, compared with 680s for BranchGRPO.

5 CONCLUSION

The paper concludes that BranchGRPO combines prefix-sharing branching, depth-wise reward fusion, and pruning to improve diffusion and flow-model alignment across image and video generation.

Across image and video generation, BranchGRPO yields faster convergence, more stable training, and higher final alignment quality under matched budgets.
Larger group sizes consistently improve performance, supporting a scaling trend for structured branching.

ETHICS STATEMENT

The experiments focus on diffusion and flow-based generative models using public resources and research-compliant data practices.

The experiments use publicly available prompt sets and reward models, with no personally identifiable information, human-subject data, or sensitive attributes collected or annotated.
The authors followed dataset and third-party model licenses and restricted use to non-commercial research when applicable.

REPRODUCIBILITY STATEMENT

The paper provides code, configurations, hyperparameters, theoretical analyses, and supplementary experiments intended to support reproduction and further evaluation.

Anonymized training and evaluation code, configuration files, exact hyperparameters, scripts, seeds, and environment specifications are released for reproduction.
The appendix includes hyperparameter settings, proofs for branch noise construction and reward fusion, and additional experiments and failure cases.
The manuscript states that an LLM was used only for language polishing and did not generate, modify, or select experimental results.
All methods use identical hyperparameters, while depth pruning and hybrid ODE–SDE use a sliding window of size 4 shifted every 30 iterations.

B THEORETICAL ANALYSIS

The theoretical analysis establishes that BranchGRPO preserves the baseline SDE rollout distribution while using fused rewards and depth-wise normalization for lower-variance, finer-grained gradient signals.

Branch noise construction uses fresh split-step randomness and correlated branch noises while maintaining the required Gaussian marginal.The child updates use branch-specific noise generated at each split step.
Branching preserves each leaf’s marginal distribution and therefore does not alter the final generator distribution.The result holds under independent noises across time and no cross-time reuse of shared components.
Fused returns estimate conditional expected returns from leaf rewards, with softmax path weighting providing a temperature-smoothed but biased estimator.The weighting reduces the influence of low-probability noisy leaves and supports stable training.
The group gradient estimator is unbiased and has strictly smaller variance than the single-trajectory estimator under the stated conditions.The comparison is stated for K siblings at a given depth and depends on the covariance condition described in the analysis.

C.1 MORE QUANTITATIVE RESULT

The supplied passages indicate that additional quantitative and ablation results are reported in Tables 3 and 4.

Table 3 reports more results for the BranchGRPO experiments.
Table 3 presents an ablation study under different BranchGRPO design choices.The caption states that all results use the same training setup as Section 4.3.

C.2 MORE TEXT2IMAGE RESULTS

The supplied passages document text-to-image prompt examples, figure references, failure-case labels, and a video-evaluation table, but do not state comparative outcomes for these materials.

Text-to-image examples: The section includes an anime flight-uniform portrait prompt with detailed digital-art and named-art-style descriptors.
Figures and evaluation: The supplied materials reference Figures 8, 9, 10, 11, and 12 without describing their encodings or outcomes.
Text-to-image examples: Another prompt specifies two women with purple hair flying in different directions against a dark background.
Text-to-image examples: Additional prompts depict a mist-surrounded raccoon in a Peaky Blinders hat, a polar bear cub with sunglasses, and a colorful anime-inspired magical composition.
Video evaluation: Table 4 reports video evaluation on vBench using mean values from 500 samples.
Failure cases: Failure-case materials include a League of Legends Arcane Jinx portrait and labels for Figures 15 and 16.

D DISCUSSION AND FUTURE WORK

The discussion reports efficiency and alignment gains while identifying tree design and reward-weighting trade-offs as areas requiring further study. Future work extends BranchGRPO toward adaptive branching, other generative paradigms, longer videos, and robotics.

Discussion: BranchGRPO improves both efficiency and alignment, while ablations show that branching schedules and pruning strategies substantially affect reward stability.
Discussion: Reward fusion provides stable gradients in practice, but its bias–variance trade-off under different weighting schemes requires further theoretical analysis.
Future Work: Future directions include dynamic branching, transfer beyond diffusion models, broader long-horizon video validation, and robotic action generation.The paper specifically notes that more extensive validation is required for high-resolution, long-duration video generation.