Модели согласованности с несколькими шагами

Multistep Consistency Models

March 11, 2024

Авторы: Jonathan Heek, Emiel Hoogeboom, Tim Salimans

cs.AI

Аннотация

Модели диффузии относительно легки в обучении, но требуют много шагов для генерации образцов. Модели согласованности гораздо сложнее в обучении, но генерируют образцы за один шаг. В данной статье мы предлагаем Многошаговые модели согласованности: объединение Моделей согласованности (Сонг и др., 2023) и TRACT (Бертелот и др., 2023), которые могут интерполировать между моделью согласованности и моделью диффузии: компромисс между скоростью генерации образцов и их качеством. Конкретно, модель согласованности на 1 шаге является обычной моделью согласованности, в то время как мы показываем, что модель согласованности на бесконечном числе шагов является моделью диффузии. Многошаговые модели согласованности работают очень хорошо на практике. Увеличивая бюджет образцов с одного шага до 2-8 шагов, мы можем более легко обучать модели, которые генерируют образцы более высокого качества, сохраняя при этом большую часть преимуществ в скорости генерации образцов. Замечательные результаты: 1,4 FID на Imagenet 64 на 8 шаге и 2,1 FID на Imagenet128 на 8 шагах с дистилляцией согласованности. Мы также показываем, что наш метод масштабируется на модель диффузии текста в изображение, генерируя образцы, близкие по качеству к оригинальной модели.

English

Diffusion models are relatively easy to train but require many steps to generate samples. Consistency models are far more difficult to train, but generate samples in a single step. In this paper we propose Multistep Consistency Models: A unification between Consistency Models (Song et al., 2023) and TRACT (Berthelot et al., 2023) that can interpolate between a consistency model and a diffusion model: a trade-off between sampling speed and sampling quality. Specifically, a 1-step consistency model is a conventional consistency model whereas we show that a infty-step consistency model is a diffusion model. Multistep Consistency Models work really well in practice. By increasing the sample budget from a single step to 2-8 steps, we can train models more easily that generate higher quality samples, while retaining much of the sampling speed benefits. Notable results are 1.4 FID on Imagenet 64 in 8 step and 2.1 FID on Imagenet128 in 8 steps with consistency distillation. We also show that our method scales to a text-to-image diffusion model, generating samples that are very close to the quality of the original model.