MobileDiffusion：移动设备上亚秒级文本到图像生成

 人工智能咨询培训老师叶梓 转载标明出处

文本到图像扩散模型在生成高质量图像方面具有卓越的能力，这些模型是多种应用的基础，包括图像编辑、控制生成、个性化内容生成、视频合成和低级视觉任务等。然而，这些大规模模型通常需要在具有强大神经计算单元的服务器上运行，在移动设备上部署大规模文本到图像扩散模型受到模型尺寸大和推理速度慢的限制。

为了克服这些限制，谷歌的研究者团队提出了MobileDiffusion，这是一种经过全面架构和采样技术优化的高效率文本到图像扩散模型。研究者们对模型架构进行了深入的审查，旨在减少冗余、提高计算效率，并在显著减少模型参数数量的同时，保持甚至提升图像生成的质量。他们还采用了蒸馏技术和扩散-GAN微调方法，分别实现了8步和1步的快速推理过程。

在本文中，研究者们还介绍了UFOGen，一种创新的单步生成模型。UFOGen通过巧妙结合对抗性训练和重建项，进一步提升了生成过程的效率，并确保了图像质量。这两种技术的融合不仅展示了在理论上的先进性，更在实际应用中证明了其强大的潜力和广泛的适用性。

MobileDiffusion 文本到图像生成的两个示例，分别是 (a) 8步生成（使用 cfg-aware distillation [23]），和 (b) 1步生成（使用 UFOGen finetuning [57]）

Mobile Diffusion

扩散模型的生成过程可以视为一个由噪声逐渐增加的过程，随后通过学习逆过程来去除噪声，恢复出清晰的图像。这个过程通常涉及以下几个步骤：

正向过程：从数据分布中采样一个干净的数据点，然后逐步加入高斯噪声，生成一系列越来越嘈杂的样本。

学习逆过程：通过训练一个神经网络，使其能够预测每一步中噪声的分布，并在逆向过程中逐步去除噪声，恢复出原始数据。

训练方法：研究者采用了特定的损失函数来训练去噪网络，使其能够准确预测每一步中的噪声，并进行有效去除。

MobileDiffusion (MD) 的模型示意图，包括卷积 (Conv.)、自注意力 (SA)、交叉注意力 (CA) 等组件

MobileDiffusion模型采用了UNet结构，这是一种常见的卷积神经网络架构，它通过对称的下采样和上采样过程来捕获图像的上下文信息和细节。在MobileDiffusion中，UNet结构经过了特别的优化，以适应移动设备上的资源限制。

下采样阶段：

Down-1和Down-2：在下采样的初始阶段，模型使用标准的卷积层（Conv.），这些层负责提取图像的初步特征。这些层可能包含深度可分离卷积（DW），以减少计算量。Down-3：在下采样的第三层，模型引入了具有自注意力（Self-attention, SA）和交叉注意力（Cross-attention, CA）的Transformer块。自注意力是可选的，而交叉注意力用于将文本特征与图像特征结合起来。这里的注意力机制通过非线性函数ϕ来计算注意力权重。

中间瓶颈阶段：

在UNet的瓶颈部分，模型进一步压缩特征图，以便于在网络的下一部分中高效地处理和上采样。瓶颈层可能包含更多的Transformer块，这些块具有深度可分离卷积和交叉注意力，以捕获更细致的特征并保持文本条件的一致性。

上采样阶段：

Up-3、Up-2和Up-1：在上采样阶段，模型逐步恢复图像的空间分辨率。这一过程涉及到使用上采样操作和卷积层来细化特征图。在Up-3中，模型再次使用具有自注意力和交叉注意力的Transformer块，但这次自注意力层可能被省略，以减少计算负担。在Up-2和Up-1中，模型继续使用卷积层和交叉注意力层来进一步细化图像的细节和结构。

门控线性层：

在Transformer块中，模型使用了门控线性层，其中包含了扩展因子e。这种层可以调节信息流，增强模型的表示能力。

非线性激活：

在整个网络中，模型使用了非线性激活函数，如ReLU或Swish，来引入非线性并增强模型的学习能力。

端到端优化：

整个MobileDiffusion模型被设计为端到端可训练，意味着从文本输入到图像输出的整个过程可以在一个统一的框架中进行优化。

与其它知名潜在扩散模型的比较，包括参数数量、卷积块数量、自注意力和交叉注意力层的组合、参数量、浮点运算次数、延迟和模型大小

Transformer架构是扩散模型中用于处理长距离依赖关系的关键组件。Transformer的优化主要集中在以下几个方面：

自注意力机制的改进：研究者探索了减少自注意力层计算复杂度的方法，例如通过降低自注意力层的分辨率或减少自注意力头的数量。

参数共享策略：为了减少模型的参数数量，研究者采用了参数共享技术，例如在自注意力层中共享键（Key）和值（Value）的投影矩阵。

