本文作者劉圳是香港中文大學(xué)（深圳）數(shù)據(jù)科學(xué)學(xué)院的助理教授，肖鎮(zhèn)中是德國馬克思普朗克-智能系統(tǒng)研究所和圖賓根大學(xué)的博士生，劉威楊是德國馬克思普朗克-智能系統(tǒng)研究所的研究員，Yoshua Bengio 是蒙特利爾大學(xué)和加拿大 Mila 研究所的教授，張鼎懷是微軟研究院的研究員。此論文已收錄于 ICLR 2025。

在視覺生成領(lǐng)域，擴(kuò)散模型（Diffusion Models）已經(jīng)成為生成高質(zhì)量圖像、視頻甚至文本的利器。然而，生成結(jié)果往往離我們所偏好的不一致：結(jié)果不美觀，圖文不符，等等。

雖然我們可以像大語言模型中的 RLHF（基于人類反饋的強(qiáng)化學(xué)習(xí)）一樣直接用傳統(tǒng)強(qiáng)化學(xué)習(xí)來微調(diào)擴(kuò)散模型，但收斂速度往往慢；而基于可微計(jì)算圖直接最大化獎(jiǎng)勵(lì)函數(shù)的方法又往往陷入過擬合和多樣性缺失的問題。

有沒有一種方法，既能保留生成樣本的多樣性，又能快速完成微調(diào)？我們基于生成流網(wǎng)絡(luò)（Generative Flow Network，GFlowNet）提出的 Nabla-GFlowNet 實(shí)現(xiàn)了這一速度和質(zhì)量間的平衡。

• 論文標(biāo)題：Efficient Diversity-Preserving Diffusion Alignment via Gradient-Informed GFlowNets
• 論文地址：https://arxiv.org/abs/2412.07775
• 代碼地址：https://github.com/lzzcd001/nabla-gfn

利用 Nabla-GFlowNet 在 Aesthetic Score 獎(jiǎng)勵(lì)函數(shù)（一個(gè)美學(xué)指標(biāo)）上高效微調(diào) Stable Diffusion 模型。

擴(kuò)散過程的流平衡視角

生成流網(wǎng)絡(luò) GFlowNet 示意圖。初始節(jié)點(diǎn)中的「流」通過向下游的轉(zhuǎn)移概率流經(jīng)不同節(jié)點(diǎn)，最后匯聚到終端節(jié)點(diǎn)。每個(gè)終端節(jié)點(diǎn)所對(duì)應(yīng)的流應(yīng)匹配該終端節(jié)點(diǎn)對(duì)應(yīng)的獎(jiǎng)勵(lì)。

在生成流網(wǎng)絡(luò)（Generative Flow Network, GFlowNet）的框架下，擴(kuò)散模型的生成過程可以視為一個(gè)「水流從源頭流向終點(diǎn)」的動(dòng)態(tài)系統(tǒng)：

• 從標(biāo)準(zhǔn)高斯分布采樣的噪聲圖像作為初始狀態(tài)，其「流量」為；
• 去噪過程 如同分配水流的管道網(wǎng)絡(luò)，把每一個(gè)節(jié)點(diǎn)的水流分配給下游每一個(gè)節(jié)點(diǎn)；
• 而加噪過程則可以回溯每一個(gè)的水流來自哪里；
• 最終生成的圖像將累積總流量。

擴(kuò)散模型示意圖

流梯度平衡條件

在 GFlowNet 框架下，前后向水流需要滿足一定的平衡條件。我們通過推導(dǎo)提出我們稱為 Nabla-DB 的平衡條件：

其中是殘差去噪過程，和分別是微調(diào)模型和預(yù)訓(xùn)練模型的去噪過程。是這個(gè)殘差過程對(duì)應(yīng)的對(duì)數(shù)流函數(shù)。

這個(gè)殘差去噪過程應(yīng)該滿足，其中 β 控制微調(diào)模型在獎(jiǎng)勵(lì)函數(shù)和預(yù)訓(xùn)練模型之間的平衡。如果 β 為零，那么殘差過程為零，也就是微調(diào)網(wǎng)絡(luò)等于預(yù)訓(xùn)練網(wǎng)絡(luò)。

稍作變換，就可以得到我們提出的 Nabla-GFlowNet 對(duì)應(yīng)的損失函數(shù) Residual Nabla-DB（其中 sg 為 stop-gradient 操作）：

前向匹配損失：

后向匹配損失：

