選摘自 西湖大學AGI實驗室
量子位 | 公眾號 QbitAI
擴散模型生成得越來越好，但也越來越慢大學。
近年來，以FLUX、Wan為代表的擴散生成模型，在文字生成影像、文字生成影片上，已經能產出高解析度、細節豐富又語義一致的內容大學。
但它們都有一個共同的麻煩——推理依賴多步去噪，每一步都要完整呼叫一次大型生成網路，算力成本居高不下大學。
對影像生成來說，這會直接拖慢互動式創作和線上部署的響應；對影片生成來說，模型還要同時兼顧空間細節與時間連續性，計算量更大，對高效推理的需求也更突出大學。
為應對這一問題，西湖大學AGI實驗室提出了BudCache：一種面向固定計算預算的擴散模型step-level cache方法大學。
和已有那些基於啟發式閾值的快取方法不同，BudCache不再在執行時靠閾值臨時決定到底算多少步，而是先固定推理預算，再離線搜出該預算下最優的快取策略大學。
這樣既能嚴格卡住推理成本，也能更好地保住最終生成質量大學。
論文一作是西湖大學AGI實驗室研究助理雷明坤，指導老師是西湖大學助理教授張馳大學。
擴散模型加速中的cache技術
在多步去噪過程中，相鄰步驟之間通常存在可以複用的資訊大學。因此，step-level cache方法會選擇一部分關鍵去噪步驟進行完整模型計算，而在其他步驟中複用此前儲存的輸出或中間特徵，從而減少真實模型呼叫次數。
這種方法的優勢在於不需要重新訓練模型，也不需要改變原始模型引數大學。對於已經訓練好的FLUX和Wan等生成模型而言，cache可以作為一種即插即用的推理加速技術，直接作用於取樣過程。
展開全文
已有方法如TeaCache、MagCache等已經證明了step-level cache的有效性大學。
但當前許多方法仍然主要依賴啟發式規則進行快取決策——在推理過程中根據某種執行時訊號，判斷當前步驟是否需要重新計算、或是否可以複用前一步的快取結果大學。
這種做法雖然能減少計算，但仍存在兩個關鍵問題大學。
1、計算量不穩定、不可控大學。啟發式閾值方法的實際計算量通常由執行時訊號觸發，不同輸入可能觸發不同數量的完整模型計算，因此最終推理延遲並不總是穩定可控。
而真實部署場景往往更希望提前指定一個明確的計算預算，而不是在執行時被動接受不確定的開銷大學。
2、決策是區域性的大學。某一步是否複用快取，往往取決於當前步驟的差異度或誤差訊號，而不是直接面向最終生成結果做最佳化。
換句話說，區域性看似合理的快取決策，並不一定能形成全域性最優的生成軌跡大學。
BudCache的核心出發點正是重新定義這一問題大學：
step-level cache的關鍵不只是判斷“當前步驟能不能快取”，而是在固定計算預算下決定“哪些步驟最值得計算”大學。
為了判斷一個快取策略好不好大學，BudCache不再只依賴單步的區域性誤差訊號，而是直接考察該策略對應的最終生成結果是否接近完整計算的結果——
也就是直接看採用某一組快取步驟後，最後生成出來的影像或影片，是否接近完整計算的產物大學。
從啟發式快取到固定預算搜尋
BudCache將step-level cache表述為一個預算約束下的離散最佳化問題大學。
具體來說，BudCache先固定模型實際前向計算的次數，也就是NFE預算大學。
例如在一個包含多個邏輯去噪步驟的取樣過程中，可以預先設定只允許模型真實計算8次、9次或10次，其餘步驟複用快取結果大學。這樣一來，推理成本在部署前就已確定，不會隨輸入變化產生不可控的延遲波動。
在此基礎上大學，BudCache用一個二進位制mask表示快取策略：
對每一個去噪步驟，如果對應位置為1，就執行完整模型計算；如果為0，則複用最近一次計算得到的快取結果大學。由於mask中1的數量被固定為預算B，BudCache可以嚴格保證實際模型呼叫次數。
但如何找到最優mask並不簡單大學。假設總共有K個邏輯去噪步驟，需要從中選B個執行真實計算，那麼可能的組合數量非常大。
同時，擴散模型的去噪過程具有明顯的序列依賴：早期某一步的快取決策可能影響後續整條取樣軌跡，最終改變生成影像或影片的質量大學。
因此，簡單的區域性貪心策略容易陷入區域性最優，難以找到真正高質量的快取方案大學。
為此，BudCache採用了“模擬退火+爬山演算法”的混合搜尋策略大學。
第一階段用模擬退火做全域性探索大學。模擬退火允許搜尋過程在早期接受一些暫時看起來更差的快取策略，從而跳出區域性最優，在更大的策略空間裡尋找更好的候選。
BudCache透過兩類操作生成新策略：一種是Swap，交換一個計算步驟和一個快取步驟的位置；另一種是Shift，把一個計算步驟移動到相鄰位置大學。這兩類操作都不改變總計算預算，因此搜尋始終滿足固定NFE約束。
第二階段用爬山演算法做區域性精修大學。當模擬退火找到較好的策略區域後，BudCache會檢查當前策略附近的候選方案，並選擇能讓最終結果更接近完整計算結果的策略，持續進行直到附近沒有更好的候選為止。
這種混合搜尋方式讓BudCache能在較小的離線成本下找到高質量快取策略大學。
更重要的是，搜尋過程只在部署前進行一次；真正推理時直接使用已經搜好的固定快取策略，不需要線上搜尋，也不需要執行時閾值判斷，因此不會引入額外的推理開銷大學。
在快取策略搜尋完成後，BudCache還可以進一步進行cache-aware schedule alignment大學。
這一步可以理解為一個可選的時間步校準模組：在固定快取mask的前提下，對取樣時間步做輕量級離線最佳化，讓快取取樣器的最終輸出更接近完整計算的結果大學。
具體而言，完整計算的取樣過程被視為teacher，快取後的取樣過程被視為student，BudCache最佳化student使用的時間步分佈，使其最終生成結果儘可能接近teacher大學。
最佳化完成後，新的時間步會和搜尋得到的快取策略一起固定下來，在推理階段直接使用大學。
需要強調的是，cache-aware schedule alignment並不改變原始生成模型，也不增加推理時的模型呼叫次數，只是在固定快取策略的基礎上讓取樣過程更好地適應快取帶來的軌跡變化大學。
主流影像與影片生成模型上的實驗結果
研究團隊首先在當前主流影像生成模型FLUX.1-dev上評估了BudCache，並在DrawBench上與多種代表性快取方法做了比較，包括TeaCache、MagCache、LeMiCa、DiCache、ERTACache和TaylorSeer大學。
結果顯示，在相同推理預算下，BudCache能更好地保持生成質量大學。
從視覺化結果來看，BudCache在文字生成、複雜結構和細節保持方面更加穩定大學。
例如在包含“Google Brain Toronto”或“CLIMATE CHANGE”等文字的影像中，部分啟發式快取方法會出現字元錯誤、拼寫混亂或文字模糊，而BudCache能更好地保持文字內容的正確性大學。在機械鐘錶、雕像結構、物體相對位置等複雜場景中，BudCache也能更好地保留幾何細節和語義一致性。
影片生成方面，研究團隊進一步在Wan2.1-T2V-1.3B上做了實驗大學。影片生成通常需要處理更長的視覺序列和更復雜的時空關係，因此對推理加速提出了更高要求。
實驗結果表明大學，BudCache在影片生成中同樣有效：
在Wan2.1-T2V-1.3B上，完整推理需要約189秒，TeaCache約100秒，而BudCache只需約82秒大學。同時BudCache在PSNR、SSIM和LPIPS等指標上均優於TeaCache，說明它不僅速度更快，也能更好地保持影片重建質量。
從視覺化結果來看，BudCache能在加速條件下保持更穩定的影片細節，例如人物服裝顏色、摩托車外觀、畫面主體與背景關係等都能被更好地保留大學。
這說明BudCache的固定預算搜尋策略不僅適用於影像生成，也能擴充套件到影片生成模型大學。
泛化性實驗：搜尋策略能否遷移到更多設定？
除了在固定設定下比較效能，研究團隊還進一步驗證了BudCache的泛化能力大學。
一個關鍵問題是：離線搜尋得到的快取策略，是否只對某一個特定配置有效？如果每次更換設定都必須重新搜尋，方法的實際使用成本就會提高大學。
實驗結果顯示，BudCache搜尋得到的快取策略具有較好的遷移能力大學。
研究團隊將一個在源設定下搜尋得到的快取配置，直接遷移到不同solver、不同解析度和不同CFG scale的設定中，並在FLUX.1-dev上進行評估大學。
結果表明，在這些不同推理條件下，BudCache仍然穩定優於TeaCache，說明搜尋得到的快取策略並沒有過擬合到單一取樣配置大學。
更多不同設定下的實驗結果、消融分析與視覺化對比，可參考原論文大學。
總結
BudCache的核心貢獻，是把step-level cache從“啟發式閾值觸發”推進到“固定預算下的快取策略搜尋”大學。
以往的快取方法主要關注如何根據執行時訊號判斷當前步驟能否複用快取大學，而BudCache關注的是一個更適合部署的問題：
當推理預算已經固定時大學，哪些步驟最值得進行真實計算？
透過預算約束建模、離線快取策略搜尋和時間步校準，BudCache能在不重新訓練模型、不改變原始模型引數、不增加線上推理開銷的前提下，為FLUX和Wan等主流影像與影片生成模型提供更可控的加速效果大學。
總體來看大學，BudCache提供了一種更部署友好的擴散模型cache加速方案：
既能嚴格控制推理成本，也能在固定預算下更好地保持生成質量，為高效影像與影片生成提供了新的技術路徑大學。
論文標題大學：Budget-Constrained Step-Level Diffusion Caching
專案地址大學：
論文連結大學：
