Large Scale Diffusion Distillation via Score-Regularized Continuous-Time Consistency

이 논문은 대규모 텍스트 - 이미지 및 비디오 생성을 위해 Jacobian-벡터 곱 연산을 병렬화하고 점수 정규화 기법을 도입하여 기존 연속 시간 일관성 모델의 한계를 극복하고, DMD2 와 유사한 품질을 유지하면서 모드 붕괴를 방지하고 다양성을 향상시킨 1~4 단계로 고충실도 샘플을 생성할 수 있는 확장 가능한 'rCM'을 제안합니다.

핵심

이 논문은 **"거대한 AI 그림/영상 생성기를 1 초 만에 그릴 수 있게 만든 획기적인 기술"**에 대한 이야기입니다.

기존의 AI 그림 그리기 기술 (확산 모델) 은 아주高质量 (고화질) 의 이미지를 만들지만, 한 장을 그리려면 50 번 이상의 복잡한 계산 과정을 거쳐야 해서 매우 느렸습니다. 이를 빠르게 만들기 위해 '지식 증류 (Distillation)'라는 기술을 쓰는데, 기존 방법들은 속도는 빠르지만 그림이 흐릿하거나, 혹은 다양성이 떨어지는 문제가 있었습니다.

이 논문은 **"rCM"**이라는 새로운 방법을 제안하며, 속도, 화질, 다양성이라는 세 마리 토끼를 모두 잡았습니다.

---

🎨 핵심 비유: "명화 복원 프로젝트"

이 기술의 원리를 이해하기 위해 명화 복원 프로젝트를 상상해 보세요.

1. 기존 방법들의 문제점 (속도 vs 화질)

• 기존의 빠른 방법 (sCM): "그림을 빨리 그리려면 붓질 횟수를 줄여야 해!"라고 생각한 기술입니다.
  • 문제: 붓질 횟수를 너무 줄이다 보니, 세부 묘사 (눈, 입, 글씨 등) 가 흐릿해지거나 모양이 뭉개지는 경우가 많았습니다. 마치 초고속으로 그린 스케치처럼요.
  • 원인: AI 가 "대충 비슷하게 그리면 돼"라는 생각 (Forward Divergence) 을 했기 때문입니다. 모든 가능성을 다 포함하려다 보니 디테일이 사라진 것입니다.
• 기존의 화질 좋은 방법 (DMD2): "화질을 잡으려면 AI 가 스스로 그림을 보고 고쳐야 해"라는 방법입니다.
  • 문제: 화질은 좋지만, 모든 그림이 똑같은 자세나 구도로 그려져 버렸습니다. (예: 모든 사람이 똑같은 방향으로 서 있음). 다양성이 사라진 것입니다.
  • 원인: AI 가 "가장 확실한 정답"만 찾다가 (Reverse Divergence), 다른 가능성을 무시해 버린 것입니다.

2. 이 논문의 해결책: "rCM (스케치 + 전문가 감수)"

이 논문은 두 가지 방법을 완벽하게 섞은 새로운 방식을 제안합니다.

• 비유:
  1. 스케치 단계 (sCM): 먼저 AI 가 빠르게 대략적인 구도와 형태를 잡습니다. (다양성은 좋지만, 디테일은 부족함).
  2. 전문가 감수 단계 (Score Regularization): 여기에 **명화 전문가 (원본 교사 모델)**가 와서 "여기 글씨는 이렇게 써야 해", "이 눈빛은 더 선명하게"라고 구체적인 코칭을 해줍니다.
  3. 결과: AI 는 빠른 스케치 능력을 유지하면서, 전문가의 코칭을 받아 디테일한 명화를 그립니다.

이 방식 덕분에 AI 는 **1~4 번의 붓질 (계산)**만으로 원본이 50 번 이상 그릴 때와媲美 (비견) 되는 고화질 영상을 만들 수 있게 되었습니다.

---

🚀 이 기술이 왜 대단한가요?

1. 거대 모델도 가능해짐 (100 억 개 이상의 파라미터):

   • 이전에는 이런 빠른 기술을 거대한 AI (영화 같은 영상을 만드는 모델) 에 적용하려면 컴퓨터 메모리가 터지거나 계산이 너무 복잡해서 불가능했습니다.
   • 이 논문은 FlashAttention-2라는 최신 기술을 활용해, 거대한 AI 의 "뇌"를 효율적으로 계산할 수 있게 만들었습니다. (마치 거대한 도서관을 관리하는 직원을 효율적으로 배치한 것과 같습니다.)

2. 화질과 다양성의 동시 달성:

   • 화질: 글씨를 정확하게 쓰거나, 사물의 모양을 왜곡 없이 그립니다.
   • 다양성: 같은 명령어 ("강아지가 공을 쫓는 영상") 를 입력해도, 매번 다른 각도, 다른 표정, 다른 배경을 만들어냅니다. (기존 빠른 방법들은 매번 똑같은 강아지 그림만 냈습니다.)

