Scalable Equilibrium Sampling with Sequential Boltzmann Generators
Authors: Charlie B. Tan
Joey Bose
Chen Lin
Leon Klein
Michael M. Bronstein
Alexander Tong
Venue: ICML 2025 poster (2025)
Keywords: Normalizing Flows
Boltzmann Generators
Annealed Importance Sampling
Primary area: applications->chemistry_physics_and_earth_sciences
TL;DR: We scale up Boltzmann Generators using large all-atom transformers and apply sequential Monte Carlo to improve sampling.

Abstract: Scalable sampling of molecular states in thermodynamic equilibrium is a long-standing challenge in statistical physics. Boltzmann generators tackle this problem by pairing normalizing flows with importance sampling to obtain uncorrelated samples under the target distribution. In this paper, we extend the Boltzmann generator framework with two key contributions, denoting our framework Sequential Boltzmann Generators (SBG). The first is a highly efficient Transformer-based normalizing flow operating directly on all-atom Cartesian coordinates. In contrast to the equivariant continuous flows of prior methods, we leverage exactly invertible non-equivariant architectures which are highly efficient during both sample generation and likelihood evaluation. This efficiency unlocks more sophisticated inference strategies beyond standard importance sampling. In particular, we perform inference-time scaling of flow samples using a continuous-time variant of sequential Monte Carlo, in which flow samples are transported towards the target distribution with annealed Langevin dynamics. SBG achieves state-of-the-art performance w.r.t. all metrics on peptide systems, demonstrating the first equilibrium sampling in Cartesian coordinates of tri-, tetra- and hexa-peptides that were thus far intractable for prior Boltzmann generators.

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
Charlie B. Tan
Avishek Joey Bose Chen Lin' Leon Klein
Michael M. Bronstein 14 Alexander Tong
Abstract
Legend
Scalable sampling of molecular states in thermo- dynamic equilibrium is a long-standing challenge in statistical physics. Boltzmann generators tackle this problem by pairing normalizing flows with importance sampling to obtain uncorrelated sam- ples under the target distribution. In this paper, we extend the Boltzmann generator framework with two key contributions, denoting our frame- work SEQUENTIAL BOLTZMANN GENERATORS (SBG). The first a highly efficient Transformer- based normalizing flow operating directly on all- atom Cartesian coordinates. In contrast to the equivariant continuous flows of prior methods, we leverage exactly invertible non-equivariant ar- chitectures which are highly efficient during both sample generation and likelihood evaluation. This efficiency unlocks more sophisticated inference strategies beyond standard importance sampling. In particular, perform inference-time scaling of flow samples using a continuous-time variant of sequential Monte Carlo, in which flow sam- ples are transported towards the target distribution with annealed Langevin dynamics. SBG achieves state-of-the of-thh-art performance w.r.t. all metrics on peptide systems, demonstrating the first equilib- rium sampling in Cartesian coordinates of tri-, tetra- and hexa-peptides that were thus far in- tractable for prior Boltzmann generators.
X
Start Sample
Energy Contour
Non-equ Inium Tansport
Figure I. SBG uses annealed Langevin dynamics transport pro- posal flow samples towards towards the target distribution.
et al., 2011), protein-ligand binding (Buch et al., 201 1), and formation of crystal structures (Parrinello & Rahman, 1980; Matsumoto et al., 2002), whose understanding can aid in problems that range from long-standing global health challenges, to efficient energy storage (Deringer, 2020).
The dominant paradigm for molecular sampling involves rumning Markov chain Monte Carlo (MCMC) molecular dynamics (MD), whereby the equations of motion are in- tegrated with finely discretized time steps. However, such molecular systems often exist in thermodynamic equilib- rium by remaining for extended periods in metastable states. Such metastable states are captured in the minima of a com- plex energy landscape, itself defining the molecular sys- tem's 's equilibrium (Boltzmann) distribution at given tem- perature. The high-energy barriers separating metastable states lead to infrequent state transitions (Wimsberger et al., 2020), presenting an obstacle for effective sampling with simulation-bart methods such as molecular dynamics or MCMC, requiring long simulation periods with small time steps on the order of femtoseconds fs 10-15g.
1. Introduction
The sampling of molecular systems at the all-atom resolu- tion is of central interest in understanding complex natural processes. These include important biophysical processes such as poottinnfoolding (Noé et al., 2009; Lindorff-Larsen
'Equal contribution Universsity of Oxford 'Mila Québec Al Institute 'Freie Universität Berlin "AITHYRA Universitté de Montréal. Correspondence to: CBT <char- lie.tan @wxeeexxeeeeooo uk>, AJB <jjccy moseo..lelmmmi.mmlllll magilmmeeil..c...... <alcxandertongdeve@ gmail.com>>.
Boltzmann generators (BG) (Noé et al., 2019) offer an al- ternative approach, in which powerful generative models, such as normalizing flows (Dinh et al., 2017; Rezende & Mohamed, 2015), are trained on existing (but assumed to biased) datasets, and leveraged as a proposal for self- normalized importance sampling (SNIS), targeting the de-
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
sired Boltzmann distribution. Boltzmann generators permit accelerated sampling through amortization the uncorre- lated proposal generation avoids the slow state transitions suffered by MD and MCMC. Despite their appeal, it re- mains challenging for existing BGs to model systems be- yond the smallest peptides (2 amino acids) in Cartesian coordinates (Klein et al., 2023b; Midgley et al., 2023a). The principal drawback inhibiting scalability stems from the lack of expressive equivariant architectures that are also exactly invertible (Bose et al., 2021; Midgley et al., 2023a), or the present over-reliance on simple E(n)- (Satorras et al., 2021) based equivariant vector fields in continuous-timg normalizing flows (Chen et al., 2018). As a result, even the most performant BGs suffer from poor target distribution overlap. leading to low sampling efficiency during SNIS.
We summarize the different aspects of our proposed SBG in comparison to other learned samplers in Table 1.
We instantiate SBG using a best-in-class, non noneequivr normalizing flow, named TarFlow (Zhai et al., 2024). TarFlow is a modernized normalizing flow architecture employing a scalable transformer backbone to parameterize an exactly invertible transformation. We demonstrate that such exactly invertible architectures, via fast and accurate log-likeldood evaluation, benefit from inference-scaling. We emphasize this is in stark contrast to continuous normalizing flows that underpin prior SOTA Boltzmann generators which require both the costly simula- tion of the 2nd order divergence operator as well as differ- entiation of an ODE solver. Furthermore, we demonstrate that enforcing equivariance softly along enables us to stably scale proposal flows in SBG, far beyond prior BGs. On theoretical front, we study novel inference-tine proposal energy adjustment to counteract the influence of training data centroid augmentation when resampling, as well as quantify the additional bias of common thresholding tricks employed to improve resampling numerical stability. Em- pirically, we observe SBG achieve state-of-the-ar results across metrics, far outperforming continuous BGs all datasets. In particular, SBG is the first uncorrelated learned sampler to scale successfully in Cartesian coordinates to tripeptides, tetrapeptides, hexapeptides, and makes signifi- cant progress towards equilibrium sampling of decapeptides.
pprpose,
Present work. In this paper, we introduce SEQUENTIAL BOLTZMANN GENERATORS (SBG) novel extension to the existing Boltzmann generator framework.' SBG makes progress on the scalability of Boltzmann generators in Cartesian coordinates along two complementary axes: (1) scalable pre-training of softly -qqiivrrqui proposal normalizing flows in BGs; and (2) inference time seal- ing via continuous-time variants of annealed importance sampling (AIS) (Neal, 2001) and sequential Monte Carlo (SMC) (Doucet et al., 2001). The use of AIS or SMC over SNIS enables more effective sampling given suboptimal proposal- overlap, enabling SBG to draw uncorrelated Murget (r) samples for peptide system up to 6 residues.
2. Background and Preliminaries
Table 1. Method overview for samplers, given biased data samples. Use Anncoling
Method
Use Bix)
Exact likeliboods
We are interested in drawing statistically independent sam- ples from the target Boltzmann distribution Maarget, with par- tition function Z, defined over Rax3,
DEM (Alaound Sudegh al., 2024)
NETS (Albergo Vandes Bijnden, 2025) (Nuc al., 2019)
SBG (Ows)
SBG scales up normalizing flows in BGs by following re- cent advances in atomistic generative modeling (Abramson et al., 2024). In particular, we remove the rigid SF((3- equivariance as an explicit architectural inductive bias in favor of softly enforcing it through simpler and more ef- ficient data augmentations. To further improve sampling we perform inference-time scaling by defining an interpo- lation between the proposal flow energy distribution (i.e., negative log density of samples) and the known target Boltz- mann energy. Crucially, simulating samples at inference via annealed Langevin dynamics may be coupled to corre- sponding time evolution of importance weights, converting naturally to continuous-tii time variants of the well-established annealed importance sampling (AIS) (Neal, 2001) and se- quential Monte Carlo (SMC) (Doucet et al., 2001). As a result, SBG can readily improve over the simple one-step importance sampling methodology used in existing BGs.
EE(x)) kig
EE(x) Kg
Mtarger (1) x exp
2 exp
dr.
J2d
The Boltzmann distribution is defined for given system and includes the Boltzmann constant kg, and a specified temperature T. Additionally, the potential energy of the system E R=*3 + R and its gradient V can be evaluated at any point x E TRax3, but the exact density (uunge (x) is not available as the partition function Z evaluation is intractable for all but the simplest systems.
In this paper, unlike pure sampling-bard based settings, we are afforded access to small biased dataset of N samples D {2 Jiz provided as an empirical distribution pD. Consequently, it is possible to perform an initial learning phase that fits a generative model pg. with parameters 6, to PD-c. by minimizing the forward KL D (pr (pellpe)-r act as a proposal distribution that can be corrected.
We open source our full codehase at https: //gp://bg/tt//////// com/ char/ 1 wewtteeetttannneee wttaaaanntttsaafeaabeaeaa saa-ssampleee samplers.
2

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
2.1. Normalizing Flows
flow-matching learning framework (Liu, 2022; Albergo & Vanden-Eijnden, 2023; Lipman et al., 2023; Tong al., 2024). Specifically, Hlw-matching regresses learnable vector field of a CNF ft, [0, 1] x ==* + the target vector field uy (x, associated to the flow 4/.. In practice, it is considerably easier to regress against a target conditional vector field uy (azel=)-whick generates the con- ditional probability path pr (x ( as we do not have closed form access to the (marginal) vector field un which gener- ates pt. The conditional flow-matching (CFM) objective can then be stated as a simple simulation-free regression,
A key desirable property needed for the correction of trained generative model Po biased dataset D is the ability to extract an exact likelihood (z)x). Normalizing flows (Dinh et al., 2017; Rezende & Mohamed, 2015) rep- resent exactly such a model class as they learn to transform an easy-to- -oa-ll---slle base density to a desired target density using a parametrized diffcomorphism. More formally, given a sample from a (prior) base density Fy ~ Po and a diffeo- morphism So: RRax3 - R=x3 that maps the initial sample to æ1 fo (zo). We can obtain expression for the log density of 21 via the classical change of variables,
I=*3 to
CCFM = E, (e (ot.elillll (l,,, x,) u (=(r1)=)1 (3)
f (2o)
log (1z) = log (x log det
The conditioning distribution q(2) can be chosen from any valid coupling, for instance, the independent coupling ((z) p(z))((())))))) ((11). We highlight that Eq. 3 allows for greater flexibility in fe,o as there is exact invertibility constraint. To generate samples and their corresponding log density according to the CNF we may solve the following flow ODE numerically with initial conditions zo (0 and c = log Po(o)), which is the log density under the prior:
Oro
In Eq. I above the log det term corresponds to the Jacobian determinant of fo evaluated at zo- Optimizing Eq. I is the maximum likelihood objective for training normalizing flows and results in fo learning P1 x Pusta There are multiple ways to construct the (flow) map fo. Perhaps the most popular approach is to consider the flow to be a composition of finite number of ele- mentary diffeomorphisms fo = far o fM--1'.. fi. resulting in the change in log density to be: (o = log Pooo) - Em= log (0fio(1;-1)/0rr We note that the construction of each fL E [M] is
d ()))
(t) -V V f(o(t,
dt [logp((a))
(4)
2.2. Boltzmann Generators
(z) and Jacobian 0f (z))(=)) computationally cheap compute.
motivated such that both the inverse f()((z)
A Boltzmann generator (Noé et al., 2019) pe pairs normal- izing flow as the proposal generative model pe. which is then corrected to obtain i.i.d. samples under Jurget using self- normalized importance sampling. More precisely, as nor- malizing flows are exact likelihood models, BG's first draw K independent samples x' ~ p)(x), e [K] and compute the corresponding (unnormalized) importance weights for each sample w(x')= = exp
Continuous normalizing flows In the limit of infinite elementary diffeomorphisms, a normalizing flow trans- forms into a continuous normalizing flow (CNF) (Chen et al., 2018). Formally, a flow is a one-parameter time- dependent diffeomorphism th+ [0, 1] x Rax3 - R that is the solution to the following ordinary differential equation (ODE): d
(-8(=)) /po(==). Leveraging the importance weights we can compute a Monte-Carlo ap- proximation to any observable o(x) of interest under Hunget using nelf-normalized importance sampling as follows:
di (z)) = u4 (z)), with initial con- ditions ())) = no. for a time-dependent vector field Me : 1] x x - Rax3 It is often desirable to con- struct the target flow by associating it to designated prob- ability path pt- [o. 1] x P3 + 33 which is a time-indexed interpolation in probability space between two distributions po: P1 E [P) In such cases, the flow 2h, is said to generate Pt if it pushes forward po pi by following Ut Pt = [ # (Po). As thy is a valid flow and satisfies an ODE the change in log density can be computed using the instantaneous change of variables:
ΣK w(x')o(x') Ex w(x+)
K
E E (n) [o(x)) = Ep. [o(z)
In addition, computing importance weights also enables resampling the pool of samples according to the collection of normalized importance weights W = {m<
3. SEQUENTIAL BOLTZMANN GENERATORS
log p(11) = logp(zo) - f(
V (x
We now present SBG, which extends and improves over classical Boltzmann generators by including an annealing process to transport proposal samples towards the target distribution. We begin by identifying the key limitation in current BGs as SNIS with a suboptimal proposal. Indeed, while the SNIS estimator is consistent, its efficacy highly
where = V/a (2o) and V. is the divergence operator.
A CNF can then be viewed as a neural flow that seeks to learn a designated target flow th for all time E [0, 1].
The most scalable way to train CNFs is to employ a
3

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
dependent on the overlap between proposal pa and target Marge: where the optimal proposal is proportional to the minimizer of the variance of 6(x²) (orrge( (Owen, 2013). Unfortunately, since po within a BG is trained on a biased dataset 'D the importance weights typically exhibit large variance, resulting in a small effective sample size (ESS).
R SO(3) and applying them to data samples R o #1 pi (a1). Secondly, the data is mean-free and has (n 1) x 3 degrees of freedom, we lift the data dimensionality back to n by adding noise to the center of mass. This allows us to easily train with non-translational equivariant prior distribution such as the standard normal Po = (0, 1). More precisely, a data sample is constructed x = Rí + c, where ( 1)x3 + R=x3 is the mean-free data point embedded in R=*3 R SO(3), and ~ |N(0, d²). At inference, the impact of this center of mass noise is that we must account for p(llel)). which follows a X3 distribution in three dimensions. Consequently, during reweighting we adjust the proposal energy to account for the impact of center of mass training augmentation as follows:
We address the need for more flexible proposals in $3.1 with modernized scalable training recipes for atomistic normaliz- ing flows. In $3.2 we outline our novel application of non- equilibrium sampling with sequential Monte Carlo (Doucet et al., 2001). We term the overall process of combining a pre- trained Boltzmann generator with inference scaling through annealing SEQUENTIAL BOLTZMANN GENERATORS.
Symmetries of molecular systems. The energy function ((x) in a molecular system using classical force fields is invariant under global rotations and translation, which cor- responds to the group SF(3) = SO(3) x (IR3+ Unfortu- nately, SE(3) is a non-compact group which does not allow for defining a prior density (0())) on R=*3 Equivariant generative models circumvent this issue by defining a mean- free prior which is projection of a Gaussian prior (0, T) onto the subspace (xx=1)x3 (Garcia Satorras al., 2021). Thus pushing forward a mean free prior with equivariant flow provably leads to an invariant proposal Pi (x1) (Köhler et al., 2020; Bose et al., 2021). We next build BGs de- parting from exactly equivariant maps by considering soft cquivariance, unlocking scalable and efficient architectures.
||c| 20
1c22
log (1x) = log (x:) +
log
3 , (5)
1/20
where r(+) is the gamma function. We empirically analyze the impact of this adjustment in $F.3.
We next outline a proposition, and prove in $B.1, that demonstrates that reweighting using the adjusted proposal provably leads to better SNIS effective sample size (ESS).
Proposition 1. Given SE(3)- (i(3)-invariant Purger (z). consider the decomposition of data point x E R into its constituent mean-free component, F E 1m=)333 - Rax3 and center of mass e E TR3, x = H + G, where c -~ -N(0, o²). Now, assume both the proposal po(x) and the adjusted proposal D6 factorize independently over the mean-free com- ponent and the center of mass. Then setting (x pa(x) 11/xxa (Ieel),, leads to the following inequality on the effective sample size in the limit of K - 0o:
3.1. Scaling Training of Boltzmann Generators
To improve proposal flows in SBG we favor scalable ar- chitectural choices that are more expressive than exactly equivariant ones. We motivate this choice by highlighting that many classes of normalizing flow models are known to be universal density approximators (Teshima et al., 2020; Lee et al., 2021). Thus, expressive enough non-equivariant flows can learn to approximate any equivariant map.
(iarger(az) po (z)
MMurger (1x) )
ESS
<
EES
(6)
Pî (z)
Soft equivariance. We instantiate SBG with a state-of- the-art TarFlow (Zhai et al., 2024) which is based on block- wise masked autoregressive flow (Papamakarios et al., 2017) based on causal Vision Transformer (ViT) (Alexey, 2021) modified for molecular systems where patches are over the particle dimension. Since the data comes mean-free we further normalize the data to unity standard deviation. Com- bined, this allows us to scale both the depth and width of the models stably as there is no tension between a hard equivariance constraint and the invertibility of the network.
3.2. Inference Time Scaling of Boltzmann Generators
Given trained BG with proposal flow pe, the self- normalized importance sampling estimator suffers from a large variance of importance weights as the dimensionality and complexity of (*)) grows in large molecular systems. We aim to address this bottleneck by proposing an inference time scaling algorithm that anneals samples of ~ ps(x) - and corresponding unnormalized importance weights w w(x') in continuous manner towards Manget
We include series of strategies to improve training of non- equivariant flows by softly enforcing First, we softly enforce equivariance to global rotations through data augmentation by sampling random rotations PESS is defined as: ESS- 1/ EKK (w²
Improved sampling via annealing. We leverage class of methods that fall under non-equilibrimm sampling to im- prove the base proposal flow samples. One of the simplest instantiations of this idea is to use annealed Langevin dy- namies with reweighting through a continuous-tine variant
4

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
of Annealed Importance Sampling (AIS) (Neal, 2001). Con- cretely, we consider the following SID that drives proposal samples towards the target Boltzmann density:
resampling during the annealing process SMC can avoid particle redundancy in which all but few particles have negligible weight. We state the full SBG sampling algo- rithm with adaptive resampling in Algorithm 1; to recover
dæ, = e, (x,)dT + 22,,WWr
(7)
the SBG AIS variant we simply set ESSphreshld = - 1.0.
where c 2 0 is a time-dependen diffusion coefficient and W, is the standard Wiener process. We distinguish т, from t used in the context of training pa. as the time variable that evolves initial proposal samples at 7 O towards the target at T = 1. The energy interpolation & is design choice, and we opt for simple linear interpolant & (1-7) and set {z) = - log.oop(pe (r))). We highlight that unlike past work in pure sampling (Mâté & Fleuret,
Algorithm 1 SBG Sampling
Require:
# particles K, # annealed distributions N, Energy annealing schedule & (x+) A - 1/N 2: for i = 1 to N do
1: Zo ~ Eg
Iy+4 f 24 - E VE, (2+ )dT + /2,,
4: log A < log u4 Ex )dr
2023; Albergo & Fijden,, 2025) which use the prior energy Ex(z) - log Po(z)) our design affords the significantly more informative proposal given by the pre- trained normalizing flow py. As such, there is no need for additional learning, with the annealing process extending the inference capabilities of the Boltzmann generator مo (z).
5:
+ + A
6:
if ESS < ESS then
7:
+ RESAMPLE(Z,, W,) W/- + 0
8:
9: end
10: end for
To resample or compute observables with the transported samples, we use the well-known and celebrated Jarzynski's equality, that enables the calculation of equilibrium statistics from non-equilibrium processes. We recall the result, origi- nally derived in Jarzynski (1997), and recently re-derived in continuous- time in the context of learned sampling algo- rithm by Vargas et al. (2024); Albergo & Handen-Fijnder (2025), that describes the importance weight time evolution.
To simulate the Langevin SDE in Equation (7), and the corresponding importance weight evolution, we require the gradient of the energy interpolant:
((xT) KBT
( (1-)) (ogg +
Proposition 2 (Albergo& & Vanden-Eijnder (2025)). Let (I,, W solve the coupled system of SDE / ODE
which requires efficient gradient computation through the log-likelihood estimation under the normalizing flow py as given by Eq. 1. This presents the first point of distinction between finite flows and CNFs. The former class of flows trained using Eq. 1 gives fast exact likelihoods especially for our scalable non-equivariant TarFlow model. In contrast, CNFs must simulate Eq. 4 and differentiate through an ODE solver to compute V log po (xx+) for each step of the Langevin SDE in Eq. 7. As a result, a TarFlow proposal is considerably cheaper to simulate and reweight with AIS than a CNF In $A we present an alternate interpolant that does not require the proposal distribution during sampling which is appealing when only samples are needed but at the cost of more expensive computation of log weights. These paths are of interest in the setting of Boltzmann emulators and other generative models, and are of independent interest, but are not considered further in the context of SBG.
du, -e, (x7) )dr + dW,
dlog w,+ = - (x,) with pa. Wy = 0
then for any test function Ø Id R we have
| o(z)), (x)dxx
E. [[w+
(8)
E E[w+]
and
2 = E[e"*] (Jarzynski's equality)
(9)
The final samples Ir-1 are then reweighted with the im- portance weights Wy=1, themselves lower variance than SNIS in conventional BGs. It is crucial to highlight that the prior is not directly constituent of this annealing process, but instead the learned proposal pa (zo) acts as the initial distribution. It is precisely this learned proposal density that d log w, evolves during the annealing process. Alnealed importance sampling can be considered a special case of se- quential Monte Carlo (SMC) (Doucet et al., 2001). In SMC, resampling can occur at arbitrary times т, typically using an ESS threshold as adaptive resampling. Intuitively by
For improved numerical stability during annealing, and to further reduce computational footprint, we propose strat- egy that eliminates the forward evolution of the initial pro- posal that already obtain high energy. Specifically, we can simulate large number of samples via Eq. 12 and thresh- old using an energy threshold 7 > 0, and evaluate the log weights of promising samples. We justify our strategy by
5

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
first remarking lower bound to the log partition function of Murget using Monte Carlo estimate,
truncation (x) == P()g(1) 2 8), for small 6 >> O. Essen- tially, this thresholding filters low probability samples under the model prior to any importance sampling. The additional bias incurred by performing such thresholding is theoreti- cally analyzed in Proposition 4 and presented in $B.4.
- E(x) exp
NOT
log 2 = log E (x))
(x)
2 E (x)(()
Ex(x)) KTT
log (og)
= log 2. (10)
4. Experiments
We evaluate SBG on small peptides using classical force- field energy functions, further experimental details are de- scribed in $E. SBG samples are generated by Algorithm 1, both with adaptive resampling (SMC) and without (AIS).
Plugging this estimate in the definition of the target Boltz- mann distribution we get an upper bound,
=(xx) KBT
log Murget (x) S log
- log E.
We consider small peptides composed of up to 6 alanine residues, with some systems additionally incorporat- ing an acetyl group and an N-methyl group. All datasets are generated from a single MD simulation in implicit solvent using classical force field. For each system, the first 1 us is used for training, the next 0.2 ps for validation, and the remainder serves as the test set. Therefore, some metastable states may not be represented in the training set. An ex- ception is alanine dipeptide, for which we use the dataset from Klein & Noé (2024). In addition to the alanine sys- tems, we also investigate the 138-atom peptide chignolin, consisting of 10 residues (GYDPETGTWG) and notable for it's formation of B-hairpin structure in water solvent Honda et al. (2004). We provide additional dataset details in $D.
Datasets.
An upper bound on Marger (x) allows us to threshold sam- ples using the energy function, E(x) 7/s of the target. Formally, this corresponds to truncating the target distri- which places zero mass on high energy conformations. Correcting flow samples with respect to this truncated target introduces additional bias into the self-normalized importance sam- pling estimate, which precisely corresponds to the differ- ence in total variation distance between the two distributions (Arargtt Msrgge)). We prove this result using an intermedi- ate result in Lemma included in $B.
bution (x) == IP (Magge (x)) 2 log z
Our next theoretical result provides a prescriptive strategy of setting appropriate threshold y as function of the num- ber of samples K and effective sample size under (targer
Baselines. For baselines, we train prior sfte-of-the-arr equivariant Boltzmann generators. Specifically, we train the exactly invertible and equivariant SE(3) (3gg)))augeeteed cou- pling flow (Midgley et al., 2023a), and the equivariant con- tinuous normalizing flow (ECNF) employed in Transferable Boltzmann Generators (Klein & Noé, 2024). We also in- clude an improved variant, denoted ECNF++, as stronger baselines; this uses a refined flow matching objective, larger network, and improved optimization hyperparameters, full details provided in $E.4 for full details. We note that both SEE()-EACH and ECNF to be equivariant E(3) and hence generate samples of both global chiralities, which we resolve by applying flip transformation as in Klein & Noć (2024), for further details and related results see $F.
Proposition 3. Given an energy threshold E(x)= > % for ry > O large and the resulting truncated target dis- tribution (sarge(()) == P (Prazee(a)) (x) 2 Took Further, assume that the density of unnormalized importance weights w.r.l to Purrger is square integrable (w(()))² Co. Given tolerance p = E/ESS and bias of the orig- inal importance sampling estimator in total variation
b = TV(us Mrarge),, then the -ttruncation threshold with K K-samples for TV(uo, )marget) is:
7 2 X
log
12pE/exp( )
Kb
+ log 2.
(11)
We report effective sample size (ESS) along with Wasserstein-2 distances on the energy distribution E-W2 and dihedral angle torus T-W2 The energy distribution is highly sensitive fine-graines details whereas the di- hedral angles encode macrostructural information such as metastable state occupancy; full metric definitions are pro- vided in $E. Additional results for Wasserstein-2 distance on time- lagged independent component analysis (TICA) pro- jections TICA-W W-22 are provided in $F. We provide energy histograms in the main text whilst Ramachandran plots (Ra- machandran et al., 1963) detailing the mode coverage via dihedral angle distributions are presented in $F.
Metrics.
The proof for Proposition 3 is located in BB.3. Proposition 3 allows us to appropriately set a energy threshold y as func- tion of tolerance p that depends on ESS. In practice, this allows us to negotiate the amount of acceptable bias when dropping initial samples that obtain high-energy before any further AIS correction. Moreover, this gives firmer theo- retical foundation to existing practices of thresholding high importance weight samples (Midgley et al., 2023b;a).
Analogous to thresholding based on E(x), we can also threshold by the probability under the proposal flow with
6

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
(a) Alanine dipeptide
(b) Trialanine
(c) Alaninc tetrapeptide
(d) Hexa-alaninc
(c) Chignolin
Figure 2. Samples generated by SBG on peptide systems ranging from to 10 residues.
Table Quantitative results on alanine dipeptide and trialanine.
Baseline methods presented with SNIS.
Alanine dipeptide
Trialanine
Algorithm +
ESS f
E-W2 4
T-W2 4
ESS f
E-W2 4
T-W2 4
SE(3)-FACF
< 10-3
108.202
2.867
ECNF
0.119
0.419
0.311
ECNF (Ours)
0.275 0.010
0.914 +: 0.122
0.189 1 0.019
0.003 0.002
2.206 + 0.813
0.962 1 0.253
SBG AIS (Ours)
0.030 0.012
0.630 + 0.249
0.418 + 0.090
0.052 0.013
0.797 + 0.094
0.450 + 0.043
SBG SMC (Ours)
0.412 + 0.125
0.430 + 0.100
0.590 + 0.267
0.455 + 0.076
0.08
MD
Propoeal
0.06
Deasity
Densiry
Density
Donsiry
SMC
0.4
0.04
0.2
0.02
40.0
20.0
0.0
40.0
20.0
0.0
-0.00
20.0
0.0
40.0
20.0
0.0
6(x)
6(x)
8(x)
6(x)
(a) (E(3)-EECE
(b) ECNF
(c) ECNF++
(d) SBG
Figure 3. Energy histograms for baseline methods and SBG alanine dipeptide dataset.
having good overlap with the MD data, we find the SNIS reweighted performance to be substantially degraded at this lower temperature when using the same 0.2% weight clip- ping threshold as the original work; see $F for analysis of more aggressive clipping thresholds. The ESS and T-W of ECNF++ outperform both SBG variants by large mar- gin, although the inverse is true for the &-W2. Furthermore, the original ECNF model trained by Klein & Noč (2024) achieves superior E-W² to our proposed ECNF++ but infe- rior ESS and T-W2- These results are further substantiated by the energy distribution histograms in Figure 3.
Proposal
Denity
Density
0.04
SNIS
0.04
0.02
0,02
-1500
-100.0
-150.0
-100.
E(x)
8(x)
Figure Energy distribution histograms for baseline ECNF+ (left) and SBG (right) on trialanine dataset.
4.1. Results
Trialanine. Despite achieving acceptable performance on alanine dipeptide, ECNF was unable to scale to trialanine and was omitted. The computational cost of (E(3) -EACF (c.f. Table 5) precluded it's consideration. The SBG vari- ants are significantly stronger in all metrics compared to ECNF++, with SBG AIS achieving higher ESS and both AIS and SMC outperforming on E-W2 and T-W2. However the performance of ECNF++ is acceptable, and constitutes the first CNF-based Boltzmann generator on tripeptide sys- tem in Cartesian coordinates. There is no notable distinction in performance between SBG variants.
We evaluate SBG and our baseline methods with quantita- tive metrics summarized Table and Table 3. Where + present three models are independently trained and sam- pled; unless otherwise stated 10a particles sampled. We provide examples of SBG generated samples in Figure 2.
Alanine dipeptide. SEA3)-EACF was originally trained on an alanine dipeptide dataset at 800 K; we retrain on our more challenging 300 K data using the original codebase of Midgley al. (2023a). Despite the proposal distribution

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
Table 3. Quantitative results on alanine tetrapeptide and hexa-alamine. ECNF++ presented with SNIS.
Datasets +
Alanine tetrapeptide
Hexa-alanine
Algorithm 4
ESS f
E-W2 4
T-W2
ESS f
E-W2
T-W2 4
ENCF+ (Ours)
0.016 + 0.001
5.638 + 0.483
1.002
0.061
0.006 + 0.001
10.668 0.285
1.902 + 0.055
SBG AIS (Ours)
0.046 + 0.014
0.883 = 0.213
0.866
0.076
0.034 + 0.015
1.021
0.239
1.431 + 0.085
SBG SMC (Ours)
1.027 = 0.465
0.888
0.114
1.189
0.357
1.444 + 0.140
SE(C)-FACC
0.06
0.06
Proposal
ECNFL
0.55
GPU Hours
SMC
Densiy
0.04
0.04
SMC
10'
0.50
0.02
0.45
Standard
10
Adjused
0.00
0.00
ALD3
A1.3
AL4
AL6
10
10
10
E(x)
E(x)
Dataset
Timesteps
Figure 5. Energy distribution histograms for baseline ECNF++ (Ieft) and SBG (right) on alanine tetrapeptide dataset.
Figure 7. Left: GPU hours (NVIDIA L40S) for sampling and reweighting 10 points. Right: T-W on trialanine as a function of Langevin timestep discretization for both standard pe (x) and center of mass adjusted proposal energy functions pa (2x).
0.10
MD
Proposai SNIS
0.04
0.06
Densiyy
0.06
Density
SMC
1.25
Proposal
0.04
Density
1.00
Donsity
SNIS
0.75
SMC
0.50
0,02
-100.1
-5001
-100.0
-50.0
0.25
E(x)
8(x)
0.00
0.00
Figure 6. Energy distribution histograms for baseline ECNF++ (left) and SBG (right) on hexa-alanine dataset.
-500.0
-450.0
-4000 -350.0
E(x)
Figure 8. SBG interatomic distance histogram (left) and en- ergy distribution histogram (right) for decapeptide chignolin (GYDPETGTWG) SNIS E-W, 12.046, SMC &-W2 = 3.571.
Alanine tetrapeptide and hexa-alanine. At this scale the ECNF++ baseline diverges from the target distribution, flected particularly in E-W, In contrast, SBG is read- ily scalable up to hexapeptides, achieving greatly reduced E-W2 on both datasets. As reweighted samples under SBG show extremely high overlap with the ground truth (targe((1) we argue that SBG successfully solves these molecular sys- tems in comparison prior BGs. This conclusion is sup- ported by the energy histograms in Figure and Figure 6, in which the SNIS reweighted SNIS does not approximate the MD data well. in contrast to the good alignment of SBG. Notably, the proposals for ECNF++ have good overlap with the target density, indicating the error to be introduced by likelihood estimation itself.
trialanine function of Langevin timestep granularity for SBG SMC both with and without the center of mass proposal energy adjustment, as stated Equation (5). When the center of mass adjusted energy is employed we observe a strong inverse relationship between time discretization steps and T-W,, however without this adjustment there no clear relationship. This evidences both the efficacy of the center of mass adjustment improving reweighting as well the potential of SBG for inference-time scaling capability not present in the standard Boltzmann generator.
4.2. Scaling to Decapeptide
Inference scaling. To illustrate the scalability of SBG in relation to other methods plot in Figure 7 the GPU hours required by each method to sample 10tt points. We observe exponential scaling of inference time for ECNF++ as the size of the system grows, whilst SBG is less sensitive to system size and over an order of magnitude faster on the hexapeptide system. We additionally plot the T-W2 on
We now apply SBG to the decapeptide chignolin. As no other method can scale to this system we report energy his- tograms and distance plots for SBG SMC only in Figure 8, We observe success of SBG at matching the interatomic dis- tance distribution. We additionally observe a strong overlap of SMC sample energy distribution despite the notably poor
8

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
proposal overlap, providing further demonstration of the vi- ability of the SBG approach to molecular system sampling. Our application of SBG to chignolin represents significant step forwards in the scalability of BGs, where prior meth- ods struggled on even alanine tetrapeptide, as observable in results for ECNF++ E-W2 presented in Table 3.
6. Conclusion
In this paper, we introduce SBG an extension to the Boltz- mann generator framework that scales inference through the use of annealing processes. Unlike past BGs, in SBG, we scale training using non-equívariant transformer-based TarFlow architecture with soft equivariance penalties to 6 peptides. In terms of limitations, using non-equilibrium sampling as presented in SBG does not enjoy easy applica- tion to CNFs due to expensive simulation, which limits the use of modern flow matching methods in a SBG context. Considering hybrid approaches that mix CNFs through dis- tillation to an invertible architecture consistency- based objectives is thus natural direction for future work. Finally, considering other classes of scalable generative models such as autoregressive ones which also permit exact likelihoods is also a ripe direction for future work.
5. Related Work
Boltzmann generators (BGs) (Noé et al., 2019) have been applied to both free energy estimation (Wirssherger et al., 2020; Rizzi et al., 2023; Schebek et al., 2024) and molecular sampling. Initially, BGs relied on pystem-speciff repre- sentations, such as internal coordinates, to achieve relevant sampling efficiencies (Noé et al., 2019; Köhler et al., 2021; Midgley et al., 2023b; Köhler et al., 2023; Dibak et al., 2022). However, these representations transferable across different systems, leading to the develop- ment of BGs operating in Cartesian coordinates (Klein et al., 2023b; Midgley et al., 2023a; Klein & Noé, 2024). While this improves transfcrability, they are currently limited in scalability, struggling to extend beyond dipeptides. Sealing to larger systems typically requires sacrificing exact sam- pling from the target distribution (Jing et al., 2022; Abdin & Kim, 2023; Jing et al., 2024a; Lewis et al., 2024). An alternative to direct sampling from Puaree (x) is to generate samples iteratively by learning large steps in time (Schreiner et al., 2023; Fu et al., 2023; Klein et al., 2023a; Diez et al., 2025; Jing et al., 2024b; Daigavane et al., 2024) to accelerate methods such as molecular dynamics via noarse-graining
generally not
Acknowledgements
The authors thank Damien Ferbach, Tara Akhound-Sace Lars Holdijk, Kacper Kapuániak, Kirill Neklyudov, Michael Albergo, and Majdi Hassan for insightful conversations and feedback. In addition, the authors thank Paul Skaluba for constructive comments on Proposition I of an older draft.
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
ca), Mila (https: /mmmmmmaaaa/aaa///../ htw.ii......eee....a.eea and NVIDIA. AJB is partially supported by an NSERC Post-doc fel- lowship. This research is partially supported by the EP- SRC Turing AI World- World-Leading Research Fellowship No. EP/X040062/1 and EPSRC AI Hub No. EP/Y028872/1.
Amortized sampling. The field of sampling has seen renewed interest with the rise of generative models. In samplers has seen
particular, the use of diffusion- rapid application with a plethora of approaches exploiting the favorable theoretical properties of mode-mixing of dif- fusion models (Berner et al., 2024; Vargas et al., 2023; Richter al., 2024; Zhang & Chen, 2022; Vargas et al., 2024). While initial approaches focused on simulation- based dynamics, including both overdamped and under- damped Langevin (Blessing et al., 2025; Chen et al., 2025), it is expected that simulation-free methods that also exploit diffusion properties et al., 2021: De Bortoli et al., 2024) are an attractive oppor- tunity to tackle larger-scale systems due to their scalability. Finally, flow-based models have also been employed for sampling with classical flows augmenting MCMC (Arbel et al., 2021; Gabrié et al., 2021; Matthews et al., 2022; Midg- ley et al., 2023b; Hagemann et al., 2023), and through CNFs that construct ODE bridges, such as linear interpolants be- tween the prior and target (Máté & Fleuret, 2023), and more general bridges that rely on satisfying the mass transport equations (Tian al., 2024; Fan et al., 2024).
Impact Statement
This work studies sampling from Boltzmann densities, a problem of general interest in machine learning and A14Science that arises both in pure statistical modeling and within applications. We highlight the training Boltzmann generators molecular tasks in turn applicable drug and material discovery. While we do not foresee immediate negative impacts of our advances in this area, we encourage due caution whilst scaling to prevent their potential misuse.
Sachound-Sadegee et al., 2024; Huang
9

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
References
Buch, I., Giorgino, T., and De Fabritiis, G. Complete re- construction of an enzyme-inhibitor binding process by molecular dynamics simulations. Proceedings of the Na- tional Academy of Sciences, (08(25)1018441011)) 2011. (Cited on page 1)
Abdin, O. and Kim, P. M. Pepflow: direct conforma- tional sampling from peptide energy landscapes through hypernetwork-condee conditioned diffusion. bioRxiv, pp. 2023- 06, 2023. (Cited on page 9)
Chen, J., Richter, L.., Berner, J., Blessing, D., Neumann, G., and Anandkumar, A. Sequential controlled langevin diffusions. In International Conference on Learning Rep- resentations (ICLR), 2025. (Cited on page 9)
Abramson, J., Adler, J., Dunger, J., Evans, R., Green, T.,
Pritzel, A., Ronneberger, O., Willmore, L., Ballard, A. J., Bambrick, J., al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, pp-
1-3, 2024. (Cited on page 2)
Chen, R. T. Q., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. Neural ordinary differential equations. Neural Information Processing Systems (NIPS), 2018. (Cited on pages 2 and 3)
Agapiou, S., Papaspiliopoulos, O., Sanz-Alonso, D., and Stuart, A. M. Importance sampling: Intrinsic dimension and computational cost. Statistical Science, pp. 2017. (Cited on page 16)
Daigavane, A., Vani, B. P., Saremi, S., Kleinhenz, J., and Rackers, J. Jamun: Transferable molecular conforma- tional ensemble generation with walk-jump sampling. arXiv, 2024. (Cited on page 9)
Aahound-Sadegee T., Bector-Brooks, J., Bose, J., Mittal, S., Lemos, P., Liu, C.-H., Sendera, M., Ravanbakhsh. S., Gidel, G., Bengio, Y., Malkin, N., and Tong, A. Iterated denoising energy matching for sampling from boltzmann densities. In International Conference on Machine Learn- ing (ICML), 2024. (Cited on pages 2 and 9)
De Bortoli, V., Hutchinson, M., Wirnsberger, P., and Doucet, A. Target score matching, arXiv, 2024. (Cited page 9)
Deringer, V. L. Modelling and understanding battery mate- rials with machine-leaminng-diieee atomistic simulations. Journal of Physics: Energy, 2(4):041003, oct 2020. doi: 10.
Albergo, M. S. and Vanden-Eijnden, E. Building normal- izing flows with stochastic interpolants. International Conference on Learning Representations (ICLR), 2023. (Cited on page?)
1. (Cited on page 1)
Dibak, M., Klein, L., Krämer, A., and Noé, F. Tem- perature steerable flows and Boltzmann generators. Phys. Rev. Res., 4:L042005, Oct 2022. doi: 10.1103/ PhysReevRseaacch44..000000
Albergo, M. S. and Vanden- E. Nets: A non- equilibrium transport sampler. In International Con- ference on Machine Learning (ICML), 2025. (Cited on pages 2, 5, and 14)
(Cited pages and 21)
Diez, J. V., Schreiner, M., Engkvist, O., and Olsson, S. Boltzmann priors for implicit transfer operators. Interna- tional Conference on Learning Representations (ICLR),
Alexey, D. An image is worth 16x16 words: Transformers for image recognition at scale. In International Confer- ence on Learning Representations (ICLR), 2021. (Cited on page
2025. (Cited on page
Dinh, L., Sohl-Dickstein, J., and Bengio, S. Density esti- mation using Real NVP. International Conference on Learning Representations (ICLR), 2017. (Cited on pages 1 and 3)
Arbel, M., Matthews, A., and Doucet, A. Annealed flow transport monte carlo. In International Conference Machine Learning, pp. 318-330. PMLR, 2021. (Cited on page 9)
Domingo-Enrich, C., Drozdzal, M., Karrer, B., and Chen, R. T. Q. Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control. In International Conference on Representation Learning (ICLR), 2025. (Cited on pages 18 and 19)
Berner, J., Richter, L., and Ullrich, K. An optimal con- trol perspective on diffusion-based generative modeling. Transactions on Machine Learning Research (TMLR), 2024. (Cited page 9)
Doucet, A., De Freitas, N., Gordon, N. J., et al. Sequential Monte Carlo methods in practice, volume 1. 2001. (Cited on pages 2, 4, and 5)
Blessing, D., Berner, J., Richter, L., and Neumann, G. Un- derdamped diffusion bridges with applications to sam- pling. In International Conference on Learning Repre- sentations (ICLR), 2025, (Cited on page 9)
Eastman, P., Swails, J., Chodera, J. D., McGibbon, R. T., Zhao, Y., Beauchamp, K. A., Wang, L.-P., Simmonett, A. C., Harrigan, M. P., Stern, C. D., et al. Openmm 7: Rapid development of high performance algorithms for
Bose, A. J., Brubaker, M., and Kobyzev, I. Equivariant finite normalizing flows. arXiv, 2021. (Cited on pages and 4)
10

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
molecular dynamics. PLoS computational biology, 13(7): e1005659, 2017. (Cited on page 21)
Hutchinson, M. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Com- munications in Statistics Simulation and Computation, 19(2):433-453 1990. doi: 10. 10.108000361091100088 (Cited on page 25)
Esser, P., Kulal, S., Blattmann, A., Entezari, R., Müller, J., Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, F., Podell, D., Dockhorn, T., English, 7., Lacey, K., Good- win, A., Marek, Y. and Rombach, R. Scaling rectified flow transformers for high-resolution image synthesis. In International Conference on Machine Learning (ICML), 2024. (Cited on page 25)
Jarzynski, C. Nonequilibrium equality for free energy dif- ferences. Phys. Rev. Lett., 78:2690-2693 Apr 1997. doi: 10.1 103/PhysRevleet 78.2660... (Cited page 5)
Jing. B., Corso, G., Chang. J., Barzilay, R., and Jaakkola, T. Torsional diffusion for molecular conformer generation. Advances in Neural Information Processing Systems, 35: 24240-24253, 2022. (Cited on page 9)
Fan, M., Zhou, R., Tian, C., and Qian, X. Path-guided particle- based sampling. International Conference on Machine Learning (ICML., 2024. (Cited page 9)
Flamary, R., Courty, N., Gramfort, A., Alaya, M. Z., Bois-
Jing. B., Berger, B., and Jaakkola, Alphafold meets flow matching for generating protein ensembles. In Interna- tional Conference on Machine Learning (ICML), 2024a. (Cited page 9)
bunon, A., Chambon, S., Chapel, L., Corenflos, A., Fatras, K., Fournier, N., Gautheron, L., Gayraud, N. T., Janati, H., Rakotomamonjy, A., Redko, I., Rolet, A., Schutz, A., Seguy, V., Sutherland, D. J., Tavenard, R., Tong, A., and Vayer, T. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1-8,, 2021. (Cited on page2) 23)
Jing, B., Stärk, H., Jaakkola, T., and Berger, B. Generative modeling of molecular dynamics trajectories. In Neural Information Processing Systems (NeurIPS), 2024b. (Cited page 9)
Fu, X., Xie, T., Rebello, J., Olsen, B., and Jaakkola, S. Simulate time-integrated poarse-grained molecular dynamics with multi-scale graph networks. Transactions Machine Learning Research, 2023. (Cited page 9)
Karczewski, R., Heinonen, M., and Garg, V. Diffusion mod- els as cartoonists! the curious case of high density regions. In International Conference on Learning Representations (ICLR), 2024. (Cited pages 14, 19, and 20)
Gabrié, M., Rotskoff, G. M., and Kanden-Eijnder E. Effi- cient Bayesian sampling using normalizing flows to assist Markov chain Monte Carlo methods. arXiv, 2021. (Cited on page 9)
Klein, L. and Noé, F. Transferable boltzmann generators. Advances in Neural Information Processing Systems, 2024. (Cited on pages 6, 7, 9, 21, 22, 25, 26, 27, and 29)
Garcia Satorras, V., Hoogeboom, E., Fuchs, F., Posner, I., and Welling, M. E(n) equivariant normalizing flows. Neural Information Processing Systems (NeurIPS), 2021.
Klein, L., Foong, A. Y., Fjelde, T. E., Mlodozeniec, B., Brockschmidt, M., Nowozin, S., Noé, F., and Tomioka, R. Timewarp: Transferable acceleration of molecular dynamics by learning time-coarsened dynamics. Neural Information Processing Systems (NeurIPS), 2023a. (Cited on pages and 21)
(Cited on page 4)
Grathwohl, W., Chen, R. Т. Q., Bettencourt, J., Sutskever, I., and Duvenaud, D. FFJORD: free-form continuous dynamics for scalable reversible generative models. In International Conference on Representation Learning (ICLR), 2019. (Cited on page 25)
Klein, L.,, Krämer, A., and Noé, F. Equivariant flow match- ing. Neural Information Processing Systems (NeurIPS), 2023b. (Cited on pages and 9)
Hagemann, P. L., Hertrich, J., and Steidl, G. Generalized normalizing flows via Markov chains. 2023. (Cited on page 9)
Köhler, J., Klein, L., and Noé, F. Equivariant flows: exact likelihood generative learning for symmetric densities. International Conference on Machine Learning (ICML), 2020. (Cited on page 4)
Honda, S., Yamasaki, K., Sawada, Y., and Morii, H. 10 residue folded peptide designed by segment statistics. Structure, 12(8):1507-1518; 2004, (Cited on page 6)
Köhler, J., Krämer, A., and Noé F. Smooth normalizing flows. In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 2796- 2809, 2021. (Cited on page 9)
Huang, J., Jiao, Y., Kang, L., Liao, x., Liu, J., and Liu, Y. Schrödinger-Föllie Föllmer sampler: sampling without ergodic- ity. arXiv, 2021. (Cited on page 9)

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
Köhler, J., Invernizzi, M., De Haan, P., and Noé, F. Rigid body flows for sampling molecular crystal structures. In- ternational Conference on Machine Learning (ICML), 2023. (Cited page 9)
Neal, R. M. Annealed importance sampling. Statisfics and computing, 1 11:125-139, 2001. (Cited on pages and 5)
Noé, F., Schütte, C., Vanden- Eijnden,, E., Reich. L., and Weikl, T. R. Constructing the equilibrium ensemble of folding pathways from short aff-equilibrium simulations. Proceedings of the National Academy of Sciences, 106 (45): (45):19011-19015) 2009. (Cited on page 1)
Lee, H., Pabbaraju, C., Sevekari, A. P., and Risteski, A. Universal approximation using well-conditioned normal- izing flows. Advances in Neural Information Processing Systems, 34:12700-12711, 2021. (Cited on page 4)
Noé, F., Olsson, S., Köhler, J., and Wu, H. Boltzmann gener- ators: Sampling equilibrium states of many-hodyy systems with deep learning. Science, 365(6457):eaaw/ I 147, 2019. (Cited on pages 1, 2, 3, and 9)
Lewis, S., Hempel, T., Jiménez Luna, J., Gastegger, M., Xie, Y., Foong, A. Y., García Satorras, V., Abdin, O., Veeling, B. S., Zaporozhets, I., et al. Scalable emulation of protein equilibrium ensembles with generative deep learning. bioRxiv, pp. 2024-12, 2024. (Cited on page 9)
Owen, A. B. Monte Carlo theory, methods and examples. 2013. (Cited on page 4)
Lindorff-Larsen. K., Piana, S., Dror, R. O., and Shaw, D. E. How fast-folding proteins fold. Science, 334(6055):517-- 520, 2011. (Cited on page 1)
Papamakarios, G., Pavlakou, T., and Murray, I. Masked autoregressive flow for density estimation. Advances in neural information processing systems, 30, 2017. (Cited on page 4)
Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. Interna- tional Conference Learning Representations (ICLR), 2023. (Cited page 3)
Parrinello, M. and Rahman, A. Crystal structure and pair potentials: A molecular- -ynamicss study. Physical review letters, (1(4))11996 1980. (Cited on page 1)
Liu, Q Rectified flow: A marginal preserving approach to optimal transport. arXiv, 2022. (Cited on page 3)
Ramachandran, G. N., Ramakrishnan, C., and Sasisekharan, V. Stereochemistry of polypeptide chain configurations. Journal of Molecular Biology, pp. 95-99, 1963. (Cited on page 6)
Liu, X., Zhang, X., Ma. J., Peng, J., and Liu, Q. Instaflow:
One step is enough for high-quality diffusion- text- to-image generation, 2024. (Cited on page 25)
Rezende, D. and Mohamed, S. Variational inference with normalizing flows. International Conference on Machine Learning (ICML). 2015. (Cited on pages and 3)
Loshchilov, and Hutter, F. Decoupled weight decay reg- ularization. In International Conference on Learning Representations (ICLR), 2017. (Cited on page 25)
Richter, L., Berner, J., and Liu, G.-H. Improved sam- pling via learned diffusions. International Conference on Learning Representations (ICLR), 2024. (Cited on page 9)
Máté, B. and Fleuret, F. Learning interpolations between boltzmann densities. Transactions on Machine Learning Research, 2023. ISSN 9335-8856. (Cited on pages and 9)
Matsumoto, M., Saito, S., and Ohmine, Molecular dynam- ics simulation of the ice nucleation and growth process leading water freezing. Nature, 416(6879):4094411 2002. (Cited page 1)
Rizzi, A., Carloni, P., and Parrinello, M. Free energies at qm accuracy from force fields via multimap targeted estimation. Proceedings of the National Academy of Sciences, 2023. (Cited page 9)
Matthews, A., Arbel, M., Rezende, D. J., and Doucet, A. Continual repeated annealed flow transport monte carlo. International Conference on Machine Learning (ICML), 2022. (Cited on page 9)
Satorras, V. G., Hoogeboom, E., and Welling, M. E (n) equivariant graph neural networks. International Con- ference Machine Learning (ICML), 2021. (Cited on page 2)
Midgley, L. I., Stimper, V., Antorán, J., Mathieu, E., Schölkopf, B., and Hernández-Lobace J. M. SE(3) equiv- ariant augmented coupling flows. Neural Information Processing Systems (NeurIPS), 2023a. (Cited pages 2, 6, 7, 9, and 24)
Schebek, M., Invernizzi, M., Noé, F., and Rogal, J. Efficient mapping of phase diagrams with conditional boltzmann generators. Machine Learning: Science and Technology. 2024. (Cited on page 9)
Schreiner, M., Winther, O., and Olsson, S. Implicit transfer operator learning: Multiple time-resolution models for molecular dynamics. In Thirty-sevntth Conference on Neural Information Processing Systems, 2023. (Cited on page 9)
Midgley, L. I., Stimper, V., Simm, G. N., Schölkopf. B., and Bormández-Lobateee J. M. Flow annealed importance sam- pling bootstrap. International Conference on Learning Representations (ICLR), 2023b. (Cited on pages and 9)
12

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
Skreta, M., Atanackovic, L.,, Bose, A. J., Tong, A., and Nek- lyudov, K. The superposition of diffusion models using the itō density estimator. In International Conference Learning Representations (ICLR), 2025. (Cited on pages 14 and 19)
Teshima, T., Ishikawa, I., Tojo, K., Oono, K., Ikeda, M., and Sugiyama, M. Coupling-basee based invertible neural networks are universal diffeomorphism approximators. Advances in Neural Information Processing Systems, 33:3362-3373, 2020. (Cited on page 4)
Tian, Y., Panda, N., and Lin, T. Liouville flow impor- tance sampler. In International Conference on Machine Learning (ICML), 2024. (Cited on page 9)
Tong, A., Fatras, K., Malkin, N., Huguet, G., Zhang, Y., Rector- Brotor.Brooks, J., Wolf, G., and Bengio, Y. Improv- ing and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research, 2024. (Cited on pages and 25)
Vargas, F., Grathwohl, W., and Doucet, A. Denoising dif- fusion samplers. International Conference Learning Representations (ICLR), 2023, (Cited on page 9)
Vargas, F., Padhy, S., Blessing, D., and Nüsken, N. Trans- port meets variational inference: Controlled Monte Carlo diffusions. International Conference on Learning Repre- sentations (ICLR), 2024. (Cited on pages 5 and 9)
Wimsberger, P., Ballard, A. Papamakarios, G., Aber- crombie, S., Racanière, S., Pritzel, A., Jimenez Rezende, D., and Blundell, C. Targeted free energy estimation via learned mappings. J. Chem. Phys., 2020. (Cited pages and 9)
Zhai, S., Zhang, R., Nakkiran, P., Berthelot, D., Gu, J., Zheng, H., Chen, T., Bautista, M. A., Jaitly, N., and Susskind, J. Normalizing flows are capable generative models. arXiv, 2024. (Cited on pages 2, 4, and 26)
Zhang, Q. and Chen, Path integral sampler: stochastic control approach for sampling. International Conference on Learning Representations (ICLR), 2022. (Cited on
page 9)
13

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
A. Alternate Paths
A.1. Proposal Free Langevin Dynamies
We can also modify the Langevin SDE in Eq. to include an additional drift term vy (r+) E Id as follows:
dæ,- - V&,(xx)dT + (xx)dd + /2e,
Under perfect drift My (T) the log weights do not change and there is no need for correction. For imperfect drift the corresponding coupled ODE time-evolution of log-weights d d'log w, needed to apply AIS was derived in NETS (Albergo & Hunden-Eijnded 2025, Proposition 3):
dw,, V My (x, )dT (x,, e, (2, )dT - 0, (x,
In contrast to learning a drift as done in NETS (Albergo & Vanden-Eijnde 2025) we now illustrate that a judicious choice of ur (x+) eliminates the need to compute the gradient of log-likelihood under the proposal. For instance, we can choose or ()(x)) = e, VE, (x+) - e, v
E(a.) kg
which by straightforward calculation gives the following SDE:
dx, -, (E)(x))dT + (x,) + /2,,W, (( KBT
= =
dr de,
(12)
This new SDE greatly simplifies the simulation of samples 2, as it independent of the proposal energy (Er())) - poggo((>)) However, the log weights ODE still requires the computation of the gradient of the proposal energy. The form Eq. 12 suggests the possibility of massively parallel simulation schemes under regular normalizing flow and CNF. However, due to simulatio the log weights remains expensive for CNFs due to the need to compute the divergence operator in Eq. 4. Furthermore, while recent advances in divergence- frree density estimation via the It density estimator (Skreta et al., 2025; Karczewski et al., 2024) might appear attractive we show that the log density under this estimator is necessarily biased and may limit the fidelity of self-normalized importance sampling incurs non noggtagotia added bias. For case of presentation, we present this theoretical investigation in $C.2 and characterize the added bias in Proposition In totality, this limits the application of continuous BG's to only the conventional IS setting, unlike finite flows like TarFlow which can benefit from non-equilibrium transport and AIS.
B. Proofs
B.1. Proof of Proposition 1
Prarge ).), consider the decomposition of a data point x. E RRxx3 into its constituent mean-free component, E 1)x-1)x3 + Rax3 and center of mass c E R3, x = H t c, where c ~ NO, o²). Now, assume both the proposal po(x)) and the adjusted proposal (x(x) factorize independently over the mean-free component and the center of mass. Then setting pl = (z) 1/oxx((lcll), leads the following inequality on the effective sample size in the limit of K >
Proposition 1.
Given an
(1)) po (z)
Miurger pg(x)
ESS
ESS
(6)
Proof. Recall the definition of effective sample size using Kish's formula, both po and the adjusted proposal (x:):
E w(xx)) 2
I K Mtarge((")) p/po))(()))))
Mraggot(x)
ESS
(x))
ΣS (w(a)))2
Σ wo(o))2
ΣN w(xt)2
Σx K w((x') )²
Σ Prangot (x') \/2 (x+))
M(arget(z))
ESS
(x(x)
Σn (1T())))22
D (xi))
ΣK K (Prange (x.) \/8 (x)))²
14

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
We rewrite the weights for the regular proposal's ESS calculation follows,
wxx') = Murgal (I' + +c)
Murgel (x)) pa (z')) ((()
po ( + c)
(xx) ) = xo(x') w(c)
where we exploited the translation invariance of [Murget to remove the center of mass e and also the independence between the mean and in the proposal. Since c (N o²) we can write p(c) p(l|e],, 8, 8, σ) in spherical coordinates to follow scaled Chi distribution ||e|| ~ (xa(]lll) with angular components that follow independent uniform distributions (0, 6) ((O)((o). Fixing canonical angular components (0, Φ) we have p(c) ( and w(c) x/oxx (Ile]).
For the adjusted proposal we set p(c) 1/0Xs which gives the following weights:
Marrger (I' +c)
Mtarper (F') (xt ( }p(c)
Mtarget (3') w/ (x²).
w w(('')
(13)
pz( +c) po
po()*)
We now seek to prove that:
1i=1 w (()))((())
K
ΣE - we wo(x)) E w(())
(14)
K - w w
Now, denote X = w('') and y w(c) random variables over the sample index i. By construction, X and Y are independent for each i. Furthermore, all (X', YY* pairs are i.i.d. for i E [K].
We may now formalize equation 14 as proving
E EL E X
(E[X])
(15)
Y
Because and independent for each sample we know:
E[ =
E = E[X B|y
(16)
Hence
E X
Y1)²
(E(X))*
(E(Y])*
(17)
E
X
Y
E[
E|Y
If Var(YY > O, then B[y?)> > ( Consequently,
())
(18)
Y
Thus, at the level of population expectations,
[E[X YJ)*
(E
EEL
[E[Y]) 2
(E[X))
(19)
E
Y
E[X2]
E
Y
E[X²]
Therefore, applying the adjustment is strictly better by ESS than unadjusted.
B.2. Proof of Lemma 1
We first prove a useful lemma that computes the total variation distance between the original distribution of the normalizing flow po and the truncated distribution fo before proving the propositions.
15

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
Lemma 1. Let pe be generative and denote ( the 6-truncated distribution such that po(x) P(po(x) for small 8 >> O. Define the constant = P(po(()) 8) the event where the truncation occurs. Then the total variation distance between the generative model and its truncated distribution is TV(ps po) B.
Proof. We begin by first characterizing the total variation distance between flow after correction with importance sampling p(æ) with truncated distribution p(r). Recall that the truncated distribution is defined as follows:
((x) = P(((x) 2 6) p(x)](p(() 2 83 [[(p(a) 2 6)p(r)dr'
(20)
where II is the indicator function. Denote the events P(X 2 8) and B P(X <6) for the random variance X ~ p(æ).. Clearly, + = 1 and a [[((() 2 6)p(())d Now consider the total variation distance between these two distributions:
TV P) sup E (()=) Ea [o(+) = 1 p(x)) deo
f(xx) dz.
(21)
where = {Φ | S 13. Next we break up the event space into two regions R₁ and R2 which correspond to the events p(x) ô and p(x) 2 respectively. Now consider the total variation distance in the region Rz whereby construction p(xr) 0,
2 //1 (p(x) f(x)) dr = 2 S. p(x)dr= = 83 2
1
(22)
similar computation R2 gives,
1 2 Jfe (o( - po(2))de
p(x) R2 a
1
1
dx ar
1)
B
p(x)
dx
p(x)
(23)
Ra
a
where exploited the fact that (x) (w) in the first equality and that Combining these results we get the full total variation distance:
SR. po(x))da in the second equality.
TV p) f
[p(x) - p(x) f
JR. (p(x) - p(x))dr 2 1p(x) p(x) B.
(24)
Thus the TV(p, p) = B and 0 in the trivial case where a and the truncated distribution are the same.
B.3. Proof Proposition 3
Proposition 3. Given energy threshold E(x) > for y large and the resulting truncated target distribution
Piarger (x)) = P (aarged (1:) 2 Tog, ) Further, assume that the density of unnormalized importance weights w.r.t. to
Pharger is square integrable (=² Given tolerance p L/ESS and bias of the original importance sampling estimator in total variation = TV(uo Htaryet), then the n-truncation threshold with K K-samples for TV{po, Pranget) is:
1 log 7 2 X
Kb
+ log E.
(11)
12pl exp( AX)
Proof. We start by recalling well-known result stating the bias of self-normalized importance sampling found in Agapiou et al. (2017, Theorem 2.1) using K samples from the proposal p(æ):
KÊ, K wxx2)22
12p K
K
sup 11 $1
IE [u (Ø) Mtarger (b)]
(e)) 5
p
(25)
ESS
Σ w(xt))
where the terms 2 (d) = Σ, w<x))((')) is the self-normalized importance estimator of Mraget with samples drawn according to x' po(o) and [[((x|] s a bounded test function.
16

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
By truncating using energy threshold (xr 7, for a large O, we truncate the support of Musget (2:) by cutting off low probability regions that constitute high-energy configurations. More precisely, we have Muarett
(uurget (x)
log 2
where log 2 is as defined in Eq. 10. Note that (*)) is absolutely continuous w.r.t. to Mtarget as the support is contained up to modulo measure zero sets. The importance sampling error incurred by using Parget can be bounded as follows:
sup 1 S1
B 100 (0) - MMarger
sup
[ [uo (Ø) (targe((()) + sup 11
E [Etarger
- Pearget ( [6]]
(26)
1| x1
12p
+ 11
(27)
K
12p
< K B1.
(28)
The first inequality follows from the triangle inequality. Here we note that P is the ESS which corresponds to using importance weights computed with respect to the truncated target Aurget rather than Murget- The constant 31 = TV (ftarget Murreet and follows from an application of Lemma 1. Further, note that 2 á since ESS must increase and thereby ğ decreases-4a as the distributional overlap between the two distributions decreases. Now observe,
31 3=P(X< log 1
where samples follow the law X (z)ge()). Then a direct application of Chernoff's inequality log
gives us (x<
< log
exp
Elexp (-XX)]. Thus the additional bias incurred is,
12p K
12p K
Ay
sup o S1
E (10 (6) Prarget [6)]
+
5
exp
log Elexp(-AX))
(29)
Where the term Elexp(-AXX) is the moment generating function. Setting b = TV Murgaa), then we have
1 log
Kb (ellleepp(AAA0)
7 x
+ loo 2.
(30)
B.4. Proof of Proposition 4
Proposition 4. Assume that the density of the model pe after importance sampling μr is absolutely continuous with respect the target Prorser- Further, assume that the density of unnormalized importance weights square integrable (wxx))) < so. Given a tolerance p p= 1/ESS of the original importance sampling estimator under pa and bias of the importance sampling estimator total variation b TV(uo Praget, then the 8-truncation for the truncated distribution po(x)= == =P(pp(x)) 5) threshold with K -samples is:
1 8 2 X
Kb
log
(31)
(22E(exp((AXX)
Proof. We aim to bound the total variation distance TV Marget) of using the truncated distribution P(pe(r) > 8) by again recalling the bias self- norfmnormalized importance sampling using K samples from (o (x:):
12p K
K ESS
K E K
w(x5)²
sup |+
E Ju (Ø) Purrer (o)]
(32)
51
Σ w(xx))
where the terms 1 (0) ΣK w(x))(('') is the self-normalized importance estimator of Maaret with samples drawn according to x² ~ (z) and [[o()) < is a bounded test function. We next characterize the error introduced by using the truncated distribution Po for importance sampling in place of po by first defining the truncated K K-sample neff-normalized
17

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
importance estimator in (0) E3 ((x)))(x²)) where x/ (z)) Specifically, bound the total variation distance:
TV()0.10) = sup 1M S1
E [0 (0) - 11 (+)]
(3.3
K
sup
Bxtrup
Σ w( (o(²)(((²))
Σ
( (o(2))
(34)
11
1
E Bginpa
Σ 2((2))
Exi mopo
Σ ii (1x7)
(35)
2
Here in the second equality, we used the fact that the test function is bounded || I 1 Next, we apply Lemma I and leverage the fact that the nelf-normllized weights are also bounded and achieve a bound on the total variation distance,
K Σ x(x')
K
1)
TV
1
Byirpo
- Exi up
Σ
in (x²²)
(36)
2
=32
(37)
where B2 the probability mass P(X- < 6) when X po(x). Like previously, the overall error can be bounded using the triangle inequality
sup |+
E
Pg (6)
Paarger (0) 5 sup
IE [ (6) Marger (+) + sup all 31
IE
[n K (Ø) - 2 [6)
(38)
$1
12p
+ 192
(39)
K
12p
K + B2.
(40)
Where the last inequality follows from the same logic as Proposition 3 where ESS goes up after truncation and therefore P P. A direct application of Chernoff's inequality gives us IP(X < 6) B2 S exp(A0))[exp(-AA))) where we used
the moment generating function of po().. Thus the additional bias incurred is.
12p
12p
sup 1
E
[ng K (Ø) Parget (0) S
B2 <
+exp(A6)Elexp(-AN)
(41)
K
Setting b TV Munge) as the bias, then we have
1 log
Kbb 1epB
8 X
(42)
-XX)]
C. Itō Filtering
C.1. Flow Matching SDE
As shown in Domingo Comingo-Enrich et al. (2025) we can write Flow Matching with Gaussian conditional paths and Diffusion models under unified SDE framework given a reference flow:
8,10 + 0gT1
(43)
where (oo)tejo.1); (3:)eto,1) are functions such that oo with linear interpolants that we consider we have:
31 = 0 and
= -30 1. 1.. In the specific case of flow matching
Ft (1 t)a0 + tx
(44)
18

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
The unified SDE for both flow matching and continuous-time diffusion models introduced in Domingo-Enrich al. (2025) is then:
( + + s(I,, t) + o,dW,, Ri
(it de = Be Cx
dx, rig2
3:
(45)
at
where s( t) is the score function estimated by the diffusion model. Thus the flow matching SDE is:
dr+ (2f - Fit +) dt + o,dWe, o, (2 t)t)
(46)
In fact, the Stein score can be estimated from the output velocity field and vice-versa:
1ff,o(t, x4) zt. fe,((, = 3t + (1 V log (24)
v (ogg)((z))=
(47)
1-t
t
Rewriting Eq. 46 in terms of the score function get,
dx: t + Volgpe(((EE) + a.dW.
(48)
C.2. Itō Filtering
Proposition 5. Assume that the density of the model pa after importance sampling ps is absolutely continuous with respect to the target Parger- Further, assume that the density of unnormalized importance weights is square integrable (w(x))) < cc. Let r(xo) be the Itô density estimator for log po(ao) of the flow matching SDE:
St + (so(t,,, + o.dWe at = dx,
(2
(49)
Given p= 1/ESS. and Ç > O which is the weight clipping threshold. Then the additional bias of using the Itô density estimator for importance sampling fr.0 with clipping is:
12p [ s
sup 11
E (Ø) Mearger (6)]
+ As + 3a.
(50)
K
where 8g- = TV(ur,0 Ho) and B4 TV(pr,0 pr.
We now recall Itô's lemma which states that for a stochastic process,
dæ, f((t + g,dW,,
(51)
and a smooth function h : Rdd + the variation of as a function of the stochastic SDE can be approximated using a Taylor approximation:
dh(t, = (00 (t, 22)
dr h(t, xf) + 0
+
dt + at Dr h(t, d.W.
(52)
Ot
where Ag is the Laplacian. We will use Itô's Lemma with h(t, xz) = et al., 2025; Karezewski et al., 2024) but for flow models
pr (x;;) to obtain the Itô density estimator (Skreta
dlogpe(<.)) = ( (xx)) + dr log (x.)) [(1, 2x) + 2 - log (x))) di + at log pe (a())dd, Or
1
(53)
0t
To solve for the change in density over time we can start from the log version of the Fokker-Plank equation:
1 of log - V. (x) (F(..0)
Øt log (x) = - (f(t,x))) +
2 ,o log
(54)
in the general case we end with:
(z)) -V. (f(t, oo (z))( + V. (oN di + V. (og dW,
(55)
19

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
We now apply this the flow-matching SDE Eq. 48 written in terms of the score function. In particular, have,
dl - V (o/V [ + It t -o, V loggoog((ar))) + == dt
+ 01 v (zx)) dW.
dlogpe((re) = -d//L oo v. (og)(()))) di + V. log dW.
(56)
The above equation makes an implicit assumption that have access the actual ground truth score function of V log, (2t) rather than the estimated one 5ø, expressed via the vector field as in Eq. 47. When working with imperfect score estimates we have the following SDE:
dr. = St + Vsoll, + o.dW t
(57)
The score estimation enror causes discrepancy in log (xx) estimates whose error is captured in the theorem from Kar- czewski et al. (2024)[Theorem 3J:
looro log Po(x)) + Y
(58)
where log rn is the bias of the log density starting time 0 of the auxiliary process that does not track xy correctly due to the estimation error of the score. Also, Y is random variable such that that bias of ro is given by:
Y = B.
10 (()))s(((, ax) - V pe
(59)
Thus the Itô density estimator forms an upper bound to the true log density, i.c. ro(zo) 2 log po(p). This allows us to form an upper bound on the normalized log weights as an expectation,
E(xo) - log Pa (x0) c
E Exo (mo [log Th (ax)] = Eurg~pu(zo)
KBT
E(za)
< (xo~ppl
KnT -ro(21o)
where O' is constant. We define log (2.) = E(zu) constants. We can now compute the additional bias self-normalized importance sampling estimator ,
KeT ro( the new normalized importance weights, module
TV
sup >
E Tp (0) H (Φ)]
(60)
sup |0 1
Eytopo
Σ Wr (xz)) }ob(zz))
Eys Apo
Σ W (+()))
(61)
Σ w(x'') 1)
K
K
E Ena
Σ Wr (2+)
- Eng mpp
(62)
2
1
E B. (ol.//e0_oilll0101.000 (oillse((. (losm(())1911
E Extrpo
exp
(63)
2
(64)
The total bias is then
20

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
12p B3.
sup |0
E
[p (6)
Margat (b)
(65)
K
Finally, when clipping weights with > 0 we induce truncated distribution Pr,o, i.e. to :- P(rozo > 9). Using Lemma 1 this creates another constant factor that contributes TV Pr.00) = B1 to the overall bias:
12p K
E [2fo(d) Mtarger |0
sup
Marger(d))] (d)]
+ As + B4.
(66)
D. Datasets
For all datasets besides alanine dipeptide we use a training set of 10 contiguous samples (1 μn simulation time) from a single MCMC chain, a validation set of the next 10a contiguous samples (0.2 us simulation time), and a test set of 10 uniformly strided subsamples from the remaining trajectory. Since these are highly multimodal energy functions, this leaves us with biased training data relative to the Boltzmann distribution; we split trajectories this way to test the model in challenging and realistic setting. We describe the datasets below and present the simulation parameters in Table 4. Ramachandran plots for the training and test data (before subsampling) are provided in Figures 9 to 12.
Table 4. Overview molecular dynamics simulation parameters.
Peptide
Force field
Temperature
Time step
Alanine dipeptide
Amber f199SBildn
300K
1fs
Trialanine
Amber 14
310K
1fs
Alanine tetrapeptide
Amber f199SBildn
300K
1fs
Hexaalanine
Amber 14
310K
1fs
Chignolin
Amber 14
310K
1fs
Alanine dipeptide. For this dataset we use the data and data split from Klein & Noč (2024). Here the training set is purposely biased with an overrepresentation an underrepresented mode, i.e. the positive & state. This bias makes easier reweight the target Boltzmann distribution. Alanine Dipeptide consist of one Alanine amino acids, acetyl group, and an N-methyl group.
Trialanine and hexa-alanine. For the peptides composed of multiple alanine amino acids, we generate MD trajectories using the OpenMM library (Eastman et al., 2017). All simulations are conducted in implicit solvent, with the simulation parameters detailed in Table 4. These systems do not include any additional capping groups, such as those present in alanine dipeptide and alanine tetrapeptide, as they are generated in the same manner as described Klein et al. (2023a). There are two peptide bonds in trialanine and five in hexa-alanine, resulting in two and five Ramachandran plots respectively.
Alanine Tetrapeptide (AD4). For this dataset as flexible. The original dataset kept all hydrogen bonds fixed, as the Boltzmann Generator was operating in internal coordinates. The MD simulation to generate the dataset is then performed as described above. Alanine Tetrapeptide consist of three Alanine amino acids, an acctyl group, and an N-methyl group. Therefore, there are four Ramachandran plots.
the same system setup in Dibak et al. (2022), but treat all bonds
Chignolin. In addition to the small peptide systems, we also investigate the small protein chignolin, consisting of ten amino acids (GYDPETGTWG). We simulate this system using the same configuration as trialanine, defined Table for 4.5 ps.
21

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
33
-4
41
4
411
(a) Train
(b) Test
Figure Alanine dipeptide Ramachandran plots and test data subsets. The lower probability state oversampled the training
data as in Klein & Noé (2024).
4
4
$
1
-7
-
198
(a) Train
(b) Test
Figure 10. Trialanine Ramachandran plots for train and test data. We observe missing mode in the first Ramachandran plot of the training subset.
M4
V1
+4
-4
2
4
(a) Train
01
IP2
(b) Test
Figure Alanine tetrapeptide Ramachandran plots for train and data subsets, We obverse an underreprcsented mode in the second Ramachandran plot of the training subset.
22

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
W1 14
Ma
14
v/
V4
4P3
4
-7
4
(a) Train
84
4P1
422
4P3
4
44
4
4P5
4
(b) Test
Figure Hexa-alanine Ramachandran plots for train and data subsets. We observe good mode coverage of the training data.
E. Experimental Details
E.1. Metrics
For computational efficiency we subsample 10a reference samples from evaluation trajectory to serve as ground truth. Similarly in cases where a method produces more than 10 samples, random subset of size 10 is selected for comparison, We quantify distributional similarity using empirical Wasserstein-2 distances between generated samples and the reference data. Given two empirical distributions, μ Σi= òz, and 1 Σド 6g the Wasserstein-2 distance is computed as
W2(1, min xell(w))
Taj c( Uz),,
(67)
where v) denotes the set of admissible transport plans and c(2, y)² is defined cost function. Optimal couplings are computed using the POT library (Flamary et al., 2021).
Energy cost. To assess energy distribution similarity, we define the following cost
==(2,y)= E(y)1
(68)
Dihedral torus cost. We evaluate macrostructural similarity in the space of backbone dihedral angles (d), v)), which encode conformational information. For molecule with L residues, define the angle vector:
Dihedrals(z) = (01, eh .i...... &L- 121-))
(69)
Due to angle periodicity, we define the cost on the resulting torus as:
2L
cz y)²
Σ
(Dihedrals(z)) Dihedrals(y)) (y); + T) mod 2 - m]²
(70)
capturing angular deviations while respecting cireular geometry.
23

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
TICA projection cost Time-lagged Independent Component Analysis (TICA) performs dimensionality reduction for identification of slow dynamical modes. From the mean-centered time series 3t, compute:
1 Coo-
Σ
S Cur T-
Σ
(71)
T T-T
1=1
and solve the generalized eigenvalue problem:
Cor w = WC
(72)
The top two eigenvectors wi and w2 define projections capturing the most autocorrelated directions. Using these, we define the TICA cost:
CTICA v)² = Σ w
wj y
(73)
Note that the TICA basis is computed from the full evaluation trajectory (un-subsampled),, while the comparison subset is restricted the 10 subset. TICA analysis is performed only heavy atom coordinates.
E.2. Timings
Sampling time calculations. For the sampling inference times in Figure 7, compute all times single NVIDIA L40S GPU, using the maximum power of two batch size possible.
Training time. For training times we compute all times on single A100 80GB GPU except for SE(3) -EA3F which is trained on a single H100. We report the total time in hours until convergence for all methods the table below.
Table 5. Training time (hours) for all methods.
Model
ALDP
AL3
AL4
AL6
Chignolin
SE(3)-EACF
160
ECNF+
9.72
12.5
17.17
76.94
SBG
16.83
24.67
41.67
57.5
427.33
E.3. SE(3)-EACF Implementation Details
Equivariant augmented coupling flow (EACF) (Midgley et al., 2023a). We adopt the original model configuration from (Midgley et al., 2023a) for our EACF baseline on ALDP. Specifically, we choose the more stable spherical-projection EACF with 20-layer configuration. Each layer has two ShiftCoM layers and two cor-ttransformation blocks. The EGNN used in the core transformation block consists of weessage-passing layers with 128 hidden states. Stability enhancement tricks like stable MLP and dynamic weight clipping on each layer's output are fully applied. The model is trained for 50 epochs with a bathh size of 20 using Adam optimizer and peak learning rate of 1 10 We use the default 20 samples for likelihood estimation.
EACF Boltzmann generator. EACF leverages augmented dimensions, and therefore to estimate the likelihood of sample æ under an EACF model, we need to use an estimate based on samples from the augmented dimension a. Specifically, for a Gaussian distributed augmented variable a, we can estimate the marginal density of an observation as
(x( a)
(xx) (11))
(74)
[[(a)))
however, this is only a consistent estimator of the likelihood and for finite sample sizes has variance. This makes EACF unsuitable for our application of large-scale Boltzmann generators, as in this setting we need to compute exact likelihoods. Variance in likelihood estimation would lead to bias in the final distribution under normanzod importance sampling or a SBG strategy. We therefore do consider EACF as a viable option for large scale Boltzmann distribution sampling.
24

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
E.4. ECNF Implementation Details
E.4.1. NETWORK AND TRAINING
Algorithm 2 ECNF flow matching training
Input: Prior Qo. Empirical samples from data 91, bandwidth o, batchsize b, initial network uo.
while Training do
2g ~ go (20); 21 ~ qı (x1) (Sample batches of size i.i.d. from the dataset)
(0, 1)
14 tan (1 1)
~ NN,,, 1²
(()) - ||v (1, x)) (x) -xo)ll
0 < Update(0, Vo((
end while
Return U/
Equivariant continuous normalizing flow (ECNF) (Klein & Noé, 2024). We use the supplied pretrained model from Klein & Noé (2024) for our ECNF baseline on alanine dipeptide. Therefore all training parameters are equivalent to, and specified in, that work. We use the specification for the model "TBG + Full" in that work.
ECNE++. We note four improvements to the ECNF, which together substantially improve scalability.
1. Flow matching loss. In Klein & Noé (2024) a flow matching algorithm with smoothing is employed which provides extra stability during training. This is depicted Alg. 2, however this smooths out the optimal target distribution (Tong et al., 2024, Proposition 3.3). ECNF uses 0.01 where we use for ECNF++. We find that causes poor molecular structures to be generated the bond lengths are not able to be controlled precisely enough. We note that 0 is used in most recent large scale flow matching models (Liu al., 2024; Esser et al., 2024).
in this case
2. Architecture size. Empirically, we find the ECNF to be underparameterized. We perform grid search over layer width and depth, finding width of 256 and depth of 5 block be a good balance between performance and speed on alanine dipeptide. We employ the same parameters for larger molecular systems.
3. Improved optimizer and LR scheduler. We find using an AdamW (Loshchilov Hutter, 2017) with moderate weight decay of 10 improves performance and stability. Prior work has found weight decay helps to keep the Lipschitz constant of the flow low and avoids stiff dynamics which enables accurate ODE solving during inference. We also use smoothly varying cosine schedule with warm-up (over 5% of iteration budget) which enables larger maximum learning rate and faster training than the two step schedule used previously. Both the start and end learning rates are 500 times lower than the defined maximum.
4. Exponential moving average. We use exponential moving average (EMA) on the weights with decay 0.999. This is standard practice in flow models, which improves performance.
These four elements together greatly improve the ECNF training, enabling larger systems to be successfully modeled, and provide a strong foundation for future Boltzmann generator training on molecular systems using equivariant continuous normalizing flows. Qualitatively, we find ECNFs quite stable to train and robust to training parameters relative to invertible architectures. However, it is very slow to compute the exact likelihoods necessary for importance sampling.
Other parameters. For inference we use a Dormand-Prin Prinnce 45 (dopri5) adaptive step size solver with absolute tolerance 10-4 and relative tolerance 40
Likelihood evaluation. Evaluating the likelihood of CNF model requires calculating the integral of the divergence, as in Eq. 2. While there exist fast unbiased approximations of the likelihood using Hutchinson's 's trace estimator (Hutchinson, 1990; Grathwohl et al., 2019), these are unfortunately unsuitable for Boltzmann generator applications where variance in the likelihood estimator leads to biased weights under nofmaaozml importance sampling. We therefore calculate the Jacobian using automatic differentiation which is both memory and time intensive. For example, on hexa-alanine, the maximum batch size that can fit on an 80GB A A100 GPU is 8. This batch takes around 2 minutes for 84 integration steps. We also use
25

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
an improved vectorized Jacobian trace implementation for all CNF which reduces memory by roughly half and time by roughly 3x over the previous implementation (Klein & Noé, 2024).
On using CNF proposal with SBG. In principle it is possible to drop in replace our NF architecture with a CNF in SBG. However, there are several drawbacks to such an approach, most notably in efficiency. As previously discussed, CNFs are extremely computationally inefficient to sample a likelihood from. We find on the order of 100 SBG steps are necessary for best performance. This would make CNFs at least two orders of magnitude slower to sample from, when they are already at the edge of tractability for the current importance sampling estimates. We leave future work to consider faster CNF likelihoods and note that our SBG algorithm could be applied there readily.
E.5. SBG Implementation Details
Architecture. We scale the TarFlow architecture for increasingly challenging datasets. As advised by Zhai et al. (2024) we scale the layers per block alongside the number of blocks. The layers / blocks / channels and resulting number of parameters presented in Table 6. We note the larger number of parameters in the TarFlow relative the ECNF+ despite the faster inference walltime, due to the lack of simulation and higher computational efficiency of the architecture.
Table 6. TarFlow configurations across different datasets.
Dataset
Layers per Block
Number Blocks
Channels
Number Parameters (M)
ALDP
4
4
256
13
AL3
6
6
256
29
AL.4
384
64
AL6
6
384
64
Chignolin
8
8
384
114
Training configuration. The training hyperparameters used closely follow those of Zhai et al. (2024), although we deviate in using a larger value of weight decay as instability was observed during training. Namely we use a learning rate of 1 x 10-4, weight decay of 4 x 10-4, Adam $1, of (0.90, 0.95). We additionally employ the same cosine decay learning rate schedule with warmup (start and end learning rate 500 times lower than maximum value) and exponential moving average decay (0.999) used in ECNF++. Training is performed for 1000 epochs. Center of mass augmentation is applied with a standard deviation of + for the number of particles, match that of the prior for given system. As non-monotonic yn
improvement was observed on validation metrics, we use early stopping on the SNIS E-W2 against the validation dataset.
Sampling hyperparameters. Whilst the TarFlow is capable of generating low-energy peptide states, also prone to generating samples of extremely high target energy. For standard importance sampling this presents issue as these samples will be assigned negligible importance weight. However, for the SBG these high energy samples were prone numerical instability during Langevin dynamics. To mitigate this issue we truncate the proposal distribution prior po based on an energy cutoff, noting that the bias introduced by this operation is bounded in Proposition 3. Similarly to EACF, we additionally remove the samples corresponding to the largest 0.2% of importance weights, in the case of SBG this is performed once, prior to Langevin dynamics. See $F.3 for an ablation on the effect of weight clipping in SBG. For alanine systems we use 100 Langevin time steps, with ESSehrssoldd = 0.5 and 1 x 10 up trialanine, and e
10-6 thereafter. For chignolin we use 500 time steps, with ESSthreshold 0.5 and 1 10-5 For ablations of these hyperparameters see $F.3.
Table 7. Overview of training configurations
Training Parameter
ECNF
ECNF++
TarFlow
Learning Rate
5 x 10-4
5 x 10-4
1 10 10-4
Weight Decay
0.0
1 10-2
4 10-4
B1, .2
0.9, 0.999
0.9, 0.999
0.9, 0.95
EMA Decay
0.000
0.999
0.999
Width
64
256
Varies
N blocks
5
Varies
Parameters
152 K
2.317
Varies
26

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
F. Additional Results
F.I Chirality
The ECNF architecture E(3) equivariant, hence cquivariant to reflections, and will generate samples of both possible global chiralities. As the energy functions are themselves invariant reflections this is not resolved by importance sampling. Having the correct global chirality is necessary match the test dataset dihedral angle distributions where only one global chirality is present in the data. The incorrect chirality can show up symmetric mode Ramachandran plots. To resolve this issue we follow Klein & Noć (2024) in detecting incorrect chirality conformations, and reflecting them. However, unlike Klein & Noé (2024) points with unresolvable symmetry (e.g mixed chirality conformations) are not dropped. The results for T-W2 before and after fixed-chirality samples are presented in Table 8. We observe a reduction in metric value (improved performance) on all configurations, which we attribute to evaluation noise. We further note that non-equivariant methods such as SBG do not suffer this effect and hence do not require any symmetry post-processing.
Table 8. T-Wa results for unprocessed and fixed-chirality samples from ECNF and ECNF+
Datasets +
Tripeptide (AL3)
Tetrapeptide (AL4)
Hexapeptide (AL6)
Algorithm
Unprocessed
Fixed
Unprocessed
Fixed
Unprocessed
Fixed
ECNF++ (Ours)
1.967 + 0.062
1.177 + 0.145
2.414 + 0.000
2.082 + 0.005
5.405 0.069
4.315 + 0.018
F2. Ramachandran Plots
In this appendix we include the Ramachandran plots for each model on each peptide system. Please note that the ground truth training and test Ramachandran plots are presented in SD.
Alanine dipeptide. In Figure 13 we present the alanine dipetide Ramachandran plots. Both ECNF and SBG cover all relevant modes. We find that that ECNF++ models the distribution well, but drops the positive yp mode. This is notable as this mode is oversampled in the training data (see Figure 9).
Trialanine. In Figure 14 we can see the Ramachandran plot for resampled points for the trialanine dataset. Comparing this to the train and test data in Figure 10 we see that the AL3 training data is missing the y1 positive mode which reflected in all of the models.
Alanine tetrapeptide. In Figure 15 we present Ramachandran plots for resampled points on the alanine tetrapeptide dataset. Comparing this to the train and test distributions in Figure 11 we observe that both ECNF++ and SBG capture the dihedral angle distribution with comparable success.
alaanne. In Figure 16 we present the Ramchandran plots for samples from ECNF++ and SBG. We find that SBG succeeds to capture the low density positive modes, albeit with a tight concentration of points as opposed to a broad range of low density. We additionally observe the negative V mode to be well captured by the SBG.
-5
-4-
.4
(F]
4
4
(a) EACF
(b) ECNF
(c) ECNF+
(d) SBG
Figure 13, Alanine dipeptide Ramachandran plots for baseline methods (SNIS) and SBG (SMC). 10 points sampled per method.
27

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
-4
-4
-4
IP1
4P2
(a) ECNE+
(b) SBG
Figure 14. Trialanine Ramachandran plots for ECNE+ (SNIS) and SBG (SMC). 10 points sampled per method.
422
ECNF++
422
IF3
(b) SBG
Figure 15. Alanine tetrapeptide Ramachandran plots for ECNF++ (SNIS) and SBG (SMC). 10 points sampled method.
4-
-4
130
148
ECNF++
4
-4
44
95
(b) SBG
Figure 16. Hcxa-alanine Ramachandran plots for ECNF+ (SNIS) and SBG (SMC). 10 points sampled per method.
28

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
E.3 Ablation Studies
Ablation of SNIS and SMC. In Tables and 10 we compare the quantitative performance of ECNF+ and SBG as proposals, with SNIS, and in the case of SBG, with SMC. This table also presents the TICA-W2 results. We almost uniformly observe a significant reduction in E-W2 using SNIS or SMC over the raw proposals, with only ECNF+ SNIS alanine tetrapeptide failing to achieve this. On alanine dipeptide we also observe large decrease in macrostructure metrics T-W2 and TICA-Wa, which is expected as the training data and subsequently the proposal distribution was intentionally biased to provide improved coverage (Klein & Noé, 2024). On all other datasets (and both models) we generally see no change or an increase in these metrics after reweighting with either SNIS or SMC, suggesting there is some tradeoff between matching the energy distribution whilst maintaining good macrostructure metrics. We lastly note that the use of SMC over SNIS does not uniformly improve performance on E-W2 for SBG, with only alanine dipeptide and trialanine exhibiting this trend. However, we believe this may be an artifact of the good proposal overlap, and draw attention to the significant reduction in E-W2 achieved on chignolin by SMC over SNIS in Figure 8.
Table Comparison of proposal, SNIS, and SMC for SBG and ECNF++ on alanine dipeptide and trialanine.
Datasets- -
Alanine dipeptide
Trialanine
Algorithm 4
E-Wa
T-W2
TICA-W2
E-W2
T-W2
TICA-W2
ECNF++ Proposal
6.675 + 0.297
1,776 + 0.018
3,920 + 0.025
5.424 + 1.595
0.277 + 0.004
0.435 + 0.009
ECNF++ SNIS
0.914 + 0.122
0.189 + 0.019
0.402 + 0.002
2.206 + 0.813
0,962 + 0.253
0.597 + 0.023
SBG Proposal
2 10 0,873 + 0.338
1.695 = 0.015
3.862 + 0.038
2 10
0.338 + 0.036
0.449 + 0.028
SBG SNIS
0.439 + 0.129
0.942 + 0.268
0.758 + 0.506
0.502 + 0.016
0.518 + 0.032
SBG AIS
0.960 + 0.617
0.430 + 0.034
0.806 + 0.166
0.754 - 0.230
0.495 + 0.033
0.476 + 0.048
SBG SMC
0.741 + 0.189
0.431 + 0.141
0.915 0.316
0.598 + 0.084
0.503 + 0.029
0.501 + 0.031
Table 10. Comparison of proposal, SNTS, and SMC for SBG and ECNF+ on alanine tetrapeptide and hexa-alanine.
Datasets
Alanine tetrapeptide
Hexa-alanine
Algorithm 4
E-W2
T-W;
TICA-W2
E-W2
T-W2
TICA-W2
ENCF++ Proposal
2.983 1: 1.266
0.576 1 0.002
0.737 + 0.013
> 10
1.136 1 0,030
0.688 1 0.066
ECNF++ SNIS
5.638 1 0.483
1.002 1 0.061
0.832 1 0.021
10.668 1 0.285
1.902 0.055
0.632 1 0.087
SBG Proposal
> 10
0.624 + 0.023
0.791 + 0.050
> 10
1.079 1 0.153
0.299 1: 0.039
SBG SNIS
1.068 0.495
0.969 + 0.067
0,879 1 0.047
1.036 = 0.534
1.473 0.114
0.452 + 0.245
SBG AIS
1.070 + 0.272
0,923 : 0.100
0.920 + 0.028
1.131 = 0.384
1.510 + 0.113
0.492 + 0.240
SBG SMC
1.007 + 0.382
1.039 = 0.069
0.904 + 0.054
1.155 = 0.635
1.517 + 0.118
0.530 + 0.198
Center of mass adjusted energy. As discussed in $3.1, the SBG proposal distribution is not mean-free due to the CoM data augmentation applied to the training data, with a centroid norm distribution ||C]| ~ 0x3- This can introduce adverse behavior, as the target energy is invariant to ||C|], and thus the importance weights will depend on ||C|| of sample Concretely, low (target) energy sample generated far from the origin (with large ||C|] will have low likelihood under po but high likelihood under leading to very large importance weight.
To provide visual intuition for this effect, we plot in Figure 17 the centroid norm distribution of the SBG proposal samples before and after SNIS. Here the empirical distribution is generated using 2 x 10" samples, to approximate the asymptotic behavior. In Figure 17a we observe that, even with this extremely large number of samples, without weight clipping or center of mass adjusted energy Equation (5) the ||C*| distribution is greatly influenced by the reweighting. In this case resampling shifts most density to high ||C'] samples, where there was very little density prior to reweighting and hence little sample diversity, and a large peak manifests resulting from single sample with very large importance weight. In contrast, in Figure 17b we see much smaller change in |[C|| distribution after reweighting, with no large peak for any given sample, after applying the CoM adjustment. In this case there is no overweighting of high |[C|| regions with limited sample diversity. Adding weight clipping to the standard proposal energy function (Figure 17c) greatly reduces the change in distribution from reweighting, although the mean remains notably shifted towards higher |[C*] samples. Applying weight clipping with the CoM adjusted proposal energy function (Figure 17d) has little effect on the ||C]| distribution beyond smoothing. We
29

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
15
2.5
Proposal
SNIS
2.0
Density
Density
1.5
1.0
5
0.5
0
0.0
0.00
0.25
0.50
0.75
1.00
1.25
0.00
0.25
0.50
0.75
1.00
1.25
||C||
||C||
(a) Weight clipping X, CoM energy x
(b) Weight clipping x, CoM energy
2.5
2.5
2.0
2.0
Density
Density
1.5
1.5
1.0
1.0
0.5
0.5
0.0
0.0
0.00
0.25
0.50
0.75
1.00
1.25
0.00
0.25
0.50
0.75
1.00
1.25
||C||
||CI|
(c) Weight clipping V, CoM energy X
(d) Weight clipping V, CoM energy /
Figure 17. Centroid norm ||C|] histograms for 2 x 10 proposal samples and reweighted proposal samples, with / without both of weight clipping (0.2%) and center mass adjusted energy.
emphasize that these plots are presented for 107 samples hence clipping may have larger still effect for both proposal energy functions for smaller sample sets.
In Figures 18 and 19 ablate the utility of performing the center of mass energy adjustment to the proposal energy. Specifically, ablate the center of mass adjustment function of number of samples used during inference and also as function of number of inference time steps. Each of these ablations is performed on trialanine. Considering Figure 18, there is little distinction between variants on E-W2 with respect to N samples, although standard energy without clipping can be identified as the least performant, and all methods improve with increased N. On T-W2 the standard energy without clipping again evidently the worst performing, with clear benefit to applying the CoM adjustment where clipping not used. The best performing variants do employ clipping, with slight but clear benefit to using the center of mass adjustment. Considering Figure 19, observe again that the center of mass adjustment improves performance greatly where clipping not employed, and notably still without clipping on T-W2
30

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
2.5
Standard (clip)
Adjusted (clip)
Standard (no clip)
Adjusted (no clip)
8-12
T-W2
4
1.5
1.0
0.5
10a
10a
10
10
10a
10a
10
10
N Samples
N Samples
Figure 18. Trialanine SBG SMC &-W2 and T-W2 performance with standard and center of mass adjusted energy, with / without weight clipping (0.2%) a variety sampling set sizes. 100 Langevin timesteps used.
1.50
1.25
E-W2
T-W2
0.75
10
10a
10a
10!
10a
10a
N Timesteps
N Timesteps
Figure Trialanine SBG SMC &-W2 and T-W; performance standard and center of mass adjusted energy, / without weight clipping (0.2%) at varierty of Langevin time discretizations. 10 samples generated.
31

---

Scalable Equilibrium Sampling with Sequential Boltzmann Generators
Ablation EACF importance weight clipping. We lastly report additional EACF results on alanine dipeptide. In the same 0.2% clipping threshold the importance weights other models for fair comparison. Nevertheless, in the resampling process, observe significant degradation in sample diversity, as evidenced by the energy histogram in Figure 3 and Ramachandran plots Figure 9. In Figure 20 we plot energy histograms and Ramachandran plots for a variety of different clipping thresholds. We observe that EACF generates highly unreliable importance weights, particularly visible in the energy histograms where there are extreme spikes and poor alignment with the true data distribution. This leads to poor resampling quality, as demonstrated fail to capture the true data distribution. While increasing the clipping threshold 10% shows some improvement, the fundamental issue of inaccurate importance weight estimation by EACF persists across different clipping ratios.
main results
corresponding Ramachandran plots where the resampled points
0.14
0.12
0.06
0.06
Progusal
0.10
SNIS
Density
Density
Density
SMC
0.08
0.04
0.04
0.06
0.04
0.02
0.02
0.02
0.00
0.00
0.00
-50.0
0.0
50.0
-50.0
0.0
50.0
-50.0
0.0
50.0
2
((x)
8(x)
Un
-4
f1
(a) Clip 0.2%
(b) Clip 2%
(c) Clip 10%
Figure 20. Energy histogram and Ramachandran Plots of SE(3) EAAF under different weight clipping ratio (0.2%, 2%, 10%).
32


Reviews:
[{"review_id": "7l9EhGyMbP", "reviewer": "Reviewer_VNPc", "summary": "This paper introduces Sequential Boltzmann Generators (SBG), an extension to the Boltzmann generator framework. By replacing conventional importance sampling with a non-equilibrium annealing process , the authors aim to transport proposal samples toward the target Boltzmann distribution. The authors also propose normalizing flow which leverages a Transformer‐based exactly invertible TarFlow that is trained with soft equivariance penalties rather than using equivariant architectures. Experimentally, the paper demonstrates increased performance over BG baselines on a series of molecular systems (ranging from dipeptides up to decapeptides such as Chignolin) in terms of effective sample size (ESS), Wasserstein distances, and inference time scaling.\n\n## Update after rebuttal\n\nMy doubts were initially about the scalability of SMC / re-weighting approaches in general being able to go beyond classical force-fields and lack of comparison to alternative diffusion samplers like iDEM (which essentially amortizes force-field evaluation in training time, not required for sampling). However, I want to acknowledge that showing efficacy of invertible NF with exact likelihoods enables some interesting downstream applications as the authors have accomplished here (to my knowledge, not currently possible with iDEM). I think this work highlights some interesting avenues in sampling and is a high-quality paper, so I want to increase my support 3->4.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Rebuttal Reviewer VNPc\n\nWe thank the reviewer for their time, feedback, and nuanced comments. We are glad that the reviewer found the non-equivariant NF an “interesting alternative” which is “necessary as we scale up to even larger systems and datasets”. We also appreciate that the reviewer recognized our use of SMC instead of IS as a “smart drop-in design choice”. Finally, we are heartened to hear that the reviewer agrees that Prop 3 “is a principled technique”. We now address their key questions raised in the review and note that additional results are included in this link: https://anonymous.4open.science/api/repo/sbg/zip\n\n## iDEM as a Baseline\n\nWe appreciate the reviewers' valuable comments regarding iDEM as an additional baseline. We would like to first politely recall the setting of iDEM and other diffusion samplers as completely data-free (i.e. there is no training set), in contrast to the BG setting which includes training on (biased) data. This allows BGs to scale much more easily than the data-free amortized samplers. Consequently, we argue that scalability claims in each setting cannot be meaningfully compared. To our knowledge only one sampling method has successfully scaled to any molecular task, which is FAB on ALDP using an intrinsic coordinate system, while molecular tasks are the focus of most works on BGs. To our knowledge, no sampling method has successfully scaled to any molecular task on cartesian coordinates, the focus of this work.\n\nTo investigate this setting, we had private correspondence with the iDEM authors, who told us that iDEM failed to scale to ALDP even with substantial effort; to their knowledge no diffusion-based sampler can scale to molecular tasks. Following the reviewer's suggestion we trained iDEM on ALDP ourselves. In this case we use a vacuum instead of the implicit solvent used in our main results to speed up training (see Fig 1 in our link). We observe that iDEM is unable to successfully sample in this easier setting with most modes missing and a poor energy distribution.\n\n## Novelty\n\nWe value the reviewer's feedback that the application of SMC, and a general-purpose Normalizing Flow may initially appear to have limited technical novelty. We would like to politely push back against this assertion, as our framework and design choices fundamentally challenge the predominant approach in BG. Prior to SBG, all modern BGs in Cartesian coordinates have resorted to equivariant CNF’s; the fact that we can omit both exact equivariance and use an exact NF—is a novel insight. Moreover, our CoM adjustment strategy is also new. We ablate its importance in new plots in the rebuttal link (Fig 2 and 3), which shows that IS reweighting significantly benefits from this adaptation strategy.\n\nIn fact Reviewer S3Vr notes, “The contribution made by this paper should significantly change the course of future work in this area.” We further wish to highlight that the paper contains numerous theoretical results on thresholding, some of which are utilized in existing papers without proper justification. Our theoretical results allow us to quantify the impact of thresholding schemes in this process. We hope that the reviewer may join us in agreeing that the design choices utilized in SBG allow for a fresh approach to building BG’s using non-equivariant components with soft penalties that demonstrably scale better to larger peptides in Cartesian coordinates—which remained an open challenge until SBG.\n\n## Additional References\n\nWe acknowledge the reviewer's comment regarding the inclusion of more non-BG based samplers. We will update the paper with a dedicated discussion on non-BG based sampling and include the references suggested by the reviewer in FAB, PIS, PDDS, and the concurrent work SCLD.\n\n## LJ Potential\n\nThe reviewer is correct to highlight that many previous samplers and BG papers also evaluate on LJ potential systems. Such systems are however less challenging than even small peptides (e.g ALDP), making them of limited interest once peptides can be successfully tackled. We thank the reviewer for their suggestion that this may help better place the work, but instead draw attention to the evident failure of iDEM on ALDP in Fig 1 as clear evidence that BG methods are superior in the data-available setting.\n\n## Computational expense\n\nWe thank the reviewer for enquiring regarding the efficiency of our method with respect to energy / force evaluations. During sampling a single force evaluation is required per-particle per timestep, hence if the force is expensive this will present a computational cost.\n\n## Closing comments\n\nWe thank the reviewer for their valuable feedback and great questions. We hope that our rebuttal fully addresses all the important points raised, and we kindly ask the reviewer to potentially upgrade their score, as they indicated, if they are satisfied with our responses. We are also more than happy to answer any further questions that arise, please do let us know!", "claims_and_evidence": "SBG claims to improve scalability of equilibrium sampling in two ways: 1. by replacing the importance sampling reweighting of proposal samples with a target-informed non-equilibrium process, 2. Increase computational efficiency of the proposal model by using a non-equivariant invertible transformer architecture, enforcing equivariance softly via data augmentation. \n\nBased on the peptide sampling experiments, the authors have shown evidence of scaling up compared to prior boltzmann generator techniques and may very well be the first to sample these larger peptides (at least in cartesian coordinates). However, there are other generative modeling methods which aim to perform scalable amortized boltzmann sampling without requiring invertible networks or additional importance weighting / SMC. One example, which is only ever mentioned in Table 1, is iDEM which has shown to scale well to high-dimensional configurations (55 particle Lennard Jones potential). If the goal is scaling up equilibrium sampling, I find it confusing why the iDEM method is cited in the table but never discussed or benchmarked. Flow-annealed Bootstrapped importance sampling (FAB) and path-integral sampler (PIS) may also be competitive here, which can fully make use of equivariant architectures.\n\nIf these concerns are clarified I would definitely reconsider my score.", "methods_and_evaluation_criteria": "I think the proposed methods are described clearly and many design choices are directly related to the goal sampling molecular configurations. As mentioned above, I find only comparing against Boltzmann generator techniques strange.", "theoretical_claims": "The theoretical claims appear correct.  Prop 2 is based on a recently established result. The energy thresholding technique suggested by Prop. 3 is interesting.", "experimental_designs_or_analyses": "(mostly reiterated from claims-evidence section) Based on the peptide sampling experiments, the authors have shown thorough evidence of scaling up compared to prior boltzmann generator techniques and may very well be the first to sample these larger peptides (at least in cartesian coordinates). However, there are other generative modeling methods which aim to perform scalable amortized boltzmann sampling without requiring invertible networks or additional importance weighting / SMC. If the goal is scaling up equilibrium sampling, I find it confusing why the iDEM method is cited in the table but never discussed or benchmarked. Flow-annealed Bootstrapped importance sampling (FAB) and path-integral sampler (PIS) may also be competitive here, which can fully make use of equivariant architectures.", "relation_to_broader_scientific_literature": "I expect to see more application of these non-equilibrium processes as a drop in replacement of self-normalized importance sampling. I think the paper could have addressed other non-BG frameworks for sampling that have been proposed recently.", "essential_references_not_discussed": "I think there definitely could be more discussion on recent generative modeling techniques for sampling outside of Boltzmann generators. In particular, there are diffusion-based samplers which all have different tradeoffs, but can possibly scale just as well here:\n\nCited but not discussed:\nFlow Annealed Importance Sampling Bootstrap (2023),\niDEM (Akhound 2024), \nTransport meets variational inference ( Vargas 2024)\n\nNot cited:\nPath-integral Sampler (Zhang 2021),\nParticle Denoising Diffusion Sampler (Phillips 2024), \nSequential Controlled Langevin Diffusions (Chen 2024) (concurrent work)", "other_strengths_and_weaknesses": "Strengths:\n\nThe paper offers some interesting alternatives for normalizing flow architectures that do not rely on explicit equivariant parameterizations. The inductive biases / symmetries for molecular sampling problems are still quite strong, but the non-equivariant architectures may be necessary as we scale up to even larger systems and datasets.\n\nThe use of NETs in place of self-normalized importance weighting is a generally smart drop-in design choice for Boltzmann samplers.\n\nProp 3 gives us a principled technique for energy thresholding.\n\nWeaknesses:\n\nFor the purpose of scaling up equilibrium sampling, I believe the paper is missing some essential baselines (in particular iDEM) and discussion of related methods listed above. It seems the paper only shows evidence of scaling over previous Boltzmann generator frameworks, but not other amortized samplers.\n\nI was a little bit disappointed in the novelty since the application of NETs and the architecture choice of the normalizing flows are somewhat orthogonal to each other. It seems they do work well together experimentally, but I believe more baselines are needed to prove its necessity for scalable equilibrium sampling.\n\nAs I said before, If these concerns are clarified I would definitely reconsider my score.", "other_comments_or_suggestions": "If you have any other comments or suggestions (e.g., a list of typos), please write them here.\n\nThe acronym EACF isnt defined until the appendix should be defined before the appendix. \n\nMany previous works evaluate on these synthetic energy functions based on the Lennard Jones potential (LJ). They are not as interesting as peptides, but they might help place the work better wrt previous work.", "questions_for_authors": "If you have any important questions for the authors, please carefully formulate them here. Please reserve your questions for cases where the response would likely change your evaluation of the paper, clarify a point in the paper that you found confusing, or address a critical limitation you identified. Please number your questions so authors can easily refer to them in the response, and explain how possible responses would change your evaluation of the paper.\n\nAlthough there is amortization here from the normalized flow, it seems that sampling still requires a lot of queries from the potential energy model. How would this technique work in situations where the potential energy is computationally expensive to evaluate?", "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": "ksdKRFKIs1", "reviewer": "Reviewer_9RzA", "summary": "The manuscript presents the Sequential Boltzmann Generators (SBG), a novel extension to the existing Boltzmann generator framework for scalable sampling of molecular states in thermodynamic equilibrium. The framework removes the SE(3)-equivariance and encodes equivariance softly via data augmentations, achieving enhanced computation efficiency. The method leverages inference-time non-equilibrium transport via Sequential Monte Carlo (SMC), progressively refining proposal samples and improving their alignment with the target Boltzmann distribution. SBG employs a Transformer-based normalizing flow for efficient likelihood computation, avoiding the costly integration required by continuous normalizing flows. The experimental results demonstrate the state-of-the-art performance of SBG, successfully scaling equilibrium sampling to larger molecular systems that were previously intractable for standard Boltzmann generators.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Rebuttal Reviewer 9RzA\n\nWe thank the reviewer for their thoughtful comments and feedback. We value that the reviewer found that most of our claims were supported by “clear and convincing evidence” and that our mathematical claims were “theoretically sound and mathematically consistent”. We are also glad that the reviewer found our use of SMC in BGs to be “novel” and allows for “scalability,” which is supported by our “insightful” comparison on computational efficiency. We next address all the salient points raised in the review and note that additional results are included in this link: https://anonymous.4open.science/api/repo/sbg/zip\n\n## (Auto) Correlation of samples\n\nWe acknowledge the reviewer's request for the inclusion of autocorrelation analysis. However, we emphasize that the samples generated by SBG can not suffer from autocorrelation. This is because only samples from $\\tau = 1$ are returned as generated samples, hence there is no time dimension on which autocorrelation could exist. This is in contrast to methods such as MD in which a single trajectory is returned as generated samples, and samples may exhibit autocorrelation over sampling time.\n\nTo bolster our study, we have included a new variant of SBG (SBG-AIS) that does not perform any resampling during the annealing process (See Tables 1 and 2). This ablation principally serves to eliminate the correlation that SMC might introduce during inference scaling. As we observe SBG-AIS, it achieves acceptable ESS and outperforms other baselines in sample-based metrics. We hope that this makes the impact of correlation during inference clearer in SBG.\n\n## Benchmarking MCMC\n\nWe thank the reviewer for the suggestion of additional baselines. We however, are unsure exactly which MCMC methods they believe to be suitable for benchmarking on the evaluation systems. We note that the goal of BGs is to amortize sampling and train a model which is able to quickly draw many uncorrelated samples. This is a fundamentally different approach with both benefits and drawbacks as compared to MCMC methods.\n\n## Evaluating SBG beyond peptides\n\nWe value the reviewer's feedback as exploring SBG’s generalizability to experimental settings beyond peptides considered in this paper is interesting. We first note that before SBG, modern BG’s utilized equivariant CNFs instead of classical normalizing flows that employed invertible architectures. This was due to the widely held belief that equivariant CNFs were more expressive and more amenable to larger-scale problems of interest. Despite this promise, these equivariant CNFs—as demonstrated in our ablations—struggle on datasets beyond AL3/4 and are much slower to evaluate due to the need for simulation. Thus, we argue that even demonstrating that a classical normalizing flow paired with an inference scaling strategy allows tackling even larger peptides in hexapeptide (AL6) is extremely interesting. These results demonstrate the expressive power of our framework in comparison to previous BGs.\n\nSuch results enabled Reviewer S3Vr to highlight in their review, “The contribution made by this paper should significantly change the course of future work in this area.” As a result, we believe the experimental validation of our introduced SBG approach is well supported in the current peptide datasets and testing on other chemically diverse systems, while extremely interesting, is beyond the scope of the current paper, but remains an exciting direction of future work. At this time we don’t anticipate any limitations specific to chemically diverse systems such as small organic molecules or metal-organic frameworks, but chose to evaluate on peptides as has been the norm in many prior BG works.\n\n## Closing comments\n\nWe thank the reviewer for their time and effort in reviewing our work, and we hope the reviewer will kindly consider a fresh evaluation of our work, given the main clarifying points outlined above. We are also eager to engage in further discussion if the reviewer has any lingering doubts.", "claims_and_evidence": "Most of the claims are supported by clear and convincing evidence. However, the following claim might be problematic to some extent:\n\n1. The paper claims that the model can generate uncorrelated samples efficiently, given the ESS results. However, while a high ESS and improved sampling efficiency are impressive, they do not directly confirm the independence of samples as an autocorrelation analysis.", "methods_and_evaluation_criteria": "1. The proposed method aligns well with the problem of scalable equilibrium sampling for molecular systems. Using a Transformer-based normalizing flow coupled with non-equilibrium transport via SMC is a reasonable approach to improve sample quality and importance weighting. \n\n1. The evaluation metrics, namely Effective Sample Size (ESS), Wasserstein distances for energy distributions and dihedral angles, and Ramachandran plots, are reasonable for the task in the paper and provide a comprehensive assessment of sampling quality. However, to demonstrate that the model can generate uncorrelated samples, autocorrelation plots might be another indicative evaluation method.\n\n1. The datasets are about different peptide systems (di-, tri-, tetra-, and hexapeptides), which are typical for molecular sampling work. However, additional benchmarks on more chemically diverse molecules could further validate the generalizability of SBG.", "theoretical_claims": "I checked propositions 1-3 and the sampling algorithms, and they are theoretically sound and mathematically consistent.", "experimental_designs_or_analyses": "1. The paper benchmarks SBG against existing SOTA methods, including SE(3)-equivariant coupling flows and equivariant continuous normalizing flows. Therefore, the comparison between SBG and baselines gives indicating evidence of SBG's performance. \n\n1. The comparison of computational efficiency is insightful, as it highlights the scalability of SBG compared to baseline models.", "relation_to_broader_scientific_literature": "The key contributions of the paper contribute to the research in this area in terms of:\n\n1. It builds upon the Boltmann generator (Noe'2019) and introduces non-equilibrium transport via Sequential Monte Carlo, improving sample quality and efficiency significantly.\n\n2. It removes the explicit constraint of equivariance but implicitly learns this via data augmentation, which is also supported by some recent work in this area like AlphaFold3 (Abramson'2024) and MCF (Wang'2024).", "essential_references_not_discussed": "All the related and essential references are discussed.", "other_strengths_and_weaknesses": "Strengths:\n\nThe paper novelly integrates Sequential Monte Carlo with Boltzmann Generators, enabling non-equilibrium transport for improved sample quality and scalability. Consequently, the normalizing flow operates directly on all-atom Cartesian coordinates, which was previously intractable for existing methods. In addition, the claims in the paper are well supported by the mathematical derivations and proofs.\n\nWeaknesses:\n\nWhile the paper claims that the model generates independent samples more efficiently, it does not analyze sample autocorrelation and compare the results to traditional MCMC methods.", "other_comments_or_suggestions": "See Strengths and Weaknesses.", "questions_for_authors": "The evaluation in the paper is primarily focused on peptide systems, whereas it is unclear how well SBG generalizes to chemically diverse systems like small organic molecules or metal-organic frameworks. Therefore, have you considered any test on different chemical systems, and if not, do you anticipate any limitations when applying SBG to those systems?", "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": "IPANTdCNOT", "reviewer": "Reviewer_HyU2", "summary": "This paper improves data-driven learning-based Boltzmann Generators (BG) with Sequential Monte Carlo (SMC), based specifically on a non-equilibrium transport method (NETS) recently proposed by (Albergo & Vanden-Eijnden (2024)). Unlike NETS whose source energy is based on a pre-defined prior, here the source energy is learned from data with a normalizing flow. Empirical results demonstrate the effectiveness of the proposed method.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Rebuttal Reviewer HyU2\n\nWe thank the reviewer for their time and effort. We are glad that the reviewer found our empirical results to “demonstrate the effectiveness” of our method SBG. We next clarify the main points raised in the review and note that additional results are included in this link: https://anonymous.4open.science/api/repo/sbg/zip\n\n## Novelty\n\nWe acknowledge the reviewer's concern that SBG can be thought of as an application of NETS-style inference on a learned proposal. We first highlight that our approach differs from NETS as we do not learn an extra drift term through an auxiliary loss e.g. A PINN objective. In NETS, such a term is crucial to reduce the variance of the importance weights during a linear interpolation. However, SBG does not need to employ this computationally intense learning objective precisely because of our learned proposal. In contrast to an uninformative proposal, which has a very small overlap with the target, a learned proposal mitigates the need to learn a drift term. Thus, we would like to politely push back against the assertion that SBG is a simple instantiation of NETS, as it in fact SMC with a computationally efficient manner to perform Langevin dynamics that is unlocked by using a normalizing flow rather than an equivariant CNF which we argue is a novel insight that we exploit in the BG context. \n \nWith regards to our framework, we again would like to politely disagree with the reviewer as our design choices fundamentally challenge the direction of BG research. More precisely, our main technical novelty lies in demonstrating the scalability of non-equivariant classical flows in contrast to the dominant trend to leverage equivariant CNFs. In addition, we include new algorithmic novelties such as CoM adjusted resampling; rebuttal Fig 2 and 3 highlight the improved stability and performance of IS / SMC when accounting for CoM augmentation. Furthermore, the overall scalability of our method to hexapeptides is a novel result, and is due to a combination of each component in SBG: 1) a non-equivariant classical flow and 2) SMC that leverages the exact energy of a classical flow. We highlight that these findings enabled Reviewer S3Vr to remark in their review, “The contribution made by this paper should significantly change the course of future work in this area.” \n\nFinally, we highlight that our paper contains several theoretical results that provide quantification of bias added through various thresholding schemes. Such schemes have routinely been employed in existing literature without justification or analysis. Our paper is the first to provide an exact characterization of the bias added—allowing practitioners to negotiate a select a problem-dependent thresholding value.\n\n## Training details\n\nWe appreciate the reviewer's comment regarding the training details for training the proposal flow. Whilst many of the training, inference and dataset details are included in Appendices D and E, we recognize the importance of including details in the main paper. We will update this in future versions of the work to include additional details in section 4 but briefly state that for all TarFlow models we train for 1,000 epochs with lr 1e-4 weight decay 4e-4 using the AdamW optimizer with a cosine lr schedule and EMA with decay 0.999. Furthermore, directly answering the reviewer's question concerning training samples, Appendix E includes a description of dataset construction from MD trajectories for each of the peptide systems. For each system we use a training set of 100k contiguous samples from a single MD trajectory. We understand that aspects of such details are important to be included directly in the experiments section, and we will revise the paper in the next draft to include the key details in the main body. \n\n## Background on CNF\n\nWe value the reviewer's feedback that the discussion of CNFs might seem ancillary to a paper that leverages classical normalizing flows as the model. The key reason to introduce CNFs is to illustrate that our inference time scaling through SMC, while theoretically possible using a CNF proposal, faces significant challenges in scalability due to the need to simulate the ODE using equation 4 to compute $\\log p_{\\theta}$ for **every energy evaluation** along the interpolation used in Langevin dynamics. In contrast, our TarFlow requires only 1 call for each Langevin step. However, we understand that this section may not have been as tightly integrated into the paper, and we will improve the clarity of presentation in the updated draft. \n\n## Closing comments\n\nWe thank the reviewer again for their questions, which allowed us an opportunity to clarify our paper. We hope that our rebuttal responses fully address all the important questions raised by the reviewer, and we kindly ask the reviewer to potentially upgrade their score if the reviewer is satisfied with our responses. We are also more than happy to answer any further questions that arise.", "claims_and_evidence": "Y", "methods_and_evaluation_criteria": "To the best of my understanding, the proposed method represents a specific instantiation of NETS, wherein a data-driven source energy is employed. Given that NETS is theoretically applicable to any choice of source energy, the proposed approach seems incremental by incorporating a pre-trained normalizing flow (NF) as the source energy.", "theoretical_claims": "Y", "experimental_designs_or_analyses": "- What's the experiment setup for learning p_\\theta? How many samples are required in $\\mathcal{D}$? Since the proposed method is a two-stage approach. Training detail should be clarified in Sec 4.", "relation_to_broader_scientific_literature": "Efficient sampling method from Boltzmann distribution could benefit AI4Science area for applications such as drug and material discovery.", "essential_references_not_discussed": "Y", "other_strengths_and_weaknesses": "See above sections", "other_comments_or_suggestions": "N/A", "questions_for_authors": "- It seems redundant to introduce CNF in the main paper, given that NF is used in practice to model log p_\\theta. What's the reason for mentioning CNF in Sec 2?", "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": "4E4cSuZKs7", "reviewer": "Reviewer_S3Vr", "summary": "This paper proposes Sequential Boltzmann Generators, consisting of two conceptual ingredients: first, that invertible normalizing flows operating on Cartesian coordinates can scale to molecules as large as hexapeptides by leveraging non-equivariant transformers and the recent TarFlow framework; second, that annealing from the flow likelihood to the target with AIS and SMC can dramatically improve sampling efficiently. The authors show across-the-board improvements over CNFs on ALDP and scale these comparisons up to chignolin.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Rebuttal Reviewer S3Vr\n\nWe would like to thank the reviewer for their time, feedback, and positive appraisal of our work. \nWe are heartened to hear that the reviewer feels that the “contribution made by this paper should significantly change the course of future work in this area.” We also thank the reviewer for acknowledging that the main claims of our work are “well supported by experiments”. We now address the questions and suggestions raised by the reviewer in order, and note that additional results are included in this link: https://anonymous.4open.science/api/repo/sbg/zip\n\n## ESS after SMC and SBG without resampling\n\nThe reviewer makes an astute observation that using SMC with adaptive resampling, ESS can be maintained to an artificially high value. We agree that such a metric in this case is not meaningful, and we had opted to include this to be in line with the broader literature, e.g. in pure sampling NETS includes ESS after SMC. We will update the paper to remove ESS after SMC and, as the reviewer suggested, include SBG without resampling (SBG-AIS) which is now included as another row in our rebuttal document Tables 1 and 2. We find that SBG with AIS outperforms the previous SOTA which is ECNF and our introduced ECNF++ for AL3-AL6 across sample-based metrics $\\mathcal{T}-\\mathcal{W}_2$ and $\\mathcal{E}-\\mathcal{W}_1$, and is slightly worse in terms of ESS to ECNF++.\n\nWe thank the reviewer for allowing us to strengthen and clarify our empirical results, and we hope this new SBG-AIS variant alleviates this particular concern.\n\n## Wasserstein distances of proposal w/o reweighting \n\nWe thank the reviewer for this insightful idea. We have included Table 3 in the rebuttal link to quantify the performance of our proposal without reweighting, with importance sampling i.e., just BG, and with SBG. We observe a drastic improvement in $\\mathcal{E}-\\mathcal{W}_1$, for IS over the proposal, and an even greater improvement with SMC. We will include these results in our updated draft.\n\n## Wasserstein in TICA space\n\nThank you for the suggestion! We have now included in the rebuttal experiments (Tables 1, 2 and 3) TICA Wasserstein metrics for all AL3+. We find SBG outperforms all baselines (ECNF/ENCF++) in this metric. We will include these results in our updated draft.\n\n## High energy structures\n\nThat is a very interesting question. Upon further investigation we found that the highest energy AL6 samples result from steric clashes, as visualized in Fig 4. In Fig 5 we show the histogram and log histogram of shortest non-bonded interatomic distance between MD (ground truth) and SBG samples. An additional source of high energy samples is covalent bonds of insufficient length (Fig 6). We will include these results in our new appendix.\n\n## Mixing energy-based training\n\nWe thank the reviewer for this great suggestion. We have experimented with mixing energy-based training with normal data-based training and find preliminary evidence that this leads to better performance on $T-\\mathcal{W}_2$ and the newly introduced TICA metric but marginally worse $\\mathcal{E}-\\mathcal{W}_1$ vs normal SBG (Tables 1-2). Given these promising initial results we will run a full set of experiments using this method and report these in the updated paper.\n\n## Transferable BG\n\nWe appreciate the reviewers' comments that operating in Cartesian coordinates more easily leads to training transferable models. In this work, our primary focus was on proving the scalability of SBG as it relied on normalizing flows and inference scaling through SMC. Consequently, we believe that extending this to transferable systems is a natural direction of future work, but out of the scope of the current paper.\n\n## Other questions\n\n> The definition of T-W2 distance is not clear.\n\nThis is a multi-dimensional Wasserstein distance on torsion angles, accounting for the torus geometry of the angle space.\n\n> 100k vs 10k samples\n\nThe subsampling is done after the final SMC step at the end of inference.\n\n> Proposition 1\n\nTo give further empirical credibility to CoM adjustment, we perform additional ablations in rebuttal Fig 2 and 3. We first find that the $||C||$ of the proposal samples does indeed follow an approximate Chi distribution, as expected given the training data augmentations. We also find that CoM adjustment is both important for stable IS reweighting (with a large but finite number of samples) and that sample metrics improve when the adjustment is employed in SMC.\n\nWe will update the proposition statement and proof to better convey our theoretical result.\n\n## Closing comments\n\nWe would like to thank the reviewer for their time and effort. We hope all our answers here allow the reviewer to continue to positively endorse our paper, and we would love to have the opportunity to clarify any lingering questions should the reviewer have them.", "claims_and_evidence": "The main claim of the work is that the proposed collection of strategies enables scaling of Boltzmann generators to systems of unprecedented size without the use of coordinate transformations. This claim is well supported by experiments — this is the first demonstration of out-of-the-box Boltzmann generators on hexapeptides or molecules of similar complexity. However, there are significant caveats to some details, discussed below.", "methods_and_evaluation_criteria": "Yes, the choice of model system and evaluation criteria is reasonable and effective at demonstrating the paper’s key claims.", "theoretical_claims": "I did not carefully check the proofs of theoretical claims.", "experimental_designs_or_analyses": "There are deficiencies with some of the experimental designs or analyses\n\n*   ***[Severe]*** For the proposed method, the authors report ESS after SMC resampling, which is a completely meaningless metric as by definition the ESS can be maintained to be arbitrarily high via resampling. ***The authors must address this point (ideally via the next suggestion); otherwise I will change the recommendation to Reject.***\n\n*   [Moderate] The authors should separately report the performance of SBG with and without resampling (i.e., only with AIS) to allow a more direct comparison of I.I.D. proposal quality and ESS.\n\n*   [Moderate] The authors should clarify (or even better, report both) the Wasserstein distances of the proposal w/o reweighting vs the samples with reweighting (which should presumably approach 0 without finite sample effects).\n\n*   [Minor] It would be great if the authors reported more fine-grained Wasserstein metrics, for example in TICA space for the larger molecules.\n\n*   [Minor] It would be great if ESS was also reported for Chignolin.\n\n*   [Minor] The proposal appears to still contain very high-energy structures. It could be interesting to analyze the types of errors exhibited in these structures.", "relation_to_broader_scientific_literature": "The paper contributes to the literature on training Boltzmann generators with access to data, where the main technical challenge is in the parameterization of the learned distribution in a way that permits exact likelihoods. The choice adopted by this paper is quite novel compared to previous works, which have generally used continuous flows or flows over internal coordinates. In particular, this work opens up the long-sought possibility of scalable, transferable, exact likelihood flows. The contribution made by this paper should significantly change the course of future work in this area.", "essential_references_not_discussed": "All essential works were discussed.", "other_strengths_and_weaknesses": "Since the model permits fast likelihood evaluation, there is a missed opportunity to explore a mix of data-based and energy-based training as done in the original Boltzmann generator paper.\n\nAlso, with the freedom from internal coordinates, there is also a missed opportunity to explore the training of transferable models.", "other_comments_or_suggestions": "The definition of T-W2 distance is not clear. Could the authors clarify? Is it a multidimensional W2 distance or an average of one-dimensional W2 distances?\n\nThe distinction between 100k vs 10k samples is not clear — could the authors clarify when this downsampling happens?\n\nI would also note that the discussion in Appendix A.1 and Appendix C seems unnecessary and might be confusing. Most readers will understand that AIS cannot be accomplished without fast likelihoods, and it would be strange to propose an additional control so as to exactly cancel out the annealing terms. The Ito filtering paper is also very new and this paper should not feel obliged to guard against misunderstandings of their key result.\n\nAlso, the writing of Proposition 1 seems to be obfuscating the fact that the additional term is just a Gaussian on the CoM and the partition function a power of $2\\pi$. Actually, could the authors clarify where the $\\log ||c^2||/\\sigma^3$ comes from in the log density of the $\\chi(3)$ distribution?", "questions_for_authors": "I have no important questions for the 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:

Meta-review: This paper introduces Sequential Boltzmann Generators, merging ideas from sequential Monte Carlo techniques with Boltzmann Generators. The approach aims to align an inference procedure of a surrogate of the Boltzmann distribution with the actual target density using a SMC scheme, instead of only doing importance sampling after inference. The approach is conceptually interesting and practical, providing state of the art results on challenging physical systems. All doubts were resolved during discussion period and the reviewers unanimously recommended acceptance.
Decision: Accept (poster)