Semplificazione, Stabilizzazione e Scalabilità dei Modelli di Coerenza nel Tempo Continuo

Simplifying, Stabilizing and Scaling Continuous-Time Consistency Models

October 14, 2024

Autori: Cheng Lu, Yang Song

cs.AI

Abstract

I modelli di consistenza (CMs) sono una potente classe di modelli generativi basati sulla diffusione ottimizzati per un campionamento rapido. La maggior parte dei CMs esistenti viene addestrata utilizzando passaggi discretizzati nel tempo, che introducono iperparametri aggiuntivi e sono soggetti a errori di discretizzazione. Sebbene le formulazioni a tempo continuo possano attenuare tali problemi, il loro successo è stato limitato dall'instabilità dell'addestramento. Per affrontare ciò, proponiamo un quadro teorico semplificato che unifica le precedenti parametrizzazioni dei modelli di diffusione e dei CMs, identificando le cause principali dell'instabilità. Sulla base di questa analisi, introduciamo miglioramenti chiave nella parametrizzazione del processo di diffusione, nell'architettura della rete e negli obiettivi di addestramento. Questi cambiamenti ci consentono di addestrare CMs a tempo continuo su una scala senza precedenti, raggiungendo 1,5 miliardi di parametri su ImageNet 512x512. Il nostro algoritmo di addestramento proposto, utilizzando solo due passaggi di campionamento, raggiunge punteggi FID di 2,06 su CIFAR-10, 1,48 su ImageNet 64x64 e 1,88 su ImageNet 512x512, riducendo il divario nei punteggi FID rispetto ai migliori modelli di diffusione esistenti al di sotto 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%.