Обучение моделей согласованности с вариационным шумовым сцеплением

Training Consistency Models with Variational Noise Coupling

February 25, 2025

Авторы: Gianluigi Silvestri, Luca Ambrogioni, Chieh-Hsin Lai, Yuhta Takida, Yuki Mitsufuji

cs.AI

Аннотация

Обучение согласованности (CT) недавно появилось как многообещающая альтернатива моделям диффузии, достигая конкурентоспособных результатов в задачах генерации изображений. Однако не дистиллированное обучение согласованности часто страдает от высокой дисперсии и нестабильности, и анализ и улучшение его динамики обучения является активной областью исследований. В данной работе мы предлагаем новый подход к обучению CT на основе концепции согласования потоков. Нашим основным вкладом является обученная схема шумообразования, вдохновленная архитектурой вариационных автокодировщиков (VAE). Обучив модель эмиссии шума, зависящую от данных и реализованную как архитектура кодировщика, наш метод может косвенно изучать геометрию отображения шума на данные, которая вместо этого фиксируется выбором прямого процесса в классическом CT. Эмпирические результаты на различных наборах изображений показывают значительное улучшение в генерации, превосходя наши модели базового уровня и достигая уровня лучших показателей (SoTA) в не дистиллированном CT FID на CIFAR-10, а также достигая FID на уровне SoTA на ImageNet с разрешением 64 на 64 пикселя в двухэтапной генерации. Наш код доступен по ссылке https://github.com/sony/vct.

English

Consistency Training (CT) has recently emerged as a promising alternative to diffusion models, achieving competitive performance in image generation tasks. However, non-distillation consistency training often suffers from high variance and instability, and analyzing and improving its training dynamics is an active area of research. In this work, we propose a novel CT training approach based on the Flow Matching framework. Our main contribution is a trained noise-coupling scheme inspired by the architecture of Variational Autoencoders (VAE). By training a data-dependent noise emission model implemented as an encoder architecture, our method can indirectly learn the geometry of the noise-to-data mapping, which is instead fixed by the choice of the forward process in classical CT. Empirical results across diverse image datasets show significant generative improvements, with our model outperforming baselines and achieving the state-of-the-art (SoTA) non-distillation CT FID on CIFAR-10, and attaining FID on par with SoTA on ImageNet at 64 times 64 resolution in 2-step generation. Our code is available at https://github.com/sony/vct .