Modelos de Emparejamiento de Flujo con Mezcla 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

Resumen

Los modelos de difusión aproximan la distribución de eliminación de ruido como una Gaussiana y predicen su media, mientras que los modelos de emparejamiento de flujo reparametrizan la media Gaussiana como velocidad de flujo. Sin embargo, su rendimiento es inferior en muestreo de pocos pasos debido al error de discretización y tienden a producir colores sobresaturados bajo guía sin clasificador (CFG). Para abordar estas limitaciones, proponemos un nuevo modelo de emparejamiento de flujo con mezcla Gaussiana (GMFlow): en lugar de predecir la media, GMFlow predice parámetros dinámicos de mezcla Gaussiana (GM) para capturar una distribución multimodal de velocidad de flujo, la cual puede aprenderse con una pérdida de divergencia KL. Demostramos que GMFlow generaliza modelos previos de difusión y emparejamiento de flujo donde se aprende una única Gaussiana con una pérdida de eliminación de ruido L_2. Para la inferencia, derivamos solucionadores GM-SDE/ODE que aprovechan distribuciones analíticas de eliminación de ruido y campos de velocidad para un muestreo preciso en pocos pasos. Además, introducimos un novedoso esquema de guía probabilística que mitiga los problemas de sobresaturación de CFG y mejora la calidad de generación de imágenes. Experimentos extensivos demuestran que GMFlow supera consistentemente a las líneas base de emparejamiento de flujo en calidad de generación, logrando una Precisión de 0.942 con solo 6 pasos de muestreo en 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.