Scale-wise Distillation of Diffusion Models

이 논문은 기존 확산 모델 증류의 효율성 한계를 극복하기 위해 중간 단계 계산을 줄이는 '스케일별 증류 (SwD)' 프레임워크와 MMD 기반의 새로운 증류 목표를 제안하여, 적은 계산 비용으로도 고품질의 2 단계 샘플링 속도와 성능을 달성함을 보여줍니다.

핵심

🎨 그림을 그리는 '스마트한 학생' 이야기: SwD (Scale-wise Distillation)

안녕하세요! 오늘 소개해 드릴 논문은 **"Diffusion Model(확산 모델)"**이라는 아주 똑똑하지만 조금 느린 그림 그리기 AI 를, 훨씬 더 빠르고 똑똑한 학생으로 만드는 새로운 방법론에 대한 이야기입니다. 이 방법은 SwD라고 불립니다.

이 복잡한 기술을 일상적인 비유로 쉽게 풀어보겠습니다.

---

1. 문제: "완벽한 그림을 그리려면 너무 오래 걸려요!" 🐢

지금까지의 AI 그림 그리기 기술 (Diffusion Model) 은 마치 아주 정교한 화가와 같습니다.

• 방식: 처음엔 완전히 하얀 천 (소음) 을 들고 시작해서, 한 번에 한 번씩 천천히 그림을 수정해 나갑니다.
• 문제: 이 화가는 그림을 완성하려면 20~50 번이나 천을 수정해야 합니다. (예: 40 번 수정)
• 결과: 그림은 정말 예쁘지만, 그걸 그리는데 시간이 너무 오래 걸려서 실생활에서 쓰기엔 답답합니다.

기존 연구자들은 "그림을 그리는 횟수 (단계) 를 줄이자!"라고 생각했습니다. 하지만 4 단계로 줄이는 건 이미 한계에 다다랐고, 1~2 단계로 줄이면 그림이 뭉개지거나 엉망이 되는 문제가 생겼습니다.

2. 새로운 아이디어: "먼저 대략적인 윤곽을 그리고, 디테일은 나중에!" 🎨

이 논문 (SwD) 의 핵심 아이디어는 **"그림을 그릴 때, 처음부터 끝까지 같은 크기의 캔버스를 쓸 필요가 없다"**는 것입니다.

• 기존 방식 (비효율적): 1 단계부터 40 단계까지 모두 **거대한 캔버스 (고해상도)**에서 그림을 그립니다. 하지만 초반에는 소음 (잡음) 이 너무 많아서 세부적인 디테일은 보이지 않습니다. 마치 거대한 캔버스에 먼지만 뿌리는 것과 같아 비효율적입니다.
• SwD 의 방식 (효율적):
  1. 초반 (소음이 많을 때): 아주 **작은 캔버스 (저해상도)**에서 그림의 대략적인 윤곽만 빠르게 그립니다. (예: 32x32 픽셀)
  2. 중반: 그림이 조금 더 선명해지면, 캔버스를 조금 더 크게 늘립니다. (예: 64x64 픽셀)
  3. 후반 (소음이 적을 때): 이제야 **최종적인 거대한 캔버스 (고해상도)**로 확대해서 세부적인 디테일을 채워 넣습니다. (예: 1024x1024 픽셀)

💡 비유:
마치 지도 앱을 사용하는 것과 같습니다.

• 처음엔 전 세계 지도를 쭉 펼쳐놓고 (고해상도) 도시 하나를 찾으려 하면 너무 느립니다.
• 대신 전체 지도를 먼저 보고 (저해상도), 원하는 지역을 **확대 (Zoom-in)**해서 들어가는 것이 훨씬 빠르고 효율적입니다. SwD 는 AI 가 그림을 그릴 때 이 '확대' 과정을 자연스럽게 활용하는 것입니다.

3. 두 가지 핵심 기술: "스마트한 학습법" 🧠

SwD 는 이 방식을 구현하기 위해 두 가지 특별한 기술을 사용합니다.

① 점진적 성장 (Scale-wise Distillation)

위에서 설명한 대로, AI 가 작은 캔버스에서 시작해 점점 큰 캔버스로 그림을 그리도록 훈련시킵니다.

• 장점: 초반에 불필요한 고해상도 계산을 하지 않으므로, 전체 속도가 2~3 배 빨라집니다.
• 결과: 4 단계만으로도 기존 AI 가 40 단계로 그리는 것과 같은 퀄리티를 냅니다.

② MMD (최대 평균 불일치) 라는 '맛보기' 테스트 🍽️

기존에는 AI 가 그리는 그림이 진짜인지 확인하기 위해 '가짜 판별기 (Discriminator)'라는 또 다른 AI 를 훈련시켜야 했습니다. 이는 매우 비싸고 복잡했습니다.

