Normalizing Flow Maps

Francesco M. Ruscio; Timothy James Hitge; Fred Zhangzhi Peng; Danyal Rehman; Zander W. Blasingame; Nicholas M. Boffi; Michael M. Bronstein; Alexander Tong

cs.LGOct 15, 2026

Normalizing Flow Maps

Authors: Francesco M. Ruscio, Timothy James Hitge, Fred Zhangzhi Peng, Danyal Rehman, Zander W. Blasingame, Nicholas M. Boffi, Michael M. Bronstein, Alexander Tong,

Organizations: AITHYRA · Imperial College London · TU Wien · Duke University ·

Abstract

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.

Figures & tables

Figure 1: Overview of NFM , which performs generative transport along a stochastic interpolant. NFM-R exploits pointwise matching at the particle level, while NFM-D leverages distributional matching.

Panel of Figure 2 — Figure 2: NFM post-training. (a) NFM-D, golden retriever ; (b) NFM-D post-trained, a single golden retriever on a sandy beach ; (c) NFM-R, bighorn sheep ; (d) NFM-R post-trained, bighorn sheep standing on a sandy beach beside the sea .

Table 1: Class-conditional generation results. Best results are in bold , and second-best are underlined .

Figure 3: Post-trained NFM-D vs. best-of-

64

and SMC baselines.

Figure 4: a) Energy histograms of MD data, the NFM-D proposal, and the resampled distribution, for a single-system model (left) and a transferable model evaluated on a $4$ - and $8$ -residue peptides (middle, right). b) TICA projections for the ground truth MD, TARFlow, and NFM-D on NHQYGSDP.

Algorithm	$\mathcal{E}$ - $\mathcal{W}_{2}\downarrow$	$\mathbb{T}$ - $\mathcal{W}_{2}\downarrow$
ArBG ( Rehman et al., 2026c )	$7.94\pm 1.05$	$\underline{2.78\pm 0.11}$
SBG SMC ( Tan et al., 2025a )	$10.82\pm 7.21$	$3.78\pm 0.44$
SBG-iTARFlow SMC	$10.87\pm 10.18$	$3.56\pm 0.92$
NFM-D IS (Ours)	$\underline{5.12\pm 0.26}$	$\mathbf{2.71\pm 0.05}$
NFM-D SMC (Ours)	$\mathbf{3.42\pm 0.27}$	$2.81\pm 0.12$

Table 2: Performance comparison on Chignolin. SBG-iTARFlow refers to our reimplementation of SBG, used to control for architectural, optimizer, and training-recipe changes. Best results are in bold ; second-best results are underlined .

	4AA (30 systems)			8AA (30 systems)
Model	$\mathcal{E}$ - $\mathcal{W}_{2}\downarrow$	$\mathbb{T}$ - $\mathcal{W}_{2}\downarrow$	TICA- $\mathcal{W}_{2}\downarrow$	$\mathcal{E}$ - $\mathcal{W}_{2}\downarrow$	$\mathbb{T}$ - $\mathcal{W}_{2}\downarrow$	TICA- $\mathcal{W}_{2}\downarrow$
TimeWarp ( Klein et al., 2023a )	$7.237$	$2.204$	$0.993$	—	—	—
BioEmu ( Lewis et al., 2025 )	$90.079$	$2.037$	$1.479$	$193.873$	$4.638$	$1.601$
UniSim ( Yu et al., 2025 )	$>10^{4}$	$2.766$	$1.733$	$>10^{3}$	$6.156$	$1.495$
ECNF++ ( Tan et al., 2025b )	$10.032$	$1.121$	$0.572$	—	—	—
Prose ( Tan et al., 2025b )	$\underline{0.932}$	$\underline{0.752}$	$\mathbf{0.367}$	$10.038$	$2.456$	$0.988$

Table 3: Transferable models comparison on unseen 4- and 8-residue peptides, using a budget of

10^{4}

energy evaluations and

T=1

. Best results are in bold ; second-best results are underlined .

	4AA (30 systems)			8AA (30 systems)
