BranchGRPO用樹形分叉與剪枝, 破局擴散模型對齊新范式

始智AI wisemodel.cn社區是源自中國的中立開放的AI開源社區。正在，歡迎加入共同成長。

在擴散/流匹配模型的人類偏好對齊中，實現高效采樣與穩定優化的統一，一直是一個重大挑戰。北京大學與字節團隊提出了名為BranchGRPO的新型樹形強化學習方法。不同順序展開的DanceGRPO，BranchGRPO通過在擴散反演過程中引入分叉（branching）與剪枝（pruning），讓多個軌跡共享前綴、在中間步驟分裂，并通過逐層獎勵融合實現稠密反饋。

該方法在 HPDv2.1 圖像對齊與 WanX-1.3B 視頻生成上均取得了優異表現。最令人矚目的是，BranchGRPO 在保證對齊效果更優的同時，迭代時間最高近 5×（Mix 變體 148s vs 698s）。BranchGRPO已上線始智AI-wisemodel開源社區，歡迎大家前去體驗。

代碼地址

https://wisemodel.cn/codes/yumingli/BranchGRPO

01.

研究背景與挑戰

近年來，擴散模型與流匹配模型憑借在圖像與視頻生成上的高保真、多樣性與可控性，已成為視覺生成的主流方案。然而，僅靠大規模預訓練并不能保證與人類意圖完全對齊：模型生成的結果常常偏離美學、語義或時間一致性的需求。

為解決這一問題，「人類反饋強化學習（RLHF）」被引入，用以直接優化生成模型，使其輸出更貼近人類偏好。

在 RLHF 體系中，「群體相對策略優化（GRPO）」被證明在圖生文、文生圖和視頻生成中具有良好的穩定性與可擴展性。然而，當 GRPO 應用于擴散 / 流模型時，依舊面臨兩大根本性瓶頸：

• 低效性：標準 GRPO 采用順序 rollout，每條軌跡必須在舊策略和新策略下獨立采樣，復雜度達到 O (N×T)（其中 T 是擴散步數，N 是組大小）。這種重復采樣帶來大量計算冗余，嚴重限制了大規模生成任務的擴展性。

• 稀疏獎勵：現有方法通常只在最終生成結果上計算單一獎勵，并將其均勻回傳至所有步。這種 “稀疏且均勻” 的反饋忽視了中間狀態中蘊含的關鍵信號，導致 credit assignment 不準確，訓練波動大、收斂不穩，甚至出現高方差梯度。

因此，一個關鍵問題被提出：如何在不破壞多樣性的前提下，既提升采樣效率，又讓獎勵信號更稠密、更穩定地作用于訓練過程？

正是在這一背景下，研究團隊提出了 BranchGRPO。通過樹形分叉、獎勵融合與剪枝機制，BranchGRPO 做到了「又快又穩、又強又準」，為大規模視覺生成對齊開辟了新路徑。

02.

如何在擴散過程中分化出樹形結構

為突破順序 rollout 的低效與稀疏獎勵瓶頸，BranchGRPO 將原本單一路徑的采樣過程，重構為一種樹形展開：

1.分叉（Branching）：在若干預設的擴散步上進行分裂，每條軌跡可以向多個子路徑擴展，前綴計算被復用，大幅減少冗余采樣。這種結構既保持了擴散過程的完整性，又讓探索更高效。

2.獎勵融合與逐層歸因（Reward Fusion & Depth-wise Advantage）：不同于將單一終末獎勵均勻分配到所有步驟，BranchGRPO 將葉子節點的獎勵自底向上傳遞，并在每一深度上進行標準化，形成逐步稠密的優勢信號，使訓練過程更穩定、更精準。

3.剪枝（Pruning）：為避免樹形結構帶來的指數級成本，BranchGRPO 設計了兩種剪枝策略：

• 寬度剪枝：僅保留關鍵葉子參與反向傳播，減少梯度計算量；

• 深度剪枝：跳過部分層的反傳（但保留前向和獎勵評估），進一步壓縮開銷。

這一系列設計使得 BranchGRPO 在效率和穩定性之間實現了統一：既能顯著加速訓練、降低迭代開銷，又能在獎勵歸因上更精細、更穩定，從而在圖像與視頻生成任務中同時提升對齊效果與收斂速度。

03.

精度、速度、穩定度

1.圖像對齊（HPDv2.1）

在圖像對齊測試中，BranchGRPO 帶來了真正的「又快又好」：

更快：

DanceGRPO (tf=1.0) 每迭代 698s；BranchGRPO 493s；剪枝版 314s；Mix 變體 148s（相對 698s 最高近 4.7× 加速）

更穩更準：

HPS-v2.1 0.363–0.369，穩定高于 DanceGRPO 的 0.360；ImageReward 1.319（DepPru） 為全表最佳。

