Modèles de correspondance de flux par mélange gaussien

Gaussian Mixture Flow Matching Models

April 7, 2025

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

cs.AI

Résumé

Les modèles de diffusion approchent la distribution de débruitage comme une Gaussienne et prédisent sa moyenne, tandis que les modèles de correspondance de flux reparamétrisent la moyenne Gaussienne en tant que vitesse de flux. Cependant, ils sous-performent dans l'échantillonnage en quelques étapes en raison de l'erreur de discrétisation et ont tendance à produire des couleurs sursaturées sous le guidage sans classifieur (CFG). Pour remédier à ces limitations, nous proposons un nouveau modèle de correspondance de flux par mélange Gaussien (GMFlow) : au lieu de prédire la moyenne, GMFlow prédit des paramètres dynamiques de mélange Gaussien (GM) pour capturer une distribution de vitesse de flux multi-modale, qui peut être apprise avec une perte de divergence KL. Nous démontrons que GMFlow généralise les modèles de diffusion et de correspondance de flux précédents où une seule Gaussienne est apprise avec une perte de débruitage L_2. Pour l'inférence, nous dérivons des solveurs GM-SDE/ODE qui exploitent des distributions de débruitage analytiques et des champs de vitesse pour un échantillonnage précis en quelques étapes. De plus, nous introduisons un nouveau schéma de guidage probabiliste qui atténue les problèmes de sursaturation du CFG et améliore la qualité de génération d'images. Des expériences approfondies démontrent que GMFlow surpasse systématiquement les modèles de référence en correspondance de flux en termes de qualité de génération, atteignant une Précision de 0,942 avec seulement 6 étapes d'échantillonnage sur ImageNet 256×256.

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.