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

TL;DR

LeDiFlow addresses the efficiency bottleneck of diffusion-model-based image generation by introducing a learned prior $P_L$ for flow matching, replacing the conventional Gaussian prior. It uses a regression-based auxiliary model conditioned on a per-image latent $Z$ and a transformer-based FM backbone, together with latent sampling to enable faster inference with fewer solver steps. The approach yields up to 3.75x speedups in pixel-space generation and a 1.32x improvement in CMMD for latent-space generation, while maintaining or improving image quality on datasets like FFHQ, ALIS, and StarGAN, and it supports controllable outputs via latent interpolation and inpainting. Limitations include reliance on the quality of the learned prior and reduced performance on very large-scale datasets; future work points to architectural scaling and extending to broader tasks such as full-text-to-image generation.

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.

Figures (12)

• Figure 1: LeDiFlow produces an easier to solve transformation for flow matching-based generative modeling via a learned prior, enabling faster inference with higher image quality. $Z$ resembles the latent space.
• Figure 2: Pipeline. An auxiliary model (middle) maps images $y \sim Q$ into latent vectors $z \sim Z$, from which the decoder defines a learned prior $P_L = \mathcal{N}(\mu(z;\theta_D), \sigma^2(z;\theta_D))$. In latent space (left), we train a flow model $\theta_L$ to match trajectories from Gaussian noise $w \sim \mathcal{N}(0,1)$ to semantic codes $z$, using supervised vectors $z - w$. In image space (right), another flow model $\theta_\mathrm{FM}$ transforms samples $x \sim P_L$ into the data distribution via flow matching towards $y \sim Q$. Dashed lines denote forward inference, solid lines are training-time supervision.
• Figure 3: Visual scatter plot showing the comparison from \ref{['tab:inf_performance']}.
• Figure 4: Samples showing unconditional generation of our method using 4 inference steps (Midpoint solver). From left to right: Karras2019ffhq, ALIS, choi2020stargan
• Figure 5: (Left) Visualization of the interpolation of samples $w \in \mathcal{N}(0,1)$ that are transformed into valid samples $z \in Z$ using $\theta_L$. (Right) Generated images from random latent vectors $z_0$ on the left and $z_1$ on the right using a constant noise seed.