Bridging Stochastic Flow Maps and Boltzmann Generators with Normalizing Flows

Louis Grenioux; RuiKang OuYang; Luhuan Wu

stat.MLSep 25, 2026

Bridging Stochastic Flow Maps and Boltzmann Generators with Normalizing Flows

Authors: Louis Grenioux, RuiKang OuYang, Luhuan Wu

Organizations: 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

Abstract

Generating independent, equilibrium samples of molecular systems at scale remains a central obstacle in computational statistical mechanics. Boltzmann Generators address this by pairing a generative model with importance sampling to obtain consistent samples from the target distribution. We introduce Normalizing Flow Flow Maps (NF $^2$ M), which combines the strengths of recent stochastic flow maps with the tractability of classic normalizing flows to build a Boltzmann Generator. Unlike most methods, which correct the generative model only at the end, NF $^2$ M reweighs each denoising transition as generation proceeds, avoiding wasted compute on trajectories that are ultimately discarded. At each denoising step, a conditional normalizing flow proposes clean configurations given the current noisy state (a simpler task than sampling directly from the target) and its exact likelihood enables correcting each proposal toward the true denoising transition of the target Boltzmann distribution. This is in contrast to most existing methods, whose likelihoods are approximate or expensive to evaluate, undermining the statistical reliability of the correction. We establish consistency of the corrected transitions and bound how approximation errors propagate through the sampling chain. We evaluate NF $^2$ M on peptide systems, demonstrating improved sampling efficiency and sample quality.

Figures & tables

Figure 1: Conceptual illustration of the NF 2 M mechanism. Given a sample

x_{t}

drawn from the marginal distribution at noise level

t

, we first draw

M

samples

\{x_{0|t}^{m}\}_{m=1}^{M}

from the learned posterior

q_{0|t}^{\theta}(\cdot|x_{t})

. We then reweight them via Eq. 12 to construct the posterior

p_{s|t}^{\theta,M}

, which proposes a sample

x_{s}

at the next noise level

s

. We repeat this procedure until reaching the target distribution.

		SNIS	NF 2 M-IS	NF 2 M-MIS	FALCON	FALCON-A	F2D2	SCALLOP
ALA-2	$\mathcal{E}$ - $\mathcal{W}_{2}$	$0.768\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.031}$	$\textbf{0.132}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.013}$	$\underline{0.242}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.080}$	$0.759\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.047}$	$0.595\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.329}$	$0.890\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.030}$	$0.874\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.033}$
ALA-2	$\mathbb{T}$ - $\mathcal{W}_{2}$	$\textbf{0.242}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.012}$	$1.050\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.011}$	$1.239\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.006}$	$0.519\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.253}$	$0.862\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 1.089}$	$0.752\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.071}$	$\underline{0.374}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.043}$
ALA-3	$\mathcal{E}$ - $\mathcal{W}_{2}$	$0.268\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.067}$	$\underline{0.164}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.011}$	$\textbf{0.157}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.009}$	$0.181\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.049}$	$0.241\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.144}$	$0.553\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.133}$	$0.703\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.186}$
ALA-3	$\mathbb{T}$ - $\mathcal{W}_{2}$	$\textbf{0.362}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.019}$	$0.490\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.016}$	$0.465\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.007}$	$\underline{0.463}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.170}$	$1.100\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.706}$	$1.049\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.476}$	$0.847\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.196}$
ALA-4	$\mathcal{E}$ - $\mathcal{W}_{2}$	$1.515\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.351}$	$\textbf{0.988}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.027}$	$\underline{1.107}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.013}$	$1.763\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.192}$	$1.985\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.737}$	$2.748\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 2.096}$	$2.981\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.970}$
ALA-4	$\mathbb{T}$ - $\mathcal{W}_{2}$	$\textbf{0.777}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.131}$	$\underline{0.824}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.041}$	$0.882\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.051}$	$0.893\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.140}$	$1.462\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 1.092}$	$2.430\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.115}$	$2.368\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.189}$

Table 1: Energy and torus Wasserstein distances on ALA-

N

and chignolin systems, from 3 random seeds, computed on the test split vs.

10^{4}

reference samples. Best bold , second underlined .

		SNIS	NF 2 M-IS	NF 2 M-MIS	FALCON	FALCON-A	F2D2	SCALLOP
ALA-2	Inf. (m)	$2.0\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.0}$	$2.3\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.2}$	$2.1\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.2}$	$24.2\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.0}$	$24.1\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.1}$	$0.8\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.0}$	$0.8\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.0}$
ALA-2	Train (h)	$3.4$	$4.3$	$4.3$	$4.8$	$5.0$	$6.3$	$5.6$
ALA-3	Inf. (m)	$2.1\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.0}$	$2.7\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.0}$	$2.5\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.2}$	$\gtrsim 129$	$\gtrsim 129$	$2.0\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.0}$	$2.0\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.0}$
ALA-3	Train (h)	$4.7$	$6.0$	$6.0$	$7.1$	$7.2$	$6.5$	$5.8$
ALA-4	Inf. (m)	$2.8\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.1}$	$3.1\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.2}$	$2.7\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.2}$	$\gtrsim 202$	$\gtrsim 201$	$2.4\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.0}$	$2.4\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.0}$
ALA-4	Train (h)	$5.8$	$5.3$	$5.3$	$7.1$	$7.1$	$5.6$	$5.1$

