編輯：編輯部 HYZ
【新智元導讀】劃時代的突破來了！來自NYU、MIT和谷歌的頂尖研究團隊聯手，為擴散模型開闢了一個全新的方向——測試時計算Scaling Law。其中，謝賽寧高徒為共同一作。
測試時計算+擴散模型，終於有人做到了！
AI大神謝賽寧團隊的新作，再次火得一塌糊塗。
眾所周知，大模型+測試時計算，是OpenAI篤定的全新Scaling Law。o3橫掃千軍，就是最好的證明。
那麼，擴散模型（DM）又如何呢？
這篇來自NYU、MIT和谷歌團隊新研究中，提出了一個創新性解決方案，通過設計通用搜索框架，從兩個維度來提升模型性能。
論文地址：https://arxiv.org/pdf/2501.09732
一是引入驗證器提供質量反饋，而是設計專門的演算法尋找更優質的雜訊候選。
簡言之，驗證器+演算法，成為了打通擴散模型測試時Scaling Law的核心要素。
謝賽寧表示，「2025年擴散模型的兩個令人興奮的方向：要麼（非常）小，要麼（非常）大」。
擴散模型，也有測試時Scaling Law
那麼，擴散模型+測試時計算，真的能夠看到全新的Scaling Law嗎？
從大佬的推文中，我們能夠獲得關鍵的一瞥。
謝賽寧發的一篇長帖表示，自己在第一次看到擴散模型時，被它們在推理過程中的擴展性震撼到了：
你在訓練時的計算量是固定的，但在測試時，你可以把計算量提高1000倍。
不過，這是在o1出現之前的事了。
然而，這種Scaling方式存在一個明顯的瓶頸——當去噪步驟增加到一定程度後，性能提升就會變得微乎其微。
根據共同一作Nanye Ma的介紹，團隊在這項研究中將推理時Scaling重新定義為採樣雜訊的搜索問題。
結果表明，增加搜索計算，可以提升生成性能，從而進一步推動擴散模型的能力。
擴散模型具有天然的靈活性，可以有在推理階段通過調整降噪步驟數量，來分配不同計算資源
然而，由於誤差累積，增加去噪步驟所帶來的性能提升通常會在幾十步後達到瓶頸。
因此，為了讓擴散模型在推理階段實現更大規模的Scaling，需要設計一個新的框架。
根據「優選效應」我們知道，在採樣過程中，一些初始雜訊會比其他的表現更好。
也就是說，可以嘗試通過投入更多的計算資源來搜索更好的雜訊，突破推理階段Scaling的極限。
那麼，問題來了：我們該如何判斷哪些採樣雜訊更好，又該如何有效地搜索這些雜訊呢？
為此，團隊提出了一種由兩個部分組成的搜索框架：驗證器提供反饋，演算法用於尋找更優的雜訊候選。
在使用SiT-XL的ImageNet上，不同驗證器與演算法的組合展現出顯著不同的Scaling特性。
接下來，團隊研究了搜索框架在文本條件生成任務中的能力。
在DrawBench上使用12B FLUX.1-dev模型時，通過結合所有驗證器進行搜索，可以提升樣本質量，但具體的改進效果在不同設置中差異很大。
這些觀察結果表明，沒有單一的搜索配置可以普遍適用；相反，每個任務都需要一個獨特的搜索設置來實現最佳的Scaling能力。
最後，團隊研究了推理時計算Scaling如何使較小的擴散模型受益。
在ImageNet上，SiT-L在有限的推理預算下優於SiT-XL；對於文本到圖像任務，0.6B PixArt-Sigma的總計算量僅為FLUX.1-dev的十分之一，卻實現了更優的表現。
這些結果表明，大量的訓練成本可以通過適度的推理時計算來部分抵消，從而更有效獲得更高質量的樣本。
接下來，就讓我們一起拜讀一下這篇大作吧。
「事半功倍」不行，那就開闢新路
正如開篇所述，全新搜索框架兩個設計軸：用於在搜索中提供反饋的驗證器，以及用於尋找更好雜訊候選項的演算法。
總而言之，最新研究的貢獻主要有三大點：
- 提出了一個用於擴散模型推理時Scaling的基礎框架。論文表明，通過搜索來Scaling函數評估次數（NFE）可以在各種生成任務和模型規模上帶來顯著改進，超越了僅增加去噪步驟的效果。
- 確定了所提出搜索框架中的兩個關鍵設計軸：提供反饋的驗證器和尋找更好雜訊候選項的演算法。通過研究了不同驗證器-演算法組合在各種任務中的表現，結果發現表明沒有一種配置是普遍最優的；每個任務反而需要特定的搜索設置才能實現最佳Scaling性能。
- 對驗證器與不同生成任務之間的對齊進行了廣泛分析。結果揭示了不同驗證器中嵌入的偏差，以及在每個不同的視覺生成任務中需要專門設計驗證器的必要性。
項目地址：https://inference-scale-diffusion.github.io/
將推理時Scaling作為搜索問題
研究人員將推理時Scaling構建為對採樣雜訊的搜索問題——具體來說，他們如何知道哪些採樣雜訊是好的，以及如何搜索它們。
站在更高層次上，他們提出考慮兩個設計軸：
1. 驗證器（Verifiers）：能夠對雜訊候選項質量提供反饋的預訓練模型；具體而言，它們接收生成的樣本和可選的相應條件作為輸入，並為每個生成的樣本輸出一個標量值作為分數。
2. 演算法（Algorithms）：基於驗證器反饋尋找更好雜訊候選項的函數。形式化定義為，演算法是函數：
該函數接收驗證器V、預訓練的擴散模型D_θ、N對生成的樣本及其對應條件，並根據雜訊和樣本之間的確定性映射輸出最佳初始雜訊。
在整個搜索過程中，f通常需要對D_θ進行多次前向傳遞。
研究人員將這些額外的前向傳遞稱為搜索成本，同樣用NFE來度量。
在下面展示的類條件ImageNet生成任務的設計流程中，研究人員使用在ImageNet-256上預訓練的SiT-XL模型，並使用二階Heun採樣器進行採樣。
他們用去噪步驟和搜索中使用的總NFE來衡量推理計算預算。去噪步驟固定為最優設置250，主要研究投入到搜索中的NFE的Scaling行為。
驗證器
驗證器方面，研究人員考慮了三種不同類型，旨在模擬三種不同的用例。
1. 預言驗證器（Oracle Verifier）：利用所選樣本最終評估的完整特權信息。
結果顯示，儘管預言驗證器很有效，但在實際場景中並不實用，因為它需要完全訪問樣本的最終評估結果。
研究人員將這些結果僅作為概念驗證，證明通過將計算資源投入到搜索中是可能實現更好的性能，並在推理時實現顯著的Scaling性能。
2. 監督驗證器（Supervised Verifier）：可以訪問預訓練模型來評估樣本的質量以及它們與指定條件輸入的對齊程度。
這裡，研究人員採用了兩個具有良好學習表示能力的模型：CLIP和DINO，並利用這兩個模型的分類視角。
在搜索過程中，他們將樣本輸入這些分類器，並選擇在生成時使用的類別標籤對應的最高logits值的樣本。
雖然與單純通過增加去噪步驟來擴展NFE相比，這種策略能有效提高樣本的IS分數，但作者使用的分類器與FID分數的目標只是部分對齊，因為它們是逐點操作的，並不考慮樣本的全局統計特性。
這可能導致樣本方差顯著降低，並且隨著計算量的增加最終表現為模式崩潰，這一點可以從不斷增加的精確度（Precision）和不斷下降的召回率（Recall）中得到證實。
3. 自監督驗證器（Self-Supervised Verifier）：使用在低雜訊水平（σ=0.4）和無雜訊（σ=0.0）樣本之間的特徵空間（分別由DINO/CLIP提取）餘弦相似度來評估初始雜訊的質量。
結果發現，這種相似度分數與DINO/CLIP分類器輸出的logits高度相關，因此可以作為監督驗證器的有效替代，如下所示。
演算法
在演算法方面，研究人員也考慮了三種不同的策略：
- 隨機搜索（Random Search）：簡單地從固定候選集中選擇最佳選項
- 零階搜索（Zero-Order Search）：利用驗證器反饋來迭代優化雜訊候選項
- 路徑搜索（Search over Paths）：利用驗證器反饋來迭代優化擴散採樣軌跡
下圖6所示，展現了這些演算法的性能，由於零階搜索和路徑搜索這兩種演算法的局部性特徵，它們都在一定程度上緩解了FID的多樣性問題，同時保持了Inception Score的縮放性能。
文本-圖像推理時Scaling
接下來，作者繼續研究搜索框架在更大規模的文本條件生成任務中的推理時Scaling能力，並研究驗證器與特定圖像生成任務之間的對齊情況。
為了對框架進行更全面的評估，研究人員使用了兩個數據集：DrawBench和T2I-CompBench。
模型方面，作者採用了新發布的FLUX.1-dev模型作為主幹網路。
分析結果：驗證器Hacking和驗證器-任務對齊
如圖8所示，並且根據LLM評分器的指示，使用所有驗證器進行搜索通常都能提高樣本質量，但具體的改進表現在不同設置下有所不同。
這證實了作者觀點：可以根據不同的應用場景專門選擇搜索設置。
在推理時使用FLUX.1-dev進行搜索的性能
從圖9中，作者還觀察到，隨著搜索預算的增加，評估指標的Scaling行為與ImageNet設置類似。
如下所示，這些驗證器在DrawBench和T2I-CompBench上的對比表現，突顯了某些驗證器可能比其他驗證器更適合特定任務。
這給作者帶來了啟發，設計更多針對特定任務的驗證器，下一步工作會繼續探討。
下表2種，作者展示了搜索演算法在DrawBench上的表現。
結果發現，這三種方法都能有效提高採樣質量，其中隨機搜索在某些方面的表現優於其他兩種方法，這是由於零階搜索和路徑搜索的局部性特徵所致。
搜索與微調兼容性
搜索和微調都旨在將最終樣本，與顯式獎勵模型或人類偏好對齊。
前者將樣本模式向特定驗證器的偏好方向偏移，而後者直接修改模型分布以與獎勵對齊。
這就引出了一個問題：在模型分布被修改後，我們是否仍然可以根據驗證器來偏移樣本模式？
作者使用DPO微調後的Stable Diffusion XL模型，並在DrawBench數據集上進行搜索。
由於該模型是在Pick-a-Pic數據集上微調的，他們用PickScore評估器替代了ImageReward。
如下表3所示，作者發現，搜索方法可以推廣到不同的模型，並且能夠提升已經對齊的模型的性能。這將成為一個有用的工具，可以用來：
- 緩解微調模型與獎勵模型產生分歧的情況
- 提高模型的泛化能力
推理計算投入的維度
由於擴散模型的迭代採樣特性，作者在搜索過程中可以在多個維度上分配計算資源。
下面列出這些維度並研究它們對搜索的影響。
- 搜索迭代次數：增加迭代次數可以使選定的雜訊更接近驗證器所認為的最優集合，作者在之前的所有實驗中都觀察到了這種行為。
- 每次搜索迭代的計算量：用NFEs/iter表示這種計算量。在搜索過程中，調整NFEs/iter可以揭示不同的計算最優區域，如下圖10所示。
計算投入的有效性
此外，研究人員還探索了在較小的擴散模型上，Scaling推理時計算量的有效性，並強調了其相對於未使用搜索的更大模型的性能效率。
對於ImageNet任務，他們使用了SiT-B和SiT-L，對於文本到圖像任務，除了FLUX.1-dev外，還使用了較小的基於Transformer的模型PixArt-ΣΣ。
由於不同大小的模型在每次前向傳播時的計算成本差異顯著，他們使用估計的GFLOPs來度量它們的計算成本，而不是使用NFEs。
如圖11所示，在ImageNet上對小型模型進行推理時計算量的Scaling可以非常有效——在固定計算預算的情況下，在推理計算資源有限的區域中，SiT-L的表現可以超過SiT-XL。
然而，這要求小型模型具有相對較強的性能基礎，SiT-B從搜索中獲得的收益不如SiT-L多，也沒有具有優勢的計算區域。
這些觀察結果也延伸到了基於文本條件的設置中，如下表4所示：僅使用1/10的計算量，PixArt-ΣΣ就超過了未使用搜索的FLUX-1.dev的性能，而使用大約兩倍的計算量時，PixArt-ΣΣ顯著超越了未使用搜索的FLUX.1-dev。
這些結果具有重要的實踐意義：在訓練時投入的大量計算資源可以通過生成時少量的計算來抵消，從而更高效地獲得更高質量的樣本。
作者介紹
Willis (Nanye) Ma
共同一作Nanye Ma是紐約大學的博士生，導師是謝賽寧，同時也是谷歌的學生研究員。此前，在NYU同時獲得了數學和計算機專業的學士學位。
他的研究方向是AI視頻生成，致力於改進潛表徵和實現長時序一致性。研究興趣則廣泛涵蓋CV領域，尤其是生成建模中的最優傳輸。
除了學術研究之外，他還喜歡打籃球、徒步、滑雪和攝影。
Shangyuan Tong
共同一作Shangyuan Tong，是麻省理工學院CSAIL的博士生，師從Tommi S. Jaakkola。此前，在UCSD獲得學士學位。
他的研究方向為深度學習和生成模型。
