擴散模型也能推理時Scaling，謝賽寧團隊重磅研究可能帶來文生圖新范式
對于 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 中使用的術語。
對于驗證器，他們考慮了三種不同的設置，這些設置旨在模擬三種不同的使用場景：
- 掌握最終評估如何進行的特權信息的場景；
- 掌握用于指導生成的條件信息的場景；
- 沒有額外信息可用的場景。
對于算法，他們研究了：
- 隨機搜索，它只是從固定的候選集中選擇最佳項；
- 零階搜索，它利用驗證器反饋來迭代改進噪聲候選項；
- 路徑搜索，它利用驗證器反饋來迭代改進擴散采樣軌跡。
作者首先在 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 行為。
這一結果對于條件信息不可用或難以獲取的用例來說（如醫學圖像生成任務）是令人鼓舞的。由于這些限制在現實場景中并不常見，本文將自監督驗證器的進一步研究留待未來工作。
搜索算法
前文的探索主要集中于簡單的隨機搜索設置，他們發現這種方法可能導致驗證器黑客行為?；谶@一認知，作者開始研究更細致的搜索算法，這些算法利用驗證器的反饋逐步優化候選樣本，每次只進行小幅調整，從而減輕過擬合風險。
具體來說，本文考慮了一種零階搜索方法：
- 從一個隨機高斯噪聲 n 作為支點開始。
- 在支點的鄰域中找到 N 個候選項。
- 通過 ODE 求解器運行候選項以獲得樣本及其對應的驗證器得分。
- 找到最佳候選項，將其更新為支點，并重復步驟 1-3。
與零階優化類似，零階搜索不涉及昂貴的梯度計算；相反，它通過在鄰域內進行多次前向函數評估來近似梯度方向。
本文還注意到，由于許多驗證器是可微的，理論上可以使用真實梯度進行一階搜索，并且在實際中已有應用。然而，這需要在整個采樣過程中進行反向傳播，通常在時間和空間復雜度上都極為昂貴，尤其是在擴展大模型時。在實踐中，本文發現一階搜索在 ImageNet 上并沒有表現出比零階搜索更明顯的優勢，盡管它的成本更高。
擴散采樣過程的迭代特性為設計局部搜索算法提供了其他可能性，沿著注入噪聲的采樣軌跡進行搜索是可行的。本文提出了路徑搜索來探索其中一種可能性。具體來說，
- 采樣 N 個初始獨立同分布噪聲，并運行 ODE 求解器直到某個噪聲水平 σ。噪聲樣本 x_σ 作為搜索起點。
- 為每個噪聲樣本采樣 M 個獨立同分布噪聲，并模擬從 σ 到 σ+Δf 的前向加噪過程，生成大小為 M 的 {x_(σ+Δf)}。
- 在每個 x_(σ+Δf) 上運行 ODE 求解器到噪聲水平 σ + Δf - Δb，并獲得 x_(σ+Δf-Δb)。在這些樣本上運行驗證器并保留前 N 個候選樣本。重復步驟 2-3，直到 ODE 求解器達到 σ=0。
- 對剩余的 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?？捎糜谶M一步擴展驗證器在不同評估方面的能力。
更多實驗設置詳情請參看原論文。
分析結果：驗證器 - 任務對齊
下面比較在不同數據集上驗證器和算法組合的結果。
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。這些結果具有重要的實踐意義：在訓練中投入的大量計算資源可以通過生成期間的一小部分計算來抵消，從而能更高效地獲得更高質量的樣本。