3. 압도적인 속도:

   • 기존에 50 초 걸리던 작업이 1~2 초 만에 끝납니다. (약 15 배~50 배 빠름).
   • 이는 마치 고속도로를 개통해서, 기존에 좁은 시골길로 1 시간 걸리던 길이 2 분 만에 도착하게 된 것과 같습니다.

💡 요약

이 논문은 **"AI 가 그림을 그릴 때, '속도'와 '화질'을 trade-off(상충) 관계로 생각하지 않고, 두 마리 토끼를 모두 잡을 수 있는 새로운 학습 방법 (rCM)"**을 개발했다는 것입니다.

앞으로 이 기술을 통해 실시간으로 영화 같은 영상을 만들거나, 수천 장의 다양한 이미지를 1 초 만에 생성하는 것이 일상적으로 가능해질 것입니다. 마치 AI 가 "명화"를 그리면서도 "스케치"만큼 빠르다는 놀라운 성과를 낸 셈입니다.

기술 요약

논문 요약: Large Scale Diffusion Distillation via Score-Regularized Continuous-Time Consistency (rCM)

이 논문은 대규모 텍스트 - 이미지 (T2I) 및 텍스트 - 비디오 (T2V) 생성 모델에 대한 **연속 시간 일관성 모델 (Continuous-Time Consistency Models, sCM)**의 확장성을 해결하고, 그 품질 한계를 극복하기 위한 새로운 프레임워크인 **Score-Regularized Continuous-Time Consistency Model (rCM)**을 제안합니다.

1. 문제 정의 (Problem)

기존의 연속 시간 일관성 모델 (sCM, MeanFlow 등) 은 이론적으로 정교하고 빠른 추론 (few-step generation) 을 가능하게 하지만, 대규모 실제 응용 (10B+ 파라미터 모델, 고해상도 비디오) 에 적용하는 데 다음과 같은 심각한 장벽이 존재했습니다.

• 인프라 및 계산적 한계: sCM 학습에는 자코비안 - 벡터 곱 (JVP, Jacobian-Vector Product) 계산이 필수적입니다. 대규모 모델 (BF16 정밀도, FlashAttention, 병렬 학습 환경) 에서 JVP 를 효율적으로 계산하는 것은 기존 PyTorch 구현으로는 불가능하거나 수치적 불안정성을 초래했습니다.
• 품질 저하 (Quality Issues): sCM 은 '모드 커버링 (Mode-Covering)' 특성을 가지며, forward divergence 기반의 목적 함수를 사용합니다. 이로 인해 미세한 디테일 (텍스트 렌더링, 객체 기하학 등) 에서 왜곡이나 흐릿함이 발생하고, 오차가 시간 단계가 커짐에 따라 누적되는 문제가 있었습니다.
• 평가의 한계: 기존 평가 지표 (FID 등) 는 대규모 T2I/T2V 작업의 정밀한 속성 (텍스트 정확도, 복잡한 프롬프트 준수) 을 충분히 반영하지 못했습니다.

2. 방법론 (Methodology)

저자들은 sCM 의 확장성을 높이고 품질 문제를 해결하기 위해 다음과 같은 기술적 혁신을 제안했습니다.

2.1 대규모 JVP 연산을 위한 인프라 구축

• FlashAttention-2 JVP 커널: 대규모 병렬 학습 (FSDP, Context Parallelism) 환경에서 JVP 를 효율적으로 계산하기 위해 FlashAttention-2 알고리즘에 JVP 로직을 통합한 커스텀 Triton 커널을 개발했습니다. 이를 통해 100 억 개 이상의 파라미터를 가진 모델과 고차원 비디오 데이터 학습을 가능하게 했습니다.
• 네트워크 구조 재구성: FSDP 와 호환되도록 각 레이어 수준에서 JVP 연산을 수행할 수 있도록 네트워크를 재구성했습니다.

2.2 Score-Regularized Continuous-Time Consistency Model (rCM)

sCM 의 품질 한계를 극복하기 위해 **Score Distillation (역방향 발산, Reverse Divergence)**을 정규화 항으로 도입했습니다.

• 하이브리드 목적 함수:
  • Forward Consistency (sCM): 실제 데이터나 교사 모델 생성 데이터 기반의 '모드 커버링' 특성을 유지하여 **다양성 (Diversity)**을 확보합니다.
  • Reverse Divergence (Score Distillation): 학생 모델이 스스로 생성한 샘플을 기반으로 한 '모드 찾기 (Mode-Seeking)' 특성을 추가하여 **품질 (Quality)**을 향상시킵니다.
  • 최종 Loss: $L_{rCM} = L_{sCM} + \lambda L_{DMD}$ (여기서 $L_{DMD}$는 분포 매칭 증류 손실).
