Simplification, Stabilisation et Mise à l'Échelle des Modèles de Cohérence en Temps Continu

Simplifying, Stabilizing and Scaling Continuous-Time Consistency Models

October 14, 2024

Auteurs: Cheng Lu, Yang Song

cs.AI

Résumé

Les modèles de cohérence (CMs) constituent une classe puissante de modèles génératifs basés sur la diffusion, optimisés pour un échantillonnage rapide. La plupart des CMs existants sont entraînés en utilisant des pas de temps discrétisés, ce qui introduit des hyperparamètres supplémentaires et est sujet aux erreurs de discrétisation. Alors que les formulations en temps continu peuvent atténuer ces problèmes, leur succès a été limité en raison de l'instabilité de l'entraînement. Pour y remédier, nous proposons un cadre théorique simplifié qui unifie les paramétrisations précédentes des modèles de diffusion et des CMs, identifiant les causes fondamentales de l'instabilité. Sur la base de cette analyse, nous introduisons des améliorations clés dans la paramétrisation du processus de diffusion, l'architecture du réseau et les objectifs d'entraînement. Ces changements nous permettent d'entraîner des CMs en temps continu à une échelle sans précédent, atteignant 1,5 milliard de paramètres sur ImageNet 512x512. Notre algorithme d'entraînement proposé, n'utilisant que deux étapes d'échantillonnage, atteint des scores FID de 2,06 sur CIFAR-10, 1,48 sur ImageNet 64x64 et 1,88 sur ImageNet 512x512, réduisant l'écart des scores FID avec les meilleurs modèles de diffusion existants à moins de 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%.