计算效率：研究者对Transformer的各个组成部分进行了细致的分析，以确定哪些部分对于保持图像质量至关重要，哪些部分可以进行简化或优化以提高计算效率。

研究者采用了深度可分离卷积技术，这种方法通过分步骤进行空间卷积和深度卷积，有效降低了模型的参数数量和计算复杂度。研究者还实施了卷积块的剪枝策略，通过消除网络中不必要的部分来简化模型结构，这不仅减少了计算负担，还提升了模型的运行速度，同时保持了图像生成的质量。

变分自编码器（VAE）是MobileDiffusion模型中用于图像压缩和重建的关键组件。研究者首先对VAE的编码器和解码器进行了改进。他们采用了修改版的VAE，其中编码器将图像的空间分辨率降低8倍，并将3通道的RGB图像编码为具有更多通道的潜在变量。这样的改进有助于提升图像重建的质量。同时，为了提高效率，研究者设计了一种轻量级的解码器架构，通过减少解码器的宽度和深度来降低模型的复杂性。

在训练策略上，研究者采用了包括L2重建损失、KL散度正则化、感知损失和对抗性损失在内的组合损失函数。他们使用了一个批量大小为256的数据集，并进行了200万次迭代的训练。为了进一步提升效率，研究者训练了一个轻量级解码器，通过在训练过程中冻结编码器，并减少对抗性损失的权重，成功地提高了模型的解码速度。

通过这些优化，研究者显著减少了VAE模型的参数数量，并在解码512×512图像时，实现了近50%的延迟改进。定量比较显示，轻量级解码器在COCO2017验证集上的质量与原始解码器相当。

UFOGen设计空间

UFOGen，即Universal Forward One-step Generation，是一种创新的文本到图像生成模型，它的核心优势在于能够通过单步过程实现图像的快速生成。在资源受限的环境中，如移动设备或边缘计算设备上，部署大型生成模型面临着重大挑战。为了克服这些限制，UFOGen采用了一种创新的单步生成策略，该策略能够在一次前向传播中快速生成图像，大大减少了与传统扩散模型相关的多步骤迭代过程。

在UFOGen的架构中，对抗性训练发挥了核心作用。通过引入一个判别器网络，模型被训练为生成高质量、逼真的图像，这些图像足以在视觉上欺骗判别器。这种对抗性训练不仅提升了图像的视觉质量，也增强了模型对文本描述的语义理解，确保生成的图像能够准确地反映文本中所蕴含的信息。

UFOGen模型中的重建项确保了生成过程的准确性和数据一致性。重建项作为训练目标的一部分，衡量了模型输出与原始数据之间的差异，引导模型学习如何生成与目标数据分布相匹配的图像。这种方法有助于提高图像的细节准确性和整体质量。

UFOGen的设计强调了在减少计算步骤的同时保持图像质量的能力。通过优化模型的架构和训练过程，研究者们展示了UFOGen在生成图像时的高效率和高保真度。这种平衡对于在计算资源受限的情况下实现快速且高质量的文本到图像生成至关重要。

UFOGen的架构还考虑了模型的灵活性和适应性。通过精心设计的训练策略，UFOGen能够适应不同的文本提示和图像风格，满足多样化的生成需求。这种灵活性使得UFOGen不仅在理论上具有吸引力，也在实际应用中展现出了广泛的应用潜力。

在UFOGen模型的训练过程中，重建项扮演着至关重要的角色，确保生成的图像与文本描述紧密对应。研究者通过精心设计的损失函数，对模型预测的图像和目标图像之间的差异进行量化，进而引导模型学习如何更准确地反映文本中蕴含的视觉信息。这种损失函数通常包括像素级的误差和结构相似性指标，促使生成图像在视觉上与原始数据分布保持一致。

为了进一步提升模型性能，研究者采用了对抗性微调技术，这种方法通过引入一个判别器网络，为生成网络提供了额外的挑战。判别器的任务是区分真实图像与生成图像，而生成网络则努力产生足以欺骗判别器的图像。这种对抗性训练不仅增强了图像的真实感，还提高了模型对细节的捕捉能力，从而在图像质量和多样性上取得了显著提升。

在实施对抗性微调时，研究者特别注重参数效率，通过选择性地更新模型参数，避免了对整个网络的大规模调整。这种策略允许模型在保持原有结构的同时，对特定任务进行快速适应。例如，通过微调网络的特定层或使用低秩矩阵调整，模型能够更加专注于学习与任务相关的特征，而不是进行全面的重新训练。

为了避免过拟合并保持模型的泛化能力，研究者采取了一系列措施。这包括在训练过程中使用dropout、引入数据增强技术，以及应用正则化方法来限制模型复杂度。这些技术的应用有助于确保模型在面对新的、未见过的数据时，仍能保持稳定的性能。

