LeDiFlow: Learned Distribution-guided Flow Matching to Accelerate Image Generation

TL;DR

LeDiFlow introduces a learned prior for flow matching in image generation, significantly reducing inference steps and improving quality, making high-quality diffusion-based image synthesis more efficient and accessible.

Contribution

The paper proposes a novel learned distribution-guided flow matching method that accelerates image generation by initializing the ODE solver with a learned prior, reducing computational complexity.

Findings

Up to 3.75x faster inference on pixel-space models.

Average 1.32x improvement in image quality measured by CMMD.

Outperforms baseline flow matching models in both speed and quality.

Abstract

Enhancing the efficiency of high-quality image generation using Diffusion Models (DMs) is a significant challenge due to the iterative nature of the process. Flow Matching (FM) is emerging as a powerful generative modeling paradigm based on a simulation-free training objective instead of a score-based one used in DMs. Typical FM approaches rely on a Gaussian distribution prior, which induces curved, conditional probability paths between the prior and target data distribution. These curved paths pose a challenge for the Ordinary Differential Equation (ODE) solver, requiring a large number of inference calls to the flow prediction network. To address this issue, we present Learned Distribution-guided Flow Matching (LeDiFlow), a novel scalable method for training FM-based image generation models using a better-suited prior distribution learned via a regression-based auxiliary model. By initializing the ODE solver with a prior closer to the target data distribution, LeDiFlow enables the learning of more computationally tractable probability paths. These paths directly translate to fewer solver steps needed for high-quality image generation at inference time. Our method utilizes a State-Of-The-Art (SOTA) transformer architecture combined with latent space sampling and can be trained on a consumer workstation. We empirically demonstrate that LeDiFlow remarkably outperforms the respective FM baselines. For instance, when operating directly on pixels, our model accelerates inference by up to 3.75x compared to the corresponding pixel-space baseline. Simultaneously, our latent FM model enhances image quality on average by 1.32x in CLIP Maximum Mean Discrepancy (CMMD) metric against its respective baseline.