谷歌研究员成功实现手机12秒一键生图

谷歌研究员成功实现手机12秒一键生图

谷歌研究人员成功运行StableDiffusion1.4，在三星手机上实现了11.5秒生成图像，内存使用量也大幅减少。论文中提出的方法通用，可以对所有扩散模型的改进。通过实验结果发现，三星S23Ultra和iPhone14Pro的整体图像生成时间分别减少了52%和33%。一部手机装下一个生成式AI模型的未来更近了。

从3080到一部手机当前，将大型扩散模型合并到任何APP中的一个关键考虑因素是，模型将在何处执行选择。在一个消费级设备上部署模型的好处是，较低的服务成本、改善扩展性、还可以离线，并且能改善用户隐私。22年，StableDiffusion刚刚发布的第一个版本，最初只能缓慢运行在RTX3080上。StableDiffusion有超过10亿的参数，DALL-E是120亿，以后随着扩散模型的发展，参数量会逐渐增加。由于设备计算和内存资源的限制，因此在运行时带来了诸多挑战。

研究人员对大型扩散模型提供了一系列实现优化，这些模型在配备GPU的移动设备上实现了迄今为止报道的最快推理延迟。在不使用INT8量化的情况下，对于一张512x512的图片进行20次迭代，StableDiffusion1.4的推理延迟低于12秒。具体是如何实现的呢？

论文中，研究人员侧重的是使用大型扩散模型，然后完成从文本描述生成图像的任务。虽说论文中，部分讨论是研究者为StableDiffusion特定结构所提出的优化建议，但这些优化可以很容易推广到其它大型扩散模型上。研究人员表示，当用文本提示进行推理时，这个过程包含根据所需的文本描述，应用额外条件来指导反向扩散。具体来说，StableDiffusion的主要组成部分包括：文本嵌入器（TextEmbedder）、噪声生成（NoiseGeneration）、去噪神经网络（DenoisingNeuralNetwork，akaUNet），以及图像解码器（ImageDecoder）。

StableDiffusion中主要组件及其相互作用的示意图。StableDiffusion中主要组件及其相互作用的示意图。StableDiffusion中主要组件及其相互作用的示意图。StableDiffusion中主要组件及其相互作用的示意图。

文本嵌入器：利用CLIP模型对文本提示y进行编码，生成一个高维嵌入向量τθ(y)，将文本提示的语义封装进去。该嵌入被当作去噪神经网络的输入，为逆向扩散的过程提供指示。噪声生成：给潜在空间提供随机噪声z，该噪声作为逆向扩散过程的起始点。去噪神经网络：该网络被设计为近似p(z|y)形式的条件分布，利用条件去噪自动编码器θ(zt,t,τθ(y))（denoisingautoencoder）。每次迭代t采用UNet架构。同时，交叉注意机制（cross-attentionmechanism）被用来操作潜在空间和文本嵌入向量，在迭代过程中预测z的去噪版本。图像解码器：逆行扩散过程在潜在空间中进行。一旦这个过程完成，图像解码器D被用来从潜在矢量中重建RGB图像。

研究人员在整个UNet架构中实现了群组归一化（Groupnormalization，GN）。这种归一化技术的工作原理是将特征图（featuremap）的pipeline划分为较小的组，并对每个组进行独立的归一化，使GN对批次大小的依赖性降低，更适合于各种大小的批次和各种网络结构。研究人员并没有依次执行上述提到的重塑、平均值、方差和归一化的所有操作，而是以GPUShader的形式设计了一个特别的程序，在一个GPU命令中执行所有这些操作，无需中间流程。

在论文中，研究人员介绍了两种可能的优化，旨在缓解这些计算瓶颈。一种是PartiallyFusedSoftmax，另一种是FlashAttention。下面仅以Softmax为例。研究人员提出了一整套优化方案，可以在各种设备上执行大型扩散模型时，共同达到了突破性的延迟数字。这些改进扩大了模型的通用性，并提高了在各种设备上的整体用户体验。

为了评估改进后的模型，研究人员分别在三星S23Ultra(Adreno740)和iPhone14ProMax(A16)进行了一组基准测试。作为去噪神经网络，UNet是计算需求最高的组件。研究人员提供了执行单次迭代的UNet所需的延迟数据，以毫秒为单位测量，图像分辨率为512x512。此外，他们记录了运行时生成的中间张量在「Tensor」列中的内存使用情况，以及为保存模型权重分配的内存在「Weight」列中的使用情况，均以兆字节为单位。请注意，内存管理器通过重用中间张量的缓冲区来优化内存占用。如表中数据显示，第一行显示了在公共Github仓库中使用内部OpenCL内核实现，但没有任何优化的结果。实现之后的结果，并且研究者在没有任何优化的情况下使用内部OpenCL内核。第2-5行，分别逐个启用每个优化：Opt.Softmax：部分融合的softmax和优化的softmax减少步骤S-GN/GELU：用于组归一化和GELU的专用内核FlashAttn.：FlashAttention实现Winograd(All)：采用Winograd卷积随着每个优化的启用，实验结果发现延迟逐步减少。与基线相比，在两种设备上都观察到了显著的总体延迟降低：三星S23Ultra降低52.2%，iPhone14ProMax降低32.9%。

此外，研究人员还评估了在三星S23Ultra进行文本到图像输出端到端延迟。进行了20次去噪迭代，生成一张512x512图像，实现了不到12秒的业界领先结果。可见，在没有数据连接或云服务器的情况下，在手机上本地运行生成式人工智能模型，将开辟了许多可能性。谷歌最新研究给出了一种全新方案。