[논문 리뷰] SnapFusion: Text-to-Image Diffusion Model on Mobile Devices within Two Seconds

Yanyu Li, Huan Wang외 7명|arXiv (Cornell University)|2023. 06. 01.

Generative Adversarial Networks and Image Synthesis인용 수 35

한 줄 요약

SnapFusion은 UNet 아키텍처, 이미지 디코더, 그리고 스텝 증류를 최적화하여 모바일 기기에서 2초 이내에 실행되는 텍스트-투-이미지 확산 모델을 제공합니다. SD-v1.5와 비교해 훨씬 적은 스텝으로도 경쟁력 있는 품질을 달성합니다.

ABSTRACT

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.

연구 동기 및 목표

• 온-디바이스 확산 모델의 병목을 식별하고 모바일 하드웨어에서 지연 원인을 정량화합니다.
• 이미지 품질을 희생하지 않으면서 계산량을 줄이기 위한 아키텍처-진화 UNet를 개발합니다.
• 메모리 및 계산 수요를 줄이기 위해 이미지 디코더를 압축하고 증류합니다.
• 더 적은 스텝으로 품질을 유지하기 위해 classifier-free guidance 정규화를 활용한 스텝 증류를 발전시킵니다.

제안 방법

• Text Encoder, UNet, VAE Decoder의 지연 병목을 찾아 Stable Diffusion v1.5를 분석합니다.
• 블록 수준 순열을 견딜 수 있도록 강인한 학습을 갖춘 아키텍처-진화 UNet를 제안하고 중복을 제거합니다.
• 합성 프롬프트를 포함한 데이터/증류 파이프라인을 사용하여 이미지 디코더를 압축하고 증류합니다.
• 추론 단계를 50에서 8로 줄이되 품질을 유지하기 위해 스텝 증류를 적용합니다.
• FID와 CLIP의 균형을 맞추기 위해 CFG-인식 스텝 증류를 CFG-가이드 손실 및 손실 혼합 스킴과 함께 도입합니다.
• 동적 감마를 통해 증류 목표를 조화시키면서 CFG-인식 증류와 원래 손실을 공동으로 사용합니다.

Figure 1: Example generated images by using our efficient text-to-image diffusion model.

실험 결과

연구 질문

• RQ1 UNet 아키텍처 중복을 온-디바이스 확산 속도 향상에 어떻게 활용할 수 있으며 품질 저하 없이 가능할까요?
• RQ2 모바일 확산을 위한 견고한 아키텍처 진화를 가능하게 하는 학습 전략은 무엇일까요?
• RQ3 압축된 이미지 디코더가 매개변수 및 MAC를 줄이면서 지각 품질을 유지할 수 있을까요?
• RQ4 스텝 증류를 통한 denoising 스텝 감소가 모바일 디바이스의 FID와 CLIP 점수에 어떤 영향을 미칠까요?
• RQ5 CFG-인식 스텝 증류가 저스텝 영역에서 CLIP 점수를 향상시키면서 FID를 보존할 수 있을까요?

주요 결과

• 8-step on-device UNet with a distilled image decoder achieves faster-than-2-second generation on mobile hardware with competitive SD-v1.5-like quality.
• Architecture evolution with robust training preserves pre-trained performance while allowing block-level pruning/removal for speed.
• Efficient image decoder achieves 3.8× fewer parameters and 3.2× speedup compared to SD-v1.5, via 50% channel pruning.
• CFG-aware step distillation improves CLIP scores for low-step models while maintaining reasonable FID, outperforming vanilla distillation in CLIP at similar FID.
• Direct 16→8-step distillation beats progressive distillation in both FID and CLIP under the same inference budget.
• On MS-COCO 2017 5K, the 8-step model achieves 24.2 FID and 0.30 CLIP, outperforming several baselines.

Figure 2: Latency (iPhone 14 Pro, ms) and parameter (M) analysis for cross-attention (CA) and ResNet blocks in the UNet of Stable Diffusion.