BranchGRPO: Stable and Efficient GRPO with Structured Branching in Diffusion Models
Authors: Yuming Li
Yikai Wang
Yuying zhu
Zhongyu Zhao
Ming Lu
Qi She
Shanghang Zhang
Venue: ICLR 2026 Poster (2026)
Keywords: Diffusion Models
Flow Matching
RLHF
GRPO
Efficient Training
Primary area: generative models
TL;DR: We introduce BranchGRPO, a tree-structured rollout framework that makes RLHF for diffusion models more efficient and stable.

Abstract: 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,as 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 HPSv2.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.7× faster than DanceGRPO without degrading alignment.
On WanX video generation, it further achieves higher motion quality reward with sharper and temporally consistent frames.

Published as a conference paper at ICLR 2026
BRANCHGRPO: STABLE AND EFFICIENT GRPO WITH STRUCTURED BRANCHING IN DIFFUSION MODELS
Yuming Li Yikai Wang2*
Yuying Zhu
Zhongyu Zhao' Ming Lu
Qi She
Shanghang Zhang't
Stttt Key Laboratory of Multimedia Information Processing,
School of Computer Science, Peking University
BBejjing Normal University 'ByteDance
shanghang@pku. cetphe@pku.uuuccmmmm edu.cn
ABSTRACT
Recent progress in aligning image and video generative models with Group Rela- tive 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, as sparse terminal rewards are uniformly propagated across timesteps, failing to capture the varying criticality of decisions during denoising. In this paper, present BranchGRPO, 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: (I) a branching scheme that amor- tizes rollout cost through shared prefixes while preserving exploration diversity: (2) 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 HPSv2.1 image alignment, BranchGRPO improves alignment scores by up to 16% over Danc while reducing per-iteration training time by nearly 55%. A hybrid variant, BranchGRPO-Min further accelerates training to 4.7x faster than DanceGRPO without degrading alignment. On WanX video genera- tion, further achieves higher motion quality reward with sharper and temporally consistent frames.
INTRODUCTION
Diffuusion and flow-matching models have advanced image and video generation with high fidelity, diversity, and controllability (Ho et al. 2020; Lipman et al. 2022 Liu et al. 2022). However, large-scale pretraining alone cannot ensure alignment with human intent, as outputs often miss aes- thetic, semantic, or temporal expectations. Reinforeement learning from human feedback (RLHF) addresses this gap by directly adapting models toward human- human-preferree preferred outcomes (Ouyang et al,
2022).
Within RLHF, Group Relative Policy Optimization (GRPO) has shown strong stability and scalabil- ity across text-to-image and text-to-video tasks (Liu et al.) 2025a Xue al. 2025). However, when applied to diffusion and flow-matching models, current GRPO variants still face two fundamental bottlenecks: (1) Inefficiency. Standard GRPO adopts sequential rollout design, where each tra- jectory must be independently sampled under both the old and new policies. This incurs O(N - T complexity with denoising steps T and group size N, leading to significant computational redun- dancy and limiting scalability in large-scale image and video generation tasks. (2) Sparse rewards. Existing methods assign a single terminal reward uniformly across denoising steps, neglecting informative signals from intermediate states. This uniform propagation leads to unreliable credit assignment and high-variance gradients, raising the central question: how can we attribute sparse outcome rewards to the specific denoising steps that truly shape final quality?
Equual contribution. Corresponning author.

---

Published as a conference paper at ICLR 2026
DescedaPo hatten-o.0
Am My y m
0.38
GRPO
Step
BranchGRPO
0.36
Sup 1-1
Reward
0.34
Gwead
0.32
100
150
300
Training Stops
Figure 1: Comparison of BranchGRPO and DanceGRPO. Left: Reward curves during training. BranchGRPO converges substantially faster, achieving up to 2.2x speedup over DanceGRPO (time fraction 1.0) and 1.5x speedup over DanceGRPO (time fraction 0.6), while ultimately sur- passing both baselines. The time fraction 0.6 variant further exhibits pronounced instability. (time fraction denotes the proportion of diffusion timesteps used during training.). Right: Illustration of rollout structures. GRPO relies sequential rollouts with only final rewards, whereas Branch- GRPO performs branching at intermediate steps and propagates dense rewards backward, enabling more efficient and stable optimization.
To overcome these limitations, propose BranchGRPO, a tree-structured policy optimization framework for diffusion and flow models. BranchGRPO replaces inefficient independent sequen- tial rollouts with branching structure, where scheduled split steps in the denoising process allow each trajectory to stochastically expand into multiple bub-trajectories while reusing shared prefixes. This design amortizes computation across common segments and aligns naturally with the stepwise nature of denoising, substantially improving sampling efficiency while reducing computational cost.
The tree structure further enables novel reward fusion with Septh-wiss advantage estimation. In- stead of uniformly propagating a single terminal reward, BranchGRPO aggregates leaf rewards and propagates them backward with depth-wise normalization, producing finer-grained step-level ad- vantages. In addition, width- and depth-pruning strategies remove redundant branches and depths during backpropagation, accelerating training and reallocating computation toward promising re- gions of the trajectory space.
We validate BranchGRPO on both text-to-image and image-to-vides alignment tasks, demonstrating its effectiveness and generality across modalities. In addition, we verify the scaling law of Branch- GRPO, larger group sizes consistently lead to better alignment performance.
Our contributions threefold:
We introduce BranchGRPO, tree-structured GRPO training paradigm. It replaces in- dependent sequential rollouts with branching during denoising, reusing shared prefixes to amortize compute and broaden exploration, thereby improving efficiency and scalability.
We propose new reward fusion and depth-wise advantage estimation method that converts sparse terminal rewards into dense, step- signals, yielding more stable optimization.
We design complementary width- and depth-pruning strategies that lower backpropagation cost and further improve alignment.
2 RELATED WORK
Diffusion models (Ho al. 2020 Rombach al. 2022) and flow matching models (Lipman al, 2022 Liu et al. 2022) have become dominant paradigms for visual generation due to their strong theoretical foundations and ability to generate high-quality content efficiently. While pretraining es- tablishes the generative prior, aligning outputs with nuanced human preferences requires reinforce- ment learning from human feedback (RLHF). In natural language processing, RLHF has proven highly successful for aligning large language models (LLMs) (Ouyang et al. 2022 Christiano al.) 2017 Lu et al.) 2025), where methods such as PPO and GRPÓ enable stable preference-driven post-training. These successes have inspired adaptation of RLHF vision.
2

---