Table 2: Inference time at the matched target-energy budget (sampling + post-hoc MALA ladder) and total training wall-clock across all required stages (mean

\pm

std, 3 seeds). Chignolin FALCON/FALCON-A ran at a reduced

10^{5}

-sample budget (vs.

10^{6}

elsewhere;

10^{4}

for NF 2 M).

Appendix figures & tables34 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 2: Values of

a_{s,t}

and

\gamma_{s,t}

along the time grid of Karras et al. (2022) .

Panel of Figure 3 — Figure 3: Candidate sharing and indexing in NF 2 M-MIS. The example uses $N=3$ sources, $M=3$ candidates per source, and $C=2$ components per mixture; two receiving trajectories are shown. Candidate $x_{0}^{j,m}$ is indexed by its source $j$ and position $m$ within that source’s batch. For receiving trajectory $i$ , a component $c\in\mathcal{C}(i)$ is itself a source index: $c=i$ always appears (the trajectory’s “own” component), and the remaining slots are drawn by posterior proximity. The middle batch, drawn once from source $j=2$ , is reused by both trajectories: it serves as component $c=2$ of trajectory $1$ ’s mixture and as the own component $c=2$ of trajectory $2$ ’s mixture, receiving weight $\Omega_{1,2,m}$ and $\Omega_{2,2,m}$ respectively — the same $M$ flow draws inform two different reverse steps. Each trajectory samples $(c,m)$ from its own categorical weights and applies the endpoint-conditioned Gaussian transition. When a source is selected more than once within the same mixture, the repeat instead draws a fresh, independent batch, indexed by the occurrence counter $r(i,c)$ introduced below.

\Omega_{i,c,m}\propto\omega_{i,c}\,\tilde{\Lambda}_{i,c,m},\qquad\tilde{\Lambda}_{i,c,m}=\frac{p_{t_{k}|0}(x_{t_{k}}^{i}\mid x_{0}^{i,c,m})\,\exp\{-\mathcal{E}(x_{0}^{i,c,m})\}}{\sum_{c^{\prime}\in\mathcal{C}(i)}\omega_{i,c^{\prime}}\,q^{\theta}_{0|t_{k}}(x_{0}^{i,c,m}\mid x_{t_{k}}^{c^{\prime}})}

Appendix

Algorithm 2 NF 2 M-MIS sampling.

Figure 4: Average effective sample size of

q^{\theta}_{0|t}(\cdot|x_{t})

against

p_{0|t}(\cdot|x_{t})

as a function of

t

, for NF 2 M-IS and NF 2 M-MIS across different numbers of importance candidates or mixture components.

M^{\star}

denotes the budget reported in Tab. 18 .

Figure 5: Inference time and distributional distances for all systems, comparing NF 2 M-IS and NF 2 M-MIS across different numbers of importance candidates or mixture components.

M^{\star}

denotes the budget reported in Tab. 18 .

	ALA-2	ALA-3	ALA-4	ALA-6	Chignolin
MAExL $\omega$ (in-loop)	$1.02{\times}10^{-2}$	$7.48{\times}10^{-3}$	$6.55{\times}10^{-3}$	$4.68{\times}10^{-3}$	$3.46{\times}10^{-3}$
MALA $\tau$ (post-hoc)	$5.27{\times}10^{-5}$	$2.75{\times}10^{-5}$	$2.14{\times}10^{-5}$	$1.10{\times}10^{-5}$	$6.01{\times}10^{-6}$

Appendix

Table 3: Per-system calibrated step sizes (Robbins–Monro, target acceptance 0.574).

		1024 MAExL steps		512 MAExL steps		0 MAExL steps
		NF 2 M-IS	NF 2 M-MIS	NF 2 M-IS	NF 2 M-MIS	NF 2 M-IS	NF 2 M-MIS
ALA-2 ( $d=66$ )	$\mathcal{E}$ - $\mathcal{W}_{2}$ $\downarrow$	$0.127\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.019}$	$0.186\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.042}$	$0.132\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.013}$	$0.242\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.080}$	$0.256\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.054}$	$0.367\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.035}$
	$\mathbb{T}$ - $\mathcal{W}_{2}$ $\downarrow$	$1.044\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.019}$	$1.241\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.010}$	$1.050\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.011}$	$1.239\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.006}$	$1.055\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.009}$	$1.257\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.023}$
	Inference (s) $\downarrow$	$129.2\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 5.8}$	$118.7\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.2}$	$137.2\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 14.5}$	$125.8\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 14.4}$	$122.8\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 14.4}$	$116.6\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 14.6}$
