Modelos de Difusão em Um Passo com Correspondência de Distribuição via f-Divergência

One-step Diffusion Models with f-Divergence Distribution Matching

February 21, 2025

Autores: Yilun Xu, Weili Nie, Arash Vahdat

cs.AI

Resumo

A amostragem de modelos de difusão envolve um processo iterativo lento que dificulta sua implantação prática, especialmente para aplicações interativas. Para acelerar a velocidade de geração, abordagens recentes destilam um modelo de difusão de múltiplos passos em um gerador estudante de passo único por meio de destilação de pontuação variacional, que iguala a distribuição das amostras geradas pelo estudante à distribuição do professor. No entanto, essas abordagens utilizam a divergência reversa de Kullback-Leibler (KL), que é conhecida por buscar modos. Neste artigo, generalizamos a abordagem de correspondência de distribuição usando uma nova estrutura de minimização de f-divergência, denominada f-distill, que abrange diferentes divergências com diferentes compensações em termos de cobertura de modos e variância de treinamento. Derivamos o gradiente da f-divergência entre as distribuições do professor e do estudante e mostramos que ele é expresso como o produto de suas diferenças de pontuação e uma função de ponderação determinada pela razão de suas densidades. Essa função de ponderação naturalmente enfatiza amostras com maior densidade na distribuição do professor ao usar uma divergência menos propensa a buscar modos. Observamos que a popular abordagem de destilação de pontuação variacional usando a divergência reversa de KL é um caso especial dentro de nossa estrutura. Empiricamente, demonstramos que f-divergências alternativas, como as divergências de KL direta e de Jensen-Shannon, superam os melhores métodos atuais de destilação de pontuação variacional em tarefas de geração de imagens. Em particular, ao usar a divergência de Jensen-Shannon, o f-distill alcança o melhor desempenho atual em geração de uma etapa no ImageNet64 e em geração de texto para imagem zero-shot no MS-COCO. Página do projeto: 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