Learning Boltzmann Generators via Constrained Mass Transport
Authors: Christopher von Klitzing
Denis Blessing
Henrik Schopmans
Pascal Friederich
Gerhard Neumann
Venue: ICLR 2026 Poster (2026)
Keywords: sampling
Boltzmann generators
annealing
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: Constrained Mass Transport (CMT) is a variational framework for Boltzmann generators that improves high-dimensional multimodal sampling by constraining KL divergence and entropy decay, reducing mode collapse, and outperforming SOTA methods.

Abstract: Efficient sampling from high-dimensional and multimodal unnormalized probability distributions is a central challenge in many areas of science and machine learning. We focus on Boltzmann generators (BGs) that aim to sample the Boltzmann distribution of physical systems, such as molecules, at a given temperature. Classical variational approaches that minimize the reverse Kullback–Leibler divergence are prone to mode collapse, while annealing-based methods, commonly using geometric schedules, can suffer from mass teleportation and rely heavily on schedule tuning. We introduce *Constrained Mass Transport* (CMT), a variational framework that generates intermediate distributions under constraints on both the KL divergence and the entropy decay between successive steps. These constraints enhance distributional overlap, mitigate mass teleportation, and counteract premature convergence. Across standard BG benchmarks and the here introduced *ELIL tetrapeptide*, the largest system studied to date without access to samples from molecular dynamics, CMT consistently surpasses state-of-the-art variational methods, achieving more than 2.5× higher effective sample size while avoiding mode collapse.

Published as a conference paper at ICLR 2026
LEARNING BOLTZMANN GENERATORS VIA CON- STRAINED MASS TRANSPORT
Christopher von Klitzing' Denis Blessing Henrik Schopmans
Pascal Friederich
Gerhard Neumann"
Autonomous Learning Robots, Karlsruhe Institute of Technology
Artificial Intelligence for Materials Sciences, Karlsruhe Institute of Technology
ABSTRACT
Efficient sampling from high-dimensional and multimodal unnormalized proba- bility distributions is central challenge in many areas of science and machine learning. We focus on Boltzmann generators (BGs) that aim to sample the Boltz- mann distribution of physical systems, such as molecules, at given temperature. Classical variational approaches that minimize the reverse Kullback-Leibler di- vergence are prone mode collapse, while annealing -bsedd methods, commonly using geometric schedules, can suffer from mass teleportation and rely heavily on schedule tuning. We introduce Constrained Mass Transport (CMT), variational framework that generates intermediate distributions under constraints on both the KI. divergence and the entropy decay between successive steps. These constraints enhance distributional overlap, mitigate mass teleportation, and counteract pre- mature convergence. Across standard BG benchmarks and the here introduced ELIL tetrapeptide, the largest system studied to date without access to samples from molecular dynamics, CMT consistently surpasses state-of-the-an art variational methods, achieving more than 2.5x higher effective sample size while avoiding mode collapse.
INTRODUCTION
We consider the problem of sampling from target probability measure E P(RR) given by p(x) = #(x)/2 where ÿ E ((R) can be evaluated pointwise but the normalization constant Z See p(x)) dz is intractable. Sampling from unnormalized densities arises in many areas, including Bayesian statistics (Gelman et al., 1995), reinforcement learning (Celik et al., 2025), and the natural sciences (Stoltz al., 2010).
promising alternative to classical Monte Carlo methods (Hammersley, 2013) is offered by varia- tional approaches (Struwe& & Struwe, 2000), which aim to minimize statistical divergence between variational probability measure E P(RR)) and the target p, commonly the reverse Kullback- (KL) divergence
arg min Dkt(g p), gEP(Ied)
whose unique minimizer is q² p. A prominent example is the variational learning of molecular Boltzmann generators (BGs) (Noé et al., 2019), for which p(x)) = exp(-E(e)/xa77) with E being an energy function, T the temperature, and kig the Boltzmann constant. BGs enable efficient sampling of thermodynamic ensembles, thereby bypassing costly molecular dynamics (MD) simulations and accelerating the exploration of rare but physically important states. However, learning BGs is challenging as the state space is typically high- high-dimensional, the target distribution is often highly multimodal, and evaluating E(x) can be very costly, especially when using accurate energies such as those from density- censity-functional theory (Argaman & Makov, 2000).
Furthermore, directly minimizing the reverse KL divergence tends to suffer from mode collapse, ignoring low-probability modes of the target (Blessing et al., 2024; Soletskyi et al., 2025). To address this, a number of recent approaches have proposed to construct a sequence of intermediate
Equal contribution. Correspondence denis. hlwissslesnnnnnee kit.edu and henrik scw.pkanchh @ kit.cdu.

---

