擴散模型也能推理時Scaling，謝賽寧團隊重磅研究可能帶來文生圖新正規化

机器之心發表於2025-01-18
模型
對於 LLM，推理時 scaling 是有效的！這一點已經被近期的許多推理大模型證明：o1、o3、DeepSeek R1、QwQ、Step Reasoner mini……

但這個說法也適用於擴散模型嗎？

近日，紐約大學謝賽寧領導的一個團隊對這一方向進行了探索。具體來說，他們藉助通用搜尋框架系統性地探索了擴散模型的推理時 scaling 情況。他們發現，推理時 scaling 對擴散模型是有效的 —— 增加推理時間計算可以顯著提高擴散模型生成的樣本的質量，並且由於影像的複雜性，可以針對不同的應用場景，對框架中的元件進行不同形式的組合。
圖片
這篇論文的署名部分還透露出了一個有趣的小細節：謝賽寧的所屬機構標記成了谷歌並且通訊作者郵箱也使用了其谷歌域名的郵箱。但我們目前還並不清楚這是否意味著這位著名 AI 研究者已經正式加入谷歌，還是在谷歌兼職（謝賽寧此前已經是谷歌研究院訪問學者）。
圖片
  • 論文標題：Inference-Time Scaling for Diffusion Models beyond Scaling Denoising Steps
  • 論文連結：https://arxiv.org/pdf/2501.09732

有趣的是，就在前幾天，來自紐約大學和哥倫比亞大學的研究者也釋出了一篇類似方向的論文。看來，scaling 推理將是改進擴散模型的一個非常有前途的方向。
圖片
論文概覽

OpenAI o1 等模型的問世已經證明，在推理階段增加計算量可以讓 LLM 的效能進一步提升。但對於擴散模型，如何有效 scaling 推理時間計算以進一步提升效能還不夠明確。

擴散模型經過訓練以去除資料中的噪聲，是一類在連續資料領域占主導地位的生成模型，如影像、音訊和影片。為了生成單個樣本，它們的生成過程通常從純噪聲開始，需要經過訓練模型的多次前向傳遞來去噪並獲得乾淨的資料。這些前向傳遞因此被稱為去噪步驟。由於去噪步驟的數量可以調整，以在樣本質量和計算成本之間進行權衡，擴散模型的生成過程自然提供了在推理時分配計算預算的靈活性。

對於生成模型，這種計算預算通常用函式評估次數（NFE）來衡量，以確保與其他使用迭代取樣過程但沒有去噪能力的模型系列進行合理比較。

經驗觀察表明，僅透過將計算投入到去噪步驟中，效能提升在達到某個 NFE 後往往會趨於平穩，限制了推理過程中增加計算的收益。因此，以前關於擴散模型的工作長期以來一直專注於在保持高效能的同時儘可能減少推理時的 NFE 以提高效率。

但是，谷歌 DeepMind 和謝賽寧等人新論文的研究方向與此相反。

與 LLM 相比，擴散模型處理的是作為初始樣本注入的噪聲或在取樣過程中注入的噪聲的顯式隨機性。已有研究表明這些噪聲並非等價，即某些噪聲會帶來更好的生成結果。這一觀察為 scaling NFE 提供了除增加去噪步驟之外的另一個維度 —— 在取樣中搜尋更好的噪聲。

新論文研究了在推理過程中透過搜尋有效利用計算資源的方法，而不是僅將 NFE 分配給去噪步驟（這往往會很快導致效能不再增長），從而在推理時改進擴散模型的效能和可擴充套件性（圖 1）。
圖片
作者主要考慮搜尋框架中的兩個設計軸：用於在搜尋中提供反饋的驗證器，以及用於尋找更好噪聲候選項的演算法，遵循 LLM 中使用的術語。

對於驗證器，他們考慮了三種不同的設定，這些設定旨在模擬三種不同的使用場景：

  1. 掌握最終評估如何進行的特權資訊的場景；
  2. 掌握用於指導生成的條件資訊的場景；
  3. 沒有額外資訊可用的場景。

