FALCON: Few-step Accurate Likelihoods for Continuous Flows
Authors: Danyal Rehman
Tara Akhound-Sadegh
Artem Gazizov
Yoshua Bengio
Alexander Tong
Venue: ICLR 2026 Oral (2026)
Keywords: Generative Models
Flow Matching
Boltzmann Generators
AI for Science
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: Few-step Flow Matching with Accurate Likelihoods for Scalable Boltzmann Generators

Abstract: Scalable sampling of molecular states in thermodynamic equilibrium is a long-standing challenge in statistical physics. Boltzmann Generators tackle this problem by pairing a generative model, capable of exact likelihood computation, with importance sampling to obtain consistent samples under the target distribution. Current Boltzmann Generators primarily use continuous normalizing flows (CNFs) trained with flow matching for efficient training of powerful models. However, likelihood calculation for these models is extremely costly, requiring thousands of function evaluations per sample, severely limiting their adoption. In this work, we propose Few-Step Accurate Likelihoods for Continuous Flows (FALCON), a method which allows for few-step sampling with a likelihood accurate enough for importance sampling applications by introducing a hybrid training objective that encourages invertibility. We show FALCON outperforms state-of-the-art normalizing flow models for molecular Boltzmann sampling and is two orders of magnitude faster than the equivalently performing CNF model. FALCON code is available at: https://github.com/danyalrehman/FALCON.

