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

Introduction

• Existing diffusion-model GRPO variants face redundant O(N · T) sequential rollouts and unreliable credit assignment from uniformly propagated terminal rewards.Uniform reward propagation neglects intermediate denoising signals and produces high-variance gradients.Source: §1 · p. 1
• BranchGRPO replaces independent sequential rollouts with scheduled branching during denoising, reusing shared prefixes to amortize computation and broaden exploration.Trajectories stochastically expand into multiple sub-trajectories from shared segments.Source: §1 · pp. 1–2
• BranchGRPO achieves up to 2.2× speedup over DanceGRPO (time fraction = 1.0) and 1.5× speedup over DanceGRPO (time fraction = 0.6), while ultimately surpassing both baselines.The time fraction = 0.6 variant exhibits pronounced instability.Source: Figure 1 · p. 2
• BranchGRPO fuses leaf rewards and propagates them backward with depth-wise normalization to produce finer-grained, step-level advantages.This converts sparse terminal rewards into dense signals for more stable optimization.Source: §1 · p. 2
• Width- and depth-pruning remove redundant branches and depths during backpropagation, reducing computation and directing it toward promising trajectory regions.The paper describes these complementary strategies as lowering backpropagation cost and further improving alignment.Source: §1 · p. 2
• BranchGRPO is validated on text-to-image and image-to-video alignment tasks, and larger group sizes consistently yield better alignment performance.The paper presents these results as evidence of effectiveness and generality across modalities and verifies a scaling law with group size.Source: §1 · p. 2

2 RELATED WORK

• Diffusion and flow-matching models dominate visual generation, and RLHF methods successful in language models have motivated preference alignment for vision.The passage cites diffusion models (Ho et al., 2020; Rombach et al., 2022), flow matching (Lipman et al., 2022; Liu et al., 2022), and language-model RLHF methods including PPO and GRPO (Ouyang et al., 2022; Christiano et al., 2017; Lu et al., 2025).Source: §2 · p. 2
• ImageReward (Xu et al., 2023) backpropagates learned rewards through denoising, while Diffusion-DPO (Wallace et al., 2024) and Videodpo (Liu et al., 2025c) avoid explicit reward modeling.Online policy-gradient methods DDPO (Black et al., 2023) and DPOK (Fan et al., 2023) explore optimization but often face stability challenges; HPS-v2.1 and VideoAlign (Liu et al., 2025b) support standardized comparisons.Source: §2 · pp. 2–3
• DanceGRPO (Xue et al., 2025) and Flow-GRPO (Liu et al., 2025a) apply GRPO across visual-generation tasks, while TempFlow-GRPO (He et al., 2025) targets sparse-reward credit assignment and MixGRPO (Li et al., 2025a) trades overhead against performance.DanceGRPO and Flow-GRPO unify diffusion and flow models via SDE reformulation; TempFlow-GRPO uses temporally aware denoising-step weights, and MixGRPO uses a mixed ODE–SDE sliding-window scheme.Source: §2 · p. 3
• BranchGRPO extends this line with branching rollouts, depth-wise reward fusion, and structured pruning, adapting tree-structured rollouts from TreePO (Li et al., 2025b) to diffusion dynamics.Source: §2 · p. 3

3.1 DOES BRANCH ROLLOUT HARM DIVERSITY?

• KID=0.0057 and MMD2=0.0067 (Inception), and KID=0.00022 and MMD2=0.0149 (CLIP ViT-B/32), support near-indistinguishable sample distributions.Across 4096 samples per method, feature-embedding distributions largely overlap; the paper concludes that branch rollouts preserve distributional and semantic diversity with negligible shifts.Source: §3.1 · p. 3

3.2 BRANCH ROLLOUT ALGORITHM

• BranchGRPO shares early denoising prefixes, then splits each selected state into K correlated children controlled by s, preserving the baseline SDE marginal distribution.A shared initial noise seeds the tree; child branches diverge through stochastic perturbations, balancing exploration diversity and stability.Source: §3.2 · p. 4

3.3 REWARD FUSION AND DEPTH-WISE ADVANTAGE ESTIMATION

• Leaf rewards are fused upward into node values, then normalized within each depth to produce dense, balanced edge advantages and finer denoising-step credit assignment.Fusion ranges from uniform averaging at β = 0 to behavior-policy weighting at β = 1: averaging retains low-probability leaves but may add variance, while path-weighting reduces variance but can over-concentrate on high-likelihood leaves. Per-depth standardization prevents late steps with smaller variance from dominating.Source: §3.3 · pp. 4–5