對於演算法，他們研究了：

  1. 隨機搜尋，它只是從固定的候選集中選擇最佳項；
  2. 零階搜尋，它利用驗證器反饋來迭代改進噪聲候選項；
  3. 路徑搜尋，它利用驗證器反饋來迭代改進擴散取樣軌跡。

作者首先在 ImageNet 類別條件生成這個相對簡單的設定中探討這些設計選擇並展示它們的有效性，為新框架提供一個具體例項。然後他們將這些設計選擇應用到更大規模的文字條件生成設定中，並評估他們提出的框架。

由於影像的複雜性質和文字條件包含的豐富資訊，需要對生成質量進行更全面的評估。因此，他們採用多個驗證器來進行搜尋中的推理時計算 scaling。這也使他們能夠探究每個驗證器具有的「偏見」，以及它們與生成任務的匹配程度。為了緩解對單個驗證器的過擬合，他們還嘗試了驗證器整合，並展示了它在不同基準測試中的良好泛化能力。

新論文的貢獻總結如下：

  • 提出了一個用於擴散模型推理時 scaling 的基礎框架。論文表明，透過搜尋去 scaling NFE 可以在各種生成任務和模型規模上帶來實質性改進，超越了僅增加去噪步驟的方法。此外，作者對推理時計算預算如何影響 scaling 效能進行了全面的實證分析。
  • 在提出的搜尋框架中確定了兩個關鍵設計軸：提供反饋的驗證器和尋找更好噪聲候選項的演算法。作者研究了不同驗證器 - 演算法組合在各種任務中的表現，他們的發現表明沒有一種配置是普遍最優的；每個任務反而需要一個獨特的搜尋設定才能實現最佳的 scaling 效能。
  • 對驗證器與不同生成任務之間的匹配度進行了廣泛分析。論文的結果揭示了不同驗證器中嵌入的偏見，以及在每個不同的視覺生成任務中需要專門設計驗證器的必要性。

如何在推理時間 scale

本文提出了一個框架，用於擴散模型的推理時間 scaling 。本文將這一挑戰表述為對取樣噪聲的搜尋問題；具體來說，本文需要知道哪些取樣噪聲是好的，以及如何搜尋這些噪聲？

在高層次上，本文考慮了兩個設計軸：

一是驗證器，用於評估候選樣本的質量。這些驗證器通常是能夠提供反饋的預訓練模型；驗證器函式形式如下：
圖片
這些函式接收生成的樣本以及可選的對應條件，並輸出一個標量值作為每個生成樣本的評分。

二是演算法，用於根據驗證器的評分找到更好的候選樣本。函式形式如下：
圖片
它接收驗證器 V、預訓練擴散模型 𝐷_θ ，以及 N 對生成的樣本及其相應的條件，並根據噪聲和樣本之間的確定性對映輸出最佳的初始噪聲。在整個搜尋過程中，f 通常會多次透過 𝐷_θ 進行前向傳遞。本文將這些額外的前向傳遞稱為搜尋成本，用 NFE 來衡量。

舉例來說，本文使用在 ImageNet 上預訓練的 SiT-XL 模型，解析度為 256×256，並使用二階 Heun 取樣器進行取樣，即除了初始噪聲外，取樣過程中沒有其他隨機性來源。此外，本文還透過去噪步驟和搜尋成本的總 NFE 來衡量推理計算預算。

本文從最簡單的搜尋演算法開始，即隨機取樣高斯噪聲，使用 ODE 生成樣本，並選擇與最佳驗證器評分對應的樣本（圖 2）。本文將這種演算法稱為隨機搜尋，它本質上是一種在所有噪聲候選上應用一次的 Best-of-N 策略。

對於驗證器，本文從最佳驗證器開始，即 Oracle 驗證器，他們假設其擁有關於選定樣本最終評估的全部特權資訊。對於 ImageNet，由於 FID 和 IS 通常被用作評估指標，因而直接將它們作為 Oracle 驗證器。
圖片
搜尋驗證器

