速度惊人，手机跑 Stable Diffusion，12 秒出图，谷歌加速扩散模型破记录

手机生图只要 12 秒？

这不是吹的，谷歌就给实现了。

最新研究中，谷歌研究人员对 4 个 GPU 分层优化，在三星手机上成功运行 Stable Diffusion 1.4。

实现了 11.5 秒生成图像，重要的是，内存使用量也大幅减少。

正所谓，Speed Is All You Need！

谷歌最新提出的方法是通用的，可以对所有扩散模型的改进，并非仅针对特定设备实现的。

通过实验结果发现，三星 S23 Ultra 和 iPhone 14 Pro 的整体图像生成时间分别减少了 52% 和 33%。

这意味着，一部手机装下一个生成式 AI 模型的未来更近了。

从 3080 到一部手机

当前，将大型扩散模型合并到任何 App 中的一个关键考虑因素是，模型将在何处执行选择。

在一个消费级设备上部署模型的好处是，较低的服务成本、改善扩展性、还可以离线，并且能改善用户隐私。

22 年，Stable Diffusion 刚刚发布的第一个版本，最初只能缓慢运行在 RTX 3080 上。

Stable Diffusion 有超过 10 亿的参数，DALL-E 是 120 亿，以后随着扩散模型的发展，参数量会逐渐增加。

由于设备计算和内存资源的限制，因此在运行时带来了诸多挑战。

在没有精心设计情况下，在设备上运行这些模型可能会导致，由于迭代去噪过程和过多的内存消耗，输出内容延迟就会增加。

此前，也有研究实现了将 Stable Diffusion 成功部署到设备上，但仅局限于特定的设备或芯片组。

对此，谷歌研究人员对大型扩散模型提供了一系列实现优化，这些模型在配备 GPU 的移动设备上实现了迄今为止报道的最快推理延迟。

在不使用 INT8 量化的情况下，对于一张 512x512 的图片进行 20 次迭代，Stable Diffusion 1.4 的推理延迟低于 12 秒。

具体是如何实现的呢？

GPU 感知优化

在论文中，研究人员侧重的是使用大型扩散模型，然后完成从文本描述生成图像的任务。

虽说论文中，部分讨论是研究者为 Stable Diffusion 特定结构所提出的优化建议，但这些优化可以很容易推广到其它大型扩散模型上。

研究人员表示，当用文本提示进行推理时，这个过程包含根据所需的文本描述，应用额外条件来指导反向扩散。

具体来说，Stable Diffusion 的主要组成部分包括：文本嵌入器（Text Embedder）、噪声生成（Noise Generation）、去噪神经网络（Denoising Neural Network，aka UNet），以及图像解码器（Image Decoder）。

如下图所示：

Stable Diffusion 中主要组件及其相互作用的示意图

下面，我们分别介绍一下这几个组成部分，各部分间的关系参照图。

・文本嵌入器：

利用 CLIP 模型对文本提示 y 进行编码，生成一个高维嵌入向量 τθ(y)，将文本提示的语义封装进去。该嵌入被当作去噪神经网络的输入，为逆向扩散的过程提供指示。

・噪声生成：

给潜在空间提供随机噪声 z，该噪声作为逆向扩散过程的起始点。

・去噪神经网络：

该网络被设计为近似 p (z|y) 形式的条件分布，利用条件去噪自动编码器 θ(zt, t, τθ(y))（denoising autoencoder）。每次迭代 t 采用 UNet 架构。

同时，交叉注意机制（cross-attention mechanism）被用来操作潜在空间和文本嵌入向量，在迭代过程中预测 z 的去噪版本。

・图像解码器：

逆行扩散过程在潜在空间

中进行。一旦这个过程完成，图像解码器 D 被用来从潜在矢量中重建 RGB 图像。

研究人员在整个 UNet 架构中实现了群组归一化（Group normalization，GN）。

