Amortized Sampling with Transferable Normalizing Flows
Authors: Charlie B. Tan
Majdi Hassan
Leon Klein
Saifuddin Syed
Dominique Beaini
Michael M. Bronstein
Alexander Tong
Kirill Neklyudov
Venue: NeurIPS 2025 poster ()
Keywords: sampling
proteins
peptides
normalizing flows
Primary area: machine_learning_for_sciences
TL;DR: we train transferable normalizing flows to sample from peptide Boltzmann distributions up to 8 residues

Abstract: Efficient equilibrium sampling of molecular conformations remains a core challenge in computational chemistry and statistical inference. Classical approaches such as molecular dynamics or Markov chain Monte Carlo inherently lack amortization; the computational cost of sampling must be paid in full for each system of interest. The widespread success of generative models has inspired interest towards overcoming this limitation through learning sampling algorithms. Despite performing competitively with conventional methods when trained on a single system, learned samplers have so far demonstrated limited ability to transfer across systems. We demonstrate that deep learning enables the design of scalable and transferable samplers by introducing Prose, a 285 million parameter all-atom transferable normalizing flow trained on a corpus of peptide molecular dynamics trajectories up to 8 residues in length. Prose draws zero-shot uncorrelated proposal samples for arbitrary peptide systems, achieving the previously intractable transferability across sequence length, whilst retaining the efficient likelihood evaluation of normalizing flows. Through extensive empirical evaluation we demonstrate the efficacy of Prose as a proposal for a variety of sampling algorithms, finding a simple importance sampling-based finetuning procedure to achieve competitive performance to established methods such as sequential Monte Carlo. We open-source the Prose codebase, model weights, and training dataset, to further stimulate research into amortized sampling methods and finetuning objectives.

Amortized Sampling with Transferable Normalizing Flows
Charlie B. Tan"
Majdi Hassan"
Leon Klein
Saifuddin Syed
Dominique Beaini 2,3,5
Michael M. Bronstein 1.6
Alexander Tong*
2,3,6
Kirill Neklyudov 2,3,7
University of Oxford
2 Université de Montréal
Mila Quebec Al Institute
FFreie Universität Berlin
Valence Labs
AITTYRRA
Institut Courtois
Abstract
Efficient equilibrium sampling of molecular conformations remains core chal- lenge in computational chemistry and statistical inference. Classical approaches such as molecular dynamics or Markov chain Monte Carlo inherently lack amortization; the computational cost of sampling must be paid in full for each system of interest. The widespread success of generative models has inspired interest towards overcoming this limitation through learning sampling algorithms. Despite performing competitively with conventional methods when trained on a single system, learned samplers have so far demonstrated limited ability to transfer across systems. We demonstrate that deep learning enables the design of scal- able and transferable samplers by introducing PROSE, 285 million parameter all-atom transferable normalizing flow trained dynamics trajectories up to residues in length. PROSE draws zero-shot uncor- related proposal samples for arbitrary peptide systems, achieving the previously intractable transferability across sequence length, whilst retaining the efficient likelihood evaluation of normalizing flows. Through extensive empirical evaluation we demonstrate the efficacy of PROSE as a proposal for a variety of sampling algorithms, finding simple importance sampling-bass
corpus of peptide molecular
fine-tuning procedure achieve competitive performance to established methods such as sequential Monte Carlo. We open-source the PROSE codebase, model weights, and training dataset, to further stimulate research into amortized sampling methods and objectives.
Introduction
Accurately sampling molecular configurations from the Boltzmann distribution is fundamental
problem in statistical physics with profound implications for understanding biological and chemical systems. Key applications include protein folding [Noé et al.. 2009, Madorfff-Larser et al., 2011]. protein-ligand binding [Buch et al., 201 1], and crystal structure prediction [Köhler et al., 2023]; processes that underpin advances in drug discovery and material science.
"Equal contribution. 'Equal advising. Correspondence to: charlie. tanSexeter.. ox. ac. uk and majdi. hassan@mila. quebec
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

MD
2.5
1.0
Prose
2.0
TICA-W>
0.8
1.5
0.6
1.0
0.4
0.5
10
10
10
10"
10
10
10
10
10
# Energy Evaluations
# Energy Evaluations
# Energy Evaluations
2.5
1.0
2.0
TICA-wz
F-12
0.8
1.5
0.6
1.0
0.4
0.5
10
10-2
10"
10
10 -2
10"
10
10-2
10"
GPU Hours
GPU Hours
GPU Hours
Figure 1: PROSE exceeds the quantitative performance of molecular dynamics on unseen peptide systems. Wasserstein-2 distances energy, dihedral torus, and TICA projection with respect to reference molecular dynamics (5 js), for a μus) molecular dynamies baseline and PROSE (with SNIS), at range of energy evaluation (ahove) and GPU walltime budgets (helow). Each value represents the mean over unseen tetrapeptide systems. PROSE outperforms the baseline with respect to energy evaluations on all metrics. Whilst comparable on E-W2 for a given time budget, the baseline is significantly inferior the T-W2 and TICA-W; macrostructure metrics, highlighting long simulation periods required traverse the separated metastable states.
Conventional approaches such as Markov chain Monte Carlo (MCMC) [Liu, 2001] and, in partic- ular, Molecular Dynamics (MD) [Leimkuhler and Matthews, 2015] seek tackle this problem by proposing a general solution, which, however, has practical limitations due to its Markov nature. To accurately integrate the corresponding Hamiltonian dynamics, MD has to be simulated with fine time-discretization (on the order of femtoseconds). which produces highly correlated samples and prevents efficient exploration of the modes of the Boltzmann density. Although running multiple chains from different initializations is possible, every chain has to be simulated for a long time to ensure proper mixing, which cannot be efficiently parallelized. Finally, the entire simulation has to be started from scratch for a new system, which bottlenecks the speed of ab initio studies.
Deep learning- baased samplers abandon the Markov chain approach drawing samples and shift the computational burden to a one-time training phase, enabling fast and inexpensive inference compared to MCMC. In the most challenging scenario, these methods consider having access only to the unnormalized density function (analogous to MC methods) [Vargas et al., 2023, Akhound-Sadege al., 2024]. Boltzmann generators (BGs) [Noé et al., 2019] consider more practical scenario where, in addition to the unnormalized density, a dataset of MD trajectory is available, which does necessarily match the target density. To mitigate the error introduced by imperfections of the model and training data, BGs train likelihood-basee models and perform self ceff-normalized importance sampling (SNIS) [Liu, 2001] at inference time. The availability of training data coupled with SNIS have enabled BGs to generalize across dipeptide systems [Klein and Noe, 2024], but they have not yet been able generalize across larger and more diverse systems of scientific interest.
In this work, we introduce PROSE, a large-scale normalizing flow which demonstrates unprecedented ability transfer previously unseen systems of varying amino acid composition, sequence length, and temperatures, outperforming MD for the same computational budget (see Fig. 1). Our approach strikingly simple and scalable, which elucidates the potential of the deep learning-based samplers for sampling applications. In particular, we outline the following series of contributions:
We introduce ManyPeptidesMMD: novel dataset of molecular dynamics trajectories for peptide systems between 2 and 8 residues. The training dataset consists of 21,700 peptide sequences simulated for 200 ns each, giving total of 4.3 ms of simulation.
Building on the recently proposed TarFlow [Zhai et al., 2024], we propose architectural modifications, which allow for better modeling peptide systems, system-transserable conditioning, and generation of peptide sequences of varying length.

---

