The Coupling Within: Flow Matching via Distilled Normalizing Flows

David Berthelot; Tianrong Chen; Jiatao Gu; Marco Cuturi; Laurent Dinh; Bhavik Chandna; Michal Klein; Josh Susskind

cs.LGMar 9, 2026

The Coupling Within: Flow Matching via Distilled Normalizing Flows

Authors: David Berthelot, Tianrong Chen, Jiatao Gu, Marco Cuturi, Laurent Dinh, Bhavik Chandna, Michal Klein, Josh Susskind,

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.

Figures & tables

	$x_{1}\neq x_{2}$		$x_{1}=x_{2}$		$x_{1}\neq x_{2}$
	$\epsilon_{1}\neq\epsilon_{2}$		$\epsilon_{1}\neq\epsilon_{2}$		$\epsilon_{1}=\epsilon_{2}$
$\eta$	$d_{x}$	$d_{z}$	$d_{x}$	$d_{z}$	$d_{x}$	$d_{z}$
$0.025$	$0.50$	$1.00$	$0.025$	$0.78$	$0.50$	$0.67$
$0.05$	$0.50$	$1.00$	$0.05$	$0.85$	$0.50$	$0.57$
$0.1$	$0.51$	$1.00$	$0.10$	$0.90$	$0.50$	$0.47$
$0.2$	$0.54$	$1.00$	$0.20$	$0.94$	$0.50$	$0.37$

Table 1:

z

-space structure for ImageNet64

	$x_{1}\neq x_{2}$		$x_{1}=x_{2}$		$x_{1}\neq x_{2}$
	$\epsilon_{1}\neq\epsilon_{2}$		$\epsilon_{1}\neq\epsilon_{2}$		$\epsilon_{1}=\epsilon_{2}$
$\eta$	$d_{x}$	$d_{z}$	$d_{x}$	$d_{z}$	$d_{x}$	$d_{z}$
$0.025$	$0.80$	$1.00$	$0.025$	$0.26$	$0.80$	$0.99$
$0.050$	$0.80$	$1.00$	$0.05$	$0.37$	$0.80$	$0.97$
$0.100$	$0.80$	$1.00$	$0.10$	$0.50$	$0.80$	$0.93$
$0.250$	$0.83$	$1.00$	$0.25$	$0.69$	$0.80$	$0.80$

Table 2:

z

-space structure for ImageNet256

Solver	NFE	FM	SD-FM	NFM
Heun( $t^{2}$ )	31	0.0864	0.0767	0.0435
Euler( $t$ )	128	0.0386	0.0289	0.0181

Table 3: Curvature

\kappa

on ImageNet64

Figure 2: FID convergence on ImageNet64

Method	NFE	FID
		$32$ MiB	$64$ MiB	$128$ MiB	$256$ MiB
FM	$31$	$14.33$	$6.11$	$3.56$	$2.57$
SD-FM	$31$	$13.83$	$6.06$	$3.63$	$2.68$
NFM	$31$	8.73	3.33	2.17	1.78
FM	$15$	$17.15$	$7.47$	$5.24$	$4.80$
SD-FM	$15$	$14.45$	$6.39$	$3.98$	$3.15$

Table 4: FID comparisons ImageNet64 with SiT-XL/4 for TarFlow teacher TF-6

\times

2+2

\times

26-1024/2 with FID=

1.98

Solver	NFE	FM	SD-FM	NFM
Heun	31	2.57	2.68	1.78
Heun	15	4.80	3.15	2.15
Heun	7	13.01	6.41	3.23
Heun	5	17.56	9.29	4.01
Euler	5	21.05	12.18	3.92
Euler	4	21.81	13.76	4.47

Table 5: FID vs NFEs comparisons for ImageNet64 with SiT-XL/4 for TarFlow teacher TF-6

\times

2+2

\times

26-1024/2 with FID=

1.98

Figure 3: FID convergence on ImageNet256 (31 NFE sampling)

Method	NFE	FID
		$32$ MiB	$64$ MiB	$128$ MiB	$256$ MiB
FM	$31$	$12.33$	$4.93$	$3.05$	$2.30$
NFM	$31$	9.68	4.17	2.78	2.29
FM	$15$	$12.40$	$5.07$	$3.37$	$2.84$
NFM	$15$	10.36	4.14	2.84	2.45
FM	$7$	$19.92$	$13.72$	$13.60$	$12.41$

Table 6: FID comparisons ImageNet256 with SiT-XL/2 for TarFlow teacher TF-6

\times

2+2

\times

26-1024/1 with FID=

3.96

Solver	NFE	FM		NFM
		$128$ MiB	$256$ MiB	$128$ MiB	$256$ MiB
Heun	31	3.05	2.30	2.78	2.29
Heun	15	3.37	2.84	2.84	2.45
Heun	7	13.60	12.41	3.93	3.43
Heun	5	48.51	48.73	6.43	5.75
Euler	5	11.60	10.48	4.72	4.04

Table 7: FID vs NFEs comparisons for ImageNet256 with SiT-XL/2 for TarFlow teacher TF-6

\times

2+2

\times

26-1024/1 with FID=

3.96

Method	FID	Latency	Speedup
TarFlow	$1.98$	$10.7965$	$1\times$
NFM(31)	$1.78$	$0.3376$	$32\times$
NFM(15)	$2.15$	$0.1585$	$68\times$
NFM(7)	$3.23$	$0.0744$	$145\times$