ALA-3 ( $d=99$ )	$\mathcal{E}$ - $\mathcal{W}_{2}$ $\downarrow$	$0.167\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.016}$	$0.151\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.030}$	$0.164\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.011}$	$0.157\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.009}$	$0.215\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.037}$	$0.212\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.033}$
	$\mathbb{T}$ - $\mathcal{W}_{2}$ $\downarrow$	$0.497\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.011}$	$0.471\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.020}$	$0.490\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.016}$	$0.465\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.007}$	$0.505\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.022}$	$0.462\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.020}$

Appendix

Table 4: MAExL ablation on ALA-

N

systems, with

N=2,3,4,6

, and on chignolin. Each column is the NF 2 M configuration of Tab. 10 , rerun with the number of in-loop MAExL steps per level given in the header; the 512-step columns are those of Tab. 1 . Results are obtained from three random seeds, reported as mean

\pm

std . Energy and torus Wasserstein distances are computed on the test split against

10^{4}

reference samples, from a proposal budget of

10^{4}

samples. Inference time covers sampling and the post-hoc MALA ladder.

Figure 6: Sampling and density evaluation share one stack. The same

L

blocks (a deep block conditioned on

y_{t}

, shallow blocks refining) run in two directions. (a) Autoregressively, from

z\sim\mathcal{N}(0,\mathrm{I})

up to a sample

x_{0}\sim p^{\theta}_{0|t}(\cdot|y_{t})

N

sequential per-atom steps per block, under a KV cache. (b) In parallel, from a given

x_{0}

down to

\log p^{\theta}_{0|t}(x_{0}|y_{t})

: one causal pass per block. Here

y_{t}

denotes the noisy configuration

x_{t}

of Sec. 3.2 , relabelled so the conditioning is never mistaken for another stage of the flow’s own

x^{0},\ldots,x^{T-1}

sequence. The green/purple arrows are the inverse/forward instance of the same per-atom update between every pair of adjacent levels,

x^{\ell}_{i}=\mu_{i}(x^{\ell}_{<i})+\sigma_{i}(x^{\ell}_{<i})\cdot x^{\ell+1}_{i}

(bottom left),

x^{\ell+1}_{i}

is literally

z_{i}

only at the very first level.

y_{t}

is shown entering as a causal prefix at the deep block only. Fig. 7 shows the cross-attention alternative and the every-block scope TarFlow uses instead ( Sec. A.5.2 ), plus the AdaLN-Zero time conditioning we add on top of both shown here only as the dashed

t

arrows reaching every block, deep and shallow alike (the modulation itself is not drawn). Both directions read the same conditioning cache, built once from

(y_{t},t)

and shared across every particle drawn at that step ( Sec. A.5.3 ). Layout follows Gu et al. (2026) ’s own figure.

Figure 7: The two conditioning mechanisms, inside one block. (a) Cross-attention:

x_{t}

is embedded once per block, and each of the

K

layers still projects its own keys and values from that shared embedding and attends to it from a query built out of

x_{0}

; every sub-layer also carries an AdaLN-Zero modulation driven by the time embedding

c

, with every modulation linear zero-initialised, so the block is the identity map at init. (b) Prefix:

x_{t}

instead joins the sequence as a causal prefix ( Eq. 82 ). A plain lower-triangular mask over the concatenated

2N

tokens already gives every target atom access to every conditioning atom while keeping the prefix invisible to

x_{0}

, so the added cost is a longer attention span and

W_{\mathrm{cond}}

, not a new sub-layer. In both variants the conditioning state is fixed across every particle sharing

x_{t}

, hence cacheable ( Sec. A.5.3 ).

		Added by conditioning
	Flow backbone	cross-attention	prefix
$\log p$	$K_{\mathrm{tot}}(N^{2}C+NC^{2})$	$\tfrac{1}{m}K_{\mathrm{cond}}(N^{2}C+NC^{2})$	$K_{\mathrm{cond}}(3N^{2}C+NC^{2})$
$x\sim p$	$K_{\mathrm{tot}}(\tfrac{1}{2}N^{2}C+NC^{2})$	$\tfrac{1}{m}K_{\mathrm{cond}}(N^{2}C+NC^{2})$	$K_{\mathrm{cond}}(2N^{2}C+NC^{2})$
additional params.	—	$\tfrac{1}{m}K_{\mathrm{cond}}\cdot 4C^{2}$	$0$

Appendix

Table 5: Cost of one call into the flow, dropping constants, with nothing precomputed: every call rebuilds the conditioning state from

x_{t}

. Notation as in the text. The flow backbone column is what every variant pays for the flow itself; the other two are what conditioning adds on top of it.