Model	$\mathcal{E}$ - $\mathcal{W}_{2}\downarrow$	$\mathbb{T}$ - $\mathcal{W}_{2}\downarrow$	TICA- $\mathcal{W}_{2}\downarrow$	$\mathcal{E}$ - $\mathcal{W}_{2}\downarrow$	$\mathbb{T}$ - $\mathcal{W}_{2}\downarrow$	TICA- $\mathcal{W}_{2}\downarrow$
Robin	$1.168$	$0.886$	$0.471$	$\underline{4.251}$	$2.325$	$\underline{0.943}$
Robin SMC	$\underline{1.079}$	$\underline{0.874}$	$\underline{0.463}$	$4.263$	$\underline{2.315}$	$0.977$
NFM-D (Ours)	$\mathbf{0.865}$	$\mathbf{0.745}$	$\mathbf{0.398}$	$\mathbf{4.148}$	$\mathbf{2.296}$	$\mathbf{0.892}$

Table 4: Performance of transferable models with temperature tuning. Best results are in bold ; second-best results are underlined .

Appendix figures & tables20 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 5 — Figure 5: Guided MNIST samples from MLE-trained iTARFlow.

System	Patch Size	Blocks	Layers	Channels	Parameters (M)
CIFAR- $10$	2	8	6	256	$~40$ M
ImageNet- $64$	2	8	8	384	$~120$ M

Appendix

Table 5: Architectural hyperparameters used for the image experiments.

Setting	Variant	FID $\downarrow$	FLD $\downarrow$	KID $\downarrow$	Prec. $\uparrow$	Rec. $\uparrow$	Invertible? $\downarrow$
No noise	Guidance	40.83	23.21	2.758	0.744	0.025	✓
No noise	No guidance	55.52	26.90	4.249	0.835	0.017	✓
0.05 noise	Guidance and Denoise	6.72	8.22	0.457	0.595	0.293	✗
	No guidance	12.04	12.00	1.277	0.627	0.340	✗
	No denoise	48.34	10.05	0.923	0.676	0.122	✓
	No guid., no denoise	71.48	15.16	2.256	0.734	0.129	✓

Appendix

Table 6: MLE iTARFlow AdamW results on conditional CIFAR-10, results are selected across the learning rates

\{10^{-4},4\times 10^{-4},10^{-3}\}

and EMAs

\{0.9999,0.9995,0.999\}

. Results computed with the default MLE iTARFlow guidance of 2.8.

Optimizer	Variant	FID $\downarrow$	FLD $\downarrow$	KID $\downarrow$	Prec. $\uparrow$	Rec. $\uparrow$	Invertible?
AdamW	Guidance and Denoise	6.72	8.22	0.457	0.595	0.293	✗
	No guidance	12.04	12.00	1.277	0.627	0.340	✗
	No denoise	48.34	10.05	0.923	0.676	0.122	✓
	No guid., no denoise	71.48	15.16	2.256	0.734	0.129	✓
Muon	Guidance and Denoise	7.54	7.69	0.394	0.622	0.275	✗
	No guidance	11.62	11.35	1.177	0.630	0.322	✗

Appendix

Table 7: MLE iTARFlow results on conditional CIFAR-10 at

0.05

input noise, comparing the AdamW and Muon optimizers. The AdamW row is selected across the learning rates

\{10^{-4},4\times 10^{-4},10^{-3}\}

; the Muon row across the same learning rates and Muon weight decays

\{0,0.1\}

. Both select over EMAs

\{0.9999,0.9995,0.999\}

and steps by FID. Results computed with the default MLE iTARFlow guidance of 2.8.

Panel of Figure 6 — Figure 6: NFM-R ( noise annealed ) CIFAR-10 FID as the CFG strength ( $\omega$ ) and the CFG interval $[t_{\rm min},t_{\rm max}]$ vary.

Figure 7: NFM-D and iTARFlow MLE guidance scale (w) vs FID on CIFAR-10

Panel of Figure 8 — Figure 8: Samples from NFM s pre- and post-REINFORCE fine-tuning with a reward favouring green frogs.

Figure 9: Post-trained NFM-R vs. best-of-

64

and SMC baselines. Prompt: bighorn sheep standing on a sandy beach beside the sea .

System	Layers	Blocks	Channels	Parameters (M)
AL3	6	6	256	29.5
AL4	6	6	256	29.5
AL6	6	6	384	66.5
Chignolin	6	6	384	66.5
Transferable	8	8	384	118.6

