Technique d'Entraînement Améliorée pour les Modèles de Cohérence Latente

Improved Training Technique for Latent Consistency Models

February 3, 2025

Auteurs: Quan Dao, Khanh Doan, Di Liu, Trung Le, Dimitris Metaxas

cs.AI

Résumé

Les modèles de cohérence sont une nouvelle famille de modèles génératifs capables de produire des échantillons de haute qualité en une seule étape ou en plusieurs étapes. Récemment, les modèles de cohérence ont démontré des performances impressionnantes, atteignant des résultats comparables à ceux des modèles de diffusion dans l'espace pixel. Cependant, le succès de la mise à l'échelle de l'entraînement à la cohérence sur des ensembles de données à grande échelle, en particulier pour les tâches de génération de texte en image et de vidéo, est déterminé par les performances dans l'espace latent. Dans ce travail, nous analysons les différences statistiques entre les espaces pixel et latent, découvrant que les données latentes contiennent souvent des valeurs aberrantes très impulsives, qui dégradent significativement les performances de l'iCT dans l'espace latent. Pour remédier à cela, nous remplaçons les pertes Pseudo-Huber par des pertes Cauchy, atténuant efficacement l'impact des valeurs aberrantes. De plus, nous introduisons une perte de diffusion aux premiers pas de temps et utilisons le couplage de transport optimal (OT) pour améliorer encore les performances. Enfin, nous introduisons le planificateur d'échelle adaptative pour gérer le processus d'entraînement robuste et adoptons la normalisation des couches non scalables dans l'architecture pour mieux capturer les statistiques des caractéristiques et réduire l'impact des valeurs aberrantes. Grâce à ces stratégies, nous entraînons avec succès des modèles de cohérence latente capables de produire des échantillons de haute qualité en une ou deux étapes, réduisant significativement l'écart de performance entre la cohérence latente et les modèles de diffusion. L'implémentation est disponible ici : https://github.com/quandao10/sLCT/

English

Consistency models are a new family of generative models capable of producing high-quality samples in either a single step or multiple steps. Recently, consistency models have demonstrated impressive performance, achieving results on par with diffusion models in the pixel space. However, the success of scaling consistency training to large-scale datasets, particularly for text-to-image and video generation tasks, is determined by performance in the latent space. In this work, we analyze the statistical differences between pixel and latent spaces, discovering that latent data often contains highly impulsive outliers, which significantly degrade the performance of iCT in the latent space. To address this, we replace Pseudo-Huber losses with Cauchy losses, effectively mitigating the impact of outliers. Additionally, we introduce a diffusion loss at early timesteps and employ optimal transport (OT) coupling to further enhance performance. Lastly, we introduce the adaptive scaling-c scheduler to manage the robust training process and adopt Non-scaling LayerNorm in the architecture to better capture the statistics of the features and reduce outlier impact. With these strategies, we successfully train latent consistency models capable of high-quality sampling with one or two steps, significantly narrowing the performance gap between latent consistency and diffusion models. The implementation is released here: https://github.com/quandao10/sLCT/