Modèles de Cohérence Multi-étapes

Multistep Consistency Models

March 11, 2024

Auteurs: Jonathan Heek, Emiel Hoogeboom, Tim Salimans

cs.AI

Résumé

Les modèles de diffusion sont relativement faciles à entraîner mais nécessitent de nombreuses étapes pour générer des échantillons. Les modèles de cohérence sont bien plus difficiles à entraîner, mais génèrent des échantillons en une seule étape. Dans cet article, nous proposons les Modèles de Cohérence Multi-étapes : une unification entre les Modèles de Cohérence (Song et al., 2023) et TRACT (Berthelot et al., 2023) qui permet d'interpoler entre un modèle de cohérence et un modèle de diffusion : un compromis entre la vitesse d'échantillonnage et la qualité des échantillons. Plus précisément, un modèle de cohérence à 1 étape est un modèle de cohérence conventionnel, tandis que nous montrons qu'un modèle de cohérence à ∞ étapes est un modèle de diffusion. Les Modèles de Cohérence Multi-étapes fonctionnent très bien en pratique. En augmentant le budget d'échantillonnage d'une seule étape à 2-8 étapes, nous pouvons entraîner des modèles plus facilement qui génèrent des échantillons de meilleure qualité, tout en conservant une grande partie des avantages en termes de vitesse d'échantillonnage. Parmi les résultats notables, on trouve un FID de 1,4 sur Imagenet 64 en 8 étapes et un FID de 2,1 sur Imagenet128 en 8 étapes avec la distillation de cohérence. Nous montrons également que notre méthode s'adapte à un modèle de diffusion texte-image, générant des échantillons très proches de la qualité du modèle 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.