Per particle
	cross-attention		prefix
	uncached	amortised	uncached	amortised
density evaluation	$\tfrac{1}{m}K_{\mathrm{cond}}(N^{2}C+NC^{2})$	$\tfrac{1}{m}K_{\mathrm{cond}}N^{2}C$	$K_{\mathrm{cond}}(3N^{2}C+NC^{2})$	$K_{\mathrm{cond}}N^{2}C$
sampling	$\tfrac{1}{m}K_{\mathrm{cond}}(N^{2}C+NC^{2})$	$\tfrac{1}{m}K_{\mathrm{cond}}N^{2}C$	$K_{\mathrm{cond}}(2N^{2}C+NC^{2})$	$K_{\mathrm{cond}}N^{2}C$
Once per distinct $x_{t}$
build	—	$\tfrac{1}{m}K_{\mathrm{cond}}NC^{2}$	—	$K_{\mathrm{cond}}(N^{2}C+NC^{2})$

Appendix

Table 6: Amortising the conditioning.

x_{t}

is fixed across every

x_{0}

drawn from the same posterior, so the conditioning state is built once per distinct

x_{t}

and read by every particle that conditions on it, in both directions of the flow. Entries are the cost added to the flow backbone; notation as in Tab. 5 . Amortising deletes the projection terms from the per-particle cost, leaving attention against the cached state; what remains is the state itself, which scales with the number of conditioning rows held live rather than with the number of particles.

	ALA-2	ALA-3	ALA-4	ALA-6	Chignolin
Channels	256	256	256	384	512
Blocks	4	6	6	6	8
Layers / block	4	6	6	6	8

Appendix

Table 7: Unconditional TarFlow architecture.

	ALA-2	ALA-3	ALA-4	ALA-6	Chignolin
Channels	256	256	256	384	512
Blocks	6	6	6	6	6
Head dim	64	64	64	64	64
Top-block depth $n$	4	14	14	14	28

Appendix

Table 8: Unconditional StarFlow architecture. Every variant (here and in Tab. 9 ) follows the pattern layers/block

=[2,2,2,2,n,n]

: the first four blocks are fixed at 2 layers, and only the last two (“top-block depth”

n

) vary.

Variant	Conditioning	ALA-2	ALA-3	ALA-4	ALA-6	Chignolin
T1	cross-attn (stride 3–4)	256/4 $\times$ 3/3	256/6 $\times$ 4/4	256/6 $\times$ 4/4	384/6 $\times$ 4/4	–
K1	cross-attn (stride 1)	256/4 $\times$ 3/1	256/6 $\times$ 4/1	256/6 $\times$ 4/1	384/6 $\times$ 4/1	–
W1	cross-attn (stride 3–4)	320/4 $\times$ 3/3	320/6 $\times$ 4/4	320/6 $\times$ 4/4	448/6 $\times$ 4/4	–
S1	prefix	256/3	256/10	256/10	384/10	512/20
S1k	cross-attn (stride 3–4)	256/3/3	256/10/4	256/10/4	384/10/4	512/20/4
S2	cross-attn (stride 4)	256/4/4	256/14/4	256/14/4	384/14/4	512/28/4

Appendix

Table 9: NF 2 M architecture sweep. TarFlow variants report channels / blocks

\times

layers-per-block / conditioning stride; StarFlow variants report channels / top-block depth

n

(layers/block

=[2,2,2,2,n,n]

) / conditioning stride. “–” marks variants that failed to train.

NF 2 M-IS
	ALA-2	ALA-3	ALA-4	ALA-6	Chignolin
Variant	W1	W1	T1	S2	S1
Architecture	320/4 $\times$ 3/3	320/6 $\times$ 4/4	256/6 $\times$ 4/4	384/14/4	512/20 (prefix)
$K$	5	8	8	5	5
Training (h)	4.3	6.0	5.3	12.8	21.7
NF 2 M-MIS

Appendix

Table 10: Selected NFM checkpoint per system, with its architecture (same notation as Tab. 9 : TarFlow variants as channels / blocks

\times

layers-per-block / conditioning stride, StarFlow variants as channels / top-block depth / conditioning stride or “prefix”).

K

is the number of inference-time levels; training time is the wall-clock of the 1000-epoch run of the checkpoint actually loaded.

	ALA-2	ALA-3	ALA-4	ALA-6	Chignolin
$d_{\text{model}}$	192	192	192	192	384
Layers	6	6	6	6	12
Heads	6	6	6	6	12
MLP ratio	4.0	4.0	4.0	4.0	4.0
Time-embedding dim	64	64	64	64	128

Appendix

Table 11: DiT backbone for the phase-1 velocity teacher.

Single-head (ESD, FALCON, FALCON-A)
	ALA-2	ALA-3	ALA-4	ALA-6	Chignolin
