Enhancing Diffusion-Based Sampling with Molecular Collective Variables

## Abstract

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.

## 1 Introduction

A central goal of statistical mechanics simulations is to estimate macroscopic thermodynamic properties, such as the energy released in a 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 $\nu(x)\propto\exp(-u(x))$, where the reduced energy $u$ encodes the chemical system, temperature, and other thermodynamic variables.
Sampling from this distribution is deceptively difficult, despite having access to its explicit unnormalized form, due to the large number of possible atomic configurations and the high energy barriers separating them.
Straightforward MD or MCMC often requires an 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-likelihood to draw independent and identically distributed (i.i.d.) samples.
Early work focused on 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 et al., 2023a; Havens et al., 2025).(^1^11More information about neural and diffusion-based samplers can be found in [Section A.2](https://arxiv.org/html/2510.11923v1#A1.SS2).)
Although training can amortize effort across multiple draws, these methods do not inherently reduce the number of energy evaluations required to obtain accurate equilibrium statistics compared to MD-based approaches (He et al., 2025).
Directly applying diffusion samplers faces a well-known pitfall of mode collapse, where training and sampling concentrate on high-probability basins while underrepresenting rare but thermodynamically important states.
This issue is not only about exploration but also about statistics, since reliable free energies and ensemble averages depend critically on assigning correct weights to rare configurations, which may be exponentially small.
Thus, successful sampling must both reach low-occupancy modes and attribute the correct equilibrium weights to them.

Figure: Figure 1: Scheme for neural enhanced sampling. (Left) Local sampling and pretraining near the reference configuration. (Middle) Energy-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.
Refer to caption: x1.png

Inspired by enhanced sampling methods, we augment the diffusion-based sampler with a bias in collective variable (CV) space.
CVs are low-dimensional functions of atomic coordinates that capture slow, chemically relevant motions and can be chosen from common system characteristics (e.g., bond lengths or torsional angles) or identified with ML (Sidky et al., 2020; Bonati et al., 2023).
During the training of a diffusion-based sampler, the sampler iteratively proposes configurations and receives 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, discouraging 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, including rare 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 are as follows:

- •
We propose the *Well-Tempered Adjoint Schrödinger Bridge Sampler* (WT-ASBS), which augments the state-of-the-art diffusion-based sampler ASBS (Liu et al., 2025) with a biasing mechanism inspired by well-tempered metadynamics (WTMetaD, Barducci et 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-based sampling of reactive landscapes, achieving shorter wall-clock times than WTMetaD. This establishes a new practical role for diffusion-based samplers in chemistry.

## 2 Background

### 2.1 Molecular Sampling

#### Problem setup

In the configurational space $\mathcal{X}\subseteq\mathbb{R}^{n\times 3}$ for $n$ atoms with elemental identities $a\in\mathcal{E}^{n}$ with a set of element types $\mathcal{E}$, the potential energy function $E:\mathcal{X}\times\mathcal{E}^{n}\to\mathbb{R}$ defines the Boltzmann distribution at inverse temperature $\beta=1/k_{\text{B}}T$ as $\nu(x;a)\propto\exp(-\beta E(x;a))$.
In molecular simulations, the central challenge is to generate configurations representative of the underlying free energy landscape that describes the relative stability and likelihood of different molecular configurations.
Three main downstream applications motivate this task: (i) obtaining a diverse and statistically meaningful set of configurations for structural analysis, (ii) 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 $\nu(x)$.

One rarely aims to sample the entirety of $\mathcal{X}$ in practice because it includes *all* configurations consistent with the chosen composition and potential energy function.
Instead, sampling is naturally restricted to an accessible component $\mathcal{A}\subset\mathcal{X}$ that is kinetically connected to a reference conformer $x_{\text{ref}}\in\mathcal{X}$ on the timescales of interest at temperature $T$ by physics.
In molecular conformational sampling, $\mathcal{A}$ implies configurations that preserve the bond topology of $x_{\text{ref}}$. If one included bond reorganizations occurring on longer timescales (hours to days), $\mathcal{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 $\nu(x)$, the exponential dependence 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 configurations are being studied.
Consequently, many effective schemes for molecular sampling generate weighted samples $\{(X_{i},W_{i})\}_{i=1}^{N}$ with $X_{i}\in\mathcal{A}$ and $W_{i}\geq 0$, obtained for example via biased simulations with reweighting.
For any bounded observable $f:\mathcal{X}\to\mathbb{R}$, the self-normalized estimator

$$ $\hat{\nu}_{n}(f)=\frac{\sum_{i=1}^{N}W_{i}f(X_{i})}{\sum_{i=1}^{N}W_{i}}\;\xrightarrow{N\to\infty}\;\mathbb{E}_{\nu}[f(X)\,|\,X\in\mathcal{A}]$ (1) $$

recovers expectations under the Boltzmann ensemble restricted to $\mathcal{A}$.
We now summarize the requirements and discuss the practical realization of these:

#### Requirement ˜ 1 : Restrained sampling

Membership in $\mathcal{A}$ is governed by the free energy barriers that separate modes, not by the stability (state free energies) of individual modes.
When sampling is performed by emulating physical dynamics (e.g., MD), high barriers suppress transitions and trajectories remain within $\mathcal{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 $\mathcal{A}$ while strongly penalizing $\mathcal{X}\setminus\mathcal{A}$.
This confines exploration to the intended modes while allowing efficient coverage within them.
Practical considerations are further discussed in [Section˜3.2](https://arxiv.org/html/2510.11923v1#S3.SS2).

#### Requirements ˜ 2 and 3 : Biasing on the CV space

High free energy barriers often align with a small set of slow coordinates, so unbiased exploration in $\mathcal{A}$ can waste effort on fast motions while undersampling rare but important transitions.
CVs address this by mapping configurations to a low-dimensional space $\xi:\mathcal{X}\to\mathcal{S}\subset\mathbb{R}^{m}$ with $m\ll n$, that capture functionally relevant degrees of freedom.
Examples include peptide backbone torsions ($\phi,\psi$) and the center-of-mass distance between a protein pocket and its ligand.
By steering sampling with a bias that depends only on $s=\xi(x)$, we enhance explorations across modes while leaving orthogonal fluctuations unchanged.
A bias $V(s)$ yields the unnormalized sampling density $\nu_{V}(x)$ and importance weights $w(x)$ as

$$ $\nu_{V}(x)\propto\exp[-\beta E(x)-\beta V(\xi(x))],$ (2) $$

$$ $w(x)\propto\frac{\nu(x)}{\nu_{V}(x)}\propto\exp[+\beta V(\xi(x))],$ (3) $$

so unbiased expectations are recovered by the estimator in [Equation˜1](https://arxiv.org/html/2510.11923v1#S2.E1) with $X_{i}\sim\nu_{V}$ and $W_{i}=w(X_{i})$, ensuring [Requirement˜2](https://arxiv.org/html/2510.11923v1#S2.I1.i2).

Biasing-based enhanced sampling approaches often construct a bias that raises the effective sampling temperature along CV space, introduced by Barducci et al. (2008) as the well-tempered target distribution.
We first define the CV marginal $\bar{\nu}(s)$ and the potential of mean force (PMF) $F(s)$ from the unbiased CV marginal (defined up to a constant)(^2^22The term PMF for $F(s)$ is used interchangeably with the free energy profile/landscape along the CV.) using the Dirac delta function $\delta$:

$$ $\bar{\nu}_{V}(s)\propto\int_{\mathcal{X}}\nu_{V}(x)\delta[\xi(x)-s]\,\mathrm{d}x,$ (4) $$

$$ $F(s):=-\frac{1}{\beta}\log\bar{\nu}(s).$ (5) $$

Now, given a bias factor $\gamma>1$ that acts as a scale factor for $T$, we define the well-tempered bias $V_{\text{WT}}(s)$ and the corresponding target distribution $\nu_{\text{WT}}(x)$ from [Equation˜2](https://arxiv.org/html/2510.11923v1#S2.E2) as

$$ $V_{\text{WT}}(s)=-(1-\frac{1}{\gamma})F(s),$ (6) $$

$$ $\nu_{\text{WT}}(x)\propto\exp[-\beta E(x)+\beta(1-\frac{1}{\gamma})F(\xi(x))],$ (7) $$

From [Equations˜4](https://arxiv.org/html/2510.11923v1#S2.E4), [5](https://arxiv.org/html/2510.11923v1#S2.E5) and [7](https://arxiv.org/html/2510.11923v1#S2.E7), the well-tempered CV marginal is $\bar{\nu}_{\text{WT}}(s)\propto\exp[-\frac{\beta}{\gamma}F(s)]=[\bar{\nu}(s)]^{1/\gamma}$, which shows that the sampled CV behaves as if at a higher effective temperature $T_{\text{eff}}=\gamma T$ ($\beta_{\text{eff}}=\beta/\gamma$), while the orthogonal conditional remains unchanged: $\nu_{\text{WT}}(x\,|\,\xi(x)=s)=\nu(x\,|\,\xi(x)=s)$.
This selective increase in sampling temperature along the CV satisfies [Requirement˜3](https://arxiv.org/html/2510.11923v1#S2.I1.i3).
To construct $V_{\text{WT}}$ on-the-fly during simulation, one can stack Gaussian kernels (Barducci et al., 2008), use a kernel density estimator (Invernizzi & Parrinello, 2020), or obtain the bias variationally (Valsson & Parrinello, 2014; Bonati et al., 2019).

#### Mode exploration

[Requirement˜2](https://arxiv.org/html/2510.11923v1#S2.I1.i2) 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-Based Neural Samplers

#### Neural samplers

Diffusion-based neural samplers learn a control vector field that steers a stochastic process to gradually transform noise into a target distribution, typically specified by an unnormalized density or energy function.
Unlike iterative generative methods trained on data from implicitly defined distributions (Sohl-Dickstein 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](https://arxiv.org/html/2510.11923v1#A1.SS2).

#### Adjoint Schrödinger Bridge Sampler (ASBS, Liu et al., 2025 )

ASBS casts Boltzmann sampling as the Schrödinger Bridge (SB) problem between source $\mu(X_{0})$ and Boltzmann target $\nu(X_{1})\propto\exp(-\beta E(X_{1}))$ that seeks the optimal control $u_{\theta}$ w.r.t. the path Kullback-Leibler (KL) objective

$$ $\displaystyle\min_{u}D_{\text{KL}}(p^{u}\|p^{\text{base}}=p^{u=0})=\mathbb{E}_{X\sim p^{u}}\left[\int_{0}^{1}\frac{1}{2}\|u_{\theta}(X_{t},t)\|^{2}\,\mathrm{d}t\right],$ (8) $\displaystyle\text{s.t.}\;\mathrm{d}X_{t}=\sigma_{t}u_{\theta}(X_{t},t)\,\mathrm{d}t+\sigma_{t}\,\mathrm{d}W_{t},\;X_{0}\sim\mu(X_{0}),\;X_{1}\sim\nu(X_{1}),$ (9) $$

where $\sigma_{t}$ is the noise schedule, and $p^{u}$ denotes the law of the SDE in [Equation˜9](https://arxiv.org/html/2510.11923v1#S2.E9).(^3^33While the original ASBS formulation is based on SDEs with general drift $f(X_{t},t)+\sigma_{t}u_{\theta}(X_{t},t)$, here we only consider the driftless base process with $f=0$.)
ASBS avoids backpropagating through the SDE during optimization by introducing the parametrized corrector $h_{\phi}(X_{1})$ that debiases the effect of endpoint coupling from base distribution to allow the SOC interpretation of [Equation˜8](https://arxiv.org/html/2510.11923v1#S2.E8), which can be converted to Adjoint Matching (AM) objective for the drift (Havens et al., 2025).
The corrector is then learned by optimizing the Corrector Matching (CM) objective:

$$ $\displaystyle\text{(AM)}\quad u$ $\displaystyle=\operatorname*{arg\,min}_{u}\mathbb{E}_{p_{t|0,1}^{\text{base}}p_{0,1}^{\bar{u}}}[\|u_{\theta}(X_{t},t)+\sigma_{t}(\beta\nabla E(X_{1})+h_{\phi}(X_{1}))\|^{2}],$ (10) $\displaystyle\text{(CM)}\quad h$ $\displaystyle=\operatorname*{arg\,min}_{h}\mathbb{E}_{p_{0,1}^{\bar{u}}}[\|h_{\phi}(X_{1})-\nabla_{X_{1}}\log p^{\text{base}}(X_{1}|X_{0})\|^{2}],$ (11) $$

where $\bar{u}:=\text{{stopgrad}}(u)$.
These paired objectives instantiate iterative proportional fitting (IPF, Kullback, 1968) and converge to the SB solution.
The expectations use closed-form base conditionals, so no target samples or importance weights are needed, which preserves scalability.
Furthermore, the control model supports approximate data-based pretraining (see [Appendix˜C](https://arxiv.org/html/2510.11923v1#A3)).

ASBS achieved strong performance on 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-seeking behavior, motivating the introduction of an enhanced sampling mechanism in this work.

## 3 Methods

Now we augment ASBS with a well-tempered bias on a small set of CVs, updated online from i.i.d. model samples during training, instantiating a neural enhanced sampler that satisfies the requirements in [Section˜2.1](https://arxiv.org/html/2510.11923v1#S2.SS1).
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 [Figure˜1](https://arxiv.org/html/2510.11923v1#S1.F1)).

### 3.1 Well-Tempered Adjoint Schrödinger Bridge Sampler

Figure: Algorithm 1 Well-Tempered Adjoint Schrödinger Bridge Sampler (WT-ASBS)

#### Setup and two-time-scale algorithm

We implement WT-ASBS with a well-tempered bias update through stacking Gaussian kernels (Barducci et al., 2008) to reach the well-tempered target [Equation˜6](https://arxiv.org/html/2510.11923v1#S2.E6).
Let $\mathcal{S}\subset\mathbb{R}^{m}$ denote a compact CV region of interest, where bias potentials $V:\mathcal{S}\to\mathbb{R}$ are defined and constructed, which induces the sampling distribution $\nu_{V}(x)\propto\exp[{-\beta(E(x)+V(\xi(x)))}]$.
WT-ASBS presented in [Algorithm˜1](https://arxiv.org/html/2510.11923v1#alg1) is executed via a two-time-scale scheme:

- 1.
Inner step (ASBS training).
For fixed $V_{k}$, train ASBS with energy $E+V_{k}\circ\xi$ until convergence so that the marginal distribution at $t=1$ equals $\nu_{V_{k}}$.
- 2.
Outer step (well-tempered bias update).
Draw a conditionally i.i.d. batch of configurations from ASBS $\{X_{1,k}^{(i)}\}_{i=1}^{N}\sim p_{1}^{u}$, project to CV space $s_{k}^{(i)}=\xi(X_{1,k}^{(i)})$, and update
$V_{k+1}(s)=V_{k}(s)+h\sum_{i=1}^{N}\exp{(-\frac{\beta}{\gamma-1}V_{k}(s_{k}^{(i)}))}K_{\sigma}(s,s_{k}^{(i)}),$
(12)
with fixed height $h>0$ and Gaussian kernel $K_{\sigma}(s,s^{\prime})=\exp(-\frac{\|s-s^{\prime}\|^{2}}{2\sigma^{2}})$.

We further establish a convergence guarantee showing that WT-ASBS provably approaches the desired well-tempered distribution under the two-time-scale update scheme (proof in [Section˜B.1](https://arxiv.org/html/2510.11923v1#A2.SS1)):

###### Proposition 3.1 (Convergence of WT-ASBS) .

When [Algorithm˜1](https://arxiv.org/html/2510.11923v1#alg1) is performed until convergence, the bias potential $V_{k}$ converges almost surely to $V^{\ast}(s)=-(1-\frac{1}{\gamma})F(s)+\mathrm{const}$, and hence the sampled distribution converges to the well-tempered target [Equation˜7](https://arxiv.org/html/2510.11923v1#S2.E7).

###### Remark 3.1 .

[Proposition˜3.1](https://arxiv.org/html/2510.11923v1#S3.Thmproposition1) ensures that, upon completion of training, the PMF along the CVs can be recovered up to a constant from the final bias as $F(s)=-\frac{\gamma}{\gamma-1}V^{\ast}(s)+\mathrm{const}$.

###### Proposition 3.1 (Convergence of WT-ASBS) .

###### Remark 3.1 .

#### Practical implementation

In practice, AM uses a replay buffer: gradient signals are not computed directly from samples of $p_{0,1}^{\bar{u}}$ with the current $u_{\theta}$, 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 [Appendix˜C](https://arxiv.org/html/2510.11923v1#A3) for details.

Although our biasing mechanism mirrors classical MD-based enhanced sampling, neural samplers like ASBS generate i.i.d. samples from the current distribution, eliminating the need to wait for decorrelation before bias deposition occurs.
This allows us to 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 or by modifying the learning objective in [Section˜A.2](https://arxiv.org/html/2510.11923v1#A1.SS2).

### 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, we consider a single reference configuration that defines the accessible component $\mathcal{A}$, corresponding to the initial conformer used to start MD or MCMC chains.
Although [Algorithm˜1](https://arxiv.org/html/2510.11923v1#alg1) does not require data samples for training, as in the original ASBS, the model can be pretrained with available samples to warm-start and improve training efficiency.
Accordingly, we perform short MD simulations (local sampling) to generate samples from the mode containing the reference configuration and apply bridge matching pretraining ([Appendix˜C](https://arxiv.org/html/2510.11923v1#A3)) to 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 a pretraining stage, high accuracy is unnecessary, and we can further reduce computational cost by using approximate energy evaluations (e.g., classical force fields), faster sampling settings, or 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 low-dimensional CVs that defines the space over which modes are explored. CVs may be selected either from chemical intuition or from topological changes induced by the underlying transformation, as demonstrated in [Section˜4](https://arxiv.org/html/2510.11923v1#S4).
For more complex systems, CVs can also be parameterized by shallow neural networks 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 $\gamma$ 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](https://arxiv.org/html/2510.11923v1#S2.I1.i1), 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 is well understood, these serve as a natural guide (Canfield et al., 2018).
We provide a detailed treatment in [Section˜A.3](https://arxiv.org/html/2510.11923v1#A1.SS3), and here illustrate the idea with chirality inversion in proteins and peptides.
The $C_{\alpha}$ atom of each amino acid is a 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-symmetric 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](https://arxiv.org/html/2510.11923v1#S3.Thmremark1) 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, to recover the full weighted conformational distribution, one can integrate [Equation˜9](https://arxiv.org/html/2510.11923v1#S2.E9) to generate samples $X_{i}$ and assign weights $W_{i}=\exp[\beta V^{\ast}(\xi(X_{i}))]$, yielding the self-normalized estimator in [Equation˜1](https://arxiv.org/html/2510.11923v1#S2.E1).
If desired, these samples can be refined prior to weighting via short MD or MCMC using $E+V^{\ast}\circ\xi$ as the energy.

## 4 Experiments

We demonstrate the efficiency of the WT-ASBS pipeline on four molecular sampling tasks involving peptides and chemical reactions.
Since [Requirements˜3](https://arxiv.org/html/2510.11923v1#S2.I1.i3) and [1](https://arxiv.org/html/2510.11923v1#S2.I1.i1) are incorporated into the design of the method, correctness with respect to [Requirement˜2](https://arxiv.org/html/2510.11923v1#S2.I1.i2) ensures that all requirements are satisfied.
We therefore assess whether the empirical PMF, $\hat{F}(s)=-\frac{1}{\beta}\log\hat{\nu}(s)$ obtained from reweighted samples [Equation˜1](https://arxiv.org/html/2510.11923v1#S2.E1) and the PMF from bias ([Remark˜3.1](https://arxiv.org/html/2510.11923v1#S3.Thmremark1)) 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_{i}=-\frac{1}{\beta}\log\int_{\mathcal{S}}\exp(-\beta F(s))\bm{1}[s\in\mathcal{S}_{i}]\,\mathrm{d}s\,+\,\mathrm{const},$ (13) $$

where $\mathcal{S}_{i}$ is the CV region of state $i$.
We report free energy differences $\Delta F_{ij}=F_{j}-F_{i}$, thereby removing the additive constant as it is the same for all $\mathcal{S}_{i}$. We compute multi-state errors with a common offset.
Although weighted sampling required by [Requirement˜3](https://arxiv.org/html/2510.11923v1#S2.I1.i3) is best assessed via PMFs and free energies, we additionally compare with diffusion samplers on alanine dipeptide in [Appendix˜E](https://arxiv.org/html/2510.11923v1#A5) 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–Nme) and alanine tetrapeptide (Ala4; Ace–Ala3–Nme), which serve as standard testbeds for evaluating molecular sampling methods.
Both peptides are modeled with a classical force field with implicit water solvation, with the chirality restraint applied to all Cα atoms (see [Appendix˜D](https://arxiv.org/html/2510.11923v1#A4)).
In this work, we focus on sampling the full Cartesian coordinates of the molecules.
Although peptides and proteins adopt a broad range of conformations, these arise largely from the combinatorial organization of local backbone torsional states, specifically the $\phi$ and $\psi$ angles ([Figure˜2](https://arxiv.org/html/2510.11923v1#S4.F2)a), as 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.

Figure: Figure 2: Sampling alanine dipeptide. (a) Molecular structure and torsional CVs ($\phi$ and $\psi$). (b) Pretraining data from a short MD simulation. (c) PMF evolution reconstructed from bias during training. (d) PMF from reweighting. (e) Reference PMF. (f) Convergence of $\Delta F$ over training.
Refer to caption: x2.png

#### Alanine dipeptide (Ala2)

We employ $\phi$ and $\psi$ torsion angles as CVs ([Figure˜2](https://arxiv.org/html/2510.11923v1#S4.F2)a).
The conformational landscape is broadly divided into two states along $\phi$ ([Figure˜2](https://arxiv.org/html/2510.11923v1#S4.F2)b).
Local sampling within the more populated state is efficient due to the absence of a significant barrier; thus, we easily sample pretraining data from a 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 conformations in the less populated state.
Over training progress, the deposition asymptotically converges to the bias potential required to reproduce the PMF ([Figure˜2](https://arxiv.org/html/2510.11923v1#S4.F2)c).
[Figure˜2](https://arxiv.org/html/2510.11923v1#S4.F2)d 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-sample PMF agree closely with the reference PMF ([Figure˜2](https://arxiv.org/html/2510.11923v1#S4.F2)e), whereas the baseline ASBS fails to reproduce the correct distribution in the less populated region ([Figure˜7](https://arxiv.org/html/2510.11923v1#A5.F7)).
Moreover, compared to the baseline enhanced sampling method (WTMetaD) with the same bias factor, WT-ASBS achieves convergence of the $\Delta F$ between the two states more accurately.

Figure: Figure 3: Sampling alanine tetrapeptide. (a) Molecular structure and torsional CVs ($\phi_{1}$, $\phi_{2}$, $\phi_{3}$). (b) Reference data distribution in torsional space. Pretraining is confined to the mode indicated by $\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\ast$. (c) State exploration during training or simulation. (d) Convergence of state free energies measured by MAE up to a constant.
Refer to caption: x3.png

#### Alanine tetrapeptide (Ala4)

As seen in the Ala2 case, the main energetic barriers separating conformational states arise from backbone $\phi$ torsions.
We therefore employ three CVs, $\phi_{1}$, $\phi_{2}$, and $\phi_{3}$.
A long 100-µs reference MD simulation shows that the system occupies eight distinct modes, corresponding to the $2^{3}$ combinations of the three $\phi$ torsional states ([Figure˜3](https://arxiv.org/html/2510.11923v1#S4.F3)b).
We initialize from a single conformation in the indicated mode and generate pretraining data from a short MD simulation confined in the mode ([Figure˜8](https://arxiv.org/html/2510.11923v1#A5.F8)).
As shown in [Figure˜3](https://arxiv.org/html/2510.11923v1#S4.F3)c, WT-ASBS explores conformational states far more efficiently than the baseline WTMetaD, discovering all eight modes within the early stages of training.
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 ([Figure˜3](https://arxiv.org/html/2510.11923v1#S4.F3)d), as measured by the mean absolute error (MAE) of free energies up to a constant $\min_{c}\frac{1}{8}\sum_{i}|F_{i}-F_{i}^{\text{ref}}+c|$, within chemical accuracy (1 kcal/mol), demonstrating the practical applicability of WT-ASBS.

### 4.2 Sampling Reactive Energy Landscapes

While peptide results highlight the efficiency of WT-ASBS in terms of the number of energy evaluations, the computational budgets in diffusion-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 accurate energetics, is 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 settings, we apply WT-ASBS to sample reactive energy landscapes involving bond formation (Trizio & Parrinello, 2021).
We employ a recent uMLIP, UMA-S-1.1 (Wood et al., 2025), which accurately reproduces DFT energies, as the underlying energy model.
For reaction sampling, we cross-check with WTMetaD results since sampling the transition region with unbiased samplers like MD is practically impossible (e.g., $\sim 10^{15}$ samples required for a 20 kcal/mol barrier at 300 K).

Figure: Figure 4: Sampling reactive landscapes. (a) SN2 reaction scheme and CVs. (b) PMF $F(d_{1},d_{2})$ for the SN2 reaction. (c) Cycloaddition reaction scheme with post-TS bifurcation. (d) Three bond-formation distances. (e) Contact CVs $c_{i}$ derived from distances $d_{i}$. (f,g) 1D $F(s_{1})$ and 2D $F(s_{1},s_{2})$ PMFs for the post-TS bifurcation.
Refer to caption: x4.png

#### Nucleophilic substitution (S N 2)

We study a prototypical SN2 reaction ([Figure˜4](https://arxiv.org/html/2510.11923v1#S4.F4)a), where a chloride ion attacks the tetrahedral carbon and another chloride ion departs, leading to interconversion between states 1 $\leftrightarrow$ 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, a single configuration (state 1) suffices for pretraining the sampler.
Using WT-ASBS, we obtain the 2-D PMF along the two bond distances ([Figure˜4](https://arxiv.org/html/2510.11923v1#S4.F4)b).
The symmetric PMF agrees well with the WTMetaD result ([Figure˜9](https://arxiv.org/html/2510.11923v1#A5.F9)), and both the location and free energy barrier of the transition state (TS) coincide with the results from standard saddle-point optimization.

#### Post-transition-state bifurcation

For simple reactions like SN2 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 landscape around the reaction progress (Truhlar et 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 [Figure˜4](https://arxiv.org/html/2510.11923v1#S4.F4)c exemplifies this difficulty: reactant 3 passes through a 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 ([Figure˜4](https://arxiv.org/html/2510.11923v1#S4.F4)d) as normalized contacts in [Equation˜55](https://arxiv.org/html/2510.11923v1#A4.E55), capturing bond formation progress ([Figure˜4](https://arxiv.org/html/2510.11923v1#S4.F4)e).
From these we define two CVs: $s_{1}=c_{1}+c_{2}+c_{3}$, which tracks overall bond formation (a soft bond count), and $s_{2}=c_{2}-c_{3}$, which distinguishes between products 4 and 5.
Details on restraints and settings are provided in the [Appendix˜D](https://arxiv.org/html/2510.11923v1#A4).
[Figure˜4](https://arxiv.org/html/2510.11923v1#S4.F4)f and g show the PMFs $F(s_{1})$ and $F(s_{1},s_{2})$, and similarly to SN2 example, the PMFs agree well with the WTMetaD result (2-D comparison in [Figure˜10](https://arxiv.org/html/2510.11923v1#A5.F10)), providing accurate energetics for complex, multimodal reactive landscapes that dominate practical challenges in chemistry.
Further discussion on reaction sampling tasks is provided in [Appendix˜E](https://arxiv.org/html/2510.11923v1#A5).

Figure: Figure 5: Convergence of PMFs. PMF MAE from the final PMF for each reaction.
Refer to caption: x5.png

#### 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 [Equation˜57](https://arxiv.org/html/2510.11923v1#A5.E57), is shown in [Figure˜5](https://arxiv.org/html/2510.11923v1#S4.F5).
For the SN2 reaction, WT-ASBS required 0.77M energy evaluations and 4.3 hours, versus 4.0M and 29 hours for WTMetaD.
For the bifurcation reaction, WT-ASBS used 2.6M evaluations and 23 hours, compared 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.

## 5 Discussion

We presented WT-ASBS, which integrates bias deposition over chosen CVs into a diffusion-based sampler, enabling exploration of unseen conformational modes and recovery of PMFs and importance sampling weights that reconstruct the Boltzmann distribution.
The method unifies diffusion-based sampling with enhanced sampling and provides a practical and efficient recipe for sampling molecular systems, which allows diverse methodological extensions, including transition path sampling (Raja et al., 2025) and conformational exploration with system-agnostic CVs (e.g., RMSD).

#### Limitations

WT-ASBS requires the use of CVs, and while we utilized predefined ones, more complex systems may benefit from ML CVs and iterative cycles that combine exploratory WT-ASBS with CV training.
A fair comparison of enhanced sampling methods is difficult because of the large number of hyperparameters involved (Hénin et al., 2022), and although WTMetaD may not be the optimal choice, we use it as a baseline since our enhanced sampling component is based on WTMetaD.
Future work will explore combinations of different enhanced sampling strategies to provide a more comprehensive analysis.
Finally, our approach does not allow amortization in the current form, while it is possible to construct the target bias variationally using neural networks (Bonati et al., 2019), which could be applied to systems that share similar slow degrees of freedom.

## Acknowledgments

The authors would like to thank Zachary Ulissi and Larry Zitnick for the helpful discussions and comments.

## References

- Akhound-Sadegh et al. (2025)
Tara Akhound-Sadegh, Jungyoon Lee, Avishek Joey Bose, Valentin De Bortoli, Arnaud Doucet, Michael M. Bronstein, Dominique Beaini, Siamak Ravanbakhsh, Kirill Neklyudov, and Alexander Tong.
Progressive inference-time annealing of diffusion models for sampling from boltzmann densities, 2025.
URL [https://arxiv.org/abs/2506.16471](https://arxiv.org/abs/2506.16471).
- Albergo & Vanden-Eijnden (2024)
Michael S Albergo and Eric Vanden-Eijnden.
Nets: A non-equilibrium transport sampler.
*arXiv preprint arXiv:2410.02711*, 2024.
- Albergo et al. (2023)
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.
- Albergo et al. (2019)
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.
- Barducci et al. (2008)
Alessandro Barducci, Giovanni Bussi, and Michele Parrinello.
Well-tempered metadynamics: a smoothly converging and tunable free-energy method.
*Physical Review Letters*, 100(2):020603, 2008.
- Berner et al. (2022)
Julius Berner, Lorenz Richter, and Karen Ullrich.
An optimal control perspective on diffusion-based generative modeling.
*arXiv preprint arXiv:2211.01364*, 2022.
- Blessing et al. (2025)
Denis Blessing, Julius Berner, Lorenz Richter, Carles Domingo-Enrich, Yuanqi Du, Arash Vahdat, and Gerhard Neumann.
Trust region constrained measure transport in path space for stochastic optimal control and inference, 2025.
URL [https://arxiv.org/abs/2508.12511](https://arxiv.org/abs/2508.12511).
- Bonati et al. (2019)
Luigi Bonati, Yue-Yu Zhang, and Michele Parrinello.
Neural networks-based variationally enhanced sampling.
*Proceedings of the National Academy of Sciences*, 116(36):17641–17647, 2019.
- Bonati et al. (2023)
Luigi Bonati, Enrico Trizio, Andrea Rizzi, and Michele Parrinello.
A unified framework for machine learning collective variables for enhanced sampling simulations: mlcolvar.
*The Journal of Chemical Physics*, 159(1), 2023.
- Canfield et al. (2018)
Peter J Canfield, Iain M Blake, Zheng-Li Cai, Ian J Luck, Elmars Krausz, Rika Kobayashi, Jeffrey R Reimers, and Maxwell J Crossley.
A new fundamental type of conformational isomerism.
*Nature chemistry*, 10(6):615–624, 2018.
- Chen et al. (2018)
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud.
Neural ordinary differential equations.
*Advances in neural information processing systems*, 31, 2018.
- Chodera et al. (2025)
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 Simmonett, Mark J. Williamson, hb0402, Jake Fennick, Sander Roet, Ian Kenney, Irfan Alibay, Benjamin Ries, Ethan Holz, Hannah Baumann, Richard Gowers, and SimonBoothroyd.
choderalab/openmmtools: 0.25.1 bugfix release, August 2025.
URL [https://doi.org/10.5281/zenodo.16783210](https://doi.org/10.5281/zenodo.16783210).
- Choi et al. (2025)
Jaemoo Choi, Yongxin Chen, Molei Tao, and Guan-Horng Liu.
Non-equilibrium annealed adjoint sampler, 2025.
URL [https://arxiv.org/abs/2506.18165](https://arxiv.org/abs/2506.18165).
- Comer et al. (2015)
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 B*, 119(3):1129–1151, 2015.
- Dama et al. (2014)
James F. Dama, Michele Parrinello, and Gregory A. Voth.
Well-tempered metadynamics converges asymptotically.
*Phys. Rev. Lett.*, 112:240602, Jun 2014.
[10.1103/PhysRevLett.112.240602](https:/doi.org/10.1103/PhysRevLett.112.240602).
URL [https://link.aps.org/doi/10.1103/PhysRevLett.112.240602](https://link.aps.org/doi/10.1103/PhysRevLett.112.240602).
- Darve et al. (2008)
Eric Darve, David Rodríguez-Gómez, and Andrew Pohorille.
Adaptive biasing force method for scalar and vector free energy calculations.
*The Journal of chemical physics*, 128(14), 2008.
- Dinh et al. (2016)
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio.
Density estimation using real nvp.
*arXiv preprint arXiv:1605.08803*, 2016.
- Eastman et al. (2023)
Peter Eastman, Raimondas Galvelis, Raúl P Peláez, Charlles RA Abreu, Stephen E Farr, Emilio Gallicchio, 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 B*, 128(1):109–116, 2023.
- Frenkel & Smit (2023)
Daan Frenkel and Berend Smit.
*Understanding molecular simulation: from algorithms to applications*.
Elsevier, 2023.
- Froitzheim et al. (2025)
Thomas Froitzheim, Marcel Müller, Andreas Hansen, and Stefan Grimme.
g-xtb: A general-purpose extended tight-binding electronic structure method for the elements h to lr (z= 1–103).
*ChemRxiv*, 2025.
[10.26434/chemrxiv-2025-bjxvt](https:/doi.org/10.26434/chemrxiv-2025-bjxvt).
- Havens et al. (2025)
Aaron J Havens, Benjamin Kurt Miller, Bing Yan, Carles Domingo-Enrich, Anuroop Sriram, Daniel S. Levine, Brandon M Wood, Bin Hu, Brandon Amos, Brian Karrer, Xiang Fu, Guan-Horng Liu, and Ricky T. Q. Chen.
Adjoint sampling: Highly scalable diffusion samplers via adjoint matching.
In *Forty-second International Conference on Machine Learning*, 2025.
URL [https://openreview.net/forum?id=6Eg1OrHmg2](https://openreview.net/forum?id=6Eg1OrHmg2).
- He et al. (2025)
Jiajun He, Yuanqi Du, Francisco Vargas, Dinghuai Zhang, Shreyas Padhy, RuiKang OuYang, Carla Gomes, and José Miguel Hernández-Lobato.
No trick, no treat: Pursuits and challenges towards simulation-free training of neural samplers, 2025.
URL [https://arxiv.org/abs/2502.06685](https://arxiv.org/abs/2502.06685).
- Ho et al. (2020)
Jonathan Ho, Ajay Jain, and Pieter Abbeel.
Denoising diffusion probabilistic models.
*Advances in neural information processing systems*, 33:6840–6851, 2020.
- Hopkins et al. (2015)
Chad W Hopkins, Scott Le Grand, Ross C Walker, and Adrian E Roitberg.
Long-time-step molecular dynamics through hydrogen mass repartitioning.
*Journal of chemical theory and computation*, 11(4):1864–1874, 2015.
- Houk et al. (1970)
KN Houk, LJ Luskus, and NS Bhacca.
Novel double [6+4] cycloaddition of tropone to dimethylfulvene.
*Journal of the American Chemical Society*, 92(21):6392–6394, 1970.
- Hénin et al. (2022)
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.
[10.33011/livecoms.4.1.1583](https:/doi.org/10.33011/livecoms.4.1.1583).
URL [https://livecomsjournal.org/index.php/livecoms/article/view/v4i1e1583](https://livecomsjournal.org/index.php/livecoms/article/view/v4i1e1583).
- Invernizzi & Parrinello (2020)
Michele Invernizzi and Michele Parrinello.
Rethinking metadynamics: from bias potentials to probability distributions.
*The journal of physical chemistry letters*, 11(7):2731–2736, 2020.
- Jacobs et al. (2025)
Ryan Jacobs, Dane Morgan, Siamak Attarian, Jun Meng, Chen Shen, Zhenghao Wu, Clare Yijia Xie, Julia H Yang, Nongnuch Artrith, Ben Blaiszik, et al.
A practical guide to machine learning interatomic potentials–status and future.
*Current Opinion in Solid State and Materials Science*, 35:101214, 2025.
- Jiao et al. (2023)
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.
- Jing et al. (2022)
Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi Jaakkola.
Torsional diffusion for molecular conformer generation.
*Advances in neural information processing systems*, 35:24240–24253, 2022.
- Karras et al. (2022)
Tero Karras, Miika 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.
- Kim et al. (2025)
Minkyu Kim, Kiyoung Seong, Dongyeop Woo, Sungsoo Ahn, and Minsu Kim.
On scalable and efficient training of diffusion samplers, 2025.
URL [https://arxiv.org/abs/2505.19552](https://arxiv.org/abs/2505.19552).
- Kollman et al. (1997)
Peter Kollman, Richard Dixon, Wendy Cornell, Thomas Fox, Chris Chipot, and Andrew Pohorille.
The development/application of a ‘minimalist’organic/biochemical molecular mechanic force field using a combination of ab initio calculations and experimental data.
In *Computer simulation of biomolecular systems: Theoretical and experimental applications*, pp. 83–96. Springer, 1997.
- Kullback (1968)
Solomon Kullback.
Probability densities with given marginals.
*The Annals of Mathematical Statistics*, 39(4):1236–1243, 1968.
- Kushner & Yin (2003)
Harold J Kushner and G George Yin.
*Stochastic approximation and recursive algorithms and applications*.
Springer, 2003.
- Laio & Parrinello (2002)
Alessandro Laio and Michele Parrinello.
Escaping free-energy minima.
*Proceedings of the national academy of sciences*, 99(20):12562–12566, 2002.
- Larsen et al. (2017)
Ask Hjorth Larsen, Jens Jørgen Mortensen, Jakob Blomqvist, Ivano E Castelli, Rune Christensen, Marcin Dułak, 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 Kolsbjerg, 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 S Thygesen, Tejs Vegge, Lasse Vilhelmsen, Michael Walter, Zhenhua Zeng, and Karsten W Jacobsen.
The atomic simulation environment—a python library for working with atoms.
*J. Phys.: Condens. Matter*, 29(27):273002, 2017.
URL [https://doi.org/10.1088/1361-648X/aa680e](https://doi.org/10.1088/1361-648X/aa680e).
- Liao et al. (2024)
Yi-Lun Liao, Brandon Wood, Abhishek Das*, and Tess Smidt*.
EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations.
In *International Conference on Learning Representations (ICLR)*, 2024.
URL [https://openreview.net/forum?id=mCOBKZmrzD](https://openreview.net/forum?id=mCOBKZmrzD).
- Lindorff-Larsen et al. (2010)
Kresten Lindorff-Larsen, Stefano Piana, Kim Palmo, Paul Maragakis, John L Klepeis, Ron O Dror, and David E Shaw.
Improved side-chain torsion potentials for the amber ff99sb protein force field.
*Proteins: Structure, Function, and Bioinformatics*, 78(8):1950–1958, 2010.
- Lipman et al. (2022)
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le.
Flow matching for generative modeling.
*arXiv preprint arXiv:2210.02747*, 2022.
- Liu et al. (2025)
Guan-Horng Liu, Jaemoo Choi, Yongxin Chen, Benjamin Kurt Miller, and Ricky T. Q. Chen.
Adjoint schrödinger bridge sampler, 2025.
URL [https://arxiv.org/abs/2506.22565](https://arxiv.org/abs/2506.22565).
- Loshchilov & Hutter (2017)
Ilya Loshchilov and Frank Hutter.
Decoupled weight decay regularization.
*arXiv preprint arXiv:1711.05101*, 2017.
- Marx & Hutter (2000)
Dominik Marx and Jurg Hutter.
Ab initio molecular dynamics: Theory and implementation.
*Modern methods and algorithms of quantum chemistry*, 1(301-449):141, 2000.
- Máté & Fleuret (2023)
Bálint Máté and François Fleuret.
Learning interpolations between boltzmann densities.
*Transactions on Machine Learning Research*, 2023.
ISSN 2835-8856.
URL [https://openreview.net/forum?id=TH6YrEcbth](https://openreview.net/forum?id=TH6YrEcbth).
- McNaught & Wilkinson (1997)
Alan D McNaught and Andrew Wilkinson.
*Compendium of chemical terminology*, volume 1669.
Blackwell Science Oxford, 1997.
- Midgley et al. (2023)
Laurence Midgley, Vincent Stimper, Javier Antorán, Emile Mathieu, Bernhard Schölkopf, and José Miguel Hernández-Lobato.
Se (3) equivariant augmented coupling flows.
*Advances in Neural Information Processing Systems*, 36:79200–79225, 2023.
- Midgley et al. (2022)
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.01893*, 2022.
- Miller et al. (2024)
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 *Proceedings of Machine Learning Research*, pp. 35664–35686. PMLR, 21–27 Jul 2024.
URL [https://proceedings.mlr.press/v235/miller24a.html](https://proceedings.mlr.press/v235/miller24a.html).
- Nam & Jung (2023)
Juno Nam and YounJoon Jung.
Enhanced sampling for free energy profiles with post-transition-state bifurcations.
*Journal of Chemical Theory and Computation*, 19(10):2735–2743, 2023.
- Nicoli et al. (2020)
Kim A Nicoli, Shinichi Nakajima, Nils Strodthoff, Wojciech Samek, Klaus-Robert Müller, and Pan Kessel.
Asymptotically unbiased estimation of physical observables with neural samplers.
*Physical Review E*, 101(2):023304, 2020.
- Noé et al. (2019)
Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu.
Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning.
*Science*, 365(6457):eaaw1147, 2019.
- Onufriev et al. (2004)
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, Function, and Bioinformatics*, 55(2):383–394, 2004.
- Paddon-Row & Warrener (1974)
Michael N Paddon-Row and Ronald N Warrener.
Concerning the periselectivity of the cycloaddition reaction between 8, 8-dimethylisobenzofulvene and tropone.
*Tetrahedron Letters*, 15(43):3797–3800, 1974.
- Peters (2017)
Baron Peters.
*Reaction rate theory and rare events*.
Elsevier, 2017.
- Raja et al. (2025)
Sanjeev Raja, Martin Sipka, Michael Psenka, Tobias Kreiman, Michal Pavelka, and Aditi S. Krishnapriyan.
Action-minimization meets generative modeling: Efficient transition path sampling with the onsager-machlup functional.
In *Forty-second International Conference on Machine Learning*, 2025.
URL [https://openreview.net/forum?id=QwoGfQzuMa](https://openreview.net/forum?id=QwoGfQzuMa).
- Ramachandran et al. (1963)
G.N. Ramachandran, C. Ramakrishnan, and V. Sasisekharan.
Stereochemistry of polypeptide chain configurations.
*Journal of Molecular Biology*, 7(1):95–99, 1963.
ISSN 0022-2836.
[https://doi.org/10.1016/S0022-2836(63)80023-6](https:/doi.org/https://doi.org/10.1016/S0022-2836(63)80023-6).
URL [https://www.sciencedirect.com/science/article/pii/S0022283663800236](https://www.sciencedirect.com/science/article/pii/S0022283663800236).
- Richter & Berner (2023)
Lorenz Richter and Julius Berner.
Improved sampling via learned diffusions.
*arXiv preprint arXiv:2307.01198*, 2023.
- Rissanen et al. (2025)
Severi Rissanen, RuiKang OuYang, Jiajun He, Wenlin Chen, Markus Heinonen, Arno Solin, and José Miguel Hernández-Lobato.
Progressive tempering sampler with diffusion.
In *Forty-second International Conference on Machine Learning*, 2025.
URL [https://openreview.net/forum?id=uBMnbCBEtZ](https://openreview.net/forum?id=uBMnbCBEtZ).
- Robbins & Monro (1951)
Herbert Robbins and Sutton Monro.
A stochastic approximation method.
*The annals of mathematical statistics*, pp. 400–407, 1951.
- Schopmans & Friederich (2025)
Henrik Schopmans and Pascal Friederich.
Temperature-annealed boltzmann generators.
In *Forty-second International Conference on Machine Learning*, 2025.
URL [https://openreview.net/forum?id=RqtRSrCbNu](https://openreview.net/forum?id=RqtRSrCbNu).
- Schreiner et al. (2023)
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, M. Hardt, and S. Levine (eds.), *Advances in Neural Information Processing Systems*, volume 36, pp. 36449–36462. Curran Associates, Inc., 2023.
URL [https://openreview.net/forum?id=1kZx7JiuA2](https://openreview.net/forum?id=1kZx7JiuA2).
- Schütt et al. (2021)
Kristof Schütt, Oliver Unke, and Michael Gastegger.
Equivariant message passing for the prediction of tensorial properties and molecular spectra.
In Marina Meila and Tong Zhang (eds.), *Proceedings of the 38th International Conference on Machine Learning*, volume 139 of *Proceedings of Machine Learning Research*, pp. 9377–9388. PMLR, 18–24 Jul 2021.
URL [https://proceedings.mlr.press/v139/schutt21a.html](https://proceedings.mlr.press/v139/schutt21a.html).
- Shuaibi et al. (2025)
Muhammed Shuaibi, Abhishek Das, Anuroop Sriram, Misko, Luis Barroso-Luque, Ray Gao, Siddharth 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://github.com/facebookresearch/fairchem](https://github.com/facebookresearch/fairchem).
- Sidky et al. (2020)
Hythem Sidky, Wei Chen, and Andrew L Ferguson.
Machine learning for collective variable discovery and enhanced sampling in biomolecular simulation.
*Molecular Physics*, 118(5):e1737742, 2020.
- Sohl-Dickstein et al. (2015)
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli.
Deep unsupervised learning using nonequilibrium thermodynamics.
In *International conference on machine learning*, pp. 2256–2265. pmlr, 2015.
- Stoltz et al. (2010)
Gabriel Stoltz, Mathias Rousset, et al.
*Free energy computations: A mathematical perspective*.
World Scientific, 2010.
- Sugita & Okamoto (1999)
Yuji Sugita and Yuko Okamoto.
Replica-exchange molecular dynamics method for protein folding.
*Chemical physics letters*, 314(1-2):141–151, 1999.
- Tabak & Turner (2013)
Esteban G Tabak and Cristina V Turner.
A family of nonparametric density estimation algorithms.
*Communications on Pure and Applied Mathematics*, 66(2):145–164, 2013.
- The PLUMED consortium (2019)
The PLUMED consortium.
Promoting transparency and reproducibility in enhanced molecular simulations.
*Nature methods*, 16(8):670–673, 2019.
- Torrie & Valleau (1977)
Glenn M Torrie and John P Valleau.
Nonphysical sampling distributions in monte carlo free-energy estimation: Umbrella sampling.
*Journal of computational physics*, 23(2):187–199, 1977.
- Tribello et al. (2014)
Gareth A Tribello, Massimiliano Bonomi, Davide Branduardi, Carlo Camilloni, and Giovanni Bussi.
Plumed 2: New feathers for an old bird.
*Computer physics communications*, 185(2):604–613, 2014.
- Trizio & Parrinello (2021)
Enrico Trizio and Michele Parrinello.
From enhanced sampling to reaction profiles.
*The Journal of Physical Chemistry Letters*, 12(35):8621–8626, 2021.
- Truhlar et al. (1996)
Donald G Truhlar, Bruce C Garrett, and Stephen J Klippenstein.
Current status of transition-state theory.
*The Journal of physical chemistry*, 100(31):12771–12800, 1996.
- Vaitl et al. (2022)
Lorenz Vaitl, Kim A Nicoli, Shinichi Nakajima, and Pan Kessel.
Gradients should stay on path: better estimators of the reverse-and forward kl divergence for normalizing flows.
*Machine Learning: Science and Technology*, 3(4):045006, 2022.
- Valsson & Parrinello (2014)
Omar Valsson and Michele Parrinello.
Variational approach to enhanced sampling and free energy calculations.
*Physical review letters*, 113(9):090601, 2014.
- Vargas et al. (2023a)
Francisco Vargas, Will Grathwohl, and Arnaud Doucet.
Denoising diffusion samplers.
*arXiv preprint arXiv:2302.13834*, 2023a.
- Vargas et al. (2023b)
Francisco Vargas, Shreyas Padhy, Denis Blessing, and Nikolas Nüsken.
Transport meets variational inference: Controlled monte carlo diffusions.
*arXiv preprint arXiv:2307.01050*, 2023b.
- Watson et al. (2023)
Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian L Trippe, Jason Yim, Helen E Eisenach, Woody Ahern, Andrew J Borst, Robert J Ragotte, Lukas F Milles, et al.
De novo design of protein structure and function with rfdiffusion.
*Nature*, 620(7976):1089–1100, 2023.
- Wood et al. (2025)
Brandon M. Wood, Misko Dzamba, Xiang Fu, Meng Gao, Muhammed Shuaibi, Luis Barroso-Luque, Kareem Abdelmaqsoud, Vahe Gharakhanyan, John R. 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 [https://arxiv.org/abs/2506.23971](https://arxiv.org/abs/2506.23971).
- Yu et al. (2017)
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–8258, 2017.
[10.1021/jacs.7b02966](https:/doi.org/10.1021/jacs.7b02966).
URL [https://doi.org/10.1021/jacs.7b02966](https://doi.org/10.1021/jacs.7b02966).
- Zhang & Chen (2021)
Qinsheng Zhang and Yongxin Chen.
Path integral sampler: a stochastic control approach for sampling.
*arXiv preprint arXiv:2111.15141*, 2021.
- Zhang et al. (2019)
Zhijun Zhang, Xinzijian Liu, Kangyu Yan, Mark E Tuckerman, and Jian Liu.
Unified efficient thermostat scheme for the canonical ensemble with holonomic or isokinetic constraints via molecular dynamics.
*The Journal of Physical Chemistry A*, 123(28):6056–6079, 2019.

## Appendix

## Appendix 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 to address this challenge. The core idea is to bias the simulation dynamics to encourage exploration of less-populated regions of phase space, while applying a 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 is umbrella sampling. Introduced by Torrie & Valleau (1977), umbrella sampling runs simulations with an added bias potential $U(s)$ along a chosen reaction coordinate or collective variable $s=\xi(x)$. By choosing a series of overlapping bias “umbrellas” that cover the range of $\xi$, one can force the system to sample even high-energy configurations in each window. Each window’s samples are then reweighted by $\exp(+\beta U(s))$ to reconstruct the unbiased free energy landscape. This method demonstrated the principle of biasing along important coordinates to access otherwise inaccessible states. However, it requires a careful tuning of windows.

Metadynamics took adaptive biasing a step further by dynamically building the bias potential during a single simulation. In the original metadynamics (Laio & Parrinello, 2002), the simulation deposits 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 $V_{t}(s)$ approximates the negative of the underlying free energy $-F(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 a finite bias that yields a modified stationary distribution rather than diverging. Specifically, well-tempering introduces a bias factor $\gamma>1$ so that the deposited bias is scaled by a 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 $\gamma$, effectively simulating the system at a higher temperature $T_{\text{eff}}=\gamma 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 flattening the landscape, avoiding infinitely large weights and ensuring the bias converges asymptotically to a finite correction.

In the enhanced sampling community, there have been continual improvements to biasing algorithms. 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 [Equation˜7](https://arxiv.org/html/2510.11923v1#S2.E7). In a 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) (Valsson & 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 a different approach: instead of depositing hills, it continuously 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 metadynamics often achieve similar outcomes: a nearly uniform sampling of $s$, but via force-balancing vs. potential-filling mechanisms.
Another family of methods is replica exchange and parallel tempering. 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 configurations, 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 (e.g., the lowest-temperature one samples the target distribution) but benefits from exchanges to traverse phase space faster. Parallel tempering effectively raises the chance of barrier crossing by leveraging high-temperature dynamics, 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 potentially 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 statistical mechanics, respectively.
These initial approaches trained discrete-time normalizing flows (Tabak & Turner, 2013) by optimizing the KL divergence between the model and the unnormalized target densities. As the use of reverse-KL objective is prone to 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 computationally expensive, yet it is a 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 a 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 a 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) optimize the variance of Radon–Nikodym derivative between the forward and backward path measures.
Albergo & Vanden-Eijnden (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 struggling to capture multiple modes in complex energy landscapes.
Schopmans & Friederich (2025) applied temperature annealing in Boltzmann Generators, and concurrently Akhound-Sadegh et al. (2025); Rissanen et al. (2025) incorporated inference-time annealing into diffusion models to sample from Boltzmann densities.
These approaches smoothly interpolate from a high-temperature (smoothed-out) version of the target distribution to the true low-temperature distribution akin to parallel tempering.
Our method biases the sampler on the fly along specific CVs, which can be seen as a 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 et 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 a diffusion-based sampler with a 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 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 taxonomy (McNaught & Wilkinson, 1997), which separates constitutional isomerism from stereoisomerism 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 are subdivided into:
–
Configurational isomers, which are separated by large barriers and are isolable on experimental 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 stereocenter is inverted, R and S related molecules are 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 𝒜 \mathcal{A}

Let $\mathcal{X}$ denote the configurational space for a fixed chemical composition. We select a reference configuration and restrict attention to the portion of $\mathcal{X}$ relevant to the molecular process under study, which is the sense in which the main text employs the term accessible component.
One can regard configurations at temperature $T$ and study time horizon $\tau$ as grouped by an equivalence relation $\sim_{\tau}$: two configurations belong to the same class if they differ only by stereochemical elements that interconvert with a half-life shorter than $\tau$, while connectivity is preserved unless the chemical transformation of interest explicitly allows bond changes.
The accessible component $\mathcal{A}$ is then the connected set, under the physical dynamics, in the quotient space $\mathcal{X}/\sim_{\tau}$ that contains the chosen reference configuration.
Because the physical dynamics are stochastic, this construction relies on a separation of timescales, which is well supported by the exponential scaling of rates with activation energies and by established isomerism taxonomies.

#### Isomerism, geometry, and order-parameter restraint

The polytope formalism in Canfield et al. (2018) provides a coordination-centric view of stereochemical change and shows how fundamental isomerization processes can be enumerated for each coordination number.
In our context, this motivates representing each stereochemical element by the smallest stable geometric feature that distinguishes its alternatives, together with a numerical order parameter defined on those features.
The order parameter then becomes the target of a light, topology-preserving restraint.
Here, we exemplify 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 as axial, planar, and helical chirality, can be handled analogously.

(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 $r_{ij}(X)=\|x_{i}-x_{j}\|$ be the instantaneous bond length between atoms $i$ and $j$.
For each bonded pair or non-bonded pair to be preserved, we apply a bond restraint or a no-bond restraint, respectively:

$$ $\displaystyle U_{ij}^{\text{bond}}(X)$ $\displaystyle=\frac{1}{2}\kappa[r_{ij}(X)-r_{\text{max}}]_{+}^{2},$ (14) $\displaystyle U_{ij}^{\text{no-bond}}(X)$ $\displaystyle=\frac{1}{2}\kappa[r_{\text{min}}-r_{ij}(X)]_{+}^{2},$ (15) $$

where $[z]_{+}:=\max(z,0)$.
Here, we set $r_{\text{max}}=1.3(r_{i}^{\text{cov}}+r_{j}^{\text{cov}})$, $r_{\text{min}}=0.7(r_{i}^{\text{vdW}}+r_{j}^{\text{vdW}})$, and set the force constant $\kappa=100$ eV/Å^2, with $r_{i}^{\text{cov}}$ and $r_{i}^{\text{vdW}}$ denoting the covalent and van der Waals radii of atom $i$, respectively.
For examples in this work, we found that no-bond restraint [Equation˜15](https://arxiv.org/html/2510.11923v1#A1.E15) could be omitted in cases where valency of atoms is well-defined by the base potential energy (see implementation details in [Appendix˜D](https://arxiv.org/html/2510.11923v1#A4)).

(2) Tetrahedral center chirality restraint:
We define the torsion (dihedral) angle of atoms $i$, $j$, $k$, and $l$, $\phi_{ijkl}(X)\in(-\pi,\pi]$ as follows:

$$ $\phi_{ijkl}(X)=\mathrm{atan2}\left(b_{2}\cdot((b_{1}\times b_{2})\times(b_{2}\times b_{3})),\,|b_{2}|(b_{1}\times b_{2})\cdot(b_{2}\times b_{3})\right),$ (16) $$

where the displacement vectors are defined as $b_{1}:=x_{j}-x_{i}$, $b_{2}:=x_{k}-x_{j}$, and $b_{3}:=x_{l}-x_{k}$.
A proper torsion is the signed angle between two planes that share a central bond, with three bonds in sequence, e.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 a bond, but rather the out-of-plane displacement of one atom relative to a reference plane defined by three others.

We now consider a tetrahedral stereocenter $i$ with three neighboring atoms $j$, $k$, and $l$ (out of its four substituents).
The corresponding improper torsion angle, $\phi_{ijkl}(X)$, exhibits a bimodal distribution: each mode corresponds to a distinct stereoisomer, and in the ideal tetrahedral geometry, the peaks are centered near $\pm 35^{\circ}$.
To preserve the stereochemical configuration, we introduce the following flat-bottom harmonic restraint potential on the improper torsion:

$$ $U_{ijkl}^{\text{imp}}(X;\kappa,\phi_{0},\phi_{\text{tol}})=\frac{1}{2}\kappa[|w(\phi_{ijkl}(X)-\phi_{0})|-\phi_{\text{tol}}]_{+}^{2}$ (17) $$

where $w(\alpha)=\mathrm{atan2}(\sin\alpha,\,\cos\alpha)\in(-\pi,\pi]$ is the $2\pi$-periodic wrapping function, $\phi_{0}$ is the reference torsion value, and $\phi_{\text{tol}}$ 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 (Cα, N, C, Cβ) is positive, whereas the improper torsion (Cα, N, C, Hα) is negative.
Based on this, we define the amino acid chirality restraint as

$$ $U_{\text{AA}}^{\text{chiral}}(X)=U_{\text{C${}_{\alpha}$,N,C,C${}_{\beta}$}}^{\text{imp}}(X;\kappa,+\phi_{0},\phi_{\text{tol}})+U_{\text{C${}_{\alpha}$,N,C,H${}_{\alpha}$}}^{\text{imp}}(X;\kappa,-\phi_{0},\phi_{\text{tol}}),$ (18) $$

where we set $\kappa=25.0$ eV/rad^2, $\phi_{0}=0.615$ rad ($35^{\circ}$), and $\phi_{\text{tol}}=0.436$ rad ($25^{\circ}$).

## Appendix B Proofs and Notes

### B.1 Proof for Proposition ˜ 3.1

See [3.1](https://arxiv.org/html/2510.11923v1#S3.Thmproposition1)

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 1: Level and driving parts of the bias and the step size.
We split the bias $V_{k}(s)$ into a level $L_{k}:=\frac{1}{|\mathcal{S}|}\int_{\mathcal{S}}V_{k}(s)\,\mathrm{d}s$ and a driving part $D_{k}(s):=V_{k}(s)-L_{k}$.
Only $D_{k}$ affects the biased distribution, and adding the scalar $L_{k}$ shifts the energy by a constant.
Now, a single step of well-tempered bias update can be split into

$$ $\displaystyle L_{k+1}$ $\displaystyle=L_{k}+\exp\left(-\frac{\beta}{\gamma-1}L_{k}\right)\Lambda[D_{k}](s_{k+1}),$ (19) $\displaystyle D_{k+1}(s)$ $\displaystyle=D_{k}(s)+\exp\left(-\frac{\beta}{\gamma-1}L_{k}\right)\Gamma[D_{k}](s,s_{k+1}),$ (20) $$

where we define $\Lambda[D](s^{\prime})$ and $\Gamma[D](s,s^{\prime})$ as follows:

$$ $\displaystyle\Lambda[D](s^{\prime})$ $\displaystyle:=\frac{1}{|\mathcal{S}|}\int_{\mathcal{S}}hK_{\sigma}(s,s^{\prime})\exp\left(-\frac{\beta}{\gamma-1}D(s^{\prime})\right)\,\mathrm{d}s,$ (21) $\displaystyle\Gamma[D](s,s^{\prime})$ $\displaystyle:=hK_{\sigma}(s,s^{\prime})\exp\left(-\frac{\beta}{\gamma-1}D(s^{\prime})\right)-\Lambda[D](s^{\prime}).$ (22) $$

Now, under mild practical conditions, the driving bias $D_{k}(s)$ stays bounded, and $\exp(-\frac{\beta}{\gamma-1}D_{k}(s^{\prime}))$ is bounded above and below by positive constants.
Also, $\int_{\mathcal{S}}hK_{\sigma}(s,s^{\prime})\,\mathrm{d}s$ is a continuous positive function of $s^{\prime}$ on a compact set, so it also has positive finite bounds.
Therefore, from [Equation˜21](https://arxiv.org/html/2510.11923v1#A2.E21), we get the following bound with positive constant $a$ and $b$:

$$ $a\leq\Lambda[D_{k}](s_{k+1})\stackrel{{\scriptstyle\lx@cref{refnum}{eq:wt_bias_level}}}{{=}}(L_{k+1}-L_{k})\exp\left(\frac{\beta}{\gamma-1}L_{k}\right)\leq b.$ (23) $$

This result implies that the asymptotic effective step size is bounded by positive multiples of $\frac{1}{k}$: $\epsilon_{k}:=\exp(-\frac{\beta}{\gamma-1}L_{k})\asymp\frac{1}{k}$.

Step 2: Stochastic recursion to a mean-field ODE.
Now, we define an auxiliary time scale $\tau$ with $\tau_{0}=0$ and $\tau_{k+1}=\tau_{k}+\epsilon_{k}$.
Then, the driving part update [Equation˜20](https://arxiv.org/html/2510.11923v1#A2.E20) can be written in stochastic approximation form

$$ $D_{k+1}(s)=D_{k}(s)+\epsilon_{k}(\Gamma[D_{k}](s)+\zeta_{k+1}(s)),$ (24) $$

where we rewrite [Equations˜21](https://arxiv.org/html/2510.11923v1#A2.E21) and [22](https://arxiv.org/html/2510.11923v1#A2.E22) with expectation over $S\sim\bar{\nu}_{D}$, the CV marginal under the bias $V=L+D$,

$$ $\displaystyle\Lambda[D]$ $\displaystyle:=\frac{1}{|\mathcal{S}|}\int_{\mathcal{S}}\mathbb{E}_{S\sim\bar{\nu}_{D}}\left[hK_{\sigma}(s,S)\exp\left(-\frac{\beta}{\gamma-1}D(S)\right)\right]\,\mathrm{d}s,$ (25) $\displaystyle\Gamma[D](s)$ $\displaystyle:=\mathbb{E}_{S\sim\bar{\nu}_{D}}\left[hK_{\sigma}(s,S)\exp\left(-\frac{\beta}{\gamma-1}D(S)\right)\right]-\Lambda[D],$ (26) $$

and $\zeta_{k+1}$ is a martingale difference with bounded second moment.
The mini-batch hill update in [Algorithm˜1](https://arxiv.org/html/2510.11923v1#alg1) provides an unbiased estimator of the expectation $\mathbb{E}_{S\sim\bar{\nu}_{D}}$.
From Step 1, $\epsilon_{k}\asymp\frac{1}{k}$, which satisfies the Robbins–Monro step condition $\sum_{k}\epsilon_{k}=\infty$ and $\sum_{k}\epsilon_{k}^{2}<\infty$ (Robbins & Monro, 1951).
Standard stochastic approximation arguments then imply that the polygonal interpolation $\hat{D}(\tau_{k})=D_{k}$, $\hat{D}(\tau)=D_{k}+\frac{\tau-\tau_{k}}{\epsilon_{k}}(D_{k+1}-D_{k})$ for $\tau\in[\tau_{k},\tau_{k+1})$ is an asymptotic pseudo-trajectory of the mean-field ODE $\frac{\mathrm{d}}{\mathrm{d}\tau}D(\tau)=\Gamma[D(\tau)]$ (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 $\tau$ window while the bias changes only slightly within that window.

Step 3: Fixed points of the asymptotic ODE.
At a fixed point $D^{\ast}$, we have $\Gamma[D^{\ast}]\equiv 0$.
From [Equation˜26](https://arxiv.org/html/2510.11923v1#A2.E26), this correspond to

$$ $\int_{\mathcal{S}}hK_{\sigma}(s,s^{\prime})\exp\left(-\frac{\beta}{\gamma-1}D^{\ast}(s^{\prime})\right)\bar{\nu}_{D^{\ast}}(s^{\prime})\,\mathrm{d}s^{\prime}=\mathrm{const.}\quad\forall s\in\mathcal{S}.$ (27) $$

For Gaussian $K_{\sigma}$ on a compact domain, the associated integral operator is positive definite.
Hence, the solution is unique and satisfies

$$ $\displaystyle\exp\left(-\frac{\beta}{\gamma-1}D^{\ast}(s)\right)\underbrace{\exp(-\beta(F(s)+D^{\ast}(s)))}_{\propto\bar{\nu}_{D^{\ast}}(s)}=\mathrm{const.}$ (28) $\displaystyle\quad\Longrightarrow\quad V^{\ast}(s)=L^{\ast}+D^{\ast}(s)=-\left(1-\frac{1}{\gamma}\right)F(s)+\mathrm{const.}$ (29) $$

Step 4. Global stability.
Now, define the tempering-reweighted instantaneous CV density

$$ $p_{w}(s;\tau)\propto\exp\left(-\frac{\beta}{\gamma-1}D(s;\tau)\right)\underbrace{\exp(-\beta(F(s)+D(s;\tau)))}_{\propto\bar{\nu}_{D}(s)}.$ (30) $$

Let $p_{w}^{\ast}$ denote the target tempering-reweighted distribution, i.e., a normalized solution of

$$ $\int_{\mathcal{S}}hK_{\sigma}(s,s^{\prime})p_{w}^{\ast}(s^{\prime})\,\mathrm{d}s^{\prime}=\mathrm{const.}$ (31) $$

For Gaussian hills on a compact $\mathcal{S}$, this solution is unique and equals the uniform density $p_{w}^{\ast}(s)=|\mathcal{S}|^{-1}$.
Consider the Lyapunov functional defined as a KL divergence between $p_{w}^{\ast}$ and $p_{w}(\cdot;\tau)$:

$$ $\mathcal{D}(\tau):=D_{\text{KL}}(p_{w}^{\ast}\,\|\,p_{w}(\cdot;\tau))=\int_{S}p_{w}^{\ast}(s)\log\frac{p_{w}(s;\tau)}{p_{w}^{\ast}(s)}\,\mathrm{d}s.$ (32) $$

We now show that $\frac{\mathrm{d}}{\mathrm{d}\tau}\mathcal{D}(\tau)\leq 0$.
First,

$$ $\displaystyle\frac{\mathrm{d}}{\mathrm{d}\tau}D(s;\tau)$ $\displaystyle=\Gamma[D(\tau)](s)$ (33) $\displaystyle\stackrel{{\scriptstyle\lx@cref{refnum}{eq:wt_bias_drive_exp}}}{{=}}\int_{\mathcal{S}}hK_{\sigma}(s,s^{\prime})\exp\left(-\frac{\beta}{\gamma-1}D(s^{\prime};\tau)\right)\bar{\nu}_{D(\tau)}(s^{\prime})\,\mathrm{d}s^{\prime}-\Lambda[D(\tau)]$ (34) $\displaystyle\stackrel{{\scriptstyle\lx@cref{refnum}{eq:temp_rew_cv_density}}}{{=}}c(\tau)\int_{\mathcal{S}}hK_{\sigma}(s,s^{\prime})p_{w}(s^{\prime};D_{k})\,\mathrm{d}s^{\prime}-\Lambda[D(\tau)]$ (35) $\displaystyle\stackrel{{\scriptstyle\lx@cref{refnum}{eq:temp_rew_cv_target}}}{{=}}c(\tau)\int_{\mathcal{S}}hK_{\sigma}(s,s^{\prime})[p_{w}(s^{\prime};\tau)-p_{w}^{\ast}(s^{\prime})]\,\mathrm{d}s^{\prime}-C(\tau),$ (36) $$

where $c(\tau)>0$ and $C(\tau)$ are constants in $s$.
Also, from [Equation˜30](https://arxiv.org/html/2510.11923v1#A2.E30), we get

$$ $\displaystyle\frac{1}{p_{w}(s;\tau)}\frac{\mathrm{d}}{\mathrm{d}\tau}p_{w}(s;\tau)=-\frac{\beta\gamma}{\gamma-1}\left(\frac{\mathrm{d}}{\mathrm{d}\tau}D(s;\tau)-\mathbb{E}_{p_{w}}\left[\frac{\mathrm{d}}{\mathrm{d}\tau}D(\cdot;\tau)\right]\right).$ (37) $$

Finally, plugging [Equations˜36](https://arxiv.org/html/2510.11923v1#A2.E36) and [37](https://arxiv.org/html/2510.11923v1#A2.E37) into the time derivative of [Equation˜32](https://arxiv.org/html/2510.11923v1#A2.E32) yields

$$ $\displaystyle\frac{\mathrm{d}}{\mathrm{d}\tau}\mathcal{D}(\tau)$ $\displaystyle\stackrel{{\scriptstyle\lx@cref{refnum}{eq:lyapunov}}}{{=}}-\int_{\mathcal{S}}p_{w}^{\ast}(s)\frac{1}{p_{w}(s;\tau)}\frac{\mathrm{d}}{\mathrm{d}\tau}p_{w}(s;\tau)\,\mathrm{d}s$ (38) $\displaystyle\stackrel{{\scriptstyle\lx@cref{refnum}{eq:temp_rew_cv_dt}}}{{=}}\frac{\beta\gamma}{\gamma-1}\left(\mathbb{E}_{p_{w}^{\ast}}\left[\frac{\mathrm{d}}{\mathrm{d}\tau}D(\cdot;\tau)\right]-\mathbb{E}_{p_{w}}\left[\frac{\mathrm{d}}{\mathrm{d}\tau}D(\cdot;\tau)\right]\right)$ (39) $\displaystyle\stackrel{{\scriptstyle\lx@cref{refnum}{eq:drive_dt}}}{{=}}-\frac{\beta\gamma}{\gamma-1}c(\tau)\int_{\mathcal{S}}\int_{\mathcal{S}}[p_{w}(s;\tau)-p_{w}^{\ast}(s)]hK_{\sigma}(s,s^{\prime})[p_{w}(s^{\prime};\tau)-p_{w}^{\ast}(s^{\prime})]\,\mathrm{d}s\mathrm{d}s^{\prime}\leq 0,$ (40) $$

since the integral operator with the Gaussian kernel $hK_{\sigma}(\cdot,\cdot)$ is positive semidefinite and $c(\tau)>0$.
The equality in [Equation˜40](https://arxiv.org/html/2510.11923v1#A2.E40) holds if and only if $p_{w}(\cdot;\tau)=p_{w}^{\ast}$, and $\mathcal{D}$ is a strict Lyapunov functional and the ODE solution converges to the unique fixed point.

Hence, the mean-field ODE $\frac{\mathrm{d}}{\mathrm{d}\tau}D(\tau)=\Gamma[D(\tau)]$ has a globally attracting equilibrium $D^{\ast}$, and one gets $D_{k}\to D^{\ast}$ and $V_{k}\to V^{\ast}$ in [Equation˜29](https://arxiv.org/html/2510.11923v1#A2.E29) almost surely by following the stochastic approximation [Equation˜24](https://arxiv.org/html/2510.11923v1#A2.E24).
$\square$

## Appendix C Implementation Details for WT-ASBS

### C.1 WT-ASBS with Replay Buffer

Here, we describe the practical implementation of [Algorithm˜1](https://arxiv.org/html/2510.11923v1#alg1), which merges the well-tempered bias update into the AM replay buffer update within the inner step.
The resulting WT-ASBS algorithm with replay buffer is shown in [Algorithm˜2](https://arxiv.org/html/2510.11923v1#alg2).
After updating the AM replay buffer $\mathcal{B}_{\text{adj}}$ (line 6), we reuse the same samples, projecting them to CV space to obtain the CV values for Gaussian kernel deposition (line 7).

Figure: Algorithm 2 WT-ASBS with replay buffer

#### Grid discretization of bias potential

Since the bias potential is accumulated over the entire training history, longer training would require summing over an increasingly large number of Gaussian kernels to evaluate $\nabla(V\circ\xi)$.
To ensure computational efficiency, we adopt the standard grid discretization approach commonly used in enhanced sampling simulations.
The bias is stored on a discretized grid with spacing chosen to be smaller than the kernel width $\sigma$, 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 $u_{\theta}$ more closely with available data samples.
In our setting, these samples are generated by short MD simulations that perform local sampling around a reference configuration $x_{\text{ref}}$.
The bridge matching objective $\mathcal{L}_{\text{pretrain}}$ is defined as follows:

$$ $\mathcal{L}_{\text{pretrain}}(\theta)=\mathbb{E}_{t\sim\mathcal{U}[0,1],X_{0}\sim\mu,X_{1}\sim q,X_{t}\sim p^{\text{base}}(\cdot|X_{0},X_{1})}\left[\hat{\lambda}_{t}\|u_{t}(X_{t})-\sigma_{t}\nabla_{X_{t}}\log p^{\text{base}}(X_{1}|X_{t})\|^{2}\right],$ (44) $$

where $q$ denotes the empirical distribution of the MD samples, and $\hat{\lambda}_{t}$ is a time-dependent weighting factor, chosen as $\kappa_{1|t}/\sigma_{t}^{2}$, where $\kappa_{1|t}$ is the conditional variance of the base process (see [Section˜C.3](https://arxiv.org/html/2510.11923v1#A3.SS3)).
This pretraining is only approximate, since the marginal distribution of the optimal $u_{\theta}^{\ast}$ matches $q(X_{1})$ exactly only when the base process is memoryless (i.e., when $X_{0}$ and $X_{1}$ are independent).
Nevertheless, for warm-start purposes, this approximation is sufficient to rapidly align the model with a distribution concentrated in a single conformational mode.

### C.2 Model Architecture

We build on PaiNN (Schütt et al., 2021), an $\mathrm{E}(3)$-equivariant graph neural network designed for molecular energy and force prediction, with minimal modifications: adding time-conditioning inputs, incorporating vector product layers (Schreiner et al., 2023) to break parity symmetry, restricting the model to $\mathrm{SE}(3)$-equivariance, and applying equivariant layer normalization (Liao et al., 2024) to improve training stability.
Each atom is assigned a unique atomic embedding, as the reference configuration used to define the accessible component ([Requirement˜1](https://arxiv.org/html/2510.11923v1#S2.I1.i1)) already specifies atom indices.
The model outputs a single vector per atom, which is scaled by $\sigma_{t}$ 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 $\sigma_{t}$ through a power-schedule parameter $\rho$ and yields an analytically integrable variance between two different times $\kappa_{t|s}$:

$$ $\displaystyle\sigma_{t}$ $\displaystyle:=\left[(1-t)\sigma_{\text{max}}^{1/\rho}+t\sigma_{\text{min}}^{1/\rho}\right]^{\rho},$ (45) $\displaystyle\kappa_{t|s}$ $\displaystyle:=\int_{s}^{t}\sigma_{\tau}^{2}\,\mathrm{d}\tau=\frac{1}{2\rho+1}\frac{\sigma_{s}^{2+1/\rho}-\sigma_{t}^{2+1/\rho}}{\sigma_{\text{max}}^{1/\rho}-\sigma_{\text{min}}^{1/\rho}}.$ (46) $$

With the reparametrized time $\gamma_{t}:=\frac{\kappa_{t|0}}{\kappa_{1|0}}\in[0,1]$, conditional base distributions for the driftless base $\mathrm{d}X_{t}=\sigma_{t}\mathrm{d}W_{t}$ are

$$ $\displaystyle p_{t|0}^{\text{base}}(X_{t}|X_{0})$ $\displaystyle=\mathcal{N}(X_{t};X_{0},\kappa_{t|0}I),$ (47) $\displaystyle p_{t|0,1}^{\text{base}}(X_{t}|X_{0},X_{1})$ $\displaystyle=\mathcal{N}(X_{t};(1-\gamma_{t})X_{0}+\gamma_{t}X_{1},\gamma_{t}(1-\gamma_{t})\kappa_{1|0}I).$ (48) $$

#### Translational invariance

As in AS (Havens et al., 2025), we enforce translational invariance by projecting the SDE onto the subspace $\mathcal{X}^{\text{CoM}}=\{x\,|\,\sum_{i=1}^{n}x^{i}=0\}\subset\mathcal{X}$, where $x^{i}\in\mathbb{R}^{3}$ is the position of $i$-th atom.
This is achieved through the projection operator $P$, which acts on any $y\in\mathbb{R}^{n\times 3}$ by subtracting the center of mass from each coordinate $y^{i}$ for $i\in[n]$.

$$ $x=Py,\qquad P=\left(I_{n}-\frac{1}{n}\bm{1}_{n}\bm{1}_{n}^{\top}\right)\otimes I_{d}$ (49) $$

$$ $dX_{t}=\sigma_{t}Pu_{t}(X_{t})\,\mathrm{d}t+\sigma_{t}P\,\mathrm{d}W_{t},\qquad X_{0}=0$ (50) $$

For the base process, the dynamics are given by

$$ $\displaystyle X_{t}$ $\displaystyle=X_{0}+\int_{0}^{t}\sigma_{s}P\,\mathrm{d}W_{s}$ (51) $\displaystyle\operatorname{Var}[X_{t}|X_{0}]$ $\displaystyle=\int_{0}^{t}\sigma_{s}^{2}PP^{\top}\,\mathrm{d}s=\kappa_{t|0}PP^{\top}.$ (52) $$

Sampling can be carried out by first drawing $Y_{t}\sim\mathcal{N}(X_{0},\kappa_{t|0}I)$ and then projecting to obtain $X_{t}=PY_{t}$, which enforces the translational invariance constraint.

Adjoint matching [Equation˜10](https://arxiv.org/html/2510.11923v1#S2.E10) and corrector matching [Equation˜11](https://arxiv.org/html/2510.11923v1#S2.E11) are modified accordingly as follows:

$$ $\displaystyle u$ $\displaystyle:=\operatorname*{arg\,min}_{u}\mathbb{E}_{p_{t|0,1}^{\text{base}}p_{0,1}^{\bar{u}}}[\|P(u_{t}(X_{t})+\sigma_{t}(\nabla E+h^{(k-1)})(X_{1}))\|^{2}],\;\bar{u}=\texttt{stopgrad}(u),$ (53) $\displaystyle h$ $\displaystyle:=\operatorname*{arg\,min}_{h}\mathbb{E}_{p_{0,1}^{u^{(k)}}}[\|P(h(X_{1})-\nabla_{X_{1}}\log p_{1|0}^{\text{base}}(X_{1}|X_{0}))\|^{2}].$ (54) $$

### C.4 Hyperparameters

**Table 1: Hyperparameters for WT-ASBS. Model and training hyperparameters for each experiment.**
| Hyperparameter | Alanine<br>dipeptide | Alanine<br>tetrapeptide | Nucleophilic<br>substitution | Post-TS<br>bifurcation |
| --- | --- | --- | --- | --- |
| Models |  |  |  |  |
| Hidden dimension | 128 | $128$ | 64 | 128 |
| Radial basis | 64 | $64$ | 32 | 64 |
| Control layers | 4 | $4$ | 3 | 4 |
| Corrector layers | 3 | $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 | $\mathcal{N}(0,1\,\text{\AA }^{2})$ | N(0, 1   Å^2) | $\mathcal{N}(0,1\,\text{\AA }^{2})$ | $\mathcal{N}(0,1\,\text{\AA }^{2})$ |
| $\sigma_{\text{min}}$ [Å] | 0.001 | $0.001$ | 0.001 | 0.001 |
| $\sigma_{\text{max}}$ [Å] | 6.0 | $6.0$ | 5.0 | 6.0 |
| EDM exponent $\rho$ | 3.0 | $3.0$ | 3.0 | 3.0 |
| Bias update |  |  |  |  |
| Initial Gaussian $h$ [eV] | $1\times 10^{-4}$ | 5 ×10^-4 | $5\times 10^{-4}$ | $2\times 10^{-4}$ |
| Gaussian width $\sigma$ | [0.2, 0.2] rad | $[0.3,0.3,0.3]rad$ | [0.2, 0.2] Å | [0.05, 0.05] |
| CV grid minimum | [$-\pi,-\pi$] rad | $[$-π,-π,-π$]rad$ | [1.0, 1.0] Å | [0.0, -1.0] |
| CV grid maximum | [$+\pi,+\pi$] rad | $[$+π,+π,+π$]rad$ | [4.0, 4.0] Å | [3.0, 1.0] |
| CV grid bins | [150, 150] | $[100,100,100]$ | [150, 150] | [300, 200] |
| Periodic CV grid | True | $True$ | False | False |
| Bias factor $\gamma$ | 8.0 | $12.0$ | 12.0 | 30.0 |
| Training |  |  |  |  |
| Pretraining learning rate | $3\times 10^{-4}$ | 3 ×10^-4 | $3\times 10^{-4}$ | $3\times 10^{-4}$ |
| Pretraining steps | 100,000 | $100,000$ | 20,000 | 100,000 |
| Training stages $K$ | 4 | $40$ | 4 | 2 |
| AM epochs $M_{\text{adj}}$ | 1,000 | $1,000$ | 500 | 1,000 |
| CM epochs $M_{\text{crt}}$ | 500 | $500$ | 250 | 500 |
| Initial learning rate | $1\times 10^{-4}$ | 1 ×10^-4 | $3\times 10^{-5}$ | $3\times 10^{-5}$ |
| Final learning rate | $3\times 10^{-5}$ | 5 ×10^-5 | $1\times 10^{-5}$ | $1\times 10^{-5}$ |
| Gradient clipping value | 100.0 | $100.0$ | 100.0 | 100.0 |

Model and training hyperparameters for each experiment are listed in [Table˜1](https://arxiv.org/html/2510.11923v1#A3.T1).
For WT-ASBS [Algorithm˜2](https://arxiv.org/html/2510.11923v1#alg2), we use $L=100$ gradient updates per AM and CM epoch, and $\alpha_{\text{max}}=100.0$ to clip the energy gradient in the terminal cost.
Models were trained with AdamW optimizer (Loshchilov & Hutter, 2017) using a weight decay of $10^{-3}$, and diffusion process was 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 a final AM stage without further bias deposition.
Learning rates followed a cosine annealing schedule during training, applied independently to AM and CM optimization.

## Appendix D Experiment Details

For small peptides, we 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 implemented in the fairchem-core 2.3.0 package (Shuaibi et al., 2025) with MD integrators in ASE 3.25.0 (Larsen et al., 2017).
All MD-based enhanced sampling (WTMetaD) simulations were carried out with PLUMED 2.9.3 (Tribello et al., 2014; The PLUMED consortium, 2019).

### D.1 Alanine Dipeptide

#### Energy model

The potential energy of the system was described by
Amber ff96 force field (Kollman et al., 1997) with GBSA-OBC implicit solvation (Onufriev et al., 2004), as implemented in OpenMM (Eastman et al., 2023).

#### Restraints

We apply the chirality restraint [Equation˜18](https://arxiv.org/html/2510.11923v1#A1.E18) to the central alanine, $U_{\text{Ala2}}^{\text{chiral}}(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^{3}$ frames at 0.1 ps intervals.
Simulations employed the Langevin middle integrator (Zhang et al., 2019) as implemented in OpenMM (Eastman et al., 2023), with a collision rate of 1 ps^-1.
In total, pretraining data generation required $2.75\times 10^{4}$ energy evaluations.

#### Reference data

We employed the test split of $10^{7}$ samples from Midgley et al. (2022), generated by replica exchange MD with 21 replicas spanning 300 K to 1300 K.
Their dataset generation required $2.3\times 10^{10}$ energy evaluations.

#### WTMetaD simulations

We performed WTMetaD simulations on the backbone torsion CVs $(\phi,\psi)$.
The dynamics were propagated with Langevin middle integrator (Zhang et al., 2019) at 300 K with a 1 fs timestep and 1 ps^-1 collision rate.
Following a 10 ps equilibration at 300 K, Gaussian 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 a grid covering $[-\pi,\pi]^{2}$ with $150\times 150$ bins.
Each production run was 10 ns in length, and simulations were repeated four times to ensure statistical reliability, yielding $4.0\times 10^{7}$ energy evaluations.

### D.2 Alanine Tetrapeptide

#### Energy model

The potential energy of the system was described by
Amber ff99SB-ILDN force field (Lindorff-Larsen et al., 2010) with Amber99-OBC implicit solvation (Onufriev et al., 2004), as implemented in OpenMM (Eastman et al., 2023).

#### Restraints

We apply the chirality restraints [Equation˜18](https://arxiv.org/html/2510.11923v1#A1.E18) to three alanine residues, $U_{\text{Ala2}}^{\text{chiral}}(X)+U_{\text{Ala3}}^{\text{chiral}}(X)+U_{\text{Ala4}}^{\text{chiral}}(X)$, during WT-ASBS training.
As in the alanine dipeptide case, these restraints are not used in MD-based sampling.

#### Pretraining data

We generated pretraining data using a 1 ns molecular dynamics simulation at 300 K with a 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 $10^{3}$ frames at 1 ps intervals.
Simulations employed the Langevin middle integrator (Zhang et al., 2019) as implemented in OpenMM (Eastman et al., 2023), with a collision rate of 1 ps^-1.
In total, pretraining data generation required $2.52\times 10^{5}$ 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 dynamics 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 hybrid 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 an energy-minimized 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 a 1 fs timestep and 1 ps^-1 collision rate.
After 100 ps of equilibration, production simulations were performed for 1.56 µs, yielding a total of 100 µs simulation time across all replicas.
Trajectories were saved every 2 ps, producing $5\times 10^{7}$ samples in total.
Reference data generation required approximately $1.01\times 10^{11}$ energy evaluations.

#### WTMetaD simulations

We performed WTMetaD simulations on the three backbone torsion CVs $(\phi_{1},\phi_{2},\phi_{3})$.
The same integrator settings as in the reference data generation were used.
Following a 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 a grid covering $[-\pi,\pi]^{3}$ with $75^{3}$ bins.
Each production run was 20 ns in length, and simulations were repeated four times to ensure statistical reliability, yielding $8.0\times 10^{7}$ energy evaluations.

Figure: Figure 6: Atom index definition for chemical reactions. (a) Nucleophilic substitution (SN2). (b) Cycloaddition reaction between tropone and dimethylfulvene with post-TS bifurcation.
Refer to caption: x6.png

### D.3 Nucleophilic Substitution

#### Energy model

The potential energy of the system was described by the UMA-S-1.1 (Wood et al., 2025) model with the DFT task label omol, a system charge of –1, and a singlet spin state (total spin 0).

#### Restraints

In WT-ASBS, three C–H bonds (1–2, 1–3, and 1–4 in [Figure˜6](https://arxiv.org/html/2510.11923v1#A4.F6)a) are maintained using the bond restraint in [Equation˜14](https://arxiv.org/html/2510.11923v1#A1.E14).
For both WT-ASBS and WTMetaD, to keep the accessible CV space $\mathcal{S}$ bounded, the same restraint form is applied to two C–Cl distances (1–5 and 1–6) with $r_{\text{max}}=3.3$ Å, preventing chloride ions from escaping the reaction complex.

#### WTMetaD simulations

We performed WTMetaD simulations using $(d_{1},d_{2})$ as the CVs.
The dynamics were propagated with a Langevin thermostat at 300 K, a 0.5 fs timestep, and a friction coefficient of 10 ps^-1.
Initial velocities were assigned from a Maxwell–Boltzmann distribution, and the total momentum was removed.
Following a 1 ps equilibration at 300 K, Gaussian biases were deposited every 250 steps with widths of 0.2 Å for both CVs.
The biasing scheme used an initial Gaussian height of 0.02 eV, a bias factor of 20, and a grid spanning $[1.0,4.0]$ Å $\times$ $[1.0,4.0]$ Å.
Each production run was 0.5 ns in length, and simulations were repeated four times to ensure statistical reliability, yielding $4.0\times 10^{6}$ energy evaluations.

### D.4 Post-Transition-State Bifurcation

#### Energy model

The potential energy of the system was described by the UMA-S-1.1 (Wood et al., 2025) model with the DFT task label omol, a system charge of 0, and a singlet spin state (total spin 0).

#### Collective variables

We define a normalized contact $c_{i}$ that maps the bond distance $d_{i}$ to a soft bond indicator,

$$ $c_{i}=\frac{1-(d_{i}/d_{\text{ref}})^{6}}{1-(d_{i}/d_{\text{ref}})^{8}}\in[0,1],$ (55) $$

which approaches 1 when $d_{i}$ is much smaller than $d_{\text{ref}}$ and decays to 0 as $d_{i}$ increases (bond dissociation).
To construct the CVs, we express the three relevant C–C distances ([Figure˜4](https://arxiv.org/html/2510.11923v1#S4.F4)d) as normalized contacts using [Equation˜55](https://arxiv.org/html/2510.11923v1#A4.E55) with $d_{\text{ref}}=1.7$ Å, thereby capturing bond formation progress ([Figure˜4](https://arxiv.org/html/2510.11923v1#S4.F4)e).
We then define two CVs: $s_{1}=c_{1}+c_{2}+c_{3}$, which measures overall bond formation, and $s_{2}=c_{2}-c_{3}$, 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, [Figure˜4](https://arxiv.org/html/2510.11923v1#S4.F4)c) are preserved using the bond restraint in [Equation˜14](https://arxiv.org/html/2510.11923v1#A1.E14).
Furthermore, to prevent tropone (upper reactant) from “teleporting” to the opposite face of dimethylfulvene (lower reactant), we enforce an inter-planar distance defined by three atoms per plane:

$$ $d(X)=\underbrace{\frac{(x^{(1)}_{1}-x^{(1)}_{0})\times(x^{(1)}_{2}-x^{(1)}_{0})}{\|(x^{(1)}_{1}-x^{(1)}_{0})\times(x^{(1)}_{2}-x^{(1)}_{0})\|}}_{\text{Unit normal of plane 1}}\,\cdot\,\underbrace{\left(\frac{1}{3}\sum_{m=0}^{2}x^{(2)}_{m}-\frac{1}{3}\sum_{m=0}^{2}x^{(1)}_{m}\right)}_{\text{Inter-centroid displacement}}.$ (56) $$

$d(X)$ is positive if the centroid of plane 2 lies on the side of plane 1 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 1: dimethylfulvene atoms 15, 13, 24; plane 2: tropone atoms 2, 3, 7 in [Figure˜6](https://arxiv.org/html/2510.11923v1#A4.F6)b) and apply a lower harmonic wall restraint $U_{\text{plane}}(X)=\frac{1}{2}\kappa[-d(X)]_{+}^{2}$ with $\kappa=100$ eV/Å^2 to enforce $d(X)>0$.

Finally, for both WT-ASBS and WTMetaD, the no-bond restraint in [Equation˜15](https://arxiv.org/html/2510.11923v1#A1.E15) is applied to atom pairs C6–C24, C3–C18, and O8–C19 to suppress sampling of irrelevant reactive pathways.

#### Pretraining data

The system was propagated with a Langevin thermostat at 300 K, a 1 fs timestep, and a friction coefficient of 1 ps^-1.
Initial velocities were assigned from a Maxwell–Boltzmann distribution 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, pretraining data generation required $1.2\times 10^{4}$ sequential energy evaluations.

#### WTMetaD simulations

We performed WTMetaD simulations using $(s_{1},s_{2})$ as the CVs.
The dynamics were propagated with a Langevin thermostat at 300 K, a 0.5 fs timestep, and a friction coefficient of 10 ps^-1.
Initial velocities were assigned from a Maxwell–Boltzmann distribution, and the total momentum was removed.
Following a 1 ps equilibration at 300 K, Gaussian biases were deposited every 250 steps with widths of 0.05 for both CVs.
The biasing scheme used an initial Gaussian height of 0.02 eV, a bias factor of 20, and a grid spanning $[0.0,3.0]\times[-1.0,1.0]$.
Each production run was 0.8 ns in length, and simulations were repeated four times to ensure statistical reliability, yielding $6.4\times 10^{6}$ energy evaluations.

## Appendix E Additional Results

### E.1 Alanine dipeptide

In this section, we extend the alanine dipeptide sampling comparison from the main text to standard distributional metrics following Liu et al. (2025).
Specifically, we report KL divergence for 1-D marginals of torsion angles, as well as $\mathcal{W}_{2}$ and Jensen–Shannon divergence (JSD) for the joint distribution of backbone torsions $(\phi,\psi)$.
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](https://arxiv.org/html/2510.11923v1#A5.T2).
We used $10^{6}$ samples to compute $D_{\text{KL}}$ and JSD, and randomly selected $10^{4}$ samples for $\mathcal{W}_{2}$ (with weighted random sampling for WT-ASBS).

**Table 2: Distributional metrics for alanine dipeptide. Comparison of diffusion samplers applied in internal or Cartesian coordinate representations of alanine dipeptide sampling. The KL divergence ($D_{\text{KL}}$) of five 1-D torsion angle marginals and the Wasserstein-2 distance ($\mathcal{W}_{2}$) and JSD on joint backbone torsion distribution are reported. ^∗Results from Liu et al. (2025). ^†Result from Choi et al. (2025). ^‡Reweighted samples are used.**
|  | $D_{\text{KL}}$ on 1-D marginal ($\downarrow$) |  |  |  |  | $(\phi,\psi)$ joint |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Method | $\phi$ | ψ | $\gamma_{1}$ | $\gamma_{2}$ | $\gamma_{3}$ | $\mathcal{W}_{2}$ ($\downarrow$) | JSD $(\downarrow)$ |
| Internal coordinates |  |  |  |  |  |  |  |
| PIS (Zhang & Chen, 2021)^∗ | 0.05$\pm$0.03 | $0.38{\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\tiny$\pm$0.49}$ | 5.61$\pm$1.24 | 4.49$\pm$0.03 | 4.60$\pm$0.03 | 1.27$\pm$1.19 | – |
| DDS (Vargas et al., 2023a)^∗ | 0.03$\pm$0.01 | $0.16{\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\tiny$\pm$0.07}$ | 2.44$\pm$0.96 | 0.03$\pm$0.00 | 0.03$\pm$0.00 | 0.68$\pm$0.09 | – |
| AS (Havens et al., 2025)^∗ | 0.09$\pm$0.09 | $0.04{\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\tiny$\pm$0.04}$ | 0.17$\pm$0.17 | 0.56$\pm$0.09 | 0.51$\pm$0.06 | 0.65$\pm$0.52 | – |
| NAAS (Choi et al., 2025)^† | 0.26$\pm$0.00 | $0.24\hphantom{{\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\tiny$\pm$0.00}}$ | 0.27$\pm$0.00 | 0.13$\pm$0.00 | 0.16$\pm$0.00 | – | – |
| ASBS (Liu et al., 2025)^∗ | 0.02$\pm$0.00 | $0.01{\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\tiny$\pm$0.00}$ | 0.03$\pm$0.01 | 0.02$\pm$0.00 | 0.02$\pm$0.00 | 0.25$\pm$0.01 | – |
| Cartesian coordinates |  |  |  |  |  |  |  |
| ASBS (Liu et al., 2025) | 0.14$\pm$0.03 | $0.09{\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\tiny$\pm$0.03}$ | 0.02$\pm$0.00 | 0.00$\pm$0.00 | 0.00$\pm$0.00 | 0.45$\pm$0.01 | 0.014$\pm$0.003 |
| WT-ASBS (this work)^‡ | 0.02$\pm$0.00 | $0.04{\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\tiny$\pm$0.01}$ | 0.02$\pm$0.00 | 0.00$\pm$0.00 | 0.00$\pm$0.00 | 0.43$\pm$0.02 | 0.005$\pm$0.001 |

Figure: Figure 7: Ramachandran plots for alanine dipeptide. Comparison of backbone torsion distributions $(\phi,\psi)$ in log densities, for reference, ASBS (internal coordinates), ASBS and WT-ASBS (Cartesian coordinates), as in [Table˜2](https://arxiv.org/html/2510.11923v1#A5.T2). All three methods reproduce the more populated regions reasonably well, but their behavior in less populated regions differs substantially. Negligibly low-density regions are omitted for visual clarity. ^∗Taken from Liu et al. (2025).
Refer to caption: x7.png

[Table˜2](https://arxiv.org/html/2510.11923v1#A5.T2) shows that WT-ASBS performs well on all metrics, with similar $\mathcal{W}_{2}$ performance to baseline ASBS (Cartesian) and slightly worse than reference ASBS (internal coordinates).
However, the Ramachandran plots in [Figure˜7](https://arxiv.org/html/2510.11923v1#A5.F7) 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˜2](https://arxiv.org/html/2510.11923v1#A5.T2) assume equal weight for all samples.
WT-ASBS, designed to satisfy [Requirement˜3](https://arxiv.org/html/2510.11923v1#S2.I1.i3), incorporates weighted sampling to achieve balanced efficiency across all relevant modes.
Since Boltzmann populations depend exponentially on free energy, focusing on more populated regions can yield good metrics without capturing the full distribution.
Therefore, we emphasize evaluations more closely aligned with the design choices discussed in the main text.

### E.2 Alanine tetrapeptide

In [Figure˜8](https://arxiv.org/html/2510.11923v1#A5.F8), we show the distribution of pretraining data and compare the PMF obtained by reweighting ([Equation˜1](https://arxiv.org/html/2510.11923v1#S2.E1)) for WT-ASBS, the bias ([Remark˜3.1](https://arxiv.org/html/2510.11923v1#S3.Thmremark1)) 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\times 10^{7}$ for the reference), underscoring the necessity of biased sampling.

Figure: Figure 8: 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 to different torsion angle pairs $(\phi_{1},\phi_{2})$, $(\phi_{1},\phi_{3})$, and $(\phi_{2},\phi_{3})$. Dotted lines at $\phi_{i}=0$ and $2$ mark the mode boundaries. All method reproduces populations within metastable modes, while biased sampling (WT-ASBS and WTMetaD) captures transition regions inaccessible to unbiased sampling.
Refer to caption: x8.png

### E.3 Sampling Reactive Energy Landscapes

To quantify the convergence of PMFs in [Figure˜5](https://arxiv.org/html/2510.11923v1#S4.F5), we define the PMF MAE between a PMF $F(s)$ and its reference $F_{\text{ref}}(s)$, both shifted to have minimum zero, as

$$ $\text{MAE}[F,F_{\text{ref}}]=\frac{\min_{c}\int_{\mathcal{S}}|F(s)-F_{\text{ref}}(s)+c|\,\bm{1}[F_{\text{ref}}(s)<F_{\text{thres}}]\,\mathrm{d}s}{\int_{\mathcal{S}}\bm{1}[F_{\text{ref}}(s)<F_{\text{thres}}]\,\mathrm{d}s},$ (57) $$

i.e., the average absolute deviation (up to an optimal constant shift $c$) over regions of CV space where $F_{\text{ref}}(s)$ is below the threshold $F_{\text{thres}}$.
We used $F_{\text{thres}}$ of 30 kcal/mol for the SN2 reaction, and 45 kcal/mol for the post-TS bifurcation reaction.

Figure: Figure 9: SN2 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 $d_{1}=d_{2}$ symmetry lines, and the purple cross indicates the transition state location.
Refer to caption: x9.png

Figure: Figure 10: Post-TS bifurcation reaction PMFs. Comparison of 1-D PMF $F(s_{1})$ and 2-D PMF $F(s_{1},s_{2})$ obtained from WT-ASBS and WTMetaD biases. Red dotted lines mark regions influenced by restraints, gray dotted lines denote the $c_{2}=c_{3}$ symmetry lines, and the purple cross indicates the transition state locations.
Refer to caption: x10.png

We compare PMFs for reactions obtained with bias from WT-ASBS and WTMetaD in [Figures˜9](https://arxiv.org/html/2510.11923v1#A5.F9) and [10](https://arxiv.org/html/2510.11923v1#A5.F10), 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, two competing mechanisms of 3 $\to$ 4 $\to$ 5 (Paddon-Row & Warrener, 1974) vs. 3 $\to$ 5 (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 5 makes it the thermodynamically favored product, consistent with previous findings that 4 forms transiently and subsequently converts to 5 (Yu et al., 2017).

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: See weaknesses.
soundness: 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.

The 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.

All 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.
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.

The 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.

The 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.
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.
- 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.
- 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.
- 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).

[1] Seong, Kiyoung, et al. “Transition Path Sampling with Improved Off-Policy Training of Diffusion Path Samplers.” arXiv preprint arXiv:2405.19961 (2024).
[2] Singh, Aditya N., Avishek Das, and David T. Limmer. “Variational path sampling of rare dynamical events.” Annual Review of Physical Chemistry 76 (2025).
[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.
[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.

---

confidence: 2
contribution: 3
initial_rating: 8
presentation: 3
questions: See weakness
soundness: 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. 

- 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.

- The paper is technically sound and well-executed.

  - The authors provide a convergence guarantee (Proposition 3.1), showing that WT-ASBS provably approaches the desired well-tempered target distribution.

  - 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.

  - 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).

- 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.
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.
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$. 

It 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.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
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?

(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.

(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.

(Q4) Could you summarize cost–benefit metrics (energy evaluations, wall-clock time, ESS, PMF error) across ASBS, WT-ASBS, and WTMetaD for standardized comparison?

(Q5) ASBS’s slow exploration (Figure 3.c) might result from strong pre-training confinement near the initial mode (marked * in Figure 3.b).
Have you explored reducing pretraining strength or fine-tuning its duration to encourage broader exploration?

(Q6) When computing the free-energy difference, which states or basins were used? Were they predefined minima or clusters in CV space?
Clarifying this would help interpret the reported PMF results.

(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.
soundness: 4
strengths: **Conceptually innovative hybridization.**
The 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.

**Clear motivation from sampling efficiency.**
The 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.

**Transparent discussion of limitations.**
The paper openly discusses the non-amortized bias, potential generalization issues, and differences from traditional WTMetaD, which reflects scientific maturity and honesty.
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.

The 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.
weaknesses: **1. Conceptual tension between learned sampler and non-learnable bias.**
Although 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.

**2. Lack of fair and direct comparison with ASBS.**
Since 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.

**3. # of force-calls inconsistency between figures.**
Figure 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.

**4. Limited diffusion-based baselines.**
Comparisons 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.

**5. Scalability and consistency.**
WT-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.

**6. Lack of statistical robustness.**
The 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.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: 1. How should practitioners determine that the two-loop WT-ASBS training has converged?

2. 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?

3. Have the authors considered coupling WT-ASBS with ML-based CV discovery?

4. WT-ASBS achieves broader exploration but sometimes slightly worse free-energy accuracy (Fig. 3d). Can the authors comment on this?

5. Can the authors explore the computational feasibility for systems with explicit solvent?
soundness: 4
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.

2. 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.

3. 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.

4. 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.
summary: This paper introduces WT-ASBS, a diffusion-based generative sampler enhanced by biasing along  CVs, inspired by WTMetaD.

The 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.
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.

2. 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.

3. 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.

4. 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.

5. 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.

6. The method’s feasibility for explicit-solvent systems remains unclear.


Decision: Accept (Poster)