生成模型近年来发展迅猛,已经表现出极强的真实感合成能力,在三维重建、AI绘画、音视频创作、可控图像生成、真实图像编辑等领域的应用广泛。例如,即便没有绘画基础,大家也可以很容易利用生成模型绘制大师级画作。但近年来SOTA生成模型的主要问题是需要大量的计算资源,这一方面是由于深度网络的框架较为复杂,另一方面是因为每次针对图像可能只是做了很小的改动,但生成模型仍然需要重新计算整张图像。
在2022 NeurIPS论文“Efficient Spatially Sparse Inference for Conditional GANs and Diffusion Models”中,CMU、MIT、斯坦福联合提出了空间稀疏推理(SSI),它可以利用编辑区域的空间稀疏性来加速深度生成模型,并且可以应用于各种生成模型!在Apple M1 Pro CPU上的推理速度加快了14倍!代码刚刚开源!
代码链接:https://github.com/lmxyy/sige
原文链接:https://arxiv.org/abs/2211.02048
2. 摘要
在图像编辑期间,现有的深度生成模型倾向于从头开始重新合成整个输出,包括未编辑的区域。这导致了计算的显著浪费,尤其是对于较小的编辑操作。在这项工作中,我们提出了空间稀疏推理(SSI),这是一种通用技术,它选择性地对编辑区域执行计算,并加速各种生成模型,包括条件GAN和扩散模型。我们的主要观察是,用户倾向于对输入图像进行渐进的改变。这促使我们缓存和重用原始图像的特征图。给定一个编辑过的图像,我们稀疏地将卷积滤波器应用于编辑过的区域,同时为未编辑的区域重用缓存的特征。基于我们的算法,我们进一步提出稀疏增量生成引擎(SIGE)来将计算减少转换为现成硬件上的延迟减少。通过1.2%的面积编辑区域,我们的方法减少了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%的区域。然而,普通的DDIM需要生成整个图像来计算新编辑的区域,在未改变的区域上浪费了80%的计算资源。解决这个问题的一个简单方法是首先分割新编辑的片段,合成相应的输出片段,并将输出与先前结果进行叠加。但这种方法很容易在新编辑和未编辑的区域之间产生明显接缝。
为解决此问题,作者提出了空间稀疏推理(Spatially Sparse Inference,SSI)和稀疏增量生成引擎(SIGE),如图2(b)所示。作者的关键思想是重用之前编辑的缓存特征图并稀疏更新新编辑的区域。
图1 在交互式编辑场景中,用户添加了一个新的建筑,它占据了9.4%的像素
如图2所示是作者与其他生成模型的对比结果。可以看出相较于DDIM和GauGAN,作者提出的方法计算量大幅降低。其中相较于DDIM,计算量降低了4~6倍,相较于GauGAN,计算量降低了15倍,当引入模型压缩方法以后,计算量进一步减少了47倍。注意一下这里的MACs指标,1 MAC等于2 FLOPs。
重要的是,这种方法可以很容易得推广到其他生成模型中!进一步降低计算量!
3.2 激活稀疏性
作者的启发灵感是,在交互式图像编辑期间,用户通常每次只编辑图像中的一部分。因此,可以为未编辑区域重用原始图像的激活。因此,可以利用编辑区域的空间稀疏性来加速深度生成模型。具体来说,给定用户输入,首先计算一个差异掩码来定位新编辑的区域。对于每一个模型中的卷积层,仅稀疏地将滤波器应用于掩蔽的区域,而对未改变的区域重复使用先前的生成模型。稀疏更新可以在不损害图像质量的情况下显著减少计算量。
此外,由于稀疏更新涉及聚集-分散过程,现有的深度学习框架会导致显著的延迟开销。为了解决这个问题,作者进一步提出了稀疏增量生成引擎(SIGE)来将算法的理论计算减少转化为在各种硬件上测量的延迟减少。
如图3所示是具体的算法原理,首先预计算原始输入图像的所有激活。在编辑过程中,通过计算原始图像和编辑图像之间的差异掩模来定位编辑区域。然后,对未编辑的区域重新使用预先计算的激活,并且通过对它们应用卷积滤波器来仅更新已编辑的区域。
其中Al表示第l层卷积层F的输入tensor,W和b分别是第l层的权重和偏置。
如图4所示,ΔAl共享了用户所做编辑中的结构化空间稀疏性,因此非零值主要聚集在编辑区域内。这样就可以直接使用原始图像和编辑后的图像来计算一个差异掩码,并用这个掩码对ΔAl进行稀疏化。
图4 左图:图像编辑示例。右图:在不同特征图分辨率下,DDIM第l层的∆Al通道平均值
3.3 稀疏增量生成引擎SIGE
一种简单的方法是为每个卷积从ΔAl中裁剪一个矩形编辑区域,并且只计算裁剪区域的特征。但作者发现这种裁剪方法对于不规则的编辑区域(图4所示的例子)效果很差。
因此,如图5所示,作者使用基于tiling的稀疏卷积算法。首先将差异掩码向下采样到不同的比例,并扩展向下采样的掩码,将ΔAl在空间上划分为多个相同大小的小块。每个块索引指的是具有非零元素的单个块。然后将非零块沿批维度进行相应的聚集,并将其馈入卷积Fl。最后,根据索引将输出块分散成零张量,以恢复原始空间大小,并将预先计算的残差计算。
内容中包含的图片若涉及版权问题,请及时与我们联系删除
评论
沙发等你来抢