雖然透過搜尋 scalingNFE 在使用 Oracle 驗證器時表現出令人印象深刻的效能（如圖 3 所示），但關鍵問題在於這種有效性是否可以推廣到使用更易獲取的預訓練模型的監督驗證器上。
圖片
為了研究這一點，本文選擇了兩個模型：CLIP 和 DINO 。對於 CLIP，本文遵循 Radford 等人的方法，使用透過提示工程生成的嵌入權重作為零樣本分類器。對於 DINO，本文直接使用預訓練的線性分類頭。在搜尋過程中，將樣本輸入分類器，並選擇與生成中使用的類別標籤對應的 logits 最高的樣本。

如圖 4 所示，與單純透過增加去噪步驟 scaling NFE 相比（圖 1），這種策略也有效提升了模型在 IS 上的效能。然而，本文注意到，由於這些分類器是逐點操作的，它們僅部分與 FID 評分的目標對齊。具體來說，它們生成的 logits 僅關注單個樣本的質量，而沒有考慮總體多樣性，這導致樣本方差顯著減少，並最終表現為隨著計算量增加而出現的模式崩潰。隨機搜尋演算法也難辭其咎，由於其無約束的搜尋空間，加速了搜尋向驗證器偏見的收斂。這種現象類似於強化學習中的獎勵駭客行為，因此將其稱為驗證器駭客。
圖片
儘管條件資訊在現實世界的生成任務中至關重要，但本文發現驗證器並不一定需要條件資訊來有效指導搜尋。如圖 5 所示，他們發現 DINO/CLIP 分類器輸出的 logits 與模型在低噪聲水平（σ=0.4）下的 x 預測值與最終生成的乾淨樣本（σ=0）之間的特徵空間（分別由 DINO/CLIP 提取）餘弦相似度之間存在強相關性。

因此，本文使用這種相似度評分作為分類 logits 的替代，並將這類驗證器稱為自監督驗證器，因為它們不需要額外的條件資訊。圖 5 再次觀察到了有效的 scaling 行為。

這一結果對於條件資訊不可用或難以獲取的用例來說（如醫學影像生成任務）是令人鼓舞的。由於這些限制在現實場景中並不常見，本文將自監督驗證器的進一步研究留待未來工作。圖片
搜尋演算法

前文的探索主要集中於簡單的隨機搜尋設定，他們發現這種方法可能導致驗證器駭客行為。基於這一認知，作者開始研究更細緻的搜尋演算法，這些演算法利用驗證器的反饋逐步最佳化候選樣本，每次只進行小幅調整，從而減輕過擬合風險。

具體來說，本文考慮了一種零階搜尋方法：

  1. 從一個隨機高斯噪聲 n 作為支點開始。
  2. 在支點的鄰域中找到 N 個候選項。
  3. 透過 ODE 求解器執行候選項以獲得樣本及其對應的驗證器得分。
  4. 找到最佳候選項，將其更新為支點，並重復步驟 1-3。

與零階最佳化類似，零階搜尋不涉及昂貴的梯度計算；相反，它透過在鄰域內進行多次前向函式評估來近似梯度方向。

本文還注意到，由於許多驗證器是可微的，理論上可以使用真實梯度進行一階搜尋，並且在實際中已有應用。然而，這需要在整個取樣過程中進行反向傳播，通常在時間和空間複雜度上都極為昂貴，尤其是在擴充套件大模型時。在實踐中，本文發現一階搜尋在 ImageNet 上並沒有表現出比零階搜尋更明顯的優勢，儘管它的成本更高。

擴散取樣過程的迭代特性為設計區域性搜尋演算法提供了其他可能性，沿著注入噪聲的取樣軌跡進行搜尋是可行的。本文提出了路徑搜尋來探索其中一種可能性。具體來說，

  1. 取樣 N 個初始獨立同分布噪聲，並執行 ODE 求解器直到某個噪聲水平 σ。噪聲樣本 x_σ 作為搜尋起點。
  2. 為每個噪聲樣本取樣 M 個獨立同分布噪聲，並模擬從 σ 到 σ+Δf 的前向加噪過程，生成大小為 M 的 {x_(σ+Δf)}。
  3. 在每個 x_(σ+Δf) 上執行 ODE 求解器到噪聲水平 σ + Δf - Δb，並獲得 x_(σ+Δf-Δb)。在這些樣本上執行驗證器並保留前 N 個候選樣本。重複步驟 2-3，直到 ODE 求解器達到 σ=0。
  4. 對剩餘的 N 個樣本透過隨機搜尋並保留最佳的一個。

