CDMD는 데이터와 교사 롤아웃 없이 확산 모델을 소수 스텝으로 증류한다

무엇인가

이 논문은 확산 및 플로우 생성 모델의 추론이 수백 번의 신경망 함수 평가를 포함하는 반복 디노이징 때문에 느리다는 문제를 다룬다. 증류는 교사의 생성 동역학을 학생이 모방해 한두 스텝 생성을 가능하게 하지만, 기존 방법은 정제된 증류 데이터셋, 비싼 교사 롤아웃, 보조 fake score 네트워크, 적대적 목적, 다단계 학습에 의존해 왔다. 또한 실제 데이터를 노이징해 중간 상태를 만들면 그 데이터가 교사를 학습시킨 분포와 다를 수 있고, 교사가 파인튜닝이나 classifier-free guidance로 분포가 이동했거나 독점 데이터로 학습됐다면 감독 상태가 교사의 확률 경로에서 벗어난다. 교사 궤적을 직접 쓰려면 역방향 시뮬레이션이 필요해 계산 부담이 다시 커진다. 저자들은 외부 데이터, 교사 롤아웃, 보조 fake score 모델 없이 교사를 소수 스텝 학생으로 증류할 수 있는가를 질문하고 CDMD를 제안한다.

어떻게 동작하나

CDMD의 핵심은 샘플 생성과 점수(속도) 추정을 하나의 학생 네트워크로 통합하는 것이다. 학생은 임의의 구간 [r,t]에서 평균 속도 u_theta(x_t,t,r)를 예측한다. 이는 (x_t - x_r)/(t-r)로 정의되는 적분된 속도이며, r을 고정하고 t로 미분하면 순간 속도를 복원할 수 있다. 한 스텝 생성은 G(z)=z-u_theta(z,1,0)로 표현되고, 같은 네트워크가 v_fake(x_t,t)≈u_theta(x_t,t,r)+(t-r)d/dt u_theta(x_t,t,r)라는 속도 대리 모델을 제공한다. 이때 총미분은 교사 속도 방향으로 평가한 하나의 JVP로 계산한다. 저자들은 Fisher divergence를 속도 형태로 다시 쓰고, L_CDMD=E||u_theta(x_hat_t,t,r)+(t-r)sg[JVP]-v_teacher(x_hat_t,t)||^2를 최소화한다. 샘플 상태와 JVP 보정에는 stop-gradient를 걸어 고차 그래디언트를 막고 학습을 안정화한다. 결과적으로 별도 프록시 네트워크 없이 교사와 학생 두 모델, 하나의 목적함수만 쓴다.

무엇과 다른가

학습 중간 상태는 외부 데이터 없이 학생 자신의 경로에서 부트스트랩한다. 먼저 사전분포 z를 학생 플로우 맵으로 이른 시간 s까지 디노이징한 뒤, 닫힌 형태의 전방 노이징으로 t까지 다시 노이징해 x_hat_t를 만든다. 식은 x_hat_t=(1-t)/(1-s)f_theta(z,1,s)+(t-s)/(1-s)epsilon이며, s=t이면 학생 경로 위 상태, s=0이면 학생 한 스텝 샘플을 노이징한 상태가 된다. 이렇게 만든 분포를 학생 확률 경로 q_t^theta라 부르고, 학생이 교사에 가까워질수록 감독 상태도 교사 경로에 가까워진다. 논문의 정리 3.1은 정규성 가정 아래 목적함수가 0으로 가면 학생 플로우 맵 푸시포워드와 교사 주변분포 사이의 시간 평균 W2 제곱 거리가 0으로 수렴하고, 허용 가능한 학생이 교사 평균 속도를 표현할 수 있으면 모든 전역 최소화점이 교사 평균 속도를 복원한다고 주장한다. 이 정리는 감독 상태가 반드시 교사 경로에서 샘플링될 필요는 없다는 점을 강조한다. 또한 JVP 보정의 분산을 분석해 교사 속도 v(x_hat_t,t)를 접선 방향으로 쓸 때 합성 데이터에서 가장 낮고 안정적인 JVP 2차 모멘트를 얻는다고 보고한다.

어떻게 쓰나

