1. 個人理解

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

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

2. 摘要

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

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

基于我們的算法，我們進一步提出稀疏增量生成引擎(SIGE)來將計算減少轉(zhuǎn)換為現(xiàn)成硬件上的延遲減少。通過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)所示，上一次編輯的生成結(jié)果已經(jīng)被計算，用戶進一步編輯其中9.4%的區(qū)域。然而，普通的DDIM需要生成整個圖像來計算新編輯的區(qū)域，在未改變的區(qū)域上浪費了80%的計算資源。解決這個問題的一個簡單方法是首先分割新編輯的片段，合成相應(yīng)的輸出片段，并將輸出與先前結(jié)果進行疊加。

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



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



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

3.2 激活稀疏性

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

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

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



圖3 稀疏卷積原理概述 具體的數(shù)學(xué)推導(dǎo)為：



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



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

3.3 稀疏增量生成引擎SIGE

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

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



圖5 基于titling的稀疏卷積

4. 實驗

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

4.1 主要結(jié)果

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

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

表1 定量質(zhì)量評估





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

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

表2 模型效率對比

4.2 消融研究

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

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

表3 模型消融實驗結(jié)果

5. 結(jié)論

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

審核編輯：劉清