Segmind가 Knowledge Distillation 기법으로 Stable Diffusion 모델을 35~55% 경량화하면서 추론 속도를 최대 100% 향상

Context

대규모 생성 AI 모델(SD 2.1, SDXL 등)의 성능은 뛰어나지만 파라미터 수가 많아 배포 비용과 추론 지연이 높다는 문제가 있었다. 이러한 모델을 실무 환경에서 활용하려면 모델 크기 축소와 추론 속도 개선이 필수적이었다.

Technical Solution

• Block-removal Knowledge Distillation 적용: 교사 모델(Realistic-Vision 4.0)로부터 학생 모델이 UNet 블록 출력을 모방하도록 3가지 손실함수(latent loss, teacher-student output loss, feature-level loss) 결합
• SD-Small 모델 개발: 기존 모델 대비 파라미터 35% 감소, LAION Art Aesthetic 데이터셋 1M개 이미지로 100K 스텝 학습
• SD-Tiny 모델 개발: 기존 모델 대비 파라미터 55% 감소, 동일 데이터셋으로 125K 스텝 학습
• 🧨 diffusers 라이브러리 기반 구현: DiffusionPipeline을 통한 표준화된 모델 사용 인터페이스 제공
• 파인튜닝 및 LoRA 지원: 경량화 모델에 대한 portrait 데이터셋 파인튜닝(7k 이미지, 131K 스텝) 및 LoRA 학습 예제 제공

Impact

• 추론 레이턴시: SD-Small과 SD-Tiny가 기존 모델 대비 최대 100% 향상
• 모델 파라미터: SD-Small 35% 감소, SD-Tiny 55% 감소
• 파라미터 수 대비 이미지 품질 유지: 경량화 모델도 기존 모델과 비슷한 수준의 이미지 충실도 달성
• LoRA 학습 속도: 경량화 모델에서 더 빠른 학습 속도 실현

Key Takeaway

Knowledge Distillation의 핵심은 블록 레벨 출력 매칭을 통해 정보 손실을 최소화하면서 모델을 축소하는 것이다. 경량화된 모델도 파인튜닝과 LoRA를 통해 특정 도메인(portrait, 추상 개념)에서 원본 모델 수준의 품질을 달성할 수 있음을 보여준다.