Consistent Sampling and Simulation: Molecular Dynamics with Energy-Based Diffusion Models
Authors: Michael Plainer
Hao Wu
Leon Klein
Stephan Günnemann
Frank Noe
Venue: NeurIPS 2025 poster ()
Keywords: Energy-based Diffusion Models
Molecular Dynamics
Coarse-Graining
Boltzmann Emulator
Fokker-Planck Equation
Mixture of Experts
Primary area: machine_learning_for_sciences
TL;DR: An energy-based diffusion model regularized with the Fokker-Planck equation to improve the consistency between iid sampling and simulation

Abstract: In recent years, diffusion models trained on equilibrium molecular distributions have proven effective for sampling biomolecules. Beyond direct sampling, the score of such a model can also be used to derive the forces that act on molecular systems. However, while classical diffusion sampling usually recovers the training distribution, the corresponding energy-based interpretation of the learned score is often inconsistent with this distribution, even for low-dimensional toy systems. We trace this inconsistency to inaccuracies of the learned score at very small diffusion timesteps, where the model must capture the correct evolution of the data distribution. In this regime, diffusion models fail to satisfy the Fokker-Planck equation, which governs the evolution of the score. We interpret this deviation as one source of the observed inconsistencies and propose an energy-based diffusion model with a Fokker-Planck-derived regularization term to enforce consistency. We demonstrate our approach by sampling and simulating multiple biomolecular systems, including fast-folding proteins, and by introducing a state-of-the-art transferable Boltzmann emulator for dipeptides that supports simulation and achieves improved consistency and efficient sampling. Our code, model weights, and self-contained JAX and PyTorch notebooks are available at https://github.com/noegroup/ScoreMD.

Consistent Sampling and Simulation: Molecular Dynamics with Energy-Based Diffusion Models
Michael Plainer 1,2,3,4
Hao Wut
Leon Klein
Stephan Günnemann
Frank Noet 1,7,8
Technische Universität Berlin Berlin Institute for the Foundations of Learning and Data 'Shanghai Jiao Tong University Technische Universität München
Freie Universität Berlin
Zuse School ELIZA
"Rice University
*Microsoft Research A14Science
Abstract
In recent years, diffusion models trained on equilibrium molecular distributions have proven effective for sampling biomolecules. Beyond direct sampling, the score of such a model can also be used to derive the forces that act on molecular systems. However, while classical diffusion sampling usually recovers the training distribution, the corresponding energy-based interpretation of the learned score is often inconsistent with this distribution, even for low-dimensional toy systems. We trace this inconsistency to inaccuracies of the learned score at very small diffusion timesteps, where the model must capture the correct evolution of the data distribu- tion. In this regime, diffusion models fail to satisfy the Fokker-Planck equation, which governs the evolution of the score. We interpret this deviation as source the observed inconsistencies and propose an energy-based based diffusion model with a Lokker-Plancc--erring regularization term to enforce consistency. We demonstrate our approach by sampling and simulating multiple biomolecular systems, including fast-folding proteins, and by introducing state-of-the-arr transferable Boltzmann emulator for dipeptides that supports simulation and achieves improved consistency and efficient sampling. Our code, model weights, and self-contained JAX and PyTorch notebooks are available at https::/ /tt///gttbb./////////////e cttpauu..oom///omm http///ogggooooorrreeee//
Introduction
Understanding biochemical systems requires modeling not only static molecular structures but also their temporal evolution and interactions. Molecular dynamics (MD) simulations offer a principled framework to obtain such information from atomistic force fields and, with recent methodological and computational advances, can reach biologically relevant timescales (Lindorff-L Wolf et al., 2020). Nonetheless, fully atomistic simulations of large systems slow conformational transitions remain computationally challenging (Kmiccik et al., 2016; Plattner et al., 2017). Coarse- graining (CG) methods address this limitation by grouping atoms into beads accelerate simulations at the cost of physical resolution. In this reduced space, interactions and forces cannot be described accurately with traditional MD methods, requiring learning- based approaches to approximate forces and recover correct dynamics (Clementi, 2008; Noid, 2013; Husic et al., 2020; Charron et al., 2025).
et al., 2011;
Diffusion models (Ho et al., 2020; Song et al., 2021) have emerged powerful generative tools capable learning high-dimensional molecular distributions (Abramson al., 2024; Watson et al., 2023; Corso et al., 2023; Plainer et al., 2023b; Lewis al., 2025). They are trained to reverse
i "Corresponding Authors: mi chael. plainerefu- berlin.de hwu810sjtu.cm edu. en, franknoo@@iicrossft.ccom
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Reference
Denoising Score Matching Epss
inconsistent
iid Sampling
Learned Density (x,
Train Diffusion
Salve revene SID Pglogp"(rr11
Denoising Score Matching
[M[HS,+ Fokker-Planck regularization Lpsx a Cep
+ -
consistent
Training diffusion models on a 2D toy example reveals inconsistencies. While classical iid diffusion sampling (i.e., denoising) correctly reproduces both modes, evaluating the score at 0 to estimate the unnormalized density yields third mode and an incorrect mass distribution. Such a diffusion model would produce incorrect dynamics while producing correct samples. model trained with Fokker-Plande regularization is self-consistent, and aligns the learned score at = 0 with the distribution recovered by diffusion sampling.
Figure 1:
a stochastic diffusion process in which data progressively perturbed with Gaussian noise over diffusion time i, starting from clean samples at / = O and approaching pure noise at / = 1. Generation then proceeds in reverse: random noise is iteratively denoised using the learned time-dependent score function Vx log p(z, t), producing molecular configurations that approximate the equilibrium training distribution. While such models can efficiently generate independent equilibrium structures (Lewis al., 2025), they lack temporal structure and cannot capture kinetic properties (Wang and Hou, 201 201200111011 la.b) or transition mechanisms (Henkelman et al., 2000; Bolhuis et al., 2002). In contrast, MD yields both, thermodynamic and kinetic properties but has high computational cost.
In principle, the score function at / 0 of diffusion model trained on molecular samples can serve as surrogate for the physical forces. In an ideal energy-based diffusion model, this correspondence would ensure thermodynamic consistency between the learned energy landscape and the equilibrium distribution. In practice, however, difflusion models rarely behave as proper energy-based models, especially as system complexity increases (Arts et al., 2023). In this work, we address these limitations and propose consistent energy-based diffusion model that can generate realistic equilibrium samples while providing physically consistent forces for simulation.
To address this, we leverage the Fokker- Planck equation, which governs the evolution of probability densities under diffusion processes (Särkkä and Solin, 2019), but is generally violated by existing diffusion models (Lai et al., 2023). By introducing a loss term during training that penalizes deviations from this equation, we improve model alignment and empirically enhance consistency between generated samples and simulated dynamics, as shown in Figure 1. In practice, we parameterize the score as the gradient of an energy function, providing access to a potential energy and conservative forces, which is the prerequisite to enable the use of established physical sampling methods (Torrie and Valleau, 1977; Laio and Parrinello, 2002) with coarse-grained diffusion models. Furthermore, by decomposing the diffusion timeline into smaller intervals and assigning a separate model to each, this regularization can be selectively applied to small diffusion timesteps (i.e., where is needed most).
Our main contributions in this work are as follows:
1. We include a Pokker-Planck-bart regularization when training energy- diffusion models, which substantially improves simulation stability and enforces consistency between diffusion sampling and molecular dynamics simulations recovered from the learned forces. This enables accurate sampling and simulation of fast-folding proteins such Chignolin and BBA.
2. We show that training on small sub-interval of the diffusion time suffices for stable simulation. Combining this with smaller models trained on complementary intervals can yield more efficient training and inference while preserving high sampling quality.
3. We develop state-of-the-art transferable Boltzmann emulator for dipeptides capable of fast high-quality independent sampling, and consistent long-time simulation.
Background
2.1 Generative Score-based Modeling
Diffusion models (Song and Ermon, 2019; Ho et al., 2020; Song et al., 2021) are self-supervised generative models that gradually corrupt the training data with noise and learn to reverse this stochastic

---

process. The forward process is defined by a stochastic differential equation (SDE)
dæ f(x 1) dl + g(() dw,
(1)
where denotes the standard Wiener process, and f and g define the drift and diffusion coefficient respectively. To generate samples, diffusion models simulate the corresponding reverse-time SDE
dæ = [f t) - Vt) log pc(x)) dt + gt) dwo
(2)
Starting from Gaussian noise at time = 1, they iteratively denoise until a sample is produced at = O. Here, (x())) denotes the density of the sample
time 1, which the model aims to approximate by learning the score v log Pt ( t). w denotes the time-reversed Wiener process.
As for f and g. the choice depends on the specific diffusion formulation. In this work, we adopt the variance preserving (VP) SDE formulation introduced by Song et al. (2021) (compare Appendix A.
Denoising score matching (Vincent, 2011; Song et al., 2021) is a common way to train diffusion models by minimizing the squared error between a time-dependent learned score function s (zxt), t) and the true score of the transition kernel pos (z(t) (0)) conditioned on the training data x (0):
0" argmin (0,1) (o) (1)) (+))
[A()) || (x) V
(0))] (3)
where X(t) is a time-dependent weighting function. For brevity, we will denote the denoising diffusion loss LDSM s (xx, t) X(t) || (x((), V Va(() log (x)(t)t) (00))|
Parameterization and instabilities. With an affine drift f. we can write the closed-form solution of Por Gaussian (Särkkä and Solin, 2019), and can efficiently train diffusion models such that
A(0) | (v(a()), (1),
l = argmin E..(i) (2) BewN
(4)
(tt)
2
where (0), t), (t) depend on the concrete choices for f and By construction, (0)) 0, which ensures correct interpolation between data and noise. Minimizing the denoising loss yields a neural network with parameters 0 that can approximate the unconditional score s sxx t) x V log (zr).
As t - 0, this parameterization introduces numerical instability, where (t) vanishes and the loss explodes. This instability makes training difficult at small timescales (Kim et al., 2022) and is typically mitigated by truncating the training interval to (e. 1) for some ε >> O. While effective for training stability, this inherent instability in training limits the model's accuracy at small 4, which is critical for applications requiring reliable scores close to the data manifold, as targeted in this work.
2.2 Boltzmann Distribution
Molecular dynamics simulations numerically integrate Newton's 's equations of motion describe the time evolution of molecular system. A widely used formulation is Langevin dynamics (Leimkuhler and Matthews, 2015), which corresponds to integrating the following set of SDEs
dx v dt,
M du - (xx) di MM dt T dwn.
M denotes the partiele masses, v the velocities, y a friction constant, kgT a constant, and w, is the standard Brownian motion. Note that here refers to the physical time instead of the diffusion time used earlier. Integration of this system requires access the forces U. However, in settings where direct force evaluation is not feasible, such as in CG models, a surrogate force function is required. In this work, we propose using the score s (xx, t = 0) for this purpose, as we describe next.
Extracting forces. After a long, equilibrated MD simulation, samples follow the Boltzmann distribu- tion (Boltzmann, 1868) such that p(z)) exp(- knT 1/Z, where U is the potential energy and Z the normalization constant. Training a diffusion model on data from such distribution yields at t = 0
8 (x, 0) V v log (z=)
(6)
U(x)
V - log exp
V. log 7
KBT
V U(x)
0.
kg?
3

---

Hence. the score at diffusion time t = O is proportional to - -v)U(=), the forces acting on the system. This equivalence implies that any diffusion model trained on be used for independent sampling but also for molecular simulation by using the learned score as force estimator with the SDEs of Equation (5). Unlike previous methods that rely on explicit force labels for training (Husic et al., 2020; Durumeric et al., 2023; Charron et al., 2025; Krämer et al., 2023), this relation allows us to learn a model directly from equilibrium samples. This is particularly useful for systems where potential energy and force labels are unavailable, in CG modeling, which is the setting we consider. The central question we will answer next is whether and how this identity can be realized in practice and if the resulting model is sufficiently accurate.
data can not only
3 Method
Diffusion models trained on Boltzmann-distribuce molecular data provide bridge between genera- tive sampling and physical energy landscapes. At the diffusion endpoint t = O, the score represents the forces acting the system (Equation (6)), suggesting that a single model can capture the sampling process and the underlying potential energy. These two perspectives correspond to reconstructing the data distribution through denoising, p(æ=), or evaluating the energy-based form at ( 0, po
p(æ) Denoising(s [s (z)
with S (xx = -V U(v) kBT
(xx)
(7)
(a( exp
KBT
Ideally, these two formulations should agree such that p(x) = Po(x). In practice, however, diffusion models often display inconsistencies (Koehler et al., 2023; Li et al., 2023; Bortoli et al., 2024). Prior attempts to use the score-force relation beyond static sampling can reproduce realistic ensembles under classical diffusion sampling p(xx) but fail to recover these distributions in simulation po (xx) (Arts et al., 2023), indicating that the Iearned model does not behave as a proper energy-based model.
We show that consistency can be recovered directly from unlabeled equilibrium data by enforcing the Pokker-Plancee equation. This equation links energy functions across different diffusion times, effectively transferring accuracy from stable, large-t regions to the small-t regime. The result is self-consistent diffusion model that learns physically meaningful forces without ever observing them.
3.1 Improving Consistency with the Fokker-Planck Equation
The Fokker-Planck equation (Øksendal, 2003; Särkkä and Solin, 2019) is partial differential equation that describes how probability densitics evolve in stochastic processes, including diffusion models. For the diffusion SDE introduced in Equation (I), the log-density formulation of the Fokker-Planck equation (Lai et al., 2023; Hu et al., 2025) can be written as
1 0, log (x) = (a 1) 5 (1)
div (V Pt) + V 0 log pol2
VS.Va log div a (f),
(8)
where div,: denotes the divergence operator div F = ur F). This equation is fundamental property of the diffusion process, and any well-trained model should satisfy Equation (8). However, as shown in Section 5, we observe that diffusion models violate this equation, particularly at small t, which is consistent with findings in prior work (Lai et al., 2023). Since the Fokker-Plande equation links the evolution of the score that of the underlying density, such violations imply that the learned score does not evolve consistently with the density over diffusion time. Consequently, small deviations can accumulate, possibly leading to inconsistencies between the distribution p(x) recovered through denoising and the instantaneous density Po (z) we want to recover.
Fokker-Planck regularization. To correct this, we introduce a regularization term that enforces compliance with Equation (8) during training, similar to Lai et al. (2023). Alongside the standard diffusion objective, we minimize the residual error
[[x t) = (2, t) 0, log log (x)]1
(9)
and define the corresponding loss as Lirllog p0 (2,, = Arr(t)D R(x, f)]l where E D
and XFP is time-dependent weighting function, which we set to be the same as X. As the regular- ization can be evaluated for all æ, t, we combine it with the score-matching objective
argmin (0-1 (o) Ex(()) ()))) [CDsM[ V. log (=x()))()))))))))))))<))) t) + a Cppllog ( (z))),
(10)

---

where a is hyperparameter that determines the regularization strength. The regularization Lep term provides an additional training signal that stabilizes learning the small-time time regime, where the denoising loss becomes ill illccondttooned and gradients are dominated by noise. By explicitly linking the temporal evolution of densities, the regularization encourages the model to remain consistent with the Fokker-Plancc dynamics throughout diffusion time.
Weak residual formulation. The exact computation of the residual Rt (and thus the loss CFP) involves costly higher-order derivatives, especially the divergence term div (V logp") can be challenging to compute for high-dimensional data. To reduce this overhead, we introduce a series of approximations and begin by using residual in the weak formulation similar to the one used by Guo et al. (2022)
R(xr, 1) = Ev R( + v,t)]
(11)
with v ~ _N(0, 12) and a small σ > 0. As a approaches O, R(x t) will be equal to R(x, t). Using the weak residual formulation, R(x,, t) can be estimated by the following unbiased estimator, which only requires the computation of first-order derivatives (see Appendix A.2 for derivation)
1 R(a, 1; v) = (0) (G)
8 (x v,t) - s² x v,t) + =1 + u, 112
(12)
20
(f(x U, v,t), S sx v, f)) div (f + v)t)) O₂ log P. (x u
where 0 Vn logr This yields an unbiased estimator of the squared residual
[i(a..)] 1)| R(x, t: u) R(x,t;v'), |
(xe( t) with v. ~ (0, yielding a computationally manageable estimate of Lrp-
(13)
We further reduce computational cost by estimating 0, log pf using finite differences (Fornberg, 1988)
log p h² log ha) (h3 - h2,) log p. (xx h2 log p. (x. hha) ho)
(14)
hh + ha)
Since both L and log pe are one-dimensional, finite differences serve robust estimate. This allows for computationally feasible approximation the loss LFp, using multiple forward passes.
3.2 Physically Consistent Model Design
To model biochemical systems accurately and support Fokker-Planck regularization, we must enforce known physical constraints through an appropriate parameterization and choice of architecture.
Conservative model parameterization. In diffusion models, it is common to parameterize the score directly as 8 = V v log P. NNET(< t) rather than defining explicit energy function and taking its gradient, 8 V. NN' 1). While energy-based formulations have been explored previously for diffusion models (Song and Ermon, 2019), they are less common in practice (Du et al., 2023), since most applications require only the score and report difference in sampling quality (Salimans and Ho, 2021). In our setting, however, we explicitly aim to construct a consistent energy- based diffusion model, ensuring the correspondence between the learned score and the underlying energy landscape as formalized in Equation (7). This gradient- based formulation provides access to the energy log Pt, which is essential for enforcing physical consistency through the Fokker-Planck equation. Moreover, for molecular simulations, a conservative formulation, where the forces are derived from well-defined energy, is critical for numerical stability and accurate force estimation (Schütt et al., 2017; Batzner et al., 2022; Arts et al., 2023), as further demonstrated in Appendix C.I.
Architecture. Our choice for the score is conservative, translation invariant, and learns SO(3) cquivariance. Similarly to Arts ct al. (2023), we use graph transformer (Shi et al., 2021), making the score permutation equivariant and achieving translation invariance by using pairwise distances instead of absolute coordinates. For the rotation equivariance, recent work (Abramson et al., 2024) has shown that the architecture itself does not need to enforce this property. Hence, apply random rotations during training so that the network learns rotational equívariance via data augmentation.
The main part of the architecture can be summarized by describing the nodes n and edges such that
eij = x Dz
n (0)
ia; t]
n (1+1) = (19) (n) ee),
(15)
5

---

