Temperature-Annealed Boltzmann Generators
Authors: Henrik Schopmans
Pascal Friederich
Venue: ICML 2025 poster (2025)
Keywords: boltzmann generators
normalizing flows
molecular systems
variational sampling
Primary area: applications->chemistry_physics_and_earth_sciences
TL;DR: We can efficiently sample molecular systems using normalizing flows with reweighting-based temperature-annealing.

Abstract: Efficient sampling of unnormalized probability densities such as the
Boltzmann distribution of molecular systems is a longstanding challenge.
Next to conventional approaches like molecular dynamics or Markov chain
Monte Carlo, variational approaches, such as training normalizing flows with
the reverse Kullback-Leibler divergence, have been introduced. However, such
methods are prone to mode collapse and often do not learn to sample the full
configurational space. Here, we present temperature-annealed Boltzmann
generators (TA-BG) to address this challenge. First, we demonstrate that
training a normalizing flow with the reverse Kullback-Leibler divergence at
high temperatures is possible without mode collapse. Furthermore, we
introduce a reweighting-based training objective to anneal the 
distribution to lower target temperatures.
We apply this methodology to three molecular systems of increasing complexity 
and, compared to the baseline, achieve better results in almost all metrics while requiring up to 
three times fewer target energy evaluations. For the largest system, our approach is the 
only method that accurately resolves the metastable states of the system.

Temperature-Annealed Boltzmann Generators
Henrik Schopmans Pascal Friederich 12
Abstract
advanced methodology, a key factor also lies in the avail- ability of abundant experimental data, such as that in the Protein Data Bank (PDB) (Burley et al., 2021).
Efficient sampling of unnormalized probability densities such as the Boltzmann distribution of molecular systems is a longstanding challenge. Next to conventional approaches like molecular dynamics or Markov chain Monte Carlo, vari- ational approaches, such as training normaliz- ing flows with the reverse Kullback-Leible di- vergence, have been introduced. However, such methods are prone to mode collapse and often do not learn to sample the full configurational space. Here, we present Boltzmann generators (TA-BG) address this challenge. First, we demonstrate that training normalizing flow with the reverse Kullback- Leibler divergence at high temperatures is possi- ble without mode collapse. Furthermore, we intro- duce reweighting- -aaee training objective to an- neal the distribution to lower target temperatures. We apply this methodology three molecular systems of increasing complexity and, compared to the baseline, achieve better results in almost all metrics while requiring up to three times fewer target energy evaluations. For the largest system, our approach is the only method that accurately resolves the metastable states of the system.
However, not all scientific domains benefit from such well- curated and extensive experimental datasets. In areas where data scarcity is a persistent challenge, computational sim- ulations play an essential role. Molecular dynamics (MD) and Markov chain Monte Carlo (MCMC) methods are the primary tools used to explore complex biochemical and physical systems and generate insights from limited experi- mental information. Despite their utility, these classical sam- pling approaches often come with significant computational costs, as they rely on iterative trajectory- -based exploration of nigh-dimensional state spaces.
ieperaturr-annnalle
As result, various approaches have been explored to speed up these methods, including integrating machine learning (MI))-based force fields (Reiser et al., 2022), enhanced sam- pling techniques (Barducci et al., 2011), and data-driven col- lective variables (Bonati et al., 2021). Furthermore, (trans- ferable) generative models have been trained on equilibrium samples from MD simulations (Noé et al., 2019; Mahmoud et al., 2022: Zheng et al., 2024; Klein & Noe, 2024).
While these advancements have significantly improved the efficiency and utility of traditional simulations, there is a growing interest in rethinking the paradigm altogether. Vari- ational sampling methods, rooted in generative modeling, offer a compelling alternative to classical MD and MCMC. These approaches aim learn the underlying probability distribution without the availability of training data, bypass- ing the need for explicit trajectory-based sampling.
1. Introduction
Machine learning, and particularly generative models, have become a transformative force across numerous domains. A prime example of this impact is in structural biology, where deep learning methods such as those in the AlphaFold family (Jumper et al., 2021; Abramson et al., 2024) have revolu- tionized our ability to predict protein structures. While a big part of AlphaFold's 's success can surely be attributed to
The most straightforward variational approach is to train a likelihood- generative model, such as a normalizing flow, using the reverse Kullback-Leibler divergence (KLD), However, this is known to yield mode collapse in many scenarios (Midgley et al., 2023b: Felardos et al., 2023). Recently, multiple variational sampling methods were de- veloped (Blessing et al., 2024), based on normalizing flows (Matthews et al., 2022; Midgley et al., 2023b), diffusion models (Zhang & Chen, 2021; Richter & Berner, 2023; Berner et al., 2023; Vargas et al., 2023; Zhang et al., 2023; Akhound-Sadege et al., 2024; Sendera et al., 2024), and flow matching (Woo & Ahn, 2024).
'Institute of Theoretical Informatics, Karlsruhe Institute of Technology. Kaiserstr. 12. 76131 Karlsruhe, Germany "Institute of Nanotechnology, Karlsruhe Institute of Technology. Kaiserstr. 12, 76131 Karlsruhe, Germany. Correspondence to: Pascal Friederich kit.edu>.
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Temperaturc-Ammealed Boltzmann Generators
Despite their promise to accelerate sampling, the applicabil- ity and scalability of variational sampling methods remain limited, and the field is in its early stages of development compared to the wealth of research on hybrid MD/ML ap- proaches. To the best of our knowledge, the only variational approach that has successfully been applied to the sampling of molecular systems with non-trivial multimodality, such as the popular benchmark system alanine dipeptide, is Flow Annealed Importance Sampling Bootstrap (FAB) (Midgley et al., 2023b).
This essentially makes the high-temperauure simulations help the lower-temperature simulations in overcoming slow energy barriers in the system.
Invernizzi et al. (2022) present variation of replica ex- change molecular dynamics using a normalizing flow that maps from the highest temperature directly to the target tem- perature. This allows direct exchanges and circumvents the need for intermediate replicas in the simulation. However, this still requires performing MD simulations at the bound- ary temperatures and it is not clear how well the method scales to larger systems.
In this work, we propose a novel and scalable flow-based framework to efficiently sample complex molecular sys- tems without mode collapse. We train a normalizing flow at increased temperature using the reverse KLDD, which we show reliably cireumvents mode collapse. Since one is typically interested in the equilibrium distribution at lower temperatures, e.g. at room temperature, we introduce reweighting-baan baasd training objective to iteratively anneal the distribution of the normalizing flow down to the target temperature. We demonstrate the capability and scalability of this methodology using three peptide systems of increas- ing complexity and achieve superior sampling efficiency and accuracy compared to baseline approaches.
Dibak et al. (2022) use a normalizing flow that is trained on samples from high-temperature molecular dynamics sim- ulations. Using special flow architecture, they show that the flow can be adapted to output low-temperature samples, even though it was only trained at the high temperature. Draxler et al. (2024) recently showed that the volume- preserving coupling layers used in that work are not univer- sal, making the approach unsuitable for complex systems.
To solve this issue, Schebek et al. (2024) propose to use normalizing flow with explicit conditioning on the tem- perature T and pressure P. The prior of the normalizing flow is formed by samples from an MD simulation at ref- erence thermodynamic state (To, Po). The normalizing flow is subsequently trained to be able to sample across a range of thermodynamic states (T, P) using the reverse KLD.
Our contribution is threefold:
We show that, in contrast to current literature, the re- verse KLD is surprisingly powerful at learning the Boltzmann distribution of molecular systems with- out mode collapse, but only at increased temperatures where barriers between different free energy minima are lower and the probability distribution maxima are interconnected.
Similarly, Wahl et al. (2025) train a cempeeatuurecondiiinn flow increased temperature using samples from MD, and the correct temperature scaling lower temperatures is ob- tained by matching the gradient of the unnormalized proba- bility density of the flow with respect to the temperature to the gradient of the known target energy function.
We introduce an iterative reweighting-base training strategy to anneal the flow distribution to arbitrary tar- get temperatures.
So far, all mentioned approaches use (a large amount of) MD samples at at least one temperature, and transfer this to a different (lower) temperature. While a multitude of variational sampling methods that do not rely on samples from MD have been proposed, to the best of our knowledge, the only approach that has so far been successfully applied to non-trivial molecular systems is Flow Annealed Importance Sampling Bootstrap (FAB) by Midgley et al. (2023b).
We introduce two complex molecular systems as new benchmarks that go far beyond the size of the typi- cally used benchmark system alanine dipeptide, and we demonstrate that our approach scales to those sys- tems without mode collapse.
Instead of the reverse KLD, FAB uses the a-divergence with a for energy-based training. The a-divergence is estimated using annealed importance sampling (AIS) (Neal,
2. Related Work
2001) from gx to p& where Hamiltonian Monte Carlo gx (HMC) (Duane et al., 1987) is used to transition between intermediate distributions. While AIS sampling is costly in terms of energy evaluations, it proved effective in learn- ing complex probability distributions. They successfully learned the Boltzmann distribution of alanine dipeptide, common benchmark molecular system, without mode col- lapse. However, how well this method scales to more com-
Leveraging the improved mixinmixig behavior when sam- pling at higher temperatures not a completely novel ap- proach, as it was introduced as accelerated sampling technique for MCMC and MD simulations before. Replica exchange Markov chain Monte Carlo (RE-MCMC) and molecular dynamics (REMD) use multiple parallel trajec- tories (replicas) at different temperatures, while allowing repeated exchanges of configurations between the replicas.
2

---