終端匹配損失：

對(duì)數(shù)流梯度參數(shù)化設(shè)計(jì)

上述損失函數(shù)需要用一個(gè)額外的網(wǎng)絡(luò)估計(jì)，但我們觀察到：如果我們對(duì) 做單步預(yù)測(cè)得到不準(zhǔn)確的去噪結(jié)果，那么是一個(gè)很好的估計(jì)，其中是一個(gè)權(quán)重常數(shù)。因此，我們提出如下參數(shù)化：

其中是用 U-Net 參數(shù)化的殘差梯度，而單步去噪估計(jì)為（是擴(kuò)散模型ε-預(yù)測(cè)參數(shù)化的網(wǎng)絡(luò)）。

方法示意圖。每條采樣路徑的每個(gè)轉(zhuǎn)移對(duì)中每張圖的上下兩個(gè)「力」需要相互平衡。

直觀解釋

如果我們只計(jì)算這一轉(zhuǎn)移對(duì)的前向匹配損失對(duì)擴(kuò)散模型參數(shù)的梯度，我們有：

其中第一項(xiàng)是基于內(nèi)積的匹配度函數(shù)（殘差擴(kuò)散模型與獎(jiǎng)勵(lì)梯度估計(jì)之間的匹配），第二項(xiàng)是讓微調(diào)模型趨近于預(yù)訓(xùn)練模型的正則化。

偽代碼實(shí)現(xiàn)

實(shí)驗(yàn)結(jié)果

我們分別用以下獎(jiǎng)勵(lì)函數(shù)微調(diào) Stable Diffusion 網(wǎng)絡(luò)：

• Aesthetic Score，一個(gè)在 Laion Aesthetic 數(shù)據(jù)集上訓(xùn)練的美學(xué)評(píng)估獎(jiǎng)勵(lì)函數(shù)；
• HPSv2 和 ImageReward，衡量指令跟隨能力的獎(jiǎng)勵(lì)函數(shù)。

定性實(shí)驗(yàn)結(jié)果表明，通過 Nabla-GFlowNet 微調(diào)，我們可以快速得到獎(jiǎng)勵(lì)更高但避免過擬合的生成圖像。

ReFL 和 DRaFT 等直接獎(jiǎng)勵(lì)優(yōu)化的方法雖然收斂速度快，但很快會(huì)陷入過擬合；而 DDPO 這一基于傳統(tǒng)策略梯度的強(qiáng)化學(xué)習(xí)微調(diào)方法由于沒有理由梯度信息，微調(diào)速度顯著劣于其他方法。

同時(shí)，我們的定量實(shí)驗(yàn)表明，我們的 Nabla-GFlowNet 可以更好保持生成樣本的多樣性。

Aesthetic Score 獎(jiǎng)勵(lì)函數(shù)上的微調(diào)結(jié)果（微調(diào) 200 步，取圖片質(zhì)量不坍塌的最好模型）。Nabla-GFlowNet（對(duì)應(yīng) Residual Nabla-DB 損失函數(shù)）方法微調(diào)的網(wǎng)絡(luò)可以生成平均獎(jiǎng)勵(lì)更高且不失自然的生成圖片。

相較于 ReFL，DRaFT 等直接獎(jiǎng)勵(lì)優(yōu)化的方法，Nabla-GFlowNet 更難陷入過擬合。

Aesthetic Score 獎(jiǎng)勵(lì)函數(shù)上的定量結(jié)果 Nabla-GFlowNet 在獎(jiǎng)勵(lì)收斂快的同時(shí)，保持更高的 DreamSim 多樣性指標(biāo)（越高代表多樣性越好）和更低的 FID 分?jǐn)?shù)（越低代表越符合預(yù)訓(xùn)練模型的先驗(yàn)）。

HPSv2 獎(jiǎng)勵(lì)函數(shù)上的微調(diào)結(jié)果

ImageReward 獎(jiǎng)勵(lì)函數(shù)上的微調(diào)結(jié)果

結(jié)語

我們利用生成流網(wǎng)絡(luò)（GFlowNet）的框架，嚴(yán)謹(jǐn)?shù)氐玫揭粋€(gè)可以更好保持多樣性和先驗(yàn)的高效的擴(kuò)散模型獎(jiǎng)勵(lì)微調(diào)方法，并且在 Stable Diffusion 這一常用的文生圖擴(kuò)散模型上顯示出相較于其他方法的優(yōu)勢(shì)。