為了確保迭代會終止，該團隊嚴格要求 Δb > Δf。另外，由於驗證器通常不適應有噪聲輸入，因此他們在步驟 3 中執行一個額外的去噪步驟，並使用乾淨的 x 預測與驗證器互動。這裡，主要的 scaling 軸是步驟 2 中新增的噪聲數量 M，在實踐中，他們研究了不同初始噪聲數量 N 的 scaling 行為。

因此該演算法被命名為 Paths-N。這兩種演算法見圖 2，從中可以看到，與隨機搜尋相比，零階搜尋和路徑搜尋都保持了很強的區域性性：前者在初始噪聲的鄰域中運作，後者則是在取樣過程的中間步驟中搜尋。

這些演算法的效能見圖 6。由於這兩種演算法的區域性性質，它們都在一定程度上緩解了 FID 的多樣性問題，同時保持了 Inception Score 的 scaling。對於零階搜尋，可以注意到增大 N 的效果很小，N = 4 似乎已經是對區域性最優的良好估計。至於路徑搜尋，可以看到不同的 N 值會導致不同的 scaling 行為，小 N 值在小生成預算下計算效率高，而大 N 值在增加更多計算時更具優勢。
圖片
在文生圖場景中的推理時 scaling

在更大規模文生圖任務中，該團隊研究了搜尋框架的推理時 scaling 能力，並研究了驗證器與特定影像生成任務之間的對齊效能。

資料集方面，該團隊採用了 DrawBench 和 T2I-CompBench 以實現更全面的評估。骨幹模型則是 FLUX.1-dev。驗證器方面，考慮到文生圖任務本質上的複雜性質，該團隊還擴充套件了監督式驗證器的選擇：Aesthetic Score Predictor、CLIPScore、ImageReward。其中 Aesthetic Score Predictor 可以預測人類對合成影像的視覺質量的評分，CLIPScore 則是透過 4 億對人工標註的影像 - 文字資料來對齊視覺和文字特徵，ImageReward 則更能體現更一般化的偏好。

此外，該團隊還創新地將這三個驗證器組合到一起，得到了第四個驗證器，稱為 Verifier Ensemble。可用於進一步擴充套件驗證器在不同評估方面的能力。

更多實驗設定詳情請參看原論文。

分析結果：驗證器 - 任務對齊

下面比較在不同資料集上驗證器和演算法組合的結果。

DrawBench。這是一個高度通用的資料集，包含來自不同類別的文字提示詞。圖 8 展示了在該資料集上的一些結果。
圖片
如圖 8 所示，根據 LLM Grader 指標，可知使用所有驗證器進行搜尋通常都會提高樣本質量，而具體的改進行為在不同設定下有所不同。這表明，可以根據不同的應用場景選擇專門的搜尋設定。

從圖 8 的左兩列中，可以看到使用 Aesthetic 和 CLIP 驗證器進行搜尋會過度擬合它們固有的偏見，對彼此產生負面影響。該團隊推測這兩個驗證器在評估上存在重大不匹配：Aesthetic Score 僅關注視覺質量，經常偏向於與文字提示詞偏離的高度風格化影像，而 CLIP 優先考慮視覺 - 文字對齊，犧牲了視覺質量。因此，如果在搜尋過程中使用一個驗證器（例如 Aesthetic Score），其偏見便會降低由另一個驗證器（例如 CLIP）評估的指標。

另外，與 ImageNet 設定類似，該團隊在這裡也觀察到了隨著搜尋預算增長而表現出的 scaling 行為。
圖片
透過觀察驗證器在 DrawBench 和 T2I-CompBench 上的行為，可以發現：驗證器的有效性取決於其標準與任務特定需求的匹配程度，某些驗證器比其他驗證器更適合特定任務。