$d_{\text{model}}$	192	192	192	192	384
Blocks	6	6	6	6	12
Heads	6	6	6	6	12
Parameters	4.1M				32M
Dual-head (LSD, F2D2, SCALLOP)

Appendix

Table 12: Flow-map backbone architecture. The single-head network serves the ESD pretraining objective, FALCON and FALCON-A; the dual-head network serves LSD’s students, F2D2 and SCALLOP, and counts its divergence-readout blocks only (the velocity reads straight off the shared trunk). Both carry a learned atom-type embedding, a bond-distance attention bias and centre-of-mass removal.

	FM	ESD / FALCON / FALCON-A	LSD / F2D2 / SCALLOP
Optimizer	AdamW	AdamW	RAdam
Learning rate	$10^{-4}$	$5{\times}10^{-4}$	$10^{-4}$
Weight decay	$10^{-4}$	$10^{-4}$	0
Schedule	sqrt-anneal	cosine	flat-then-sqrt
Warmup	2.5%	5%	none
Grad-norm clip	10.0	10.0	10.0

Appendix

Table 13: Flow-map optimizer hyperparameters, per phase. Each pretraining objective takes the optimizer settings of the phase it initializes, so a warm start lands its student in a basin that student’s own optimizer is already tuned for. All three columns use

\beta=(0.9,0.999)

\varepsilon=10^{-8}

, an EMA of decay

0.999

and batch size 1024. “Flat-then-sqrt” holds the rate at its maximum for an absolute number of steps and then decays as

1/\sqrt{\text{step}}

; the onset is 10,000 steps on ALA-2 and ALA-4 and 20,000 elsewhere, read off the reference’s molecular configurations. RAdam’s variance rectification is why that schedule carries no warmup.

	ESD	FALCON	FALCON-A
Self-consistency weight $\lambda_{\text{avg}}$	1.0 †	1.0 †	1.0 †
Cycle-consistency weight $\lambda_{r}$	–	10.0	10.0
Split point $\lambda$	–	–	0.5
Diagonal share	0.75 †	0.75 †	1.0
Time pair	unordered square
Time conditioning	pair	pair	pair

Appendix

Table 14: ESD / FALCON / FALCON-A loss hyperparameters, the single-head family. “Diagonal share” is the fraction of each batch drawn at

s=t

and given to the instantaneous term, its complement carrying the self-consistency term; it is an exact-count split of every batch rather than a Bernoulli mask. FALCON-A is the exception at 1.0: forcing

s=t

degenerates its additivity target to

0=0

rather than to flow matching, so it keeps an explicit instantaneous term and evaluates both on every sample. Values marked

\dagger

are not stated numerically by any of the three papers and were never swept. “–” marks a term the objective does not have.

	LSD	F2D2	SCALLOP
Self-consistency weight $\lambda_{\text{avg}}$	1.0 †	1.0	1.0
Diagonal share	0.5	0.5	0.5
Time pair	unordered square
Time conditioning	anchor + step	anchor + step	anchor + step
Average-velocity derivative	live	live	live
Divergence derivative	–	stopped	stopped

Appendix

Table 15: LSD / F2D2 / SCALLOP loss hyperparameters, the dual-head family, which alone carries a divergence head. Same conventions as Tab. 14 :

\dagger

marks values neither paper states numerically, “–” marks a term the objective does not have.

Phase 1 — velocity teacher
	ALA-2	ALA-3	ALA-4	ALA-6	Chignolin
Velocity	177,510	162,378	139,389	206,125	117,661
Phase 1.5 — flow-map pretraining
ESD	130,768	226,158	208,892	362,830	293,385
LSD	68,937	70,578	55,023	77,465	270,018
Phase 2 — distillation

Appendix

Table 16: Training steps per system and phase, at batch size 1024. Phase 1 is one velocity teacher per system, shared by all four methods; phase 1.5 is one pretraining run per objective, shared by the two students that read it (ESD

\to

FALCON, FALCON-A; LSD

\to

F2D2, SCALLOP). Each entry is what the reported checkpoint was actually trained for.

Seconds per step at batch 1024 (one H100)
	ALA-2	ALA-3	ALA-4	ALA-6	Chignolin
FM	0.0383	0.0380	0.0416	0.0581	0.1453
ESD	0.0332	0.0341	0.0361	0.0470	0.1040
LSD	0.0284	0.0320	0.0339	0.0472	0.1130
FALCON	0.0390	0.0430	0.0460	0.0560	0.1180
FALCON-A	0.0260	0.0450	0.0510	0.0790	0.1890

Appendix

Table 17: Measured per-step cost of every flow-map objective, and the two wall-clock quantities the step budget is solved against. Costs are seconds per optimizer step at batch size 1024 on one H100, each measured at the compilation mode its production run uses; chignolin is measured on its own 4-rank configuration at 256 rows per rank, so its column is directly comparable to the others per step.

	ALA-2	ALA-3	ALA-4	ALA-6	Chignolin
