空間稀疏推理(SSI)加速深度生成模型

1. 個人理解

生成模型近年來發(fā)展迅猛，已經表現出極強的真實感合成能力，在三維重建、AI繪畫、音視頻創(chuàng)作、可控圖像生成、真實圖像編輯等領域的應用廣泛。例如，即便沒有繪畫基礎，大家也可以很容易利用生成模型繪制大師級畫作。但近年來SOTA生成模型的主要問題是需要大量的計算資源，這一方面是由于深度網絡的框架較為復雜，另一方面是因為每次針對圖像可能只是做了很小的改動，但生成模型仍然需要重新計算整張圖像。

在2022 NeurIPS論文“Efficient Spatially Sparse Inference for Conditional GANs and Diffusion Models”中，CMU、MIT、斯坦福聯合提出了空間稀疏推理(SSI)，它可以利用編輯區(qū)域的空間稀疏性來加速深度生成模型，并且可以應用于各種生成模型！在Apple M1 Pro CPU上的推理速度加快了14倍！

2. 摘要

在圖像編輯期間，現有的深度生成模型傾向于從頭開始重新合成整個輸出，包括未編輯的區(qū)域。這導致了計算的顯著浪費，尤其是對于較小的編輯操作。在這項工作中，我們提出了空間稀疏推理(SSI)，這是一種通用技術，它選擇性地對編輯區(qū)域執(zhí)行計算，并加速各種生成模型，包括條件GAN和擴散模型。

我們的主要觀察是，用戶傾向于對輸入圖像進行漸進的改變。這促使我們緩存和重用原始圖像的特征圖。給定一個編輯過的圖像，我們稀疏地將卷積濾波器應用于編輯過的區(qū)域，同時為未編輯的區(qū)域重用緩存的特征。

基于我們的算法，我們進一步提出稀疏增量生成引擎(SIGE)來將計算減少轉換為現成硬件上的延遲減少。通過1.2%的面積編輯區(qū)域，我們的方法減少了7.5倍的DDIM和18倍的GauGAN的計算，同時保持視覺保真度。通過SIGE，我們加速了3.0倍在RTX 3090上的DDIM和6.6倍在蘋果M1 Pro CPU上的推理時間，以及4.2倍在RTX 3090上的GauGAN和14倍在Apple M1 Pro CPU上的推理時間。

3. 算法分析

3.1 效果對比

話不多說，先看效果！ 如圖1(a)所示，上一次編輯的生成結果已經被計算，用戶進一步編輯其中9.4%的區(qū)域。然而，普通的DDIM需要生成整個圖像來計算新編輯的區(qū)域，在未改變的區(qū)域上浪費了80%的計算資源。解決這個問題的一個簡單方法是首先分割新編輯的片段，合成相應的輸出片段，并將輸出與先前結果進行疊加。

但這種方法很容易在新編輯和未編輯的區(qū)域之間產生明顯接縫。 為解決此問題，作者提出了空間稀疏推理(Spatially Sparse Inference，SSI)和稀疏增量生成引擎(SIGE)，如圖2(b)所示。作者的關鍵思想是重用之前編輯的緩存特征圖并稀疏更新新編輯的區(qū)域。



圖1 在交互式編輯場景中，用戶添加了一個新的建筑，它占據了9.4%的像素 如圖2所示是作者與其他生成模型的對比結果?？梢钥闯鱿噍^于DDIM和GauGAN，作者提出的方法計算量大幅降低。其中相較于DDIM，計算量降低了4~6倍，相較于GauGAN，計算量降低了15倍，當引入模型壓縮方法以后，計算量進一步減少了47倍。注意一下這里的MACs指標，1 MAC等于2 FLOPs。



圖2 作者提出方法的計算量對比 重要的是，這種方法可以很容易得推廣到其他生成模型中！進一步降低計算量！

3.2 激活稀疏性

作者的啟發(fā)靈感是，在交互式圖像編輯期間，用戶通常每次只編輯圖像中的一部分。因此，可以為未編輯區(qū)域重用原始圖像的激活。因此，可以利用編輯區(qū)域的空間稀疏性來加速深度生成模型。具體來說，給定用戶輸入，首先計算一個差異掩碼來定位新編輯的區(qū)域。

對于每一個模型中的卷積層，僅稀疏地將濾波器應用于掩蔽的區(qū)域，而對未改變的區(qū)域重復使用先前的生成模型。稀疏更新可以在不損害圖像質量的情況下顯著減少計算量。 此外，由于稀疏更新涉及聚集-分散過程，現有的深度學習框架會導致顯著的延遲開銷。

