Упрощение, стабилизация и масштабирование моделей непрерывного времени согласованности.

Simplifying, Stabilizing and Scaling Continuous-Time Consistency Models

October 14, 2024

Авторы: Cheng Lu, Yang Song

cs.AI

Аннотация

Модели согласованности (CM) представляют собой мощный класс генеративных моделей на основе диффузии, оптимизированных для быстрого сэмплирования. Большинство существующих CM обучаются с использованием дискретизированных временных шагов, что вводит дополнительные гиперпараметры и подвержено ошибкам дискретизации. В то время как формулировки в непрерывном времени могут смягчить эти проблемы, их успех ограничен нестабильностью обучения. Для решения этой проблемы мы предлагаем упрощенную теоретическую основу, объединяющую предыдущие параметризации моделей диффузии и CM, выявляя коренные причины нестабильности. Основываясь на этом анализе, мы вносим ключевые улучшения в параметризацию процесса диффузии, архитектуру сети и цели обучения. Эти изменения позволяют нам обучать CM в непрерывном времени на беспрецедентном масштабе, достигая 1.5 миллиарда параметров на ImageNet 512x512. Наш предложенный алгоритм обучения, использующий всего два шага сэмплирования, достигает значений FID 2.06 на CIFAR-10, 1.48 на ImageNet 64x64 и 1.88 на ImageNet 512x512, сокращая разрыв в значениях FID с лучшими существующими моделями диффузии до 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%.