PaperlayerView exact revision on arXiv

Source-linked AI summary

Flow Matching in Latent Space

Quan Dao, Hao Phung, Binh Nguyen, Anh Tran

arXiv:2307.08698v1cs.CVcs.LG

TL;DR

Prior flow-matching methods incurred expensive computation and many solver evaluations in pixel space, while latent flow matching remained underexplored and high-resolution flow matching was not yet ready. The paper applies flow matching in pretrained autoencoder latent spaces, adds conditional inputs with classifier-free velocity guidance, and derives a Wasserstein-2 bound. Across image-synthesis tasks, it reports competitive performance and improved efficiency, while noting risks from misinformation and substantial greenhouse-gas emissions during training.

  • Problem

    Prior flow-matching methods faced expensive computation and many function evaluations in pixel space, while high-resolution flow matching remained underdeveloped.

  • Method

    The paper performs flow matching in pretrained autoencoder latent spaces and redesigns the velocity field to support conditional inputs with classifier-free guidance.

  • Results

    The method attains competitive performance across image-synthesis tasks, with extensive quantitative and qualitative evaluations and a Wasserstein-2 upper bound.

  • Takeaways & Limitations

    Latent flow matching offers an approach toward more efficient and scalable high-resolution and conditional image generation.

  • Takeaways & Limitations

    Training the method emits substantial greenhouse gases and poses potential risks related to misinformation dissemination.

Abstract

from arXiv · show

Flow matching is a recent framework to train generative models that exhibits impressive empirical performance while being relatively easier to train compared with diffusion-based models. Despite its advantageous properties, prior methods still face the challenges of expensive computing and a large number of function evaluations of off-the-shelf solvers in the pixel space. Furthermore, although latent-based generative methods have shown great success in recent years, this particular model type remains underexplored in this area. In this work, we propose to apply flow matching in the latent spaces of pretrained autoencoders, which offers improved computational efficiency and scalability for high-resolution image synthesis. This enables flow-matching training on constrained computational resources while maintaining their quality and flexibility. Additionally, our work stands as a pioneering contribution in the integration of various conditions into flow matching for conditional generation tasks, including label-conditioned image generation, image inpainting, and semantic-to-image generation. Through extensive experiments, our approach demonstrates its effectiveness in both quantitative and qualitative results on various datasets, such as CelebA-HQ, FFHQ, LSUN Church & Bedroom, and ImageNet. We also provide a theoretical control of the Wasserstein-2 distance between the reconstructed latent flow distribution and true data distribution, showing it is upper-bounded by the latent flow matching objective. Our code will be available at https://github.com/VinAIResearch/LFM.git.

1 Introduction

The paper motivates flow matching as a faster, easier-to-train alternative to diffusion models, then moves flow matching into pretrained autoencoder latent spaces for efficient high-resolution synthesis. It also adds conditional generation and a Wasserstein-2 theoretical bound.

  • Diffusion models can require long sampling times, slow convergence, and sub-optimal probability paths that hinder large-scale performance.
  • Flow matching learns an ordinary differential equation from a source distribution to a target distribution and can be easier to train and faster to sample than diffusion models.
  • Latent flow matching applies path matching in pretrained autoencoder latent spaces to improve computational efficiency and scalability for high-resolution image synthesis.
  • The redesigned velocity field supports class labels, segmentation masks, and images for conditional generation tasks.The paper presents this as extending flow matching beyond prior unconditional tasks.
  • The paper provides a Wasserstein-2 bound showing that the distance between reconstructed latent flow and true data distributions is upper-bounded by the latent flow matching objective.
  • The contributions include latent-space flow matching, conditional inputs, theoretical analysis, and experiments across tasks and datasets.The listed applications include label-conditioned generation, image inpainting, and mask-to-image generation.

2 Related works

Prior generative approaches face computational or scalability limitations, while flow matching had not yet been developed for high-resolution or class-conditional image generation. This work addresses both gaps through latent representations and classifier-free guidance for velocity-based sampling.

  • Diffusion models often require hundreds to thousands of function evaluations per sample, motivating latent-space methods that reduce compute through compact autoencoder representations.
  • Diffusion models can also suffer from training-convergence and sub-optimal-trajectory issues that may affect training time and overall performance.
  • Earlier continuous normalizing flow frameworks were computationally prohibitive and difficult to scale because they required solving ODEs at every training iteration.
  • Current flow-matching methods were not yet ready for high-resolution image synthesis.
  • The paper studies latent representations in flow matching to enhance scalability and performance for high-resolution generation.
  • Flow-based models had not yet been explored for class-conditional generation, so the paper incorporates classifier-free guidance for latent flow matching.Unlike traditional frameworks that estimate noise, flow matching drives sampling with velocity.

