PaperlayerView exact revision on arXiv

Source-linked AI summary

SnapFusion: Text-to-Image Diffusion Model on Mobile Devices within Two Seconds

Yanyu Li, Huan Wang, Qing Jin, Ju Hu, Pavlo Chemerys, Yun Fu, Yanzhi Wang, Sergey Tulyakov, Jian Ren

arXiv:2306.00980v3cs.CVcs.AIcs.LG

TL;DR

Text-to-image diffusion models are too computationally expensive for fast, private mobile inference. SnapFusion combines an efficient UNet and image-decoder distillation with improved CFG-aware step distillation, achieving sub-two-second generation while matching or exceeding Stable Diffusion v1.5 quality.

  • Problem

    Text-to-image diffusion models require substantial computation, making mobile inference slow and motivating costly cloud execution with privacy concerns.

  • Method

    SnapFusion identifies UNet redundancy, evolves an efficient architecture, distills the image decoder, and improves step distillation with training strategies and CFG-aware regularization.

  • Results

    1.84 seconds: SnapFusion runs on-device with image quality similar to or better than Stable Diffusion while being significantly faster.

  • Takeaways & Limitations

    SnapFusion brings large-scale text-to-image diffusion to mobile devices for faster, more private content generation.

  • Takeaways & Limitations

    Latency analysis is conducted mostly on the more powerful iPhone 14 Pro, and the model still has relatively many parameters.

Abstract

from arXiv · show

Text-to-image diffusion models can create stunning images from natural language descriptions that rival the work of professional artists and photographers. However, these models are large, with complex network architectures and tens of denoising iterations, making them computationally expensive and slow to run. As a result, high-end GPUs and cloud-based inference are required to run diffusion models at scale. This is costly and has privacy implications, especially when user data is sent to a third party. To overcome these challenges, we present a generic approach that, for the first time, unlocks running text-to-image diffusion models on mobile devices in less than $2$ seconds. We achieve so by introducing efficient network architecture and improving step distillation. Specifically, we propose an efficient UNet by identifying the redundancy of the original model and reducing the computation of the image decoder via data distillation. Further, we enhance the step distillation by exploring training strategies and introducing regularization from classifier-free guidance. Our extensive experiments on MS-COCO show that our model with $8$ denoising steps achieves better FID and CLIP scores than Stable Diffusion v$1.5$ with $50$ steps. Our work democratizes content creation by bringing powerful text-to-image diffusion models to the hands of users.

1 Introduction

SnapFusion addresses the cost, latency, and privacy barriers that limit text-to-image diffusion on mobile devices. It combines efficient architecture design with improved step distillation to generate images in under two seconds while maintaining or improving quality.

  • Motivation: Cloud-based inference requires high-end GPUs, incurs high costs, and raises privacy concerns when private data is sent to third parties.
  • Motivation: Existing mobile diffusion methods reach 11.5s on a Samsung Galaxy S23 Ultra, which does not support a seamless user experience.Prior studies also lacked systematic quantitative analysis of on-device generation quality.
  • Approach: SnapFusion targets two bottlenecks: the slow UNet and the large number of denoising steps needed for generation.The paper identifies UNet architecture redundancy and the quality impact of directly reducing denoising steps.
  • Approach: The proposed system uses an evolving training framework for an efficient UNet, data distillation for the image decoder, and CFG-aware step-distillation regularization.The training strategy also explores teacher-student paradigms for the on-device model.
  • Outcome: Less than 2 seconds and image quality similar to Stable Diffusion v1.5 are achieved for 512 × 512 text-to-image generation on mobile devices.The result is attributed to improved step distillation and network architecture development.

2 Model Analysis of Stable Diffusion

The analysis explains Stable Diffusion’s latent-space denoising pipeline and locates its mobile-inference bottlenecks in the UNet and image decoder. It further distinguishes parameter concentration in downsampled stages from latency concentration in high-resolution stages.

  • Stable Diffusion Pipeline: DDIM repeatedly feeds the denoised latent at a previous time step back into the diffusion model until the process reaches t′ = 0.
  • Stable Diffusion Pipeline: Latent diffusion denoises in latent space, decodes the resulting latent into an image with a VAE, and uses text embeddings as conditioning.Stable Diffusion v1.5 is used as the baseline, with classifier-free guidance controlling conditional strength and the quality-diversity trade-off.
  • Latency Bottlenecks: The text encoder contributes 8 ms, the VAE decoder 369 ms, and the denoising UNet 1.7 seconds per inference analysis.The UNet is both computationally intensive and repeatedly executed during generation.
  • UNet Structure: The time-conditional UNet combines cross-attention, which integrates text embeddings, with ResNet blocks, which capture locality.Cross-attention uses noisy-data queries and text-condition keys and values.
  • UNet Bottlenecks: UNet parameters concentrate in downsampled middle stages, whereas latency is highest in high-resolution input and output stages because spatial cross-attention scales quadratically with token count.This micro-level breakdown guides architecture redesign and search.

