通过分支结构提升扩散模型对齐效率,训练速度提升近55%且性能更优。
BranchGRPO: Stable and Efficient GRPO with Structured Branching in Diffusion Models
- 将生成过程改造成共享前缀的分支树,减少重复计算。
- 将稀疏奖励转化为每步信号,提升决策关键性识别能力。
- 剪枝策略降低梯度计算开销,适合高效训练图像视频生成模型。
近期基于组相对策略优化(GRPO)的图像与视频生成模型对齐进展显著,但现有方法因序列化采样和大量采样步数导致效率低下,且信用分配不可靠:稀疏的终端奖励被均匀传播至所有时间步,无法捕捉去噪过程中各决策的重要性差异。本文提出BranchGRPO,将回放过程重构为具有共享前缀的分支树结构,通过共享计算和剪枝策略消除低价值路径与冗余深度。主要贡献包括:(1) 分支机制通过共享前缀摊薄回放成本,同时保持探索多样性;(2) 奖励融合与逐层优势估计器将稀疏终端奖励转换为密集的步骤级信号;(3) 剪枝策略在不影响前向推理与探索的前提下,削减梯度计算量。在HPDv2.1图像对齐任务中,BranchGRPO相比DanceGRPO提升对齐分数达16%,每轮训练时间减少近55%。混合变体BranchGRPO-Mix进一步实现4.7倍于DanceGRPO的加速,且不损失对齐效果。在WanX视频生成任务中,其生成视频具有更高的Video-Align得分,画面更清晰且时序更一致。代码已公开于BranchGRPO。
原文摘要 · Abstract (English)
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}.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。