Дистилляция EM для моделей диффузии одного шага

EM Distillation for One-step Diffusion Models

May 27, 2024

Авторы: Sirui Xie, Zhisheng Xiao, Diederik P Kingma, Tingbo Hou, Ying Nian Wu, Kevin Patrick Murphy, Tim Salimans, Ben Poole, Ruiqi Gao

cs.AI

Аннотация

Хотя модели диффузии могут изучать сложные распределения, для выборки требуется вычислительно затратный итерационный процесс. Существующие методы дистилляции обеспечивают эффективную выборку, но имеют заметные ограничения, такие как деградация производительности при очень небольшом количестве шагов выборки, зависимость от доступа к обучающим данным или оптимизация поиска моды, которая может не захватить полное распределение. Мы предлагаем EM Дистилляцию (EMD), подход на основе максимального правдоподобия, который дистиллирует модель диффузии до модели генератора с одним шагом с минимальной потерей качества восприятия. Наш подход разработан через призму ожидания-максимизации (EM), где параметры генератора обновляются с использованием выборок из совместного распределения априори учителя диффузии и выведенных латентов генератора. Мы разрабатываем перепараметризованную схему выборки и технику отмены шума, которые вместе стабилизируют процесс дистилляции. Мы также раскрываем интересную связь нашего метода с существующими методами, минимизирующими KL по поиску моды. EMD превосходит существующие одношаговые генеративные методы по показателям FID на ImageNet-64 и ImageNet-128, а также сравнительно с предыдущими работами по дистилляции моделей диффузии текста в изображение.

English

While diffusion models can learn complex distributions, sampling requires a computationally expensive iterative process. Existing distillation methods enable efficient sampling, but have notable limitations, such as performance degradation with very few sampling steps, reliance on training data access, or mode-seeking optimization that may fail to capture the full distribution. We propose EM Distillation (EMD), a maximum likelihood-based approach that distills a diffusion model to a one-step generator model with minimal loss of perceptual quality. Our approach is derived through the lens of Expectation-Maximization (EM), where the generator parameters are updated using samples from the joint distribution of the diffusion teacher prior and inferred generator latents. We develop a reparametrized sampling scheme and a noise cancellation technique that together stabilizes the distillation process. We further reveal an interesting connection of our method with existing methods that minimize mode-seeking KL. EMD outperforms existing one-step generative methods in terms of FID scores on ImageNet-64 and ImageNet-128, and compares favorably with prior work on distilling text-to-image diffusion models.