• 작동 원리: sCM 은 오차 누적을 유발할 수 있는 단계를 건너뛰는 (long-skip) 방식으로 작동하고, Score Distillation 은 이를 보정하는 정규화제 (regularizer) 역할을 하여 시각적 품질을 개선하면서도 다양성을 유지합니다.

2.3 안정적인 시간 미분 계산

JVP 계산 시 발생하는 수치적 불안정성 (특히 BF16 환경) 을 해결하기 위해 두 가지 전략을 사용했습니다.

• Semi-Continuous Time: 유한 차분 (Finite Difference) 을 사용하여 시간 미분을 근사화.
• High-Precision Time: 10B+ 모델 및 비디오 작업에서는 시간 임베딩 레이어만 FP32 정밀도로 강제하여 JVP 계산의 안정성을 확보했습니다.

3. 주요 기여 (Key Contributions)

1. 대규모 확산 증류의 첫 시도: 100 억 개 이상의 파라미터 (Cosmos-Predict2 14B, Wan2.1 14B) 와 5 초 길이의 비디오 생성에 연속 시간 일관성 모델을 성공적으로 확장한 최초의 연구입니다.
2. JVP 기반 대규모 학습 인프라: FlashAttention-2 와 병렬 학습을 지원하는 JVP 커널을 개발하여 대규모 모델 학습의 기술적 장벽을 제거했습니다.
3. rCM 프레임워크 제안: Forward (sCM) 와 Reverse (Score Distillation) 발산을 결합하여, sCM 의 품질 저하를 해결하고 DMD2 와 같은 기존 최첨단 방법보다 우수한 다양성을 확보한 새로운 증류 프레임워크를 제시했습니다.
4. 실용적인 성능: GAN 튜닝이나 복잡한 하이퍼파라미터 탐색 없이, 1~4 스텝으로 고품질 샘플을 생성하여 기존 확산 모델 대비 15 배~50 배의 추론 속도 향상을 달성했습니다.

4. 실험 결과 (Results)

• 텍스트 - 이미지 (T2I): Cosmos-Predict2 모델 (0.6B ~ 14B) 에서 rCM 은 4 스텝 추론 시 SOTA 모델 (FLUX.1-schnell 등) 과 유사하거나 더 나은 GenEval 점수를 기록했습니다. 특히 작은 텍스트 렌더링 등 미세한 디테일에서 sCM 의 왜곡 문제를 해결했습니다.
• 텍스트 - 비디오 (T2V): Wan2.1 (1.3B, 14B) 및 Cosmos-Predict2 모델에서 rCM 은 4 스텝 추론 시 VBench 점수에서 교사 모델 (Teacher) 을 능가하거나 동급의 성능을 보였습니다.
• 다양성 vs 품질: DMD2 는 품질은 높지만 모드 붕괴 (Mode Collapse) 로 인해 다양성이 떨어지는 경향이 있는 반면, rCM 은 높은 품질과 높은 다양성을 동시에 달성했습니다. (Figure 1 참조: rCM 은 객체의 위치/방향/운동이 교사 모델과 유사하면서도 DMD2/SiD 보다 다양한 생성 결과를 보입니다.)
• 추론 효율성: 14B 파라미터 모델 기준, 4 스텝 생성 시 14.4 FPS(14B Wan2.1) 의 처리량을 보여주며, 1~2 스텝에서도 경쟁력 있는 성능을 입증했습니다.

5. 의의 및 결론 (Significance)

이 연구는 대규모 생성 모델 분야에서 **이론적으로 정교한 연속 시간 일관성 모델 (sCM)**을 실제 산업 수준의 대규모 모델에 적용할 수 있는 길을 열었습니다.

• 패러다임의 전환: 단순히 'Forward' 또는 'Reverse' 발산 중 하나를 선택하는 것이 아니라, 두 가지 접근법을 결합하여 품질과 다양성이라는 상충되는 목표 (Trade-off) 를 동시에 최적화할 수 있음을 증명했습니다.
• 실용성: 복잡한 GAN 기반 튜닝 없이도 1~4 스텝으로 고품질 비디오/이미지를 생성할 수 있어, 실시간 인터랙티브 애플리케이션 및 대규모 콘텐츠 생성에 매우 실용적인 솔루션을 제공합니다.
• 미래 전망: rCM 의 패러다임은 자기 강제 (Self-Forcing) 기반의 오토레거시 비디오 확산 모델 등 차세대 생성 모델 연구에도 중요한 영감을 줄 것으로 기대됩니다.

요약하자면, rCM은 대규모 확산 모델의 증류 (Distillation) 를 위한 새로운 표준을 제시하며, 빠른 추론 속도와 높은 생성 품질/다양성을 모두 만족시키는 실용적이고 이론적으로 탄탄한 프레임워크입니다.