Enhancing Diffusion-Based Sampling with Molecular Collective Variables
Authors: Juno Nam
Bálint Máté
Artur P. Toshev
Manasa Kaniselvan
Rafael Gomez-Bombarelli
Ricky T. Q. Chen
Brandon M. Wood
Guan-Horng Liu
Benjamin Kurt Miller
Venue: ICLR 2026 Poster (2026)
Keywords: diffusion sampler
generative modeling
conformational sampling
enhanced sampling
collective variables
free energy methods
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: Diffusion-based samplers learn to sample complex, high-dimensional distributions using energies or log densities alone, without training data. Yet, they remain impractical for molecular sampling because they are often slower than molecular dynamics and miss thermodynamically relevant modes. Inspired by enhanced sampling, we encourage exploration by introducing a sequential bias along bespoke, information-rich, low-dimensional projections of atomic coordinates known as collective variables (CVs). We introduce a repulsive potential centered on the CVs from recent samples, which pushes future samples towards novel CV regions and effectively increases the temperature in the projected space. Our resulting method improves efficiency, mode discovery, enables the estimation of free energy differences, and retains independent sampling from the approximate Boltzmann distribution via reweighting by the bias. On standard peptide conformational sampling benchmarks, the method recovers diverse conformational states and accurate free energy profiles. We are the first to demonstrate reactive sampling using a diffusion-based sampler, capturing bond breaking and formation with universal interatomic potentials at near-first-principles accuracy. The approach resolves reactive energy landscapes at a fraction of the wall-clock time of standard sampling methods, advancing diffusion-based sampling towards practical use in molecular sciences.

Published as a conference paper at ICLR 2026
ENHANCING DIFFUSION-BASED SAMPLING WITH MOLECULAR COLLECTIVE VARIABLES
Juno Nam Bálint Máté
Artur P. Toshev Manasa Kaniselvan
Rafael
Ricky T. Q. Chen' Brandon Wood'
Guan-Horng Liu l. Benjamin Kurt Miller
'FAIR at Meta
MIT University of Geneva *TUM 'ETH Zurich
"Core contributors TWork done during internship at FAIR
ABSTRACT
Diffusion-basee samplers learn to sample complex, high-dimensional distribu- tions using energies or log densities alone, without training data. Yct, they remain impractical for molecular sampling because they often slower than molecular dynamics and miss thermodynamically relevant modes. Inspired by enhanced sampling, we encourage exploration by introducing sequential bias along bespoke, information-rich, low-dimensional projections of atomic coordi- nates known as collective variables (CVs). We introduce repulsive potential centered on the CVs from recent samples, which pushes future samples towards novel CV regions and effectively increases the temperature in the projected space. Our resulting method improves efficiency, mode discovery, enables the estimation of free energy differences, and retains independent sampling from the approximate Boltzmann distribution via reweighting by the bias. On standard peptide confor- mational sampling benchmarks, the method recovers diverse conformational states and accurate free energy profiles. We are the first to demonstrate reactive sampling using a diffusion- based sampler, capturing bond breaking and formation with uni- versal interatomic potentials solves reactive energy landscapes at a fraction of the wall-clock time of standard sampling methods, advancing diffusion-based sampling towards practical use in molecular sciences.
saau-iist-principlee accuracy. The approach re-
INTRODUCTION
A central goal of statistical mechanics simulations is to estimate macroscopic thermodynamic prop- erties, such as the energy released in chemical transformation. These simulations typically employ molecular dynamics (MD) or Markov chain Monte Carlo (MCMC), which propagate an atomic structure forward in time or along a chain using energies from quantum mechanical calculations or their approximations (Marx & Hutter, 2000; Frenkel Smit, 2023; Jacobs et al., 2025). Over long times, the distribution of states converges to the Boltzmann distribution v(x) ε exp(-u(x))) where the reduced energy u encodes the chemical system, temperature, and other thermodynamic vari- ables. Sampling from this distribution is deceptively difficult, despite having access its explicit unnormalized form, due to the large number possible atomic configurations and the high en- ergy barriers separating them. Straightforward MD MCMC often requires impractically large number of sequential energy evaluations to capture conformational changes or chemical reactions.
Machine learning (ML)-based samplers attempt to bypass dynamics by reformulating sampling as statistical inference on the Boltzmann distribution, using its log- likelihoddo to draw independent and identically distributed (i.i.d)) samples. Generally, models that approximate thermodynamic ensem- bles are known as Boltzmann Generators. The first work in this area applied normalizing flows (Noé et al., 2019), while recent approaches employ diffusion-based samplers that steer stochastic differential equations (SDEs) whose trajectories yield samples consistent with the target distribution (Zhang& & Chen, 2021; Vargas al., 2023a; Havens et al., 2025). Although training can amortize
More information about neural and di diffusion-based samplers can he found in Section A.2.

---

Published as a conference paper at ICLR 2026
effort across multiple draws, these methods do not inherently reduce the number of energy evalua- tions required to obtain accurate equilibrium statistics compared to MD-based approaches (He et al., 2025). Directly applying diffusion samplers faces well-known pitfall of mode collapse, where training and sampling concentrate on high-probability basins while underrepresenting rare but ther- modynamically important states. This issue is not only about exploration but also about statisties, since reliable free energies and ensemble averages depend critically on assigning correct weights to rare suniiuuatioos, which may be exponentially small. Thus, successful sampling must both reach low-oceupancy modes and attribute the correct equilibrium weights to them.
Data-Based Pretraining
Energy-Based
Sampling Free Energies
Diffusion Sampler Training
Bolizaman ensembile from weighted samples
Samples
Reweighting
Energy
Free energy landscape (PMF) from bias
landscape
Sampling
Well-Tempered Update
Figure Scheme neural enhanced sampling. (Left) Local sampling and pretraining near the reference configuration. (Middle) Energy-basee based sampler training with a CV-space bias that is updated online via well- tempered deposition. (Right) After convergence, the final bias yields the potential of mean force (PMF) along CVs; also, reweighting recovers the Boltzmann ensemble, including samples and free energy differences.
Inspired by enhanced sampling methods, we augment the diffusion- sampler with bias in collective variable (CV) space. CVs are low-dimensional functions of atomic coordinates that cap- ture slow, chemically relevant motions and can be chosen from common system characteristics (e.g., bond lengths or torsional angles) or identified with MI (Sidky al., 2020; Bonati al., 2023). Dur- ing the training of a diffussion-based sampler, the sampler iteratively proposes configurations and re- ceives feedback from their energies. An additional online repulsive potential is maintained over the CVs: regions visited more often accumulate higher bias, which raises their effective energy, discour- aging repeated visits. This flattens free energy barriers over training progress and mitigates collapse onto dominant modes. At inference, importance weights remove the effect of the bias, ensuring that ensemble estimates match the target Boltzmann distribution. The result is broader exploration, in- cluding modes, while preserving statistically correct estimates after reweighting. Furthermore, compared to MD-based enhanced sampling, the diffusion process mixes more efficiently within the biased distribution, enabling faster exploration. Our contributions
follows:
We propose the Well-Tempered Adjoint Schrödinger Bridge Sampler (WT-ASBS), which aug- ments the state-of-the-art diffusion-based sampler ASBS (Liu et al., 2025) with biasing mech- anism inspired by well-tempered metadynamics (WTMetaD, Barducci al., 2008), achieving substantially improved exploration and accurate estimation of free energy differences.
We develop a protocol for applying WT-ASBS to molecular systems, demonstrating (i) accurate sampling of peptide conformations in Cartesian coordinates, where the baseline method fails; and (ii) the first diffusion- sampling of reactive landscapes, achieving shorter wall-clock times than WTMetaD. This establishes a new practical role for diffusion-bar samplers in chemistry.
2 BACKGROUND
2.1 MOLECULAR SAMPLING
Problem setup In the configurational space X' C Rax3 for n atoms with elemental identities Ed with set of element types E, the potential energy function E X E" R defines the Boltzmann distribution at inverse temperature 3 = 1/kg as vxx; a) x exp(-BE(x; a)). In molec- ular simulations, the central challenge is generate configurations representative of the underlying
2

---

