Modelos de Consistencia de Múltiples Pasos

Multistep Consistency Models

March 11, 2024

Autores: Jonathan Heek, Emiel Hoogeboom, Tim Salimans

cs.AI

Resumen

Los modelos de difusión son relativamente fáciles de entrenar, pero requieren muchos pasos para generar muestras. Los modelos de consistencia son mucho más difíciles de entrenar, pero generan muestras en un solo paso. En este artículo proponemos los Modelos de Consistencia Multietapa: una unificación entre los Modelos de Consistencia (Song et al., 2023) y TRACT (Berthelot et al., 2023) que puede interpolar entre un modelo de consistencia y un modelo de difusión: un equilibrio entre la velocidad de muestreo y la calidad de las muestras. Específicamente, un modelo de consistencia de 1 paso es un modelo de consistencia convencional, mientras que demostramos que un modelo de consistencia de ∞ pasos es un modelo de difusión. Los Modelos de Consistencia Multietapa funcionan muy bien en la práctica. Al aumentar el presupuesto de muestras de un solo paso a 2-8 pasos, podemos entrenar modelos más fácilmente que generan muestras de mayor calidad, conservando gran parte de los beneficios en velocidad de muestreo. Resultados destacados incluyen un FID de 1.4 en Imagenet 64 en 8 pasos y un FID de 2.1 en Imagenet 128 en 8 pasos con destilación de consistencia. También demostramos que nuestro método se escala a un modelo de difusión de texto a imagen, generando muestras muy cercanas en calidad al modelo original.

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.