$M$	14	12	14	20	20

Appendix

Table 18: NF 2 M-IS candidate budget

M

for the selected (variant,

K

) of Tab. 10 .

	ALA-2	ALA-3	ALA-4	ALA-6	Chignolin
$C$	4	4	2	3	2
$M$	2	3	7	5	6

Appendix

Table 19: NF 2 M-MIS selected mixture configuration.

C

and

M

are the validation-selected mixture size and per-particle candidate count in every system.

	ALA-2	ALA-3	ALA-4	ALA-6	Chignolin
NFE	4	8	8	16	16

Appendix

Table 20: Flow-map inference steps (NFEs), shared by all four methods. The four ALA values are the ones both papers report; chignolin appears in neither and reuses the largest.

SNIS: Unconditional TarFlow/StarFlow
	ALA-2	ALA-3	ALA-4	ALA-6	Chignolin
Selected flow	TarFlow	TarFlow	StarFlow	StarFlow	TarFlow
Sample	131072	131072	65536	32768	16384
Log-prob	131072	131072	32768	32768	32768
NF 2 M (measured knee, bf16)
Selected variant (IS / MIS)	W1	W1	T1	S2 / S1k	S1

Appendix

Table 21: Evaluation chunk sizes, in rows, on an A100-80GB. The SNIS and NF 2 M architectures are selected per system based on validation metrics. NF 2 M variant labels (T1, W1, S1, S1k, S2) are defined in Tab. 9 .

		Velocity	FALCON	FALCON-A	F2D2	SCALLOP
ALA-2 ( $d=66$ )	$\mathcal{E}$ - $\mathcal{W}_{2}$ $\downarrow$	$0.765\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.097}$	$0.787\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.091}$	$0.966\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.285}$	$1.132\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.018}$	$1.110\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.014}$
	$\mathbb{T}$ - $\mathcal{W}_{2}$ $\downarrow$	$0.274\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.032}$	$0.524\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.212}$	$0.892\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.932}$	$0.774\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.059}$	$0.391\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.029}$
	Inference (s) $\downarrow$	$1375.3\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 10.4}$	$1477.6\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 2.0}$	$1461.3\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 11.9}$	$70.6\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.1}$	$71.5\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.6}$
	Training (h) $\downarrow$	$1.9$	$4.8$	$5.0$	$6.3$	$5.6$
ALA-3 ( $d=99$ )	$\mathcal{E}$ - $\mathcal{W}_{2}$ $\downarrow$	$0.410\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.219}$	$0.250\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.031}$	$0.445\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.410}$	$0.648\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.211}$	$0.596\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.448}$
	$\mathbb{T}$ - $\mathcal{W}_{2}$ $\downarrow$	$0.677\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.168}$	$0.468\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.144}$	$1.102\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.540}$	$1.072\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.370}$	$0.946\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.117}$

Appendix

Table 22: Flow-map samplers against the velocity field they are distilled from, on ALA-

N

systems with

N=2,3,4,6

. Results are obtained from three random seeds, reported as mean

\pm

std . Energy and torus Wasserstein distances are computed on the test split against

10^{4}

reference samples.

In-loop MAExL (512 steps) and post-hoc MALA (128 steps)
	ALA-2	ALA-3	ALA-4	ALA-6	Chignolin
NF 2 M-IS	6 (4.1%)	14 (8.6%)	17 (9.4%)	23 (5.0%)	108 (3.8%)
NF 2 M-MIS	9 (7.0%)	14 (9.5%)	17 (10.7%)	36 (8.8%)	112 (5.0%)
Budget-matched post-hoc MALA only ( Tab. 25 )
TarFlow/StarFlow	8 (7.2%)	14 (10.6%)	17 (10.1%)	35 (9.3%)	166 (6.8%)
Velocity	435 (24.2%)	322 (9.8%)	470 (9.9%)	499 (5.0%)	–

Appendix

Table 23: Langevin refinement cost at a matched target-energy budget: 512 MAExL steps for NF 2 M and the budget-matched post-hoc MALA ladder for all other methods. Each entry gives the wall-clock time spent in the Langevin kernels and its share of the sampler’s total inference time (means over 3 seeds, A100-80GB).

		SNIS	NF 2 M-IS	NF 2 M-MIS	FALCON	FALCON-A	F2D2	SCALLOP
