English

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

Computer Vision and Pattern Recognition 2025-09-30 v5 Artificial Intelligence Machine Learning

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: 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}.

Keywords

Cite

@article{arxiv.2509.06040,
  title  = {BranchGRPO: Stable and Efficient GRPO with Structured Branching in Diffusion Models},
  author = {Yuming Li and Yikai Wang and Yuying Zhu and Zhongyu Zhao and Ming Lu and Qi She and Shanghang Zhang},
  journal= {arXiv preprint arXiv:2509.06040},
  year   = {2025}
}

Comments

12 pages, 6 figures

R2 v1 2026-07-01T05:25:06.619Z