Het vereenvoudigen, stabiliseren en opschalen van continue-tijd consistentiemodellen.

Simplifying, Stabilizing and Scaling Continuous-Time Consistency Models

October 14, 2024

Auteurs: Cheng Lu, Yang Song

cs.AI

Samenvatting

Consistentiemodellen (CM's) zijn een krachtige klasse van diffusie-gebaseerde generatieve modellen geoptimaliseerd voor snelle bemonstering. De meeste bestaande CM's worden getraind met behulp van gedisciplineerde tijdstappen, die extra hyperparameters introduceren en gevoelig zijn voor discretisatiefouten. Hoewel continue-tijd formuleringen deze problemen kunnen verminderen, is hun succes beperkt gebleven door trainingsinstabiliteit. Om dit aan te pakken, stellen we een vereenvoudigd theoretisch kader voor dat eerdere parameterisaties van diffusiemodellen en CM's verenigt, waarbij de oorzaken van instabiliteit worden geïdentificeerd. Op basis van deze analyse introduceren we belangrijke verbeteringen in de parameterisatie van diffusieprocessen, netwerkarchitectuur en trainingsdoelstellingen. Deze veranderingen stellen ons in staat om continue-tijd CM's op een ongekende schaal te trainen, met 1,5 miljard parameters op ImageNet 512x512. Ons voorgestelde trainingsalgoritme, met slechts twee bemonsteringsstappen, behaalt FID-scores van 2,06 op CIFAR-10, 1,48 op ImageNet 64x64, en 1,88 op ImageNet 512x512, waarbij het verschil in FID-scores met de beste bestaande diffusiemodellen tot binnen 10% wordt verkleind.

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%.