Potential Score Matching: Debiasing Molecular Structure Sampling with Potential Energy Guidance
Authors: Liya Guo
Zun Wang
Chang Liu
Junzhe Li
Pipi Hu
Yi Zhu
Tao Qin
Venue: Accepted by TMLR ()

Abstract: The ensemble average of physical properties of molecules is closely related to the distribution of molecular conformations, and sampling such distributions is a fundamental challenge in physics and chemistry. Traditional methods like molecular dynamics (MD) simulations and Markov chain Monte Carlo (MCMC) sampling are commonly used but can be time-consuming and costly. Recently, diffusion models have emerged as efficient alternatives by learning the distribution of training data. Obtaining an unbiased target distribution is still an expensive task, primarily because it requires satisfying ergodicity. To tackle these challenges, we propose Potential Score Matching (PSM), an approach that utilizes the potential energy gradient to guide generative models. PSM does not require exact energy functions and can debias sample distributions even when trained on limited and biased data. Our method outperforms existing state-of-the-art (SOTA) models on the Lennard-Jones (LJ) potential, a commonly used toy model. Furthermore, we extend the evaluation of PSM to high-dimensional problems using the MD17 and MD22 datasets. The results demonstrate that molecular distributions generated by PSM more closely approximate the Boltzmann distribution compared to traditional diffusion models.

Published in Transactions on Machine Learning Research (08/2025)
Potential Score Matching: Debiasing Molecular Structure Sampling with Potential Energy Guidance
Liya Guo"
Myg22@mmill.tsinguua..co
You Mathematical Sciences Center, Tsinghua University Deportment of Mathematical Sciences, Tsinghua University
Zun Wang
zunuang@mmsccrooofl.com
Microsoft Research A14Science, Beijing, China
Chang Liu'
StunigLiu@micccooff..co
Microsoft Research A1fScience, Beijing, China
Junzhe Li School of Computer Science, Peking University
kjunahhe1028@sstupuuu....u
Pipi Hu
pisquare@micerooofl.com
Microsoft Research AIfScience, Beijing, China
Yi Zhu
yishu@tsinghua..de edu.cn
You Matheratical Sciences Center, Tsinghua Universily Yangi Lake Beijing Institute of Mathematical Sciences and Applications
Tao Qin
taoginn@microsoft.com
Microsoft Research Al4Science, Beijing, China
Reviewed on OpenReview: https: //
/ppere/v/pe openreview. net/ forum? id= hada://aaaa///oossssstttttoo
Abstract
The ensemble average of physical properties of molecules closely related to the distribution of molecular conformations, and sampling such distributions is a fundamental challenge in physics and chemistry, Traditional methods like molecular dynamics (MD) simulations and Markov chain Monte Carlo (MCMC) sampling commonly used but can be time- consuming and costly. Recently, diffusion models have emerged as efficient. alternatives by lcarning the distribution of training data. Obtaining an unbiased target distribution is still an expensive task, primarily because it requires satisfying ergodicity. To tackle these challenges, we propose Potential Score Matching (PSM), an approach that utilizes the potential energy gradient. to guide generative models. PSM does not require exact energy functions and can debias sample distributions even when trained on limited and biased data. Our method outperforms existing stale-of-the-art (SOTA) models the Lennard- Jones (LJ) potential, commonly used toy model. Furthermore, we extend the evaluation of PSM to high-dimensional problems using the MD17 and MD22 datascts. The results
WWrkk done during internship at Microsoft Research AI4Scicnce. 'Correspondng author.
Work done during internship Microsoft Research A14Science. $ CCorespooning author.
Corresponding author.

---

Published in Transactions on Machine Learning Research (08/2025)
demonstrate that molecular distributions generated by PSM more closely approximate the Boltzmann distribution compared to traditional diffusion models.
1 Introduction
Physical quantities of interest, such as free energy, are often determined by ensemble averages and are intrinsically linked to the distribution of molecular conformations Alder & Wainwright, (1959); ; Stoltz et al.
(0100). Traditional techniques for quantifying these physical quantities, notably Markov Chain Monte Carlo
(MCMC) sampling and Molecular Dynamics (MD), are well-ostablishdd Gelman & Rubin (9966) McCammon jet
(1977); ; Noé et al. (2019). However, these methods are computationally intensive, particularly for systems with high-dimensional molecules.
Unlike the inherently sequential sampling of MD and MCMC, generative models, espocially diffusion models, have emerged as efficient alternatives for generating independent and identically distributed (i.i.d.) samples Song et. al. (2020b); Ho et al. (2020); Song et. al. (2020a); Phillips et al. (2024); De Bortoli et. al. (2022)
Wu & Li (2023) Xu et al. (002 ; Hoogeboom et al. (2022): Woo & Ahn (2024). One such approach is
Denoising Score Matching (DSM) Song et al. (2020b). These models apply score function to learn the data distribution, which particularly useful in molecular systems where the equilibrium discribution adheres to the Boltzmann distribution. This distribution can be expressed as p(æ) e-E(a)/(kaTT) with kB denoting the Boltzmann constant, T the temperature, and E(x) the energy function of the system. Constructing a dataset that follows the Boltzmann distribution remains a challenging task, and failure to provide unbiased training data can result. in DSM overfitting to a biased distribution, leading to inaccurate observables.
Recent advances generative modeling have sought to incorporate the principles of the Boltzmann distri- bution for improved molecular sampling Wu al. (2024): Woo Ahn (2024): Bortoli al. (2024); Chen et al. 2024): ; Chung et al. 2023). Techniques such the Denoising Diffusion Sampler (DDS) and the
Path Integral Sampler (PIS) Zhang & Chen [002]; Vargas et al. (0023) have been developed to amortize the computational expense of traditional MCMC and MD methods. facilitating learning processes that do not necessitate equilibrium data. These methods typically rely on integration paths derived from ordinary differential equations (ODEs) and stochastic differential equations (SDEs), which, despite being innovative, still involve substantial computational time. The Iterated Denoising Energy Matching (iDEM) framework Akhound-Sadegh et al. (2024) represents another stride forward, employing energy functions for data gen- eration and introducing energy-guided sampling independent of the initial distribution. The complexity of its iterative loops and the absence of predefined initial data increase the computational overhead in high dimensional space, and at larger time scales, iDEM's sampling elficiency diminishes, demanding a greater number of samples to obtain precise outcomes. A more recent approach, Target Score Matching (TSM) Bortoli et al. (2024), integrates energy information with DSM to enable sampling from designated energy functions under the assumption of data unbiasedness.
Table Comparison of molecular sampling methods and their properties.
Inference
Doesn't Require
Efficiency in High Dim.
Doesn't Require
i.i.d. Samples
Exact Boltzmann
Energy Function
Data for Training
(Only Energy Labels)
MD HD/MCMC
DSM
X
DDS/PI
X
iDEM
V
TSM
V
X
PSM (ours)
In this work, we introduce the Potential Score Matching (PSM) mothod, which incorporates potential energy derivatives into generative models to more closely align the sample distribution with the Boltzmann distri- bution. Furthermore, PSM requires only force labels for the reference structures and obviates the need for
2

---

