Better Source, Better Flow: Learning Condition-Dependent Source Distribution for Flow Matching

TL;DR

This paper introduces a method to learn condition-dependent source distributions for flow matching in text-to-image generation, leading to faster convergence and improved quality by optimizing the source distribution.

Contribution

It proposes a novel approach to design and learn source distributions conditioned on input, addressing stability issues and enhancing flow matching performance in generative models.

Findings

Up to 3x faster convergence in FID scores.

Improved stability with variance regularization.

Enhanced performance across multiple benchmarks.

Abstract

Flow matching has recently emerged as a promising alternative to diffusion-based generative models, particularly for text-to-image generation. Despite its flexibility in allowing arbitrary source distributions, most existing approaches rely on a standard Gaussian distribution, a choice inherited from diffusion models, and rarely consider the source distribution itself as an optimization target in such settings. In this work, we show that principled design of the source distribution is not only feasible but also beneficial at the scale of modern text-to-image systems. Specifically, we propose learning a condition-dependent source distribution under flow matching objective that better exploit rich conditioning signals. We identify key failure modes that arise when directly incorporating conditioning into the source, including distributional collapse and instability, and show that appropriate variance regularization and directional alignment between source and target are critical for stable and effective learning. We further analyze how the choice of target representation space impacts flow matching with structured sources, revealing regimes in which such designs are most effective. Extensive experiments across multiple text-to-image benchmarks demonstrate consistent and robust improvements, including up to a 3x faster convergence in FID, highlighting the practical benefits of a principled source distribution design for conditional flow matching.