Appendix

Table 8: Architectural hyperparameters used for NFM .

	Tri-alanine (AL3)		Tetrapeptide (AL4)		Hexa-alanine (AL6)
Algorithm	$\mathcal{E}$ - $\mathcal{W}_{2}\downarrow$	$\mathbb{T}$ - $\mathcal{W}_{2}\downarrow$	$\mathcal{E}$ - $\mathcal{W}_{2}\downarrow$	$\mathbb{T}$ - $\mathcal{W}_{2}\downarrow$	$\mathcal{E}$ - $\mathcal{W}_{2}\downarrow$	$\mathbb{T}$ - $\mathcal{W}_{2}\downarrow$
ArBG ( Rehman et al., 2026c )	$0.271\pm 0.113$	$0.311\pm 0.009$	$0.886\pm 0.076$	$0.593\pm 0.008$	$0.722\pm 0.063$	$1.085\pm 0.069$
ECNF++ ( Tan et al., 2025a )	$2.206\pm 0.813$	$0.962\pm 0.253$	$5.638\pm 0.483$	$1.002\pm 0.061$	$10.668\pm 0.285$	$1.902\pm 0.055$
RegFlow ( Rehman et al., 2026b )	$0.853\pm 0.105$	$1.577\pm 0.140$	$3.277\pm 0.546$	$2.342\pm 0.102$	—	—
FALCON-A ( Rehman et al., 2026a )	$1.385\pm 0.182$	$0.343\pm 0.004$	$2.929\pm 0.068$	$1.094\pm 0.034$	$1.211\pm 0.105$	$\mathbf{1.163\pm 0.112}$
FALCON ( Rehman et al., 2026a )	$\underline{0.544\pm 0.013}$	$0.452\pm 0.011$	$\mathbf{0.686\pm 0.047}$	$\mathbf{0.858\pm 0.077}$	$\underline{0.892\pm 0.311}$	$\mathbf{1.256\pm 0.132}$

Appendix

Table 9: Benchmark results for continuous flows on AL3, AL4, and AL6. We report ArBG in gray as an autoregressive model reference. Bold values indicate the best result between SBG and NFM-D .

Figure 10: Energy histograms for the true ground truth MD data, the proposal, and the resampled distribution for AL3, AL4, and AL6.

Figure 11: Ramachandran plots for tri-alanine (AL3). Left: reference MD data; right: angular distributions generated by NFM-D.

Figure 12: Ramachandran plots for tetra-alanine (AL4). Left: reference MD data; right: angular distributions generated by NFM-D.

Figure 13: Ramachandran plots for hexa-alanine (AL6). Left: reference MD data; right: angular distributions generated by NFM-D.

	Raw		Resampled
Algorithm	$\mathcal{E}$ - $\mathcal{W}_{2}\downarrow$	$\mathbb{T}$ - $\mathcal{W}_{2}\downarrow$	$\mathcal{E}$ - $\mathcal{W}_{2}\downarrow$	$\mathbb{T}$ - $\mathcal{W}_{2}\downarrow$
NFM-R	$>10^{9}$	$\mathbf{0.31}$	$0.54$	$0.51$
NFM-D	$>10^{9}$	$0.32$	$\mathbf{0.15}$	$\mathbf{0.40}$

Appendix

Table 10: Comparison of the best NFM-R and NFM-D models on AL3 using raw and resampled metrics. Best results are in bold .

Algorithm	$\mathcal{E}$ - $\mathcal{W}_{2}\downarrow$	$\mathbb{T}$ - $W_{2}\downarrow$	Parameters (M)
SBG SMC	$10.819\pm 7.206$	$3.778\pm 0.440$	114
SBG-iTARFlow IS	$17.015\pm 11.781$	$2.845\pm 0.069$	$66.5$
SBG-iTARFlow SMC	$10.874\pm 10.179$	$3.558\pm 0.918$	$66.5$
NFM-D IS (Ours)	$\underline{5.118\pm 0.262}$	$\mathbf{2.712\pm 0.050}$	$66.5$
NFM-D SMC (Ours)	$\mathbf{3.419\pm 0.268}$	$\underline{2.809\pm 0.119}$	$66.5$

Appendix

