[논문 리뷰] BK-SDM: A Lightweight, Fast, and Cheap Version of Stable Diffusion

Bo-Kyeong Kim, Hyoung‐Kyu Song외 2명|arXiv (Cornell University)|2023. 05. 25.

Generative Adversarial Networks and Image Synthesis인용 수 5

한 줄 요약

이 논문은 U-Net 아키텍처에서 잔차 및 어텐션 블록을 프루닝한 후, 단지 13일치 A100 연산 자원과 소규모 데이터셋을 사용해 특징 수준의 지식 정련 재학습을 수행함으로써 모델 크기를 최대 51% 감소시키고 지연 시간을 43% 향상시킨 경량형, 빠르고 비용 효율적인 Stable Diffusion의 변종인 BK-SDM을 소개한다. 이는 경쟁적인 제로샷 텍스트-이미지 생성 성능을 달성하며, 4초 미만의 추론 시간을 확보해 엣지 장치에 효율적으로 구현 가능하다.

ABSTRACT

Text-to-image (T2I) generation with Stable Diffusion models (SDMs) involves high computing demands due to billion-scale parameters. To enhance efficiency, recent studies have reduced sampling steps and applied network quantization while retaining the original architectures. The lack of architectural reduction attempts may stem from worries over expensive retraining for such massive models. In this work, we uncover the surprising potential of block pruning and feature distillation for low-cost general-purpose T2I. By removing several residual and attention blocks from the U-Net of SDMs, we achieve 30%~50% reduction in model size, MACs, and latency. We show that distillation retraining is effective even under limited resources: using only 13 A100 days and a tiny dataset, our compact models can imitate the original SDMs (v1.4 and v2.1-base with over 6,000 A100 days). Benefiting from the transferred knowledge, our BK-SDMs deliver competitive results on zero-shot MS-COCO against larger multi-billion parameter models. We further demonstrate the applicability of our lightweight backbones in personalized generation and image-to-image translation. Deployment of our models on edge devices attains 4-second inference. Code and models can be found at: https://github.com/Nota-NetsPresso/BK-SDM

연구 동기 및 목표

• 대규모 텍스트-이미지 확산 모델(예: Stable Diffusion)의 계산 및 훈련 비용을 성능 저하 없이 감소시키는 것.
• 비용이 많이 들기 때문에 거의 탐색되지 않은 대규모 확산 모델에서 잔차 및 어텐션 블록 프루닝을 통한 아키텍처 압축을 탐색하는 것.
• 제한된 데이터와 계산 자원으로도 대규모 사전 훈련된 모델의 지식을 작고 프루닝된 변종으로 효과적으로 전달할 수 있음을 입증하는 것.
• Jetson AGX Orin 및 iPhone 14와 같은 엣지 장치에 하위 4초 이내 추론 속도로 효율적인 작동이 가능한 작고 강력한 확산 모델의 구현을 가능하게 하는 것.
• 연구자들과 개발자들이 현실적인 자원 예산으로 소형이지만 강력한 확산 모델을 훈련하고 배포할 수 있는 실용적이고 접근 가능한 길을 제공하는 것.

제안 방법

• Stable Diffusion의 U-Net 아키텍처에 블록 프루닝을 적용하여 모델 크기와 MACs를 줄이기 위해 다수의 잔차 및 어텐션 블록을 제거한다.
• 성능 유지 목적으로 특징 수준의 지식 정련을 사용해 프루닝된 모델을 재학습한다. 이는 원본 대규모 모델의 지식을 작고 압축된 모델로 전달한다.
• 정련 과정은 0.22M개의 LAION 쌍 데이터셋과 단지 13일치 A100 연산 자원을 사용하여 재학습 비용을 크게 감소시킨다.
• 이 방법은 SDM-v1.4 및 SDM-v2.1-base 모두에 적용되었으며, 다양한 모델 변종에서 일관된 성능 향상을 보였다.
• 프루닝 및 정련된 모델은 제로샷 텍스트-이미지 생성, 개인화된 생성, 이미지-이미지 번역, 엣지 배포 등에 대해 평가되었다.
• 얼굴 생성을 위한 조건부 LDM(celebA-HQ)에 대해 검증되었으며, 이는 텍스트-이미지 작업을 넘어서 일반화 가능성을 확인한다.

실험 결과

연구 질문

• RQ1 Stable Diffusion의 U-Net에서 잔차 및 어텐션 블록을 아키텍처 프루닝함으로써 성능 저하 없이 모델 크기와 추론 지연 시간을 크게 줄일 수 있는가?
• RQ2 제한된 데이터와 계산 자원으로도 대규모 사전 훈련된 확산 모델의 능력을 작고 프루닝된 변종으로 효과적으로 전달할 수 있는가?
• RQ3 경량 확산 모델이 다수의 억억 파라미터를 가진 모델들과 비교해 경쟁적인 제로샷 텍스트-이미지 생성 성능을 달성할 수 있는 정도는 어느 정도인가?
• RQ4 프루닝 및 정련된 모델은 더 적은 피니팅 비용으로도 개인화된 텍스트-이미지 생성 및 이미지-이미지 번역에 효과적으로 사용될 수 있는가?
• RQ5 이러한 작고 강력한 모델을 실시간 추론 속도로 엣지 장치에 배포하는 것이 현실적으로 가능한가?

주요 결과

• BK-SDM는 CPU 및 GPU에서 최대 51%의 모델 크기 감소와 43%의 지연 시간 향상을 달성했으며, MACs와 추론 시간 모두 30–50% 감소하였다.
• 정련된 모델은 단지 13일치 A100 연산 자원과 0.22M개의 LAION 쌍 데이터셋을 사용해 MS-COCO 30K에서 FID 14.61을 달성했으며, 훨씬 더 많은 자원을 사용해 훈련된 더 큰 모델들과 견줄 만한 성능을 보였다.
• BK-SDM-Small는 Jetson AGX Orin과 iPhone 14에서 모두 512×512 이미지를 4초 미만으로 생성하여 엣지 배포 가능성에 강력한 증거를 제시했다.
• DreamBooth를 통한 개인화된 생성에서 BK-SDM은 원본 SDM의 주제 및 프롬프트 충실도를 95–99% 유지하면서도 피니팅 비용을 줄였다.
• 이 방법은 다른 확산 모델로도 일반화 가능하며, 187M 파라미터를 가진 조건부 LDM을 성공적으로 압축하여 고품질의 얼굴 생성을 달성했다.
• 파라미터와 훈련 데이터가 감소했음에도 불구하고 BK-SDM 모델들은 경쟁적인 CLIP 및 IS 점수를 유지하여 강력한 생성 품질을 보였다.