[论文解读] MobileDiffusion: Instant Text-to-Image Generation on Mobile Devices

Yang Zhao, Yanwu Xu+3 more|arXiv (Cornell University)|Nov 28, 2023

Generative Adversarial Networks and Image Synthesis被引用 4

一句话总结

MobileDiffusion 通过架构剪枝和先进的采样技术，引入了一种高度高效的文本到图像扩散模型，专为移动端部署而优化，在 iPhone 15 Pro 上实现了 512×512 图像生成的亚秒级推理（仅 0.2 秒），同时保持了高质量输出，开创了设备端生成的新 SOTA。

ABSTRACT

The deployment of large-scale text-to-image diffusion models on mobile devices is impeded by their substantial model size and slow inference speed. In this paper, we propose extbf{MobileDiffusion}, a highly efficient text-to-image diffusion model obtained through extensive optimizations in both architecture and sampling techniques. We conduct a comprehensive examination of model architecture design to reduce redundancy, enhance computational efficiency, and minimize model's parameter count, while preserving image generation quality. Additionally, we employ distillation and diffusion-GAN finetuning techniques on MobileDiffusion to achieve 8-step and 1-step inference respectively. Empirical studies, conducted both quantitatively and qualitatively, demonstrate the effectiveness of our proposed techniques. MobileDiffusion achieves a remarkable extbf{sub-second} inference speed for generating a $512 imes512$ image on mobile devices, establishing a new state of the art.

研究动机与目标

• 解决大规模文本到图像扩散模型在移动端部署时因模型过大和推理缓慢带来的挑战。
• 突破扩散模型在边缘设备部署中高参数量与迭代采样双重瓶颈。
• 为基于 UNet 的扩散模型设计一个整体性的架构效率框架，聚焦于组件级优化，而非简单的模块移除。
• 通过结合架构效率与最先进采样加速技术（如蒸馏与 GAN 微调），实现实用的设备端推理。
• 展示轻量化、高性能文本到图像模型在真实移动端应用中的可行性，包括可控生成与微调能力。

提出的方法

• 对 UNet 架构进行全面的架构分析，识别并移除冗余组件，降低参数量并提升计算效率。
• 通过注意力机制与特征图动态的洞察，系统性地剪枝与重组模块，设计出参数少于 4 亿的轻量化 UNet。
• 应用配置感知蒸馏，训练一个 8 步推理的蒸馏模型，在保持质量的同时减少采样步数。
• 利用 UFOGen 微调实现 1 步推理，仅造成轻微质量损失，借助基于 GAN 的蒸馏实现样本高效。
• 将优化后的架构与现有采样加速技术整合，实现在移动端硬件上的端到端亚秒级推理。
• 使用原生推理工具在 iPhone 15 Pro 上进行性能基准测试，测量文本编码器、VAE 解码器、UNet 及整体流水线的延迟。

实验结果

研究问题

• RQ1 能否通过架构优化从零开始设计一个高度高效的移动端部署文本到图像扩散模型？
• RQ2 在轻量化架构上结合蒸馏与基于 GAN 的微调，采样速度可降低到何种程度而不牺牲图像质量？
• RQ3 架构效率与采样加速技术的结合，如何影响移动端真实场景下的推理延迟？
• RQ4 该模型能否在设备端有效支持下游应用，如可控生成与 LoRA 微调？
• RQ5 在速度、质量与模型大小方面，设备端文本到图像生成的性能上限是什么？

主要发现

• MobileDiffusion 在 iPhone 15 Pro 上实现 512×512 图像生成的亚秒级推理（0.2 秒），显著优于先前工作。
• 蒸馏后的 8 步模型在 MS-COCO 上取得 9.01 的 FID 分数，与 SD 复现基线（FID 8.43）相当，尽管模型更小 50% 且速度快 50%。
• UFOGen 微调后的 1 步模型取得 11.67 的 FID 分数，证明单步推理在质量损失极小的情况下具备可行性。
• 两种模型在 CLIP-ViT-g/14 上的得分均约为 0.33，表明与 CLIP 嵌入空间具有强对齐性。
• 设备端延迟测量显示，MobileDiffusion 将整体推理时间减少至 1.23 秒（8 步）与 2.19 秒（1 步），优于先前方法如 SnapFusion。
• 该模型成功支持设备端微调，可用于可控生成任务，包括 LoRA 与适配器方法，展现出在真实移动端应用中的实际可用性。