Temperaturc-Ammealed Boltzmann Generators
(a)
(b)
reweight
M
AA
2 ~ qz
rev. KLD
sample
fwd. KLD
5
Eval.
INN
fwd.. KLD
fwd. KLD
sample
sample
2 2 px- exp (- E (KBT))
ewwight
AA
reweight
(c)
Alanine Dipeptide
Alanine Tetrapeptide
Alanine Hexapoptide
Figure 1. (a) Illustration of Boltzmann generators based on normalizing flows. The goal is to learn the equilibrium Boltzmann distribution of the 3D conformations of molecular system. We facus on data-free training, where only the unnormalized prohability density is known. (b) Illustration of workflow. (1) To avoid mode collapse, we first train the flow high temperature with the the learned distribution is iteratively annealed to the target temperature by sampling at the current temperature, reweighting these samples subsequently performing forward KLD training at this lower temperature. This is repeated multiple times until the desired target temperature is reached, (5) In the end, samples are drawn from the flow for evaluation. (c) Visualization of the three peptide molecular systems used work.
KLD. (2-4) Then,
a lower temperature,
plex systems is unclear, and we will address this in this work when comparing our approach to FAB.
identical (see Figure in the appendix for an illustration), If the elementwise transformation is invertible (monotonic), the whole transformation becomes invertible. Furthermore, the Jacobian matrix of such coupling transform is lower triangular and can be efficiently computed (Durkan al., 2019).
3. Preliminaries
3.1. Normalizing Flows
Normalizing flows latent distribution (z)),, typically Gaussian uniform distribution, which is transformed to the target space using an invertible transformation
Training with Samples. The key property of normalizing flows that the likelihood (Equation directly available, typically the cost of forward pass. This allows data- based maximum likelihood training (forward KLD);
g(2; 0), = f(:; 0) = و 'x;; 0) (Figure la).
The transformed density of the flow can be expressed using change variables formula (Dinh al., 2015):
KL lxx|lxx px log qx (x; dr
(x(x;0) = qz (f(x; 0)) (det Jansz 0f(x; 0)
(1)
=C -C -e-pps [log qz + log |det Ja
with the Jacobian Jm
drT
Training by Energy. Next data-based training, a nor- malizing flow can be trained by energy if only the target density px (a) is known. In case of physical systems, such as the molecules studied in this work, this is the Boltzmann distribution px (z) ~ exp (PE((x))(kBT)) Here, x the 3D configuration the molecule, for example, the Carte- sian coordinates of all atoms, kB the Boltzmann constant,
The most common approach to parameterize the invertible function is to use invertible coupling layers. In each cou- pling layer, the input FT.D is split into two parts End and Id+1:D. The first part is transformed elementwise condi- tioned on the second part, while the second part is kept
3

---

Temperaturc-Ammealed Boltzmann Generators
T' is the temperature, and E(x) is the energy of the given configuration, evaluated either using quantum mechanics, e.g. with density functional theory, or, as in this work, using a parameterized force field. While other lata-free training objectives exist (such as the a = 2 divergence used in FAB), the reverse KLD is the most straightforward objective to fit the distribution of the flow to the target density, using samples from the flow itself:
representation based on bond lengths, angles, and dihedral angles represent the molecular conformations. This in- corporates the symmetries of the potential energy, which is invariant to translations and rotations of the whole molecule. For all experiments, we use a normalizing flow built from 16 monotonic rational-quadratie spline coupling layers (Durkan et al., 2019) with fully connected parameter networks the couplings. Dihedral angles are treated using circular splines (Rezende et al., 2020) to incorporate the correct topology. Details can be found in Sections A and B of the appendix.
KL [qx [px]] = KLo [92]|pz]
(4)
= C - / qz(2) (0: 0) dz
4.2. Temperature-Anneale Boltzmann Generators
(5)
Our approach to learn the Boltzmann distribution of molecu- lar systems can be separated into two phases (see Figure Ib); First, we learn the distribution at a high temperature using the reverse KLD (step 1 in Figure lb). Due to decreased barrier heights, this can be done without mode collapse. Sec- ondly, the distribution is iteratively annealed using the mass- covering forward KLD with aepportacc-samplle datasets to obtain the distribution at the target temperature. We now explain both steps in detail and discuss why they avoid the challenges described above.
= C - Rengz [log (g(2; 0)) + log det Jz.sal]
(6)
Importance Sampling Since normalizing flows provide the likelihood of the generated samples, one can perform importance sampling to the true distribution px using the importance weights wxx))= = ox (m:0)- When estimating an
(x(x)
expectation value of an observable h(x) using samples T, from the flow distribution qx, this offers asymptotically
unbiased estimates (Martino et al., 2017; Noé et al., 2019):
Training by Energy: Avoiding Mode Collapse
w(xx) h (In) N- 1i w(xn)
f h())px(()d (7)
The mode-seeking behavior of the reverse KLD has been discussed and observed in multiple previous publications (Midgley et al., 2023b; Felardos et al., 2023; Soletskyi et al., 2024). Once the flow collapsed to a mode, meaning that some remaining modes of the target distribution are not within the support of the flow, it will generally not escape this collapsed state if the remaining modes are too far sepa- rated from the collapsed mode. This is not surprising, since the reverse KLD is evaluated using an expectation value with samples from the flow distribution qx itself, which will only cover the collapsed modes.
While diffusion models and continuous normalizing flows can provide likelihoods by computing divergences of the involved vector field, this is often prohibitively expensive in practice, already for relatively small systems (Klein & Noe, 2024).
While Equation 7 theoretically allows unbiased estimates, this is limited in practice by the actual overlap between the flow distribution qx and the target distribution px. A helpful measure, here, is the effective sample size (ESS), defined as the number of independent samples needed from the target distribution px to achieve the same variance in estimating expectation values as when using the flow dis- tribution qx (Martino et al., 2017). The reverse ESS is an approximation of the ESS, where samples from the flow are used (see Section E in the appendix). The ESS is thus only estimated within the support of the flow, meaning that high ESS can still be achieved if only parts of the true distribution are covered. Therefore, the reverse ESS value needs to always be interpreted together with other metrics.
At the typical target temperature for molecular systems, i.e. 300 K, modes are too far separated to be successfully covered with the reverse KLD. However, when sampling at increased temperature, the modes become more connected. Eventually, one can use the reverse KLD to efficiently learn the distribution. In this work, we performed the reverse KLD experiments at 1200 K, which allows a relatively small batch size and number of gradient descent steps to be used without mode collapse (see Section F.3 in the appendix for an extended ablation and discussion regarding choosing the starting temperature).
Using any loss function that directly includes the target en- ergy of a molecular system can be challenging. If two atoms overlap sufficiently, the repulsive van der Waals energy di- verges, leading to unstable training. Following previous work (Midgley et al., 2023b), we thus use a regularized energy function for training (see Section C in the appendix).
4. Methods
4.1. Flow Architecture
Analogous to previous works (Midgley et al., 2023b; Schop- mans & Friederich, 2024), we use an internal coordinate
4

---

Temperature-AAmnalll Boltzmann Generators
1200K (GT)
1200K (rev KLD)
882K
556K
408K
300K
300K (GT)
10
free energy Ng
Dipeptide V4
11
d1
Φ1
41
41
d1
@1
Tetrapeptide
10
free energy Kg
6
$1
41
61
61
41
free energy Rg
Hexapeptide
10
d2
d2
42
42
d2
Figure 2. Visualization of the iterative annealing process, showing the free energy F= - In p(ós, thi) of backbone dihedral angles (Ramachandran plots) in each iteration. After learning the distribution at 1200 K using the reverse KLD, the distribution is annealed step by step to the target temperature 300 K. Note that not all annealing iterations are shown. Since the tetrapeptide has three pairs of backbone dihedral angles, and the hexapeptide five, selected pair of dihedral angles for this illustration. The amnealing of all pairs of backbone dihedrals can be found in Figures 13 and the appendix. We used x 10? samples for the Ramachandran plots at 300 and 10 samples for the rest.
This avoids very high values in the loss function and sta- bilizes training. Furthermore, analogous to previous work (Schopmans & Friederich, 2024), found that removing small fraction of the largest energy values in the loss contri- butions of each batch stabilizes training.
According to these importance weights, resample a dataset Wi+1 with replacement from D;
4. Perform forward KLD training on WH+1-
Throughout this annealing workflow, we keep updating the flow parameters, without reinitialization. this way, we anneal the distribution of the flow step by step toward the target temperature Tuager (see steps 2-4 in Figure lb). Since we use the mass-covering forward KLD objective based on importance-samppalia datasets, mode collapse is not problem during the annealing.
Reweighting- Annealing
As explained, the reverse KL.D first step step learn the Boltzmann distribution increased temperature. To obtain the distribution at lower target temperature, here 300 K, we utilize importance sampling (Equation 7). While one could do importance sampling directly from Taun 1200 to Tasga 300 K. this will yield bad overlap and sampling efficiency (see Figure in the appendix).
For experiments, we chose 9 temperature annealing itera- tions. To ensure a similar overlap between two consecutive distributions, we choose the temperatures T; using geomet- rie progression between Tatart and Turga (Sugita & Okamoto, 1999) (see Section F.3 in the appendix for comparison with a linear temperature schedule):
Instead, perform importance sampling using multiple temperatures Tj Tamn: T2, .a.... ......... T;, where Ti 2 T2 2 T3 2 2 Tk. In one annealing iteration, we perform the following steps:
Turgat
1. Sample dataset D, { of N samples z, (x(x;; 0) from the flow the current temperature T.
Trrret
#
T, Tstart
(8)
T
2. Calculate importance weights w(z,) px.r=7;41 (xy) for each sample v, in D, to transition T
Furthermore, we added a final fine-tuning iteration, where we sample 300 K and reweight to 300 K (Tx-1 Tk
ax (x,:0)
5

---

Temperaturc-Amnealed Boltzmann Generators
Turge). Since the temperature is not lowered in this fine- tuning step, the effective sample size of the training dataset is higher, which empirically improves the final metries ob- tained at 300 K. For the hexapeptide system, we added such a fine-tuning iteration with T++1 = Ti after each anneal- ing iteration (see Equation 14 in the appendix). While this increases the total number of target potential energy evalua- tions, it improves the obtained results substantially. For the two less complex systems, intermediate fine-tuning was not necessary. We refer to Section F.3 of the appendix for extended discussion and ablations regarding the fine-tuning iterations.
their challenging potential energy surface, TA-BG can also be applied to other sampling tasks. We refer to the appendix, Section K, for additional experiments on a Gaussian mixture system.
5.1. Baseline Methods
To judge the performance of our approach, we compare it to baseline methods. First, we trained a normalizing flow with the forward KLD using MD data from the target distribution. While this is not a variational sampling approach, it serves as good baseline to show the expressiveness of the flow if data is available. Next, we trained a normalizing flow with the reverse KI.D, targeting the Boltzmann distribution 300 K. As the final and most powerful baseline, we trained Flow Annealed Importance Sampling Bootstrap (FAB) (Midgley et al., 2023b). As already discussed, to the best of our knowledge, this is the only method that so far has shown success without mode collapse on our smallest test system alanine dipeptide. It therefore serves as a strong baseline.
Variations
We note that our buffered reweighting approach is not the only option to anneal the temperature of a normalizing flow. As discussed in Section 2, a concurrent study ( (Wahl et al., 2025) uses a temperature-conditione normalizing flow with a temperature scaling loss to learn the Boltzmann distribu- tion at the target temperature. While this approach achieves promising results, it requires the repeated estimation of the partition function Z using importance sampling with the flow distribution. Obtaining a low-variance estimate of Z can be computationally expensive, especially for high- dimensional systems such as the hexapeptide studied here. However, a systematic comparison of different temperature scaling approaches is an interesting avenue for future work.
5.2. Metrics
To evaluate the distribution of the normalizing flow at the target temperature 300 K, we use combination of multiple metrics. First, we use the negative log-likelihood (NL.)) calculated on the ground truth test dataset. This is good overall measure of the learned distribution. We note that since the metastable regions of our systems form only a small part of the ground truth test datasets, seemingly minor differences in the NLL can be decisive, especially when assessing mode collapse or the quality of the description of the metastable region.
We further experimented with variations of our reweighting approach. For example, we tried training a temperature- conditioned flow with a reweighting-based objective contin- uously on the whole temperature range. This is described in more detail in Section 11 of the appendix. In practice, we found the iterative buffered annealing workflow to be superior, both in terms of accuracy and sampling efficiency.
To better assess potential mode collapse, we additionally evaluate the free energy F = - 7 In p(c,, vhx) of the back- bone dihedral angles (Ramachandran plots). Since these are the main slow degrees of freedom of the peptide sys- tems, mode collapse will be directly visible here. To assess the quality of the Ramachandran plot, we calculate the for- ward KLD between the probability distribution given by the Ramachandran plot of the ground truth and the one of the flow distribution (RAM KLD). Since the tetrapeptide and hexapeptide have multiple pairs of backbone dihedral angles, we report the mean of their RAM KLD values. Fur- thermore, we repeat the calculation of the RAM KLD also using the Ramachandran plots obtained from importance sampling.
5. Experiments
We now describe the conducted experiments to evaluate our cmpperature-annealiin approach. The objective is to learn the Boltzmann distribution of three molecular systems, in- creasing in complexity (see Figure Ic). The first molecule is alanine dipeptide, a popular system that previously served as a benchmark for variational sampling (Midgley et al., 2023b), but also other related tasks such as transition path sampling (Holdijk et al., 2023; Scong et al., 2024). We fur- ther evaluate on two higher- higher-dimensionall and more complex systems, alanine tetrapeptide and alanine hexapeptide. All three systems have complex metastable high-energy regions that make up only small fraction of the entire state space, which makes them suitable hard objectives for benchmark- ing.
Last, to evaluate the sampling efficiency, we report the re- verse effective sample size (ESS) (see Section 3).
6. Results
Figure 2 shows how the Ramachandran plots of each of the three systems are annealed the target temperature.
While the focus of our work is on molecular systems due to
6

---

Temperaturc-Annealed Boltzmann Generators
showing four exemplary steps of the annealing workflow. Both the distribution at 1 1200 K learned with the reverse KLD and the distribution at the target temperature 300 K match the ground truth obtained from MD.
7. Discussion
To summarize, TA-BG achieves better results in most met- ries compared to the baselines and thus learns the Boltzmann distribution of the investigated systems more accurately, while requiring significantly fewer target energy evaluations. We note that force field evaluations in our current setup are relatively inexpensive. With more accurate and com- putationally costly target evaluations, such as ML-based foundation models or density functional theory, the energy evaluation cost becomes dominant. In such cases, the higher sampling efficiency of our method will translate into signif- icantly reduced computational cost. We refer to Section J in the appendix for detailed analysis of the computational cost of TA-BG in comparison with FAB.
We now compare the obtained distribution at 300 K with that from the baseline methods by evaluating the introduced metrics (Table 1). Furthermore, Figure shows the Ra- machandran plots at 300 K obtained by all methods side-by- side. Figure 12 in the appendix shows the same comparison, but with importance sampling to the target distribution.
We start with the smallest system, alanine dipeptide. All methods, except for the reverse KLD, are able to learn the distribution without mode collapse (Figure 3). While the reverse KLD covers the high-energy region to some extent, partial mode collapse is visible. In terms of metrics, FAB and TA-BO obtain comparable results. Our method achieves better NLL and ESS values, while FAB achieves slightly lower RAM KLD values. Our approach only uses approxi- mately a third of the target energy evaluations of FAB. We further note that the metrics we obtained with FAB on ala- nine dipeptide are slightly better but mostly comparable to those in the original publication.
Next, we would like to point out the surprisingly good re- sults we obtained by simply training with the reverse KLD at 300 K for alanine dipeptide (Figure 3). Even though at 300 K the metastable states are only connected to the global minimum by very low-probability transition region, only partial mode collapse is observed. Next to our results of being able to successfully learn the distribution at high temperatures with the reverse KLD, this further shows the effectiveness of simple reverse KLD training.
Similar results can be observed for the tetrapeptide system. The reverse KLD training now collapses almost fully to the main mode, missing most of the metastable region (Fig- ure 3). Both TA-BG and FAB achieve a good match with the ground truth distribution, fully resolving the metastable region. The metrics of TA-BG and FAB are again close, our approach achieves slightly better NLL and RAM KLD values, while FAB has a lower reweighted RAM KLD value and slightly higher ESS. Our approach again uses approxi- mately a third of the target energy evaluations of FAB.
Furthermore, previous work showed that not only training normalizing flow from scratch with the reverse KLD is prone to mode collapse, but also fine-tuning pre-trained flow to target distribution with the reverse KLD typically collapses (Felardos et al., 2023). This is not well under- stood phenomenon and has only recently been investigated theoretically (Soletskyi et al., 2024). While Felardos et al. (2023) proposed a solution to this problem, it is not clear how well it scales larger systems. As described before, after our annealing workflow reaches 300 K, we fine-tune the flow distribution at 300 K by performing forward KLD training with a buffered dataset from the flow distribution, reweighted to the target distribution. This offers a simple yet effective solution to the problem of fine-tuning pre- trained flows and can also be used in other scenarios. For exam- ple, one can train flow on biased MD simulation data that is not properly equilibrated, and then fine-tune with our buffered reweighting approach to obtain an unbiased flow distribution.
In the case of the most complex investigated system, alanine hexapeptide, the distribution of the reverse KLD training again almost entirely collapses. While FAB is partially able to resolve the metastable states, they differ in shape compared to the ground truth distribution. In contrast, our approach covers all metastable states without mode collapse and resolves them accurately with only small imperfections. This is also reflected and quantified by the metrics: Com- pared to FAB, our approach achieves better results in all metrics, while requiring 3.08 x 108 target energy evalua- tions compared 4.2 x 10x used by FAB.
Furthermore, while our method generally yielded better results with fewer target energy evaluations for the hexapep- tide, the results we obtained using FAB are still relatively close in terms of accuracy, especially for the other two sys- tems. This establishes it as powerful method for variational sampling. Therefore, combination of our temperature- annealing approach with FAB is an interesting avenue for future work. Instead of using the reverse KLD for train- ing at increased temperatures, FAB can be used and the distribution can be subsequently annealed with our TA-BG
A tradeoff exists between the accuracy of the obtained dis- tribution and the number of target evaluations. This is es- pecially true for FAB, where the number of intermediate AIS distributions and the number of HMC steps can be var- ied. We present corresponding variations in Table 11 in the appendix. For FAB applied to the hexapeptide, even when using almost 3 times as many target evaluations compared to our approach, we still achieve a lower NLL value.
7

---

Temperature-AAmealee Boltzmann Generators
Ground truth
FWD KLD
TA-BG (ours)
FAB
REV KLD
10
free energy KoT
Dipeptide
0
61
41
Tetrapeptide
10
free energy/ kg
6n
61
Hexapeptide
10
free energy Kg
d2
Figure 3. Comparison of the free energy F - kaT In p( thx) of the backhone dihedral angles (Ramachandran plots) 300 K obtained by the different methods. Since the tetrapeptide has three pains backhone dihedral angles, and hexapeptide five, we selected the pair with the largest visible deviation among the methods. The Ramachandran plots for all pairs of backbone dihedral angles be found in the appendix. Figures 14 and 17.
workflow. If and when this can be better than simple reverse KLD training at high temperatures needs to be investigated.
8. Conclusion
We introduced temperature-anne annealed Boltzmann generators, technique that uses combination of high-temperature reverse KLD pre-training and a subsequent annealing work- flow to efficiently sample the Boltzmann distribution of molecular systems at room temperature. On the molecular systems investigated, our approach achieves better results in almost all metrics, while requiring up to three times fewer target energy evaluations compared to the baselines. Further- more, it the only variational approach that accurately resolved the metastable region of the most complex sys- tem studied, demonstrating its scaling capabilities from toy examples to application-relevan systems.
Our results indicate that other variational sampling ap- proaches, such as those based on diffusion models, might benefit from sampling at higher temperatures and anneal- ing the temperature afterwards. Other objectives will likely also benefit from the reduced mode separation elevated temperatures.
In this work, we used a relatively simple flow architecture based on neural splines and internal coordinates. Recent improvements in flow architectures (Zhai et al., 2024; al., 2025) can be leveraged in the future to improve the obtained results and scale larger systems. Furthermore, the use of internal coordinate representation can be limitation, since it is not transferable between systems and it does not incorporate the permutation invariance of iden- tical particles. A transferable equivariant normalizing flow has recently been proposed (Midgley et al., 2023a). Using the TA-BG methodology
Similar to how replica exchange molecular dynamics is an established method in the toolbox of computational sci- entists,
confident that high-temperature sampling with temperature-annealing a powerful approach that will move the field of variational sampling forward.
transferable setting is an in- teresting avenue for future work, which may also allow hybrid approach where data-free training is combined with Rata-based training (Lewis et al., 2024).
8

---

Temperature-Anmealed Boltzmann Generators
Table Comparison of metrics obtained for all three peptide systems. The table shows the number of potential energy evaluations (PE EVALS), the negative log likelihood (NLL), the reverse effective sample size (ESS), the (mean) Ramachandran KLD (RAM KLD), and the (mean) Ramachandran KLD with reweighting (RAM KLD RW). metrics are reported as mean and standard error obtained from four independent experiments. The best-performing variational method for each metric is highlighted in bold.
SYSTEM
METHOD
PE EVALS
NLL J
ESS 1
RAM KLD J
RAM KLD w. RW +
FORWARD KLD
x 10"
-213.581 # 0.000
(82.16 + 0.09) %
(2.21 # 0.05) x 10"3
(1.99 + 10"3
REVERSE KLD
2.56 x 10a
-213.609 = 0.006
(94.11 = 0.21) %
(1.75 x 10
(1.65 0.29) 10
DIPEPTIDE
FAB
2.13 10
-213,653 = 0.000
(94.81 :: 0.04) %
(1.50 = 0.03) 10
(1,25 + 0.01) 10"3
TA-BG (OURS)
x.56 10"
-213.665 = 0.002
(95.61 + 0.14) %
(1.97 0.08) x 10
(1.39 + 0.03) 10
FORWARD KLD
4.2x10'
-330.069 + 0.001
(45.29 0.08) %
(2.26 10
(2.50 0.03) 10
TETRA- PEPTIDE
REVERSE KLD
2.56 10
-329.191 = 0.122
(74.88 3.65) %
(3.00 10
(2.87 0.440) 10-
FAB
2.13 10
-330.100 # 0.002
(63.59 +: 0.23) %
(6.89 # 0.25) 10
(1.25 # 0.01) x 10
TA-BG (OURS)
7,56
-330.114 0.004
(62.36 + 0.46) %
(2,61 = 0.200) x 10-2
(1.94 0.12) 10-*
FORWARD KLD
x.2x 10"
-501,598 + 0.005
(10.97 + 0.11) %
(4.16 0.26) 10
(7.69 a: 0.03) 10
HEXA- PEPTIDE
REVERSE KLD
2.56 x 10a
497.378 = 0.277
(22.22 # 1.44) %
(5.41 = 0.38) x 10
(5.32 +.388 10-1
FAB
4.2 10"
-501.268 + 0.008
(14.64 0.08) %
(2.09 2 0.02) 10
(1.12 a: 0.02) 10
TA-BG (OURS)
3.08 x 10
-501.511 = 0.013
(14.71 = 0.18) %
(8.63 = 0.66) x 10
(8.77 # 0.47) 10
Software and Data
The source code to reproduce our experiments can found https: : / /gg thub. com/a a imat aab/ Saa//TAA-P TT--/GAAC BG. Fur- thermore, the ground truth datasets from MD simulations are provided at https: /ww./iiooo//// www./0...............o 5281 15526429.
Acknowledgments
The authors would like to thank the anonymous reviewers for their valuable comments and suggestions. H.S. acknowl- edges financial support by the German Research Foundation (DFG) through the Research Training Group 2450 "Tailored Scale-Bridging Approaches to Computational Nanoscience". P.F. acknowledges funding from the Klaus Tschira Stiftung gGmbH (SIMPLAIX project) and the pilot program Core- Informatics of the Helmholtz Association (KiKIT project). Parts of this work were performed on the HoreKa super- computer funded by the Ministry of Science, Research and the Arts Baden-Wiinne Württemberg and by the Federal Ministry of Education and Research. This work is supported by the Helmholtz Association Initiative and Networking Fund on the HAICORE@ KIT partition.
Impact Statement
This paper presents work whose goal advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.

---

Temperaturc-Ammealed Boltzmann Generators
References
for basic and applied research and education in funda- mental biology, biomedicine, biotechnology, bioengineer- ing and energy sciences. Nucleic Acids Research, 49 (D1):D437-D451, January 2021. ISSN 0305-1048. doi: 10. 10.1093/na//kkaa0003
Abramson, J., Adler, J., Dunger, J., Evans, R., Green, T.,
Pritzel, A., Ronneberger, O., Willmore, L.,, Ballard, A. J., Bambrick, J., Bodenstein, S. W., Evans, D. A., Hung. C.-
C., O'Neill, M., Reiman, D., Tunyasuvunakool, K., Wu, Z. Zemgulyte, A., Arvaniti, E., Beattie, C., Bertolli, O., Bridgland, A., Cherepanov, A., Congreve, M., Cowen- Rivers, A. I., Cowie, A., Figurnov, M., Fuchs, F. B., Gladman, H., Jain, R., Khan, Y. A., Low, C. M. R., Per- lin, K. Potapenko, A., Savy, P., Singh, S., Stecula, A., Thillasundaram, A., Tong, C., Yakneen, S., Zhong, E. D., Zielinski, M., Zidek, A., Bapst, V., Kohli, P., Jaderberg, M.. Hassabis, D., and Jumper, J. M. Accurate structure prediction of biomolecular interactions with AlphaFold Nature, 630 43008006):493-5000 June 2024. ISSN 1476-4687. doi: 10. 10.10388/441586--44--0444
Conor Durkan, Artur Bekasov, lain Murray, and George Pupamakarios. Nilows: Normalizing flows in PyTorch, November 2020.
D.A. Case, H.M. Aktulga, K. Belfon, I.Y. Ben-Shalom, J.T. Berryman, S.R. Brozell, D.S. Cerutti, T.E. Cheatham, III, V.W.D. Cruzeiro, T.A. Darden, N. Forouzesh, G. Gi- ambasu, T. Giese, M.K. Gilson, H. Gohlke, A.W. Goetz, J. Harris, Izadi, S.A. Izmailov, K. Kasavajhala, M.C. Kay- mak, E. King, A. Kovalenko, T. Kurtzman, T.S. Lee, P. Li, C. Lin, Liu, T. Luchko, R. Luo, M. Machado, V. Man, M. Manathunga, K.M. Merz, Y. Miao, O. Mikhailovskii, G. Monard, H. Nguyen, K.A. O' O'Hearn, A. Onufriev, F. Pan, S. Pantano, R. Qi. A. Rahnamoun, D.R. Roe, A. Roit- berg, C. Sagui, S. Schott- Verduuo, A. Shajan, Shen, C.L. Simmerling, N.R. Skrynnikov, J. Smith, J. Swails, R.C. Walker, J. Wang, J. Wang, H. Wei, X. Wu, Y. Wu, Y. Xiong, Y. Xue, D.M. York, S. Zhao, Q Zhu, and P.A. Kollman. Amber 2023. University of California, San Francisco, 2023.
Akhound-Sadegh, T., Rector-Brooks, J., Bose, J., Mittal, S., Lemos, P., Liu, C.-H., Sendera, M., Ravanbakhsh, S., Gidel, G., Bengio, Y., Malkin, N., and Tong, A. Iterated Denoising Energy Matching for Sampling from Boltz- mann Densities. In Forty-First International Conference on Machine Learning, June 2024.
Barducci, A., Bonomi, M., and Parrinello, M. Metady- namics. WIREs Computational Molecular Science, 1(5): 2011. ISSN 1759-0884, doi: 10. 1002/wems.31.
Dibak, M., Klein, L., Krämer, A., and Noé, F. Temper- ature steerable flows and Boltzmann generators. Phys. Rev. Res., 444)0004 October 2022. doi: 10.1103/ PhyaRewwesearch.44L0000000
Berner, J., Richter, L., and Ullrich, K. An optimal con- trol perspective on diffusion-based generative modeling. Transactions on Machine Learning Research, October 2023. ISSN 2835-8856.
Dinh, I.,, Krueger, D., and Bengio, Y. NICE: Non-linear Independent Components Estimation, April 2015.
Dinh, I.,, Sohl-Dickstein, J., and Bengio, S. Density estima- tion using Real NVP, February 2017.
Blessing, D., Jia, X., Esslinger, J., Vargas, F., and Neumann, G. Beyond ELBOs: A Large-Scale Evaluation of Varia- tional Methods for Sampling. In Forty-First International Conference on Machine Learning, June 2024.
Draxler, F., Wahl, S., Schnoerr, C., and Koethe, U. On the Universality of Sluume-Prserrvigg and Coupling-Based Normalizing Flows. In Forty-First International Confer- ence on Machine Learning, June 2024.
Bonati, L., Piccini, G., and Parrinello, M. Deep learning the slow modes for rare events sampling. Proceedings of the National Academy of Sciences, (18(44))21113333111 November 2021. doi: 10.1073/pnas.22 13533118.
Duane, S., Kennedy, A. D., Pendleton, B. J., and Roweth, D. Hybrid Monte Carlo. Physics Letters B, 195(2):216- 222, September 1987. ISSN 0370-2693. doi: 10.1016/ 8370-2693887/111977X
Burley, S. K., Bhikadiya, C., Bi, C., Bittrich, S., Chen, L., Crichlow, G. V., Christie, C. H., Dalenberg, K., Di Costanzo, L.. Duarte, M., Dutta, S., Feng, Z., Gane- san, S., Goodsell, D. S., Ghosh, S., Green, R. K., Gura- nović, V., Guzenko, D., Hudson, B. P., Lawson, C. L., Liang, Y., Lowe, R., Namkoong, H., Peisach, E., Per- sikova, I., Randle, C., Rose, A., Rose, Y., Sali, A., Segura, J.. Sekharan, M.. Shao, C., Tao, Y.-P., Voigt, M.. West- brook, J. D., Young, J. Y., Zardecki, C., and Zhuravleva, M. RCSB Protein Data Bank: Powerful new tools for exploring 3D structures of biological macromolecules
Durkan, C., Bekasov, A., Murray, I., and Papamakarios, G. Neural Spline Flows. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Eastman, P., Galvelis, R., Peláez, R. P., Abreu, C. R. A., Farr, S. E., Gallicchio, E., Gorenko, A., Henry, M. M., Hu, F., Huang, J., Krämer, A., Michel, J., Mitchell, J. A., Pande, V. S., Rodrigues, J. P., Rodriguez- Guerraa J., Sim- monett, A. C., Singh, S., Swails, J., Turner, P., Wang,
10

---

Temperature-Anmealel Boltzmann Generators
Y., Zhang, I., Chodera, J. D., De Fabritiis, G., and Mark- land, T. E. OpenMM 8: Molecular Dynamics Simulation with Machine Learning Potentials. J. Phys. Chem. B, 128(1):109-116, January 2024. ISSN 1520-6106. doi: 10. 1a110//acs.ipcc.3c6666.
Mahmoud, A. H., Masters, M., Lee, S. and Lill, M. A. Accurate Sampling of Macromolecular Conformations Using Adaptive Deep Learning and Coarse-Grained Rep- resentation. J. Chem. Inf. Model., 62(7):1602-1661, April 2022. ISSN 1549-9596. doi: 10.102
Felardos, L., Hénin, J., and Charpiat, G. Designing losses for data-free training of normalizing flows on Boltzmann distributions, January 2023.
Martino, L.,, Elvira, V., and Louzada, F. Effective Sample Size for Importance Sampling based discrepancy mea- sures. Signal Processing. 131:386-400 February 2017. ISSN 01651684, doi: 10. siggro.22.....88...0 2016.08......
Goshtasbpour, S., Cohen, V., and Perez-Cruz, F. Adap- tive Annealed Importance Sampling with Constant Rate Progress. In Proceedings of the 40th International Con- ference on Machine Learning, pp. 11642-1111 (1642-1668,, PMLR, July 2023.
Matthews, A., Arbel, M., Rezende, D. J., and Doucet, A. Continual Repeated Annealed Flow Transport Monte Carlo. In Proceedings of the 39th International Con- ference on Machine Learning, pp. 15196-15219, PMLR, June 2022.
Holdijk, L.., Du, Y., Hooft, F., Jaini, P., Ensing, B., and Welling, M. Stochastic Optimal Control for Collective Variable Free Sampling of Molecular Transition Paths. Advances in Neural Information Processing Systems, 36: 79540-79556, December 2023.
Midgley, L. I., Stimper, V., Antorán, J., Mathieu, E., Schölkopf, B., and Hernández-Lohann J. M. SE(3) Equiv- ariant Augmented Coupling Flows. In Thirty-Sevent Conference on Neural Information Processing Systems, August 2023a.
Invernizzi, M., Krämer, A.. Clementi, C., and Noć, F. Skip- ping the Replica Exchange Ladder with Normalizing Flows. J. Phys. Chem. Lett., 13(50): 1643-11166611111 1649, De- cember 2022. doi: 10.1022/aass.... 10.1002//aesspccccttt22032222
Midgley, L. I., Stimper, V., Simm, G. N. C., Schölkopf, B., and Hernández-Lobato, J. M. Flow Annealed Impor- tance Sampling Bootstrap. In The Eleventh International Conference on Learning Representations, 2023b. doi:
Jumper, J., Evans, R., Pritzel, A., Green, T., Figurnov, M.. Ronneberger, O., Tunyasuvunakool, K., Bates, R., Zidek, A., Potapenko, A., Bridgland, A., Meyer, C., Kohl, S. A. A., Ballard, A. J., Cowie, A., Romera-Paredes, B., Nikolov, S., Jain, R., Adler, J., Back, T., Petersen, S., Reiman, D.. Clancy, E., Zielinski, M., Steinegger, M., Pacholska, M., Berghammer, T., Bodenstein, S., Silver, D., Vinyals, O., Senior, A. W., Kavukcuoglu, K., Kohli, P., and Hassabis, D. Highly accurate pro- tein structure prediction with AlphaFold. Nature, 596 (7873):583-589, August 2021. ISSN 1476-4687. doi: 10.1038/41586-021-033311
10-48550/a2Xi.2228..011
Neal, R. M. Annealed importance sampling. Statistics and Computing, 11(2):125-139 April 2001. ISSN 1573-1375 doi: 10. 10.1023/A:10008922155000
Noé, F. Bgflow. A14Science group, FU Berlin (Frank Noé and co-workers), December 2024.
Noé, F., Olsson, S., Köhler, J., and Wu, Boltzmann gen- erators: Sampling equilibrium states of many-body sys- tems with deep learning. Science, 365(6457) 147, September 2019. doi: 10. 1..1126/sience.aam 147.
Kingma, D. P. and Ba, J. Adam: A Method for Stochastic Optimization, January 2017.
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Köpf, A., Yang, E., DeVito, Z., Rai- son, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S. PyTorch: An Imperative Style, figh-Perffrmance Deep Learning Library, Decem- ber 2019.
Klein, L. and Noe, F. Transferable Boltzmann Generators. In The Thirty-eighth Annual Conference on Neural Infor- mation Processing Systems, November 2024.
Lewis, S., Hempel, T., Jiménez-Luna, J., Gastegger, M., Xie, Y., Foong, A. Y. K., Satorras, V. G., Abdin. O., Veeling, B. S., Zaporozhets, I., Chen, Y., Yang, S., Schneuing, A., Nigam, J., Barbero, F., Stimper, V., Campbell, A., Yim, J. Lienen, M., Shi, Y., Zheng, S., Schulz, H., Munir, U., Clementi, C., and Noé, F. Scalable emulation of pro- tein equilibrium ensembles with generative deep learning, December 2024.
Reiser, P., Neubert, M., Eberhard, A., Torresi, L., Zhou, C., Shao, C., Metni, H., van Hoesel, C., Schopmans, H., Sommer, T., and Friederich, P. Graph neural networks for materials science and chemistry. Commun Mater, 3(1): 1-18, November 2022. ISSN 2662-4443. doi: 10.1038/ $43246-022 6--446-02200031556

---

Temperature-AAmealld Boltzmann Generators
Rezende, D. J., Papamakarios, G., Racaniere, S., Albergo, M., Kanwar, G., Shanahan, P., and Cranmer, K. Normal- izing Flows Tori and Spheres. In Proceedings of the 37th International Conference on Machine Learning, PP. 8083-8092. PMLR, November 2020.
Wahl, S., Rousselot, A., Draxler, F., and Koethe, U. TRADE: Transfer of Distributions between External Conditions with Normalizing Flows. In The 28th International Con- ference on Artificial Intelligence and Statistics, February 2025.
Richter, L., and Berner, J. Improved sampling via learned diffusions.. In The Twelfth International Conference Learning Representations, October 2023.
Woo, D. and Ahn, S. Iterated Energy-based Flow Matching for Sampling from Boltzmann Densities, August 2024.
Zhai, S., Zhang, R., Nakkiran, P., Berthelot, D., Gu, J., Zheng, H., Chen, T., Bautista, M. A., Jaitly, N., and Susskind, Normalizing Flows Capable Generative Models, December 2024.
Schebek, M., Invernizzi, M., Noé, F., and Rogal, J. Ef- ficient mapping of phase diagrams with conditional Boltzmann Generators. Mach. Learn.: Sci. Technol., 5(4):045045, November 2024. ISSN 2632-2153. doi: 10.1088/2633--1555/888484
Zhang, D., Chen, R. T. Q., Liu, C.-H., Courville, A., and Bengio, Y. Diffusion Generative Flow Samplers: Im- proving learning signals through partial trajectory opti- mization. In The Twelfth International Conference Learning Representations, October 2023.
Schopmans, H. and Friederich, P. Conditional Normalizing Flows for Active Learning of Coarse- Grasine Molecular Representations. In Forty-First International Conference on Machine Learning, June 2024.
Zhang, Q. and Chen, Y. Path Integral Sampler: A Stochas- tic Control Approach For Sampling. In International Conference on Learning Representations, October 2021.
Sendera, M., Kim, M., Mittal, S., Lemos, P., Scimeca, L., Dector-Brooks, J., Adam, A., Bengio, Y., and Malkin, N. Improved off-policy training of diffusion samplers. In The Thirtt-eighth Annual Conference on Neural Information Processing Systems, November 2024.
Zheng. S., He, J., Liu, C., Shi, Y., Lu, Z., Feng, W., Ju, F,, Wang. J., Zhu, J., Min, Y., Zhang, H.. Tang. S., Hao, H., Jin, P., Chen, C., Noé, F., Liu, H., and Liu, T.-Y. Predict- ing equilibrium distributions for molecular systems with deep learning. Nat Mach Intell, 665):558-567, May 2024. ISSN 2522-5839. doi:
Seong, K., Park, S., Kim, S., Kim, W. Y., and Ahn, S. Transi- tion Path Sampling with Improved Off-Policy Training of Diffusion Path Samplers. In The Thirteenth International Conference on Learning Representations, October 2024.
Soletskyi, R., Gabrié, M., and Loureiro, B. A theoretical perspective mode collapse in variational inference, October 2024.
Stimper, V., Midgley, L. I., Simm, G. N. C., Schölkopf, B., and Hernández-LLobato, J. M. Alanine dipeptide in implicit solvent 300K. August 2022. doi: 10.5281/ 1enodo.6993124.
Sugita, Y. and Okamoto, Replica-exchange molecular dynamics method for protein folding. Chemical Physics Letters, 314(1): November 1999. ISSN 0009- 2614. doi: 22.006680009-221199900111111
Tan, C. B., Bose, J., Lin, C., Klein, I.., Bronstein, M. M., and Tong, A. Scalable Equilibrium Sampling with Se- quential Boltzmann Generators, February 2025.
Vargas, F., Grathwohl, W. S., and Doucet, A. Denoising Diffusion Samplers. In The Eleventh International Con- ference on Learning Representations, September 2022.
Vargas, F., Padhy, S., Blessing, D., and Nüsken, N. Trans- port meets Variational Inference: Controlled Monte Carlo Diffusions. In The Twelfth International Conference on Learning Representations, October 2023.
12

---

Temperature--Ammealel Boltzmann Generators
A. Internal Coordinate Representation
As discussed the main text, we use an internal coordinate representation based on bond lengths, angles, and dihedral angles to represent the conformations of the molecular systems. As discussed in the next section, we use splines as the invertible transformations in our coupling blocks. These splines are only defined for mappings from the interval [0, BJ to [0, B). Therefore, we scale all internal coordinates fit in this range (here, B 1).
To achieve this, we divide all dihedral angles by 2m. Furthermore, bond lengths and angles are transformed as
1/3 (1) 77smmn) 1 0.5.
Here, Mmin is the value of the corresponding degree of freedom from minimum energy structure obtained from minimizing the initial structure with the force field. a was empirically chosen as 0.07 nm for bond length dimensions and 0.5730 rad for angle dimensions.
For the peptides studied in this work, two chiral forms (mirror images) exist. In nature, one almost exclusively finds only one of the two (L-form). However, since the potential energy is invariant to the mirror symmetry, there no preference given by the energy model itself. Previous work (Midgley al., 2023b; Schopmans & Friederich, 2024) simply filtered the "wrong" R-chirality during training. In contrast, directly constrain generation to the L-chirality by restricting the output bounds of the splines that generate the dihedrals of the hydrogens at the chiral centers to the range [0.5, 1.0). For this, we transform the corresponding dimensions as T + 7/2 0.5 + 0.5 after the flow generated them. This entirely removes the R-chirality from the space the flow can generate.
Furthermore, there is no preference given for the permutation of hydrogens in CH₃ groups. However, since the ground truth molecular dynamies simulations start from a given starting configuration, preference does exist in the ground truth. Therefore, restrict the generated distribution of the flow this preference, by constraining the spline output range of the respective dihedral angles, analogous to how we constrain the chirality (see above).
B. Architecture
Spline Trafo
Input
Split dims
Merge dims
Output
NN1
Figure Illustration of normalizing flow coupling layer.
For the normalizing flow architecture, & Friederich, 2024). As the invertible transformation in the coupling layers, we use monotonic rational -uuadraatq splines (Durkan et al., 2019) that map the interval [o, 1] to [0, 1] using monotonically increasing rational-quadratic functions with K bins.
architecture similar to previous works (Midgley et al., 2023b; Schopmans
We use 8 pairs of neural spline coupling layers. In each pair, we use randomly generated mask to decide which dimensions transform and which dimensions to condition the transformation on (see Figure 4). In the second coupling of each pair, the inverted mask is used. The dimensions of the dihedral angles are treated using circular splines (Rezende et al., 2020) to incorporate the correct topology. After each coupling layer, we add random (but fixed) periodic shift to the dihedral angle dimensions.
The latent distribution 9z the normalizing flow is a uniform distribution the range [o, 1] for the dihedral angle dimensions, and a Gaussian distribution with o 0.5 and o = 0.1, truncated to the range [o, 1), for the bond length and angle dimensions.
13

---

Temperature--Anmalle Boltzmann Generators
As the conditioning network 1NN; in Figure 4), we fully connected neural network with hidden dimensions 2,56,256,25,,256,25555, and ReLU activation functions. Previous works used fully connected neural network with skip connection (Midgley et al., 2023b; Schopmans & Friederich, 2024), however, we found no benefit in this and therefore did not use skip connection. To incorporate their periodicity, dihedral angles are represented as (cos a sin e)) T before being passed as input to the neural network.
We used the same architecture for all experiments of all methods. The number of parameters the architecture of each system can be found in Table 2. All experiments used the Adam optimizer (Kingma Ba, 2017) to train the flow.
Table 2. Number of parameters in the normalizing flow architecture for cach system.
ALANINE DIPEPTIDE
ALANINE TETRAPEPTIDE
ALANINE HEXAPEPTIDE
NUMBER OF PARAMETERS
7421 512
9752.376
12 624
To implement the normalizing flow models and the internal coordinate representations, we used the bgflow (Noé, 2024) and nflows (Conor Durkan et al., 2020) libraries with PyTorch (Paszke et al., 2019).
C. Molecular Systems
Table 3. Overview of the molecular systems. The number of constrained bonds is given in brackets.
NAME ALANINE DIPEPTIDE
SEQUENCE NO. ATOMS
NO. HYDROGENS
NO. BONDS
NO. ANGLES
NO. TORSIONS
ACE-ALA-NMO
22
12
21
20
19
ALANINE
TETRAPEPTIDE
ACE-3-ALA-NME
42
22
19 (+ 22)
40
39
ALANINE HEXAPEPTIDE
ACE-S ALA-NME
62
32
29 (+ 32)
60
59
To avoid diverging van der Waals energies due to atom clashes, we train with regularized energy function (Midgley et al., 2023b):
Eveg (E)
E,
S Enight
log(E- Enigh + 1) Enight
Emigh < < S Emax,
(10)
log(Fmax Fnigh + 1) + Fnigh:
if E > Fmax-
For all systems, we used the energy regularization parameters Enigh 10 and Emax x 2000
Throughout this work, use visualizations of the free energy F = kky of the systems (Ramachandran plots). For all Ramachandran plots, we used I x 107 samples if not otherwise specified, Furthermore, the axes are in scaled internal coordinates in the range [0, 1].
I' In p(o,, vh,) of the backbone dihedral angles
D. Force Field and Ground Truth Simulations
All ground truth simulations have been performed with OpenMM 8.0.0 (Eastman et al., 2024) using the CUDA platform. Ground truth energy evaluations during training have been performed with 18 workers in parallel using the OpenMM CPU Platform.
Details on the performed simulations and force field parameters for each system can be found in Table 4. For all systems, we used variants of Amber force fields (D.A. Case et al., 2023). For alanine dipeptide, the parameters are identical to the ones used in the FAB publication (Midgley et al., 2023b). We use the dataset made available by Stimper et al. (2022) as our test dataset. In addition to this ground truth test dataset, we performed another molecular dynamics simulation for alanine
14

---

Temperaturc-AAmealld Boltzmann Generators
dipeptide 300 K (see Table 4). We used 50 for equilibration and production simulation time of 5 us. The small time step of 1 fs was chosen since this system does not hydrogen bond length constraints. This additional simulation was used for training the forward KLD experiments and to create separate validation dataset.
The force field parameters of the alanine tetrapeptide system match those used in the temperature steerable flow publication (Dibak et al., 2022). However, no public dataset for this system was available, which is why we performed two replica exchange molecular dynamics (REMD) simulations to obtain 300 ground truth data. The hexapeptide system was, to the best of our knowledge, not used in previous publications, so we also here performed REMD simulations obtain a ground truth dataset for evaluation. All REMD simulations used 200 equilibration without exchanges, 200 ns equilibration with exchanges, and ju for the production simulation. For both the tetrapeptide and hexapeptide, used one simulation for the ground truth test dataset and the other simulation to form the training dataset for the forward KLD experiments and to create a separate validation dataset.
Additionally to the simulations at 300 K, we performed high-temperature simulations at 1200 to create validation datasets for the reverse KI.D pre-training (see Table 4).
All datasets were subsampled randomly from the total production MD trajectories. The ground truth test datasets at 300 contain 1 10? samples, the datasets used for the forward KLD experiments contain x 10 samples. The additional validation datasets at 300 K and 1200 K contain x 10 samples.
Table Overview the molecular dynamics simulations performed obtain the ground truth datasets. We only specify the production simulation time. In the case of REMD simulations, this is the simulation time for cach replica.
SYSTEM
FORCE FIELD
CONSTRAINTS
T / K
SIM. TIME / us
TIME STEP /
AMBER H=96 WITH OBCI IMPLICIT SOLVATION
ALANINE DIPEPTIDE
NONE
300
5.0
1.0
1200
ALANINE
AMBER99SB- ILDN WITH AMBER99 OBC IMPLICIT SOLVATION
HYDROGEN BOND LENOTHS
300, 332, 368, 408, 451, 500 (REMD)
1.0
2.0
TETRAPEPTIDE
1200
2.5
1.0
ALANINE HEXAPEPTIDE
AMB BIIBER99SB-ILLOE WITH AMBER99 OBC
HYDROGEN ROND LENGTHS
300, 332, 368, 408, 451, 500 (REMD)
1.0
2.0
1200
2.5
1.0
IMPLICIT SOLVATION
E. Metrics Details
obtain comparable results to those the original FAB publication (Midgley et al., 2023b), we evaluated the forward KLD of the Ramachandran plots the same way. First, we calculated the probability density of the Ramachandran plot of the ground truth and that of the flow distribution on 100x100 grid, using 1 10 samples for both. Then, the forward KLD is calculated between the two distributions.
RAM KLD
RAM KLD W. RW We repeat the same procedure to assess the obtained Ramachandran plot after reweighting to the target distribution. Here, analogous to Midgley et al. (2023b), we clipped the 0.01 % highest importance weights to the lowest value among them. This is necessary because of outliers in the importance weights due to flow numerics.
ESS
The ESS was calculated according to the following equation (Midgley et al., 2023a):
Norrv
(11)
N
N WEN w (x:)²
w (wx))
with
I ~ qx (zz;)) (xx)) ΣN w(x;)
Also here, we clipped the 0.01 % highest importance weights to the lowest value among them. Furthermore, the ESS was calculated with respect to the regularized energy function (Equation 10).
15

---

Temperaturc-Anmealld Boltzmann Generators
While one also calculate the forward ESS, which uses samples from the ground truth (Midgley al., 2023a), in practice found this metric yield spurious results, depending heavily on the chosen importance weight clipping value. Therefore, we chose to only use the reverse ESS, even though it does not capture mode collapse.
F. TA-BG
E.. Workflow Variations
As described in the main text, next to our buffered iterative annealing workflow, we also tried training temperature- conditioned normalizing flow on the whole continuous temperature range. To fix the distribution 1200 K the distribution learned by the reverse KL.D, used split architecture to train with lemperatuur--onditioonnng:
A base flow generates samples 1200 K. This was trained with the reverse KLD at 1200 and the model parameters of this base were frozen afterwards.
A lemperature-condiitonne head flow is added to the base flow. which transforms the highttemperature samples to lower temperatures. The spline couplings of the head flow are scaled in such way that they always output the identity for T = 1200 K.
In each batch, we sampled a support temperature T support and a target temperature T reweight and performed reweighted forward KLD training:
E(x) exp 2
Lreweight
E xmai (a:: T)
log 9x (x; T reweight 8)
(12)
9x (I;; T supports 0)
Stopped gradients
We also used this training objective with self-normalized importance weights within each batch, and with resampling each batch according to the importance weights.
The training objective in Equation 12 is similar to the one introduced by Wahl al. (2025), but does not require the estimate of the partition function. practice, found the iterative annealing workflow yield more accurate results, while also being more sampling efficient compared using Equation 12.
A systematic comparison of using Equation 12 and using the training objective introduced by Wahl et al. (2025) can be explored in future work.
F.2. Hyperparameters of Main Experiments
Table Hyperparameters of the TA-BG experiments, annealing from 1200 300 The cosine annealing learning scheduler applied within each annealing iteration, so the learning rate resets in the beginning of each new annealing iteration. "Annealing iteration" here also refers to the fine-tuning iterations with T Ti.
DIPEPTIDE
TETRAPEPTIDE
HEXAPEPTIDE
GRADIENT DESCENT STEPS PER ANNEALING ITERATION
30000
20000
20000
LEARNING RATE
5 10
1 10 4096
x 10 2048
BATCH SIZE
2048
LR SCHEDULER
COSINE ANNEALING
BUFFER SAMPLES DRAWN PER
5 10
5 10"
1 x 10"
ANNRALING ATION
BUFFER RESAMPLED TO
2 10
2 10"
2 10
In all TA-BG experiments, we used the following annealing iterations:
16

---

Temperaturc-Ammealed Boltzmann Generators
1028.69K 1200K-
8 K 648.04 K 5
T1
T2
- K + K 3 + 300.0K
3
(13)
TK-1
Tk Torn
As described in the main text, the hexapeptide additionally used intermediate fine-tuning iterations after each annealing iteration, not only in the very end:
K200 + K028.69K - 1028.69K + 881.84K 8 - 755.95 K
Ti
T2
648.04 - 648.04K - K - K 476.22 476.22 - 40 - 408.24 K
3.3.96 K 349.96 - +300,K 300.0K
(14)
TK-1
Tk TT=Tup
As described in the main text, in each iteration of the annealing workflow, use buffered dataset for training, resampled according to the importance weights. Similarly to how we clipped the importance weights to calculate the forward KLD of the reweighted Ramachandrans, also here we clipped the highest 0.01 % of the importance weights to the lowest value among them. This can prevent overemphasizing outliers in the importance weights, though the effect is minimal.
We further note that instead of using only the training buffer dataset of the current annealing iteration, one can also reuse older samples that correspond to previous annealing iterations by reweighting them the current target temperature. To simplify the workflow, we did not pursue this option in our experiments, but it might further increase sample efficiency.
F.3. A Ablation Experiments
In this section, we discuss several ablation experiments to analyze the impact of different hyperparameter choices. To keep the computational cost down, performed each ablation experiment only once (except for the ablation experiments on starting temperatures). All ablation experiments start from the hyperparameters chosen in the main experiments, changing only the hyperparameters specified.
When ablating choices in the annealing workflow, we start all experiments from the same checkpoint pre-trained with reverse KLD at high temperature, to remove variations in the pre-training.
Direct Importance Sampling Figure 5 shows the result when directly performing importance sampling from 1200 K to 300 without using the iterative annealing workflow. As one can see, the overlap between the two distributions is very small, resulting in very noisy result due the small effective sample size. This shows that multistep annealing workflow is necessary.
Starting Temperature As discussed in the main text, perform high-temperature reverse KLD training in the beginning of our workflow to pre- trait the flow distribution. The starting temperature Ti at which reverse KLD training is performed needs to be chosen high enough such that minima are less separated and mode collapse is avoided.
In Figure 6, we analyze the propensity for mode collapse for each of the three studied systems. For each temperature, we performed four reverse KLD experiments and display the fraction of experiments with mode collapse among them (mode collapse is defined here by manual visual inspection of the Ramachandran plots). As one can see, for each system, a critical temperature exists above which mode collapse does not oceur.
We note that the propensity for mode collapse not only depends on the temperature, but also on the chosen batch size and the number of gradient descent steps. The experiments in Figure 6 were performed with a batch size of 1024 and 250 000 gradient descent steps. Generally, when increasing the target temperature beyond the so-obtained critical temperature, one can use a smaller batch size and fewer gradient descent steps without obtaining mode collapse. This allows more efficient training in terms of target energy evaluations. For our main experiments, we thus chose a relatively large initial temperature
17

---

Temperaturc-AAmealled Boltzmann Generators
REV KLD 1200K
IS 1200K
300K
Ground truth (300K)
10
free energy ke
Dipeptide
5
+:
91
Tetrapeptide
free energy Kg
41
61
10
Hexapeptide
free energy
d3
48
63
Figure 5. Free energy F - p(és,, dis) of dihedral angles (Ramachandran plots). reweighted directly from 1200 K. to 300 We used 10 samples importance sampling (middle).
T = K. A tradeoff exists: Increasing Ti allows cheaper reverse KLD training without mode collapse, but also requires more annealing steps to reach the desired target temperature.
Temperature Schedule As discussed in Section 4, we used geometric progression between T and Tusget as the temperature schedule for our main experiments. We compare this choice with linear temperature schedule for alanine dipeptide in Figure 7, and for the hexapeptide in Figure 8. As one can see, the geometric temperature schedule is able keep an approximately constant buffer ESS and therefore overlap between two consecutive distributions, while the buffer ESS drops down significantly for the linear schedule. In the future, also adaptive schedules can be explored (Goshtasbpour et al., 2023) to further improve the transition from one temperature level to the next.
Final Fine-Tuning Tuning As described in Section 4, all TA-BG experiments used final fine-tuning iteration with Tk- Tk- Importance sampling to the same temperature (without lowering the target temperature) yields higher overlap between the two distributions, as is evident from the increased buffer ESS the final iteration shown in Figure We found empirically that such final fine-tuning iteration with increased buffer ESS helpful to improve the final metrics for systems, as summarized in Table
Intermediate Fine- Tine-Tuning For the hexapeptide system, we additionally added intermediate fine-tuning iterations after each annealing iteration, instead of only in the very end. We compare the impact of such intermediate fine-tuning for all systems in Table 7. As we see, the impact the largest for the hexapeptide system, where including intermediate fine-tuning iterations significantly boosts the final ESS (see also Figure 9). Since such intermediate fine-tuning costs additional target energy evaluations, we did not include them for the two smaller systems.
NO. Annealing Iterations For our main experiments, we annealed the temperature times for all systems. We summarize the impact of the number of annealing iterations on the final metrics for alanine dipeptide in Table 8, Choosing the number of annealing iterations is tradeoff between improving the final metrics and lowering the number of target evaluations.
18

---

Temperaturc-Ammealed Boltzmann Generators
NO. Samples Per Annealing Iteration Also the number of samples drawn form the training buffer datasets in each annealing iteration forms tradeoff between improving the final metrics and lowering the number of target evaluations. This is summarized in Table 9 for alanine dipeptide, where we draw the specified number of samples and resample this dataset to the same size for training in each annealing iteration.
100
Dipeptide
Tetrapeptide
Hexapeptide
80
Fraction collapeed 1%
60
40
20
100
600
800
1000
1200
Temperature
Figure Ablation of starting temperature Ti: Fraction of reverse KLD experiments with mode collapse a given temperature.
Buffer ESSy, >x,, (lincar temp schedule)
0.9
Buffer ESSy (geometric temp schedule)
0.8
0.7
Buffer BSS
0.6
0.5
0.4
0.3
0.2
1200
1000
800
600
400
T. K
Figure 7. The ESS of the training buffer datasets W used anneal from T. T.-1, comparing linear and geometric temperature schedule for alanine dipeptide. The increase of the huffer ESS in the end is due to the final fine-tuning iteration with T,+1 T;
19

---

Temperature-Ammealed Boltzmann Generators
0.35
Butler ES (linear temp schedule)
Bufler E (geometxic temp schedule)
0.25
Bufla ESS
0.20
0.10
800
T: K
Figure 8. The ESS of the training buffer datasets W used to anncal from T. to T.-1, comparing a linear and geometric temperature schedule for alanine hexapeptide. Note that the hexapeptide system includes intermediate fine-tuning iterations (T.41 = T,), where the buffer ESS is increased due to better overlap compared to the iterations where the temperature is lowered (T;+1 < T.).
Table 6. Impact of final fine-tuning step on final metrics per system.
SYSTEM
INTERMED, FINE-TUNING
FINAL FINE-TUNING
NLL
ESS %
ALANINE DIPEPTIDE
WITHOUT
WITH
-213.667
95.85
WITHOUT
WITHOUT
-213.666
95.55
ALANINE
WITHOUT
WITH -330.128
62.57
TETRAPEPTIDE
WITHOUT
WITHOUT
-330.087
57.57
ALANINE HEXAPEPTIDE
WITH
WITH -501.510
15.38
WITH
WITHOUT
-501.461
13.37
Table 7. Impact of intermediate fine-tuning on final metrics per system.
SYSTEM
INTERMED, FINE-TUNING
NLL
ESS %
ALANINE DIPEPTIDE
WITH
-213.671
96.46
WITHOUT
-213.667
95.85
ALANINE TETRAPEPTIDE
WITH
-330.176
68.93
WITHOUT
-330.128
62.57
ALANINE HEXAPEPTIDE
WITH -501.510
15.38
WITHOUT
-499.580
7.56
Table 8. Impact of number of annealing iterations on final metrics for alanine dipeptide.
NO. ANNEALING ATERATIONS
NLII
ESS /
-213.647
92.27
-213.663
95.06
7
-213.666
95.57
9
-213.667
95.85
11
-213.668
96.05
20

---

Temperatuurc-Ammalle Boltzmann Generators
0.35
Buffer ESS (without intermed. finetuning)
0.30
Buffer ESSy (with intermed. finetuning)
Buffez ESS
0.20
0.15
0.05
0.00
1200
1000
800
600
400
T: /
Figure The ESS of the training buffer datasets W used anneal from to Ti+1, comparing the case with and without intermediate fine-tuning iterations for the hexapeptide system.
Table 9. Impact of number of samples drawn per annealing iteration on final metrics for alanine dipeptide.
NO. SAMPLES
NLL
ESS %
500 000
221.5556
84.49
1 000
-213.649
92.47
2 000
-213.663
95.12
000
-213.670
96.43
100000000
-213.673
96.86
21

---

Temperaturc-Ammealled Boltzmann Generators
F.4. Scaling to Higher Dimensions
Our proposed annealing workflow relies on importance sampling (IS) to anneal the distribution iteratively from T; to T.+1 Since importance sampling does not scale well when increasing the dimensionality N of the problem, want shortly discuss how annealed importance sampling (AIS) (Neal, 2001) can be used in place of IS in the future, and how this can mitigate potential problems encountered with IS when scaling to larger systems.
The main problem of IS is that small local bias, e.g., the possibility of two atoms clashing locally in the proposal, gets amplified due to the dimensionality the problem. We consider a molecular system with N independent neighborhoods, each with chance for clash under the proposal distribution. A sample's importance weight if no clashes occurred,
and otherwise. The ESS is ESS (E m.) 1 M
EM so E(ESS) = (1 n) Thus, the ESS drops exponentially,
- w
which captures the scaling problem of IS.
To demonstrate how AIS addresses this, we move to a continuous model (a similar example and analysis can be found in (Midgley et al., 2023b))) The proposal distribution is N-dimensional Gaussian with 1.1, l 0, and the target distribution is an N-dimensional Gaussian with 1.0, O. This is simplified model of single annealing iteration. To perform AIS in this toy example, we use single step of HMC with 5 leap-frog steps to transition between two intermediate distributions. We further scale the number of intermediate distributions I linearly with the dimensionality N, we chose T 5 in this example.
We visualize the results of this toy example Figure 10. As one can see, the ESS of vanilla drops exponentially. However, when using AIS and scaling the number of intermediate distributions linearly with N, one can obtain an approximately constant ESS. While this is, of course, simplified setup, it still captures the essence of why IS fails for higher dimensions and how AIS can mitigate this issue. We also refer to (Neal, 2001) for a theoretical analysis. While for molecular systems, exact factorizability of the distribution is of course not given, AIS can still remove atom clashes, etc., which are often local effects. Thus, by using AIS instead of IS, potentially arising scaling problems can be avoided in the future.
100
Empirical ESS (IS)
Empirical ESS (AI8, no scaling of intermed. dist.)
Empirical ESS (AIS, lneer scaling intermed. dist.)
60
ISS 1%
40
100
200
300
400
500
Dimensions (N)
Figure 10. Comparison of the empirical ESS as function the number dimensions N for Gaussian toy system, as obtained using vanilla IS, AIS without increasing the number of intermediate distributions, and AIS when lincarly scaling up the number of intermediate distributions with N.
22

---

Temperature-AAmealed Boltzmann Generators
G. FAB
G.I. Hyperparameters Main Experiments
For the FAB experiments, we started from the hyperparameters reported in the original publication. To make the original FAB hyperparameters a good starting point, we scaled our internal coordinates from the range of the splines [0.1] to [0, 10], which is the range used in the original FAB implementation. With this, used initial HMC step size of 0.05 for all experiments. Since FAB obtained significantly better results by using prioritized replay buffer, and were able to reproduce this finding, we chose the same replay buffer used originally by FAB for all experiments.
Table 10. Hyperparameters of the FAB experiments at 300 K. All experiments used cosine annealing learning rate scheduler with single cycle.
DIPEPTIDE
TETRAPEPTIDE
HEXAPEPTIDE
GRADIENT DESCENT STEPS
50000
50000
50000
LEARNNNNG RATE
x 10
1 10
1 10
BATCH SIZE
1024
1024
1024
GRAD NORM CLIPPING
1000.0
1000.0
1000.0
1000
1000
1000
WEIGHT DECAY (1.2)
x 10-
1 x 10
1 10
NO. INTERMED. DIST.
8
8
8
NO. INNER HMC STEPS
8
G.2. Ablation Experiments
As with most sampling approaches, number target evaluations the accuracy of the obtained distribution form a tradeoff. For the hexapeptide, FAB was not able to resolve the metastable high-energy region accurately. Therefore, we performed additional experiments where we varied the number of intermediate distributions and the number of HMC steps. This improves the results slightly, while requiring significantly more target energy evaluations (see Table 11).
For the smaller systems, alanine dipeptide and alanine tetrapeptide, we further performed experiments with smaller batch sizes while using the same number of gradient descent steps. This lowers the number of required target energy evaluations. However, as one can see from Table 11, this comes at the cost of further increasing the NLL.
Table Impact hyperparameter choices on NLL FAB. In are hyperparameters were used main experiments, chosen trade-off between computational cost and accuracy. Only the main experiments were performed four times, all other experiments were performed once.
NO. INTERMED.
NO. INNER
SYSTEM
BATCH SIZE
PE EVALS +
NLL +
DIST.
HMC STEPS
1024
8
2.13 x 10*
-213.653 + 0.000
ALANINE
DIPEPTIDE
512
1.07 10
-213.643
256
8
4
5.33 x 10
-213.623
1024
8
4
2.13 10
-330.100 + 0.002
TETRAPEPTIDE
ALANINE
512
1.07 10
-330.019
256
8
4
5.33 x 10²
-329.874
1024
8
4
2.13 10
-501.157
512
1.07 x 10
-500.857
ALANINE HEXAPEPTIDE
1024
16
4
4.20 10
-501.255
1024
8
4.20 x 10
501.268 + 0,008
1024
16
8.34 108
-501.327
23

---

Temperature-AAmealeed Boltzmann Generators
H. Reverse KLD
H.1. Hyperparameters Main Experiments
Table 12. Hyperparameters the reverse KLD experiments. This includes the 300 K experiments, but also the 1200 K experiments used start the cemperature-annealin workflow, All experiments used cosine annealing learning rate scheduler with single cycle.
DIPEPTIDE K
TETRAPEPTIDE 300
HEXAPEPTIDE 300
DIPEPTIDE 1200
TETRAPEPTIDE
HEXAPEPTIDE
T
1200
1200 K
GRADIENT DESCENT
250000
250000
250000
100
100000
250
STEPS
LEARNING RATE
x 10
1 10
1 10
1 x 10
I x 10
I x 10
BATCH SIZE
1024
1024
1024
256
256
512
GRAD NORM CLIPPING
100.0
100.0
100.0
100.0
100.0
100.0
LR LINEAR WARMUP
1000
1000
1000
1000
1000
1000
STEPS
WEIGHT DECAY (L2) NO. HIGHEST ENERGY
x 10
1x 10
x 10
1 x 10
1 x 10
1 x 10
40
40
40
10
10
20
VALUES REMOVED
I. Forward KLD
L. Hyperparameters of Main Experiments
Table Hyperparameters of the forward KLD experiments 300 K. All experiments used cosine annealing learning rate scheduler with a single cycle.
DIPEPTIDE
TETRAPEPTIDE
HEXAPEPTIDE
GRADIENT DESCENT STEPS
100000
100000
120000
LEARNING RATE
5 10
5x 10
5 10
BATCH
1024
1024
1024
24

---

Temperaturc-AAnnalled Boltzmann Generators
J. Comparison of Computational Cost
In this section, we compare the computational cost TA-BG and FAB, summarized in Table 14. We point out that neither of the two implementations was optimized for speed, thus the wall times are only approximately indicative the actually achievable performance.
As already discussed, TA-BG is more efficient compared to FAB in terms of target energy evaluation. However, this increased sampling efficiency comes with the cost of more flow evaluations, since we train for an extended period of time on large buffer datasets. The force field evaluations of the benchmark systems in this work are relatively inexpensive. Thus, the total wall time of FAB is currently slightly lower compared to TA-BG for the systems investigated here.
However, we note that while the force field evaluations are inexpensive, they are also not very accurate. When applying our approach systems with more expensive target evaluations, such as ones based on foundation model force fields or density functional theory, target evaluations become dominant, favoring methods with improved sampling efficiency, such as TA-BG.
Table 14. Computational cost of TA-BG and FAB, comparing the number of target evaluations, the number of flow evaluations (batched), and the total wall time (excluding evaluation). To determine the wall times, we ran experiments in parallel on a compute node with x AMD EPYC Rome 7402 CPU (24 cores each) and NVIDIA A100 GPU.
SYSTEM
METHOD
NO. TARGET EVALS
NO. FLOW EVALS (BATCHES)
WALL TIME
FAB
x.13 10
2.06 x 10 (AIS UPDATE BUFFER. BS 1024)
18.0h
5 10 (TRAINING, 1024)
2.56 10
ALANINE DIPEPTIDE
TA-BG
x 10 (REV, KLD PRE-TRAINING)
1 x 10S (REV. KLD PRE-TRAINING, BS 256)
22.2h
+ 5 x 10" (ANNEALING)
5 10? / 4096 (BUFFER CREATION, BS 4096)
7.56 x 10
3 10 (FWD. KLD ANNEALING, 2048)
4.12 10
FAB
2.13 10"
2.06 10 (AIS UPDATE BUFFER, 1024)
19.67 h
5 10 (TRAINING, 1024)
2.56 10
ALANINE TETRAPEPTIDE
TA-BG
2.56 10 (REV. KLD PRE-TRAINING)
1 X 10S (REV. KLD PRE-TRAINING, BS 256)
23.3h
+ 5 x 10' (ANNEALING)
10 / 4096 (BUFFER CREATION, 4096)
7.56 x 10"
2 10 (FWD. KL.D ANNEALING, 4096)
3.12 10
FAB
4.20 10
4.06 10 (AIS
BUFFER, 1024)
41.4h
x 10 (TRAINING, 1024)
= 4.56 x 10
ALANINE HEXAPEPTIDE
TA-BG
1.28 x 10 (REV. KLD PRE-TRAINING)
2.5 x 10S (REV. KLD PRE-TRAINING, BS 512)
52.2h
1.8 x 10" (ANNEALING)
1.8 10 4096 (BUFFER CREATION, 4096)
3.08 x 10
10 (FWD. KLD ANNEALING, BS 2048)
= 6655x 10
25

---

Temperature-Anmealll Boltzmann Generators
K. 2D GMM System
Next to the molecular systems, we additionally consider the Gaussian mixture density dimensions as introduced by Midgley et al. (2023b). The ground truth distribution, as well as the results obtained with TA-BG and FAB, are visualized in Figure 11.
Ground truth
TA-BG (ours)
FAB
50
50
50
50
-50
50
50
-50
50
Figure Comparison of obtained distributions for the GMM system. visualize 1000 samples from ground distribution. TA-BG, and FAB. We additionally visualize the contour lines of the ground truth log probability.
We report the obtained NLL and ESS in Table 15. This table additionally contains results for diffusion-based samplers: Path Integral Sampler (PIS, Zhang Chen (2021)), Denoising Diffusion Sampler (DDS, Vargas et al. (2022)), and Iterated Denoising Energy Matching (iDEM, Akhound Sahhundd-Sadehh al. (2024)). For these methods, we show the results reported by Aahound-Sadege et al. (2024).
FAB (REAL NVP) in Table 15 refers to the original FAB results reported by Akhound-Sadeen al. (2024), which based the Real NVP coupling flow architecture (Dinh et al., 2017). We repeated experiments with FAB based on the more expressive neural splines coupling flows (Durkan et al., 2019), reported as FAB (NEURAL SPLINES) in Table 15, The same architecture was used for our TA-BG experiments.
As one can see both from Figure 11 and Table 15, TA-BG and FAB obtain strong results on the 2D GMM system, outperforming or matching the considered diffusion based baselines. The results of TA-BG and FAB are practically identical and likely close the optimal solution possible with the given architecture.
We note that the 2D GMM system is very simple target, thus the inclusion here mostly serves illustrative purposes to show that our methodology can also be applied to more traditional sampling tasks. Since the informative value of systematic comparison on this system is limited, we did not tune the results of TA-BG and FAB with respect to the number of target energy evaluations.
TA-BG can, of course, also be applied to other sampling tasks considered related literature. However, we purposely set the focus of this work on molecular systems, where the potential energy is more correlated and complex compared many of the toy problems eurrently considered in the variational sampling literature. Symmetrie target densities with identical particles, such DW-4, LJ-13, LJ-55 (Akhound-Sadee Sadegh et al., 2024), should idcally be handled with equivariant flow architecture, thus we did not consider them here.
K.1. Details
Here, we specify hyperparameters and other details for TA-BG and FAB applied to the GMM system. We only specify hyperparameters that differ from the ones used in our main experiments on the molecular systems.
Contrary to the experiments on the molecular systems, we did not use a regularized energy function for the GMM system. Also, we did not clip importance weights when calculating the ESS. Following previous work, the ESS was calculated using 1000 samples from the model, the NLL. was calculated using 1000 ground truth samples.
26

---

Temperature-AAmealee Boltzmann Generators
Table 15. Comparison of metrics the 2D GMM system. The best-performing method for each metric is highlighted in bold. For experiments (FAB (NEURAL SPLINES) and TA-BG (NEURAL SPLINES)), specify the mean and standard deviation of the metrics across four independent runs, the rest reported with the mean and standard deviation over three independent runs.
METHOD
NLL.
ESS 1
FAB (REAL NVP)
7.14
(65.3 # 1.7) %
FAB (NEURAL SPLINES)
6.92
(97.09 0.58) %
TA- BG (NEURAL SPLINES)
6.91
(96.89 = 0.48) %
PIS
0.03
(29.5 # 1.8)
DDS IDEM
7
(68.7 = 20.8)
6,96 + 0.07
(73.4 9.2) %
Architecture For experiments on the GMM system, we used monotonically increasing rational-quadratic splines with K 16 bins in the range [050, 50]. For the parameter networks, we used fully connected neural networks with hidden dimensions [120, 120] and ReLU activation functions. We used 13 coupling layers, swapping the two dimensions after each coupling. As the latent distribution qz. we used 2D normal distribution with [-50, 50).
0 and 10.0, truncated to the range
FAB Hyperparameters We used a learning rate of I x 10 a batch size of 8192, and 000 gradient descent steps. We scaled the coordinates from [-50, 50] to [-5, 5] to perform AIS. Contrary to the experiments on the molecular systems, we set both the number of intermediate AIS distributions and the number of inner HMC steps to one. We further used 0.05 as initial HMC step size.
TA-BG Hyperparameters Reverse KLD pre-training was performed batch size of 128, and 50 gradient descent steps. Contrary to the reverse KLD experiments the molecular systems, here did not remove the largest energy values in the loss contributions of each batch. We also did not learning rate scheduler.
30.0, using learning rate of x 10
For the annealing, we used learning rate of 5 x 10-5, a batch size of 8192, and no learning rate scheduler. We used the following annealing iterations, following geometric temperature schedule:
30.0K + 18.45K + 11.35K - 6.98 + 4.30K + 2.64K + 1.63K + LOK TK-1
LOK
(15)
Ti
T2
Tk
Ty=Tag
For each annealing iteration, we draw 200 000 samples, resample 2 2000 000 samples using the importance weights (without clipping), and perform 20 000 forward KLDD gradient descent steps.
27

---

Temperature-Ammealed Boltzmann Generators
L. Additional Figures
L.1. All Systems (Reweighted)
Ground truth
FWD KLD
TA-BG
FAB
REV KLD
10
free energy 1g
Dipeptide
61
41
Tetrapeptide
10
free energy Kg.
61
41
61
Φ1
Hexapeptide
10
free energy Ke
5
92
42
62
92
Figure 12. Reweighted version of Figure Comparison of the free energy machandran plots) at 300 K, reweighted to 300 K.
kp p(ds, eh) of selected dihedral angles (Ra-
28

---

Temperature-AAmealee Boltzmann Generators
L.2. Alanine Tetrapeptide
1200K (GT)
1200K (rev KLD)
882K
556K
408K
300K
300K (GT)
fooo enege AiBegg Ro T
11
$1
41
41
10
42
42
02
92
free energy /o T
10
d
40
4g
dg
Figure Visualization the iterative annealing process for tetrapeptide, showing the free energy - kp Tnn p(bi, 1hs) backbone dihedral angles (Ramachandran plots) in each iteration. Note that all annealing iterations are shown. We used I x 10? samples for the Ramachandran plots at 300 K and x 10 samples for the rest.
29

---

Temperaturee-Amnallel Boltzmann Generators
Ground truth
TA-BG
FAB
REV KLD
FWD KLD
10
free energy Fop
6
91
41
61
61
10
free energy kg
42
62
42
62
62
10
free energy ko
6
dg
43
dg
ds
Figure 14, Comparison of the free energy F - In p(pi, tht) the backbone dihedral angles (Ramachandran plots) the tetrapeptide at 300 K.
30

---

Temperaturce-Amnallel Boltzmann Generators
Ground truth
TA-BG
FAB
REV KLD
FWD KLD
10
free energy Fip
91
41
61
91
10
free energy kg
6
42
62
40
62
62
10
free energy ke
193
6
ds
43
d3
di
Figure 15. Reweighted version Figure 14. Comparison of the free energy (Ramachandran plots) of the tetrapeptide at 300 K, reweighted to 300 K.
T In p( 1h)) of the backbone dihedral angles
31

---

Temperature-AAmealed Boltzmann Generators
L.3. Alanine Hexapeptide
1200K (GT)
1200K (rev KLD)
882K
556K
408K
300K
300K (GT)
10
free anergy Ro
11
41
61
61
free energy Kg T
42
d2
/2
freo energy 1g
10
69 dg
49
69
dg
Hee snenn oo T
5
de
64
10
free energy
d5
ds
d5
48
ds
ds
Figure 16. Visualization of the iterative annealing process for the hexapeptide, showing the free energy F Th dis) of backbone dihedral angles (Ramachandran plots) in each iteration. Note that not all annealing iterations are shown. We used x 10? samples for the Ramachandran plots 300 K and 10 samples for the
32

---

Temperature-Ammealed Boltzmann Generators
Ground truth
TA-BG
FAB
REV KLD
FWD KLD
10
free energy Kg.
d1
41
41
41
44
10
free energy Kg
d2
62
10
free energy KoT
5
43
63
10
free energy Kg
Via
64
44
de
64
10
free energy No
Figure 17. Comparison of the free energy F = ky T In p(pt, whs of the backbone dihedral angles (Ramachandran plots) of the hexapeptide 300 K.
33

---

Temperature-Ammealed Boltzmann Generators
Ground truth
TA-BG
FAB
REV KLD
FWD KLD
10
free energy To.
d1
41
4.
61
61
10
free energy
d2
62
10
free energy Kg
5n
43
63
10
kg
free energy
Via
64
44
de
64
10
free energy
6
d5
Figure 18. Reweighted version of Figure 17. Comparison of the free energy F = ky In p(di, vhr) of the backhone dihedral angles (Ramachandran plots) of the hexapeptide at 300 K, reweighted 300 K.
34


Reviews:
[{"review_id": "xl6XzUr5Sz", "reviewer": "Reviewer_SzWG", "summary": "This paper proposes temperature annealed Boltzmann generators. The proposed idea is train a normalizing flow with reverse KL at some high temperature, then train a series of models down to the target temperature using forward KL using generated samples reweighted with importance sampling from a higher temperature model.\n\nIt is shown that this method scales up to Alanine Hexapeptide in an internal coordinate system favourably compared to flow annealed bootstrapping and reverse KL training directly at the target temperature.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed feedback and the many helpful suggestions. We address them in the following:\n\nWe performed ablation studies to investigate the impact of hyperparameters, which we will include in our revised manuscript:\n* We ran ablation experiments to determine suitable starting temperatures for each of the systems. At a given temperature, we performed 4 experiments to determine how many of these experiments show mode collapse in the Ramachandran plots when pre-training with the reverse KLD:  \nhttps://ibb.co/5XgHtZYM  \nThere is a large range of temperatures where reverse KLD training is possible without mode collapse. This ablation was performed with a batch size of 1024 and 250k gradient descent steps. Increasing the temperature beyond the “critical temperature” where the experiments start not collapsing anymore allows us to pre-train cheaper with a smaller batch size and fewer gradient descent steps without mode collapse. A tradeoff exists: When increasing the temperature beyond the “critical temperature”, one can pre-train cheaper, but then the annealing will be more expensive.  \nThis ablation nicely extends our mode collapse discussion in the paper. The fact that alanine dipeptide can be sampled without mode collapse with the reverse KLD at only slightly increased temperature (375 K) is likely interesting and surprising to the sampling community in itself.\n* We ran experiments to justify the geometric temperature schedule. We compare the geometric schedule with a linearly spaced temperature schedule. In the following figures, we show the ESS of each buffer dataset we train on during the annealing:  \nhttps://ibb.co/TDfMwLbs (alanine dipeptide, one final fine-tuning step)  \nhttps://ibb.co/gbBHb66W (hexapeptide, with intermediate fine-tuning steps)  \nThe buffer ESS “spikes” visible for alanine dipeptide in the end and for the hexapeptide in-between come from the fine-tuning steps, where $T_{i+1}=T_i$, yielding better overlap. One can see that the geometric temperature schedule yields an approximately constant buffer ESS, whereas the ESS drops down significantly toward the end for the linear schedule (since the temperature steps are too large!).\n* Regarding details on the fine-tuning steps, these are performed with the same parameters as the annealing steps. The only difference is that we do not decrease the temperature, but rather reweight to the same temperature the flow was annealed to in the previous step, so $T_{i+1}=T_i$. We will make this clearer in the revised version of our manuscript.  \nWe further created an ablation to show the impact of fine-tuning. You can find details on this in our reply to reviewer 1DKq. To summarize, the final fine-tuning step is important for all systems, while the intermediate fine-tuning helps keep the ESS up for the hexapeptide.\n* For alanine dipeptide, we ran ablation experiments to show the impact of the number of temperature steps on the final metrics. We summarize the results in the following table: https://ibb.co/NgPHGG9J  \nAs you can see, a tradeoff between increasing the number of intermediate temperatures (which yields more target evaluations) and improving the final metrics exists.  \nIn the following figure, we further show the buffer ESS over the course of training: https://ibb.co/Z6VGsLFj  \nOne can see that with smaller temperature steps, the buffer ESS during training increases.\n* For alanine dipeptide, we ran ablations to show the impact of the number of samples drawn in each annealing iteration (https://ibb.co/jxwVY4M). Again, a tradeoff is visible between the number of target evaluations and the final NLL / ESS.\n\nFurthermore, thank you for suggesting the additional baseline of reweighting directly from 1200 K to 300 K. You can see what this looks like (using $10^6$ samples) in the following figure: https://ibb.co/7JHtBb2v  \nWe will include this additional baseline in our main text.\n\nIn general, the ablations show that a drop in ESS can be counteracted by using more intermediate temperatures, drawing more samples, or with additional fine-tuning steps. We are thus confident that we can scale our method to even larger and more complex systems, but we believe that already scaling to the hexapeptide is a significant milestone that we want to share with the community.\n\nIn addition to the molecular systems studied in our work, we now additionally benchmarked against the 2D GMM system (see our rebuttal to reviewer CidM), where TA-BG outperforms FAB and all diffusion-based baselines.\n\nAgain, thank you for your helpful comments, we believe to have addressed all suggestions and concerns of your review. The suggested ablations shed light on interesting dependencies present in our method. We are confident that our results, including the suggested ablations, are interesting for the sampling community and provide a foundation for future work. We hope that you can raise your score accordingly.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "No theoretical claims", "experimental_designs_or_analyses": "Yes", "relation_to_broader_scientific_literature": "This paper proposes using reverse KLD at a higher temperature then annealing which is a complementary approaches to previous works. While many works use temperature and many works use normalizing flows in this field I have not seen work that combines these two pieces together in this way.\n\nI agree with the authors assertion that to the best of my knowledge too the only ML approach that has been successfully applied to ALDP is FAB.", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "Strengths:\n* I think this is a very strong work providing a useful method and observations. I also found a number of tricks in the appendix that I personally had not thought about in parametrization. I think this will be a useful work for the community to build off of.\n* I found the exposition and figures very clear, although I do think that it could be made slightly clearer to the broader ML community why these are the interesting figures to look at.\n* The main results appear strong with adequate baselines.\n\nWeaknesses:\n* This work could be made much stronger with additional ablations around hyper parameter choices. I think this paper proves that there exists a setting which scales to at least AL6 using this method, but not a systematic understanding of what it takes to get there or further. In particular, there are a number of hyper parameters which seem to appear without justification and seem like they might be extremely important to the overall performance of the method. Specifically,\n  * Initial target temperature (set to 1200K here). It's unclear where this number came from or how it was decided. It seems like this would have a significant effect on results, but it is unclear how much of an effect it has. It would be very reassuring to see that this method also works within a reasonable range of initial target temperatures, or that there is a reasonable method to selecting the initial target temperature.\n  * Number of intermediate temperature annealing steps. As far as I can tell all that is mentioned is \"For all experiments, we chose 9 temperature annealing steps.\" \n  * Intermediate fine tuning. Very little detail is given on how this is done, what parameters are used, and what effect it has. What is the effect of this empirically? \n  * Number of samples drawn for each intermediate distribution. \n  * Number of steps trained for / stopping criterion at each temperature.\n  * Temperature annealing schedule", "other_comments_or_suggestions": "None", "questions_for_authors": "1. I think the missing baseline is training with reverse KL at the high (1200K) temperature then directly importance sampling down to 300K (and either training a flow or not). I think this would emphasize the need for annealing and the benefit of training many flow models.\n\n2. Would it be possible to include additional ablations I think the most interesting would be on starting temperature and number of temperature scales? I think this would greatly improve the practical applicability of this work.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "t5vPUYsrNw", "reviewer": "Reviewer_1DKq", "summary": "The authors propose an iterative training approach for learning normalizing flows to approximate unnormalized densities, such as Boltzmann distributions for physical systems.   The method proceeds to first train a normalizing flow using the mode-seeking KL divergence match a high-temperature target density.   Using samples drawn at the previous temperature reweighted according to the next (lower) temperature, the method proceeds to train on these examples using the mass-covering KL divergence.   The procedure is iterated until reaching the target temperature.\n\n\nThe authors demonstrate that this approach can mitigate mode-seeking behavior of directly training a normalizing flow at the target temperature.   The proposed method is competitive with or outperforms the previous gold standard method (FAB) on tetrapeptide and hexapeptide systems which are more difficult that standard baselines in the literature.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive, detailed, and constructive feedback.\n* We address your points following the same numbering as in your review:\n   2. We will more clearly outline the advantages of our method. For example, explicitly stating the mass-covering nature of the forward KLD and why this is helpful in our case is a great suggestion. We will generally work on making our introduction of forward / reverse KLD more intuitive.\n   3. Thank you for these suggestions to improve the presentation of our method. We will work on these points to revise our manuscript. We will include the equation $T_i = T_\\text{start} \\left( \\frac{T_\\text{target}}{T_\\text{start}} \\right)^{\\frac{i-1}{K-1}}$ to clarify the “geometric temperature progression”. Further, we will improve the explanation of the annealing procedure, including the fine-tuning and its motivation.\n   4. Choosing hyperparameters for a fair comparison of the different methods is not straightforward, as the final metrics and the total number of target evaluations form a tradeoff. Thus, you are right that saying our method “requires” a certain number of evaluations is imprecise; we will correct this.  \nTo better make the tradeoff visible, we included FAB hyperparameter variations in Appendix H of our original submission. Furthermore, we now performed ablation studies to make this tradeoff better visible for our method, see our answer to reviewer SzWG, and the fine-tuning ablation study discussed below.\n   5. One of the reasons why using Cartesian coordinates might make the problem harder is that with internal coordinates, we can constrain the angles and bond lengths to reasonable bounds. This is not possible with Cartesian coordinates, where even completely different chemical graphs can be formed. We suspect this is partly why diffusion-based sampling methods have not yet been successful in modeling molecular systems, since handling periodic torsions is less straightforward with diffusion than with flows, and thus Cartesian coordinates are often used.\n   6. Not clipping the importance weights results in numerical outliers visible in the reweighted Ramachandran plots. Therefore, we clipped the weights when determining the KLD of the Ramachandran plots, analogous to FAB.\nWe further ran ablation experiments for alanine dipeptide to determine the impact of clipping the importance weights when resampling the buffer datasets during the annealing. The final NLL and ESS stayed within our error bounds, so clipping has little impact here. We will mention this in the revised manuscript.\n* Further responses to your additional suggestions and questions:\n   7. We performed an ablation study to show the impact of fine-tuning. First, we look at the impact of the intermediate fine-tuning for the hexapeptide system. In the following figure, we show the ESS of the training buffer dataset in each iteration for the case with and the case without intermediate fine-tuning (both experiments have a final fine-tuning step): https://ibb.co/XxZJvr0r  \nThe buffer ESS of the fine-tuning steps that follow each annealing step is much higher, which is reasonable since we here reweight to the same temperature the flow was previously annealed to ($T_{i+1}=T_i$). One can further see that the buffer ESS generally drops down when not using intermediate fine-tuning. We summarize the impact of intermediate fine-tuning on the final metrics for all systems in this table (https://ibb.co/hjXk96b). As you can see, the impact is the largest for the hexapeptide, which is why we included intermediate fine-tuning only in this system for our main experiments.  \nWe further summarize the impact of the final fine-tuning step for all systems in the following table: https://ibb.co/XfcbxGb3\n   8. Testing other objectives (such as the α=2 divergence) for the annealing is a great idea. However, we believe this goes beyond the main claims we want to cover. We will add a short discussion of other possible objectives to our manuscript.\n   9. Thank you for pointing out the concurrent work [1]. We will mention it in our revised manuscript.\n   10. We will mention more advanced flow architectures as a path for the future [2,3].\n   11. Thank you for your additional suggestions to improve our manuscript. We will include them in the revised version.\n\nIn addition to the molecular systems studied in our work, we now additionally benchmarked against the 2D GMM system (see our rebuttal to reviewer CidM), where TA-BG outperforms FAB and all diffusion-based baselines.\n\nOverall, your comments were very helpful in improving our paper. Our approach outperforms the current SOTA (FAB) in this area and demonstrates scalability to more complex systems. We believe that this will be very helpful for the community to advance in this task, and we would appreciate it if you could revise your score based on our revision.\n\n[3] Tan et al. 2025, “Scalable Equilibrium Sampling with Sequential Boltzmann Generators”", "claims_and_evidence": "1) \nThe motivation for training on successively lower temperatures is natural and well presented.   \n- the proposed method conveniently uses mass-covering / maximum-likelihood KL divergence for training after the first iteration\n\n2) \nThis benefit of the algorithm is not properly emphasized throughout the paper\n- Below eq. 3, the transition to `we focus on the case' leads the reader to question the introduction of Eq. 2-3 and suggests it may not play a key role\n- The mass-covering nature of the KL used in training steps after the first likely plays a key role in the lack of mode dropping. \n- This change in loss, facilitated by iterative sampling, should be emphasized in \"Training by Energy\" and the introduction to Sec. 4.2\n\n3) \nThe presentation of the method could be improved:\n- clarify that the same parameters are used to retrain at each step (i.e. $\\theta_{T_i}$ used to initialize $\\theta_{T_{t+1}}$ and then discarded)\n- an inline equation specifying the temperature schedule would be appreciated\n- it is not 100% obvious what is the meaning of a 'fine-tuning step' in Lines 254-260 R.   It would be beneficial to highlight the \"off-by-one\" nature of the proposal samples in the stated algorithms, how the \"fine-tuning\" corrects this, and that hexapeptide uses e.g. 2x the number of steps\n\n\n\nMost additional findings are empirical rather than theoretical.", "methods_and_evaluation_criteria": "The method and evaluation measures make sense for the problem at hand.\n\n\nI might be interested to see how the ESS of the sampled datasets changes over iterations (with and without 'fine-tuning').   \n\n\nThe authors appear to use a large number of training steps at the initial mode-seeking KL training run.   This makes sense since later sampling results depend crucially on good initial learning, but could be emphasized.\n\n\n\n4) \nI am slightly confused by the philosophy behind setting training parameters.   In line 334-336 L, \"our approach achieves better results... while requiring $3.08 x 10^8$ target evaluations,\"  it is somewhat unclear what is meant by \"requires\".   How is the number of gradient steps decided?   Eventually, the authors approximately match the number of function evaluations for forward KL (MD steps), FAB (AIS steps), and the proposed method (resampling), but it might be good to emphasize that the proposed training framework allows for many more gradient steps for the same # PE Evals than comparison methods.", "theoretical_claims": "None given.", "experimental_designs_or_analyses": "see above", "relation_to_broader_scientific_literature": "The paper is well-positioned in relation to the literature.   \n\nIn Related Work (Lines 342-348 R), the points about diffusion models should be moved to a new paragraph.   \n\"Our results *suggest* (the possibility) that... *might* benefit\" would be better wording.\n\n\"computing exact likelihoods for these models is prohibitively expensive\" \n- this claim should be expanded upon in Lines 191-192 L and/or Lines 345-365 R.    (PF ODE, need for divergences, possible approximations)", "essential_references_not_discussed": "*Not essential*, but I noticed concurrent work also proposing to train a sampler from a higher temperature target distribution [1].  Their approach is not iterative and involves Sequential Monte Carlo resampling to reweight back to the target temperature along a diffusion trajectory.\n\nThe authors might also consider advances in normalizing flow architectures demonstrated in [2]\n\n[1] Skreta et. al 2025, \"Feynman-Kac Correctors\"\n[2] Zhai et. al 2024, \"Normalizing Flows are Capable Generative Models\"", "other_strengths_and_weaknesses": "The experimental results are promising, but the presentation can be greatly improved.", "other_comments_or_suggestions": "I personally dislike the \"forward\" and \"reverse\" KL nomenclature and don't intuitively follow it when reading papers.   This further emphasizes the need for the authors to clearly state the difference in training objectives in (i) their first step and existing methods, versus (ii) training steps at subsequent temperatures.", "questions_for_authors": "The FAB method applies just as easily for α=1 (mass-covering KL) as it would for α=2 (perhaps requiring fewer PE Evals for shorter annealing).   Since the proposed method also optimizes this loss, it may be an interesting further benchmark.\n\n5) \nDo the authors have any comment on the use of internal coordinates versus Cartesian coordinates?  I see that Midgley et. al 2023a train from maximum-likelihood on replica exchange MD simulation.    It seems that methods in the field are just \"not quite there yet\"\n\n6)\nHow sensitive is the method to clipping the highest importance weights?   The authors throw out the \"forward ESS\" values due to sensitivity to clipping, but this begs the question of its impact on training and/or Ramachandran plots (although for the latter, the authors helpfully provide reweighting-free evaluations)", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DphW21UyLm", "reviewer": "Reviewer_CidM", "summary": "This paper considers the problem of off-policy sampling from the unnormalized density distributions and proposes a novel method (TA-BG) based on a normalizing flow architecture (like FAB) that is less prone to mode collapse. In fact, the authors present a way of training a normalizing flow in this setting with the reverse KLD without mode collapse. They start from training a normalizing flow at high temperatures and then annealing the learned distribution into lower temperatures, which prevent collapsing. Moreover, the paper proposes two novel benchmark problems - similar to alanine dipeptide modeling, but harder.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed and helpful feedback. We address your questions and concerns below:\n* While Table 1 suggests TA-BG mainly improves in terms of the number of target evaluations, it’s important to note that the high-energy metastable regions constitute a small part of the state space. Small NLL differences can thus have significant implications. For example, in the hexapeptide system, TA-BG and FAB show similar NLLs, but Ramachandran plots (Fig. 3) reveal FAB fails to resolve the high-energy metastable region, while TA-BG succeeds. Therefore, quantitative metrics should be complemented with system-specific analyses like Ramachandran plots.  \nIn summary, TA-BG achieves better NLL across all systems, needs significantly fewer target evaluations, and is the only method to capture the complex interactions in the hexapeptide.\n* FAB is, indeed, the only method successfully applied to learning the Boltzmann distribution of alanine dipeptide without prior knowledge (as also confirmed by Reviewer SzWG in his review). Of the works you cited:\n    * [1] and [2] discuss transition path sampling between points A and B, which is a different task than sampling the full Boltzmann distribution.\n    * [3] relies on data from the target distribution; we cover the data-free case.\n    * [4] depends on a low-dimensional collective variable, typically non-trivial to obtain.  \nWe will cite these papers and discuss that they are related but ultimately solve a different task.\n* While diffusion samplers have gained popularity, they have not yet been successfully applied to alanine dipeptide. Most diffusion sampling studies use synthetic benchmarks with less correlated transitions between minima than the molecular systems we examine. In our reply to reviewer 1DKq we further discuss why the necessity of using Cartesian coordinates might make applying diffusion models to sample molecular systems harder, but we can mostly speculate here. As reviewer 1DKq points out, diffusion models are \"not quite there yet\", but we hope to see a successful publication on this topic soon. FAB and our method remain the only approaches that scale to system complexities like alanine dipeptide, which we now significantly extended to the tetrapeptide and hexapeptide.\n* We now include the 2D GMM system (as introduced by FAB). Since diffusion methods have been successfully applied here, this allows a direct comparison. We pre-trained the flow at T=30.0 with reverse KLD, annealed to T=1.0 in 7 steps, and performed one fine-tuning step. Results are shown in this figure (neural splines: https://ibb.co/7dkKCr76) and table (https://ibb.co/bgrDkVyD), using values reported by iDEM [5]. We note that we used two different flow architectures, REAL NVP (as originally used in FAB) and an improved neural splines architecture. We did not have enough time in the rebuttals to repeat the FAB experiments with neural splines, but we will prepare this for the revised manuscript (and also include error bars for our method). As shown, TA-BG outperforms FAB and all diffusion baselines on the 2D GMM system. We will further explore other sampling tasks, such as DW-4, and include them in the appendix.  \nWe emphasize our focus on molecular systems, which present more challenging and realistic benchmarks than commonly used synthetic systems. Our goal is to motivate the application of sampling methods, including diffusion, to these more difficult and application-relevant domains, supported by our release of the two new benchmark systems.\n* We approximated training time (excluding evaluation) for FAB and TA-BG on the alanine dipeptide system. On our hardware, FAB takes ~18.2h, TA-BG requires ~17.1h. Thus, runtimes are comparable. Furthermore, inference time is identical, since the same model is used for all experiments.  \nForce field evaluations in our current setup are relatively inexpensive. With more accurate and computationally costly forces, such as ML-based foundation models or DFT, the evaluation cost becomes dominant. In such cases, the higher sampling efficiency of our method will translate into significantly reduced computational cost.\n* Thank you for pointing out the missing standard sampling and diffusion references; we will revise the manuscript accordingly.\n* We have performed several new ablation studies (see responses to reviewers SzWG and 1DKq).\n\nWith this, we hope that we were able to convince you that our method and results are significant achievements and interesting to the sampling community. We hope to see more sampling approaches, including diffusion models, being applied and benchmarked on the complex and more application-relevant molecular systems proposed in our work.\n\n[5] Akhound-Sadegh et al. 2024, “Iterated Denoising Energy Matching for Sampling from Boltzmann Densities”  \n[6] Durkan et al. 2019, “Neural Spline Flows”", "claims_and_evidence": "The paper demonstrates that training a normalizing flows with reverse KLD does not necessarily end with mode collapse when we consider higher temperatures. Moreover, the authors present a novel annealing scheme from higher temperatures to lower. The general statements regarding problems with mode collapsing in previous models are reasonable.\n\nThe main problem is with the evidence supporting the claims for TA-BG superiority. Firstly, the experimental setup is limited to three datasets, similar to each other. Moreover, the only baselines are FAB, forward KLD, and reverse KLD, even though, FAB is not the only sampler that was tried in alanine dipeptide setting. Finally, the presented results for NLL and ESS for TA-BG and FAB are similar, without any significant difference. More on improving evaluation you can find in the following sections.", "methods_and_evaluation_criteria": "The evaluation setting is limited in the sense of number of considered energies, their variability, and the considered baselines. However, the proposed novel benchmarks might be beneficial for the community.", "theoretical_claims": "I haven’t found any obvious flaws in the theoretical properties of TA-BG.", "experimental_designs_or_analyses": "I think that the experimental design is very limited and should be significantly improved:\n\n- Firstly, FAB is not the only sampling method successfully used for alanine dipeptide, please see, e.g.,  the following papers: [1], [2], [3], and [4]. At least part of them and other samplers should be considered as baselines also (e.g., DDS, DIS, PIS, GFlowNets, or iDEM).\n- The evaluation setup is limited to three energies, which are in some sense similar. I will suggest considering also typical sampling problems, e.g., log-cox, DW-4, and LJ potentials.\n- A comparison of training costs for the considered baselines would be beneficial, regarding the presented (not significantly better) results of TA-BG.\n- Finally, the presented results in Tab. 1 (especially NLL and ESS) don’t support the claims in a convincing way.\n\n\n**References:**\n\n[1] Holdijk, Lars, Yuanqi Du, Ferry Hooft, Priyank Jaini, Berend Ensing, and Max Welling. \"Stochastic optimal control for collective variable free sampling of molecular transition paths.\" Advances in Neural Information Processing Systems 36 (2023): 79540-79556.\n\n[2] Seong, Kiyoung, Seonghyun Park, Seonghwan Kim, Woo Youn Kim, and Sungsoo Ahn. \"Transition Path Sampling with Improved Off-Policy Training of Diffusion Path Samplers.\" arXiv preprint arXiv:2405.19961 (2024).\n\n[3] Petersen, Magnus, Gemma Roig, and Roberto Covino. \"Dynamicsdiffusion: Generating and rare event sampling of molecular dynamic trajectories using diffusion models.\" (2023).\n\n[4] Phillips, Dominic, and Flaviu Cipcigan. \"MetaGFN: Exploring distant modes with adapted metadynamics for continuous GFlowNets.\" arXiv preprint arXiv:2408.15905 (2024).", "relation_to_broader_scientific_literature": "This paper is positioned within the sampling community, focused on sampling from molecular energies like alanine dipeptide. However, the proposed method seems to be novel, the paper doesn’t compare (theoretically and empirically) against related baselines.", "essential_references_not_discussed": "**Missing references for standard sampling methods:**\n\n[1] Duane, Simon, Anthony D. Kennedy, Brian J. Pendleton, and Duncan Roweth. \"Hybrid monte carlo.\" Physics letters B 195, no. 2 (1987): 216-222.\n\n[2] Hoffman, Matthew D., and Andrew Gelman. \"The No-U-Turn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo.\" J. Mach. Learn. Res. 15, no. 1 (2014): 1593-1623.\n\n[3] Halton, J. H. Sequential Monte Carlo. In Mathematical Proceedings of the Cambridge Philosophical Society, volume 58, pp. 57–78. Cambridge University Press, 1962 \n\n[4] Chopin, N. A sequential particle filter method for static models. Biometrika, 89(3):539–552, 2002. \n\n\n**And novel neural diffusion-based samplers like:**\n\n[1] Vargas, Francisco, Shreyas Padhy, Denis Blessing, and Nikolas Nüsken. \"Transport meets Variational Inference: Controlled Monte Carlo Diffusions.\" In The Twelfth International Conference on Learning Representations.\n\n[2] Sendera, Marcin, Minsu Kim, Sarthak Mittal, Pablo Lemos, Luca Scimeca, Jarrid Rector-Brooks, Alexandre Adam, Yoshua Bengio, and Nikolay Malkin. \"Improved off-policy training of diffusion samplers.\" Advances in Neural Information Processing Systems 37 (2024): 81016-81045.\n\n[3] Zhang, Qinsheng, and Yongxin Chen. \"Path Integral Sampler: A Stochastic Control Approach For Sampling.\" In International Conference on Learning Representations.", "other_strengths_and_weaknesses": "**Strengths:**\n\n[1] Proposing a way of annealing the target distribution to lower temperatures.\n\n[2] Showing that training a NF with reverse KLD at high temperatures is possible without mode collapse.\n\n[3] Introducing two novel benchmarks.\n\n\n\n**Weaknesses:**\n\n[1] Missing related references and methods (see above).\n\n[2] Limited experimental setup, not supporting the claims.\n\n[3] Lack of comparison with other than FAB baselines, especially worth to consider are diffusion-based samplers.", "other_comments_or_suggestions": "Overall, I think that the current state of the experimental setting is not enough for the general work in sampling community.\n\nFor other comments, please refer to the previous sections.", "questions_for_authors": "**Questions:**\n\n[1] Since the proposed method needs an annealing procedure from high to low temperatures, it seems to be computational heavy. Could you present the computational cost of TA-BGs and compare against baselines? I think that considering both training and inference time might be beneficial.\n\n[2] I think that comparing against diffusion-based samplers would be beneficial also from the perspective of mode collapse issue. If I remember correctly, it is not such an issue as for Normalizing Flow-based samplers?\n\nFor additional questions, please refer to the previous sections.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "z5EaQaHcTF", "reviewer": "Reviewer_BtBN", "summary": "In this paper, the authors present a temperature-annealing strategy to train normalizing flows to match unnormalized probability densities (in this case focusing on the equilibrium Boltzmann distribution of high-dimensional molecular systems). The training is done with the reverse KL divergence, assuming no access to any samples from the true density, and only the true energy function. To avoid the mode collapse issue associated with training with the reverse KL divergence objective, the authors propose to first train at a high temperature at which the sample density is closer to a uniform distribution and energy barriers are less difficult to overcome, followed by iterative reweighting and retraining at slightly lower temperatures, until the target temperature (e.g room temperature) is reached. The proposed method is demonstrated on three small protein systems, increasing in complexity from alanine dipeptide to alanine hexapeptide. The method avoids mode collapse in larger systems and successfully reproduces the ground truth distributions obtained from MD simulations.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your detailed and helpful feedback. We will address your points in the following:\n* Comparison to FAB in accuracy and metrics: While Table 1 suggests TA-BG mainly improves in terms of the number of target evaluations, it’s important to note that the high-energy metastable regions constitute only a small part of the state space. Small NLL differences can thus have significant implications. For example, in the hexapeptide system, TA-BG and FAB show similar NLLs, but Ramachandran plots (Fig. 3) reveal that FAB fails to resolve the high-energy metastable region, while TA-BG succeeds. Therefore, quantitative metrics should be complemented with system-specific analyses, such as Ramachandran plots.\n* Comparison to FAB in computational cost: We want to briefly discuss why our method is substantially more efficient than FAB. FAB uses annealed importance sampling (AIS) with intermediate Hamiltonian Monte Carlo (HMC) transitions to evaluate the α=2 divergence loss. While this yields accurate results (except for the hexapeptide), the AIS costs a significant amount of potential energy evaluations. Our method simply uses the reverse KLD at high temperature, which is significantly less costly to evaluate. While the annealing adds some additional cost, it is still very efficient since large buffers are used, where training samples are reused several times. Therefore, in total, our method is significantly cheaper than FAB.\n* Scaling to larger systems: First, we would like to emphasize that scaling to the hexapeptide is already a big achievement in itself, as the previous SOTA method FAB was the only method that worked even on the smallest of our three systems studied. Despite this, we are confident that our method scales to larger systems. While it is true that the ESS drops down for the larger systems, this can be counteracted with a more expressive architecture, e.g., the one recently proposed in [3]. At the same time, even with the current flow architecture, one can counteract the drop in ESS with other measures:\n    * Increase in the number of temperature steps or the number of drawn samples per annealing step, see our ablations in the answer to reviewer SzWG.\n    * Additional fine-tuning steps, see our ablations in the answer to reviewer 1DKq.\n    * One can use AIS instead of IS to estimate the loss in the annealing steps. While this comes at an additional cost, it can keep the variance of the loss low when increasing the intermediate AIS steps linearly with the dimensionality (see [4] for a theoretical analysis).\n* Your question about data-driven vs data-free sampling of Boltzmann distributions is very interesting. First, we want to point out that the task of sampling the equilibrium Boltzmann distribution of proteins is significantly harder than just predicting folded structures, as is done in AlphaFold. While the references [1] and [2] that you cited tackle the task of equilibrium sampling, the results are only rough approximations of the true equilibrium distribution. Even though large MD datasets have become available recently, we believe that much more data is necessary to achieve good transferability that can replace direct sampling through MD or variational methods, such as the one proposed in our work.  \nAnother recent publication [5] trained transferable Boltzmann generators on a large dataset of dipeptide MD simulations, which is a very narrow chemical space. While the transferability was successfully shown, also here the transfer to new systems only resulted in a rough estimate of the true equilibrium distribution.  \nWe thus believe that a hybrid data / variational approach could be pursued in the future. In our method, MD data can be used to supplement the high-temperature pre-training, which will make finding new modes easier and may lead to quicker convergence. We agree that this discussion is very important and we will include a discussion of these aspects in our revised manuscript.\n* The reported ESS in Table 1 is the ESS when sampling from the flow distribution (which has been annealed through multiple iterations to 300K) in the very end, reweighting to the 300K target.\n\nWe further note that we have now performed several new ablation studies, as detailed in our responses to reviewers SzWG and 1DKq. Lastly, we now additionally benchmarked against the 2D GMM system (see our rebuttal to reviewer CidM), where TA-BG outperforms FAB and all diffusion-based baselines.\n\nWith this, we hope you agree with us that our work is a significant achievement that the sampling community can build upon in the future. We hope to see more methods applied to the complex benchmark systems proposed in our work, while we are eager to scale our method to even larger systems next.\n\n[3] Zhai et al. 2024, \"Normalizing Flows are Capable Generative Models\"  \n[4] Midgley et al. 2023, “Flow Annealed Importance Sampling Bootstrap”  \n[5] Klein et al. 2024, “Transferable Boltzmann Generators”", "claims_and_evidence": "The authors claim that their proposed annealing approach prevents mode collapse by starting at a higher temperature at which the sample density is closer to uniform, followed by gradual reweighting. This claim is generally supported by the results in Figures 2, 3 and Table 1, which shows that training with reverse KL without temperature annealing often results in mode collapse. \n\nThe results of the proposed TA-BG approach are only marginally better than the strongest baseline, FAB, with most of the improvement appearing in the number of potential energy evaluations (typically much lower than FAB). Can the authors discuss more in depth what the aspect of their method is that enables them to learn more efficiently in this way? This wasn’t particularly intuitive to me.", "methods_and_evaluation_criteria": "I am concerned about the scalability of the approach to higher-dimensional systems. Specifically, I noted that the effective sample size (ESS) resulting from reweighting quickly becomes smaller (from 95% to 15%) as the system complexity increases from alanine dipeptide to alanine hexapeptide. This makes intuitive sense, as the overlaps between distributions become smaller in higher dimensions. As the authors correctly note, importance sampling between these distributions can become ineffective. I would like the authors to demonstrate, either empirically or theoretically, some scaling properties of their proposed method, perhaps as a function of system size and/or number/spacing of reweighting temperatures. I can suggest the fast-folding proteins from D.E. Shaw simulations [1] as a natural testbed for some slightly larger proteins with interesting conformational dynamics and metastable states. \n\n[1] Lindorff-Larsen, Kresten, et al. \"How fast-folding proteins fold.\" Science 334.6055 (2011): 517-520.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "See above.", "relation_to_broader_scientific_literature": "As acknowledged by the authors, the idea of temperature annealing is similar to replica-exchange molecular dynamics, in which higher temperatures are used to circumvent high energy barriers and speed up sampling.\n\nIn relation to other variational sampling methods, the paper compares most closely with Flow Annealed Importance Sampling Bootstrap (FAB). The authors demonstrate that their method matches the ground truth distribution more closely than FAB for the systems considered, and requires fewer potential energy evaluations.\n\nRegarding other generative modeling methods for molecular systems, the main other line of work is training deep, generative models given large-scale data from, e.g. equilibrium MD simulations [1, 2]. The main benefit of the presented method is that no ground truth data, which can be expensive to collect, is needed. However, I would like a more detailed justification from the authors as to why they believe their approach can remain relevant compared to this line of work, particularly as the availability of structural and simulation databases like the PDB and ATLAS [3] steadily grow. Relatedly, can the authors’ method make use of this growing data availability to improve the efficiency of learning with their method on new/unseen systems? I think some discussion regarding this would be useful to better situate the contribution in the current landscape of protein generative modeling. \n\n[1] Lewis, Sarah, et al. \"Scalable emulation of protein equilibrium ensembles with generative deep learning.\" bioRxiv (2024): 2024-12.\n\n[2] Zheng, Shuxin, et al. \"Predicting equilibrium distributions for molecular systems with deep learning.\" Nature Machine Intelligence 6.5 (2024): 558-567\n\n[3] Vander Meersche, Yann, et al. \"ATLAS: protein flexibility description from atomistic molecular dynamics simulations.\" Nucleic acids research 52.D1 (2024): D384-D392.", "essential_references_not_discussed": "While not directly related since it uses data to learn generative models, the authors should consider citing this recent paper: Lewis, Sarah, et al. \"Scalable emulation of protein equilibrium ensembles with generative deep learning.\" bioRxiv (2024): 2024-12.", "other_strengths_and_weaknesses": "Strength: The paper is very well-written and easy to follow.", "other_comments_or_suggestions": "See above.", "questions_for_authors": "Is the reported ESS metric the result of reweighting from the original temperature (1200K) to the final temperature (300K), or from some combination of all of the intermediate steps?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper introduces a new training objective for reverse Kullback-Leibler divergence based-training of Boltzmann Generators based on temperature annealing. The work has strong empirical results and the discussion phase strengthen both the ablation studies and comparisons against baselines.
Decision: Accept (poster)