ALA-2 ( $d=66$ )	$\mathcal{E}$ - $\mathcal{W}_{2}$ MALA $\downarrow$	$0.77\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.03}$	$\mathbf{0.13}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.01}$	$\underline{0.24}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.08}$	$0.76\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.05}$	$0.60\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.33}$	$0.89\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.03}$	$0.87\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.03}$
	$\mathcal{E}$ - $\mathcal{W}_{2}$ raw $\downarrow$	$1.07\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.15}$	$\mathbf{0.20}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.05}$	$\underline{0.24}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.07}$	$0.64\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.07}$	$1.50\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 1.05}$	$2.84\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.42}$	$2.77\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.36}$
	$\mathbb{T}$ - $\mathcal{W}_{2}$ MALA $\downarrow$	$\mathbf{0.24}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.01}$	$1.05\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.01}$	$1.24\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.01}$	$0.52\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.25}$	$0.86\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 1.09}$	$0.75\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.07}$	$\underline{0.37}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.04}$
	$\mathbb{T}$ - $\mathcal{W}_{2}$ raw $\downarrow$	$\mathbf{0.25}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.01}$	$1.05\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.01}$	$1.24\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.01}$	$0.52\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.20}$	$0.92\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.97}$	$0.78\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.06}$	$\underline{0.40}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.03}$
	Inference MALA (s) $\downarrow$	$117\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0}$	$137\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 15}$	$126\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 14}$	$1453\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 2}$	$1448\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 5}$	$46\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0}$	$46\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0}$
	Inference raw (s) $\downarrow$	$109\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0}$	$137\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 15}$	$126\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 14}$	$1444\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 2}$	$1440\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 4}$	$37\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0}$	$38\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0}$

Appendix

Table 24: Results on ALA-

N

systems, with

N=2,3,4,6

, and on chignolin, with and without post-hoc MALA refinement, at a matched target-energy budget Tab. 25 . Results are obtained from three random seeds, reported as mean

\pm

std . Energy and torus Wasserstein distances are computed on the test split against

10^{4}

reference samples. The proposal budget differs by method:

10^{4}

samples for NF 2 M,

10^{5}

for FALCON and FALCON-A on chignolin, and

10^{6}

for every other column. Inference time covers sampling, plus the post-hoc ladder for the MALA rows. Best in bold , second best underlined , within each row.

Figure 8: Energy Wasserstein distance against the matched target-energy budget, on the four ALA systems and chignolin. Every point on a panel spends the same number of target-energy calls. The three budgets are the 0-, 512- and 1024-MAExL-step settings of Tab. 4 for NF 2 M, and the equalizing post-hoc MALA ladders of Tab. 25 for the SNIS-based samplers. Means over three seeds.

Figure 9: Torus Wasserstein distance against the matched target-energy budget, same protocol and seeds as Fig. 8 .

Figure 10: Inference wall-clock against the matched target-energy budget, same protocol and seeds as Fig. 8 . FALCON and FALCON-A are omitted, their per-step Jacobian log-determinant putting them off the scale of every panel.

		MAExL steps
System	Methods	0	512	1024
ALA-2 ( $d=66$ )	SNIS, FALCON, FALCON-A, F2D2, SCALLOP	114	4212	8308
ALA-3 ( $d=99$ )		127	4231	8327
ALA-4 ( $d=126$ )		135	4239	8335
ALA-6 ( $d=189$ )		134	4238	8334
Chignolin ( $d=414$ )	SNIS, F2D2, SCALLOP	129	4228	8324

Appendix

Table 25: Number of post-hoc MALA steps run for the SNIS-based samplers at each matched budget. Each step count is chosen so that the SNIS-based total,

P

energy calls for the importance weights plus

n\,(S+1)

score calls on

n=10^{4}

chains with one score call counted as two energy calls, matches the budget of the corresponding NF 2 M setting. All SNIS-based samplers use

P=10^{6}

proposals, except FALCON and FALCON-A on chignolin, which use

P=10^{5}

and therefore take longer ladders for the same total.

		SNIS	NF 2 M-IS	NF 2 M-MIS	SBG-SMC StarFlow	SBG-SMC TarFlow
ALA-2 ( $d=66$ )	$\mathcal{E}$ - $\mathcal{W}_{2}$ $\downarrow$	$0.752\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.027}$	$\mathbf{0.132}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.013}$	$\underline{0.242}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.080}$	$0.737\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.142}$	$0.686\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.042}$
	$\mathbb{T}$ - $\mathcal{W}_{2}$ $\downarrow$	$\mathbf{0.245}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.014}$	$1.050\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.011}$	$1.239\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.006}$	$0.387\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.075}$	$\underline{0.385}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.021}$
	Inference (s) $\downarrow$	$128.8\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.4}$	$137.2\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 14.5}$	$125.8\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 14.4}$	$242.1\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 9.7}$	$263.2\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 11.4}$
	Training (h) $\downarrow$	$3.4$	$4.3$	$4.3$	$3.6$	$3.4$
ALA-3 ( $d=99$ )	$\mathcal{E}$ - $\mathcal{W}_{2}$ $\downarrow$	$0.225\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.047}$	$\underline{0.164}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.011}$	$\mathbf{0.157}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.009}$	$2.715\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 3.857}$	$0.524\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.072}$
	$\mathbb{T}$ - $\mathcal{W}_{2}$ $\downarrow$	$\mathbf{0.371}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.016}$	$0.490\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.016}$	$\underline{0.465}\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.007}$	$1.256\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 1.256}$	$0.470\color[rgb]{0.5,0.5,0.5}{\scriptscriptstyle\pm 0.044}$