Published as a conference paper at ICLR 2026
FALCON: FEW-STEP ACCURATE LIKELIHOODS FOR CONTINUOUS FLOWS
Danyal Rehman" Tara
Tara Akhound-Sadegh" 1.4 Artem Gazizov,, Yoshua Bengio 1,26 B,.
Alexander Tong3
Mila Québec Al Institute, UUniversité de Montréal, 3 AITHYRA, McGill University, 'Harvard University, 'CIFAR Senior Fellow
ABSTRACT
Scalable sampling of molecular states in thermodynamic equilibrium is long- standing challenge statistical physics. Boltzmann Generators tackle this problem by pairing generative model, capable of exact likelihood computation, with importance sampling to obtain consistent samples under the target distribution. Current Boltzmann Generators primarily use continuous normalizing flows (CNFs) trained with flow matching for efficient training of powerful models. However, likelihood calculation for these models is extremely costly, requiring thousands of function evaluations per sample, severely limiting their adoption. In this work, we propose FEW-STEP ACCURATE LIKELIHOODS FOR CONTINUOUS FLOwS (FALCCON,, a method which allows for few-step sampling with likelihood accurate enough for importance sampling applications by introducing hybrid training objective that encourages invertibility. We show FALCON outperforms state-of-the-ar normalizing flow models for molecular Boltzmann sampling and is over two orders of magnitude faster than the equivalently performing CNF model. FALCON code is available at https:/ /ttt:/gg..tt...///// thub. com/ comm/annallrehmaan FALCO..
INTRODUCTION
Sampling molecular configurations from the Boltzmann distribution p(z) x exp(-E(x)) where (x)) is the potential energy of a configuration æ, a foundational and long-standing challenge in statistical physics. The ability to generate samples according to this distribution is the foundation for determining many other observables such as free energies and heat capacities real-world behaviour. Consequently, efficient Boltzmann sampling is essential for progress in large range of areas, from characterizing the function of biomolccules, to accelerating drug design, and discovering novel materials (Frenkel & Smit, 2023; Liu, 2001; Ohno et al., 2018; Stoltz et al., 2010).
govern
The difficulty of this task arises from the structure of the energy for molecules of interest, The energy landscape is dighndimensional and non-smooth with many local energy minima. These rugged energies severely challenge classical simulation-based methods like Molecular Dynamics (MD) (Leimkuhler & Matthews, 2015) and Monte Carlo Markov Chains (MCMC) as they become easily trapped in local minima, requiring computationally inaccessible number of steps to mix between modes. These samplers generate many correlated samples, creating large inefficiencies. where an ideal sampler would generate i.i.d. samples from the underlying data distribution, p(z).
Boltzmann Generators (BGs) have emerged as way to address this inefficiency by amortizing the cost through the training of a generative model to learn to sample from (=(x) close to p(x). These samples can then corrected to p(r) using self-normalized importance sampling (SNIS) (Noé et al., 2019). SNIS requires efficient access to (x())) and E(x) for every sample drawn x ~ (x() for the correction step, but guarantees statistical consistency of the corrected samples, as illustrated in Fig. I (Noé et al., 2019; Liu, 2001).
The main design choice in BGs is which type of generator to use. Modern BGs (Klein et al., 2023; Klein & Noe, 2024) primarily make use of generators based on continuous normalizing flows (CNFs) (Chen et al., 2018; Grathwohl et al., 2019) due to their expressive power, ease of training,
"Correspondence danyal. rehman@mi la. quebeec, atong@ai thyra.t

---

Published as a conference paper at ICLR 2026
True Density
True Density
Self-romalized Importance Sempling
Fi
21
24h (N(1))
((Ts; 8)
20
Source
Biased Target
Debiased Target
Figure I: Flow map learns from biased data, with SNIS re-weighting generated samples consistent with the Boltzmann distribution, approaching equality with infinite samples under mild regularity conditions.
and flexibility of parameterization (Köhler al., 2020) (see Table 1). However, while it is possible to access the (=( of CNF, it is extremely computationally costly to approximate pa( with sufficient accuracy. Two primary reasons contribute to this cost: (1) Approximate estimators are not sufficiently accurate, making full Jacobian calculations necessary for each step along the flow; and (2) Many steps are necessary to control discretization error for sufficient performance (Fig. 2).
Recently, there have been significant advancements in few-step generation using flow models (Song et al., 2023; Boffi et al., 2025a; Frans et al., 2025; Guo et al., 2025; Sabour et al., 2025; Geng et al., 2025a). These models are extremely powerful few-
16
FALCON
SBG (Tan al.)
Torus Wasserstein-W2 -2, F-W;
2.0
DIT euler (exact)
DIT dopris (exact)
DIT dopris (Hutch)
step generators with flexible architectures and effi- cient simulation-free training; however, these few- step samplers do not natively admit efficient esti-
1.5
mators of the likelihood, making them unsuitable for the high-precision demands of importance sam- pling and scientific applications such Boltzmann Generation (Rehman et al., 2025).
1.0
In this work, we investigate how to design a generative model that combines the best of both worlds: the training efficiency and architectural freedom of simulation-free flow models with the fast sampling and likelihood evaluation of discrete-time invertible models. We propose
0.0
Lud
10d
10-3
10-2
10 1
10a
10a
Evaluation Time, GPU Hours
Figure Performance-ifference time comparison between NFs and CNFs for 10* dipeptide samples.
FEW-STEP ACCURATE LIKELIHOODS FOR CONTINUOUS FLOWS (FALCON). a flow-based model that enables few-step sampling while providing likelihood estimate that is both fast to compute and accurate enough for importance sampling applications. FALCON leverages hybrid training objective that combines a regression loss for stable and efficient few-step generation with a cycle-consistency term to encourage invertibility prior to convergence. Our main contributions are:
We introduce FEW-STEP ACCURATE LIKELIHOODS FOR CONTINUOUS FLOWS (FALCON): new continuous flow-based generative model for Boltzmann sampling that is invertible, trainable with a regression loss, and supports free-form architectures, while enabling both few-step generation and efficient likelihood evaluation.
Orthogonally, we introduce a simple and scalable, softly equivariant continuous flow architecture that significantly improves over the current state-of-the-an equivariant flow model architecture.
We show that FALCON two orders of magnitude faster than CNF-based Boltzmann Generators for equivalent performance (Fig. 2), drastically reducing the computational cost of CNFs, and taking significant strides towards real-world large-scale molecular sampling applications.
We show that FALCON outperforms the current state-of-the-a normalizing flow-based Boltzmann Generator across metrics, even when FALCON is given 250 x fewer samples (Tan et al., 2025a).
2

---

Published as a conference paper at ICLR 2026
2 BACKGROUND AND PRELIMINARIES
We are interested in drawing statistically independent samples from a target Boltzmann distribution Pharget with partition function Z, defined over IR:
Puarget (I) exp (-)(x), 2 2= exp (-E(x)) da
(I)
where : IRd + R. is the energy the system, which we can efficiently compute for any x. this work we do not require the energy to be differentiable. Unlike in the pure sampling setting (Akhound- Sadegh et al., 2024; Havens et al., 2025; Akhound-Sadegh et al., 2025; Midgley et al., 2023; Zhang & Chen, 2022; Vargas et al., 2023), we also assume access to a small biased dataset D = {x² ) of N samples (Noé et al., 2019). This makes it possible to perform an initial learning phase that fits a generative model with parameters 6, producing proposal distribution, p=)) (Noé et al., 2019).
Boltzmann Generators. (BGs) (Noé et al., 2019) combine generative models capable of exact likelihoods, with a target energy function and self-normalized importance sampling (SNIS) to re- weight generated samples to the target Boltzmann distribution. The model is first trained on possibly biased dataset ID as close as possible to Purget- BGs then draw K independent samples xf pf. [K] and compute the corresponding unnormalized importance weights such that w(xx)) -
Given these importance weights, we can then compute a consistent Monte-Carlo estimate of any observable o(x) under Ptarget using SNIS (Liu, 2001) as:
Σx ((<x))((>)) Σ w(x*)
E Puner [o(x:)] = [o(x)(()(
This allows for inference-time scaling as the Monte- Carlo estimate of any observable converges in probability to the correct value as the number of samples grows.
Flow Matching Models. Flow matching models (Lipman et al., 2023; Albergo & Vanden- 2023: Liu, 2022; Peluchetti, 2021) are probabilistic generative models that learn a continuous interpolation between an easy-to-sample distribution po = Pnoise and the data distribution P1 = Pdata in R. Let I, sæ1 + (1 s) )ro be point at time s E [0, 1] between two points zg ~ Po and T1 P1. The flow matching objective is then Bxgo-po.. ~ps -pp.-s--tt w() (E) s) (xz - xn) 1|
for some parameterized vector field vo, and weighting function w [0, 1] + Rt. We can then sample using an ordinary differential equation (ODE) 2 Jo (0 r)dr with the initial condition To Po and are guaranteed (under some mild assumptions) that (f x pi (.1).
Furthermore, the density p can be computed using the instantaneous change of variables for- by solving a single d d+1 dimensional ODE:
mula (Chen al., 2018)
8 &lggpla.) Gs
-tr
Oua
Up (x T) Oue Dir.
Tg log
-
dr, with initial condition
To log po (2o)
(3)
where the integral is discretized into T steps and the trace can either computed exactly in (dT)) function evaluations approximated using Hutchinson's 's trace estimator tr(J) [e Jc] for some noise vector c E IRd in O(T) function evaluations (Hutchinson, 1990). practice, this a major bottleneck because large number of steps needed control discretization error (see Fig. 2).
Few-step Flow Models. Flow matching models can require hundreds of steps to accurately approximate Pt. To speed up
Table 1: Related method overview
Method
Invontble
Regression- loss
Few Step FreeForm Arch.
BioEmu
Flow/Maps
TBG
sampling, few-step flow models such as consistency models (CMs) (Song al., 2023), optimal transport- based methods (Pooladian et al., 2023; Tong et al., 2024a;b; Shi et al., 2023), and flow map models (Boffi
RegHlow
Prose
Fiflloss
X
FALCON (Ours)
al., 2025a; Sabour et al., 2025; Geng et al., 2025a; Frans et al., 2025; Guo al., 2025) attempt to
3

---

Published as a conference paper at ICLR 2026
train model that generates high quality samples in many fewer steps. Recently, efficient models that take not only the current sample time, but also the target sample time have shown flexibility and effectiveness in the one- to few-step regimes. In these models, uo is augmented with an additional input I, which denotes the target time to capture the average velocity (Geng et al., 2025a) t as:
ulIs,, 8, ( ((xxr;))e
to minimize the average velocity objective:
[ - I o(rr, T r)d |1]
E n, w(s,t) uo(
(5)
where the average velocity ue is parameterized by a neural network as depicted in Fig. and is the vector field of the probability flow ODE that transports samples from the noise distribution Po to the data distribution pt. Once the average velocity is learned, then samples can be drawn using any discretization of the time interval [0, 1]. such that to 0, 14.... Ir as Lt, -t,,1 + (4: t4) for i E T.. However, thus far, few-step flow models have only been applied for fast generation, and, as we show, do not natively guarantee efficient access to likelihoods in realistic settings as the learned average velocity map ug is not guaranteed to be invertible before the training objective is perfectly minimized, making the standard change-off of-variabless formula inapplicable. These models and their relationship to FALCON are summarized in Table 1.
( t4-1,t4)
3 FEW-STEP ACCURATE LIKELIHOODS FOR CONTINUOUS FLOWS
We now introduce FEW-STEE ACCURATE LIKELIHOODS FOR CONTINUOUS FLOWS (FALCON), a novel flow-based generative model designed to address the inherent efficiency limitations of using continuous flow models Boltzmann Generators. Our method departs from traditional continuous normalizing flow (CNFs) by training a flow map that operates in a few discrete steps, while simultaneously achieving invertibility ensure fast and accurate likelihood computation for Boltzmann Generation. This is achieved through a hybrid training objective, which, by enabling stable few-step generation, dramatically reduces the inference cost. This efficiency allows us train more expressive architectures (Vaswani et al., 2017; Peebles & Xie, 2023; Ma et al., 2024) that were previously computationally infeasible to scale the Boltzmann Generator setting.
Flow Maps are Flawed Boltzmann Generators. The core FALCON is a generative process that learns an invertible map from a simple base distribution Po to the target molecular distribution Pr in a small number of steps. We first examine the suitability of the existing few-step flows for importance sampling applications, concluding that, their own, they are not sufficient. We first define the continuous map with respect to a vector field v as:
f(o(trr (o(x):T)de Ioe
Xu t)
(6)
and note that under mild regularity conditions on this map is always invertible up to discretization error. For any invertible map, we can compute the change in density with respect to the input using the change of variables formula (which requires the Jacobian), the cost of d function evaluations, and the determinant, which, while an ((33) operation, is in practice negligible compared to the NFE cost (See $D.2). This property also holds for flow maps at the optima, which address in Proposition 1:
Proposition Let uy be a minimizer of Eq. 5 with respect to some Also, define the Jacobian of X as Jx ax, and the discrete flow map: Dz.'
XulIs, 8, = J. + (1 s)ug(r,, Then, for sufficiently smooth uy and v and for any (s, E [0, 1|2 1. Xu(,,,,) is an invertible map everywhere,
(7)
2. log pt (xx) log p3 (xx)) log Idet J Jx_ (xx)) almost everywhere.
4

---

Published as a conference paper at ICLR 2026
We provide precise statement and proofs for all propositions in $A. This means that optimal flow maps are, in some ways, ideal Boltzmann Generators in that they have relatively efficient access to both samples and likelihood; however, this property only holds the optima, and in the case that (,,,,,1) ( 8, 1) for all I, which in practice is extremely challenging to satisfy.
In practice, for standard flow map models, (=(-, 8, t) + Xo(", and we have guarantee that X, will be invertible, making efficient likelihood calculation all but impossible. However, we note that this condition is actually much stronger than we need for FALCON, For our uses, while would like (,(,, 8, 1) to be close to (- 8, (), for accurate and efficient likelihood computation, we only require that X, invertible, not that matches the particular invertible map defined by X This leads us to define an additional invertibility loss:
Lim Es,t,a, I|= Xu (Xu s,t), s)l
(8)
be used in conjunction with the average velocity objective and flow matching objectives, Letm for a final loss comprised of three components:
((0) (0) >wvgCarg(0) + Aun (0),
with variants Lavg proposed below. Minimizing this loss has a less strict requirement for the correct- ness of the Boltzmann Generator specifically:
Proposition 2. Let ug be a minimizer of Lim (Eq. (8)) with respect to some u. Then, for sufficiently smooth ui and v and for any (s.t) t) E [0, 1]2 Xu(", s,t) is an invertible map everywhere, and logp? (z,) p" (Ix) log det J x., (xr,) almost everywhere.
Thus, minimizing the invertibility loss is sufficient for valid Boltzmann Generation, even without exactly reproducing the continuous-time flow. Note that the proposition provides a constructive guarantee of invertibility, and, in practice, only require the existence of inverse, not its explicit form. This condition ensures that FALCON acts as a consistent generator of the target energy distribution, E(x), while benefiting from fast inference-time scalability.
FALCON Enables Scalable Architectures. Previous Boltzmann Generators based on continuous normalizing flows for molecular applications utilize small equivariant architectures (Klein al., 2023; Klein & Noe, 2024; Tan et al., 2025a; Aggarwal et al., 2025). These models are limited in their scale due to the high cost of inference with multi-step adaptive step size samplers, which are needed to control the error in the likelihood calculation. FALCON, by enabling relatively cheap few-step sampling, can greatly improve performance. Specifically, we use a standard diffusion transformer (DiT) (Peebles & Xie, 2023) with an additional time embedding head. We also use combination of data augmentation to enforce soft SO(3) (rotation) equivariance and subtraction of the mean to enforce translation invariance following Tan et al. (2025a;b).
Formulations for Lave in the context of Boltzmann Generators. Many forms of Lave have been explored in the context of fast generation (Geng et al., 2025a; Guo et al., 2025; Boffi et al., 2025a; Sabour et al., 2025); We discuss these losses in $B. In this work, we focus on the following loss, which is equivalent the MeanFlow loss of Geng al. (2025a), but also explore additive losses (Boffi et al., 2025a) similar to shortcut (Frans et al., 2025), split-MeanFlow (Guo et al., 2025). This choice based on the superior performance of this loss in image experiments (Sabour al., 2025), as well as its potential for efficient implementation, which we discuss below.
Lavg
E Est. uo (,,,,, 8, sg e(x,, s) - (t - (u((,, s)O2, + 0,u0)
(10)
Note that since x. = s + (1 s) we can directly use v(2,, 8) =1 Fo-
Efficient Implementation. As noted in multiple previous works, Lavg can be efficiently imple- mented using single Jacobian vector product (JVP) call using forward automatic differentiation. Specifically, we have:
dug jvp(ue, (I., s,t,, (v,, 1,)))
(z s, s,t), ds
where the jvp function takes a callable function, inputs, and a vector (the vector part of the JVP).
5

---

Published as a conference paper at ICLR 2026
Additionally, for this loss specifically, can combine Leim and Lavg if we implement 5( ug(Zz, 8, s), i.c. passing the same time to u, representing the instantaneous veloc- ity. Specifically,
Algorithm 1: Training FALCON
Input: Sampleable Po and pl, regularization weight A,. network ug Output: The trained network ug
can implement the of the two losses by changing the distribution of s,t in Eq. 10 to include some percentage of the time when s = t to correspond with some fraction of Lem loss.
while training do (Io, .1) ~ po x PPi 2s f- s21 (1 s)zo:
Us x- σ1 -u;
Oug - jvp uo: (Iz, t), (V,, 1,0));
Our method is the first to require flow maps both in the forward and backward directions. We therefore need to consider the parameterization ug(x,, L, 8). i.e. the backwards direction, specif- ically at the discontinuity when
us. ds
Uug Us [L 8) Bue.
0s
ft Is (t s)ui::
Gg t &t + (s t)uol t, s);
t. When t - st, then (, t) = v(x,, s). but when t - sT. then uo (Zzg. s, t) = -v( s). To ad- dress this, we parameterize our flow map ug such that up (1r, 8, sign(l ha (3ts; 8, l).
( + ||u sg(ung) 1|2 + X || 1::1:
0 update 0, VoL(0))
return uo:
EXPERIMENTS
In this section, we first demonstrate that FALCON achieves more scalable performance over state-of-the-art continuous flows across both global and local metrics on tri-alanine, alanine tetrapeptide, and hexa-alanine (Table 3). Next, we empirically demonstrate that FALCON flows exceed the performance of state-of-the-art discrete NFs, even when they are given vastly larger sampling budgets (Fig. 4). Then, we elucidate the importance invertibility and aiding generative performance (Fig. 6). Finally, we ablate inference schedules and show their impact on performance across metrics as a function of sampler choice (Fig. 7).
regularization in achieving
4.1 EXPERIMENTAL SETUP
FALCON setup. We evaluate FALCON with two losses, one based on MeanFlow (Geng et al., 2025a) (FALCON) as described in $3, and the other an additivity loss (FALCON-A) (SB). We default to the first loss as it performed best or second best on all metrics on larger systems (Table 3).
Datasets. We evaluate the performance of FALCON on equilibrium conformation sampling tasks, focusing on alanine dipeptide (ALDP), tri-alanine (AL3), alanine tetrapeptide (AL4), and hexa-alanine (AL6). Datasets are obtained from implicit solvent molecular dynamics (MD) simulations with the amber-14 force field, as detailed in $D.3. We train biased data and test on a held-out unbiased dataset, using nerf-normalized importance sampling (SNIS) and force-field energies to compute log- ligelhoosss and re-sample molecules from the target Boltzmann density.
Baselines. We benchmark FALCON against both discrete and continuous normalizing flows We include four discrete normalizing flow baselines: (1) SB(3)-EACF (Midgley et al., 2023); (2) RegFlow (Rehman et al., 2025); (3) SBG (Tan et al., 2025a) with standard SNIS (SBG IS); and (4) SBG with SMC sampling (SBG SMC), as well as three continuous flows: (1) ECNF (Klein et al., 2023); (2) ECNF++ (Tan et al., 2025a); and (3) BoltzNCE (Aggarwal et al., 2025), recent -uuuivaqui architecture leveraging geometric vector perceptrons (GVPs) (Jing et al., 2021) on alanine dipeptide. For all continuous flows, samples and likelihoods generated by integrating over the vector field using the Dornand-Prince 4(5) integrator with atol = rtol = 10-3 (Dormand & Prince, 1986) to ensure a fair comparison between methods. More details architectures and parameters are covered in I.
We report Effective Sample Size (ESS), and the 2- Wasserstein distance on both the energy distribution (E-W2), and dihedral angles (T-W2). The full definitions of the metrics included in $E. The energy captures local details, as minor atomic displacements yield large variations in the energy distribution, while T-W2 captures global structure via mode coverage across metastable states. We
Metrics.

---

Published as a conference paper at ICLR 2026
include energy histograms in the main text, with Ramachandran plots relegated to $F.4. For robustness, all quantitative experiments are performed on three seeds the model and reported as mean + standard deviation in the tables and figures. For all benchmarks, in cases where dashes are present, data was unavailable, except for SBG SMC (Tan et al., 2025a), where ESS is not a valid metric.
4.2 FALCON OUTPERFORMS STATE-OF-THE-ARE METHODS
Superior Scalability Over Continuous Flows. Computing likelihoods in CNFs is
Table 2: Results on alanine dipeptide. Best results are bolded, with second-best underlined.
computationally prohibitive, limiting their scalability in the Boltzmann Generator set- ting. Although the current state-of-the-art, ECNF++, performs exceptionally well on ESS and T-W for alanine dipeptide (see Ta- ble 2) (Tan et al., 2025a), it fails to scale to larger molecules, as seen in Table 3. In con- trast, for larger cystems-tri-alante tetrapeptide, and hexa- maxx--a alanine-FALCOM substantially outperforms ECNF++ across all metrics, demonstrating superior scalability to larger molecular systems. The true MD energy distributions, learned proposals, and re-sampled energies for alanine dipeptide, tri-alanine, alanine tetrapeptide, and hexa-alanine are all shown Fig. 3.
Alanine dipeptide (ALDP)
Algorithm 4
ESS
E-W/ 4
T-W;
BolzNCE
0.27 0.02
0.57 0.00
SR(3)-FACF
< 10--
108.202
2,867
RegFlow
0.036
0,519
0.958
ECNT+
0.275 4. 0,010
0,914 : 0.122
0.189 =: 0.019
SBG
0.030 0.012
0.873
1:733=00 0.338
0,439
0.129
SBG SMC
0.741
0.189
0.431
0.141
alanine
FALCON-A (Ours)
0.097 1: 0.007
0.512 0.038
0.180 0.005
FALCON (Ours)
0.013
0.402
Table Quantitative results on tri-alanine (AL3), alanine tetrapeptide (AL4), and hexa-alanine (AL6). Baseline methods presented with SNIS, unless stated otherwise. Evaluations are conducted over x 10 samples.
Trimbbiine (AL3)
Terapopeale (AL44
Hesn-slaiin: (AL66)
Algenihi
ES5
8-Wo
T-315
ESS
E-W9
T-31%
ESS
8-Woy
T-W
4CNF+
0003 0,02
2,506 a 0.513
0.962 = 0.253
0.015 è
56688 s 0.485
10002
0.061
0,006
0.001
10,568 # 0.285
1.002 #: 0.055
ReaElow
0.029
1.051
1,612
0.010
6,277
3.476
SBK
0,53 0013
0.788
0.506
0.52
0016
0065 0014
1,08
0.495
0.999
006)
0.034 0.015
1,11 0.29
1431 0.05
SRG
0.508
0:044
050
1.077
0.382
10%
0006
(Dans)
IL104
1.38
IL343
0.994
0.007
1,163
FALCON (Ours)
0,013
0,047
0.011
1.256
Alanine dipeplide (4 steps]
Tri-elanine (8I steps)
Alanine Alnnine tetrapeptide (8 steps)
Hexe-wlinine (3 steps)
0.14
0.10
0.00
Tnue data
0.08
0.10
Propooal pansigtted)
Prypocal
Nortallbod Dindday
0.06
0.00
Mormaalbea 62
0.07
0.05
0.06
0.04
0.05
Mornaltod
0.64
0.03
0.04
0.00
0.04
0.03
0.02
0.02
0.02
0.61
0.01
0.00
-10
-20
20
0.00
-00
-130
-100
0.00
100
0.00
-100
Figure 3: True MD energy distribution with best FALCON unweighted and re-sampled proposals alanine dipeptide (left), tri-alanine (center left), and alanine tetrapeptide (center right), and hexa-alanine (right).
Improved Sample Quality Over Discrete Flows. Discrete NF' have recently been shown to be highly performant Boltzmann Generators (Rehman et al., 2025; Tan et al., 2025a;b). SBG (Tan et al., 2025a), based on the TARFlow architecture (Zhai et al., 2024), outperforms all previously reported methods across both global and local metrics (see Table 3). Here, we demonstrate that FALCON, even in a few steps, can outperform SBG across all metrics. We make two main claims to assert FALCON's competitive advantage in comparison with discrete NFs: (1) Discrete NFs, despite being fast one-step generators, consistently underperform compared to FALCON across all global and local metrics (Table 3); and (2) Increasing the number of samples can partially close this gap; however, even with 5 x 10 samples-250. x more than those used to evaluate FALCON-SBOE performance on E-W, remains significantly worse than that of 4-step FALCON Flow (as demonstrated in Fig. 4).
7

---

Published as a conference paper at ICLR 2026
4.3 ANALYSIS OF COMPUTATIONAL EFFICIENCY
1.2
5BG (Tan FALCON (Ours)
Training Efficiency Compared to Discrete NFs. Discrete NFs benefit from fast inference, but slow and unstable to train due to the maximum likelihood
Energy-W W2. E WW
objective (Xu & Campbell, 2023; Andrade, 2024). By
0.6
contrast, CNFs trained with flow matching objective trade more stable and faster training for slower infer- ence. When considering both training and evaluation time together (Table 4), we see that FALCON-despp being marginally slower at inference than the discrete NFs for the same number samples samples-achiee achieves faster cumulative training inference times for superior per- formance due to the expedited training objective.
0.2
9
10a
10a
Number of Samples Generated, M
Figure 4: Performance with additional samples.
Inference Efficiency Compared to CNFs. A primary contribution of FALCON is the dramatic reduction in the computational cost required to achieve high-quality samples with accurate likelihoods. Fig. 2 directly illustrates this advantage. To reach comparable level of performance metric for alanine dipeptide, a traditional CNF requires inference times that are two orders of magni- tude longer than FALCON. This efficiency gain what enables the use of more expressive alternative architectures and makes large-scale molecular sampling practical. For additional discussion, see $F.3.
the T-W2
Table 4: Cumulative training inference time across flows. 2 10 samples evaluated with (atol = rtol = 10 for our CNF and 4-step FALCON). All experiments were conducted on NVIDIA L40S with batch size 1024.
ECNF++
SBG
DiT CNF (Ours)
FALCON (Ours)
Alanine dipeptide
12.52
16.83
9.56
7.65
Tri-alanine
19.59
24.67
17.54
11.45
Alanine tetrapeptide
32.17
41.67
24.10
18.84
Hexa-alanine
137.4
57.50
82.10
25.76
Table Quantitative results alanine dipeptide, tri-alanine, alanine tetrapeptide, and hexa-alanine compared our DopriS-integrater CNFs. Evaluations were conducted over x 10 points across methods.
System
Algorithm
ESS f
E-W2 4
T-W2 4
NFB
FALCON-Dope -opris
0,264 0.058
0.442 1 0.048
0.218 1 0.023
257
Alanine dipeptide
FALCON
0.067 0.013
0.225 1 0.104
0.402 1 0.021
FALCON-Dophl
0.125 0.034
0.382 + 0.053
0.370 = 0.093
265
Tri-alanine
FALCON
0.077 0.004
0.544 + 0.013
0.452 + 0.011
8
FAL.CO--ooprS
0.129 0.015
0.665 +: 0.047
0.640 + 0.093
200
Alanine tetrapeptide
FALCON
0.055 + 0.003
0.686 + 0.047
0.858 = 0.077
8
FAI.CON-Dopris
0.128 + 0.031
1.013 + 0.115
1.320 + 0.201
207
Hexa-alaning
FAI.CON 0.060 + 0.017
0.892 + 0.311
1.256 + 0.132
16
Inference vs. Accuracy Trade-off. In FALCON, by using a flow map formulation, we can trade off performance for faster evaluation by adjusting the number of inference steps post-hoc. In Table 5, we show that hig,-NFE, adaptive step solver achieves superior performance the state-of-the -ttth-rrt continuous time ECNF++, as well as a few-step FALCON Flow; however, we also demonstrate that in this few-step regime, FALCON still outperforms every method considered (see Table 3) using two orders of magnitude fewer function evaluations (ranging between 4- 16 steps, depending on the dataset). Depending on the compute budget and goal, we demonstrate Fig. how FALCON is able to interpolate between slow and accurate sampling with fast but less accurate sampling.
4.4 ABLATION STUDIES AND DESIGN CHOICES
Verifying FALCON's Invertibility. As CNFs are only invertible at convergence, we introduce a regularization term in the loss to promote numerical invertibility in the few-step regime. In Fig.
8

---

Published as a conference paper at ICLR 2026
1:step
2-step
4-step
8-step
0.14
Thu Sata
0.12
0.12
0.12
0.12
Normaizes Denetty
trow grtee)
Normalteed Density
Normaited Dendty
0.10
0.10
0.00
3.08
0,08
0.08
Normalieed
0.06
0.04
0.04
0.04
0.04
3.02
0.02
0,22
40
-20
0.00
20
-20
20
8(x)
((r)
8(n)
81-)
Figure 5: Improved proposal and re-weighted sample energies with increased steps for alanine dipeptide.
we demonstrate the trade-off from this term on both the ESS and Wasserstein distance on dihedral angles: weak regularization leads to poor invertibility and degraded performance, whereas strong regularization enforces flow invertibility, albeit at the cost of reduced sample quality. We fix the regularization constant to 10.0 for all experiments performed, unless stated otherwise to balance performance and numerical invertibility.
We also directly prove that in- verse exists for our trained flow, by training auxiliary network to in- vert a frozen FALCON Flow in the forward direction. We find that FAL- CON achieves invertibility errors on the order of 10 which is the same order of magnitude as the in- vertibility of discrete and continuous NFs. Additional details can be found $F2 and Fig.
2etep
2.5
64tep
Estop
Torus W2..
2.0
16step
Enteerive Samppe
10-2
1.5
1.0
2stes
10-4
30a
10'
0.0
10F
10'
10
10'
Requbrizztton..
Figure 6: Performance trade-off with increasing regularization.
Impact of Inference Schedules. In the few step regime, performance can be significantly impacted by the choice of inference schedule.
Torus Wasserstein-W. T-WW(() W3(4)
0.8
We run ablations on various schedules for ala-
Energy-W,,F E ME(
nine dipeptide with 8-steps, summarizing the re-
0.6
sults in Fig. 7. We note that EDM substantially outperforms all other schedulers, in agreement with observations from the diffusion literature (Karras et al., 2022); sampling more points near the data distribution proves beneficial in aiding generative performance as the variance of the flow field higher closer the target distribu- tion. For all reported results, we use the EDM scheduler. In D.2.2, we provide additional de- tails regarding scheduler definitions, inference setups, and parameter selection for EDM.
0.4
10a
0.2
0.0
in
geam. cas.
choby
edm.
fn
geam. cos
cheby
edm.
Figure 7: Performance vs. choice of inference schedule.
RELATED WORK
Boltzmann Generators. Boltzmann Generators (BGs) (Noé et al., 2019) are used sample molecular conformations (Klein et al., 2023) and enable consistent estimates of thermodynamic observables (Wirnsberger et al., 2020; Rizzi al., 2023; Schebek et al., 2024). While traditionally BGs are based on discrete normalizing flows, more recent work in machine learning makes use of more powerful continuous normalizing flow architectures for invariance (Köhler et al., 2020; Köhler et al., 2023) and expressive power (Klein al., 2023; Klein & Noe, 2024). A few other works have explored the usage of approximate likelihoods (Draxler et al., 2024; Sorrenson et al., 2024; Aggarwal al., 2025), but have until been unable scale. Rehman et al. (2025) also proposes more

---

Published as a conference paper at ICLR 2026
efficient BG using new-regression- baseed objective to train discrete normalizing flow architectures, but the approach necessitates an invertible architecture limiting scalability and performance.
Few-step Flows. Diffusion and flow matching methods are now central to domains from vision (Song et al., 2021; Lipman et al., 2023) to scientific applications in material and drug discovery (Abramson al., 2024; Noé al., 2019). Scalable regression- based losses make these models fast to train, yet inference remains costly due to the need for numerous vector field integrations, motivating efforts to reduce computational expense. One- and few-step methods, like consistency models (Song et al., 2023; Song & Dhariwal, 2024; Geng et al., 2025b), shortcut models (Frans et al., 2025), and flow maps (Boffi et al., 2025a;b; Sabour et al., 2025; Guo et al., 2025) have gained recent attention. Concurrent work demonstrates that approximate likelihoods can be obtained through direct regression and are useful in guiding few step flows (Ai et al., 2025). Although there have been numerous efforts in improving generative performance in image settings, developing invertible few-step flows for scientific applications has seen substantially less interest. Our work, to the best of our knowledge, is the first demonstration of an invertible few-step flow with tractable likelihoods.
CONCLUSION
In this work, introduced FEW-STEP ACCURATE LIKELIHOODS FOR CONTINUOUS FLOWS, a novel few-step flow-based generative model designed to address the long-standing challenge of scalable and efficient Boltzmann distribution sampling. Our approach successfully combines the expressiveness and training efficiency of modern flow-based models with a few-step sampling capability for fast yet accurate likelihood estimation. By leveraging hybrid training objective, FALCON provides a practical solution for the computationally expensive likelihood evaluations that have historically limited the widespread adoption of Boltzmann Generators.
Our empirical results demonstrate that FALCON not only outperforms the existing state-of-ther discrete normalizing flow models, but also provides a significant leap in computational efficiency over previous continuous flow models. We showed that our model is two orders of magnitude faster than an equivalently performing Cas- Boltzmann Generator, making real-world, molccular sampling applications significantly more feasible. This represents critical step toward unlocking the potential of Boltzmann Generators in fields ranging from drug discovery to materials science.
Limitations. Despite its advancements, FALCON has several key limitations that crucial to acknowledge. First, while our results demonstrate that the computed likelihoods are empirically good enough for practical applications, we cannot efficiently guarantee their theoretical correctness. This represents a trade-off between computational efficiency and absolute theoretical certainty. Additionally, while theoretically possible, achieving true one-step generation remains out of reach for our current models, and we believe further architectural improvements and training methodologies are necessary to fully realize this potential.
Finally, our current research has primarily focused on the application of FALCON to Boltzmann Generation in molecular conformation sampling. Future work will explore applying our approach to Bayesian inference, robotics, and other domains where rapid and accurate likelihood estimation critical. We also see potential in models with structured Jacobians (Rezende & Mohamed, 2015; Dinh et al., 2017; Zhai al., 2024; Kolesnikov et al., 2025) to facilitate even faster sampling.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENTS
DR received financial support from the Natural Sciences and Engineering Research Council's (NSERC) Banting Postdoctoral Fellowship under Funding Reference No. 198506. The authors acknowledge funding from UNIQUE, CIFAR, NSERC, Intel, and Samsung. The research was enabled in part by computational resources provided by the Digital Research Alliance of Canada (https: pettp::/alllaaneeaas ca), Mila (https: / /mila. quebec), AITHYRA (https: //www. www..//wwwwwwww..oooooooooo.oo.. oea.. ac. at/aithyra), and NVIDIA.
ETHICS STATEMENT
Our work is primarily focused on theoretical algorithmic development for faster and more accurate generative models for sampling from Boltzmann densities, with reduced focus experimental implementation. However, we recommend that future users of our work exereise appropriate caution when applying it to domains that may involve sensitive considerations.
REPRODUCIBILITY STATEMENT
We undertake multiple measures to ensure the reproducibility of our work. A dedicated section in $F.1 outlines the setup required to generate each of our reported figures. Further, we provide comprehensive information the MD datasets used to train the training, validation, and test splits used. We also include a separate section detailing model configu- rations, learning rate schedules, optimizer settings, hyperparameter choices, and other relevant aspects to facilitate reproduction in $D.1. We will also release all developed code publicly upon acceptance.
models, including simulation parameters well as

---

Published as a conference paper at ICLR 2026
REFERENCES
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew J. Ballard, Joshua Bambrick, Sebastian W. Bodenstein, David A. Evans, Chia-Chun Hung, Michael O'Neill, David Reiman, Kathryn Tunyasuvunakool, Zachary Wu, Akvilé Zemgulyte, Eirini Arvaniti, Charles Beattie, Ottavia Bertolli, Alex Bridgland, Alexey Cherepanov, Miles Congreve, Alexander I. Cowen-Rivers, Andrew Cowie, Michael Fig- urov,, Fabian B. Fuchs, Hannah Gladman, Rishub Jain, Yousuf A. Khan, Caroline M. R. Low, Kuba Perlin, Anna Potapenko, Pascal Savy, Sukhdeep Singh, Adrian Stecula, Ashok Thillaisun- daram, Catherine Tong. Sergei Yakneen, Ellen D. Zhong, Michal Zielinski, Augustin Zidek, Victor Bapst, Pushmeet Kohli, Max Jaderberg, Demis Hassabis, and John M. Jumper. Accurate structure prediction of biomolecular interactions with AlphaFold 3. Nature, 2024, ISSN 1476-4687. doi: 10.1038/s41586-00 024-00487-0-
Rishal Aggarwal, Jacky Chen, Nicholas M Boffi, and David Ryan Koes. BoltzNCE: Learning likelihoods for boltzmann generation with stochastic interpolants and noise contrastive estimation. In Neural Information Processing Systems (NeurIPS), 2025.
Xinyue Ai, Yutong He, Albert Gu, Ruslan Salakhutdinov, J Zico Kolter, Nicholas Matthew Boffi, and Max Simchowitz. Joint distillation for fast likelihood evaluation and sampling in flow-based models. arXiv, 2025.
Tara Akhound-Sadegh, Jarrid Rector-Brooks, Avishek Joey Bose, Sarthak Mittal, Pablo Lemos, Cheng-Hao Liu, Marcin Sendera, Siamak Ravanbakhsh, Gauthier Gidel. Yoshua Bengio, Nikolay Malkin, and Alexander Tong. Iterated Denoising Energy Matching for Sampling from Boltzmann Densities. In International Conference on Machine Learning (ICMI.), 2024.
Tara Akhound Sadeeh, Jungyoon Lee, Avishek Joey Bose, Valentin De Bortoli, Arnaud Doucet, Michael M. Bronstein, Dominique Beaini, Siamak Ravanbakhsh, Kirill Neklyudov, and Alexander Tong. Progressive inference-time annealing of diffusion models for sampling from boltzmann densities. In Neural Information Processing Systems (NeurIPS), 2025.
Michael S. Albergo and Eric Vanden-Eijnden, Building normalizing flows with stochastic interpolants. International Conference on Learning Representations (ICLR), 2023.
Daniel Andrade. Stable training of normalizing flows for nigh-dimensional variational inference. arXiv, 2024.
Nicholas M. Boffi, Michael S. Albergo, and Eric Vanden- How to build a consistency model: Learning flow maps via self-distillation. In Neural Information Processing Systems (NeurIPS). 2025a.
Nicholas Matthew Boffi, Michael Samuel Albergo, and Eric Vanden-Eijnden. Flow map matching with stochastic interpolants: mathematical framework for consistency models. Transactions on Machine Learning Research (TMLR), 2025b.
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Neural Information Processing Systems (NIPS), 2018.
Laurent Dinh, Jascha Sohl-Dickstein and Samy Bengio. Density estimation using Real NVP. International Conference on Learning Representations (ICLR), 2017.
John R Dormand and Peter J Prince. Runge-kutta triples. Computers & Mathematics with Applications, 1986.
Felix Draxler, Peter Sorrenson, Lea Zimmermann, Armand Rousselot, and Ullrich Köthe. Free- form flows: Make any architecture a normalizing flow, In International Conference on Artificial Intelligence and Statistics (AISTATS), 2024.
Kevin Frans, Danijar Hafner, Sergey Levine, and Pieter Abbeel. One step diffusion via shortcut models. In International Conference on Learning Representations (ICLR), 2025.
Daan Frenkel and Berend Smit. Understanding Molecular Simulation: from Algorithms to Applica- fions. Elsevier, 2023.
12

---

Published as a conference paper at ICLR 2026
Zhengyang Geng, Mingyang Deng, Xingjian Bai. J. Zico Kolter, and Kaiming He. Mean flows for one-step generative modeling. In Neural Information Processing Systems (NeurIPS), 2025a.
Zhengyang Geng, Ashwini Pokle, William Luo, Justin Lin, and J Zico Kolter. Consistency models made easy. In International Conference on Learning Representations (ICLR), 2025b.
Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, llya Sutskever, and David Duvenaud. FFJORD: free-form continuous dynamics for scalable reversible generative models. In International Confer- ence on Representation Learning (ICLR), 2019.
Yi Guo, Wei Wang, Zhihang Yuan, Rong Cao, Kuan Chen, Zhengyang Chen, Yuanyuan Huo, Yang Zhang, Yuping Wang, Shouda Liu, and Yuxuan Wang. Splitmeantlow: Interval splitting consistency in few-step generative modeling. arXiv, 2025.
Aaron Havens, Benjamin Kurt Miller, Bing Yan, Carles Domingo- Enrich,, Anuroop Sriram, Brandon Wood, Daniel Levine, Bin Hu, Brandon Amos, Brian Karrer, Xiang Fu, Guan- Horng Liu, and Ricky T. Q. Chen. Adjoint sampling: Highly scalable diffusion samplers via adjoint matching. In International Conference on Machine Learning (ICML), 2025.
M.F. Hutchinson. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in Statistics Simulation and Computation, 1990. doi: 10. 10.1080/ 0361091900881
Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael JL. Townshend, and Ron Dror. Learning from protein structure with geometric vector perceptrons. In International Conference on Learning Representations (ICLR), 2021.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. In Neural Information Processing Systems (NeurIPS), 2022.
Leslie Kish. Confidence intervals for clustered samples. American Sociological Review, 1957.
Leon Klein and Frank Noe. Transferable boltzmann generators. In Neural Information Processing Systems (NeurIPS), 2024.
Leon Klein, Andreas Krämer, and Frank Noé. Equivariant flow matching. Neural Information Processing Systems (NeurIPS), 2023.
Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: exact likelihood generative learning for symmetric densities. International Conference on Machine Learning (ICML), 2020.
Alexander Kolesnikov, André Susano Pinto, and Michael Tschannen. Jet: A modern transformer- based normalizing flow. Transactions Machine Learning Research (TMLR), 2025.
Jonas Köhler, Michele Invernizzi, Pim De Haan, and Frank Noe. Rigid Body Flows for Sampling Molecular Crystal Structures. In International Conference on Machine Learning (ICML), 2023.
Benedict Leimkuhler and Charles Matthews. Molecular Dynamics: With Deterministic and Stochastic
Numerical Methods. Springer, 2015.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. International Conference on Learning Representations (ICLR), 2023.
Jun Liu. Monte Carlo Strategies in Scientific Computing. Springer, 2001.
Qiang Liu. Rectified flow: A marginal preserving approach to optimal transport. arXiv, 2022.
Nanye Ma, Mark Goldstein, Michael S. Albergo, Nicholas M. Boffi, Eric Vanden- Eijnden, and Saining Xie. Sit: Exploring flow and diffusion-bar based generative models with scalable interpolant transformers. In Computer Vision ECCV 2024, 2024.
Laurence Illing Midgley, Vincent Stimper. Gregor NC Simm, Bernhard Schölkopf, and José Miguel Hernández- Lobato. Flow annealed importance sampling bootstrap. International Conference on Learning Representations (ICLR), 2023.
13

---

Published as a conference paper at ICLR 2026
Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 2019.
Kaoru Ohno, Keivan Esfarjani, and Yoshiyuki Kawazoe. Computational Materials Science: From Ab Initio Monte Carlo Methods. Springer, 2018.
William Peebles and Saining Xie. Scalable Diffusion Models with Transformers. In 2023 IEEE/CVF International Conference on Computer Vision (ICCV), Paris, France, 2023. IEEE. ISBN 9798350307184, doi: 10. 10.11090CCV511000.223..0000 00087..
Stefano Peluchetti. Non-denoising forward-time diffusions, 2021.
Aram-Alexandia Pooladian, Heli Ben- Hamu, Carles Domingo-Enrich Brandon Amos, Yaron Lipman, and Ricky Q. Chen. Multisample flow matching: Straightening flows with minibatch couplings. In International Conference on Machine Learning (ICML), 2023.
Danyal Rehman, Oscar Davis, Jiarui Lu, Jian Tang, Michael Bronstein, Yoshua Bengio, Alexander Tong. and Avishek Joey Bose. Efficient regression- based training normalizing flows for boltzmann generators. arXiv, 2025.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. International Conference on Machine Learning (ICML), 2015.
Andrea Rizzi, Paolo Carloni, and Michele Parrinello. Free energies at qm accuracy from force fields via multimap targeted estimation. Proceedings of the National Academy of Sciences, 2023.
Amirmojtaba Sabour, Sanja Fidler, and Karsten Kreis. Align your flow: flow map distillation. In Neural Information Processing Systems (NeurIPS), 2025.
Scaling continuous-timg time
Maximilian Schebek, Michele Invernizzi, Frank Noé, and Jutta Rogal. Efficient mapping of phase diagrams with conditional boltzmann generators. Machine Learning: Science and Technology.
2024.
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion schrödinger bridge matching. In Neural Information Processing Systems (NeurIPS), 2023.
Yang Song and Prafulla Dhariwal. Improved techniques for training consistency models. In Interna- lional Conference on Learning Representations (ICLR), 2024.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference Leurning Representations (ICLR), 2021.
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. In International Conference on Machine Learning (ICML), 2023.
Peter Sorrenson, Felix Draxler, Armand Rousselot, Sander Hummerich, Lea Zimmermann, and Ullrich Köthe. Lifting architectural constraints of injective flows. In International Conference on Learning Representations (ICLR), 2024.
Gabriel Stoltz, Mathias Rousset, and Tony Lelièvre. Free Energy Computations: A Mathematical Perspective. World Scientific, 2010.
Charlie B. Tan, Avishek Joey Bose, Chen Lin, Leon Klein, Michael M. Bronstein, and Alexander Tong. Scalable Equilibrium Sampling with Sequential Boltzmann Generators. In International Conference on Learning Representations (ICLR), 2025a.
Charlie B. Tan, Majdi Hassan, Leon Klein, Saifuddin Syed, Dominique Beaini, Michael M. Bronstein, Alexander Tong, and Kirill Neklyudov. Amortized sampling with transferable normalizing flows. In Neural Information Processing Systems (NeurIPS), 2025b.
Alexander Tong, Kilian Fatras, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research (TMLR), 2024a.
14

---

Published as a conference paper at ICLR 2026
Alexander Tong, Nikolay Malkin, Kilian Fatras, Lazar Atanackovic, Yanlei Zhang, Guillaume Huguet, Guy Wolf, and Yoshua Bengio. Simulation-fie frreeeon--eeeeee schrödinger bridges via score and flow matching. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2024b.
Francisco Vargas, Will Grathwohl, and Arnaud Doucet. Denoising diffusion samplers. International Conference on Learning Representations (ICLR), 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Neural Information Processing Systems (NIPS), 2017.
Peter Wirnsherger, Andrew J Ballard, George Papamakarios, Stuart Abererombie, Sébastien Racanière, Alexander Pritzel, Danilo Jimenez Rezende, and Charles Blundell. Targeted free energy estimation via learned mappings. J. Chem. Phys., 2020.
Zuheng Xu and Trevor Campbell. Embracing the chaos: analysis and diagnosis of numerical instability in variational flows. Neural Information Processing Systems (NeurIPS), 2023.
Shuangfei Zhai, Ruixiang Zhang, Preetum Nakkiran, David Berthelot, Jiatao Gu, Huangjie Zheng, Tianrong Chen, Miguel Angel Bautista, Navdeep Jaitly, and Josh Susskind. Normalizing Flows are Capable Generative Models, 2024.
Qinsheng Zhang and Yongxin Chen. Path integral sampler: stochastic control approach for sampling. International Conference on Learning Representations (ICLR), 2022.
15

---

Published as a conference paper at ICLR 2026
APPENDIX
A THBORY
Proposition 1. Let ug be minimizer Eq. 5 with respect to Also, define the Jacobian of X as J 0x,'
OX and the discrete flow map:
Xu( s, = xa + (t s)e s, Then, for sufficiently smooth vi and v and for any (s, E [0, 1]2 Xu(,,,, i) an invertible map everywhere,
(7)
logp" (xx) log p (ar)) log |dct x., (2x) almost everywhere.
We first define what we mean by sufficiently smooth in both propositions. Assumption 1. ug is C almost everywhere.
This assumption on ug allows the application of the change of variables formula almost everywhere, which necessary compute the log likelihood under u. We note that this is satisfied by most sufficiently expressive modern architectures including those using ReLU type activations, which are not O everywhere, but are almost everywhere.
Assumption 2. is uniformly Lipschitz continuous in x and continuous in
This assumption on necessary in order to satisfy the Freard-Lindeli Theorem. note that the Lipschitz condition does not hold then the Peano existence theorem implies that the initial value problem (IVP) of the ODE may not be invertible. If, however, we have both continuity x and as then by the Picard-Lindelöf Theorem, we have existence and uniqueness of the IVP. We first recall Picard-Lindelöss
well as Lipschitz condition for
Theorem 1 (Picard-Lindelêf), Let D TRd TR. be closed rectangle with (No, (o) int D, the interior of D. Let D + Rd be function that is continuous in t and Lipschitz continuous in (with Lipschitz constant independent from t). Then there exists some such that the initial value problem is:
dr di -( x(1o) 20
(11)
has unique solution æ(t) the interval [to e, to e).
Under these conditions on o we are now able prove the proposition.
Proof. Recall that if ug minimizes Eq.
1
L: u(x,, u(x,,,))de r)
Butra w(s ()llul
8,1) (1),
then we have that (2 s, - Se u(z,, T)dT for all s, 1].. Furthermore we have
v. + +t (t-s) (d)s)up(C,,
(12)
x. + (t ()-s) [ u((r,
(13)
x. + L: T)d
(14)
which by application of the Picard-Lindelöf theorem invertible for all 8, proving part 1 of the proposition.
To prove part 2, we that Eq. 14 differentiable with respect v, and X, is invertible, therefore by the change-of-varidee
formula and Assumption 1, we arrive at part 2 of the proposition. O
Proposition Let up be a minimizer of Lim (Eq. (8)) with respect to some e. Then, for sufficiently smooth un and and for any (s, E 1J3 Xa(,,,, s,t) t) is an invertible map everywhere, and
log r. (xx) log ph (xr.) log /det. J (x.)) almost everywhere.
16

---

Published as a conference paper at ICLR 2026
Proof. First recall Eq. 8:
Linv = Es,t., I|,,
(15)
In this case, the minimizer of Liwv that is also in C' is pointwise invertible by definition (u(,,, the inverse (',,, (', 8, t) as Liv - 0 everywhere.
For the change of variables to apply, we also need Assumption apply almost everywhere.
We note that this loss ensures pointwise invertibility. Additional restrictions are needed to ensure that the log likelihood can be calculated efficiently.
B OTHER FORMULATIONS FOR LAvG
There are various formulations of Lavg that have been explored in the recent literature. We note that FALCON directly benefits from any advancements in this rapidly-evolving research direction. Specifically, the following three different Lavg losses can be used for training flow maps.
1. L1 A Estt.r, (n((,,,2.)) sg v(x,,))) (1-s)))(s)s)))) s)O,, ua (u,0) which is equivalent
to the MeanFlow loss of Geng et al. (2025a), as well as the ESD objective in Boffi et al. (2025a). Note that since Vg = $z + -s) s) we can directly use v(z,, s) 20.
2. L2 A E.,.., uo(s v.) ) sg Auo 8, + (X ( r),
for As
(1- A)t, is equivalent to the SplitMeanFlow loss of Guo al. (2025), as well as the scaled PSD objective in Boffi et al. (2025a). In Boffi et al. (2025a), two ways of sampling the intermediate time r are explored. Deterministic mid-point sampling, which sets À = 0.5, as well as uniform sampling of X. Guo al. (2025) only explores uniform sampling. This is also generalization of Frans et al. (2025), which additionally samples and in a binary-tree fashion. We denote this loss when integrated in our setting as FALCON-A. We it was most stable.
0.5 for all experiments this
setting as we found
3. L3 E.,.a, (I,, sg up (,,,, s,t), (,)
is the Lagrangian objective pre- sented in Boffi et al. (2025a). We were not able to get this objective to train stably our setting.
C EXTENDED RELATED WORK
C.I RELATION FREE-FORM FLOWS
fore-form Flows (FFF) enable arbitrary architectures to function normalizing flows by jointly training an auxiliary network to approximate the inverse the forward model (Draxler et al., 2024). The forward model maps data to latents, while the auxiliary network regularizes this mapping by learning a tractable estimator for the Jacobian term the change of variables. As correctly highlighted by Draxler et al. (2024), although the loss encourages invertibility, this guarantee only holds when the reconstruction error sufficiently small-a condition that can be difficult to meet in practice. In our setting, FALCON is trained with entirely different loss function- albeit with cycle-consistency term to promote invertiblity-simile to that observed by Draxler et al. (2024). In addition, we explicitly validate invertibility for FALCON, in line with their observations through a dedicated experiment.
C.2 RELATION TO REGFLOWS
RegFlows (Rehman et al., 2025) enable a more efficient training process for standard discrete normal- izing flows by avoiding the hiitmmm--ikelhhoe (MLE) objective, whose unstable training dynamics often make optimization a challenge; instead, RegFlows distill the knowledge of a predefined invert- ible coupling-ein either using a pre-trained continuous normalizing flow pre-computed optimal
17

---

Published as a conference paper at ICLR 2026
transport map- soup-nnooo discrete normalizing flow via regression objective. While this sidesteps issues associated with MLE and yields an efficient training pipeline, it also imposes a strict requirement on the existence of an invertible reference map, which must be provided in advance.
Further, RegFlows are constrained to inherently invertible architectures, which limits their expressiv- ity and ultimately their scalability. As show in Table 2 and Table 3, even when comparing against the strongest RegFlow variant (NSF), FALCON consistently outperforms it across all metrics, with pronounced gains on larger peptide systems. In contrast, FALCON is the first few-step Boltzmann Generator to employ a fully free-form architecture, surpassing state-of-the-art inherently invert- ible discrete flows (and even continuous normalizing flows) while avoiding the strict architectural constraints imposed by invertibility. Note, although FALCON is only guaranteed to be invertible either at convergence or via additional regression objective, this relaxation enables substantially higher expressivity and delivers significant empirical performance gains. Lastly, we demonstrate the practical invertibility of FALCON through experiments performed in $F.2.
D EXPERIMENTAL DETAILS
All training experiments are run on a heterogeneous cluster of NVIDIA H100 and L40S GPUs using distributed data parallelism (DDP). All models were trained with three random seeds, and reported values are averages across runs. Additional training and inference details are included in the following. For benchmarks with existing methods, in cases where dashes are present, data was unavailable, except for the case with SBG SMC (Tan et al., 2025a), where ESS is not a valid metric.
D.1 TRAINING DETAILS
Architecture. We adopt Diffusion Transformer (DiT) backbone for FALCON, with the same model size used for all peptides. The details of the backbone configuration are in Table 6 below.
Table 6: Overview of FALCON configurations across datasets.
Dataset
Hidden Size
Blocks
Heads
Cond. Dim.
Parameters (M)
Alanine dipeptide
192
64
3.2
Tri-alanine
192
3.2
Alanine tetrapeptide
192
3,2
Hexa-alaming
3.2
Training Configuration All models were trained with an exponential moving average (EMA) on the weights using decay rate of 0.999. Logit values energy regularization disabled. For evaluation, we generated 10 proposal samples, and used the same number for re-sampling and computing all metries. Lastly, center-of-mass augmentation was applied with a standard deviation of 1/ /T, where is the number of particles.
clipped 0.002, with compositional
Hyperparameters Optimization was performed using AdamW with learning rate 10-1, = (0.9, 0.999), = 10-8 and weight decay 10-1, A cosine annealing learning rate schedule with a warm-up phase covering 5% of the training iterations was also used.
D.2 INFERENCE DETAILS
In this section clarify the inference strategies used in this work. Both in terms of strategies used for likelihood computation, as well as how to schedule evaluations of FALCON.
D.2.1 LIKELIHOOD COMPUTATION
Likelihood computation of CNFs generally uses the continuous-timg time formulation in Eq. (3), repeated below for clarity:
Vo(
Ze log (x.)) =
-tr
Bigg ikr.
dr, with initial condition
log po (xo)
18

---

Published as a conference paper at ICLR 2026
There are two main degrees of freedom in approximating log for some generated sample, the integrator, and the trace computation or approximation.
Following previous work use the Dormand-Prince-44 (Dopri5) adaptive step size ODE integrator for all continuous-time models. To our knowledge, all Boltzmann Generators that use base continuous normalizing flow generator use the Dopris solver as the importance sampling is otherwise too biased. Fixed-step size integrators, while effective for generation, are known systematically overestimate likelihood for flow matching models.
Integrator:
Trace Computation: There two main strategies used in practice to estimate the trace
(uu 80=-
Either using a direct exact computation, or using the Hutchinson trace approximation with Monte-Carlo estimate of size 1. We find the exact computation necessary for Boltzmann Generator applications to maintain stability of the importance sampling.
In practice, this implemented using d + dimensional ODE where the trace computation is computed in a vectorized manner using reverse mode automatic differentiation.
For FALCON, the continuous-time formulation no longer applies, and we need to use the exact change of variables formula at each discrete time-step as detailed in Algorithm 2.
Algorithm 2: Inference for FALCON
Input: Sampleable po. trained network up, time schedule {4
Output: Samples IN ~ pe with associated log likelihoods p.v ar Po;
Po < loggoo((o)) (xo);
for NJ do
(s, t) +- (1i-1, f;);
/* where (x s. xg + (t -ss-)))e (Eu0
2; + (x t):
+ Pi-1 - log |det. J x. (xi-1));:
return PNN
Computational Cost Cost-wise, this computation turns out to be bottle-necked by function evalua- tions, and requires approximately Ti- function evaluations where Ti is the number of integration- steps chosen by the adaptive step-size solver for the likelihood computation, and d is the dimension of the system. This the number of integration steps chosen in this generation-only case. In practice, we find Ti Tg. and often up to 2-3x larger.
contrast to the T, function evaluations required to generate samples where T
FALCON the other hand has a computational cost of N function evaluations where N the chosen number of steps to evaluate FALCON. In practice we set N << Tí, e.g. Table where Ti 200 and N 8. We note that technically there is an additional O(d3) term calculate the determinant. However, in practice this is not bottleneck compared to the cost of function evaluations.
D.2.2 INFERENCE SCHEDULERS
We evaluated five inference schedules to assess their impact on generative performance. The linear baseline distributes steps uniformly across the trajectory, while the geometric, cosine, Chebyshev, and EDM schedules bias step allocation to emphasize regions where the data distribution more sensitive. The mathematical definitions and parameter settings for each schedule are provided below.
Linear. The linear schedule spaces is distributed uniformly between and 0, with N E N denoting the total number of inference steps such that:
t; =
N i (O....,N) N}.
(16)
19

---

Published as a conference paper at ICLR 2026
Geometric. The geometric schedule exponentially allocates more resolution to early steps. We let a E IR and a > 1, denote the geometrie base. For all experiments conducted, we set a 2:
N- E (O...,,,N
(17)
a
Cosine. The cosine schedule follows a sauared-cosin law, concentrating more steps near This has been used in diffusion models for improved stability. With N E N:
i iE (0, ................................ NJ, iN O.
cos
(18)
2 N
Chebyshev. The Chebyshev schedule, derived from the nodes of Chebyshev polynomials of the first kind, minimizes polynomial interpolation error, and distributes steps more densely near boundaries:
4
cos
0(5+-0.5) N
+ +1), E (O,..,,,,, IN = O.
(19)
The EDM schedule (Karras et al., 2022) parameterizes the noise level a using power-law with exponent p RR The interpolation is performed p-space allocate more steps where the generative process is most sensitive. For all experiments performed, we use the same parameters proposed by Karras al. (2022), with p = 7, Fmin 10-3 and F 1, and min, max E R
EDM.
max + N 0i = (o
2/p min
o
1/p maxx
)), E (O...,,,NN N),
(20)
t
0.
(21)
max
min
This formulation ensures a smooth transition between max (high noise) and m (low noise, near- deterministic refinement).
D.3 DATASET DETAILS
For the datasets used, we follow the same training, validation, and test split used by (Tan al., 2025a), where single MCMC chain is decomposed into 105, 2 10 and 10 samples for training, validation, and test. The training and validation data are each taken as contiguous regions of the chain simulate the realistic scenario where you have generated an MCMC trajectory and would like to use Boltz- mann Generator to continue generating samples. The data is split so that (after warmup) the first 103 samples are for train, the next 2 x 1 are for validation, and the test samples uniformly strided sub- samples from the remaining MCMC chain. Earlier parts of the trajectory undersample certain modes enabling biased training set that attempt debias through access the energy function and SNIS. MD simulations were all run for 1 μs, with timestep fs, at temperatures 300K, 3 310K, and 300K for alanine dipeptide, tri-alanine, and alanine tetrapeptide, respectively, in line with those conducted by Klein & Noe (2024). The force fields used for all three molecules, in the same order, were the Amber ff99SBildn, Amber 14, and Amber ff99sBi 1dn, in line with prior work by Tan et al. (2025a).
Alanine Dipeptide. The Ramachandran plots for the training and test split are provided in Fig. 8. Following Klein et al. (2023), we up sample the lowest sampled mode (middle right) to make the problem easier for BGs. This also has the added benefit of testing what happens to FALCON when the training dataset is biased relative to the true distribution.
Tri-alanine. Similarly to alanine dipeptide, the Ramachandran plots for the training and test split in Fig. 9. As there are two pairs of torsion angles that parameterize the system, there two sets of Ramachandran plots included for each training and test. In tri-alanine, we can see that the training set actually misses mode entirely (1/ π r3),, and undersamples this mode for sha relative the test set. This is great test of finding new mode (in ei)) and correctly weighting a mode (in 4/2).
Alanine Tetrapeptide. For the tetrapeptide, there are three sets of Ramachandran plots each for training and test given the three pairs of torsions angles that parameterize the molecule, in Fig. 10.
20

---

Published as a conference paper at ICLR 2026
MD Data (Training Set)
MD Data (Held-out Test Set
Figure 8: Left: Training data for alanine dipeptide; Right: Test data for alanine dipeptide.
NO (training Set)
(training Set)
MD (Helld-ot Test Seti
(Held-out Test Set
Figure 9: Left and left center: Training data for tri-alanine; Right center and right: Test data for tri-alanine.
Dasa Thairing
Data (hainin Selt
(Paining Set)
Data (hairing
Data
Figure 10: First three: Training data alanine tetrapeptide; Last three: Test for alanine tetrapeptide.
MD Data Training Set
MD Data Training Set
MD Data Traininn Set
MD Dota Training Set
MD Data Training Set
MD Data (Hald-out Tast Sat
Data (Held-out Test Sat)
MD Data (Hold-out Tast Sen)
MD Data (Held-aut Tast Set
MD Data (Held-out Test Set
Figure 11: First five: Training data for hexa-alanine; Last three: Test data for hexa-alanine.
21

---

Published as a conference paper at ICLR 2026
Hexa-alanine. For hexa-alanine, also include Ramachandran plots in Fig. 11. The first row shows training data, while the second row shows a held-out test set used to evaluate performance.
E
METRICS
Effectivv Sample Size (ESS), To quantify sampling efficiency, we compute the effective sample size (ESS) following Kish's definition (Kish, 1957). Given N N generated partieles with unnormalized importance weights {wr)i= R,, the ESS is normalized by as:
FSS((w)))
Σ m
(22)
Zi=1 W
The ESS reflects how many independent, gqually-weighter samples would provide equivalent statisti- cal power to the weighted sample. Higher ESS is desirable for more performant models.
2- Wasserstein Energy Distance (E-WV)) To compare energy distributions, we compute the Wasserstein distance between generated and reference samples from our MD dataset. For distributions p.9 E (TR) over energy values, and [T q) the set of admissible couplings, we define:
/
E-W2(p.9)
min *ell(p.4)
x v1 dr ().
(23)
RXB
The E-W2 measures how closely the generated energy histogram aligns with that of the reference, with high sensitivity to structural accuracy due bond-length -eeegddee energies. Since small perturbations in local structure induce large fluctuations in the energy distribution, this metric captures this variance, with lower values of E-W2 being more favourable generative models.
Torus 2- Wassersstii Distance (T-W2). To assess structural similarity in torsional space, compute a 2- Wasserstein distance over dihedral angles. For molecule with L E residues, define the dihedral vector as:
Dihedrals(z) = (ø1, Vi...., ΦL-1, v2-1) E [0, 2x 2(2-1)
(24)
The cost on the torus accounts for periodicity of angles:
2(L
Tt (x.
[DDihedrals(a): Dihedrals(y): + x) mod 2nr -r]² m]²
(25)
The torus Wasserstein distance between two distributions p.q E P([O, ((+))( is then:
f v) dx (z.
T-W2(p,0)
min
(26)
=EIl(p-a)
This captures macroscopic conformational differences while respecting angular periodicity. The T-W2 provides more global assessment performance, for instance revealing missed conforma- tional modes that would not be captured by the &-W,.
F ADDITIONAL EXPERIMENTS
F.1 DETAILS FOR FIGURE GENERATION
In Fig. 2, we compare performance on the T-W2 metric for alanine dipeptide between FALCON and our continuous-time DIT CNF For FALCON, accuracy is controlled by the number of inference steps (1-8). For the DiT CNF, we consider three settings: (1) adaptive step Dormand-Prince 4(5) with exact Jacobian trace evaluation, varying atol/rtol from 10 to 10 10-5; (2) the same tolerances but instead using the Hutchinson trace estimator, trading performance for faster likelihoods with higher variance; and (3) fixed-step Euler integration with step sizes from 4 to 256 with exact Jacobian traces, where the upper bound chosen to roughly match the number of function evaluations needed for Dopri5. The models and training configurations used are presented in $D.1.
22

---

Published as a conference paper at ICLR 2026
In Fig. 3, we demonstrate the energy distributions for unweighted and re-weighted samples for our most performant FALCON Flows. For alanine dipeptide, tri-alanine, alanine tetrapeptide, and hexa-alanine, 4, 8, 8, and 16 steps were used for figure generation. Energy distributions reveal microscopic detail, as marginal changes in local atomic position can have significant impacts on total energy. The best models for each molecular system were used, with pertinent details model size, training configurations, and hyperparameters detailed in S. Similarly, in Fig. 5, we show the proposal and re-weighted sample estimates for alanine dipeptide, and demonstrate how an increasing number of steps, improves the energy distribution.
In Fig. we compare the performance between SBG's discrete NFs and FALCON, illustrating that despite more samples increasing performance across metrics, the approach is still unable to reach FALCON's performance. For SBG, we specifically take their best model weights for the TarFlow architecture from: https: / thub. com/trans ansstrrrnrrraaa ferable- -arrraaeamlleeerreeeeeeeeae httlllsssstttnneee transf frrasnarrabaa wereerralleeeeeeee samplers and draw N₂ E {10 22 x 10a, 5 10 10 2 10 5 x 105, 10 2 10 x 10%) samples three different times (the error bars are representative of the three draws) and evaluate & E-W2 for each set.
In Fig. 6, we investigate how the strength of regularization impacts performance on ESS and E-W2. Specifically, we demonstrate that small amounts of regularization enable generative performance but impede invertibility, while too much regularization detrimentally impacts sample quality. This trade-off leads to an optimum on both metries. We investigate Xp {10", 10 10 10 10 10 }. and conclude upon A = 10 To ascertain the optimal X,, we ran these experiments on alanine dipeptide, using the same model details and configurations highlighted in $1D.1.
Fig. 7 demonstrates the sensitivity the inference schedule plays on generative performance. For this figure, we took best models-trained with the details presented in Snd ran inference using all three trained seeds to generate uncertainties for each inference schedule. We used our 8-step system, as inference scheduler choice less important the fewer steps are used.
F.2 PROOF OF INVERTIBILITY
In our loss, we use cycle-consistency term that regularizes training to promote numerical invert- ibility. We see in Fig. 6, the introduction of this modified loss aids generative performance; however, the forward-backward reconstruction yields errors on the order of 10-2, indicating an approximately, but not entirely invertible model. To prove that the flow is invertible (but the in- verse challenging to discover during train- ing), we train an auxiliary FALCON Flow in the reverse direction. We use the same network to parameterize the reverse flow as described in $D.1, with the same training configuration and hyperparameter set. Next, we freeze the for- ward model (which goes from latents - data), generate synthetic data (2 x 10 prior-target ple pairs) and train the auxiliary model on these reflow targets to learn the mapping from data back to latents. In Fig. 12, we illustrate the loss curve of the trained auxiliary FALCON. To evaluate invertibility, we draw i..d. samples from our prior distribution, pass them through the frozen forward model, and test the auxiliary reverse model on these unseen latents. We evaluate an Ez error on the recovered latents to find reconstruction within 10-4 after 1000 epochs matching the reconstruction accuracy of discrete NFs that are invertible by design. These results support the claim that the learned flow is indeed invertible.
10
Reversible (Raw)
Reversible Loss Curve (0.99 EMA)
Training Loss
10-1
10-2
1
200
400
600
800
1000
Epoch
Auxiliary model loss from reflow target train- ing forward flow.
Figure 12:
23

---

Published as a conference paper at ICLR 2026
F.3 PERFORMANCE AGAINST NUMBER OF FUNCTION EVALUATIONS
Fig. 13 highlights the efficiency of FALCON in
2.5
LI 11
ELECH FALCON
terms of number of function evaluations. We
DIT euler (exact)
see that FALCON achieves the same torus 2- Wasserstein performance as our DiT CNF with
DIT depriS (exact)
2.0
Torus Torus Wasserstein-W2. x
dopriS (Hutch)
Dopri5, while requiring over two orders of magni- tude fewer function evaluations. Whereas Fig. quantified efficiency in terms of inference time, here we measure the number of function evalua- tions against both fixed-step solvers (Euler with 4-266 steps) and adaptive solvers (Dopri5 run until reaching target tolerances atol rtol
1.5
1.0
{10-2, 2,022.10 1003 10 4.10 1100g $} for Hutchinson and
0.5
(10-2, 10-3 10-1 10-5, 10-6) for exact).
0.0
Although Hutchinson's trace estimator yields a speedup relative exact Jacobian computation for CNFs, as seen in Fig. 2, the number of function evaluations needed is higher than the exact Jacobian computation. In either setting, however, FALCON still remains substantially faster. Even with 4-step solver, FALCON matches the accuracy of DiT CNFss using Hutchinson at atol = rtol = 10-5, with an 8-step solver nearly matching the performance of DiT CNF with Dopris set to an atol rtol 10-6.
10a
10a
Number of Function Evaluations,
Performance of FALCON our DiT CNFss as a function of NFEs.
Figure 13:
F.4 RAMACHANDRAN PLOTS
Alanine Dipeptide. We demonstrate FALCON's capacity to learning global features through the Ramachandran plots in Fig. 14 for alanine dipeptide. We include both the held-out test set and the learned model's map, showcasing its ability to debias the data and capture the true MD distribution. Specifically, the undersampled ø mode in the training data is correctly upweighted in the learned model's predictions, indicating accurate likelihood estimates from the learned flow.
MD Data (Held out Test Set)
Model Predictions
4,0
Free onergy Hes
2.0
0.0
Figure 14: Left: Test data for alanine dipeptide; Right: FALCON's angular predictions for alanine dipeptide.
Tri-alanine. Similarly, show the Ramachandran plots for tri-alanine in Fig. 15 exhibiting similar behaviour. Most conformations are correctly captured, with some modes being underweighted.
Alanine Tetrapeptide.
Next, we show the Ramachandran plots for alanine tetrapeptide Fig. 16.
Hexa-alanine Finally, we show the Ramachandran plots for hexa-alanine in Fig. 17.
24

---

Published as a conference paper at ICLR 2026
MD Data (Teat 5et)
MD Data (Test Set)
Model Predictions
Model Predictions
free energy May
2.3
(P2
PPa
Figure 15: Left and left center: Test data for tri-alanine; Right and right center: FALCON's angular predictions for tri-alanine.
Figure 16: First three: Test data for alanine tetrapeptide; Last three: FALCON's angular predictions for alanine tetrapeptide.
Figure 17: First five: Test data hexa-alanine: Last five: FALCON's angular predictions for hexa hexa-alanine. alanine.
25


Reviews:
[{"review_id": "07w6CClQvi", "reviewer": "Reviewer_ExRz", "summary": "In this work, the authors tackle the problem of accurate likelihood estimation for the family of Flow Map models, an extension of standard CNFs. While efficient in terms of the number of steps needed for accurate sampling, Flow Maps cannot model the required change of variables needed for accurate density estimation.\n\nThe authors tackle this problem by introducing an additional regularization term that enforces the flow map to be approximately invertible (Eq. 8). If strong enough, this regularization term makes it possible to calculate the change in density using the log-determinant of the flow map’s Jacobian, similar to what is used in time-discrete normalizing flows. The authors formally show that the exactly invertible architecture is the minimizer of the proposed loss.\n\nIn their experimental evaluation, the authors show how a Flow Map trained using their proposed regularization technique can be used within the context of equilibrium sampling of peptide systems using importance sampling, similar to how standard CNFs and discrete-time NFs are used in Boltzmann generators. The experimental section focuses on showing improved scalability over CNFs and improved sample quality over discrete-time NFs. Furthermore, the paper studies the efficiency of the method (comparing training against discrete-time NFs and inference against CNFs) and studies the inference–accuracy trade-off. Throughout all experiments, the approach performs better than or on par with current state-of-the-art methods.", "questions": "See weaknesses", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 1, "strengths": "Overall, the paper is well executed. The background and preliminaries sections are well written and contain sufficient depth; the problem setting is clear, and the proposed solution is well described. Furthermore, the proposed solution itself seems intuitive (although novelty might be limited, as noted below). Finally, the experimental section is extensive and provides clear evidence in favor of the proposed approach.\n\nAs such, I am generally in favor of accepting the paper for publication based on the content alone. However, as highlighted below, I have a few concerns regarding the novelty of the paper that need to be answered before I can vote for acceptance with absolute confidence.", "weaknesses": "My primary objection to voting for acceptance of the paper at this time is the similarity of the proposed method to the one presented in Rehmann et al., 2025. This paper has a very similar objective: combining the benefits of CNFs and discrete-time NFs by enforcing approximate invertibility. Additionally, both papers share a very similar set of experiments (Rehmann et al. is, however, slightly more extensive). While the method presented here has, in theory, sufficient novelty to warrant publication, further clarification about the similarities and a comparison across methods are needed.\n\nWithout this clarification included in the paper, I do not feel confident in accepting the paper for publication.\n\nIn addition, I have included a short list of comments/questions below that I would appreciate the authors’ response to or see addressed in the paper:\n- Fig. 1: This is very nit-picky, but the debiased target should completely overlap with the true density.\n- Fig. 2: This should also include a discrete-time NF.\n- Tab. 3: A key contribution of the paper is that, with the Flow Map–based approach instead of discrete-time NFs, FALCON should scale better. Table 3, however, shows that when considering large systems (e.g., Hexa-Alanine), the difference between the methods shrinks quite significantly, with SBG-IS and SBG-SMC even outperforming FALCON in terms of Wasserstein distance over energy distributions. This should be discussed in the paper.\n- Tab. 4: It’s difficult to establish a clear trend in these results. ECNF is worse for smaller systems but has better training + inference time for Alanine Tetrapeptide. Conversely, the largest system is again considerably slower. SBG, on the other hand, is slower than FALCON for all system sizes except the largest. This should be discussed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JmkJcUAXJA", "reviewer": "Reviewer_cmBL", "summary": "This paper describes FALCON (Few-step Accurate Likelihoods for Continuous Flows), which does what the title says. The addition of an invertability loss supports the training of a few-step flow map from which accurate likelihoods can be computed efficiently.  The results demonstrate this advances the state-of-the-art for Boltzmann Generators.", "questions": "What do you mean by combining L and L_{avg}? This is unclear.\n\nHow could your approach be adapted to a transferable model (not single molecule system)?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The results (in particular Figure 2) are impressive and advance the field.  Inference time is orders of magnitude faster than traditional CNFs.\n\nScalable Boltzmann generators can enable new direction in molecular modeling - this is an important advance.\n\nThe paper is well written and the approach is well motivated, both intuitively and mathematically.", "weaknesses": "Figure 4 is hard to interpret. What are the x's? Does FALCON really have a constant W2 regardless of the number of samples generated? I don't think this figure is clearly illustrating the data as intended.\n\nAlthough the ability to use a larger model is attributed as a strength of the approach, for comparison purposes it would be nice to train a model of similar size to comparative approaches to factor the effect of model size on performance.\n\nL_1 is used without introduction (defined in appendix, but not main paper).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NW2RpOU6gq", "reviewer": "Reviewer_vNfw", "summary": "This paper introduces FALCON, a framework designed to overcome a long-standing bottleneck of flow matching and related continuous-time generative models: although training is efficient, evaluating accurate likelihoods remains computationally expensive due to the need to integrate the divergence term along the ODE trajectory.\n\nFALCON proposes a theoretically grounded and computationally efficient approach that enables few-step yet accurate log-likelihood estimation by optimizing a surrogate network to approximate the evolution of the log-density under flow dynamics. The method provides an effective approximation that avoids explicit computation of the ODE while maintaining high accuracy in the estimated likelihoods. Theoretical analyses support its consistency and error bounds, and extensive experiments demonstrate that FALCON achieves comparable or better accuracy than full integration methods at a fraction of the cost.\n\nOverall, the paper offers a mathematically elegant and practically useful solution to a central problem in flow-based models.", "questions": "1. If the Jacobian determinant becomes close to zero in some regions (i.e., the mapping is nearly non-invertible), how does this affect the numerical stability and bias of the estimated likelihood? Can the authors provide practical guidance or regularization strategies to mitigate this issue?\n\n2. What specific smoothness or Lipschitz continuity assumptions are required on u_theta for the theoretical guarantees to hold? Would the results still apply to models with ReLU-based, piecewise-linear velocity fields that are not globally C^1?", "limitations": "", "rating": 10, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The paper addresses a core unresolved issue in the flow-matching and continuous normalizing flow (CNF) literature — accurate and efficient likelihood computation — using a very novel and conceptually clean idea. The formulation of few-step likelihood estimation through continuous optimization is both original and impactful.\n\n2. The theoretical development is rigorous, with solid mathematical grounding that explains why the surrogate optimization scheme yields unbiased or low-bias likelihood estimates.\n\n3. Empirically, the experiments are thorough and convincing, covering a range of flow architectures and datasets. The results consistently demonstrate that FALCON significantly reduces computational overhead without sacrificing estimation accuracy.\n\nIn summary, this contribution is potentially highly influential for both the theoretical and applied machine-learning communities. It removes a practical barrier that has limited the deployment of flow-based models in real-world density estimation. FALCON may become a standard component for efficient probabilistic inference with continuous flows.", "weaknesses": "Overall, the work is strong and technically solid, but a few aspects could benefit from further clarification or discussion:\n\n1. Since the method still depends implicitly on the invertibility and numerical conditioning of the local Jacobian, it would be important to analyze the behavior when the Jacobian determinant approaches zero, i.e., when the approximate mean velocity field u_theta becomes locally near-singular. Could such regions lead to unstable or biased likelihood estimates? Most practical architectures cannot guarantee strict invertibility.\n\n2. Theoretical analysis likely assumes certain smoothness properties of the vector field to ensure the validity of the continuous approximation. These assumptions should be clearly stated and their necessity discussed, especially for large-scale neural networks that are not strictly smooth.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "g9EEzvKm2U", "reviewer": "Reviewer_tZEt", "summary": "This paper proposes to train a mean-flow model with an additional soft constraint on the flow invertibility by optimising the $\\mathcal{L}_{inv}$ (Equation 8). Once the optimum of this inv loss is achieved, the log-prob of a generated sample can be calculated by change-of-variable. The proposed model, FALCON, could generate samples with few network calls while having access to model density, which can be applied as a Boltzmann generator.\n\nThe method is clearly written, and the empirical results show a clear gap between other baselines on either ESS or sample quality related metrics when scaling to complex tasks. However, the main concern lies in the low ESS of FALCON on ALDP, which is significantly lower than one of the baseline ECNF++.\n\nOverall, the reviewer recommends a weak accept.", "questions": "Please see weakness.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The method is clearly written\n\n2. The proposed method is simple and effective\n\n3. The experiments show the scalability of the proposed method, which outperforms other baselines", "weaknesses": "The main weakness of the paper lies in the experimental result on ALDP, where the ESS of FALCON is much lower than that of ECNF++. It would be great to elaborate more on why this happens. On the other hand, to showcase the effectiveness of proposed method, one could use the same network architecture as ECNF++ and then train FALCON. Such setting would be fairer to illustrate the effectiveness of the invertibility loss.\n\nOn the other hand, it is not clear how invertible the trained network is. In particular, the author could elaborate more on the invertibility of $X(x, s, t)$ by using a 2-dimensinonal heatmap. Intuitively speaking, for farer s and t, the invertibility is more difficult to maintain, and therefore increasing the number of NFE could improve ESS.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)