Cross-Resolution Distribution Matching for Diffusion Distillation

이 논문은 로그 신호대잡음비 (logSNR) 기반 매핑과 분포 정합을 통해 해상도 간 격차를 해소하고, 예측 노이즈 재주입 메커니즘을 도입하여 고화질 생성을 유지하면서 SDXL 및 Wan2.1-14B 모델에서 최대 33.4 배의 추론 가속화를 이루는 '교차 해상도 분포 정합 증류 (RMD)' 프레임워크를 제안합니다.

핵심

🎨 1. 문제점: "왜 처음부터 4K 화질로 그리면 안 될까?"

기존의 AI 그림 그리기 기술 (확산 모델) 은 아주 정교한 그림을 만들려면 수백 번의 반복 작업이 필요합니다. 마치 아티스트가 4K 화질의 거대한 캔버스에 처음부터 마지막 붓터치까지 모든 디테일을 한 번에 그리려다 지쳐버리는 상황과 비슷합니다.

• 기존의 시도: "그림을 빨리 그리려면 붓질 횟수 (단계) 를 줄이자!"라고 했지만, 횟수를 너무 줄이면 그림이 뭉개지거나 엉망이 됩니다.
• 새로운 아이디어: "그럼 작은 캔버스 (저해상도) 로 먼저 전체 구도 (얼굴, 배경 위치 등) 를 대충 그리고, 그걸 바탕으로 큰 캔버스 (고해상도) 로 디테일을 채워보자!"

하지만 여기서 큰 문제가 생깁니다. 작은 캔버스에서 그린 그림과 큰 캔버스에서 그린 그림은 '분위기'나 '스타일'이 다릅니다. (예: 작은 스케치에서는 눈이 작게 보이지만, 크게 확대하면 눈동자 디테일이 달라져서 어색해짐). 이걸 **해상도 간 차이 (Distribution Gap)**라고 합니다.

🌉 2. 해결책: RMD 의 마법 같은 다리

이 논문은 이 '작은 캔버스'와 '큰 캔버스' 사이의 괴리를 메우는 **마법 같은 다리 (RMD)**를 만들었습니다.

① "소금물 농도"로 구간 나누기 (LogSNR 기반 구간 분할)

그림을 그릴 때, 어떤 단계에서 어떤 크기의 캔버스를 쓸지를 아주 정교하게 결정합니다.

• 초반 (소음 많을 때): 전체적인 구도 (산이 어디에, 사람이 어디에) 를 잡는 단계입니다. 이때는 **작은 캔버스 (저해상도)**로 빠르게 그립니다. (계산 비용이 적게 듭니다.)
• 후반 (소음 적을 때): 눈썹 하나, 옷 주름 하나 같은 디테일을 채우는 단계입니다. 이때는 **큰 캔버스 (고해상도)**로 정교하게 그립니다.

② "스타일 통일" 주문하기 (분포 매칭)

작은 캔버스에서 그린 구도가 큰 캔버스에 옮겨질 때, AI 가 "아, 이건 작은 그림이니까 이런 스타일로 그렸지"라고 착각하지 않게 합니다.

• 비유: 작은 스케치를 큰 캔버스에 옮길 때, 작은 스케치의 '느낌'을 그대로 유지하면서 크지만, 큰 캔버스의 '고화질 스타일'과 똑같아지도록 AI 를 훈련시킵니다. 이를 분포 매칭이라고 합니다.

③ "예상되는 실수"를 미리 고쳐주기 (노이즈 재주입)

작은 그림을 크게 늘리면 (업샘플링), 그림이 흐릿해지거나 찌그러질 수 있습니다.

• 비유: 확대할 때 AI 가 "아, 여기가 흐릿해지면 안 되겠네, 내가 예측한 디테일을 미리 넣어줄게!"라고 예상되는 노이즈 (실수) 를 다시 주입해서 그림이 뭉개지지 않게 안정적으로 만듭니다.

---

🚀 3. 결과: 얼마나 빨라졌을까?

이 방법을 쓰면 어떻게 될까요?

• 속도: 기존에 그림을 그리는데 100 걸음 걸렸다면, 이 방법은 33 배나 빨라져서 3 걸음 만에 그립니다. (SDXL 기준 33.4 배 속도 향상)
• 화질: 속도가 빨라졌다고 그림이 뭉개진 건 아닙니다. 고화질 그대로의 퀄리티를 유지합니다.
• 영상에도 적용: 정지된 그림뿐만 아니라, **움직이는 영상 (비디오)**을 만들 때도 같은 원리가 적용되어 훨씬 빠르게 고퀄리티 영상을 만들어냅니다.