where are the coarse- graniee positions, are the edge features, are atom features, t is the diffusion time, n (1) are the node embeddings of layer I, and & is one layer of the graph transformer. When training on single molecule, we use one-hot atom types; for the transferable model, we use the atom identity, atom number, residue index, and amino acid type, following Klein and Noé (2024).
Finally, to achieve conservativeness, we map the last node embeddings n!( E R to scalar energies via vi R + I,, and compute the score as V Σ, (nn Overall, this yields a translation- invariant, approximately iitation-qqivarinnn conservative architecture that also avoids issues caused by mirror symmetrics (Trippe et al., 2023; Klein and Noć, 2024).
3.3 Mixture of Experts
Classical diffusion models are trained across the full diffusion timeline to enable generative sampling. In some applications, learning the score only potential. To support this, we introduce time- time-based mixture of experts (MoE) approach, inspired by ideas from the image domain to improve efficiency (Balaji et al., 2023; Ganjdanesh et al., 2025). Rather than training a single model across the entire range / E (0, 1). we partition the interval into disjoint subintervals II..I... with U, I, = (0, 1), and assign a separate expert 8f to each interval. The overall score can then be written by evaluating the correct model
can be sufficient, as it gives access to the
s sx t) = sf (x, for E.
(16)
Thus, only one expert is active at any given t, which simplifies memory management and allows independent and parallel training of all experts. For simulation, only the expert corresponding to t= 0 I is needed, while iterative diffusion sampling sequentially loads every expert as t decreases.
This design provides several advantages. For simulation, training only on the relevant small-f range (e.g., (0, 0.1)) avoids wasting resources on unnecessary diffusion times that are never used. For generative sampling, dividing the timeline still offers clear benefits. Applying the Fokker-Planck loss from Equation (10) across all diffusion times can lead to overregularization at large t, degrading iid sampling performance. Since large-t models do not require Fokker-Planck regularization or conservative parameterization, using simpler unconstrained models improves generative quality and prevents unnecessary constraints. Moreover, the experts specialize in distinct temporal regions: small-t experts capture fine structural details, while large-t experts focus coarse features (Ganj- danesh et al., 2025). This specialization stabilizes training and enhances the consistency between sampling and simulation beyond what could be achieved by simply increasing model capacity (see Appendix C.2). Finally, because only one expert is loaded into memory at a time, we can overall train more parameters without exceeding GPU limits. We will explore these benefits in Section 5.
4 Related Work
In recent years, a variety of deep learning methods have been proposed to enhance or replace molecular simulation. The work most closely related to ours is that of Arts et al. (2023), who employ an energy-based diffusion model for coarse-grained systems. However, their approach fails to maintain consistency between sampling and the learned score. They mitigate some of the inconsistencies we describe by evaluating the diffusion model at larger timesteps & > 0. which introduces additional noise and reduces structural fidelity. We compare against this model in Section 5. Raja et al. (2025) also use the score-force relation but for transition path sampling and could benefit from our approach to have more accurate forces. Daigavane et al. (2025) use the inherent noise of the diffusion process to converge to an equilibrium distribution faster, but they do not learn accurate scores either. Several works instead learn coarse-grained force fields via a force-matchin (Husic et al., 2020; Köhler et al., 2023; Charron et al., 2025; Durumeric et al., 2024). Rather than training a model to represent the data distribution, these methods approximate the target forces. However, they typically need system-specific energy priors, relying heavily on domain knowledge.
objective
Other approaches bypass MD sampling entirely, generating Boltzmann distrrbuted configurations either sequentially, by conditioning each sample on its predecessor (Dibak et al., 2022; Plainer et al., 2023a; Schreiner et al., 2023; Tamagnone et al., 2024), or completely independently (Noé et al., 2019; Wirnsberger et al., 2020; Köhler et al., 2020; Midgley et al., 2023; Klein et al., 2023b; Abdin and Kim, 2024; Kim et al., 2024; Wu and Noé 2024; Schebek et al., 2024; Diez et al., 2024; Tan et al., 2025). These methods frequently leverage diffusion- based architectures. In the latter

---

case, for all-atom systems with known energy functions, they can guarantee asympiotically unbiased sampling via reweighting or integration into an MCMC scheme. Although this is often a desirable property, extending it to larger systems remains challenging, as CG is not directly possible.
The inaccurate behavior of the score function has also been studied in low-dimensional settings by Koehler et al. (2023); Li et al. (2023), who demonstrated inconsistencies and derived error bounds. Bortoli al. (2024) propose to improve the score at small diffusion timesteps, when the probability distribution of the samples is known. However, this is not directly compatible with coarse-graining, as the potential and the derived forces are not known. Similarly to us, Lai et al. (2023) also proposed a regularization; however, their method applies a higher-order regularization to the score itself improve iid sample quality, rather than enforcing consistency through the potential to improve the learned score. Daras et al. (2023) propose a regularization to ensure average consistency of the score, wheras our objective aims to ensures concistency of the energy over complete trajectories through the Fokker-Planck equation. Relatedly, Hu et al. (2025) propose a score-based solver for high-dimensional Fokker- Planck equations, focusing on solving general SDE forward problems, and Du et al. (2024) use the Bokker-Plancke equation to describe MD as series of Gaussians that can be integrated but do not recover the full unconditional Boltzmann distribution.
seirr-PPanc--insss
5 Experiments
In this section, we evaluate the consistency of diffusion models by comparing samples obtained through classical denoising (iid) with those generated by sequential Langevin simulations (sim) using forees derived with the score-force relation from Equation (6). In other words, we verify the align- ment of Equation (7) for energy-based models and show superior performance. We demonstrate our approach (ScoreMD) on three biomolecular systems- alanine dipeptide, Chignolin, and BBA-and introduce transferable model that generalizes across dipeptides, improving over existing state-of-the- art Boltzmann generators (Klein and Noé, 2024). The code, model weights, and self-contained note- books in JAX and Py Torch are publicly available at https: : / //g gitggtt....... com www.o//nogoup///ooooooooooo Sww.grrru//S//ree0/ /ScorMDD
Metries. As the main metric of interest, we compare the 2D free energy surfaces of the equilibrium distributions obtained by different methods. For dipeptides, we project the data onto the dihedral angles 9 and v) while for proteins we perform time-lagged independent component analysis (TICA) (Pérez-Hernández et al., 2013) on bond distances and dihedral angles to recover two representative coordinates. To quantify differences between free energy surfaces, we report the potential of mean force (PMF) error (Durumeric et al., 2024), which measures the squared distance between the negative logarithms of the sampled and reference densities in the projected space. This metric places higher weight low-density regions compared to alternatives such as the Jensen-Shannon (IS) divergence. Additional details are provided in Appendix B.2.
Baselines. We train a conservative Diffusion model and re-implement the Two For One method (Arts et al., 2023) within the same continuous- time diffusion framework to ensure fair comparison. Both models use identical architectures and training procedures and hence produce the same iid samples (up to numerical errors); their only difference lies in the diffusion time used during evaluation. Specifically. Two For One performs simulation at a finite diffusion time i > 0 to enhance stability. whereas Diffusion evaluated a = 10-5, For transferability, we re-train the Transferable Boltzmann generator (BG) model (Klein and Noé, 2024) with coarse- grariiigg and use it without reweighting.
5.1 Alanine Dipeptide and Fast-Fo Proteins
Dataset. For alanine dipeptide, we use 50k samples from an MD simulation in implicit solvent (Köhler et al., 2021), coarse-grained to five atoms: [C, N. CA, C, NJ. For Chignolin and BBA, we use the dataset from Lindorff-Larsen et al. (2011), coarse-grainee to one bead per amino acid (10 and 28 residues, respectively), and use 80% of the samples for training. For the iid setting, we generate the same number of samples as in the training set. For sim, we initialize 100 parallel simulations from random training conformations, and simulate with timestep of 2 fs for total length of 1.2 us for alanine dipeptide, 855.6 ns for Chignolin, and 200 ns for BBA. All trajectories are downsampled to match the number of training samples for consistency with the iid setting.
(In-) non)-consittet sampling. Figure 2 compares the free energy landscapes in the projected space for iid sampling and Langevin simulation (sim). While all methods reproduce the training distribution under iid sampling, simulation quality varies and reveals inconsistencies. For the small system alanine dipeptide, the equilibrium distributions appear similar across methods. However, as system size
7

---

Reference
Diffusion iid Sampling
+
Two For One Simulation
Ours
Ours Simulation
iid Sampling
Alanine Dipeptide
Chignolin
me.
nic
TC!
mc
BBA
TIC
T
TCC
TIC
Figure Comparison of equilibrium distributions obtained by iid sampling and Langevin simulation (sim) across different systems and methods. While classical iid sampling recovers the reference equilibrium distribution, performing simulation with the learned score reveals inconsistencies when models are not trained with Fokker- Planck regularization, i... p(x) # (2)). Regularized models achieve consistent behavior across systems.
Method
ALDP
Chignolin
BBA
Teo for Ona
PMF ())
PMF ())
PMF (L)
Diffusion
0.098 + 0.006
0,027 = 0.000
0,034 = 0.000
Ours
0.097 0.008
0.035 0.001
0.234 0.001
Two For One Ours
sim
0.206
0.004
1.438
= 0.019
1.624
0.107
0.091 0.004
0.038 0.006
0.254 0.005
10
12
Table 1: Quantitative comparison of sampling and simulation across different methods for alanine dipeptide (ALDP), Chignolin, and BBA. Ours yields significantly more consistent results between iid and sim.
Figure 3: Density of C-N bond
length in Á for simulation of alanine dipeptide using different models.
increases, artifacts become more pronounced: for Chignolin, the Two For One simulation exhibits strong noise, and for BBA, the system separates into unphysical modes, yielding an incorrect equi- librium distribution. In contrast, our Fokker-Plande regularization enforces consistency, producing matching ensembles for iid sampling and Langevin simulation (sim), i.e., p(ar) 2 po(a)).
These trends can be seen in Table 1, which shows that Ours substantially improves simulation quality. We also observe modest reduction in iid sampling performance, reflecting tradeoff between generative fidelity and consistency, goverened by the regularization strength a. Larger values promote stronger alignment between sampling and simulation but may slightly degrade iid sampling accuracy.
Noisy simulation. Two For One improves consistency by evaluating the model at a larger diffusion time t >> during simulation, which introduces substantial noise that degrades structural fidelity. While iid sampling remains unaffected, simulations exhibit pronounced deviations, as seen for Chignolin in Figure 2. For alanine dipeptide, the free energy surface obtained with Two For One appears reasonable, but because the method relies on "noisy" forces, the observables are also affected by this noise. For instance, bond-length distributions will be incorrect, as depicted in Figure 3.
Dynamics. A key advantage of consistent energy-based diffusion model is that provides access not only to generative diffusion sampling but also to physically meaningful energy function usable for downstream tasks such as MD simulation. As such, we can recover kinetic information and temporal behavior, as studied in Figure 4 for BBA, which is not accessible with classical diffusion sampling. Additional analysis of the dynamics, such as transition probabilities between states and the corresponding results for Chignolin, is provided in Appendix C.6.4.
Force matching. CG models are typically trained using force matching rather than diffusion-based objectives. In Appendix C.5.4, we benchmark our method against recent force-matching approaches for alanine dipeptide. Although these baselines often rely on explicit physical priors and force labels, our fiffusion-based model achieves ensembles that more closely match the reference distribution.
8

---

(a)
(b)
(c)
TCC
TAC 1
30000 2000
30000
4000
5000
6000
(Coarse-Grained) Timesteps
Figure 4: Recovering and analyzing the dynamic behavior of BBA. (a) 3D molccular structures from a single MD trajectory generated by our model. (b) Trajectory projected onto the first two TIC coordinates. (c) Illustrating the TIC coordinates over the coarse-grained (and subsampled) timesteps.
5.2 Transferability Across Dipeptides (Two Amino Acids)
Dataset. We use the dataset introduced by Klein et al. (2023a), which contains 49k samples from implicit-solvent simulations of 1 jus covering all possible combinations of the 20 standard amino acids. From these, we use 175 dipeptides for training. Each dipeptide is coarse-grained by retaining the atoms [N, CA, CB, C, O] for each amino acid, a common CG resolution (Charron et al., 2025). With this, up 10 atoms draw 49k independent samples, while for sim we perform 30 ns of Langevin simulation per dipeptide, initialized from 10 random conformations with timestep of 0.5 fs downsampled to 49k frames.
retained per molecule, as seen in Figure 5 (a). For iid sampling,
Models. We ablate the individual contributions of the Pokker-Planck regularization, the impact of MoE, and the effect of evaluating models at larger diffusion timesteps t >> 0. Similar ablation studies for previous results are provided in Appendices C.5 and C.6. Specifically, we compare the Diffusion and Two For One baselines with three variants: Mixture denotes the MoE scheme, consisting of three experts trained on the intervals (0, 0.1). [0.1, (0.1,0.6) 0.6), and [0.6, 1.0). The models were robust to moderate changes in these ranges, provided the smallest-time expert was trained on sufficiently large subinterval. All three experts are combined for iid sampling, while only the smallest- model is used for simulation. The experts corresponding to larger timescales are reduced in size and complexity. Pokker-Plancee refers to a single model trained with the loss from Equation (10). Both combines these two approaches, where the regularization is applied only to the smallest-timescale expert within the MoE setup. This corresponds to the model referred to previously as Ours.
Overdispersion. As shown in Figure 5 (b), all models can generate independent samples that closely match the reference distribution. However, samples from the Transferable BG model exhibit increased noise and over-explore wwwprobabillity states. In the original work, this behavior is corrected through sample reweighting, which we omit here and instead treat the model simply as a Boltzmann emulator. While the resulting differences are within standard deviations, the model yields higher mean with comparable variance (see Table 2). Morcover, note that this model does not support simulation. When using the score for simulation, Two For One produces broader, overdispersed distributions, as shown in Figure 5 (b). This overdispersion also propagates to structural features such as inter-atomic distances, consistent with the behavior observed in previous systems (see Appendix C.8.2). More metrics and evaluation for more dipeptides can be found in Appendix C.8.
Ad Advantages of mixture. The MoE architecture can significantly improve simulation quality compared to Diffusion. For the specific dipeptide shown in Figure (c), Both further enhances performance over Fokker-Planck, improving both iid and simulated distributions. While the degree of improvement varies across dipeptides-winn with some cases showing comparable results (see Appendix C.8.3)-MoF generally increases simulation stability, reflected in the lower PMF errors reported in Table 2.
Since Mixture and Both contain more parameters (see Appendix B.5), part of the improvement can be attributed to increased capacity. However, MoE also offers clear computational advantages: by applying regularization only at small diffusion times and using smaller, unconstrained models for larger timescales, it reduces training and inference cost, cutting sampling time by more than 50% in this case (see Appendix C.3) and prevents overregularization at large diffusion timesteps. Moreover, additional ablation studies in Appendix C.2 show that the benefits of MoE extend beyond parameter scaling, demonstrating that training on smaller temporal subranges improves robustness and that independently trained experts can be seamlessly combined into coherent model.

---

