적은 단계의 확산 모델 증류를 위한 연속 시간 분포 매칭

기존의 확산 모델 증류 방식은 고정된 이산적 시간 단계에 의존하여 이미지의 세부 사항이 뭉개지거나 아티팩트가 발생하는 한계가 있었다. 이 논문은 분포 매칭 증류를 연속 시간 영역으로 확장하여, 추가적인 GAN이나 보상 모델 없이도 단 몇 번의 추론만으로 고품질의 이미지를 생성할 수 있는 새로운 패러다임을 제시한다.

확산 모델의 추론 속도를 높이기 위한 기존의 증류 기법들은 주로 교사 모델이 밟았던 특정 시간 지점(Discrete Timesteps)들을 학생 모델이 그대로 따라 하도록 학습시킨다. 하지만 이 방식은 학습 시 정해진 지점 사이의 정보를 놓치게 되어, 실제 추론 시 큰 보폭으로 이동할 때 궤적을 이탈하거나 이미지가 흐릿해지는 현상이 발생한다.

CDM은 이를 해결하기 위해 학습 과정을 '연속적인 시간' 개념으로 전환한다. 학생 모델이 특정 지점만이 아니라 전체 시간 축의 어느 지점에서든 교사 모델의 분포를 닮도록 학습하며, 특히 현재 예측한 속도(Velocity)를 바탕으로 다음 단계로 이동했을 때 발생할 수 있는 오차를 미리 예측하여 교정하는 원리를 적용한다.

결과적으로 학생 모델은 단순히 교사의 출력을 흉내 내는 수준을 넘어, 스스로 궤적 이탈을 수정할 수 있는 능력을 갖추게 된다. 이는 단 4번의 연산만으로도 수백 번 연산한 교사 모델에 필적하는 선명한 질감과 세밀한 디테일을 복원할 수 있게 만든다.

CDM은 기존 DMD(Distribution Matching Distillation)를 연속 시간 영역으로 확장한다. 먼저 Dynamic Time Schedule을 도입하여, 각 학습 반복마다 역방향 시뮬레이션 길이 N을 무작위로 샘플링(N ~ U{1, Nmax})하고 연속적인 시간 시퀀스를 생성한다. 이를 통해 학생 모델은 고정된 앵커에 과적합되지 않고 전체 시간 도메인에서 매끄러운 속도 필드를 학습한다.

핵심 손실 함수인 LCDM은 궤적 외 정렬(Off-trajectory Alignment)을 수행한다. 현재 잠재 변수 xti와 예측 속도 vti가 주어질 때, xt'i = xti + (t'i - ti)vti 연산을 통해 1차 오일러 외삽을 수행하여 새로운 지점 xt'i를 얻는다. 여기서 [현재 위치와 속도 → 시간 변화량 곱셈 및 덧셈 → 외삽된 위치] 과정을 거쳐 실제 샘플링 시 발생할 오차 지점을 시뮬레이션한다.

이후 외삽된 지점 xt'i에서 학생 모델의 예측값 x^0을 구하고, 이를 다시 노이즈를 섞어(Re-noising) 교사 모델이 평가하게 함으로써 속도 필드의 재료 미분(Material Derivative)을 직접적으로 규제한다. 최종 손실 함수는 CFG Augmentation(LCA), Distribution Matching(LDM), 그리고 제안된 CDM 손실(LCDM)의 합으로 구성되어 텍스트 정렬과 이미지 품질을 동시에 최적화한다.

SD3-Medium 및 Longcat-Image 모델을 대상으로 한 실험에서 CDM은 단 4회의 함수 평가(NFE)만으로 기존 SOTA 모델들을 압도하는 성능을 보였다. SD3-Medium 기준 Aesthetic Score 6.075, HPSv3 9.561을 기록하며 DMD2(8.419)나 D-DMD(9.176)보다 월등히 높은 시각적 충실도를 증명했다.

특히 4-NFE CDM 학생 모델은 일부 지표(DPGBench, HPSv3)에서 100-NFE를 수행한 교사 모델의 성능을 추월하는 결과를 보였다. 이는 연속 시간 최적화가 교사의 출력을 단순히 복제하는 것을 넘어, 더 효율적이고 정교한 생성 경로를 찾아냈음을 시사한다.

Ablation Study 결과, 제안된 세 가지 손실 함수가 상호 보완적으로 작용함이 확인됐다. LCA는 구조적 정렬을, LDM과 LCDM은 각각 궤적 내외의 분포 정렬을 담당하며, 특히 LCDM이 제거될 경우 이미지의 세부 디테일과 선명도가 크게 저하되는 현상이 관찰됐다.

CDM은 DMD 프레임워크의 이산적 한계를 수학적으로 분석하고 이를 연속 시간 제어로 해결한다. 오일러 샘플링의 국소 절단 오차가 속도 필드의 재료 미분(Material Derivative)에 비례한다는 점에 착안하여, LCDM 손실이 이 미분 값을 직접적으로 억제하도록 설계되었다.

학습 시 학생 모델의 가중치 θ는 LCA, LDM, LCDM의 결합 그래디언트를 통해 업데이트되며, 가짜 교사(Fake Teacher) 역할을 하는 온라인 모델 ψ는 별도의 옵티마이저로 업데이트되어 학생 모델의 현재 분포를 추적한다. SD3-Medium 증류 시 16개의 A100 GPU에서 약 24시간의 학습으로 최적의 성능에 도달했다.

이론적으로 CDM은 Tweedie's formula를 통해 스코어 매칭과 KL 발산 최소화 사이의 연결 고리를 유지하면서도, 연속 시간 스케줄링을 통해 전체 시간 축에 걸친 암시적 분류기(Implicit Classifier) 그래디언트를 정규화하는 효과를 거둔다.

동적 연속 스케줄과 CDM 손실 도입으로 인해 학습 시 반복당 연산 비용이 기존 D-DMD 대비 약 1.8배 증가한다. 또한 증류 프레임워크 특성상 교사 모델의 성능이 상한선으로 작용하며, 교사가 제대로 처리하지 못하는 복잡한 구도는 학생 모델에서도 개선하기 어렵다.