3.4 PRUNING STRATEGIES

• Width and depth pruning reduce backpropagation cost after reward fusion and normalization, while preserving all forward rollouts and reward evaluations.Width pruning uses Parent-Top1 or Extreme selection; depth pruning skips gradients at selected timesteps and shifts its pruning window toward later steps during training.Source: §3.4 · p. 5

4.1 EXPERIMENT SETUP

• Image alignment is evaluated on HPDv2.1 using FLUX.1-Dev, with DanceGRPO and MixGRPO as baselines and efficiency and preference metrics reported.The evaluation uses 103k training prompts and 400 balanced test prompts; metrics include NFE, iteration time, HPSv2.1, PickScore, ImageReward, and Unified Reward.Source: §4.1 · p. 6

4.2 IMPLEMENTATION DETAILS

• The image experiments use depth 20, branch factor 2, and 16 leaves per rollout before pruning, with the Dense branching schedule as default.Training runs for 300 optimizer steps on 16 NVIDIA H200 GPUs, with other GRPO hyperparameters held identical across methods.Source: §4.2 · p. 6

4.3 MAIN RESULTS

• BranchGRPO-DepthPruning raises HPS-v2.1 from 0.360 for DanceGRPO to 0.369 and achieves PickScore 0.231, ImageReward 1.625, and Unified Reward 3.404.BranchGRPO-WidthPruning and BranchGRPO-Mix reduce iteration time to 314s and 148s, respectively, with marginal quality drops; qualitative comparisons show sharper details and better semantic alignment than Flux and DanceGRPO.Source: §4.3 · pp. 6–7, Table 1 · p. 7, Figure 4 · p. 7

4.4 ABLATION STUDIES

• With branch correlation s = 4.0, training reaches the highest reward with stable convergence; early, denser splits accelerate reward growth, while later or sparser schedules grow more slowly.Very large correlation (s = 8.0) destabilizes early training, whereas smaller values limit exploration; all tested split densities eventually converge to similar reward levels.Source: §4.4 · pp. 7–8, Figure 5 · p. 8
• Path-probability weighting produces higher, more stable rewards than uniform averaging, while depth pruning gives the best final reward and Parent-Top1 width pruning the smoothest, lowest-variance curve.Depth pruning skips losses and gradients in a sliding four-step window that shifts every 30 iterations while leaving forward sampling unchanged; the hybrid ODE–SDE scheme reaches 148s versus 289s for MixGRPO and 469s for DanceGRPO.Source: §4.4 · p. 8, Figure 5 · p. 8

4.5 SCALING WITH BRANCH

• Scaling both branch factor and number of branching steps raises reward growth and final performance, with gains becoming more pronounced as the rollout tree expands.At 81 rollout samples, a GRPO step takes over 3500s for DanceGRPO versus 680s for BranchGRPO.Source: §4.5 · pp. 8–9, Figure 6 · p. 9

4.6 VIDEO GENERATION RESULTS

• On Wan2.1-1.3B, BranchGRPO produces sharper, more temporally coherent frames and faster convergence with higher final rewards than DanceGRPO.Each iteration takes about 8 minutes for BranchGRPO versus 20 minutes for DanceGRPO; the reward is Video-Align motion quality.Source: §4.6 · p. 9, Figure 7 · p. 9

5 CONCLUSION

• Across image and video generation, BranchGRPO achieves faster convergence, more stable training, and higher final alignment quality under matched budgets.The conclusion presents structured branching as a practical, scalable path for RLHF in diffusion and flow models.Source: §5 · p. 9
• BranchGRPO replaces sequential rollouts with prefix-sharing branching and depth-wise reward fusion, augmented by lightweight pruning for compute reallocation.Source: §5 · p. 9
• Larger group sizes consistently improve performance, establishing a scaling trend for BranchGRPO.Source: §5 · p. 9

ETHICS STATEMENT

• 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 usage terms, restricting use to non-commercial research where applicable.Source: ETHICS STATEMENT · p. 10
• To address possible bias amplification, the authors report failure cases and stability and variance ablations, and avoid safety-critical deployment.Source: ETHICS STATEMENT · p. 10
• The authors plan to release code and configuration files to support scrutiny and responsible reuse.Source: ETHICS STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

• The supplementary materials provide code, configurations, exact hyperparameters, scripts, seeds, and environment specifications for reproducing all experiments.The released training and evaluation code is anonymized.Source: REPRODUCIBILITY STATEMENT · p. 10

APPENDIX