💡 요약: 한 문장으로 정리하면?

"처음부터 거대한 캔버스에 모든 디테일을 다 그리느라 지친 AI 에게, '먼저 작은 스케치로 전체 구도를 잡고, 그걸 바탕으로 큰 캔버스에 디테일을 채우는' 새로운 방식을 가르쳐서, 화질은 그대로 유지하면서 속도는 30 배 이상 빠르게 만든 기술입니다."

이 기술은 앞으로 우리가 AI 로 그림이나 영상을 만들 때, 기다리는 시간을 획기적으로 줄여주면서도 멋진 결과물을 얻을 수 있게 해줄 것입니다.

기술 요약

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

• 확산 모델의 비효율성: 최신 확산 모델 (Diffusion Models) 은 고화질 이미지 및 비디오 생성에서 뛰어난 성능을 보이지만, 수백 번의 반복적인 디노이징 (denoising) 단계가 필요하여 계산 비용과 지연 시간이 매우 큽니다.
• 기존 증류 (Distillation) 방법의 한계:
  • 단계 수 축소 (Step Reduction): 기존 증류 기법들은 생성 단계를 줄여 (예: 48 단계) 추론 속도를 높이지만, 단계를 과도하게 줄이면 (13 단계) 성능이 급격히 저하되는 '성능 포화' 현상이 발생합니다.
  • 해상도 불일치 (Resolution Mismatch): 추론 효율을 높이기 위해 초기 단계에서 저해상도로 생성을 시작하는 다중 해상도 캐스케이드 전략은 유망하지만, 기존 모델은 저해상도와 고해상도 데이터 분포 간의 간극 (Distribution Gap) 이 존재합니다.
  • 핵심 문제: 동일한 텍스트 프롬프트라도 모델이 학습된 해상도 (저해상도 vs 고해상도) 에 따라 데이터 분포가 달라, 단순히 해상도를 낮추고 디노이징을 수행하면 생성 품질이 크게 떨어지고 구조적 일관성이 깨집니다 (Fig. 1 참조).

2. 제안 방법: RMD (Methodology)

저자들은 Cross-Resolution Distribution Matching Distillation (RMD) 을 제안하여, 저해상도 생성기의 분포와 고해상도 교사 (Teacher) 모델의 분포 간의 간극을 해소하고 few-step(소수 단계) 다중 해상도 캐스케이드 생성을 가능하게 합니다.

가. 로그 SNR 기반 해상도 구간 분할 (Resolution Trajectory Division)

• 로그 SNR (logSNR) 활용: 노이즈 스케줄은 해상도에 따라 다르게 작용합니다. RMD 는 logSNR 곡선을 기준으로 디노이징 경로를 여러 구간으로 나눕니다.
• 해상도별 시간 간격 매핑: 각 구간마다 다른 해상도 ($r_1 < r_2 < \dots < r_K$) 를 할당합니다. 저해상도 구간에서는 전역 구조 (Global Structure) 를 빠르게 복원하고, 고해상도 구간으로 넘어가면서 세부 디테일 (Fine-grained Details) 을 정제합니다.
• 정합 (Alignment): 교사 모델의 logSNR 임계값을 기반으로 학생 모델의 시간 단계 (timestep) 를 조정하여, 서로 다른 해상도에서도 동일한 디노이징 상태 (denoising state) 에서 증류가 이루어지도록 합니다.

나. 교차 해상도 분포 정합 (Cross-Resolution Distribution Matching)

• 목표: 저해상도에서 생성된 잠재 변수 (latent) 를 고해상도 공간으로 업샘플링한 후, 교사 모델의 분포와 정합시킵니다.
• 손실 함수: KL 발산 (Kullback-Leibler Divergence) 을 최소화하여 학생 모델의 출력 분포가 교사 모델의 분포와 일치하도록 학습합니다.
  • $L(\theta) = \sum \lambda_{r_i} \mathbb{E}_{t_i} [ \mathrm{KL}( p_\theta(x_{t_i}^{(r_i)}) \parallel p_\varphi(x_{t_i}) ) ]$
• 가짜 스코어 모델 (Fake Score Model): 직접적인 스코어 추정이 어렵기 때문에, DMD 나 TDM 과 유사하게 '가짜 확산 모델'을 도입하여 학생 모델의 스코어를 추정하고 증류합니다.

