DiffusionOPD：復(fù)旦聯(lián)合通義萬相提出擴散模型在線策略蒸餾新范式

擴散模型在單一任務(wù)上的強化學(xué)習(xí)已經(jīng)取得了顯著進展，例如提升文字生成質(zhì)量、增強構(gòu)圖準(zhǔn)確性，或優(yōu)化畫面美感等。但當(dāng)這些能力需要同時集成到同一個模型中時，訓(xùn)練往往會變得十分困難：不同任務(wù)之間容易產(chǎn)生相互干擾，訓(xùn)練目標(biāo)也會變得復(fù)雜而不穩(wěn)定。

近期，來自復(fù)旦大學(xué)與阿里巴巴通義萬相的研究團隊對此提出了新的思考。他們認為，多任務(wù)強化學(xué)習(xí)不應(yīng)被視為一個統(tǒng)一優(yōu)化問題，而應(yīng)該解耦為兩個彼此獨立的過程：單任務(wù)的在線策略探索 & 多任務(wù)能力整合。

基于這一觀點，他們提出了DiffusionOPD，為 diffusion 領(lǐng)域的 On-Policy Distillation 提供了一個統(tǒng)一視角，并建立了相應(yīng)的理論與實驗框架。

DiffusionOPD 的核心思路，是先針對不同任務(wù)分別訓(xùn)練各自的「專家教師」模型；隨后，再通過在線策略蒸餾，將這些教師模型的能力統(tǒng)一蒸餾到同一個學(xué)生模型中，實現(xiàn)多任務(wù)能力整合。最終，一個統(tǒng)一的 student model 便能夠同時兼顧構(gòu)圖、OCR、美學(xué)等多項能力。

• 論?標(biāo)題：DiffusionOPD: A Unified Perspective of On-Policy Distillation in Diffusion Models
• 論文地址：https://arxiv.org/abs/2605.15055
• 項?主?：https://quanhaol.github.io/DiffusionOPD-site/
• 代碼鏈接 https://github.com/ali-vilab/DiffusionOPD

多任務(wù)強化學(xué)習(xí)方法

過去常?的多任務(wù)強化學(xué)習(xí)?法主要有兩類。

聯(lián)合多任務(wù) RL (Joint Multi-Task Optimization) ：使用現(xiàn)有的 RL 算法例如 DiffusionNFT, GRPO 去聯(lián)合優(yōu)化多個任務(wù)。這種范式會撞上兩個問題： 1 獎勵沖突：不同任務(wù)的優(yōu)化?向往往存在相互干擾； 2 任務(wù)失衡：簡單任務(wù)會主導(dǎo)訓(xùn)練過程，導(dǎo)致復(fù)雜任務(wù)難以充分學(xué)習(xí)。

級聯(lián) RL (Cascade RL)：按階段依次訓(xùn)練不同任務(wù)。雖然能夠緩解任務(wù)沖突，但是訓(xùn)練流程復(fù)雜，需要分別調(diào)整各階段的超參數(shù)與訓(xùn)練策略，而且容易產(chǎn)生災(zāi)難性遺忘，后續(xù)任務(wù)訓(xùn)練的時候會削弱已有能力。

圖 1：（a）相比所有多任務(wù)強化學(xué)習(xí)基線方法，DiffusionOPD 展現(xiàn)出顯著更快的收斂速度以及更高的性能上限。（b）在包括 GenEval、OCR 與美學(xué)在內(nèi)的多個任務(wù)領(lǐng)域中，DiffusionOPD 均優(yōu)于所有基線方法。

DiffusionOPD: 單任務(wù)探索 + 多任務(wù)整合

DiffusionOPD 給出的答案?脆利落：多任務(wù)強化學(xué)習(xí)不應(yīng)被視為一個統(tǒng)一優(yōu)化問題，而應(yīng)該解耦為兩個彼此獨立的過程：單任務(wù)的在線策略探索 & 多任務(wù)能力整合。

整體訓(xùn)練過程可分為兩個階段

