Simplificando, Estabilizando e Dimensionando Modelos de Consistência em Tempo Contínuo.

Simplifying, Stabilizing and Scaling Continuous-Time Consistency Models

October 14, 2024

Autores: Cheng Lu, Yang Song

cs.AI

Resumo

Os modelos de consistência (CMs) são uma classe poderosa de modelos generativos baseados em difusão, otimizados para amostragem rápida. A maioria dos CMs existentes é treinada usando passos de tempo discretizados, que introduzem hiperparâmetros adicionais e são propensos a erros de discretização. Enquanto formulações em tempo contínuo podem mitigar esses problemas, seu sucesso tem sido limitado pela instabilidade de treinamento. Para abordar isso, propomos um arcabouço teórico simplificado que unifica as parametrizações anteriores de modelos de difusão e CMs, identificando as causas raiz da instabilidade. Com base nessa análise, introduzimos melhorias-chave na parametrização do processo de difusão, arquitetura de rede e objetivos de treinamento. Essas mudanças nos permitem treinar CMs em tempo contínuo em uma escala sem precedentes, alcançando 1,5 bilhão de parâmetros no ImageNet 512x512. Nosso algoritmo de treinamento proposto, usando apenas dois passos de amostragem, alcança pontuações FID de 2,06 no CIFAR-10, 1,48 no ImageNet 64x64 e 1,88 no ImageNet 512x512, reduzindo a diferença nas pontuações FID com os melhores modelos de difusão existentes para dentro 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%.