Distilação Aprimorada de Correspondência de Distribuição para Síntese Rápida de Imagens

Improved Distribution Matching Distillation for Fast Image Synthesis

May 23, 2024

Autores: Tianwei Yin, Michaël Gharbi, Taesung Park, Richard Zhang, Eli Shechtman, Fredo Durand, William T. Freeman

cs.AI

Resumo

Abordagens recentes têm mostrado promessas ao destilar modelos de difusão em geradores eficientes de uma única etapa. Entre elas, a Distilação por Correspondência de Distribuição (DMD) produz geradores de uma etapa que correspondem à distribuição de seus modelos professores, sem impor uma correspondência um-para-um com as trajetórias de amostragem de seus professores. No entanto, para garantir um treinamento estável, o DMD requer uma perda de regressão adicional calculada usando um grande conjunto de pares ruído-imagem gerados pelo professor com muitas etapas de um amostrador determinístico. Isso é custoso para a síntese de texto para imagem em larga escala e limita a qualidade do modelo estudante, vinculando-o muito de perto aos caminhos de amostragem originais do professor. Introduzimos o DMD2, um conjunto de técnicas que supera essa limitação e melhora o treinamento do DMD. Primeiro, eliminamos a perda de regressão e a necessidade de construção de conjuntos de dados caros. Mostramos que a instabilidade resultante se deve ao crítico falso não estimar com precisão a distribuição das amostras geradas e propomos uma regra de atualização em duas escalas de tempo como solução. Segundo, integramos uma perda GAN no procedimento de destilação, discriminando entre amostras geradas e imagens reais. Isso nos permite treinar o modelo estudante com dados reais, mitigando a estimativa imperfeita da pontuação real do modelo professor e melhorando a qualidade. Por último, modificamos o procedimento de treinamento para permitir amostragem em múltiplas etapas. Identificamos e resolvemos o problema de incompatibilidade de entrada entre treinamento e inferência nesse cenário, simulando amostras do gerador durante o tempo de treinamento. Juntas, nossas melhorias estabelecem novos benchmarks na geração de imagens de uma única etapa, com pontuações FID de 1,28 no ImageNet-64x64 e 8,35 no COCO 2014 zero-shot, superando o professor original apesar de uma redução de 500X no custo de inferência. Além disso, mostramos que nossa abordagem pode gerar imagens de megapixels ao destilar o SDXL, demonstrando qualidade visual excepcional entre métodos de poucas etapas.

English

Recent approaches have shown promises distilling diffusion models into efficient one-step generators. Among them, Distribution Matching Distillation (DMD) produces one-step generators that match their teacher in distribution, without enforcing a one-to-one correspondence with the sampling trajectories of their teachers. However, to ensure stable training, DMD requires an additional regression loss computed using a large set of noise-image pairs generated by the teacher with many steps of a deterministic sampler. This is costly for large-scale text-to-image synthesis and limits the student's quality, tying it too closely to the teacher's original sampling paths. We introduce DMD2, a set of techniques that lift this limitation and improve DMD training. First, we eliminate the regression loss and the need for expensive dataset construction. We show that the resulting instability is due to the fake critic not estimating the distribution of generated samples accurately and propose a two time-scale update rule as a remedy. Second, we integrate a GAN loss into the distillation procedure, discriminating between generated samples and real images. This lets us train the student model on real data, mitigating the imperfect real score estimation from the teacher model, and enhancing quality. Lastly, we modify the training procedure to enable multi-step sampling. We identify and address the training-inference input mismatch problem in this setting, by simulating inference-time generator samples during training time. Taken together, our improvements set new benchmarks in one-step image generation, with FID scores of 1.28 on ImageNet-64x64 and 8.35 on zero-shot COCO 2014, surpassing the original teacher despite a 500X reduction in inference cost. Further, we show our approach can generate megapixel images by distilling SDXL, demonstrating exceptional visual quality among few-step methods.