Modelos de Fluxo de Correspondência com Mistura Gaussiana

Gaussian Mixture Flow Matching Models

April 7, 2025

Autores: Hansheng Chen, Kai Zhang, Hao Tan, Zexiang Xu, Fujun Luan, Leonidas Guibas, Gordon Wetzstein, Sai Bi

cs.AI

Resumo

Os modelos de difusão aproximam a distribuição de remoção de ruído como uma Gaussiana e preveem sua média, enquanto os modelos de correspondência de fluxo reparametrizam a média Gaussiana como velocidade de fluxo. No entanto, eles apresentam desempenho inferior na amostragem com poucos passos devido ao erro de discretização e tendem a produzir cores super-saturadas sob orientação sem classificador (CFG). Para abordar essas limitações, propomos um novo modelo de correspondência de fluxo com mistura Gaussiana (GMFlow): em vez de prever a média, o GMFlow prevê parâmetros dinâmicos de mistura Gaussiana (GM) para capturar uma distribuição de velocidade de fluxo multimodal, que pode ser aprendida com uma perda de divergência KL. Demonstramos que o GMFlow generaliza modelos anteriores de difusão e correspondência de fluxo, onde uma única Gaussiana é aprendida com uma perda de remoção de ruído L_2. Para inferência, derivamos solucionadores GM-SDE/ODE que aproveitam distribuições analíticas de remoção de ruído e campos de velocidade para uma amostragem precisa com poucos passos. Além disso, introduzimos um novo esquema de orientação probabilística que mitiga os problemas de super-saturação da CFG e melhora a qualidade da geração de imagens. Experimentos extensivos demonstram que o GMFlow supera consistentemente as linhas de base de correspondência de fluxo em qualidade de geração, alcançando uma Precisão de 0,942 com apenas 6 passos de amostragem no ImageNet 256x256.

English

Diffusion models approximate the denoising distribution as a Gaussian and predict its mean, whereas flow matching models reparameterize the Gaussian mean as flow velocity. However, they underperform in few-step sampling due to discretization error and tend to produce over-saturated colors under classifier-free guidance (CFG). To address these limitations, we propose a novel Gaussian mixture flow matching (GMFlow) model: instead of predicting the mean, GMFlow predicts dynamic Gaussian mixture (GM) parameters to capture a multi-modal flow velocity distribution, which can be learned with a KL divergence loss. We demonstrate that GMFlow generalizes previous diffusion and flow matching models where a single Gaussian is learned with an L_2 denoising loss. For inference, we derive GM-SDE/ODE solvers that leverage analytic denoising distributions and velocity fields for precise few-step sampling. Furthermore, we introduce a novel probabilistic guidance scheme that mitigates the over-saturation issues of CFG and improves image generation quality. Extensive experiments demonstrate that GMFlow consistently outperforms flow matching baselines in generation quality, achieving a Precision of 0.942 with only 6 sampling steps on ImageNet 256times256.