3 Background

Flow matching learns an ODE-based velocity field that transports a source distribution to a target distribution. Linear interpolation and constant-velocity paths reduce trajectory curvature and can improve training and sampling efficiency.

  • Flow matching framework: Flow matching estimates a coupling between empirical data x0 ∼ p0 and noise x1 ∼ p1, then models their evolution with an ordinary differential equation.The velocity field drives the flow over t ∈ [0, 1].
  • Flow matching framework: The ODE gives a Lagrangian view of point-cloud dynamics, while the equivalent continuity equation describes the evolution of the measure p_t.The divergence operator appears in the continuity-equation formulation.
  • Probability flow ODE: Probability flow ODEs use the score function and diffusion-process drift and diffusion coefficients to define a deterministic generative trajectory.This formulation connects SDE-based diffusion processes with ODE sampling.
  • Probability flow ODE: ODE sampling with the variance-preserving path can reduce sampling costs compared with discretizing the corresponding diffusion SDE.The cited comparison concerns the ODE sampler using the path in Equation (6).
  • Constant velocity ODE: Constant-velocity flow matching replaces nonlinear interpolation with x_t = (1 − t)x0 + tx1, yielding velocity v_t = x1 − x0.The linear path avoids unnecessary curvature that can reduce training and sampling efficiency.
  • Constant velocity ODE: The method in this paper uses linear interpolation for training and the ODE in Equation (1) for sampling.This adopts the constant-velocity framework described above.

4 Methodology

Latent Flow Matching encodes images into pretrained-autoencoder latents, learns a velocity field from Gaussian noise to latent data, and decodes integrated samples into images. The framework also supports conditional inputs and provides assumptions and a Wasserstein-2 bound linking latent flow estimation to reconstructed-data quality.

  • Training and sampling: At sampling time, numerical ODE integration transports random noise toward the target latent distribution, after which a pretrained decoder produces the output image.The implementation tests Euler and Runge-Kutta variants with fixed or adaptive steps.
  • Training and sampling: Input images are encoded into compact latent codes, where a velocity network learns a path from z1 ∼ N(0, I) to the latent data distribution.The compact latent space reduces the dimensionality used for velocity-network optimization.
  • Conditional generation: Conditional flow matching supplies conditional information alongside the latent state and supports class labels, segmentation masks, and images.The paper applies these inputs to conditional generation tasks including inpainting and semantic-to-image synthesis.
  • Conditional generation: Classifier-free guidance combines conditional and unconditional velocity fields within one network, avoiding integration of a pretrained classifier.The unconditional field is trained by using an empty condition token with a certain probability.
  • Theoretical analysis: The theoretical analysis assumes a Gaussian encoder, a deterministic Lipschitz decoder, a regular estimated velocity field, and bounded reconstruction error.These assumptions define the conditions used for the Wasserstein-2 bound.
  • Theoretical analysis: Theorem 4.1 links the latent flow-matching objective to the Wasserstein distance between the true data distribution and reconstructed samples.The analysis also indicates that encoder and decoder capacities influence sample quality, exposing a speed–quality tradeoff.

5 Experiments

Experiments evaluate latent flow matching across unconditional and conditional image-generation tasks, datasets, architectures, and ODE solvers. The method is tested in latent space at multiple resolutions and supports image, semantic, and class conditioning.

  • Experimental setup: All experiments operate in the latent space of a pretrained Stable Diffusion VAE, whose encoder downsamples images by a factor of 8.For an RGB image of shape h × w × 3, the latent representation has shape h/8 × w/8 × 4.
  • Experimental setup: Experiments cover CelebA-HQ, FFHQ, LSUN Bedrooms, LSUN Church, and higher-resolution CelebA-HQ at 512 × 512.Both CNN-based ADM and transformer-based DiT architectures are evaluated using FID, Recall, running time, and function evaluations.
  • Conditional generation: The conditional-generation evaluation includes class-conditioned ImageNet, image inpainting, and semantic-to-image synthesis.Inpainting conditions the velocity network on encoded masked images and resized masks, while semantic-to-image uses an encoded semantic layout.
  • Conditional generation: Classifier-free velocity guidance substantially improves conditional ImageNet FID, reducing DiT B/2 from 20.38 to 4.46 at cfg = 1.5.For ADM, FID decreases from 16.71 to 8.58 at cfg = 1.25.
  • Conditional generation: Figure 4 compares conditional inputs, generated results, and ground-truth images for each image-conditional task.The layout provides a qualitative comparison across task-specific conditioning settings.
  • ODE solvers: Heun’s fixed-step solver achieves comparable performance to the adaptive solver with only 50 steps but is slightly slower because it performs an additional corrective forward pass.The experiments therefore select the adaptive solver by default.