對比其他方法：

MixGRPO 雖然也能壓縮時間到 289 秒，但對齊分數略有下降，并且 MixGRPO 訓練常常不穩定；相比之下，BranchGRPO-Mix 在極致加速的同時，依舊保持了與原始 BranchGRPO 相當的對齊效果和穩定的訓練，展現出驚人的性價比。

2.視頻生成(WanX-1.3B)

在視頻生成任務中，BranchGRPO 同樣展現了強大的優勢：

更清晰：

不使用 RLHF 的基礎模型常出現嚴重的閃爍和變形；DanceGRPO 雖有所改善，但畫面依舊模糊、不夠穩定。相比之下，BranchGRPO 生成的視頻幀更銳利，細節更豐富，角色和物體在時間維度上保持一致，真正實現了「流暢不掉幀」的觀感。

更快：

在相同硬件條件下，DanceGRPO 每次迭代大約需要 近 20 分鐘；而 BranchGRPO 僅需約 8 分鐘 就能完成一次迭代，訓練效率直接翻 2 倍以上。

3.消融實驗

從消融實驗可以看到：適中的分支相關度、早期更密集的分裂能加快獎勵提升；路徑加權的獎勵融合讓訓練更穩；深度剪枝帶來最佳最終效果；而混合 ODE–SDE 調度則在保持穩定的同時達到最快訓練速度。

4.多樣性保持

分叉并未削弱樣本分布，MMD2≈0.019，幾乎與順序采樣一致。

5.擴展性（Scaling Law）

得益于 BranchGRPO 的高效性與訓練穩定性，研究團隊能夠輕松擴大分支規模而不崩潰：無論是增加分支因子還是分支次數，性能都持續提升。比如在 81 個樣本規模下，DanceGRPO 每次迭代要花 2400 秒，而BranchGRPO 只需 680 秒，真正把大規模對齊訓練變得可行。

04.

總結與展望

BranchGRPO 通過樹形分叉、獎勵融合與輕量剪枝，創新性地融合了效率與穩定，獎勵從「終點一錘子」變「全程有信號」—— 在速度、穩定與對齊效果上全面提升（HPDv2.1 最高近 5×，視頻生成更清晰更一致）。成為視覺生成對齊的新一代解決方案。

未來，若引入自適應分裂 / 剪枝策略，并拓展至多模態與更大規模生成任務，BranchGRPO 有望成為擴散 / 流模型 RLHF 的核心方法，為高效、穩定的人類偏好對齊提供新的范式。

該項目主要由來自北京大學、北京師范大學、字節跳動的師生聯合研究，作者包括李聿明、王一凱等，通訊作者為北京大學仉尚航。

編輯：成蘊年

----- END -----

wisemodel相關：

系列模型：

關于wisemodel更多

1

歡迎持續關注和支持

開源社區建設需要長期堅持和投入，更需要廣大用戶的積極參與、貢獻和維護，歡迎大家加入wisemodel開源社區的志愿者計劃和開源共創計劃。期待更多開發者將開源成果，包括模型、數據集和代碼等發布到 wisemodel.cn 社區，共建中立、開放的AI開源社區生態。歡迎掃碼添加wisemodel微信，申請加入wisemodel社群，持續關注wisemodel.cn開源社區動態。

2

歡迎加盟wisemodel開源社區

始智AI wisemodel社區自2023年9月上線以來，逐漸成為影響力日益擴大的中立開放的AI開源社區，為了加快公司發展，我們長期需要技術、運營等人才加盟，技術側重在AI infra、后端開發，熟悉K8S、模型訓練和推理等技術， 以及熟悉開發者生態運營的成員，歡迎感興趣的朋友加盟，可以通過添加wisemodel微信，或者將簡歷投遞到郵箱：liudaoquan@wisemodel.cn

3

歡迎投稿優質內容

歡迎投稿分享人工智能領域相關的優秀研究成果，鼓勵高校實驗室、大企業研究團隊、個人等，在wisemodel平臺上分享各類優質內容，可以是AI領域最新論文解讀、最新開源成果介紹，也可以是關于AI技術實踐、應用和總結等。投稿可以發郵件到liudaoquan@wisemodel.cn，也可以掃碼添加wisemodel微信。

4

關于wisemodel開源社區

始智AI wisemodel.cn開源社區由清華校友總會AI大數據專委會副秘書長劉道全創立，旨在打造和建設中立開放的AI開源創新社區，將打造成“HuggingFace”之外最活躍的AI開源社區，匯聚主要AI開源模型、數據集和代碼等，歡迎高校科研院所、大型互聯網公司、創新創業企業、廣大個人開發者，以及政府部門、學會協會、聯盟、基金會等，還有投資機構、科技媒體等，共同參與建設AI開源創新生態。

向上滑動查看