Published as a conference paper at ICLR 2026
distributions that transport probability mass from tractable base distribution qu to the target (Arbel et al., 2021; Matthews et al., 2022: Vargas et al., 2023; Tian et al., 2024; Albergo Vanden-Eijnden 2024). This idea, which dates back more than two decades annealed importance sampling (Neal, 2001), is most often realized through a geometric annealing path, which is defined as a sequence of (q:);= which follows & 1--8, po, where the corresponding annealing schedule (8.))-1 ensures that qı = p. Despite its simplicity, geometric annealing can suffer from mass teleportation, where large portions of the probability mass shift to disjoint regions between successive steps, complicating mass transport Lakhhund-Sadege
al., 2025; Maurais et al., 2025; Chehab Korba, 2024; Béreux et al., 2024). Moreover, its performance critically depends on the choice of annealing schedule (Syed ct al., 2024). By "mass teleportation", we refer to the failure mode of geometric annealing in which probability mass abruptly shifts to regions where the current intermediate has negligible density, blocking effective transport. This differs from the definition of Mâté & Fleuret (2023), who focus on preserving relative weights (often termed "mode switching" (Phillips et al., 2024)).
To counteract this, we build on ideas from reinforcement learning (Schulman et al., 2015) and use trust-region constraint that modifies (1) by bounding the KL divergence between successive distributions, which results in the geometric annealing path with automatic schedule tuning. Adapting these ideas sampling problems, we further introduce constraint that explicitly controls the rate which the entropy of the variational distribution decreases along the transport path. This added degree of freedom enables deviations from the standard geometric annealing path, mitigating issues such as mass teleportation and premature convergence, while fostering greater overlap between consecutive distributions. These contributions culminate in Constrained Mass Transport (CMT), a general framework for transporting variational densities along well-defined annealing paths. To highlight its practical utility, we instantiate the framework with normalizing flows and demonstrate that consistently outperforms state-of-the-at approaches, often by a substantial margin when learning molecular Boltzmann generators directly from energy evaluations, without reliance on additional MD samples.
Contributions.
Our contributions be summarized as follows:
We introduce a general framework for addressing sampling problems through a sequence of constrained variational problems, considering trust-region, entropy, and hybrid constraints.
We establish connection between these sequences and annealing paths, which interpolate between tractable prior and the target distribution.
We instantiate our framework in practice by employing normalizing flows (Papamakarios et al., 2021) to learn molecular Boltzmann generators (Noé et al., 2019).
We show that our method, Constrained Mass Transport (CMT), consistently surpasses state-of- the- art approaches, often by a significant margin when learning molecular Boltzmann generators solely from energy evaluations, without relying on additional MD samples.
We introduce new benchmark, the ELIL tetrapeptide, which, to the best of our knowledge, is the largest system studied to date under the setting of learning exclusively from energy evaluations.
2 CONSTRAINED MASS TRANSPORT
Here, denote by P ((Rd)) the space of probability measures Id that are absolutely continuous with respect to Lebesgue measure and admit smooth densities. We approach the sampling problem by dividing (1) into a sequence of constrained optimization problems that result in an annealing path of intermediate densities ( that bridge between a tractable prior qo and the target
Trust-region constraint. Trust regions aim at dividing the problem (I) into subproblems by constraining the updated density to be close to the old density in terms of KL divergence. Formally, this is given by the iterative optimization scheme
arg min DKl(g s.t. DKL(glla) Eas (R4)
4++1
Too ensure that & ((R))) need an additional constraint f q(z)dx= = 1 which we omitted in the main part of the paper for readability.
2

---

Published as a conference paper at ICLR 2026
Geometric AP
Geometric AP via (2)
Tempered AP via (7)
Geometrie-tempered AP via (9)
linear schedule
Figure 1: Illustration of the annealing paths (AP) obtained by solving the variational problems (2), (7), or (9). Trust-region-based optimization (2) mitigates the irregularities of naive schedules (e.g.. the linear schedule), but resulting geometric AP suffers from mass teleportation as the right mode of the target distribution p emerges without overlap with carlier intermediate densities. Constraining the entropy decay between successive densities (7) prevents mass teleportation, yet fails to guarantee sufficient overlap between the initial distribution go and subsequent intermediate densities. In contrast, combining both constraints (9) yields APs that both maintain overlap between successive densities and avoid mass teleportation.
for i E N. trust-region bound Eir O and some lo E P(RR). Due to the convexity of the KL divergence, we can show that in all but the last step we actually have an equality constraint in (2); Appendix A. Thus, there exists assumptions, we can approach the above constrained optimization problem using a relaxed Lagrangian formalism, i.e..
N such that qr q" (= p). Under suitable regularity
(g. A) DeL (gll) + A (DRl(alla) ar) where À 2 O is a Lagrangian multiplier, and solve the saddle point problems
c
max min A20 (EP(
c(f) (q, A).
We note that ( convex in by convexity of the KL divergence and the dual function
(++) = min 9ir
(A)
cl (4,
gEP(Rá)
concave in since the pointwise minimum of a family of linear functions of X. Thus, (4) has unique optima which denote by qi+1 and X, respectively. Indeed, (2) admits an analytical solution which is characterized by Proposition 2.1. We refer to Appendix A for a proof and further details on problem (2).
Proposition 2.1 (Optimal intermediate trust-region densities). The intermediate optimal densities that solve (2) satisfy
(0t1x, 9: (2) p(x)) with Z (A)= (()))))((()))))))))))))))))))) p(ar) ((x))txxx da,
(5)
1 (A)
where 9++1 are the unique optima of the Lagrangian corresponding to (2).
The optimal multiplier X; that solves (2) is obtained by plugging qi+ 1++1(A) in the Lagrangian (3) to obtain the dual function fhr E C(R, R) given by
(++ (A) = c(i+ (Q+1 = -(1 +1+ A) log Z (A) A
9tr
(6)
Assuming access (A) one can solve d arg max
Iu (i+
(A) to obtain the optimal
q E P((d)) that solves (2) as gi+1 == (Ai).
Entropy constraint. In similar fashion to (2), we avoid premature convergence by regulating the entropy decay of the model by constructing a sequence of intermediate densities whose change in entropy is constrained. Formally, aim to solve the following problem
arg min DKl(gllp) (EP(34)
s.t. H(g)) H()) < Eents
4++1
where H(g) = f ((x) log (z is the Shannon entropy and Eent 0 the entropy bound. We can again approach (7) using Lagrangian formalism by introducing a Lagrangian multiplier 7) 2 0. The analytical solution (7) is characterized by Proposition 2.2 whose proof can be found in Appendix A.
3

---

Published as a conference paper at ICLR 2026
Proposition 2.2 (Optimal intermediate densities for entropy constraint). The intermediate optimal densities that solve (7) satisfy
p(x) rin Z
qi-1 n)
with Z++1(n) - 0(xx) =o,
(8)
where qi+1 the unique optima of the Lagrangian corresponding (7).
Despite the potential of (7) for counteracting premature convergence, we identify two challenges depending on the entropy of the initial density (go): First, H (qu) < Hp) then the constraint is inactive resulting in nu
reducing (7) the optimization problem stated in (1). Second, if (20) > (p) then the KI. divergence between go and 9 p1/1+% can be arbitrarily large and therefore could cause instabilities due to lack of overlap between the successive densities. While the former challenge can typically be addressed by initializing qo with large entropy, the second can be more intricate. In the following, we explain how this challenge can be addressed by combining the trust-region and entropy constraint.
Combining both constraints. One can straightforwardly combine the constraints in (2) and (7) into a single iterative optimization scheme defined as
[KL((ll))) S Etri {
4++1 = arg min (KL(g PEP(33)
s.t.
(H(a) - H(4) & Sent-
In analogy to the previous section, we introduce Lagrangian multipliers and for the trust-region and entropy constraint, respectively. Indeed, one can again obtain an analytical expression for the evolution of the optimal densities, see Proposition 2.3 and Appendix A for proof.
Proposition 2.3 (Optimal intermediate densities for entropy and trust-region constraint). The inter- mediate optimal densities that solve (9) satisfy
(0, ), = (())) ाτेा Zi+1(A,))
with Z1+1(An)) = (1()) = p(x)) 14xto (x)
(10)
where qi+1 are the unique optima of the Lagrangian corresponding to (9).
Clearly, if (qu)) > Hp)). the trust-region constraint ensures that the KL divergence between qu and Qi is most Str and, therefore, for suitable choice of ětr ensures that two consecutive densities
have sufficient overlap. Lastly, the Lagrangian dual function firent E C(R2, RR) corresponding to (9), that
Iur-ent (A,n) == -(1 A 17) log (A. dE n (q.) Eent):
(11)
requires solving a two-dimensional convex optimization problem to obtain X, 77, which can be done efficiently in practice; Section for additional details.
Connection to annealing paths. Iteratively solving (2). (7) (9) induces an annealing path, that is, a sequence of densities (q: Jex that interpolates between go and We characterize these paths in Theorem 2.4; See Appendix for a proof.
Theorem 2.4 (Annealing paths). Letp E P(RR) be the target density and qo E (((R)) some initial density. The intermediate optimal densities that solve (2), (7) and (9) satisfy
4s x 40 1-B P'7 9s Thh (i2 1), and α 1-8: (77%)
(12)
respectively, with 3 and being functions of the corresponding Lagrangian multipliers. Moreover, the sequences (as) and (Bi)seNo take values in [0, 1], satisfy ag = Bu = and aj = B1 1 for some I E N and (Bi)ieN, is monotonically increasing.
In what follows, refer to the annealing paths in (12) as geometric (G), tempered (T), and geometric- tempered (GT) annealing paths, respectively. We further refer to Figure for an illustration of the impact of the introduced constraints on the annealing paths.
4

---

Published as a conference paper at ICLR 2026
3 LEARNING THE INTERMEDIATE DENSITIES
general recipe. So far, we discussed how one can construct a sequence of intermediate measures (9:)ee using our constrained mass transport formulation. However, despite having access to the analytical form of 9;. it is typically not possible to sample from it directly. As such, we approximate each %. by distribution from tractable class Q P(R))) that permits efficient sampling and density evaluation. Given an approximation family Q, we select d: E Q to approximate qi by solving
arg min D(gi, q),
(13)
TEQ
where D is arbitrary statistical divergence between probability measures. This formulation is general: the choice of Q and D determines the trade-off between expressivity, computational cost, and statistical properties such mode coverage robustness.
Practical algorithm. In this work, we choose Q to be a normalizing flow family constructed via push-forwards of simple base measure (Rezende & Mohamed, 2015; Dinh al., 2016; Durkan et al., 2019; Kingma & Dhariwal, 2018; Gabrié et al., 2022; Kolesnikov et al., 2024; Zhai et al., 2024). Let ds P(Rd)) be an say-to-sample base measure (e.g., standard Gaussian), and let JF be class of smooth invertible maps f Fd + Rd. We define
det 05 (x)
QNF {5#0= If FF,, with (f#9=)(2) 9= 9x
(14)
Oz
with fush-forward f= To fit within this family, take D to be the importance-weighte forward KI. divergence
((t1(2) q((x)
9i+1 (++1(2))
]
g++1= = =arg min (KL(0+1ll4) with DKL (a+11l)) = Bggg
log
(15)
q(z)
qEQNr
This choice offers several advantages. First, forward KL strongly penalizes underestimating the support of qi+1, encouraging mode coverage and reducing the risk of mode collapse. Second, because 9--1 is available closed form from the constrained transport updates (see Proposition 2.1, 2.2 and 2.3), the importance weights (x11<x/a;;x)) can be computed solely from and P. Third, the importance- weighted formulation allows us to reuse samples drawn from q. enabling a seamless integration of replay buffers, resulting in increased sample efficiency. Lastly, the trust-region constraint controls the variance of the importance weights, keeping approximately constant, independent of the problem dimension d (see Appendix C.3), resulting in highly scalable algorithm. Details regarding this specific form CMT are provided Appendix C.2.
Lagrangian dual optimization. Maximizing the concave dual function (11) requires evaluating intermediate normalization constants Z1+1- This can be done efficiently by expressing Z,+1 as an expectation under qi and using Monte Carlo estimation. For instance, the expression for Z,+1 in (10) can be estimated as
Z (A, Bizng ( Ø(x)
(L)))) (
5(xx)) N (xx
I Σ Pa
11xtr
(16)
We note that, in contrast to estimating the normalization constant of the target Z, estimation of the normalization constant of the intermediate distributions Z,41 be performed with low variance since the trust-region constraint ensures sufficient overlap with the next intermediate distribution. Furthermore, samples In ~ qi and the corresponding evaluations (xx) and p(xx) are typically already computed when solving (13), so the additional cost of determining the Lagrangian multipliers negligible (see Appendix D. D.4). For example, on alanine dipeptide, it accounts for only about 0.01%
of the total training time. Details of the dual optimization procedure are provided in Appendix C.I, including a code example. Lastly, we refer to Algorithm for an algorithmic overview of the trained measure transport method.
4 RELATED WORK
Boltzmann generators. Learning molecular Boltzmann generators (Noé et al., 2019) purely from energy evaluations has been explored with both flow-based methods (Stimper al., 2022b; Midgley
5

---

Published as a conference paper at ICLR 2026
Algorithm 1 Constrained mass transport
Require:
Initial measure Qu. target measure D. divergence D, approximation family Q, buffer size N for i f- o,..., / - 1 do
Draw N samples
i,, evaluate (x(.))
Initialize buffer (()) = (an, qi (an), b(an)))=
Compute multipliers A; VJs = arg max, mER!! grem (A, 1)) using B
Compute q++1 = Q++1 = arg m D(g+t1, q) using B
return qr =
al., 2022; Schopmans & Friederich, 2025) and diffusion- methods (Liu et al., 2025; Choi et al., 2025; Kim et al., 2025). While flow-based approaches have demonstrated strong performance, their diffusion-based counterparts remain less competitive on molecular systems, often struggling with mode collapse, even on relatively small systems.
Next to purely energy-based approaches, recent work showed success in leveraging samples from a higher temperature, which are typically easier to obtain, and transferring the distribution to lower target temperature (Dibak et al., 2022a; Wahl et al., 2025; Schopmans & Friederich, 2025; Rissanen et al., 2025; Akhound-Sant Sadegh et al., 2025).
Alternative methods train on samples from molecular dynamics (Klein et al., 2023; Midgley et al., 2023; Tan et al., 2025a; Peng & Gao, 2025), allowing for amortized sampling due transferability to unseen systems (Jing et al., 2022; Abdin & Kim, 2023; Klein & Noé, 2024; Jing et al., 2024; Lewis al., 2025; Tan et al., 2025b).
Constrained optimization. Trust-region methods have long history robust optimization algo- rithms that iteratively minimize an objective within an adaptively sized "trust region"; see Conn et al. (2000) for overview. Beyond classical optimization, these methods have been extended to operate over spaces of probability distributions, with applications in reinforcement learning (Peters et al., 2010; Schulman al., 2015: 2017; Achiam et al., 2017; Pajarinen et al., 2019; Akrour al., 2019; Yang et al., 2020: Otto et al., 2021; Xu et al.. 2024; Wu et al., 2017; Abdolmaleki et al., 2018b;a; Meng et al., 2021). black-box optimization (Sun et al., 2009; Wierstra et al., 2014: Abdolmaleki et al., 2015), variational inference (Arenz et al., 2020; 2022), and path integral control (Gómez et al., 2014; Thalmeier et al., 2020). The first explicit link between trust-region optimization and geometric annealing paths was established by Blessing et al. (2025) for path space measures in the setting of stochastic optimal control. Entropy constraints, often introduced as entropy regularization, have also been studied in policy optimization and reinforcement learning, either in the form of soft constraints (Ahmed et al., 2019; Mnih et al., 2016; O' D'Doooguue et al., 2016) or hard constraints (Ab- dolmaleki et al., 2015; Pajarinen al., 2019; Akrour et al., 2016; 2018; 2019). However, prior work typically constrains the absolute entropy value, which is problematic for inference tasks, since it requires prior knowledge of the target density's 's entropy. To the best of our knowledge, such methods have not yet been extended to sampling problems. Furthermore, the connection between entropy-constrained optimization and annealing paths has not previously been established.
Improved annealing paths. Research on improving annealing paths (APs) has largely focused on geometric APs in the context of annealed importance sampling (AIS) (Neal, 2001) and their extensions to sequential Monte Carlo (SMC) (Del Moral et al., 2006); see Jasra et al. (2011); 1); Goshtasbpour et al. (2023); Chopin et al. (2023): Syed et al. (2024). Beyond the standard geometric AP, altermative constructions have been proposed, such as the moment- aaverging path for exponential family distributions (Grosse et al., 2013) and the arithmetic mean path (Chen et al., 2021). The geometric path itself can be interpreted as a quasi-arithmetic mean (Kolmogorov 1930) under the natural logarithm, which motivated Brekelmans et al. (2020) to propose APs based on the deformed logarithm transformation. A variational characterization of these paths was later analysed by Brekelmans & Nielsen (2024). Related work also explores improved schedules for parallel tempering Surjanovic et al. (2022); Syed et al. (2021) and for the diffusion coefficient in ergodic
Castelnuovo,
processes used to train denoising diffusion models (Ho et al., 2020;

---

Published as a conference paper at ICLR 2026
Song et al., 2020); see, e.g., Nichol & Dhariwal (2021); Williams et al. (2024); Benita et al. (2025); Zhang (2025).
5 NUMERICAL EVALUATION
Table 1: Results for all systems of varying dimensionality d. Evaluation criteria include the number of target evaluations (TaRGET EVALS), the evidence upper bound (EUBO), the reverse effective sample size (ESS) and the average total variation distance to the Ramachandran plots generated from molecular dynamics samples (RAM TV). Details on all metrics can be found in Appendix D.3. Each value is shown as the mean + standard error over four independent runs. An exception is TA-BG on the ELIL tetrapeptide, for which only two runs were successful due numerical instabilities. The best results are highlighted in bold, except for the forward KL and reverse KL. Reverse KL is prone to mode collapse, which makes ESS values not directly comparable, and forward KL is trained from samples rather than from energy.
SYSTRM METHOD
TARGET EVALS
EUBO
[
RAM TV
FORWARD REVERSE KL.
6x 10
174.92 0.00
(82.14++ 0.08)
(1.09 0.01) 10
ALANINE DIPEPTIDE
2.56 10*
174.96 0.00
(94.13 4: 0:21)
(1.36 0.05)
FAB
2,13 x 10"
-174.88 + 0.00
(91.8 0.04) %
(1.03 : 0:01) 10
(d 60)
1174.99 0.00
(95.76
(1.21
10-2
CMT (OURS)
1
-175.00 = 0.00
(97.69 + 0.08) %
(9.43 : 0.08) x 10
ALANINE
FORWARD REVERSE KL
10"
-353.79 0.00
0.11)
(1.47
10
2.56 10"
-332.66 a: 0.13
(75.06
3.50)
(2.89 0.02) 10-2
PEPTIDE
FAB
2,13 x 10
3333.93 0.00
(63.59
0.25) %
(3.10 0.04) 10
(d = 120)
x
333.99
0.00
(65.81
0.24)
(1.53 0.09) 10-2
CMT (OURS)
1 10
334.00 +: 0.00
(68.60
+ 0.21) %
(1.443 +: 0.03) x 10-2
ALANINE PEPTIDE
FORWARD REVERSE KI.
4.2 10"
558.16 0.01
(10.98 6.11)
(L.88 0.01) 10
2.56 108
529.26 0.26
(21.83 1.30) %%
(7.73 + 0.63) 10-
FAB
42 10a
-532 98 0.01
(11.35
0.05) %
(6.43 0.03) 10-2
(d = 180)
TA-BG
4 10
-333 6333-4-- 0.00
(18.22
(2.50 0.03) 10-2
CMT (OURS)
4 10
533.51 :: 0.01
(29.63 + 0.08) %
(2.48 +: 0.02) x 10-2
ELIL
FORWARD REVERSE KL
10"
276.76 0.00
(5.85 + 0.03) 1%
(1.58 0.01)) 10-2
TETRA-
2.56 10
-262.34 # 3.48
[1/26 + 0.53) 9
(2:61 + 0.27) x 10
PEPTIDE
FAB
8.43 10"
226..6 1: 0.01
(7.21 1 0.08) %
(7.54
10-2
(d 219)
TA-BG
80"
-277.40 =: 0.06
(13.75 11:::11449 1.42)
(2.54 0.15) x 10-2
CMT (OURS)
8 10
-277.83 + 0.00
(26.06 + 0.26) %
(3:13 + 0.03) 10
In this section, we compare our approach against state-of-the-an methods on four challenging molecular systems. We provide a brief overview of the experimental setup here, with full details in Appendix D. Additional experimental results are provided in Appendix B, including extended performance metrics, an ablation study on the effect of both constraints, and an analysis of different trust-region bounds across systems of different dimensionality.
5.1 EXPERIMENTAL SETUP
Benchmark problems. Our evaluation covers a range of molecular systems, beginning with the well-studied alanine dipeptide (d 60) (Dibak et al., 2022b; Stimper et al., 2022b; Midgley et al., 2022; Tan et al., 2025a), and extending to the larger alanine tetrapeptide (d 120) and alanine hexapeptide (d 180), which have only recently been addressed using variational methods (Schopmans & Friederich, 2025). addition, we introduce a new benchmark, the ELIL tetrapeptide (d = 219), which is digher-dimennional and which contains more complex side chain interactions compared to the alanine hexapeptide. To the best of our knowledge, this represents the largest and most complex molecular system investigated using variational approaches to date. A detailed description of all benchmark systems is provided in Appendix D.2, and visualizations of all systems can be found next Table 1.
While using Lagrangian multipliers as a stopping criterion is possible (as À = 0 implies satisfied constraints), we use a fixed number of annealing steps to strictly control the computational budget for fair benchmarking (see Algorithm 2).
7

---

Published as a conference paper at ICLR 2026
Baseline methods. Our main baselines are Flow Annealed Importance Sampling Bootstrap (FAB) (Midgley et al., 2022) and Semperature-Annalle Boltzmann Generators (TA-BG) (Schopmans & Friederich, 2025), which currently define the state of the art for variational sampling of molecular systems. For reference, we also include reverse and forward KI. training; the latter leverages ground truth samples obtained from molecular dynamics (MD) simulations (see Appendix D.2). To ensure fair comparison, all methods use neural spline flows (Durkan et al., 2019) and identical architectures.
Performance criteria. We evaluate methods primarily using three criteria. First, the evidence upper bound (EUBO), computed with ground truth MD samples. Up to an additive constant, the EUBO corresponds to the forward KL divergence and is therefore well suited for detecting mode collapse (Blessing et al., 2024). Second, we consider the effective sample size (ESS), defined as ESS( 4. [mzy~ ((t)) -1 ESS is common measure of sample quality, but it is known to be less reliable for assessing mode collapse (Blessing et al., 2024).
Finally, we consider Ramachandran plots as a qualitative criterion for assessing mode collapse. These plots visualize the 2D log-densit of the joint distribution of a pair of dihedral angles in a peptide's ' backbone. This is low-dimensional representation of important molecular configurations, making possible to assess whether the generated samples capture all relevant modes of the distribution or fail to represent certain regions of the state space. For more details on Ramachandran plots, we refer to Appendix D.3 and Schopmans & Friederich (2025). To assess the quality of the Ramachandran plots, we use the total variation distance (Ram TV) between the sampled and ground -oout (MD) Ramachandran histograms, the TV distance is symmetric and more naturally reflects the bidirectional nature of matching generated and target Boltzmann distributions, thereby also penalizing overestimation of density by the model.
For details on all metrics, we refer to Appendix D.3. Since evaluating the target density of molecular systems is typically expensive, we also report the number of target evaluations required by each method.
5.2 RESULTS
Main results. The main findings are summarized in Table 1. Across all systems and metries, our method outperforms the baselines while requiring the same or fewer target evaluations. It produces samples closer the ground-trutt distribution (EUBO), allows more efficient importance sampling (ESS), and provides superior mode coverage and resolution of metastable high-energy regions (RAM TV). While the performance gap between our method and the baselines is less pronounced for smaller systems, it widens substantially for the larger ones. In particular, on alanine hexapeptide and FI IL.II. tetrapeptide, our method attains approximately twice the ESS of competing approaches, while also avoiding mode collapse, reflected in improved EUBO and Ram TV values. In contrast, the reverse KL objective exhibits significant mode collapse, as evidenced by the widening gap in EUBO and Ram TV relative the other methods, with the most pronounced discrepancy observed on the largest system, ELIL tetrapeptide.
Taken together, the consistency of these trends across metrics and systems highlights the robustness of our method, particularly in challenging high-dimensional systems.
Ablation study for constraints. Additionally, investigate the effect of different constraint choices on the performance of the alanine hexapeptide system. Specifically, we compare four settings: using both constraints, each constraint individually, and no constraint (which corresponds to importance-weighted forward KL minimization). The results are summarized in Figure 2 and Figure Figure 2a shows that omitting the trust-region constraint causes entropy to decrease rapidly, which leads to mode collapse during training. Moreover, using only the entropy constraint yields unstable training, as evidenced by violations of the prescribed linear entropy decay. In contrast, incorporating a trust-region constraint stabilizes training, as reflected in Figure 2b, where it produces a substantially higher ESS between successive intermediate densities. Figure 3 shows Ramachandran plots of alanine hexapeptide with the constraints selectively enabled or disabled. Visible signs of mode collapse appear in all cases except for the tempered (7) and ceometric-temperee (9) variants, with the most accurate Ramachandran plot observed in the latter. Overall, our findings indicate that both constraints are necessary to achieve high ESS values while simultaneously avoiding mode collapse.
8

---

Published as a conference paper at ICLR 2026
No constraint
Geumetric via (2)
Tempered via (7)
Geometric-teripered via
581
000
531.4
531
CR2
Targtt Exsiuations 16
Tirget
-10"
(a) (G))
(b) ESS(
[%]]
(c) -UBO
(d) ESS( [%]
Figure 2: Impact of the trust-region and entropy constraint visualized on alanine hexapeptide. Figure 2a visualizes the model entropy over the course of the training. Analogously. Figure 2b shows the wmportance-weight effective sample size (ESS) between successive intermediate densities. Figures 2c and 2d depict the final log-likelihood and ESS to the target density, respectively. The variants in Figure 2d marked with "*" exhibit visible mode-collapse on the Ramachandran plots; see Figure 3. The ESS is therefore not directly comparable to the other methods. We denote by q the variational approximation of the intermediate density q:
Alanine Hexapeptide
Ground truth
constraint
Geometric
Tempered
Geometric-tempered
Fo0 Soneeg 16
Figure 3: Ramachandran plots for alanine hexapeptide with trust-region and entropy constraints selectively enabled or disabled. Using a single or no constraint leads to mode collapse, whereas combining both constraints avoids it. See Appendix D.3 for details.
CONCLUSION
We have introduced Constrained Mass Transport (CMT), variational framework for constructing intermediate distributions that transport probability mass from a tractable base measure to a complex, unnormalized target. By enforcing constraints both the KL divergence and the entropy decay between successive steps, CMT balances exploration and convergence, thereby mitigating mass teleportation, reducing mode collapse, and promoting smooth distributional overlap. Our empir- ical evaluation across established Boltzmann generator benchmarks and the here proposed ELIL letrupeptide, learned purely from energy evaluations without access to molecular dynamics samples, demonstrates that CMT consistently outperforms existing annealing-bass and variational baselines, achieving over 2.5x higher effective sample size while preserving mode diversity.
Promising directions for future work include exploring alternative approximation families Q and divergences D) for learning intermediate densities, which may lead further performance improve- ments. Applying our method in Cartesian coordinate representations also presents an interesting avenue, facilitates transferability across different molecular systems (Klein & Noé, 2024; Tan et al., 2025b). A key limitation of the current approach is the large number of gradient updates needed to approximate each intermediate target during training. Future work could investigate using more efficient loss functions, such as the log-variance loss (Richter et al., 2020), to reduce computational cost and improve scalability.
REPRODUCIBILITY STATEMENT
The source code for all experiments is available https://githut: com/ htt.s/ppppeeeootiiiii///ett----.
The ground- -ruth molecular dynamics data have been made publicly accessible at https: h/tp://doooooooooo wtooooorg/1..... 5281 /zeoodo. 188224 442244445 45.

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
The authors acknowledge support by the state of Maden-Württemeert through bwHPC. This work is supported by the Helmholtz Association Initiative and Networking Fund on the HAICORE@ KIT partition. D.B. acknowledges support by funding from the pilot program Core Informatics of the Helmholtz Association (HGF). H.S. acknowledges financial support by the German Research Foundation (DFG) through the Research Training Group 2450 "Tailored Scale-Bridgeh Approaches Computational Nanoscience". P.F. acknowledges funding from the Klaus Tschira Stiftung gGmbH (SIMPLAIX project) and the pilot program Core-Informatio of the Helmholtz Association (KiKIT project).
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Osama Abdin and Philip M Kim. Pepflow: direct conformational sampling from peptide energy landscapes through hypermetort-coodttonee diffusion. bioRxiv, pp. 2023-06, 06, 2023.
Abbas Abdolmaleki, Rudolf Lioutikov, Jan R Peters, Nuno Lau, Luis Pualo Reis, and Gerhard Neumann. Model-based relative entropy stochastic search. Advances in Neural Information Processing Systems, 28, 2015.
Abbas Abdolmaleki, Jost Tobias Springenberg. Jonas Degrave, Steven Bohez, Yuval Tassa, Dan Belov, Nicolas Heess, and Martin Riedmiller. Relative entropy regularized policy iteration. arXiv preprint arXiv: httx:v:///2.002255 2018a.
Abbas Abdolmalcki, Jost Tobias Springenberg, Yuval Tassa, Remi Munos, Nicolas Heess, and Martin Riedmiller. Maximum posteriori policy optimisation. arXiv preprint arXiv:1806.000 2018b.
Joshua Achiam, David Held, Aviv Tamar, and Pieter Abbeel. Constrained policy optimization. In International conference on machine learning, pp. 22-31. PMLR, 2017.
Zafarali Ahmed, Nicolas Le Roux, Mohammad Norouzi, and Dale Schuurmans. Understanding the impact of entropy on policy optimization. In International conference on machine learning, pp. 151-160. PMLR, 2019.
Tara Akhound-Sadegh. Jungyoon Lee, Joey Bose, Valentin De Bortoli, Arnaud Doucet, Michael M. Bronstein, Dominique Beaini, Siamak Ravanbakhsh, Kirill Neklyudov, and Alexander Tong. Progressive inference-time annealing of diffusion models for sampling from boltzmann den- sities. In ICML 2025 Generative Al and Biology (GenBio) Workshop, 2025. URL https: //openreview net wwwttttttooeeeeeeee forum? id=uo024Lveut
Riad Akrour, Gerhard Neumann, Hany Abdulsamad, and Abbas Abdolmaleki. Model-free trajectory optimization for reinforcement learning. In International Conference on Machine Learning, pp- 2961-2970. PMLR, 2016.
Riad Akrour, Abbas Abdolmaleki, Hany Abdulsamad, Jan Peters, and Gerhard Neumann. Model-free trajectory-based policy optimization with monotonic improvement. Journal of machine learning research, 19(14):1-25, 2018
Riad Akrour, Joni Pajarinen, Jan Peters, and Gerhard Neumann. Projections for approximate policy iteration algorithms. International Conference on Machine Learning, pp. 181-190. PMLR, 2019.
Michael S Albergo and Eric Vanden-Fijnden, Nets: non-equilibrium transport sampler. arXiv preprint arXiv:2410.02711,, 2024.
Michael Arbel, Alex Matthews, and Amaud Doucet. Annealed flow transport monte carlo. In International Conference on Machine Learning, pp. 318-330. PMLR, 2021.
Oleg Arenz, Mingjun Zhong, and Gerhard Neumann. Trust-region variational inference with gaussian mixture models. Journal of Machine Learning Research, 21(163): 2020.
Oleg Arenz, Philipp Dahlinger, Zihan Ye, Michael Volpp, and Gerhard Neumann. A unified per- spective on natural gradient variational inference with gaussian mixture models. arXiv preprint arXiv:2209. 11533, 2022.
Nathan Argaman and Guy Makov, Density functional theory: An introduction. American Journal of Physics, 68(1):69-77) January 2000. ISSN 1943-2909. doi: 10.1119/1.19375 URL http: //xxx doi org/ 10. 11111199111 19375.
Roi Benita, Michael Elad, and Joseph Keshet. Spectral analysis of diffusion models with application to schedule design. arXiv preprint 2xXi:22502.00180 2025.
Nicolas Béreux, Aurélien Decelle, Cyril Furtlehner. Lorenzo Rosset, and Beatriz Scoane. Fast training and sampling of restricted boltzmann machines. arXiv preprint arXiv:2405. 15376, 2024.

---

Published as a conference paper at ICLR 2026
Denis Blessing, Xiaogang Jia, Johannes Esslinger, Francisco Vargas, and Gerhard Neumann. Beyond ELBOs: A large-scale evaluation of variational methods for sampling. In Ruslan Salakhutdi- nov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), Proceedings of the 41st International Conference on Machine Learning, vol- ume 235 of Proceedings of Machine Learning Research, pp. 4209-4229. PMLR, 21-27 Jul 2024. URL https:://///mmm http:////ccceeeeeeeeeeee/e mlr httpp/eeeeeee/ee//eeeseeiieeeee .tt.iinnn2aa...........
Denis Blessing, Julius Bemnr, Lorenz Richter, Carles Domingo-Enrich, Yuanqi Du, Arash Vahdat, and Gerhard Neumann. Trust region constrained measure transport in path space for stochastic optimal control and inference, 2025, URL /ttp::///axxvvoooo/ org/ abs/ 2500//50.... 12511.
Stephen Boyd and Lieven Vandenberghe. Convex Optimization. Cambridge University Press, Cambridge, UK, 2004. ISBN 9780521833783, Seventh printing with corrections, 2009.
Rob Brekelmans and Frank Nielsen. Variational representations of annealing paths: Bregman information under monotonic embedding. Information Geometry, 7(1):193 228, 2024.
Rob Brekelmans, Vaden Masrani, Thang Bui, Frank Wood, Aram Galstyan, Greg Ver Steeg, and Frank Nielsen. Annealed importance sampling with q- q-paths. arXiv preprint arXiv:001.007823, 2020.
Richard P Brent. Algorithms for minimization without derivatives. Courier Corporation, 2013.
Onur Celik, Zechu Li, Denis Blessing, Ge Li, Daniel Palenicek, Jan Peters, Georgia Chalvatzaki, and Gerhard Neumann. Dime: Diffusion-based maximum entropy reinforcement learning. arXiv preprint wrXiv:2502.023166 2025.
Omar Chehab and Anna Korba. A practical diffusion path for sampling. arXiv preprint arXiv:2406. 14040, 2024.
Junya Chen, Danni Lu, Zidi Xiu, Ke Bai, Lawrence Carin, and Chenyang Tao. Variational inference with holder bounds. arXiv preprint arXiv:21 2021.
Jaemoo Choi, Yongxin Chen, Molei Tao, and Guan-Horng Liu. Non-equilibrium annealed adjoint sampler. arXiv preprint arXiv:2506. 18/65, 2025.
Nicolas Chopin, Omiros Papaspiliopoulos, et al. An introduction to sequential Monte Carlo, volume 4. Springer, 2020.
Nicolas Chopin, Francesca Crucinio, and Anna Korba. A connection between tempering and entropic mirror descent. arXiv preprint arXiv:2310. 11914, 2023.
Andrew R Conn, Nicholas IM Gould, and Philippe L Toint. Trust region methods. SIAM, 2000.
Thomas M Cover. Elements of information theory. John Wiley & Sons, 1999.
Pierre Del Moral, Arnaud Doucet, and Ajay Jasra. Sequential monte carlo samplers. Journal of the Royal Statistical Society Series B: Statistical Methodology, 68(3):41 1-436, 2006.
Manuel Dibak, Leon Klein, Andreas Krämer, and Frank Noé. Temperature steerable flows and boltzmann generators. Physical Review Research, 4(4):1.04205, 2022a.
Manuel Dibak, Leon Klein, Andreas Krämer, and Frank Noé. Temperature steerable flows and Boltzmann generators. Physical Review Research, 4(4):L042005, October 2022b. doi: 10.1
PhysRerRessarch.4110400000
Laurent Dinh, Jascha Sohl-Dickstein. and Samy Bengio. Density estimation using real nvp. arXiv preprint arXiv: 08803, 2016.
Simon Duane, A. D. Kennedy, Brian J. Pendleton, and Duncan Roweth. Hybrid Monte Carlo. Physics Letters B, 195(2):216-2255 September 1987. ISSN 0370-2693, doi: 10.1016/0370-269366 91 91197-X.
Conor Durkan, Artur Bekasov, lain Murray, and George Papamakarios. Neural spline flows. Advances in neural information processing systems, 32, 2019.
12

---

Published as a conference paper at ICLR 2026
Conor Durkan, Artur Bekasov, lain Murray, and George Papamakarios, nflows: normalizing flows in PyTorch, November 2020. URL https: / doi. org/ 10. 5281/zenodo. 4296287.
Peter Eastman, Raimondas Galvelis, Raúl Peláez, Charlles R. A. Abreu, Stephen E. Farr, Emilio Gallicchio, Anton Gorenko, Michael M. Henry, Frank Hu, Jing Huang. Andreas Krämer, Julien Michel, Joshua A. Mitchell, Vijay S. Pande, João PGLM Rodrigues, Jaime Rodriguez-Guent Andrew C. Simmonett, Sukrit Singh, Jason Swails, Philip Tumnrr, Yuanqing Wang, Ivy Zhang, John D. Chodera, Gianni De Fabritiis, and Thomas E. Markland. OpenMM 8: Molecular Dynamics Simulation with Machine Learning Potentials. J. Phys. Chem. B, 128(1):109-116) January 2024.
ISSN 1520-6106. doi: 10. 101110//ccc wptl//cs..jcb..00666
Marylou Gabrié, Grant M Rotskoff, and Eric Vanden-Eijnden. Adaptive monte carlo augmented with normalizing flows. Proceedings of the National Academy of Sciences, 119(10):2109420119 19, 2022.
Andrew Gelman, John B Carlin, Hal S Stern, and Donald B Rubin. Bayesian data analysis. Chapman and Hall/CRC, 1995.
Vicenç Gómez, Hilbert Kappen, Jan Peters, and Gerhard Neumann. Policy search for path integral control. In Machine Learning and Knowledge Discovery in Databases: European Conference, ECML PKDD 2014, Nancy, France, September 15-19, 2014. Proceedings, Part I 14, pp. 482-497. Springer, 2014.
Shirin Goshtasbpour, Victor Cohen, and Fernando Perez-Cruz Adaptive annealed importance sampling with constant rate progress. In International Conference on Machine Learning, pp. 11642-11658, PMLR, 2023.
Roger B Grosse, Chris J Maddison, and Russ R Salakhutdinov. Annealing between distributions by
averaging moments. Advances in Neural Information Processing Systems, 26, 2013.
John Hammersley, Monte carlo methods. Springer Science & Business Media, 2013.
Jonathan Ho. Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
Ajay Jasra, David A Stephens, Arnaud Doucet, and Theodoros Tsagaris. Inference for lévy-driven stochastic volatility models via adaptive sequential monte carlo. Scandinavian Journal of Statistics, 38(1):1-22, 20
Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi Jaakkola. Torsional diffusion for molecular conformer generation. Advances in neural information processing systems, 35:24240 242533, 2022.
Bowen Jing, Bonnie Berger, and Tommi Jaakkola. Alphafold meets flow matching for generating protein ensembles. arXiv preprint 204.::2002004445, 2024.
Minkyu Kim, Kiyoung Scong. Dongyeop Woo, Sungsoo Ahn, and Minsu Kim. On scalable and efficient training of diffusion samplers. arXiv preprint vrXiv:2505.19533 2025.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, January 2017.
Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1x1 convolutions. Advances in neural information processing systems, 31, 2018
Leon Klein and Frank Noé. Transferable boltzmann generators. Advances in Neural Information Processing Systems, 37:45281 114455511-53144 2024.
Leon Klein, Andreas Krämer, and Frank Noé. Equivariant flow matching. Advances in Neural Information Processing Systems, 36:59886-59910, 2023.
Alexander Kolesnikov, André Susano Pinto, and Michael Tschannen. Jet: A modern transformer- based normalizing flow. arXiv preprint arXiv:2412. 15129,, 2024.
Andreï Nikolaevich Kolmogorov and Guido Castelnuovo. Sur la notion de la moyenne. G. Bardi, tip. della R. Accad. dei Lincei, 1930.
13

---

Published as a conference paper at ICLR 2026
Sarah Lewis, Tim Hempel, José Jiménez-Luna, Michael Gastegger, Yu Xie, Andrew YK Foong, Victor García Satorras, Osama Abdin, Bastiaan Veeling, Iryna Zaporozhets, et al. Scalable emulation of protein equilibrium ensembles with generative deep learning. Science, pp. eadv9817, 2025.
Guan- Horng Liu, Jaemoo Choi, Yongxin Chen, Benjamin Kurt Miller, and Ricky TQ Chen. Adjoint schr\"" odinger bridge sampler. arXiv preprint arXiv:2506. 2025.
Bálint Máté and François Fleuret. Learning interpolations between boltzmann densities, 2023. URL hilps://arxiv. org/abss/301.07388 2023.
Alex Matthews, Michael Arbel, Danilo Jimenez Rezende, and Arnaud Doucet. Continual repeated annealed flow transport monte carlo. In International Conference on Machine Learning, pp. 15196-15219, PMLR, 2022.
Aimee Maurais, Bamdad Hosseini, and Youssef Marzouk. Learning paths for dynamic measure transport: A control perspective. arXiv preprint arXiv:2511.0379) 2025.
Wenjia Meng, Qian Zheng, Yue Shi. and Gang Pan. An off off-polliy trust region policy optimiza- tion method with monotonic improvement guarantee for deep reinforcement learning. IEEE Transactions Neural Networks and Learning Systems, 33(5):2223 3355)2223--235 2235, 2021.
Laurence Midgley, Vincent Stimper, Javier Antorán, Emile Mathieu, Bernhard Schölkopf, and José Miguel Hemúndezz-Lbato. Se (3) equivariant augmented coupling flows. Advances in Neural Information Processing Systems, 36:79200-79225, 2023.
Laurence Illing Midgley, Vincent Stimper, Gregor NC Simm, Bernhard Schölkopf, and José Miguel Hernández-Lobate Flow annealed importance sampling bootstrap. arXiv preprint arXiv:2208.0189) 2022.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In International conference on machine learning, pp. 1928-1937. PmLR, 2016.
Radford M Neal. Annealed importance sampling. Statistics and computing, 1 (2(2):225-133) 2001.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International conference on machine learning, pp. 8162-8171. PMLR, 2021.
Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(6457) 2019.
Frank Noć and co workers. bgflow: Boltzmann generators and normalizing flows in pytorch, 2025. URL /tttp:://githiie com/ hotm//oeeroupp. wbgp flw..w.. ow. (ba25ffccccass22200dd44777777779
Commit
Brendan O' O'Donoghue, Rémi Munos, Koray Kavukcuoglu, and Volodymyr Mnih. Pgq: Combining policy gradient and q-learning CoRR, 2016.
Fabian Otto, Philipp Becker, Ngo Anh Vien, Hanna Carolin Ziesche and Gerhard Neumann. Dif- ferentiable trust region layers for deep reinforcement learning. arXiv preprint arXiv:2101.09207 2021.
Joni Pajarinen, Hong Linh Thai, Riad Akrour, Jan Peters, and Gerhard Neumann. Compatible natural gradient policy search. Machine Learning, 108:1443-1466, 2019.
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. Journal of Machine Learning Research, 22(57):
2021.
14

