Улучшенное сопоставление распределений для быстрой синтеза изображений

Improved Distribution Matching Distillation for Fast Image Synthesis

May 23, 2024

Авторы: Tianwei Yin, Michaël Gharbi, Taesung Park, Richard Zhang, Eli Shechtman, Fredo Durand, William T. Freeman

cs.AI

Аннотация

Недавние подходы показали перспективы сжатия моделей диффузии в эффективные генераторы одного шага. Среди них Distribution Matching Distillation (DMD) создает генераторы одного шага, соответствующие своему учителю по распределению, не требуя однозначного соответствия с траекториями выборки своих учителей. Однако для обеспечения стабильного обучения DMD требует дополнительной потери регрессии, вычисленной с использованием большого набора пар шум-изображение, сгенерированных учителем с помощью множества шагов детерминированного сэмплера. Это затратно для синтеза текста в изображение в крупном масштабе и ограничивает качество студента, слишком тесно связывая его с исходными траекториями выборки учителя. Мы представляем DMD2, набор техник, которые преодолевают это ограничение и улучшают обучение DMD. Во-первых, мы устраняем потерю регрессии и необходимость в дорогостоящем построении набора данных. Мы показываем, что возникшая нестабильность обусловлена тем, что фальшивый критик недостоверно оценивает распределение сгенерированных образцов и предлагаем правило обновления на двух временных шкалах в качестве лекарства. Во-вторых, мы интегрируем потерю GAN в процедуру дистилляции, различая сгенерированные образцы и реальные изображения. Это позволяет обучать модель студента на реальных данных, смягчая недостаточную оценку реального балла от модели учителя и улучшая качество. Наконец, мы модифицируем процедуру обучения для возможности многократной выборки. Мы выявляем и решаем проблему несоответствия ввода обучения и вывода в этой ситуации, симулируя образцы генератора во время вывода во время обучения. В целом, наши улучшения устанавливают новые показатели в генерации изображений одного шага, с оценками FID 1,28 на ImageNet-64x64 и 8,35 на zero-shot COCO 2014, превосходя исходного учителя несмотря на 500-кратное снижение стоимости вывода. Кроме того, мы показываем, что наш подход может генерировать мегапиксельные изображения, дистиллируя SDXL, демонстрируя исключительное визуальное качество среди методов с небольшим количеством шагов.

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.