실험은 2D 나선 합성 데이터, ImageNet 256×256, SANA-1.6B 텍스트-투-이미지로 구성된다. 2D 나선에서는 정확한 속도장을 교사로 두고 CDMD와 FreeFlow를 비교했는데, 두 방법은 목표까지의 엔드포인트 W2 거리를 비슷한 속도로 줄였다. 그러나 교사 평균 속도 MSE는 CDMD만 계속 줄었고 FreeFlow는 더 높은 값에서 정체했다. 저자들은 이것이 엔드포인트 주변분포 매칭이 교사의 수송 맵 학습을 의미하지 않는다는 증거라고 해석한다. ImageNet 256×256에서는 SiT-XL/2, VA-REPA로 학습한 SiT-XL/2, VA-VAE 오토인코더를 쓴 LightningDiT-XL/2를 각각 교사로 증류했다. CDMD는 세 교사 모두에서 1-NFE FID가 가장 좋았고, 각각 3.47, 2.93, 2.04였다. 더 강한 교사로 바꿀 때 CDMD는 FID가 0.54 개선됐지만 FreeFlow는 0.41, DMD2는 0.07 개선에 그쳤고 MeanFlow는 오히려 나빠졌다. 재학습 없이 같은 학생으로 2-NFE에서 VA-REPA 2.60, LightningDiT-XL/2 1.62, 4-NFE에서 1.37을 기록해 교사의 1.35에 근접했다. 데이터-프리 설정으로 적응시킨 원스텝 베이스라인들은 CDMD보다 나빴다.

전제와 한계

CDMD는 임의 구간의 평균 속도를 학습하므로 추론 시 스텝 수를 늘릴 수 있다. FreeFlow는 노이즈 t=1에 고정돼 부분 구간에 적용할 수 없지만, CDMD는 여러 구간에 같은 맵을 적용해 few-step 샘플링을 한다. 실험에서 guidance w=1.5일 때는 스텝이 늘수록 FID가 계속 개선됐고, w=2.0일 때는 매우 적은 스텝에서 FID가 가장 좋고 스텝이 늘면 나빠졌다. 저자들은 강한 guidance가 원스텝 생성에는 도움이 되지만 다단계 샘플링에서 학생-교사 불일치를 증폭한다고 설명한다. 학습 상태 보간 비율 m은 중간값이 좋아서 FID는 m=0.5에서 가장 낮았고 IS는 m=0.875에서 가장 높았으며 수렴 속도는 비슷했다. t=r 쌍의 비율은 10%일 때 1-NFE FID가 좋아지지만 다단계 샘플링이 나빠졌고, 50%일 때 few-step 성능이 더 안정적이었다. 텍스트-투-이미지에서는 SANA-1.6B를 LAION-Aesthetics 6.5+에서 MFD 프로토콜로 증류했고, 학습된 4-NFE 방법 중 Aesthetic Score, PickScore, ImageReward가 가장 좋았으며 FID는 27.83으로 MFD의 28.77보다 낮았다. CDMD는 4-NFE 교사보다 여섯 지표 모두에서 좋았고 20-NFE 교사보다 Aesthetic Score, ImageReward, FID에서 앞섰지만 HPSv2와 CLIPScore에서는 뒤졌다. 저자들은 이 결과가 통계적 우월성을 입증하는 것은 아니라고 명시한다. 학습 비용은 약 7 GPU 시간으로 Diff-Instruct 22, MFD 25, DMD2 30보다 낮았고, consistency distillation은 약 4 GPU 시간으로 더 저렴하지만 다섯 포인트 지표에서 모두 낮았다.

개발자 관점에서 CDMD는 이미 강한 확산 또는 플로우 교사 체크포인트가 있고, 원 학습 데이터에 접근하기 어렵거나 교사 롤아웃과 보조 점수 네트워크 학습 비용을 줄이고 싶을 때 유용하다. 두 모델과 하나의 옵티마이저로 끝나며 적대적 학습이 없고, 학생이 임의 구간을 다루므로 지연 시간과 품질 사이에서 NFE를 조절할 수 있다. 다만 실제 도입 전에는 교사 품질, guidance 스케일, NFE와 샘플러 조합, JVP 항의 안정성, 학생 초기화 전략을 확인해야 한다. 논문의 FID와 선호도 지표는 연구 벤치마크이므로 도메인 데이터와 라이선스, 독점 데이터 접근 제약, 프로덕션 컴퓨팅 예산에 맞춰 별도 검증이 필요하다.

저자들이 밝힌 한계와 전제도 분명하다. 이론은 정규성, 커버리지, 적분 가능성 가정과 교사 평균 속도를 표현할 수 있는 학생이라는 조건에 기대며, 학생 부트스트랩 상태는 학생 경로가 교사 경로에 가까울 때 원칙적인 대체물이 된다. JVP 보정은 접선 방향에 민감하고, 강한 guidance는 다단계 샘플링에서 학생-교사 불일치를 증폭할 수 있다. 텍스트-투-이미지 결과는 HPSv2와 CLIPScore에서 뒤지고 통계적 우월성을 주장하지 않는다. 향후 과제로 JVP 항의 안정성 개선, 모델 가중치가 없을 때의 초기화 전략, 데이터 희소성이 큰 대규모 생성 과제로의 확장을 제시한다.