实验

训练过程使用了专有的数据集，这个数据集包含了从公共网络收集的1.5亿个图像-文本对。大多数图像的分辨率超过256×256像素，其中4000万像素图像的分辨率超过512×512像素。为了训练，研究者遵循了一致的预处理方法，包括裁剪图像以匹配所需的分辨率。

模型采用AdamW优化器进行训练，配置了学习率为0.001、β1为0.9、β2为0.999和权重衰减为0.01。训练采用了渐进式策略：首先在256×256分辨率下进行训练，批量大小为4096，训练75万步以捕获高级语义，然后在512×512分辨率下继续训练25万步，批量大小为2048，以细化特征。

在搜索和选择优秀候选模型时，研究者依赖于训练步骤30K时报告的FID和CLIP分数来决定继续进行的候选模型。这大约需要使用32个TPU（每个具有16GB内存）计算4小时。研究者设计了一个停止机制，当FID和CLIP分数与之前的候选相比显示出较慢的改进迹象时，实验将停止。这个策略在过程中效果合理。完成网络搜索大约消耗了512个TPU，历时15天。

研究者使用MS-COCO数据集作为评估的主要来源。他们展示了零样本FID-30K场景的结果，其中从COCO验证集中随机选择了3万个标题。这些标题作为图像合成过程的输入。研究者计算FID分数来衡量生成样本与3万个参考真实图像之间的差异。

不同模型在零样本MS-COCO 2014验证集（30K）上的性能

 在表2中，研究者对不同的文本到图像生成模型进行了全面比较。对于DDIM（Denoising Diffusion Implicit Models），通过调整cfg尺度至大约3.0，达到了最低的FID（Fréchet Inception Distance）得分。由于UFOGen模型没有cfg，因此直接报告了FID得分。研究者还比较了与他们模型相近的其他模型的CLIP得分。MobileDiffusion模型在保持模型尺寸紧凑的同时，实现了比SnapFusion和UFOGen等先前解决方案更好的指标。

MobileDiffusion模型与其他几个模型在人类偏好评分（Human Preference Score, HPS）上的表现比较

不同模型在iPhone上的延迟测量结果，包括文本编码器、解码器、UNet步骤以及整体推理步骤的延迟

研究者在iPhone 15 Pro上使用特定工具对他们提出的MD（MobileDiffusion）进行了基准测试。如表4所示，MD在多个方面展现出了卓越的效率，包括文本编码器、VAE解码器、UNet每步推理，以及最终的整体延迟。具体来说，MD在单步推理时的整体延迟显著低于Stable Diffusion 1.5和SnapFusion，这表明MobileDiffusion非常适合在资源受限的移动设备上运行。

对抗性微调设置的不同变体的消融研究结果

研究者在表5中报告了消融研究的FID得分。包括没有使用LoRA时的得分，以及使用扩散损失（Diffusion loss）和蒸馏损失（Distill loss）的得分。例如，不使用LoRA时的FID得分是11.67，而使用LoRA时的FID得分是13.45。使用蒸馏损失时，不使用LoRA的FID得分是11.20，而使用LoRA时是13.24。FID指标显示使用蒸馏损失有轻微的性能优势，原始的扩散损失在应用中表现最佳。

应用

模型在图像编辑、风格转换和个性化内容生成等方面的能力，显示了MobileDiffusion模型在实际使用中的灵活性和有效性。

MobileDiffusion 在控制文本到图像生成中的应用示例，包括边缘到图像的翻译和 LoRA 微调

如图4所示，该模型在可控生成、个性化生成和修复等任务上始终能够即时生成视觉上吸引人的结果。条件生成模型是在预训练的扩散模型上训练的，并且它们能够无缝适应单步扩散模型而无需重新训练。这种灵活性使得MobileDiffusion模型能够兼容更多潜在的应用场景。

具体应用示例包括：

1. Canny边缘：使用Canny边缘检测作为条件，模型能够根据边缘信息生成完整的图像。
2. 深度图：利用深度信息作为条件，模型可以生成具有正确深度感知的图像。
3. 风格LoRA：通过微调LoRA（Low-Rank Adaptation）技术，模型能够适应不同的风格并生成具有特定艺术风格的图像。
4. 对象LoRA：使用LoRA技术对模型进行微调，以生成包含特定对象或特征的图像。
5. 修复（Inpainting）：在图像中指定区域进行修复，模型能够智能填充缺失部分，生成连贯且视觉上合理的图像。

这些技术的结合不仅为移动设备上的图像生成提供了强大的工具，也为未来的研究和应用开辟了新的可能性。随着人工智能技术的不断进步，MobileDiffusion模型有望在更广泛的领域中发挥其潜力，为用户带来更加丰富和个性化的视觉体验。

论文链接：https://arxiv.org/abs/2311.16567

​