演算法。表 2 展示了三種搜尋演算法在 DrawBench 上的表現。
圖片
可以看到，這三種方法全都能有效提高取樣質量，其中隨機搜尋在某些方面優於其他兩種方法。同樣，該團隊認為這種行為的原因是零階搜尋和路徑搜尋的區域性性質（圖 2）。由於這裡展示的所有驗證器和指標都是基於每個樣本進行評估的，隨機搜尋將大大加速向驗證器偏見的收斂，而其他兩種演算法需要對次優候選項進行改進。

搜尋與微調是相容的

搜尋和微調的目標都是將最終樣本與顯式獎勵模型或人類偏好對齊。前者將樣本模式向特定驗證器的偏見移動，後者是直接修改模型的分佈以與獎勵對齊。於是問題來了：在模型分佈被修改後，是否仍然可以根據驗證器移動樣本模式？

在這裡研究過的所有微調方法中，作為 RLHF 方法的一個更高效和更簡單的替代方案，DiffusionDPO 已被廣泛用於對齊大規模文生圖模型。為了回答這個問題，該團隊採用論文《Diffusion model alignment using direct preference optimization》中的 DPO 微調的 Stable Diffusion XL 模型，並在 DrawBench 資料集上進行搜尋。由於該模型是在 Pick-a-Pic 資料集上微調的，於是他們也用 PickScore 評估器替換 ImageReward。結果見表 3。
圖片
可以看到，搜尋方法可以泛化到不同的模型，並可以改進已對齊模型的效能。看起來，這個工具很有用，可以緩解微調模型與獎勵模型不一致的情況，並提高它們在其他指標上的泛化能力。

投入推理計算上不同維度的不同效果

由於擴散模型本質上的迭代取樣性質，因此透過搜尋來 scaling NFE 的維度有很多。這篇論文也對此展開了一些討論。

搜尋迭代的次數。直觀地說，增加搜尋迭代次數可以使選定的噪聲相對於驗證器更接近最優集，從而可以顯著提高效能。

每次搜尋迭代的計算量。在每次搜尋迭代中，去噪步驟數量都可以調整。這裡將模型採用的去噪步驟次數記為 NFEs/iter。儘管僅增加去噪步驟時模型效能很快達到平穩（圖 1），但該團隊觀察到在搜尋過程中，調整 NFEs/iter 可以揭示不同的計算最優區域，如圖 10 所示。在搜尋過程中，較小的 NFEs/iter 能實現高效收斂，但最終效能較低。相反，較大的 NFEs/iter 會導致收斂較慢但產生更好的效能。
圖片
最終生成的計算量。儘管可以自由調整最終生成的去噪步驟，但該團隊的做法是始終使用最優設定以獲得最佳的最終樣本質量。在 ImageNet 中，他們將去噪預算固定為 250 NFEs，在文字到影像設定中使用 30 步取樣器，因為進一步 scaling 時的效能將很快趨平。

投資計算的有效性

該團隊探索了在較小的擴散模型上進行推理時計算 scaling 的有效性，並強調了相對於它們不進行搜尋的大型模型的效能效率。結果見圖 11。
圖片
可以看到，在 ImageNet 上，對小型模型進行推理時計算 scaling 會非常有效。在固定的推理計算預算下，對小型模型進行搜尋可以超過不進行搜尋的大型模型。舉個例子，在推理計算量有限時，SiT-L 相對於 SiT-XL 存在優勢。但是，將 SiT-B 與其他兩個模型進行比較表明，這種有效性取決於小型模型相對強大的基線效能。當小型模型的基線效能明顯落後時，scaling 的收益是有限的，會導致結果次優。

在基於文字的設定中也出現這種現象，如表 4 所示。
圖片
可以看到，僅使用十分之一的計算量，PixArt-Σ 就超過了不進行搜尋的 FLUX-1.dev。而如果使用大約兩倍的計算量，PixArt-Σ 還能大幅超過不進行搜尋的 FLUX-1.dev。這些結果具有重要的實踐意義：在訓練中投入的大量計算資源可以透過生成期間的一小部分計算來抵消，從而能更高效地獲得更高質量的樣本。

相關文章