(a)
(b)
Refromnc
Trazaferabie BO Gid
For One (ai))
Misture (idl
Friker Pland (id
Beth rid)
Cousse groined
(c)
Jeen Eoergy Projection (sim)
Reference
Difiasion
Mistre (aim)
Flankk-Planck (sim)
Both Gam)
Figure 5: Comparison of methods on testset dipeptide KS.. (a) The coarse graining scheme. (b) Comparison of the Ramachandran plots of different methods for iid sampling and Langevin simulation. (c) The projection of the free energy surface and differences along the dihedral angle yp for samples generated with simulation.
Method
iid PMF (1)
sim. PMF (L)
Transferable BG
0.230
0.119
Diffusion
0.206
0.159
6,515 1 3,175
Two For One
0.203
0.149
0.741
+ 0.319
Mixture
0,200
0.127
0.658
1 0.407
Fokker-Planck
0.241
0.105
0.368 0.267
Both
0.199 0.127
0.203 0.104
2.6
2.6
38
Table 2: Comparison of metrics across all testset dipeptides with PMF error. compute mean and standard deviation, we have averaged the results across all evaluated dipeptides. Lower values are better.
Figure 6: Comparing the Fokker- Planck error for log of multiple models.
Fokker-Planck error. Figure shows the deviation from the Fokker-Planck equation, quantified as [z. 0 log p (z) I1 plotted on a log scale. For models using MoE, this error is evaluated only up to f = 0.1, since only the small-timescall model is energy- energy-based. Across all methods, the error largest near = O. Applying the Pokker-Planck regularization significantly reduces this error, correlating with the improved sampling-simulation simulation consistency observed earlier.
Interestingly, while Mixture improves consistency, its Fokker-Planck error at 0 remains com- parable to that of unregularized models. This suggests that Bokker-Pllanck regularization and MoE improve consistency through different mechanisms, which explains why combining them outperforms either approach on its own, making Both again clearly the best model (compare Table 2).
6 Conclusion, Limitations and Future Work
In this work, we investigated energy-based diffusion models and analyzed the discrepancy between the density recovered through denoising and the learned energy at diffusion time = 0. We showed that diffusion models are generally inconsistent, and these two densities do not coincide. To address this, we introduced Nokker-Planckk-aseee regularization on the model's energy, demonstrating that reducing deviations from the Fokker-Planck equation substantially improves model consistency. Furthermore, we found that focusing training on smaller diffusion times further enhances simulation quality. With these improvements, we obtain a physically consistent energy-based model that can generate independent equilibrium samples and recover realistic molecular dynamics from the same learned potential, providing access to both static and kinetic properties. This is particularly useful in coarse-grained settings where direct force information not available. We validated these findings across multiple systems, trained on fast-folding proteins, and generalized across dipeptides.
Despite the theoretical motivation behind our approach, the results presented are primarily empirical. While our findings indicate that reducing the Fokker-Planck deviation improves consistency, this unlikely be the only source of simulation error. In fact, due to the fundamental differences between diffusion sampling and Langevin simulation, perfect aligmment may not be achievable without limiting model expressivity. Moreover, evaluating the Tokker-Planch residual introduces computational overhead, which we mitigate through a weak residual formulation, though it still requires multiple forward passes during training and thus remains computationally demanding. Future work could extend this framework to larger molecular systems and explore transfer learning across biomolecular families. Another promising direction is to fine-tune existing diffusion models with the proposed regularization to explicitly correct Pokker-Planck deviations and enable simulation.
10

---

Acknowledgments and Disclosure of Funding
The authors would like to thank Aleksander Durumerie, Yaoyi Chen, Maximilian Schebek, Klara Bonneau, Winfried Ripken, Marcel Kollovieh, Nicholas Gao, Matej Mezera, Yuanqi Du, Jiajun He, and Jungyoon Lee for the fruitful discussions and their helpful input. The work of Michael Plainer was supported by the Konrad Zuse School of Excellence in Learning and Intelligent Systems (ELIZA) through the DAAD program Konrad Zuse Schools of Excellence in Artificial Intelligence, sponsored by the Federal Ministry of Education and Research. The work of Hao Wu was supported by the National Natural Science Foundation of China, Grant No. 12171367. Moreover, we gratefully acknowledge support by the Deutsche Forschungsgemeinnchaaft (SFB 1 14, Projects No. A04 and No. B08) and the Berlin Mathematics center MATH+ (AAI-10). We gratefully acknowledge the computing time made available to them the reghpeefformanc computer "Lise" at the NHR Center NHR @ ZIB. This center is jointly supported by the Federal Ministry of Education and Research and the state governments participating in the NHR (www phr verein. deuusser-- parerrr-artter)..

---

References
Batzner, S., Musaelian, A., Sun, L.,, Geiger, M., Mailoa, J. P., Kornbluth, M., Molinari, N., Smidt, Abdin, O. and Kim, P. M. (2024). Direct conformational sampling from peptide energy landscapes diffusion. Nature Machine Intelligence, 6(7):775-786.
through hypernetoor-ccoddiionn
Abramson, J., Adler, J., Dunger, J.. Evans, R., Green, T., Pritzel, A., Ronneberger, O., Willmore, L., Ballard, A. J., Bambrick, J., Bodenstein, S. W., Evans, D. A.. Hung, C.-C., O' O'Neill, M., Reiman, D.. Tunyasuvunakool, K., Wu, Z., Zemgulyté, A., Arvaniti, E.,, Beattie, C., Bertolli. O., Bridgland, A., Cherepanov, A., Congreve, M., Cowen-Rivers A. I., Cowic, A., Figurnov, M., Fuchs, F. B., Gladman, H., Jain, R., Khan, Y. A., Low, C. M. R., Perlin, K., Potapenko, A., Savy, P., Singh, S., Stecula, A., Thillaisundaram, A., Tong. C., Yakneen, S., Zhong, E. D., Zielinski, M., Zídek, A., Bapst, V., Kohli, P., Jaderberg, M., Hassabis, D., and Jumper, J. M. (2024). Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, 630(8016):493-55)))
Albergo, M. S. and Vanden-Eijnden E. (2025). NETS: A non-equilibrium transport sampler. In Singh, A., Fazel, M., Hsu, D., Lacoste- S., Berkenkamp, F, Maharaj. T., Wagstaff, K., and Zhu, J., editors, Proceedings of the 42nd International Conference on Machine Learning, volume 267 of Proceedings of Machine Learning Research, pages 1026-1055 PMLR.
Arts, M., Garcia Satorras, V., Huang, C.-W., Zügner, D., Federici, M., Clementi, C., Noé, F., Pinsler, R., and van den Berg, R. (2023). Two for one: Diffusion models and force fields for coarse-grained molecular dynamies. Journal of Chemical Theory and Computation, 19(18):6151-6)
Balaji, Y., Nah, S., Huang, X., Vahdat, A., Song, J., Zhang, Q., Kreis, K., Aittala, M., Aila, T., Laine, S., Catanzaro, B., Karras, T., and Liu, M.-Y. (2023). ediff-i: Text-to-image diffusion models with Stii::2211001222.
an ensemble of expert denoisers.
T. E., and Kozinsky, B. (2022). E (3) (3)equivvriiant graph neural networks for data-efficient and accurate interatomic potentials. Nature communications, 13(1):2453.
Bolhuis, P. G., Chandler, D., Dellago, C., and Geissler, P. L. (2002). Transition path sampling: Throwing ropes over rough mountain passes, in the dark. Annual review of physical chemistry, 53(():291-318.
Boltzmann, L. (1868). Studien üùer das Gleichgewicht der lebendigen Kraft zwischen bewegten materiellen Punkten: vorgelegt in der Sitzung am 8. October 1868. k. und k. Hof- und Staatsdr.
Bortoli, V. D., Hutchinson, M., Wirnsberger, P., and Doucet, A. (2024). Target score matching. arXiv:2402.08667.
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman- S., and Zhang, Q. (2018). JAX: composable transformations of Nuthonn+uumPy programs. http:/github..om/jiat
Charron, N. E., Bonneau, K., Pasos Trejo, A. S., Guljas, A., Chen, Y., Musil, F., Venturin, J., Gusew, D., Zaporozhets, I., Krämer, A., Templeton, C., Kelkar, A., Durumeric, A. E. P., Olsson, S., Pérez, A., Majewski, M., Husic, B. E., Patel, A., De Fabritiis, G., Noé, F., and Clementi, C. (2025). Navigating protein landscapes with machine-learned transferable coarse-grained model. Nature chemistry, pages 1-9.
Clementi, C. (2008). Coarse-grained models of protein folding: toy models predictive tools? Current opinion in structural biology, 18(1):10-15.
Corso, G., Stärk, H., Jing, B., Barzilay, R., and Jaakkola, T. (2023). DiffDook: and turns for molecular docking, In International Conference on Learning Representations.
Diffusion steps, twists,
Daigavane, A., Vani, B. P., Davidson, D., Saremi, S., Rackers, J. A., and Kleinhenz, J. (2025). Jamun: Bridging smoothed molecular dynamics and score-based learning for conformational ensemble generation. In ICML 2025 Generative Al and Biology (GenBio) Workshop.
Daras, G., Dagan, Y., Dimakis, A., and Daskalakis, C. (2023). Consistent diffusion models: Mitigating sampling drift by learning to be consistent. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S., editors, Advances in Neural Information Processing Systems, volume 36, pages 42038 42063. Curran Associates, Inc
12

---

Dibak, M., Klein, L.. Krämer, A., and Noé, F. (2022). Temperature steerable flows and Boltzmann generators. Phys. Rev. Res., 4:L042005.
Diez, J. V., Atance, S. R., Engkvist, O., and Olsson, S. (2024). Generation of conformational ensembles of small molecules via surrogate model-assisted molecular dynamics. Machine Learning: Science and Technology, 5(2):025010.
Du, Y., Durkan, C.. Strudel, R., Tenenbaum, J. B., Dieleman, S., Fergus, R., Sohl-Dickstein J.. Doucet, A., and Grathwohl, W. (2023). Reduce, reuse, recycle: Compositional generation with energy-based diffusion models and MCMC. In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J., editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 8489-8510, PMLR.
Du, Y., Plainer, M., Brekelmans, R., Duan, C., Noé, F., Gomes, C. P., Aspuru-Guzik, A., and Neklyudov, K. (2024). Doob's lagrangian: A sample-efficient variational approach to transition path sampling. In Globerson, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U., Tomezak, J., and Zhang, C., editors, Advances in Neural Information Processing Systems, volume 37, pages 65791-65822. Curran Associates, Inc.
Durumeric, A. E., Charron, N. E.. Templeton. C., Musil, F., Bonneau, K., Pasos-Trejo, A. S., Chen, Y.. Kelkar, A., Noé, F,, and Clementi, C. (2023). Machine learned coarse-grained protein force-fields Are we there yet? Current opinion structural biology, 79:102533.
Durumeric, A. E. P., Chen, Y., Noé, F., and Clementi, C. (2024). Learning data efficient coarse-grained molecular dynamics from forces and noise. arXiv:2407.01286
Eastman, P., Swails, J., Chodera, J. D., McGibbon, R. T., Zhao, Y., Beauchamp, K. A., Wang, L.-P., Simmonett, A. C., Harrigan, M. P., Stern, C. D., Wiewiora, R. P., Brooks, B. R., and Pande, V, S. (2017). OpenMM 7: Rapid development of high performance algorithms for molecular dynamics. PLOS Computational Biology, 13(7):e
Fornberg, B. (1988). Generation of finite difference formulas on arbitrarily spaced grids. Mathematics of Computation, 51(184):699-706.
Ganjdanesh, A., Kang, Y., Liu, Y., Zhang, R., Lin, Z., and Huang, H. (2025). Mixture efficient diffusion experts through automatic interval and sub-network selection. In Leonardis, A., Ricci, E., Roth, S., Russakovsky, O., Sattler, T., and Varol, G., editors, Computer Vision - ECCV 2024, pages 54-71, Cham. Springer Nature Switzerland.
Guo, I.,, Wu, H., and Zhou, T. (2022). Normalizing field flows: Solving forward and inverse stochastic differential equations using physics-informed flow models. Journal of Computational Physics, 461:111202.
Heek, J., Levskaya, A., Oliver, A., Ritter, M., Rondepicrre, B., Steiner, A., and van Zec, M. (2024). Flax: A neural network library and ecosystem for JAX. http://gtthub.oomgggglleaa
Henkelman, G., Uberuaga, B. P., and Jónsson, H. (2000). A climbing image nudged elastic band method for finding saddle points and minimum energy paths. The Journal of chemical physies, (13(22):9901-9900
Ho, J., Jain, A., and Abbeel, (2020). Denoising diffuuion probabilistic models. Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H., editors, Advances in Neural Information Processing Systems, volume 33, pages
Curran Associates, Inc.
Hoffmann, M., Scherer, M. K., Hempel, T., Mardt, A., de Silva, B., Husic, B. E., Klus, S., Wu, H., Kutz, J. N., Brunton, S., and Noé, F. (2021). Deeptime: a python library for machine learning dynamical models from time series data. Machine Learning: Science and Technology.
Hu, Z., Zhang, Z., Karniadakis, G. E., and Kawaguchi, K. (2025). Score-based physics-informed neural networks for high-dimensional fokker-planck equations. SIAM Journal on Scientific Computing, C680
13

---

Husic, B. E., Charron, N. E... Lemm. D., Wang, J., Pérez, A., Majewski, M.. Krämer, A., Chen, Y., Olsson, S., de Fabritis, G., Nocé, F., and Clementi, (2020). Coarse graining molecular dynamics with graph neural networks. The Journal of Chemical Physics, 153(19):19101.
Kim, D., Shin, S., Song, K., Kang, W., and Moon, L.-C. (2022). Soft truncation: A universal training technique of score-based diffusion model for high precision score estimation. In Chaudhuri, K., Jegelka, S., Song, I.., Szepesvari, C., Niu, G., and Sabato, S., editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 11201-11228, PMLR.
Kim, J. C.,, Bloore, D., Kapoor, K., Feng, J., Hao, M.-H., and Wang, M. (2024). Scalable normalizing flows enable boltzmann generators for macromolecules. International Conference on Learning Representations (ICLR).
Klein, L.,, Foong, A., Fjelde, T., Mlodozeniec, B., Brockschmidt, M., Nowozin, S., Noe, F., and Tomioka, R. (2023a). Timewarp: Transferable acceleration of molecular dynamics by learning time- coarsened dynamics. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S., editors, Advances in Neural Information Processing Systems, volume 36, pages 52863-52883. Curran Associates, Inc.
Klein, L., Kelkar, A., Durumeric, A., Chen, Y., Clementi, C., and Noé, F. (2025). Operator forces for coarse-grained molecular dynamics. The Journal of Chemical Physics, 163(10).
Klein, L., Krämer, A., and Noć, F. (2023b). Equivariant flow matching. Advances in Neural Information Processing Systems, 36:59886-59910.
Klein, L. and Noé, F. (2024). Transferable Boltzmann generators, In Globerson, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U., Tomezak, J., and Zhang, C., editors, Advances in Neural Information Processing Systems, volume 37, pages 45281-45314. Curran Associates, Inc.
Kmiecik, S., Gront, D., Kolinski, M., Wieteska, L., Dawid, A. E., and Kolinski, A. (2016). Coarse- grained protein models and their applications. Chemical reviews, 1 116(14)7798--796.
Koehler, F., Heckett, A., and Risteski, A. (2023). Statistical efficiency of score matching: The view from isoperimetry. In International Conference on Learning Representations.
Köhler, J., Chen, Y., Krämer, A., Clementi, C., and Noé, F. (2023). Flow-matching: Efficient coarse- graining of molecular dynamics without forces. Journal of Chemical Theory and Computation, 19(3):942-92..
Köhler, J., Klein, L., and Noé, F. (2020). Equivariant Flows: Exact Likelihood Generative Learning for Symmetric Densities. In Proceedings of the 37th International Conference on Muchine Leurning.
Köhler, J., Krämer, A., and Noé, F. (2021). Smooth normalizing flows. In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W., editors, Advances in Neural Information Processing Systems, volume 34, pages 2796-2809, Curran Associates, Inc.
Krämer, A., Durumerie, A. E., Charron, N. E., Chen, Y., Clementi, C., and Noé, (2023). Statistically optimal force aggregation for coarse-graining molecular dynamics. The Journal of Physical Chemistry Letters, 17777)3970-3979.
Lai, C.-H., Takida, Y., Murata, N., Uesaka, T., Mitsufuji, Y., and Ermon, S. (2023). FP-Diffusion: Improving score-based diffusion models by enforcing the underlying score Bokker-Planck equation. In Krause, A., Brunskill, F., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J., editors, Proceed- ings of the 40th International Conference on Machine Learning. volume 202 of Proceedings of Machine Learning Research, pages 18365-18398 PMLR.
Laio, A. and Parrinello, M. (2002). Escaping free-energy minima. Proceedings of the national academy of sciences,
Leimkuhler, and Matthews, C. (2015). Molecular Dynamics: With Deterministic and Stochastic Numerical Methods. Springer International Publishing.
14

---

Lewis, S., Hempel, T., JiménezzLuna, J., Gastegger, M., Xie, Y., Foong, A. Y.. Satorras, V. G.. Abdin, O., Veeling, B. S., Zaporozhets, I., Chen, Y., Yang. S., Foster, E., Schneuing, A., Nigam, J., Barbero, F., Vincent, S., Campbell, A., Yim, J., Lienen, M., Shi, Y., Zheng, S., Schulz, H., Munir, U., Sordillo, R., Tomioka, R., Clementi, C., and Noé, (2025). Scalable emulation of protein equilibrium ensembles with generative deep learning. Science, page eadv9817.
Li, P., Li, Z.. Zhang, H., and Bian, J. (2023). On the generalization properties of diffusion models. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S., editors, Advances in Neural Information Processing Systems, volume 36, pages 2097-2127. Curran Associates, Inc.
Lindorff-Larsen K., Piana, S., Dror, R. O., and Shaw, D. E. (2011). How fast-folding proteins fold. Science, 334(6055):517-5200
Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., and Le, M. (2023). Flow matching for generative modeling. In International Conference on Learning Representations (ICLR).
Lipman, Y., Havasi, M., Holderricth, P., Shaul, N., Le, M., Karrer, B., Chen, R. T. Q.. Lopez- D., Ben-Hamu, H., and Gat, I. (2024). Flow matching guide and code. arXiv:2412.06264.
Loshchilov, I. and Hutter, F. (2019). Decoupled weight decay regularization. In Energy Based Models Workshop.
McGibbon, R. T., Beauchamp. K. A.. Harrigan, M. P., Klein, C., Swails, J. M., Hernández, C. X., Schwantes, C. R., Wang, L.-P., Lane, T. J., and Pande, V. S. (2015). Mdtraj: A modern open library for the analysis of molecular dynamics trajectories. Biophysical Journal, 109(8):1528 1532.
Midgley, L. I., Stimper, V., Simm, G. N. C., Schölkopf. B., and Hemándezz-Lbato,, J. M. (2023). Flow annealed importance sampling bootstrap. In The Eleventh International Conference on Learning Representutions.
Müller, K. and Brown, L. D. (1979). Location of saddle points and minimum energy paths by constrained simplex optimization procedure. Theoretica Chimica Acta, 53(1):75-92
Noé, F., Olsson, S., Köhler, J., and Wu, H. (2019). Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(6457))eaww1477
Noid, W. (2013). Perspective: Coarse-grained models for biomolecular systems. The Journal of chemical physics, 139(9).
Øksendal, B. (2003). Stochastic Differential Equations, pages 65-84. Springer Berlin Heidelberg.
Pérez-Hermnndez, G., Paul, F., Giorgino, T., De Fabritiis, G., and Noé, F. (2013). Identification of slow molecular order parameters for markov model construction. The Journal of chemical physics, 139(1).
Plainer, M., Stärk, H., Bunne, C., and Günnemann, (2023a). Transition path sampling with boltzmann generator-based meme moves. In Generative Al and Biology Workshop.
Plainer, M., Toth, M., Dobers, S., Stärk, H., Corso, G., Marquet, C., and Barzilay, R. (2023b). Diffocc--Pekett Diffusion for peeket-level docking with sidechain flexibility. In Machine Learning in Structural Biology.
Plattner, N., Doerr, S., De Fabritiis, G., and Noé. F. (2017). Complete protein protein association kinetics in atomic detail revealed by molecular dynamics simulations and markov modelling. Nature chemistry, 1000):00051100
Raja, S., Sipka, M., Psenka, M., Kreiman, T., Pavelka, M., and Krishnapriyan, A. S. (2025). Action- minimization meets generative modeling: Efficient transition path sampling with the onsager- machlup functional. In Singh, A., Fazel, M.. Hsu, D., Lacoste-Julien. S., Berkenkamp. F., Maharaj. T.. Wagstaff, K., and Zhu, J., editors, Proceedings of the 42nd International Conference on Machine Learning, volume 267 of Proceedings of Machine Learning Research, pages 50972-51000 51008. PMLR.
Salimans, and Ho, (2021). Should EBMs model the energy or the score? In Energy Based Models Workshop.
15

---

