Segmind SD-Small 및 SD-Tiny 지식 증류 출시

TL;DR

Segmind가 블록 제거 지식 증류(block-removal knowledge distillation)를 통해 압축된 확산 모델인 SD-Small 및 SD-Tiny의 학습 코드와 사전 학습된 체크포인트를 공개했습니다. 이 모델들은 파라미터 수를 35%–55% 줄이면서도 베이스 모델과 유사한 이미지 충실도를 유지하며, 최대 100% 빠른 추론 속도를 달성합니다.

지식 증류 방법론

공개된 모델들은 On Architectural Compression of Text-to-Image Diffusion Models (Shinkook et al.)에 기술된 Block-Removal Knowledge Distillation 기술을 사용하여 학습되었습니다.

  • Teacher model: 고품질 Stable Diffusion 체크포인트인 Realistic-Vision 4.0.
  • Student architectures: 레이어가 제거된 UNet 변형 모델로, 파라미터 수가 각각 35%(SD-Small) 및 55%(SD-Tiny) 적습니다.
  • Loss composition:
    1. 타겟 이미지와 생성된 이미지의 잠재 표현(latent representations) 사이의 표준 확산 손실(diffusion loss).
    2. 학생 모델이 생성한 잠재 표현을 교사 모델의 잠재 표현과 일치시키는 잠재 수준 손실(Latent-level loss).
    3. 교사 모델과 학생 모델의 각 UNet 블록 출력을 매칭하는 특징 수준 손실(Feature-level loss) (가장 핵심적인 구성 요소).
  • Training data: 이미지 점수 > 7.5로 필터링된 LAION Art Aesthetic 데이터셋.
  • Training schedule: 1M 이미지를 사용하여 SD-Small은 100k 스텝, SD-Tiny는 125k 스텝 동안 학습.

전체 증류 파이프라인은 [segmind/distill-sd](https://github.com/segmind/distill-sd) 저장소에서 확인할 수 있으며, 사전 학습된 체크포인트는 Hugging Face의 segmind 네임스페이스에 호스팅되어 있습니다.

🤗 Diffusers를 이용한 모델 사용법

두 모델 모두 🤗 Diffusers 라이브러리의 DiffusionPipeline을 사용하여 직접 로드할 수 있습니다:

from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained(
    "segmind/small-sd", torch_dtype=torch.float16
)
prompt = "Portrait of a pretty girl"
negative_prompt = (
    "(deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, "
    "cartoon, drawing, anime:1.4), text, close up, cropped, out of frame, "
    "worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, "
    "mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn "
    "face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, "
    "extra limbs, cloned face, disfigured, gross proportions, malformed limbs, "
    "missing arms, missing legs, extra arms, extra legs, fused fingers, too many "
    "fingers, long neck"
)
image = pipeline(prompt, negative_prompt=negative_prompt).images[0]
image.save("my_image.png")

모델 식별자를 segmind/tiny-sd로 바꾸면 동일한 API가 적용됩니다.

추론 속도 향상

동일한 하드웨어에서의 벤치마킹 결과, 증류된 모델은 기존 베이스 체크포인트에 비해 **최대 2배 낮은 지연 시간(latency)**을 보여줍니다. 이러한 측정에 사용된 추론 스크립트는 저장소(inference.py)에 포함되어 있습니다.

알려진 제한 사항

  • 이 모델들은 초기 단계의 출시 버전입니다. 시각적 품질이 아직 상용 수준의 확산 모델과 일치하지 않을 수 있습니다.
  • 범용 생성에 최적화되어 있지 않으며, 복잡한 구도 프롬프트나 여러 개념을 처리하는 데 어려움을 겪을 수 있습니다.
  • 더 높은 충실도를 달성하려면 특정 도메인 데이터에 대해 미세 조정(fine-tune)하거나 LoRA를 적용하는 것이 권장됩니다.

초상화 데이터셋을 이용한 SD-Tiny 미세 조정

Segmind는 Realistic-Vision 4.0으로 생성된 7k 이미지 초상화 세트를 사용하여 SD-Tiny를 미세 조정한 사례를 선보였습니다. 학습 하이퍼파라미터:

  • Steps: 131,000
  • Learning rate: 1e-4
  • Batch size: 32 (gradient accumulation = 4)
  • Image resolution: 768 px
  • Mixed-precision: fp16

결과물은 55%의 파라미터 감소를 유지하면서도 원래의 교사 모델의 품질에 근접합니다.

증류된 모델에 대한 LoRA 학습

SD-Tiny에 저차원 적응(LoRA)을 적용하면 모델 크기가 줄어들었기 때문에 더 빠른 LoRA 수렴이 가능합니다. 추상적 개념에 대해 학습된 LoRA 체크포인트 예시는 저장소(lora_training.py)에 제공됩니다.

커뮤니티 초대

Segmind는 개발자들이 프로젝트에 기여하고, 이슈를 보고하며, 미세 조정된 체크포인트를 공유할 것을 권장합니다. 소통 채널로는 Discord 서버와 GitHub 저장소가 있으며, Star와 Pull Request를 환영합니다.

Sources