• SwD 의新方法: 새로운 AI 가 그리는 그림의 **특징 (맛)**을 분석하는 MMD라는 수학적 도구를 사용합니다.
• 비유: 미식가가 두 요리의 '맛'을 비교할 때, 요리사 (AI) 를 따로 고용할 필요 없이, **요리 자체의 재료와 향 (특징)**을 분석해서 "이게 진짜 같은 맛이야!"라고 판단하는 것입니다.
• 장점: 추가적인 AI 를 훈련시킬 필요가 없어 계산 비용이 적게 들고, 그림의 질이 더 좋아집니다.

4. 실제 성과: "기존의 거장보다 10 배 빠르고 똑똑해!" 🚀

이론만 좋은 게 아닙니다. 실험 결과 놀라운 성과가 나왔습니다.

• 속도: 기존 최고의 AI 들보다 10 배 이상 빠릅니다. (예: 40 초 걸리던 게 4 초로 줄어듦)
• 퀄리티: 같은 시간 (계산 비용) 을 썼을 때, 기존 방식보다 훨씬 더 선명하고 아름다운 그림을 그립니다.
• 영상까지: 이미지뿐만 아니라 동영상을 만들 때도 적용되어, 영상 생성 속도를 3 배나 높였습니다.

5. 결론: "효율의 새로운 시대" ✨

이 논문 (SwD) 은 **"그림을 그릴 때, 처음부터 끝까지 같은 힘으로 일할 필요는 없다"**는 통찰을 주었습니다.

• 과거: 무조건 큰 캔버스에서 천천히 그렸다. (비효율적)
• SwD: 작은 캔버스에서 시작해 점점 키워가며 그렸다. (초효율적)

이 기술은 앞으로 우리가 AI 로 그림을 그리거나 동영상을 만들 때, 더 빠르고 저렴하게 고퀄리티 콘텐츠를 즐길 수 있게 해주는 열쇠가 될 것입니다. 마치 스마트폰이 처음엔 느렸지만, 이제는 순식간에 고화질 영상을 편집하듯, AI 그림 그리기도 이제 '초고속' 시대로 진입하게 된 셈입니다.

기술 요약

1. 연구 배경 및 문제 제기 (Problem)

• 현재의 한계: 대규모 텍스트 조건부 이미지 및 비디오 생성을 위한 확산 모델 (Diffusion Models, DMs) 은 고품질 생성을 위해 일반적으로 2050 단계의 순차적 샘플링이 필요합니다. 최근 연구들은 이를 4 단계 이하로 줄이는 '디스틸레이션 (Distillation)' 기법을 통해 속도를 개선했으나, 12 단계로 더 줄이는 것은 여전히 매우 어렵습니다.
• 비효율성: 기존 few-step 모델들은 전체 생성 과정 동안 고정된 해상도 (Full-resolution) 에서 작동합니다. 그러나 확산 과정의 초기 단계 (고잡음, High noise) 에서는 고주파수 정보가 억제되어 저해상도에서도 데이터 신호를 충분히 표현할 수 있음에도 불구하고, 불필요하게 고해상도 계산을 수행하여 계산 자원을 낭비하고 있습니다.
• 과제: 기존 디스틸레이션 방법들은 모델 아키텍처나 입력 해상도 등 다른 자유도를 고정시킨 채 단계 수만 줄이는 데 집중해 왔습니다. 효율성을 극대화하기 위해 해상도 (Scale) 라는 새로운 축을 활용한 접근이 필요합니다.

2. 핵심 방법론 (Methodology)

이 논문은 Scale-wise Distillation (SwD) 프레임워크를 제안하며, 이는 두 가지 주요 기술로 구성됩니다.

A. 스펙트럼 분석 및 점진적 생성 (Spectral Analysis & Progressive Generation)

• 잠재 공간 스펙트럼 분석: 저자들은 VAE 잠재 공간 (Latent Space) 에서 확산 과정을 분석한 결과, 자연 이미지와 마찬가지로 확산 과정이 저주파수에서 고주파수로 점진적으로 예측되는 '스펙트럼 자기회귀 (Spectral Autoregression)' 특성을 가짐을 확인했습니다.
• 고잡음 시 저해상도 작동 가능: 고잡음 (High noise) 단계에서는 고주파수 성분이 노이즈에 가려지므로, 이 단계에서는 낮은 잠재 공간 해상도로 모델을 작동시켜도 정보 손실 없이 생성이 가능함을 증명했습니다.
• SwD 파이프라인:
  1. 해상도 스케줄링: 생성 단계마다 잠재 공간의 해상도를 점진적으로 증가시킵니다 (예: 32x32 → 64x64 → 128x128).
  2. 업샘플링 전략: 기존 방식 (노이즈를 먼저 넣고 업샘플링) 은 분포 왜곡을 일으킵니다. SwD 는 먼저 저해상도에서 깨끗한 샘플 ($\hat{x}_0$) 을 예측한 후, 이를 업샘플링하고 다음 단계의 노이즈를 주입하는 방식을 사용합니다. 이는 고해상도에서의 노이즈 통계를 올바르게 유지합니다.
  3. 단일 모델 통합: 캐스케이드 (여러 모델 사용) 방식이 아닌, 단일 few-step 모델이 다양한 해상도에서 점진적으로 작동하도록 학습시킵니다.