Särkkä, S. and Solin, A. (2019). Applied Stochastic Differential Equations. Cambridge University Press.
Schebek, M., Invernizzi, M., Noé, F., and Rogal, (2024). Efficient mapping of phase diagrams with conditional boltzmann generators. Machine Learning: Science and Technology.
Schreiner, M., Winther, O., and Olsson, S. (2023). Implicit transfer operator learning: Multiple time- resolution models for molecular dynamics. In seirth- Conference Neural Information Processing Systems.
Schüit,, K., Kindermans, P.-.., Sauceda Felix, H. E., Chmiela, S., Tkatchenko, A., and Müller, K.-R. (2017). Schnet: A continuous-fiteer convolutional neural network for modeling quantum interactions. In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R., editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc.
Shi, Y., Huang, Z., Feng, S., Zhong. H.,, Wang, W., and Sun, Y. (2021). Masked label prediction: Unified message passing model for semi-supervised classification. In Zhou, Z.-H., editor, Pro- ceedings of the Thirtieth International Joint Conference on Artificial Intelligence, IJCAI-21, pages 1548-1554 International Joint Conferences on Artificial Intelligence Organization.
Song, Y. and Ermon, S. (2019). Generative modeling by estimating gradients of the data distribution. In Wallach, H., Larochelle, H., Beygelzimer, A., d' d'Alché-Bit -ucc, F., Fox, E., and Garnett, R., editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. (2021). Score-based generative modeling through stochastic differential equations. In International Conference Learning Representations.
Tamagnone, S., Laio, A., and Gabrié, M. (2024). Coarse-grained molecular dynamics with normaliz- ing flows. Journal of Chemical Theory and Computation, 20(18);7796-7805.
Tan, C. B., Bose, A. J., Lin, C., Klein, L., Bronstein, M. M., and Tong, A. (2025). Scalable equilibrium sampling with sequential boltzmann generators. arXiv:2502.18466 18462..
Torrie, G. M. and Valleau, J. (1977). Nonphysical sampling distributions in monte carlo free-energy estimation: Umbrella sampling. Journal of computational physics, 23(2):187-199
Trippe, B. L., Yim, J., Tischer, D., Baker, D., Broderick, T., Barzilay, R., and Jaakkola, T. (2023). Diffusion probabilistic modeling of protein backbones in 3d for the motif-scaffolding problem. In International Conference on Learning Representutions.
Vincent, P. (2011). A connection between score matching and denoising autoencoders. Neural Computation, 23(7))1661-166
Wang, J. and Hou, T. 1a). Application of molecular dynamics simulations in molecular property prediction. density and heat of vaporization. Journal of chemical theory and computation, 7(7):1151-2155.
Wang, J. and Hou, T. (2)1 Application of molecular dynamics simulations in molecular property prediction ii: diffusion coefficient. Journal of computational chemistry, 32(16)::505-3511
Wang, J., Olsson, S., Wehmeyer, C., Pérez, A., Charron, N. E.,, Fabritiis, G. D., Noé, F,, and Clementi, C. (2019). Machine learning of coarse-graines molecular dynamics force fields. ACS Cent. Sci., 5:755-767.
Watson, J. L., Juergens, D., Bennett, N. R., Trippe, B. L., Yim, J., Eisenach, H. E., Ahern, W., Borst, A. J.. Ragotte, R. J., Milles, L. F., Wicky, B. I. M., Hanikel, N., Pellock, S. J., Courbet, A., Sheffler, W., Wang, J., Venkatesh, P., Sappington, I., Torres, S. V. Lauko, A., De Bortoli, V., Mathieu, E., Ovchinnikov, S., Barzilay, R., Jaakkola, T. S., DiMaio, F., Back, M., and Baker, D. (2023). De novo design of protein structure and function with rfdiffusion. Nature, 620(7976):1089 1100.
Wirnsberger, P., Ballard, A. Papamakarios, G., Abercrombie, S., Racaniere, S., Pritzel, A., Jimenez Rezende, D., and Blundell, C. (2020). Targeted free energy estimation via learned mappings. The Journal of Chemical Physics, 154(142:44411.
16

---

Wolf, S., Lickert, B., Bray, S., and Stock, G. (2020). Multisecond ligand dissociation dynamics from atomistic simulations. Nature Communications, 11(1).
Wu, H. and Noć, F. (2024). Reaction coordinate flows for model reduction molecular kinetics. The Journal of Chemical Physics, 160(4).
17

---

A Proofs, Derivations, and Formalisms
A.1
Diffusion
We have opted to use VP diffusion (Song et al., 2021) throughout the paper, and thus, the drift and diffusion coefficient can be written
f(x, 1) --1()) s( = VB(1),
(17)
where
((1) Smin + (Bmax -miin)
(18)
with the hyperparameters from Song et al. (2021) such that (8min m (0.1, For this noise schedule to be suitable for molecules, it is important that we normalize the data to have unit variance.
With this specific choice for f and Solin, 2019) with moving mean and standard deviation such that
can write the transition kernel as Gaussian (Särkkä and
(x( (=))) = (=((() ((<(() (()
(19)
(xx(t)) 16 s(s)ds (0), (1 e- (((6
(20)
A.2 Residual Loss
In this section, we prove Equation (12) and show that R(x, t; v) can be used to get an unbiased estimation of R(x, 1). For simplicity of notation, let us express the log Fokker-Planck equations from Equation (8) as
2 g()) div s (x. r0 (x. t) = 0,
(21)
where (x, Va log p (z), and vo involves only the first-order gradient of pg Here define the "weak" residual (Guo et al., 2022) of the above equations for each (z, t)
1 (1) div s - v,t) T0 + v,
R(x, 6) -ee
(22)
where a > is a small number. It can be seen that residuals are zero the two parts of the equations exactly equal. We now aim to get the unbiased estimation of the above residual, without calculating high-order derivatives.
As such, we can show that for an arbitrary
v exp
20
Ep [dvv s + v,t)]
div & 80(x + v. t)d
(23)
(2 &
1
(V.cos) exp
(-
vTv
s0 +,,)
do (24)
(2TO2)
4
20
TT v
exp
(x
+
(25)
(2
20²
(x
V.
Eg
(26)
g²
U
8
+ v.t)
v 0 x u
E
(27)
2
0²
0²
9
(x
+ v,t) t)
8 (x v2
E.
(28)
2
where E DR DR. V
D and E I.
18

---

Based on this, we can obtain
R(x, t) = Eu R(2, [:
(29)
1 (1)
8 (x + v. p) s 2,t) 70( + v, + 7o (a - u,
= Ev 2
20
2
Hence, R(x, 1; v) is an unbiased estimation of R(x, 1) by drawing single sample v ~ NN,
In practice, further reduce the computational overhead by only using single approximation for Ne. and defining
8 + v, 50 u.t) + (0 v,
R(æ, t; - -ー
2
(30)
2x
We found 10
4 be an effective choice throughout our experiments.
A.2.1 Relation with Hutchinson Trace Estimation
In fact, our treatment the divergence term in the weak residual formulation is closely related to the Hutchinson's 's trace estimator used by Albergo and Vanden-Eijin
(2025) for computing the PINN objective. Both approaches utilize Gaussian perturbations of x to obtain unbiased estimates of the divergence. However, it is important to note a key difference in how the residual is defined. As described in Equation (29), in our weak residual R(x, t), we apply the same Gaussian perturbation not only to the divergence term but also to other terms in the residual. This means that when the strong residual R(x, 1) is exactly zero, the weak residual R(æ, () E[R(æ + ()] also remains exactly zero, regardless of the perturbation variance o², avoiding truncation errors. This distinction allows our weak residual formulation maintain consistency even under large perturbations.
A.3 Finite Difference Approximation
To approximate 0 log po, relied on a finite difference approximation Fornberg (1988), stated in Equation (14). For this estimation, we have followed the work of Lai et al. (2023), and used the hyperparameters that they suggested (ha, ha) (0.001, 0.0005).
A.4 Connection with Flow Matching
The score-force relation introduced Equation (6) not specific to diffusion models. It also applies to other generative models, such as flow matching (Lipman et al., 2023), and in general to any model that can estimate v log (t=) In standard Gaussian flow matching, the learned vector field o relates to the score through
1 V logp((x) = 1 1- (t, 0 (x, -x) x)
(31)
as discussed by Lipman et al. (2024). Thus, forces can also be obtained from flow-matching models via reparameterization. However, in our experiments, we found that flow-based models perform worse near x 0, likely because the stochasticity inherent to diffusion models improves generalization in this regime.
A.5 Formalization Coarse-Craining
In coarse-gre graaiiing, we aim to reduce the number of dimensions of our system by combining multiple atoms into individual beads. Given non-CG samples æ, the Boltzmann distribution of CG samples can be recovered by
U(x) KgT
p(z) x
exp
6(E(x) 2) dæ.
(32)
which defines the CG potential up to a constant. 5 the Dirac delta function.
B Implementation Details and Experimental Setup
In this section, we will discuss additional details for the main experiments presented in the paper.
19

---

B.1 Langevin Integrator
We perform NVT dynamics with the Langevin integrator as implemented in openMM (Eastman et al., 2017) version 8.2.0. The update reads
v a av
VoU(T)M + (kaT(1-a)MM
(33)
1+ x' At 2to,
with a c cxp(( Here, and At index the simulation timesteps (not the diffusion time), y is the friction coefficient, M the mass matrix, R ~ (0, 1) standard normal random vector, kg the Boltzmann constant, and 7 the temperature.
B.2 Metrics
To compute the JS divergence and the PMF error, we first project our data onto either the dihedral angles the first two TIC coordinates (Pérez-Hermáne Herendderd et al., 2013) and then discretize the observed free energy into binned histograms. For the IS divergence. we then compute the JS distance between the two probability vectors (we flatten the 2D histograms). To prevent discontinuities, we assume that in each bin there is at least one observation by adding
As for the PMF error, we discretize into 64 bins for the dipeptides and into 100 bins for the proteins and compute the proportion of samples in each bin. These are then transformed by taking the log in each bin and then computing the square loss, which is averaged over all bins. Similarly, we have ensured that each bin contains some data and have added 10 6 as a baseline proportion. The approach and implementation
analogous to Durumeric et al. (2024).
B.3 Alanine Dipeptide
Dataset. The alanine dipeptide datasets is available part of the public bgmol (MIT licence) repository here: https://gi tttbbb.t........ com/ com/noogroup//eee bogmoollpppggl.... The dataset was generated with an MD) simulation, using the classical Amber ff99SBildn force-field at 300K for implicit solvent for duration of 1 ms Köhler et al. (2021) with the оpenМM library (Eastman et al., 2017). For training, we have selected 50k random samples from this simulation.
Architecture. For alanine dipeptide we have used quite small architecture, where the hyperparame- ters are listed in Table 3. When multiple parameters are listed for the same model, this means that they are used for the corresponding MoE model. Note that when using MoE, we have mostly used the same model architecture, except that only the Fokker-Planne regularized model is conservative. As for the optimizer, we have used Adam W (Loshchilov and Hutter, 2019).
Parameter
Diffusion
Mixture
Tikker-Planck
Both
Epochs
10000
7000, 2000, 1000
10000
7000, 2000, 1000 3. 10-4
Max Learning Rate
3 10-4
3. 10-4
3. 10-4
Min Learning Rate
10-5
10-5
10-8
10
BS
1024
1024
1024
1024
Attention Heads
8
Feature Dim
16
Model-Ranges
(0, 1)
(0,0.1), 1. 0.6). [0.6. 1.0)
(0.
(0,0.1), (0.1. 0.6). [0.6, 1.0) Yes. No. No
Conservative
Yes
Yes. No, No
Yes
0
0. 0, 0
5 10-4
10-3, 0,
Hidden Dimension
96
96, 96, 96 2, 2,
96
96, 96, 96 2, 2,
Layers
Table 3: Alanine dipeptide. This table contains the hyperparameters for the different models shown.
Simulation. To perform Langevin simulation, have extracted the forces from the model via Equation (6) at t = 10-5 for all models except for Two For One, where we chose = 0.02, which is the same value presented in Arts et al. (2023) for the same number of training samples.
Evaluation. To compute the standard deviation in the experiments, we have trained the same model with three different seeds.
20

---

B.4 Fast-Folding Proteins
Dataset. We evaluate our model on the proteins Chignolin (10 amino acids) and BBA (28 amino acids) using the data from Lindorff-Larsen et al. (2011). The reference data was simulated in explicit solvent with step size of 2 at 340K for Chignolin and 325K for BBA. total 106 jss and 223 jus were simulated for Chignolin and BBA. For training, we have randomly selected 80% of the samples which accounts to 427, 794 and 1, 300, 156 samples for Chignolin and BBA respectively.
Architecture. We have used the same architecture as for alanine dipeptide, with the only difference being that we use larger networks. Compare the hyperparameters in Table for Chignolin and Table for BBA. Note that for Two For One we used the same model as Diffussion but evaluated it at different for simulation. The hyperparameters slightly deviate from the choices in (Arts al., 2023) to produce consistent results with our method. As optimizer, we used AdamW (Loshchilov and Hutter, 2019).
Parameter
Diffusion
Mixture
Fokker-Planck
Both
Epochs
2000
1700, 200, 100
2000
1700, 200, 100 3 10
Max Learning Rate
3 10-
3 10
10
Min Leaminng Rate
10-5
10-5
10-5
10-5
1024
1024
1024
1024
Attention Hcads
Feature Dim
16
16
16
16
Model-Ranges
(0, 1)
(0,0.1), [0.1, 0.6). [0.6, 1.0)
(0.
(0,0.1), 10.1. 0.6). [0.6, 1.0)
Conservative
Yes
Yes. No, No 0, 0
Yes
Yes. No. No
a
2 10
2 10 3, 0, 128, 96,
Ilidden Dimension
128, 96, 3. 2,
128
Layers
2
3, 2,
Table 4: Chignolin. This table contains the hyperparameters for the different models shown.
Parameter
Diffusion
Mixture
Fokker-Planck
Both
Epochs
1200, 400, 200
1500
1200, 400, 200
Max Learning Rate
3 10
3 10-
10-4
10-4
Min Learning Rate
10-5
10-5
10-6
10-5
BS
1024
1024
1024
1024
Attention Heads
8
8
8
Feature Dim
16
16
16
16
Model-Ranges
(0, 1)
(0,0.1), (0.1, 0.6). (0.6.100)
(0,
(0,0.1), (0.1, 0.6). (0.6, Yes. No. No
Conservative
Yes, No. No
Yes
0, 128, 128, 128
10-2
2 10 0.
Hidden Dimension
128
128
128, 128, 128 3, 3,
Layers
3
3,
3
Table 5: BBA. This table contains the hyperparameters for the different models shown.
Simulation. To perform I Langevin simulation, we have extracted the forces from the model via Equation (6) t 10 for all models except for Two For One, where chose 0.02 and t = 0.005 for Chignolin and BBA respectively, which is the same value as presented in Arts et al. (2023) for the same number training samples.
Evaluation. To compute the standard deviation in the experiments, we have used the same models but three different seeds for inference.
B.5 Dipeptides (2AA)
Dataset. The original dipeptide dataset (2AA) was introduced in Klein et al. (2023a) (MIT License) and available here: https:/ //uggs///uugggfaaae ace cadaaaasss/iiccoooootttteeeeeeee
As this includes a lot of intermediate saved states and quantities, like energies, there is smaller version made available by Klein and Noé (2024) (CC BY 4.0): https://ost. iooii////// //vz2//vveee _nnl 00000011778aa..... For comprehensive overview of the simulation details, refer to Klein et al. (2023a). All dipeptides were simulated implicit solvent with classical amber-14 14 force-field at T = 310K. The simulation of the training and validation peptides were run
21

---

for 50ns, while the test peptides were simulated for lus. All simulation were performed with the оpenMM library (Eastman et al.. 2017).
Note that we have removed dipeptides containing Glyeine from our dataset to ensure that all dipeptides have the same number of is not a technical limitation of our architecture. It is split into 175 train, 75 validation, and 92 test dipeptides, of which have used 15 for the results presented in the paper (also the metrics) to reduce inference time. The fifteen test dipeptides are: AC, AP, AT, ET, HP, HT, IM, KQ, KS, LW, NF, NY, RL, RV, TD.
atoms. This made it easier to handle it in the code, but it
Architecture.
The hyperparameters are listed Table 6. When multiple parameters are listed for the same model, this means that they are used for the corresponding MoE model. Note that when using MoF, we have used smaller networks for larger diffusion times, and only the Fokker-Planck regularized model is conservative. As optimizer, used AdamW (Loshchilov and Hutter, 2019).
Parameter
Diffusion
Mixture
Fokker- Planck
Both
Epochs
120
100, 20, 10
120
100, 20, 10
Max Learning Rate
3 10
3 10-
3 10
3 10
Min Learning Rate
10-5
10-5
10-5
10-5
BS
1024
1024
1024
1024
Attention Heads
8
8
8
Feature Dim
16
16
Model-Ranges
(0, 1)
(0.0.1), (O.L. 0.6). (0.6, 1.0)
(0,
(0.0.1), (0.1, 0.6). (0.6,
Conservative
Yes
Yes, No, No
Yes
Yes, No. No
0, 0, 0
5 10
10 0,
Hiiden Dimension
128
128, 96, 96 2,
128
128, 96, 96 3, 2,
Layers
3
Table Dipeptides. This table contains the hyperparameters for the different models shown.
Simulation. To perform Langevin simulation, have extracted the forces from the model via Equation (6) at t = 10-5 for all models except for Two For One. As this system has not been tested by Arts et al. (2023), opted to use the same t for alanine dipeptide, namely t 0.02.
Evaluation. To compute the standard deviation in the experiments, we compute the values across all testset dipeptides.
B.6 Compute Infrastructure
We have used a single RTX 3090 GPU for the toy systems, an A100 with 80GB memory for alanine dipeptide, two A 100 80GB GPUs for the dipeptide dataset and Chignolin, and four for BBA. Depending on availability, have also used H100 for some experiments. For inference, we use single GPU.
B.7 Software Licences
In our code, have used jax (Bradbury al., 2018) (Apache-2.0) and the accompanying machine learning library flax (Heek et al., 2024) (Apache-2.0), For the graph transformer architecture, have extended code from Arts al. (2023) (MIT) and have re-implemented the code from wrwrrppppttrassfoorr-------oooee
https://gt thhtbbb. thub
(MIT) in jax.
For the free-energy plots of the Müller-Brown potential, we used Hoffmann et al. (2021) (L.GPL-3.0). For trajectories and simulations, we have used openMM (Fastman et al., 2017) (MTT) and mdtraj (McGibbon et al., 2015) (LGPL-2.I).
C Further Studies and More Experiments
C.1
Comparing Conservative and Score-based Models
Previous work (Arts et al., 2023) suggested that conservative models improve the quality of the diffusion process. However, this effect was not observed for image data (Salimans and Ho, 2021). In Figure 7, we compare these approaches in practice. For iid sampling, conservative models provide
22

---

slight improvement. In contrast, for simulation, we were unable to train stable score- score-based models without a conservative parameterization, aligning with knowledge from the force-field literature (Schütt et al., 2017; Batzner et al., 2022). Using conservative model yields much more stable forces, making simulation feasible. We attribute this to the smoother behavior of the conservative parameterization, which prevents sudden changes in the score. Since the impact on iid sampling is negligible, we consider the conservative parameterization most relevant for small timescales, where model training is more sensitive. Consequently, in our MoE architecture, we apply the conservative parameterization only to the small-time diffusion model, achieving comparable or even superior iid sampling performance.
Scote-besse Dulusiion Model (id)) Compuryaat ive uagg Model (id)
iid
Scom-btesed Diffielon Model (lm)
Conpencat ve Diffustin Model is
sim
Figure 7: Alanine dipeptide. We compare a conservative diffusion model with a score-based model. We can see that around the low-density regions, the conservative parameterization generates better iid samples. As for simulation. score-based model exhibits stability issues, and the simulation becomes unstable after few thousand steps.
C.2 Mixture vs. More Parameters
In our experiments, using the MoE scheme naturally increases the total number of parameters. To assess whether the observed improvements stem from this additional capacity or from the mixture design itself, we perform ablation studies on alanine dipeptide. Specifically, we compare the original results from the paper against three variants: (Deep), and one using three instances of the conservative diffusion model combined into a Mixture of experts. The results in Table 7 show that Mixture outperforms both larger single models and achieves more stable simulation behavior, and requires the least time for training and inference (see also Appendix C.3). This demonstrates that the advantages of the MoE design extend beyond simple parameter scaling and can yield runtime improvements.
with wider layers (Wide). one with more layers
Method Parameters
Training (1)
iid JS (1)
sim. JS (L)
iid PMF (L)
sim. PMF ())
Diffusion
647585
50min
0.0881 = 0.0003
0,0695 = 0.0517
0.095 1 0.003
1.047 0.924
Wide
1964321
67min
0.0082
0.0003
0.0406 1. 0.0236
0.096 1 0.003
0.467
0.245
Deep
1968562
111min
0.0078 1 0.0003
0.0376
1: 0.0087
0.091 1: 0.002
0.478
0.038
Mixture
1957567
50min
0.0079
0.0003
0.0264 :: 0.0085
0.093 : 0.007 0.325
0.113
Table 7: Alanine dipeptide. Ablation study comparing the MoE scheme with parameter scaling. Reported are
the IS divergence and PMF error for iid sampling and simulation (sim). Lower values indicate better performance.
The MoE achieves the most stable and accurate simulations while maintaining the same training cost as the baselinc.
C.3 Runtime Comparison
We compare the training and inference runtimes of different approaches in Table 8. While the MoE setup could, principle, trained in parallel, our current implementation does not distribute the experts across devices and is therefore not fully optimized. For the low-dimensional systems, this
23

---

can even introduce small additional overhead. As a result, runtime improvements are only visible for larger systems. In contrast, models trained with Fokker-Plande regularization show a clear computational overhead due to the additional residual evaluation, which requires multiple forward passes per training step. During inference, however, they perform the same as their counterparts.
Dataset
Task
Diffusion
Mixture
Fokker-Planck
Both
Alanine Dipeptide
Training
50min
50min
4h 39min
3h 59min
iid Sampling
3min
4min
3min
4min
Training
4h 58min*
3h Omin
22h 36min
21h 33min
Chignolin
idd Sampling
43min
14min
43min
14min
Training
Id 20h
1d 19h*
7 2h
6d 15h
BBA
iid Sampling
12h 57min
5h 56min
12h 57min
5h 56min
Training 4h Smin
3h 50min
28h 39min
27h Smin
Dipeptides
iid Sampling
Smin
4min
Smin
4min
Table 8: Training and inference runtimes for all models. Methods marked with were trained using half the number of GPUs (one for Chignolin and two for BBA). The MoE model reduces inference time substantially modular design, while Lokker-Plande regularization introduces additional computational overhead from multiple forward evaluations during training.
In terms of simulation, the main computational advantage of our approach arises from its ability to perform simulations in parallel and to leverage searse-rrannigg In classical molecular dynamics, simulating protein typically requires modeling the solvent explicitly, which adds thousands of water molecules and substantially increases computational cost. Coarse-grainin eliminates these degrees of freedom, thereby accelerating computation and improving parallel scalability.
For instance, in the dipeptide dataset, where each molecule contains only few atoms and no solvent to coarse-grain, our model achieves 150 parallel simulations at 830 steps/s each, corresponding to roughly 125k steps/s on a single NVIDIA A A100. In comparison, conventional force-field simulation on an NVIDIA V40 reaches around 10k steps/s. Even larger gains can be expected for systems with explicit solvent, where coarse-graide grariiig removes many more particles. As an example, our BBA protein simulation (100 parallel runs of IM steps each) completes in approximately one hour.
C.4 Toy System
Dataset. We have used the Müller- potential (Müller and Brown, 1979) to demonstrate the capabilities of our approach in two dimensions. For this, we have used the following potential
U(x v)
200 exp (-(x 1) 10y²)
100 exp
(-x² 10 (y 0.5))
(34)
170 exp (-6.5- (0.5 + 11 (r+ 0.5) (y - 6.5 (y - (y-115)
+ 15 exp (0.7 (1 ()+ + (x + (v 0.7 ()-1)²)
To generate training samples from this potential, we have performed a Langevin simulation (compare Equation (5)). For this, have performed 5M steps with kB 23, di 0.005, M 0.5 where we only store every 50th sample to generate 100k training samples.
Architecture and training. For the toy systems, we have used simple multi- llyer perception with the hyperparameters presented in Table 9.
Parameter
Diffusion
Mixture
Fokker-Planck
Both
Parameters
17849
17263
17849
17263
BS
128
128
128
128
Model-Ranges
(0, 1)
(0,0.1). [0.1, 0.6). [0.6, 1.0)
(0.1)
(0,0.1). 0.6). [0.6, 1.0)
Epochs
180
120, 30, 30
180
120. 30, 30
Hidden Dimension
02
64, 64, 54
92
64, 64, 54
Layers
3
3, 2.2 0
3
3, 2
0
5- 10
5 10 0
Table 9: Müller- Browln This table contains the hyperparameters for the different models shown.
24

---

(a)
Course graind
(c)
Pize Energy Projection (SUI)) bicerne boste
Pokker-Planck
Figure 9: Alanine dipeptide. (a) The coarse-graining scheme. (b) Comparison of the Ramachandran plots of different methods for iid sampling and Langevin simulation. (c) The projection of the free energy surface and differences along the dihedral angle for samples generated with simulation.
Evaluation. In Figure 8, we can see the free energy plots the different methods. All methods produce iid samples that match the Reference. However, when using the learned score for Langevin simulation, the standard Diffusion model fails to reproduce the correct distribution and undersamples the low-probability state, highlighting the inconsistency between sampling and simulation. Although having roughly the same number of parameters, the Mixture model partially improves this, but consistency is only achieved with Fokker-Planck regularization. Combining Both approaches further improves performance. Numerical results can be seen in Table 10.
Method
iid JS
sim JS
iid PMF
sim PMF (L)
Reference
0.01 19 0.0004
0.087
0.002
Diffusion
0.0122 + 0.0013
0.0448
0.0125
0.111 + 0.006
0.504 + 0.150
Mixture
0.0109 1 0.0007
0.0254 1 0.0109
0.097 1. 0.004
0.247 1 0.113
Fokker-Planck
0.0130 + 0.0010
0.0166
0.0009
0.122 + 0.006
0.163 + 0.008
Both
0.0110 0.0007 0.0108 0.0008
0.098 0.003
0.099 0.004
Table 10: Müller-Brown. Comparison of methods based on JS Divergence and the PMF error. Lower values better. To compute the standard deviation, have trained different models and performed sampling/simulation with them. As for the reference, have started multiple simulations with different seed on the same ground-truth potential. This serves as a reference what could optimally achieved.
idd
sim
Reference
Diffusion
Mixture
Fooker-PPanne
Both
Figure 8: Müller-Brown Comparing free energy plots of different models for classical diffusion sampling (iid) and Langevin simulation (sim). The energies should align with the training data (Reference).
C.5
Alanine Dipeptide
In this section, we report some further results and plots on alanine dipeptide.
C.5.1 Free Energies
Inconsistent sampling. Figure (b) compares the free energies of the sampled dihedral angles for iid sampling and Langevin simulation (sim). While methods can match the training distribution under lid sampling, simulation quality varies, and existing models show inconsistencies. Standard Diffussion fails to recover the low-probability mode (i.e., & > 0) completely, even when starting simulation from these regions. Mixture generally improves the results, but still does not find the other
25

---

Method
iid JS (U)
sim. JS (1)
iid PMI (U)
sim. PMF ())
Ditfusion
0.0081 + 0.0003
0.0695 +: 0.0517
0.095 + 0.003
1.047 1 0.924
Two For One
0.0081 = 0.0003
0.0158 + 0.0002
0.098 + 0.006
0.206 + 0.004
Mixture
0,0080 1 0.0004
0.0353 1 0.0117
0.092 1 0.007
0.388 1 0.109
Fokket-Planck
0.0082 = 0.0002
0.0090 + 0.0006
0.098 + 0.003
0.104 + 0.004
Both
0.0082 = 0.0004
0.0080 0.0002
0.097 + 0.008
0.091 = 0.004
Table 11: Alanine dipeptide. Comparison of JS divergence and PMF emor. To compute the mean and the standard deviation, we have trained and evaluated three models with three different seeds.
energy projection
energy projection
175-
Difision
Diflucoon
For
150-
Tan
Mictars
FaccerPanck
125-
Foboann bock
Mizzure
Buth
1000
Both
iid
75--
5.-
25-
00-
Free energy projection of
Free energy projection of
Tacharanss
Rela Dimelan
fiflution
sim
15.
Bot
Envegy
Enegy
12
Figure Alanine dipeptide. Comparing free energy along the dihedral angles yp, d for iid sampling simulation across different models.
numerical results in Table 11, where Mixture achieves lower means and smaller variance, but simulation emrrrs remain noticeable. We attribute this behavior to the smaller time range, which focuses the model's attention, allowing it to learn a better, more stable optimum.
mode, This reflected
Consistent models. Fokker-Planck regularization enables the model to recover the missing states without modifying the diffusion time, and thus preserving structural accuracy. Table shows that the regularization substantially improves consistency between iid and simulation, although iid performance slightly declines in favor of improved simulation performance. Combining MoE with Fokker- regularization for Both, enhances simulation quality further while barely mitigating drop iid performance. The resulting model achieves close alignment between iid and simulation, and captures the free energy landscape in simulations accurately (see Figure 9 (c)). The similar iid performance of Both and Tokker-Plancl suggests that applying regularization introduces additional constraints and restrictions the model and hence degrade generative quality.
Detailed free energy. In Figure 10 compare the free energies along the dihedral angles y, bb for iid sampling and simulation. We can see that the results from the main paper persist.
C.5.2 Fokker-Planck Residual
In Figure 11 the Fokker-Plande residual error || (x. t) Ô log rf (x((())) results are similar to what was reported in Figure 6. We can note that the Fokker- error of Mixture is lower than Diffusiion, indicating that MoE can improve the model's consistency.
(z) ||2
is reported. Overall, the
C.5.3 Bonds
Figure 12 shows the immediate bond lengths of the coarse -aarineer molecule for iid sampling and Langevin simulation. Since Two For One does not evaluate the model at = 0, it introduces noise across all bonds. We can also see this behavior by looking at the Wasserstein distance of the bond- lengths the reference data as seen in Table 12. In this table, Mixture provides the best bonds during simulation, but does not recover all modes.
26

---

Two For One
Bokker-Planck Error
Mixture
Folkar-Planck
Both
10.
10"
0.6
1.0
Diffusion time
Figure 11: Alanine dipeptide. Comparing Fokker Planck error for log pd of multiple models.
CA C bond langen
100
121
123
idd
Lraghin
10 ureelidd
1.0
C- langhh
N4 bond length
CA-C bond langth
C- band langth
Ratue
sim
12
Ennghin
Figure Alanine dipeptide. This illustration shows direct bond lengths sampled by the different methods.
Method
iid relative W1 (J)
sim relative W1 (1)
Diffusion
0.88 i. 0,70
0.83 1: 0.15
Two For One
0.88 1 0.58
23,51 4: 1.91
Mixture
1.13 + 0.46
0.47 + 0.12
Cokker-Planck
1.65 0.37
1.23 1 0.16
Both
1.00 0.00
1.00 0.00
Table 12: Comparison of methods on alanine dipeptide based on the Wasserstein distance of the C-N bond lengths to the reference data. We have divided all entries by the Wasserstein I distance of Both that the numbers are easier to compare. other words, numbers larger than mean that the bonds are worse than Both.
C.5.4 Force Matching
We further compare model against prior work force matching, which either employs explicit forces from the training set (Husic et al., 2020) or infers forces implicitly from noise the data (Durumeric al., 2024; Klein al., 2025). All these approaches share the same underlying architecture, the CGSchNet model (Husic et al., 2020), which combines a trainable SchNet component (Schütt al., 2017) with fixed energy terms accounting for bonded interactions such bonds, angles, and dihedrals. In contrast, our models operate without any such energy priors, and thus require significantly less prior knowledge about the investigated systems. As these previous studies commonly evaluated their models on six-bead representation of alanine dipeptide (Wang et al., 2019; Husic al., 2020; Klein et al., 2025), retrain our models as well the Two for One model on the same coarse- -raiine system. The corresponding training and evaluation trajectory was generated via MD for 1 μs explicit solvent (Wang et al., 2019), and the training set for all methods consists of 100k samples. For the force-maten matching baselines, we report the results from Klein et al. (2025). The comparison, shown Table 13, evaluates again the PMF between the simulated and target trajectories. We only evaluate the simulations, the force-matching methods do not permit independent sampling. Overall, our models achieve superior performance compared to all force matching baselines on this dataset, despite not using any force information, unlike Husic al. (2020).
27

---

Method
sim. JS (1)
PMF (4)
Force Matching
0.0243 0.0015
0.322
0.005
Noise Force Matching
0,0402
0.864 0.146
Operator Forces
0.0214
0.0041
0.301
0.043
Operator Forces
0.0166 0.0009
0.282 0.031
Diffusion
0.0928 0.0165
1.135
0.205
Two for One
0.0151 0.0006
0.166 0.009
Mixture
0,0506 0.0370
0.545 0.429
Fokker-Planne
0.0092 0.0017
0.107
0.028
Both
0.0075 0.0024
0.090 0.049
Table 13: Six-bead alanine dipeptide. Comparison of different models for learning forces. The table reports the divergence the PMF error between simulated and target distributions (lower is better). Our Fokker-Planck- regularized models outperform all force- morce-matching methods despite being trained without any force labels.
C.6 Proteins
We now provide additional comparisons and ablation studies for the independent contributions of MoE and Fokker-Planck regularization for Chignolin and BBA as well as analysis of dynamics.
C.6.1 Free Energy Surfaces
We present the same free energy surfaces as in the main text, but distinguish between Mixture, Fokker-Planck, and Both (referred to as Ours in the main text). In Figure 13, we show the free energy along the first TIC coordinates by projecting the data into 100 bins, and in Figure 14, we show the corresponding kernel density estimates along each axis. Beyond the results discussed in the main text, we observe that simulations using Diffasson (i.e., evaluating the model / 0) unstable. The MoF. variant in Mixture improves stability but does not fully resolve these issues. In contrast, introducing the Fokker-Plande regularization (with or without mixture) yields stable simulations, and the resulting free energy surfaces from diffusion sampling and simulation become consistent. In Figure 14, the models appear nearly identical for iid sampling but differ for sim. The qualitative metrics are similar between Bookker-Plande with and without MoE, showing only minor differences.
The same trends observed for BBA in Figures 15 and 16. The Diffusion model fails recover a meaningful free energy landscape, while Mixture using MoE substantially improves the results but remains insufficient. Once the Fokker-Planck regularization is applied, the simulations become consistent with diffusion sampling. As shown in Figure 16, the iid results of several models are nearly indistinguishable in the one-dimensional projections, making it difficult determine which performs best, For sim, it is also unclear whether Fokker-Planck or Both yields better performance. quantitative comparison is provided Appendix C.6.3.
Reference
Diffusion (16
For
Micure
Parck (1
Both
iid
TC3
TC3
TC3
TC3
TC3
TC3
Befeeence
Fffiam (Nm))
Tho (smn)
Mictur
Teblen-iande [4ml
Dath (sm)
sim
Figure 13: Chignolin. We compare energy projected onto first two axes presented methods for iid sampling and Langevin simulation (sim). The aim is to match the Reference (left) and have self-consin consistent model where iid and sim closely aligns.
28

---

energy projection
energy projection
50
One
Emergy (ol
60-
Misture Rosker Farce Sozn
iid
Eneeryy
Free energy projection of TIC 0
Free energy projection TIC
60
toloronee
Ditialon
Energy (el
sim
Emergy
20
0055
0.0
0.5
1.5
3.5
2.0
25
-1
215---00----- 0.00
025 a.0
T1C
TIC
Figure 14: Chignolin. Comparing the free energy along the first two TIC axes for iid sampling and simulation across different models with the reference in black.
Reforence
Diffusion (id
Tao For One Ridd
Micure Hid)
Folioor Parck dish
Both (id)
iid
Refreence
fifusiin (sl))
Teo One Isin
Micuum 6im)
Ralden-Pandd fsmn
Both (sm)
sim
Figure 15: BBA. We compare the free energy projected onto the first two TIC axes for all presented methods for iid sampling and Langevin simulation (sim). The aim match the Reference (left) and have self-consistent model where iid and sim closely aligns.
Free energy projection of TFC 0
Free energy projection of TIC
Ralstune
D Dhidden
Two fy One
Nisure
Poiden-Parck
Both
iid
10
5
1.0
555
0.2
6.5
1.5
29
TCC
T1C 1
energy projection
energy projection
25
For
20
Naturs
Forke Both
sim
Energy
0.5
15
25
Figure 16: BBA. Comparing the free energy along the first two TIC axes for iid sampling and simulation across different models with the reference in black.
29

---

C.6.2 Contact Maps
To analyze structural consistency beyond local distances, we compute pairwise contact maps between all atoms. such, they provide a compact way representing molecular structure and indicate which atom pairs are in spatial proximity. Specifically, for each generated ensemble, we compute the distances between all atom pairs and record whether they fall below threshold of 10 Å. Averaging these contacts over all generated samples yields probability map indicating how frequently each atom pair interacts. We then take the log of these to smooth out the range and improve visualization.
This is visualized for Chignolin and BBA in Figures 17 and 18. This representation captures the overall spatial organization of molecule, such as which regions are folded together or form stable secondary structures. Comparing contact maps from different models provides intuitive way to assess whether the generated ensembles reproduce realistic atomic arrangements.
These figures align with the overall trends in the paper showing that: tid works well across all models, Diffusson and Mixture is unstable for simulation, and Paaker-PlanckBBotte match the reference and the sampling methods are consistent. However, when examining these specific metries, the additional noise introduced by Two For One is less apparent. We attribute this to the binary cutoff used to determine whether two atoms in contact, which reduces sensitivity small structural perturbations. This effect is not present in the pairwise distance matrix analysis Tables 14 and 15.
iid
Rnference
Two For One
Vistre
Nokko: Pbrci
hett
sim
Figure Chignolin. Pairwise contact maps comparing different models and the two sampling methods. Two atoms are considered in contact if their distance is below 10 Á. Bright colors (corresponding to small values) indicate low probability of close proximity. The maps symmetric by construction.
Teo One
Mieturt
Plance
S4th
iid
Tao ForOne
Fakkr Panic
sim
Figure 18: BBA. Pairwise contact maps comparing different models and the sampling methods. Two atoms are considered in contact their distance is below 10 Á. Bright colors (corresponding to small values) indicate a low probability of close proximity. The maps are symmetric by construction.
C.6.3 Quantiative Results
In addition to the results presented in the main paper, we provide further quantitative analyses in this section. Alongside the PMF error, we also report the JS divergence. While the PMF error evaluates energy differences, the JS divergence compares probability densities, which can make model differences less pronounced. Both metries are computed in 2D TICA space.
We additionally report the IS divergence computed on the pairwise distance (PWD) matrix. For this, we first calculate all interatomic distances and extract the upper triangular part of the full PWD matrix to avoid redundancy. We then remove the diagonal elements and exclude pairs separated by fewer than three bonds to focus on nonlocal structural correlations. In other words, we compare distances between atoms that are not close neighbors along the chain. For each atom pair, we then construct normalized histograms of their distance distributions for both the generated and reference ensembles and compute the IS divergence between them. Averaging these divergences across all pairs yields
30

---

single scalar value that quantifies how well the generated ensemble reproduces the overall spatial organization of the molecule. Lower values correspond higher structural fidelity.
The resulting metrics are summarized in Tables 14 and 15 for Chignolin and BBA, respectively. The trends are consistent with those discussed in the main text: for Chignolin, observe only slight decrease in iid sampling quality when introducing Fokker-Planck regularization or the MoE model. However, for the larger and more complex BBA system, it becomes evident that achieving consistent energy-based model requires trade-off's in iid sampling quality, In contrast, our models excel in the sim setting, achieving the best results in both energetic and structural metrics, particularly for the PWD-but based measure of structural fidelity.
Method
tid JS (
sim. JS (1)
iid PMF (1)
sim. PMF ())
iid PWD JS (3)
sim. PWD JS (L)
Dittusion
0.0036 + 0.0001
0.4351 + 0.0141
0.027 + 0.000
63.804 + 0.372
0.0001 + 0.0000
0.3817 = 0.0009
Two For One
0.0036 0.0001
0.1023 + 0.0008
0.027 = 0.000
1.488 + 0.019
0.0001 = 0.0000
0.0082 = 0.0000
Mixture
0.0042 0.0001
0.4336 1 0.0075
0.033 0.000
11.185
0.430
0.0003 0.0000
0.2045
0.0004
Fokker Planck
0.0048 + 0.0001
0.0050 + 0.0001
0.037 + 0.000
0.039
+ 0.001
0.0004 + 0.0000
0.0008
0.0000
Both
0.0045 0.0001
0.0050 0.0008
0.035 0.001
0.038 0.006
0.0003 + 0.0000
0.0012 0.0005
Table 14: Chignolin. Quantitative comparison of all methods based on JS divergence, PMF error (both computed in TICA space), and the mean JS divergence of the pairwise distance (PWD) distributions of atoms. Lower values indicate better agreement with the reference data. The iid metrics assess equilibrium sampling quality, while the sim metrics evaluate the consistency of the simulated dynamics.
Method
iid ())
sim. JS ())
iid PMF (1)
sim. PMF
iid PWD JS (J)
sim. PWD JS (L)
Diffusion
0.0043
0,0000
0.2014 0.003
0,034 + 0.0I0)
5.387 + 0.144
0.0IM16 + 0.0000
0.1003 + 0.0050
Two For One
0.0043 + 0.0000
0.1162 = 0.0021
0.034 + 0.000
1.624 + 0.107
0.0006 + 0.0000
0.0240 + 0.0008
Mixture
0.0091 0. 0,00000
0.1964 + 0.0113 0.0113
0.081 + 0.000
5.970 - 0.278
0.0012 + 0.0000
0.1016 + 0.0020
Fokker Planck
0.0130
0.0000
0.0292
0.0035
0.117 1: 0.001
0.275
4 0.033
0.0018 1: 0.0000
0.0078 : 0.0014
Both
0.0238 0.0001
0.0238 + 0.0001
0.234 + 0.001
0.254 + 0.005
0.0023 + 0.0000
0.0042 + 0.0001
Table 15: BBA. Quantitative comparison of all methods based on JS divergence, PMF error (both computed in TICA space), and the mean IS divergence of the pairwise distance (PWD) distributions of atoms. Lower values indicate better agreement with the reference data. The metrics assess equilibrium sampling quality, while the sim metrics evaluate the consistency of the simulated dynamics.
C.6.4 Dynamics
The main advantage of our method is that we have access to a consistent and physically meaningful energy function that allows us to perform simulation and recover dynamic properties. To illustrate this benefit, we analyze the recovered dynamies and provide further evaluations.
As in the main text for BBA, we visualize in Figure 19 single simulated trajectory for Chignolin. The system exhibits folding and unfolding transitions between metastable states, showing realistic conformational changes over time. For this, we employ the Both model, which combines Fokker- Planck regularization with the MoE scheme.
(a)
(b)
(c)
TC1
0.2
-05
2000
4000
5000
5000
10000
12000
(Coarse-Grainee Granse- Timesteps
Figure 19: Chignolin. Recovering and analyzing the dynamic behavior. (a) 3D molecular structures from a single MD trajectory generated by our model. (b) Trajectory projected onto the first two TIC coordinates. (c) Illustrating the TIC coordinates over the coarse goasieer (and subsampled) timesteps.
31

---

To further assess the long-time MD simulations, we follow the work of Arts et al. (2023) and cluster the conformations into three metastable states for Chignolin and four for BBA. State assignments are obtained via k-means clustering the TIC coordinates, from which we estimate Markov model and compute the transition probabilities between states. The state definitions and transition matrices are visualized in Figures 20 and 21. Again, this analysis is only applicable to models that recover continuous dynamics rather than independent samples.
Although the visual differences between transition matrices are subtle, the Posker-Planckk-reulliizze models clearly provide the most accurate estimates of transition probabilities when compared the reference simulation. To quantify this, we compute the mean JS divergence between the transition probabilities of each model and the reference, reported in Table 16. These results confirm that Fokker-Plaan Planck regularization improves kinetic consistency, while the MoE scheme leads to slightly less accurate transitions; both when used independently and combination with Tokker-Plancke
Reference
Ore
Micoue
Folkenfunee
States
TC2
Nokke-fanch
2och
nace
AM
0.193
1134
0.99
0212
0017
E223
Probs.
0216
1222
0.22
6/06
CAR
3661
8,00
6.
CAH
0000
Figure 20: Chignolin Each conformation is assigned to one of three states based on the TIC coordinates (top row). shown State (purple). State 2 (green), and State 3 (yellow). Using these assignments, we estimate the corresponding transition probabilities between states.
Reference
Ore
Victur
Felder
States
TC
Saforwaee
For
Sath
Probs.
nEen
6.00
0.89
DETY
0.00
Figure 21: BBA. Each conformation is assigned to one of four states based on the TIC coordinates (top row), shown as State 1 (purple). State 2 (blue), State 3 (green), State 4 (yellow). Using these assignments, we estimate the corresponding transition probabilities between states.
Method
Chignolin JS ())
BBA JS (L)
Diffusion
3.4 10
10
Two for One
2.3 10
7.1 -.1-10-3
Mixture
3.0 10
4.0-10 10-2
Fokker-Planck
4.6 10-4
2.5 10-3
Both
2.1 10
4.2 12-1
Table 16: Mean JS divergence between the transition probabilities estimated from simulation and those of the reference trajectories (compare Figures 20 and 21). Lower values indicate better agreement.
C.7 Umbrella Sampling
Similarly to Appendix C.6.4, access to an -bsed diffusion model enables further applications such as molecular simulation and umbrella sampling (Torrie and Valleau, 1977). umbrella sampling, several simulations are carried out with restraints applied to specific regions of configuration space (e.g., along selected dihedral angles), and the resulting biased ensembles are reweighted by the bias potential to recover unbiased free energy estimates.
32

---

While such enhanced sampling techniques lie beyond the main scope of this work, we conducted preliminary experiments with our model. We found that, its current form, umbrella sampling remains challenging. The model tends to perform poorly when forced into high-energy or previously unseen regions of configuration space. For example, Chignolin, umbrella sampling along the first TIC coordinate failed to adequately explore the extremes of the coordinate, leading to unreliable free-energy estimates in those regions. Nonetheless, the overall shape the free energy profile could still be recovered with significantly fewer simulation steps. We believe that the proposed Fokker-Planck regularization could help address this limitation. Since the loss can be evaluated for arbitrary configurations æ, it may encourage the model generalize better to out-of-distee -itt-ii-ittiittiionn noisier samples, thereby improving its robustness in biased or enhanced sampling settings.
C.8 Transferability Across Dipeptides (two amino acids)
C.8.1 Complete Metries
In the main text, reported the PMF enror across the test set, and in Table 17, additionally include the JS divergence for the same experiments. Both metrics show consistent trends, although IS divergence is less sensitive to small differences, as it measures discrepancies between densities rather than energies.
Method
iid JS
sim.
iid PMF
sim. PMF
Transferable BG
0.0183 + 0.0970
0.230 0.119
Diffusion
0.0155 1: 0.0083
0.2256 0.1304
0.206 0.159
6.515 3.175
Two For One
0.0153 + 0.0080
0,0466 = 0.0114
0.203 0.149
0.741 0.319
Mixture
0.0155 1 0.0078
0.0444 0.0237
0.200)
0.127
0.658 0.407
Bokker-Plannee
0.0200 + 0.0071
0.0254 + 0.0119
0.241 + 0.105
0.368 + 0.267
Both
0.0157 + 0.0078
0.0177 = 0.0084
0.199 = 0.127
0,203 + 0.104
Table 17: Dipeptides. Comparison of metrics across all testset dipeptides with JS divergence and PMF To compute the mean and standard deviation, we have averaged the metries across the dipeptides from the test set.
C.8.2 Lysine-Serine (KS)
In Figure 22, we present extended results for the dipeptide analyzed in the main text (KS). The trends remain consistent with previous findings: using Two For One for simulation leads to incorrect observables and distorted equilibrium distributions. In addition, we show the free energy surfaces along the dihedral angles and v separately. The Fokker- Planck regularization improves these free energy profiles, and the Both model produces particularly accurate and well-aligned surfaces.
C. distance
Free energy projection of
Free energy projection of
Mastus Wum Turahhaade
Energyy Ao
Energy NxT
iid
Distance A
distance
Free energy projection of
Free energy projection of
Tetennce Melum wit
Energ ATT
sim
Eneege
2/
Figure 22: KS: We compare further metrics between iid sampling and Langevin simulation. We compare the
C.. CC, distance for the dipeptides and also the free energy projections along the dihedral angles wb.
33

---

C.8.3 Transferability: Results on More Dipeptides
In this section, show additional dipeptides from the test set to evaluate the performance of all models and their combinations. While individual systems exhibit minor variations, the overall trends remain consistent. We present the following dipeptides: AC Figures 23 and 24, AP Figures 25 and 26, ET Figures 27 and 28, HP Figures 29 and 30, NY Figures 31 and 32, RV Figures 33 and 34, and TD Figures 35 and 36.
Notably, Figures 25, 29 and 35, the Diffusion model fails to produce stable simulations, resulting in divergence and visible artifacts. This instability appears across several systems, not only those containing proline. While we found in our experiments that introducing dropout can help to stabilize the simulation, chose not to include it since all other methods remain stable without it. In particular, these cases highlight that restricting training to a smaller diffusion cittssioonimeeeeee subrange, as is done with Mixture, improves simulation stability.
telounss
be
Nodde Namcs
iid
sim
Figure 23: AC: We compare the free energy plot the dihedral angles y,, tb for all presented methods for iid sampling and Langevin simulation.
distance
Free energy projection of
Free energy projection of
Energy NNT
iid
Turahhaade
Distance A
distance
Free energy projection of
Free energy projection of
Tetennce Malum wit
Energe ATT
sim
Enwegee
2/
Figure 24: AC: We compare further metrics between iid sampling and Langevin simulation. We compare the
C. C., distance for the dipeptides and also the free energy projections along the dihedral angles wb.
34

---

iid
Reference
emisenn
Tue for cne
Muaue
sim
Figure 25: AP: We compare the free energy plot on the dihedral angles y, v for all presented methods for iid sampling and Langevin simulation.
C, distance
Free energy projection of
Free eneryy projection of
Netteene Tex
109
Enorgy Ae?
Enoogy Fay
iid
Dath
Distance in.A
Ge Cu distance
Free energy projection of
Free eneryy projection of
tor
Energr hat
Energg chr
sim
Soth
Distance 1n
Figure 26: AP: We compare further metrics between iid sampling and Langevin simulation. We compare the C.-Ca distance for the dipeptides and also the free energy projections along the dihedral angles wb.
No One
Mudure
iid
Leference
Mieure
Tecoe Plancs
sim
Figure 27: ET: We compare the free energy plot the dihedral angles yp, v for all presented methods for iid sampling and Langevin simulation.
35

---

Co C. distance
energy projection
energy projection
Oillaiin
Two One
Energy Ned
iid
Distance bn
C, C. distance
Free enerpy projection
eneray projection
hetenncc
Tuo For One
Energe AAT
sim
Distance i.
Figure 28: ET: We compare further metrics between iid sampling and Langevin simulation. We compare the
Ca Ca distance for the dipeptides and also the free energy projections along the dihcdral angles yo, 2b.
idd
Reference
prbicen
We for car
Muntre
Bmm
sim
Figure 29; HP: We compare the free energy plot on the dihedral angles şo, d for all presented methods for iid sampling and Langevin simulation.
C. distance
Free energy projection
Free energy projection of
139
Energy NoT
iid
Wut Thrahhaade
Energy
Distance A
distance
Free energy projection
Free energy projection of
Teternce
139
Malum wit
100
Energe AUT
sim
2/0
Figure 30: HP: We compare further metrics between iid sampling and Langevin simulation. We compare the
Cu- C.. distance for the dipeptides and also the free energy projections along the dihedral angles wb.
36

---

iid
seterecce
sim
Figure 31: NY: We compare the free energy plot on the dihedral angles y, v for all presented methods for iid sampling and Langevin simulation.
C, distance
Free energy projection of
Free energy projection
Anory//AuT
iid
Dath
Distance 1n.A
Ga Cu distance
Free energy projection of
Free energy projection
for
Energy NNT
sim
Soth
Distance 1n
Figure 32: NY: We compare further metrics between iid sampling and Langevin simulation. We compare the C.-Ca distance for the dipeptides and also the free energy projections along the dihedral angles tb.
Nelorence
Forske pesttrgaoo
iid
Reference
The
Midure
Tecoe Planes
sim
Figure 33: RV: We compare the free energy plot the dihedral angles y, v for all presented methods for iid sampling and Langevin simulation.
37

---

Co C. distance
energy projection
energy projection
Oillaiin
Two One
Energy Ne?
iid
7
Distance 1n
Cu C. distance
Free enerpy projection
eneray projection
hetenncc
Tuo For One
Bnerps NoT
sim
G
Distance i
Figure 34: RV: We compare further metrics between iid sampling and Langevin simulation. We compare the
Ca-Ca Ca distance for the dipeptides and also the free energy projections along the dihedral angles y, 2b.
idd
Reference
Dificen
We for
Manurr
sim
Figure 35: TD: We compare the free energy plot on the dihedral angles şo, ai for all presented methods for iid sampling and Langevin simulation.
C. distance
Free energy projection of
Free energy projection of
Energyihe
iid
Thrahhrabee
Distance A
distance
Free energy projection of
Free energy projection of
Tetennce Malum wit
Everpe ArT
sim
Eneege
2/
Figure 36: TD: We compare further metrics between iid sampling and Langevin simulation. We compare the
Cu- C.. distance for the dipeptides and also the free energy projections along the dihedral angles wb.
38

---

D
Societal Impact
Our work focuses on improving the efficiency of molecular sampling and simulation. We consider this research foundational, with the potential to accelerate applications such as drug and material discovery. While we do not identify any immediate risks, the technology could be misused, for example, in the development of biological weapons. Furthermore, our method currently does not provide formal guarantees, which poses a risk of misleading downstream research if the method produces incorrect or biased results.
39

---

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

---

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

---

The authors should provide scripts to reproduce all experimental results for the new reproducible, they 5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: We have used synthetically generated data that explain in detail how to produce (compare Appendix C.4), publicly available datasets (compare Appendices B.3 and B.5), publicly available libraries (compare Appendix B.7) and make code available at https:: http//ggttttbt httpg/gthb..c.oooogggeeeee conn/noogroup/Scoeet
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (kttps://nips.cm
guuides/Codeuummi
for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:
/ //wips.cc///////// httpp//cccpbblcc/guuttccpee gtaebiiigggidis////eeeoeo Stte://ooeSSSemieeeeeeeeee wttpbbmmissioooooooyy. for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide all technical details for the experiments either in the main text in Section 5 or in Appendix B. Further, we make our code available at attps://githuu: httaa//co/ooogoo////rreeeee which contains all settings and hyperparame- ters.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yss]
Justification: We report the mean and standard deviation in the experiments where suitable (in the main text compare Table and Table 2).
Guidelines:
The answer NA means that the paper does not include experiments.
42

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative eтror rates).
show in tables or
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We list the computational resources we used for the experiments in Ap- pendix B.6 and the runtime in Appendix C.3.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: http:///rrpp.pccppl/////e/e////e
Answer:
[Yes]
Justification:
We believe that our research adheres to the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: We have discussed societal impacts in Appendix D.
43

---

Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (c.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymemt of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
We do not believe that any safe guards needs to be in place, as our approach does not pose any direct risks. As for the data, only widely used data and commonly used resources have been used and, hence, does not pose risk for misuse.
Justification:
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (c.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We have properly credited all works, for the datasets compare Appendices B.3 and B.5 and software licenses are listed in Appendix B.7.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
44

---

The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if it has changed) should provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We make it clear in Section 5 that our code is publicly available at https: https://githbb com www./ooo//eoeeeeeeee /ttp///ooooo//ree/// and provide information in README how to use the code.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
No crowdsoureing or human subjects are involved in this research.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: No human subjects are involved in this research.
45

---

Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: LLMs have been used to implement standard methods, and for writing, editing, and formatting purposes of this document and hence no further declaration is required.
Guidelines:
The answer NA means that the core method development in this research does involve LLMs as any important, original, or non non-standard components.
Please refer to our LLM policy pottts://neur/e cc/ /ct.:///oooeeeeccc////////// for what should or should not be described.
46


Reviews:
[{"review_id": "dweCJcqLep", "reviewer": "Reviewer_GH1q", "summary": "This paper explores the connection between the score of diffusion models and the potential energy surface of a molecular system. The authors show that there is a mismatch between samples generated via reverse diffusion and Langevin dynamics using the score as a force field. To fix this, the authors augment the standard denoising score matching loss with a Fokker–Planck residual penalty that biases the learned energy/time evolution to satisfy the Fokker–Planck equation. They further reduce computational cost by partitioning diffusion time into segments and training a Mixture-of-Experts. They show empirically that they can train consistent Boltzmann Generators.", "questions": "1.\tCould you provide a more detailed analysis of the non-smoothness of the resulting PES? (see above for more detail on the question) How smooth are geometry optimisations for example? \n2.\tCould you test a more diverse dataset, potentially of small molecules? With more diversity in the functional groups present. Does the FP regularisation truly transferrable learn more consistent scores? This is of interest as there is no guarantee of the score satisfying the FP equation from an architectural perspective.", "limitations": "-\tOverall adequately addressed, however the problems associated with a non-smooth PES do not seem adequately investigated.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "-\tThe clarity of the paper is very strong as is the overall quality of the text.\n-\tFigure 2 nicely illustrates the need for the added FP regularisation. \n-\tThe results are convincing and the free energy plots informative. \n-\tThe transferability of the approach could be probed more. Is it transferrable across diverse molecular datasets? Rather than just similar dipeptides? \n-\tThe authors mention that the MoE approach is more expensive, but a detailed cost comparison of multiple inferences is not provided. \n-\tThe use of only a single MoE and the associated non smoothness of the potential energy surface is not thoroughly investigated and appears to me as a significant limitation of the approach. The artifacts may appear in kinetic observables and may produce artifacts in other relevant observables not investigated in this approach. Could the authors quantitatively investigate the non-smoothness of the potential energy surface, by probing the forces as the model changes between mixture of experts? How significant can these differences be? \n-\tHow does the proposed method compare to recently proposed long-stride MD simulations such as FlashMD, BoostMD, TrajCast, Timewarp?", "was_revised": true, "final_justification": "The rebuttal phase addressed many of my questions. It is clear that the paper has the potential for high impact in at least one sub-area of AI. I particularly appreciate the more direct comparisons to standard force-matching approaches. I would encourage the authors to further benchmark their simulation model as a force field, particularly by examining the smoothness of the potential energy surface. Their response that such investigations will be pursued in future work is reasonable.", "rebuttal": "We thank the reviewer for their insightful review and questions. We appreciate their judgment that our paper is a significant contribution to this conference. We now address their questions and concerns individually.\n\n> The authors mention that the MoE approach is more expensive, but a detailed cost comparison of multiple inferences is not provided.\n\nThe MoE approach introduces essentially *no* computational overhead during training or inference, as only a *single* model is active at any given diffusion time $t$. In fact, Appendix C.3 shows that MoE can reduce training and significantly reduce inference time by up to 50% in our experiments. These performance gains are achieved by assigning simpler, smaller models to larger diffusion times with a lot of noise, where high precision is not required.\n\nThis design differs from conventional MoE approaches that combine multiple experts per input or rely on overfitting specialized models. Our MoE is a time-partitioned scheme, optimized for computational efficiency, and ensures consistency across different $x$ at the same diffusion time.\n\nWe will make this clearer in the next revision.\n\n> The use of only a single MoE and the associated non smoothness of the potential energy surface is not thoroughly investigated and appears to me as a significant limitation of the approach. The artifacts may appear in kinetic observables and may produce artifacts in other relevant observables not investigated in this approach. Could the authors quantitatively investigate the non-smoothness of the potential energy surface, by probing the forces as the model changes between mixture of experts? How significant can these differences be?\n\n> Could you provide a more detailed analysis of the non-smoothness of the resulting PES? (see above for more detail on the question) How smooth are geometry optimisations for example?\n\nWe thank the reviewer for raising this important point. As noted above, for each diffusion time $t$, only a single model is trained and used. This means that at diffusion time $t=0$, which is required for simulation, only a *single* independent expert model is used. This model is essentially a small-time expert and was trained on the interval $(0, 0.1)$. Since the entire simulation trajectory remains within this interval, only one model is used throughout, and thus no model switching occurs during simulation or training. As a result, the potential energy surface (PES) and corresponding forces remain smooth within this regime, and the MoE architecture does not introduce discontinuities that could affect kinetic observables.\n\nDiscontinuities in the score can occur during iid sampling, where samples are generated across the full diffusion timeline by stitching together different experts. In this case, the score is not guaranteed to be continuous at the expert boundaries. We experimented with smoothing transitions between models by interpolating scores or jointly training experts with overlapping time intervals, but observed no qualitative improvements in sample quality or training stability. Thus, we opted for the simpler approach of independently training each expert.\n\nWe acknowledge that this distinction between simulation and sampling use cases may not have been made sufficiently clear in the current version of the manuscript and will revise the text accordingly.\n\n> How does the proposed method compare to recently proposed long-stride MD simulations such as FlashMD, BoostMD, TrajCast, Timewarp?\n\nWe thank the reviewer for highlighting these related approaches. Methods like FlashMD, TrajCast, and TimeWarp all accelerate MD by directly predicting updated configurations at larger time steps. In contrast, our approach learns a (coarse-grained) force field, derived as the gradient of a learned energy function, and then uses that force field to drive simulations. By explicitly modeling the energy, our method can be combined with enhanced-sampling techniques such as umbrella sampling and metadynamics.\n\nBoostMD proposes a framework to speed-up simulations by evaluating ML force fields only every N simulations steps, and evaluating a cheaper surrogate model for the intervening updates. By contrast, our method requires no force-field information during training, only configuration data. In principle, this allows us to drive simulations directly with our learned model, obviating any explicit ML force field. Whether such purely data-driven trajectories achieve sufficient accuracy in these domains remains to be seen.\n\nDirect comparisons are challenging, since each method has been validated on different systems. \n\nWe thank the reviewer for the suggestions and will add these in the related work section.\n\n> The transferability of the approach could be probed more. Is it transferrable across diverse molecular datasets? Rather than just similar dipeptides?\n\n> Could you test a more diverse dataset, potentially of small molecules? With more diversity in the functional groups present. Does the FP regularisation truly transferrable learn more consistent scores? This is of interest as there is no guarantee of the score satisfying the FP equation from an architectural perspective.\n\nWe thank the reviewer for this suggestion. While evaluating broader chemical diversity, such as small molecules with diverse functional groups, is indeed of interest, the computational cost of extending to such a dataset is significant and unfortunately beyond the scope of what we can address during the rebuttal period.\n\nThat said, we agree that demonstrating scalability beyond the dipeptide dataset is important. To this end, we extended our evaluation to two larger and more complex systems: **Chignolin** (10 amino acids) and **BBA** (28 amino acids), both of which are well-studied fast-folding proteins. These systems go significantly beyond the small coarse-grained dipeptides and serve as a step toward testing generalization to more diverse and larger molecular systems.\n\n**Chignolin Results:**\n|Method|iid JS|sim JS|iid PMF|sim PMF|\n|---|---|---|---|---|\n|Diffusion|**0.0036 ± 0.0001**|0.4351 ± 0.0141|**0.027 ± 0.000**|63.804 ± 0.372|\n|Two For One|**0.0036 ± 0.0001**|0.1023 ± 0.0008|**0.027 ± 0.000**|1.438 ± 0.019|\n|Mixture|0.0042 ± 0.0001|0.4336 ± 0.0075|0.033 ± 0.000|11.185 ± 0.430|\n|Fokker-Planck|0.0048 ± 0.0001|**0.0050 ± 0.0001**|0.037 ± 0.000|0.040 ± 0.001|\n|Both|0.0045 ± 0.0001|**0.0050 ± 0.0008**|0.035 ± 0.001|**0.038 ± 0.006**|\n\n**BBA Results:**\n|Method|iid JS|sim JS|iid PMF|sim PMF|\n|---|---|---|---|---|\n|Diffusion|**0.0043 ± 0.0000**|0.2014 ± 0.0043|**0.034 ± 0.000**|5.387 ± 0.144|\n|Two For One|**0.0043 ± 0.0000**|0.1162 ± 0.0021|**0.034 ± 0.000**|1.624 ± 0.107|\n|Fokker-Planck|0.0070 ± 0.0000|**0.0427 ± 0.0053**|0.060 ± 0.000|**0.438 ± 0.074**|\n\nIn both systems, we observe a substantial improvement in sampling-simulation consistency when applying Fokker-Planck regularization. These effects are also visible in the corresponding free energy surfaces: while the Two For One [1] model tends to overdisperse and fragment the distribution across multiple modes, our method better preserves the structure of the equilibrium landscape and recovers the correct metastable states. To our knowledge, our method is the first that can accurately sample and simulate the dynamics of BBA.  \n\nThese results support the conclusion that Fokker-Planck regularization improves sampling-simulation consistency even in substantially larger systems, without relying on architectural constraints. While these proteins are still relatively small compared to the full range of molecular diversity in chemistry, they demonstrate the scalability and consistency benefits of our approach in higher dimensions. We will include these additional results and dynamic analyses in the revised manuscript.\n\n> ### Limitations:\n\n> Overall adequately addressed, however the problems associated with a non-smooth PES do not seem adequately investigated.\n\nWe thank the reviewer for pointing this out. As clarified in our response above, the Mixture-of-Experts (MoE) design does not introduce non-smoothness in the potential energy surface (PES) during simulation, since only a single expert trained over a narrow diffusion time interval is used throughout. This ensures that the PES remains smooth along simulation trajectories. We acknowledge that this was not made sufficiently clear in the original submission and will revise the manuscript to better explain this point.\n\n-------\n\n[1] Arts et al. (2023). Two for one: Diffusion models and force fields for coarse-grained molecular dynamics"}, {"review_id": "RK28nNpt7j", "reviewer": "Reviewer_c2SZ", "summary": "The authors look at the problem of molecular dynamics, in particular sampling from boltzmann distributions, through the use of diffusion models and langevin sampling. They highlight that current score-based diffusion models, while immensely successful at modeling images, are not as efficient at modeling MD simulations which can be consequently improved with a combination of (a) energy parameterization as opposed to score parameterization, (b) mixture-of-experts formulation for the architecture, and (c) Fokker-planck regularization to enable consistency. Their experiments indicate that the combination of these approaches leads to better simulation results, while still maintaining good iid sampling.", "questions": "- It was unclear how the authors successfully train energy-based parameterization since it would require taking the gradient of the model w.r.t the state as well as consequent gradient w.r.t the parameters of the model. Is this operation efficient or are there any approximations leveraged?\n- Are equations 10-12 somehow related to the hutchinson trace estimator or the estimator used in [1]?\n- What is the purpose of using a mixture-of-experts system instead of just increasing the model size/capacity with number of parameters?\n\n[1] Albergo, Michael S., and Eric Vanden-Eijnden. \"Nets: A non-equilibrium transport sampler.\" arXiv preprint arXiv:2410.02711 (2024).", "limitations": "No additional limitations.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 1, "originality": 1, "strengths_and_weaknesses": "**Strengths**\n\n- The authors provide tractable means of leveraging fokker-planck regularization through finite-difference based estimation of divergence and derivative terms.\n- They showcase that such regularization leads to improved performance for both simulation and iid sampling, when combined with mixture-of-experts architecture.\n\n**Weaknesses**\n\n- The authors conjecture that learning a score function through denoising score matching leads to unreliable estimates closer to the data distribution. However, this has been a well known property in literature, and is in particular showed in detail in [1].\n- The authors aim to solve the main problem of simulation which requires score at $t=0$ for which they leverage the learned score. However, typically for these applications energy, and thus forces, are available at $t=0$ so what is the aim for using the learned score from denoising score matching? I believe an insight into this was lacking in the current draft.\n- While it is evident that enforcing the fokker-planck equation at intermediate time points would lead to consistent sampling, what is the reason for why it would lead to better estimation of score at $t=0$ since it is only the boundary point?\n\n[1] De Bortoli, Valentin, et al. \"Target score matching.\" arXiv preprint arXiv:2402.08667 (2024).", "was_revised": true, "final_justification": "The authors have resolved all my concerns regarding both the use case as well as the importance of their work. Their additional experiments demonstrate the use case of simulation in scientific applications and their proposed method of including a regularization from the FPE directly leads to better results. I am happy to recommend accept.", "rebuttal": "We thank the reviewer for their comments and will use the feedback to further clarify the motivation behind our approach. Below, we address the questions in detail.\n\n> The authors conjecture that learning a score function through denoising score matching leads to unreliable estimates closer to the data distribution. However, this has been a well known property in literature, and is in particular showed in detail in [1].\n\nWe agree that the limitations is well known, and we reference multiple works in our manuscript [A,B,C] that highlight this issue. We do not present this as a novel observation, but address it empirically by showing that reducing deviation from the Fokker-Planck equation improves the accuracy of the learned score, particularly near the data distribution.\n\nImportantly, approaches like target score matching [1] are not applicable here, as they require access to force labels, which we explicitly avoid (see next answer). We will gladly cite the suggested literature [1] as further support for the known limitations.\n\n> The authors aim to solve the main problem of simulation which requires score at $t=0$ for which they leverage the learned score. However, typically for these applications energy, and thus forces, are available at $t=0$ so what is the aim for using the learned score from denoising score matching? I believe an insight into this was lacking in the current draft.\n\nOur work specifically focuses on settings where force labels at $t=0$ are not available and forces cannot be computed, and hence are not available for simulation or training.\n\nThis situation arises frequently in coarse-grained modeling, where the atomistic resolution is reduced and the underlying all-atom fine-grained potential (and thus the forces) are no longer well-defined. In our approach, we learn from samples alone, which allows for arbitrary levels of system abstraction (see Sec 1), similar to prior work [C].\n\nEven in systems where forces could be computed, practical limitations often prevent this. E.g., obtaining accurate forces using quantum methods such as DFT can be prohibitively expensive, while collecting configurations is often much cheaper. Also, in many experimental settings, one can measure samples, but not the forces acting on the system.\n\nAs our approach does not require forces, it can also be used in mixed-data settings where we want to train e.g., using simulation and real-world data.\n\nWe agree that this distinction could have been made clearer in the manuscript and will revise the introduction to better motivate the broader applicability of learning from samples alone.\n\n> While it is evident that enforcing the fokker-planck equation at intermediate time points would lead to consistent sampling, what is the reason for why it would lead to better estimation of score at $t=0$ since it is only the boundary point?\n\nWhile the Fokker-Planck equation is indeed enforced over intermediate time points, its inclusion improves the estimation of the score at $t=0$ through the following: When $t$ is small, the target energy (i.e., $-\\log p_t$) or score function can be non-smooth, and denoising score matching (DSM) becomes numerically unstable (as discussed in Sec 2.1). In contrast, for larger $t$, the energy function becomes smoother, and DSM tends to yield more stable and accurate estimates.\n\nThe Fokker-Planck equation provides a principled way to relate energy functions across time. By minimizing the Fokker-Planck residual during training, accurate estimates at large $t$ can be *propagated back* toward $t=0$, improving the estimation near the data distribution. Indeed, we observe in our numerical experiments that this leads to improved consistency and accuracy of the energy and score functions around $t=0$. This is reflected in our experiments in Sec 5. Fig 6 shows that our approach reduces the deviation from the Fokker-Planck equation at $t=0$ by nearly two orders of magnitude.\n\nTechnically, our models are evaluated at a small $t > 0$ to ensure numerical stability and define gradients cleanly. In practice, we observe no difference when using $t=0$ during simulation, suggesting the learned score remains well-behaved at the boundary.\n\nEmpirically, we can show that this regularization leads to better learned scores at $t = 0$ (see Sec 5). We also applied this approach to the protein BBA (28 amino acids), where to our knowledge, we are the first to provide consistent sampling and simulation [C] (see table below). We believe that this further strengthens our point that Fokker-Planck regularization improves the consistency even at the boundary.\n\n**BBA**\n|Method|iid JS|sim JS|iid PMF|sim PMF|\n|---|---|---|---|---|\n|Diffusion|**0.0043 ± 0.0000**|0.2014 ± 0.0043|**0.034 ± 0.000**|5.387 ± 0.144|\n|Two For One|**0.0043 ± 0.0000**|0.1162 ± 0.0021|**0.034 ± 0.000**|1.624 ± 0.107|\n|Fokker-Planck|0.0070 ± 0.0000|**0.0427 ± 0.0053**|0.060 ± 0.000|**0.438 ± 0.074**|\n\nWe thank the reviewer for raising this point and will incorporate a more detailed explanation of this effect in the manuscript, including a discussion of the boundary behavior and its implications for score estimation.\n\n> It was unclear how the authors successfully train energy-based parameterization since it would require taking the gradient of the model w.r.t the state as well as consequent gradient w.r.t the parameters of the model. Is this operation efficient or are there any approximations leveraged?\n\nIn Sec. 3.2, we note that an energy-based parameterization requires computing gradients of the model output with respect to the input. This is implemented by performing a backward pass through the network. While this introduces additional computational overhead, roughly a factor of 2, it remains tractable and has been used before [C,D].\n\nAs shown in App C.2, the energy-based parameterization improves stability and simulation quality and is crucial for obtaining reliable forces for MD simulation. Similar behavior has been reported in [C].\n\nWe will revise the manuscript to point to App C.2 and include the computational overhead introduced by the energy-based parameterization.\n\n> What is the purpose of using a mixture-of-experts system instead of just increasing the model size/capacity with number of parameters?\n\nRather than increasing the size of a single model, we chose a time-based mixture-of-experts (MoE) approach for several practical reasons (see Sec 3.3).\n\n1. Increasing model capacity increases training time, memory consumption, and inference cost. In contrast, our MoE setup trains separate, smaller models on disjoint time intervals where only one expert is active at a time. We effectively gain the benefits of a larger model without the overhead.\n\n2. MoE allows us to match model complexity to the specific demands of each diffusion time range. Near $t=0$, the score function must resolve fine-grained structure, motivating the use of larger models. At higher $t$, the data is noisy and can be modeled effectively with smaller networks. \n\n3. Dividing a time-dependent learning task into subproblems helps avoid optimization issues due to sharp transitions or local minima, and allows each model to focus on a simpler subproblem. This approach is common in the numerical PDE literature [E], where a separate network is trained per time segment.\n\n4. The performance benefits of MoE go beyond scaling. Below, we compare models trained on alanine dipeptide using either a larger model (\"Wide\" with more hidden units and \"Deep\" with more layers) or MoE with three experts of equivalent total parameter count. MoE outperforms all single-model baselines in simulation accuracy, while not introducing any overhead.\n\n|Method|Training|# Parameters|iid JS|sim JS|iid PMF|sim PMF|\n|---|---|---|---|---|---|---|\n|Diffusion|**50 min**|650k|**0.0081 ± 0.0003**|0.0695 ± 0.0517|0.095 ± 0.003|1.047 ± 0.924|\n|Wide|67 min|2M|0.0082 ± 0.0003|0.0406 ± 0.0236|0.096 ± 0.003|0.467 ± 0.245|\n|Deep|111 min|2M|**0.0078 ± 0.0003**|0.0376 ± 0.0087|**0.091 ± 0.002**|0.478 ± 0.038|\n|MoE|**50 min**|2M|**0.0079 ± 0.0003**|**0.0264 ± 0.0085**|**0.093 ± 0.007**|**0.325 ± 0.113**|\n\nWe will add this ablation and further clarify the benefits of MoE in the revised manuscript, including a pointer to this connection with time-splitting approaches for solving time-dependent PDEs.\n\n> Are equations 10-12 somehow related to the hutchinson trace estimator or the estimator used in [1]?\n\nOur treatment of the divergence term in the weak residual formulation is closely related to the Hutchinson's trace estimator used by Albergo and Vanden-Eijnden [1] for computing the PINN objective. Both approaches utilize Gaussian perturbations of $x$ to obtain unbiased estimates of the divergence. However, it is important to note a key difference in how the residual is defined. As described in App. A.2, in our weak residual $\\tilde R(x,t)$, we apply the same Gaussian perturbation not only to the divergence term but also to other terms in the residual (see Eq. (30)). This means that when the strong residual $R(x,t)$ is exactly zero, the weak residual $\\tilde R(x,t) = \\mathbb{E}[R(x+v,t)]$ also remains exactly zero, regardless of the perturbation variance $\\sigma^2$, avoiding truncation errors. This distinction allows our weak residual formulation to maintain consistency even under large perturbations. We will clarify this point further in the revised manuscript and cite the work [1] appropriately.\n\n-------\n\n[A] Kim et al. (2022) Soft truncation: A universal training technique of score-based diffusion model for high precision score estimation\n\n[B] Koehler et al. (2023) Statistical efficiency of score matching: The view from isoperimetry\n\n[C] Arts et al. (2023) Two for one: Diffusion models and force fields for coarse-grained molecular dynamics\n\n[D] Song et al. (2019) Generative modeling by estimating gradients of the data distribution\n\n[E] Wight et al. (2020) Solving Allen-Cahn and Cahn-Hilliard equations using the adaptive physics informed neural networks"}, {"review_id": "uvXmrHMTXB", "reviewer": "Reviewer_TwQZ", "summary": "The paper addresses the inconsistency between molecular dynamics (MD) simulations and traditional sampling in diffusion models by introducing a Fokker–Planck-based regularization to the diffusion loss. To reduce computational costs, the authors use efficient approximations and a time-based mixture of experts architecture. Experiments on toy systems, alanine dipeptide, and transferable dipeptide settings demonstrate that their approach improves consistency between sampling and simulation. In short, the contribution of this paper lies in solving the inconsistency, i.e., mismatch between how we generate samples from a model and how we use it to simulate time-dependent behavior.", "questions": "1. Mixture intervals\n\nA simple question in Section 5. The Mixture of Experts (MoE) model is trained on diffusion time intervals (0, 0.1), [0.1, 0.6), and [0.6, 1.0). Could the authors clarify the intuition or practical rationale behind choosing these specific cutoffs? Were they based on empirical findings, model behavior (e.g., Fokker–Planck error), or heuristic choices?\n\n2. High variance of Mixture of experts in simulations results\n\nIn Tables 1 and 2, the Mixture of Experts model shows relatively high variance in the simulation metrics compared to other methods. This seems like an interesting and potentially important observation. Could the authors offer an explanation or intuition for this behavior?\n\n3. Experiment metric (JS)\n\nFrom section 5, it seems that the JS divergence is computed for the energy distribution. Could the authors state why this is done, instead of JS divergence on dihedral angles / TIC(time-lagged independent component) or pairwise distance as in Two For One? A simple results on JS divergence as in Two For One or why they have chosen energy would be good be enough to resolve my question.", "limitations": "yes, the authors adequately addressed the limitations and potential negative societal impact of their work.", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\n1. (Clarity) clear and organized presentation\n\nThe paper is clearly written and well organized, making it easy to follow. Most experiments are well designed and support the authors’ claims across various settings.\n\n**Weakness**\n\n1. (Quality) Quantitative transferable dipeptides results\n\nFor the transferable dipeptide, quantitative and qualitative results only for dipeptide AC has been presented in the main paper. Though there are qualitative results on other test dipeptides in the appendix, I do not see any quantitative results on other peptides.\n\n(Minor) Moreover, additional qualitative results for dipeptides could be helpful. Specially, most dipeptide show similar landscape to the Ramachandran plot in Figure 5 and 12, while dipeptides containing proline and glycine shows some different Ramachandran plot trends. Though there is one in Figure 14, one or two additional test peptides results with different Ramachandran plot trends could strengthen the claim for the generalization of the proposed method in the transferable dipeptide settings.\n\n2. (Quality) Efficiency of training and inference\n\nWhile the authors have emphasized that the proposed method achieves efficient training and inference, I feel that the experiment results fall short on this. The author have claimed efficiency in many aspects, such as \n\n- line 58 (contribution 2) achieve efficient training and inference\n- line 80 efficiently evaluate the loss\n- line 150 efficient approximation of the loss\n- line 183 MOE approach to allocate model capacity more efficiently\n\nThere are not results supporting this claim in the main paper, while there is table 7 in the appendix reports the training and inference time for the proposed method. However, methods including Fokker-Planck regularization terms shows several longer training time than the original model. \n\n3. (Significance) Dynamic results\n\nThe authors have experimented only to dipeptide system scales, while a major baseline in the paper “Two For One” has done experiments to fast folding protein system scales (chignolin, BBA, etc). Experiments on systems with this scales allows dynamics analysis as in Figure 5 of the Two For One paper. Testing two or three fast folding protein could strengthen the paper’s claim a lot, since this dynamics analysis is closely related to why we run simulations.", "was_revised": true, "final_justification": "The authors have well resolved almost all of my concerns and questions, therefore I modify the score accordingly.", "rebuttal": "We thank the reviewer for their thoughtful review and agree that results on larger systems would strengthen the paper. We now include fast-folding proteins, which further support our main claims. In the following, we address all points raised by the reviewer.\n\n> For the transferable dipeptide, quantitative and qualitative results only for dipeptide AC has been presented in the main paper. [...] I do not see any quantitative results on other peptides.\n\nThe quantitative results in Table 2 are in fact computed across the *entire test set*, not just the dipeptide AC. The table caption is mislabeled in this regard. The rest of the caption is accurate \"To compute the mean and standard deviation, we have averaged the metrics across the dipeptides from the test set.\" We thank the reviewer for pointing this out and we will fix this.\n\n> (Minor) Moreover, additional qualitative results for dipeptides could be helpful. [...]. Though there is one in Figure 14, one or two additional test peptides results with different Ramachandran plot trends could strengthen the claim for the generalization [...]\n\nThe dipeptides shown in App. C5 were selected randomly from the test set, as we found them to be representative in terms of model performance. We agree that including additional examples, particularly those that exhibit more diverse behavior, could better support our generalization claims. We are happy to add such results to the appendix.\n\n> While the authors have emphasized that the proposed method achieves efficient training and inference, I feel that the experiment results fall short on this. The author have claimed efficiency in many aspects, [...]\n\n> There are not results supporting this claim in the main paper, while there is table 7 in the appendix reports the training and inference time for the proposed method. However, methods including Fokker-Planck regularization terms shows several longer training time than the original model.\n\nWe thank the reviewer for this observation and agree that our use of the term efficient could be more consistent and better contextualized. To clarify, we do *not* claim efficiency in terms of wall-clock training time for models with Fokker-Planck regularization, which we acknowledge introduces overhead (see conclusion/limitations in Section 6 and Appendix C.3). Below, we provide additional context for each instance and will revise the manuscript to make these distinctions clearer:\n> line 58 (contribution 2) achieve efficient training and inference\n\nThis is found in our second contribution, where we introduce the mixture of experts (MoE) scheme as an independent contribution. MoE reduces cost by training smaller models for most of the diffusion timeline.\n> line 80 efficiently evaluate the loss\n\nThis statement appears in the background section and refers to the standard denoising score matching (DSM) loss, which has a closed-form expression for the VP SDE and can be computed efficiently.\n> line 150 efficient approximation of the loss\n\nHere, we refer to the *approximations* to make the Fokker-Planck regularization tractable. The original loss involves second-order derivatives, which would be intractable in high dimensions. Here we refer to the weak formulation and finite-difference approximation, which make the regularization feasible. That said, we agree that calling this \"efficient\" without qualification may be misleading, and we will revise the wording accordingly.\n> line 183 MOE approach to allocate model capacity more efficiently\n\nHere, efficiency refers to how model capacity is used, not training speed. When using a single model across all timesteps, capacity must be shared uniformly. In contrast, the MoE approach allows us to tailor model complexity to the difficulty of the subtask, improving how the model capacity is allocated. \n\nWe appreciate the reviewer's attention to this detail and will clarify these aspects in the final version.\n> [...]  a major baseline in the paper \"Two For One\" has done experiments to fast folding protein system scales (chignolin, BBA, etc). Experiments on systems with this scales allows dynamics analysis as in Figure 5 of the Two For One paper. Testing two or three fast folding protein could strengthen the paper's claim a lot, since this dynamics analysis is closely related to why we run simulations.\n\nWe thank the reviewer for this valuable suggestion. We agree that demonstrating scalability to larger systems would further strengthen our claims, especially in the context of dynamical consistency.\n\nFollowing this suggestion, we prepared results on **Chignolin** and **BBA**. As with smaller systems, we observe that diffusion models produce unstable simulations, and that \"Two For One\" introduces notable noise in the trajectories, or in the case of BBA, tends to fragment the distribution across multiple modes. Our Fokker-Planck-regularized models yield significantly improved simulation results, with better alignment between sampling and simulation.\n\n**Chignolin**\n|Method|iid JS|sim JS|iid PMF|sim PMF|\n|---|---|---|---|---|\n|Diffusion|**0.0036 ± 0.0001**|0.4351 ± 0.0141|**0.027 ± 0.000**|63.804 ± 0.372|\n|Two For One|**0.0036 ± 0.0001**|0.1023 ± 0.0008|**0.027 ± 0.000**|1.438 ± 0.019|\n|Mixture|0.0042 ± 0.0001|0.4336 ± 0.0075|0.033 ± 0.000|11.185 ± 0.430|\n|Fokker-Planck|0.0048 ± 0.0001|**0.0050 ± 0.0001**|0.037 ± 0.000|0.040 ± 0.001|\n|Both|0.0045 ± 0.0001|**0.0050 ± 0.0008**|0.035 ± 0.001|**0.038 ± 0.006**|\n\n**BBA**\n|Method|iid JS|sim JS|iid PMF|sim PMF|\n|---|---|---|---|---|\n|Diffusion|**0.0043 ± 0.0000**|0.2014 ± 0.0043|**0.034 ± 0.000**|5.387 ± 0.144|\n|Two For One|**0.0043 ± 0.0000**|0.1162 ± 0.0021|**0.034 ± 0.000**|1.624 ± 0.107|\n|Fokker-Planck|0.0070 ± 0.0000|**0.0427 ± 0.0053**|0.060 ± 0.000|**0.438 ± 0.074**|\n\nWe will include these results in the revised manuscript, along with a dynamic analysis akin to Figure 5 in Two For One, to support the evaluation of kinetic consistency in larger systems.\n\n> [...] The Mixture of Experts (MoE) model is trained on diffusion time intervals (0, 0.1), [0.1, 0.6), and [0.6, 1.0). Could the authors clarify the intuition or practical rationale behind choosing these specific cutoffs? Were they based on empirical findings, model behavior (e.g., Fokker-Planck error), or heuristic choices?\n\nThe choice of intervals was based primarily on empirical findings. We experimented with different numbers and placements of intervals and found that performance was generally robust; especially for the larger diffusion times, where the data is more heavily corrupted and the model behavior is less sensitive to the exact range.\n\nThat said, for the smallest interval used during simulation, we observed that setting the upper bound too low (e.g., < 0.1) made the model less robust to noise and led to poor generalization. Conversely, setting it too high introduces high variability in the data, making it harder for the model to focus on the detailed structures required for simulation.\n\nThe chosen partition seems to strike a balance between specialization and stability, and also regularizes the region with the highest Fokker-Planck error. We will include a discussion of this design choice in the revised manuscript to better motivate the interval selection.\n\n> In Tables 1 and 2, the Mixture of Experts model shows relatively high variance in the simulation metrics compared to other methods. This seems like an interesting and potentially important observation. Could the authors offer an explanation or intuition for this behavior?\n\nThe higher variance of the MoE model can be best explained with the alanine dipeptide example. Neither the MoE nor the single-model baselines consistently recover the low-probability region in simulation. However, while the single model tends to systematically miss this mode, the MoE model occasionally captures it. \n\nWe attribute this to the increased capacity of MoE, which allows it to \"remember\" more of the training distribution. However, since it does not explicitly address the underlying inconsistency at small diffusion times, the model may or may not recover rare modes depending on initialization or training noise, leading to the observed variance.\n\nThis further motivates our combination of MoE with Fokker-Planck regularization. We thank the reviewer for highlighting this observation, and we will clarify this point in the revised manuscript.\n\n> From section 5, it seems that the JS divergence is computed for the energy distribution. Could the authors state why this is done, instead of JS divergence on dihedral angles / TIC(time-lagged independent component) or pairwise distance as in Two For One? A simple results on JS divergence as in Two For One or why they have chosen energy would be good be enough to resolve my question.\n\nWe thank the reviewer for raising this point. We would like to clarify that we do *not* compute JS divergence on energy distributions as this would not be feasible in our coarse-grained setting, where the true energy function is unknown. Instead, all reported metrics are based on 2D projections such as dihedral angles or TICA coordinates, as suggested by the reviewer, consistent with the approach taken in Two For One.\n\nWe agree that pairwise distance (PWD)-based metrics provide a useful complementary view. Following the suggestion, we have computed JS divergence on the PWD histograms which can be found below. We will include these results in the revised manuscript and clarify the basis of all reported metrics.\n\n**Chignolin**\n|Method|iid PWD JS|sim PWD JS|\n|---|---|---|\n|Diffusion|**0.0001 ± 0.0000**|0.3817 ± 0.0009|\n|Two For One|**0.0001 ± 0.0000**|0.0082 ± 0.0000|\n|Mixture|0.0003 ± 0.0000|0.2045 ± 0.0004|\n|Fokker-Planck|0.0004 ± 0.0000|**0.0008 ± 0.0000**|\n|Both|0.0003 ± 0.0000|**0.0012 ± 0.0005**|\n\n**BBA**\n|Method|iid PWD JS|sim PWD JS|\n|---|---|---|\n|Diffusion|**0.0006 ± 0.0000**|0.1003 ± 0.0050|\n|Two For One|**0.0006 ± 0.0000**|0.0240 ± 0.0008|\n|Fokker-Planck|0.0011 ± 0.0000|**0.0104 ± 0.0010**|"}, {"review_id": "8GweuhFmYq", "reviewer": "Reviewer_ad79", "summary": "The paper proposes a regularization loss for diffusion model that is based on Fokker–Planck-equation (FPE) to enforce consistency between sampling and simulation modes in molecular dynamics (MD) simulations. The authors identify that diffusion models, when trained on equilibrium molecular distributions, exhibit inconsistencies between diffusion reverse process (backward inference) and use the first score function for Langevin simulation. To address this, they introduce a computationally efficient regularization method that ensures consistency. They validate their approach on a simple toy system, alanine dipeptide, and a transferable Boltzmann emulator for dipeptides demonstrating improved consistency and efficient sampling.", "questions": "1. For the definition of $L_{FP}$ (line 132), which does the D^{-2} denote or indicate? \n2. How do you obtain the $\\log p^{\\theta}(x, t)$ in equation 13 as your practical estimation? This term is not trivial for most of the diffusion models. Which is the concrete form of residual error (eq (8))?\n3. In equation (6), i do not quite understand the derivation in the following sense: the EBM that correlates p(x) and U(x) has its support over the whole configuration space of x (exp is non-negative). However, the learned score at t=0 $s^{\\theta}(x)$ does not necessarily capture the gradient field with arbitrary x, especially the region far from the data manifold. I am hesitant about the correctness of such derivation. Following such derivation, one can even simply train the network via denoising score matching to estimate the \"force\" around the data manifold. Training diffusion models on multi-tasking across multiple time steps yet throwing away the score network for time t>0 seems weird for sampling. Maybe you can elaborate or justify it furthermore.\n4. What does it mean by $NNET(x,t)$ in line 155 and line 156?\n5. In practical concern, why do we actually need to do “sim” instead of just “iid” for diffusion sampling? (\"sim\", \"iid\" are following the notation in the experimental section)", "limitations": "Yes.", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strength:\n- Introducing Fokker-Planck equation (FPE) based regularization loss for diffusion models to learn to sample consistently with the energy-based model formulation.\n- The experiments demonstrate that with the proposed loss, the diffusion works well by generating consistent sample using langevin dynamics (LD) simulator, on MB potential, ALDP and dipeptide tasks.\n\nWeakness:\n- While the paper empirically demonstrates improvements on toy systems, the theoretical underpinning connecting regularization to broader implications of diffusion models is insufficiently developed, which may not support the main claim of diffusion models in the paper.\n- The writing is not clear enough and seems to lack of necessary explanation (see equations).\n- The significance of the methodological advance seems unclear. I am not convinced how incorporating the FP loss improves the diffusion sampling from the experimental results.", "was_revised": true, "final_justification": "Summary. The work is well-motivated and original albeit the limited scale of systems studied. The major concerns hover around 1) the motivation of using \"simulation\" manner for sampling while the generative models (diffusion) are trained by i.i.d samples (via the score matching objective). It is hard to straight conclude whether the learned score function s(x, t) actually reflects the underlying physical kinetic like the physical force fields rather than hypothetical hallucination; 2) The experiments only involves toy peptides systems and BBA while the main baseline in this study (\"two for one\") had tested for several fast folding proteins. \n\nRegarding evaluation. A good motivation of using diffusion model in replace of traditional MD simulation is the prohibitive scaling law of system size with force field and exhaustive simulation. For the small system, the MD simulation (+enhanced sampling methods) can already work well, while the authors do not show promise of this method to work for larger systems. That being said, the regularization according to FP equations is very interesting and of good significance. After taking rebuttal and responses into consideration, I lean toward borderline rejecting this paper mainly due to the limited evaluation for properly justification of the proposed regularization.", "rebuttal": "We thank the reviewer for their feedback and concrete questions, which help us clarify the contributions of the paper. It seems that some aspects of the application context, particularly the importance of simulation, may not have been introduced with sufficient clarity. We therefore briefly revisit the motivation and have reordered the questions to address them more coherently.\n\n\n> In practical concern, why do we actually need to do \"sim\" instead of just \"iid\" for diffusion sampling? (\"sim\", \"iid\" are following the notation in the experimental section)\n\nThe goal of our work is to build diffusion models that support not only independent (iid) sampling but also molecular dynamics (MD) simulation (sim) by leveraging the learned score through Langevin dynamics. Previous work [1] has failed to recover the underlying forces, and the simulation produced different results from sampling. \n\nWhile iid sampling allows us to generate equilibrium configurations, simulation enables access to kinetic properties and dynamical information, which is crucial for many scientific applications. For example, to study how a molecule binds to a protein, we need temporally coherent trajectories; something that iid samples cannot provide. Whole fields of biochemistry (e.g., Transition Path Sampling) are interested in studying such dynamics. Furthermore, sim enables targeted exploration of rare states. Starting from a single low-probability conformation, we can simulate the system and uncover long-timescale behavior without the inefficiency of repeatedly sampling iid until the desired state reappears. \n\nIn addition, the ability to estimate forces and the potential from the score opens up opportunities for advanced sampling techniques (e.g., metadynamics, umbrella sampling), which would otherwise be inaccessible.\n\nWe agree that our introduction could better highlight these practical motivations, and we thank the reviewer for pointing this out. We will revise the introduction to make these use cases more explicit and will add supporting literature.\n\n> In equation (6), i do not quite understand the derivation in the following sense: [...] the learned score at t=0 $s^{\\theta}(x)$ does not necessarily capture the gradient field with arbitrary x, especially the region far from the data manifold. [...]\n\nWe would like to clarify that Eq 6 is a well-established identity that has been discussed in prior work [1] and follows directly from the Boltzmann distribution under the assumption that the model accurately learns the score (see Sec 2.2).\n\nThe more relevant question, that we explicitly address, is whether this identity can be realized in practice. As discussed in Sec 2.1, the score function at $t=0$ can be non-smooth, and training via denoising score matching (DSM) becomes numerically unstable. This is a known limitation of diffusion models and has been documented in the literature [1,2,3].\n\nWe improve the accuracy of the learned score at small $t$ with a Fokker-Planck regularization. Importantly, the introduced loss can be evaluated even on samples farther from the data distribution, and offers a mechanism to improve generalization in regions where DSM alone provides limited guidance. While no model can be expected to yield perfect gradients far from the training data, our method helps mitigate this limitation.\n\nWe will introduce a discussion of this in a revision.\n\n> While the paper empirically demonstrates improvements on toy systems, the theoretical underpinning connecting regularization to broader implications of diffusion models is insufficiently developed [...]\n\nOur work specifically targets diffusion models trained on molecular systems. In this setting, the connection between the model's score and the physical force field is well established, and simulation via Langevin dynamics becomes meaningful. While the Fokker-Planck regularization is derived from first principles, our claims remain scoped to molecular systems where these conditions apply. \n\nEmpirically, we show that our approach improves consistency between diffusion sampling (iid) and energy-based simulation (sim). These improvements are not limited to low-dimensional toy systems: we have extended our evaluation to larger molecules, including proteins such as BBA (28 amino acids). These systems go well beyond the dipeptide examples and exhibit the same trends:\n\n**BBA Results:**\n|Method|iid JS|sim JS|iid PMF|sim PMF|\n|---|---|---|---|---|\n|Diffusion|**0.0043 ± 0.0000**|0.2014 ± 0.0043|**0.034 ± 0.000**|5.387 ± 0.144|\n|Two For One|**0.0043 ± 0.0000**|0.1162 ± 0.0021|**0.034 ± 0.000**|1.624 ± 0.107|\n|Fokker-Planck|0.0070 ± 0.0000|**0.0427 ± 0.0053**|0.060 ± 0.000|**0.438 ± 0.074**|\n\nTo our knowledge, we are the first to provide consistent sampling and dynamics for such high-dimensional systems, where previous attempts [1] failed. These results reinforce our main claim: our regularization improves the consistency between sampling and simulation for molecules, even in high-dimensional settings. \n\n> The significance of the methodological advance seems unclear. I am not convinced how incorporating the FP loss improves the diffusion sampling from the experimental results.\n\nWe want to clarify that our goal is not to improve standard iid sampling but to enable consistent simulation for diffusion models by learning accurate scores for small $t$. This distinction is important: while we maintain comparable performance to state-of-the-art methods for iid generation, our focus is on enabling *simulation*, which traditional diffusion models do not support reliably [1].\n\nThe core idea behind the Fokker-Planck regularization is to improve the score estimation at small diffusion times $t$, where DSM becomes unstable. As discussed in Sec 2.1, these scores can be non-smooth and ill-conditioned, making accurate learning difficult. In contrast, at larger $t$, the score is smoother and easier to learn.\n\nThe Fokker-Planck equation provides a consistent relation between scores across time. By minimizing the Fokker-Planck residual, we leverage the more stable learning at large $t$ and *propagate this information back* to smaller $t$, leading to improved estimates near the data manifold. This can be seen in Figure 6, where the deviation from the Fokker-Planck equation at $t=0$ is reduced by nearly two orders of magnitude with our regularization.\n\nWe will revise the paper to make this distinction and motivation clearer.\n\n> [...] one can even simply train the network via denoising score matching to estimate the \"force\" around the data manifold. Training diffusion models on multi-tasking across multiple time steps yet throwing away the score network for time t>0 seems weird for sampling. Maybe you can elaborate or justify it furthermore.\n\nWe do not discard any of the models. For iid sampling, we use all expert models across the full timeline, recovering standard diffusion behavior. For simulation, we only use the expert trained on small $t$, which provides the most accurate estimate of the score near the data manifold. In the case where we only need simulation, we can train only this expert. \n\nNote that we cannot train a model purely at $t=0$, as the DSM loss is ill-defined there (see Sec 2.1). We need to train over an interval (which we do in MoE), and even then, this region is difficult to learn accurately. We show in Sec 5 that the regularization improves performance in this regime and enables MD simulations that standard diffusion training cannot support.\n\nWe will improve the paper by more clearly distinguishing the objectives of sampling and simulation, and by better motivating the role of MoE and regularization in each.\n\n> The writing is not clear enough and seems to lack of necessary explanation (see equations).\n\nWe thank the reviewer for their feedback. We agree that the paper would benefit from clearer motivation to reach a wider audience. We will revise the introduction to better explain why learning forces from data is important (see first answer). To improve clarity, we will include an algorithm outlining training and inference steps.\n\nNext, we will answer the questions about the equations.\n\n> For the definition of $L_{FP}$ (line 132), which does the D^{-2} denote or indicate? \n\nAs denoted in line 132: \"and define the corresponding loss as $L_{FP}[\\log p^{\\theta}](x, t) = \\lambda_{FP}(t) D^{-2} \\left\\Vert R(x, t) \\right\\Vert^2_2$, where $x \\in \\mathbb{R}^D$\", $D$ is the dimension of the data. \n\n> How do you obtain the $\\log p^{\\theta} (x, t)$ in equation 13 as your practical estimation? This term is not trivial for most of the diffusion models. Which is the concrete form of residual error (eq (8))?\n\nAs discussed in Section 3.2, computing $\\log p^{\\theta}(x,t)$ is non-trivial for standard diffusion models. However, we use an energy-based parameterization where the score is defined as the gradient of a scalar neural network output $s^{\\theta}(x,t) = \\nabla_x \\log p^{\\theta}(x,t)$, which gives us access to the energy.\n\nFor the concrete form of the residual used in training, we refer to Eq 11.\n\n> What does it mean by $NNET(x,t)$ in line 155 and line 156?\n\nNNET denotes any neural network. In this context, the question likely refers to the distinction between two parameterizations: modeling the score directly with a neural network, versus modeling the score as the gradient of an energy function. \n\nWe adopt the latter, energy-based formulation (see Sec 3.2), which gives us access to the unnormalized log-density required for computing the Fokker-Planck residual. We will try to make the notation clearer and will revise it to better highlight this distinction.\n\n-------\n\n[1] Arts et al. (2023). Two for one: Diffusion models and force fields for coarse-grained molecular dynamics\n\n[2] Kim et al. (2022). Soft truncation: A universal training technique of score-based diffusion model for high precision score estimation\n\n[3] Koehler et al. (2023). Statistical efficiency of score matching: The view from isoperimetry"}]

Author rebuttal:

Meta-review: The work introduces a Fokker–Planck-based regularization to the standard denoising score matching loss of a diffusion model. This improves the alignment between sampling from Boltzmann distributions using the score of diffusion and Langevin dynamics for molecular dynamics simulations.

The reviews have generally found the work potentially highly impactful. Particularly after the rebuttal, where the authors presented a broader experimental evaluation and stronger motivation for the approach, the consensus is that this paper provides a valuable contribution to the field.
Decision: Accept (poster)