f-발산 분포 매칭을 활용한 단일 단계 확산 모델

One-step Diffusion Models with f-Divergence Distribution Matching

February 21, 2025

저자: Yilun Xu, Weili Nie, Arash Vahdat

cs.AI

초록

확산 모델에서의 샘플링은 느린 반복적 과정을 수반하며, 이는 특히 상호작용형 애플리케이션에서의 실질적인 배포를 방해합니다. 생성 속도를 가속화하기 위해, 최근 접근법들은 다단계 확산 모델을 변분 점수 증류(variational score distillation)를 통해 단일 단계 학생 생성기로 증류하여, 학생이 생성한 샘플의 분포를 교사 모델의 분포와 일치시키는 방법을 사용합니다. 그러나 이러한 접근법들은 모드 탐색(mode seeking) 성향이 있는 역방향 쿨백-라이블러(Kullback-Leibler, KL) 발산을 분포 매칭에 사용합니다. 본 논문에서는 모드 커버리지와 학습 분산 간의 다양한 절충점을 제공하는 새로운 f-발산 최소화 프레임워크인 f-distill을 통해 분포 매칭 접근법을 일반화합니다. 우리는 교사와 학생 분포 간의 f-발산의 그래디언트를 유도하고, 이가 그들의 점수 차이와 밀도 비율에 의해 결정되는 가중치 함수의 곱으로 표현됨을 보입니다. 이 가중치 함수는 덜 모드 탐색적인 발산을 사용할 때, 교사 분포에서 더 높은 밀도를 가진 샘플을 자연스럽게 강조합니다. 우리는 역방향 KL 발산을 사용한 인기 있는 변분 점수 증류 접근법이 우리 프레임워크 내의 특수한 경우임을 관찰합니다. 실험적으로, 우리는 순방향 KL 및 젠센-섀넌(Jensen-Shannon) 발산과 같은 대안적 f-발산들이 이미지 생성 작업에서 현재 최고의 변분 점수 증류 방법들을 능가함을 입증합니다. 특히, 젠센-섀넌 발산을 사용할 때, f-distill은 ImageNet64에서의 단일 단계 생성 성능과 MS-COCO에서의 제로샷 텍스트-이미지 생성에서 현재 최첨단 성능을 달성합니다. 프로젝트 페이지: https://research.nvidia.com/labs/genair/f-distill

English

Sampling from diffusion models involves a slow iterative process that hinders their practical deployment, especially for interactive applications. To accelerate generation speed, recent approaches distill a multi-step diffusion model into a single-step student generator via variational score distillation, which matches the distribution of samples generated by the student to the teacher's distribution. However, these approaches use the reverse Kullback-Leibler (KL) divergence for distribution matching which is known to be mode seeking. In this paper, we generalize the distribution matching approach using a novel f-divergence minimization framework, termed f-distill, that covers different divergences with different trade-offs in terms of mode coverage and training variance. We derive the gradient of the f-divergence between the teacher and student distributions and show that it is expressed as the product of their score differences and a weighting function determined by their density ratio. This weighting function naturally emphasizes samples with higher density in the teacher distribution, when using a less mode-seeking divergence. We observe that the popular variational score distillation approach using the reverse-KL divergence is a special case within our framework. Empirically, we demonstrate that alternative f-divergences, such as forward-KL and Jensen-Shannon divergences, outperform the current best variational score distillation methods across image generation tasks. In particular, when using Jensen-Shannon divergence, f-distill achieves current state-of-the-art one-step generation performance on ImageNet64 and zero-shot text-to-image generation on MS-COCO. Project page: https://research.nvidia.com/labs/genair/f-distill