B. 최대 평균 불일치 (MMD) 기반 디스틸레이션 목적 함수

• 새로운 목적 함수: 기존 분포 일치 (Distribution Matching) 기법을 보완하기 위해, 사전 학습된 확산 모델의 특성 공간 (Feature Space) 에서 패치 (Patch) 수준의 최대 평균 불일치 (Maximum Mean Discrepancy, MMD) 를 최소화하는 목적 함수 ($L_{MMD}$) 를 도입했습니다.
• 특징:
  • 추가적인 학습 가능한 모델 (Discriminator 등) 이 필요하지 않아 계산 효율적입니다.
  • 생성된 샘플과 타겟 샘플의 공간적 토큰 (Spatial tokens) 평균을 비교하여 분포를 정렬합니다.
  • 기존 DMD 나 GAN 손실과 결합하거나, 단독으로도 강력한 성능을 발휘합니다.

3. 주요 기여 (Key Contributions)

1. Scale-wise Distillation (SwD) 프레임워크: 기존 few-step 모델을 변형하여 생성 단계마다 해상도를 점진적으로 높이는 새로운 디스틸레이션 방식을 제안했습니다.
2. MMD 기반 패치 레벨 목적 함수: 별도의 판별자 없이 사전 학습된 모델의 특징을 활용하여 분포 정렬을 수행하는 간단하지만 효과적인 새로운 손실 함수를 제시했습니다.
3. 스펙트럼 분석을 통한 이론적 근거: 확산 모델의 잠재 공간에서 고잡음 단계는 저해상도로 모델링할 수 있음을 스펙트럼 분석을 통해 입증하고, 이를 효율적인 생성 전략으로 연결했습니다.

4. 실험 결과 (Results)

논문은 SDXL, SD3.5, FLUX (이미지) 및 Wan2.1 (비디오) 등 최첨단 모델에 SwD 를 적용하여 검증했습니다.

• 생성 속도 및 효율성:
  • 이미지: 기존 최첨단 few-step 모델 대비 약 2 배 빠른 추론 속도를 달성했습니다 (동일한 계산 예산 하에서).
  • 비디오: Wan2.1 기반 모델에서 약 3 배 빠른 속도를 달성했습니다.
  • 비교: 4 단계 SwD 모델은 2 단계 풀해상도 모델보다 더 빠르면서도 품질이 우수하거나 동급인 것으로 나타났습니다.
• 생성 품질:
  • 자동 지표 (FID, HPSv3, ImageReward 등): 대부분의 모델 그룹에서 SOTA 성능을 기록하거나 기존 디스틸레이션 모델 (Turbo, DMD2 등) 을 능가했습니다.
  • 인간 평가: 텍스트 관련성, 이미지 미학, 복잡성, 결함 (Defects) 등 다양한 기준에서 인간 평가자들이 SwD 를 선호했습니다. 특히 2 단계 풀해상도 모델 대비 결함이 현저히 적고 이미지 복잡도가 높았습니다.
• MMD 손실의 효과: MMD 손실만으로도 경쟁력 있는 few-step 모델을 생성할 수 있으며, 기존 손실 함수와 결합 시 수렴 속도와 성능을 크게 향상시킵니다.

5. 의의 및 결론 (Significance)

• 효율성의 새로운 패러다임: 확산 모델의 가속화를 위해 단순히 '단계 수 줄이기'에 집중하던 기존 접근에서 벗어나, 해상도 스케일링 (Scale-wise) 을 활용한 효율성 증대를 제시했습니다.
• 실용성: 추가적인 학습 가능한 모델 없이 기존 디스틸레이션 파이프라인에 쉽게 통합 가능하며, 텍스트 - 이미지 및 텍스트 - 비디오 생성 모두에 적용 가능합니다.
• 미래 전망: 제안된 MMD 기반 목적 함수는 자체적으로 강력한 성능을 보여주어, 향후 추가적인 학습 모델 없이도 효율적인 디스틸레이션 파이프라인을 구축하는 데 중요한 기여를 할 것으로 기대됩니다.

요약하자면, SwD는 확산 모델이 고잡음 단계에서는 저해상도로 작동해도 무방하다는 통찰을 바탕으로, 해상도를 점진적으로 높이는 생성 방식과 새로운 분포 정렬 손실 함수를 결합하여, 품질 저하 없이 생성 속도를 획기적으로 개선한 획기적인 연구입니다.