• Stage 1?單任務(wù)?師獨?訓(xùn)練：針對不同任務(wù)（如 GenEval、OCR、Aesthetic 等），分別使用現(xiàn)有的 diffusion RL 方法訓(xùn)練對應(yīng)的「專家教師」模型。其中，GenEval 任務(wù)采用 DiffusionNFT，OCR 與美學(xué)任務(wù)采用 GRPO-Guard。由于每個教師僅負責(zé)單一任務(wù)，因此能夠避免跨任務(wù)干擾。
• Stage 2?在線策略蒸餾多任務(wù)能力到學(xué)?模型：隨后，從一個預(yù)訓(xùn)練擴散模型初始化統(tǒng)一的學(xué)生模型，并通過在線策略蒸餾整合多任務(wù)能力。在訓(xùn)練過程中，學(xué)生模型針對不同任務(wù)，基于自身策略生成去噪軌跡；隨后，在學(xué)生生成的每個去噪狀態(tài)上，由對應(yīng)任務(wù)的教師模型提供監(jiān)督信號。因此，學(xué)生模型無需重新對所有任務(wù)進行從零探索，而是能夠直接學(xué)習(xí)各任務(wù)教師的策略與能力，從而實現(xiàn)高效的多任務(wù)能力融合。

Diffusion 領(lǐng)域 OPD 的?標(biāo)函數(shù)推導(dǎo)

在 LLM 中，OPD 的做法很自然：學(xué)生模型先按照自己的策略生成 token，隨后教師模型在學(xué)生訪問到的每一個 token 狀態(tài)上提供監(jiān)督。由于語言模型本身是離散 token 分布，因此可以直接對每一步的 token distribution 做 KL 蒸餾。

但 diffusion model 不一樣。它不是離散 token 序列，而是一個連續(xù)狀態(tài)的去噪過程。

因此作者首先把 diffusion 的去噪過程重新視作一個 continuous-state Markov chain（連續(xù)狀態(tài)馬爾可夫鏈）。在這個視角下，每一步去噪 transition 都對應(yīng)一個 Gaussian transition kernel；學(xué)生模型和教師模型分別定義自己的 transition distribution：

接著，論文進一步推導(dǎo)發(fā)現(xiàn)：由于 student 和 teacher 的transition covariance 是相同的，于是整個擴散版 OPD ?標(biāo) reverse KL，就被寫成了?個完全解析、? Monte-Carlo ?差的均值匹配損失：

作者進一步指出這一框架同時統(tǒng)一了stochastic SDE sampler 與 deterministic ODE sampler。在 ODE 情況下目標(biāo)會退化成均值之間的 L2 匹配。

與 PPO-style policy gradient 的比較

另一個一非常自然的想法是：把老師當(dāng)作「過程獎勵模型」，把 KL 損失看作每一個去燥步的 dense reward 然后計算 advantage，最后套一個 PPO 的損失函數(shù)。

DiffusionOPD 論文里嚴(yán)格證明了直接閉式 KL 與 PPO -style policy gradient在期望意義下梯度完全相等。但 PPO 的梯度里會多出一項 score-function 項，它與高斯噪聲成正比，期望為零但方差不為零。也就是說，PPO 估計天然比閉式 KL 更「吵」。

更關(guān)鍵的是 PPO 形式離不開 logprob 與 ratio 的計算，因此它在 ODE 確定性采樣器下根本無法定義，僅僅支持 SDE sampler。

實驗結(jié)果

1.與多任務(wù)強化學(xué)習(xí)方法的對比

定量效果對比：

訓(xùn)練曲線對比：

定性效果對比：

圖 2：與多任務(wù)強化學(xué)習(xí)方法以及單任務(wù)教師模型的定性對比結(jié)果。每個案例分為兩行展示：第一行從左到右依次為 DiffusionOPD（本文方法）、Multi-Task GRPO-Guard、Multi-Task NFT 和 Cascade NFT；第二行從左到右依次為輸入文本、Aesthetic Teacher、GenEval Teacher 和 OCR Teacher 的生成結(jié)果。

2.蒸餾方法消融：

作者還做了一組很有意義的對照實驗：固定同一批專家老師，分別用 DiffusionOPD、DMD、TDM、SFT 蒸餾到同一個學(xué)生，控制變量后對比誰更適合「多任務(wù)能力整合」這個場景。

訓(xùn)練曲線對比：

曲線表示同樣的老師、同樣的采樣設(shè)置，DiffusionOPD 在收斂速度和上限上都明顯更好。

定性效果對比：

圖 3：與不同蒸餾方法的定性對比結(jié)果。從左到右依次為：DiffusionOPD（本文方法）、DMD、TDM 和 SFT。

3.Loss 形式以及 Sampler Type 消融

圖 6：關(guān)于損失函數(shù)形式與采樣器噪聲水平的消融實驗。當(dāng)噪聲水平設(shè)為 0 時，SDE sampler 將退化為 ODE sampler。實驗結(jié)果表明，PPO-style policy gradient 的表現(xiàn)遜于同樣 noise level 的 closed-form KL objective；此外，更低的噪聲水平能夠帶來更快的收斂速度和更高的性能上限。