다양성 보존 디스트릴레이션으로 빠른 이미지 생성

초록

본 논문은 기존 디스트릴레이션 방식인 DMD가 역 KL 손실 때문에 발생하는 모드 붕괴 문제를 해결하고자, 첫 번째 스텝을 다양성 유지용 목표 예측(v‑prediction)으로 전용하고 이후 스텝은 기존 DMD 손실만 적용하는 역할 분리(distillation) 프레임워크인 DP‑DMD를 제안한다. 추가적인 퍼셉추얼 네트워크나 판별기 없이도 텍스트‑투‑이미지 작업에서 샘플 다양성을 크게 회복하면서도 품질은 기존 최첨단 수준을 유지한다.

상세 분석

본 연구는 확산 모델을 몇 단계만 사용해 빠르게 샘플링하고자 하는 디스트릴레이션 분야에서, 특히 Distribution Matching Distillation(DMD)이 역 KL(Kullback‑Leibler) 손실을 최소화하면서 “모드‑시킹” 경향을 보이는 점을 핵심 문제점으로 지적한다. 역 KL는 확률 질량이 높은 영역에만 집중하도록 유도해, 학습이 진행될수록 다양한 모드가 사라지고 결과적으로 이미지 다양성이 급격히 감소한다. 기존 해결책은 퍼셉추얼 손실(LPIPS 등)이나 GAN‑기반 판별기를 도입해 추가적인 정규화를 제공하지만, 이는 GPU 메모리와 연산량을 크게 늘리고, 특히 대규모 텍스트‑투‑이미지 모델에서는 학습 불안정성을 초래한다.

DP‑DMD는 이러한 문제를 “역할 분리”라는 간단하면서도 효과적인 설계로 극복한다. 첫 번째 디스틸 단계는 teacher 모델이 특정 노이즈 레벨(K 단계)에서 생성한 중간 상태를 이용해, v‑prediction(속도 예측) 목표를 설정한다. 이는 흐름 매칭(flow‑matching) 손실 L_Div = ‖v_θ(ε,1) − v_target‖² 로 구현되며, 여기서 v_target은 teacher‑derived 중간 샘플의 실제 흐름을 선형 보간식으로 역산한 값이다. 이 단계는 고노이즈 구간에서 전역 구조와 객체 배치를 결정하므로, 다양성을 보존하는 데 결정적인 역할을 한다.

그 후의 N‑1 단계는 기존 DMD 손실 L_DMD만을 사용해 품질을 미세 조정한다. 중요한 점은 첫 번째 단계의 출력(z₁)을 stop‑gradient 처리해, 이후 단계에서 역 KL 손실이 첫 단계에 역전파되지 않도록 차단한다. 이렇게 하면 역 KL이 초기에 다양성을 억제하는 영향을 완전히 차단하고, 후속 단계는 순수히 이미지 디테일과 텍스처를 개선하는 데 집중한다.

손실 함수는 L = L_DMD + λ·L_Div 로 결합되며, λ는 두 목표 사이의 균형을 조절한다. 실험에서는 λ=5×10⁻², K=5, M=5(가짜 모델 업데이트 주기) 등 적절한 하이퍼파라미터가 선택되었다.

DP‑DMD의 가장 큰 장점은 추가적인 퍼셉추얼 백본, 판별기, 혹은 외부 라벨 이미지가 전혀 필요 없다는 점이다. 모든 연산은 latent space에서 이루어지며, 기존 DMD와 동일한 NFEs(네트워크 함수 평가 횟수)만 사용하면서도 다양성 지표(DINO, CLIP 기반 코사인 유사도)에서 현저히 높은 점수를 기록한다. 특히 4 NFEs(≈4 스텝) 조건에서도 SD3.5‑M과 SDXL 교사 모델에 비해 다양성 손실이 최소화되며, 시각적 품질은 기존 DMD와 거의 차이가 없다는 것이 실험 결과로 입증된다.

이러한 설계는 “초기 단계는 전역 구조와 다양성을, 후기 단계는 세부 품질을 담당한다”는 확산 과정의 단계적 특성을 명확히 활용한 것으로, 향후 다른 흐름 기반 생성 모델이나 비디오 디스트릴레이션에도 일반화 가능성이 크다.