We study the use of PROSE as a proposal distribution for different Monte Carlo algorithms, finding the learned proposal be sufficiently powerful for accurate sampling with standard SNIS, which does not require tuning of sampling parameters. Furthermore, resampled generations can used for efficient fine-tuning of PROSE on previously unseen systems.
Finally, we empirically demonstrate that PROSE achieves satte-of-thearte performance when sampling from the equilibrium distribution on previously unseen peptide systems of length up to 8 residues surpassing the continuous normalizing flow-based transferable Boltzmann generator [Klein and Noe, 2024] whilst generating proposals 4 103 times faster.
We open source our codebase https: /ttgtgggtbbb/t//// com com/t:a e samplers/ aransferablee-ampll.
ManyPeppiddesMD dataset https haaps://huggnngeee datasets teassetssttrasseeebble wwwpee/www...s..aarraaaaaeeaeeee / //huggingface. hoo/transferaait samplers/model wwwpsr/s//moowwwwwttt.......
co/ md and model weights https:
Background
2.1 Normalizing flows
The fundamental challenge of probabilistic modeling is designing a density model from which samples can be efficiently generated. Normalizing flows [Rezende and Mohamed, 2015] approach this challenge by defining a diffeomorphism: a differentiable invertible function with a differentiable inverse. Namely, given simple prior density (=(=) and parameterized flow (diffeomorphism) fI()), one can define the push-forward distribution as the map of samples from the prior distribution ~ qu (z) via the inverse flow I = fo (z) ~ go (z) with learnable parameters 0. The density of the push-forward distribution can then be computed via the change-of-variables formula
go(x) = -f (z(0(( - (ffa))) (fo(z))
(z)
(I)
where (z)()))O is the Jacobian determinant the map fo- In practice, one has to be able to efficiently evaluate fo (z) for sample generation, and (z))())00r)) for likelihood evaluation.
Autoregressive normalizing flows [Kingma et al., 2016, Papamakarios et al., 2017, Zhai et al., 2024] define family of invertible maps with tractable Jacobian as sequence of composed transformations x is defined autoregressively. In the case of the TarFlow [Zhai al., 2024], this is autoregressive affine update defined over blocks of latent variable corresponding to image patches 2, e R. That is, the i-th latent block is
fa fr o... o...o So. where each transformation z = =f((z))
i= 00 ( (2 (zrl: i)) exp(-as ())) i [1,N 1),
[ 2z [i)
2t+1 [
(2)
where we adopt slicing notation denoting the i-th block as z|i] and blocks up to the i-th (exclusive) as 2|: i]. Notably, the autoregressive structure allows for efficient evaluation of the Jacobian de- terminant due to its nower-triangulat structure log [0f(zz)/0zz = N EnD0 ar(2]]: i)), Furthermore, the autoregressive affine updates are invertible with inverse 2, f. (2441) given by
[2+1/1] 22411 [i] c exp(ar(z1l: 4) + (z () i [1, 1].
i= 0,
2/ =
However, clearly, such transformations leave the leading dimension 2, untouched, hence must be interleaved with permutations n, over the dimensions fo = म fr o o... o To o fo- For example, Zhai et al. [2024] use simple inversions of the latent block sequence for all n, across the entire model.
2.2 Boltzmann generators
Despite normalizing flows allowing various forms of training supervision, such as variational inference for unnormalized densities [Rezende and Mohamed, 2015], or maximum likelihood from empirical distributions [Kingma and Dhariwal, 2018], errors present the parameterized distribution prevents accurate evaluation within scientific applications requiring high precision, e.g. free energy estimation.
3

---

Boltzmann generators address specifically this challenge by performing self- neff-normalized importance sampling (SNIS) at inference time. Namely, evaluate the expectation of a statistic ((x) w.r.t. the target Boltzmann density p(ar) one can use the following consistent Monte Carlo estimator
[p())((1)) ~
WC
((:)), W
wi
p(x+) go
32 ~ go(r),
(4)
where qo (z) is the density of the learned normalizing flow. The SNIS estimator converges, for n - 00, to the true value (())) Furthermore, Tan et al. [2025] extended the Boltzmann generator framework to more general Monte Carlo algorithms, in particular a continuous-time formulation of Sequential Monte Carlo [Jarzynski, 1997, Albergo and Vanden-Fijnden, 2025].
Transferable Boltzmann generator (TBG) [Klein and Noe, 2024] made first attempt of learning sampler that generalizes across target densities corresponding to different peptide systems. TBG parametrizes the proposal distribution as a continuous normalizing flow (CNF) [Chen et al., 2018b], where the vector field is defined by an equivariant graph neural network [Satorras al., 2021, Klein et al., 2023c]. Crucial to the method is the system dependent conditioning of N atoms
h = [A,, R, P.),
(5)
where atom type A,, residue type R, and residue position P. are each encoded as one-hot vectors. Training on set of MD trajectories for dipeptide systems [Klein et al., 2023b] with this system- conditional encoding enables TBG to generate a proposal for previously unseen dipeptides.
However, despite successful generalization across dipeptides, the TBG architecture introduces significant bottlenecks for inference and fine-tuning. Indeed, the learned CNF requires accurate integration of the vector field and computationally expensive evaluation of its divergence for evaluating the learned density model. For instance, the implementation of [Klein and Noe, 2024] requires GPU-days to produce 3 10 samples with their corresponding proposal likelihoods for a single dipeptide system. Furthermore, the expensive evaluation of density makes it infeasible to train or finetune TBG via the reverse KL-divergence create replay buffer of substantial size.
3 Scalable transferable normalizing flows as Boltzmann generators
3.1 Architecture of PROSE
PROSE builds on the TarFlow architecture [Zhai al., 2024], which parametrizes a sequence of autoregressive affine transformations via blocks of transformer layers. The expressivity and favorable scalability of the transformer layers enables TarFlow effectively model high dimensional data, whilst the affine autoregressive flow parameterization ensure fast and accurate energy evaluation. With minimal modifications TarFlow capable of successfully modeling high-dimensional molecular data (Tan et al., 2025]. Here we describe our design choices that make transferability possible.
Transferability across system dimensions. We extend TarFlow to support concurrent training on sequences of variable length. Whilst transformers natively support sequences of arbitrary length, special consideration is required within normalizing flow such as TarFlow that is defined for fixed input and output dimensions. We therefore define appropriate masking to the affine sequence updates and log- determinant aggregation prevent padding tokens influencing either computation, under arbitrary sequence permutations: further details of the tokenization and masking are provided in Appendix We additionally replace the fixed-length learnable position embedding with the more extrapolation-rreenlly sinusoidal embedding. This design enables PROSE to efficiently train across distribution systems by maximizing the normalized likeliodl
max E. d(s)
E Ba-n(a ()) log (0(x:)
(6)
where d(s) is the size of the system [Klein and Noe, 2024]. This extended architecture allows for parallel processing of data dimensions, enabling transferability and scalability across lengths.
Adaptive system conditioning. The standard TarFlow employs simple additive conditioning for
class-conditional image generation. Whilst we find this to be sufficient to define a teasfellr normalizing flow, we follow large-scale atomistic transformer architectures in applying conditioning

---

Ivenntatioo
Rp-1
Saaul 8 (x/6 9 Ju 4J
mep -
14
Autoregresive Flow Block
xT
Transformee
B2
R2
A
R
Barklaan: Ponsatation
Coordinates
Cond. Features
Figure 2: All-atom block-wise autoregressive normalizing flow based the TarFlow [Zhai et al., 2024]. Peptides are encoded via atom types A, residue types R, residue position P. and sequence length L. Atom positions in 3D Cartesian coordinates define the system state. The embedding of the peptide is applied as conditioning to the coordinates such that PROSE achieves transferability between systems. Within each block the sequence 2u is permuted and passed transformer. defining an autoregressive affine update. In the backbone permutation the backbone [N,, Ca.s, C., 0.1/20 of residues (with associated hydrogens) updated before any sidechains, providing additional diversity to causal attention for global structure modeling.
through adaptive layer normalization, adaptive scaling, and SwiGLU transition blocks [Abramson et al., 2024, Geffner et al., 2024]. The system conditioning features are constructed from atom types A, residue types R, residue positions P. and sequence lengths L. Atom and residue types are embedded using lookup-table embedding layers, whilst sinusoidal embeddings are employed for the naturally ordered sequence position and sequence length. See Appendix A for further details.
Chemistry-aware sequence permutations. In the image setting. Zhai et al. [2024] employ only an identity and flip permutation the sequence of image tokens. Similarly, when applying TarFlow to peptide systems Tan et al. [2025] employ only an identity and flip permutation on the ordering defined per-residue starting with backbone atoms followed by sidechain atoms. Whilst simple identity and flip may be appropriate for the regular grid of image data, we argue this to be suboptimal for the diversity of geometric interactions present in molecular systems. This motivates our introduction of chemistry-an
sequence permutations, defined to promote effective peptide modeling. We define the backbone permutation. such that the backbone atoms [N; Ca.i. C, O,1 (with associated hydrogens) for all residues are located the start of the sequence, and followed by the sidechains. By processing the coordinates of the backbone atoms at the start of the sequence, the model refines the global structure of the peptide as contiguous sequence. Crucially, the sidechain positions are subsequently able to causally attend to the full backbone structure, hence enabling local updates to be influenced by global structure. We further employ a backbone-flip permutation to provide additional diversity to the autoregressive modeling.
3.2 Inference and fine-tuning of PROSE
The practical applicability of learned sampling methods depends significantly on their inference-time throughput, as well as their transferability to unseen systems. this section, we describe how one can employ PROSE for inference-time importance sampling, importance sampling- based fine-tuning, as well annealing of the learned proposal to different target temperatures.
Importance sampling, At the inference time, one can use PROSE to estimate the expectation of statisties sx(x) w.r.t. the target Boltzmann density p(x) via self-normalized importance sampling (SNIS) estimator. Namely, we consider standard SNIS, discrete-time sequential Monte Carlo (SMC) [Neal, 2001, Doucet al., 2001], and continuous-time SMC [Jarzynski, 1997, Albergo and Vanden- Eijnden, 2025]. All these estimators are of the form
E (xx))((x)) ~
W:
p(x)), w; p(x;)) g(2,) Ei g(x),
Σ
プラ= Wj
5

---

where the only difference between them is the proposal density q(z). Note that these estimators can be interpreted as the expectation over the empirical distribution, i.c.
((x;) (x)))((), p(x) j=1 W
Wi
8( x;)
(8)
In practice, we compare the true density p(x) with our generated distribution p(z) instead of measur-
ing statistics pzx). For completeness, we describe all the considered estimators in Appendix D.
demonstrate the ability fine-tune PROSE using a self-improvement strategy. Namely, we iteratively generate the empirical distribution p(x: {8) by resampling the samples from the model qo (x: proportionally p(a s) and these samples for fine-tuning. Note that this is different from classical fine-tuning as true samples from the target are available. We update the parameters by maximizing the likelihood on the resampled proposal, i.e.
Self-improvement. For unseen system
p(z; s) 20 (1g (8)
max Ep(x (a) log (2 [s), p(x1s)
8(x xz) detach Wj
(9)
This is akin to the energy-based training of Jing et al. [2022], in which samples are proposed by ODE-integration of a diffusion model, resampled, and then used in the score-matching objective.
Temperature transfer. Temperature is fundamental to molecular simulation, with significant influ- ence on conformational dynamics and statistic expectations. learned sampler may transfer across temperature without retraining. Formally, we aim to change the temperature T = 1/3 of the learned density model when generating samples, i.e.
therefore highly desirable that
(x) 0z
log qo (f (2)) = (1)gg=(2) , log
(10)
Note that, for measure-presserving flows log (2)= o. one simply has to change the temper- ature of the prior distribution (i.e. sample 2 ~ 92 (2) instead of ~ (z)) to change the temperature of the density model, which is a standard technique in the normalizing flow literature [Kingma and Dhariwal, 2018, Dibak et al., 2022]. Whilst PROSE a non-volume preserving flow [Dinh et al., 2017], hence violating this assumption, found that simply scaling the prior temperature log qu (2) results in a suitable proposal for the Boltzmann density with the corresponding temperature.
4 Experiments
To establish the performance of PROSE, we first introduce a new dataset of peptide molecular dynamics. We employ this dataset to train PROSE and prior methods, and evaluate using metrics computed against reference molecular dynamics trajectories. We additionally evaluate PROSE as a proposal for variety of sampling algorithms, and in the cemeeraturettranseer setting.
4.1 Molecular dynamics trajectory dataset
We introduce ManyPepttidesMDD a novel dataset of peptide MD trajectories for sequences ranging from 2 to 8 residues in length' Following Klein et al. [2023b] all simulation is performed using OpeNMM [Fastman et al., 2017] with the Amber14 forcefield [Case et al., 2014). For training, a total of 21,700 uniformly sampled sequences are simulated for 200 ns. For evaluation, 30 sequences of length 2, 4, and 8 are randomly sampled such that all amino acids are represented equally, and simulated for
5 us. Further details on dataset collection and MD configuration provided in Appendix B.
Table 1: Number of sequences used per peptide length for training and evaluation.
Sequence length
5
6
7
8
Training Evaluation
200
1,000
1,500
2,000
3,000
4,000
10,000
30
30
30
Available attppsr/wngggigo..ooooottttt/tts http://///eeeetooeeee//e//e/r//a

---

4.2 Experimental configuration
We train the first Boltzmann generators transferable across peptide sequence length. We train the PROSE architecture defined in Section 3.1, unmodified TarFlow [Zhai et al., 2024] as in SBG [Tan et al., 2025], and the equivariant CNF of Klein and Noe [2024], with the improved training recipe of Tan al. [2025]. denoted as ECNF++. All models are trained for x 10 iterations with batch size 512. Both PROSE and TarFlow are suitably scalable to long sequences and are trained on the full dataset detailed in Section 4.1. However, generating 8 residue sequences with likelihoods for ECNF++ was found to be prohibitively expensive, hence the training data was limited to sequences up and including length 4. Comprehensive training details provided in Appendix C.
The primary evaluation metrics are the Wasserstein-2 distance on: (i) the energy distribution E-W2, (ii) the dihedral angle torus distribution T-W2, (iii) the first 2 TICA component projections T1CA-Wa The energy distribution is highly sensitive to perturbation in bond length and angle, hence E-Wa measures accuracy fine-grained details. The dihedral angle tori and TICA projection describe macrostructure, hence T-W2 and TICA-Wa measure accuracy in terms of metastable state coverage. We additionally report effective sample size (ESS); the variance of the importance weights. For metric definitions and further details on sampling evaluation procedure please refer to Appendix E.
4.3 Scale transterability of PROSE
To establish the performance of PROSE sampler proposal distribution, first evaluate the trained flows in the Boltzmann generator setting. Here we generate a set of proposal particles {z; L. evaluate model likelihoods go (I,) and reweight using SNIS as in Eq. (7). In addition the trained models we benchmark against the following pretrained baselines; (i) the TBG model trained by Klein and Noe [2024], denoted ECNF (ii) Warp [Klein et al., 2023a] (iii) BioEmu [Lewis et al., 2024] (iv) Unisim Yu et al., 2025], For all methods we permit a budget of 10 energy evaluations. For the Boltzmann generator methods (ECNF, ECNF++, TarFlow, PROSE) this corresponds to 10 SNIS particles; for further information on the budget allocation of non-BG methods (Time BioEmu, UniSim) see Appendix F. We additional provide results for the unweighted proposal distributions in Appendix F. We note the Warp dataset to lack any sequences containing Proline at the N-Terminal, hence neither Warp nor ECNF were evaluated on such sequences.
Table 2: Quantitative results for baseline methods, and flows with self-normalized importance sampling on peptide systems up to residues. All methods evaluated budget 10 energy evaluations. Best values in bold. Not evaluated on sequences with N-terminal proline due to absence training data.
Sequence length-
2AA (huaan
BAA 03ymmns
Model |
PSS 1
R-W, 1
T-W,
FSS
F-W- J
T-W:
TICA-W2 J
FSS f
F-WWa J
T-W²
TICA-Wa
TameWarp
4.532
7.237
2.204
BioEmu
45.313
1.208
90.079
2.037
1.499
193,873
4.638
L.601
UniSim
10
1.289
10"
2.766
1.733
10'
6.156
1.495
ECNF"
0.086
0.894
0.488
ECNT++
0.024
3.470
0.302
0.008
10,032
1.121
0.572
TurFoow
0.134
0.452
0.193
0.045
1.260
0.924
0.492
0.008
11.298
2.733
1.087
PROSE
0.191
0.371
0.210
0.071
0.932
0.752
0.367
0.011
10.038
2,456
0.988
We present metrics for PROSE and baseline methods in Table 2. PROSE achieves the strongest performance on all metrics aside from dipeptide T-W,, where it is marginally outperformed by TarFlow, confirming to be a strong SNIS proposal for peptide systems of varying sequence length. ECNF+ performs very poorly on E-W2 on both dipeptides and tetrapeptides, seemingly unable to learn an effective vector field when trained on tetrapeptides. Warp the strongest non-Bolzzmann generator baseline, with both BioEmu and UniSim attaining high values of E-W, However, we note the training data for these pretrained models does not correspond exactly with our evaluation data and hence they are not directly comparable. Fig. 1 further confirms the success of PROSE with SNIS as an amortized sampler, surpassing the performance of a baseline MD trajectory on the critical T-W2 and TICA-Wa describing metastable state coverage w.r.t. both energy evaluations and GPU walltime. We additionally present qualitative results on the unseen octapeptide DGVAHALS Fig. 3, demonstrating the unprecedented scalability of PROSE; further results are provided in Appendix F.
7

---

0.06
MD
Proposal
Density
0.04
SNIS
TIC₁
TIG
0.00
-300
-250
-200
TICo
TICo
E(x)
Figure 3: PROSE accurately samples from the Boltzmann distributions unseen octapepitde system. Empirical results for sampling from DGVAHALS peptide system, not present in training data. Energy histogram (left) for reference MD data, PROSE proposal and PROSE reweighted using SNIS, demonstrate fine-grat ine-grained detail accuracy. TICA plots for MD (center) and SNISSreweighted PROSE (right) illustrate mode coverage.
4.4 Architecture ablation study
We proceed to ablate the architectural variations applied in PROSE, as described in Section 3.1; (i) the adaptive system conditioning blocks in the transformer layers (ii) the backbone permutations interleaved into our permutation sequence. We train ablation models using an identical training con- figuration to that described in Section 4.2. We present metrics for these modifications in Table 3. We observe a significant improvement in all metrics across scales of peptide sequence length, confirming the efficacy of these modifications for atomistic modeling, notably the backbone permutations which introduce negligible runtime complexity over the standard TarFlow architecture.
Table Ablation results for PROSE architecture components. SNIS performed with x 100h energy evaluations.
Sequence
2AA chyeres
4AA cppeae
8AA obyees
Model
PSS
R-WS |
T-Wo
FS8
F-W6 J
T-W J
TICA-Ws
FSS
F-WWs J
T-Wh
TICA-W:
PROSE
0.191
0,282
0.177
0.071
0.646
0,607
0,349
0.011
9.360
2.019
0.960
w/o Backbone-lic first
0.170
0.295
0.152
0.051
0.816
0.697
0.421
0.009
10.261
2,275
1.044
weo Trunsition
0.152
0.322
0.282
0.054
0.880
0.691
0.384
0.009
11.384
2.209
1.012
4.5 Sampling algorithms
Having established the unmatched performance of PROSE in the standard Boltzmann generator framework, we now consider alternative sampling algorithms made tractable by its efficient likelihood. We evaluate SNIS, SMC in continuous time, SMC in discrete time, and the simple instantiation of self-improvement defined in Section 3.2. All methods are permitted a budget of 10t energy evaluations, further details on method configurations are provided in Appendix E. Metric results are presented in Table 4. These results reveal the surprising result that, given a suitably strong proposal distribution, SNIS is competitive with both SMC variants, despite requiring no tuning. While SMC discrete achieves the best value of E-W2 on octapeptides, both SMC variants introduce a notable SNIS. Furthermore, the
deterioration of macrostructure metrics at this scale when compared performance of SNIS with self- improvement at improving the tetra- and octapeptide E-W, provides strong evidence in favor proposal fine-tuning within sampling methods, as an alternative to resource allocation solely on annealing- based methods. We provide results for the unseen RLMM system in Fig. 4, illustrating the superior mode coverage of PROSE with SNIS over a MD baseline given an allocation of 10th energy evaluations, further evidence of successful amortized sampling with PROSE.
Table 4: Results for samplers using PROSE proposal. Methods provided with hudget of 10 energy evaluations. Best values bolded. 'Not evaluated sequences with N N-terminal proline due absence training data.
Sequence length
2AA
4AA
8AA
Algorishm I
1
E-W, 1
T-W,
E-W, 4
T-W5
TICA-W.
E-W,
7-W,
TICA-
Timewarp
2.551
0.580
4.125
1.600
01813
0.271
0.070
0.665
9.386
0.964
SMC Continanus
0.318
0.177
0.764
0.688
0.338
10.563
2.642
1.049
SMC Diserete
0.249
0.147
0.653
0.721
01-00))
7.672
2.524
1,086
Self- Improvee
0.265
0.070
0.345

---

TIC
TIC1
TIC
TICo
TICo
TICo
Figure 4: By drawing uncorrelated proposal samples, PROSE achieves greater metastable state coverage than molecular dynamies for the same number of energy evaluations. TICA projection plots for unseen tetrapeptide system (RLXM). After 5 10 energy evaluations the reference molecular dynamies (left) traversed four distinct metastable states, taken to be ground truth. However, with an energy evaluation budget of 10 molecular dynamics explores only a single metastable state (center). highlighting the limitations of simulation- based sampling methods for mode exploration. PROSE with SNIS (right) samples all 4 states given the same budget of energy evaluations, indicating successful amortization of the mode exploration problem.
4.6 Inference-time temperature transfer
We lastly evaluate the scaled prior (SP) technique for inference-time temperature transfer introduced in Section 3.2. We collect additional 1 jus MD trajectories for the RLMM unseen tetrapeptide at temperatures defined by geometric series between the base model temperature of 310 K, and 800 K. We then perform SNIS using 2 x 10t energy evaluations from PROSE, both naively and with the scaled prior inference method. Results are presented in Fig. 5, with scaled prior universally outperforming naive SNIS. We emphasize that scaled prior does require any fine-tuning and introduces negligible increase in complexity at inference. These results thus demonstrate that PROSE is transferable not only in system, but also in temperature, opening variety of avenues of further exploration.
0.6
Prose
1.4
Prose (SP)
0.5
40
1.2
TICA->
E-W-
0.4
20
0.3
0.8
0.2
300
400
500
600
700
800
300
400
500
600
700
800
300
400
500
600
700
B00
Temperature (K)
Temperature (K)
Temperature (K)
Figure 5: Scaled prior greatly improves the ability of PROSE to accurately reweight to arbitrary tempera- tures. Metries for PROSE RLMM unseen tetrapeptide, targeting temperatures up 800 K. Naively applying SNIS to the larget temperature leads to rapid degradation in energy distribution, and to lesser extent the dihedral angle distribution. Applying prior scaling (PROSE SP) leads to a significant improvement in energy distribution at high temperatures and moderate improvement in dihedral angles. Notably. the TICA distribution improves higher temperatures irrespective of scaled prior usage, although scaled prior remains more effective.
5
Related Work
Normalizing flows. Normalizing flows [Rezende and Mohamed, 2015, Dinh et al., 2017, Kingma and Dhariwal, 2018, Durkan et al., 2019] fell from favor as generative adversarial networks (GANs) [Goodfellow et al., 2014], diffusion models [Ho et al., 2020, Song et al., 2021], and continuous normalizing flows [Chen et al., 2018a, Liu, 2022, Albergo and Vanden-Eijnden 2023], demonstrated superior empirical generative quality. However, they have still found relevance in scientific applications where efficient likelihood calculations are necessary. Furthermore, the recent introduction of Transformer Kolesnikov et al., 2024] has enabled previously intractable data distributions be modeled whilst retaining efficient likelihood evaluation, and brought renewed research attention to this area.
generative models as
normalizing flows [Zhai et al., 2024,

---

Boltzmann generators. Boltzmann generators are machine learning based samplers that train likelihood-basee models and employ inference-time SNIS achieve consistent sampling of the target density (Noé et al., 2019]. A major limitation of standard Boltzmann generators is the need to train the proposal model on a dataset of true density samples, motivating methods that transfer between molecular systems. Accordingly. Jing et al. [2022] propose a transferable Boltzmann generator operating on the torsion angles of small molecules, and Klein and Noe [2024] develop a Boltzmann generator operating on Cartesian coordinates that transfers between dipeptide systems. However, the scalability of such methods has remained limited due to the difficulty of designing expressive generative models that possess efficient and accurate likelihood evaluations. In particular the use of continuous normalizing flows implies a large cost to proposal likelihoods due to the need to integrate the vector field divergence [Grathwohl et al., 2019]. Schopmans and Friederich [2025] replace the single-step SNIS of standard Boltzmann generators with a normalizing flows, performing SNIS with each flow sample from a given target distribution.
sequence of
Approaches machine learning-based sampling. The widespread empirical success of generative modeling has inspired many approaches to machine learning- -baed sampling, Boltzmann emulators, like Boltzmann generators, seek uncorrelated sampling of the target density, but forgo efficient likelihood evaluation in favor of scalable generative modeling on large pre-collected datasets [Abdin and Kim, 2024, Wayment-Steed -ttele et al., 2024, Lewis et al., 2025]. In this case the lack of efficient likelihood evaluation precludes the use of Monte Carlo estimators such as SNIS. Diffusion samplers, propose novel objectives for training diffusion models in the absence of an empirical data distribution, both simulation- based [Berner et al., 2024, Vargas et al., 2023, Richter et al., 2024, Zhang and Chen, 2022, Vargas et al., 2024] and simulation-fir De Bortoli et al., 2024]. Notably, Havens al. [2025] develop diffusion sampler that transferable across small molecules. Time coarseners are another family of model, in which ML is used to predict large time transitions for simulation [Schreiner al., 2023, Fu et al., 2023, Klein al., 2023b, Daigavane et al., 2024, Yu et al., 2025], whereas methods like MDGen apply generative modeling to both the spatial and temporal dimensions of MD data [Jing et al., 2024]. Lastly, several works integrate normalizing flows with classical Monte Carlo methods [Albergo et al., 2019, Arbel et al., 2021, Gabrié et al., 2021, Matthews al., 2022, Midgley et al., 2023b, Hagemann al., 2023].
[Akhound-Sadegh et al., 2024, Huang et al., 2021,
Conclusion
We develop PROSE, demonstrating that deep learning-based samplers can efficiently transfer to previously unseen systems at unprecedented scale. PROSE outperforms learned baseline methods, as well as molecular dynamics, at a variety of energy evaluation and walltime budgets. Notably, PROSE demonstrates cate-off-the-an performance whilst retaining many simple design choices; thus, leaving many directions for further development. The competitive performance of SNIS compared to SMC invites further investigation into the merits of annealing- based samplers given a proposal with good coverage of the target density. Naturally, annealing- based samplers have been enhanced beyond the simple instantiations we explore; careful tuning of SMC may yield further improvements [Syed al., 2024]. We lastly note the self- ceff-improvement strategy discussed to not be restricted to SNIS, the integration of advanced Monte Carlo methods presents an avenue for future work.
I Limitations. Whilst conventional Monte Carlo algorithms make no assumption on the target density, transfcrable learned samplers, including PROSE, rely on the assumption that the system belongs to structured space of energy functions, in our case the chemical space of peptides. To achieve greater practical relevance it will be necessary to consider a more diverse chemical space, such as the recent OMol25 [Levine et al., 2025] dataset. We lastly comment that, despite the scaled prior method demonstrating surprising abilities to sample from the higher temperatures, we believe that precise transfer to lower temperatures would require further algorithmic development.
10

---

Acknowledgments
This research partially supported by the EP- SRC Turing AI World-Leading Research Fellowship No. EP/X040062/1 and EPSRC AI Hub No. EP/Y028872/1. The authors acknowledge funding from UNIQUE, CIFAR, NSERC, Intel, and Samsung. The research was enabled in part by computational resources provided by the Digital Research Alliance of Canada (https: / /atla//allii ancecan. ca), Mila (uttps://mila,
quebbee quebec), and NVIDIA. The authors additionally thank HuggingFace for hosting the ManyPeptidesMDD dataset. KN was supported by IVADO and Institut Courtois.
References
Osama Abdin and Philip M. Kim. Direct conformational sampling from peptide energy landscapes through sxpernetwork-condttioo diffusion. Nature Machine Intelligence, 6(7):775-786) July
2024. ISSN 2522-5839. doi: 10. 10.1038//42256-0244 022-0560-00-0000--44444
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew J. Ballard, Joshua Bambrick, Sebastian W. Bodenstein, David A. Evans, Chia-Chun Hung, Michael O' O'Neill, David Reiman, Kathryn Tunyasuvunakool, Zachary Wu, Akvilé Zemgulyte, Eirini Arvaniti, Charles Beattie, Ottavia Bertolli, Alex Bridgland, Alexey Cherepanov, Miles Congreve, Alexander Cowen-Rivers, Andrew Cowie, Michael Fig- urnov, Fabian B. Fuchs, Hannah Gladman, Rishub Jain, Yousuf A. Khan, Caroline M. R. Low, Kuba Perlin, Anna Potapenko, Pascal Savy, Sukhdeep Singh, Adrian Stecula, Ashok Thillaisun- daram, Catherine Tong, Sergei Yakneen, Ellen Zhong, Michal Zielinski, Augustin Zídek. Victor Bapst, Pushmeet Kohli, Max Jaderberg, Demis Hassabis. and John M. Jumper. Accurate structure prediction of biomolecular interactions with AlphaFold 3. Nature, 2024.
Tara -aauun, Jarrid Rector-Brooke Joey Bose, Sarthak Mittal, Pablo Lemos, Cheng-Hao Liu, Marcin Sendera, Siamak Ravanbakhsh, Gauthier Gidel, Yoshua Bengio, Nikolay Malkin, and Alexander Tong. Iterated denoising energy matching for sampling from boltzmann densities. In International Conference Machine Learning (ICML), 2024.
Michael S. Albergo and Eric Vanden- Building normalizing flows with stochastic interpolants. International Conference on Learning Representations (ICLR), 2023.
Michael S Albergo and Eric Vanden-Eijnden. Nets: A non-equilibrium transport sampler. In International Conference on Machine Learning (ICML), 2025.
Michael S Albergo, Gurtej Kanwar, and Phiala Shanahan. Flow-based generative models for markov chain monte carlo in lattice field theory. Physical Review D, 2019.
Michael Arbel, Alex Matthews, and Arnaud Doucet. Annealed flow transport monte carlo. In International Conference on Machine Learning, 2021.
Julius Berner, Lorenz Richter, and Karen Ullrich. An optimal control perspective on diffusion-based generative modeling. Transactions on Machine Learning Research (TMLR), 2024.
Ignasi Buch, Toni Giorgino, and Gianni De Fabritiis. Complete reconstruction of enzyme- -iinibiin binding process by molecular dynamics simulations. Proceedings of the National Academy of Sciences, 2011.
D. A. Case, V. Babin, J. T. Berryman, R. M. Betz, Q. Cai, D. Cerutti, T. E. Cheatham, T. A. Darden, R. E. Duke, H. Gohlke, A. W. Goetz, S. Gusarov, N. Homeyer, P. Janowski, J. Kaus, I. Kooossváry, Luchko, Luo, B. Madej, K. M. Merz, Paesani, D. R. Roe, A. Roitberg, C. Sagui, R. Salomon-Ferrer, G. Seabra, C. L. Simmerling, W. Smith, Swails, R. C. Walker, J. Wang, R. M. Wolf, X. Wu, and A. Kollman. AMBER 14. University of California, San Francisco, 2014.
A. Kovalenko, T. S. Lee, S. LeGrand,
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Neural Information Processing Systems (NIPS), 2018a.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 2018b.

---

Ameya Daigavane, Bodhi P Vani, Saeed Saremi, Joseph Kleinhenz, and Joshua Rackers. Jamun: Transferable molecular conformational ensemble generation with walk-jump sampling, arXiv, 2024.
Valentin De Bortoli, Michael Hutchinson, Peter Wirnsbegee, and Arnaud Doucet. Target score matching. arXiv, 2024.
Pierre Del Moral. Mean field simulation for monte carlo integration. Monographs on Statistics and Applied Probability, 2013.
Manuel Dibak, Leon Klein, Andreas Krämer, and Frank Noé. Temperature steerable flows and Boltzmann generators. Phys. Rev. Res., 2022.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using Real NVP. International Conference on Learning Representations (ICIR), 2017.
Arnaud Doucet, Nando De Freitas, Neil James Gordon, et al. Sequential Monte Carlo methods in practice. 2001.
Conor Durkan, Artur Bekasov, lain Murray, and George Papamakarios, Neural spline flows. Advances in neural information processing systems, 2019.
Peter Eastman, Jason Swails, John D. Chodera, Robert T. McGibbon, Yutong Zhao, Kyle A. Beauchamp, Lee-Ping Wang. Andrew Simmonett, Matthew P. Harrigan, Chaya D. Stern, Rafal P. Wiewiora, Bernard R. Brooks, and Vijay S. Pande. OpenMM 7: Rapid development of high performance algorithms for molecular dynamics. PLoS computational biology, 2017.
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Lactitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T. H. Gayraud, Hicham Janati, Alain Rakotomamonjy, levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong. and Titouan Vayer. POT: Python Optimal Transport. Journal of Machine Learning Research, 2021.
Xiang Fu, Tian Xie, Nathan J Rebello, Bradley Olsen, and Tommi S Jaakkola. Simulate time- integrated coarse-grained molecular dynamics with multi-scale graph networks. Transactions on Machine Learning Research, 2023.
Marylou Gabrié, Grant M Rotskoff, and Eric Vanden-Eijnden Efficient Bayesian sampling using normalizing flows to assist Markov chain Monte Carlo methods. arXiv, 2021.
Tomas Geffner, Kieran Didi, Zuobai Zhang, Danny Reidenbach, Zhonglin Cao, Jason Yim, Mario Geiger, Christian Dallago, Emine Kucukbenli, Arash Vahdat, and Karsten Kreis. Proteina: Scaling Flow-based Protein Structure Generative Models. In International Conference on Learning Representations (ICLR), 2024.
Ian J. Goodfellow, Jean Pouget-Abadia Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. In NeurIPS, 2014.
Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. Ffjord: Free-form continuous dynamics for scalable reversible generative models. International Conference on Learning Representations (ICLR), 2019.
Paul Lyonel Hagemann, Johannes Hertrich, and Gabriele Steidl. Generalized normalizing flows via Markoy chains. 2023.
Aaron Havens, Benjamin Kurt Miller, Bing Yan, Carles Domingo-Enrien Enrich, Anuroop Sriram, Brandon Wood, Daniel Levine, Bin Hu, Brandon Amos, Brian Karrer, Xiang Fu, GGan-Hormp Liu, and Ricky T. Q. Chen. Adjoint sampling: Highly scalable diffusion samplers via adjoint matching, 2025.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models, 2020.
Shinya Honda, Kazuhiko Yamasaki, Yoshito Sawada, and Hisayuki Morii. 10 residue folded peptide designed by segment statisties. Structure, 2004.
12

---

Jian Huang, Yuling Jiao, Lican Kang, Xu Liao, Jin Liu, and Yanyan Liu. Schrödinger-Föllmer sampler: sampling without ergodicity. arXiv, 2021.
Christopher Jarzynski. Nonequilibrium equality for free energy differences, Physical Review Letters, 1997.
Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi Jaakkola. Torsional diffusion for molecular conformer generation. Advances in Neural Information Processing Systems, 2022.
Bowen Jing, Hannes Stärk, Tommi Jaakkola, and Bonnie Berger. Generative modeling of molecular dynamics trajectories. In Neural Information Processing Systems (NeurIPS), 2024.
Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1x1 convolutions. Advances in neural information processing systems, 2018.
Durk Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Improved variational inference with inverse autoregressive flow. Advances in neural information processing systems, 2016.
Leon Klein and Frank Noe. Transferable Boltzmann Generators. 2024.
Leon Klein, Andrew Y. K. Foong, Tor Erlend Fjelde, Bruno Kacper Mlodozeniec, Marc Brockschmidt, Sebastian Nowozin, Frank Noe, and Ryota Tomioka. Timewarp: Transferable Acceleration of Molecular Dynamics by Learning Time-Cor Coae-eCaree Dynamics. In Neural Information Processing Systems (NeurIPS), 2023a.
Leon Klein, Andrew YK Foong, Tor Erlend Fjelde, Bruno Mlodozeniec, Mare Brockschmidt. Sebastian Nowozin, Frank Noć, and Ryota Tomioka. Timewarp: Transferable acceleration of molecular dynamics by learning time-coarsened dynamics. Neural Information Processing Systems (NeurIPS), 2023b.
Leon Klein, Andreas Krämer, and Frank Noe. Equivariant flow matching. Advances in Neural Information Processing Systems, 2023c.
Alexander Kolesnikov, André Susano Pinto, and Michael Tschannen. Jet: modern transformer- based normalizing flow. arXiv preprint arXiv:2412.15129 2024.
Jonas Köhler, Michele Invernizzi, Pim De Haan, and Frank Noe. Rigid Body Flows for Sampling Molecular Crystal Structures. In Proceedings of the 40th International Conference on Machine Learning, 2023.
Benedict Leimkuhler and Charles Matthews. Molecular Dynamics: With Deterministic and Stochastic Numerical Methods. 2015.
Daniel S. Levine, Muhammed Shuaibi, Evan Walter Clark Smotte-Smith, Michael G. Taylor, Muham- mad R. Hasyim, Kyle Michel, Ilyes Batatia, Gábor Csányi, Misko Dzamba, Peter Eastman, Nathan C. Frey, Xiang Fu, Vahe Gharakhanyan, Aditi S. Krishnapriyan, Joshua A. Rackers, San- jeev Raja, Ammar Rizvi, Andrew S. Rosen, Zachary Ulissi, Santiago Vargas, C. Lawrence Zitnick, Samuel M. Blau, and Brandon M. Wood. The open molecules 2025 (omol2.) dataset, evaluations, and models, 2025.
Sarah Lewis, Tim Hempel, José Jiménez Luna, Michael Gastegger, Yu Xie, Andrew YK Foong, Victor Gareía Satorras, Osama Abdin, Bastiaan S Veeling, Iryna Zaporozhets, et al. Scalable emulation of protein equilibrium ensembles with generative deep learning. bioRxiv, 2024.
Sarah Lewis, Tim Hempel, José Jiménez-Luna, Michael Gastegger, Yu Xie, Andrew Y. K. Foong, Victor García Satorras, Osama Abdin, Bastiaan S. Veeling, Iryna Zaporozhets, Yaoyi Chen, Soojung Yang, Arne Schneuing, Jigyasa Nigam, Federico Barbero, Vincent Stimper, Andrew Campbell, Jason Yim, Marten Lienen. Yu Shi, Shuxin Zheng, Hannes Schulz, Usman Munir, Ryota Tomioka, Cecilia Clementi, and Frank Noé. Scalable emulation of protein equilibrium ensembles with generative deep learning, bioRxiv, 2025.
13

---

Kresten Sandorff-Larsen Stefano Piana, Ron O. Dror, and David E. Shaw. How Fast-Folding Proteins Fold. Science, 2011.
Jun Liu. Monte Carlo Strategies in Scientific Computing. 2001.
Qiang Liu. Rectified flow: A marginal preserving approach to optimal transport. arXiv, 2022.
Iyy Loshchilov and Frank Hutter. Decoupled Weight Decay Regularization. 2018.
Alex Matthews, Michael Arbel, Danilo Jimenez Rezende, and Arnaud Doucet. Continual repeated annealed flow transport monte carlo. International Conference on Machine Learning (ICML), 2022.
Laurence Midgley, Vincent Stimper, Javier Antorán, Emile Mathieu, Bernhard Schölkopf, and José Miguel Hernández-Lobato. SB(3) Equivariant Augmented Coupling Flows. Advances in Neural Information Processing Systems, 2023a.
Laurence Illing Midgley, Vincent Stimper, Gregor NC Simm, Bernhard Schölkopf, and José Miguel Hemández-Lobato Flow annealed importance sampling bootstrap. International Conference on Learning Representations (ICLR), 2023b.
Radford M Neal. Annealed importance sampling. Statistics and computing, 2001.
Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 2019.
Frank Noé, Christof Schütte, Eric Vanden -ijedenn, Lothar Reich, and Thomas R. Weikl. Constructing the equilibrium ensemble of folding pathways from short off-equilibrium simulations. Proceedings of the National Academy of Sciences, 2009.
George Papamakarios, Theo Pavlakou, and lain Murray. Masked autoregressive flow for density estimation. In Neural Information Processing Systems (NeurIPS), 2017.
Sanjeev Raja, Martin Sipka, Michael Psenka, Tobias Kreiman, Michal Pavelka, and Aditi Krish- napriyan. Action-minimization meets generative modeling: Efficient transition path sampling with the onsager-machlup functional. Forty-second International Conference on Machine Learning, 2025. URL https: // hptpr/veeeeeeeeeeeeeeeeeeeeeeeee
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. International Conference on Machine Learning (ICML), 2015.
Lorenz Richter, Julius Berner, and Guan-Horm Liu. Improved sampling via learned diffusions. International Conference on Learning Representations (ICLR), 2024.
Gareth O Roberts and Richard L Tweedie. Exponential convergence of langevin distributions and their discrete approximations. Bernoulli, 1996.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks. International Conference Machine Learning (ICML), 2021.
Henrik Schopmans and Pascal Friederich. Temperature- Annealed Boltzmann Generators, 2025.
Mathias Schreiner, Ole Winther, and Simon Olsson. Implicit transfer operator learning: Multiple time- resolution models for molecular dynamics. In Thirty-sevnnth Conference Neural Information Processing Systems, 2023.
Noam Shazeer. Glu variants improve transformer, 2020. URL U...ittts://aavivoorrr/bbb// 05202.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations, 2021.
Saifuddin Syed, Alexandre Bouchard-Côté. Kevin Chern, and Arnaud Doucet. Optimised annealed sequential monte carlo samplers. arXiv preprint a0x 2024.
14

---

Charlie B. Tan, Avishek Joey Bose, Chen Lin, Leon Klein, Michael M. Bronstein, and Alexander Tong. Scalable Equilibrium Sampling with Sequential Boltzmann Generators, 2025.
Francisco Vargas, Will Grathwohl, and Arnaud Doucet. Denoising diffusion samplers. International
Conference on Leurning Representations (ICLR), 2023.
Francisco Vargas, Shreyas Padhy, Denis Blessing, and Nikolas Nüsken. Transport meets variational in- ference: Controlled Monte Carlo diffusions. International Conference on Learning Representations (ICLR), 2024.
Hannah K. Wayment-Steck Adebayo Ojoawo, René Otten, al. Predicting multiple conformations via sequence clustering and alphafold2. Nature, 2024.
Xiongwu Wu and Bernard R. Brooks. Beta-hairpin folding mechanism of a nine-residue peptide revealed from molecular dynamics simulations in explicit water. Biophysical Journal, 2004.
Ziyang Yu, Wenbing Huang, and Yang Liu. Unisim: A unified simulator for time-coarsened dynamics of biomolecules, 2025.
Shuangfei Zhai, Ruixiang Zhang, Preetum Nakkiran, David Berthelot, Jiatao Gu, Huangjie Zheng, Tianrong Chen, Miguel Angel Bautista, Navdeep Jaitly, and Josh Susskind. Normalizing Flows are Capable Generative Models, 2024.
Qinsheng Zhang and Yongxin Chen. Path integral sampler: a stochastic control approach for sampling. International Conference on Learning Representations (ICLR), 2022.
15

---

Appendix
PROSE architecture Details
17
A.I Tokenization and masking
17
A.2 Permutations
17
A.3 Adaptive layer norm and transition
17
Dataset
18
C Training configuration
18
C.1 Computational resources
19
D Importance sampling variants
19
D.I Self-normalized importance sampling
19
D.2 Continuous time annealed importance sampling
20
D.3 Discrete-time annealed importance sampling
20
D.4 Sequential Monte Carlo
21
Evaluation configuration
E.1
Proposal sampling and likelihood evaluation
21
E,2
Metries
22
E,3 Sampling algorithm contigurations
24
E,4
Additional baseline configurations
24
E.5
Computational resources
25
Supplementary experimental results
25
F1
Unweighted proposal performance
25
F2
Additional ablations
25
F3
Dataset scaling
26
F.4
Sequence-length extrapolation
27
F.5
JSD evaluation
27
F66
Effective sample per second
F7
Octapeptide Ramachandran plots
28
F88
Temperature plots
29
16

---

A PROSE architecture Details
A.1 Tokenization and masking
As in SBG [Tan et al., 2025], the token sequence is constructed using single atom per token xri] R². As different peptide sequences contain varying numbers atoms, we zero-pad atom sequences to fixed maximum sequence length and introduce padding mask m[i] or padded (m]ii] = 0) tokens. In the context of a causal transformer, the implementation greatly simplified by ensuring all padding tokens are placed at the end of the token sequence, irrespective of the permutations applied. We may therefore state the following block update rules, as only a minor adaptation of Eq. (2)
indicating valid (m[i) 1)
== 0,
[
(2 (z+(: a)({i) o exp(-or(<)) M)))
m m 1, m 0,
(1I)
0,,0,
with log- determmittmn of Jacobian given by
ff(()))
N D-1 Σ
log
Σ
mii
ax(2 d]
(12)
Ōz.
The inverse transformation is correspondingly minor adaptation of Eq. (3)
Z++ [i]
i= 0, m|i mli 0.
2,
2++1|d [i] @ exp(ar(<]) D)) + pe(zal: e]]]i [0,0, 0J.
(13)
A.2 Permutations
All permutations are defined from the N-terminal the C-terminal. In the residue-by-residue permutation the atoms are ordered such that each residue forms a contiguous sequence, with sidechain atoms immediately following the corresponding backbone atoms. In the backbone-firs first permutation the entire sequence of backbone atoms is placed at the start of the sequence before any sidechain atoms. Where constituent residues possess branch or ring we also introduce variant in which the branch ordering is flipped or ring traversal is inverted. The flip permutations are a simple inversion of the permutation, where padding tokens and not moved from their position at the end of the sequence. The specific sequence of permutations employed in PROSE is presented in Table 5.
Table 5: Autoregressive permutation order used the eight transformation blocks.
Permutation
Description
To
Backbone first
T
Residue- (flip)
TT2
Backbone-first (flip)
TT
Residue-by-reiduee
Backbone-first (variant)
TT
Residue-by- (variant, flip)
TT
Backbone-first (variant, flip)
Residue-by-residue (variant)
A.3 Adaptive layer norm and transition
PROSE integrates the adaptive layer normalization and SwiGLU-based [Shazeer, 2020] transition modules employed by Geffner et al. [2024] into the transformer blocks the TarFlow architecture [Zhai et al., 2024]. The positions of the latent vector z, are encoded using sinusoidal positional encoding, which are added directly to z.. The conditional embedding is used in the adaptive layer normalization and adaptive scale components.
17

---

Counce
Litror
Spiit
P
Litene-
Splts
Dviral
Magaive
Linou
Spli
Good. Eiatediting
Figure 6: Adaptive Layer Norm and Transition. The transformer block is modified to incorporate conditional
information using adaptive layer normalization and transition block. Figure adapted from Geffner et al. [2024].
B Dataset
Sequence sampling. Training sequences are collected for all peptide lengths by uniformly sampling the 20 standard amino acids. For the 8-residue test data, a sequence length 30 8 240 constructed by concatenating 12 of each amino acid. This sequence is then randomly permuted and split into peptides of length 8, ensuring that each amino acid is represented uniformly. A similar process is performed for length 4 but was not possible at length 2. In both training and test sets, the N- and C-terminal residues are protonated to form the zwitterionic state of the peptides. Initial structure files (PDB format) are generated using AmberTools' tleap.
Molecular dynamics simulation. Local energy minimization is performed with the Limited-memory Broyden-Fletcher Dletcher-Goldfarbb Soetfarr-Shanno (L-BGFS) algorithm. Energy minimization is followed by burn-in simulation of length ps, after which samples are collected every 1 ps (train) or 10 ps (test) until the simulation budget is exhausted. Full MD simulation parameters are provided in Table 6.
Table OpenMM simulation parameters.
Force field
Amber14
Integration time step
fss
Friction coefficient
0.3 ps
Temperature
310K
Nonbonded method
CutoofflonPeriodic
Nonbonded cutoff
2 mm
Integrator
LangevinnhiddlleIntegrraa
Table Training and evaluation dataset parameters.
Train
Test
Burn-in period
50 ps
50 ps
Sampling interval
1 ps
10
Simulation time
200
5
C
Training configuration
All models are trained for 105 iterations using a batch size of 512 with the AdamW optimizer [Loshchilov and Hutter, 2018]. We employ cosine learning rate schedule in which the initial and final learning rates are a reduction of the maximal value by factor of 500, as well as exponential moving average with decay of 0.999. No overfitting was observed hence no early stopping was required. Given the large size of the training trajectories, we subsample to 10 ps per frame. Samples normalized using values approximating the standard deviation of the 8AA data (o 0.35), if
18

---

absent the 4AA data (o = 0.28), noting that single value must be shared across systems of different dimensionality. An overview of all training configurations is provided in Table
Continuous Normalizing Flows. We use the +CNF++ training recipe defined by Tan et al. [2025]; this entails learning rate of 5 10 and weight decay of 1 10 with default AdamW hyperparam- eters of AdamW 31, .2 of (0.9, 0.999). In contrast the ECNF of Klein and Noe [2024] was trained without weight decay exponential weight averaging. The channel width and layer depth both models is defined in Table 9.
learning rate of 1 10-4, weight decay of 10 10-4, and AdamW B1, B2 of (0.9, 0.95). Data augmentation is applied as random rotations and Gaussian center of mass augmentation, in which every the entire system conformation translated by a vector e ~ ~N [a). The o² value is chosen to match that of the prior, which has center of mass variance of N where N is the number of atoms. Given N is in variable for single model trained on multiple systems, this augmentation is applied-per system before padding applied. The architecture width and depth is provided Table 9.
TarFlows. Following Zhai al. [2024] and Tan et al. [2025]
Table 8: Overview of training configurations.
BCNF
+CNF++
TarFlow / PROSE
Learning Rate
10
50
10
Weight Decay
0.0
1-10-2
4 10-4
31. 13
0.9, 0.999
0.9, 0.999
0.9, 0.95
EMA Decay
0.0
0.999
0.999
Table Overview of model scaling parameters. For TarFlow variants depth corresponds number parame- terized transformations, ECNF variants this simply the number of graph neural network layers.
ECNF
ECNF+
TarFlow
PROSE
Channels
128
256
384
384
Depth
9
9
8
8
Layers per block
N/A
N/A
8
Parameters (M)
1
4
115
285
C.1 Computational resources
All training experiments are run NVIDIA H100 GPUs using distributed data parallelism. The training throughput for each model is presented in Table 10.
Table 10:
Training throughput for models presented in Table 2. We highlight ECNF++ to be trained only on sequences up to length whereas TarFlow and PROSE are trained on sequences up to length
ECNF+
TarFlow
PROSE
Training iterations H100 hour
960
1132
260
D Importance sampling variants
D.1 Self-normalized importance sampling
Self-normalized importance sampling (SNIS) corresponds to the following estimator
Wi
p(x)), w; = p(xx)
E B (x)
q(z,)* x: ~ q(x),
(14)
Σ wy
where uses the learned density model of PROSE gz( q(x).
19

---

D.2 Continuous-time annealed importance sampling
Below, we repeat the derivations from [Jazyyskki, 1997, Albergo and Vanden- -inden-Eijnden, 2025]. Namely, we consider the continuous family of marginal densities
(x))) Ze f fdr exp(-Fe(xx)
q (3)
(15)
Z
The PDE describing the time-evolution of this density is dq, (x)
(x)
(x)
(x))
+ + (a) (a)
(16)
Ot
0t
δt
[E((z)
(z)
+
V. (x)
v log (xr)
+q((z))
+ EE (x))
(17)
2
01.
al
of V. (x))
+ of Aq((z) + (((())
(B.(z)
(E.())
v log
(x)
2
Öt
Ot
(18)
This is Feynman- PDE which can be simulated [Del Moral, 2013] as the following SDE on the extended space of states and weights Wt
of
dx. = (E(())dt + o,dW,, Iz=0 ~ go(x)
2
(19)
d dog 1h = (E((x)
dt, w²=0 =1.
Øt
The expectation the statistics (0x) w.r.t. the density qr(x) then can be estimated using SNIS as follows
wi
E (n) (x))
(((t),
(20)
Ej=1 wa
where (x w+) are the solutions of the SDE Eq. (19).
For the inference time of PROSE, we define the continuous family of marginals as
qe (x) exp (1 (1-1) log (x()) + (lgg(x)) I [0,1],
(21)
Ex(r))
where (z()) is the learned density of PROSE. Thus, Eq. (19) becomes
00 ((1 - t) V log + tV logp(x))dt + a,dW,, Iz=o ~ (x(x)
dx, 2
(22)
log W/ = (logp(r) - log go (x) dt,, w1= =1
D.3 Discrete-time annealed importance sampling
Consider sequence of marginal densities
Go (x) x exp(-Fo(z)).... qK (x) x ex)(-Fx(x)))
(23)
Let's denote by (a 1 the kernel that satisfies the detailed balance w.r.t. de (xt), i.c.
ge (xt-1 (x (m+-1)
qe (x (x)).
(24)
Then, one can write importance sampling estimator for the final marginal as
f dux (x(xx)<(zx) =
(x - (x)
(25)
dæ do dox (x æ - ) (x (I
(26)
gK( [x((KK1)
Egx 1xxx-
QK- (IK- (p(k)
(27)
20

---

Clearly, we can repeat the trick but now for 9K-1(xx-1). Thus, applying this trick recursively to different marginals, have
K-1
/ dax qu (x(x))p(<k) Boo(zo) Bes,.....
II ==
(+1l p((x),
(28)
(e(x))
where 21...., JK hn (x1 .. ( - 2 (T xx--1).
(29)
Thus, have the following SNIS estimator
i dx qu (x))(())
E Σx E, w²k p((k))
wk
(30)
where 2f ~ le x4-1) = L, I....,KK K, 2o ~ go(xo)
(31)
log : - Ei +
+ log wi
(32)
Note that there is of flexibility for the choice of (3r (x-1)) because we not use the densities of the transition kernel for the weights. In particular, the Metropolis Hettopolis-Hastingss algorithm with any proposal yields a reversible kernel (satisfies the detailed balance), which result a consistent final estimator. Furthermore, compared to the continuous-time AIS, discrete-time AIS does not introduce the time-discretization error.
At the inference step of PROSE, we choose
1
(e( x exp (1 - 1) log (x()) + (og))))) 1=
(33)
-E()))
and Metropolis-Adjuste Langevin Dynamics the transition kernel (,, 12)-1) [Roberts and Tweedie, 1996].
D.4 Sequential Monte Carlo
Sequential Monte Carlo (SMC) [Doucet et al., 2001] can be understood annealed importance sampling equipped with adaptive resampling. As in AIS, particles and weights are propagated in coupled system. However, SMC tracks the effective sample size (ESS)
(ΣN
i==l wi
ESS( ,...,,. W
(34)
Σ (2
and performs resampling step whenever ESS falls below a chosen threshold. This in contrast to AIS, in which resampling only occurs in the final timestep. This adaptive resampling seeks to mitigate weight degeneracy, which large compute allocations are required AIS to propogate very low weight particles. While the continuous-time and discrete-time time AIS variants are both suitable for SMC, only the continuous-time variant was considered by Tan et al. [2025].
E Evaluation configuration
E.1 Proposal sampling and likelihood evaluation
Equivariant continuous normalizing flows. Sampling from a continuous normalizing flow (CNF) involves solving the ODE defined by the parameterized vector field u, [0,
R== + R
dx. dt
(xt))
TD Po
(35)
The corresponding likelihoods can be obtained using the instantaneous change of variables formula
logpr(s1) = (og))( - ./ (x
(36)
21

---

where V. is the divergence operator. In practice both Eq. (35) and Eq. (36) can be integrated simultaneously with an ordinary differential equation (ODE) solver. We use the Dormand-Pring 5 (dopri5) adaptive solver in all ECNF experiments. Given the E(3) equivariance of the ECNF, samples are generated in both possible global chiralities. Following Klein and Noe [2024] we check for incorrect global sample chirality and flip samples appropriately to match the L-amino acids present in the evaluation data. Unlike Klein and Noe [2024] we do not omit any samples with unresolvable chirality. We additionally apply logit clipping, removing the samples with the 0.2% highest importance weights before resampling Midgley al. [2023a].
TarFlow variants. As discussed in Section 2.1, samples are generated from a normalizing flow simply by applying fo to prior samples (O, D)) Model likelihoods are obtained using the change of variables formula (Eq. (1)). Given the lack of translation equivariance in the TarFlow, and the data augmentation applied during training, samples are generated with an approximate scaled X3 distributiin over centroid norm ||e]| = ll Σi=1 Li,=ll ~ 0x3. This leads to adverse behavior when resampling with finite samples, hence apply the center of mass adjustment of [Tan al., 2025], in which the Xs probability density function is divided out of the proposal likelihoods
||l|
|||l 202
log p (x) log Po (z) log
+
log
211'
(37)
where I' is the gamma function. This adjustment seeks to account for the radial component introduced by translation non-equivariance. We additionally apply the same weight clipping threshold as the ECNF when performing SNIS, before SMC.
E.2 Metrics
We report both effective sample size and variety of Wasserstein-2 distances as evaluation metrics. For the Wasserstein distances subsample 10 samples are randomly sampled from the evaluation trajectory as ground truth. Similarly, at most 10 generated samples are employed; if a method has generated more samples a random subset is drawn without replacement.
Effective sample size. We compute the effective sample size (ESS) using Kish's formula, normalized by the number of samples generated
Σ i= W/
ESS( ((sw))))))))
(38)
N
w/
Empirical Wasserstein distance. We compare generated samples to ground truth data, collected as defined in Appendix B. using empirical Wasserstein-2 distances. Given empirical distributions J Ei= 8z, and v m Σラ= 8pp- the empirical Wasserstein-2 distance is defined
W₂ (2( min nETT(RP)
Σ
Ty c(ai, 12)
(39)
where I[(,, v)) denotes the set of couplings with marginals j and v, and c(x, y)² is a defined cost function. Different choices of c(æ, y)² define different measures of dissimilarity. We the POT [Flamary et al., 2021] linear optimal transport solver to compute the optimal couplings.
Energy cost. The energy of a sample E(x) is sensitive to both bonded forces and non-bonded forces. For the energy Wasserstein-2 distance E-W2 the cost function simply
(=( [E(z) - 2
(40)
Dihedral torus cost. The and 1/ backbone dihedral angles of a peptide conformation encode essential information regarding secondary and tertiary structure. We compare generated and ground truth samples in angle space by defining the dihedral angle vector
22

---

Dihedrals(z)) (Φ1, 1h Φ2 ........ 01.1......
(41)
where the number of residues. Given the torus geometry implied by angle periodicity di (-T, T.. natural cost function is the minimal signed angle difference
2L cT(,,)) Σ
(Dihhedrabs()); Dihedrals(y); + +) mod 2 - -m]².
(42)
This metric captures the geometric dissimilarity in dihedral angle space, respecting periodicity.
Time-lagged independent component analysis cost. The time-lagged independent component analysis (TICA) projection of time-series data captures directions along which the data exhibits maximal autocorrelation. Within molecular dynamics, TICA is commonly used to detect distinct metastable states. Given mean-free time series data Tt the instantaneous (zero-lag) empirical covariance and time-lagged empirical covariance matrix (at lag time 7) computed as
T- Σ
Coo = - Σ a Cur = T -
(43)
TICA seeks linear projection vectors E R" that maximize autocorrelation at lag
T TC w max T Coow
(44)
The solution to which is obtained by solving the generalized eigenvalue problem
Cor = XCoo
(45)
where the eigenvalue À measures the autocorrelation of the projected component at lag and the eigenvector w defines the corresponding slow mode. To define the TICA Wasserstein-22 distance TICA-Wa we take the full evaluation trajectory without subsampling and solve Eq. (45) to obtain the first two TICA projection vectors W1, W2. We may then define the following cost function
2
CHICA (x, u)²
Σ
w x w
2
(46)
defining similarity TICA projection space. practice, we compute the TICA projection for the heavy-atom pairwise distances and dihedral angles. We also emphasize that the TICA projection must be computed on the full 5 us evaluation trajectory (such that the slowest transitions may be detected), but that the samples y used in the Wasserstein metric are restricted to the 10 subset.
Jenson-Shannon divergence metries. Following Raja et al. [2025]. on the reference MD trajectory we run k-means (k = 20) on the first two time-lagged independent components (TICs) fitted to the pairwise distances and dihedral angles of the peptide confirmations. Using these clusters, we obtain the occupancy distributions for the reference MD and the samples generated from PROSE, and report their Jenson-Soe Jenson-Shannon Shannon divergence (TICA-K-JSD) defined as
JSD(P.Q) = 1 - KL (P) KL M), where KL(P,M) Σ P. log
P. M,
(47)
where P and Q are discrete distributions, and M = (P 1 procedure using k-means only on the dihedral angles, and report this as T-k-JSD. In both cases we fit the k-mmans clustering on the features (TICA projections or dihedrals) from full reference MD, compute metries of generated samples against a subsampling of the reference trajectory as previously discussed. Note that JSD depends on the arbitrary choice of clustering (here h 20) and can be sensitive to binning resolution, whereas Wasserstein distances avoid discretization errors by operating directly the continuous distributions.
Q). Moreover, we repeat the same
23

---

E.3 Sampling algorithm configurations
In this section we define configurations for the sampling algorithms presented in Table 4, as well as the molecular dynamics baseline used in Fig. I and Fig. 4. In particular the allocation of the 10 energy evaluations within the method defined.
Molecular dynamics baseline. We follow the same procedure used for collecting the main datasets defined in Appendix B, where the parameters defined in Table 6 are unchanged. We apply logarithmically decaying frame interval with appropriate reweighting to obtain accurate resolution across many orders of magnitude. The simulation is run for (1 jus) using 10 energy evaluations.
Sequential Monte Carlo. For both variants we generate a proposal set of 10 samples. For the discrete variant, we perform 50 annealing steps, requiring two energy evaluations per step: to update the samples and one for the Metropolis Hastiiggs correction. For the continuous variants, we perform 100 annealing steps. In both cases, resampling is performed every step. For the continuous variant, Langevin dynamics is used with a step size o, of 10 for dipeptides and tetrapeptides, and T of 10-8 for octapeptides. Details of the formulation are found in Appendix D.2. For the discrete variant, we apply Langevin dynamics with step size of 10-5, followed by a Metropolis-Hastinne step accept or reject proposals. The step size is adaptively updated to maintain acceptance rate of approximately 60%, under the assumption of sufficient smoothness in the intermediate densities. Further details of discrete-time AIS are found in Appendix D.3.
Self-improvement. We perform 4 rounds self-improvement. In each round, we spend a portion of the budget to generate 2- 10t samples and reweight using SNIS. The resulting reweighted samples are then used to finetune the model for 250 gradient steps with a batch size of 256. Notably, once the buffer is established, the finetuning does not further expend the allocated budget it does not involve energy evaluations. After the final round, the remaining computational budget allocated to generate final set of 2 10s samples, which are again SNIS reweighted yield the empirical distribution p(zls) for a given system s. We found it beneficial to introduce an L2 regularization term between the log-densities of the current model and "teacher" model initialized from the pre-trained weights.
Density and samples
NF
SNIS
Training samples
Figure 7: Self-improvement procedure. A pre-trained PRose is finetuned at inference-time by iteratively generating samples, reweighting them using SNIS, and training on the reweighted samples.
E.4 Additional baseline configurations
Warp. The original codebase and model weights of Klein al. [2023b] were sampled using the asymptotically unbiased MCMC variant (with Metropolis-Hasting acceptance).
BioEmu. The inference code of Lewis et al. [2024] was employed. BioEmu does not directly model all-atom resolution hence hpacker is employed to introduce the side chains before energy minimization. The codebase of [Lewis al., 2024] was adapted to use the same Amber14 forcefield as ManyPeptidesMMD, on which the 10 energy evaluation budget was spent on per-sample minimization. We additionally experimented with equilibration target the larger 10 energy budget but could not achieve superior results to minimization-on hence this was not included. The adapted codebase for this baseline is provided https: //tit//gtttttttt cww/ htmm//taasfeeabbllcm samplers/ //ii/BBiommi
UniSim. The UniSim model trained on the PepMD dataset of Yu al. [2025] was evaluated. UniSim applies energy minimization following proposal step; simulation
until the 10 energy budget
24

---

was expended. Increasing the simulation time to include the larger 10t energy evaluation budget was not found to improve performance and was omitted. The adapted codebase for this baseline is provided at https: /ttp://ggtttt//t///
samplers/ unisim.
E.5 Computational resourees
All evaluation experiments are run on heterogeneous cluster of NVIDIA L40S and RTX8000 GPUs. ECNF++ sampling is parallelized across multiple nodes with unique seeds to reduce sequential runtime. All evaluation timings are recorded using NVIDIA L40S GPUs. The sampling time required for 10 samples for each model is presented in Fig. 8.
10a
ECNF+ +
TarFlow
10
GPU Hours
Prose
10-1
10
Sequence Length
Figure 8: Sampling time for 10 samples NVIDIA L40S GPU for Boltzmann generators presented in Table 2.
F
Supplementary experimental results
F.1
Unweighted proposal performance
We compare the performance of ECNF++, TarFlow, and PROSE before and after SNIS (with 10 samples) Table 11. Evidently, the strongest proposal distribution given by ECNF++, with both TarFlow and PROSE having large values of E-W2 due to high-energy samples. Notably, ECNF+ deteriorates on all metrics after reweighting, suggesting error accumulation in the divergence integration. TarFlow and PROSE improve significantly on E-W2 after reweighting; however, only PROSE achieves a reduction macrostructure metrics through reweighting. Whilst the unweighted ECNF++ proposal achieves stronger T-W2 and TICA-W, on dipeptides and tetrapeptides than SNIS-reweigtted PROSE, this must be considered with the higher E-W2 indicating poor local detail.
Table Quantitative results for flows comparing the proposal performance and performance after importance sampling on peptide systems up to 8 residues. SNIS performed with a budget of 10* energy evaluations.
Sequence length
+
2AA pienno
4AA qunco
8AA pi
Model
B-W J
T-W, 1
E-Wa 4
T-W L
TICA- 1
E-W; 1
T-W+1
TICA-
Proposal
1.958
0.143
5.006
0.582
0.335
ECNF+
SNIS
3.470
0.302
10.032
1.121
0.572
Proposal SNIS
>> 10
0.178
10
0.882
0.384
10
2.475 2.733
1.026
TarFlow
0.452
0.193
1,260
0.924
0.492
11.298
1.087
Proposal SNIS
> 10" 0.371
0.261
> 10"
0.916
0.546
> 10"
2.456
1.081
PROSE
0.210
0.932
0.752
0.367
10.038
2.456
0.988
F.2 Additional ablations
We consider three further ablations of the PROSE architecture. Firstly, we replace the adaptive conditioning and transition blocks with (i) deeper transformation blocks in which the 8 transformer layers are increased 20 (ii) wider transformer block in which the dimension is increased from 384 to 576. In both cases the increased depth / width was defined to approximately match the parameter count of PROSE. We additionally ablate the use lookahead conditioning, which the atom token =|i] is conditioned not only on [A,, R2 P,, Lg] but also [Az+1, Ri+1, Pe+1, Le+1]. with conditioning information for this pair of indexes mixed using a small ML.P This was motivated by the observation that given naive conditioning the causal masking implies updates to 2[i] are computed without
25

---

knowledge of [A,, Pr, L.]. Preliminary experiments up to tetrapeptides suggested look-ahead conditioning to be beneficial hence it was included in the final PROSE model.
In Table 12 we present results for these ablation models, trained using the same procedure as PROSE. We observe PROSE marginally outperform both w/o transition variants, while removing the look- ahead conditioning is in fact beneficial, particularly on the shorter sequences. These results invite further research into the optimal allocation of parameters, and advanced conditioning techniques for transferable autoregressive flows.
Table 12: SNIS is performed with a fixed hudget of 2 x 10 energy evaluations.
Soquence longth
2AA otvermo
4AA ct reams
8AA Orogesn
Model
ESS 1
E-W;
7-W6 J
ESS 1
E-Wg
7-W1 J
TICA-W;
ESS
E-W, 4
T-W
TICA-Wy J
PROSE
0.191
0.071
0.349
9.360
0.960
w/o transition (doep)
0.166
0,290
0.154
0.060
0.643
0.613
0.338
0.010
9.257
2.123
0.961
we transition (wide)
0.187
0.291
0.158
0.065
0.637
0.634
0.354
0,010
9.426
2,121
0.988
w/o lookahead
0.212
0.270
0.158
0.074
0.591
0.623
0.372
0.011
9.319
2,070
0,962
F33 Dataset scaling
We explore the effect of dataset size on the performance of PROSE. The full ManyPepti dataset contains 21,700 sequences simulated for 200 ns each. We train a further four models wherein (i) the trajectories are limited to 50 ns and 100 ns respectively, and (ii) the set of sequences is reduced to 25% and 50%. For the 25% and 50% sequence reduction variants sequence reduction is applied non-linearly with simple estimate of relative cost, to avoid excessively removing the relatively cheap simulations of short sequences. In Fig. 9 we present metrics for these models. We note the E-Wa to lack any meaningful trend, and the T-Wa have limited sensitivity to the dataset variation, despite a slight trend towards larger data subsets improving the metrics. Most interesting of these plots is the TTCA-W2 with clear trend indicating the full 200 trajectories to be beneficial, whilst using only 50% of the sequences is in fact superior to the full dataset. We caveat these results with subsampled
the observation that PROSE may indeed be operating in the compute-bound regime; training dataset contains sufficient data for over 10a
to 10 ps per frame the ManyPeptidesMi training iterations (batch size 512) without repeating single data sample, in excess of the training budget of 5 10a permitted to the PROSE models in this paper.
E-Wi
T-W2
TICA-N W2
9.708
2.085
0.988
9.574
2.082
0.981
9.441
0.975
9.308
2.078
0.969
50ns
100ns
200ns
S0ns
1 100ns
200ns
50ns
100ns
200ns
Training trajectory lengths
E-W2
7-W2
TICA- W22
10.011
2.088
0.982
9.885
2,082
0.968
9.760
2.075
0.953
25%
50%
100%
25%
50%
100%
25%
50%
100%
Training sequences used
Figure Wasserstein-distance metrics for PROSE trained on variants of ManyPeptidesMDD. Upper row: ManyPep- tidesMD contains 200 ns trajectories for all training sequences, we train on a variant with 50 ns and 100 ns respectively. Lower: ManyPeptidesMID contains 21,700 sequences, train using 25% and 50% of the total. Evaluation metrics computed using SNIS 30 octapeptide sequences.
26

---

F.4 Sequence-length extrapolation
Recall that PROSE is trained exclusively on peptide sequences up length eight. We explore the capacity of the model to generalize in sequence length beyond its training distribution by evaluating on the nine-residue sequence YQNPDGSQA described by Wu and Brooks [2004] and the well-studied ten-residue small protein Chignolin GYDPETGTWG [Honda et al., 2004). We additionally evaluate BioEmu and UniSim baselines on these systems.
Table 13: Comparison of metrics for YQNPDGSQA (9AA) and Chignolin GYDPETGTWG (10AA).
YQNPDGSQA (9AA)
Chignolin (I0AA)
ESS
E-W,
T-W,
TICA-Wa
ESS
E-W;
T-W,
TICA-W3
UniSim
10
6.00
0.84
267.68
6.48
0.20
BioEmu
160.52
4.52
1.14
198.90
5.14
0.65
PROSE PROSE
Proposal
10
3,91
1.65
10
3.63
0.96
SNIS
0.0049
23.79
3.85
1.94
0.0001
832.59
4.35
1.25
PROSE (self-improve)
Proposal SNIS
10
3.73
1.95
>> 100
3.94
1.13
PROSE (self-improve)
0.0123
18.85
3.79
1.95
0,0002
275.87
4.43
1.20
E,5 JSD evaluation
To further assess distributional alignment between generated samples and samples from the reference MD, we follow Raja et al. [2025] and compute the TICA projections and backbone torsion angles; for more details on this metric see Appendix E.2
divergence (ISD) across both
Table Quantitative results flows comparing JSD performance on TICA projections and torus angles before and after importance sampling. SNIS performed with a budget of 10 energy evaluations.
Sequence length
2AA (10mmamm))
4AA (8lygtenn)
8AA (Mayssm
Model
T-k-JSD 4
T-k-JSD 4 TICA-k-ISD 1 T-k-JSD TICA-k-ISD J
Time Warp
0,280
0.460
0.415
BioEmu
0.329
0.245
0.315
0.371 0.879
0.403
UniSim
0.381
0.586
0.376
0.609
Proposal
0.007
ECNF
SNIS
0.031
Proposal
0.002
0.004
0..04
ECNF++
SNIS
0.020
0.051
0..522
Proposal
0.006
0.034
0.017
0.104
0.098
TarFlow
SNIS
0.005
0.022
0.019
0.139
0.124
Proposal SNIS
0.006
0.027
0..23
0.095
01.091
PROSE
0.004
0.01
0..099
0.109
0..82
F.6
Effective sample per second
We report the effective sample size per second (ESS/s) for PROSE and baseline Boltzmann generators, evaluated using SNIS with 10 energy evaluations on an NVIDIA L40s GPU.
Table 1.: Effective sample size per second (ESSIs)
2AA
4AA
8AA
ECNF
1.59 10 10-2
ECNF++
4.76
10-3
2.78
10-4
TarFlow
1.48 10
2.50 10
1.31 10a
PROSE
7.58 10
1.64 10
9.86 10
27

---

E.7
Octapeptide Ramachandran plots
Reference MD
Proposal
SNIS
-2
Figure Ramachandran plots DGVAHALS unseen octapeptide system. Reference molecular dynamics (left column), PROSE proposal (center column), PROSE SNIS with 10 samples (right column).
28

---

F.8
Temperature plots
We present TICA plots in Fig. and energy distributions in Fig. 12 across range of temperatures (310K, 393K, 498K, 631K, 800K). At each temperature, we generate 2 10 samples by scaling the prior with the inverse temperature 8, sampling from (O, 1) For SNIS, we use the energy the corresponding temperature to reweight the samples.
T-310K
T 393K
T-498K
T- 631K
T-800K
True
TIG,
Proposal TIC
Resampled
TIG:
TICa
T1Cb
TICo
TICa
TIC:
Figure 11: TICA plots for RLMM at different temperatures. Reference molecular dynamics (top row). PROSE proposal (middle row), PROSE SNIS (bottom row) with 5 10 samples.
T=310K
T=393K
T=498K
T=631K
T=800K
WD
Propocs
0.03
5NI5
0.03
0.03
0.03
0.03
0.02
0.02
0.02
0.02
600
4.00
200
400
200
300
2.00
100
200
100
150
100
30
E(x)
E00
EGK
Figure 12: Energy histogram plots for RLMM different temperatures.
29

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Yes, the claims are supported in Section and Section 4.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Yes, we discuss limitations in Section 6.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
30

---

Justification:
The paper does not include theoretical results as it is an empirical study.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We disclose the full details in Section 4 and within the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
31

---

Answer: [Yes]
Justification:
We provide open access to data and code in Section 1.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. cc/ Sttpc//ggiees//ooeuuuee
wwwssssssspoolileooeooo ooiicylcyyy for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / //nips. wtt.cc/ppuu. wtte//ublccggggiileeeeeee
Sttpa///ooeeSbiiiieeee wwwsssinooliooooooo for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Yes, we detail the experimental settings and hyperparameters in Section 4 and in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Computing error metrics would be computationally prohibitive, given the large number of evaluation systems employed.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
32

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπхо rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Yes, details are provided in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: pestps:/nnerrae httpppp....p.ppcc/cccieeeeeeeeee
Answer: [Yes]
Justification:
We follow the the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This paper presents work whose goal is to advance sampling methods for sampling from the Boltzmann distribution. There may be potential societal consequences of our work, but none that we feel must be specifically acknowledged.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
33

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e-g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI...
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: This paper presents work whose goal is advance sampling methods for sampling from the Boltzmann distribution. There may be potential risks for misuse of our work, but none that feel must be specifically safeguarded.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
We appropriately cite all creators and original owners of work referenced in this manuscript, and display the original license for any third-party code employed in our codebase.
Justification:
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswi wwwode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
34

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: Yes, the work introduces new dataset and codebase, both with accompanying documentation.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moddl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
35

---

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods in this research? Note that the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality the research, declaration is not required.
Answer: [NA]
Justification:
The core method development does not involve LLMs.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non non-standard components.
Please refer to our LLM policy (https / htt.pppp..//cooofeeece////////// for what should or should not be described.
36


Reviews:
[{"review_id": "ZdcYyDA4Vp", "reviewer": "Reviewer_VbSw", "summary": "This paper addresses the challenge of efficiently sampling molecular conformations. It introduces ENSEMBLE, a large-scale, all-atom transferable normalizing flow model trained on molecular dynamics data of peptides ranging from 2 to 8 residues. ENSEMBLE enables fast, zero-shot sampling across different peptide sequences and lengths, surpassing traditional methods. ENSEMBLE is demonstrated through extensive experiments and various systems, while the trajectory data is also provided.", "questions": "Most of my questions reading the papers has been resolved in the supplementary tool, I thank the authors for the detailed appendix.", "limitations": "yes,  the authors adequately addressed the limitations and potential negative societal impact of their work", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Strengths**\n\n1. **(Significance) Scalable and transferable approach**\n    \n    Unlike prior works that focus on small systems [1] or transferability within fixed-length peptides like dipeptides [2], the proposed method scales up to peptides of up to 8 residues and demonstrates transferability across varying lengths. This represents a meaningful step forward in general-purpose, all-atom peptide modeling.\n    \n2. **(Clarity) Clear and well-organized presentation**\n    \n    The paper is well structured and easy to follow. The inclusion of comprehensive experiments across multiple peptide systems, as well as ablation studies, effectively supports the authors’ claims.\n    \n\n**Weaknesses**\n\n1. **(Quality, minor) Performance on 2-residue (2AA) systems**\n    \n    As acknowledged by the authors, the method underperforms on 2AA systems compared to other lengths. While this is a minor issue and does not detract significantly from the overall contributions, it may indicate areas where the model’s generalization could be improved.\n    \n2. **(Quality, minor) Efficiency claim could be better quantified**\n    \n    While the method outperforms baselines under fixed energy evaluation budgets, adding a simple experiment showing performance as a function of energy evaluations (e.g., performance-vs-budget curves on one system) would strengthen the claim of efficiency and allow for better interpretability of trade-offs.\n    \n    **Minor issues**\n    \n- There is a missing parenthesis “(” after “4AA” in the first line of Table 3.\n- Bold formatting appears to be missing for “TarFlow base” and “backbone permutation” in Table 3, which may reduce clarity.", "was_revised": true, "final_justification": "My minor concerns have been resolved, and I keep my score accordingly.", "rebuttal": "# Response to Reviewer VbSw\n\nWe thank the reviewer for their detailed feedback and their time spent. We appreciate their positive evaluation, in particular, the reviewer describes the proposed approach as \"scalable and transferable\", the empirical study as \"extensive\", and the presentation as \"clear and well-organized\". In what follows, we provide requested clarifications.\n\n> (Quality, minor) Performance on 2-residue (2AA) systems. As acknowledged by the authors, the method underperforms on 2AA systems compared to other lengths. While this is a minor issue and does not detract significantly from the overall contributions, it may indicate areas where the model's generalization could be improved.\n\nThe fact that Ensemble slightly underperforms on 2-residue systems is the consequence of ECNF being trained only on 2 residues. Indeed, ECNF++ -- our attempt to scale this approach to 4 residues (following the improved training techniques from [1]), demonstrates worse performance across both 2 and 4 residues.\n\n> (Quality, minor) Efficiency claim could be better quantified. While the method outperforms baselines under fixed energy evaluation budgets, adding a simple experiment showing performance as a function of energy evaluations (e.g., performance-vs-budget curves on one system) would strengthen the claim of efficiency and allow for better interpretability of trade-offs\n\nIn Fig. 1 (page 2), we report the performance (across 30 unseen 4-residue systems) as a function of energy evaluations and GPU walltime. Notably, the proposed method outperforms Molecular Dynamics (MD) simulations for the same evaluation budget, as the latter fails to capture all metastable states (see the 3rd plot in the top row).\n\n> Minor issues\n\nWe thank the reviewer for a thorough reading of the manuscript. We will fix those typos.\n\nWe thank the reviewer again for the detailed feedback and constructive suggestions on the improvement of the manuscript. We hope we address their main concerns, and we remain available throughout the discussion period for further clarifications and suggestions for improvement.\n  \n[1] Tan, Charlie B., Avishek Joey Bose, Chen Lin, Leon Klein, Michael M. Bronstein, and Alexander Tong. \"Scalable equilibrium sampling with sequential boltzmann generators.\" arXiv preprint arXiv:2502.18462 (2025).**"}, {"review_id": "z2G45vQRt1", "reviewer": "Reviewer_EUPF", "summary": "This paper addresses the sampling of molecular conformations at equilibrium, and notably in a computationally efficient manner that does not require repeated sampling (e.g. MCMC or running molecular dynamics). The authors propose a normalizing flow generative architecture, building from the recent TarFlow, that notably now operates at a scale and performance that generalizes to different molecules and systems. Evaluations are conducted on peptide systems up to octapeptides.", "questions": "1. What is the impact of each of the proposed modifications for the resulting architecture? Which components had the most impact for the empirical results?\n2. What practical obstacles arise from either extending to or evaluating on systems beyond octopeptides?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Writing and presentation in general is quite clear, and the order is logical. The empirical work is strong: improvements from the relevant baselines are notable, multiple sampling schemes are analyzed. Metrics for results are modern and appropriate; compute budgets and hyper-parameters are documented and reproducible. Generalization on up to length 8 peptide systems is also significant.\n\nThe primary concern for the work is academic novelty compared to previous methods like TarFlow. More detailed experimental analysis on all of the modifications in Section 3.1 and each of their roles would help further readers to understand the roles of each individual component of the proposed architecture.", "was_revised": true, "final_justification": "While empirical in nature, the paper does present a notable empirical jump in learning-based conformation sampling, and the authors have done a good job with comparisons and ablations for future readers to understand the true role of their work.", "rebuttal": "# Response to Reviewer EUPF\n\nWe thank the reviewer for their time and effort in reviewing our work. We are glad that the reviewer describes the scale and generalization to other molecular systems as \"notable\", the empirical evaluation as \"strong\" and \"significant\", and the presentation as \"clear\". In what follows, we address the concerns raised and provide suggested evaluations.\n\n> More detailed experimental analysis on all of the modifications in Section 3.1 and each of their roles would help further readers to understand the roles of each individual component of the proposed architecture. What is the impact of each of the proposed modifications for the resulting architecture? Which components had the most impact for the empirical results?\n\nIn Section 4.3, we describe the ablation study of all the architectural changes introduced in our work. In particular, Table 3 compares the results of the original TarFlow architecture with the proposed architecture that includes the \"adapt + transition\" blocks as well as the proposed chemistry-aware permutations. We observe that both changes provide a performance boost; notable in particular given the chemistry-aware permutations imply no increase in training or inference wall-time. The walltime increase of \"adapt + transition\" is only pronounced on the smaller systems and is not significant on the 8AA dataset (Figure 22, supplementary material).\n\n> What practical obstacles arise from either extending to or evaluating on systems beyond octopeptides?\n\nFollowing the reviewer's suggestion, in the table below, we report the performance of our method on the 9 residue peptide introduced by [1] and the well-studied small protein Chignolin. The reported metrics demonstrate that the proposed method outperforms all of the baselines. Note that the training data only contains systems of length up to 8 amino acids, which suggests that the method generalizes across systems of previously unseen length. \n\nWe further explore an additional step in the self-refinement procedure for Chignolin, where an initial set of 1000 samples from Ensemble undergo constrained energy minimization with restraints on the macrostructure. The resulting configurations are then used to finetune our model before running the self-refinement training. \n\n|  | YQNPDGSQA (9AA) |  |  |  | Chignolin (10AA) |  |  |  |\n| :- | - | :- | :- | :- | - | :- | :- | :- |\n|  | ESS | E-W2 | Torus-W2 | TICA-W2 | ESS | E-W2 | Torus-W2 | TICA-W2 |\n| Ensemble (SNIS) | 0.0062 | 24.78 | 3.73 | 1.39 | 0.0001 | 829.59 | 4.21 | 0.73 |\n| Ensemble (self-refine + SNIS) | 0.0082 | 21.34 | 3.77 | 1.39 | 0.0003 | 447.72 | 4.64 | 1.002 |\n| Ensemble + constrained energy min + self-refine | \\- | \\- | \\- | \\- | 0.0002 | 407.13 | 4.52 | 0.81 |\n| BioEmu | - | 1629.55 | 3.85 | 1.03 | - | 9.74e9 | 3.74 | 0.89 |\n| UniSim | - | 1630.42 | 5.65 | 1.05 | - | 3.21e9 | 6.06 | 1.26 |\n\nWe thank the reviewer again for the detailed feedback and constructive suggestions on the improvement of the manuscript. We hope we address their main concerns, and we remain available throughout the discussion period for further clarifications and suggestions for improvement.\n\n[1] Xiongwu Wu and Bernard R Brooks. \"Beta-hairpin folding mechanism of a nine-residue peptide revealed from molecular dynamics simulations in explicit water.\" Biophys J. (2004)"}, {"review_id": "5qlOus1FPY", "reviewer": "Reviewer_7hoH", "summary": "The paper demonstrates that transformer-based normalizing flows, trained on MD simulations of thousands of peptides, can generalize to new peptides and serve as exact-likelihood importance samplers of their Boltzmann distributions.", "questions": "See above.", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Strengths**\nThe paper is the first demonstration of a Boltzmann generator that is transferable, evaluated on unseen peptides, and has fast and tractable likelihoods so as to permit efficient importance sampling. This will stand as a important milestone in the field.\n\n**Weaknesses**\nWhile the result is conceptually significant, the paper leaves a lot of questions unanswered and it is hard to judge the performance of the method in a rigorous context. It is the omission of these details in itself that detracts from the paper, not the results - the authors should optimize for transparency rather than hiding results that qualify the significance.\n* Figure 1. If test peptides have been simulated for up to 1us, then the MD results should be available for wall-clock times up to O(10) GPU-hrs, but authors have concealed this part of the curve.\n* It is very hard to conceptualize TICA and T-W2s. Authors should use JSDs or build markov state models and report recovery of markov states.\n* With very small ESS, it is unclear if SNIS is effective at estimating observables better than the unweighted proposal with small sample sizes. Authors should also report results without reweighting.\n* It seems that the model still generates many high-energy samples. Authors should report statistics - what are the mean, median, 95th percentile energies of the unweighted samples? How does this change from tetrapeptides to octopeptides? This is important to understand the further scalability of the method.\n* It would be helpful to understand the impact of scaling the training data from 1k to 10k training peptides. There is a missed opportunity to explore scaling laws wrt both longer simulations and more training peptides.\n* Authors should show results on chignolin - the only peptide of nontrivial length considered to have been simulated to convergence - since it is a common reference point with many other works.", "was_revised": true, "final_justification": "Authors provided new experimental details to transparently convey the performance of the method.", "rebuttal": "# Response to Reviewer 7hoH\n\nWe thank the reviewer for their time and effort in reviewing our work. We are pleased that the reviewer describes this paper as \"the first demonstration of a transferable Boltzmann generator\" and as \"an important milestone in the field.\" In what follows, we address the concerns raised and provide suggested evaluations.  \n\n> Figure 1. If test peptides have been simulated for up to 1us, then the MD results should be available for wall-clock times up to O(10) GPU-hrs, but authors have concealed this part of the curve.\n\nFigure 1 does not contain MD trajectories up to 1$\\mu$s; instead, the baseline MD trajectories there are simulated up to a fixed energy evaluation budget ($10^6$ energy evaluations). Note that simulating 1$\\mu$s would require $10^9$ energy evaluations. The long MD simulations (1$\\mu$s long) are used only as ground truth samples for metrics evaluation, and use a less frequent frame storing interval to prevent excessive data accumulation (making them unsuitable for use as a baseline across many orders of magnitude of energy evaluation budget). To avoid potential confusion, in the text, we will refer to different types of MD trajectories as train MD, eval MD, and ground true MD, which are used, respectively, as training data, as a baseline, and as the ground true distribution of samples.\n\nHowever, we agree with the reviewer that the \"stopping\" of the MD baseline in Figure 1 is a missed opportunity for comparison between our method and MD at larger wall-time budgets. Hence we have rerun the MD baseline using a decaying frame storing interval up to 1e8 energy evaluations. Given we cannot communicate an updated figure, in the below table we present the metrics at various wall time budgets, in comparison to ENSEMBLE, where a spline was fit to enable arbitrary wall times to be queried.\n\n| L40 Hours |      E-W2      |           |     Torus-W2    |           |     TICA-W2    |           |\n|--|--|--|--|-- |-- |-- |\n|           | MD Baseline    | Ensemble  | MD Baseline     | Ensemble  | MD Baseline    | Ensemble  |\n| 4.41e-4 | 6.764          | 4.400       | 2.734           | 3.252     | 1.192          | 1.025     |\n| 3.09e-3  | 2.490           | 2.585     | 2.044           | 2.663     | 1.238          | 0.769     |\n| 2.95e-2 | 1.445          | 1.356     | 1.605           | 2.057     | 1.166          | 0.571     |\n| 2.94e-1  | 0.861          | 0.964     | 1.556           | 1.672     | 0.983          | 0.494     |\n| 2.93e0 | 0.611          | 1.323     | 1.56            | 1.673     | 0.825          | 0.558     |\n\n> It is very hard to conceptualize TICA and T-W2s. Authors should use JSDs or build markov state models and report recovery of markov states.\n\nFollowing the reviewer's suggestion, we're currently running the evaluation using JSD and will share the updated results during the discussion period.\n\n> With very small ESS, it is unclear if SNIS is effective at estimating observables better than the unweighted proposal with small sample sizes. Authors should also report results without reweighting.\n\nThe suggested comparison is provided in Table 10 (Appendix C.1). This comparison demonstrates that the proposal samples multiple outliers (unphysical structures having very high values of energy), which are efficiently filtered out via SNIS.\n\n> It seems that the model still generates many high-energy samples. Authors should report statistics - what are the mean, median, 95th percentile energies of the unweighted samples? How does this change from tetrapeptides to octopeptides? This is important to understand the further scalability of the method.\n\nWe thank the reviewer for the excellent point. We provide the requested statistics for the unweighted proposal samples for sequence lengths of 4 and 8.\n\n| seq_length | mean          | median        | percentile_95  |\n|------------|---------------|---------------|----------------|\n| 4          | 8.337811e+11  | 2.436593e+05  | 2.531143e+11   |\n| 8          | 1.154441e+18  | 9.880774e+13  | 2.782726e+18   |\n\nThe results highlight a notable increase in energy with longer sequences, reflecting the greater complexity of the larger system. We believe this further motivates the need for reweighting strategies at higher sequence lengths.\n\n> It would be helpful to understand the impact of scaling the training data from 1k to 10k training peptides. There is a missed opportunity to explore scaling laws wrt both longer simulations and more training peptides.\n\nWe thank the reviewer for their suggestion. We agree that this could be quite useful for future work. Given the computational requirement of proper scaling plots in this domain, we will add this experiment in the camera-ready version of the paper.\n\n> Authors should show results on chignolin - the only peptide of nontrivial length considered to have been simulated to convergence - since it is a common reference point with many other works.\n\nFollowing the reviewer's suggestion, in the table below, we report the performance of the method on the 9-residue peptide YQNPDGSQA introduced by [1] and Chignolin. The reported metrics demonstrate that the proposed method outperforms all of the baselines. Note that the training data only contains systems of length up to 8 amino acids, which suggests that the method generalizes across systems of previously unseen length. \n\n|  | YQNPDGSQA (9AA) |  |  |  | Chignolin (10AA) |  |  |  |\n| :- | - | :- | :- | :- | - | :- | :- | :- |\n|  | ESS | E-W2 | Torus-W2 | TICA-W2 | ESS | E-W2 | Torus-W2 | TICA-W2 |\n| Ensemble (SNIS) | 0.0062 | 24.78 | 3.73 | 1.39 | 0.0001 | 829.59 | 4.21 | 0.73 |\n| Ensemble (self-refine + SNIS) | 0.0082 | 21.34 | 3.77 | 1.39 | 0.0003 | 447.72 | 4.64 | 1.002 |\n| Ensemble + constrained energy min + self-refine | \\- | \\- | \\- | \\- | 0.0002 | 407.13 | 4.52 | 0.81 |\n| BioEmu | - | 1629.55 | 3.85 | 1.03 | - | 9.74e9 | 3.74 | 0.89 |\n| UniSim | - | 1630.42 | 5.65 | 1.05 | - | 3.21e9 | 6.06 | 1.26 |\n\nWe thank the reviewer again for the detailed feedback and constructive suggestions on the improvement of the manuscript. We hope we address their main concerns, and we remain available throughout the discussion period for further clarifications and suggestions for improvement."}, {"review_id": "1unyZa5BVA", "reviewer": "Reviewer_qKai", "summary": "This work introduces ENSEMBLE, an all-atom normalizing flow model designed to be transferable across peptide sequences of varying lengths. The architecture builds upon TarFlow, with key modifications to enable dimensional transferability, adaptive system conditioning, and chemistry-aware sequence permutations. With extensive experiments on a newly curated dataset comprising molecular dynamics trajectories of approximately 16,000 peptide systems, ENSEMBLE consistently outperforms existing deep learning baselines on efficiency and efficacy, and exhibits strong capability for accurate sampling with Monte Carlo algorithms.", "questions": "1. In Section 3.1, the description of the model architecture is rather vague. In particular, the concepts of transferability across system dimensions and adaptive system conditioning require clearer and more precise explanations. It would be helpful if the authors could elaborate on these design components in more concrete terms, possibly with supporting equations or schematic illustrations to clarify how they are implemented within the model.\n2. The main experimental results presented in Table 2 are not sufficiently comprehensive. It is unclear why the results of ECNF on the 4AA and 8AA test sets, as well as ECNF++ on the 8AA test set, are missing. Even if the model is not inherently adaptable to variable sequence lengths, it is still possible to train it from scratch on datasets with fixed sequence lengths to obtain corresponding evaluation results. This would allow for a fairer and more complete comparison.\n3. The authors claim that ENSEMBLE achieves state-of-the-art performance on unseen peptide systems. However, many recent deep learning methods developed for molecular simulation and conformation sampling, such as Timewarp [1], ConfDiff [2], MDGEN [3], AlphaFlow [4], UniSim [5], etc., are not considered. A comparison with these methods is necessary to substantiate the claim.\n4. Table 8 shows that ENSEMBLE has significantly more model parameters than the other baselines. This raises the question of whether the observed performance gains are due to increased model size rather than architectural improvements. A more fair comparison, ideally under similar model capacities, would help clarify this point.\n5. Table 2 uses ESS as the evaluation metric for sampling efficiency. However, ESS alone does not fully capture inference efficiency, as higher ESS may come at the cost of significantly longer inference time. Following the practice in [1], please report the ESS per second (ESS/s) for each method (including MD trajectories), that is, ESS divided by the wall-clock time required for inference. This would provide a more comprehensive and fair comparison of sampling efficiency across methods.\n6. When applying importance sampling using Equation (6), the method requires computing the unnormalized Boltzmann density $p(x)$. Could the authors clarify which force field was used in this computation? Furthermore, please discuss whether the evaluation of $p(x)$ significantly impacts the model’s inference speed. It would be helpful to include relevant statistics to quantify this overhead.\n7. It would significantly strengthen the paper to include zero-shot or fine-tuning results of ENSEMBLE on real well-studied protein systems, such as Chignolin. If available, including these results would greatly enhance the practical value and impact of the work.\n\n**Reference**\n> [1] Klein, L., Foong, A., Fjelde, T., Mlodozeniec, B., Brockschmidt, M., Nowozin, S., ... & Tomioka, R. (2023). Timewarp: Transferable acceleration of molecular dynamics by learning time-coarsened dynamics. Advances in Neural Information Processing Systems, 36, 52863-52883.\n\n> [2] Wang, L., Shen, Y., Wang, Y., Yuan, H., Wu, Y., & Gu, Q. Protein Conformation Generation via Force-Guided SE (3) Diffusion Models. In Forty-first International Conference on Machine Learning.\n\n> [3] Jing, B., Stark, H., Jaakkola, T., & Berger, B. Generative Modeling of Molecular Dynamics Trajectories. In The Thirty-eighth Annual Conference on Neural Information Processing Systems.\n\n> [4] Jing, B., Berger, B., & Jaakkola, T. AlphaFold Meets Flow Matching for Generating Protein Ensembles. In Forty-first International Conference on Machine Learning.\n\n> [5] Yu, Z., Huang, W., & Liu, Y. (2025). UniSim: A Unified Simulator for Time-Coarsened Dynamics of Biomolecules. In Forty-second International Conference on Machine Learning.", "limitations": "As noted in the paper, ENSEMBLE still faces challenges when sampling from lower temperatures. Moreover, the lack of validation on larger and well-studied molecular systems undermines the credibility of the practical applicability of the method.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n- Compared to the backbone model TarFlow, ENSEMBLE introduces improvements in modeling peptide systems, particularly in terms of its ability to generalize across sequences of varying lengths.\n- The model supports accurate re-sampling through Monte Carlo algorithms such as Self-Normalized Importance Sampling (SNIS), which also facilitates efficient fine-tuning on previously unseen peptide systems.\n- ENSEMBLE demonstrates strong performance in equilibrium sampling tasks on unseen peptide systems, suggesting good generalization capabilities.\n\n**Weaknesses** \n- In terms of technical novelty, the contributions appear incremental. The method largely builds on existing architectures (e.g., TarFlow) with relatively minor modifications. The degree of innovation may not be sufficient to justify a major advancement without stronger empirical or theoretical justification.\n- The experimental evaluation is insufficient. Some baseline results are incomplete or missing, and the selection of baseline models is limited. A broader and more representative set of comparisons is necessary to convincingly demonstrate the method’s advantages.\n- The manuscript lacks clarity in several sections, particularly in the architectural description of ENSEMBLE. Key design choices and their motivations are not always clearly explained, which may hinder reproducibility and understanding by readers from outside the immediate subfield.\n- While Transformer-based models like ENSEMBLE must explicitly address issues arising from variable sequence lengths, GNN-based methods are inherently adaptable to sequences of arbitrary length. A more rigorous comparison with representative GNN-based approaches would be important to contextualize the claimed benefits of the proposed architecture.", "was_revised": true, "final_justification": "My main concerns were:\n\n- Insufficient baselines: a fair comparison with more recent MD trajectory generation models was needed.\n\n- The generalization ability of ENSEMBLE to larger molecular systems.\n\n- The necessity of using and extending a transformer-based architecture instead of a GNN-based one.\n\nThese points were addressed in a fairly detailed manner in the rebuttal. However, the manuscript still lacks clarity in the description of certain implementation details. I suggest the authors revise Section 3.1 to provide a clearer explanation of the model implementation.  Therefore, I will raise my score to 4.", "rebuttal": "# Response to Reviewer qKai\n\nWe would like to thank the reviewer for the time and effort they spent on reviewing our work. We are appreciative of the fact that the reviewer found ENSEMBLE demonstrates \"strong performance in equilibrium sampling tasks on unseen peptide systems\", and suggests \"good generalization capabilities\". We now address the concerns raised by the reviewer.\n\n## Baselines\n\n> However, many recent deep learning methods developed for molecular simulation and conformation sampling, such as Timewarp [1], ConfDiff [2], MDGEN [3], AlphaFlow [4], UniSim [5], etc., are not considered. A comparison with these methods is necessary to substantiate the claim.\n\nWe thank the reviewer for suggesting these additional baselines on conformation sampling. We include additional baselines in the table below, using the same 1e4 energy evaluation budget as in Table 1 of the main paper.\n\n|           | (2AA)          |             | (4AA)          |             |             | (8AA)         |             |             |\n|-----------|----------------|-------------|----------------|-------------|-------------|---------------|-------------|-------------|\n|           | E-W2           | Torus-W2    | E-W2           | Torus-W2    | TICA-W2     | E-W2          | Torus-W2    | TICA-W2     |\n| Ensemble  | 0.523          | 0.379       | 1.565          | 2.051       | 0.522       | 12.201        | 5.354       | 1.047       |\n| TimeWarp  | 5.765          | 1.515       | 6.436          | 3.448       | 1.203       |       --        |     --        |        --     |\n| UniSim    | 176882.028     | 1.052       | 30498.199      | 2.860        | 2.112       | 36559.357     | 4.813       | 1.704       |\n| BioEmu    | 2622.421       | 0.959       | 519.781        | 2.538       | 1.168       | 159.524       | 4.98        | 1.14        |\n\nIn each case we use the pretrained model weights and sampling configurations provided by the original authors. In the case of Timewarp the model only supports peptides of lengths 2AA and 4AA hence no results are presented for 8AA.  Following the authors procedure, for BioEmu the sidechains are added using hpacker followed by energy minimization. Additionally, for UniSim we employ the energy minimization step during the trajectory generation. In both cases a sweep was performed to optimize the energy minimzation budget allocated per-sample with the number of samples generated. \n\nWe would like to emphasize that only Ensemble (with SNIS), ECNF (with SNIS), and Timewarp (with Metropolis-Hastings acceptance) are able to draw asymptotically unbiased samples from the target distribution. Empirically we observe an a significant improvement of ENSEMBLE over baselines with respect to all metrics, except the macrostructure (Torus-W2 and TICA-W2) metrics on the largest systems (8AA). We however, emphasize that the E-W2 achieved by Ensemble to be significantly lower on the 8AA systems than the next best method.\n\n## Practical applicability \n\n> Moreover, the lack of validation on larger and well-studied molecular systems undermines the credibility of the practical applicability of the method.\n\nWe would like clarify that these are by far the largest systems that have been studied in a transferable setting for Boltzmann generator type models - a 4-fold increase given the only prior transferable Boltzmann generator operated on dipeptides only. We believe our work therefore significantly enhances the practical value and impact of this model class. While Boltzmann Generators as a model class are inherently less scalable than vanilla generative models, they offer the key advantage of exact likelihoods for rigorous statistical re-weighting. We argue that solving the fundamental challenge of length-transferability for peptides up to 8 residues is a critical and necessary step forward, establishing the foundation needed before this rigorous approach can be scaled to larger, canonical protein systems. Our work pushes the practical frontier for this entire class of models, making the future application to larger systems a more tangible research goal.\n\n## Evaluation\n\n> please report the ESS per second (ESS/s) for each method (including MD trajectories) \n\nIn the following table we report effective samples per second for ENSEMBLE and baseline methods. We will report ESS/s for MD trajectories and the additional baseline methods during the discussion period.\n\n|          | 2AA       | 4AA       | 8AA       |\n|--|--|--|--|\n| ECNF     | 3.20E-02  | --         | --         |\n| ECNF++   | 3.89E-03  | 3.17E-04  | --         |\n| TarFlow  | 8.44E+01  | 9.44E+00  | 5.56E-01  |\n| Ensemble | 2.20E+01  | 8.06E+00  | 5.00E-01  |\n  \n> This raises the question of whether the observed performance gains are due to increased model size rather than architectural improvements.\n\nWe thank the reviewer for their concern. We would like to clarify that in our view the model size increase is actually a substantial benefit of ENSEMBLE over previous Boltzmann generator type approaches. Specifically, because our likelihood computation is efficient, we can scale the model significantly more without computational difficulties that might severely limit sampling efficiency. We further draw attention to the sampling wall-time results presented in Figure 22 of the supplementary material, which indicate ENSEMBLE to be significantly faster to sample from than the ECNF, with diminishing increase in wall-time over the standard TarFlow as the system size is increased.\n\n## Chignolin\n\nFollowing the reviewer's suggestion, in the table below, we report the performance of the method on the 9-residue peptide YQNPDGSQA introduced by [1] and the well-studied Chignolin small protein. The reported metrics demonstrate that the proposed method outperforms all of the baselines. We further explore an additional step in the self-refinement procedure where the first round of proposal samples are not reweighed using SNIS, but instead undergo constrained energy minimization before they are used for self-refinement training. The self-refined model is then used to generate proposal samples which are reweighed using SNIS. Note that the training data only contains systems of length up to 8 amino acids, making clear the ability of Ensemble to extrapolate beyond it's training data distribution, despite the baselines being trained on larger systems.\n\n|  | YQNPDGSQA (9AA) |  |  |  | Chignolin (10AA) |  |  |  |\n| :- | - | :- | :- | :- | - | :- | :- | :- |\n|  | ESS | E-W2 | Torus-W2 | TICA-W2 | ESS | E-W2 | Torus-W2 | TICA-W2 |\n| Ensemble (SNIS) | 0.0062 | 24.78 | 3.73 | 1.39 | 0.0001 | 829.59 | 4.21 | 0.73 |\n| Ensemble (self-refine + SNIS) | 0.0082 | 21.34 | 3.77 | 1.39 | 0.0003 | 447.72 | 4.64 | 1.002 |\n| Ensemble + constrained energy min + self-refine | \\- | \\- | \\- | \\- | 0.0002 | 407.13 | 4.52 | 0.81 |\n| BioEmu | - | 1629.55 | 3.85 | 1.03 | - | 9.74e9 | 3.74 | 0.89 |\n| UniSim | - | 1630.42 | 5.65 | 1.05 | - | 3.21e9 | 6.06 | 1.26 |\n  \n## Novelty\n\nWe thank the reviewer for this crucial feedback. We argue our technical novelty lies not in creating a flow from scratch, but in solving a long-standing challenge: **enabling a single, all-atom model with efficient likelihood calculation to transfer zero-shot across variable-length peptides**\n\nThis required novel architectural solutions for **dimensional transferability and adaptive conditioning**, which are non-trivial extensions that allow a Transformer-based flow to generalize across system sizes for the first time. The significance of this breakthrough is validated by our empirical results, where ENSEMBLE unlocks a previously unattainable capability and achieves a >4000x speedup over transferable baselines. We will revise the paper to frame our contribution more clearly around solving this specific, challenging problem.\n\n## Clarity\n\n> Could the authors clarify which forcefield was used in this computation?\n\nThe amber-14 forcefield was used through the OpenMM package. For details on the simulation please see Table 5 in the appendix.\n\nYes, energy computation has a significant impact on the model's inference speed, this is why we set a fixed number of energy evaluations budget across all methods. While we use a relatively efficient energy function here, we note that our method can easily be adapted to any energy function, some of which may be significantly more expensive than even our fairly large ENSEMBLE model to run. Given the favourable performance of ENSEMBLE with respect to energy evaluations demonstrated in Figure 1, any increase in energy function complexity can be anticipated to translate into further-still wall-time performance advantages.\n\n## Concluding remarks\n\nWe thank the reviewer again for their time. We believe we have answered, to the best of our ability, all the great questions raised by the reviewer. We hope our answers allow the reviewer to consider potentially upgrading their score if they see fit. We are also more than happy to answer any further questions.\n\n[1] Xiongwu Wu and Bernard R Brooks. \"Beta-hairpin folding mechanism of a nine-residue peptide revealed from molecular dynamics simulations in explicit water.\" _Biophys J._ (2004)"}]

Author rebuttal:

Meta-review: The paper introduces a normalizing flow framework for sampling conformers, with the aim of improving generalization across molecules. This is an important problem and the present approach is well-grounded, and obtains good empirical results on a peptide dataset. The consensus from reviewers was generally positive reviewers highlighting both the potential impact of improved sampling methods and the practicality of the proposed framework. While some weaknesses about novelty and experimental validation were raised, the positives outweigh the negatives, therefore I recommend acceptance.
Decision: Accept (poster)