Date: January 31, 2024

Developers have been yearning for the ability to run powerful generative AI directly on a phone. Until now, the sheer size of state‑of‑the‑art diffusion models—often billions of parameters—has forced us to rely on cloud servers. The new MobileDiffusion research from Google AI demonstrates that a well‑engineered latent diffusion model can produce a 512×512 image in under 0.5 seconds on modern iOS and Android devices. This post breaks down the key architectural decisions, the one‑step sampling trick, and practical takeaways you can apply to your own on‑device ML pipelines.

---

Why Text‑to‑Image Diffusion Was Historically Mobile‑Unfriendly

Diffusion models generate images by repeatedly denoising a noisy latent. Each denoising step requires a full forward pass through a large UNet, and the number of steps (often 50‑100) multiplies the computational load. Two problems arise for mobile hardware:

1. Iterative denoising – every step invokes the same heavy network, inflating latency.
2. Massive model footprints – classic pipelines such as Stable Diffusion hover around 1‑2 B parameters, exceeding the memory and compute budget of a smartphone.

Past work has mainly tackled the first issue with clever solvers or distillation, reducing the number of function evaluations (NFEs) to single digits. However, even a handful of steps remains costly when the network architecture itself is not optimized for mobile constraints.

---

MobileDiffusion Overview

MobileDiffusion adopts the latent diffusion paradigm but re‑imagines every component for on‑device efficiency:

Component	Choice	Params
Text encoder	CLIP‑ViT/L14	125 M
Diffusion UNet	Custom lightweight UNet (see below)	~380 M
Image decoder (VAE)	Pruned decoder	9‑40 M

The total model size is ≈ 520 M parameters, a fraction of typical desktop‑grade models, yet it still delivers high‑fidelity results.

---

Re‑Designing the Diffusion UNet

1. Transformer placement

Traditional diffusion UNets interleave self‑attention (SA), cross‑attention (CA), and feed‑forward blocks throughout the spatial hierarchy. Attention scales quadratically with the sequence length, making it expensive at high resolutions. MobileDiffusion follows the UViT insight: push most transformer blocks to the bottleneck where the spatial dimension is smallest. This reduces the overall quadratic cost while preserving the model’s ability to fuse textual cues.

2. Convolutional efficiency

Residual blocks remain at every level for feature extraction, but regular convolutions are replaced with depthwise‑separable convolutions in deeper stages. The lightweight kernels keep FLOPs low without sacrificing perceptual quality.

3. FLOP/parameter comparison

Benchmarks against other diffusion UNets show MobileDiffusion achieves the lowest FLOP count for a comparable parameter budget, directly translating to faster inference on mobile CPUs/NPUs.

---

Slimming the Image Decoder

The VAE decoder converts the 8‑channel latent back to an RGB image. Google’s team pruned both width and depth of the original Stable Diffusion decoder, resulting in three variants:

• Full decoder – 39.3 M params, PSNR 30.0 dB, SSIM 0.83
• Lite decoder – 9.8 M params, PSNR 30.2 dB, SSIM 0.84 (better quality despite being ~¼ the size)

Latency improves by roughly 50 %, and visual fidelity surpasses the baseline Stable Diffusion decoder.

---

One‑Step Sampling with DiffusionGAN

Even with an efficient UNet, iterating 50‑100 times is still too slow for sub‑second experiences. MobileDiffusion integrates a DiffusionGAN hybrid that learns to perform the entire denoising process in a single forward pass.

Training flow:

1. Initialize both generator (the diffusion UNet) and discriminator with weights from a pre‑trained diffusion model.
2. Feed a noisy latent to the generator; compute a reconstruction loss against the ground‑truth image.
3. Add noise to the generator’s output and present it to the discriminator; apply a GAN loss to enforce realistic texture and semantics.
4. Fine‑tune for < 10 K iterations.

The result is a model that needs only one function evaluation per image, dramatically cutting latency while preserving the richness of diffusion‑based synthesis.

---

Real‑World Performance Numbers

Device	Runtime Optimizer	Latency (seconds)
iPhone 15 Pro (Apple Neural Engine)	Core ML GPU	0.46
Pixel 8 Pro (Google Tensor)	MediaPipe GPU	0.48
Samsung Galaxy S24 Ultra (Exynos)	NNAPI	0.52

All measurements are for generating a 512×512 image, confirming that the sub‑second target is comfortably met on flagship hardware.

---

How to Bring MobileDiffusion into Your App

1. Export the model – The authors provide a TensorFlow Lite / Core ML checkpoint. Convert it to the appropriate on‑device format using the TensorFlow Lite Converter or Apple’s Core ML tools.
2. Quantize – Post‑training dynamic range quantization reduces size by ~30 % with negligible quality loss. Full integer quantization can be explored if your target hardware supports INT8 acceleration.
3. Integrate the tokenizer – MobileDiffusion relies on CLIP‑ViT text embeddings. Bundle the tokenizer (e.g., SentencePiece) and run it on‑device before feeding the embedding to the UNet.
4. Cache the VAE encoder/decoder – Reuse latent representations across frames if you build an interactive drawing UI.
5. Profile – Use Android Studio Profiler or Xcode Instruments to ensure the model stays within the desired latency envelope on your target devices.

---

Key Takeaways

• Sub‑second on‑device text‑to‑image generation is feasible with a 520 M‑parameter latent diffusion model.
• Placing most attention layers at the UNet bottleneck and swapping regular convolutions for separable ones yields major FLOP reductions.
• A lightweight VAE decoder (as small as 9.8 M params) can outperform the original Stable Diffusion decoder on PSNR, SSIM, and LPIPS.
• DiffusionGAN enables one‑step sampling, turning a multi‑step diffusion process into a single forward pass.
• Real‑world latency on flagship iOS/Android devices stays under 0.5 s for 512×512 output.

---

Conclusion

MobileDiffusion proves that with thoughtful architectural pruning and a GAN‑augmented sampler, high‑quality text‑to‑image generation can move from the cloud to the palm of your hand. As mobile hardware continues to improve, we can expect even larger generative models to run locally, opening doors for privacy‑preserving AI, offline creative tools, and instant visual feedback in every app.

---

Source: MobileDiffusion: Rapid text-to-image generation on-device