这种归一化技术的工作原理是将特征图（feature map）的 pipeline 划分为较小的组，并对每个组进行独立的归一化，使 GN 对批次大小的依赖性降低，更适合于各种大小的批次和各种网络结构。

应用公式①，每个特征值

被归一化为其所属组的组均值

和方差

。

研究人员并没有依次执行上述提到的重塑、平均值、方差和归一化的所有操作，而是以 GPU Shader 的形式设计了一个特别的程序，在一个 GPU 命令中执行所有这些操作，无需中间流程。

这里先介绍一下 Gaussian Error Linear Unit（GELU）。

GELU 作为模型中普遍存在的激活函数，包含许多数值计算，如乘法、加法和高斯误差函数，如公式②所示。

研究人员弄了一个专门的 Shader 来整合这些数值计算及其伴随的分割和乘法操作，使其在一次绘图调用中完成执行。

稳定扩散中的文本 / 图像变换器有助于对条件分布 P (z|τθ(y)) 进行建模，这对文本到图像的生成任务至关重要。

然而，自我 / 交叉注意力机制在处理长序列时遇到了困难，因为它们的时间和内存复杂性是平过方的。在论文中，研究人员介绍了两种可能的优化，旨在缓解这些计算瓶颈。

一种是 Partially Fused Softmax，另一种是 FlashAttention。

下面仅以 Softmax 为例。

上图是在注意力模块中，优化过的 softmax 实现。

虚线以上的流程图描述的是直接在矩阵

中应用 softmax 的初始实现。

虚线以下的展示的则是修正后的模块（红色部分）。

总而言之，论文中研究人员提出了一整套优化方案，可以在各种设备上执行大型扩散模型时，共同达到了突破性的延迟数字。

这些改进扩大了模型的通用性，并提高了在各种设备上的整体用户体验。

12 秒，业界领先

为了评估改进后的模型，研究人员分别在三星 S23 Ultra (Adreno 740) 和 iPhone 14 Pro Max (A16) 进行了一组基准测试。

作为去噪神经网络，UNet 是计算需求最高的组件。

研究人员提供了执行单次迭代的 UNet 所需的延迟数据，以毫秒为单位测量，图像分辨率为 512x512。

此外，他们记录了运行时生成的中间张量在「Tensor」列中的内存使用情况，以及为保存模型权重分配的内存在「Weight」列中的使用情况，均以兆字节为单位。

请注意，内存管理器通过重用中间张量的缓冲区来优化内存占用。

如表中数据显示，第一行显示了在公共 Github 仓库中使用内部 OpenCL 内核实现，但没有任何优化的结果。

实现之后的结果，并且研究者在没有任何优化的情况下使用内部 OpenCL 内核。

第 2-5 行，分别逐个启用每个优化：

Opt. Softmax：部分融合的 softmax 和优化的 softmax 减少步骤

S-GN / GELU：用于组归一化和 GELU 的专用内核

FlashAttn.：FlashAttention 实现

Winograd (All)：采用 Winograd 卷积

随着每个优化的启用，实验结果发现延迟逐步减少。

与基线相比，在两种设备上都观察到了显著的总体延迟降低：三星 S23 Ultra 降低 52.2%，iPhone 14 Pro Max 降低 32.9%。

此外，研究人员还评估了在三星 S23 Ultra 进行文本到图像输出端到端延迟。

进行了 20 次去噪迭代，生成一张 512x512 图像，实现了不到 12 秒的业界领先结果。

可见，在没有数据连接或云服务器的情况下，在手机上本地运行生成式人工智能模型，将开辟了许多可能性。

谷歌最新研究给出了一种全新方案。

参考资料：

• https://arxiv.org/abs/2304.11267

• https://www.reddit.com/r/MachineLearning/comments/12zclus/d_google_researchers_achieve_performance/

本文来自微信公众号：新智元 （ID：AI_era）

广告声明：文内含有的对外跳转链接（包括不限于超链接、二维码、口令等形式），用于传递更多信息，节省甄选时间，结果仅供参考，IT之家所有文章均包含本声明。