CVPR 2024에서 MIT(Massachusetts Institute of Technology)와 Adobe Research에서 공동으로 연구 및 발표한 "분포 매칭 증류를 이용한 원스텝 확산 모델(One-step Diffusion with Distribution Matching Distillation)"은 확산 모델을 통해 이미지 생성의 효율성을 대폭 개선하였습니다. 특히, 기존의 StableDiffusion v1.5와 비교해보았을 때 비슷한 이미지 품질을 30배 빠른 속도로 달성했다고 합니다.
DMD 소개
DMD(Distribution Matching Distillation) 연구는 확산 모델을 사용하여 불가능해 보이는 이미지를 실제와 더 가까운 확률 밀도 영역으로 유도하는 방법에 초점을 맞춥니다. 기존의 확산 모델과 비교하여, 연구팀은 실제 분포뿐만 아니라 가짜 분포의 점수 함수도 추정하는 두 개의 확산 모델을 훈련시킵니다. 이를 통해 생성된 이미지를 더욱 현실적으로 만들며, 훈련의 안정성을 높이고 빠른 속도로 이미지를 생성할 수 있게 됩니다.
확산 모델의 속도를 높이기 위한 기존의 시도는 주로 샘플링 과정을 최적화하거나 생성 능력을 더 효율적인 아키텍처로 증류하는 데 중점을 두었습니다. 일부 접근 방식은 이미지 품질을 희생하거나 전체 계산 시간을 크게 줄이지 못하는 등의 한계가 있습니다.
DMD 구조 및 동작
DMD(Distribution Matching Distillation)의 핵심은 생성 과정을 모델링하는 접근 방식에 있습니다. 기존 방식이 반복 과정 자체를 가속화하는 데 초점을 맞춘 것과 달리, DMD는 멀티스텝 확산 모델의 최종 출력 분포를 단일 생성 단계에서 재현하려고 합니다. 이 방법은 확산 모델이 생성하는 이미지와 유사한 품질의 이미지를 단 하나의 단계로 생성함으로써, 이미지 생성 속도를 대폭 향상시키는 것을 목표로 합니다. 다음은 DMD 방식의 동작 과정을 설명합니다:
1. 원스텝 생성기(One-step Generator) 학습
- 원스텝 생성기 Gθ의 핵심 목표는 임의의 노이즈 z를 실제와 유사한 이미지로 변환하는 것입니다. 이 과정은 확산 모델이 여러 단계에 걸쳐 샘플링하여 생성하는 출력과 일치하도록 설계되었습니다.
2. 노이즈-이미지 쌍 사전 계산
- 노이즈-이미지 쌍의 사전 계산은 이 방법의 핵심입니다. 확산 모델로부터 생성된 이미지와 해당 이미지를 생성하기 위해 사용된 노이즈의 쌍을 미리 계산하고 저장합니다. 이러한 쌍은 학습 과정 중에 원스텝 생성기의 입력으로 간헐적으로 사용됩니다.
3. LPIPS 회귀 손실 (regression loss)
- LPIPS(Learned Perceptual Image Patch Similarity) 회귀 손실은 원스텝 생성기가 생성한 이미지와 확산 모델의 출력 간의 유사성을 강제합니다. 이는 이미지 간의 지각적 유사성을 측정하는 데 사용되며, 생성된 이미지의 리얼리즘을 향상시키는 데 중요한 역할을 합니다.
4. 분포 일치 그라디언트 (distribution matching gradient)
- 분포 일치 그라디언트 ∇θ DKL은 생성된 가짜 이미지에 추가된 노이즈를 바탕으로 합니다. 이 그라디언트는 이미지의 리얼리즘을 더욱 강화하기 위해 제공됩니다.
5. 노이즈 주입 및 확산 모델 사용
- 생성된 가짜 이미지에 임의의 양의 노이즈를 주입한 후, 두 개의 확산 모델을 통과시킵니다. 하나는 실제 데이터에 대해 사전 학습된 모델이고, 다른 하나는 가짜 이미지에 대해 지속적으로 학습되는 모델입니다. 이를 통해 가짜 이미지의 노이즈 제거된 버전을 얻습니다.
6. 더 리얼리즘을 위한 방향 지시
- 노이즈 제거 점수(평균 예측으로 시각화됨)는 이미지를 더 리얼리스틱하게 만들기 위한 방향을 나타냅니다. 두 모델 간의 차이는 더 많은 리얼리즘과 더 적은 가짜성을 향한 방향을 나타내며, 이는 원스텝 생성기로 역전파됩니다.
이 과정을 통해 DMD는 실시간 또는 대화형 애플리케이션에서 사용할 수 있는 속도와 효율성을 갖춘, 고품질의 이미지 생성을 가능하게 합니다. 확산 모델의 느린 이미지 생성 속도를 극복하고, 실제와 유사한 퀄리티를 유지하면서 이미지 생성 과정을 가속화하는 혁신적인 접근 방법입니다.
결과 비교
Stable Diffusion과 비교하여 본 연구의 원스텝 확산 모델은 이미지 생성 속도에서 현저한 차이를 보입니다. 예를 들어, Stable Diffusion은 50단계 과정을 거쳐 이미지를 생성하는 데 약 2590ms가 걸린 반면, 원스텝 확산 모델은 단 90ms 만에 비슷한 품질의 이미지를 생성할 수 있습니다. 이는 AI 기반 이미지 생성 분야에서 획기적인 속도 향상을 의미합니다.
더 읽어보기
프로젝트 홈페이지
DMD 논문
DMD 소개 슬라이드
이 글은 GPT 모델로 정리한 글을 바탕으로 한 것으로, 원문의 내용 또는 의도와 다르게 정리된 내용이 있을 수 있으니 원문도 함께 참고해주세요! 읽으시면서 어색하거나 잘못된 내용을 발견하시면 덧글로 알려주시기를 부탁드립니다. 
파이토치 한국 사용자 모임
이 정리한 이 글이 유용하셨나요? 회원으로 가입하시면 주요 글들을 이메일
로 보내드립니다! (기본은 Weekly지만 Daily로 변경도 가능합니다.)
아래
쪽에 좋아요
를 눌러주시면 힘이 됩니다~ 