[Paper Review] 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 Synthesis4 citations

TL;DR

MobileDiffusion introduces a highly efficient text-to-image diffusion model optimized for mobile deployment through architectural pruning and advanced sampling techniques, achieving sub-second inference (0.2 seconds on iPhone 15 Pro) for 512×512 image generation while maintaining high-quality outputs, setting a new state of the art in on-device generation.

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.

Motivation & Objective

• Address the challenge of deploying large-scale text-to-image diffusion models on mobile devices due to high model size and slow inference.
• Overcome the dual bottleneck of high parameter count and iterative sampling in diffusion models for edge deployment.
• Develop a holistic architectural efficiency framework for UNet-based diffusion models, focusing on component-level optimization beyond simple block removal.
• Enable practical on-device inference by combining architectural efficiency with state-of-the-art sampling acceleration techniques (distillation and GAN finetuning).
• Demonstrate the viability of lightweight, high-performance text-to-image models for real-world mobile applications, including controllable generation and fine-tuning.

Proposed method

• Conduct a comprehensive architectural analysis of the UNet architecture to identify and remove redundant components, reducing parameter count and improving computational efficiency.
• Design a lightweight UNet with fewer than 400 million parameters through systematic pruning and reorganization of blocks, informed by insights from attention mechanisms and feature map dynamics.
• Apply config-aware distillation to train a distilled 8-step inference model that preserves quality while reducing sampling steps.
• Utilize UFOGen finetuning to enable 1-step inference with minimal quality drop, leveraging GAN-based distillation for sample efficiency.
• Integrate the optimized architecture with existing sampling acceleration techniques, enabling end-to-end sub-second inference on mobile hardware.
• Benchmark performance on iPhone 15 Pro using native inference tools, measuring latency across text encoder, VAE decoder, UNet, and overall pipeline.

Experimental results

Research questions

• RQ1 Can a highly efficient text-to-image diffusion model be designed from the ground up for mobile deployment through architectural optimization?
• RQ2 To what extent can sampling speed be reduced without sacrificing image quality using distillation and GAN-based finetuning on a lightweight architecture?
• RQ3 How does the combination of architectural efficiency and sampling acceleration impact real-world inference latency on mobile devices?
• RQ4 Can the resulting model support downstream applications like controllable generation and LoRA fine-tuning effectively on-device?
• RQ5 What is the performance ceiling of on-device text-to-image generation in terms of speed, quality, and model size?

Key findings

• MobileDiffusion achieves sub-second inference time of 0.2 seconds for 512×512 image generation on iPhone 15 Pro, significantly outperforming prior work.
• The distilled 8-step model achieves a FID score of 9.01 on MS-COCO, comparable to the SD-replicated baseline (FID 8.43), despite being 50% smaller and 50% faster.
• The UFOGen-finetuned 1-step model achieves a FID of 11.67, demonstrating that single-step inference is viable with minimal quality degradation.
• The CLIP-ViT-g/14 score for both distilled and 1-step models is approximately 0.33, indicating strong alignment with CLIP embedding space.
• On-device latency measurements show that MobileDiffusion reduces overall inference time to 1.23 seconds (8 steps) and 2.19 seconds (1 step), outperforming prior methods like SnapFusion.
• The model successfully supports on-device fine-tuning for controllable generation tasks, including LoRA and adapter-based methods, demonstrating practical utility in real-world mobile applications.