Published in Transactions on Machine Learning Research (08/2025)
explicit energy function. PSM leverages the characteristics of DSM to facilitate the generation of i..d. sam- ples, offering computationally efficient alternative to MD simulations. We provide a comparative analysis of PSM with related methodologies in Table I We also present a series of theoretical proofs demonstrating that PSM provides a more accurate estimation of the score function in the vicinity of = when training data biased and results in reduced variance near = The performance of PSM validated on both simple toy models, such as Lennard-Jones (IJJ) potentials, and more complex, high-dimensional physical datasets, including MD17 and MD22. Our experimental results consistently indicate that PSM outperforms baseline models.
2 Preliminaries
2.1 Gcore-Based Generative Modeling
Diffusion models operate through a two-step process Ho et al. [2020] Song et al. (2000a): (1) a forward diffusion process that incrementally adds noise to the data To until it converges to a Gaussian distribution at time 7, and (2) a reverse denoising process that reconstructs samples using a learned score function. The forward process defined by 2t = +0,6, where ai and Ør are noise schedule coefficients, and represents standard Gaussian noise. Two widely used frameworks are the variance- preeeccireeerviiggg (VP-SDE) and variance- exploding (VE-SDE) diffusion processes Song et al. (2020b). In VE-SDE, O2 and o mmin (Fmisx//m///
while in VP-SDE, a = = à S 8.ds and 0, = Vi-e S. B. Both VP and VE can be abstracted as dæ; = f(x,,) t) dt + g() dW, Following this, the reverse process, solves the reverse SDE from T back to 0 is
dæ, = g log (z:)) dt + g(t) dW.
(1)
The term Væ, log (c),, referred to the "score function' time t, unknown and is approximated by training a neural network 8p parameterized by 0, using the objective function ||sg - Va, log gi (-) How- ever, since the marginal distribution (xx)) intractable Vincent (2011). later works propose simplified and equivalent loss function based on the conditional distribution |8 Vz, log = where the
conditional distribution is modeled as a Gaussian transition kernel. In this case, we can leverage the identity
Va, log (e xz-1) de to instead train a network to predict the noise e directly. This leads to the commonly used loss function:
min Bi [2(1) Iee () + or Va. log (a:)12) < min Ey,,.. (X(() [|eg (x,, - <112] 0
(2)
where the expectation E is taken over time sampled from a uniform distribution U((0, 1)), the noised data points Hy, sampled from the distribution q(z,|<0), and the noise e, sampled from a standard Gaussian distribution. The function A(t) represents a weighting function that adjusts the importance of different time steps during the optimization process. In molecular systems, D, often denotes atomic positions or, in some contexts, discrete atomic types. Diffusion models have become increasingly prevalent in the generation of molecular structures and the prediction of their properties De Bortoli et al. (2022).
2.2 Boltzmann Distribution
Molecular systems at equilibrium are typically characterized by the Boltzmann distribution, with the target distribution given by p(zg) e- (((ma)/(kkT)) A profound connection exists between the score function utilized in diffusion models and the system's potential energy, V p(zo) x -V((o))((kg)) Nevertheless, har- nessing the Boltzmann distribution for generative modeling poscs several challenges. For molecular systems with analytically specified energy functions, the exact distribution is the normalized form of the Boltzmann factor. The normalization constant, typically expressed as an integral over all possible configurations, is often intractable and cludes a closod-form expression. Recent rescarch has integrated energy considerations into scsre-based model variants to improve molecular prodictions. Innovations include the introduction of an equivariant energy-guided SDE Bao et al. (2022) and novel score functions Janner et al. (022): Durumeric
3

---

Published in Transactions on Machine Learning Research (08/2025)
al. 2024): Phillips et al. (2024) Wang et al. 2024p; Akhound-Sadegh et al. (2024). Additionally, attain-
ing equilibrium presupposes ergodicity in the system, which is ensured by simulating molecular trajectories over extended periods. This requirement poses computational and temporal demands. In our approach, we address these issucs by harnessing the derivatives of energy to facilitate this process.
3 Potential Score Matching in Molecular Systems
In this section, we introduce our Potential Score Matching (PSM) method, which leverages the derivatives of potential energy to efficiently approximate the Boltzmann distribution. Consider a training sample æ0 drawn from the dataset, subjected to noise injection process defined a, a, + 0,C, where o, is a time-dependent noise schedule. We formalize the representation of the score function for molecules that adhere to the Boltzmann distribution as follows:
Theorem 1. In molecular system where molecules follow the distribution p(æo) x e eE/(kkoT) given æ, 0g + oee, the score function at time t. is an expectation of the force,
1 - Em
[va K
1 ate
Va, log p(a,))
Emo la
(3)
Ott
KBT
where & and F represent the potential energy and force of æo, respeclively. Thus, the PSM loss is defined
1
Va &
Lsmmadel Ba
Euz
X(t)
s0(2,,
(-
a
KBT
(4)
1 F
X(t) 80(2
Og KBT
where 1(0,1) denotes uniform distribution, and p(a))/00) is the conditional probability distribution of the noise-injecled sample given the 27-
The proof of Theorem is provided in Appendix B.1 The loss function equation also known the "score
loss", can be further expressed in terms of the "xo loss" and "e loss", as detailed in Luo (2022), and the
relationship among these three losses is explained in Appendix B.2
of V E(o)) Do xz
mog-madel = Bm Emo
A(t)
(5)
at
KBT T
ot F 2
Le-model 10, o(2 Euso
X(t) ca (x
(6)
a₁ KBT
where Dg and Eg are neural networks designed to approximate the original data point Dy and the noise term e, respoctively. Rocall that the formulas for the VESDE forward process that a, = 1, the loss can be written For simplicity, denote KBT = 1 in this section.
L - -EE..uro...Beerrpteelle...
((tt) -
F [|]
Related ideas are explored in Target Score Matching (TSM) Bortoli et al. (0244), which utilizes an explicit energy function to model the score. Building on this concept, we link energy derivatives to the score function and present PSM, which embodies the aforementioned score representation. By utilizing force labels as a proxy to the Boltzmann distribution, our approach reduces computational costs. Our method extends TSM from boy models with analytically defined energy functions to realistic, high-dimensional molecular systems. Unlike prior approaches that rely energy- energy-based MLPs, incorporate with a novel time embedding tailored for molecular data, later introduced in the the experimental section. Moreover, our theoretical findings assert that PSM remains robust even when trained on data that does not adhere to the Boltzmann distribution, thereby correcting biases in the training dataset and yielding samples that more closely resemble the true distribution.
preeeervin architectures
Now we denote that p(z,) the Boltzmann distribution and g(x))) the data distribution. Ideally, p(æv) should coincide with (())); however, when starting with biased distribution, these two may not be equal.
4

---

Published in Transactions on Machine Learning Research (08/2025)
We demonstrate that PSM can provide superior training process in terms of data fidelity, regardless of whether the training data is biased or unbiased. Since Vay logp(a;) = p( log pa the optimal solution satisfies:
arg min (p(a)) Ep(a s a(--)
s(x,, () Væ, log p()p( = Ep( (ea)) V log p(zo))
(7)
Thus, the ground truth is given by Ep(mgle)) [V log (xo)] and the Denoising Score Matching (DSM) method learns this same expectation. The following theorem establishes that at small time steps, PSM more closely approximales the Boltzmann distribution than DSM. A detailed explanation is provided in Appendix B..
Theorem 2. For the Boltzmann distribution p(æo), the data distribution g(Zo), and smoll t, PSM learns more accurate score function than DSM,
E glao| x VVoog( (x0)] (p(aolee)) (w.)) IV log (x))
B Bg(aoler)) (V log (aro)]
Bp(aoll.) (V log (20))
PSM
Ground Truth
DSM
Ground Truth
(8)
Furthermore, denoting the right-hand side as these two terms satisfies ||22|| - 11 2 11 +0(t",,
[ and the left-hand side as [|21, the difference between where 13 = -1((0,, xx,) > 0 in neighborhood of = 0.
Recent works Yang et al. (023): Phillips et al. 2024g have indicated that during the training of traditional diffusion models, there is a tendency to encounter large Lipschitz constants and significant. variance in relation to the time variable near = O, can be seen in Lemma [2 and 3 in Appendix Appendix [B.4] which have the potential to destabiliie the training process. These highlight the need for improved training strategies in the small-t regime. Encouragingly, Theorem [2] shows that PSM offers theoretically grounded debiasing mechanism by promoting convergence toward the Boltzmann distribution in molecular systems. This suggests the potential of applying PSM labels in the small-t region of diffusson models. We propose weighted combination of both losses, referred to as "Piecewise Loss" and "Piecewise Weighted Loss". These methods prioritize PSM for small t while favoring DSM at larger time values. We denote the noise adding time range of PSM loss as fpsm, and similarly denote the time of DSM label as fdsm-
Piecewise loss. Assuming kB = 1 and given a chosen time point lps we use the force loss (equation exclusively for E tpl. while employing DSM for t > tp as follows. Here, tpsm - [0, tpl and fasm - [tp.
Emolas (v lag)( ift < lp. Eanglr lan VV, log (x: zo)], if e [tp. L].
s (xt, t) =
a:
(9)
Piccewise Weighted loss. Since s (xt; t) = Emgla
LV æv logp (2o)
= Exollar (n. log (x, xo)] =
Emg jm
20-20 we consider another form of the score
of
(1ie, 1) Feao la,,
Wi -v
log (2t) + (1 (1-w))
Fo 2c
(10)
of
and the the loss is L 20, ((t) /| ca(()).)
we F (1 w/
We choose Wi time-varying
de
Or
function which satisfies that wo = 1 and W1 = O. For example, Wa = sigmoid(50 - 0.05)), f E [0, 0.1] and W. = when / E [0.1, 1]. We can see that this loss is a generalization of "Piecewise loss". With a random noise e, the "Piecewise loss" and the "Piecewise Weighted loss" can be unified as
CPSM = Er,.... [I t) - ((1 - WieE + w1e2)1 where e = e[tasm] E2 = - -F[tpem] ot faam E [0, 1] (tpsm)) (11)
where random noise, tpsm refers to the noise adding time range using F as label, and tasm is the element in its complement. The training process is shown in Algorithm 1
5

---

Published in Transactions on Machine Learning Research (08/2025)
Algorithm Potential Score Matching
Require:
Network Ea, total iteration N, data zo, forces F, weight function wh, noise schedule o,, and 0;; Select random noise
Determine the time when the labels F and random noise e act, fpsm and ddss-
while n < N do
(0, 1), I; ~ NN(aa of), e e(tasml 62 = Flymee 0t: LPSM = [,, ()..t)--o - (w)-e)) + 21 < Up VoLPsM);
end while
Our approach incorporating force information is designed to be sample-efficient and does not impose constraints on the sampling methodology. It compatible with common techniques employed in diffusion models, including the Euler method, Prediction-Correttion (PC) method, and EDM Song et al. (2020b):
Karras et al. 20022)
4 Experiments
Datasets. We assess the performance of our proposed Potential Score Matching loss across various settings, including toy models like the Lennard-Jones (LJ) potential Köhler et al. (2020), and more intricate molecular systems such as MD17 and MD22. The LJ configurations, LJ-13 and LJ-55, consist of 13 and 55 atoms, respectively, arranged in a three-dimensional space. Historically, energy- baaeey molecular sampling has largely concentrated on toy models Midgley et al. [A AAkhound-Sadege et al. (2024); Woo & Ahn (2024); Bortoli jet al. (2024) but our approach exlends this to nigher-dinensional real-world dalasets, such as MD17 and MD22. MD17 contains molecular trajectories for different molecules, with 9 to 21 atoms, sampled at 500 K, where the thermal energy T in the Boltzmann distribution equals 1 kcal/mol. In contrast, MD22 poses more demanding challenge due its greater complexity, featuring molecular dynamics (MD) trajectories from small peptide with 42 atoms to double-walled nanotube with 370 atoms Chmiela et al. sampled 400 to 500 K with time resolution of 1 fs.
To demonstrate our method's capability to correct biased training distributions and ensembles, we use Denoising Score Matching (DSM) Song et al. (020b) as a baseline and compare our results with those from iDE and Flow AIS bootstrap (FAB) in toy model contexts Sichound-Sadege et al. (2024); Midgley et al. (2022). Our experiments primarily utilize biased training sets, employing the first 1, 000 frames for LJ data and the first 5, 000 frames for MD reference trajectories. Further investigation into the impact of dataset bias is conducted in the ablation study (Section 4.2)
Diffusion settings and network. We evaluate our method using DSM, Piccewise, and Piccewiso
Weighted approaches, employing VESDE for denoising with equal time weight functions, X(I) = 1. In the Piecewise approach, PSM is applied for S 0.05, while DSM used for t [0.05, 1]. In the Piecewise Weighted setting, the label is combined as w; x psm +arget+((1ww)) x target, where wt = I=exp(5((--005)) During training, we employ the "c-model" loss function to learn the noise
Molecular systems often represent molecular structures as graphs, with atomic coordinates, numbers, and bonds defining a molecular graph. Symmetry and equivariance in these graph representations are crucial for accurate molecular interactions. We modify the Equíformer-v2 network based Liao et. al. (2024) by
adding the time embedding, more details about this network can be seen in Appendix A.2.2
For sampling process, we use the Prediction-Correctiin (PC) sampler Song et al. (020bb where each pre- diction step are followed by one correction step. Additional experimental details, including hyperparameter settings, are available in Appendix C
Evaluation metrics. To evaluate the physical plausibility of the generated molecular conformations, we assess the stability of molecules in the MD17 dataset, following the method in Fu et al. (022). A molecule

---

Published in Transactions on Machine Learning Research (08/2025)
is deemed unstable at time T if:
max [[x - =z(T)]l big > 0.5, (6.)
(12)
where bij is the equilibrium bond length between atoms and j. Tf no unstable molecules are detected, the dataset is considered stable. Moreover, we evaluate whether generated samples show an unbiased distribution by analyzing the statistical distribution of interatomic distances (h(r)), defined as:
1 h(r) N(N 1
5(r - Iz
(13)
where is the interatomic distance and 5 is the Dirac delta function. This metric indicates whether the generated samples follow unbiased molecular distribution. The mean absolute error (MAE) between the sampled h(r) and reference data h(r) also serves as numerical comparison, providing more intuitive results.
For further numerical insights, we provide the Total Variation Distance (TVD) to measure the maximum distribution distance between sampled and reference data. For probability distributions P and Q, TVD is defined as:
R2 1 PP(x) Q(x)) da, TVD(P,Q)=
equivalently for discrete distributions: TVD(P,Q) EnEx P(xx) Q(x)). For datasets with available analytical energy expressions, additionally provide energy distribution visualizations.
4.1
Main Results
4.1.1
Lennard-Jones Potential
The Lennard-Jones (LJ) potential encapsulates the fundamental principles of interatomic interactions: repul- sive forces dominate al short distances, while altractive forces prevail al longer ranges. It. is mathematically expressed as:
12
EWع (x)
2T
Σ
I'm
dig
-
diz
I'
(14)
Σ 8o3e (x) = 2
x a meaan
?
grot = = (w((22)) (z) + gose (x).
(15)
Here, Tm and F are constants that characterize the potential. In our experiments, we use the standard values I'm = and 1, in line with previous studies Köhler et al. (2020]; Akhound-Sadegh et al. (2024). The term gose represents the harmonic potential energy associated with particle displacements relative to the system's center of mass, Emeen
As indicated by equation 14 the function magnitude significantly increases when any interatomic distance dy approaches zero, presenting substantial challenges, particularly in high-dimensional systems. To examine the effect of biased training data, we plot the histograms of interatomic distance distributions, r djj. for LJ-13 and LJ-55, denoted as h(r), in Figure 1 samples when trained on biased distributions.
The results demonstrate that PSM effectively debiases
Compared to Akhound-Sadegh et al. (2024): Midgley et al. (2022), we also report the corresponding Wasserstein-2 (W-2) distance and lotal variation distance (TVD) based on three sets of samples obtained by sampling with different random seeds. Table 2 presents these metric comparisons, where we test the LJ potential using a random 10% subset of the reference data as the training set. Our experimental results show that applicable to the toy model and has advantages in high-dimensional situations. From another aspect, PSM also significantly reduces sampling time.

---

Published in Transactions on Machine Learning Research (08/2025)
(a)2.0
2.0
2.0
Ref data
Ref data
Ref data
DSM Sample
Piecewise Sample
Piecewise Weighted
1.5
1.5
1.5
h())
1.0
h(r
1.0
hir)
1.0
0.5
0.5
0.5
0.0
0.0
2
6
4
0
2
4
r (A)
(A)
(A)
(b)1.0
1.0
1.0
Ref
Ref
0.8
Sample
0.8
Piecewise Sample
0.8
Piecewise Weighted
0.6
0.6
0.6
M0
0.4
0.4
0.4
0.2
0.2
0.0
0.0
0.0
0
4 (A)
(A)
(A)
Figure 1: The distribution of interatomic distances h(r) for LJ potential using biased training data with sample size of 500. (a) Comparison of h(r) for DSM, Piecewise, and Piecewise Weighted losses (from left to right) against reference data for LJ-13; (b) Comparison of h(r) for DSM, Piecewise, and Piecewise Weighted losses for L.-55.
Table 2: Comparisons of PSM with FAB Midgley el al. (2022) and iDEM Akhound-Sadegh et al. (2024) on LJ-13 and LJ-55 results. Metrics include Wasserstein distance and atomic distance TVD, evaluated on three different seeds.
LJ-13 (39D)
LJ-55 (165D)
Sample W-2
Distance TVD
Sample W-2
Distance TVD
FAB iDEM
4.35 + 0.001
0.252 = 0.002
18.03 + 1.21
0.24 + 0.09
4.26 + 0.03
0.044 + 0.001
16.128 + 0.071
0.09 + 0.01
Piecewise
4.287 + 0.003
0.0582
0.001
15.894 + 0.003
0.047 + 0.000
Piecewise- Weighted
4.278 + 0.001
0.0585 0.001
16.054 + 0.008
0.023 + 0.002
4.1.2 Molecular Dynamical Data
Unlike other studies, our model proves effective on higher-dimensionall datasets, specifically testing on MD17 and MD22 datasets. In MD17, we consider molecules such as Uracil (12 atoms), Naphthalene (10 atoms), Aspirin (21 atoms), Salicylic Acid (16 atoms), Malonaldehyde (9 atoms). Ethanol (9 atoms), and Toluene (15 atoms). This dataset provides molecules with properties like Cartesian coordinates (in A), atomic numbers, total energies (in keal/mol), and atomic forces (in scal/mol/ A) at 500 K, where the thermal energy 7T the Boltzmann distribution corresponds 1 kkal/mol Additionally, we test the MD22 benchmark dataset, which includes Ac-Ala3-NHMe (42 atoms), Docosahexacnoic Acid (DHA) (56 atoms), Stachyose (87 atoms), DNA base pair (AT-AT) (60 atoms), DNA base pair (AT-AT-CG-CG) (118 atoms), Buckyball catcher (148 atoms), and Double-walled nanotube (370 atoms).
MD17 dataset. We first evaluale the stability of generated molecular configurations across DSM, Piece- wise, and Piecewise Weighted methods, as proposed in Fu et al. (022).. Stability is assessed on 1, 000 samples generated at the epoch with the lowest training loss over 1, 0.0 epochs. All samples meet the stability crite- rion. With no unstable molecules in any of the sampled data, we plot the distribution of interatomic distances for these MD17 datasets, as shown in Figure g The figure indicates that at the peak, DSM tends to capture
8

---

Published in Transactions on Machine Learning Research (08/2025)
(a)
[6]175
[c]1.2
(0) 0.8
1.25
0.8
0.6
0.75
1.00
0.6
0.4
0.50
0.25
0.2
0.00
(A)
(A)
(AA)
(g)².0
0.8
0.6
Reference
Reeceisee Sample
0.2
Sample
Plecewise Weighted Sample
0.0
0.0
0.0
6
6
(A)
(A)
(A)
Figure 2: The distribution of interatomic distances h(r) of r(A) for (a) aspirin, (b) ethanol, (c) malonalde-
hyde, (d) naphthalene, (e) salicylic acid, (f) toluene, and (g) uracil in MD17 dataset. The insets display the ball-and-stick representations of these molecules learned by Piccewise method. The sample number is 000.
greater fluctuations, and PSM can more accurately learn the neighboring structure. Table summarizes the mcan absolute errors (MAEs) of the sampled interatomic distances and total variation distance compared to the reference molecular data. These results demonstrate that, despite the inherent bias in the training data, PSM successfully debiases molecular samples.
Table 3: Comparison MAE interatomic distance and Total Variation Distance (TVD) for different
MD17 molecules with sample size 0,00
MAE h(r)
TVD
MD 17
DSM
Piecewise
Piecewise Weighted
DSM
Piecewise
Piecewise Weighted
Aspirin
0.1240
0.0570
0.0510
0.0662
0.0373
0.0290
Ethanol
0.1193
0.0845
0.0911
0.0663
0.0508
0.0501
Malonaldehyde
0.1003
0.0843
0.0630
0.0515
0.0492
0.0518
Naphthalene
0.1257
0.0701
0.1164
0.0716
0.0365
0.0592
Salicylic Acid
0.0799
0.0554
0,0584
0.0450
0.0424
0.0337
Toluene
0.0970
0.0769
0.0948
0.0536
0.0428
0.0490
Uracil
0.0747
0.0705
0.0560
0.0178
0.0421
0.0126
MD22 dataset. Similarly, Figure and Table present l(r)) and the numerical metrics of the MD22
dataset, showing improved performance with the PSM method. The sample size is 128.
Table 4: Comparison of MAF of M() and Total Variation Distance (TVD) for different MD22 molecules.
MAE of h(r)
TVD
MD
DSM
Piecewise
Piecewise Weighted
DSM
Piecewise
Piecewise Weighted
Ac-Ala3-NHMe
0.0873
0.0505
0.0486
0.0415
0.0224
0.0301
Stachyose
0.0384
0.0333
0.0424
0.0174
0.0161
0.0198
Docosahexaenoic acid
0.0602
0.0444
0.0518
0.025-
0.0221
0.0228
AT--A
0.0767
0.0620
0.0611
0.0336
0.0292
0.0285
AT-AT-CG-CG
0.0706
0.0655
0.0677
0.0341
0.0324
0.0329
Buckyball catcher
0.0978
0.0350
0.0567
0.0417
0.0158
0.0264
Dw-nanotube
0.1527
0.0919
0.0484
0.09416
0.0597
0.0256

---

Published in Transactions on Machine Learning Research (08/2025)
(8)
0.17
[6)
0.20
0.20
0.150
0.125
0.15
0.10
0.10
0.075
0.05
0.03
0.025
0.000
14
0.0
2.5
5,0
7510.012.11.0.7..6 5...0...550.0
0.0
25 5.0
012 520.0.0000
0.0 2.5 5.0 7.5 7.5.000022 022.515 720.00 (A)
r (A)
/ (A)
(A)
b.14
(gP.12
0.10
0.08
0.08
0.06
0.06
0.04
Reference
0.02
0.05
Piecewise Sample
0.00
DSM Sample
2
10 (A)
12
24
16
0.0
2.5 5.0
7.510.022 515.007
$20.0
10
15 (A)
20
25
Plecewise Weighted Sample
(A)
Figure 3: The distribution of interatomic distances (r(A)) for (a) Ac-Ala3-NHMe, (b) DNA base pair (AT-
AT), (c) DNA base pair LAT-AT-CG-CO) (d) Docosahexaenoe acid, (e) Stachyose, (f) Buckyball catcher,
(g) Dw nanotube in MD22 dataset. The insets display the ball-and-stick representations of these molecules.
comana com
24:255
29.45C
31.0%
10
25
15
10
UMA cimension
USAP cimorsion
UNAA Urioonsion
Figure 4: UMAP of SOAP feature for reference data, PSM, and DSM.
Debiasing conformations using SOAP descriptors. In molecular dynamics, to illustrate that PSM cxtends boyond recovering ensemble averages such as h(r) and achieves debiased results in terms of structural accuracy, we utilize the Smooth Overlap of Atomic Positions (SOAP) descriptor. SOAP describes the local atomic environment around central atom a as a smoothly varying atomic density function:
(r( Σ fo(riag(rra (om (fra):
(16)
where Tia = Ir:-ral I is the distance between the central atom a and its neighbor i. The function fe(r) ensures locality, s(r) is a radial basis function, and Yım (f) are spherical harmonics describing angular dependence. SOAP maintains a high-dimensional feature representation invariant under rotations. To visualize SOAP results, we apply UMAP, a nonlinear dimensionality reduction technique, for intuitive comparison.
To demonstrate PSM's capability in debiasing, we construct a deliberately biased training dataset informed by ablation study (Section 4.2 which reveals that the first 1,0 simulation frames exhibit structural deviation. Specifically, approximately 1, 000 frames) with these biased carly frames to introduce controlled bias into the training data. We then compare the resulting molecular distributions generated by PSM, DSM, and the reference data using UMMAP-prjeecto for the reference, PSM, and DSM are (28.2%, 71.8%), (29.4%, 70.6%), and (31.0%, 68.4%), respoctively. The close alignment between PSM and the reference indicates that PSM more effectively captures the underlying atomic environments and debiases the molecular distribution.
augment randomly selected subset (comprising 0.5% of the full trajectory,
SOAP features clustered via DBSCAN, As shown in Figure 4 the cluster proportions
10

---

Published in Transactions on Machine Learning Research (08/2025)
4.2 Ablation Study
In this section, we conduct ablation experiments on the debias capability of PSM and the amortized MD computation cost to illustrate the efficiency of PSM.
Debias ability. First, we investigate the impact of biased versus unbiased training data on the performance of DSM and PSM. To emphasize the effect of debiasing, we conduct experiments on both biased and unbiased datasets. Specifically, we use two different training sets: (1) a biased dataset consisting of the first 1, 000 frames of the trajectory, and (2) an unbiased dataset created by randomly selecting 10% of the reference data. Figure 5 (a) confirms that these datasets exhibit distinet bias characteristics. For this study, use the ethanol molecule example.
(a)
1b/2.00
(c)2.00
1.75
Fandom 0.1 percent
Reference data
1.75
Seferetce
Hecendee Gamplh
Referetce
Peoee famp
1.50
Aist 1000 data
1.50
Flccedine Ueighted Samp
Sample
1.50
DSM Sample
Picna se Waighted Sample
1.25
1.25
1.25
1.00
1.00
0.75
0.75
0.75
0.50
0.50
0.50
0.25
0.25
0.00
0.00
0.00
3
(A)
(A)
(A)
Figure 5: The distribution of interatomic distance comparisons. (a) Distribution of interatomic distances in
the reference dataset, the randomly selected subset, and the first 1, 000 frames. (b) Comparison of DSM, Piccewise, and Piccowise Weighted methods trained on the first 1,000 frames. (c) Comparison of the three methods when trained on randomly selected 10% subset of the reference data.
Figure 5 compares the sample distributions generated by DSM and PSM when trained on the first 1,000 frames. The results show that DSM mainly reflects the characteristics of the biased training data, whereas PSM successfully corrects the bias, producing samples that better align with the unbiased distribution. In contrast, Figure 5 (c) demonstrates that when the training data is initially unbiased, both DSM and PSM yield accurate estimates. These findings underscore the importance of data quality in DSM training and highlight PSM's capability to generating samples that more accurately approximate the Boltzmann distribution.
Amortized MD computation cost. We further compare PSM with molecular dynamics simulations on the aspirin molecule. For MD, we adopt the Quin force field Wang et al. (2023). following the original training protocol. We measure both the force field training time and the time required to generate sufficiently long, converged trajectories (treated as the sampling time). For PSM, we train the diffusion model for 500 epochs until convergence and mcasure the sampling time using the Euler discretization scheme. All experiments are conducted on a single NVIDIA A100 GPU. Our results show that PSM completes training in approximately 8 hours and requires only 10 minutes to generate 1,000 batches, each consisting of 1,000 time steps. In contrast, training the MD force field to a comparable accuracy takes around 24 hours, while achieving state-of-the-ar performance may require up to one week of training. Moreover, the MD simulation phase takes about 31 hours and 18 minutes to generate 210, 000 steps-tte the trajectory length used in the reference MD17 aspirin dataset. These results demonstrate that PSM offers substantially improved sampling efficiency over MD.
5
Conclusions
In this study, we introduce Potential Score Matching (PSM), a method designed to mitigate biases in non-Boltzmann distributions by incorporating force labels. Traditional molecular dynamics simulations are computationally expensive. While diffusion models offer an alternative by learning a score function generate samples that satisfy the training data distribution, obtaining training data that accurately follows
11

---

Published in Transactions on Machine Learning Research (08/2025)
the Boltzmann distribution is challenging. PSM relaxes these constraints. Theoretical analysis indicates that PSM provides training with lower variance and more accurate score estimations in the vicinity of I = 0, which allows for a concentrated training effort around this point. This has led to the development of two novel loss formulations: Piecewise and Piccewise Weighted losses. Empirical assossments validate the performance of PSM over Denoising Score Matching by elfectively debiasing data towards the Boltzmann distribution. Furthermore, our research extends the application of PSM to high-dimensional real-world datasets, such as MD22 and MD17, demonstrating its capability in handling complex molecular modeling challenges.
Additionally, flow matching has proven to be an effective method for generating molecular conformations and predicting properties Lipman et al. (2022); : Chen & Lipman (2023); y Miller et al. (2024). Progress in this area Woo & Ahn (204)); Domingo-Enriin Enrich et al. (2024) has resulted in unified frameworks that integrate diffusion models and flow matching techniques. Considering the established relation between the velocity field and the
+ ( - - s(x, where æ; atyi + 7:390 0g + 7xe, and you y1 are sampled from the noise and data distributions respectively, our PSM framework is well- positioned for extension to flow matching. This prospective expansion is earmarked for future exploration, while further theoretical insights and formulations for VPSDE and VESDE are delineated in Appendix D
score function, as expressed by v(z, = fi at
References
Tara Skhound-Sadegh, Jarrid Rector- Brotor,Brooks, Avishek Joey Bose, Sarthak Mittal, Pablo Lemos, Cheng-Hau Liu, Marcin Sendera, Siamak Ravanbakhsh, Gauthier Gidel, Yoshua Bengio, et al. Iterated denoising energy matching for sampling from boltzmann densities. arXiv preprinl arXiv:2402.0611 001112,,, 2024.
Berni Alder and Thomas Everett Wainwright. Studics in molecular dynamics. i. general method. The
Journal of Chemical Physics, (1()):559-48) 1959.
Fan Bao, Min Zhao, Zhongkai Hao, Peiyao Li, Chongxuan Li, and Jun Zhu. Equivariant energy-guided sde for inverse molecular design. In The eleventh international conference on learning representations, 2022.
Valentin De Bortoli, Michael Hutchinson, Peter Wirnsberger, and Arnaud Doucet. Target score matching, 2024. URL /tttps://axxi..cot
08667
Ricky TQ Chen and Yaron Lipman. Riemannian flow matching on general geometries. arXiv preprint arXiv:2302.0m 03660, 2023.
Wenlin Chen, Mingtian Zhang, Brooks Paige, José Miguel Hornándcz-Lobato, and David Barber. Diffusive gibbs sampling. In Forty-first International Conference on Machine Learning, 2024.
Stefan Chmicla, Valentin Vassilev-Galindo, Oliver T. Unke, Adil Kabylda, Huziel E. Saucoda, Alexandre Tkatchenko, and Klaus-Robert Müller. Accurate global machine learning force fields for molecules with hundreds of atoms. Science Advances, =ead00737
Hyungjin Chung, Jeongsol Kim, Michael Mccann, Marc L Klasky, and Jong Chul Ye. Diffusion poste- rior sampling for general noisy inverse problems. Tn The Fleventh International Conference on Learning Representations, ICLR 2023. The International Conference on Learning Representations, 2023.
Valentin De Bortoli, Fmile Mathieu, Michael Hutchinson, James Thornton, Yee Whye Teh, and Arnand Doucet. Riemannian score-based generative modelling. Advances in Neural Information Processing Sys- lems, 35:2406 2422, 2022.
Carles Domingo-Enrich, Michal Drozdzal, Brian Karrer, and Ricky TQ Chen. Adjoint matching: Finc- (uning flow and diffusion generative models with memoryless stochastic optimal control. arXie preprint arXiv:2409.0SS 08861, 2024.
Aleksander EP Durumeric, Yaoyi Chen, Frank Noé, and Cecilia Clementi. Learning data efficient coarse- grained molecular dynamics from forces and noise. arXio preprint arXiv:2407.0 01 010000000006000000 2024.
Lawrence C Evans. An introduction stochastic differential equations, volume 82. American Mathematical Soc., 2012.
12

---

Published in Transactions on Machine Learning Research (08/2025)
LawrenceCraig Evans. Measure theory and fine properties of functions. Routledge, 2018.
Xiang Fu, Zhenghao Wu, Wujic Wang. Tian Xie, Sinan Keten, Rafacl Comez-Bombarclli, and Tommi Jaakkola. Forces are not enough: Benchmark and critical evaluation for machine learning force fields with molecular simulations. arXie preprint or Xrxx:222000022 072377 2022.
Andrew Gelman and Donald B Rubin. Markov chain monte carlo methods in biostatistics. Statistical methods in medical research, 5(4):339-333) 1996.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Aduances in neural information processing systems, 33:6840 6851, 2020.
Emicl Hoogcboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. International conference on machine learning, pp. 8867 8887. PMLR, 2022.
Michael Janner, Yilun Du, Joshua Tenenbaun, and Sergey Levine. Planning with diffussio for flexible behavior synthesis. In International Conference on Machine Learning, pp. 90009990299999 9915. PMLR, 2022.
Tero Karras, Mika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-basse generative models. Advances in neural information processing systems, 14225655-65777 2022.
Jonas Köhler, Leon Klein, and Frank No6. Equivariant flows: exact likelihood generative learning for symmetric densities. In International conference on machine learning, pp. 5361 5370. PMLR, 2020.
Tuan Le, Frank No'e, and Ajnrk-Arné Clevert. Equivariant graph attention networks for molecular prop- erty prediction. ArXiv, 0ts/2202.09888 2022. URL https://api www.stiiscoolleeeeeeee wwwpatttiiiccooaaooorooororouooe Tngg/CorpuII::
247011158
Yi-Lum Liao, Brandon Wood, Abhishek Das*, and Tess Smidt* EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations. In International Conference Learning Rep- resendations (ICLR), 2024. URL httpps://openreviex http://wwwweeoooomm////ododdoneo
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. orXiv preprinl arXiv:2210.024 020777 2022.
Calvin Luo. Understanding diffusion models: A unified perspective. arXiv preprint ar Xri::22008.113 11970, 2022.
Andrew McCammon, Bruce Gelin, and Martin Karplus. Dynamics of folded proteins. nature, (56(2)) 585 590, 1977.
Laurence Illing Midgley, Vincent Stimper, Gregor NC Simm, Bernhard Schölkopf, and José Miguel Hernández-Tobate Flow annealed importance sampling bootstrap. arXio preprint arXiv:2208.01800 2022.
Benjamin Kurt Miller, Ricky TQ Chen, Amuroop Sriram, and Brandon M Wood. Flowmm: Generating materials with riemannian flow matching. arXiv preprint 000,240660471,, 2024.
Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron Owen, Mordechai Kornbluth, and Boris Kozinsky. Learning local equivariant representations for large-scale atomistic dynamics. Nature Commanicaaiinss, 14(1):579, 2023.
Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equilibrium states many-body systems with deep learning. Science, 365(6447) 2019.
Angus Phillips, Hai-Dang Dau, Michael John Hutchinson, Valentin De Bortoli, George Deligiannidis, and Arnaud Doucet. Particle denoising diffusion sampler. In Forty-first International Conference on Machine Learning, 2024.
Chence Shi, Shitong Luo, Minkai Xu, and Jian Tang. Learning gradient fields for molecular conformation generation. In International conference on machine learning, pp. 9558 9568. PMTRR 2021.
13

---

Published in Transactions on Machine Learning Research (08/2025)
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Represenlations, 2020a.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2020b.
Gabriel Stoltz, Mathias Rousset, et al. Free eneryy compulations: A mathematical perspective. World Scientific, 2010.
Francisco Vargas, Will Sussman Grathwohl, and Arnaud Doucet. Denoising diffusion samplers. In In- ternational Conference on Learning Representations, 2023. URL http:s//opparerreee.eeeeerrrrree @pvnfTAbulf
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural compulation, 23 (7):1661-167,, 2011.
Yan Wang, Lihao Wang, Yuning Shen, Yiqun Wang, Huizhuo Yuan, Yue Wu, and Quanquan Gu. Protein conformation generation forco-guided se (3) diffusion models. arXiv preprint arX.v:2409.1 14088, 2024.
Zun Wang, Guoqing Liu, Yichi Zhou, Tong Wang, and Bin Shao. Efficiently incorporating quintuple inter- actions into geometric deep learning force fields. Advances in Newral Information Processing Systems, 36: 77043-77055, 2023.
Dongyeop Woo and Sungsoo Ahn. Iterated energy based flow matching for sampling from boltzmann den- sities. Xi preprint or Xrx:::2408. 16249, 2024.
Fang Wu and Stan Z Li. Diffin:: geometric diffusion model for molecular dynamics simulations. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp. 5321-5329, 2023.
Kevin Wu, Kevin K Yang, Rianne van den Berg, Sarah Alamdari, James Y Zou, Alex X Lu, and Ava P Amini. Protein structure generation via folding diffusion. Nalure Commeniications, 15(1):1059, 2024.
Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. In International Conference Learning Repre- sentations, 2022. URL https:://ppeneev//// http://wwwwwwwneeeeeeee http:///ooooo////////////tt////
Zhantao Yang, Ruili Feng, Han Zhang, Yujun Shen, Kai Zhu Lianghua Huang, Yifei Zhang, Yu Liu, Deli Zhao, Jingren Zhou, et al. Lipschitz singularities in diffusion models. In The Twelfth International Conference on Learning Representations, 2023.
Qinsheng Zhang and Yongxin Chen. Path integral sampler: stochastic control approach for sampling. In International Conference on Learning Representations, 2022.
A Preliminary
A.1 Diffusion Models
A diffusion model typically consists of forward noise addition process and reverse denoising process. In the forward process, noise is iteratively added to the data E (at t 0), gradually transforming the distribution into an approximation of a Gaussian distribution at time T, independent of the initial state æo- The reverse process then generates samples that conform to the data distribution by leveraging information obtained from the forward process.
We introduce the Prrrancc-rrrervving Stochastic Differential Equation (VPSDE) and the Variance-Explodinng Stochastic Differential Equation (VESDE) [Song et al.) (2020b), summarized in the following Table Here, random noise, i (I, .N} is of noising steps.
14

---

Published in Transactions on Machine Learning Research (08/2025)
Table 5: VESDE and VPSDE.
VESDE
VPSDE
Noise addition
x; x,-1 + V o 1 2
x; = 0,N0 + O,
SDE Formulation
dx; (1)dW,
dx, 11 + VB,dW;
Parameterization
(t) min max/ min)
B, 3min + (3max Mmin) )
- 3 9.ds (1 - -e- So Pada Pado)I 1I
Transition Distribution
N (zr;: To. min ( maxx Baaux//Omiin min I)
N
2 The
We sample from the diffusion model follows Predietor-Corredlle (PC) scheme. The predictor step follows the time-reversed SDE, while the corrector step applies Langevin dynamics to refine the sample. Given a time step At, the predictor step updates the sample as:
f(x, ((()))))(+(()))))
(17)
where 2, ~ (O(, Gaussian noise, and and 9 are drift and diffusion coefficients. The corrector step refines -t using Langevin dynamics Xt- X + Vog (x)) /2rz'' where rj is the step size, and x' ~ (O, I). This two-step process improves the quality of generated samples and enhances convergence to the data distribution.
A.2 Equiformer-v2
A.2.1 Equivariance and Irreducible Representation
In molecular systems, molecular structures are typically represented as graph, where atomic coordinates, atomic numbers, and bonded connections define a molecular graph. Ensuring properties such as symmetry and equivariance these graph representations is crucial for capturing molecular interactions accurately.
Graph representation. A molecular graph is defined as G E), where V {vi, (2,,,, repre- sents the set of atoms (nodes), and E {ei (i.j) C V x V} denotes the set of edges capturing atomic interactions. The number of atoms in molecule is denoted by its nuclear charge and 3D coordinate r; E R Our goal is to design a generative model that learns to generate molecular distributions while preserving the underlying chemical and spatial properties.
N [V]. Each node vs V is characterized
Graph neural networks. Graph Neural Networks (GNNs) are widely used to process graph-structured data by propagating information across nodes and edges. At each iteration node features h are updated using a sassage-passing mechanism:
((+1)
Σ (E)(U)
o(h!",h)
h((+1) (h("), m! m(+1)),
(18)
where (v)) represents the neighboring nodes of u, cou denotes the edge feature between nodes u and u, and o(-) and v-)) are learnable functions, implemented as networks.
Equivariance and irreducible reproscntations. Equivariance serves as a. fundamental property neural networks, providing strong prior knowledge that enhances data efliciency in molecular modeling. A function f X + Y is equivariant under transformation group G if, for any input z e X, output y Y, and group element E G, f (Dx (g)x) = Dy (x) holds, where Dx and (g) are transformation matrices parameterized by in X and Y. In 3D atomistic graphs, molecular structures must maintain equivariance under the Euclidean group F(3), which includes translations, rotations, and reflections.
For the Euclidean group E(3), scalar quantitics remain invariant under rotations, whereas vector quantitics transform accordingly. To enforce translation symmetry, computations are performed on relative positions. Since rotation and inversion transformations commute, any representation of the special orthogonal group
15

---

Published in Transactions on Machine Learning Research (08/2025)
SO(3) can be decomposed into irreducible representations (irreps), which serve fundamental transfor- mation components. Equivariant neural networks leverage these irreducible representations to construct features that remain equivariant to 3D rotations.
For probabilistic models, the equivariance property extends to probability distributions. A probability distribution p(x:)) is equivariant under the special Fuclidean group SF(3) if, for any transformation T3 corresponding to g SE(3), then p(æ,) (x))( and p(æ;-1 xx) p(Tg((x-11 (xx))) A similar definition applics for E(3).
A.2.2 Equiformer-v2 Network
In our research, we harness the capabilities of the Equiformer-v2 network Liao al. (2024) to maintain (B)
a critical aspect for accurately modeling molecular configurations. This advanced architecture builds the principles of irreducible reprosentations, ensuring that its operations and features are equivariant and robust. Other works have also incorporated designs ensuring molecular invariance and
equivariance Xu et al. 2022 Shi et al. [2021): Le et al. 20222; Musaelian et al. 20233)
Building upon the Equiformer, the Equiformer-v2 introduces sophisticated equivariant graph attention mech- anisms, supplanting standard Translormer operations with their SE(3)/E(3)) (e3i)-equvaa counterparts and utilizing tensor products for higher-degree representations. Node embeddings are crafted by concatenating channel-dimension features and applying rotations based relative positions or edge orientations.
To enhance computational efficiency, Equiformer-v2 forgoes separate depth-wiss wrpph-wise tensor products and linear layers in favor of single SO(2) linear layer that processes scalar and irreps features distinctly. It encodes edge distance information using radial basis functions, integrating them into the node embeddings.
The attention mechanisms benefit from extra normalization and separable S activation for improved sta- bility. The network's food-forward section employs a two-layer MLP with SiLU activation. The output head either aggregales scalar predictions or computes atom-wise forces, utilizing equivariant graph attention to boost expressiveness and scalability.
ESCN convolution. The ESCN convolution refines equivariant tensor products by substituting conven- tional SO(3) operations with SO(2) linear ones. Typically, SO(3) convolutions combine input irreps features with spherical harmonics of relative positions using SabschhGGoddal coefficients. ESCN simplifies this by aligning the relative position vector with fixed axis through rotation, reducing dependencies and limiting interactions to cases where my 0, which streamlines computations while preserving equivariance.
Attention re-normalization. To maintain stability at higher Lmax, we introduce an additional layer normalization step before non-linear transformations, ensuring well-scaled scalar features r.6 Attention weights are computed through a leaky ReLU and a linear layer, promoting numerical stability in operations like softmax, which aids in training convergence and model robustness.
Separable 2 activation The separable S activation processes degree-O and higher-degree features sepa- rately. Degree-O vectors are partially activated with SILU,, while the remainder is blended with higher-degree vectors for 52 activation. This selective processing minimizes cross-degree interference, stabilizes gradients, and enhances expressivenes, particularly in FFNs.
Separable layer normalization. Separable layer normalization (SLN) extends traditional equivariant normalization by independently normalizing degree-O and higher-degree features. By computing their means and standard deviations scparately, SLN preserves inter-degree dynamics, enhancing stability and expres- siveness, especially in high Luxx scenarios.
Our network architecture predicated on Equiformer-v2. which yields both one-dimensional energy and three-dimensional force terms. For the c-net output, we utilize only the energy head. Furthermore, we adapt the model to consider time an additional input feature to align with the diffusion model's time-dependent noise addition.
16

---

Published in Transactions on Machine Learning Research (08/2025)
B Derivation of Loss Function
B.1 Proof of Theorem
Proof. Assume that p(zo) & e the derivation of the loss function follows from the transformation of scores using the Gaussian transition probability density function and the integration by parts. Since V Va,p(a;;((a)) ae Vagp(an
Isy 1my-ocmoll* Re
S (((o Var
dæy
V a. I p(a Do 2o)
Ve, logp(z;)
p(x+)
p(x))
p(aro)
Vange
mo
znz
daco
|m mm 207
az
S 1
dro
p(x+)
p(ace)
(19)
1
(p(3r) V mg log log)())))
daro
at
p(ar,)
/ p(xo xe)) Vao log p(oo))do
1
at
-vn
Euo aa
KgT
O2
KBT
where & is the energy function, and F represents the force derived from V V. Specifically, for the VESDE,
0g 1, then Vz, log p(x;) Bx
E
KgT
For simplicity, let kgT =. Combining these results, the potential score matching loss can be derived as:
1
Lecore model (E- Bur (m Ezw ||s [ t) (-
V E)12
(20)
att
1 F112.
Fae Exp (1o
(21)
The reason why Væ, log p(+z)
= Ex le | Vv, (objective equation (equation 20 lies in Theorem According to the above expansion, label Van lox;( can be rewritten as
19 is the solution of the loss function
Vm, (p((o))((aa (o) p(xx)
f p(ao)
21 Pi-20 af
)p(x ddoo
Va, a, logp(a;)
p p(xx))
(22)
- / p(zoo=)))
xt 20 )dxo = Emolan
2pp 27
of
of
which is consistent with the writing of "xo-model". Similarly, we can rewrite the loss function with a network representing Zn-
ED Et V Va, E(zo) = Emn
ZD Et of E(xo)
Tu -DD
Egro we
Enolas
(23)
of
o²
Lxs model EE uplae (mn) Ex, [Do wt of E
(24)
Theorem 3. Let X be integrable random variable. Then for each oalgebra V and Y E V, Z = B(XIV)
solves the least square problem Evans
2012
2 X
min J|Y - XI1,
YEV
where ||V|| (J
17

---

Published in Transactions on Machine Learning Research (08/2025)
To prove the theorem, we first introduce lemma.
Lemma 1. If V-measurable, and f measurable function the sense that its domain and codomain oppropriately aligned with the aalgebbras, then f(Y) will also be V-measurable Evans 2018).
Proof of Theorem Let X be an integrable random variable, be a o-algebra, and Y E V. We need to show that Z V minimizes the least square problem 1 X| minyevlYY-a
First, recall the property of conditional expectation: B(ZIV) Z for any Z E V. Consider the difference X --Y for any Y V.V. We can write this difference as:
X -x-=(X-E(X>))) (B(XIV) - Y)
Note that B(X]V) E V, so the second term (B(X(V)-Y) V. Using the property of conditional expectation, have (V) - Y V = EX Y.
Now, let's calculate the conditional expectation of the product of (X E((())) and (E(XPV - Y):
B((X - ((XX)((())(())))))) Y)I = -(XX-- =0
The last equality follows from the fact that the product of the two terms is uncorrelated, and their expectation is zero. This implies that (x R(X[V) and (E(X Y) are orthogonal.
Now, we can use the Pythagorean theorem for Hilbert spaces:
X - YP= X -(XM)) [E([V) Y
Notice that the right-hand side is minimized when [E(XIV) Y is minimized. This is true because X E(X is constant and non-negative, Therefore, the least square problem is minimized when Y = E(XIV), i.e., Z = (V..
Hence, we have proved that Z E(X[V) solves the least square problem:
7
X
min Y
X|.
YEv
B.2 Relationships Among Three Diffusion Losses
In this section, we give the forms of three commonly used losses, e-loss, s-loss and Øg-loss in the PSM setting. Assume a = and KB for simplicity, since
V mi logp(æ:) Va. p(2 (( /axo
p(xx))
1
J p(zo)
-Vaee
207
dxy
Va,
[ log log
= Emglm
(a)
p(xx)
K
)p (o) daro
of
Ct 200 dxx / p( ((xullx))(-
Emg mi
TDo of
(b),
p(xz)
of
(25)
where E is the energy of 2o. When using a score network so to learn the gradient of log probability Van log ((x)( the s-loss can be derived by (a) that
Lscore- model (E Be ppt Iss (a0 t) (- (-Va,E))
(26)
18

---

Published in Transactions on Machine Learning Research (08/2025)
and accordingly, the e-loss is
Le- model (e Beo (10 t) - os Va. 211
(27)
if take oes which is independent of St and 0, out of the norm. Similarly by (b), the formula + Vage(z)) is expected to be zero, as the optimal score function satisfies that Va, log p(z;)
m00 E mt o?
Exg le, V-,,, (xxo).. Since Bmolas las function with network representing 20,
2o 211 of
V 22 Exx)) Bay lay
can the loss mo moomy m. to2 Van ((mo)
"
Da, to learn the diffusion.
Lag model EE,U((, B., mpla Eun IDo 2t V.V mn Elx0)l
(28)
Therefore, e-loss and To-loss are different objects, but are essentially derived from the same formula.
B.3 Proof of Theorem 2
Proof. Let In be the formula on the left side (o(a))))) right one, By(xolees) (V log (ax))] Ep(wolx)) V log (20)]
[ log (2ro) E Ep(eo me) Vllog (xo)), and Iz be the
1 logp (2o) (4 (ao æx)) p (xo xx)) dxo
12 / vV (2ogg( æx)) V log (xx) p (at xx)) dxy,
(29)
I3 = / (Vlog (xx)
v Vlogp (xo)) (zo xz dæy-
Then I1 + I3.
We only need to prove that for near 0,
/ V log p (2o) (g (Zo xx) ( xx)) dxy
|
1
(30)
/ V log (2o) q (wo xe) (()) p (aro xx)) doo
-|/211
By ||2112 = ||/3 + 1 11 = 1/ 2 + 11 + 21, Ia, we need to explain that I3 (Ia + 211) > 0.
For simplicity, we assume the reverse SDE satisfies that dæ, -h(t) log (xx)) dt + /2h(t)dW, Here h(tt) min omnt )²n and We Brownian motion. For function f. [0, 1] x Re R, the Ito formula of f (xx) is expanded as:
df (at)=== (a:) day, Sd (ar+) (da.)²
=f' (ax) | Vn log (z,, + 2h(t)aW; + 555" (æx) 22(()dt
1
(a:) (x: logp f" (z.) h(t)) dt + 2h()) (z.) dw.
For reverse transition p (xx 2x), where < s, assume that p (ax xx) the form of N (g (xx), (PPa), g(xx) is function of Øs and Br, is time function related to I, 8. Then
Ft (xt.)
dp (xe) x.) =
(x, xx) l(t) logp(x.))
x gg
2 p (x. xx) h(t)
dt + (- ((((
x. (xxx)
/2h(tjaw.
+
p (xt xx))
82
3r,
87,
19

---

Published in Transactions on Machine Learning Research (08/2025)
Therefore, @p(a:/a=)) p (exx xx)) lu(t)
(ar)
+
(1-g(ax)) Pr.
27
and
dt
Pi
1/a
Zu - (xx)) ( [x) p20 [s)) - p p x2s) h(0))
Vx log (xo)
1
+
To
(F+)
+ 0
Bu.
(Eo g (x+s
=p (xx)) h1(0))
(g (xxz) m0) Var log (2:0) - 1
8 8 0(82)
30.a
x.) "e, (lo(zz)) - x0) aa log (x0) - -1+ (at,) + Ci (20g, x2.) (
[0- o(at)) ?) )))
Let C1 (xo. v.) = -h(0pppp (xo
then p (20) =
fo,-
I3 / (Vlogg (x0) - V logg p (x))) [p (xx) + C1 (zo, x (t)] ((]) doo
= (2r) / (V (Vlo (x00) log (x))) dxy +t / (V 1oggg (xo) - v Flogp (xo)) C1 (zo, xz) doy + 0 (ta).
(c(.) xo) Van logg (mr) 1 200 (000((e))) )r) then
where C2 (zo, = -h(0)g (aro x..) Pe.
p (orolæ.)) (x:)) + Ci (wo, v.) . 8 (
We then expand 11 == f V log (x0) (g (zo xx) (wr xz)) dæg. From above, we know that,
11 / v (og) la (xx) + C2 (xo, xt) p (xx) Ci (xo, xx) 0
= (1())) - (xx1) / V (ogg don
V Flogp(xx) (C2 (zo, - Ci (xx) xz) dxy + +0 -1²)
21g 1 2p (xx)) p (ax)) (ʃ (Vlogg(Zo)) Vlogp dra) )
T "(
V log p (ara) daw )
(wo. (
(Vlgg log (mo) - V log (mo)) daro) )
(/ log (mo) dxo
+ 21 (a (#x) p (mi))
( ((VV (Vlog (xo) - V log (zu)) dzo)
+ (xx)
- C1 (zo, xz)) V logp (#o) dxoo) T
+ +
(f
((V V logp(ao)) C1 (Iv. xx) dxy)
( V Vlog (zo) (C2 (Io. xx) C1 (xo. xx)) dxy
-
+ +0 (19)
Since / VVlop( dxo = 0,
(((Voogg(am)) - V (og)( C1 (Zo, x() dxo) )
(ʃ V log (zx) (C2 (xo. x.) - C1 (xo xx)) doy
21 11 2t²
+0
Consider the second term of the last equality,
C2 (Zo, xt)) O (xo, æx) = hh00)
[[g (xx) - 20) (p x.)) Van logp (x0) - q -4(00 x.)) Van log g (go))
x g (2s) )² (p
1
+ h())
-1 -1 +
(xu (p(zu)
xx)) 9 (Iu xx))
20

---

Published in Transactions on Machine Learning Research (08/2025)
Then
/ v logp (zv) - (C2 (Zo, xt)
Ci (2o, xx)) dru h(0)- / no (are)
)
(p (2ca xx) - (2ol2x)) V log (xo) dx
f-(())) zo y (x:)
2
p(20)
q (2r)
=k())
p (2)(()))
log (zo) dæo
p p(x))
q (x)))
(t b(0) / WB0
(are)
p (x: zo) V log (at) dxa
(6)
The (1) above is the assumption p(a.) ~ 1, which is result from that for p(z, (mo) =N(xr, o) where or
min Then
Omax mmnn
and expanding p(ax)) around wy is p(ax)) p(aro) + 3o O(()) Hence, p(ao)
p(ae)
O()).
21 20 (j
Ci (wo, (u). (V logg log pog ) ((())
(
20 - g(
(x) wo) pogg((p)) dac
Therefore,
1] (Is + 211) = 1 + 0(t2)
= =
-( (V log (200 - V log (2r0) (2o xx)) dwo )² " '²
(V (Vlogg(zn) v log (xo)) ) (20 xx) dann) + 0(10).
If small, (z))) 6(z0 fxx)), where f is a function of æe. Since V Vlog (xx) * (logp(xo)) then [3 (Ts 2h)) > 0 when & is small.
B.4 Singularity and Variance Problem
This section explains the role of PSM in the neighborhood of 0 from two aspects: (1) PSM has smaller training variance near 0 and the training is more stable; (2) PSM can alleviate the problem that the Lipschitz constant of DSM is too large at O.
(1) Variance of PSM. Let X be a sample from the data distribution, and let Y be a sample from the noisy process 2; = 0g + O2E.
Lemma 2. When the noise schedule close near t = 0, and the variance of the PSM target Varx (V logp(X)) smaller than DSM when near
close lo zero, the variance of the DSM largel salisfies Verxjy (V IY > 0 if
Proof. Let X the sample from the data distribution and Y is the noised data, Y ax + orE, we can prove the variance of PSM target Varx (V log p(X)) smaller than the DSM target V Varxly (V log p(yla)) near t = 0 while larger when t is larger. The proof strategy here is similar to that of TSM Bortoli et al. (202h;; however, their work only considers the Gaussian case, whereas address more complex distributions. Nevertheless, the corresponding Boltzmann distribution or data distribution can be approximated by mixture Gaussian distribution, and thus it suffices to establish the desired properties under the Gaussian setting for simplicity. Let be the dimension the problem. Assume that px (x) = (p, Par),, and the conditional probability
Prix (y x) DV (aur,
(31)
By convolution, py())= N (a.B(X), of). Then, by Bayes' formula, Px y (x (v) (x(x)) Py
N (( (07)
Py(2)
(x-14)
(y - axx)
Px(x) PrIx(U) x) x exp
20
20
tar
21

---

Published in Transactions on Machine Learning Research (08/2025)
[o)2 + u²² - uw)) + olar (y² +
= - -2awzy)
exp
20 2
tax
[[a [aafotar oar + +o,) x² (oN} + asyorar) x + 1²² oi + y saall
= exp
20702ar
Hence, (1)²
oio aj our to?
and
N NV o(00))
(x- 22)) (2))2))
(y (x)))
(x (3)
Py
exp
(0(
2 (0] + arotar)
1
x (x-1(1) v((l)) 2
o(1)²
(v ax)(2)))
exp
2(0
oi and
1
(x- C (n1t), ,01) M, ar ofr B(x))
exp
2(o(1)
where C() is a constant decided by :. Therefore, (arx|x(x) = =
o. ajo top
Therefore,
(ax - y (ax-yy
Varxly (Velogrrix = x))
Varx _0,2 o,2
Var xaxx
o,
(32)
2
0
Otar
1=1
Ct 2
2 o
tar
+
=
da
0:
of tar + of
and
Varx (V, logp(z))
Σ
Varxiy
d
01
1
(33)
otar
a
otar
of far
+
1=1
(2) The Lipschitz problem of DSM.
Lemma 3. Given the forward process æ, og + o.e, the noise schedule satisfies or - 0 as t - 0, then the Lipschitz constant of the neural network eg satisfies lom sup
(t) Ot + 00.
Proof. By 6g = oe 80 RM oe Vz l and the chain rule, drg assume that the v log (xt) is smooth. Since x; = ax + O
dey
V, log (xx.)
0V. log ot
0z. We
at di
(i) VESDE: aa 1, 0o =: where do min; σ maxx, then da. ao
In o2 In
dr
do
Therefore, de mentioned that, do approaches 0 theoretically, but in practice, do 0.01 is taken to avoid singular values.
das V. log (xx.) has the order of O(In z z), where is the random noise. As Song al. 2020b 5u
(ii) VPSDE: a So Bod, or = V1- e- -S p, then da. dt
= Bee (-e Jo a,de) ) (1-c
- -So (0 -
8,0 So So nods
+00 if
Vi-e- - -S" Saaaaa fada
It shows the Lipschitz singularity property.
In summary, our theoretical analysis (Lemma)) 2 Lemma and Thcorem 2 demonstrates that PSM roduces variance and corrects estimation bias in the small-t regime, providing strong motivation for its use when is close to zero.
22

---

Published in Transactions on Machine Learning Research (08/2025)
C
More Experimental Details
Dataset and settings. The MD17 dataset and the MD22 dataset can be downloaded from http://
quantum-machiie org/gdm///atta/pp the LJ13 and LJ55 data are from ptttps://osf.i///e/// http....///r/ogf..e....
osfstorage Hyperparameters for experiments are recorded in Table i and 7
Table 6: Hyperparameters for different datasets with biased trainset
Dataset
Data Split
Batch Size
Max
Radius of Graph
(train, val, test)
(train, val)
Epochs
LJ13
First 1k, 0.1, remaining
(64, 64)
2000
LJ55
First 1k, 0.1, remaining
(64, 64)
2000
5
MD17
First 5k, 0.1, remaining
(64, 64)
1000
5
MD22 (Dw nanotube)
First 5k, 0.1, remaining
(8,
3000
4 7 (AT-AT)
(AT-AT-CG-CG)
MD22 (others)
First 5k, 0.1, remaining
(16, 16)
500
6 (Ac-Ala3-NHIMe)
6 (Docosalexaenoic Acid)
(Stachyose, Buckyballeatcher)
Additional Settings
Learning rate: 0.0002
Lr scheduler: Cosine
Seed: 42
Optimizer: Adam
Weight decay: 10
All diffusion parameters: Omin = 0.1, Fmax = 5, time weight function X(t) =
Sampling timesteps: 1, 0,000
Sampling method:
Predict-corrector sampler
Table Hyperparameters compared with previous work Table
Dataset
Data Split
Batch Size
Max Epochs
Radius of Graph
Noise Schedule
(train, val, test)
(train, val)
LJ13
0.1, 0.1, remaining
(64, 64)
5000
4
min 0.01, T
LJ55
0.1, 0.1, remaining
(64, 64)
5000
5
0min = 0.01, Tmax
Additional Settings are the same as before in Table 6
Toy models in 1D and 2D. We conduct experiments using PSM alone in 1D and 2D toy distributions, and the according force label is (-4|| [|1 10)x. It indicates that PSM and Piecewise perform comparably to DSM in ID and 2D case, and it can be seen that Piecwise and PSM can better learn the sparse area in the middle. Previous work also addresses the challenges for learning high-dimensional molecular problems using only an energy label or log-density gradient Woo & Ahn (024h; Bortoli et al. (2024). The constraint of force labels F become weak due to the projection of time-zero information onto highly noisy Je for large is smaller, which unstable for dynamical learning. Refs Yang et al. (0233 and Bortoli al. (2024) have shown that DSM suffers from high variance and singularities near and encouragingly, our analysis further demonstrates that PSM achieves lower variance and corrects bias in the small-t regime (Lemma 2 and Theorem 2), motivating its use specifically for small t. Therefore, for high dimensional molecular system, we use PSM in small t range, and use DSM for large t.
shown in Figure where p(x:) -5||=||++||=] ee
Lennard-Jones potential. Here we additionally supplement the LJ potential with biased data training to obtain the energy comparisons in Figure
23

---

Published in Transactions on Machine Learning Research (08/2025)
2
0.8
Reference
0.6
Dim
0.4
-1
0.2
0.0
-4
-
-2
-1
0
0
X
Dim 1
&
(a)
(b)
0.8
0.6
0.6
0.6
0.4
0.4
0.4
0.2
0.2
0.0
0.0
0.0 4
3
-2
-1
1
2
4
4
-3
-2
1
0
1
3
-3
-2
-1
2
4
X
X
(8)
(b)
(c)
Dim
mm
2 -2
2
-2
-1
-1
1
2
-2
-1
1
2
Dim
Dim
Dim
Figure Comparisons of (a) DSM, (b) Piecewise, (c) PSM one-dimensional (the second row) and
two-dimennional (the third row) examples. The first row is the ground truth for 1d and 2d cases where the exact density function is p(x) = e +||a]]++|a]]]
24

---

Published in Transactions on Machine Learning Research (08/2025)
DSM
Piecewise
Piecewise Weighted
(a)
Ref Energy
Ref Energy Piecewise
Ref Energy Piecewise
0.06
Weighted
Frequency
Frequency
Frequency
0.02
0.00
40
-20
-60
Energy
Energy
Energy
DSM
Piecewise
Piecewise Weighted
(b)
0.03
Ref Energy DSM
Ref Energy Piecewise
Ref Energy
Plecewise Weighted
Frequency
Frequency
Frequency
0.02
0.01
12
M
-400
-300
-200
-100
-400
-300
-200
-100
-400
-300
-200
-100
Energy
Energy
Energy
Figure 7: Comparisons of energy for (a) LJ-13 and (b) LJ-55 potential training with first 1, 000 data.
D
Relationship with Flow Matching
Flow matching is a powerful tool for generating molecular conformations and predicting molecular properties
Lipman et al. (002); Chen & Lipman 2023); Miller et al. 2024).. It enables faster training and sampling
while achieving superior generalization performance. Motivated by recent advancements Woo & Ahn (2024):
Domingo- Domingo-Enrich et al. (2024), unified frameworks have been developed to describe both diffusion models
and flow matching mothods comprehensively.
Let yo and y1 be samples from the noise distribution and the data distribution, respectively, in the flow matching framework. As demonstrated in Domingo-Ennnn Enrioh et al. (2024). the relationship between the velocity field and the score function is given by:
di x +
v( 1)
t/ s(x0, 1),
04
where xe = a:y1 + Y: = 0yzo + YE. Based on this relationship, our method can be extended to flow matching. In our setting, the coefficients are defined as follows:
a ()) o0
VP: a - Aeds 7t (1 (1-e So ((Rads) ).ar))
VE: a 1, == op
(34)
2 In s(x,t) for VESDE, and v(x,t)
Therefore, the vector field u(z,t) nis(L,,t) of
- 3,, + " a in) s(x, 1) for VPSDE. Building on the interplay between energy and the score func- tion, we propose investigate an energy- iergg-informedd flow matching method. This energy-based extension will form the basis of our future research endeavors.
25


Reviews:
[{"review_id": "NM75VRmzII", "reviewer": "Reviewer_oAeN", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The paper is nice because it recognizes the unique setting of molecular systems, and finds a bespoke diffusion strategy for it and exploits it effectively.\n2. The numerical simulations are quite comprehensive, and clearly demonstrate the efficacy of the strategy.\n\nWeaknesses:\n1. The method itself is not really novel, considering [Bortoli, et al, 2024].    The authors do propose some variants of the approach which are novel, combining classical Diffusion Score Matching losses with the new PSM loss, but at no point is it really explained why such approaches are beneficial or even necessary.   This is the main weakness of the paper, in my opinion.\n2. The paper has some issues with clarity and motivation in some areas - which I will detail in the next section.\n3. The paper is a bit inconsistent in the mathematical exposition.   Some quite technical aspects (which are arguably well known in the diffusion literature) are completely glossed over (e.g. the relationship between the score and denoiser, etc), while very basic aspects of probability theory are expanded on in great detail.   \n4. The notation is challenging at times", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces Potential score matching,  a score-based diffusion model variant which is able to leverage knowledge about the target distribution to improve the behaviour and performance of the model at small times (i.e. when close to the data distribution).\n\nThe strategy is based on an alternative representation of the score-function in terms of the gradient of the log target density, which results in a loss function which requires access to the grad log target density at the data-points (thus can be considered additional label information).    Similar strategies have been explored before, e.g. in [Bortoli et al, 2024] targeted score matching paper, but the authors propose a novel way of integrating the psm score with dsm through interpolation.\n\nThis setting plays very well in the context of molecular dynamics, where you would typically have access to the energy and force values, as the training data would have arisen from MD simulations, etc.     To this end, the authors demonstrate the efficacy of their approaches (combining PSM with classical DSM) on a number of benchmark molecular problems, showing that their approach is effective.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Can the authors please explain in more detail:\n1. Why the need for piecewise and piecewise-weighted losses,  e.g. why not just use PSM all the way through.\n2. Why the specific form for the piecewise-weighted loss, i.e. the use of t_dsm and t_psm are completely unclear to me.\n\nThe section defining these two needs to be substantially rewritten and clarified.\n\nThe authors should be a bit more explicit about going from the score-based approach to the denoiser approach -- why do this, what's the advantage in this setting?    Also, why and what the x_0 loss and epsilon loss, (5) and (6) respectively, needs to be explained in more detail,  for example, what is $\\mathcal{D}\\_\\theta$ a function of?    They do cite Luo 2022, but it would be better to provide clarity directly within the text.   Similarly, jumping from denoiser to score is a bit 'sloppy' at times, for example in Page 2, it is not really correct to say \"leveraging the relationship $s_{\\theta}$ = ....    The relationship is between the losses - one can choose to optimise either score or denoiser loss, but not simply prescribe equality between the two networks.\n\n\nPage 15: typo \"Repressentation\".  \n\nLemma 1 and Theorem 3 -- these are largely undergraduate textbook proofs.   I would simply cite the relevant sources here, and dedicate the space to explaining the other aspects (denoising vs score based approaches in far more detail).", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "jbV9shILQB", "reviewer": "Reviewer_whEG", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n- The paper is mostly nicely written and easy to follow. \n\n- The method is straightforward an easy to implement. Compared to existing approaches like iDEM, PIS or DDS, it can be easily scaled to high dimensions. It can also easily utilize biased training data. Furthermore, theorem 2 demonstrates theoretical PSM's advantage over DSM. \n\n- The overall empirical results seem convincing. \n\nWeakness:\n\n- Limited novelty: PSM is largely built upon TSM. In particular, the original TSM paper mentioned the option of combining TSM and DSM (in their page 3). And in principle, TSM can be applied to biased training data as well (contradicting the presentation in table 1). I wonder if the authors can clarify these points and highlight the novelty of the proposed method?\n\n- Lacking theoretical guarantee: there is not theoretical guarantee that the proposed method would generate exact samples from the Boltzmann distribution. Hence i wonder if the \"debiasing\" claim should be reframed. \n\n- Lacking comparison to MD: In high dimensional tasks, while it is reasonable to exclude comparison to other diffusion-based methods (as they won't easily scale), I feel it'd be nice to compare to classical MD approaches, e.g.  the number of force evaluations given the same performance, or convergence speed, or the amortization advantage.\n\n- Lacking comparison to using only PSM: I wonder if the authors considered using only PSM, instead of som combination of DSM and PSM. Is DSM necessary in high-t regime?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes Potential Score Macthing (PSM) to sample from Boltzmann distribution using biased training data and the potential energy gradient. The key idea of PSM is to leverage the target score matching (TSM, Bortoli 2024) identity, which translates the score at time t to some expectation of the \"clean\" target score -- equivalent to the force field in the molecular simulation context. PSM can be further combined with Denoising Score Matching (DSM, the standard way of training diffusion models) to improve the empirical performance \n \nThe paper benchmarks PSM against existing methods on low-dimensional molecular dynamics (MD) systems, demonstrating competitive performance. In high-dimension MD systems, the combination of PSM and DSM outperforms using only DSM.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "For most of the requested changes, see Weakness section. \n\nAdditionally, I wonder what \"energy labels\" in table 1 refer to? Are they just the force (or equivalently the clean score function)? In that case, shouldn't some of the MD and MCMC methods satisfy the requirement that they don't require energy function but only the energy labels? For example, langevin dynamics only utilize the gradient information, i.e. the score function.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "WOupkeIm4p", "reviewer": "Reviewer_VW23", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths)**\n- Simple modifications in the loss function, but exhibits superior performance\n- Conduct experiments on several high-dimensional benchmarks.\n\n**Weakness)**\n- I have a question regarding the choice of $t_psm$ and also the weighted function $\\omega_t$. While it seems that they are fixed across all experiments, it would be better to conduct some ablation studies regarding those components.\n\n**Questions)**\n- As I am not an expert on molecular dynamics. I have some questions regarding the experiments. First, for Figure 1, it seems that the suggested method more clearly approximates the reference data distribution. However, according to the captions, they are trained on the biased dataset. Could you give me more elaboration on what the figure indicates exactly?\n- Similar question regarding Figure 4. It seems that all the figures look very similar. Could you give me more elaboration on what the figure indicates exactly?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors present Potential Score Matching, a novel method for sampling molecular conformations from the Boltzmann distribution. A main component of PSM is a piecewise loss and a piecewise weighted loss, which prioritize PSM at small $t$ and DSM at large $t$. \n\nThe authors conduct several experiments on molecular dynamical datasets, including MD17 and MD22, which are relatively high-dimensional compared to commonly used benchmarks such as the LJ potentials.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please see above.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)