Table 11: Comparison of SBG and NFM-D on Chignolin. SBG-iTARFlow refers to our reimplementation of SBG, used to control for architectural, optimizer, and training-recipe changes. Best results are in bold ; second-best results are underlined .

Sequence length $\rightarrow$	4AA ${}_{\text{(30 systems)}}$			8AA ${}_{\text{(30 systems)}}$
Model	$\mathcal{E}$ - $\mathcal{W}_{2}$ $\downarrow$	$\mathbb{T}$ - $\mathcal{W}_{2}$ $\downarrow$	TICA- $\mathcal{W}_{2}$ $\downarrow$	$\mathcal{E}$ - $\mathcal{W}_{2}$ $\downarrow$	$\mathbb{T}$ - $\mathcal{W}_{2}$ $\downarrow$	TICA- $\mathcal{W}_{2}$ $\downarrow$
MLE-TARFlow (original)	$1.260$	$0.924$	$0.492$	$11.298$	$2.733$	$1.087$
MLE-TARFlow (rerun)	$1.414$	$0.933$	$0.585$	$12.126$	$2.808$	$1.107$
MLE-TARFlow †	$1.362$	$0.843$	$0.424$	$12.723$	$2.654$	$0.994$
MLE-iTARFlow	$\underline{1.068}$	$\underline{0.757}$	$\mathbf{0.378}$	$\underline{8.569}$	$\underline{2.379}$	$\underline{0.944}$
NFM-D (Ours)	$\mathbf{0.907}$	$\mathbf{0.750}$	$\underline{0.387}$	$\mathbf{7.396}$	$\mathbf{2.315}$	$\mathbf{0.890}$

Appendix

Table 12: Quantitative results on unseen 4- and 8-residue peptides, with a budget of

10^{4}

energy evaluations at

T=1

. Rows marked with † denote our reimplementation, used to control for differences in optimizer and training settings. Best results are in bold ; second-best results are underlined .

Figure 14: TICA projections for the ground truth MD, and NFM-D on six unseen octapeptides (reported in the titles).

Figure 15: Per-sequence

\mathcal{E}

\mathcal{W}_{2}

\mathbb{T}

\mathcal{W}_{2}

, and TICA-

\mathcal{W}_{2}

on the 8AA ManyPeptidesMD test set. Bars compare TARFlow and NFM-D; lower values indicate better agreement with the reference distribution.

Figure 16: Per-sequence

\mathcal{E}

\mathcal{W}_{2}

\mathbb{T}

\mathcal{W}_{2}

, and TICA-

\mathcal{W}_{2}

on the 4AA ManyPeptidesMD test set. Bars compare TARFlow and NFM-D; lower values indicate better agreement with the reference distribution.

Normalizing Flows

Normalizing flows are invertible neural transformations that model complex probability distributions with tractable likelihoods and sampling.

See all 70 papers

The Coupling Within: Flow Matching via Distilled Normalizing Flows
Mar 9, 2026David Berthelot, Tianrong Chen, Jiatao Gu +6Flow-Based Generative Modeling Flow Matching
Bridging Stochastic Flow Maps and Boltzmann Generators with Normalizing Flows
Sep 25, 2026Louis Grenioux, RuiKang OuYang, Luhuan WuBoltzmann Sampling Molecular Generation
Center for Computational Mathematics, Flatiron Institute, New York, NY, USA · Department of Engineering, University of Cambridge, Cambridge, UK · Department of Applied Mathematics and Statistics, Johns Hopkins University, Baltimore, MD, USA
Normalizing Flows with Iterative Denoising
Apr 21, 2026Tianrong Chen, Jiatao Gu, David Berthelot +2Normalizing Flows Autoregressive Image Generation
Apple Machine Learning Research.

Mar 9, 2026J/K move · Enter open · S save

The Coupling Within: Flow Matching via Distilled Normalizing Flows

David Berthelot, Tianrong Chen, Jiatao Gu, Marco Cuturi+5

Flow-Based Generative Modeling Flow Matching Normalizing Flows

Open paper AskPDF arXiv

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.

Abstract

Normalizing Flows

Explore similar work

The Coupling Within: Flow Matching via Distilled Normalizing Flows

Bridging Stochastic Flow Maps and Boltzmann Generators with Normalizing Flows

Normalizing Flows with Iterative Denoising