The Coupling Within: Flow Matching via Distilled Normalizing Flows

Abstract: Flow models have rapidly become the go-to method for training and deploying large-scale generators, owing their success to inference-time flexibility via adjustable integration steps. A crucial ingredient in flow training is the choice of coupling measure for sampling noise/data pairs that define the flow matching (FM) regression loss. While FM training defaults usually to independent coupling, recent works show that adaptive couplings informed by noise/data distributions (e.g., via optimal transport, OT) improve both model training and inference. We radicalize this insight by shifting the paradigm: rather than computing adaptive couplings directly, we use distilled couplings from a different, pretrained model capable of placing noise and data spaces in bijection -- a property intrinsic to normalizing flows (NF) through their maximum likelihood and invertibility requirements. Leveraging recent advances in NF image generation via auto-regressive (AR) blocks, we propose Normalized Flow Matching (NFM), a new method that distills the quasi-deterministic coupling of pretrained NF models to train student flow models. These students achieve the best of both worlds: significantly outperforming flow models trained with independent or even OT couplings, while also improving on the teacher AR-NF model.

• The paper proposes Normalized Flow Matching (NFM) which leverages a distilled normalizing flow to create a structured, data-dependent noise coupling during training.
• The methodology replaces heuristic and optimal transport pairings with deterministic noise assignment derived from a pretrained normalizing flow, leading to smoother training trajectories.
• Empirical results on ImageNet benchmarks demonstrate lower loss, improved FID scores, and up to 32x faster generation, even allowing student models to outperform their teachers.

The Coupling Within: Flow Matching via Distilled Normalizing Flows

Introduction and Background

This work introduces Normalized Flow Matching (NFM), a generative modeling technique that unifies Flow Matching (FM) with distilled Normalizing Flow (NF) models as a means for optimal data-noise coupling. The authors address a foundational issue in FM-based Neural ODE training: how to efficiently and effectively pair data and noise samples during model training. Standard independent coupling is suboptimal in practice, while recent approaches leverage data-dependent schemes, notably those based on Optimal Transport (OT). However, such approaches are fundamentally geometric or heuristic.

The proposition of this paper is to replace such heuristics by directly leveraging the bijective mapping established by a pretrained NF. Since these models inherently learn an invertible mapping between data and Gaussian noise via maximum likelihood, their internal coupling can be "distilled"—i.e., used as a teacher map in the FM training of student models. The approach is agnostic to the student's architecture and is not limited by the invertibility constraint of the NF teacher.

Figure 1: Schematic illustration of FM (independent coupling) vs. NFM training leveraging a distilled NF for structured coupling.

Flow Matching and the Importance of Coupling

Flow Matching constructs a vector field that learns to transport Gaussian noise onto data via linear interpolants in sample space, typically using regression on randomly paired $(x, \epsilon)$. The coupling method—how data and noise are assigned as pairs—is central, as it determines the learning and generalization properties of the resulting generative flow.

While independent coupling is convenient, it does not guarantee optimal training dynamics, particularly with high-dimensional and complex data. OT-based pairings improve matters but are still limited by the chosen cost function and label handling. The approach here replaces all such heuristics with the quasi-deterministic coupling produced by a normalizing flow teacher. This leverages both the data-adaptiveness and the marginal-preserving bijections of NFs, particularly modern autoregressive NFs such as TarFlow.

Methodology

NFM distills a teacher NF (trained via maximum likelihood, e.g., TarFlow) to guide the FM training of a student model. The critical difference during FM training is that, rather than sampling noise $\epsilon \sim \mathcal{N}(0,I_n)$ independently, the noise is deterministically assigned via the output of the NF's forward pass on $x$ (optionally with small additive noise for regularization):

$$z_{\epsilon'} = f_\text{NF}(x + \eta\epsilon', c)/\sigma_f$$

This "pseudonoise" is then linearly interpolated with data, and the student predicts velocities targeting $z_{\epsilon'} - x$. The authors provide a comprehensive analysis of the $z$-space structure, empirically showing that the induced "noise" is not a simple isometric embedding and that the nearest-neighbor structure is not preserved, but that these properties lead to lower loss variance, improved trajectory smoothness, and accelerated convergence in FM.

Empirical Results

The paper provides extensive experiments on class-conditional ImageNet64 and ImageNet256 benchmarks. Key metrics include Fréchet Inception Distance (FID), loss convergence, and generation latency. The authors compare vanilla FM, SD-FM (semidiscrete OT-based FM), and their proposed NFM.

NFM consistently yields marked improvements in training convergence (lower loss in fewer epochs), FID scores, and—in a striking result—outperforms even the teacher NF in FID despite leveraging a potentially less-expressive (non-invertible) student. This effect is substantial at low NFE (number of function evaluations, i.e., sampling steps), indicating that the teacher's coupling leads to both more efficient training and higher sample quality even in fast generation settings.

Figure 2: FID convergence curves on ImageNet64 showing NFM's superior convergence and sample quality across training epochs.

Figure 3: FID convergence on ImageNet256 (31 NFE), highlighting NFM's lead over FM baselines in large-scale, high-resolution settings.

Beyond sample quality, the paper reports dramatic latency improvements. A distilled NFM student with 31 NFE achieves a 32x speedup in generation compared to the AR-NF teacher, with even greater acceleration for fewer steps, without degrading sample quality.

Analysis of Latency and Teacher Properties

The authors analyze the impact of teacher NF quality (measured by NLL), architecture, and regularization noise $\eta$ on the performance of the distilled student. The FID improvements generally correlate with the teacher's NLL, up to the onset of overfitting. The choice of $\eta$ directly affects both the structure of the $z$-space and student performance. The best $\eta$ for teacher sampling is aligned with optimal teacher-student distillation.

Discussion and Implications

This paradigm reframes the utility of NF models, suggesting that their core contribution may be the ability to establish high-quality, structured data-noise couplings useful for distillation and accelerated sampling, rather than final deployment as generative models per se. It also highlights that the geometry of the NF-induced $z$-space is nontrivial and empirically beneficial for flow model training.

The discovery that NFM-trained students can outperform their teachers in FID and generation speed suggests new theoretical questions about knowledge distillation and the role of inductive biases. The combination of student/teacher architectural diversity and the coupling quality is central. The authors hypothesize the existence of further improved couplings—potentially by hybridizing SD-FM with NFM—that better align with the student's inductive bias.

Latent encoding in NFs offers an analogy to the use of autoencoders for downstream tasks. The potential for foundation models built as NFs whose distilled couplings are re-usable across vision, language, or multimodal domains is a significant direction. Additionally, there is excitement about further reducing the required number of integration steps (mean flows), which could bring truly one-step generation closer to practice.

Figure 4: Comparison of FM, SD-FM, and NFM training with classifier-free guidance (CFG), illustrating sample quality alignment.

Conclusion

The paper makes the case that flow matching with distilled normalizing flows solves a key open problem in generative modeling—the selection of an effective, flexible data-noise coupling—yielding superior convergence, sample quality, and efficiency relative to both FM and OT-based methods. The experimental and theoretical analyses suggest that the structured latent space imposed by the NF teacher both regularizes and accelerates FM learning, enabling high-quality, low-latency generation. Future directions include hybrid coupling strategies, deeper exploration of $z$-space geometry, transfer to other data modalities, and extension to mean flow models.

For complete references, see "The Coupling Within: Flow Matching via Distilled Normalizing Flows" (2603.09014).