Published on

擴散模型推理擴展新範式

Authors
  • avatar
    Name
    Ajax
    Twitter

擴散模型推理擴展：新範式探索

近年來，大型語言模型 (LLMs) 在推理階段的擴展性已展現出驚人的成效。諸如 o1, o3, DeepSeek R1, QwQ 和 Step Reasoner mini 等模型都證明，增加推理時的計算量能顯著提升效能。這引發了一個問題：這種原則是否也適用於擴散模型？

紐約大學謝賽寧 (Xie Saining) 領導的團隊深入探討了這個問題。他們透過使用通用搜尋框架進行系統性研究，發現推理時的擴展確實對擴散模型有效。增加推理時的計算量能顯著改善生成樣本的品質。此外，圖像的複雜性允許框架內不同組件的組合，以適應各種應用場景。

主要研究發現

  • 推理時擴展對擴散模型有效： 在推理期間分配更多計算資源可以產生更高品質的樣本。
  • 組件組合的靈活性： 該框架允許不同的組件配置，以滿足各種應用需求。
  • 超越去噪步驟： 研究表明，在採樣過程中尋找更好的雜訊是擴展 NFE 的另一個維度，而不僅僅是增加去噪步驟。
  • 兩個設計軸： 該框架側重於兩個關鍵設計軸：
    • 驗證器 (Verifiers)： 在搜尋過程中提供回饋。
    • 演算法 (Algorithms)： 尋找更好的雜訊候選。

研究方法

研究團隊模擬了三種不同的驗證器情境，以模擬各種用例：

  1. 擁有最終評估特權資訊的情境。
  2. 擁有條件資訊以引導生成的情境。
  3. 沒有額外資訊可用的情境。

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

  1. 隨機搜尋 (Random Search)： 從一組固定的候選中選擇最佳的。
  2. 零階搜尋 (Zero-Order Search)： 使用驗證器回饋迭代改進雜訊候選。
  3. 路徑搜尋 (Path Search)： 使用驗證器回饋迭代改進擴散採樣軌跡。

該研究最初在相對簡單的 ImageNet 類別條件生成設置中探索了這些設計。隨後，他們將這些設計應用於更大規模的文本條件生成，並評估了他們提出的框架。

推理時間的擴展

該論文提出了一個擴展擴散模型推理時間的框架，將挑戰定義為搜尋最佳採樣雜訊。這個過程涉及兩個核心組件：

  • 驗證器 (Verifiers)： 這些是預先訓練的模型，用於評估生成樣本的品質。它們接收生成的樣本以及可選的條件，並輸出一個標量分數。
  • 演算法 (Algorithms)： 這些演算法使用驗證器分數來尋找更好的候選樣本。該函數接收一個驗證器 (V)、一個預先訓練的擴散模型 (_θ) 以及一組生成的樣本和條件，並輸出最佳的初始雜訊。

總推理預算由函數評估總數 (NFE) 衡量，包括去噪步驟和搜尋成本。

搜尋驗證器

研究人員首先使用 Oracle 驗證器，該驗證器具有關於選定樣本最終評估的完整資訊。對於 ImageNet，這包括 FID 和 IS 等指標。然後，他們探索了更容易取得的預先訓練模型作為監督驗證器，例如 CLIP 和 DINO。這些模型用於對樣本進行分類，選擇與類別標籤對應的最高 logits 的樣本。

然而，他們觀察到，這些逐點運算的分類器僅部分符合 FID 分數的目標。這導致隨著計算量的增加，樣本變異減少和模式崩潰。這種被稱為「驗證器駭客」的現象，因隨機搜尋演算法不受約束的搜尋空間而加速。

有趣的是，研究發現驗證器不一定需要條件資訊來有效地引導搜尋。他們觀察到來自 DINO/CLIP 分類器的 logits 與低雜訊水平下的 x 預測和最終乾淨樣本之間的特徵空間的餘弦相似度之間存在很強的相關性。這導致使用自我監督驗證器，這些驗證器不需要額外的條件資訊，並且仍然表現出有效的擴展行為。

