Treinando Modelos de Consistência com Acoplamento de Ruído Variacional

Training Consistency Models with Variational Noise Coupling

February 25, 2025

Autores: Gianluigi Silvestri, Luca Ambrogioni, Chieh-Hsin Lai, Yuhta Takida, Yuki Mitsufuji

cs.AI

Resumo

O Treinamento de Consistência (CT) surgiu recentemente como uma alternativa promissora aos modelos de difusão, alcançando desempenho competitivo em tarefas de geração de imagens. No entanto, o treinamento de consistência sem destilação frequentemente sofre com alta variância e instabilidade, e a análise e melhoria de sua dinâmica de treinamento é uma área ativa de pesquisa. Neste trabalho, propomos uma abordagem de treinamento CT inovadora baseada no framework de Correspondência de Fluxo. Nossa principal contribuição é um esquema de acoplamento de ruído treinado inspirado na arquitetura de Autoencoders Variacionais (VAE). Ao treinar um modelo de emissão de ruído dependente dos dados implementado como uma arquitetura de codificador, nosso método pode aprender indiretamente a geometria do mapeamento de ruído para dados, que é em vez disso fixado pela escolha do processo direto no CT clássico. Resultados empíricos em diversos conjuntos de dados de imagens mostram melhorias generativas significativas, com nosso modelo superando as bases e alcançando o estado-da-arte (SoTA) de CT sem destilação FID no CIFAR-10, e atingindo FID em nível com SoTA no ImageNet em resolução de 64 por 64 em geração de 2 etapas. Nosso código está disponível em 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 .