BranchGRPO用樹形分叉與剪枝, 破局?jǐn)U散模型對齊新范式

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

在擴散/流匹配模型的人類偏好對齊中，實現(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）。BranchGRPO已上線始智AI-wisemodel開源社區(qū)，歡迎大家前去體驗。

代碼地址

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

01.

研究背景與挑戰(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ù)雜度達(dá)到 O (N×T)（其中 T 是擴散步數(shù)，N 是組大小）。這種重復(fù)采樣帶來大量計算冗余，嚴(yán)重限制了大規(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ī)模視覺生成對齊開辟了新路徑。

02.

如何在擴散過程中分化出樹形結(jié)構(gòu)

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

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

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

3.剪枝（Pruning）：為避免樹形結(jié)構(gòu)帶來的指數(shù)級成本，BranchGRPO 設(shè)計了兩種剪枝策略：

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

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

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

03.

精度、速度、穩(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 秒，但對齊分?jǐn)?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)嚴(yán)重的閃爍和變形；DanceGRPO 雖有所改善，但畫面依舊模糊、不夠穩(wěn)定。相比之下，BranchGRPO 生成的視頻幀更銳利，細(xì)節(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)定的同時達(dá)到最快訓(xùn)練速度。

4.多樣性保持

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

5.擴展性（Scaling Law）

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

04.

總結(jié)與展望

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

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

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

編輯：成蘊年

----- END -----

wisemodel相關(guān)：

系列模型：

關(guān)于wisemodel更多

1

歡迎持續(xù)關(guān)注和支持

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

2

歡迎加盟wisemodel開源社區(qū)

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

3

歡迎投稿優(yōu)質(zhì)內(nèi)容

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

4

關(guān)于wisemodel開源社區(qū)

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

向上滑動查看