---

Published as a conference paper at ICLR 2026
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vito,, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang. Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Asso- ciates, Inc., 2019. URL httpps::/prooeeednssseeerrrrpppp wwtpparwisss.pppppppemmmmmm wc/paper_raeeeeeeeeeeeeeeeeeeese fi firerr_fflleeeeeeeeeeeeeeeeeeeel eapaprr//////// 2019 9/19// / file e www.bddbbceeee. 2tt.:aaeeeee.2222222200000222222
pdf.
Xin Peng and Ang Gao. Flow perturbation to accelerate boltzmann sampling. Nature Communications, 16(1):6604, 2025.
Jan Peters, Katharina Mulling, and Yasemin Altun. Relative entropy policy search. In Proceedings of the AAAI Conference Artificial Intelligence, volume 24, pp. 1607-1612, 2010.
Angus Phillips, Hai-Dang Dau, Michael John Hutchinson, Valentin De Bortoli, George Deligiannidis, and Arnaud Doucet. Particle denoising diffusion sampler. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.). Proceedings of the l International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 40688 40724. PMLR, 21 21-27 Jul 2024. URL https:/ proceedings mwr.mlr.pree m.r..pres// vv3///235////22/ /hhpppiiiii ips aa html.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International conference on machine learning, pp. 1530-1538. PMLR, 2015.
Danilo Jimenez Rezende, George Papamakarios, Sebastien Racaniere, Michael Albergo, Gurtej Kanwar, Phiala Shanahan, and Kyle Cranmer. Normalizing Flows on Tori and Spheres. In Proceedings of the 37th International Conference on Machine Learning, pp. 8083-8092. PMLR, November 2020.
Lorenz Richter, Ayman Boustati, Nikolas Nüsken, Francisco Ruiz, and Omer Deniz Akyildiz. Vargrad: low-variance gradient estimator for variational inference. Advances in Neural Information Processing Systems, 33:13481-,34,,, 2020.
Severi Rissanen, RuiKang Yang, Jiajun He, Wenlin Chen, Markus Heinonen, Arno Solin, and José Miguel HernándezzLobats Progressive Tempering Sampler with Diffusion. In Forty-Second International Conference on Machine Learning, June 2025.
Henrik Schopmans and Pascal Friederich. Conditional Normalizing Flows for Active Learning of Coarse-Grained Molecular Representations. In Forty-First International Conference on Machine Learning. June 2024.
Henrik Schopmans and Pascal Friederich. Temperature- boltzmann generators. In Forty- second International Conference on Machine Learning, 2025, URL / net/ forum? RccRR.
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International conference on machine learning, pp. 1889-1897. PMLR, 2015.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov, Proximal policy optimization algorithms. arXiv preprint arXiv: 1rXiv:1707.06347 2017.
Roman Soletskyi, Marylou Gabrié, and Bruno Loureiro. theoretical perspective mode collapse in variational inference. Machine Learning: Science and Technology, 6(2):025056, 2025.
Yang Song, Jascha Sohl- Dickseicks,,,,, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint 1rXiv:2011.1334 2020.
Vincent Stimper, Laurence I. Midgley, Gregor N. C. Simm, Bernhard Schölkopf, and José Miguel Hernández-Lobate Alanine dipeptide in an implicit solvent at 300K. August 2022a. doi: 10.528 1288//enooo.69933124
15

