Gaussian Mixture Solvers for Diffusion Models

TL;DR

This paper introduces Gaussian Mixture Solvers (GMS), a novel SDE-based approach for diffusion models that improves sample quality by better approximating the transition kernel, addressing efficiency and effectiveness issues in existing methods.

Contribution

The paper proposes GMS, which estimates moments and optimizes Gaussian mixture kernels during sampling, significantly enhancing diffusion model performance.

Findings

GMS outperforms existing SDE solvers in image quality.

GMS improves stroke-based synthesis results.

Empirical validation across various diffusion models confirms effectiveness.

Abstract

Recently, diffusion models have achieved great success in generative tasks. Sampling from diffusion models is equivalent to solving the reverse diffusion stochastic differential equations (SDEs) or the corresponding probability flow ordinary differential equations (ODEs). In comparison, SDE-based solvers can generate samples of higher quality and are suited for image translation tasks like stroke-based synthesis. During inference, however, existing SDE-based solvers are severely constrained by the efficiency-effectiveness dilemma. Our investigation suggests that this is because the Gaussian assumption in the reverse transition kernel is frequently violated (even in the case of simple mixture data) given a limited number of discretization steps. To overcome this limitation, we introduce a novel class of SDE-based solvers called \emph{Gaussian Mixture Solvers (GMS)} for diffusion models. Our solver estimates the first three-order moments and optimizes the parameters of a Gaussian mixture transition kernel using generalized methods of moments in each step during sampling. Empirically, our solver outperforms numerous SDE-based solvers in terms of sample quality in image generation and stroke-based synthesis in various diffusion models, which validates the motivation and effectiveness of GMS. Our code is available at https://github.com/Guohanzhong/GMS.