Diffusion約2倍無損加速！訓練——推理協同的快取學習框架來了｜HKUST&北航&商湯

HarmoniCa團隊投稿量子位 | 公眾號 QbitAI

面對擴散模型推理速度慢、成本高的問題，HKUST&北航&商湯提出了全新快取加速方案——HarmoniCa：訓練-推理協同的特徵快取加速框架，突破DiT架構在部署端的速度瓶頸，成功實現高效能無損加速。

△HarmoniCa整體壓縮框架

由於現有指標並不能完全反映影像效果優劣，因此該團隊研究人員提供了大量視覺化效果對比圖，更多對比請看原論文。

△PIXART-影像生成效果圖

該工作已被ICML 2025接收為Poster，並開源專案程式碼。

Diffusion 加速難在哪？

Diffusion Transformer（DiT）作為高解析度影像生成主力架構，在推理階段仍面臨「重複計算多」「耗時嚴重」的現實瓶頸。例如，使用PIXART-α生成一張2048×2048影像即需14秒，嚴重影響落地效率。

近期「特徵快取（Feature Caching）」成為加速新思路，但已有方法普遍存在兩大關鍵問題：

前序時間步無感知：訓練階段忽略快取歷史，推理時則高度依賴先前結果，二者邏輯斷裂；

訓練目標錯位：訓練對準中間噪聲誤差，推理關注最終影像質量，最佳化方向南轅北轍；

這兩大錯配，導致已有快取學習方法加速有限、影像失真明顯。快取機制的基本工作原理如下：

△快取機制

HarmoniCa快取學習框架

一句話總結：目標一致、路徑同步，訓練與推理真正協同最佳化

該工作提出的HarmoniCa框架透過兩個關鍵機制，從根本上解決了以往學習型特徵快取方法中的訓練-推理脫節問題：

一、Step-Wise Denoising Training（SDT）

逐步去噪訓練，模擬推理全流程，誤差不再層層積累。

傳統方法在訓練時僅取樣某個時間步，快取是空的，完全跳過了“歷史快取影響”，而推理時，快取是從頭累積的，訓練和推理根本不是一回事。

進而該工作提出 SDT 來打破這一不一致：

1）構建完整的 T 步去噪過程，與推理一致；2）教師-學生結構：學生使用快取進行去噪，教師不使用快取作為“理想輸出”；3）每一時間步的Router都被獨立更新，顯式對齊多輪快取路徑下的輸出誤差；4）學生模型每步將自己的輸出作為下一個輸入，使得誤差傳播機制貼近真實推理軌跡。

效果：SDT顯著降低了時間步間誤差積累，提升最終影像清晰度與穩定性。

△SDT 有效抑制誤差蔓延（紅色為舊方法，藍色為 SDT）

二、Image Error Proxy Objective（IEPO）

一句話總結：不是“中間好”，而是“最後影像好”，最佳化目標就是結果本身。

以往方法訓練時只對齊每一步的噪聲誤差，而推理的目標是最終影像質量，兩者目標嚴重錯配，導致快取Router學出來“看似合理”但效果很差。

該工作提出 IEPO 機制，核心思想是：

透過代理項 λ(t) 來估算“使用快取 vs 不使用快取”在時間步 t 對最終影像 x₀ 的影響；

越關鍵的時間步，其 λ(t) 越大，引導 Router 減少該步快取複用，保留精度；

每隔若干輪重新生成一批影像，動態更新 λ(t)，保證目標始終貼合訓練狀態。IEPO 的最佳化目標為：

即在“影像質量”與“加速率”之間實現可控權衡。

實驗結果

該工作在兩個典型任務場景中驗證了HarmoniCa的有效性：

分類條件生成（DiT-XL/2@ImageNet）
文字生成影像（PIXART-α@COCO，多解析度）

對比方法包括當前最佳的快取學習方法 Learning-to-Cache (LTC)、啟發式快取方法 FORA / ∆-DiT，以及多種加速器設定（DDIM 步數縮減、量化剪枝等）。

分類條件生成（DiT-XL/2 256×256）

重點結論：

在高壓縮率場景（10步推理）下，HarmoniCa保持影像質量優勢，FID比LTC更低、IS更高
同時達成更高快取利用率，提升實際加速效果

文字生成影像（PIXART-α 256×256）

重點結論：

即使在2K高解析度下，HarmoniCa仍保持1.69×實際加速
在CLIP語義匹配、FID等主流指標上均超過FORA

量化/剪枝VS HarmoniCa

除了與主流快取方法的對比，該工作也評估了HarmoniCa相比剪枝和量化等壓縮技術的表現。在統一的 20 步取樣設定下，傳統方案如 PTQ4DiT、EfficientDM等雖然模型更小，但實際加速依賴硬體支援，特別是一些定製CUDA核心在H800等新架構上表現並不穩定。更重要的是，量化模型在小步數取樣時往往精度下降嚴重，PTQ4DiT就出現了明顯的效能下滑。而HarmoniCa不依賴底層魔改，無需專用硬體，在各種主流取樣器和裝置上都能穩定提速，保持影像質量，是當前更通用、更穩妥的部署選擇。

△與量化/剪枝方法的比較

與量化結合

該工作還驗證了HarmoniCa與模型量化技術的高度相容性。在 PIXART-α 256×256 場景下，將HarmoniCa應用於4bit量化模型（EfficientDM），推理速度從1.18×提升至1.85×，FID僅略增0.12，幾乎無感知差異。說明HarmoniCa不僅可獨立提速，也能作為“加速外掛”疊加於量化模型之上，進一步釋放效能潛力。未來，該工作也計劃探索其與剪枝、蒸餾等技術的組合能力，為DiT模型的輕量部署開闢更多可能。