---

Published as a conference paper at ICLR 2026
Vincent Stimper. Bernhard Schölkopf, and Jose Miguel Hernandez-Lobato. Resampling Base Distributions of Normalizing Flows. In Proceedings of The 25th International Conference on Artificial Intelligence and Statistics. pp.
PMLR, May 2022b.
Gabriel Stoltz, Mathias Rousset, et al. Free energy computations: A mathematical perspective. World Scientific, 2010.
Michael Struwe and M Struwe. Variational methods, volume 991. Springer, 2000.
Yuji Sugita and Yuko Okamoto. Replica-exchange molecular dynamics method for protein folding. Chemical Physics Letters, 314(1):141-151, November 1999. ISSN 0009-2614. doi: 10.1016/
50009-2614(99)0122244
Yi Sun, Daan Wierstra, Tom Schaul, and Jürgen Schmidhuber. Efficient natural evolution strategies. In Proceedings of the 11th Annual conference on Genetic and evolutionary computation, pp. 539-546, 2009.
Nikola Surjanovic, Saifuddin Syed, Alexandre Bouchard-Côré, and Trevor Campbell. Parallel tempering with a variational reference. Advances Neural Information Processing Systems, 35: 565-577, 2022.
Saifuddin Syed, Vittorio Romaniello, Trevor Campbell, and Alexandre Bouchard-Côté. Parallel tempering on optimized paths. In International Conference on Machine Learning, pp. 10033- 10042. PMLR, 2021.
Saifuddin Syed, Alexandre Bouchard-Côté, Kevin Chern, and Amaud Doucet. Optimised annealed sequential monte carlo samplers. arXiv preprint arXiv:2408. /2057, 2024.
Charlie B. Tan, Joey Bose, Chen Lin, Leon Klein, Michael M. Bronstein, and Alexander Tong. Scal- able Equilibrium Sampling with Sequential Boltzmann Generators. Forty-Second International Conference on Machine Learning, June 2025a.
Charlie B Tan, Majdi Hassan, Leon Klein, Saifuddin Syed, Dominique Beaini, Michael M Bronstein, Alexander Tong, and Kirill Neklyudov. Amortized sampling with transferable normalizing flows. In ICML. 2025 Generative Al and Biology (GenBio) Workshop, 2025b.
Dominik Thalmeier, Hilbert J Kappen, Simone Totaro, and Vicenç Gómez. Adaptive smoothing for path integral control. Journal of Machine Learning Research, 21(191):1-37, 2020.
Yifeng Tian, Nishant Panda, and Yen Ting Lin. Liouville flow importance sampler. arXiv preprint (.Xiv:2405.0662. 2024.
Francisco Vargas, Shreyas Padhy, Denis Blessing, and Nikolas Nüsken. Transport meets variational inference: Controlled monte carlo diffusions. arXiv preprint arXiv:2307.01050, 2023.
Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy, David Cournapeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, Stéfan J. der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nelson, Eric Jones, Robert Kern, Eric Larson, C Carey, Îlhan Polat, Yu Feng. Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, Antônio H. Ribeiro, Fabian Pedregosa, Paul van Mulbregt, and SciPy 1.0 Contributors. SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods, 17:261-272. 2020. doi: 10908///159--0--066666
Stefan Wahl, Armand Rousselot, Felix Draxler, and Ullrich Koethe. TRADE: Transfer of Distributions between External Conditions with Normalizing Flows. In The 28th International Conference Artificial Intelligence and Statistics, February 2025.
Daan Wierstra, Tom Schaul, Tobias Glasmachers, Yi Sun, Jan Peters, and Jürgen Schmidhuber. Natural evolution strategies. The Journal of Machine Learning Research, 15(1):949-980 2014.
Christopher Williams, Andrew Campbell, Arnau Doucet, and Saifuddin Syed. Score-optimal diffusion schedules. Advances in Neural Information Processing Systems, 37:107960-107983, 2024.
16

---

Published as a conference paper at ICLR 2026
Yuhuai Wu, Elman Mansimov, Roger B Grosse, Shun Liao, and Jimmy Ba. Scalable trust-region method for deep reinforcement learning using cronecker-factoree approximation. Advances in neural information processing systems, 30, 2017.
Haotian Xu, Junyu Xuan, Guangquan Zhang, and Jie Lu. Trust region policy optimization via entropy regularization for kullback-leibler divergence constraint. Neurocomputing, 589:127716, 2024.
Tsung-Yes Yang, Justinian Rosca, Karthik Narasimhan, and Peter J Ramadge. Projection-bais based constrained policy optimization. arXiv preprint arXiv:2010.03/52, 2020.
Shuangfei Zhai, Ruixiang Zhang, Preetum Nakkiran, David Berthelot, Jiatao Gu, Huangjic Zheng, Tianrong Chen, Miguel Angel Bautista, Navdeep Jaitly, and Josh Susskind. Normalizing flows capable generative models. arXiv preprint arXiv:24/2.06332) 2024.
Leo Zhang. The cosine schedule is fisher-rao-optimal for masked discrete diffusion models. arXiv preprint arXiv:2508. 04884, 2025.
Ciyou Zhu, Richard Byrd, Peihuang Lu, and Jorge Nocedal, Algorithm 778: L-bfgs-b: Fortran subroutines for large-scale bound-constrained optimization. ACM Transactions mathematical software (TOMS), 5354):500-560, 1997.
17

---

Published as a conference paper at ICLR 2026
APPENDIX
Proofs
19
Extended numerical evaluation
22
C Further details constrained mass transport
27
C.1 Dual optimization in practice
27
C.2 CMT with normalizing flows
29
C.3 Connection to weportance-weight variance
29
C.4 Guidance for hyperparameter selection
31
C.5 Annealing path visualizations.
32
Experimental setup
32
D.1 Architecture
32
D.2 Target densities
33
................................
D.3 Metrics
34
D.3.1 Task-agnostic metrics
34
D.3.2
Task-specific metrics
35
D.4
Computational cost
36
D.5
Hyperparameters
18

---

