Модели согласования потоков на основе гауссовских смесей

Gaussian Mixture Flow Matching Models

April 7, 2025

Авторы: Hansheng Chen, Kai Zhang, Hao Tan, Zexiang Xu, Fujun Luan, Leonidas Guibas, Gordon Wetzstein, Sai Bi

cs.AI

Аннотация

Диффузионные модели аппроксимируют распределение удаления шума как гауссово и предсказывают его среднее значение, тогда как модели согласования потоков перепараметризуют гауссово среднее как скорость потока. Однако они демонстрируют более низкую производительность при выборке за небольшое количество шагов из-за ошибок дискретизации и склонны создавать избыточно насыщенные цвета при использовании классификатор-независимого управления (CFG). Чтобы устранить эти ограничения, мы предлагаем новую модель согласования потоков на основе гауссовой смеси (GMFlow): вместо предсказания среднего значения GMFlow предсказывает динамические параметры гауссовой смеси (GM) для захвата многомодального распределения скорости потока, которое может быть обучено с использованием функции потерь на основе дивергенции Кульбака-Лейблера. Мы показываем, что GMFlow обобщает предыдущие диффузионные модели и модели согласования потоков, где обучается одна гауссова функция с использованием L_2-потери для удаления шума. Для вывода мы разрабатываем решатели GM-SDE/ODE, которые используют аналитические распределения удаления шума и поля скорости для точной выборки за небольшое количество шагов. Кроме того, мы представляем новую схему вероятностного управления, которая смягчает проблемы избыточной насыщенности CFG и улучшает качество генерации изображений. Многочисленные эксперименты демонстрируют, что GMFlow последовательно превосходит базовые модели согласования потоков по качеству генерации, достигая точности 0,942 всего за 6 шагов выборки на наборе данных 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.