Published as a conference paper at ICLR 2026
DanceGRPO Sequential Rollout
Sampled Images Distribution
BranchGAPO
Rollout
BezedinGROO DanaCarg
Figure 2: Comparison of sequential and branch rollouts. Left/Right: example generations from DRPceGRPO and BranchGRPO, respectively. Middle: distribution of sampled images projected into 2D feature space, where red and blue dots correspond DanceGRPO and BranchGRPO.
In the visual domain, RLHF for diffusion has been developed along two main directions hima et al. 2025). Reward-modd boae---eddel approaches such as ImageReward (Xu et al. (223)) back-
propagate learned rewards through the denoising process. Direct Preference Optimization (DPO)
(Rafailov et al., 2023 Liu et al. (222c) has also been extended to diffusion, leading to Diffusion- DPO (Wallace et al. 2024) and Videodpo (Liu et al., 2025c), which achieve competitive alignment
without explicit reward modeling. Policy-gradient formulations such DDPO (Black et al. 2023)
and DPOK (Fan et al. 2023) further explore online optimization but often face stability challenges. Meanwhile, standardized reward models including HPS-v2.1 (Wu et al. 2023). VideoAlign (Liu
ct al. 2025b) enable systematic comparison of alignment algorithms on image and video tasks.
More recently, Group Relative Policy Optimization (GRPO) (Shao et al. 2024) has been intro- duced as a scalable alternative to PPO for preference optimization. DanceGiPO (Xue et al. 2025) and Flow-GRPO (Liu et al., 2025a) pioneers the application of GRPO to visual generation, unify- ing diffusion and flow models via SDE reformulation and demonstrating stable optimization across text-to-image, text-to-video, and image-to-video tasks. TempFlow- ther highlights the limitation of sparse terminal rewards with uniform credit assignment, proposing temporally-an
(He et al. 2025) fur-
weighting across denoising steps. MixGRPO (Li et 2025a) further enhances bodding-window efficiency with a mixed ODE-SDE sliding window scheme, though it still faces trade-offs between overhead and performance. Our work continues this line by introducing BranchGRPO, which lever- ages branching rollouts, depth-wise reward fusion, and structured pruning to improve both stability and efficiency; while related TreePO in LLMS (Li at. 2025b), our method adapts rtr-sstuetured rollouts specifically to diffusion dynamics.
3 METHOD
3.1 DOES BRANCH ROLLOUT HARM DIVERSITY?
natural concern with branch rollouts that reusing shared prefixes might reduce sample diver- sity. To investigate this, we generate 4096 samples each from DanceGRPO and BranchGRPO and evaluate their distributions across multiple feature spaces. Figure 2 provides both qualitative and quantitative evidence: the left/rght panels show representative generations, while the middle panel visualizes the sampled distributions in 2D feature embedding, where DanceGRPO and Branch- GRPO points largely overlap.
Quantitatively. in the Inception feature space, the distributions remain close, with KID (Binkowski ett al. 2018 20188000055 =00057 and MMD (Gretton et al. 2012 In the CLIP feature space (ViT- B/32 Radford et al., 2221)). the similarity is even stronger: KID=0.00022 and MMD = 0.0149, both indicating that the two distributions are almost indistinguishable the semantic level.
3

---

Published as a conference paper at ICLR 2026
Branch Rollout
Reward Fusion
Depth-wise Adv. Pruning
Pruned node
1:35
Figure 3: Left: branch rollout process. Middle: leaf rewards are fused upward. Right: depth-wise normalization and pruning yield dense advantages and reduce cost.
Taken together, Figure and these statistics demonstrate that branch rollouts preserve distributional and semantic diversity, introducing at most negligible shifts across different feature spaces.
3.2 BRANCH ROLLOUT ALGORITHM
Preliminaries. Given prompt, BranchGRPO reformulates denoising into tree-structured pro- cess. We align terminology with tree search: (i) depth T denotes the number of denoising steps; (ii) width w is the number of completed trajectories (leaves); (iii) branching steps indicate split timesteps; (iv) branch correlation controls the diversity among child nodes; (v) branching factor K is the number of children per split.
Branch sampling. Unlike prior GRPO variants such as DanceGRPO and FlowGRPO that rely on sequential rollouts, where each trajectory is sampled independently from start to finish, Branch- GRPO reorganizes the process into tree-structured rollout (Figure D For each prompt, we ini- tialize a root node with a same initial noise 2u ~ NNO,, I) and then denoise step by step along the reverse SDE. At designated split steps B, the current state expands into K children, producing multiple tujetrajectorrie that share early prefixes but diverge afterward. The branching is achieved by injecting stochastic perturbations into the SDE transition, with hyperparameter controlling the diversity strength among child nodes. This mechanism balances exploration diversity and sta- bility while keeping the marginal distribution unchanged. The rollout continues until reaching the maximum depth T', at which point all leaves are collected for reward evaluation.
Formally, following Xue et al. (2025), the reverse-time dynamics can be written in the SDE form:
da,
[ee
its gi log) di + StDt dwy;
where e, controls stochasticity.
At split step i B with step size h4 - 1,41, instead of sampling single successor we generate K correlated children:
2 = (2) + It. / Ths Eo.
(b)
E
Eo + s f/b
1,....,, K,
vI 1
Lid. (O, 1). with En shared across branches and ?/s denoting branch-specific innova- tions. The parameter s x20 0 tunes inter-branch correlation: small yields highly correlated but stable branches, while large s makes branches nearly independent. By construction, each &a ~ (), (b) so every child 2441 has the same marginal distribution as the baseline SDE step.
where fo. Tb
3.3 REWARD FUSION AND DEPTH- WISE ADVANTAGE ESTIMATION
Branch rollouts form trajectory tree with shared ancestral prefixes, allowing internal node values to be expressed by descendant rewards and enabling backward propagation of leaf signals. However, existing GRPO variants use a single terminal reward at every step, ignoring intermediate states and yielding high- high-variance, unreliable credit assignment (Fig. 3. BranchGRPO addresses this by propagating leaf rewards upward and converting them into dense step-level advantages via path- probability fusion and depth-wise normalization.
4

---

Published as a conference paper at ICLR 2026
Reward fusion. We design dynamically adjustable reward fusion scheme that aggregates leaf rewards into internal node values through soft weighting mechanism. For an internal node n with descendant leaves [(n).
exp(Bse)
F(n) = Σ
(n) We re,
oe = EjE(()) exp(Bs;)) 88 = log Poet n2).
(3)
BEL(N)
Here Pbeh is the behavior policy and 8 controls concentration. _ = 0 reduces to uniform averag- ing; when = 1, the fusion reduces to weighting by the behavior policy distribution poeh (e n)- Uniform averaging is robust to log-prob calibration errors and encourages exploration by retaining low-probability leaves, but introduces variance when branches are many. Path-weighting reduces variance and stabilizes updates, though it may over-concentrate on high-likelihood leaves in deep trees. We empirically compare both variants Sec. 4.4
Depth-wise normalization. Nodes the same depth share the same noise level and thus di- rectly comparable, while rewards across depths vary drastically due to changing noise states. To balance gradient contributions, we normalize aggregated rewards within each depth d:
Aa(n) = F(n) pd +
Na m, F(n), ad = stdn,N,N,
(4)
where Na denotes all nodes at depth d. Each edge advantage A(e) inherits from its child node and is optionally clipped to [-AA A-Amax: Amax]. This per-depth standardization prevents late denoising steps with smaller variance from dominating, yielding process-dense and balanced credit signals. Compared to GRPO's prompt-level normalization, which broadcasts a single terminal reward, our scheme produces stable gradients and finer credit assignment to the denoising steps that matter.
We optimize the standard clipped GRPO loss over tree edges:
J(0) = E
Σ 333
min (p)(O)((e), clip(p. - 1-,, (,),,)))(((())
(5)
where an edge denotes transition (st, at) at depth f, & is the set of such edges in behavior tree, and (f( To (a; | st).
3.4 PRUNING STRATEGIES
While branch rollouts improve efficiency and provide dense process rewards, an excessive number of branches may induce exponential growth in trajectory count, leading to prohibitive backpropa- gation cost. To further accelerate training, introduce the context of BranchGRPO: width pruning, which reduces the number of leaves used for gradient updates, and depth pruning, which skips unnecessary denoising steps.
complementary pruning strategies in
Importantly, pruning is applied only after reward fusion and depth-wise normalization, and affects backpropagation but not forward rollouts reward evaluation. This design ensures that all trajecto- ries contribute to reward signals, while gradients are computed only for the selected subset.
Width Pruning. After computing rewards and advantages for all leaves L, we restrict gradient up- dates to a subset of them. We investigate two modes. The first, Parent-Topl, retains the child with the higher reward from each parent the last branching step. This strategy roughly halves gradient computation while ensuring coverage of all parents, yielding stable but slightly less diverse updates. The second, Extreme selection, preserves both the globally best and worst ò leaves by reward score. This explicitly maintains strong positive and negative signals, which may enhance exploration but also increase variance.
Depth Pruning. Branch rollouts generate dense rewards across all denoising steps, but computing gradients at every depth remains costly. To improve efficiency, we introduce depth pruning, which skips gradient computation at selected timesteps Concretely, maintain set of pruned depths D and ignore gradients from nodes at these layers. To prevent permanently discarding certain steps, we adopt a sliding window schedule: the pruned depths gradually shift toward later timesteps training progresses, until reaching a predefined maximum depth. Formally, pruning is active throughout training, and at fixed intervals the window slides one step deeper until the stop depth is reached.
5

---

Published as a conference paper at ICLR 2026
Algorithm BranchGRPP Training Process
Require: dataset C; policy Te; behavior policy Tew reward models {R denoising steps T;
1: for iteration branching steps B; branching factor K; branch correlation to .M do
2: TeaM f- Te
Sample batch C. C C
4: for prompt E C, do
5: Sample root noise 20 (0, Build rollout tree T with T for / = 7
6:
7:
8:
if I E 8 then
9:
Branch into K children with correlation s
10:
else
11:
Sinple-step denoising
12:
end if
13:
end for
Evaluate rewards for leaves C(7)
15:
Reward fusion: aggregate leaf rewards upward (path-prob. weights)
Depth-wise normalization: standardize per depth, assign edge advantages A(e)
17:
Pruning: select nodes for backprop only
Form edge set (c)) from tree T
19:
Compute J(O) (clipped-GRPO over e Ec)), averaged by E(())
Update policy: 0 + J()
21:
end for
22:
end for
EXPERIMENTS
4.1 EXPERIMENT SETUP
We evaluate BranchGRPO on HPDv2.1 (Wu et al.] 2023) (103k training and 400 balanced test
prompts). The backbone is FLUX. 1-Dev (Black Forest Labs 2024), SD3.5-M Esser al. (2024) baselines include DanceGRPO and MixGRPO under identical settings. We report efficiency (NFE. iteration time) and quality (HPS-v2.1, PickScore (Kirstain et al. 2023). ImageReward (Xu et al, 2023 Unified Reward (Wang et al. 2025
4.2 IMPLEMENTATION DETAILS
We set the tree depth to d 20 and the branch factor to K 2, yieldingg 16 leaves per rollout before pruning. The branching steps B use three presets: Dense 30,3,6, 6, as the default, Mixed (0, 4, 4, 12), and Sparse (0, 10, 15). The branch correlation sweeps for 300 optimizer steps with gradient accumulation y 12 and per-GPU batch size 2, on 16x NVIDIA H200 GPUs Optimization with bél6 precision and EMA weights stored on CPU. All GRPO-related hyperparameters are kept identical across methods, with full details deferred the supplementary material.
21, 4) Training runs
A (learning rate 10-3, weight decay 1 10-4)
4.3 MAIN RESULTS
summarizes efficiency and alignment performance. BranchGRPO consistently outperforms baselines across cuuman-prefereee metrics while offering favorable compute trade-offs. In particu- lar, SamnchGRPO-DeptPPruttti
Table
achieves the best overall alignment, raising HPS-v2.1 from 0.360 (DanceGRPO) to 0.369 and delivering the highest PickScore (0.231), ImageReward (1.625), and further reduce itera-
Unified Reward (3.404). BranchGRPO- SiaecGGGPO-WiddPPuunnnnn WidthPruning and BranchGRPO-MM tion time to 314s and 148s, respectively, with only marginal drops in quality quality-making them highly practical for lcrge-scale training. Compared with MixGRPO (289s, SPPSvv.110.33 Unified Re- ward=3.380), BranchGRPO variants yield both stronger alignment and more flexible scaling.

---

Published as a conference paper at ICLR 2026
inspined by greenn pose
Cartaonish ilustration machire shop shipping containge.
onarge standing its legs
with open.
DarceG
DunceSRPO
BranchGRPO(Curs
Frontal portrait of queen elizobeth.
standing forest, stylized.
Have a tortoise deliver your package for you, don't be surprised if sings "Hoppy Birthd back
a tiny mouse on tunnel.
SD3.5-M
FlowGRPO
BranchGRPO
SD3.5-M
FlowGRPO
BranchGRPO
Figure 4: Qualitative comparison of generations. Top: Flux, DanceGRPO, and our BranchGRPO. Bottom: SD3.5-M, FlowGRPO, and our BranchGRPO.
Reward curves in Figure confirm these findings: (fnc)G suffers from instability, while the full-timestep variant converges smoothly but at high cost. BranchGRPO achieves faster early reward growth, smoother convergence, and higher final rewards. Qualitative compar- isons in Figure further show that method produces sharper details and better semantic align- ment than Flux and DanceGRPO.
Table 2 further validates the generality of BranchGRPO on the SD3.5-M backbone. When plugged into GRPOO BranchGRPO reduces the total GPU-hours (from 2000 to 1460) while keeping the same number of training steps. Despite using less than half the compute, BranchGRPO consistently improves all aligmment metrics HPS-vPPS22--- 1,.1, PickScore, ImageReward, and GenEval. These results show that BranchGRPO transfers cleanly across different backbones and GRPO-style pipelines, providing both higher training efficiency and stronger alignment quality under the same or even lower- compute budgets.
Reward-KL Efficiency. Beyond final re- ward and training speed, we also examine the reward-KI. tradeoff, which measures how efficiently RL policy converts KL diver- gence into reward gains. As shown on the right, BaanncGRRO-DeephPPPi lies above the DanceGRPO frontier across the full KL range. This indicates that BranchGRPO extracts more reward per unit KL, reflecting more stable credit assignment better opti- mization efficiency.
3.2
Fusildiit Deya PPhutin
consistently
3.3
Figure 5: Reward-KI. curves comparing BranchGGRPO-DepthPPunnng and DanceGRPO.

---

Published as a conference paper at ICLR 2026
Table 1: Efuiciency-quality comparison. The best and second-best results in each column are high- lighted in bold and underline, respectively. NFE denotes the number of function evaluations of the denoiser. For branching methods, we report the average per-sample NFE, computed as the total function evaluations in the tree divided by the number of final samples.
Method
NF
NFE.
Iteration Time (8)
HPS-v2.11
Pick Scoref
Image Rowardf
Unified Roward
FLUX
0.313
0.227
1.112
3.370
DanceGRPOrd=-L0)
20
20
698
0.360
0.234
1.612
3,388
DanceCiPPQUE=066)
0.353
1,517
3.362
MixGRPO (20.5)
20
289
0.359
0.228
1.594
3.380
BranchGRPO
13.68
13.68
493
0.363
0.229
1.603
3,386
BranchGRPO-WidPP
13.68
8.625
0.364
0.231
1.609
3.383
BranchGRPO-Dee DepPru
13.68
8.625
314
0.369
0.235
1.625
3.404
Mix
13.68
3.384
Table 2: Generalization on SD3.5-M and integration into GRPO-style training pipelines. Branch- GRPO consistently improves alignment quality and training efficiency.
Method
GPU Hours
HPS-x2.1
Pick Scoref
Image Reward
GenEvall
SD3.5-M
0.204
20.51
0.85
0.63
FlowGRPO
2000
0.316
23.50
1.29
0.86
FlowGRPO
1000
0.280
22.41
0.95
0.73
BranchGRPO
1460
0.323
23.58
1.32
0.89
4.4 ABLATION STUDIES
We conduct a series of ablation studies better understand the design choices in BranchGRPO. Unless otherwise stated, all are carried out under the same training setup as in Section The fol- lowing analyses highlight how different branching configurations and aggregation strategies affect efficiency, reward quality, and stability.
Branch Correlation. Figure o) shows the effect of varying the branch correlation s. Smaller values (s = 1.0, 2.0) limit exploration and lead to slower reward growth, while very large values (s 8.0) destabilize early training. A moderate setting (s = 4.0) achieves the best trade-off, reaching the highest reward and stable convergence, confirming that stochastic branching is necessary but should be carefully tuned.
Branching Steps. We next vary the positions of split timesteps (Figure b(b)). Early splits such as (0, 3, 6, promote faster reward increase in the early stage, whereas later splits like (9, 12, 15, 18) delay exploration and yield lower rewards. Intermediate schedules (e.g., (3, 6, 9, 12)) balance effi- ciency and reward quality, suggesting that early splits are generally more effective for exploration.
Branch Density. Finally, we compare different densities of split points while keeping the overall horizon fixed (Figure 6 c)). Although all configurations eventually converge to similar reward levels, denser splits (e.g.. (0,3, 6, 9)) accelerate early training, while sparser configurations converge more slowly. This indicates that increasing the density of branching in the early phase improves sample efficiency without harming stability.
Branch Rollout Diversity. A potential con- cern is that Branch Rollout may reduce diver- sity. Beyond the distribution-level checks Sec. 3.1, also evaluate prompt-conditioned diversity using LPIPS-MPD and TCE follow- ing prior work Kim al. (2025). As shown on the right, BranchGRPO remains very close DanceGRPO across branching schedules, in- dicating that sample diversity is essentially un- changed.
Table 3: Prompt-conditioned diversity under different branching schedules (new).
Method
LPIPS-MPD 1
TCE 1
DanceGRPO
0.723
4.45
BranchGRPO (0,2,4,8)
0.719
4.44
BranchGRPO (0,3,6,9)
0.713
4.42
BranchGRPO (0.4.8.12)
0.704
4.39
Reward Fusion Strategies. Figure d) compares uniform averaging (8 = 0) with rath-probability weighting (3 1) under identical training settings. Uniform averaging shows higher variance and
8

---

Published as a conference paper at ICLR 2026
(e) Branch Correlation
(3) Branch Steps
(c) Branch Delsity
0.40
-.40
188810
Cor
puandigg 11
2004 002000011
Son
Buaddng -A16/A1
Tewed
0.39
0.36
0.33
032
0.32
0.32
100
155
200
100
200
250
100
150
200
250
Traiiin Steps
Tra Tragiing Steps
Triniing Stepp
(d) Revand Fiusion Stretage
(c) Pruning Strotogics
(0) Iylrnd (331-600
wnter slong coney
Asccaerre (191)
C33
0.30
0.32
Roward
0.54
0.34
0.34
6.32
0.90
0.30
0.30
130
30
300
130 200
Trising
Figure 6: Ablation studies of BranchGRPO. Moderate branch correlation, early and denser splits improve reward growth; path-weighted fusion enhances stability; depth pruning achieves the best final reward; and the hybrid ODE-SDE provides the fastest training speed while remaining stable.
clear late-stage plateau, whereas path- weighted fusion delivers consistently higher and more stable rewards throughout training. This empirically supports Sec. 3.3] uniform averaging is exploration- friendly but noisy, while path weighting aligns credit assignment with the behavior distribution and improves convergence at no extra cost.
Pruning Strategies. Figure compares pruning methods applied after depth-wise normalization and only during backpropagation. For depth pruning, we adopt a Blddngg-inndww schedule over denoising steps, the window is initialized at the last split point, has fixed size of 4, and shifts by one denoising step every 30 training iterations. GRPO losses and gradients in the active window are skipped, while forward sampling remains unchanged. This schedule yields the best final re- ward and reveals substantial redundancy at late timesteps. For width pruming -oore,tTTop)), we retain only the locally better child at each branch for gradient updates, effectively halving updates and producing the smoothest, lowest-variance curve, though with slightly lower final reward than depth pruning. Width pruning (Extreme-b) keeps both the globally best and worst 6 leaves, injecting stronger positive/negative signals and remaining competitive at the end, but with higher variance.
Hybrid ODE-SDE. To further explore depth pruning, inspired by MixGRPO (Li et al. 2025a), design a hybrid ODE-SDE schedule: all branching steps are preserved as SDE, while a sliding window determines additional SDE steps, with the remaining updates replaced by ODE. Figure f) shows that this scheme achieves the fastest speedups (148s vs. 289s for MiXGRPO vs. 469s for DanceGRPO while maintaining stable and fast reward growth.
4.5 SCALING WITH BRANCI
DanceGRPO scales poorly: one GRPO training step with 81 rollout samples takes over 3500s, whereas BranchGRPO achieves the same scale only 680s, making large-scale scaling feasible. We investigate two settings: scaling the branch factor (K=2, 3, 4 yielding 16, 81, and 256 leaves) and scaling the number of branching steps (3, 4, splits yielding 8, 16, and 32 leaves under K=2).
As shown in Figure 7 scaling along both dimensions leads to clear and substantial gains in reward growth and final performance. Larger branch factors and more branching steps consistently push the reward curves higher, with improvements becoming increasingly pronounced as the rollout tree expands. This demonstrates that BranchGRPO can effectively generate additional samples, making large-scale scaling both practical and beneficial without compromising stability.

---

Published as a conference paper at ICLR 2026
Method
Group Size
GPU Hours
Reward
Ausily han cosse
DaxeGRPO
16
928
0.360
BranchGRPO
BrancbGRPO
32
368
0.373
BranchGiPO
787
0.381
BranchGiPPO
906
0.387
BrancbCRPPO
256
3072
0.404
Tratnng 2o00
Tea 00g Ste5e
Table 4: Here use BranchGRPO-Min Under similar GPU-hours, BranchGRPO with group size 81 achieves reward of 0.387, substantially outperforning Dance- GRPO (0.360).
Figure 7: Impact of scaling branch rollouts in BranchGRPO. Larger branch factors (a) and more branching steps (b) consistently improve reward, fol- lowing a clear scaling law.
Do und arg
10 Manue 2IP
13
150
Figure 8: Video generation results on Wan2.1-1.3B. Left: qualitative frame comparisons across three settings. Right: reward curves showing faster convergence and higher final rewards with BranchGRPO compared to DanceGRPO.
4.6 VIDEO GENERATION RESULTS
We further evaluate BranchGRPO on video generation using Wan2.1-1.3B (Wan al. 2025). with the Aldeo-Align's motion quality (Liu et al. 2025b) as reward.
As shown in Figure 8 the base model exhibits severe temporal flickering and deformation, while DanceGRPO improves consistency but still produces blurry details. BranchGRPO generates sharper and more coherent frames across time, and the reward curves demonstrate faster convergence and higher final rewards compared to DanceGRPO.
These results highlight that branching rollouts particularly effective for video generation, where reward sparsity and temporal coherence are especially challenging. In practice, BranchGRPO also improves efficiency: each iteration takes only about 8 minutes, compared 20 minutes for Dance- GRPO. Additional visual examples are provided in the supplementary material.
5 CONCLUSION
We introduced BranchGRPO, tree-structured GRPO paradigm that replaces sequential rollouts with prefix-sharing branching and depth-wise reward fusion, augmented by lightweight pruning for compute reallocation. Across image and video generation, BranchGRPO yields faster convergence, more stable training, and higher final alignment quality under matched budgets. These results estab- lish structured branching as a practical, scalable path for RLHF in diffusion and flow models, and we further verify a scaling trend in which larger group sizes consistently improve performance.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
This work was supported by the National Natural Science Foundation of China (No. 6247601 and No. 62576043), and by the Beijing Natural Science Foundation (No. L252060).
ETHICS STATEMENT
This work focuses on improving optimization efficiency and stability for diffusion/flow-based gen- erative models. Our experiments rely on publicly available prompt sets and reward models cited in the paper. No personally identifiable information, human subjects data, or sensitive attributes were collected or annotated by the authors. We followed the licenses and usage terms of all datasets and third-party models; when applicable, we restricted use to non-commercial research. Given that generative models may amplify biases in reward models or datasets, we report failure cases and abla- tions that probe stability and variance, and we avoid deploying the models in safety-critical settings. We will release code and configuration files to facilitate scrutiny and responsible reuse.
REPRODUCIBILITY STATEMENT
We release anonymized training and evaluation code, configuration files, and exact hyperparameters in the supplementary materials. All experiments can be reproduced using the provided scripts, seeds, and environment specifications.
LLM USAGE
An LLM was used only for language polishing of the manuscript (c.g.. grammar, wording, and minor clarity edits). No experimental results were generated, modified, selected by the LLM;; all technical content, claims, ideas, and analyses are the authors' own. No confidential or non-public data were shared with the LLM. The authors take full responsibility for the final content.

---

Published as a conference paper at ICLR 2026
REFERENCES
Mikotaj Binkowski, Danica J Sutherland, Michael Arbel, and Arthur Gretton. Demystifying mmd gans. arXiv preprint arXiv:1801 01 00111 2018.
Kevin Black, Michael Janner, Yilun Du, llya Kostrikov, and Sergey Levine. Training diffusion models with reinforcement learning. arXiv preprint arXiv:2305.133011 2023.
Black Forest Labs. Flux.
2024.
Paul F Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. Advances in neural information processing sys- lems, 30, 2017.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. Forty-first international conference on machine learning, 2024.
Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, and Kimin Lee. Dpok: Reinforcement learning for fine-tuning text-to-image diffusion models. Advances Neural Information Processing Systems, 36:79858-79885 2023.
Arthur Gretton, Karsten M Borgwardt, Malte Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. The journal of machine learning research, 13(1):723777, 2012.
Xiaoxuan He, Siming Fu, Yuke Zhao, Wanli Li, Jian Yang, Dacheng Yin, Fengyun Rao, and Bo Zhang. Tempflow-grpor When timing matters for grpo in flow models. arXiv preprint 1rXiv:2508.04324 2025.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-663 6851, 2020.
Sunwoo Kim, Minkyu Kim, and Dongmin Park. Test-time alignment of diffusion models without reward over-optimization. arXiv preprint arXiv:2501.05800 2025.
Yuval Kirstain, Adam Polyak, Uriel Singer, Shahbuland Matiana, Joe Penna, and Omer Levy. Pick- a-pic: An open dataset of user preferences for text-to-image generation. 2023.
Junzhe Li, Yutao Cui, Tao Huang. Yinping Ma, Chun Fan, Miles Yang, and Zhao Zhong, Mixgrpo: Unlocking flow-based grpo efficiency with mixed ode-sde. arXiv preprint arXiv:2507.21802, 2025a.
Yizhi Li, Qingshui Gu, Zhoufutu Wen, Ziniu Li, Tianshun Xing, Shuyue Guo, Tianyu Zheng, Xin Zhou, Xingwei Qu, Wangchunshu Zhou, et al. Treepo: Bridging the gap of policy optimiza- tion and efficacy and inference efficiency with heuristic tree-based modeling. arXiv preprint arXiv:2508.11 17445, 2025b.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.027477 2022.
Jie Liu, Gongye Liu, Jiajun Liang, Yangguang Li, Jiaheng Liu, Xintao Wang, Pengfei Wan, Di Zhang, and Wanli Ouyang. Flow-grpo: Training flow matching models via online rl. arXiv preprint 2tt.:2255505477.7 2025a.
Jie Liu, Gongye Liu, Jiajun Liang, Ziyang Yuan, Xiaokun Liu, Mingwu Zheng, Xiele Wu, Qiulin Wang, Wenyu Qin, Menghan Xia, et al. Improving video generation with human feedback. arXiv preprint arXiv:2501.13918 2025b.
Runtao Liu, Haoyu Wu, Ziqiang Zheng, Chen Wei, Yingqing He, Renjie Pi, and Qifeng Chen. Videodpo: Omni- preferecce alignment for video diffusion generation. In Proceedings of the Computer Vision and Pattern Recognition Conference, pp. 8009-8019, 2025c.
12

---

Published as a conference paper at ICLR 2026
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning generate and transfer data with rectified flow. arXiv preprint arXiv:2209.0300 2022.
Fanbin Lu, Zhisheng Zhong, Shu Liu, Chi-Wing Fu, and Jiaya Jia. Arpo: End-to-end policy opti- mization for gui agents with experience replay. arXiv preprint N.Ki:22005.6288. 2025.
Yuta Oshima, Masahiro Suzuki, Yutaka Matsuo, and Hiroki Furuta. Inference-time text-to-video alignment with diffusion latent beam search. arXiv preprint arXiv:2501.19255 2025.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to fol- low instructions with human feedback. Advances in neural information processing systems, 35: 27730-27744, 2022.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp. PmLR, 2021.
Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly reward model. Advances in neural information processing systems, 36:53728-53747 2023.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and pattern recognition, pp. 10684-10695, 2022.
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song. Xiao Bi. Haowei Zhang, Mingchuan Zhang, YK Li, Yang Wu, et al. Deepseekmath: Pushing the limits of mathemati- cal reasoning in open language models. arXiv preprint arXiv:2402.03300, 2024.
Bram Wallace, Meihua Dang, Rafael Rafailov, Linqi Zhou, Aaron Lou, Senthil Purushwalkam, Stefano Ermon, Caiming Xiong, Shafiq Joty, and Nikhil Naik. Diffusion model alignment using direct preference optimization. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pp. 8228-8238, 2024.
Team Wan, Ang Wang, Baole Ai, Bin Wen, Chaojie Mao, Chen- Bein-Wei Xie, Di Chen, Feiwu Yu, Haiming Zhao, Jianxiao Yang. Jianyuan Zeng, Jiayu Wang, Jingfeng Zhang. Jingren Zhou, Jinkai Wang, Jixuan Chen, Kai Zhu, Kang Zhao, Keyu Yan, Lianghua Huang, Mengyang Feng. Ningyi Zhang, Pandeng Li, Pingyu Wu, Ruihang Chu, Ruili Feng. Shiwei Zhang, Siyang Sun, Tao Fang, Tianxing Wang. Tianyi Gui, Tingyu Weng. Tong Shen, Wei Lin, Wei Wang. Wei Wang, Wenmeng Zhou, Wente Wang, Wenting Shen, Wenyuan Yu, Xianzhong Shi, Xiaoming Huang, Xin Xu, Yan Kou, Yangyu Lv, Yifei Li, Yijing Liu, Yiming Wang. Yingya Zhang, Yitong Huang, Yong Li, You Wu, Yu Liu, Yulin Pan, Yun Zheng, Yuntao Hong, Yupeng Shi, Yutong Feng, Zeyinzi Jiang, Zhen Han, Zhi-Fan Wu, and Ziyu Liu. Wan: Open and advanced large-scale video generative models. arXiv preprint arXiv: 20132.20344... 2025.
Yibin Wang, Yuhang Zang, Hao Li, Cheng Jin, and Jiaqi Wang. Unified reward model for multi- modal understanding and generation. arXiv preprint arXiv:2503.05236, 2025.
Xiaoshi Wu, Yiming Hao, Keqiang Sun, Yixiong Chen, Feng Zhu, Rui Zhao, and Hongsheng Li. Human preference score v2: A solid benchmark for evaluating human preferences of text-to- image synthesis, 2023.
Jiazheng Xu, Xiao Liu, Yuchen Wu, Yuxuan Tong, Qinkai Li, Ming Ding, Jie Tang, and Yuxiao Dong. Imagereward: learning and evaluating human preferences for text-to-image generation. In Proceedings of the 37th International Conference on Neural Information Processing Systems, pp. 15903-15935, 2023.
Zeyue Xue, Jie Wu, Yu Gao, Fangyuan Kong, Lingting Zhu, Mengzhao Chen, Zhiheng Liu, Wei Liu, Qiushan Guo, Weilin Huang, et al. Dancegrpo: Unleashing grpo on visual generation. arXiv preprint ar0iv:2505. 07818, 2025.
13

---

Published as a conference paper at ICLR 2026
APPENDIX
This appendix provides additional theoretical proofs, implementation details, and experimental re- sults that complement the main paper. The contents are organized as follows:
Section A
Hyperparameter settings used in all experiments.
Section Theoretical analysis, including proofs for branch noise construction, reward fusion, and variance reduction.
Section C Additional experiments, including ablations, more text-to-image and image-to-
video results. Failure Cases: Qualitative examples where BranchGRPO fails, highlighting current limitations.
Section D
Discussion and future work.
A HYPERPARAMETER SETTINGS
Table 5 summarizes the detailed hyperparameter configuration used in our experiments. All hyper- parameters kept identical across all methods, including DanceGRPO and MixGRPO, to ensure a fair comparison. For depth pruning and hybrid- ODE-SDE, we follow the design of MixGRPO and adopt sliding window of size 4, which shifts one step deeper every 30 iterations.
Table 5: Hyperparameter settings used in all experiments.
Parameter
Value
Parameter
Value
Random seed
42
Learning rate
x 10
Train batch size
2
Weight decay
10-4
SP size
Mixed precision
bfloat16
SP batch size
2
Grad. checkpointing
Enabled
Dataloader workers
4
Max grad norm
0.01
Grad. accum. steps
12
Warmup steps
Checkpoint steps
40
Use TF32
Yes
Resolution
720 x 720
Sampling steps
16
Eta
0.3
Sampler seed
1223627
Num. generations
12
Shift (branch offistt)
Use group reward
Yes
Ignore last step
Yes
Clip range
1 x 10
Adv. clip max
5.0
Use EMA
Yes
EMA decay
0.995
Init same noise
Yes
Table 6: Recommended defauutt berpaarreeeter settings for BranchGRPO. These values form a unified and robust configuration that works across all tasks (image/video) and backbones (Flux, SD3.5-M).
Component
Parameter
Defanlt
Stable Range
Branching
Branch points
(0,3.6.9)
Branching factor
(2,3,4)
Correlation
Correlation scale (s)
(3.0-5.0)
Reward Fusion
Temperature (8)
or
Pruning
Pruning type
depth pruning
depth width pruning
Depth window size
Hybrid ODE-SDE
Mix ratio
30%
To further guide practitioners, identify optimal configurations based resource priorities:
14

---

Published as a conference paper at ICLR 2026
Scenario 1: Maximum Performance (Scaling Mode).
If the goal is to push the upper bound of alignment performance, we recommend using BranchGRPO-Mix and scaling the group size to N 81 (c.g., Branch Factor K 3, Number of Branching Steps 4).
Benefit: Due to the efficiency of BranchGRPO-Mix, this configuration maintains an iteration time and total GPU-hour budget similar to standard DanceGRPO/FlowGRPO (at . = -6).
Outcome: Under this similar compute budget, it yields significant performance gains (HPS 0.360 + 0.387) by leveraging massive exploration that is computation- ally prohibitive for sequential baselines.
Scenario 2: Balanced Efficiency & Quality (Universal Default).
If the goal is accelerate training while maintaining slightly improving quality, recommend BranchGRPO-Depthrruin
with the default settings in Table
Benefit: It delivers a 2.2 x speedup (314s vs. 698s per iteration) compared to the
baseline.
Outcome: It achieves strictly better alignment (HPS 0.360 0.369) without harming the original model's diversity or capacity.
15

---

Published as a conference paper at ICLR 2026
B THEORETICAL ANALYSIS
B.1 BRANCH NOISE CONSTRUCTION AND BOUNDARY DISTRIBUTION PRESERVATION
We use reverse-time grid to Consider the reverse SDE discretized by Euler-Maruyama:
t so that = t++ >
2++1 (1( t.) + go 7ii
Th ~ ~0,11,, gi :- 1==g(ti)V/
(6)
At split step E B, we construct K branch noises
E
Eo + 878 v1 + 82
3d4 Eu. Th
N I),
(7)
with fresh (Eo, {mo,) drawn at each split step and no cross-time sharing. Then &o ~ N(0, for each b, and Cov (v) 14.2 for b + b'.
Each child branch then updates as
++1 = Mo(2), + 91 Ea-
(b)
(8)
Lemma (Single-step marginal preservation). For any fixed parent 24 we have
(b)
po(zi,ta) + galhs
Th ~ N(0,1.
Lemma 2 (Leaf marginal preservation). Assuming independent noises across time steps and no cross-time reuse of the shared component. conditioned on prefix (zo, To.....T--11) each branch AN generated by the split rule has the same distribution buseline SDE sample Theorem 1 (Boundary distribution invariance). Under the branching construction above, for any of split steps B, each leaf 2x has the same marginal law baseline SIDE rollout. Hence, branching does not alter the final generator distribution.
(b)
B.2 REWARD FUSION: UNBIASEDNESS AND VARIANCE REDUCTION
Let T.(n) be the leaf set node n. Each leaf has reward re = r(2)9)). Define the conditional expected return
((m) = [[[((N) n].
Uniform fusion.
F(n) = [T(n)] (EL(N)
Σ
(9)
Then
(2(n)
E[F(n) n] = ((n).
Var(f(n) n) [L(n]
where d²(n)= = Var(re n2.
Path-probability weighted fusion.
If leaves are drawn from proposal g(E n) with weights
we Poch (E n)
q( n)
then the IS estimator
4rs (n) =
1
Σ
Wyfe
[[m)]
(EL(n)
unbiased: E(ris(n) [n] = Vnn).. The self-normalized IS estimator
E, WeFe
fsNis(n))
Ee wi
16

---

Published as a conference paper at ICLR 2026
is consistent with variance O(1/E8S), where
(2 we)2 Ee wz
ESS
Remark. practice (Eq. (4) in the main text), we adopt a softmax weighting we x exp(3st) based on path log- prg--obbabilitiss This can be interpreted as a temperature-smootine variant of SNIS, where controls the sharpness importance weights. Although no longer strictly unbiased, this form provides stable training and reduces the influence of low-probability noisy leaves.
B.3 DEPTH-WISE BASELINE (CONTROT. VARIATES)
For K siblings at depth i, let fused returns f(b), and group mean F; = 1 K
Es F(h) Define
1 = ((b) - Ti
Σ (b) =
(10)
Let 9; (0) Ve log Pa (branch b at depth i). Then
1
Σ 9,
(h)
group
K
unbiased gradient estimator with strictly smaller variance than V single 1. E, =()) unless Cov(F(M) (b) O.
(b)
gj
B.4 CONTINUOUS REWARDS AND CONCENTRATION
Assume r is L-Lipschitz and the SDE flow M4+1- is K,-Lipschitz. Then
|m(=/²) r(=N )1 LK.gills - Eor
Thus r(=x0) sub-Gaussian with parameter (T, Averaging over [[()]] leaves gives
Pr (IF(n) - V(n) = < 2 exp (-c [ (1223k))))
C ADDITIONAL EXPERIMENTS
C. MORE QUANTITATIVE RESULT
Table 7: Effect of applying BranchGRPO only to the rollout procedure in DiffusionNFF.
Iter. Time(s)
GPU Hours
Method
HPSSv2.11
PickScoret
DiffusionNFT
159
706
0.331 0.328
23.80
DiffusionNFT + BranchGRPO (rollout only)
497
23.74
Table reports more results.
we conducted small human A/B study (100 prompts, 5 annotators). As shown in Table ??, Branch- GRPO is preferred over DanceGRPO 47% of comparisons, consistent with its higher HPSv2.1 score. This verifies that the observed improvements are reflected in actual human preference rather than reward-model.
17

---

Published as a conference paper at ICLR 2026
Table 8: Ablation study of BranchGRPO under different design choices. Best and second-best per column are in bold and underline. All results are obtained under the same training setup as in Section
Configuration
NFEx Paid
NFE.
HPS-v2.1
Pick Scoret
Imuge Reward
Time (8)
Scoret
Branch Density
(0,3,6,0 (0,4, 4, 12)
13.68
13.68
493
0.363
0.229
1.603
0.374
11.56
11.56
416
0.359
0.229
1.594
0.374
(0, 10,
9.44
9.44
340
0.354
0.228
1.565
0.366
Reward Fusion
Uniform Fosion Path Weighted Fasion
13.68
13.68
493
0.361
0.229
1.595
0.368
13.68
13.68
493
0,363
0.229
1.603
0.374
Pruning Strotegy
Width Pruning Depth Pruning
0.364 0.369
0.231
1,009
0.374
13.68
8.625
314
0.235
1.625
0.381
Table 9: Human preference and Best-of-N scaling results.
(a) Human preference evaluation on 100 prompts
(b) Best-of-N (BoN) Reward Scaling on HPS-v2.1.
with 5 annotators.
Method
N-1
No2
No4
N-8
Nw16
Method
Flux
DanceGRPO
BranchGRPO
DaneeGRPO
0.3524
0.3581 0.3666
0.3626 0.3706
0.3663
0.3670
User Study
19%
33%
48%
BranchGRPO
0.3613
0.3737
0,3742
Table 10: GenEval Result.
Overall
Single Obj.
Two Obj f
Counting f
Colors
Pocition
SD3.5.M
0.63
0.98
0.78
0.50
0.81
0.24
FlowGRPO
FlowGRPO
0.73
0.99
0.85
0.77
0.81
0.67
1000
BranchGRPO
0.90
Table 11: BranchGRPO on Qwen-Image.
We report GPU-hours, per-iteration time, and four
alignment metrics.
Method
GPU
HPPS f
PickScore
ImageReward
GenEEval
Qwen- Image (Base) +DanceGRPO
0.263
22,21
0.87
0.87
389
584
0.318
23,03
1.13
0.89
+BramcGGRFO-Deeetu
184
276
0.321
23.17
1.19
0.90
Table 12: CLIPScore analysis under HPS-only and multi-ohjective training. All HPS-only GRPO variants reduce prompt adherence; simple multi-objective setup partially recovers CLIPScore while maintaining strong HPS.
Method
Reward Model
HPSv2 (1)
CLIP Score (t)
Flux (Base Model)
0.313
0.405
DanceGRPO(tE=1.0))
HPS-v2.1 only
0.360
0.371
DanccGRPO((E-0666
HPS-v2.1 only
0.353
0.375
BranchGRPOY(.33669)
HPS-v2. only
0.363
0.369
BranchGRPPO0.48..12)
HPS-v2.1 only
0,359
0.371
BranchGRRPrrO..10.155
HPS-v2.1 only
0.354
0.376
DanceGRPO((=-066)
HPS-v2.1 + CLIP Score
0.334
0.398
BranchGRPO(0336699
HPS-v2.1 + CLIP Score
0.339
0.392
18

---

Published as a conference paper at ICLR 2026
C.2 MORE TEXT2IMAGE RESULTS
An anime man in flight uniform with hyper detailed digital artwork and an art style inspired by Klimt, Nixeu, lan Sprigger, Wlop, and Krenz Cushart.
Flux
DanceGRPO
BranchGRPO(Ours)
A raccoon riding oversized fox through forest in furry art anime still.
Flux
DanceGRPO
BranchGRPO(Ours)
Totem pole made out of cats.
Flux
DanceGRPO
BranchGRPO(Ours)
Figure 9
19

---

Published as a conference paper at ICLR 2026
A portrait of two women with purple hair flying in different directions against a dark background.
Flux
DanceGRPO
BranchGRPO(Ours)
mmnl
Flux
DanceGRPO
BranchGRPO(Ours)
A cute anime schoolgirl with sad face submerged in dark pink and blue water, portrayed in oil painting style.
Flux
DanceGRPO
BranchGRPO(Ours)
Figure 10
20

---

Published as a conference paper at ICLR 2026
The image of a raccoon wearing Peaky Blinders hat, surrounded by swirling mist rendered with fine detail.
Flux
DanceGRPO
BranchGRPO(Ours)
Portrait of an anime princess in white and golden clothes.
Flux
DanceGRPO
BranchGRPO(Ours)
A cute little anthropomorphic Tropical fish knight wearing a cape and crown in short, pale blue armor.
Flux
DanceGRPO
BranchGRPO(Ours)
Figure
21

---

Published as a conference paper at ICLR 2026
A white polar bear cub wearing sunglasses sits meadow with flowers.
Flux
DanceGRPO
BranchGRPO(Ours)
photo of a mechanical angel woman with crystal wings, the sci-fi style of Stefan Kostic, created by Stanley Lau and Artgerm.
Flux
DanceGRPO
BranchGRPO(Ours)
Close-up shot of person running on a treadmill with worn running shoes under dramatic lighting and a comic book-style painting effect.
Flux
DanceGRPO
BranchGRPO(
Figure 12
22

---

Published as a conference paper at ICLR 2026
Family assembling missile living room.
Flux
DanceGRPO
BranchGRPO(Ours)
The image depicts alien flowers and plants surrounded by visceral exoskeletal formations in front mythical mountains with dramatic contrast lighting, created with surreal hyper detailing 3D render.
Flux
DanceGRPO
BranchGRPO(Ours)
A colorful tin toy robot runs a steam engine on a path near a beautiful flower meadow in the Swiss Alps with mountain panorama in the background, captured in long shot with motion blur and depth field.
M1I
Flux
DanceGRPO
BranchGRPO(Ours)
Figure 13
23

---

Published as a conference paper at ICLR 2026
C.3 MORE IMAGE2VIDEO RESULTS
Table 13: Video evaluation on vBench. We report mean values on 500 samples.
Method
Aesthetic Quality
Background Consistency
Dynamic Degree
Imaging Quality
Motion Smoothness
Iteration Time (s)
Base Model
0.5206
0.9588
0.5150
71.92
0.9784
DanceGRPO
0.5178
0,9647
0.4992
71.94
0,9899
1352
BranchGRPO
0.5190
0.9659
0.5000
71.94
0.9912
493
Witneut G27g
BranchGRPO
(a) Case
Wienoue GA20
DanceGAPO
(b) Case 2
Wethout Garg
DanoscGPo
Branchttr
(c) Case
24

---

Published as a conference paper at ICLR 2026
Without GRPO
DanceGRPO
BranchGRPO
(a) Case
Without
GRPO
DEA
141
S7A
DanceCRPO
BranchGRPO
(b) Case
25

---

Published as a conference paper at ICLR 2026
C.4 FAILURE CASES
A colorful digital painting with front view and anime-inspired vibes featuring magical composition.
Flux
DanceGRPO
BranchGRPO(Ours)
A one-eyed dwarf wizard holding flagon clean cel shaded vector art.
Flux
DanceGRPO
BranchGRPO(Ours)
Australian soldiers surrendering to emu.
Flux
DanceGRPO
BranchGRPO(Ours)
Figure 16: Failure
26

---

Published as a conference paper at ICLR 2026
Head and shoulders portrait Jinx from League of Legends of Arcane animated Series.
Flux
DanceGRPO
BranchGRPO(Ours)
The image is of Pixel Art Huggy Wuggy performing a jumpscare.
aL
Flux
DanceGRPO
BranchGRPO(Ours)
A 3D render volcanic icon on a rocky background, in isometric perspective and darkly lit.
Flux
DanceGRPO
BranchGRPO(Ours)
Figure 17: Failure case
27

---

Published as a conference paper at ICLR 2026
D DISCUSSION AND FUTURE WORK
While our results demonstrate the stability and efficiency benefits of BranchGRPO, several open directions remain.
Discussion. BranchGRPO introduces structured branching and pruning into GRPO training, which we have shown to improve both efficiency and alignment. Our ablations suggest that the choice of branching schedule and pruning strategy can substantially affect reward stability, high- lighting the importance principled tree design. Moreover, reward fusion provides stable gradients in practice, but its bias-variance tradeoff under different weighting schemes warrants further theo- retical analysis.
Future Work. Several promising directions extend beyond the present scope. (1) Dynamic branching. Instead of fixed hyperparameters, one can design adaptive policies that adjust branch fac- tor, correlation, or pruning windows on-the-fly based on sample difficulty or intermediate rewards, enabling more efficient rollouts.
(2) Beyond diffusion models. The branching framework could naturally transfer to other generative paradigms, including diffusion- LLMs and multimodal foundation models. (3) Scaling to long-horizon video. While initial experiments on WanX-1.3B 12V show benefits, more extensive validation on high-resolution, long-duration video generation tasks is required. (4) Robotics and action generation. Tree-structured rollouts are naturally suited to robotics, where intermediate states provide dense and verifiable rewards (e.g., task success signals). Extending BranchGRPO to robotic action generation and embodied video generation learning could open a promising direction for embodied AI.
28


Reviews:
[{"review_id": "UGTLDJI3EH", "reviewer": "Reviewer_UFyj", "summary": "This paper redesigns GRPO for diffusion models into a branching rollout (BranchGRPO) that amortizes computation via shared prefixes, converts sparse terminal rewards into dense stepwise signals through leaf-reward fusion with depth-wise advantages, and reduces only the backprop cost via width/depth pruning. In experiments on image and video alignment, it trains faster and more stably than prior work (e.g., DanceGRPO), improves metrics such as HPS v2.1, and substantially shortens iteration time while maintaining or improving diversity and temporal consistency.", "questions": "Please address the above concerns with concrete clarifications.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well organized and easy to follow: the problem setup, methodology, and objectives are presented coherently, with clear sectioning and notation that make the technical ideas accessible.\n- By redesigning sequential rollouts into branching rollouts, it amortizes computation via shared prefixes while preserving search breadth, which has novelty.\n- With selective backpropagation (width/depth pruning) that keeps forward generation and evaluation intact while reducing only the learning compute, it enables flexible speed–quality trade-offs, so it has usefulness.", "weaknesses": "The proposed method introduces multiple new degrees of design freedom, e.g., branching positions/counts, branching factor, branch-noise correlation, reward-fusion temperature, pruning choice, depth window and its sliding schedule, and the Hybrid ODE–SDE mixing ratio. Although the paper includes ablations of these components, it does not provide practical guidance for selecting them when tasks or compute budgets change. As a result, there is a concern about benchmark-specific overfitting and variability in reproducibility.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "F373HaO5k7", "reviewer": "Reviewer_gDg6", "summary": "BranchGRPO introduces a new reinforcement learning alignment framework tailored for Diffusion and Flow-Matching models, addressing the instability and inefficiency of prior approaches like DanceGRPO. Its key innovation, “Structured Branching,” enables multiple branches of sampled trajectories to share common prefixes, allowing computation reuse and faster convergence. Combined with techniques such as Reward Fusion and lightweight Pruning, BranchGRPO achieves significant gains in training efficiency, optimization stability, and alignment quality, while maintaining scalability across larger branching factors and deeper training configurations.", "questions": "* In Section 3.1, the authors address the question “Does Branch rollout harm diversity?” and provide quantitative results. Specifically, they project the generated samples into both CLIP and Inception feature spaces and compare the distribution similarity between BranchGRPO and DanceGRPO using KID and FID metrics. While the authors report specific numerical values, it remains unclear whether these values are sufficiently small to substantiate that the two distributions are indeed close. Can the authors provide further justification supporting this claim?\n\n* Building upon the first point, it seems intuitive that the diversity of generated samples may depend on the branching stage. In particular, if branching occurs earlier, the resulting samples might exhibit greater diversity, whereas later branching could limit variation. Have the authors conducted any analysis or observed patterns that validate this hypothesis?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper is well-written and clearly structured, with a well-motivated approach that effectively addresses the inefficiency and instability challenges in existing diffusion model alignment methods.\n\n* The structured branching design achieves approximately a 4.7× training speedup and smoother optimization through depth-wise normalization and reward fusion, leading to superior alignment stability.\n\n* The proposed method surpasses DanceGRPO in both image and video generation quality (e.g., higher HPS-v2.1, ImageReward, and Unified Reward scores) and demonstrates strong scalability, maintaining predictable improvements under larger branch factors and depths.", "weaknesses": "* The current validation is primarily conducted on the FLUX.1-Dev model. To demonstrate the generality and robustness of the proposed framework, it would be valuable to include experiments on additional architectures (e.g., Qwen-Image or other text-to-image backbones).\n\n* The evaluation primarily relies on HPS-v2.1, PickScore, and ImageReward. Incorporating more comprehensive and multidimensional benchmarks such as GenVal could provide deeper insights into alignment quality, compositional reasoning, and spatial understanding.\n\n* The manuscript contains several typos (e.g., redundant brackets in Line 67). These should be carefully reviewed and corrected in the revised version.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OI1b5oBUNV", "reviewer": "Reviewer_gfya", "summary": "BranchGRPO restructures GRPO for diffusion models into a branching rollout to simultaneously remove inefficiency and unstable credit assignment caused by sparse terminal rewards. By branching at intermediate denoising steps while sharing prefixes, it preserves exploration diversity with far less redundant computation, and by propagating leaf rewards upward into dense, step-level advantages, it stabilizes optimization. Lightweight width/depth pruning further reduces backprop cost without affecting forward exploration. Experiments show faster convergence and higher or comparable alignment than DanceGRPO across image and video tasks, improving both efficiency and final quality.", "questions": "Please see the weakness section.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper is clearly written and well-presented, making the motivation, problem formulation, and key ideas easy to follow. The proposed BranchGRPO achieves a favorable efficiency–performance trade-off: it reduces computation relative to the baseline DanceGRPO (as evidenced by lower NFE and shorter iteration time) while maintaining or even improving alignment metrics. Finally, the method is validated not only on image generation but also on video generation, demonstrating its effectiveness and generality across modalities.", "weaknesses": "1) Concern about generational diversity.  \nAlthough Sec. 3.1 presents figures and KID/MMD distances (Inception/CLIP space) to claim that branching does not harm diversity, these metrics only measure distributional proximity, not conditional diversity nor modality-level mode coverage across prompts. Since branching may suppress global structure formation at high-noise regimes (where large-scale semantics are determined), the risk of mode collapse or reduced coverage cannot be ruled out. A direct and prompt-conditioned evaluation of diversity (rather than distributional closeness alone) would strengthen the claim. For evaluating diversity more directly, the evaluation methods in [1] may serve as a helpful reference.\n\n2) Questions on evaluation validity.  \n- For image experiments, no human evaluation is reported, making it unclear whether the improvements reflected by proxy rewards correlate with gold human preference. \n- For video experiments, the observed improvement might reflect reward hacking rather than a genuine multi-objective gain. Prior work has shown that video quality metrics often exhibit non-trivial trade-offs [2, 3], demonstrating that consistent improvement along multiple independent axes (e.g., visual quality, text alignment, dynamic degree) would be significant. \n\n3) Limited experimental coverage\n- Despite Sec. 2 acknowledging the existence of many alternative alignment approaches, the experiments primarily compare against DanceGRPO and MixGRPO. It would be informative to include other competitive baselines. \n- The tested backbones are also limited; validating the method on more recent or diverse models such as Qwen-Image [4] or SD3.5 [5] would support claims of generality. \n- Finally, almost all rewards used are human-preference proxies. It remains unclear whether the method is applicable to other classes of rewards (e.g., object count or visual text rendering, as explored in [6]).\n\nReferences \n\n[1] Kim et al., Test-time Alignment of Diffusion Models without Reward Over-optimization, ICLR 2025.  \n[2] Liu et al., VideoDPO: Omni-Preference Alignment for Video Diffusion Generation, CVPR 2025.  \n[3] Oshima et al., Inference-Time Text-to-Video Alignment with Diffusion Latent Beam Search, NeurIPS 2025.  \n[4] Wu et al., Qwen-Image Technical Report, arXiv:2508.02324.  \n[5] Esser et al., Scaling rectified flow transformers for high-resolution image synthesis, ICML 2024.  \n[6] Liu et al., Flow-GRPO: Training Flow Matching Models via Online RL, NeurIPS 2025.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NSLbutGuyN", "reviewer": "Reviewer_H6yd", "summary": "The paper proposes a dynamic branching and pruning approach that seems to amortize roll-out costs, and improve efficiency. Authors provide detailed mathematical formulation of the proposed procedure as well as a suite of quantitative and qualitative results.", "questions": "- Nowhere in the draft I can figure out which Diffusion model has been used? SDXL? \n- Possible to add video frames of DanceGRPO in Section C.3 as well for consistency?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper is reads well, and is well structured offering a coherent narrative and useful insights. \n- The idea of dynamic branching and pruning is an interesting perspective applied to GRPO, and it seems to have some impact in terms of run-time efficiency.\n- The results seem to be promising, even though the impact might not be as significant. More on this in the next section.", "weaknesses": "- The impact of the proposed idea feels slightly inflated. The results in Table 1, Figure 4 and 7 confirm this point. Could you please elaborate how/why you find the impact (beyond relative decrease in run-time) significant? How should one read and interpret an improvement of $1.594 \\rightarrow 1.625$ or $3.59 \\rightarrow 3.69$? I don't think it also makes sense to accentuate on the iteration time $148$ while the arguments regarding performance are not based upon BranchGRPO-Mix ... Speaking of the results, if possible during rebuttal, I would advise comparing against one or two extra recent baselines (2024/5); as it stands most of the argument is built around outperforming DanceGRPO as the key baseline. Same goes with Section C.3 Table 4 (vBench) the performance gain (over DanceGRPO) is almost negligible. \n\n- On the qualitative results, a similar argument applies, in my view. I agree that more elaborate details (of BranchGRPO) might lead to higher preference or aesthetics scores, but in practice this might also come at the cost of divergence in prompt adherence. Fig. 4 right hand side is an example, I would argue Flux dinosaur and Gundam adhere much better with the prompt. To be more specific, I would argue the Gundam of BranchGRPO has little to no connection to Shrek anymore, and its Dinosaur does not look like a toy but more so an aesthetically elaborate 3D graphics. Extending on this, in Section C.2 (Appendix) top figure I would argue BranchGRPO does not look like an Anime anymore, same goes with DanceGRPO. In contrast, Flux seems to adhere to that part of prompt (at least) better. Please consider elaborating on this point. \n\n- I think pareto-front type of results can be way more insightful for your audience. A first step in that direction would be to consider adding prompt adherence using CLIPScore (text part --- please see https://arxiv.org/pdf/2104.08718) or something along those lines to your tables. Going beyond that, given that you propose an alignment or guidance approach, I would suggest considering paret-front reward vs. KL divergence plots, if possible. In one way or another you need to establish when your approach offers a better rewards (say HPS), in a faster fashion, it does not diverge drastically from the base model or does not sacrifice prompt adherence. I would suggest taking a look at the approach as well as the plots illustrated on this topic in the following two refs for inspiration: https://arxiv.org/pdf/2502.00968 & https://arxiv.org/pdf/2501.05803v2", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)