Stable Diffusion 3 Medium과 Diffusers 통합
Hugging Face는 Stability AI의 2B 파라미터 모델인 Stable Diffusion 3 Medium을 diffusers 라이브러리에 통합했습니다. 이번 릴리스에서는 Multimodal Diffusion Transformer (MMDiT) 라는 새로운 아키텍처와 rectified flow‑matching 학습 목표를 활용해 고품질 텍스트‑투‑이미지 및 이미지‑투‑이미지 합성을 수행할 수 있게 되었습니다.
Stable Diffusion 3의 아키텍처 혁신
Stable Diffusion 3 (SD3)는 텍스트와 이미지 데이터 간에 양방향 흐름을 이용함으로써 기존 텍스트‑투‑이미지 아키텍처와 차별화됩니다.
Multimodal Diffusion Transformer (MMDiT)
SD3는 MMDiT 모델을 새롭게 도입하고, 세 가지 텍스트 인코더인 CLIP L/14, OpenCLIP bigG/14, T5‑v1.1‑XXL을 함께 사용합니다. 또한 Stable Diffusion XL에 사용된 것과 유사한 16채널 AutoEncoder도 활용합니다.
이전 모델들이 고정된 텍스트 표현을 사용해 cross‑attention을 적용했던 것과 달리, MMDiT 블록은 텍스트 입력과 픽셀 라텐트를 임베딩 시퀀스로 처리합니다. 각각의 시퀀스는 별도의 가중치를 통해 동일 차원으로 임베딩된 뒤 연결(concatenate)되어, modulated attention과 MLP를 거칩니다. 이를 통해 두 모달리티가 attention 연산 중에 서로 영향을 주고받을 수 있습니다.
Rectified Flow Matching
SD3는 조건부 flow‑matching 목표를 사용해 학습됩니다. 이 접근법은 데이터와 노이즈 분포를 직선으로 연결하는 "rectified flow"를 전방 노이징 과정으로 정의합니다.
이를 지원하기 위해 Hugging Face는 FlowMatchEulerDiscreteScheduler를 도입했습니다. 이 스케줄러는 Euler 방법 단계들을 구현하고, 해상도에 따라 타임스텝 스케줄을 이동시키는 shift 파라미터를 포함합니다. 2B 모델의 경우, 높은 해상도에서 노이즈 스케일링을 더 잘 처리하기 위해 shift=3.0을 권장합니다.
메모리 최적화 기법
SD3는 대형 T5‑XXL 텍스트 인코더(4.7B 파라미터)를 사용하기 때문에, VRAM이 24GB 미만인 GPU에서는 실행이 어려울 수 있습니다. Hugging Face는 메모리 사용량을 줄이기 위한 여러 최적화 전략을 제공합니다:
- Model Offloading:
pipe.enable_model_cpu_offload()를 사용하면 사용되지 않을 때 모델 구성 요소를 CPU로 옮겨 VRAM 사용량을 감소시키지만, 지연 시간이 증가합니다. - Dropping the T5 Encoder: 추론 시
text_encoder_3=None으로 T5‑XXL 인코더를 제거하면 메모리 요구량이 크게 줄어들고 성능 저하도 미미합니다. - 8-bit Quantization:
bitsandbytes라이브러리를 이용해 T5‑XXL 모델을 8비트 정밀도로 로드하면 메모리 소비를 추가로 낮출 수 있습니다.
메모리 벤치마크 비교
A100 GPU(80GB VRAM)에서 fp16 정밀도와 PyTorch 2.3을 사용해 수행한 벤치마크 결과는 다음과 같습니다:
| Technique | Inference Time (secs) | Memory (GB) |
|---|---|---|
| Default | 4.762 | 18.765 |
| Offloading | 32.765 | 12.0645 |
| Offloading + no T5 | 19.110 | 4.266 |
| 8bit T5 | 4.932 | 10.586 |
성능 및 파인튜닝
torch.compile()을 활용한 추론 가속
torch.compile()을 사용해 VAE와 트랜스포머 컴포넌트의 계산 그래프를 최적화하면 큰 속도 향상을 얻을 수 있습니다. 80GB A100 한 대에서 수행한 벤치마크에서는 평균 추론 시간이 0.585초로 감소했으며, 이는 eager execution 대비 약 4배 빠른 속도입니다.
DreamBooth 및 LoRA 학습
Hugging Face는 SD3용 DreamBooth 파인튜닝 스크립트를 공개했으며, 여기서는 Low‑Rank Adaptation (LoRA) 를 활용합니다. 이 스크립트를 통해 효율적인 모델 커스터마이징이 가능하며, rectified flow 기반 학습 파이프라인 구현을 위한 참고 자료로 활용될 수 있습니다.