Gaussian Mixture Flow Matching Models

📄 arXiv: 2504.05304v3 📥 PDF

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

分类: cs.LG, cs.CV

发布日期: 2025-04-07 (更新: 2025-08-30)

备注: ICML 2025. Code: https://github.com/Lakonik/GMFlow

💡 一句话要点

提出高斯混合流匹配模型(GMFlow)，提升少步采样质量并缓解图像生成中的色彩过饱和问题。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱四：生成式动作 (Generative Motion)

关键词: 生成模型 流匹配 高斯混合模型 少步采样 图像生成 无分类器指导 KL散度

📋 核心要点

  1. 现有扩散模型和流匹配模型在少步采样时存在离散化误差，且在无分类器指导下易产生色彩过饱和问题。
  2. GMFlow通过预测动态高斯混合模型参数来捕捉多模态流速分布，并使用KL散度损失进行学习，从而改进了采样质量。
  3. 实验表明，GMFlow在生成质量上优于现有流匹配方法，并在ImageNet 256x256数据集上取得了显著的精度提升。

📝 摘要（中文）

扩散模型通常将去噪分布近似为高斯分布并预测其均值，而流匹配模型则将高斯均值重新参数化为流速。然而，由于离散化误差，它们在少步采样中表现不佳，并且在无分类器指导(CFG)下容易产生过饱和的颜色。为了解决这些限制，我们提出了一种新的高斯混合流匹配(GMFlow)模型：GMFlow 不预测均值，而是预测动态高斯混合(GM)参数，以捕获多模态流速分布，这可以通过 KL 散度损失来学习。我们证明 GMFlow 推广了先前的扩散和流匹配模型，在这些模型中，单个高斯分布通过 L2 去噪损失来学习。对于推理，我们推导了 GM-SDE/ODE 求解器，它们利用解析去噪分布和速度场进行精确的少步采样。此外，我们引入了一种新的概率指导方案，可以减轻 CFG 的过饱和问题并提高图像生成质量。大量的实验表明，GMFlow 在生成质量方面始终优于流匹配基线，在 ImageNet 256x256 上仅使用 6 个采样步骤即可实现 0.942 的精度。

🔬 方法详解

问题定义：现有的扩散模型和流匹配模型在图像生成任务中，尤其是在需要快速生成（少步采样）的情况下，存在两个主要问题。一是由于离散化误差，少步采样的质量不高。二是当使用无分类器指导（CFG）时，容易生成色彩过饱和的图像。这些问题限制了这些模型在实际应用中的效果。

核心思路：GMFlow的核心思路是将传统的单高斯去噪分布扩展为高斯混合模型（GMM）。不再直接预测高斯分布的均值（如扩散模型），而是预测GMM的参数，从而能够更好地捕捉复杂的多模态流速分布。这种方法允许模型学习更丰富的去噪过程，从而提高少步采样的质量，并缓解CFG导致的色彩过饱和问题。

技术框架：GMFlow的整体框架包括训练和推理两个阶段。在训练阶段，模型学习预测动态高斯混合模型的参数，使用KL散度作为损失函数。在推理阶段，论文推导了基于高斯混合模型的SDE/ODE求解器（GM-SDE/ODE），利用解析去噪分布和速度场进行精确的少步采样。此外，还引入了一种新的概率指导方案，以减轻CFG的过饱和问题。

关键创新：GMFlow最重要的创新点在于使用高斯混合模型来表示去噪分布，而不是传统的单高斯分布。这使得模型能够捕捉更复杂、更精细的流速分布，从而提高了生成质量和采样效率。与现有方法的本质区别在于，GMFlow不再局限于预测单一的均值，而是学习整个分布的参数，从而能够更好地适应数据的复杂性。

关键设计：GMFlow的关键设计包括：1) 使用KL散度作为损失函数，用于学习高斯混合模型的参数。2) 推导了GM-SDE/ODE求解器，用于在推理阶段进行高效的少步采样。3) 引入了一种新的概率指导方案，用于减轻CFG的过饱和问题。这些设计共同作用，使得GMFlow能够在生成质量和采样效率方面取得显著的提升。

📊 实验亮点

实验结果表明，GMFlow在图像生成质量方面显著优于现有的流匹配模型。在ImageNet 256x256数据集上，仅使用6个采样步骤，GMFlow就达到了0.942的精度。这表明GMFlow在少步采样方面具有很强的竞争力，能够以更少的计算资源实现更高的生成质量。

🎯 应用场景

GMFlow模型在图像生成领域具有广泛的应用前景，例如高质量图像合成、图像编辑、超分辨率重建等。其高效的少步采样能力使其在需要快速生成图像的场景中具有优势，例如实时渲染、移动设备上的图像处理等。此外，GMFlow还可以应用于其他生成建模任务，例如音频生成、视频生成等。

📄 摘要（原文）

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 256$\times$256.