加速近5倍！北大與字節(jié)團隊提出BranchGRPO，用「樹形分叉 + 剪枝」重塑擴散模型對齊
在擴散 / 流匹配模型的人類偏好對齊中，實現(xiàn)高效采樣與穩(wěn)定優(yōu)化的統(tǒng)一，一直是一個重大挑戰(zhàn)。
近期，北京大學(xué)與字節(jié)團隊提出了名為 BranchGRPO 的新型樹形強化學(xué)習(xí)方法。不同于順序展開的 DanceGRPO，BranchGRPO 通過在擴散反演過程中引入分叉（branching）與剪枝（pruning），讓多個軌跡共享前綴、在中間步驟分裂，并通過逐層獎勵融合實現(xiàn)稠密反饋。
該方法在 HPDv2.1 圖像對齊與 WanX-1.3B 視頻生成上均取得了優(yōu)異表現(xiàn)。最令人矚目的是，BranchGRPO 在保證對齊效果更優(yōu)的同時，迭代時間最高近 5×（Mix 變體 148s vs 698s）。
- 論文鏈接: https://arxiv.org/pdf/2509.06040
- 項目主頁:
https://fredreic1849.github.io/BranchGRPO-Webpage/ - 代碼鏈接:
https://github.com/Fredreic1849/BranchGRPO - PKU HMI 實驗室主頁：https://pku-hmi-lab.github.io/HMI-Web/index.html
- 單位：該項目主要由來自北京大學(xué)、北京師范大學(xué)、字節(jié)跳動的師生聯(lián)合研究，作者包括李聿明、王一凱等，通訊作者為北京大學(xué)仉尚航。
研究背景與挑戰(zhàn)
近年來，擴散模型與流匹配模型憑借在圖像與視頻生成上的高保真、多樣性與可控性，已成為視覺生成的主流方案。然而，僅靠大規(guī)模預(yù)訓(xùn)練并不能保證與人類意圖完全對齊：模型生成的結(jié)果常常偏離美學(xué)、語義或時間一致性的需求。
為解決這一問題，「人類反饋強化學(xué)習(xí)（RLHF）」被引入，用以直接優(yōu)化生成模型，使其輸出更貼近人類偏好。
在 RLHF 體系中，「群體相對策略優(yōu)化（GRPO）」被證明在圖生文、文生圖和視頻生成中具有良好的穩(wěn)定性與可擴展性。然而，當(dāng) GRPO 應(yīng)用于擴散 / 流模型時，依舊面臨兩大根本性瓶頸：
低效性：標(biāo)準(zhǔn) GRPO 采用順序 rollout，每條軌跡必須在舊策略和新策略下獨立采樣，復(fù)雜度達到 O (N×T)（其中 T 是擴散步數(shù)，N 是組大?。?。這種重復(fù)采樣帶來大量計算冗余，嚴重限制了大規(guī)模生成任務(wù)的擴展性。
稀疏獎勵：現(xiàn)有方法通常只在最終生成結(jié)果上計算單一獎勵，并將其均勻回傳至所有步。這種 “稀疏且均勻” 的反饋忽視了中間狀態(tài)中蘊含的關(guān)鍵信號，導(dǎo)致 credit assignment 不準(zhǔn)確，訓(xùn)練波動大、收斂不穩(wěn)，甚至出現(xiàn)高方差梯度。
因此，一個關(guān)鍵問題被提出：如何在不破壞多樣性的前提下，既提升采樣效率，又讓獎勵信號更稠密、更穩(wěn)定地作用于訓(xùn)練過程？
正是在這一背景下，我們提出了 BranchGRPO。通過樹形分叉、獎勵融合與剪枝機制，BranchGRPO 做到了「又快又穩(wěn)、又強又準(zhǔn)」，為大規(guī)模視覺生成對齊開辟了新路徑。
BranchGRPO如何在擴散過程中分化出樹形結(jié)構(gòu)
為突破順序 rollout 的低效與稀疏獎勵瓶頸，BranchGRPO 將原本單一路徑的采樣過程，重構(gòu)為一種樹形展開：
- 分叉（Branching）：在若干預(yù)設(shè)的擴散步上進行分裂，每條軌跡可以向多個子路徑擴展，前綴計算被復(fù)用，大幅減少冗余采樣。這種結(jié)構(gòu)既保持了擴散過程的完整性，又讓探索更高效。
- 獎勵融合與逐層歸因（Reward Fusion & Depth-wise Advantage）：不同于將單一終末獎勵均勻分配到所有步驟，BranchGRPO 將葉子節(jié)點的獎勵自底向上傳遞，并在每一深度上進行標(biāo)準(zhǔn)化，形成逐步稠密的優(yōu)勢信號，使訓(xùn)練過程更穩(wěn)定、更精準(zhǔn)。
- 剪枝（Pruning）：為避免樹形結(jié)構(gòu)帶來的指數(shù)級成本，BranchGRPO 設(shè)計了兩種剪枝策略：
- 寬度剪枝：僅保留關(guān)鍵葉子參與反向傳播，減少梯度計算量；
- 深度剪枝：跳過部分層的反傳（但保留前向和獎勵評估），進一步壓縮開銷。
這一系列設(shè)計使得 BranchGRPO 在效率和穩(wěn)定性之間實現(xiàn)了統(tǒng)一：既能顯著加速訓(xùn)練、降低迭代開銷，又能在獎勵歸因上更精細、更穩(wěn)定，從而在圖像與視頻生成任務(wù)中同時提升對齊效果與收斂速度。
精度、速度、穩(wěn)定度
1.圖像對齊（HPDv2.1）
在圖像對齊測試中，BranchGRPO 帶來了真正的「又快又好」：
更快：
DanceGRPO (tf=1.0) 每迭代 698s；BranchGRPO 493s；剪枝版 314s；Mix 變體 148s（相對 698s 最高近 4.7× 加速）
更穩(wěn)更準(zhǔn)：
HPS-v2.1 0.363–0.369，穩(wěn)定高于 DanceGRPO 的 0.360；ImageReward 1.319（DepPru） 為全表最佳。
對比其他方法：
MixGRPO 雖然也能壓縮時間到 289 秒，但對齊分數(shù)略有下降，并且 MixGRPO 訓(xùn)練常常不穩(wěn)定；相比之下，BranchGRPO-Mix 在極致加速的同時，依舊保持了與原始 BranchGRPO 相當(dāng)?shù)膶R效果和穩(wěn)定的訓(xùn)練，展現(xiàn)出驚人的性價比。
2.視頻生成（WanX-1.3B）
在視頻生成任務(wù)中，BranchGRPO 同樣展現(xiàn)了強大的優(yōu)勢：
更清晰：
不使用 RLHF 的基礎(chǔ)模型常出現(xiàn)嚴重的閃爍和變形；DanceGRPO 雖有所改善，但畫面依舊模糊、不夠穩(wěn)定。相比之下，BranchGRPO 生成的視頻幀更銳利，細節(jié)更豐富，角色和物體在時間維度上保持一致，真正實現(xiàn)了「流暢不掉幀」的觀感。
更快：
在相同硬件條件下，DanceGRPO 每次迭代大約需要 近 20 分鐘；而 BranchGRPO 僅需約 8 分鐘 就能完成一次迭代，訓(xùn)練效率直接翻 2 倍以上。
3.消融實驗
從消融實驗可以看到：適中的分支相關(guān)度、早期更密集的分裂能加快獎勵提升；路徑加權(quán)的獎勵融合讓訓(xùn)練更穩(wěn)；深度剪枝帶來最佳最終效果；而混合 ODE–SDE 調(diào)度則在保持穩(wěn)定的同時達到最快訓(xùn)練速度。
4.4.多樣性保持
分叉并未削弱樣本分布，MMD2≈0.019，幾乎與順序采樣一致。
5.擴展性（Scaling Law）
得益于 BranchGRPO 的高效性與訓(xùn)練穩(wěn)定性，我們能夠輕松擴大分支規(guī)模而不崩潰：無論是增加分支因子還是分支次數(shù)，性能都持續(xù)提升。比如在 81 個樣本規(guī)模下，DanceGRPO 每次迭代要花 2400 秒，而 BranchGRPO 只需 680 秒，真正把大規(guī)模對齊訓(xùn)練變得可行。
總結(jié)與展望
BranchGRPO 通過樹形分叉、獎勵融合與輕量剪枝，創(chuàng)新性地融合了效率與穩(wěn)定，獎勵從「終點一錘子」變「全程有信號」—— 在速度、穩(wěn)定與對齊效果上全面提升（HPDv2.1 最高近 5×，視頻生成更清晰更一致）。成為視覺生成對齊的新一代解決方案。
未來，若引入自適應(yīng)分裂 / 剪枝策略，并拓展至多模態(tài)與更大規(guī)模生成任務(wù)，BranchGRPO 有望成為擴散 / 流模型 RLHF 的核心方法，為高效、穩(wěn)定的人類偏好對齊提供新的范式。
如果您在研究中使用BranchGRPO，歡迎引用我們的工作：
@article{li2025branchgrpo,
title={BranchGRPO: Stable and Efficient GRPO with Structured Branching in Diffusion Models},
author={Li, Yuming and Wang, Yikai and Zhu, Yuying and Zhao, Zhongyu and Lu, Ming and She, Qi and Zhang, Shanghang},
journal={arXiv preprint arXiv:2509.06040},
year={2025}
}
