打開網(wǎng)易新聞 查看精彩圖片

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

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

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

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

打開網(wǎng)易新聞 查看精彩圖片

論文標(biāo)題:Efficient Diversity-Preserving Diffusion Alignment via Gradient-Informed GFlowNets

論文地址:https://arxiv.org/abs/2412.07775

代碼地址:https://github.com/lzzcd001/nabla-gfn

打開網(wǎng)易新聞 查看精彩圖片

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

擴散過程的流平衡視角

打開網(wǎng)易新聞 查看精彩圖片

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

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

打開網(wǎng)易新聞 查看精彩圖片

流梯度平衡條件

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

打開網(wǎng)易新聞 查看精彩圖片
打開網(wǎng)易新聞 查看精彩圖片

前向匹配損失:

打開網(wǎng)易新聞 查看精彩圖片

后向匹配損失:

打開網(wǎng)易新聞 查看精彩圖片

終端匹配損失:

打開網(wǎng)易新聞 查看精彩圖片

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

打開網(wǎng)易新聞 查看精彩圖片
打開網(wǎng)易新聞 查看精彩圖片

直觀解釋

打開網(wǎng)易新聞 查看精彩圖片
打開網(wǎng)易新聞 查看精彩圖片

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

偽代碼實現(xiàn)

打開網(wǎng)易新聞 查看精彩圖片

實驗結(jié)果

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

Aesthetic Score,一個在 Laion Aesthetic 數(shù)據(jù)集上訓(xùn)練的美學(xué)評估獎勵函數(shù);

HPSv2 和 ImageReward,衡量指令跟隨能力的獎勵函數(shù)。

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

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

同時,我們的定量實驗表明,我們的 Nabla-GFlowNet 可以更好保持生成樣本的多樣性。

打開網(wǎng)易新聞 查看精彩圖片

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

打開網(wǎng)易新聞 查看精彩圖片

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

打開網(wǎng)易新聞 查看精彩圖片

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

打開網(wǎng)易新聞 查看精彩圖片

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

打開網(wǎng)易新聞 查看精彩圖片

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