Appendix

Table 26: Results on ALA-

N

systems, with

N=2,3,4,6

, and on chignolin. Results are obtained from three random seeds, reported as mean

\pm

std . Energy and torus Wasserstein distances are computed on the test split against

10^{4}

reference samples. The proposal budget differs by method:

10^{4}

samples for NF 2 M and SBG-SMC, and

10^{6}

for SNIS. The SNIS column reports, per system, the better of the TarFlow and StarFlow checkpoints; the SBG-SMC columns are reported per architecture. Inference time covers sampling and the post-hoc MALA ladder; training time is the wall-clock of the corresponding model’s training run. Best in bold , second best underlined .

Explore similar work

Mar 15, 2026

Sampling Boltzmann distributions via normalizing flow approximation of transport maps

In a celebrated paper \cite{noe2019boltzmann}, Noé, Olsson, Köhler and Wu introduced an efficient method for sampling high-dimensional Boltzmann distributions arising in molecular dynamics via normalizing flow approximation of transport maps. Here, we place this approach on a firm mathematical foundation. We prove the existence of a normalizing flow between the reference measure and the true Boltzmann distribution up to an arbitrarily small error in the Wasserstein distance. This result covers general Boltzmann distributions from molecular dynamics, which have low regularity due to the presence of interatomic Coulomb and Lennard-Jones interactions. The proof is based on a rigorous construction of the Moser transport map for low-regularity endpoint densities and approximation theorems for neural networks in Sobolev spaces. Numerical simulations for a simple model system and for the alanine dipeptide molecule confirm that the true and generated distributions are close in the Wasserstein distance. Moreover we observe that the RealNVP architecture does not just successfully capture the equilibrium Boltzmann distribution but also the metastable dynamics.

Zia Ur Rehman, Gero Friesecke

Department of Mathematics, School of Computation, Information and Technology, Technical University of Munich, Boltzmannstrasse 3, 85748 Garching, Germany

Normalizing Flows Molecular Dynamics Simulation Optimal Transport

Ask PDF

Jun 27, 2026

Few-Step Boltzmann Generators via Scalable Likelihood Flow Maps

Recent progress in flow-based generative modeling has led to models that output high-quality samples while using only a small number of function evaluations. However, at present, there is a lack of similar advances in estimating the model likelihood. In particular, most existing methods either rely on restrictive architectures that enable exact calculations, or use stochastic approximations such as Hutchinson's trace estimator that introduce substantial variance. In this work, we introduce SCAlable LikeLihood distillation of flOw maPs (SCALLOP). SCALLOP builds on the recently proposed F2D2, a likelihood flow map model that can generate samples and their densities in a small number of function evaluations. While F2D2 uses Hutchinson's estimator during training, we introduce an alternative and more scalable likelihood distillation objective that is Hutchinson-free and admits a vectorized formulation. Empirically, we demonstrate the effectiveness of SCALLOP as a Boltzmann generator in molecular science, and further validate its benefit on image datasets. SCALLOP significantly reduces both training variance and training time while consistently improving performance compared to F2D2, and is competitive with the state-of-the-art while achieving up to 10x inference speedup over the fastest baseline.

RuiKang OuYang, Hanlin Yu, Xinyue Ai +7

University of Cambridge · University of Helsinki · Peking University +3

Flow Map Learning Generative Modeling Molecular Generation Flow-Based Generative Modeling

Ask PDF

Sep 26, 2026

FUND: Density Flow for Sampling Unnormalised Distributions

Efficient sampling from Boltzmann distributions is central to modelling complex physical systems. Markov Chain Monte Carlo (MCMC) methods suffer from critical slowing down, high autocorrelation, and poor mode-mixing, limiting their scalability. Recent advances, like Boltzmann Generators, offer a promising alternative but remain constrained by costly MCMC-based training, inefficient sampling, and poor ergodicity. We introduce an algorithm for learning Boltzmann distributions that does not require any true samples for training. Our approach draws inspiration from flow matching but departs fundamentally from sample-trajectory matching to distribution-trajectory matching. The algorithm iteratively reshapes the target distribution, using model generated samples to guide learning and ensure comprehensive mode coverage. We validate our method on standard benchmarks, including a 2D Gaussian mixture, Many-Well distributions, and high-dimensional scalar

φ^4

theory. The proposed approach not only improves sampling performance and accuracy over traditional MCMC and flow-based baselines but also establishes a new method for sample-free learning of complex physical distributions.

Vikas Kanaujia, Vipul Arora

Department of Electrical Engineering IIT Kanpur · Department of Electrical Engineering KU Leuven

Flow-Based Generative Modeling Flow Matching Markov Chain Monte Carlo

Ask PDF