12秒内AI在手机上完成作画！谷歌提出扩散模型推理加速新方法

明敏源于奥菲寺。

量子比特|微信官方账号QbitAI

只需要12秒，仅仅依靠手机本身的计算能力，就可以生成一个扩散稳定的图像。

而且是已经完成20次迭代的那种。

要知道，现在的扩散模型基本都超过10亿个参数。要快速生成一张图片，要么基于云计算，要么本地硬件足够强大。

随着大规模模型应用的普及，在个人电脑和手机上运行大规模模型很可能是未来的新趋势。

因此，谷歌的研究人员带来了这项新成果，名为Speed is all you need:通过GPU优化，加快设备上大规模扩散模型的推理速度。

三步优化加速该方法针对稳定扩散进行优化，但也可以适应其他扩散模型。面向的任务是从文本生成图像。

具体优化可以分为三个部分：

设计一个专门的内核来提高注意力模型Winograd卷积加速的效率首先看专门设计的内核，它包含了分组归一化和GELU激活函数。

在整个UNet架构中实现了分组规范化。这种归一化的工作原理是将特征映射的通道分成更小的组，并对每个组进行独立的归一化，这样组归一化对批量的依赖性更小，可以适应更大范围的批量和网络架构。

研究人员以GPU shader的形式设计了一个独特的内核，它可以在单个GPU命令中执行所有内核，而无需任何中间张量。

GELU激活函数包含了大量的数值计算，比如罚函数，高斯误差函数等等。

这些数值计算以及伴随的除法和乘法运算由一个特殊的着色器集成，因此这些计算可以放在一个简单的绘制调用中。

Draw call是CPU调用图像编程接口，命令GPU渲染的操作。

接下来，为了提高注意力模型的效率，本文介绍了两种优化方法。

一种是softmax函数的部分融合。

为了避免整个softmax在大矩阵A上的计算，设计了一个GPU着色器来计算L和S向量以减少计算量，最终得到一个大小为N2的张量。然后结合softmax的计算和矩阵V的矩阵乘法。

这种方法大大降低了中间程序的内存占用张量和整体延迟。

需要强调的是，从A到L和S的计算映射的并行性是有限的，因为结果张量中的元素数量远小于输入张量A中的元素数量.

为了增加并行性并进一步减少延迟，将A中的元素分组到一个块中，并将归约操作分成几个部分。

然后在每个block上执行计算，然后将其简化为最终结果。

利用精心设计的线程和内存缓存管理，可以在多个部分实现使用单个GPU命令降低延迟。

另一种优化方法是FlashAttention。

这是去年火起来的IO感知精确注意力算法，具体有两种加速技术：按块递增计算即平铺、并在后向传递中重新计算注意力，将所有注意力操作融合到CUDA内核中。

相较于标准Attention，这种方法能减少HBM（高带宽内存）访问，提高整体效率。

不过FlashAttention内核的缓存器密集度非常高（register-intensive），所以该团队是有选择性地使用这一优化方法。

他们在注意力矩阵d=40的Adreno GPU和Apple GPU上使用FlashAttention，其他情况下使用部分融合softmax函数。

第三部分是Winograd卷积加速。

它的原理简单来说就是使用更多的加法计算来减少乘法计算，从而降低计算量。

但弊端也很明显，将会带来更多的显存消耗和数值错误，尤其是在tile比较大的情况时。

Stable Diffusion的主干非常依赖33卷积层，尤其是在图像解码器方面，这里90%的层都是由33卷积层构成的。

研究人员分析后发现，在使用44大小的tile时，是模型计算效率和显存利用率的最佳平衡点。

实验结果为了评估提升效果，研究人员先在手机上进行了基准测试。

结果表明，两部手机在使用了加速算法后，生成图片的速度都明显提升。

其中三星S23 Ultra的延迟降低了52.2%，iPhone 14 Pro Max上的延迟降低了32.9%。

在三星S23 Ultra上端到端从文本生成一张512512像素的图片，迭代20次，耗时在12秒以内。

论文地址：https://arxiv.org/abs/2304.11267— 完—量子位QbitAI 头条号签约关注我们，第一时间获知前沿科技动态