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

1. 個(gè)人理解

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

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

2. 摘要

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

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

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

3. 算法分析

3.1 效果對(duì)比

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

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



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



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

3.2 激活稀疏性

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

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

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



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



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



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

3.3 稀疏增量生成引擎SIGE

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

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



圖5 基于titling的稀疏卷積

4. 實(shí)驗(yàn)

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

4.1 主要結(jié)果

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

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

表1 定量質(zhì)量評(píng)估





圖6 所提出方法的定性對(duì)比

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

表2 模型效率對(duì)比

4.2 消融研究

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

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

表3 模型消融實(shí)驗(yàn)結(jié)果

5. 結(jié)論

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

審核編輯：劉清