분포형 디퓨전 모델을 소수 스텝 이미지 생성에 실용화한다

무엇인가

이 논문은 확산 모델의 디노이저를 어떻게 학습할 것인가라는 문제를 다룬다. 표준 확산 모델은 조건부 평균을 예측하는 mean-prediction 디노이저를 쓰지만, Distributional Diffusion Models(DDM)은 이를 분포형(distributional) 디노이저로 바꾸고 scoring rule 목적함수로 학습한다. 즉 p(x1 | xt)의 조건부 평균이 아니라 그 분포 자체에 대한 확률적 근사를 학습하는 것이다. 문제는 이 접근을 현대적인 이미지 생성 규모로 확장할 때 두 가지 장애물이 생긴다는 점이다. 첫째, 멀티파티클 학습은 파티클 수에 비례해 오버헤드가 커진다. 둘째, DDM은 scoring rule 하이퍼파라미터를 전역적으로 고정하기 때문에 샘플링 예산 전반에 걸쳐 단 하나의 트레이드오프만 강제된다.

어떻게 동작하나

제안 방법은 이 두 한계를 각각 다른 방식으로 완화한다. 파티클 확장을 트랜스포머의 늦은 레이어로 미루어(deferring particle expansion to late transformer layers) 멀티파티클 학습의 부담을 줄이고, 하이퍼파라미터 트레이드오프 문제는 시간 의존적(time-dependent) scoring rule 스케줄을 도입해 해결한다. 이 스케줄은 Biroli 등(2024)이 제시한 동역학적 영역(dynamical regimes)에서 착안한 것이다. 여기에 DiT 기반의 latent 셋업을 결합한 것이 전체 레시피의 골격이다.

무엇과 다른가

결과적으로 이 변경들은 DDM 학습을 클래스 조건부 ImageNet-256² 규모에서 실용적으로 만든다. DiT-XL/2 기준으로 4스텝에서 4.48 FID, 50스텝에서 2.38 FID를 기록했고, 이는 teacher도 self-distillation도 JVP도 없이 한 단계(one stage)에서 처음부터 학습한 단일 모델에서 나온 수치다. 즉 별도의 증류 파이프라인이나 추가 모델 없이 하나의 모델을 학습해 얻은 결과라는 점이 핵심이다.

어떻게 쓰나

특히 주목할 부분은 샘플링 예산에 대한 내성이다. 이 모델은 확률적(stochastic) few-step 생성기이며, 샘플링 예산이 4 NFE에서 50 NFE로 늘어나도 FID가 나빠지지 않는다. 일반적으로 스텝 수를 바꾸면 품질과 속도의 트레이드오프가 다시 조정되어야 하지만, 여기서는 그렇지 않다는 뜻이다. 같은 레시피는 텍스트-투-이미지 생성으로도 전이된다고 저자들은 밝힌다.

전제와 한계

개발자 입장에서 이 논문이 의미하는 바는 추론 예산을 유연하게 다룰 수 있는 생성기라는 점이다. 서비스 환경에서 지연 시간이 중요한 요청은 4스텝으로, 품질이 중요한 요청은 50스텝으로 같은 모델을 그대로 쓸 수 있다는 주장이기 때문이다. DiT 기반 latent 확산 파이프라인을 이미 쓰고 있다면 구조 변경 없이 디노이저 학습 방식만 바꾸는 형태로 접근할 수 있다. 다만 실제 도입 전에는 파티클 수와 시간 의존 스케줄 하이퍼파라미터가 자신의 데이터 규모에서 어떤 비용을 만드는지, 그리고 공개된 코드와 사전학습 모델이 자신의 해상도·도메인에 맞는지 확인해야 한다.

한계는 저자들이 명시한 전제에서 드러난다. 애초에 이 연구가 겨냥한 문제 자체가 멀티파티클 학습 오버헤드와 전역 고정 하이퍼파라미터라는 두 제약이며, 논문은 이를 제거했다고 주장하는 것이지 완전히 새로운 학습 패러다임을 제안하는 것은 아니다. 또한 초록이 근거로 제시하는 정량 결과는 클래스 조건부 ImageNet-256²와 DiT-XL/2 조합에 집중되어 있고, 텍스트-투-이미지는 같은 레시피가 전이된다는 수준으로 언급된다. 그 밖의 데이터셋이나 모델 크기에서의 거동은 초록만으로는 확인할 수 없다.