Diffusion Models专栏文章汇总：入门与实战
前言：部署扩散模型面临着两个棘手的挑战：参数过大和推理时间过长，因此目前想在手机端用上扩散模型看似“奢不可求”。最近谷歌研究院的最新一项研究研究了如何把端侧部署大型扩散模型的梦想变成现实，这篇博客就和大家一起学习一下。
目录
贡献概述
方法详解
stable diffusion的组成
优化手段一：Group Norm and GELU
优化手段二：提高Attention Module的效率
优化手段三：Winograd Convolution
实验结果
论文和代码
论文地址
代码地址
个人思考
贡献概述
常见的大型扩散模型具有超过10亿个参数，并且由于设备上有限的计算和内存资源而带来挑战。作者提出了一系列针对大型扩散模型的实现优化，这些模型在配备gpu的移动设备上实现了最快推理延迟(对于512 × 512图像，在没有INT8量化的情况下，Stable diffusion 1.4在12秒内实现了20次迭代)。具体是通过设计专门的内核、提升Attention模型效率、用Winograd卷积加速三种方式实现的。
方法详解
stable diffusion的组成
（1）Text Embedder：利用CLIP对文本提示y进行编码，产生一个高维嵌入向量
，它封装了输入提示的语义，为反向扩散过程提供条件指导。这种条件扩散模型还不太清楚的同学可以看看我之前的博客：条件DDPM：Diffusion model的第三个巅峰之作_沉迷单车的追风少年的博客-****博客
（2）Noise Generation：在隐空间z中提供随机噪声，作为反向扩散过程的起始点。这个不需要什么运算资源，忽略。
（3）Denoising Neural Network：这是一个带有cross-attention结构的Unet，不熟悉的朋友们可以看：为什么Diffusion Models钟爱U-net结构？_diffusion网络结构_沉迷单车的追风少年的博客-****博客
（4）Image Decoder：stable diffusion源于latent diffusion，他们都是在latent域上操作的，不熟悉的朋友们可以参看：详细解读Latent Diffusion Models：原理和代码_沉迷单车的追风少年的博客-****博客
优化手段一：Group Norm and GELU
归一化手段优化：Group Norm用于diffusion models应该最早是《Diffusion Models Beat GANs on Image Synthesis》中提出的，这种归一化技术的工作原理是将特征图的通道划分为更小的组，并独立地归一化每个组，使GN更少地依赖于批量大小，更适合广泛的批量大小和网络架构。在工程实现上，作者设计了一个独特的内核，以GPU shared的形式，在单个GPU命令中执行所有这些操作，没有任何中间张量（降低访存）。
激活函数优化：之前diffusion models里面激活函数用的是SiLU，这篇论文作者改成了GELU-高斯误差线性单元(Gaussian Error Linear Unit, GELU)。我们实现了一个专用的shader来整合这些数值计算及其伴随的分割和乘法操作，使它们能够在一次single draw call中执行。
优化手段二：提高Attention Module的效率
stable diffusion models中的text/image transformer有助于对条件分布P进行建模(其实指的是text encoder和unet中都有attention结构)。然而，由于自/交叉注意力机制的二次时间和记忆复杂性，在长序列上遇到了困难。作者提出了Partially Fused Softmax和FlashAttention两种方法来解决这一挑战。
（1）Partially Fused Softmax
attention中有一次softmax和一次大矩阵运算：
softmax里面的
可分为两个步骤:1)归约操作;2)元素级操作。归约操作是指计算A中每一行的最大值及其修正后的指数和S，如公式4所示。随后，利用向量L和S对A中的值进行元素操作归一化。
为了避免在大矩阵A上执行整个softmax计算，作者用一个GPU shared 用于计算L和S向量的归约操作，从而得到大小为N × 2的张量（因此降低了矩阵的维度达到降低计算量的目的）。然后将元素级softmax计算与以下涉及矩阵V的矩阵乘法融合。这种方法大大减少了中间张量的内存占用和整体延迟。重要的是要强调从A到L的计算映射的并行性;S是有限的，因为结果张量中的元素数量远远小于输入张量A中的元素数量。
为了提高并行性并进一步降低延迟，通过将A中的元素分组为块，将归约操作划分为多个阶段。对每个块进行计算，然后简化为最终结果。通过精心设计的线程和内存缓存管理，该多阶段实现可以用一个GPU命令完成，并带来额外的延迟减少。
（2）FlashAttention
FlashAttention是一种IO感知的精确注意力算法，利用分块来减少GPU高带宽内存(HBM)和片上SRAM之间的内存读/写。这种方法导致比标准attention更少的HBM访问，使其适合于一系列SRAM大小，并提高整体效率。
详细的信息可以看一下原论文和解读：
- 科学家提出注意力算法FlashAttention，可执行更少的读写，且运行速度提高2-4倍 - 知乎
- /abs/2205.14135
摘要机翻速览：transformer在长序列上速度慢且内存消耗大，因为自注意力的时间和记忆复杂度在序列长度上是二次的。近似注意力方法试图通过权衡模型质量来降低计算复杂度来解决这个问题，但往往无法实现整体加速。本文认为，缺失的一个原则是使注意力算法具有io意识——考虑GPU内存级别之间的读写。本文提出FlashAttention，一种IO-aware的精确注意力算法，使用分块来减少GPU高带宽内存(HBM)和GPU片上SRAM之间的内存读写次数。分析了FlashAttention的IO复杂度，表明它比标准注意力需要更少的HBM访问，并且对于一系列SRAM大小是最优的。将FlashAttention扩展到块稀疏注意力，产生了一种近似注意力算法，比任何现有的近似注意力方法都快。FlashAttention训练transformer的速度比现有基线更快:BERT-large (seq。长度512)与MLPerf 1.1的训练速度记录相比，在GPT-2 (seq。长度1K)，在远程竞技场(seq.竞技场)上加速2.4倍。长度1 k-4k)。闪光注意力和块稀疏闪光注意力使transformer具有更长的上下文，产生更高质量的模型(GPT-2上的困惑度提高了0.7，长文档分类上的提升点提高了6.4)和全新的能力。长度16K，准确率61.4%)和路径256 (seq。长度64K，准确率63.1%)。
优化手段三：Winograd Convolution
Winograd卷积将卷积操作转换为一系列矩阵乘法。Winograd的关键在于，通过仔细选择变换矩阵，可以减少许多所需的乘法，从而提高计算效率。然而，它也会增加内存消耗和数值错误，特别是在使用较大的分块尺寸时。stable diffusion的主干在很大程度上依赖于3× 3卷积层，特别是在图像解码器中，它们包含了超过90%的层。
作者进行了分析，如表2所示，以探索在3 × 3内核卷积上使用不同分块大小的Winograd的潜在好处。根据发现，我们选择了4 × 4的分块大小，因为它在计算效率和内存利用率之间提供了最佳平衡。此外，我们基于启发式规则战略性地应用Winograd，仅在它能产生有利结果的地方，以进一步最大化其功效。
实验结果
作者是做了两款手机的结果，看上去每一种优化手段的效果都比较明显。其中三星S23 Ultra的延迟降低了52.2%，iPhone 14 Pro Max上的延迟降低了32.9%。在三星S23 Ultra上端到端从文本生成一张512×512像素的图片，迭代20次，耗时在12秒以内。
论文和代码
论文地址
/pdf/2304.
代码地址
暂无
个人思考
- 现在大部分端侧的推理部署还是以CPU为主(因为不同手机的GPU型号差异很大，很难做适配)，这里的优化手段主要针对GPU的。恐怕大规模落地(例如在CPU上取得一秒出图的效果)还有一段路要走。
- 作者从算法层面、工程层面都进行了优化。但是都算是比较“常规”手段吧，例如激活函数优化很早就被用在diffusion中了，attention相关的优化手段近年来层出不穷，卷积优化更是部署时候的标配。所以我个人觉得创新度上并没有太多新颖的东西，这篇论文目前的版本应该也很难中很好的会议。
- 不过这篇论文的意义在于作者分享了很多stable diffusion部署优化方面的工程性经验，避免了我们后人踩坑，非常值得一读。