Ajustement de la stabilité de la cohérence : Compréhension et amélioration des modèles de cohérence

Stable Consistency Tuning: Understanding and Improving Consistency Models

October 24, 2024

Auteurs: Fu-Yun Wang, Zhengyang Geng, Hongsheng Li

cs.AI

Résumé

Les modèles de diffusion parviennent à une qualité de génération supérieure mais souffrent d'une lenteur de génération en raison de la nature itérative du débruitage. En revanche, les modèles de cohérence, une nouvelle famille générative, obtiennent des performances compétitives avec une vitesse d'échantillonnage significativement plus rapide. Ces modèles sont entraînés soit par distillation de cohérence, qui exploite des modèles de diffusion pré-entraînés, soit par entraînement/ajustement de cohérence directement à partir des données brutes. Dans ce travail, nous proposons un nouveau cadre pour comprendre les modèles de cohérence en modélisant le processus de débruitage du modèle de diffusion comme un Processus de Décision Markovien (MDP) et en encadrant l'entraînement du modèle de cohérence comme l'estimation de la valeur par l'Apprentissage Temporel de Différence (TD). Plus important encore, ce cadre nous permet d'analyser les limitations des stratégies actuelles d'entraînement/ajustement de cohérence. Basé sur l'Accord de Cohérence Facile (ECT), nous proposons l'Accord de Cohérence Stable (SCT), qui intègre un apprentissage à variance réduite en utilisant l'identité de score. SCT conduit à des améliorations significatives des performances sur des références telles que CIFAR-10 et ImageNet-64. Sur ImageNet-64, SCT atteint un FID en 1 étape de 2,42 et un FID en 2 étapes de 1,55, un nouveau record pour les modèles de cohérence.

English

Diffusion models achieve superior generation quality but suffer from slow generation speed due to the iterative nature of denoising. In contrast, consistency models, a new generative family, achieve competitive performance with significantly faster sampling. These models are trained either through consistency distillation, which leverages pretrained diffusion models, or consistency training/tuning directly from raw data. In this work, we propose a novel framework for understanding consistency models by modeling the denoising process of the diffusion model as a Markov Decision Process (MDP) and framing consistency model training as the value estimation through Temporal Difference~(TD) Learning. More importantly, this framework allows us to analyze the limitations of current consistency training/tuning strategies. Built upon Easy Consistency Tuning (ECT), we propose Stable Consistency Tuning (SCT), which incorporates variance-reduced learning using the score identity. SCT leads to significant performance improvements on benchmarks such as CIFAR-10 and ImageNet-64. On ImageNet-64, SCT achieves 1-step FID 2.42 and 2-step FID 1.55, a new SoTA for consistency models.