• Appendix Section A provides hyperparameter settings used in all experiments.The detailed configuration is summarized in Table 2.Source: APPENDIX · p. 13, Table 2 · p. 13
• The appendix includes theoretical proofs covering branch noise construction, reward fusion, and variance reduction.It also supplies additional theoretical proofs and implementation details that complement the main paper.Source: APPENDIX · p. 13
• The appendix reports ablations and additional text-to-image and image-to-video results, shows failure cases, and includes discussion of future work.The failure examples highlight current limitations.Source: APPENDIX · p. 13

LLM USAGE

• An LLM was used only to polish manuscript language, and did not generate, modify, or select experimental results.No confidential or non-public data were shared; the authors take responsibility for the final content and state that technical content, claims, ideas, and analyses are their own.Source: LLM USAGE · p. 13

A HYPERPARAMETER SETTINGS

• All methods use identical hyperparameters, and depth pruning and hybrid-ODE-SDE use a size-4 sliding window that shifts one step deeper every 30 iterations.Table 2 summarizes the detailed settings; the passage specifies that the shared configuration includes DanceGRPO and MixGRPO.Source: A HYPERPARAMETER SETTINGS · p. 13, Table 2 · p. 13, APPENDIX · p. 13

B.1 BRANCH NOISE CONSTRUCTION AND BOUNDARY DISTRIBUTION PRESERVATION

• At each split, fresh shared and branch-specific noise creates correlated sibling noises while retaining the single-step marginal distribution.The construction specifies no cross-time sharing of the shared component.Source: §B.1 · p. 14
• Each branching leaf has the same marginal law as a baseline SDE rollout, so branching preserves the final generator distribution.This result holds for any set of split steps under the stated construction and assumptions about independent noise across time and no cross-time reuse of the shared component.Source: §B.1 · p. 14

B.2 REWARD FUSION: UNBIASEDNESS AND VARIANCE REDUCTION

• Path-probability-weighted importance sampling gives an unbiased estimate of a node’s conditional expected return.The self-normalized importance-sampling estimator is also introduced, but the supplied passage does not specify its bias properties.Source: §B.2 · p. 14
• In practice, softmax weights based on path log-probabilities trade strict unbiasedness for stable training and less influence from low-probability noisy leaves.The temperature parameter controls the sharpness of the importance weights.Source: §B.2 · p. 14

B.3 DEPTH-WISE BASELINE (CONTROL VARIATES)

• The depth-wise group-mean baseline yields an unbiased gradient estimator with strictly smaller variance than the single-sample gradient estimator, except under the stated covariance condition.The supplied passage fragment does not include the full condition.Source: §B.3 · p. 15

B.4 CONTINUOUS REWARDS AND CONCENTRATION

• For continuous rewards, the analysis states a sub-Gaussian property and considers averaging rewards over a node’s leaves.The supplied fragments do not provide the complete parameter or resulting concentration bound.Source: §B.4 · p. 15

• Table 3 presents BranchGRPO ablations across design choices, with all results using the Section 4.3 training setup.Source: §C.1 · p. 15, Table 3 · p. 15

C.2 MORE TEXT2IMAGE RESULTS

• Figure 8 is identified as an additional text-to-image result, but its outcome is not described in the supplied passage.Source: Figure 8 · p. 16

C.3 MORE IMAGE2VIDEO RESULTS

• Table 4 evaluates video on vBench, reporting mean values over 500 samples; the supplied passage does not include the values.Source: Table 4 · p. 21

C.4 FAILURE CASES

• Figure 15 is labeled a failure case, but the supplied passage does not specify what fails.Source: Figure 15 · p. 23
• Figure 16 is labeled a failure case, but the supplied passage does not specify what fails.Source: Figure 16 · p. 24

Discussion

• The authors report stability and efficiency benefits; branching schedules and pruning affect reward stability, while reward fusion’s weighting-dependent bias–variance tradeoff needs theoretical analysis.Reward fusion provides stable gradients in practice.Source: D DISCUSSION AND FUTURE WORK · p. 25

Future Work

• Future work proposes adaptive branching policies that tune branch factor, correlation, or pruning windows using sample difficulty or intermediate rewards.The goal is to make rollouts more efficient.Source: D DISCUSSION AND FUTURE WORK · p. 25
• Future work calls for transfer beyond diffusion, stronger validation on high-resolution long-duration video, and extensions to robotics and embodied video generation.Initial WanX-1.3B I2V experiments show benefits, but broader video validation is needed; robotics could leverage intermediate states with dense, verifiable task-success rewards.Source: D DISCUSSION AND FUTURE WORK · p. 25