6 Conclusions

The paper concludes that latent flow matching achieves competitive image-synthesis performance while extending flow matching to classifier-free class conditioning and several conditional tasks. It also identifies misinformation and greenhouse-gas emissions as potential risks.

  • Conclusions: Latent flow matching attains competitive performance across various image-synthesis tasks and narrows the gap between flow matching and diffusion models.The conclusion presents this as the central empirical outcome of the approach.
  • Conclusions: Classifier-free velocity guidance enables class conditioning in flow matching and improves class-conditional image generation.The conclusion frames this as a key capability added by the method.
  • Limitations and societal impact: The authors identify misinformation dissemination and substantial greenhouse-gas emissions as potential risks associated with the method.These concerns are stated in the paper’s limitation and societal-impact discussion.

A Additional background and proofs

The appendix introduces Wasserstein distance and the continuity-equation result used to analyze reconstructed latent-flow distributions. The proof relies on regularity assumptions and bounds the terminal W2 distance between true and reconstructed distributions.

  • Optimal-transport background: The appendix defines squared 2-Wasserstein distance through couplings between probability measures with finite second moments.The distance is framed through joint distributions whose marginals are the compared measures.
  • Optimal-transport background: The Kantorovich formulation expresses optimal transport through an optimization problem over couplings, potentially with entropic regularization.This provides the optimal-transport background for the theoretical analysis.
  • Continuity-equation analysis: Given a velocity field, the appendix defines the reconstructed density as the solution of an initial-value problem governed by the continuity equation.The result connects a velocity field to the evolution of the reconstructed flow distribution.
  • Continuity-equation analysis: If the velocity field is continuously differentiable and uniformly Lipschitz in space, the terminal squared W2 distance between true and reconstructed distributions admits a bound.The regularity assumption supplies the conditions for the proposition used in the proof.
  • Proof strategy: The proof combines the Wasserstein definition with reconstruction-error and Lipschitz-continuity arguments to control the final distributional discrepancy.The appendix also notes that differing base/target notation in prior work does not prevent applying its result.

B Analysis of fixed-steps ODE solvers

The solver analysis compares fixed-step Euler and Heun methods through FID and sampling time, while the sampling algorithms implement unconditional and classifier-free conditional latent-flow generation. Heun improves FID across tested settings but costs more sampling time.

  • Solver comparison: Fixed-step experiments vary the number of ODE steps and measure both FID and sampling time.The comparison is conducted on CelebA-HQ 256.
  • Solver comparison: Heun consistently obtains better FID scores than Euler across the tested fixed-step settings, with slightly higher sampling time.The trade-off favors quality for Heun and speed for Euler.
  • Unconditional generation: Unconditional training interpolates between encoded data latents and Gaussian noise, then minimizes squared velocity-prediction error.The training procedure uses the latent interpolation z_t and a velocity estimator.
  • Conditional generation: Class-conditional training randomly drops the class condition with probability p_u, enabling classifier-free velocity guidance during sampling.The sampling algorithm combines conditional and unconditional velocity estimates using guidance scale γ.
  • Sampling procedure: Euler sampling updates the latent state through successive velocity steps and decodes the final latent with the VAE decoder.The same latent-flow sampling structure supports unconditional generation, while conditional sampling adds the class condition.

D Additional experiments

On ImageNet 256 × 256, the proposed DiT-based approach with classifier-free velocity guidance consistently outperforms the original LDM and DiT-B/2 alternatives while remaining competitive with established methods.

  • The approach uses a DiT variant with classifier-free velocity guidance for ImageNet generation.
  • It consistently outperforms the original LDM and DiT-B/2 alternatives.
  • The method remains competitive with established methods despite smaller model sizes.

D.2 More qualitative results

The paper provides additional visual examples across face, scene, and ImageNet datasets, including multiple ImageNet guidance scales.

  • Additional visual examples cover CelebA-HQ 256, FFHQ, LSUN Church, LSUN Bedroom, and CelebA-HQ 512.
  • ImageNet examples are provided across Figures 11 to 20.
  • The ImageNet examples use guidance scales of 4.0, 2.0, and 1.5.

E Implementation details

The experiments use UNet and transformer-based DiT architectures, with configurations and training hyperparameters reported across the ADM and DiT networks.

  • Network configuration: Most experiments use either UNet-based ADM or transformer-based DiT architectures.
  • Network configuration: DiT-L/2 is used for unconditional 256 × 256 generation, while DiT-B/2 is used for class-conditional ImageNet generation.
  • Network configuration: Tables 6 and 7 report DiT network sizes and ADM configurations across datasets.
  • Training hyper-params: Tables 8 and 9 report ADM and DiT training hyperparameters, including estimated training days for each setting.
Loading 2307.08698v1…