Würstchen이 42배 공간 압축을 달성해 Stable Diffusion XL 대비 이미지 생성 속도 향상 및 GPU 학습 비용 16배 감소

Context

기존 확산 모델들은 이미지 생성 시 높은 계산 비용이 발생했으며, 4x~8x 정도의 공간 압축만 가능했다. 16x 이상의 압축을 시도하면 상세한 이미지 재구성에 실패했고, 1024×1024 이미지 학습에는 막대한 GPU 시간이 필요했다.

Technical Solution

• 두 단계 압축 구조(Stage A+B) 도입: VQGAN(Stage A)과 Diffusion Autoencoder(Stage B)를 순차적으로 적용해 42배 공간 압축 달성
• 고도로 압축된 잠재 공간에서 텍스트 조건부 확산 모델(Stage C, Prior) 학습: 기존 픽셀 공간 대신 32×32 크기의 압축 공간에서 모델 훈련
• PyTorch 2.0의 torch.nn.functional.scaled_dot_product_attention(SDPA) 자동 적용: Flash Attention 기법을 기본으로 활성화하여 주의 메커니즘 메모리 효율성 향상
• torch.compile 적용 옵션 제공: Prior 및 Decoder 모델에 컴파일링(mode="reduce-overhead", fullgraph=True)을 적용해 초기 2분의 일회성 컴파일 비용으로 지속적인 추론 속도 향상
• Model offload, Sequential CPU offload, xFormers flash attention 지원: Diffusers 라이브러리 통합으로 메모리-성능 트레이드오프 옵션 제공

Impact

• 학습 비용: Würstchen v1은 9,000 GPU 시간(Stable Diffusion 1.4 대비 16배 감소), Würstchen v2는 1536×1536 해상도에서 24,602 GPU 시간(512×512 기준 Stable Diffusion 1.4 대비 6배 감소)
• 추론 성능: 배치 크기별 비교에서 Stable Diffusion XL보다 빠른 생성 속도 달성
• 메모리 효율: SDXL 대비 메모리 사용량 대폭 감소
• 지원 해상도: 1024×1024~1536×1536 범위에서 학습되었으며 1024×2048에서도 양호한 출력 관찰

Key Takeaway

극단적 공간 압축(42배)을 달성함으로써 확산 모델의 학습 및 추론 비용을 획기적으로 감소시킬 수 있음을 입증했다. 엔지니어 리소스가 제한된 조직도 대규모 텍스트-이미지 생성 모델을 학습할 수 있는 진입장벽이 낮아졌다.