Table 8: Latency vs FID ImageNet64 with SiT-XL/4 students and TF-6

\times

2+2

\times

26-1024/2 teacher.

Panel of Figure 4 — Figure 4: ImageNet64 same-seed random samples. Rows 1,2,3 are resp. from TF teacher, FM(31) student and FM(7) student.

Teacher	Size	NLL	FID (T/S)
TF-6 $\times$ 2+2 $\times$ 26-768/4	$460$ K	-2.1252	$3.93/2.67$
TF-6 $\times$ 2+2 $\times$ 26-768/2	$460$ K	-2.1410	$2.14/1.80$
TF-6 $\times$ 2+2 $\times$ 26-1024/2	$823$ K	-2.1500	$1.98/1.78$
TF-6 $\times$ 2+2 $\times$ 26-1536/2	$1839$ K	-2.1691	$2.21/1.84$
TF-8 $\times$ 8-768/2	$460$ K	-2.1353	$3.42/2.41$

Table 9: TarFlow teacher impact on FM student for ImageNet64, T/S columns for FID contain Teacher and Student FIDs

$\eta$	NLL	FID (T/S)
$0.025$	-2.6223	$2.40/1.85$
$0.05$	-2.1500	$1.98/1.78$
$0.1$	-1.6040	$1.95/1.84$
$0.2$	-1.0040	$3.67/2.31$

Table 10: TF-6

\times

2+2

\times

26-1024/2 teacher noise impact on FM student for ImageNet64, T/S columns for FID contain Teacher and Student FIDs.

Figure 5: FID convergence on ImageNet64 without CFG

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 6 — Figure 6: FM loss convergence

Panel of Figure 7 — Figure 7: ImageNet64 samples for class 12: house finch.

Panel of Figure 8 — Figure 8: ImageNet64 samples for class 207: golden retriever.

Panel of Figure 9 — Figure 9: ImageNet64 samples for class 309: bee.

Panel of Figure 10 — Figure 10: ImageNet64 samples for class 417: balloon.

Panel of Figure 11 — Figure 11: ImageNet64 samples for class 554: fireboat.

Panel of Figure 12 — Figure 12: ImageNet64 samples for class 985: daisy.

Panel of Figure 13 — Figure 13: ImageNet256 samples for class 12: house finch.

Panel of Figure 14 — Figure 14: ImageNet256 samples for class 207: golden retriever.

Panel of Figure 15 — Figure 15: ImageNet256 samples for class 309: bee.

Panel of Figure 16 — Figure 16: ImageNet256 samples for class 417: balloon.

Panel of Figure 17 — Figure 17: ImageNet256 samples for class 554: fireboat.

Panel of Figure 18 — Figure 18: ImageNet256 samples for class 985: daisy.

Flow-Based Generative Modeling

Flow-based generative modeling learns probability transports from simple noise to data for fast sampling, density estimation, and controllable generation.

See all 538 papers

Normalizing Flow Maps
Oct 15, 2026Francesco M. Ruscio, Timothy James Hitge, Fred Zhangzhi Peng +6Normalizing Flows Flow Map Learning
AITHYRA · Imperial College London · TU Wien +5
One-Sided Quantile Coupling for Flow Matching
Aug 2, 2026Jin-Young Kim, So-Yoon Cho, Hyun-Gyoon KimFlow Matching Image Generation
Independent Researcher · Division of Finance & AI, Hankuk University of Foreign Studies · Dept. of Financial Engineering, Ajou University
Normalizing Flows with Iterative Denoising
Apr 21, 2026Tianrong Chen, Jiatao Gu, David Berthelot +2Normalizing Flows Autoregressive Image Generation
Apple Machine Learning Research.

Oct 15, 2026J/K move · Enter open · S save

Normalizing Flow Maps

Francesco M. Ruscio, Timothy James Hitge, Fred Zhangzhi Peng, Danyal Rehman+5

AITHYRA · Imperial College London · TU Wien+5

Normalizing Flows Flow Map Learning Conditional Generative Modeling

Open paper AskPDF arXiv

Scalable approaches to learning neural dynamical systems have placed diffusion, flow matching, and few-step generators in flow maps at the forefront of generative modeling. Such state-of-the-art models have limited access to explicit model densities despite evolving from classical generative modeling approaches that leveraged tractable likelihoods, e.g., Normalizing Flows and Autoregressive models. In this paper, we combine these threads of few-step generation and tractable sample likelihoods and introduce Normalizing Flow Maps (NFM), a class of exactly invertible pushforward maps that generalizes normalizing flows and flow maps. More precisely, NFMs learn the integrator of an underlying dynamical system, but crucially, they provide efficient exact likelihood evaluation by construction. To enable scalable training of NFMs, we derive two complementary axes of consistency inherent to NFMs that yield (self) distillation objectives: (1) pointwise matching, which unlocks

\ell_2

-based flow-map style regression losses, and (2) distributional matching, which matches pushforward distributions to the marginals of a prescribed forward process. On conditional image generation on CIFAR-10 and ImageNet-64, NFMs achieve state-of-the-art sample quality among parameter-matched normalizing flows, while supporting likelihood-based RLHF. For equilibrium sampling in Boltzmann Generators, NFMs outperform prior state-of-the-art normalizing flows on both single-system and transferable peptide benchmarks.

Abstract

Flow-Based Generative Modeling

Explore similar work

Normalizing Flow Maps

One-Sided Quantile Coupling for Flow Matching

Normalizing Flows with Iterative Denoising