Simplificando, Estabilizando y Escalando Modelos de Consistencia en Tiempo Continuo.

Simplifying, Stabilizing and Scaling Continuous-Time Consistency Models

October 14, 2024

Autores: Cheng Lu, Yang Song

cs.AI

Resumen

Los modelos de consistencia (CMs) son una potente clase de modelos generativos basados en difusión optimizados para un muestreo rápido. La mayoría de los CMs existentes se entrenan utilizando pasos de tiempo discretizados, lo que introduce hiperparámetros adicionales y es propenso a errores de discretización. Si bien las formulaciones en tiempo continuo pueden mitigar estos problemas, su éxito ha sido limitado por la inestabilidad en el entrenamiento. Para abordar esto, proponemos un marco teórico simplificado que unifica las parametrizaciones previas de modelos de difusión y CMs, identificando las causas fundamentales de la inestabilidad. Basándonos en este análisis, introducimos mejoras clave en la parametrización del proceso de difusión, la arquitectura de la red y los objetivos de entrenamiento. Estos cambios nos permiten entrenar CMs en tiempo continuo a una escala sin precedentes, alcanzando 1.5B de parámetros en ImageNet 512x512. Nuestro algoritmo de entrenamiento propuesto, utilizando solo dos pasos de muestreo, logra puntuaciones FID de 2.06 en CIFAR-10, 1.48 en ImageNet 64x64 y 1.88 en ImageNet 512x512, reduciendo la brecha en las puntuaciones FID con los mejores modelos de difusión existentes a menos del 10%.

English

Consistency models (CMs) are a powerful class of diffusion-based generative models optimized for fast sampling. Most existing CMs are trained using discretized timesteps, which introduce additional hyperparameters and are prone to discretization errors. While continuous-time formulations can mitigate these issues, their success has been limited by training instability. To address this, we propose a simplified theoretical framework that unifies previous parameterizations of diffusion models and CMs, identifying the root causes of instability. Based on this analysis, we introduce key improvements in diffusion process parameterization, network architecture, and training objectives. These changes enable us to train continuous-time CMs at an unprecedented scale, reaching 1.5B parameters on ImageNet 512x512. Our proposed training algorithm, using only two sampling steps, achieves FID scores of 2.06 on CIFAR-10, 1.48 on ImageNet 64x64, and 1.88 on ImageNet 512x512, narrowing the gap in FID scores with the best existing diffusion models to within 10%.