本文作者劉圳是香港中文大學（深圳）資料科學學院的助理教授，肖鎮中是德國馬克思普朗克-智慧系統研究所和圖賓根大學的博士生，劉威楊是德國馬克思普朗克-智慧系統研究所的研究員，Yoshua Bengio 是蒙特利爾大學和加拿大 Mila 研究所的教授，張鼎懷是微軟研究院的研究員。此論文已收錄於 ICLR 2025。
在視覺生成領域，擴散模型（Diffusion Models）已經成為生成高質量影像、影片甚至文字的利器。然而，生成結果往往離我們所偏好的不一致：結果不美觀，圖文不符，等等。
雖然我們可以像大語言模型中的 RLHF（基於人類反饋的強化學習）一樣直接用傳統強化學習來微調擴散模型，但收斂速度往往慢；而基於可微計算圖直接最大化獎勵函式的方法又往往陷入過擬合和多樣性缺失的問題。
有沒有一種方法，既能保留生成樣本的多樣性，又能快速完成微調？我們基於生成流網路（Generative Flow Network，GFlowNet）提出的 Nabla-GFlowNet 實現了這一速度和質量間的平衡。
-
論文標題：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 獎勵函式（一個美學指標）上高效微調 Stable Diffusion 模型。
擴散過程的流平衡視角
生成流網路 GFlowNet 示意圖。初始節點中的「流」透過向下遊的轉移機率流經不同節點，最後匯聚到終端節點。每個終端節點所對應的流應匹配該終端節點對應的獎勵。
在生成流網路（Generative Flow Network, GFlowNet）的框架下，擴散模型的生成過程可以視為一個「水流從源頭流向終點」的動態系統：
-
從標準高斯分佈取樣的噪聲影像 作為初始狀態，其「流量」為
；
-
去噪過程 如同分配水流的管道網路，把每一個節點的水流分配給下游每一個
節點；
-
而加噪過程 則可以回溯每一個的水流來自哪裡；
-
最終生成的影像 將累積總流量。
擴散模型示意圖
流梯度平衡條件
在 GFlowNet 框架下，前後向水流需要滿足一定的平衡條件。我們透過推導提出我們稱為 Nabla-DB 的平衡條件：
其中
是殘差去噪過程，
和
分別是微調模型和預訓練模型的去噪過程。
是這個殘差過程對應的對數流函式。
這個殘差去噪過程應該滿足
，其中 β 控制微調模型在獎勵函式和預訓練模型之間的平衡。如果 β 為零，那麼殘差過程為零，也就是微調網路等於預訓練網路。
稍作變換，就可以得到我們提出的 Nabla-GFlowNet 對應的損失函式 Residual Nabla-DB（其中 sg 為 stop-gradient 操作）：
前向匹配損失：
後向匹配損失：
終端匹配損失：
對數流梯度引數化設計
上述損失函式需要用一個額外的網路估計
，但我們觀察到：如果我們對
做單步預測得到不準確的去噪結果
，那麼
是一個很好的估計，其中
是一個權重常數。因此，我們提出如下引數化：
其中
是用 U-Net 引數化的殘差梯度，而單步去噪估計為
（
是擴散模型ε-預測引數化的網路）。
方法示意圖。每條取樣路徑的每個轉移對
中每張圖的上下兩個「力」需要相互平衡。
直觀解釋
如果我們只計算
這一轉移對的前向匹配損失對擴散模型引數的梯度，我們有：
其中第一項是基於內積的匹配度函式（殘差擴散模型與獎勵梯度估計之間的匹配），第二項是讓微調模型趨近於預訓練模型的正則化。
虛擬碼實現
實驗結果
我們分別用以下獎勵函式微調 Stable Diffusion 網路：
-
Aesthetic Score，一個在 Laion Aesthetic 資料集上訓練的美學評估獎勵函式； -
HPSv2 和 ImageReward，衡量指令跟隨能力的獎勵函式。
定性實驗結果表明，透過 Nabla-GFlowNet 微調，我們可以快速得到獎勵更高但避免過擬合的生成影像。
ReFL 和 DRaFT 等直接獎勵最佳化的方法雖然收斂速度快，但很快會陷入過擬合；而 DDPO 這一基於傳統策略梯度的強化學習微調方法由於沒有理由梯度資訊，微調速度顯著劣於其他方法。
同時，我們的定量實驗表明，我們的 Nabla-GFlowNet 可以更好保持生成樣本的多樣性。
Aesthetic Score 獎勵函式上的微調結果（微調 200 步，取圖片質量不坍塌的最好模型）。Nabla-GFlowNet（對應 Residual Nabla-DB 損失函式）方法微調的網路可以生成平均獎勵更高且不失自然的生成圖片。
相較於 ReFL，DRaFT 等直接獎勵最佳化的方法，Nabla-GFlowNet 更難陷入過擬合。
Aesthetic Score 獎勵函式上的定量結果 Nabla-GFlowNet 在獎勵收斂快的同時，保持更高的 DreamSim 多樣性指標（越高代表多樣性越好）和更低的 FID 分數（越低代表越符合預訓練模型的先驗）。
HPSv2 獎勵函式上的微調結果
ImageReward 獎勵函式上的微調結果
結語
我們利用生成流網路（GFlowNet）的框架，嚴謹地得到一個可以更好保持多樣性和先驗的高效的擴散模型獎勵微調方法，並且在 Stable Diffusion 這一常用的文生圖擴散模型上顯示出相較於其他方法的優勢。
© THE END
轉載請聯絡本公眾號獲得授權
投稿或尋求報道：[email protected]