3 Architecture Optimizations

SnapFusion redesigns the diffusion architecture by evolving a robust, efficient UNet and distilling a smaller image decoder. The approach preserves generative performance while reducing on-device computation.

  • Robust Training: Stochastic block skipping during robust training makes the UNet tolerant to architecture permutations.Each cross-attention and ResNet block can execute normally or be replaced by an identity mapping.
  • Evaluation and Architecture Evolving: The robustly trained network enables stable architectural evolution while preserving the performance of the pre-trained UNet.The method supports accurate block assessment during online network changes.
  • Efficient UNet: Architecture evolution removes redundant UNet blocks or adds valuable ones while evaluating latency and CLIP impact.Actions target cross-attention and ResNet blocks at specific stages and positions.
  • Efficient Image Decoder: The efficient image decoder has 3.8× fewer parameters and is 3.2× faster than the Stable Diffusion v1.5 decoder.It results from 50% uniform channel pruning and has approximately one-quarter the original size and MACs.

4 Step Distillation

SnapFusion combines step distillation with classifier-free-guidance-aware training to reduce denoising iterations while addressing quality degradation. Its pipeline progresses from Stable Diffusion teachers to an 8-step efficient UNet.

  • Distillation Overview: Step distillation reduces inference iterations by transferring a teacher’s denoising behavior to a student with fewer steps.For example, distilling from 32 steps to 16 steps provides 2× speedup against the teacher.
  • Vanilla Step Distillation: The model is fine-tuned for v-prediction before distillation, with velocity defined from the clean latent, noise, and time-dependent coefficients.The target velocity is v ≡ α_tϵ − σ_tx.
  • Distillation Pipeline: The final pipeline directly distills a 32-step Stable Diffusion model to 16 steps, creates a 16-step efficient UNet, then distills it to 8 steps.Progressive distillation was empirically slightly worse than direct distillation.
  • Vanilla Step Distillation: The teacher performs two DDIM denoising steps, while the student predicts the teacher’s resulting noisy latent in one step.The student objective is formulated in x-space using the teacher’s two-step trajectory.
  • CFG-Aware Step Distillation: CFG-aware distillation applies classifier-free guidance to teacher and student outputs before calculating the distillation loss.This remedy targets the CLIP-score degradation observed with vanilla step distillation.
  • CFG-Aware Step Distillation: Randomly sampling the CFG scale from the default range [2, 14] provides a training-time trade-off between FID and CLIP score.The CFG scale w is sampled uniformly during training.
  • Total Loss Function: Mixing vanilla and CFG distillation losses combines their observed benefits: lower FID from vanilla loss and higher CLIP score from CFG loss.A predefined probability p selects CFG loss per training iteration, while the original denoising loss is also weighted by γ.
  • Comparison with Prior Work: The proposed CFG-aware scheme differs from w-condition distillation by targeting image quality without adding CFG-scale input parameters.It also enables training-time diversity-quality trade-offs for models favoring quality or diversity.

5 Experiment

Experiments show that SnapFusion combines robust architecture evolution with CFG-aware step distillation to achieve strong image quality at substantially lower inference cost. On MS-COCO, its 8-step model matches or outperforms Stable Diffusion baselines across several settings.

  • Text-to-Image Generation: 0.004−0.010 higher CLIP score is achieved by the 8-step, 230ms-per-step UNet than SD-v1.5 at the same FID level on 30K MS-COCO samples.With eight DDIM steps, SnapFusion also outperforms another distillation method by 2.7 FID with on-par CLIP score.
  • Text-to-Image Generation: 0.002−0.007 higher CLIP score is achieved by CFG-aware distillation than SD-v1.5 under similar FID, using an 8-step student distilled from a CFG-distilled 16-step teacher.The CFG-aware objective consistently improves the 8-step model’s CLIP score across tested configurations.
  • Text-to-Image Generation: SnapFusion’s efficient UNet uses robust training to tolerate architecture permutations and block ablations while preserving pretrained performance.This supports evaluating and evolving candidate architectures without fine-tuning each altered network to recover performance.
  • Ablation Analysis: Direct distillation to 8 steps outperforms progressive distillation on both FID and CLIP while requiring only one teacher transition.The adopted direct path is 16 → 8, compared with progressive 64 → 32 → 16 → 8.
  • Ablation Analysis: The vanilla distillation loss achieves the lowest FID, whereas CFG distillation achieves the highest CLIP score.Mixing the losses is motivated by combining these complementary metric strengths.
  • Ablation Analysis: Using the original denoising loss lowers FID without reducing CLIP score, while a 16-step teacher provides the best CLIP score among tested teacher settings.More teacher steps produced worse student CLIP scores despite similar lowest FID values.
  • Ablation Analysis: The CFG-aware 8-step distilled model achieves performance comparable to the 50-step SD-v2 model when applied to SD-v2.The same hyperparameters used for SD-v1.5 were reused, and further tuning might improve results.