搜尋演算法

為了緩解驗證器駭客的問題，研究人員探索了更精細的搜尋演算法，這些演算法逐漸優化候選樣本。這包括零階搜尋方法：

  1. 從隨機高斯雜訊作為基準點開始。
  2. 在基準點的鄰域中找到 N 個候選者。
  3. 通過 ODE 求解器運行候選者以獲得樣本和驗證器分數。
  4. 使用最佳候選者更新基準點並重複步驟 1-3。

他們還研究了一種路徑搜尋演算法，該演算法探索沿著採樣軌跡搜尋的可能性：

  1. 採樣 N 個初始雜訊樣本，並將 ODE 求解器運行到雜訊水平 σ。
  2. 向每個樣本添加雜訊，並模擬前向加噪過程。
  3. 在每個加噪樣本上運行 ODE 求解器，並根據驗證器分數保留前 N 個候選者，重複直到 ODE 求解器達到 σ=0。
  4. 隨機搜尋剩餘的 N 個樣本並保留最佳的一個。

與隨機搜尋相比，零階和路徑搜尋演算法都保持了很強的局部性。

在文本到圖像場景中的擴展

研究團隊檢驗了搜尋框架在更大規模的文本到圖像任務中的擴展能力。他們使用 DrawBench 和 T2I-CompBench 數據集進行評估，並以 FLUX.1-dev 模型作為基礎模型。他們還擴展了監督驗證器的選擇，包括 Aesthetic Score Predictor、CLIPScore 和 ImageReward。此外，他們還通過組合這三個驗證器創建了一個驗證器集成。

分析：驗證器與任務的對齊

該研究比較了不同數據集上各種驗證器-演算法組合的結果。在 DrawBench 上，他們發現使用所有驗證器通常可以提高樣本品質。然而，他們觀察到，單獨使用 Aesthetic 和 CLIP 驗證器可能會導致過度擬合其偏差，從而對彼此產生負面影響。這源於它們的評估重點不匹配：Aesthetic Score 側重於視覺品質，通常偏愛高度風格化的圖像，而 CLIP 優先考慮視覺文本對齊，有時會犧牲視覺品質。

他們指出，某些驗證器更適合特定任務，並且驗證器的有效性取決於其與任務要求的對齊性。

演算法效能

三種搜尋演算法（隨機、零階和路徑）都有效地提高了 DrawBench 上的採樣品質。然而，由於其他兩種方法的局部性質，隨機搜尋在某些方面表現更好。隨機搜尋更快地收斂到驗證器偏差，而其他兩種演算法需要改進次優候選者。

與微調的相容性

研究團隊調查了他們的搜尋方法與微調模型的相容性。他們使用了一個 DPO 微調的 Stable Diffusion XL 模型，發現搜尋方法可以推廣到不同的模型，並提高已經對齊模型的效能。

推理計算不同維度的影響

該研究探討了推理計算的不同方面如何影響結果：

  • 搜尋迭代次數： 增加迭代次數可以使雜訊更接近最佳值。
  • 每次搜尋迭代的計算量： 調整每次迭代的去噪步驟數量可以揭示不同的計算最佳區域。
  • 最終生成計算： 該團隊使用最終去噪步驟的最佳設置，以確保最終樣本的最高品質。

計算投資的有效性

研究人員探討了在較小的擴散模型上進行推理時擴展的有效性。他們發現，對於 ImageNet，擴展較小的模型可能非常有效。在某些情況下，在較小的模型上進行搜尋可以勝過沒有搜尋的較大型模型。然而，有效性取決於較小模型的基準效能。

在基於文本的設置中，PixArt-Σ 僅使用一小部分的計算量，就優於 FLUX-1.dev。這些結果表明，在訓練期間花費的大量計算資源可以通過在生成期間使用較少的計算量來抵消，從而更有效地產生更高品質的樣本。