為了解決這個問題，作者進一步提出了稀疏增量生成引擎(SIGE)來將算法的理論計算減少轉化為在各種硬件上測量的延遲減少。 如圖3所示是具體的算法原理，首先預計算原始輸入圖像的所有激活。在編輯過程中，通過計算原始圖像和編輯圖像之間的差異掩模來定位編輯區(qū)域。然后，對未編輯的區(qū)域重新使用預先計算的激活，并且通過對它們應用卷積濾波器來僅更新已編輯的區(qū)域。



圖3 稀疏卷積原理概述 具體的數學推導為：



其中Al表示第l層卷積層F的輸入tensor，W和b分別是第l層的權重和偏置。 如圖4所示，ΔAl共享了用戶所做編輯中的結構化空間稀疏性，因此非零值主要聚集在編輯區(qū)域內。這樣就可以直接使用原始圖像和編輯后的圖像來計算一個差異掩碼，并用這個掩碼對ΔAl進行稀疏化。



圖4 左圖:圖像編輯示例。右圖:在不同特征圖分辨率下，DDIM第l層的?Al通道平均值

3.3 稀疏增量生成引擎SIGE

但是如何利用結構化稀疏性來加速Wl*ΔAl呢？ 一種簡單的方法是為每個卷積從ΔAl中裁剪一個矩形編輯區(qū)域，并且只計算裁剪區(qū)域的特征。但作者發(fā)現這種裁剪方法對于不規(guī)則的編輯區(qū)域(圖4所示的例子)效果很差。 因此，如圖5所示，作者使用基于tiling的稀疏卷積算法。

首先將差異掩碼向下采樣到不同的比例，并擴展向下采樣的掩碼，將ΔAl在空間上劃分為多個相同大小的小塊。每個塊索引指的是具有非零元素的單個塊。然后將非零塊沿批維度進行相應的聚集，并將其饋入卷積Fl。最后，根據索引將輸出塊分散成零張量，以恢復原始空間大小，并將預先計算的殘差計算。



圖5 基于titling的稀疏卷積

4. 實驗

作者分別在三個模型上進行實驗，包括擴散模型和GAN模型：DDIM、Progressive Distillation (PD)、GauGAN。使用LSUN Church數據集和Cityscapes數據集進行實驗。在評價指標方面，使用PSNR、LPIPS、FID來評估圖像質量。對于Cityscapes數據集還是用了語義分割中的mIoU這一指標。

4.1 主要結果

表1所示是作者方法應用于DDIM、Progressive Distillation (PD)和GauGAN的定量結果，并在圖6中顯示了定性結果。對于PSNR和LPIPS來說，對于DDIM和Progressive Distillation (PD)來說，作者方法始終優(yōu)于所有基線，并獲得與原始模型相當的結果。當由于全局上下文不足而編輯的區(qū)域很小時，補片推理失敗。

盡管作者方法僅將卷積濾波器應用于局部編輯區(qū)域，但是可以重用存儲在原始激活中的全局上下文。因此，作者的方法可以像原始模型一樣執(zhí)行。對于GauGAN，作者的方法也比GAN Compression執(zhí)行得更好，MACs減少更多。當應用于GAN Compression時，進一步實現了大約40倍MACs的減少，性能略有下降，同時超過了0.19 GauGAN和GAN Comp。

表1 定量質量評估





圖6 所提出方法的定性對比

在模型模型效率方面，作者測試了了圖6所示的編輯實例在4個設備上的加速比，包括RTX 3090、RTX 2080Ti、Intel Core i9-10920X CPU和Apple M1 Pro CPU，并且設置batch size為1來模式真實應用。對于GPU設備，首先執(zhí)行200次預熱運行，并測量接下來200次運行的平均耗時。對于CPU設備，首先執(zhí)行10次預熱運行和10次測試運行，重復此過程5次并報告平均耗時。結果如表2所示。

表2 模型效率對比

4.2 消融研究

表3顯示了消融研究結果。 內存使用：原始圖像預先計算的激活需要額外的存儲量，但作者所提出的方法僅將DDIM、PD、GauGAN和GAN Compression的單次轉發(fā)的峰值內存使用量分別增加了0.1G、0.1G、0.8G和0.3G。表3(a)所示是在RTX 2080Ti上為DDIM添加的每個內核優(yōu)化的有效性。

簡單地應用基于tiling的稀疏卷積可以將計算量減少7.6倍。表3(b)是在TensorRT上進行了模型部署，TensorRT進一步加快了模型的運行效率。

表3 模型消融實驗結果

5. 結論

在2022 NeurIPS論文“Efficient Spatially Sparse Inference for Conditional GANs and Diffusion Models”中，CMU、MIT、斯坦福聯合提出了空間稀疏推理(SSI)和稀疏增量生成引擎(SIGE)。這種算法減少了現有深度生成模型的計算資源浪費問題，對于生成模型的落地和應用具有重要意義。重要的是，算法已經開源，并且可以應用于各種生成模型，包括條件GAN和擴散模型！

審核編輯：劉清