Published as a conference paper at ICLR 2026
free energy landscape that describes the relative stability and likelihood of different molecular con- figurations. Three main downstream applications motivate this task: (i) obtaining diverse and statistically meaningful set of configurations for structural analysis, (ïi) accurately computing free energy differences between relevant states, and (iii) uncovering transition mechanisms that govern rare events that are nevertheless thermodynamically essential. For notational convenience, we omit elemental identities after this paragraph and write E(x)) and v(x)).
One rarely aims sample the entirety of X in practice because includes all configurations con- sistent with the chosen composition and potential energy function. Instead, sampling is naturally restricted an accessible component A C X that kinetically connected to reference conformer Jee E X' on the timescales of interest at temperature T by physics. In molecular conformational sampling, A implies configurations that preserve the bond topology of Trer- If one included bond reorganizations occurring on longer timescales (hours to days), A would have to grow to include these states. This restriction ensures that the ensemble reflects the connected set of metastable states and transition pathways relevant to the molecular process of interest, while excluding kinetically inaccessible or chemically irrelevant configurations, such as incorrect molecular chirality.
While basic sampling algorithms aim to produce i.i.d. samples from v(r), the exponential depen- dence of probabilities on E(x) leads to significant undersampling of low-population modes even at moderate energy differences. This is especially important when high-energy transition configu- rations are being studied, Consequently, many effective schemes for molecular sampling generate weighted samples {(x;, W with X; A and W; 2 O, obtained for example via biased simu- lations with reweighting, For any bounded observable f X + IR the self- norm-norma estimator
(a( Σn W.f(X.) No+00 (x((()) A
(1)
W;
recovers expectations under the Boltzmann ensemble restricted to A. We now summarize the re- quirements and discuss the practical realization of these:
Restricted Support: Sampling is confined to the relevant metastable state and transitions.
Configurations in X' \ A are not produced receive weight.
2. Consistency: The reweighted ensemble converges to the Boltzmann ensemble on A, i.e., convergence in Eq. (1) holds for all bounded f.
Local Efficiency: Within any chemically relevant region of A, the procedure yields suffi-
cient sample diversity to estimate observables and capture transition mechanisms.
Requirement 1: Restrained sampling Membership in A is governed by the free energy barriers
that separate modes, not by the stability (state free energies) of individual modes. When sampling performed by emulating physical dynamics (e.g.. MD), high barriers suppress transitions and trajectories remain within A by construction. For sampling procedures that do not generate time- correlated frames, one should enforce the restriction explicitly using restraint potentials that are flat on AA while strongly penalizing X' \ A. This confines exploration to the intended modes while al- lowing efficient coverage within them. Practical considerations are further discussed in Section 3.2.
Requirements 2 and 3: Biasing on the CV space High free energy barriers often align with small set of slow coordinates, so unbiased exploration in A can waste effort on fast motions while undersampling rare but important transitions. CVs address this by mapping configurations a low- dimensional space & X + S C ! with m < n, that capture functionally relevant degrees of firedom. Examples include peptide backbone torsions (ø, v) and the center-of-mass distance between a protein pocket and its ligand. By steering sampling with bias that depends only on {(x), enhance explorations across modes while leaving orthogonal fluctuations unchanged. A bias (s) yields the unnormalized sampling density (I) and importance weights w w(x) as
(x) x expl-3F(xc) - (V(+)))
(2)
w(x)
v(w) (v( & exp(+BV((((z)),
(3)
so unbiased expectations are recovered by the estimator in Eq. (1) with X, ~ by and W; = w(X)), ensuring Requirement 2.
Biasing- based enhanced sampling approaches often construct a bias that raises the effective sam- pling temperature along CV space, introduced by Barducci al. (2008) as the well-tempered target
3

---

Published as a conference paper at ICLR 2026
distribution. We first define the CV marginal D(s) and the potential of mean force (PMF) F(s) from the unbiased CV marginal (defined up to a constant) using the Dirac delta function d:
(s) & fx IV (x))(((())) s) dx,
(4)
F(s) = - log (s)).
(5)
Now, given a bias factor 7 1 that acts scale factor for T, we define the well-tempered bias (sr(s) and the corresponding target distribution pwr(z) from Eq. (2) as
(sT(s) - - (s))
(6)
(2) x exp(- + 3(1 4) (7)
From Eqs. (4), (5) and (7), the well-tempered CV marginal is (s) expl- = [/(
which shows that the sampled CV behaves as if at higher effective temperature Tan = 7/T (Bet 8/2), while the orthogonal conditional remains unchanged: (WT [(xr) = 8) = v( [(x) *)- This selective increase in sampling temperature along the CV satisfies Requirement 3. To construct VwT on-the-fly during simulation, one can stack Gaussian kernels (Barducci et al., 2008), use a kernel density estimator (Invemiizzi & Parrinello, 2020), or obtain the bias variationally (Valsson & Parrinello, 2014; Bonati et al., 2019). These requirements also extends to a variety of molecular enhanced sampling algorithms, as discussed in Section A.1.
Mode exploration Requirement 2 is sometimes relaxed so that we generate samples uniformly across modes without weighting. This is a practically relevant setting that emphasizes capturing the diversity of metastable modes rather than reproducing their exact Boltzmann probabilities. Mode exploration is the usual objective of data-driven generative models for molecular conformations (Jing et al., 2022), protein structures (Watson et al., 2023), and crystalline materials (Jiao et al., 2023; Miller et al., 2024), which are trained to approximate the empirical distribution of representative structures extracted from exploratory simulations or experimental data without Boltzmann weights.
2.2 DIFFUSION-BSSED NEURAL SAMPLERS
Neural samplers Diffusion- based neural samplers learn a control vector field that steers stochas- tic process to gradually transform noise into target distribution, typically specified by an unnormal- ized density or energy function. Unlike iterative generative methods trained on data from implicitly defined distributions
et al., 2015; Ho et al., 2020; Lipman et al., 2022; Albergo et al., 2023), neural samplers can generate samples by querying the energy alone. This ability is especially relevant for Boltzmann distributions (Zhang & Chen, 2021; Berner et al., 2022; Vargas et al., 2023a; Havens et al., 2025). More information about neural samplers is provided in Section A.2.
Adjoint Schrödinger Bridge Sampler (ASBS, Liu et al., 2025) ASBS casts Boltzmann sampling as the Schrödinger Bridge xxp(-BE(X,)) that seeks the optimal control up w.r.t. the path Kullback-Leible (KL) objective
(SB) problem between source (xXp) and Boltzmann target v(X,) x
min, DKL (p" = =p"=0) = Ex-p= [ 1 Ellua (Xl,t)l²dd
(8)
s.t. dx2, = o, (X,,t) dt + o, dW,, Xo ~ (X)o), X1 ~ v(Xi),
(9)
where or is the noise schedule, and p" denotes the law of the SDE in Eq. (9)² ASBS avoids back- propagating through the SDE during optimization by introducing the parametrized corrector (X,) that debiases the effect of endpoint coupling from base distribution to allow the SOC interpretation of Eq. (8), which can be converted to Adjoint Matching (AM) objective for the drift (Havens al., 2025). The corrector is then learned by optimizing the Corrector Matching (CM) objective:
(AM)
arg min. E Eprer 18. [0) t) + (BVE((X, + has (X1))
(10)
(CM) h arg min; Epg, (X,, Vx, log (X, 11
(11)
where := stopgrad(u)) These paired objectives instantiate iterative proportional fitting (IPF, Kullback, 1968) and converge to the SB solution. The expectations use closed-form base condition- als, so no target samples or importance weights are needed, which preserves scalability. Further- more, the control model supports approximate data-based pretraining (see Section C).
The term PMF for F(x) is used interchangeahly with the free energy profilefaandscape along the CV. While the original ASBS formulation is hased SDEs with general drift f(Xe,t) t) + (xuo t), here we only consider the driftless hase process with f 0.
4

---

Published as a conference paper at ICLR 2026
ASBS achieved strong performance synthetic potentials and amortized conformer sampling (mode exploration setting) compared to previous diffusion- based samplers. However, it missed less populated modes in alanine dipeptide sampling due to its mode- mode-seeking behavior, motivating the introduction of an enhanced sampling mechanism this work.
METHODS
Now we augment ASBS with a well-tempered bias on a small set of CVs, updated online from ii.d. model samples during training, instantiating neural enhanced sampler that satisfies the quirements in Section 2.1. We first present the WT-ASBS setup and mechanism together with a convergence guarantee. We then detail practical strategies with data-based pretraining. WT-ASBS training, and post-training reweighting to recover the Boltzmann ensemble (scheme in Fig. 1).
3.1 WELL-TEMPERED ADJOINT SCHRÖDINGER BRIDGE SAMPLER
Algorithm Well- Temppredee Adjoint Schrödinger Bridge Sampler (WT-ASBS)
Require: prior μ. potential energy F(ar), CV E(x), control uo(2; (z, f), corrector ha())
I: Initialize (s(s) t 0
2: for Outer step k = 0, 1, do
3: for Inner step 1 = 0, 1....... do
4- arg ming 1 [ll (Xe, (L(VV(E V2 ho)) Adjoint Matching
5: Φ arg ming Epg. (X1)
Vx, log (X11Xo)11F)
Corrector Matching
6: Sample i.i.d. {x REI ~ pp,
Sk
(0)
E((
7: (1) - (s) + h EE exp(-
A Va(ol() exp( 2
Well-Tempered Bias Update
8:
output Final control up (zx t). final bias (s)
Setup and two-time-scale algorithm We implement WT-ASBS with a well-tempered bias update through stacking Gaussian kernels (Barducci et al., 2008) reach the well-tempered target Eq. (6). Let S C 1 denote a compact CV region of interest, where bias potentials V : S + IR are defined and constructed, which induces the sampling distribution (v x expl-3(E(x)) (E())))) WT-ASBS presented in Algorithm is executed via a two-time-scale scheme:
1. Inner step (ASBS training). For fixed V&, train ASBS with energy V4 of until conver- gence so that the marginal distribution at t = equals vv.,
2. Outer step (well-tempered bias update). Draw a conditionally i.i.d. batch of configurations
from ASBS (x 1,kj== PÝ, project to CV space 8 = <(x 1,6), and update
(8) (i)
(8 (0)), (s) = (s) + h + exp ( - Vi (0) (s,
3
(12)
||s-s' 12
with fixed height h > and Gaussian kernel (s,s, s') exp(
20
We further establish a convergence guarantee showing that WT-ASBS provably approaches the de- sired well-tempered distribution under the two-time-scale update scheme (proof in Section B.1):
Proposition 3.1 (Convergence of WT-ASBS), When Algorithm is performed until convergence, the bias potential Va converges almost surely to V* = -(1 F()) + const, and hence the sampled distribution converges to the well-tempered targel Eq. (7).
Remark 3.1. Proposition 3.1 ensures that, upon completion of training, the PMF along the CVs can be recovered up to a constant from the final bias F(s) - V8) const.
Practical implementation In practice, AM uses replay buffer: gradient signals are not computed directly from samples of P8,1 with the current ug, but rather from subsets drawn from stored samples in the buffer. Assuming the bias deposition per step is small and AM training is efficient enough for generated samples to follow the current bias closely, we can merge the outer and inner sample generation steps by updating the bias during AM training. See Section C for details.
5

---

Published as a conference paper at ICLR 2026
Although our biasing mechanism mirrors classical MD-based enhanced sampling, neural samplers like ASBS generate i.d. samples from the current distribution, eliminating the need to wait for decorrelation before bias deposition occurs. This allows us exploit the bias deposition mechanism more efficiently to explore the regions of configurational space. We further discuss recent works that promote exploration by shaping the annealing path by modifying the learning objective in Section A.2.
3.2
RECIPE FOR EFFICIENT SAMPLING OF MOLECULAR SYSTEMS
Here, we outline three practical ingredients for deploying WT-ASBS to molecular systems.
Warm-start with localized pretraining In this work, consider single reference configura- tion that defines the accessible component A, corresponding to the initial conformer used to start MD or MCMC chains. Although Al Algorithm 1 does not require data samples for training, in the original ASBS, the model can be pretrained with available samples to warm-start and improve train- ing efficiency. Accordingly, we perform short MD simulations (local sampling) to generate samples from the mode containing the reference configuration and apply bridge matching pretraining (Sec- tion C) initialize the control model. During early-stage local conformational exploration without energetic barriers, correlated samplers such as MD or MCMC are more efficient because they move locally on the energy landscape, unlike diffusion samplers that must transform noise to physical samples through hundreds of model evaluations for every sample. Since this is only pretraining stage, high accuracy is unnecessary, and we can further reduce computational cost by using approx- imate energy evaluations (e.g.. classical force fields), faster sampling settings, transfer operator surrogates (e.g., Schreiner et al., 2023) to create the pretraining data.
WT-ASBS training with CVs and restraints Upon adopting ASBS, the algorithm introduces several additional hyperparameters governing training. The first is the choice of www-dimesiional CVs that defines the space over which modes are explored. CVs may be selected either from chemi- cal intuition or from topological changes induced by the underlying transformation, as demonstrated in Section 4. For more complex systems, CVs can also be parameterized by shallow neural net- works and trained from available conformations (Sidky et al., 2020; Bonati et al., 2023). Other hyperparameters follow conventions from prior enhanced sampling studies. However, the improved mode-mixing behavior of diffusion samplers permits the use of smaller ry values, i.e., the effective temperature along the CVs can be lower. Similarly, the Gaussian height h must be reduced, since diffusion samplers generate uncorrelated samples and thus deposit bias more frequently.
To satisfy Requirement 1, we use restraint potentials. The specific combination of restraints is system dependent, but their application can be formalized through the lens of chemical isomerism. Since which types of isomerism are suppressed under physical conditions well understood, these serve natural guide (Canfield et al., 2018). We provide a detailed treatment in Section A.4, and here illustrate the idea with chirality inversion in proteins and peptides. The Ca atom of each amino acid is chiral center: while both L- and D-configurations exist, nearly all natural amino acids only adopt the L-configuration. As the energy function does not inherently distinguish mirror-symmetre configurations, we introduce flat-bottom harmonic potentials on improper torsional angles defining chirality, thereby biasing sampling toward the physically relevant L-configurations.
Sampling and refinement Remark 3.1 implies that the PMF along the CVs can be recovered directly from the final bias. While samples stored in the buffer may be analyzed to characterize configurations associated with specific regions in CV space, recover the full weighted con- formational distribution, one can integrate Eq. (9) to generate samples X; and assign weights W: = (xp](("(((xx)))) yielding the self-normalized estimator in Eq. (1). Note that this exact only in the limit of perfectly trained sampler. In principle, the path importance weights in ASBS could be used to debias an imperfect sampler with the gradient-based parametrization of control and corrector (Liu et al., 2025), though these weights may suffer from high variance. The learned bias also enables a practical refinement strategy: short MD or MCMC trajectories are initiated from the generated samples under the biased energy E + V E. and the resulting configurations reweighted with V*. This restores asymptotic correctness while keeping the procedure computa- tionally efficient thanks to the smoother energy landscape along the CVs.

---

Published as a conference paper at ICLR 2026
4 EXPERIMENTS
We demonstrate the efficiency of the WT-ASBS pipeline on four molecular sampling tasks involving peptides and chemical reactions. Since Requirements and 3 are incorporated into the design of the method, correctness with respect to Requirement 2 ensures that all requirements are satisfied. We therefore assess whether the empirical PMF, F(s) = -3 log ((s) obtained from reweighted samples Eq. (I) and the PMF from bias (Remark 3.1) agree with those from reference densities. Accuracy is computed from the PMF by measuring the free energy of state i in conformational states as
F = - log Is
E S.] ds + const,
(13)
where S, is the CV region of state i. We report free energy differences AF, F; F;, thereby removing the additive constant as is the same for all S,. We compute multi-state errors with a common offset. Although weighted sampling required by Requirement is best assessed via PMFs and free energies, we additionally compare with diffusion samplers on alanine dipeptide in Section F using standard distributional metrics. Our main comparisons are with ASBS and WTMetaD.
4.1 CONFORMATIONAL SAMPLING FOR SMALL PEPTIDES
We first benchmark the WT-ASBS method on two widely studied peptide systems, alanine dipeptide (Ala2; Ace-Ala All--me) and alanine tetrapeptide (Ala4; Ace-Alag-Nmer which serve as standard testbeds for evaluating molecular sampling methods. Both peptides are modeled with classical force field with implicit water solvation, with the chirality restraint applied to all C. atoms (see Section D). In this work, we focus on sampling the full Cartesian coordinates of the molecules. Al- though peptides and proteins adopt broad range of conformations, these arise largely from the com- binatorial organization of local backbone torsional states, specifically the Φ and zi angles (Fig. 2a), classically illustrated by Ramachandran et al. (1963). Building on this insight, we demonstrate that energy- based neural samplers can be made effective for sampling peptides directly in Cartesian space when augmented with biasing along these torsional angles.
Pretralning data (100 os MD. 23k E evals)
d PMF from reweighting
[IM samples)
Reference PMF
AF convergence
Fid. 10 (ecat/mol)
WT-AS85 WTMettD
x/2
5.0
A.
0.5
2.
4,0
n//2
3.5
Reference
-n/2
m/2
-
-n/2
m/2
-m/2
n/2
0.0
255
5.0
7.5
evaluations [x10']
PMF from
PMF fram bits 0.35M evais
PMF from biss 0.64M
PMF from biss 1.6M
PMF 8.34M
0.15M Evaals
PMF Ría. 60 Bcca/moll
x/2
-20
4/2
72/2
x/2
-8
-2/2
21/2
-2
-272
0
2
-3 /2/
1/2
Figure 2: Sampling alanine dipeptide. (a) Molecular structure of alanine dipeptide and definition of the torsional collective variables & and V. (b) Distribution of & and obtained from short unbiased MD trajectory used pretraining sampler. (c) Evolution
PMF reconstructed from the accumulated bias during training, illustrating progressive recovery of the major basins in the free energy landscape. The dotted box delineates the two states, one with è < 0 or & 2 and the other with á E [0, 2). (d) PMF estimated by reweighting 10 samples generated by the trained sampler using the final bias. (c) Reference PMF computed from long unbiased MD simulation for comparison. (f) Convergence of free energy differences AF between two states as function number of energy evaluations during training.
Alanine dipeptide (Ala2) We employ o and torsion angles as CVs (Fig. 2a). The conforma- tional landscape is broadly divided into two states along ó (Fig. 2b). Local sampling within the more populated state is efficient due to the absence of significant barrier; thus, we easily sample

---

Published as a conference paper at ICLR 2026
pretraining data from short MD simulation in this region. The pretrained model produces samples restricted to this state. However, once WT-ASBS training begins, the bias deposition progressively shifts the sampling distribution away from the explored region, enabling the discovery of conforma- tions in the less populated state. Over training progress, the deposition asymptotically converges to the bias potential required reproduce the PMF (Fig. 2c). Fig. 2d shows the PMF obtained from the generated samples of the fully trained model using weights from the final bias. Both the bias-derived PMF and the reweighted-sampl PMF agree closely with the reference PMF (Fig. 2e), whereas the baseline ASBS fails to reproduce the correct distribution in the less populated region (Fig. 14). Moreover, compared to the baseline enhanced sampling method (WTMetaD) with the same bias factor, WT-ASBS achieves convergence of the A/ between the two states more accurately.
To quantify the robustness of WT-ASBS and clarify practical choices, we performed ablations of the WT bias hyperparameters and training mechanisms, and reported detailed results in Section E. Varying the Gaussian height h, kernel width a, and bias factor 7 over wide ranges leaves the final AF unchanged except at the most extreme h, while larger h and broader mainly accelerate early exploration and barrier crossing. The training mechanism ablations show that overly frequent replay buffer refreshes slow convergence because the sampler does not have time to adapt to the evolving bias, whereas moderate update intervals and localized pretraining lead to a faster approach to the reference AF and reduce the total number of energy evaluations. Furthermore, Section F.4 demon- strates that WT-ASBS can also operate with an automatically learned MI. CV (Trizio & Parrinello, 2021), recovering accurate PMF and AF after reweighting, which indicates that the method could still be utilized when suitable www-iienniooal CVs are not known a priori.
Reference data
State exploration
Free energy convergence
8-
3.0
WT-A585
WS-ASBS
MAE (moll
2.5
WT MetaD
2.-
Explored
WTTMtaD
1.5
chemical acouracy
1.0
ASBS
2
0.5
MD
0
0
2
10
20
30
40
per run [x109]
Total [x10%
Figure 3: Sampling alanine tetrapeptide. (a) Molecular structure of alanine tetrapeptide with torsional collec- tive variables Φ1, 02, and Φ3. (b) Reference distribution in torsional space. Each torsion is partitioned into two regions, d; E [0,2] and d; < 0 or o; > 2. producing eight metastable states corresponding to all combinations of these regions across the three di. Pretraining is restricted to the single state indicated by so the sampler initially encounters only one of the eight possible regions. (c) Accumulated number of explored states during the initial phase training or simulation (up to 2M energy evaluations), showing how the sampler progres- sively discovers additional states as the bias is deposited. (d) Convergence of the free energies of the eight states, reported as the MAE up to an additive constant. For WT-ASBS, the curve extends up to 20M energy evaluations where the hias has converged.
Alanine tetrapeptide (Ala4) As seen in the Ala2 case, the main energetic barriers separating con- formational states arise from backbone é torsions. We therefore employ three CVs, Φ1, d2, and Φ3. A long 100- reference MD simulation shows that the system occupies eight distinct modes, corre- sponding the 2 combinations of the three à torsional states (Fig. 3b). We initialize from single conformation in the indicated mode and generate pretraining data from short MD simulation con- fined in the mode (Fig. 15). As shown in Fig. 3c, WT-ASBS explores conformational states far more efficiently than the baseline WTMetaD, discovering all eight modes within the early stages of train- ing. This advantage arises from its ability to generate uncorrelated samples from the biased potential, enabling exploration of multiple directions of conformational change simultaneously, whereas MD- based enhanced sampling explores states sequentially in simulation time. Consequently, WT-ASBS also achieves effective convergence of free energies (Fig. 3d), as measured by the mean absolute error (MAE) of free energies up to a constant mine 1 E; FF. F.eed + cl, within chemical accuracy (1 kcal/mol), demonstrating the practical applicability of WT-ASBS. We note that it does not yield lower free energy MAE than WTMetaD, likely because MD already provides efficient local mixing within each basin once barriers along the chosen CV are crossed, whereas the diffusion sampler must learn the full high- dige--imensional intra-basin distribution. Measure- based global sam- pling is most beneficial for crossing large barriers, whereas MDD-based local sampling approaches
8

---

Published as a conference paper at ICLR 2026
can still perform well when refining physical observables, which further suggests that combining the two strategies may be advantageous for more complex systems.
4.2 SAMPLING REACTIVE ENERGY LANDSCAPES
While peptide results highlight the efficiency of WT-ASBS in terms of the number of energy eval- uations, the computational budgets in liffusion-based sampling in these cases are dominated by SDE integration with a control network, since classical force fields are extremely fast to evaluate. However, classical force fields are insufficient for accurate conformational sampling, particularly in cases involving bond reorganization, while density functional theory (DFT), the standard for accu- rate energetics, computationally prohibitive. Recent advances in semiempirical quantum methods (Froitzheim et al., 2025) and universal ML interatomic potentials (uMLIPs) demonstrate that near- DFT accuracy can be achieved at a fraction of the compute, and open the door to scalable and accurate sampling. To illustrate this direction and compare the computational cost in practical set- tings, we apply WT-ASBS to sample reactive energy landscapes involving bond formation (Trizio & Parrinello, 2021). We employ recent uMLIP, UMA-S-L.I 1-11 (Wood et al., 2025), which accurately reproduces DFT energies, as the underlying energy model. For reaction sampling, we cross-check with MTtte results since sampling the transition region with unbiased samplers like MD is prac- tically impossible (e.g.. 10 samples required for 20 kcal/mol barrier at 300 K).
Fik. (rallmoll
25
WT AS85
4
05
b
Past-TS buturcation
0.5
-36
PMF R5y <l Rccl/moll
30
20
0.0
E
16
1,0
10
12
2.0
0.5
-05
12
1.5
0.0
1.5
2.0
2.5
3.0
3.5
0.0
2.5 [A]
Figure 4: Sampling reactive landscapes. (a) Sx2 reaction scheme with nucleophilic substitution between states 1 and 2, along with the two bond distance CVs and d2 that describe approach and departure of the chloride ions. (b) PMF F(1,, d2) d2) for the Sx2 system, showing the reactant and product basins and the transition-state region along the two-dimensional distance coordinates. (c) Cycloaddition reaction scheme ex- hibiting post-TS bifurcation, highlighting the competing pathways that originate from the same transition structure. (d) Definition of the three forming bond distances di, d2, and da used parameterize progress along bifurcating reaction pathways. (e) Contact CVs c; obtained from distances de, used to construct smoother CVs for sampling. (f) One-dimensional PMF F(si) along the CV # e + c2 + c3, comparing eg ca, resolving the post-TS bifurcation and showing the distinet product channels, with TS locations from saddle point optimization.
WT-ASBS and WTMetaD. (g) Two-dimensional PMF F(s1,5) 82) with
Nucleophilic substitution (Sx2) We study prototypical Sx2 reaction (Fig. 4a), where chloride ion attacks the tetrahedral carbon and another chloride ion departs. leading to interconversion be- tween states 1 + 2. We set two C-Cl bond distances as CVs and apply restraints to prevent chloride ions from diffusing away, ensuring the CV space remains bounded. Given the small system size, single configuration (state 1) suffices for pretraining the sampler. Using WT-ASBS, we obtain the 2-D PMF along the two bond distances (Fig. 4b). The symmetric PMF agrees well with the WTMetaD result (Fig. 16), and both the location and free energy barrier of the transition state (TS) coincide with the results from standard saddle-point optimization.
Dost-transition-state bifurcation For simple reactions like Sx2 with a single reaction channel between reactants and products, the process can be well characterized by identifying the TS and evaluating single-point energies, with harmonic approximations describing the local energy land- scape around the reaction progress (Truhlar al., 1996; Peters, 2017). However, many realistic challenges in chemistry involve branched or multimodal energy landscapes that cannot be captured by such approximations and instead requires explicit sampling. The reaction in Fig. 4c exemplifies

---

Published as a conference paper at ICLR 2026
this difficulty: reactant 3 passes through single TS region that bifurcates into two products 4 and 5. Because the TS is not uniquely associated with a single product (i.e., the TS is ambimodal), characterizing this reaction requires sampling the full reactive landscape.
To construct CVs, we represent the three relevant C-C distances (Fig. 4d) as normalized contacts in Eq. (55), capturing bond formation progress (Fig. 4e). From these define two CVs: 81 e, + + C3. which tracks overall bond formation (a soft bond count), and $2 = C2 Cg, which dis- tinguishes between products 4 and 5. Details on restraints and settings are provided in the Section D.
Fig. 4f and g show the PMFs (s))) and (( 82), and similarly to Sx2 example, the PMFs agree well with the WTMetaD result (2-D
comparison in Fig. 17), providing accurate energetics for complex,
PMF MAE
[kcal/mol]
Sx2
WTMetaD WT-ASB5
multimodal reactive landscapes that dominate practical challenges in chemistry. Further discussion on reaction sampling tasks is pro-
vided in Section
PMF MAE [kcal/mml]
10
Post-TS
Bifurcation
Computational efficiency We compared wall-clock time on four A100 80GB GPUs, where the MAE of the current bias relative to the final one, defined in Eq. (57), is shown in Fig. 5, For the Sx2 reaction, WT-ASBS required 0.77M energy evaluations and 4.3 hours, versus 4.0M and 29 hours for WTMetaD. For the bifurca- tion reaction, WT-ASBS used 2.6M evaluations and 23 hours, com- pared to 6.4M and 48 hours. Although WT-ASBS integrates SDEs with a control model, it needs fewer energy evaluations, which can be batched during training, effectively distilling a lighter sampler model from the heavier uMLIP.
30 40 50 Wall-clock time [hours]
10
20
Figure 5: PMFss PMF MAE from the final converged PMF for each reaction and method, compared over wall-clock time.
Convergence of
DISCUSSION
We presented WT-ASBS, which integrates well-tempered bias deposition over chosen CVs into a diffusion-based sampler, enabling exploration of unseen conformational modes and recovering PMFs and importance sampling weights that reconstruct the Boltzmann distribution. By embedding enhanced sampling directly into the sampler, WT-ASBS links diffusion- based generative samplers with classical rare event techniques and provides a practical recipe for sampling molecular systems with controllable bias strength and straightforward statistical reweighting. We have also demon- strated that WT-ASES can be combined with ML CVs (Section F.4), illustrating that the method is compatible with data-driven representations and can be used as a downstream sampler once informa- tive low-dimensional coordinates have been learned. Looking ahead, we expect the most effective workflows to combine energy-based sampling with data-based generative models and mix global moves (diffusion-basse proposals in configuration or CV space) with local sampling (short MD or MCMC) in a problem-dependent way, for example, using global diffusion moves for slow solute rearrangements and local samplers for fast solvent motions. Finally, although this work focuses on sampling molecular configurations, din-dimensional latent variables have proven useful in discrete lattice sampling (Swendsen & Wang, 1987; Wang & Landau, 2001) and dynamic Bayesian networks (Doucet et al., 2000), suggesting that extending the WT-ASBS-like approach beyond molecular sys- tems would be fruitful direction.
Limitations WT-ASBS requires the specification of CVs, and we primarily utilized predefined CVs and extended to ML CVs where metastable states are known. More complex systems will benefit from ML CVs and iterative cycles that interleave exploratory WT-ASBS runs with CV train- ing updates. A fair comparison of enhanced sampling methods remains challenging due to the large number of algorithmic and gystem-specifii hyperparameters (Hénin et al., 2022), and although WTMetaD may not be the optimal baseline for every system, we use it here as natural baseline because our bias deposition scheme is inherited from WTMetaD. Future work will explore combina- tions of different enhanced sampling strategies to provide a more comprehensive analysis. Finally, WT-ASBS does not provide amortized sampling across systems or thermodynamic conditions, and extending it with variationally parameterized bias (Bonati et al., 2019) shared control networks across related systems is a promising direction for transferring biases between systems that share similar slow modes.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Tara Akhound-Sadepa Jungyoon Lee, Avishek Joey Bose, Valentin De Bortoli, Arnaud Doucet, Michael M. Bronstein, Dominique Beaini, Siamak Ravanbakhsh, Kirill Neklyudov, and Alexan- der Tong. Progressive inference-time annealing of diffusion models for sampling from boltzmann densities, 2025. URL
16471. 16411. (cited on page 19)
Michael S Albergo and Eric Vanden-Eijnden. Nets: A non-equilibrium transport sampler. arXiv preprint arXiv:2410.02711, 2024, (cited on page 19)
Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797 2023. (cited pages 4 and 19)
MS Albergo, G Kanwar, and PE Shanahan. Flow-based generative models for markov chain monte carlo in lattice field theory. Physical Review D, 100(3):034515, 2019. (cited on page 18)
Alessandro Barducci, Giovanni Bussi, and Michele Parrinello. Well-tempered metadynamics: smoothly converging and tunable free-energy method. Physical Review Letters, 100(2):020003, 2008. (cited on pages 2, 3, 4, and 18)
Zakarya Benayad and Guillaume Stimemann. Hamiltonian replica exchange augmented with diffusion-based generative models and importance sampling to assess biomolecular conforma- tional basins and barriers, 2025. URL https:tttt//tee or..org////ooo abs///bs//// 2505. 08357. (cited on page 20)
Julius Berner, Lorenz Richter, and Karen Ullrich, An optimal control perspective on diffusion- -based generative modeling. arXiv preprini 22Xiv:2211.01364 2022. (cited pages and 19)
Denis Blessing, Julius Berner, Lorenz Richter, Carles Domingo-Enin Enricc,, Yuanqi Du, Arash Vah- dat, and Gerhard Neumann. Trust region constrained measure transport in path space for stochastic optimal control and inference, 2025. URL https://arxiv. orw. orgoorg///aa hwwpg//bb//e...... 12511. (cited on page 19)
Luigi Bonati, Yue-Yr Yu Zhang, and Michele Parrinello. Neural networks-based variationally en- hanced sampling. Proceedings ofe National Academy of Sciences. 116(36):17641117677, 2019. (cited on pages 4 and 10)
Luigi Bonati, Enrico Trizio, Andrea Rizzi, and Michele Parrinello. unified framework for ma- chine learning collective variables for enhanced sampling simulations: mlcolvar. The Journal of Chemical Physics, 159(1), 2023. (cited on pages 2, and 36)
Peter J Canfield, Iain M Blake, Zheng-Li Cai, Ian J Luck, Elmars Krausz, Rika Kobayashi, Jeffrey R Reimers, and Maxwell J Crossley. new fundamental type of conformational isomerism. Nature chemistry, 10(6):615-66) 624, 2018. (cited on pages 6 and 20)
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differ- ential equations. Advances in neural information processing systems, 31, 2018. (cited on page 19)
John Chodera, Andrea Rizzi, Levi Naden, Kyle Beauchamp, Patrick Grinaway, Mike Henry, Iván Pulido, Josh Fass, Alex Wade, Gregory A. Ross, Andreas Kraemer, Hannah Bruce Macdonald, jaimergp. Bas Rustenburg, Dominic Rufa, Ivy Zhang, David W.H. Swenson, Andy Simmon- ett, Mark J. Williamson, hb0402, Jake Fennick, Sander Roet, Ian Kenney, Irfan Alibay, Ben- jamin Ries, Ethan Holz, Hannah Baumann, Richard Gowers, and SimonBoothroyd. choder- alab/openmmtools: 0.25.1 bugfix release, August 2025. URL https:///oo http://doooorr//..22/2/. 5281/ zenodo. 16783210. e. (cited on page 28)
Jaemoo Choi, Yongxin Chen, Molei Tao, and Guan- Liu. Non-equilibrium annealed adjoint sampler, 2025. URL https: /ttp::///axiie http:///wxxvv..aaaaabbaa
18165. (cited on page 34)
Jeffrey Comer, James C Gumbart, Jérôme Hénin, Tony Lelièvre, Andrew Pohorille, and Christophe Chipot. The adaptive biasing force method: Everything you always wanted to know but were afraid to ask. The Journal of Physical Chemistry 119(3):1129) 129-11112-11151 2015. (cited page 18)

---

Published as a conference paper at ICLR 2026
James F. Dama, Michele Parrinello, and Gregory A. Voth. Well-tempered metadynamics converges asymptotically. Phys. Rev. Lett., 112:240602, Jun 2014, doi: 10. 10.1103/PhysRevLett. 12.240602. (cited on page 22)
URL https: aps wwt.ooo//o//oooo/oooooo//o/// http///.1111////eseeettttt.t.t
Eric Darve, David Rodrz, and Andrew Pohorille. Adaptive biasing force method for scalar and vector free energy calculations. The Journal of chemical physícs, 128(14), 2008. (cited page 18)
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real nvp. arXiv preprint arXiv: (60.::6000.00) 2016. (cited on page 19)
Joseph L Doob. Conditional brownian motion and the boundary limits of harmonic functions. Bul- letin de la Société mathématique de France, 85:431-458, 1957. (cited on page 19)
Arnaud Doucet, Nando de Freitas, Kevin Murphy, and Stuart Russell. Rao-blackwellised particle filtering for dynamic bayesian networks. In Proceedings of the Sixteenth Conference on Uncer- tainty Artificial Intelligence, UAT000, pp. 176-183, San Francisco, CA, USA, 2000. Morgan Kaufmann Publishers Inc. ISBN 1558607099. (cited on page 10)
Yuanqi Du, Michael Plainer, Rob Brekelmans, Chenru Duan, Frank Noe, Carla P Gomes, Alan Aspuru-Guzik, and Kirill Neklyudov. Doob's lagrangian: A sample-efficient variational approach to transition path sampling. Advances in Neural Information Processing Systems, 37:65791-
65822, 2024. (cited on page 19)
Peter Eastman, Raimondas Galvelis, Raúl Peláez, Charlles RA Abreu, Stephen E Farr, Emilio Gal- licchio, Anton Gorenko, Michael M Henry, Frank Hu, Jing Huang, et al. Openmm 8: molecular dynamics simulation with machine learning potentials. The Journal of Physical Chemistry 128 (1):109-116, 16, 2023. (cited on pages 27 and 28)
Daan Frenkel and Berend Smit. Understanding molecular simulation: from algorithms to applica- fions. Elsevier, 2023. (cited on page 1)
Thomas Froitzheim, Marcel Müller, Andreas Hansen, and Stefan Grimme, g-xtb: A general-purpose electronic structure method for the elements h to Ir (z= 1-103). ChemRxiv, extended 22264444/chemmxii22255----0 bjxvt. (cited on page 9)
2025. doi: 10.
Aaron Havens, Benjamin Kurt Miller, Bing Yan, Carles Domingo-Enrie Anuroop Sriram, Daniel S. Levine, Brandon M Wood, Bin Hu, Brandon Amos, Brian Karrer, Xiang Fu, Guan- Horng Liu, and Ricky Q. Chen. Adjoint sampling: Highly scalable diffusion samplers via adjoint matching. In Forty-second International Conference on Machine Learning, 2025. URL https://
httt/foruuiidd=EE 0mmmg2.. (cited on pages 1, 4, 19, 25, and 34)
Jiajun He, Yuanqi Du, Francisco Vargas, Dinghuai Zhang, Shreyas Padhy, RuiKang Yu Carla Gomes, and José Miguel Hernández-Lobato No trick, no treat: Pursuits and challenges towards simulation- -fret training of neural samplers, 2025. URL https: /ttp///axvvv////////////////.avv wwtpxxivvorr//.oom atgg/bs//285 26620/665........... 06685. (cited page 2)
Lukas Herron, Kinjal Mondal, John S Schneekloth Jr, and Pratyush Tiwary. Inferring phase transi- tions and critical exponents from limited observations with thermodynamic maps. Proceedings of the National Academy of Sciences, (21(52)ee231197112 121, 2024. (cited page 20)
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-68 6851, 2020. (cited on pages and 19)
Lars Holdijk, Yuanqi Du Ferry Hooft, Priyank Jaini, Berend Ensing, and Max Welling. Stochastic optimal control for collective variable free sampling of molecular transition paths. Advances in Neural Information Processing Systems, 36:79540-79556 2023. (cited on page 19)
Chad W Hopkins, Scott Le Grand, Ross C Walker, and Adrian E Roitberg. Long- time-step molccular repartitioning. Journal of chemical theory and computation, 11
dynamics through hydrogen (4): (0):18641874 2015. (cited on page 27)
KN Houk, LJ Luskus, and NS Bhacca. Novel double [6+4] cycloaddition of tropone to dimethylful- vene. Journal of the American Chemical Society, (112))633266649944 1970. (cited page 35)
12

---

Published as a conference paper at ICLR 2026
Jérôme Hénin, Tony Lelièvre, Michael R. Shirts, Omar Valsson, and Lucie Delemotte. Enhanced sampling methods for molecular dynamics simulations [article v1.0]. Living Journal of Computational Molecular Science, 4(1):1583, Dec. 2022. doi: 10.33011/ Awvecoms4.11.1585 URI. https: httess////ccoomsouraaall org/index. www.iiieewweeeeeooeeeeeepe php/Liivecoms/rtiicl.. view/v4i1e1588, (cited on page 10)
Michele Invernizzi and Michele Parrinello. Rethinking metadynamics: from bias potentials to probability distributions. The journal of physical chemistry letters, 11(7):273-2236,, 2020. (cited pages 4 and 18)
Ryan Jacobs, Dane Morgan, Siamak Attarian, Jun Meng, Chen Shen, Zhenghao Wu, Clare Yijia Xie, Julia H Yang, Nongnuch Artrith, Ben Blaiszik, et al. practical guide to machine learning interatomic potentials- status and future. Current Opinion in Solid State and Materials Science, 35:101214, 2025. (cited on page 1)
Rui Jiao, Wenbing Huang, Peijia Lin, Jiaqi Han, Pin Chen, Yutong Lu, and Yang Liu. Crystal structure prediction by joint equivariant diffusion. Advances in Neural Information Processing Systems, 36:17464-17497, 2023. (cited on page 4)
Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi Jaakkola. Torsional dif- fusion for molecular conformer generation. Advances in neural information processing systems, 35:24240 24253, 2022. (cited on page 4)
Tero Karras, Miik Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. Advances in neural information processing systems, 35:26565-26577, 2022. (cited on page 25)
Minkyu Kim, Kiyoung Scong, Dongyeop Woo, Sungsoo Ahn, and Minsu Kim. On scalable and efficient training of diffusion samplers, 2025. URL (cited on page 19)
19552.
Peter Kollman, Richard Dixon, Wendy Cornell, Thomas Fox, Chris Chipot, and Andrew Poho- rille. The dexeloomeen//aplliaation of 'minimalist' oraiimaaissoggaanc/bboieoeo organic/biochemical molecular mechanic force field using a combination of ab initio calculations and experimental data. In Computer sim- ulation of biomolecular systems: Theoretical and experimental applications, pp. 83-96. Springer, 1997. (cited on page 27)
Solomon Kullback. Probability densities with given marginals. The Annals of Mathematical Statis- fics,
1968. (cited on page
Harold J Kushner and G George Yin. Stochastic approximation and recursive algorithms and appli- cations. Springer, 2003. (cited on page 22)
Alessandro Laio and Michele Parrinello. Escaping free-energy minima. Proceedings of the national academy of sciences, 99(20):12522-125555 2002. (cited on page 18)
Ask Hjorth Larsen, Jens Jørgen Mortensen, Jakob Blomqvist, Ivano E Castelli, Rune Christensen, Marein Dutak, Jesper Friis, Michael N Groves, Björk Hammer, Cory Hargus, Eric D Hermes, Paul C Jennings, Peter Bjerre Jensen, James Kermode, John R Kitchin, Esben Leonhard Kols- bjerg, Joseph Kubal, Kristen Kaasbjerg, Steen Lysgaard, Jón Bergmann Maronsson, Tristan Maxson, Thomas Olsen, Lars Pastewka, Andrew Peterson, Carsten Rostgaard, Jakob Schiøtz, Ole Schütt, Mikkel Strange, Kristian Thygesen, Tejs Vegge. Lasse Vilhelmsen, Michael Wal- ter, Zhenhua Zeng, and Karsten W Jacobsen. The atomic simulation environment- python library for working with atoms. Phys.: Condens. Matter, 29(27):273002, 2017. URL https://doi. org/10. 10. 1088/1333 1188/116611 648X/ aa680e. (cited on page 27)
Yi-Lun Liao, Brandon Wood, Abhishek Das*, and Tess Smidt*. EquiformerV2: Improved Equiv- ariant Transformer for Scaling Higrer-Degree Representations. In International Confer- ence on Learning Representations (ICLR), 2024. URL https:://peenevvii net/forummid mCOBBZmrzZ. (cited on page 25)
Kresten Lindorff-Larsen Stefano Piana, Kim Palmo, Paul Maragakis, John L Klepeis, Ron O Dror, and David Shaw. Improved side-chain torsion potentials for the amber ff99sb protein force field. Proteins: Structure, Function, and Bioinformatics, 1958, 2010. (cited on page 27)
13

---

Published as a conference paper at ICLR 2026
Yaron Lipman, Ricky TQ Chen, Heli Ben- Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling, arXiv preprint arXiv:2210.02747, 2022.. (cited on pages 4 and 19)
Guan-Horng Liu, Jaemoo Choi, Yongxin Chen, Benjamin Kurt Miller, and Ricky Q. Chen. Adjoint schrödinger bridge sampler, 2025. URL https://arxiv.org/ wwwpivv..oggbbb////ee abgg/bb2.0...... 2255/2255.....5.... (cited on pages 2, 4, 6, 24, 27, 34, and 35)
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv:1711.05101, 2017. (cited on page 25)
arXiv preprint
Dominik Marx and Jurg Hutter. Ab initio molecular dynamics: Theory and implementation, Modern methods and algorithms of quantum chemistry, 1(301 449)):111, 2000. (cited on page 1)
Bálint Máté and François Fleuret. Learning interpolations between boltzmann densities. Transac- tions on Machine Learning Research, 2023. ISSN 2835-8856. URL https://openrevieet net/
Forumeidd-T6YYEEchhh (cited on page 19)
Alan D McNaught and Andrew Wilkinson. Compendium of chemical terminology, volume 1669.
Blackwell Science Oxford, 1997. (cited on page 20)
Laurence Midgley, Vincent Stimper, Javier Antorán, Emile Mathieu, Bernhard Schölkopf, and José Miguel Herúánde--Lbbto. Se (3) equivariant augmented coupling flows, Advances in Neu- ral Information Processing Systems, 36:79200-79225, 2023. (cited on page 19)
Laurence Illing Midgley, Vincent Stimper, Gregor NC Simm, Bernhard Schölkopf, and José Miguel Hernández-Lobato. Flow annealed importance sampling bootstrap. arXiv preprint arXiv:2208.0 01893, 2022. (cited on pages 19 and 27)
Benjamin Kurt Miller, Ricky T. Q Chen, Anuroop Sriram, and Brandon M Wood. FlowMM: Generating materials with Riemannian flow matching. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Pro- ceedings of Machine Learning Research, pp. 35664-35686. PMLR, 21-27 Jul 2024. URL html. (cited on page 4)
https://proceedings www.eeeddinnsm....rreesesee pttpsperes////mmlsl/ls
Juno Nam and YounJoon Jung. Enhanced sampling for free energy profiles with post-transition- state bifurcations. Journal of Chemical Theory and Computation, (700))2773527333 2023.
(cited on page 29)
Kim A Nicoli, Shinichi Nakajima, Nils Strodthoff, Wojeiech Samek, Klaus-Robert Müller, and
Pan Kessel. Asymptotically unbiased estimation of physical observables with neural samplers. Physical Review E, 101(2)002334, 2020. (cited on page 18)
Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equi- librium states of many-body systems with deep learning. Science, 365(6457)xeaaw 147, 2019. (cited pages and 18)
Alexey Onufriev, Donald Bashford, and David A Case. Exploring protein native states and large- scale conformational changes with a modified generalized born model. Proteins: Structure, Func- fion, and Bioinformatics, 55(2))383-39) 2004. (cited page 27)
Michael N Paddon-Row and Ronald N Warrener. Concerning the periselectivity of the cycloaddition reaction between 8, d-dimetttliisobenzofflvveee and tropone. Terrahedron Letters, 15(43):3797- 3800, 1974. (cited on page 35)
Baron Peters. Reaction rate theory and rare events. Elsevier, 2017. (cited on page 9)
Yunrui Qiu, Richard John, Lukas Herron, and Pratyush Tiwary. Latent thermodynamic flows: Uni- fied representation learning and generative modeling of lemperature-dependen behaviors from limited data. arXiv preprint 1oXXi:25077.3174 2025, (cited on page 20)
Sanjeev Raja, Martin Sipka, Michael Psenka, Tobias Kreiman, Michal Pavelka, and Aditi S. Krish- napriyan. Action cation-minimizationn meets generative modeling: Efficient transition path sampling with the onsager-machlug functional. In Forty-second International Conference on Machine Learning, 2025. URL httt:::///pppeeiiieeee www.:ee/oooooiodooonooooooooo
(cited on page 20)
14

---

Published as a conference paper at ICLR 2026
G.N. Ramachandran, C. Ramakrishnan, and V. Sasisekharan. Stercochemistry of polypeptide chain configurations. Journal of Molecular Biology, 7(1):95 99, 1963. ISSN 0022-2836. doi: https:// (...:g/0000166600222233666666600 arttccleppi//50222333330000
URL
com/ http///eetttccccceeeeeeeeeeeeeee science/
(cited on page
Lorenz Richter and Julius Bemnr. Improved sampling via learned diffusions.
arXiv preprint
arXiv:2307.01198, 2023. (cited on page 19)
Severi Rissanen, RuiKang Ou Yang, Jiajun He, Wenlin Chen, Markus Heinonen, Arno Solin, and José Miguel Hernández-Lobat Progressive tempering sampler with diffusion. In Forty-second International Conference Machine Learning. 2025. URL http://pppnnrveeeeeeeeeeeeeeeeee id=uBMnbCBEtZ. (cited on page 19)
Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathemati- cal statistics, pp. 400-407, 1951. (cited on page 22)
Henrik Schopmans and Pascal Friederich. Temperature- boltzmann generators. In Forty- second International Conference on Machine Learning, 2025. URL https::tt//s......m http::::////p/eeeeeeeeee net/
Forum?id-RqgRSrCbNN. (cited on page 19)
Mathias Schreiner, Ole Winther, and Simon Olsson. Implicit transfer operator learning: Multiple time-resolution models for molecular dynamics. In A. Oh, T. Naumann, A. Globerson, K. Saenko, Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 36449-36462. Curran Associates, Inc., 2023. URL httpp:///ppereevie wtwpeeewweeeeeerrure/// net/forum?id= 1kZx7JiuA2. (cited on pages 6 and 25)
Kristof Schütt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the pre- diction of tensorial properties and molecular spectra. Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference Machine Learning, volume 139 of Pro- ceedings of Machine Learning Research, pp. 9377-9388. PMLR, 18-24 Jul 2021. URL https: html. (cited page 25)
//proceedings.mlr.. htt...ppees//13//scuueeeeee
Kiyoung Seong, Seonghyun Park, Seonghwan Kim, Woo Youn Kim, and Sungsoo Ahn. Transition path sampling with improved off-policy training of diffusion path samplers, 2025. URL https: //arxiv. wwtpriiw.ogggb.///////aap 2400. 19961. (cited on page 19)
Muhammed Shuaibi, Abhishek Das, Anuroop Sriram, Misko, Luis Barroso-Luque, Ray Gao, Sid- dharth Goyal, Zachary Ulissi, Brandon Wood, Tian Xie, Junwoong Yoon, Brook Wander, Adeesh Kolluru, Richard Barnes, Ethan Sunshine, Kevin Tran, Xiang, Daniel Levine, Nima Shoghi, Ilias Chair, Janice Lan, Kaylee Tian, Joseph Musielewicz, clz55, Weihua Hu, Kyle Michel, willis, and vbttchr. FAIRChem, 2025. URL https://gite thw.ggtthhh...... com wttp://aacebbooeeeerrrraaaaaaaaa (cited on page 27)
Hythem Sidky, Wei Chen, and Andrew I. Ferguson. Machine learning for collective variable dis- covery and enhanced sampling in biomolecular simulation. Molecular Physics, 11377)2 2020. (cited on pages 2 and 6)
Aditya Singh, Avishek Das, and David Limmer. Variational path sampling of rare dynamical events. Annual Review of Physical Chemistry, 76, 2025. (cited on page 20)
Jascha Sohl-Dicksteen, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learn- ing, pp. 2256-2265. pmlr, 2015. (cited on pages and 19)
Gabriel Stoltz, Mathias Rousset, et al. Free energy computations: A mathematical perspective. World Scientific, 2010. (cited on page 28)
Yuji Sugita and Yuko Okamoto. Replica-exchange molecular dynamics method for protein folding. Chemical physics letters, 114(1-1): 1999. (cited on page 18)
Robert H Swendsen and Jian-Sheng Wang. Nonuniversal critical dynamics in monte carlo simula- tions. Physical review letters, 58(2):86, 1987. (cited page 10)
15

---

Published as a conference paper at ICLR 2026
Esteban G Tabak and Cristina V Turner. A family of nonparametric density estimation algorithms. Communications on Pure and Applied Mathematics, 66(2):145-16) 2013. (cited on page 19)
The PLUMED consortium. Promoting transparency and reproducibility in enhanced molecular sim- ulations. Nature methods, 16(8):670-61) 2019. (cited on page 27)
Glenn M Torrie and John Valleau. Nonphysical sampling distributions monte carlo free- energy estimation: Umbrella sampling. Journal of computational physics, 23(2):187-199, 1977. (cited on page 17)
Gareth A Tribello, Massimiliano Bonomi, Davide Branduardi, Carlo Camilloni, and Giovanni Bussi. Plumed 2: New feathers for old bird. Computer physics communications, 185(2):604-133, 2014. (cited on page 27)
Enrico Trizio and Michele Parrinello. From enhanced sampling reaction profiles. The Journal of Physical Chemistry Letters, 12(35):8621-88) 8626, 2021. (cited on pages 8, 9, and 36)
Donald G Truhlar, Bruce C Garrett, and Stephen Klippenstein. Current status of transition-stee state theory. The Journal of physical chemistry, 100(31)112711-1280 1996. (cited page 9)
Lorenz Vaitl, Kim A Nicoli, Shinichi Nakajima, and Pan Kessel. Gradients should stay path: bet- ter estimators of the reverse-and forward divergence for normalizing flows. Machine Learning: Science and Technology, 3(4):045006, 2022. (cited on page 19)
Omar Valsson and Michele Parrinello. Variational approach to enhanced sampling and free energy calculations. Physical review letters, 113(9))090001, 2014. (cited on pages 4 and 18)
Francisco Vargas, Will Grathwohl, and Arnaud Doucet. Denoising diffusion samplers. arXiv preprint arXiv:2302.138..2 2023a. (cited on pages I, 4, 19, and 34)
Francisco Vargas, Shreyas Padhy, Denis Blessing, and Nikolas Nüsken. Transport meets varia- tional inference: Controlled monte carlo diffusions. arXiv preprint arXiv:2307.01050, 2023b. (cited on page 19)
Fugao Wang and David Landau. Efficient, multiple-range random walk algorithm calculate the density of states. Physical review letters, 86(10):2050, 2001. (cited on page 10)
Yihang Wang, Lukas Herron, and Pratyush Tiwary. From data to noise to data for mixing physics across temperatures with generative artificial intelligence. Proceedings of the National Academy of Sciences, 119632):e2203656113 2022. (cited on page 20)
Joseph L Watson, David Juergens, Nathaniel Bennett, Brian L Trippe, Jason Yim, Helen E Eise- nach, Woody Ahern, Andrew Borst, Robert J Ragotte, Lukas F Milles, et al. De novo de- sign of protein structure and function with rfdiffusion. Nature, 12897976):009-11000 2023. (cited page 4)
Brandon M. Wood, Misko Dzamba, Xiang Fu, Meng Gao, Muhammed Shuaibi, Luis Barroso- Luque, Kareem Abdelmaqsoud, Vahe Gharakhanyan, John Kitchin, Daniel S. Levine, Kyle Michel, Anuroop Sriram, Taco Cohen, Abhishek Das, Ammar Rizvi, Sushree Jagriti Sahoo, Zachary W. Ulissi, and C. Lawrence Zitnick. Uma: A family of universal models for atoms, 2025. URL hnttps://arat htttp::/aarxi..ogg/abb//////////
/2t////a//255.... 23971. (cited on pages 9, 27, 28, and 29)
Peiyuan Yu, Tiffany Q. Chen, Zhongyue Yang, Cyndi Qixin He, Ashay Patel, Yu-hong Lam, Ching- Yang Liu, and K. N. Houk. Mechanisms and origins of periselectivity of the ambimodal [6 + 4] cycloadditions of tropone to dimethylfulvene. Journal of the American Chemical Society, 139 (24):8251-8558. 2017. doi: (..122/jjas.7b02296 URL https: httts://doi.orr org//....org/11. 1021 jacs. 7b82966. (cited on page 35)
Qinsheng Zhang and Yongxin Chen. Path integral sampler: a stochastic control approach for sam- pling. arXiv preprint 2rXiv:2111.15141 2021. (cited on pages 4, 19, and 34)
Zhijun Zhang, Xinzijian Liu, Kangyu Yan, Mark E Tuckerman, and Jian Liu. Unified effi- cient thermostat scheme for the canonical ensemble with holonomic or isokinetic constraints via molecular dynamics. The Journal of Physical Chemistry A, 123((8):6056-60))) -076, 2019. (cited on pages and 28)
16

---

Published as a conference paper at ICLR 2026
APPENDIX
A Additional Background
17
A.1 Enhanced Sampling
17
A.2 Neural Samplers
18
A.3 Extended Related Works
19
A.4 Molecular Isomerism
20
B Proofs and Notes
22
B.1 Proof for Proposition 3.1
22
Implementation Details for WT-ASBS
24
C.1 WT-ASBS with Replay Buffer
24
C.2 Model Architecture
25
C,3 Base Process
25
C.4 Hyperparameters
25
D Experiment Details
27
D.I Alanine Dipeptide
27
D.2 Alanine Tetrapeptide
................................
27
D.3 Nucleophilic Substitution
D.4
Dost--rranition-SSte Bifurcation
29
Ablation Studies
30
E.1
Well-Tempered Bias Update
30
E.2
Training Mechanism
30
F
Additional Results
34
E.I
Alanine Dipeptide
34
E2
"eep............................
35
E3
Sampling Reactive Energy Landscapes
35
F44
MI. Collective Variables
35
Large Language Model Statement
37
A ADDITIONAL BACKGROUND
A.1 ENHANCED SAMPLING
Efficiently sampling molecular configurations requires overcoming high free energy barriers that cause rare but important states to be visited infrequently. Enhanced sampling methods are a class of techniques developed in the computational chemistry community address this challenge. The core idea is to bias the simulation dynamics to encourage exploration of less-populated regions of phase space, while applying compensating weight to recover correct thermodynamic averages. Unlike unbiased MD or MCMC (which can get trapped in deep energy wells), enhanced sampling methods add external forces or modify the energy surface so that the system escapes metastable states more readily.
One foundational approach umbrella sampling. Introduced by Torrie & Valleau (1977), umbrella sampling runs simulations with added bias potential (s) along chosen reaction coordinate or collective variable æ = E(x). By choosing a series of overlapping bias "umbrellas" that cover the range of force the system to sample even high-energy configurations in each window.
17

---

Published as a conference paper at ICLR 2026
Each window's samples are then reweighted by cx)(++3Uss) to reconstruct the unbiased free en- ergy landscape. This method demonstrated the principle of biasing along important coordinates to access otherwise inaccessible states. However, it requires careful tuning of windows.
Metadynamics took adaptive biasing a step further by dynamically building the bias potential during single simulation. In the original metadynamics (Laio & Parrinello, 2002), the simulation de- posits small Gaussian hills on the free energy landscape along the chosen CVs every few timesteps. These hills accumulate in wells (low free energy regions), gradually filling them and effectively raising the energy floor so the system escapes. Over time, the bias Vi(s) approximates the negative of the underlying free energy (s). flattening the CV landscape. Well-tempered metadynamics (Barducci et al., 2008) refined this approach by gradually reducing the height of added hills as the bias grows. This means early in the simulation, the bias grows quickly (promoting exploration), but later it converges towards finite bias that yields a modified stationary distribution rather than diverging- Specifically, well-temn introduces a bias factor ? > so that the deposited bias is scaled by factor that depends on the current bias (preventing overfill). The resulting steady-state is the well-tempered distribution in CV space, in which the free energy barriers are lowered by a factor of 7 effectively simulating the system at higher temperature Tar >T along those CVs. Importantly, even with bias, samples can be reweighted to recover the true Boltzmann distribution. Well-tempered metadynamics thus balances exploration and accuracy by never completely flatten- ing the landscape, avoiding infinitely large weights and ensuring the bias converges asymptotically to finite correction.
In the enhanced sampling community, there have been continual improvements to biasing algo- rithms. Notably, on-the-fly probability enhanced sampling (OPES) by Invernizzi & Parrinello (2020) is a performant metadynamics variant where the bias is constructed by directly targeting a pre- defined distribution (often the well-tempered distribution) at each step rather than depositing fixed- shape hills. This results in smoother bias updates and often faster convergence. Although we do not explicitly implement OPES, our algorithm likewise targets the well-tempered distribution Eq. (7)- In sense, WT-ASBS could incorporate OPES-style updates (adjusting the bias based on deviation between current CV histogram and target distribution), which could be an interesting direction for future work. Additionally, techniques like variationally optimized enhanced sampling (VES) (Vals- son & Parrinello, 2014) cast bias finding as an optimization problem of V(s) to minimize the KL. divergence between biased and unbiased ensembles. Such approaches could also be combined with our generative sampler by using the model's output to estimate that KL objective. In our current work, we use the simpler hill-based scheme inspired by WTMetaD.
Beyond metadynamics, numerous enhanced sampling strategies exist. The adaptive biasing force (ABF) method (Darve et al., 2008) takes different approach: instead of depositing hills, it contin- uously estimates the mean force along the CV and subtracts it, effectively flattening the free energy gradient. ABF and its extensions (e.g., extended ABF in Comer et al. (2015)) guarantee convergence by driving the observed force toward zero in well-sampled regions. In practice, ABF and metady- namics often achieve similar outcomes: a nearly uniform sampling of s, but via force-balancing vs. potential-fillin mechanisms. Another family of methods is replica exchange and parallel temper- ing. In replica exchange MD (Sugita & Okamoto, 1999), multiple copies of the system are simulated at different temperatures (or Hamiltonians) in parallel. These replicas occasionally swap configura- tions, allowing a cold replica to receive a configuration from a hot replica that has overcome barriers. Over many swaps, each replica samples from its designated ensemble (c.g., the lowest- temperttre one samples the target distribution) but benefits from exchanges to traverse phase space faster. Par- allel tempering effectively raises the chance of barrier crossing by leveraging high-temperature dy- namics, albeit at the computational cost of running many simulations. It is a powerful method to explore complex landscapes, though it requires careful selection of the temperature ladder and has high computational overhead if many replicas are needed.
A.2 NEURAL SAMPLERS
Historical developments As opposed to MCMC and MD methods, generative models can poten- tially be used to generate samples from multimodal distributions without getting trapped in local minima or requiring lengthy equilibration times. Using generative models to sample probability distributions specified by a potential function was originally proposed by Noé et al. (2019); Albergo et al. (2019), and Nicoli et al. (2020) in the contexts of many-body physics, lattice field theory, and
18

---

Published as a conference paper at ICLR 2026
statistical mechanics, respectively. These initial approaches trained discrete- time normalizing flows (Tabak & Turner, 2013) by optimizing the KL divergence between the model and the unnormal- ized target densities. As the use of reverse-KL objective is prone mode collapse, the community explored alternative training objectives for flow-based sampling (Vaitl et al., 2022; Midgley et al., 2022; Máté & Fleuret, 2023). However, using flow-based methods with these approaches presents challenges. Computing the likelihood of continuous flow models (Chen et al., 2018) is computation- ally expensive, yet
necessary component of many of these potential function-based learning objectives. Meanwhile, designing expressive discrete-time normalizing flows is complex due to the restrictions on their coupling layers (Dinh et al., 2016; Midgley et al., 2023).
Thinking of the sampling task as problem of measure transport opens up other lines of research which make use of ideas introduced by denoising diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020) and related, more general frameworks for generative modeling (Lipman et al., 2022; Albergo et al., 2023). In data-based settings, these approaches use intermediate distributions to transport samples to the target distribution, guided by local quadratic regression ("matching") objective that ensures the correct change of measure. The simulation-free nature of these approaches, combined with their straightforward training objectives, makes them more scalable and easier to train than previous methods, establishing them as the leading generative paradigm today. In the data-free setting, however, since samples are not available, trajectories must be sampled from the model. The central goal of recent works has thus been to design training objectives that (i) lead to a correct transport of measures and (ii) do not require backpropagation through the learned dynamics.
Vargas et al. (2023a); Berner et al. (2022) optimize the KL divergence in path space and establish connections between generative modeling and SOC. Zhang & Chen (2021); Havens et al. (2025) also approach the problem from the viewpoint of SOC by minimizing the control objective with the terminal cost defined by the target energy. Richter & Berner (2023) and Vargas et al. (2023b) opti- mize the variance of Radon-Nii -ikkodym derivative between the forward and backward path measures. Albergo & Vanden-Fijnden (2024) generalizes annealed importance sampling (AIS) by learning an SDE that minimizes the impact of unbiasing weights of AIS.
Related works Recent methods attempt to alleviate the challenge of energy-based training strug- gling to capture multiple modes in complex energy landscapes. Schopmans & Friederich (2025) applied temperature annealing in Boltzmann Generators, and concurrently Akhound-Sadee Sadegh et al. (2025); Rissanen et al. (2025) incorporated inference-time annealing into diffusion models to sam- ple from Boltzmann densities. These approaches smoothly interpolate from tegprrrrper (smoothed-out) version of the target distribution to the true low-temperature distribution akin to par- allel tempering. Our method biases the sampler on the fly along specific CVs, which can be seen as more targeted form of annealing that raises the effective sampling temperature only along selected important degrees of freedom.
Another line of related work focuses on directly encouraging exploration during training. Kim al. (2025) improved training efficiency by incorporating MCMC with bias from random network distillation during the search stage, steering exploration toward configurations that would otherwise be rarely visited, conceptually similar to our goal of amplifying exploration. In recent work, Blessing et al. (2025) formulate diffusion-basee based sampler with trust-region constrained optimal control perspective, which effectively adds Lagrange multiplier terms to enforce constraints on the path distribution. Our approach is unique in that it incorporates an adaptive bias during training that specifically targets low-dimensional CVs, which is highly effective for exploring rare-event regions and further enables reconstruction of the PMF
A.3 EXTENDED RELATED WORKS
MI.-based transition path sampling While our method is designed for equilibrium sampling with a CV-based biasing inside an energy-based diffusion sampler, closely related ideas have been developed for sampling rare dynamical trajectories by learning control forces on trajectory space, in the context of transition path sampling (TPS). Holdijk et al. (2023) cast molecular TPS as an SOC problem and learn CV-free policies that steer trajectories between metastable states. Du et al. (2024) propose Doob's Lagrangian, a variational formulation of Doob's h-transform (Doob, 1957) that learns conditional dynamics for generating transition paths between prescribed endpoints. Scong et al. (2025) introduce diffusion-hased based TPS with off-policy training based log-variance di-
19

---

Published as a conference paper at ICLR 2026
vergence between biased MD and target path measures. Different variational approaches path sampling of rare events are reviewed in Singh et al. (2025), with focus on trajectory ensemble reweighting and optimal control approaches. Finally, Raja et al. (2025) repurpose pretrained gen- erative models for TPS by interpreting paths as trajectories of an induced SDE and finding most- probable transition paths via Onsager-Machlup action minimization.
Data-based diffusion sampling in chemistry Complementary to our energy-based sampler that only assumes access to the potential energy and its gradients, several works train diffusion flow- based generative models directly on molecular simulation or thermodynamic data to build data- driven equilibrium samplers. Wang et al. (2022) train diffusion models on ensembles multiple thermodynamic conditions to generate molecular configurations across temperature or parameter space. Herron et al. (2024) develop thermodynamic maps, a diffusion-like generative approach to model the joint distribution of configurations and auxiliary thermodynamic variables to recon- struct phase behavior and other macroscopic properties from sparse data. Qiu et al. (2025) develop latent thermodynamic models that jointly learn now-dimensional representations and temperature- dependent equilibrium distributions. Benayad & Stirnemann (2025) combine replica exchange sim- ulations with diffusion models and importance sampling refine conformational ensembles and explore free-energy landscapes along chosen CVs.
A.4 MOLECULAR ISOMERISM
This appendix specifies how we operationalize chemical "sameness" and "difference" in terms of minimal geometric features, their coordination numbers at the atoms that define each stereochemical element, and the order parameter restraints we apply during sampling. We adopt the IUPAC tax- onomy (McNaught & Wilkinson, 1997), which separates constitutional isomerism from stereoiso- merism and further divides stereoisomerism into configurational and conformational classes.
Constitutional isomers share an elemental formula but differ in bond connectivity and therefore are different compounds. In our default setup, they are out of scope, except when the declared reaction of interest explicitly includes bond formation or breaking.
Stereoisomers have the same connectivity but differ in three-dimensional arrangement. They subdivided into:
Configurational isomers, which are separated by large barriers and are isolable on exper- imental time scales. This category includes absolute configuration at stereogenic centers (R/S), geometric isomerism such as E/Z, and helical forms. When only a single stereocen- ter is inverted, R and S related molecules enantiomers; with multiple stereocenters, the possibilities include diastereomers.
Conformational isomers, which interconvert over low barriers, for example, through rotation about single bonds or ring flips.
This taxonomy is paired with barrier and half-life delineations that mark fluxional, practically isolable, and configurationally stable regimes (Canfield et al., 2018).
Accessible component A Let X' denote the configurational space for fixed chemical composi- tion. We select a reference configuration and restrict attention to the portion of X' relevant to the molecular process under study, which is the sense in which the main text employs the term accessi- ble component. One can regard configurations at temperature T and study time horizon 7 as grouped by an equivalence relation ~y; two configurations belong the same class if they differ only by stereochemical elements that interconvert with half-life shorter than т, while connectivity is pre- served unless the chemical transformation of interest explicitly allows bond changes. The accessible component AA is then the connected set, under the physical dynamics, in the quotient space / ~e, that contains the chosen reference configuration. Because the physical dynamics are stochastic, this construction relies separation of timescales, which is well supported by the exponential scaling of rates with activation energies and by established isomerism taxonomies.
Isomerism, geometry, and parameer restraint The polytope formalism in Canfield et al. (2018) provides a peerdination-centriy view of stereochemical change and shows how fundamen- tal isomerization processes can be enumerated for each coordination number. In context, this
20

---

Published as a conference paper at ICLR 2026
motivates representing each stereochemical element by the smallest stable geometric feature that distinguishes its alternatiees, together with a numerical order parameter defined on those features. The order parameter then becomes the target of light, aopology-preservve emplify this with two main types of restraints considered in this work, bond restraint and tetrahedral center chirality restraint. While not discussed here, other types of isomerism, such axial, planar, and helical chirality, can be handled analogously.
restraint. Here,
(1) Bond restraint: The preservation of molecular topology, i.e., the fixed bond graph of the system, should be encoded explicitly as a restraint whenever bond connectivity is not supposed to change. The exceptions are when the simulation potential already contains topological information (e.g,, non-reactive classical force fields) or the scientific question explicitly involves bond making or bond breaking. Let riy (X) = ||ar; - 2ty be the instantaneous bond length between atoms and j. For each bonded pair or non-bonded pair to be preserved, we apply bond restraint or no-bond restraint, respectively:
Utie (X) kl krrig (rri(()) Tmax]
(14)
2
Uig
ano-bond (X)
= 2 x) x//min - -rt)
(15)
where [a]+ == max(z,0) 0). Here, we set Tmax = 1.3(r + ) Tmin = 0.7 + r rydw), and set the force constant k 100 eV/A², with r! and Mpبd denoting the covalent and van der Waals radii of atom i, respectively. For examples in this work, we found that no-bond restraint Eq. (15) could be omitted in cases where valency of atoms is well-defined by the base potential energy (see implementation details in Section D).
(2) Tetrahedral center chirality restraint: We define the torsion (dihedral) angle of atoms i, and 1, (X) E (-T1, m] as follows:
Oiyza (X) = atan2 (b2 ((b) x b2) x (b2 x ba)), (621(61 x 62) (b2 x ba)).
(16)
where the displacement vectors are defined as by := Ty by := I F. and by == 54 Tg. A proper torsion is the signed angle between two planes that share central bond, with three bonds in sequence, c.g., i-j-k-l, with the torsion angle describing the relative rotation around the bond j-k. In contrast, an improper torsion does not describe rotation around bond, but rather the out-of-plant plane displacement of one atom relative to a reference plane defined by three others.
We now consider tetrahedral stereocenter i with three neighboring atoms j. k. and (out of its four substituents). The corresponding improper torsion angle, Φijua( ()), exhibits a bimodal distribution: each mode corresponds to distinct stereoisomer, and in the ideal tetrahedral geometry, the peaks are centered near +355. To preserve the stereochemical configuration, we introduce the following flat-bottom harmonic restraint potential on the improper torsion:
> Eyke (X; Po, Pwo)) [ - Φb)1 croll
(17)
where w(a) atan2(sin cos a:) E (-T,, x] is the 27-periodic wrapping function, dn is the reference torsion value, and Φuat is the flat-bottom tolerance.
For clarity, we illustrate this construction using the C., stereocenter of an amino acid. In naturally occurring L-amino acids, the improper torsion defined by (Ca, N, C, C3) is positive, whereas the improper torsion (Ca, N. C, H.) is negative. Based on this, define the amino acid chirality restraint as
Jhhirall (X) Urimp
C..N.CC (X;K, + Pia)) + Timp U.N.C.H. (X;K, 151,6-000 -oo Piot),
(18)
where we H 25.0 eVrad, do 0.615 rad (35%), and Pua 0.436 rad (25%).
21

---

Published as a conference paper at ICLR 2026
B PROOFS AND NOTES
B.1 PROOF FOR PROPOSITION 3.1
Proposition 3.1 (Convergence of WT-ASBS). When Algorithm is performed until convergence, the bias potential V& converges almost surely to V* = -(1 Fs) + const, and hence the sampled distribution converges to the well-tempered target Eq. (7).
Proof. We follow Dama et al. (2014) for the bias decomposition and stochastic approximation framework, extending it with WT-ASBS-specific derivations and clearer notation.
Step Level and driving parts of the bias and the step size. We split the bias V&(s) into level
T. = [s] Is (s) ds and a driving part Dy(s) = Vs(s) L.- Only Da affects the biased distribution, and adding the scalar La shifts the energy by constant. Now, single step of well- tempered bias update can be split into
L&+1 Ls + exp
(- 7
B L&
A(Ds](8++1),
(19)
8
D (5) = Dk((s) exp
1-1 Lx
r[Dal(s,
(20)
where we define (s") and [()] 80) as follows:
A[D](6)= = [&] / hKa (8, (-10 D(6')) - ds,
(21)
r[Di(s, = hKo (s. exp (-7=1((s)) A[D](s').
D(5')
(22)
Now, under mild practical conditions, the driving bias (s((s) stays bounded, and exp(- Ds((s))
is bounded above and below by positive constants. Also, Is h.o (s, s') ds is continuous positive function s' on a compact set, so also has positive finite bounds. Therefore, from Eq. (21), we get the following bound with positive constant a and b:
a S A[Dx](sM+1) (19) (La41 - Lb) exp
B
S b.
(23)
This result implics that the asymptotic effective step size is bounded by positive multiples of :
<x = exp(-
-w) -x)
Step 2: Stochastic recursion to meare-fiedd ODE. Now, we define auxiliary time scale with To = and Tk+1 = Th + cg. Then, the driving part update Eq. (20) can be written in stochastic approximation form
(s) = (s) + (T(DD)]((s) + (s+1(s)),
(24)
where we rewrite Eqs. (21) and (22) with expectation over S ~ Dp, the CV marginal under the bias V = L 1+10 D,
(- 7- (>=D((s)) (-
A[D] = [ Es-ep (K, (- 3
ds,
(25)
[sD = ES-En hKo exp (-,² B ((s)) - A(D)),
(26)
and (k+1 martingale difference with bounded second moment. The mini-batch hill update in Algorithm provides an unbiased estimator of the expectation Es-Pp. From Step 1, e, x k. which satisfies the Sobbins-Monne step condition Σe C = oo and Σx ot < co (Robbins & Monro, 1951). Standard stochastic approximation arguments then imply that the polygonal interpolation ((T)) Dy: ()T) D2 + I-Th (Dx+1 D₂) for J [Tk: Tk+1) is an asymptotic pseudo- trajectory of the mean mean-field ODE d D(r) r[D(r)] (Kushner & Yin, 2003). Physically, this corresponds to the long-time limit equation in Dama et al. (2014), obtained by aggregating many small hill updates over a fixed window while the bias changes only slightly within that window.
22

---

Published as a conference paper at ICLR 2026
Step 3: Fixed points the asymptotic ODE. At fixed point D*, we have r[D*] = 0. From Eq. (26), this correspond to
[ (hoo exp ( 7-
(-)))) PD- (5') ds' = const. Vs E S.
(27)
For Gaussian Ko on a compact domain, the associated integral operator is positive definite. Hence, the solution is unique and satisfies
exp (-7- B
( D*()) exp((B(F(s) + D (s)))
const.
(28)
xfi. (8)
(())=='+ (1- F(s) + const.
(29)
Step 4. Global stability. Now, define the tempering-reweighte instantancous CV density
(-,-1 D(s:))
(s & exp
(30)
xVD(s)
Let Pie denote the target tempering-reweighnel distribution, i.e., a normalized solution of
/ (x, 8 ) (s') ds' const.
(31)
For Gaussian hills on compact S, this solution is unique and equals the uniform density pt (t)
[S] -1. Consider the Lyapunov functional defined as KL divergence between pa and Pw (:TT): T):
S pi (3) log (8; ds.
D(r) := (P ll Pw (sT))
(32)
Py (s)
We now show that a (()) $ 0. First.
d
(33)
dr
26
[
(- (s)r) - A[D(T)]
(34)
00 c(F)) / hK hKa (so D&) ds' - A(((TT)
(35)
@n c()) / hKo ('),s))))) pi, (s') ds' C(T),
(36)
where c(7) 0 and C(7) constants in Also, from Eq. (30), get By (- D(s:)) Ep..
1 d
dr P()
d
(s r) dr Pw (s:
(37)
7 dr
Finally, plugging Eqs. (36) and (37) into the time derivative of Eq. (32) yields d dr
D(T)
P (x)
(8)
(s; T) dr Pew (s;
(38)
By
D(::))
Byg
dr D(::) - Bpp..
(39)
dr
(66)
By
c(T) / [muo(osT)
- vio (8)) (s, (pm (io(') dads' < 0,
(40)
since the integral operator with the Gaussian kernel hKa is positive semidefinite and c(r) > 0. The equality in Eq. (40) holds if and only if (::r)) Pte. and D is strict Lyapunov functional and the ODE solution converges to the unique fixed point.
Hence, the mean-field ODE d D(T) I[D(7)) has globally attracting equilibrium D* and one
dr gets Dr + D* and Ve + in Eq. (29) almost surely by following the stochastic approximation Eq. (24).
23

---

Published as a conference paper at ICLR 2026
C IMPLEMENTATION DETAILS FOR WT-ASBS
C.I WT-ASBS WITH REPLAY BUFFER
Here, describe the practical implementation of Algorithm 1, which merges the well-tempered bias update into the AM replay buffer update within the inner step. The resulting WT-ASBS algo- rithm with replay buffer is shown in Algorithm 2. After updating the AM replay buffer Budy (line 6). we reuse the same samples, projecting them to CV space to obtain the CV values for Gaussian kernel deposition (line 7).
Algorithm WT-ASBS with replay buffer
Require: prior μ.. potential energy (xx), restraint (x)). CV E(z), control (æ t). corrector h4(2). replay buffers Badj and But
I: Initialize ((s) + 0, h +-
> Bias and IPF initialization
2: for Stage k 2, K do
3: for AM epoch
1,
Mag do
P Adjoint Matching with WT bias update where (A) = s stopgrad(u") (1))
4: Sample from model {(xb) (0) x ) 20())
5:
Compute adjoint target (1) + stopgrad(clip(V +U = V {) amm) + h2 (x(1))
6:
Update replay buffer Baij Buy u ((X)) x(()
)
7:
Update bias with sampled CVs (fi) E(x(()):
V(s) V(s) + h EN exp( 7-T V(s(t)) exp( 116-00113 2n
(41)
8:
Take L gradient steps VaLAM w.r.t. the AM objective:
(0) E4-000.3,, Bag.
|e (1,
orael??
(42)
9:
for CM epoch m 1,2, Man do
Corrector Matching
Sample from model ((x))) x(1)) pa(a) where ((4)) = stopgrad (k)
10:
Update replay buffer Ben Ben U ((X) x D
11:
12:
Take L gradient steps Co w.r.t the CM objective:
[cM( <- E(xo.x))BBn
[1ng (xi) Vx, log base (x11X)01
(43)
13: output Final control (a (z. 1). final bias (s)
Grid discretization of bias potential Since the bias potential is accumulated over the entire train- ing history, longer training would require summing over an increasingly large number of Gaussian kernels to evaluate V(V o E). To ensure computational efficiency, we adopt the standard grid dis- cretization approach commonly used enhanced sampling simulations. The bias stored on discretized grid with spacing chosen to be smaller than the kernel width a, and both energies and forces are obtained through interpolation between grid points.
Pretraining for WT-ASBS As proposed in ASBS (Liu et al., 2025), we warm-start the training process with the approximate bridge matching objective, which aligns the control model uy more closely with available data samples. In our setting, these samples are generated by short MD sim- ulations that perform local sampling around a reference configuration zee-- The bridge matching objective Epectrii is defined as follows:
Ceretrain (0) =
(X.)) as Vx, log base (X1
(44)
where q denotes the empirical distribution the MD samples, and 7 is a time-dependent weight- ing factor, chosen as K1jz/of, where Kıje is the conditional variance of the base process (see Sec- tion C.3). This pretraining is only approximate, since the marginal distribution of the optimal ug matches (Xx) exactly only when the base process is memoryless (i.e., when Xo and X, are inde- pendent). Nevertheless, for warm-start purposes, this approximation is sufficient to rapidly align the model with distribution concentrated in single conformational mode.
24

---

Published as a conference paper at ICLR 2026
C.2 MODEL ARCHITECTURE
We build on PaiNN (Schütt et al., 2021), an E(3)-equivariant graph neural network designed for molecular energy and force prediction, with minimal modifications: adding time-conditioning in- puts, incorporating vector product layers (Schreiner et al., 2023) to break parity symmetry, restricting the model to -qquivarrree,, and applying equivariant layer normalization (Liao et al., 2024) to improve training stability. Each atom is assigned a unique atomic embedding, as the reference con- figuration used to define the accessible component (Requirement 1) already specifies atom indices. The model outputs single vector per atom, which is scaled by de to produce the final control and corrector outputs.
C.3 BASE PROCESS
We employ the base distribution in Karras et al. (2022) (EDM), which defines o through power- schedule parameter p and yields an analytically integrable variance between two different times Kejss
= (1 t)o 1/0 + to 1/p]p o dr
(45)
min
ds 2+1/p 2+1/p
Kela
(46)
2p +
T max
0 min
With the reparametrized time 7 == dx2, = o,dW, are
Maja 6lj0
E [0, 1]. conditional base distributions for the driftless base
base (X, (Xx;; Xo, 1) Pejo
(47)
buse Prjo.1 (Xt X1) = [Xx,; (1 Xo + 7x 7a (17))))111)))
(48)
Translational invariance As in AS (Havens et al., 2025), enforce translational invariance by projecting the SDE onto the subspace XCCOM = {x Σx x² 0} C X, where æf E R is the position of i-th atom. This is achieved through the projection operator P. which acts on any E R= by subtracting the center of mass from each coordinate y for E [n].
111-1,11
Py,
P =
In
1,1
@ la
(49)
dx, = o. Pu,(X) (X,) dt PWW, For the base process, the dynamics are given by
Xo =0
(50)
X, = No + / o,PaW,
(51)
Var[X,|Xo] /' PP ds Kejo PPT
(52)
Sampling can be carried out by first drawing Y, N kajol) and then projecting to obtain
X, PY, which enforces the translational invariance constraint.
Adjoint matching Eq. (10) and corrector matching Eq. (11) are modified accordingly follows:
: arg min Eptrr: PG. [l),((u(X,)) (VVV (R- (Xx))) stopgrad(u), (53)
arg min E P(((X)) Vx, log P
base (X [7)
(54)
C.4 HYPERPARAMETERS
Model and training hyperparameters for each experiment are listed in Table 1. For WT-ASBS Al- gorithm 2, we use 100 gradient updates per AM and CM epoch, and amax = 100.0 to clip the energy gradient in the terminal cost. Models were trained with AdamW optimizer (Loshchilov
25

---

Published as a conference paper at ICLR 2026
Table Hyperparameters for WT-ASBS. Model and training hyperparameters for each experiment.
Alanine dipeptide
Alanine tetrapeptide
Nucleophilic substitution
Post-TS bifurcation
Hyperparameter
Models
Hidden dimension
128
128
128
Radial basis
64
64
32
64
Control layers
Corrector layers
3
3
3
Radial cutoff [Å]
8.0
8.0
8.0
8.0
Batch and replay buffer
AM/CM buffer size
16.384
4,096
4,096
4,096
Buffer samples per epoch
2,048
512
512
512
Batch size
2,048
512
512
512
Base SDE
Source distribution
((0, 1A*
(O 1A*) 0.001
N(O,, 1A²) 0.001
W((0, L 0.001
[Å]
0.001
mmax [Å]
6.0
6.0
5.0
6.0
EDM exponent P
3.0
3.0
3.0
3.0
Bias update
Initial Gaussian h [eV]
5 x 10-4
x 10
x 10
2 x 10
Gaussian width
[0.2, 0.2] rad
[0.3, 0.3] rad
[0.2, 0.2] A
[0.05, 0.05]
CV grid minimum
[- -w] rad
[-n, -,, - rad [+ t. rad [100, 100, 100]
[1.0, 1.0] Á
[0.0, -100]
CV grid maximum
[+m,, [m] rad
[4.0, 4.0J Á
[3.0, 1.0]
CV grid bins
[150, 150]
[150, 150]
[300, 200]
Periodic CV grid
True
Truc 12.0
False
False
Bias factor
8.0
12.0
30.0
Training
Pretraining learning rate
3 10-4
3 x 10-4
x 10
3 10-4
Pretraining steps
100,000
100,000
20,000
100,000
Training stages K
4 1,000
40
AM epochs Madj
1,000
500
1,000
epochs Man
500
500
250
500
Initial learning rate
10
1 x 10
x 10
3 10
Final learning rate
3 x 10
5 x 10
1 10
1 10
Gradient clipping value
100.0
100.0
100.0
26

---

Published as a conference paper at ICLR 2026
& Hutter, 2017) using weight decay of 10-3, and diffusion process integrated with Euler- Maruyama method using 250 steps. For the first AM stage, we applied a 50-epoch learning rate warmup, after which bias deposition began. After completing K training stages, we added fi- nal AM stage without further bias deposition. Learning rates followed a cosine annealing schedule during training. applied independently to AM and CM optimization.
Several hyperparameters affecting training duration were selected based on the convergence of Al- gorithm 2. The number of gradient updates was chosen ensure that the controller parameters could track the outer-loop bias updates, as examined further in the ablation studies in Section E. The AM and CM epochs per stage (Mad and Mar) were largely adopted from previous work (Liu al., 2025), and we found them sufficient for the AM and CM losses to stabilize within each training stage. The number of alternating training stages K was set according to the convergence of the bias or PMF, assessed by whether newly deposited bias became negligible or whether the estimated free energy differences along the selected CVs had stabilized.
D EXPERIMENT DETAILS
For small peptides, performed classical force field simulations with OpenMM 8.3.1 (Eastman et al., 2023), while reactions were modeled using UMA-S-1.1 uMLIP (Wood et al., 2025) as im- plemented in the fairchem-core 2.3.0 package (Shuaibi al., 2025) with MD) integrators in ASE 3.25.0 (Larsen et al., 2017). All MD)-based enhanced sampling (WTMetaD) simulations were car- ried out with PLUMED 2.9.3 (Tribello al., 2014: The PLUMED consortium, 2019).
D.1 ALANINE DIPEPTIDE
Energy model The potential energy the system was described by Amber f196 force field (Koll- man et al., 1997) with GBSA-OBC implicit solvation (Onufriev et al., 2004), implemented in OpenMM (Eastman et al., 2023).
Restraints We apply the chirality restraint Eq. (18) to the central alanine, Ushiza (X), during WT- ASBS training. This restraint is not used in MD-based sampling because chirality inversion does not occur under physical dynamics.
Pretraining data We generated pretraining data using a 100 ps molecular dynamics simulation at 300 K with a 4 fs time step, enabled by hydrogen bond constraints and hydrogen mass repartitioning (HMR: Hopkins et al. (2015)) with hydrogen masses set to 1.5 amu. The system was equilibrated for 10 ps before sampling 10 frames at 0.1 ps intervals. Simulations employed the Langevin middle integrator (Zhang et al., 2019) implemented in OpenMM (Fastman et al., 2023), with a collision rate of ps 11, In total, pretraining data generation required 2.75 10a energy evaluations.
Reference data We employed the test split of 10" samples from Midgley et al. (2022), gener- ated by replica exchange MD with 21 replicas spanning 300 K to 1300 K. Their dataset generation required 2.3 x 10 energy evaluations.
WTMetaD simulations We performed WTMetaD simulations on the backbone torsion CVs (ó, v6). The dynamics
propagated with Langevin middle integrator (Zhang et al., 2019) 300 K with a fs timestep and ps collision rate. Following a 10 ps equilibration at 300 K, Gaus- sian biases were deposited every 500 steps with widths of 0.2 rad for both variables. The bias update employed an initial Gaussian height of 1.2 kJ/mol, a bias factor of 8, and grid covering [- π]² with 150 x 150 bins. Each production run was 10 in length, and simulations were repeated four times to ensure statistical reliability, yielding 4.0 x 10" energy evaluations.
D.2 ALANINE TETRAPEPTIDE
Energy model The potential energy of the system was described by Amber f199SB-ILDD force field (Lindorff-Lard et al., 2010) with Amber99- OBCC implicit solvation (Onufriev et al., 2004), as implemented in OpenMM (Fastman et al., 2023).
27

---

Published as a conference paper at ICLR 2026
Restraints We apply the chirality restraints Eq. (18) to three alanine residues, Ala (X) U ) + X (X), during WT-ASBS training. As in the alanine dipeptide case, these restraints are not used in MD-based sampling.
(X) +
Alud
Pretraining data We generated pretraining data using a ns molecular dynamics simulation at 300 K with 4 fs time step, enabled by hydrogen bond constraints and HMR with hydrogen masses set to 1.5 amu. The system was equilibrated for 10 ps before sampling 103 frames at ps inter- vals. Simulations employed the Langevin middle integrator (Zhang al., 2019) as implemented in OpenMM (Eastman et al., 2023), with a collision rate of ps- In total, pretraining data generation required 2.52 10t energy evaluations.
Reference data We generated unbiased alanine tetrapeptide conformations by first performing parallel tempering simulations to sample initial states, followed by long unbiased molecular dy- namics runs for data collection. Parallel tempering simulations were carried out with OpenMM and OpenMMTools (Chodera et al., 2025) using MPI parallelization across 8 replicas, with temperatures geometrically spaced between 300 K and 600 K. Replicas were propagated with the generalized hy- brid Monte Carlo (GHMC) integrator (Stoltz et al., 2010) using a 2 fs timestep (with hydrogen bond constraints), 1 ps-1 collision rate, and 1000 steps between exchange attempts (2 ps). Simulations were initialized from energy-minimize structure, equilibrated for 200 ps, and then run for 100 ns of production with replica exchanges attempted every 2 ps.
From the 300 K replica, 64 independent unbiased MD simulations were seeded from randomly selected frames. Each simulation employed a Langevin middle integrator (Zhang et al., 2019) at 300 K with simulations were performed for 1.56 us, yielding a total of 100 js simulation time across all replicas. Trajectories were saved every ps, producing 10" samples total. Reference data generation required approximately 1.01 x 10 energy evaluations.
timestep and ps collision rate. After 100 ps equilibration, production
WTMetaD simulations We performed WTMetaD simulations on the three backbone torsion CVs (Ø1, 62, da). The same integrator settings as in the reference data generation were used. Following 10 ps equilibration at 300 K. Gaussian biases were deposited every 500 steps with widths of 0.3 rad for each CV. The biasing scheme employed an initial Gaussian height of 1.2 kJ/mol, a bias factor of 12, and grid covering [-x, T with 753 bins. Each production was 20 ns in length, and simulations were repeated four times to ensure statistical reliability, yielding 8.0 x 10 energy evaluations.
b
14
tropone
016
11
d220
23
9
d2
dg
02
31
18
32
25
17
19
29
dimethylfulvene
Figure 6: Atom index definition for chemical reactions. (a) Nuclcophilic substitution (Sx2). (b) Cycloaddi- tion reaction between tropone and dimethylfulvene with post-TS TS bifurcation.
D.3 NUCLEOPHILIC SUBSTITUTION
Energy model The potential energy of the system was described by the UMA-S-1.1 (Wood al., 2025) model with the DFT task label omol, system charge of -1, and singlet spin state (total spin 0).
Restraints In WT-ASBS, three C-H bonds (1-2, 1-3, and 1-4 in Fig. 6a) are maintained using the bond restraint in Eq. (14). For both WT-ASBS and WTMetal), keep the accessible CV space
28

---

Published as a conference paper at ICLR 2026
S bounded, the same restraint form is applied to two C-Cl distances (1-5 and 1 1-6) with Fmax = 3.3 A. preventing chloride ions from escaping the reaction complex.
WTMetaD simulations We performed WTMetaD simulations using (d1, d2) the CVs. The dynamics were propagated with a Langevin thermostat at 300 K, 0.5 fs timestep, and a friction coefficient of 10 ps -. Initial velocities were assigned from a Maxwell-Boltzmanl distribution, and the total momentum was removed. Following a ps equilibration at 300 K. Gaussian biases were deposited every 250 steps with widths of 0.2 A for both CVs. The biasing scheme used an initial Gaussian height of 0.02 cV. a bias factor of 20, and grid spanning [1.0, 4.0) Á x [1.0, 4.0] A. Each production run was 0.5 ns length, and simulations were repeated four times to ensure statistical reliability, yielding 4.0 x 10 energy evaluations.
D.4 POST-TRANSITION-SAAA BIFURCATION
Energy model The potential energy the system was described by the UMA-S-L.1 (Wood et al., 2025) model with the DFT task label omol, a system charge of 0, and singlet spin state (total spin 0).
Collective variables We define normalized contact e₁ that maps the bond distance de to a soft bond indicator,
1 (dg drer)
4 1 - (de/ )er) [0,
(55)
which approaches when d; is much smaller than dnet and decays to 0 as de increases (bond disso- ciation). To construct the CVs, we express the three relevant C-C distances (Fig. 4d) as normalized contacts using Eq. (55) with deef 1.7 A, thereby capturing bond formation progress (Fig. 4e). We then define two CVs: 81 Ci + 02 + C3. which measures overall bond formation, and =
which distinguishes between products 4 and 5. A similar approach has been applied to enhanced sampling of mechanistically related reactions (Nam & Jung, 2023).
Restraints In WT-ASBS, bonds present in the reactant configuration (3, Fig. 4c) are preserved using the bond restraint in Eq. (14). Furthermore, to prevent tropone (upper reactant) from "tele- porting" to the opposite face of dimethylfullvnee (lower reactant), we enforce an inter-planar distance defined by three atoms per plane:
d(x)= = (a (1)
2 (1)
x (12 F0 (1)
1
2
(I) 2
(56)
[1xx]
(1)
20 x 32
3
m==
3
m==
Unit normal of plane
Inter-centroid displacement
d(X) is positive if the centroid of plane 2 lies on the side of plane pointed to by its normal vector, and negative otherwise. The magnitude is the perpendicular distance between the two centroids along the normal direction of the first plane. We evaluate d(X) using three pairs of bond-forming atoms (plane I: dimethylfulvene atoms 15, 13, 24; plane 2: tropone atoms 2, 3, 7 in Fig. 6b) and
apply a lower harmonic wall restraint Uplane (X) = 2 x[ d with k = 100 eV/A² to enforce d(X) 0.
Finally, for both WT-ASBS and WTMMeaDD, the no-bond restraint in Eq. (15) is applied to atom pairs C6-C24, C3-C18, and O8-C19 suppress sampling of irrelevant reactive pathways.
Pretraining data The system was propagated with a Langevin thermostat at 300 K, a fs timestep, and friction coefficient of ps -1 Initial velocities were assigned from a tribution and the total momentum was removed. After 2 ps of equilibration (2000 steps), production dynamics were carried out for 10 ps (10k steps), saving configurations every 10 fs. In total, pretrain- ing data generation required 1.2 10 sequential energy evaluations.
dis-
WTMttaD simulations We performed WTMetaD simulations using ($1, S2) as the CVs. The dynamics were propagated with a Langevin thermostat at 300 K, 0.5 fs timestep, and a friction coefficient of 10 p0 Initial velocities were assigned from a Maxwell-Boltzmam the total momentum removed. Following ps equilibration 300 K, Gaussian biases were
Boltzmann distribution, and
29

---

Published as a conference paper at ICLR 2026
deposited every 250 steps with widths of 0.05 for both CVs. The biasing scheme used initial
Gaussian height of 0.02 eV, a bias factor of 30, and grid spanning [0.0, 3.0] x [-1.00, 1.0]. Each production was 0.8 ns in length, and simulations were repeated four times to ensure statistical reliability, yielding 6.4 x 10 energy evaluations.
E ABLATION STUDIES
E.1 WELL-TEMPERED BIAS UPDATE
Here, we examine how the well-tempered bias hyperparameters (Gaussian height h, width o, and bias factor 7) influence bias deposition, the intermediate sample distribution, and the convergence of free energy differences, using alanine dipeptide as a representative example.
Effect of
Effect of
Effect of
1.5e-3
0.3
16
5e-4
0.2
AF (kcat/mol)
1.5c-4
1.5e-5
0.0
2.5
5.0
7.5
0.0
2.5
5.0
7.5
0.0
2.5
5.0
7.5
Energy evaluations [x109]
Energy evaluations [x10%)
Energy evaluations [x10%]
Figure 7: Ablation of well-tempered bias updates on free energy convergence. The free energy difference AF between states of alanine dipeptide is shown function of the number of energy evaluations during training. Each panel varies h, a, and 7, and the default values used in the main text are indicated in gray background. and a are in units of eV and radians, respectively.
The convergence of AF shown in Fig. 7 across different well-tempered bias hyperparameters indi- cates that the method is robust over a broad range of settings. For the Gaussian height, values from 1.5 X 10-5 to 10-4 4 eV yield nearly identical converged results despite spanning almost a factor of 30, while only the largest value, h 1.5 x 10 eV, shows slight deviation. As seen in Fig. 8, increasing h enhances early exploration of new modes and leads to a faster initial approach to the target AF, but excessively large h can impair final convergence. We note that, across different experiments, h values on the order of 10 eV perform well. This is roughly two orders mag- nitude smaller than the h values commonly used in standard WTMetaD simulations, since multiple Gaussians are deposited simultaneously during WT-ASBS training.
For a, larger values promote faster initial exploration by producing broader bias early in training (Fig. 9), while their impact on AF convergence is comparatively modest (Fig. 7). For 7, we observe that small value such as y = 4 leads to much slower convergence (Fig. 7) and reduced exploration (Fig. 10), as the effective temperature along the CVs is lower. For both a and y, we observed that values commonly used in WTMetaD simulations provide good performance.
E.2 TRAINING MECHANISM
Now we perform analogous ablations on two training mechanisms central to the practical implemen- tation of WT-ASBS: the replay buffer and data-based pretraining. Although several replay buffer hyperparameters, such as AM/CM buffer sizes, the number of buffer samples per epoch, and the number of gradient updates between buffer refreshes, can influence training dynamics, we use the number of AM gradient updates (Algorithm T. in line 8) as representative control variable to assess the effect of the replay buffer. Because the replay buffer update is coupled to the bias update, increasing makes both updates less frequent.
First, the AF convergence in Fig. 11 shows that the replay buffer has substantial impact on training efficiency: using only a small number of AM steps per replay buffer update slows convergence significantly. The bias deposition patterns in Fig. 12 further indicate that overly frequent buffer and bias updates prevent the sampler parameters from keeping pace with the evolving bias, causing excessive bias accumulate in the initial basin and leading slow exploration of the next state. This
30

---

Published as a conference paper at ICLR 2026
h= 1.5e-3 0.15M evals
h= 1.5e-3 0.35M evals
h= 1.5e=3 0.64M evals
h= 1.5e-3
== 1.5e- 8.34M evals
1.66M eealls
40 (kcak/mol)
10
0-
-2/2
h= Se-4
h= 5e-4 0.35M evals
h- Se-4 0.64M evals
5--4 1.66M Evaals
h-
0.15M evals
8.34M E evals
Via M).0) (kcal/mal)
10
-2/
2
1.5ee4 0.15M evals
1.5e-4 1.66M Eevals
1----
0.35M E evals
0.64M evals
8.34M E evals
VieVVie [acalimal]
110
6
n/2
-1
5e-5 1.66M evals
0.15M E evalls
0.35M E evals
0.64M E evals
8.34M E evals
40 Ikca/moll
10
8
5
m//
77
h- 1.5e-5
h- 1-15 0.35M E evals
h- 1.5e-5
1.5e- 1.66M Eevals
15-5 8.34M E evals
0.15M Eevals
0.64M E evals
wll [kcal/moll
8/2
10
8
6
mm/2
2
11
-n/2
-n/2
0
11
-n/2
m/2
Figure 8: Bias deposition and projected samples during training for different values. row shows, for selected h, how the deposited bias V(á, w)) (shifted to have a minimum of zero) evolves with increasing numbers of energy evaluations, along with the samples produced by the process at the corresponding training step, projected onto the &- space. h values reported units of eV.
illustrates that sufficient parameter updates between buffer updates help bridge the gap between the ideal woo-iime-scat WT-ASBS in Algorithm 1, where ASBS training converges at each fixed bias, and the practical WT-ASBS in Algorithm 2. where parameter and bias updates proceed concurrently. The replay mechanism must therefore provide enough time for parameter updates to keep up with bias deposition.
Next, Fig. 11 shows that A convergence is substantially faster when training begins from localized pretraining. The sample distributions in Fig. clarify this behavior: without pretraining, the model must identify low-energy modes from scratch, which requires large number of energy gradient evaluations during the early stages of training. While the well-tempered bias update can converge to the correct PMF under mild conditions regardless of the initial sampling distribution, these results show that beneficial to exploit the initial structure of the chemical system and the efficiency of local, dynamics-batt exploration.
31

---

Published as a conference paper at ICLR 2026
0.15M evals
0.35M evals
0.64M evals
1.66M Eevals
8.34M evals
Bias [kcal/mol)
0
-/2
a-
0.35M evals
1.66M
Vio. [kcal/mol]
8/2
-10
0-0.1 0.64M evals
0.15M evals
0.35M evals
1.66M eevals
8.34M evals
Bia Ve. w0 [kcal/mol)
x/2
0-
-n/2
-1/2
0
m/2
-n/2
m/2
-n/2
0
m/2 -17
-2/2
n/2
-n/2
0
n/2
Figure 9: Bias deposition and projected samples during training for different values, Each row shows, for selected a, how the deposited bias V (é, w) (shifted to have minimum of zero) evolves with increasing numbers energy evaluations, along with the samples produced by process the corresponding training step. projected onto the Φ space.
y- 16 0.15M evals
y=16 0.35M eeals
yo 0.64M evals
y= 16 1.66M Eevals
y-16
8.34M evals
Bias di Ikcal/moll
0-
-//2
y-8 0.15M
0.35M y-B
y-8 1.66M
y-8 8.34M
0.64M evals
Vio d0 [kcal/mol]
8/2
n/2
4-4 0.35M evals
y-4 0.64M evals
y-4 1.66M eevals
y-4 8.34M evals
0.15M Eevals
Vi. ui [kcal/mol)
m/2
0-
mn/
-1/2
7/2
m/2
m/2
-//2
0
n/2
-77
-2/2
n1/
-n/2
7/2
Figure 10: Bias deposition and projected samples during training for different values. Each row shows, for selected 7, how the deposited bias V ((, (shifted to have minimum of zero) evolves with increasing numbers energy evaluations, along with the samples produced by process the corresponding training step. projected onto the Φ- space.
32

---

Published as a conference paper at ICLR 2026
Effect of the replay buffer (train steps per epoch)
Effect of pretraining
pretrain
AF [kcal/mmll
pretrain
25
8-
100
6-
3
0
2
8
2
8
Energy evaluations [x109)
Energy evaluations [x10f]
Figure 11: Ablation of training mechanisms on free energy convergence. The free energy difference AF between two states of alanine dipeptide is shown function of the number of energy evaluations during train- ing. Each panel varies the number of steps per epnch (replay buffer update frequeney) and whether pretraining used.
steps/epach 0.15M E evals
steps/epoch 0.35M evals
steps/epoch 0.64M E evals
steps//pochh 1.66M E evals
steps/epach 8.34M E evals
Ikca/mmol
0-
steps/epach
steps/epoch 0.35M
steps/epach 1.66M
steps/epech 8.34M
Ikcal/moll
2/2
-10
6
m//
steps/spoch 25 0.15M E evals
stops/spoch = 25 0.35M evals
stops/epoch 25 0.64M E evals
stops/spach = 25 1.66M E evals
staps/spoch 25
8.34M E evals
Bias Ikcal/moll
10
m2/
stops/epach 100
steps/epoch 100 0.35M
steps/epoch 100 0.64M
steps/epach 100 1.66M
stops/spach 100 8.34M
Dias a) [kcal/moll
x/2
10
6
4
n/2
2
77
2
/
-n/2 0
0
2
6
Figure 12: Bias deposition and projected samples during training for different number of AM steps per epoch. Each row shows, for selected number steps per epoch, how the deposited bias V x) (shifted to have minimum of zero) evolves with increasing numbers of energy evaluations, along with the samples produced by the process at the corresponding training step, projected onto the d-v space.
33

---

Published as a conference paper at ICLR 2026
No
pretrain 0.35M evals
pretrain 0.64M evals
pretrain 1.66M evals
pretrain
0. 0.15M evals
8.34M E evals
Bias Vid () (kcal/moll
R/2
10
-n/2
-Xt
Pretrain 0.15M E evals
Pretrain 0.35M E evals
Pretrain 0.64M evals
Pretrain 1.66M evals
Pretrain 8.34M E evals
kkal/molll
-10
-8
-
m/2
-
2
m/2
-n/2 0
6 -11
-n/2
0
Figure 13: Bias deposition and projected samples during training with or without pretraining. The upper row shows results without pretraining and the lower row shows results with pretraining. Each row illustrates how the deposited bias V 1 (shifted to have minimum of zero) evolves with increasing numbers of energy evaluations, along with the samples produced by the process the corresponding training step, projected onto the d- space.
F ADDITIONAL RESULTS
ALANINE DIPEPTIDE
In this section, extend the alanine dipeptide sampling comparison from the main text standard distributional metrics following Liu et al. (2025). Specifically, we report KL divergence for 1-D marginals of torsion angles, well as W2 and Jensen-Shannn Shannon divergence (ISD) for the joint distribution of backbone torsions (á, 16). Results for previous diffusion samplers using internal coordinate representations (bond distances, angles, and torsions) and for models trained in this work using Cartesian coordinates are summarized in Table 2. We used 10t samples to compute DxL and ISD, and randomly selected 104 samples for W2 (with weighted random sampling for WT-ASBS).
Table 2: Distributional metrics for alanine dipeptide. Comparison of diffuusio samplers applied in internal or Cartesian coordinate representations alanine dipeptide sampling. The divergence (Dxt. five I-D torsion angle marginals and the Wasserstein-2 distance are reported. "Results from Liu et al. (2025), Resull from Choi et al. (2025). 'Reweighted samples are used.
(W2) and JSD on joint backbone torsion distribution
DEL on I-D marginal (L)
(Ø, 2b) joint
Method
m
72
7
Wa (L)
JSD (L)
Internal coordinates
PIS (Zhang & Chen, 2021)*
0.05 10.00
0.38
1:049
5.61 1:124
4.49 10.03
1.60
1.27 1:19
DDS (Vargas et al., 2023a)*
4.03
0.16
2.44
0.03
0.03
0.68 =0,09
AS (Havens et al., 2025)*
0.09
0.04
0.17:1017
0.562om
0.51
0.65
NAAS (Choi et al., 2025)t
0.26
0.24
0.27
0.13
0.16
ASBS (Liu et al., 2025)*
0.02 +0.00
0.01
0.03 0:00
0.02
0.02 +:000
0.25
Cartesian coordinatey
ASBS (Liu et al., 2025)
0.14
4.09
0.02-1000
2.00
0.00
0.45-m.01 n.00
0.014=0.0m
WT-ASBS (this w
0.02
0.04
0.02:1000
0.00
0.00
0.43
0.005
Table 2 shows that WT-ASBS performs well on all metrics, with similar W2 performance to base- line ASBS (Cartesian) and slightly worse than reference ASBS (internal coordinates). However, the Ramachandran plots in Fig. 14 reveal that WT-ASBS better reproduces the overall density, while the other two methods either collapse into highly populated regions or leak into sparsely populated ones. This discrepancy arises because the metrics in Table assume equal weight for all samples.
34

---

Published as a conference paper at ICLR 2026
Reference
ASBS [internal)*
(Cartesian)
WT-ASBS (Cartesian)
8/2
Log density
0
11
-7/2
-n/2
Figure 14: Ramachandran plots for alanine dipeptide. Comparison backbone torsion distributions (ó, v)) in log densities, for reference, ASBS (internal coordinates), ASBS and WT-ASBS (Cartesian coordinates), as in Table 2. All three methods reproduce the more populated regions reasonably well, but their behavior in less populated regions differs substantially. Negligibly low-density regions omitted for visual clarity. "Taken from Liu al. (2025).
WT-ASBS, designed satisfy Requirement 3, incorporates weighted sampling achieve balanced efficiency across all relevant modes. Since Boltzmann populations depend exponentially free energy, focusing on more populated regions can yield good metries without capturing the full distri- bution. Therefore, emphasize evaluations more closely aligned with the design choices discussed in the main text.
F.2 ALANINE TETRAPEPTIDE
In Fig. 15, show the distribution of pretraining data and compare the PMF obtained by reweight- ing (Eq. (1)) for WT-ASBS, the bias (Remark 3.1) for WTMetaD and WT-ASBS, and the reference PMF. The populations within each mode are well reproduced, while biased sampling methods (WT- ASBS and WTMetaD) capture the energetic features of the transition region between modes. This region is inaccessible even with a huge number of unbiased samples (5 107 for reference), underscoring the necessity of biased sampling.
F.3 SAMPLING REACTIVE ENERGY LANDSCAPES
To quantify the convergence of PMFs in Fig. 5, we define the PMF MAE between a PMF F(s) and its reference (s) both shifted have minimum zero,
MAE[F, F.el min. Ss [F()) Fre((s + 1(Frer()) < Fitnes) ds IS 1[Fir(()) < Fihrcs]
(57)
ie., the average absolute deviation (up to optimal constant shift e) over regions of CV space where (s) is below the threshold Fines- We used Fitres of 30 kcal/mol for the Sx2 reaction, and 45 kcal/mol for the post-TS bifurcation reaction.
We compare PMFs for reactions obtained with bias from WT-ASBS and WTMetaD in Figs. 16 and 17, and confirm that the final PMFs from both methods agree well. For the post-TS bifurcation, valuable chemical insights into the mechanistic understanding of reactions can be obtained from the free energy profiles. For example, while 5 is the experimentally observed product,
competing mechanisms of 3 - 4 + (Paddon-Row & Warrener, 1974) vs. 3 + (Houk et al., 1970) have been proposed. The 2-D PMF reveals that the initial reactive flux is biased toward product 4. but the deeper free energy basin of makes it the thermodynamically favored product, consistent with previous findings that 4 forms transiently and subsequently converts to 5 (Yu et al., 2017).
F.4 ML COLLECTIVE VARIABLES
To illustrate the generalizability of the method non non-hand non-hand-crafted CVs, present initial exper- iments coupling WT-ASBS with an ML CV learned automatically from unbiased samples of the metastable states of alanine dipeptide.
35

---

Published as a conference paper at ICLR 2026
Pretralninn data
PMF from reweighding (WT-1585, samples)
PMF from WVT-A8B5)
PMF from (VTMe::D
Reference PMF (100 MD)
5. 2 p00 14. (hraimol
n/2
1/2
A/2
6'
0-
0
0
-1
-11
-5
1/2
4. 5 MM: Ná. henhhmoll
n/2
n/2
N/2
-07
-072
-1/2
M
-5
./2
x/2
-7
2/
a/2
-7
1/2
-1
n/2
1/2
1/2
6. PNI 1ib Acaiimall
1/2
n/2
n/2
n/2
x/2
x/2
M/2
a/2
n/2
22
2
2
m/2 2/2
m/2
n/
n/2 -2/
n/2
Figure 15: Alanine tetrapeptide PMFs. Columns show pretraining data, PMFs from reweighting (WT-ASBS), PMFs from WT-ASBS and WTMetaD biases, and the reference PMF. Rows correspond different torsion angle pairs (2). (é), Φs). and (d2, Φs). Dotted lines at di = and 2 mark the mode boundaries. All method reproduces populations within metastable modes, while biased sampling (WT-ASBS and WTMetaD) captures transition regions inaccessible unbiased sampling.
WTMetaD
WT-ASBS
28
3.5
PMF Fid, dg) [kcal/mll]
24
3.0
3.0
20
2.5
2.5
16
12
2.0
2.0
8
1.5
1.5
1.5
2.0
2.5
3,0
3.5
1.5
2.0
2.5
3.0
3.5
d, [A]
4 [A]
Figure 16: Sx2 reaction PMFs. Comparison of PMFs obtained from WT-ASBS and WTMetaD biases. Red dotted lines mark regions influenced by restraints, gray dotted lines denote the di = d2 symmetry lines, and the purple cross indicates the transition state location.
Starting from 5,000 unbiased configurations collected from each state (Fig. 18a), we trained a Deep- TDA (Trizio & Parrinello, 202 1) CV that maps the samples to a one-dimensional space in which each state corresponds
single Gaussian distribution. We used the default settings of the micolvar package (Bonati et al., 2023), employing a multilayer perceptron CV model that takes pairwise heavy-atom distances as inputs and includes hidden layers of sizes 24 and 12. We assign states and 2 to Gaussians (-700 1.0²) and /V 1.0²), respectively. The resulting MI. CV values for long unbiased reference configurations (including transitions) are shown Fig. 18b: configurations from the two metastable states cluster around 77 and 7, while transition configurations lie smoothly between them. We then perform WT-ASBS training using the same settings as the main-text alanine dipeptide experiment, except for smaller Gaussian height h 1 x 10 eV to account for the reduced CV dimensionality.
The final 1-D PMF along the CV in Fig. 18c shows good agreement with the reference PMF, and AF approaches the reference value smoothly (Fig. 18f). The 2-D PMF on (Ø. e)) obtained by reweighting 1M samples with their bias values (Fig. 8d) also agrees reasonably well with the reference 2-D PMF
36

---

Published as a conference paper at ICLR 2026
WTMetaD
WT-ASBS
Fisyl [kcal/mol]
42
PMF FMF. [kcal/mol]
0.5
36
5 Cs
30
0.0
24
18
-0.5
12
6
1
Figure 17: Post-TS bifurcation reaction PMFs. Comparison of I-D PMF F(s)) and 2-D PMF F(s1, 82) obtained from WT-ASBS and WTMetaD biases. Red dotted lines mark regions influenced by restraints, gray dotted lines denote the ca symmetry lines, and the purple cross indicates the transition state locations.
a
ML CV training data
b
ML values
PMF along CV
Reference WT-ASBS
State 2
PMF |kcak/mol)
TDA CV
6
0
-m/2
-m/2
-7
-7
-27
-n/2
n/2
-7
-R/2
m/2
-10
10
PMF from reweighting
d
(1M samples)
e
Reference PMF
AF convergence
PMF Fig. vi [kcal/nal]
WT-AS85
n/2
m/2
&F [kcal/mol]
5-
0
4-
-n/2
-n/2
Reference AF
n
-1
-7
R/2
n/2
-7
n/2
n/2
Energy ovaluations (x109
Figure 18: WT-ASBS for alanine dipeptide with MI CV. (a) CV training data from each state in the (á, 1)) space. (b) Trained MI. CV values for reference equilibrium configurations. One-dimensional PMF along the Deep-TDA CV obtained with WT-ASBS and from reference unbiased configurations. (d, e) Two-dimmesiional PMF on (ø, +) from (d) WT-ASBS with reweighting and (e) reference unbiased configurations. (f) Conver- gence of AF between the two states for WT-ASBS as function of the number of energy evaluations.
(Fig. 18e), though the modes appear broadened along the s direction, which is less well resolved by the ML CV. This behavior is consistent with the difference between baseline ASBS and WT- ASBS (Fig. 14), where WT-based weights help recover structure sparsely sampled regions by downweighting configurations that are distinguishable in the CV space. These preliminary results suggest that WT-ASBS can be combined with ML CV pipelines to enable enhanced exploration when good CV is not known priori.
G LARGE LANGUAGE MODEL STATEMENT
We made limited use of large language models (LLMs) to assist with language refinement and proofreading. No content, ideas, or analyses were produced by these tools.
37


Reviews:
[{"review_id": "8JC6fviJPp", "reviewer": "Reviewer_xH8L", "summary": "In this work, the authors present a collective-variable-based approach to enhance diffusion-based samplers. By incorporating a biasing force that is gradually deposited, similar to metadynamics, the approach improves the exploration ability of diffusion-based samplers. This addresses a consistent issue with diffusion-based samplers being susceptible to mode collapse.\n\nThe specific diffusion-based sampler enhanced using collective variables (CVs) is the Adjoint Schrödinger Bridge Sampler (ASBS), although the CV-based approach does not appear to be necessarily specific to this instance of diffusion-based sampler. The biasing along the CV is approached similarly to most biased enhanced-sampling methods: an additional bias potential is added to the base potential energy. This bias potential is then updated using the same rule as in well-tempered MetaD. This results in a two-step optimization process: an inner step, in which ASBS is performed (resulting in mode collapse), and an outer step, in which the bias potential is updated. In addition, a few practical improvements are implemented, such as the use of a replay buffer and a warm start.\n\nThe presented method is evaluated on a number of interesting systems. First, the evaluation considers two peptides, Alanine Dipeptide and Alanine Tetrapeptide, and, following this, the authors study reactive pathways in the form of a simpler SN2 reaction and a more complex reaction with multiple products.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The authors present an interesting and seemingly successful approach to alleviate the well-known issue of mode collapse in diffusion-based samplers within the domain of molecular conformations. The focus on ASBS as the diffusion-based sampler to enhance, and the choice of the well-tempered bias update, also seems the most sensible approach. Furthermore, the paper is well written and does an excellent job highlighting the core considerations that go into an enhanced-sampling method.\n\nThe experimental evaluation is relatively in-depth and focuses on two interesting problems by considering both simple peptides and reactions. While quantitative results are limited, the presented results suggest a significant improvement over standard WTMetaD.\n\nAll in all, the paper presents an excellent approach to enhance diffusion-based samplers to make them more appropriate for the equilibrium sampling of a molecular system. As such, I vote to accept the paper for publication.", "weaknesses": "- Looking at Figure 2, I’m surprised to see roughly uniform sampling from the second outer-loop iteration onwards (assuming that this is what the dots in panel c represent). I would have expected the samples to be located around the two smaller modes at this point, or still located at the original modes. Only at the later iterations would I expect to see roughly uniform sampling. Notably, this is also directly in contrast with the intuition presented in Figure 1. An author response to this inconsistency is required for me to increase my score beyond a borderline accept.\n- There are a number of important components to the training setup that would be good to see studied more in-depth, such as the use of a replay buffer and the warm start. Most importantly, I am interested in seeing how these components influence the number of energy evaluations needed.\n- Due to the reliance on collective variables, it is difficult to envision the applicability of the presented approach outside domains where CVs are known and provide a clear boundary to the sampling domain. For example, in areas such as the study of protein folding dynamics, it is hard to define CVs that sufficiently restrict the important sampling domain.\n- As the work could potentially spark new interest in enhanced sampling, it would be good to see a more in-depth discussion of work in this area, such as umbrella sampling and adaptive biasing force methods. This could potentially be combined with the discussion of the requirements in Section 2.1, whose current purpose is not entirely clear. Additionally, it would be good to see additional works discussed, such as [1–4] (not a complete list), all of which use a similar biasing force to that proposed here in the context of enhanced sampling (but are more focused on transition path sampling).\n\n[1] Seong, Kiyoung, et al. “Transition Path Sampling with Improved Off-Policy Training of Diffusion Path Samplers.” arXiv preprint arXiv:2405.19961 (2024).\n[2] Singh, Aditya N., Avishek Das, and David T. Limmer. “Variational path sampling of rare dynamical events.” Annual Review of Physical Chemistry 76 (2025).\n[3] Holdijk, Lars, et al. “Stochastic optimal control for collective-variable-free sampling of molecular transition paths.” Advances in Neural Information Processing Systems 36 (2023): 79540–79556.\n[4] Du, Yuanqi, et al. “Doob’s Lagrangian: A Sample-Efficient Variational Approach to Transition Path Sampling.” Advances in Neural Information Processing Systems 37 (2024): 65791–65822.", "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": "Wrj8zKQ3rf", "reviewer": "Reviewer_5dYX", "summary": "This paper introduces the Well-Tempered Adjoint Schrödinger Bridge Sampler (WT-ASBS), a novel method that significantly enhances the exploration capabilities of energy-based diffusion models for molecular systems. The core innovation is integrating an adaptive bias, inspired by well-tempered metadynamics (WTMetaD), into the Adjoint Schrödinger Bridge Sampler (ASBS) training loop. This bias is deposited along pre-defined Collective Variables (CVs), which are low-dimensional projections of molecular coordinates. The bias effectively increases the sampling temperature in the CV space, mitigating the mode collapse pitfall often seen in standard diffusion samplers. Crucially, the method retains the ability to recover the correct Boltzmann ensemble through importance reweighting. The authors demonstrate WT-ASBS's efficacy on conformational sampling benchmarks (alanine dipeptide and tetrapeptide) and, notably, achieve the first demonstration of reactive sampling using a diffusion-based model on $S_N2$ and post-transition-state bifurcation reactions, showing significant efficiency gains over traditional WTMetaD in terms of wall-clock time and energy evaluations.", "questions": "See weakness", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The reviewer here acknowledges that I am familiar with MCMC methods and general energy-based sampling, but is less knowledgeable about the specifics of modern molecular sciences applications. \n\n- The work presents a highly original and timely unification of two distinct fields: enhanced sampling (metadynamics) and generative diffusion modeling. By incorporating the well-tempered biasing mechanism directly into the iterative proportional fitting (IPF) scheme of the ASBS, the authors solve the critical problem of mode collapse for diffusion models applied to complex, high-dimensional, and multi-modal free energy landscapes. The successful application of reactive sampling, which is extremely challenging for standard diffusion models, is a significant first.\n\n- The paper is technically sound and well-executed.\n\n  - The authors provide a convergence guarantee (Proposition 3.1), showing that WT-ASBS provably approaches the desired well-tempered target distribution.\n\n  - The experiments are convincing. On the alanine tetrapeptide, WT-ASBS discovers all eight metastable states much faster than WTMetaD (Figure 3c). The ability to accurately resolve complex free energy landscapes (PMFs) for both conformational and reactive systems, with close agreement between PMF from bias and PMF from reweighting, strongly supports the method's correctness and efficiency.\n\n  - For the complex reactive systems using universal ML interatomic potentials (uMLIPs), WT-ASBS achieves convergence with significantly fewer energy evaluations and less wall-clock time compared to WTMetaD (Figure 5).\n\n- The paper is well-structured and clearly written, making complex concepts accessible. The methodological details, especially the two-time-scale algorithm and the practical implementation with a replay buffer (Algorithm 2), are laid out logically. Figure 1 clearly illustrates the overall scheme, and the experimental figures (e.g., Figure 2) effectively show the evolution of the PMF during training.", "weaknesses": "The effectiveness of WT-ASBS hinges on the accurate selection of several critical hyperparameters inherited from the enhanced sampling domain, such as the initial Gaussian height $h$, the Gaussian width $\\sigma$, and the bias factor $\\gamma$. \n\nIt seems that the paper lacks a systematic ablation study to investigate the sensitivity of the method to these parameters (e.g., how the convergence speed or final accuracy changes with different $\\gamma$ values). Given that WT-ASBS generates uncorrelated samples, the optimal choice for parameters like $h$ might deviate significantly from WTMetaD conventions. It requires clearer guidelines for practitioners.", "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": "I4G1KttiEB", "reviewer": "Reviewer_2tNT", "summary": "This paper proposes WT-ASBS (Well-Tempered Adjoint Schrödinger Bridge Sampler), a method that combines a neural diffusion-based sampler (ASBS) with a well-tempered bias potential defined in collective-variable (CV) space. The goal is to accelerate exploration in molecular systems by gradually flattening the free-energy landscape while retaining thermodynamic consistency through reweighting.\n\nThe authors claim that this hybrid design achieves faster and broader sampling compared to standard ASBS and conventional enhanced sampling methods such as WTMetaD. Empirical results on alanine dipeptide, alanine tetrapeptide, and chemical reaction benchmarks suggest improved exploration and free-energy reconstruction.", "questions": "(Q1) If $V_{WT}(\\xi)$ is not parameterized by a neural network, how does the model ensure that the learned sampler contributes beyond a standard MetaD bias? Could the authors clarify whether WT-ASBS should be interpreted as (a) a neural sampler guided by an analytic bias, or (b) a learnable bias parameterized through the neural transport model?\n\n(Q2) In Figure 8, the reweighted and bias-derived PMFs are nearly identical, suggesting that the sampling in CV space is almost uniform. If so, what ensures that the sampled structures remain physically valid? Could the authors provide representative molecular structures sampled from different regions of the CV space (e.g., high- vs low-free-energy regions) to illustrate what kind of configurations the model actually produces? This would clarify whether the method truly explores meaningful conformations or simply performs uniform random exploration over CVs.\n\n(Q3) Could you extend Figure 3.c to the same energy-evaluation range ($\\approx$10–40 M) as Figure 3.d to show late-stage exploration and better quantify the WT bias’s long-term effects? Please include an ASBS PMF accuracy curve for Ala2 system.\n\n(Q4) Could you summarize cost–benefit metrics (energy evaluations, wall-clock time, ESS, PMF error) across ASBS, WT-ASBS, and WTMetaD for standardized comparison?\n\n(Q5) ASBS’s slow exploration (Figure 3.c) might result from strong pre-training confinement near the initial mode (marked * in Figure 3.b).\nHave you explored reducing pretraining strength or fine-tuning its duration to encourage broader exploration?\n\n(Q6) When computing the free-energy difference, which states or basins were used? Were they predefined minima or clusters in CV space?\nClarifying this would help interpret the reported PMF results.\n\n(Q7) Are all reported curves from single runs, or averaged over multiple simulations? If single, could you comment on the run-to-run variability? Showing variance bands or repeated trials would help establish statistical reliability.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "**Conceptually innovative hybridization.**\nThe integration of Schrödinger bridge–based diffusion sampling with a well-tempered metadynamics bias is novel. It bridges probabilistic transport and enhanced sampling, offering a unified view of energy-driven exploration in molecular systems.\n\n**Clear motivation from sampling efficiency.**\nThe work addresses a concrete challenge in molecular generative modeling — slow mode discovery and poor coverage of rare events — and proposes a pragmatic biasing mechanism that is both simple and theoretically interpretable.\n\n**Transparent discussion of limitations.**\nThe paper openly discusses the non-amortized bias, potential generalization issues, and differences from traditional WTMetaD, which reflects scientific maturity and honesty.", "weaknesses": "**1. Conceptual tension between learned sampler and non-learnable bias.**\nAlthough the method is presented as a neural diffusion sampler, the core bias $V_{WT}$ is manually accumulated (non-neural). This design choice blurs whether the sampler truly learns the target distribution or merely follows a hand-crafted MetaD bias.\n\n**2. Lack of fair and direct comparison with ASBS.**\nSince WT-ASBS is an ASBS variant augmented with a well-tempered bias, the most meaningful benchmark is ASBS itself under identical conditions. However, ASBS only appears in Fig. 3c (explored states) and is not evaluated for PMF accuracy, making the incremental benefit of WT unclear.\n\n**3. # of force-calls inconsistency between figures.**\nFigure 3.c shows up to $\\approx$2 M evaluations, while Figure 3.d extends to $\\approx$40 M. Because WT-ASBS reaches chemical accuracy after $\\approx$4–5 M evaluations, truncating Fig. 3c prevents fair comparison at later stages.\n\n**4. Limited diffusion-based baselines.**\nComparisons are restricted to WTMetaD, which mainly highlights conceptual differences rather than validating the WT mechanism across diffusion methods. Adding at least one diffusion-based baseline (annealed, energy-guided, or score-guided) would strengthen the evaluation.\n\n**5. Scalability and consistency.**\nWT-ASBS excels on alanine dipeptide (Figure 2.f) but underperforms WTMetaD on tetrapeptide (Figure 3.d), suggesting potential sensitivity to CV dimensionality or generalizability to complex system.\n\n**6. Lack of statistical robustness.**\nThe reported results appear to be based on a single simulation run for WT-ASBS. Given the stochastic nature of both diffusion and metadynamics sampling, multiple independent runs (with mean and variance) are crucial to assess convergence stability and reproducibility. Without them, it is difficult to judge whether the reported trajectories and PMF curves reflect consistent behavior or a favorable random seed.", "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": "kerXrS3T4p", "reviewer": "Reviewer_mrdW", "summary": "This paper introduces WT-ASBS, a diffusion-based generative sampler enhanced by biasing along  CVs, inspired by WTMetaD.\n\nThe method replaces the local MD step in WTMetaD with a diffusion-based neural sampler (ASBS) that learns to generate molecular configurations directly from the potential energy function. A repulsive potential is iteratively constructed in CV space to encourage exploration of new regions and flatten free-energy barriers, with subsequent reweighting to recover the unbiased Boltzmann ensemble. Experiments demonstrate that WT-ASBS improves sampling efficiency and mode discovery for molecular systems.", "questions": "1. How should practitioners determine that the two-loop WT-ASBS training has converged?\n\n2. Can the authors clarify whether the proposed reweighting scheme guarantees unbiased Boltzmann expectations, or whether the residual approximation in the learned diffusion process introduces systematic bias? Could additional refinement (e.g., short MCMC runs or importance sampling) ensure asymptotic correctness similar to Boltzmann Generators?\n\n3. Have the authors considered coupling WT-ASBS with ML-based CV discovery?\n\n4. WT-ASBS achieves broader exploration but sometimes slightly worse free-energy accuracy (Fig. 3d). Can the authors comment on this?\n\n5. Can the authors explore the computational feasibility for systems with explicit solvent?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. The idea of integrating diffusion-based neural samplers with enhanced-sampling techniques from molecular simulation is elegant and promising. While both ASBS and WTMetaD exist independently, their combination is novel and creates a bridge between machine-learning-based and physics-based sampling paradigms.\n\n2. The methodology is technically well-founded, with a clear derivation and an explicit algorithmic description (Algorithm 1). The paper includes a convergence proposition showing that the bias potential approaches the well-tempered limit, ensuring theoretical consistency.\n\n3. The paper is clearly written, logically organized, and visually well-supported by figures. The workflow in Fig. 1 and the schematic diagrams for peptide systems help the reader follow the training and bias-updating procedure.\n\n4. From an application perspective, the work demonstrates that diffusion-based samplers can efficiently handle physically meaningful molecular systems. The results on peptides and reactive landscapes show large efficiency gains over WTMetaD. This points toward real practical potential in molecular modeling.", "weaknesses": "1. The method mainly replaces the MD propagation in WTMetaD with ASBS without introducing new learning objectives or architectures. The innovation lies more in application and system integration than in core ML algorithmic development. For ICLR, where the focus is typically on advancing machine-learning methodology, the contribution might appear not so significant. The work may find a more natural home in computational chemistry or physics-oriented venues.\n\n2. Algorithm 1 involves two nested loops (outer bias update over k, inner ASBS optimization over l). In addition, an iid sampling procedure is involved. However, the paper does not specify when or how to stop the iterations in practice. ASBS itself is self-consistent and lacks a single scalar loss whose decrease guarantees convergence. The absence of heuristic or empirical convergence criteria (e.g., stabilization of bias potential or free-energy estimates) makes reproducibility difficult.\n\n3. All experiments use well-known CVs such as torsional angles or bond distances. For realistic, high-dimensional systems, CV discovery is non-trivial. It would strengthen the paper to evaluate the method with ML-discovered CVs (e.g., TICA, SPIB, or FMRC). Without this, the method’s applicability to problems where good CVs are not known remains uncertain.\n\n4. Boltzmann Generators and related approaches can guarantee asymptotically correct estimates via MCMC or importance-sampling refinement. WT-ASBS claims reweighting via the bias but does not discuss whether such reweighting ensures unbiased expectations in the presence of neural-sampler approximation error.\n\n5. In Fig. 3d, the free-energy mean-absolute-error of WT-ASBS is slightly larger than that of WTMetaD. It would be valuable to analyze the reason to understand when diffusion-based sampling may underperform.\n\n6. The method’s feasibility for explicit-solvent systems remains unclear.", "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)