6 Related Work

Related work improves diffusion efficiency through faster sampling and on-device optimization, but prior mobile methods remain too slow for seamless use. Step distillation reduces denoising steps, while guided-distillation methods address classifier-free guidance with different trade-offs.

  • Efficiency Directions: Diffusion models require high-end GPUs because large networks and tens or hundreds of iterative denoising steps make inference computationally expensive.Efficiency research therefore focuses primarily on improving sampling and enabling on-device solutions.
  • Step Distillation: Progressive distillation gradually reduces sampling steps from a pretrained teacher, while w-condition reduces classifier-free-guidance evaluation cost by 2× through extra student parameters.SnapFusion instead uses step distillation with different objectives aimed at improving image quality during distillation.

7 Discussion and Conclusion

SnapFusion combines efficient architecture and improved step distillation to achieve fast on-device generation with image quality comparable to Stable Diffusion. The paper notes remaining model-size and device-coverage limitations, alongside broader-impact concerns about misuse.

  • 1.84 seconds: SnapFusion performs denoising on-device with image quality on par with Stable Diffusion.
  • The approach analyzes UNet redundancies, develops an evolving-training framework, and adds CFG-aware distillation loss.
  • The model still has a relatively large number of parameters, limiting compatibility with some edge devices.
  • Latency analysis focuses mainly on iPhone 14 Pro, which has more computational power than many other phones.
  • The authors caution that content-generation models should be applied carefully to avoid malicious applications.

A Efficient UNet

The efficient UNet is designed from a computation and parameter analysis of Stable Diffusion, while the broader pipeline benchmarks latency across phones and a server GPU. The text encoder remains uncompressed because its inference latency is negligible relative to the UNet and VAE decoder.

  • For 512 × 512 images, latent-space UNet processing uses an input and output resolution of 64 × 64.
  • Latency is benchmarked on iPhone 12 Pro Max, iPhone 13 Pro Max, and iPhone 14 Pro.
  • Over 12× speedup and 46% lower running memory are achieved versus SD-v1.5 on an Nvidia A100 40G GPU.
  • The A100 latency analysis uses TensorRT in single precision (FP32).
  • The released pipeline does not compress the text encoder because its 4ms inference latency is negligible compared with the UNet and VAE decoder.

B.2 VAE Decoder

The compressed VAE decoder is trained by distilling teacher-generated latent-image pairs, reducing decoder computation while preserving generation quality across common CFG scales. It is evaluated against the original decoder and SD-v1.5 on MS-COCO.

  • The compressed decoder uses teacher-generated latent-image pairs and mean squared error to match the teacher's output images.The student receives the teacher model's latent as input and is optimized against the teacher image.
  • 116ms: The compressed decoder provides comparable generation quality while reducing decoder latency.
  • The distillation pipeline needs no paired text-image samples and can generate unlimited data on the fly.
  • Across CFG scales from 7 to 9, the compressed decoder performs similarly to the original decoder and better than SD-v1.5.The comparison uses 6K samples from the MS-COCO 2014 validation set.

C Detailed Derivations of Step Distillation

The derivation contrasts two teacher DDIM denoising steps with one student step. It defines the intermediate teacher latents and the student prediction used to form the distillation objective with truncated-SNR weighting.

  • The teacher UNet takes timestep t, noisy latent z_t, and text embedding c, then performs DDIM steps through t′ and t′′.
  • The first teacher step predicts v̂_t, clean latent x̂_t, noise ε̂_t, and intermediate noisy latent z_t′.
  • The second teacher step repeats the process from t′ to t′′ to obtain z_t′′.
  • The student UNet performs one DDIM denoising step and is trained to predict the teacher's noisy latent after the skipped steps.
  • The resulting loss objective uses truncated SNR weighting coefficients defined as ϖ(λ_t) = max(...).

D Different Teacher Options for Step Distillation

The authors evaluate several teacher models for distilling the final 8-step efficient UNet, finding both self-distillation and SD-v1.5 distillation effective.

  • Self-distillation from the 16-step efficient UNet effectively boosts the final 8-step model.
  • Distillation from the 16-step Stable Diffusion v1.5 baseline also effectively boosts the final 8-step model.
  • CFG-aware distillation from a stronger 16-step teacher is investigated as a possible further performance improvement.

E Additional Qualitative Results

Additional examples show SnapFusion synthesizing diverse content with high aesthetics, including realistic objects, scenery, artistic styles, and cartoons.

  • SnapFusion synthesizes diverse content with high aesthetics, including realistic food and animals, scenery, artistic styles, and cartoon styles.
  • Figure 9 presents additional generated images from the efficient text-to-image diffusion model.
Loading 2306.00980v3…