다. 예측 노이즈 재주입 메커니즘 (Predicted-Noise Re-injection)

• 문제: 단순한 업샘플링과 가우시안 노이즈 주입은 고해상도 교사의 ODE 경로를 따르지 못해 구조적 일관성이 깨질 수 있습니다.
• 해결: 업샘플링 시 순수 가우시안 노이즈 대신 예측된 노이즈 (Predicted Noise) 와 확률적 노이즈 (Gaussian Noise) 를 혼합하여 주입합니다.
  • $\epsilon_{t_i} = \alpha \cdot U(\epsilon_\theta) + \beta \cdot \epsilon$
  • 해상도 간격이 클수록 확률적 노이즈의 비중 ($\beta$) 을 높여 분포 불일치를 보완하고, 해상도가 가까워질수록 예측 노이즈 ($\alpha$) 를 강조하여 교사의 경로를 따르도록 합니다.

라. 학습 및 추론 전략

• Warm-up Training: 초기 저 logSNR(의미론적 구조) 구간을 먼저 증류하여 안정적인 초기화를 제공합니다.
• Multi-Resolution Cascaded Inference: 추론 시 저해상도에서 시작하여 단계적으로 해상도를 높이며, 해상도가 변경될 때마다 업샘플링과 노이즈 재주입을 수행하여 매끄러운 전환을 보장합니다.

3. 주요 기여 (Key Contributions)

1. 새로운 증류 프레임워크: 단계 수 축소만 의존하는 기존 방식의 한계를 넘어, 해상도 축소와 분포 정합을 결합하여 고화질 few-step 생성을 가능하게 한 RMD 프레임워크 제안.
2. 해상도 간 분포 간극 해소: logSNR 기반의 시간 단계 매핑과 예측 노이즈 재주입 메커니즘을 통해 저해상도 생성 분포와 고해상도 교사 분포 간의 불일치를 효과적으로 해결.
3. 범용성 및 확장성: 이미지 생성 (SDXL, PixArt-α, SD3.5) 과 비디오 생성 (Wan2.1-14B) 모두에서 검증되었으며, 다양한 백본 모델에 적용 가능.

4. 실험 결과 (Results)

• 이미지 생성 (Text-to-Image):
  • SDXL: 33.4 배의 추론 속도 향상 (Speedup) 을 달성하면서도 HPS(Human Preference Score), Aesthetic Score, CLIP Score 등 모든 지표에서 기존 증류 방법 (SDXL-Turbo, DMD2, TDM 등) 보다 우수한 성능을 보임.
  • PixArt-α 및 SD3.5: 각각 21.0 배, 32.0 배의 속도 향상과 함께 최상위 품질 유지.
• 비디오 생성 (Text-to-Video):
  • Wan2.1-14B: 25.6 배의 속도 향상 (기존 6 단계 대비 3+3 단계 전략) 을 달성.
  • 품질: VBench 및 T2V-CompBench 에서 경쟁 모델 (DMD2, TDM) 보다 시간적 일관성과 운동 디테일에서 우위를 점함.
• 생성 품질: 2+2(저해상도 2 단계 + 고해상도 2 단계) 구성이 효율성과 품질 간의 최적 균형을 제공함을 실험을 통해 입증.

5. 의의 및 결론 (Significance)

• 효율성 병목 현상 극복: 기존 확산 모델의 추론 속도를 획기적으로 개선하면서도 고화질 생성 품질을 유지하는 새로운 패러다임을 제시합니다.
• 실용성: 실시간 애플리케이션이나 리소스가 제한된 환경에서도 고품질 이미지/비디오 생성을 가능하게 하여, 생성형 AI 의 상용화 장벽을 낮춥니다.
• 미래 지향성: 단순한 단계 축소를 넘어, 해상도와 시간적 스케일을 통합적으로 최적화하는 증류 기법의 중요성을 부각시켰으며, 차세대 대규모 생성 모델 가속화의 핵심 기술로 평가됩니다.

요약하자면, RMD는 "저해상도로 빠르게 구조를 잡고, 고해상도로 디테일을 다듬는" 캐스케이드 방식에 분포 정합 (Distribution Matching) 기술을 접목하여, 기존 증류 방법의 한계를 극복하고 초고속 고화질 생성을 실현한 획기적인 연구입니다.