Published as a conference paper at ICLR 2026
A PROOFS
Proof of Propositions 2./ to 2.3. We divide the proof into two parts and start with the most general formulation using both constraints, referring Proposition 2.3. We will then derive the solution with just the trust-region (Proposition 2.1) and just the entropy constraint (Proposition 2.2), as they can be interpreted special cases of the more general formulation.
General formulation.
Consider the constrained optimization problem
H - H(g) S Cents / |dq= (17)
arg min (xr(all)) (EP(R*) and its corresponding Lagrangian
gi+11
s.t. (xl(alla)) S Etr,
C 1) (g. (.A.n.w) (kl() + A(DRL (gl14:) - Etr)
+ ( -H(g) Eent)
(18)
+ w 1ww(ddg dg dg
Using the convexity of the KI. divergence in its arguments, the convexity of the negative Shannon entropy, and that the integral is a linear functional, the objective from Equation (17) and its Lagrangian, given by Equation (18), convex in q. Using that P(I is continuous, there always exists measure # qi with (7) < Etr and H H (q) < Eent that satisfies the inequality constraints strictly. Following Boyd & Vandenberghe (2004) (Sec. 5.2.3), Slater's condition holds, implying strong duality. We can therefore instead solve the dual problem.
We start by setting up the Euler- Lagrange equation, given by
① ((+1)
(g.2, = 0, fr
0q
using
(( (1+
(q, X, w) = a( (1=) ( 7)) log ((x) log p(x) Alog (ogg:(6) (x) + w)
-eer + 17( (q:) Eent) -
and solve for q. Hence, we get
0 C
(
(g, A, w. = (1 + m) + - logg Alog + = =0
0g
(1+x+++1) +1)
w
+ logg = log qi
(19)
Using this, we can further determine w using
/ dg = / q:(2) (1))))+ 1+ dæ/ exp
1+t+1 +1) =
W
+ w = 1 (log Z.+1 (A. - 1), with 3 (), 1) = q((x)) i+ito 1नwo (x)d
Substituting w back into (19) and simplifying the fraction using p = 2p yields
(41122,A.) A, (x) Ιズスي IPAT with Z / (x)) )イ(+)p(x) +.. (x)dz,
(2,1(Am)
which uses the unnormalized target p. proving Proposition 2.3.
Special cases. Setting Eent = so or = effectively deactivates the respective constraint, yielding the trust-region objective (2) or the entropy objective (7), respectively. This is equivalent to setting the Lagrangian multipliers ' = 0 or X = 0, proving Proposition 2.1 and Proposition 2.2. respectively.
Proof of Theorem 2.4. We divide the proof into three parts and start with the most general formulation using both constraints. The first two parts will show form and monotonicity, while part three will derive the special cases with just the trust-region and just the entropy constraint.
19

---

Published as a conference paper at ICLR 2026
General form (part 1).
Given are the sequences of Lagrangian multipliers (As)eno 2 0 and (B)eEN, 2 O. We now aim to proof that the sequence (q:)ieNo> given by
40
= 0
q:
(20)
9--1
+
p putd-itt
i2 1
takes the form
i 0000%
(pa0aa with
A; =
Aj
1 A3 7
j=0
= 0 i2
and
02
E
m
IT
3,
+
14Ay+mi
We use the common convention that empty products evaluate to one. Base case (i 0): The simplest case
10 = 1- Bo (7 80
holds due to 30 = 0 (using empty product convention).
Inductive step: We start from Equation (20) and transform it using the assumption that q. 40 1- 3 (fox) )d: holds for some arbitrary but fixed E No. yielding
-+4,+ peamiiatthe p2ta/tt
Ø++1 = qi
= (gb 1 (pai) Itdetm piiao
--ti par inbctrr + t1aett
du
Using
X
A
B:
+ A +
1 A +
1+ 1 +
1=0
1 1 77
1
1 de + 1/:
1+
5=0
1 + 1//
B
1 A₁ + 7i
we now can rewrite the exponent of p yielding
ag3,
1 1 A +
1+ A +
B1
7k
II 1=8+1
Xj
A
A TK
1 A3 + ml
1
A 7h
1 A + T
k=0
1 + 7
T A
A
1
B:+1
1
+ 1/i
+ 1J
1 Xÿ + +
1
+
j=b+1
-Bi+1
A 1K
1
k=o
j=k+1
044181+2
again using the convention that an empty product evaluates to one. It directly follows qu" 1- 9:1
1 1
l++1 q
which completes the induction.
20

---

claimed. Published as a conference paper at ICLR 2026
Monotonicity (part 2). It remains show that (a:)sENo and (Bo)seN, take values in [0, 1] with ag = 30 = and aj = 31 = 1 for some E N4 and (Bi)seN, is monotonically increasing.
The first case (og 3 0) holds by definition. Assuming that there exists N such that 11-1 = 1 = 0,
3, =1 j=0
I
1
+
and
1
II j=k+1
B
A
1
K=0
follow directly for all 2 I. Both the trust-region and entropy constraints become inactive at optinum and do not prevent (g:)ieNo from reaching the unique optimum (etr, Sent > 0). Consequently, both Lagrangian multipliers will eventually vanish, motivating the existence of some e N+, such that A = fl1-1
We now show monotonicity of (8.),E,, using (A.)sEN, N and (m:)EN, 2 0. Let i E No be arbitrary. As a direct consequence
B141 B
Th A3
Ag
m
j=0
1-1
A
(1- i +
1 Xj +
A 2g 20 BjENo
1 1 M 1+
+
sequence (B:),en, must be monotonically increasing.
Special cases (part 3). We now consider the special cases in which only the trust-region constraint or the entropy constraint is active by setting the Lagrangian multiplier sequence of the other constraint to zero
We first consider only the trust-region constraint (2), which corresponds to setting the Lagrangian multiplier of the entropy constraint to zero, i.e., Y = 0 for all E No- In this scenario, ag simplifies ag 0 and a; for all i > 1. Consequently, and using So the iterates take the form
1 1- 88 i No.
gi
40
Analogously, the trust-region constraint can be rendered inactive by setting A; = for all E No- leaving only the entropy constraint active, corresponding to Equation (7). In this case, 30 0 and 3; for all 2 1, yielding
go
i 0,
9: α 9,
i 2 1,
which concludes the proof.
Proof of uniqueness and rightness of the trust-region solution. Closely following Blessing al. (2025), we now establish the uniqueness of the trust-region solution and show that the trust-region constraint is tight for all but the final step. Specifically, we show
DKL (q:l(p) < Ecr = qi =
q: arg min DKL (glp)) s.t. DkL(gllg--1) $ Etr is unique
21

---

Published as a conference paper at ICLR 2026
If PKL(9:l(p) Eur. the KKT conditions imply that the Lagrangian multiplier satisfies X; = 0, so the constraint is inactive. Consequently, 9, must solve the strictly convex unconstrained problem
min DKL(allp), (EP()))
which has the unique minimizer Since is feasible (Dxl(pllp) = S etc). it follows that qi =
Uniqueness of q; further follows from the convexity of the feasible set {a (lDKk( S Ece together with the strict convexity of the objective in when is fixed.
B EXTENDED NUMERICAL EVALUATION
Ground truth
CMT (ours)
TA-BG
FAB
Reverse KL
Alanine Dipeplide
Alanine Tetrapeptide
10
1000 enalgg
9.
Alanine Hexapeptide
92
92
92
ELIL Tetrapeptide
Ag (eess sag
Figure Comparison of Ramachandran plots of backbone dihedral angle pairs obtained with different methods. See Appendix D.3 for details on Ramachandran plots.
Complementing the results of Section 5, this section reports additional metrics for the main method comparison in Table (see also the corresponding Ramachandran plots in Figure 4), ablation study on the effect both constraints (see Table and Figure 5), and an ablation study on the trust-region constraint and its effect bounding importance-welde weight variance across different system sizes and trust-region bounds Etr (see Figure 6).
Main results. We begin with Table 2, which introduces several additional metrics, such as Ram KI. and Ram KI. w. RW, to quantify the discrepancy between the ground truth and method-geneeatee Ramachandran plots (see Appendix D.3 for details on all metrics). Substantially elevated Ram KL values serve as robust indicators of mode collapse, exemplified by the results for reverse KL training, where the Ram KL values are consistently at least an order of magnitude higher than those observed for other methods. Corresponding Ramachandran plots for the different methods shown in Figure 4. We include Ram KL our evaluation, despite its asymmetry, to ensure comparability with previous work (Midgley al., 2022; Schopmans & Friederich, 2025). For the same reason,
22

---

Published as a conference paper at ICLR 2026
Table 2: Comparison of metries obtained for all four peptide systems. To remain comparable FAB (Midgley et al., 2022) and TA-BG (Schopmans & Friederich, 2025), we first report the number of target evaluations (TARGET EVALS), the negative log-liielhhood (NLL), the reverse effective sample size (ESS), the average forward KL divergence to the Ramachandran plots (RAM KL) generated from molecular dynamics (MD) samples, and its importance weighted version (RAM KL W. RW). We then report the evidence lower bound (ELBO), the estimated target log normalization constant (log 2), the evidence upper bound (EUBO), the average total variation distance to the Ramachandran plots from MD samples and its importance-weighhee version, and the torus Wasserstein-2 2 distance (RAM CEAMTWW3) T- The torus Wasserstein-2 distance was computed using only 10 samples due to its high computational cost, making it less reliable. All values are presented as the mean + standard error across four independent experiments. An exception is TA-BG on the ELIL tetrapeptide, for which only two runs were successful due to numerical instabilities. The best-performing variational method for each metric is highlighted in bold. Reverse KL was excluded, as it tends to suffer from mode collapse, making the corresponding metries incomparable. The ELBO is marked with "*" because rare extreme log- impooraaenee weights can dominate its mean, making it unreliable for comparing Boltzmann generators (see Appendix D.3 for details).
SYSTEM
METHOD
TARGET RVALS J
NEL
ESS [%]] 7
RAM KI.
RAM KL. RW 1
FORWARD KI.
6x 10
-214.358 0.000
(82.14 + 0.08) %
(2.17 + 0.03) x 100-
(1.8++0.06) 10-3
ALANINE DIPEPTIDE
REVERSE
10a
-214,396
0.001
09..3 %
(1.66 0:25) 10-2
(1.37 10-2
(d = 60)
FAB
2,13 x 10
-214.410 0.000
(91.80 0.04) %
(1.22 x 10"3
(1.23 0.01) x 10
TA-BG
1x 10
-214.419 1
(95.76 a: 0.13) %
(1.92 1: 0.07) 10-3
(1.37 (1.37 10-3
CMT (OURS)
1 10
-214.429 0.000
(97.69 + 0.03) %
(1.48 + 0.03)
10
(1.35 x 10-2
FORWARD KL. REVERSE
4.2x 10"
332,020 + 0.001
(45.20 + 0.11) %
(2.33 + 0.06) x 10-3
(2.33 0.04) 10-3
ALANINE
TETRA-
2.56 10a
-331,199 0.126
(75.06 %
(3.05 à 0.35) x
(2.96 10
PEPTIDE
FAB
2.13 10
-352.158 +: 0.002
(63.59 = 0.23) %
(7:00 + 0.25) x 10-3
(1.13 0.01) x 10-3
(d 120)
TA-BG
10"
332.219 0.003
(61.81 0..44 %
(1.92 0.08) 10-3
(1.56+0.00) 10-3
CMT (OURS)
1 102
332.231
0.002
(68.60 + 0.21) %
(2.17 + 0.09)
x 10-3
(1.57:00000) 10-3
FORWARD KL REVERSE
4.2x 10"
304,528 # 0.005
(10.96 =: 0.11) %
(4.09 + 0.23) x 10-3
(7.49 1.00) 10-3
ALANINE
10a
-300,225 0.265
(21.53 1: 1.30) %
(5.14 0.36) x 10
(5.07 0.355 x10
PEPTIDR
FAR
4.2x 10*
204.346 + 0.000
(14.55 + 0.05) %
(2.16 + 0.02) x 10-2
(1.06+ 0.02) x 10-22
(d 180)
TA-BG CMT (OURS)
4x 10
-04,792 0.005
(18.22 :: 0.15) %
(6.33 0.12) 10-3
(6.27 1: 0.10) 10-3
4 10"
504.875 0.008
(29.63 +: 0.08) %
(1.18 + 0.00) x 10
(1.15+00.01) 10-2
FORWARD KL REVERSE
4.2 10P
197 2: 0.005
(5.85 2: 0.08) %
(2.68 4 0.05) x 10-3
(8.98 2 0.07) x 10-3
TETRA- PEPTIDE (d 219)
ELIL
10
-386,78
3.185
[1.26 11.26:00.58)
(1.32 0.31 10
(1.27 2: 0.38) 10
FAB
8.43 10a
001.105
0.009
(7.21 0.08) %
(2.49 + 0.10)
x 10-2
(1.25 0.10) 10-2
TA-BG (OURS)
8 10
-001.400 0.064
(13.75 +: 1.42) %
(6.24 + 0.84) 10-3
(6.22 0.51) 10-8
8.x
602.272
0.004
(26.06 1 0.26) %
(7:95 0.11)
10-0
(4.79 :: 0.06) x 10-a
SYSTEM
METHOD
ELBO *
log 2
EUBO
RAM TV
RAM TV W. RW
FORWARD KL
117.99 0.37
-175.010
0.000
-176.92 00000
(1.09:0001)
10
(9.13 10
ALANINE DIPEPTIDE
REVERSE KL.
175.34 +: 0.10
175.000 4: 0.000
-174.96
0.00
(1336:0.06) x 10
[9.36=0.05) x 10
(4 60)
TA-BG
PAB
-33..76 175---23.0 21.27
-1.5.01.0.00
-175.009
-175.010
0.000
-174.99
-174.96
0000
(1:05++001) 10"3
(1211 10
[8.67 10,670010) 10-0
(8.44 x 10
CMT (OURS)
115.06 : 0.01
175.009 4: 0.000
-17.000
= 0.00
(9.43 4: 0.06) x 10-3
(8.51 4 0.06) x 10-3
ALANINE
TETRA-
REVERSE KI.
FORWARD KL
-525.18 44.77
-334.40 0.01
-334.193 0.000
334.233 0.007
-333.79
332.96
0.00
0.13
(1.47 ( 0.03) 10-2
(2.89 0.02) 10
(1.01=001) x 10-2
x 10-2
PRPFINE
FAB
-17658.25 1: 3801.78
-331
0.000
-339.9
0.00
(3,10 x 10"2
(8.85 = 0.044 x 10
(4 120)
TA-BG
333.77 0.07
334.193
333.90
0.00
(1.33 0.09) 10
(0.29 10
CMT (OURS)
3300.444 15.76
334.183
0.000
333.00 = 0.00
(1.43 + 0.03) x
x 10-4
FORWARD KL
-80802.8 13 1275.91
-531.112
0.002
533,
0.01
(1.88 1 0.01) x 10-3
(1.8 x 10
ALANINE
HIXA-
REVERSE KL.
335.46 0.20
634.547 0.000
529.26
0.26
(7.73 0.63) 1
(6.41=0411) 10
PEPTIDE
FAB
-11788137
111.2. 83
-534.120 1: 0.003
-532.96
0329996-21 0.01
(6.43 x 10-2
(2.140001) x 10-2
[d 180)
TA-BG
-51042.60
49233 UB
-534-416
0.002
0.00 5333.43
(2.50 :: 0J33 10
(1.85 = 0.01) 10
CMT (OURS)
-56239.16
5690.08
-531.428 0.001
5533.51
(2.48 0.02) x 10-2
(1.97 10
ELIL TETRA- PEPTIDE (d 219)
REVERSE KL
FORWARD KL.
75532.96
-282.53 4: 0.50
9545.42
-279.070
278.483 + 0.000
0.156
-262.34
276.76-0.00 0.00
3.48
(1.38 + 0.01) x 10-2
(2.61 10-1
[1.67 0.27) x 10-1
(2.72-0.09) x 10
FAB
544538.24
10288.
278.479
276.67
0.01
(7.54 10
[3.38
10
TA-BG
22859158
4220. 44
-278.480
0.002
227.400
0.06
2.64410.11)) x 10-2
(2.18-011) x 10
CMT (OURS)
900.75 + 164.10
278.481 0.001
277.83
0.00
(3.13 0.03) 10
(1.68 0.01) 10
also report the negative log-likelihood (NLL), which is equivalent to the EUBO and forward KL up to an additive constant. Further details on all metrics are provided in Appendix D.3.
23

---

Published as a conference paper at ICLR 2026
Table 3: Performance of CMT with the trust-region and entropy constraints selectively enabled or disabled.
TRUST-RAGIUN ENTROry
-175.01 11:00
(1.17 (.00)
000)
(8:13 (00))
(.06)
175.00
(1.45 0.03)
=002)
0/0)
24 (0)
10
(135
ALANINE
1
183.01 0.21
11.66:24.2222
10-3
[T.12 =
10-2
(1.70 0.27) 10-2
(1:45:2:002
-0-1
3
158.9 2
(2.16 0.05)
10
(1.24
0.02) 10-
(0.56 0.14) 10--
12.40 (.00)
139)
ALANIME BEXA-
10F
581
0.211
0.57 10-1
12.82 (1,5) 10-1
03:3 0144) 10-3
(3485: 0.09) 10-1
8.. 10
3553
0.37
(1.20 1.01)
10
[C11 1.57)
10
(2 0.791
(7:35
1072
13.35
2.1.0
10:1
=1.30
-.12) 12) 10-7
1
0:03) 10-7
(2.20 (.00))
30-1
189)
13
-533.19 0..01
(1.25 0.05)
10"2
(1:21 0.02) 10
(2.48 22
(002)
x (1.07 0.00) 10m2
No constraint
Geometric via (2)
Tempered via (7)
Geometric-temperee via (9)
Model Enkmpy
-100
500
200
300
Gradiait Noma
200
100
100
100
0
0-
Target Evmilnations
-10"
Target
-10*
Target Evaluations
-10
(a) Alanine Dipeptide
(b) Alanine Tetrapeptide
(c) Alanine Hexapeptide
Figure 5: Effect of trust-region and entropy constraint on the model entropy (top row) and the gradient norm (bottom row) across different molecular systems.
Ablation study for constraints. Table presents the performance of our method under different configurations, with the trust-region and entropy constraints selectively enabled disabled. In addition to the alanine hexapeptide results shown in the main paper, we also report results for alanine dipeptide and alanine tetrapeptide. The absence of both constraints effectively corresponds seporaance-wiightee forward KI. training. Considering the EUBO, which serves as a forward metric, becomes clear that both constraints are necessary achieve optimal performance. Variants of the method without the entropy constraint showed at least partial mode collapse, making the reverse ESS mostly incomparable. For this reason, in addition to reporting the EUBO, also provide Ramachandran blsed metrics to better assess mode collapse.
Figure 5 depicts the evolution of model entropy and the gradient norm (prior to clipping) during training across different systems. Training with only the entropy constraint yields an approximately linear decay of entropy for both alanine dipeptide and alanine hexapeptide. hexapeptide, however, the entropy constraint is noticeably violated, likely due to the system's 's higher dimensionality and the pronounced discrepancy between the initial model distribution go and the first intermediate distribution QI Larger system sizes also tend to increase the gradient norm, most prominently in alamine hexapeptide. The combination of the trust-region and entropy constraints produces the most stable gradient norms, while the approximately linear entropy decay indicates that the entropy constraint is effectively enforced, thereby enabling its practical application even in the case of alanine hexapeptide. By contrast, the trust-region constraint alone leads to rapid
the case of alanine
24

---

Published as a conference paper at ICLR 2026
entropy collapse, which reduces exploration and ultimately limits the algorithm's final performance in practice.
Sr 0.01
6tz 0.05
Etr 0.1
Etc
0.2
Etr
0.3
Et.
100
100
Impertuseewwiggh ESS
50
0
0
0
0.2
0.4
0.6
0.8
1
0
0.2
0.4
0.6
0.8
3
Targel Evaluations
100
Targel Evaluationn
-10"
Targel Evaluations
-10
(a) Alanine Dipeptide
(b) Alanine Tetrapeptide
(c) Alanine Hexapeptide
Figure 6: suportance-weightt variance between successive intermediate distributions, shown in terms of effective sample size (ESS), for different trust-region bounds and system sizes. Each trust-region bound Str defines an approximate lower bound on the ESS, indicated by dashed lines.
Ablation study on the trust-region bound. Figure 6 illustrates the seroortanceweiihhe variance of CMT across different trust-region bounds and system sizes, highlighting the approximate direct relationship between the trust-region bound and the variance of importance weights between con- secutive intermediate distributions. Importance- weight variance expressed in terms of effective sample size (ESS). In the absence of trust-region constraint (Etr = oo). the ESS decreases with increasing system size. By contrast, finite trust-region bounds yield a substantially larger and more stable ESS, with the approximate lower bound on the ESS becoming increasingly well realized as the trust-region bound Etr decreases. Notably, this approximate lower bound is independent of the problem's 's dimensionality, property that is empirically supported.
Comparison to non-adaptive schedules To assess the importance of adaptivity, we performed two additional experiments on alanine hexapeptide comparing CMT non-adaptive alternatives. We omit numerical results in this section, as they would not provide meaningful insight: both non-adaptive baselines exhibit substantially lower performance across all relevant metrics.
Constant trust-region multiplier. In the first experiment, we replace the adaptive trust-region schedule by a constant multiplier A; 43.3 for all i and remove the entropy constraint. The corresponding schedule is visualized in Figure 7. In practice, the non-adaptive variant performs significantly worse: training becomes unstable and exhibits substantial mode collapse, failing to adequately cover the target distribution.
TA-BG with more annealing steps. In the second experiment, we compare CMT to TA-BG with fixed (non-adaptive) annealing schedule and varying numbers of annealing steps. TA-BG consists of a pre-training phase followed by a main training phase with annealing. To ensure fair comparison with CMT (which uses 400 anncaling steps and total of 800000 gradient updates), account for the fact that TA-BG uses 250 (100 steps to pre-training and 550 000 steps to annealed training. Matching the total number of updates yields effective number of annealing steps
550 000
400 275.
250 000 + 550 000
We therefore tested both 275 and 400 annealing steps. Across both settings, TA-BG performs significantly worse.
Energy histograms and TICA projections and systems shown in Figure 8.
The target energy histograms for the different methods
We further report the Wasserstein-2 distance relative to the ground- truth MD data, using 10 samples for both the energy histograms and the TICA projections, shown in Table 4.
25

---

Published as a conference paper at ICLR 2026
CMT adaptive (Ai)s
CMT constant (A:):
10
62 0.5
10P
100
200
3000
2000
300
400
Annosling
Annealing Steps i
Figure 7: Visualization of constant Lagrangian multiplier schedule (Az): compared an adaptive one using CMT on alanine hexapeptide. Both variants were performed without an entropy constraint.
Ground truth
Predicted
Predicted with reweighting
CMT (ours)
TA-BG
FAB
Alanine Dipeptide
Alanine Tetrapeptide
6:0-
-80
-8 40
Alanine Hesapepride
d0
Tetrapeptide
0.2
Energy KgT
Figure 8: Comparison of target energy histograms shown as normalized density for different molecular systems (rows) and methods (columns).
We also ask the reader focus Figure 9, which shows several TICA projections of ELIL tetrapeptide. We chose this peptide because it most clearly illustrates the issue with using only 104 samples for the Wasserstein metrics. Figure qualitatively demonstrates that the Wasserstein-2 metric on the TICA projections using 10tt samples cannot reliably assess their quality. To further support this, we also report the TICA KL divergence between the MD TICA histogram and the predicted TICA histogram computed using 10" samples on ELIL tetrapeptide in direct comparison to the TICA Wasserstein distance using 10a samples (see Table 5).
26

---

Published as a conference paper at ICLR 2026
Table 4: Wasserstein-2 distances for the Ramachandran plots (Ram-T-Wy, Ram-T-W; RW), energy histograms (E-W2 RW), and TICA projections (TICA-W2, TICA-W; RW), evaluated against the generated ground- truth MD data using 10 samples. We report the Wasserstein-2 distance on the energy histograms only after reweighting (RW) to the target distribution, both to avoid domination by rare extreme target-energy values and to stay consistent with related work (Tan et al., 2025b;a). We also wish to emphasize that the Wasserstein metrics often exhibit large standard errors due to the comparatively small sample size used (10*). For the same reason the Wasserstein-2 distance does not seem to be reliable metric for comparing TICA projection quality. Further details are provided in Figure 9.
SYSTEM
METHOD
RAM-T-W;
RAM-T-W, RW
E-W; RW
TICA-Wa
TICA-W; RW
FAB
0.078 = 0.007
0.085 0,012
0.232 = 0.028
0.025 = 0.004
0.016 + 0.003
ALANINE
TA-BG
0.064 # 0.002
0.073 + 0.002
0.202 + 0.031
0.015 + 0.001
0.021 + 0.003
DIPEPTIDE
CMT (OURS)
0.059 + 0.003
0.068 + 0.002
0.139 + 0.026
0.015 0.002
0.012 0.001
ALANINE
FAB
0,617 = 0.010
0.506 + 0.005
0.193 + 0.018
0.080::000
0.029 0.005
TETRA-
TA-BG
0.456 0.006
0,517 0.008
0.263 + 0.053
0.038 0.006
0.045 - 0.004
PEPTIDE
CMT (OURS)
0.492 = 0.006
0.517 1417
0.235 + 0.035
0.054 + 0.003
0.048 + 0.007
ALANINE
FAB
1.089 + 0.001
1.053 + 0.015
0.475 + 0.053
0.183 + 0.005
0.097 + 0.006
HEXA-
TA-BG
0.838 + 0.006
1.088 + 0.036
0.416 + 0.017
0.078 1 0.003
0. + 0.026
PEPTIDE
CMT (OURS)
0.833 0.003
+ 0.153
0.817 + 0.314
0.006 0.087
0.114
ELIL
FAB
0.813 0.011
0.829 + 0.077
1.133 0.418
0.108
0.010
0.184 0.034
TETRA-
TA-BG
0.586 + 0.024
0.569 + 0.000
0,450 + 0.067
0.150 0.013
0.070 + 0.002
PEPTIDE
CMT (OURS)
0.631 = 0.008
0.628 + 0.048
0.384 + 0.053
0.153 = 0.003
0.084 + 0.011
Table 5: TICA KL divergences (TICA KL and TICA KL RW) evaluated against the ground-trutt MD data using 10" samples, shown alongside the TICA Wasserstein distances computed with 10* samples for comparison. Unlike the Wasserstein metric, which is less sensitive to fine details, the TICA KL with a larger sample size is more sensitive, as reflected in this table.
SYSTEM
METHOD
TICA KL
TICA KL RW +.
TICA-W2
TICA-W, RW
ELIL
FAB
(6.82
10-2
(4.37 0.37)
10-2
0.108 + 0.010 0.184 + 0.034
TETRAPEPTIDE
CMT
(8.58 1 0.22) X 10-3
(1.05 1 0.05) X 10-2 0.153 1 0.003
0.084 1 0.011
Ablation different entropy bounds We visualize the final performance of CMT alanine hexapeptide for different entropy bounds Sent in Figure 10. While the optimal entropy bound Sent varies slightly between systems, we note that they are all in the same order of magnitude. Most hyperparameters of CMT remain constant across systems, making the entropy bound one of the few hyperparameters that requires some tuning. However, tuning Eent straightforward in practice and described in detail in Appendix C.4.
Comparison of CMT to MD To compare CMT to MD, we analyze the mean absolute error (MAE) of the free energy profile along the & dihedral angle, which is the slowest mode of the system, over the course of an MD simulation. We visualize the MAE as a function of target energy evaluations in Figure and compare with the MAE obtained by CMT (using x 10 target evaluations). MD requires approximately one order of magnitude more target evaluations to yield a similarly well-equilibrated free energy profile as CMT.
C FURTHER DETAILS ON CONSTRAINER MASS TRANSPORT
In this section, we provide some additional details on constrained mass transport, including practical tips for hyperparameter selection, implementation and the connection to importance-weight variance.
C.I DUAL OPTIMIZATION IN PRACTICE
The concavity of the dual functions permits the use of any suitable nonlinear optimization algo- rithm. For one-dimensional dual optimization, we employ the bounded Brent method (Brent, 2013),
27

---

Published as a conference paper at ICLR 2026
10 samples
10 samples RW
107 samples
10a samples RW
Ground truth
FAB
CMT (ours)
Figure 9: TICA projections for ground truth (MD), FAB, and CMT on ELIL tetrapeptide using 10a and 10" samples with and without reweighting to the target distribution. The TICA projection of CMT with 10" samples is visibly more accurate, in contrast to FAB's, which appears slightly blurred or smoothed. For computational reasons, the Wasserstein metrics can only be computed using 10 samples. The corresponding TICA projections shown in the first column and clearly lack sufficient detail to accurately assess TICA projection quality.
-004.7
135 [X]
-504.8
CMT
TA-BG
-504.9
1.5
bonund
Entropy: bound Cust
Figure 10: NLL and ESS performance with different entropy bounds for CMT alanine hexapeptide.
implemented via scipy. aptimiiee library's default ID algorithm due its robustness and efficiency. A minimal working example on how a Lagrangian multiplier is estimated is given in Code Example 1. For 2D duals, we use sci scipy oecipy.ottiim iptimi ze. ze.mi miniiizze with the L-BFGS-B algorithm (Zhu et al., 1997), one SciPy's default quasi-Newton algorithms. There, additionally passed the dual gradient function, which we obtained through automatic differentiation. Due to the constraints X, 77 2 0, and to avoid numerical overflow, scipy- opt imi ze mi minniiizi ze requires an initial guess, which we set to x 10 20, value chosen be close to the lower bound.
(Virtanen et al., 2020), which is the
bound both optimizers to stay within the interval |0, 10 The method
We note that the pure dual optimization is very fast in practice. For alanine dipeptide, dual optimization time makes up approximately 0.01 % of the total training time. Details can be found in Appendix D.4.
28

---

Published as a conference paper at ICLR 2026
0.150
MD
0.125
CMT (ours)
MAE Free Energy
0.100
0.075
0.050
0.025
10
Number of target evaluations
Figure 11: Comparison of CMT with molecular dynamics for alanine dipeptide. We show the mean absolute error of the free energy profile along the & dihedral angle, which is the slowest mode of the system, as a function of the number of target evaluations.
C.2 CMT WITH NORMALIZING FLOWS
Algorithm 1 allows for multiple variants. Algorithm 2 presents a specific instance of Algorithm 1, in which the KL-divergence employed to update the normalizing flow model to the next intermediate distribution at each annealing step by adjusting its parameters. The dual function Peu, is defined as in Equation (11), with Equation (16) providing the sample-based approximation. As is observable
(i+1)
in Algorithm& 2, the optimization of the Lagrangian multipliers and the computation of importance weights only add negligible computational and memory overhead, as they reuse the same buffer samples subsequently employed to update the flow model itself.
C.3 CONNECTION TO SEOETANNEEWEIIGHT VARIANCE
In this section, we show that using the trust-region constraint yields an approximate lower bound for the effective sample size between any consecutive distributions qi and 9:-1- This approximate lower bound only depends on Str-
This idea parallels common practice in sequential Monte Carlo (SMC). As stated by Chopin al. (2023), line search on the >-ddivergence between consecutive distributions is commonly used to control the ESS. By restricting the -diierrgeeeee SMC ensures that the importance weights do not degenerate, effectively enforcing a minimum ESS between intermediate distributions. Similarly, our trust-region constraint guarantees conservative lower bound on ESS, providing principled way to adaptively select the sequence of distributions.
The variance of the importance weights
(x)) (x)
((2))
with trust-region constraint (2)
(Az)
(z))))
p(x)
with trust-region entropy constraint (9)
(Asii))
between normalized consecutive distributions is closely connected the effective sample size via
1
ESS(gi, (i+1)
+ Varg.
0i (a:)
(=(x))
also explained in Appendix D.3. The relation Var (a++1 (t))/((z)) = x² (2((+11))) (Chopin et al., qs
2020) and the well-known Taylor approximation x (gi+1 (4.) ~ 2DxL(q+1 (g:) (Cover, 1999) lets us rewrite the effective sample size in terms of the Kullback-Leible (KL) divergence between 9++1 and yielding
ESS (++1) 2
1 2D (1:1
29

---

Published as a conference paper at ICLR 2026
Code Example I: Minimal working example of the dual optimization for objective (2).
import numpy np import torch
from scipy.optimize import minimize_scalar
def estimate_log_är
model_log_praie Torch.Tenser
target_log_prob: torch.Tensor
tr_mul: float,
TorchhTensor.
"""Estimate log-partition function of next intermediate density""
log_ torch.tensors (target_Log prob. shape .log) .log()
log iw (target_ log prob preb model poodel_log__ii prob) / (1 tr tr_mul)
log_2 torch. torch.logsumee logsumexp (iog_iw, dim=0) 1og_N
return log_2
def Bet_beet__k_mmutipplee
model_log_pran torch.Tensen target_Log_pre Corch.Tenser. eps float, max multiplier: float le10,
float:
"*"Finds the best Lagrangian multiplier by maximizing the dual""M
define dual function (dependent Lagrangian multiplier) def dual float):
log_2 estimate__og__
model_log,
tr_mul=tr_mul,
dual_value (1 + tr_mul) * log_ tr_mul eps_tr
return dual lue. u it
neg qual lambda mul: -uual (mul) concave
res
neg_dual,
bounds= ( 0, max_mlltiplerr)
method="Bounded"
best_tr_mjj float (res.x) return Dest_tr_mij
as an approximation for the effective sample size. This approximation is justified under the assumption that a++1 is close to q, condition that is satisfied by the design of the problem for small trust- region bound Eur > Due to gi+1 being the optimal solution to an objective with the constraint (Kl(gll4) Etr, the constraint must also hold for q 4++1 resulting in the approximate lower bound
ESS(g; (1+1) 2
(21)
1 25tr
for the effective sample size of the importance weights with equality all but the last step.
This upproximate lower bound justifies the use of Monte Carlo approximations in Section 3, helping to stabilize training independent of the problem's dimensionality. Empirical results on this property provided in Appendix B.
30

---

Published as a conference paper at ICLR 2026
Algorithm 2 CMT with normalizing flows
Require: Initial normalizing flow - unnormalized target density ý, buffer size N, batch size B, chosen number of annealing steps 7, number of iterations K per annealing step, trust-region bound Etz, entropy hound Eent
for <- O.... O......- do
Draw N samples In qe, (4),, evaluate 4a, P(an)
Initialize buffer g = (mns q0: (In),, p(an))n==
Compute multipliers As ./. = arg max, meRt Perrfont (A, 1j) using B
Define 90 98,
Update flow to fit 9-+1= for l <- O,..., K do
// Compute importance weights from R
Retrieve mini-batch by (In, qo, (Im), ()xx.))) from (3()
Compute (g++1 (2,, )) from by and multipliers A 7J+
Compute importance weights (w.))== wn (zx)(z))/00, (xn) for 1... .B
// Update flow model:
Define loss 4 Σ w., log q0,, (æn) // loss derived from DeL. (91.(99+114000. Update parameters from 02 to 0,, using lx
Define qp: = 4
return q0, R P
C.4 GUIDANCE FOR HYPERPARAMETER SELECTION
Owing to the relationship between the trust-region constraint and the variance of importance weights, it may appear sensible to choose very small trust-region bound. While doing so indeed improves the quality of the importance weights by lowering their variance, it simultancously increases the number of intermediate annealing steps required for convergence. In practice, we observed this to constitute trade-off between training stability and computational efficiency. For this reason, we adopted fixed trust-region bound of Sir 0.3, which consistently yielded the best performance across all systems under fixed number of target evaluations.
Assuming the target entropy were known and the number of intermediate densities fixed to e N+, one could select the entropy bound
H(go)
Hp)
Eent
where the inequality allows for the possibility that the entropy constraint becomes inactive earlier during training. practice, however, the target entropy rarely available. The design of our entropy constraint mitigates this issue by restricting the search domain for Sent to R This restriction makes straightforward to identify suitable range for Sent- Due to the approximately linear entropy decay, empirically validated in Figure 5, the relationship between the iteration at which the entropy constraint becomes inactive and the entropy bound Eent is likewise approximately linear, which makes tuning Eent straightforward.
Across all molecular systems, suitable values for Eent fall within the range [0.8, 1.8). We aimed for the entropy constraint become inactive after roughly half of the total training iterations, this consistently yielded the best final performance. The entropy bound was tuned only to the first decimal place, with most other hyperparameters remaining fixed across systems.
31

---

Published as a conference paper at ICLR 2026
Str 0.01
Str 0.05
T, 1/o,
Etr = 0.1
Str = 0.2
εtr = 0.3
1000
2.00
300
400
100
200
300
400
Anmealing Stopes
Anncaling Steps
Eent 1.0
Eent 1.1
Eent = 1.2
02.
0.5
Eent 1.3
Cent 1.4
0
Eent 1.5
0
100
200
500
100
100
2000
-000
400
Aunesling Stepe
Annealing Stepe:
Figure 12: Visualization of 3; and T, = 1//0; over the annealing steps on alanine hexapeptide for different trust-region bounds (first row) and entropy bounds (second row), each evaluated while the other bound kept fixed (Et, = 0.3 and Sent 1.4). Varying the trust-region bound visibly changes the convergence speed of both (B:); and (Ta); along the annealing path, whereas varying the entropy bound seems to lead similar behavior for (8,), but strongly affects the temperature behavior (Ti);- We further emphasize that the displayed 3, does not accurately represent the actual "closeness" of q. to the target distribution p during training. Each intermediate density is only approximated via l. arg min, D(g; q). the sequence (4:) y does not follow the analytical annealing path exactly. As a consequence, the calculated evolution of 83 does not provide a reliable absolute measure of how close yi is p in terms of KL divergence.
C.5 ANNEALING PATH VISUALIZATIONS
To complement Figure 1, we provide visualizations of the sequences (as); and (6,), along the annealing path
1
in Figure 12, illustrating the effects of different trust-region and entropy bounds.
D EXPERIMENTAL SETUP
D.1 ARCHITECTURE
Our normalizing flow architecture closely follows the one used in previous works (Midgley et al., 2022: Schopmans & Friederich, 2025; 2024). We represent the conformations of the studied molecular systems using internal coordinates based on bond lengths, angles, and dihedral angles.
We 8 pairs neural spline coupling layers based on monotonic rational-quadratic splines (Durkan et al., 2019). The splines map from [0, 1] to [o. 1] using 8 bins. We use random mask to select transformed and conditioned dimensions in the first coupling of each pair, and the corresponding inverted mask for the second coupling. The dihedral angle dimensions are modeled with circular splines (Rezende al., 2020) to respect their topology, with random (fixed) periodic shift applied after each coupling layer. The parameter networks that calculate the spline parameters in each coupling are fully connected neural networks with hidden dimensions [256, 256, 256, 256, 256] and ReLU activation functions. To capture their periodicity, dihedral angles v/, are encoded as (cos v/, sin thx) when passing them to the parameter network.
As the base distribution of the normalizing flow, angles and a Gaussian truncated to [0, 1] with mean = 0.5 and standard deviation or = 0.1 for the bond lengths and angles.
uniform distribution in [0, 1] for the dihedral
32

---

Published as a conference paper at ICLR 2026
Table 6: Overview of the molecular systems and corresponding force field parametrization.
NAME
SEQUENCE
NO. ATOMS
FORCE FIELD
CONSTRAINTS
AMBER FF96
ALANINE DIPEPTIDE
ACE-ALA-NME
22
WITH OBC1
NONE
IMPLICIT SOLVATION
AMBER99SB-ILDN
ALANINE TETR APEPTIDEE
HYDROGEN BOND LENGTHS
ACE-3-ALA-NMA
WITH AMBER99 OBC
IMPLICIT SOLVATION
AMBER99SB-ILDD
ALANINE HEXAPEPTIDE
ACE-5-ALA-NME
62
WITH AMBER99 OBC
HYDROGEN BOND LENGTHS
IMPLICIT SOLVATION
AMBER99SB-ILDN
ELIL TETRAPEPTIDE
HYDROGEN BOND LENGTHS
GLU-LEU-ILE-LEU
75
WITH AMBER99 OBC
IMPLICIT SOLVATION
We follow Schopmans Friederich (2025) to map the internal coordinates to the range [0. 1] of the spline transformations: Dihedral angles are divided by 2T. Bond lengths and angles are shifted and scaled as 1// = (1/ 7//,minn/ a + 0.5, where T/mmin is obtained from a minimum energy structure after energy minimization. was set to 0.07 nm for bond lengths and 0.5730 for angle dimensions.
The studied molecular systems have two chiral forms (mirror images), L- and R-chirality, while in nature, one almost only finds the L-chirality. To constrain the generated molecular configurations to L-chirality, we constrain the spline output ranges of the relevant dihedral angles (Schopmans Friederich, 2025). Similarly, some atoms and groups (such the hydrogen atoms in CH3 groups) permutation invariant in the force field energy parametrization, but have a preference in the ground truth molecular dynamics data due to very large barriers. Similarly to the chirality constraints, constrain the splines such that only the permutation found in the ground truth data can be generated (Schopmans & Friederich, 2025).
D.2 TARGET DENSITIES
The goal of all our experiments is to sample molecular systems at 300 An overview the studied molecular systems, including their force field parametrization, is given in Table 6. We explicitly note that the largest studied system, ELIL tetrapeptide, does not contain capping groups, in contrast to the other three systems.
The energy evaluations during training were performed with the OpenMM 8.0.0 (Eastman et al., 2024) CPU platform, using 18 workers in parallel.
Following previous work (Midgley et al., 2022; Schopmans & Friederich, 2025), we use a regularized energy function avoid large van der Waals energies due atom clashes:
E, log( E Enigh + 1) Enight logí Emux
S Enigh:
Ereg. (E)-
if Enigh < < Emax.
(22)
Brigh + 1) Enight
if E > Emax-
We set Frigh
x 10 and Emux 1 x 1020 (Midgley et al., 2022).
DATASETS
We employed extensive molecular dynamics simulations to generate training and validation datasets with 10 samples and test datasets with 107 samples to calculate the metrics reported in Table 1. For each system, one molecular dynamics simulation was used to construct the training and validation dataset, and separate simulation was performed for the test dataset. We used slightly updated simulation parameters compared to Midgley et al. (2022) and Schopmans & Friederich (2025) to obtain better-equilibrated data, which became especially relevant for alanine dipeptide, where
33

---

Published as a conference paper at ICLR 2026
noticed small bias in the energy histogram of the MD data provided by Midgley al. (2022) (Stimper et al., 2022a) and used by Schopmans & Friederich (2025).
1. For alanine dipeptide, molecular dynamics simulations LangevinMiddle integrator with a time step of 1 fs. The system was initially equilibrated for 200 ns, after which production run of 5 us was performed.
carried out using the OpenMM
For alanine tetrapeptide, alanine hexapeptide, and ELII tetrapeptide, we conducted replica- exchange molecular dynamics (Sugita & Okamoto, 1999) simulations at temperatures of 300.0 K, 332.27 K, 368.01 188.011 K, 407.60 K, 451.44 K, and 500.0 K using the OpenMM LangevinMiddle integrator with time step of fis. Each replica was first equilibrated independently for 200 ns without exchanges, followed by an additional 200 ns of equilibration with exchanges. Production simulations of us per replica were then carried out. For dataset construction, we used the trajectory from the 300.0 replica.
D.3 METRICS
In this section, we present metries used for both theoretical analysis and experimental evaluation. Unless stated otherwise, all metrics were evaluated using 10" samples.
D.3.1 TASK-AGNOSTIC METRICS
We begin with the forward metrics evidence upper bound (EUBO) and the related negative log likelihood (NLL). We then turn to the reverse metric evidence lower bound (ELBO), discuss its limitations in the context of Boltzmann generators, and conclude with the effective sample size (ESS) measure of sample quality. For details on these metrics, refer to Blessing al. (2024).
Evidence upper bound (EUBO) and negative log-likelihood (NLL). EUBO and NLL are both forward metrics, each equivalent the forward KL divergence up an additive constant, They defined via
p(x)
DKL(pll)) Ep()) log
log Z const.
((x)
BUBO
= ppp( [log ((x)]
(())
NLL
const. W.r.l.
Evidence lower bound (ELBO). The ELBO is defined
ELBO Eg(e)) log log q(x)),
where g is the variational distribution and p the target distribution. The ELBO provides a lower bound on log Z, while the EUBO provides upper bound on log Z. Typically, both bounds are reported, as the difference between them, the EUBO-ELBE gap, gives a useful measure of how tight the variational approximation is. A small gap indicates that the variational distribution closely approximates the true target, while a large gap signals that the approximation may be poor.
Important note ELBO. In experiments, found that the ELBO is not reliable metric for comparing molecular Boltzmann generators. The Boltzmann distribution of molecular systems contains extremely steep regions due the repulsive term of the Lennard-Jones interaction, which grows with 1/r12 when the distance r between two atoms becomes small. Boltzmann generators typically cannot fully capture this behavior, resulting in a small number of samples with extremely small log importance weights. These outliers dominate the mean, making it unrepresentative of actual model performance. This issue does not affect metrics such as reverse ESS (and thus importance sampling performance), which are computed directly on the importance weights rather than their logarithms, nor forward metrics evaluated under the support of the ground truth data (e.g., EUBO). To mitigate this problem, clamped the lowest 0.01% 0 of log importance weights to the highest observed value among them when computing the ELBO. While this approach works reasonably well for alanine dipeptide. it is insufficient for the other systems, which likely require However, we decided to keep the same threshold for all systems for consistency. Consequently, we caution against using the ELBO for quantitative comparisons molecular Boltzmann distributions.
aggressive clamping.
34

---

Published as a conference paper at ICLR 2026
Ground truth
Reverse
Alanine Hexapeptide
10
free energy Kg?
422
d2
Figure 13: Illustration of mode collapse in Ramachandran plots. The figure highlights high-energy regions where mode collapse often occurs (e.g., under reverse KL training) but which remain important for accurately capturing of the peptide backbone distribution. Mode collapse is visible as the absence of density in these regions.
Effective sample size (ESS). The ESS defined as
ESS b)
a E P(RR))
Vara(=)
b(x)
a(x)
Closely following the notation of Blessing et al. (2024), the reverse ESS
Z,
p(x))
ESS(g.p)
with
Z, Bg(e))
2
[q(x)
E
gx(x)
p(m)
((m))
can be directly estimated via Monte Carlo using samples from the model q and the unnormalized target p.
Following Midgley al. (2022); Schopmans & Friederich (2025), for reverse ESS, we clipped the top 0.01% importnce-weightss, setting them to the smallest value among them for numerical reasons. Furthermore, ESS computed using the regularized energy function, defined in Equation (22).
Although forward ESS could be computed using samples from the target distribution, Schopmans Friederich (2025) found it to be extremely sensitive to the chosen clipping threshold and prone to instability. Consequently, only the reverse ESS was used, even though it may not fully capture phenomena such as mode collapse.
D.3.2 TASK-SPECIFIC METRICS
Ramachandran plots. A Ramachandran plot visualizes the 2D log-density the joint distribution of a pair of dihedral angles in peptide's 's backbone. For more details, refer to Schopmans & Friederich (2025). These plots are used to visualize peptide's ' main degrees of freedom and are likely to show mode collapse if it occurs. A Ramachandran plot effectively a histogram of the occurrence of dihedral angle values and computed solely from model or ground-truth samples.
How mode collapse be detected and where the high-energy density region is located in Ramachandran plot illustrated in Figure 13.
For alanine tetrapeptide, alanine hexapeptide, and the ELIL tetrapeptide, which contain multiple backbone dihedral angle pairs, we always show the pair exhibiting the most pronounced deviation from the ground truth, which the same across methods. Among the four runs made per method in Figure 4, we selected the one with the lowest Ram KL value. For Figure 3, we always selected the run with the highest Ram KL value illustrate that variations with fewer constraints are more likely to exhibit mode collapse.
Ramachandran KL (Ram KL) and Ramachandran KL with reweighting (Ram KL w. RW). Ram KL and Ram KL RW can be used to quantify deviation from the ground truth Ramachandran histograms. Following Midgley et al. (2022); Schopmans Friederich (2025), we computed the forward KI. divergence between the Ramachandran histograms from ground truth and model samples
35

---

Published as a conference paper at ICLR 2026
Table Computational cost (in hours) of FAB, TA-BG, and CMT on different molecular systems, reporting pre-training, training and total time. For all systems, we performed four parallel runs single node equipped with 4x NVIDIA A100-SXM4-40GB and 2x AMD EPYC Rome 7402 CPU with 48 cores in total.
SYSTEM
METHOD
PRE-TRAINING
TRAINING
TOTAL
FAB
18.1h
18.1 h
ALANINE DIPEPTIDE
TA-BG
6.2h
16.9h
23.2h
CMT (OURS)
20.4h
20.4
FAB
19.5
19.5
ALANINE TETRAPEPTIDE
TA-BG
7.2h
24.3h
31.5h
CMT (OURS)
21.7h
21.7h
FAB
41.0h
41.0h
ALANINE HEXAPEPTIDE
TA-BG
22.1
43.7
65.8
CMT (OURS)
64.3h
64.3h
FAB
56.6h
56.6h
ELIL TETRAPEPTIDE
TA-BG BG
22.3
116.1
138.4
CMT (OURS)
138.2h
138.2h
(Ram KL). For this, we used 100 x 100 bins and 1 x 10 samples. Additionally, we also computed a reweighted version of the metric (Ram KL RW) where the model samples were first reweighted to the target distribution before generation of Ramachandran histograms. For the larger systems, where more than Ramachandran plot exists, reported the average.
Ramachandran total variation (Ram TV) The Ram TV and Ram TV RW (with reweighting) distance complement the existing Ram KL and Ram KL w. RW metrics. While the KL-based metrics are asymmetric and computed with respect to the support of the ground-truth distribution, the total variation metrics are symmetric, allowing them to capture discrepancies both within and outside the support of the ground truth. This provides a more balanced assessment how accurately the model represents the overall distribution. The total variation distance on Ramachandran plots, which are 2D discrete normalized histograms, is defined as
TVD(P, Q = Σ Poj - Qiylaig
where P and Q are the Ramachandran plots of the target and the model q. respectively. The factor az corresponds to the area of bin (i, in our setting, all bins have equal area dij = (2Tr (2x/1100) 2
TICA KL Similarly to the Ramachandran KL, we also computed the TICA KL on the TICA histograms, using the range from the test dataset define the domain.
Wasserstein distance We report several Wasserstein-2 distances for evaluation: on the Ramachan- dran plots (Ram T-W2), the ground-truth energies of the predicted samples (E-W),, and on the TICA vectors (TICA-W)). Details on the Wasserstein metric and its implementation can found in Tan et al. (2025a). Loosely speaking. the Wasserstein-2 distance quantifies how much probability mass must be transported to transform one distribution into another. We apply it to the periodic dihedral angle pairs forming the Ramachandran plots by defining the distance the resulting torus, thus respecting the periodicity of the angles. For details, see Tan et al. (2025a) (App. E. For computational reasons, also only used 10a samples, as the Wasserstein distance does not scale well to substantially larger sample sizes, which limits sensitivity to mode collapse in high-energy metastable regions.
D.4 COMPUTATIONAL COST
Total training time
To provide an indication of computational cost, we report extrapolated training times for FAB, TA-BG, and CMT across different molecular systems in Table 7. Because training
36

---

Published as a conference paper at ICLR 2026
Table 8: Dual optimization time of CMT on different systems per annealing step and in total.
SYSTEM
PER ANNEALING STEP
TOTAL
ALANINE DIPEPTIDE
53.31 ms
10.6s
ALANINE TETRAPEPTIDE
52.48 ms
10.50
ALANINE HEXAPEPTIDE
61.25 ms
24.50s
ELIL TETRAPEPTIDE
77.17 ms
61.74
Table 9: Number of flow parameters for each system. The number of parameters is completely determined by molecular system's 's size, as the architecture the same across all systems.
ALANINE DIPEPTIDE ALANINE TETRAPEPTIDE ALANINE HEXAPEPTIDE
ELII
NUMBER OF PARAMETERS
7 512
952 376
12 124616
13 727 952
molecular Boltzmann generators is computationally expensive, the reported values are based on extrapolation and reflect pure training time only, excluding evaluation steps and other overhead.
Compared to Schopmans Friederich (2025), we increased the total number of target evaluations and gradient descent steps for TA-BG to match those used for CMT, ensuring fair comparison. While this modification improves the final performance of TA-BG, also raises its computational cost to level comparable to CMT. Since both CMT and TA-BG rely on similar training routines, their overall computational demands are generally comparable.
The comparatively high computational cost observed for CMT may be reduced through further optimization, for example by refining the loss function or decreasing the number gradient descent steps and target evaluations. Its adaptive schedule and use of importance weights rather than resampling may also improve robustness when using smaller buffer sizes compared to TA-BG. We further emphasize, that our implementation represents specific realization of broader algorithmic framework, and additional efficiency improvements may be possible.
For FAB, report training time using updated number of target evaluations but retain the original number of gradient descent steps from (Midgley et al., 2022), thereby providing a low-computational- cost baseline.
The dual optimization time all methods is negligible and summarized in Table 8.
Inference time The inference time of all methods is identical, as they share the same underlying normalizing flow architecture, with its size varying only between systems but not between methods. Details on the architecture are provided in Appendix D.1.
D.5 HYPERPARAMETERS
Hyperparameters play a crucial role in the performance of all models. Common hyperparameters include the choice of optimizer, learning rate, batch size, gradient steps, and weight decay. Below, we provide a description of the hyperparameters for each method, emphasizing any method-specific choices.
All experiments employed the Adam optimizer (Kingma & Ba, 2017). Our implementation builds on the Python packages bgflow (Noé & со workers, 2025), nflows (Durkan et al., 2020), and PyTorch (Paszke et al., 2019), The number parameters the normalizing flow architecture for each system summarized in Table 9.
CMT
We refer to Table 10 for the general and -peccfc- hyperparameters CMT. The number of annealing steps / defined implicitly by the total number of gradient descent steps and the number of gradient descent steps per annealing step, and is summarized in Table 11.
37

---

Published as a conference paper at ICLR 2026
Table 10: Hyperparameter settings for CMT (general and method-specific) for all systems.
HYPPEPAARMEETES
TETRAPEPTIDE
HEXAPEPTIDE
ELIII
BATCH SIZE
1000
1000
2000
2000
LEARNING RATE
4 10->
S 10->
5 X 10
5 10
LR SCHEDULER
COSINE
GENERAL
GRADIENT DESCENT STEPS
1000000
100000
8000000
6000
WEIGHT- DECAY LR LINEAR WARMUP STEPS
1000
1000
1000
1000
100.0
100.0
RRUST-REGION BOUND
0.3
0.3
D.3
0.3
METHOD-
ENTROPY BOUND
0.8
1.8
1.4
0.7
BUFFER
000)
500
0000
1000 00U
GRADIENT DESCENT STEPS
2000
2000
2000
2000
PER ANNEALING STEP
Table 11: Fixed number annealing steps / for CMT for each system. The number is implicitly defined by the total number of gradient descent steps and the number gradient descent steps per annealing step. We keep it fixed strictly control the computational budget for fair benchmarking.
ALANINE DIPEPTIDE ALANINE TETRAPEPTIDE
ALANINE HEXAPEPTIDE
ELII
NUMBER OF ANNEALING STEPS
200
200
400
800
TA-BG
Table 12 summarizes the hyperparameters for the pre-training of TA-BG (Schopmans Friederich, 2025) using the reverse KL divergence.
After pre-training, the temperature is annealed with a geometrically decaying temperature sequence and the hyperparameters summarized in Table 13. The TA-BO TA- experiments alanine dipeptide and alanine tetrapeptide used the geometric temperature annealing sequence
1200 K + 1028.0 69 K - 881.84K + K - 648.04K + 555.52K
476.22 4 349.96 K 300.00 K 300.00 K.
Including additional finetuning step per temperature, TA-BG employs the temperature sequence
1200 K K > 1028.69 K 881.84 881.84K K
755.95 K - 648.04 K + 648.04 + 555.52 K + 555.52 K + 476.22
- 476.22K + 408.24K - 408.24K. + 349.96 K + 349.96 + 300.00 -+ 300.00
on alanine hexapeptide. On the ELIL tetrapeptide. reverse KL pre-training suffers from mode at 1200 K. Therefore, the temperature annealing starts at 3000 K, resulting in the temperature sequence
3 2 + 2573.09 2573.09 - 2206.93K 2206.93
+ 2206.93 > 1892.88 K + 1892.88 K > 1892.88 > 1623.52 + 1623.52
1623.52 K 1392.49 K + 1392.49 1392.49K 1194.33 + 1194.33 K
+ 1194.33 K + 1024.37 K + 1024.37K + 1024.37 - 878.60 K + 878.60
K 753.57 K 753.57 753.57K 646.34 K 646.34
- 646.34K + 554.36 K + 554.36 + 554.36 + 475.48K + 475.48K
475.48K 407.81 K 407.81 407.81K 349.78K 3
349.78 - 300.00 - 3.0.00
FAB
The used hyperparameters for FAB (Midgley al., 2022) can be found in Table 14. Furthermore, used a step size of 0.05 for the Hamiltonian Monte Carlo (Duane et al., 1987) transitions. For details on the method and its hyperparameters, refer to Midgley al. (2022).
38

---

Published as a conference paper at ICLR 2026
Table 12: Hyperparameter settings for TA-BG pre-training for all systems.
HYPERRARRMETERS
TETRAPEPTIDE
TARGET TEMPERATURE
1200
1200
1200K
3000
BATCH SIZE
256
256
512
512
LEARNING
10-
1 10
10
10
LR SCHEDULER
COSINE
COSINE
COSINE
COSINE
GENERAL GRADIENT DESCENT STEPS
100.000
100000
250
250000
WEIGHT DECAY
10-3
1 10
10
10
LINEAR WARMUP STEPS
1000
1000
1000
MAX GRAD NORM
100.0
100.0
100.0
100.0
NO. HIGHEST ENERGY
10
10
20
20
VALUES REMOVED
Table 13: Hyperparameter settings for TA-BG (general and method-specific) for systems.
HYPERPARAMETERS
ALANINE DIPEPTIDE ALANINE TETRAPEPTIDE
ALANINE HEXAPEPTIDE
ELIL
BATCH SIZE
2048
4096
2048 5 10
2048 5 10-6
GENERAL
LEARNING RATE
5 10-5
x 10
LR SCHEDULER
BUFFER
74400
7 74400000
15 111 111
100 0CU
METHOD-
BUITER RESAMPLED
20000000
2000
2
100000000
SPECIFIC
GRADIENT DESCENT STEPS
30
PER ANNEALING STEP
Table 14: Hyperparameter settings of FAB (general and method-specific) for all systems.
HYPERRARAMETERS
ALANINE DIPEPTIDE ALANIEE TETRAPEFTIDE
ÁLANINE HEXAPEPTIDE
ELIL
BATCH SIZE
1021
1021
1021
2018
LR SCHEDULER
COSINE
COSINE
COSINE
COSINE
GENERAL
GRADIENT DESCENT
50,000
25:000
WEIGIT-DECAY LR LINEAR WARMUP STEPS
10
1 10-5
1 10-5
x 10-3
1000
1000
1000
1000
MAX GRAD NORM
1000.0
1000.0
1000.0
1000.0
METTIOD-
No. INTERMED. DIST. HMC STEPS
16
FORWARD AND REVERSE KL
This section reports the used hyperparameters for training with the forward KL divergence on MD data (Table 15) and the hyperparameters for training with the reverse KL divergence (Table 16). A description of how the MD data was obtained can be found Appendix D.2.
39

---

Published as a conference paper at ICLR 2026
Table 15: Hyperparameter settings of forward KL training using MD data for all systems.
HYPERPARAMETERS
ALANINE DIPEPTIDE ALANINE TETRAPEPTIDE ALANINE HEXAPEPTIDE
ELII
BATCH SIZE
1021 10
1021
1024
1921
LEARNING
5 10-5
5 10-3
10
LR SCHEDULER
COSINE 100000
COSINE
COSINE
COSINE
STEPS
Table 16: Hyperparameter settings of reverse KL training for all systems.
HYPERPARAMETORS
ALANINE DIPEPTIDE ALANINE TETRAPEPTIDO
ALANINE HCXAPEPTIDE
ELIL
BATCH SIZE
1024
1024
1024
1024
LEARNING RATE
1 10
1 10
1 x
1
SCHEDULER
GENERAL
GRADIENT DESCENT STEPS
STEPS
250
250000
250
250000
WEIGHT-DECAY
LR LINEAR WARMUP STEPS
1000
1000
1000
1000
MAX GRAD NORM
100.0
100.0
100.0
100.0
NO. HIGHEST ENERGY
VALUES REMOVED
40


Reviews:
[{"review_id": "PXc4TwIG5u", "reviewer": "Reviewer_VDn6", "summary": "This paper tackles the problem of sampling from unnormalized densities with Boltzmann generators. It proposes Constrained Mass Transport (CMT), which introduces trust-region and entropy constraints to effectively learn an adaptive annealing schedule that interpolates between the prior and the target distributions. The proposed method is evaluated on several molecular systems (Alanine dipeptide, tetrapeptide, hexapeptide, and ELIL tetrapeptide).", "questions": "See questions/weaknesses above.", "limitations": "", "rating": 2, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "* The idea of learning an adaptive schedule for the annealing path is conceptually interesting and well-motivated.\n* The paper is well-written and easy to follow. The mathematical details are clearly presented, and the proofs appear sound as far as I could verify.", "weaknesses": "* **Unclear specification of the number of annealing steps.** The number of annealing steps ($I$) appears to be a hyperparameter but is never explicitly stated. From Table 7, it seems to be roughly 200 steps for ALDP (400k total / 2000 per step ≈ 200).\n    * This number is quite large, and it is unclear whether the reported gains over TA-BG stem from the adaptive schedule itself or simply from using many more steps. Ablations comparing CMT to TA-BG with an equivalent number of steps would clarify this.\n    * Additionally, the conceptual role of $I$ is somewhat unclear to me, once the trust-region parameter $\\epsilon_{\\mathrm{tr}}$ is introduced. Since $\\epsilon_{\\mathrm{tr}}$ already constrains the KL divergence between successive distributions and, as a result,  implicitly determines how many steps are required, fixing both $I$ and $\\epsilon_{\\mathrm{tr}}$ doesn’t make sense to me.\n\n* **Metric design and presentation are confusing and seem cherry-picked.** Several of the most representative metrics appear only in the appendix (Table 2). In the main table, the authors report only Rama TV. Why was this specific metric chosen? Even if one chooses to report total variation, shouldn’t the reweighted Rama TV be reported, given that the process is not explicitly pinned to the target distribution? Similarly, the reweighted Rama $\\mathbb{T}\\text{-}W_2$ metric is omitted. The rationale for which metrics appear in the main results versus the appendix should be clarified.\n\n* **Mismatch in optimization budgets across methods.** The number of gradient steps differs substantially between CMT and TA-BG (see Tables 7 vs 8–9 in the appendix). While the authors mention equalizing the number of target evaluations, comparisons should also be made with the same number of gradient steps. Otherwise, it isn’t clear how this plays a role in the reported performance improvements.\n\n* **Hyperparameter selection seems ad-hoc and system-specific.** The trust-region bound $\\epsilon_{\\mathrm{tr}}$ is fixed globally, but the entropy bound $\\epsilon_{\\mathrm{ent}}$ varies per molecular system with no principled tuning rule. Appendix C.4 even notes that entropy tuning is difficult. This per-system manual adjustment undermines the claim of a “learned schedule.” The authors should provide clearer guidance or automatic criteria for selecting $\\epsilon_{\\mathrm{ent}}$, and ablations showing its impact on training stability and performance.\n\n* **Limited intuition about the learned schedule.** The paper emphasizes that CMT learns the annealing path, yet it never visualizes or analyzes what that path looks like. Providing intuition about the learned schedule and explaining why it may be advantageous over those in TA-BG would make the results much more interpretable.\n\n* **No analysis of $\\log Z$ along the annealing path.** Since the learned intermediates depend on Monte Carlo estimates of the partition function, it is important to assess how accurate those estimates are throughout the path. The paper only reports overall $\\log Z$ bounds (ELBO/EUBO) after training, with no diagnostics of these estimates or variance of the weights along the trajectory. Showing how $\\log Z$ (or ESS) evolves over the annealing path would substantiate the claim of the papers and make it easier to evaluate the effectiveness of the proposed adaptive schedule.\n\n* **Missing justification for how CMT avoids mass teleportation.** My understanding of “mass teleportation” follows from [1], where the linear interpolation path results in an ill-conditioned and exploding vector field, which results from the fact that the relative weight of the modes isn’t preserved. However, by this definition, I do not understand how constraining the entropy decay helps. Can the authors clarify what they mean by mass teleportation?\n\n\n[1] Máté, B., & Fleuret, F. (2023). Learning Interpolations between Boltzmann Densities. arXiv preprint arXiv:2301.07388. https://arxiv.org/abs/2301.07388", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2TwLlaBz3X", "reviewer": "Reviewer_hjaT", "summary": "The work “Learning Boltzmann Generators via Constraned Mass Transport” deal with the learning of a Boltzmann distritbution throuhg the minimization of the reverse KL divergence, sometimes named Boltzmann Generators. The authors proposes a variational approach based on a annealing procedure, where the optimization scheme is further constrain by (i) contraining the entropy between successing temperatures and (ii) the divergence between two successives temperature. The author shows that their result present better results than other approach, and show in their experiment that their protocol doesn’t seem to suffer from mode collapse.", "questions": "1. I think it would be a strong added value to be able to apply the method on a toy model suffering from mass teleportation to understand how the optimization scheme managed to mitigate this problem.\n\n    2. It is not discuss if the additional constraint add make the method slower w.r.t. other approaches, and if yes how much slower ? That would be a valuable information.\n\n    3.It looks to me the manuscript overlooked some references which seems closely related to the present work, which might be worth discussing, concerning respectively mode collapse when using KL divergence,( Soletskyi, Gabrié, Loureiro, B. (2025). “A theoretical perspective on mode collapse in variational inference”.); the use of normalizing flows (Gabrié, Rotskoff, Vanden-Eijnden, E. PNAS (2022). “Adaptive Monte Carlo augmented with normalizing flows”); and the use of a well adapted annealing path with a different protocol to avoid mass teleportation (“Fast training and sampling of Restricted Boltzmann Machines”, Béreux et al, ICLR 2025).", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The papers proposes a reasonably simple protocol to optimize over the KL divergence adding further constraints. The authors derive the optimal distribution of their optimization constraint, very similar to the classical reference Neal 2001. Their experiments seem quite convincing that their method seem to perform better than the others without suffering from mode collapse.", "weaknesses": "To my opinion, the authors do not provide a clear explanation on why their method should work best w.r.t other annealing scheme. At least, it was not clearly explained why constraining the entropy and the KL divergence should tackle the mass teleportation. It would be a strong added value to have a analytical or intuitive explanation, possibly on a toy example, on why this new optimization schemes avoid this problem.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "aqUrHw2qYd", "reviewer": "Reviewer_jMV6", "summary": "The paper tackles sampling by learning normalizing flows along an annealing path defined by a sequence of constrained variational problems. At each step, it minimizes $\\text{KL}(q\\|p)$ to the target $p$ subject to two constraints: (i) a trust-region bound on $\\text{KL}(q\\|q_{\\text{prev}})$ and (ii) a limit on entropy decay. The first constraint helps preserve overlap between iterates and the second helps avoid mass teleportation. They evaluate the approach on several high-dimensional molecular systems, including the ELIL tetrapeptide, which they introduce, and show consistent gains, particularly in effective sample size, over recent variational baselines.", "questions": "1) Can the authors clarify or provide intuition for why the entropy constraint mitigates mass teleportation? Is there a theoretical justification?\n\n2) How robust is performance to the optimization of the Lagrange multipliers? Normalizing constant estimation can be tricky, which suggests possible issues, but one could also imagine that errors just \"shift the annealing path\" without hurting final performance.\n\n3) How expensive is the constrained optimization relative to the rest of the training pipeline?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The paper is clearly written and easy to follow. The motivation for combining trust-region and entropy constraints is well explained, and the authors do a good job connecting these ideas to the practical issues of maintaining overlap and avoiding mode collapse during annealing. While the individual components are known, their combination in this specific setting is original and well justified. The experimental section is strong and demonstrates clear empirical gains. Introducing the ELIL tetrapeptide as a new benchmark is also a nice contribution that could help future work in this area.", "weaknesses": "From a technical standpoint, the contribution is somewhat incremental. The main novelty is combining two existing constraints (trust-region and entropy regularization) within a single variational framework, rather than introducing a new theoretical ingredient. This does not detract from the paper given the strong empirical results and clear exposition.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RY2rtBv5CA", "reviewer": "Reviewer_QZLu", "summary": "This paper describes a mathematically elegant approach for imposing both trust-region and entropy constraints on a mass transport problem.  This results in improved learning of the target distribution.", "questions": "How does CMT balance exploration and convergence?\n\nWill there be issues with exploration as systems scale up to proteins with rugged energy landscapes? Or will a uniform distribution on internal coordinates be sufficient?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The paper is well written and motivated. The model is trained using energy evaluations, not data samples, which is a challenging and largely (for useful systems) unsolved problem.  The evaluations are sound and target larger than typical systems for this field (including a more complex peptide data set that will be contributed). Meaningful ablations are performed.", "weaknesses": "If there is a discussion of the computational efficiency of training and inference I can't find it - especially compared to other approaches.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sZpLG1NDUf", "reviewer": "Reviewer_dNiK", "summary": "This paper introduces Constrained Mass Transport (CMT) for learning Boltzmann Generators via a sequence of constrained variational updates. The method proposes to minimize $\\mathrm{KL}(q||p)$ subject to (1) a trust-region constraint $\\mathrm{KL}(q||q_i)≤\\varepsilon_{\\text{tr}}$ and/or (2) a bound on entropy decay $H(q_i)−H(q)≤ \\varepsilon_{\\text{ent}}$. The authors derive closed-form optimal intermediate densities and show these induce geometric, tempered, or geometric-tempered annealing paths when (1), (2), or both are active, respectively. The authors demonstrate gains in performance under matched target evaluation budgets of the method on conformational sampling for Alanine peptide systems (dipeptide, tetrapeptide, and hexapeptide) and a new ELIL tetrapeptide system. CMT improves EUBO, ESS, and Ramachandran TV under matched target-evaluation budgets. Ablations indicate both constraints are needed to avoid mode collapse and stabilize ESS across steps.", "questions": "1. Have you tried Cartesian coordinates? Does CMT’s path construction help maintain overlap in higher dimensions and improve stability in training?\n2. It would be interesting to see how this approach compares to sequential boltzmann generators that are trained on MD data to see if these energy-only methods are approached data-driven. [1]\n3. Any special initialization for $q_0$​? How often do you observe divergence or mode collapse? \n4. For alanine dipeptide and tripeptide, how does this compare to PITA (which is run on all-atom coordinates) [2]\n\n[1] Tan, Charlie B., et al. \"Scalable equilibrium sampling with sequential boltzmann generators.\" arXiv preprint arXiv:2502.18462 (2025).\n[2] Akhound-Sadegh, Tara, et al. \"Progressive Inference-Time Annealing of Diffusion Models for Sampling from Boltzmann Densities.\" arXiv preprint arXiv:2506.16471 (2025).", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-written and easy to understand. \n\n2. The method is conceptually clean. The paper leverages known existing ideas in reinforcement learning and cleverly adapt them to the sampling problem. The constrained objectives yield analytic intermediates and a unified view of annealing paths.\n\n3. Strong empirical results across peptide systems and metrics under fixed energy budgets.\n\n4. Useful ablations highlighting the complementary roles of the trust-region and entropy constraints.", "weaknesses": "1. It would strengthen the paper to report Energy-W2 and TICA-W2 alongside the current metrics. Rama-TV (and Torus-W2 / TICA-W2) are informative at the macro level, but a model can score well on such macro metrics while still missing fine-grained details e.g., bond-length, that Energy-W2 would better capture. TICA-W2 adds a complementary macro view tied to slow collective modes. Including these would provide a more complete picture of both global conformational coverage and local physical accuracy. [1] \n2. Providing non-ML baselines (MD and SMC) for comparison under the same energy evaluation budget would better contextualize the contribution.\n3. Authors provide an analysis for $\\varepsilon_{\\text{tr}}$ vs (ESS lower bounds) but not for $\\varepsilon_{\\text{ent}}​$. It would be interesting to see this.\n\n[1] Tan, Charlie B., et al. \"Scalable equilibrium sampling with sequential boltzmann generators.\" arXiv preprint arXiv:2502.18462 (2025).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)