Energy Loss Functions for Physical Systems
Authors: Sékou-Oumar Kaba
Kusha Sareen
Daniel Levy
Siamak Ravanbakhsh
Venue: NeurIPS 2025 poster ()
Keywords: deep learning
AI for science
physics
physics-informed
diffusion
molecule
generative
Primary area: machine_learning_for_sciences
TL;DR: We use energy differences as loss functions for physics applications

Abstract: Effectively leveraging prior knowledge of a system’s physics is crucial for applications of machine learning to scientific domains. Previous approaches mostly focused on incorporating physical insights at the architectural level. In this paper, we propose a framework to leverage physical information directly into the loss function for prediction and generative modeling tasks on systems like molecules and spins. We derive energy loss functions assuming that each data sample is in thermal equilibrium with respect to an approximate energy landscape. By using the reverse KL divergence with a Boltzmann distribution around the data, we obtain the loss as an energy difference between the data and the model predictions. This perspective also recasts traditional objectives like MSE as energy-based, but with a physically meaningless energy. In contrast, our formulation yields physically grounded loss functions with gradients that better align with valid configurations, while being architecture-agnostic and computationally efficient. The energy loss functions also inherently respect physical symmetries. We demonstrate our approach on molecular generation and spin ground-state prediction and report significant improvements over baselines.

Energy Loss Functions for Physical Systems
Sékou-Oumar Kaba", Kusha Sareen", Daniel Levy, Siamak Ravanbakhsh
McGill University
Mila Quebec Artificial Intelligence Institute
Abstract
Effectively leveraging prior knowledge of a system's physics is crucial for appli- cations of machine learning to scientific domains. Previous approaches mostly focused on incorporating physical insights at the architectural level. In this paper, we propose framework to leverage physical information directly into the loss function for prediction and generative modeling tasks on systems like molecules and spins. We derive energy loss functions assuming that each data sample is thermal equilibrium with respect to an approximate energy landscape. By using the reverse KI. divergence with Boltzmann distribution around the data, we obtain the loss as an energy difference between the data and the model predic- tions. This perspective also recasts traditional objectives like MSE as energy- based, but with physically meaningless energy. In contrast, our formulation yields physically grounded loss functions with gradients that better align with valid configurations, while being architecture-agnostie and computationally ef- ficient. The energy loss functions also inherently respect physical symmetries. We demonstrate approach molecular generation and spin ground-state pre- diction and report sigmificant improvements over baselines. Code is available at https://gi /ithhtttttuuu// com www.auuuseeeeeeeeeee seeehhasareee/eeerryy energyeerooooor _loss
Introduction
A key challenge in applications of machine learning to the physical sciences is that data often be searce and expensive to generate. However, we often have some prior knowledge of the physics the system of interest, which can be used to design useful inductive biases. A common learning problem involves training a machine learning model to predict configurations of physical systems based on data collected close equilibrium such protein folding (Noé et al. 2020 Jumper al. 202T Abramson et al. 2024). crystal structure prediction [Ryan et al. 201 8 liao et al. [0023 Zeni et al. 2025], calculation ground states given Hamiltonian parameters [Carrasquilla and Melko 2017] or generative modeling of physical systems
Bomnez-Bombarell et al. 2018 (Sanceez-Lenglling and Aspuru-Guzik 2018]. A significant body of work has focused implementing physical inductive
biases, such as equivariance at the level of architectures (see e.g. Zhang et al. 2023 for a review).
This work explores complementary direction: embedding physical principles directly into the loss function. The fundamental question we ask is: can loss functions grounded in physical principles pro- vide more effective training signals and yield models that better reflect physically valid configurations compared to generic losses such as the mean-squared error (MSH) and the cross-entropy loss?
As response, we propose framework for deriving energy loss finctions tailored for physical systems in the thermal equilibrium regime. This is motivated by the fact that loss functions can be obtained from a distribution representing the uncertainty around each prediction or data sample. For physical systems thermal equilibrium, the sensible choice the Boltzmann distribution. Employing the reverse Kullback-Leible (KL) divergence leads to loss functions that take the form of approximate
Equal contribution. Correspondence: kabaseko kabaseko@mila.coom caa.kkmmmmila.uubbb....
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

X]
X2
%
&
X
(a) MSE loss
(b) Configuration with same MSE loss as (a)
(c) Energy loss
Figure 1: Energy interpretation of loss functions. Ground truth positions are denoted in green and predictions in blue. (a) The MSE loss function for particle positions corresponds to quadratic potential energy centered on the data. (b) This choice however physically unsound and leads to penalizing the model for configurations that are correct, i.e. related by rigid motion to the target. (c) more accurate choice would be to use a loss function based on physically sound energy, which would not suffer from the aforementioned problem.
energy differences between data and predictions. This allows for a more principled quantification of the errors made by the model, which we hypothesize provides better gradients for learning.
Our framework is general in the sense that encompasses many existing loss functions and allows us to interpret them energies. The energy loss functions also naturally capture relevant symmetries if the underlying energy approximation does. Specifically, they make so that no loss incurred by the model for predicting configurations that related to the data by symmetry. Loss functions that have this property have been suggested for atomistic systems, but they require expensive alignment or minimization procedures [Klein et al. 2023] which our framework does not require. This framework finds broad applicability to systems in thermal equilibrium, from direct regression tasks to generative modeling with diffusion models (Sohl-Dickstein et al., 2013 Ho et al., 2020 [ong al., (021). Note that we consider tasks that can be framed as regression and classification problems with data, sampling problems where are given ground-truth energy (like for example in Boltzmann generators [Noé et al. 2019)).
Contributions: (i) Methodology for deriving loss functions grounded in physical principles by minimizing the reverse KL divergence between a prediction and a Boltzmann distribution centered around data (ii) Instantiation of this framework for atomistic systems that yield distance-based loss functions and an analysis the invariance properties of these losses (iii) Applications to diffusion models and analysis of the resulting systems (v) Empirical evaluation on a range of tasks, showing consistent improvement over baselines.
estimator (iv) Instantiation of this framework for spin
2 Background
2.1 Forward and reverse KL loss functions
We first consider a regression setting. Consider the empirical distribution pn associated with the IID dataset D {x(1) (0), (+) 'iE and parametric model fo Rd R ý associated with the family of conditional distributions p (y Jo (x)). We take fe (x) to be the model prediction of the target; the usual assumption that conditional distribution is parametrized by a location parameter (like the mean for Gaussian), and the model is trained to maximize the likelihood of the data C (0) Σx log (y()) fo (x(+)). The Gaussian assumption for the model results in the Mean Squared Error (MSE) loss function. For n-way classification, the model predicts the logits a categorical distribution and maximum likelihood yields the cross-entropy loss function. Maximizing the likelihood is equivalent to minimizing the Kullback-Leiller (KL) divergence.
In this work, we will consider reverse KL divergence objective. In the regression case, this amounts to taking the model as deterministic and instead accounting for the uncertainty at the level of the data samples. For regression, we then have PD (x = Σn + (x - x:) (y y() and q (x.y) Σ: 1 specifies the uncertainty around each target. The reverse KL. objective is then
(x- x0) (y (x("))), where p (y (()) a conditional distribution that
N
DKL(g Pp) E, flog (x, y) log (x, = Σ logp (Fo (x(()) sr(i))
(1)

---

For classification, the model distribution is still categorical distribution parametrized by logits to ensure differentiability, but the distribution associated with data samples can be general.
In general, the reverse KI. divergence is not equal to the forward KI. divergence. Instead, it gives the likelihood of the prediction given uncertainty model for targets. However, it is exactly equal when the sample point and the parameter can be swapped in the distribution p. It is for example the case when p is chosen as Gaussian. Our general goal will be define more appropriate distributions p( y) for the loss function. As we will see, the reverse KL formulation is convenient since it enables defining these distributions only around each data sample.
2.2 Diffusion models
We also consider generative modeling with diffusion models another case for more informed energy loss functions. This class of generative models has proven powerful. as they can efficiently learn interpolations between a prior distribution and the data distribution [Albergo al. 2023]. The objective is typically formulated as a noise prediction task
[Ho et al. 2020
[ = J(- (xx-p(x)errrppee)) [walls ||- 20l1"] dt
where the noise prediction is the output of a neural network éØ = fo (xe, t). x, a+x + diE, oz, Og define the noise schedule and wy is weighting factor. In practice, the expectation is estimated by Monte Carlo. The objective also admits an interpretation as denoising score matching [Vincent 2011). with the optimal noise prediction satisfying e* (x, t) -o, x, log (x,).
The loss can equivalently seen as prediction of the data sample, with appropriate reweighting, see e.g. Kingma and Gao 22024]. With the sample prediction defined as Xa [x)_ox)a(x.. we have:
I (0) = / (xx-p(x),)
wiaa o?
x x01 dt
yielding regression-type objective with the MSE loss.
Energy Loss Functions
In Section 2.1 we saw that loss functions can be obtained through reverse KL formulation with respect to a conditional distribution p (y y) centered on the data. Importantly, the conditional distribution is always uncertainty model; as such, there is not necessarily a true one.
We will define the conditional distribution p as a Boltzmann distribution
p( yy) = exp ((E((yy)) y)TT)
7(y, T)
where E R x IR* > RR is related to the physical potential energy of the system around the data point y, is the temperature and Zyy T) is the partition function. We assume the system is observed in physically likely configurations; hence, each data point is modeled as an approximate local minimum in the energy landscape. The use of Boltzmann distributions to model the uncertainty around such configurations is natural and can be motivated from first principles (Jaynes] 1957, Pathria 2017]. the steady-stan state distribution of a system undergoing stochastic dynamics contact with reservoir at temperature 7' (see derivation in Appendix A.2)
Assuming general Boltzmann distribution, the reverse KL divergence Equation 1 loss function we obtain for the continuous case is
E (0))
(0) - Σ log (8)0))
+ log (,) T),
T
where the log- pege-pattiion function does not depend on the parameters. The model is penalized for errors by an amount given by the approximate increase in energy with respect to the data sample.
This picture allows for obtaining physical interpretation of different conditional distributions and loss functions depending on the choice of energy E(y.y). The Gaussian conditional distribution is
3

---

obtained with T = 2a² and isotropic harmonic potential energy centered around y:
E(g.y) = 1 yll
(6)
We can now justify our choice of the reverse KL estimation over maximum likelihood estimation. First, in the reverse KL case, the partition function, which could be challenging to evaluate for some energies, does not depend the model parameters 6. Second, functions around each data sample y()) rather than around each prediction. This is a significant advantage, as we can expect some predictions to be poor, leading to configurations that are not approximate equilibria and to nonsensical energies.
only need to define potential
3.1 Desiderata for energy functions
There is considerable freedom in the choice of the energy function. One fundamental criterion is agreement with the system's underlying physics, but this is not the only one. An appropriate energy should, in addition, satisfy the following desiderata:
The minimizer of the energy function Ely, y) should be the data (and its symmetry equivalents). Many tasks require regressing to the data even if not the minimum of the true energy landscape.
Minimized at the data and symmetries:
2. Optimization stability: The gradient of the energy function Vy Vy should be smooth and bounded to ensure well-behaved optimization with gradient-based methods.
3. Fast evaluation: Evaluation of the energy and its derivative should be efficient and compati- ble with automatic differentiation.
Based on this, we argue that one should not often use the true energy function, even if is known, since it may violate all the criteria. The energy landscapes of systems of interest typically admit multiple local minima and are highly rugged [Mézard et al., 1987 Frauenfelder et al. 1991 Wales et al. 2000]. The cost of evaluating the energy can also be prohibitive [Schuch and Verstracte 2009
4 Energies for Atomistic Systems
We consider energies associated with the positions of n atoms in d dimensions, such that y,y E Rax The potential energy Equation (6) leading the Gaussian distribution poorly motivated from the physical point of view. It describes the effect of an external force bringing back particles to position y. However, a realistic potential energy should model interactions between particles (see Figure Ic].
Many approximations exist for the potential energies of physical systems around equilibrium. For atomic systems, the Morse potential [Morse] 1929] and the Lennard-Jones potential [Lennard-lones 1931] are examples of popular models. However, using these potentials for the loss Equation (5) can pose challenges for optimization, as they have highly nonlinear gradients that can explode or vanish, A simple approximation that avoids this issue and that is much more principled than the MSE potential use quadratic pair potential of the form
Ely.y) = Σ 5 Kig (y) (lly: - yy 119 - -)1 1,j
where the indices j.j are taken over particles. This is the general form of a second-order Taylor approximation in pairwise distances of an interaction potential (see Appendix AA3)). Motivated by the fact that coordinate regression can lead to poor realism due to inconsistencies with bond lengths, this type of distance-dependent loss function has been used heuristically
regularizer in some applications [Peng et al. 2023 Yang and G6mme-BBombreell 2023 Abramson et al. 2024), but to the best of our knowledge, not primary objective. Note that this is different from directly predicting the distances [Simm and Hernández-
2019 Nesterov et al. 2020 Xu et al. 2021
There is significant freedom in the choice of the coefficients k.j (y). We propose simple heuristics to set these coefficients. First, we consider setting the coefficients are set to a constant value kij (y)
Note that this can be obtained from Taylor approximation of the Morse potential (see Appendix A.3)
Second, we consider setting the coefficient as a decreasing function of the distance between two atoms Kij (y) f(lly: yzll) to capture the fact that interactions between particles decrease at long

---

range. We consider inverse, unverse-squared and exponential decay dependence of f on the distance. Taylor approximation of the Lennard-Jones potential yields inverse squared distance dependence (see Appendix A.3) Other possibilities can be considered: in general, given an interaction potential between particles, the coefficients can be obtained by second-order Taylor expansion.
4.1 Invariance properties
An important property of energy loss functions is that they respect the symmetries of the associated physical energy. We formalize this in the following way:
Definition 4.1 (Invariant loss function). A loss function Z: IR x R - R between prediction and a target is invariant to the action of the group G on R if
(g y. (v.g.-9.
().y), Vg GY
(8)
An invariant loss function essentially compares input and targets up to transformations in G. An example of a common S (En(3-invariint loss function
apply the Kabsch algorithm [Kabsch] 1976)
to find an optimal alignment between a predicted structure and a target, and to use the MSE after
applying the alignment (Klein et al. 2023). J. It has been shown that in cases where there are multiple possible symmetry -reanttterrreetedd predictions for given input- so-called symmetry senmettr--breainn predictions (Smidt et al. 2022| Kaba and Ravanbakhsh] 2023 non-invariant loss functions exhibit pathological behaviour Xie and Smidi |024] Jing ct al. 2024 Lawrence et al.) [025]. For example, the MSE is minimized when the prediction is the mean of the possible targets rather than for any of them. There are multiple ways to define these invariant losses, which are analogous to the different ways in which invariant neural networks can be designed (see Appendix[AA4]fee more discussion).
It is easy to see that the energy loss Equation (5) is invariant to G = E(d) if kig (y) is invariant,
since it then only depends on invariant distances. This is analogous how invariant functions be built from scalars Villar et al. (202T]. These, however, are not the only symmetries of the loss. The energy loss function is additionally invariant to permutations that correspond to the symmetries of the ground-truth truth distance matrix. Denote the distance matrix of the data by Ayt = ||y; - y, || and the automorphism group of a matrix m _ as Aut (m) S S, where the automorphisms act on the matrix by conjugation. We then have the following:
Proposition 4.2. The loss function Equation is invariant to the group
G E(d) (Aut (y)) n Aut (Ay)).
All the proofs follow in Appendix A. This allows us to characterize the family of loss minimizers: Corollary 4.3. For any y E Raax and kis (y) > 0,
arg min E( {g y gEG). seRuxd
(10)
The loss landscape, therefore, presents family of global minimizers associated with symmetries. We hypothesize that the symmetry is beneficial for learning, since it allows the model to regress any target that is equivalent to the data by symmetry, as shown in the example in Figure 2 As our experimental results show, the benefits of invariance in the loss function are different and complementary to that of equivariance of the architecture. Equivariance guarantees that the output changes predictably under transformations of the input. However, it does not guarantee that the correct output will be learned. Invariant loss functions make the learning task easier by allowing to regress to any symmetry related configurations, which equivariance with a non-invariant loss does not allow.
4.2 Diffusion models with distance-based loss functions
The training objective of diffusion models involves the prediction of data sample from noisy latent one. We suggest that the energy loss functions can be used as a straightforward replacement for the MSE in these objectives. Similar distance-based objectives have been previously used in the context of diffusion models (Yang and Gömez-Bombares Bombbrelll 2023 Abramson et al.) 2024 Cognolato et al. 2025]. However, is not immediately clear that using such objectives results in learning correct score estimates. Here, we show that this is indeed the case under some conditions.
5

---

y1f
data
Ve
prediction
71
Yo
Yo
(a) MSE loss
(b) Energy loss
Figure Loss landscapes. The model has to predict the positions of two particles one dimension. The prediction for the first particle Yo is closer to the ground-truth for the second particle yı and vice- versa. (a) The MSE minimizes the forward KL divergence between a Gaussian model distribution (blue) and the data distribution (green). It does not capture the symmetry. (b) The energy loss is obtained via the reverse KL with the pair energy and admits a family of minimizers associated with symmetries. It results in gradient that points towards the closest correct configuration.
Consider the energy loss function Equation with constant coefficients kay (x) = k. The loss function computes the MSE between distance matrices for the data and the sample prediction, given
by Arg (x) || - xyll- Denote the Jacobian this function by J(x) = Vxd (x) E Rdnxx" We will seek to characterize the minimizers of diffusion model trained using this loss
Xx o.ē (x,, (s,,t) t)
é' E arg min (xxp(xx))e dERaxd
x
(11)
ax
The following result allows us to obtain an approximation of this set, valid for small noise scales: Proposition 4.4. Let p (xx) be continuously differentiable, (En()-iivariinn density, Assume |é] bounded. For small dt.
= lxg +v, v ker (J (xx))
(12)
In addition, the minimum norm minimizer Edist is given by
xx (x:)
(13)
The set of minimizers is therefore given by the true score, up to a translation in the direction of rigid motions. The second fact follows since for invariant measure, the score is orthogonal to the Lie algebra generators. We can also show that due to its invariance, the distance- based loss function offers reduction in variance with respect the MSE:
Proposition 4.5. Let p (xt) be continuously differentiable, SE(d) (E(d)-invariant density. Denote by dist and EaSE the minimum norm minimizers of the Monte-Carle estimators of the energy loss and MSF loss, respectively. For small
Bias [E Rd 0,
Var (E S Var (éWse)
(14)
Note that surprisingly, even though these results are in principle only valid for the energy loss function with constant coefficient ky (x), our results of Section 6.2, the variants using more physically motivated coefficients still performed better empirically. We hypothesize that this is because the evaluation of the model assesses the physical plausibility of the samples rather than the agreement between the learned and data distributions. Incorporating physical information in the loss function can therefore be beneficial, even though (or because) it biases the score estimate.
4.3 Linear scaling and rigidity theory
One potential downside of using the energy loss Equation (7) that it has quadratic number of terms in the number of particles N, in contrast to the linear number of terms in typical losses such as MSE loss. While in many architectures- such as transformers or densely connected graph neural networks- the quadratic cost of operations in the network makes this non-issue, the feasibility of linear scaling may prove valuable for applications involving a large number of particles, e.g., modelling macromolecules crystals with large unit cells.

---

Fortunately, a solution is provided by rigidity theory. Results in rigidity theory Laman 1970 Asimow and Roth 1978 provide the conditions for recovering the coordinates of a point cloud from linear number of pairwise distances. In this work, we consider construction for sparse rigid graphs. reducing the computational cost of energy loss without affecting its global optima (see Appendix A.5 for more background and Appendix C.4 for wall-times of different loss calculations).
5 Energy Loss for Discrete Systems
The energy loss formulation can be leveraged for other types of systems. We derive version for the discrete case, which can replace the cross-entropy loss function. Denote logits predictions as (i)
Z0.j and the associated categorical distribution q y z The reverse KL between the model distribution and a Boltzmann distribution around the data Equation
(4) is given by
N
( = T _
E ( -
E (5. - TS a
) + Tog Z (y(0,))
(15)
(5) is the model prediction for the logits associated with class j and Slg] is the entropy of q. The loss therefore proportional to the free energy difference. The last term is the negative free energy at the data, and does not depend on the parameters. The loss function, therefore, simply reduces to the variational free energy of the prediction.
where 20j
5.1 Application to spin systems
We consider modeling systems of spins as an application of the discrete formulation. Predicting configurations of these systems with machine learning models is a problem of high interest in physics [Carrasquilla and Melko 2017 Pahnng and Brenner 2020] and in combinatorial optimization (Fu and | 1986]. We will be interested specifically in systems
square lattice A such that
2 Σe Jıyyry, where the coupling -1 Jay < 1 is non-zero only for neighboring sites in the lattice A, but does not necessarily exhibit any symmetry. Systems with unstructured couplings are known as spin glasses [Mézard al. 1987] and often exhibit large number of local energy minima.
y,ye e 11 We consider Ising-type Hamiltonians of the form E (y) =-
Energy loss functions of the form Equation (15) can be used for classification of spin configurations. We suggest to use an approximate local energy around the data defined as
y( Σ ME (y)
(16)
where the local field is given by h!P E (Jg The local field energy captures the change energy from flipping a spin in the configuration y. It therefore provides an appropriate way to quantify deviations from that configuration: large values of local field are associated with spins that result in large increases energy and that should be weighted more importantly in the loss.
An alternative would be to use the true energy instead. The objective Equation (15) would then be interpreted as hurropyrregllarizer energy minimization. This would be expected to perform well in strict terms of minimizing the energy. However, the true energy is not classification objective, since it does not make use of the data. In addition, can exhibit large number of local minima. By contrast, the local energy loss Equation (16) is convex, due to the linear dependence in ý. It also admits the data point y as its unique minimum if No > 4 (see Appendix C.3). the data is a ground state of the true energy, ht > 0 is sufficient. The energy loss is therefore a proper classification objective.
6 Experiments
6.1 Regular shape prediction
Experiment Setup. To develop an understanding of energy loss functions, we propose a simple task where the goal is to generate regular shapes in two dimensions. Given a radius, model is tasked
7

---

Shape Generation
Energy Loss: Linear Sca ling
Qualie 0u
Quann
Energy iCompletel Energy (Sparee)
MSF Algn
1
Energy
MSE
n//4
n/2
3T/4 T
2
25
2
2
Augmentation Angle
Number Vertices
Figure 3: Regular shape prediction results. (a) Typical samples from optimal models trained with MSE and energy loss when oag = т. (b) The impact of aug sample quality. We can see as aagg increases, MSE performance deteriorates but the invariant losses (Energy and Kabsch Align) remain performant. (c) As the number of shape vertices scales, sparse version of the energy loss remains equally performant as a complete-edge energy loss using only O(N) operations.
with predicting the N vertices of a regular polygon of that radius. The dataset is constructed by sampling regular polygons of a radius r ~ 0.3,, 5| and then applying an augmentation by randomly rotating the shape by an angle in Ul-00mg Oagg]. Prediction is performed using two hidden-layer MLP. We compare standard MSE loss with the atomic energy loss using exponential coefficients, an SE(2)-invariant loss using the Kabsch algorithm to align points, and a version of the energy loss using sparse rigid graphs. We empirically confirm that the sparse graphs in Appendix A.5) To evaluate, we introduce a quality metric based on the regularity of the angular differences and the radial variation in given shape. Intuitively, for regular shape the angular difference variation 0 Auugte and the radial variation radius across points should be small. A full definition follows in Appendix C.1 C.
globally rigid w.h.p.
Results. Figure shows that the energy loss and other invariant losses continue to produce high- quality shapes when rotation augmentation is applied whereas MSE fails. Additionally, the sparse energy loss maintains nearly the same performance as the number of vertices N increases, while reducing computation by O(N) operations. Interestingly, models trained with an invariant loss automatically learn to produce canonical orientations of shapes.
6.2 Molecule generation
Table Metrics for GDM-aug on GEOM-Drugs.
Experiment Setup. First, we train diffusion models uncon- ditionally generate molecules in
Loss
Mol. stab, (%)
Atom stab, (%)
Valid, (%)
Unique (%)
MSE Energy
0.8
85.6
94.8
100
the QM9 dataset [Ramakrishnan) et al. 2014]. We evaluate the performance of the energy loss
24.6
96.0
89.7
when training EGNN diffusion models (EDM) [Hoogeboom et al. 2022], GNN diffusion models with and without data augmentation (GDM and GDM-aug) and near sfate-of-the-art joint 2D & 3D diffusion models (JODO) Huang et al. [2023]. As baselines, compare the convergence properties to models trained with MSE and a Kabsch-aligned MSE [Kabsch] 1976). Exponential coefficients are chosen for the energy loss. Additionally, compare with version of the energy loss using sparse rigid graphs.
Table Evaluation metries for GDM-aug on QM9,
Loss
Molecule stability (%)
Atom stablity (%)
Validity (%)
Uniqueness (%)
GDM-aug
MSE
83.7 + 2.3
98.3 + 0.004
93.6 +: 1.7
100.0 +: 0.0
Kabsch align
82.3
0.5
97.8 + 0.004
90.8 + 2.0
100.0 +: 0.0
Energy
89.8
2.8
99.3 + 0.3
97.7
1.4
99.9 + 0.002
Energy (sparse)
89.1 0.9
99.0 99.0 + 0.1
97.4 2.5
100 + 0.0
EDM
MSE
82.4 3.4
98.8 + 1.7
93.0
2.5
99.89 0.32
We also generate large molecules with GDM and GDM-aug using the GEOM-Drugs dataset Axelrod and Kromez-Bombare 2022), comparing the MSE and energy loss. similar evaluation setup
to Satorras et al. 2022 Hoogeboom et al. 2022 used for GDM and EDM while JODO uses
8

---

Energy
Align
Energy (Sparse)
GEOM-DRUGS: Molecule Stability
GEOM-Drugs: Atom Stabilitt
QM9: Dataset Size
QM9: Molecule Stability
Molecule Stability (%)
Molecule Stability (%)
Molecule Stability (96]
100
20
Atom Stabilty (%)
80
80
90
15
60
60
80
10
40
40
5
20
60
0
6
0.5
1.0
2
Step
Step
Dataset Fraction
Step
165
Figure 4: Molecule generation results. (Left) We observe a dramatic improvement on stability metrics for the GEOM-Drugs dataset, demonstrating the scalability of our approach. (Right) On QM9, energy loss improves metrics over all baselines. This is especially present in the low data regime where energy loss gives +10% molecule stability over MSE.
Table 3: 3D and alignment metrics for JODO variants.
Metric-3D
Metrie-Align
Model
At. stab
Mol. stab, (%)
Val. (%)
Compl. (%)
FCD
Bond
Angle
Dihedral
JODO (paper)
99.2
93.4
0.885
0.1475
0.0121
6.29c-4
JODO (ours)
99.2
92.8
95.6
95.5
0.854
0,1218
0.0110
5.91e-4
JODO + Energy (Inv.)
99.4
94.3
97.1
97.0
0.892
0.1125
0.0046
4,95e-4
JODO Energy (Exp.)
99.6
96.6
98.4
98.4
1.495
0.0928
0.0142
4.97c-3
broader set of 3D and align metrics Huang et al. (2023). Since the MSE is longer the optimization objective, we no longer report the ELBO, which depends on the MSE. Instead, we evaluate the method using several desirable features of generated molecules relevant to the drug discovery pipeline: atom stability, molecule stability, validity, and uniqueness. Additionally, for JODO, we compute the Maximum Mean Discrepancy (MMD) for bond lengths, bond angles and dihedral angles against the data distribution as well as the Fréchet ChemNet Distance (FCD) Preuer et al. [2018). For all settings, conduct exhaustive sweeps for learning rate and the weighting between the loss on positions and atom types. All comparisons are compute-matched. matched.
Results. Figure 4) shows the energy loss results in faster convergence and better optima over base- lines" In addition, we observe that energy loss is much more data efficient than baselines, allowing for the training of capable molecular generative models, producing over 75% stable molecules using only 50% of the training set (50K samples). Table | contains results on the GEOM-Drup data. Sablee]shows results on the QM9 data with GDM-aug model and and its equivariant variant EDM, with comprehensive results in Appendix C.2 Importantly, Table 2 shows that energy loss with non-equivariant architecture results in more improvement than using an equivariant architecture, at negligible computational cost.
The results using the JODO model are reported in Table 3 We observe that using energy loss with JODO is able to improve all align metrics, and nearly all 3D metrics, with comparable FCD, compared the default Kabsch-aligned loss. This suggests that energy loss can push the state of the art and offers complementary benefits to equivariant architectures.
Ablation. We conduct an ablation over the form of the spring coefficients kij (y) in the energy loss (Table We consider the following func-
Table 4: Ablation for coefficients Ki (y).
Coeff.
Mol. stah. (5)
Atom stah, (%)
Valid. (%)
tional forms: constant, inverse distance, inverse square distance, and exponential decay. thorough sweep over learning rates was conducted. With EDM/GDM, we find exponential decay to give the best empirical results. However, inverse distance coefficients work well for JODO and Appendix C.4 shows a less stark decay works better for the
Exp. Dist.
89.8 +: 2.8
99.3 0.3
97.7 #: 1.4
Ins. Sq. Dist.
84.6 1 1.8
98.9 0.2
96.6 1: 1.5
Dist.
1.5
Constant
83.6 1: 1.5
98.7 0.1
93.6 1: 0.7
We note that our results with MSE are better than those reported in Hoogeboom et al. 2022 and attribute this exhaustive learning rate tuning.

---

sparse loss function on large molecules. This suggests it is necessary to ablate these coefficients on new tasks.
6.3 Spin ground state prediction
Experimental Setup We consider the task of predicting the ground states of the spin Hamiltonian and compare the effectiveness of different loss functions. We construct a dataset of 10,000 training
and test apin-glass Hamiltonians, each with couplings uniformly sampled from [-1, 1]. We consider
grids of size 16 16, which offer challenging problem. The target ground-states are obtained by solving the associated integer linear program [Billionnet and Elloumi 2007]. A convolutional neural network (CNN) trained to take as input coefficients Jij and predict the ground state configuration ý; More details on the architecture and training setup are provided in Appendix C..
We compare training with the energy loss function of Equation (16) to the cross-entropy loss function and the margin loss function, another commonly used Ioss function for classification. The evaluation metric we consider is the energy of the predicted configuration. We also compare with direct minimization of the true energy baseline, despite it not being classification objective.
Results The results in Ta- ble [ show that using the local energy leads to
Table 5: Results on ground-stan state prediction.
Loss
Cross-entropy
Margin loss
Local energy
True energy
lower configuration en- ergies than the cross- entropy loss function and
Test energy
58.8 + 0.8
49.87 + 1.5
45.6 + 1.6
14.6 + 0.3
the margin loss function. As expected, minimizing the true energy still leads to lower overall energy, despite not using the data. The seel-field loss also requires fewer training epochs to converge than the cross-entropy loss. The results support the hypothesis that directly embedding physical insights through the local-field formulation effectively guides the learning process toward physically meaningful predictions.
7 Conclusion
We demonstrated a new approach to designing loss functions for machine learning tasks in physical systems based on the system's energy. When applied to both continuous and discrete settings, we found that replacing a simple MSE or cross-entropy loss with our energy loss functions leads to improved predictions across experiments. We further demonstrate the suitability of this loss for diffusion models and analyze its symmetry-invariance properties and scalability.
Limitations and future work Some limitations remain, which also point to directions for future work. First, when energy loss functions are used for diffusion models, the correct score is recovered at low noise levels; exact recovery at higher noise levels would require an explicit correction, which we leave for later study. Second, while offer more principled approach to designing loss functions, some choices are still ad hoc. Looking ahead, richer surrogate energies that capture torsional angles could be investigated. The approach could also potentially fruitfully extended a broader class of systems, including crystalline materials and proteins.
Acknowledgments
We are thankful to Mohsin Hasan, Simon Blackburn, Bruno Rousseau and Simon Verret for helpful discussions. This research was supported by the CIFAR AI Chairs program, Intel AI Labs and NSERC Discovery. S.-O. K.'s research is also supported by IVADO and FRQNT, and D.L. research additionally supported by FRQNT. Mila and Compute Canada provided computational resources.
References
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew Ballard, Joshua Bambrick, et al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, pages 1-3, 2024.
10

---

Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint vrXiv:2303.08797, 2023.
Leonard Asimow and Ben Roth. The rigidity of graphs. Transactions of the American Mathematical Society, 245:279-289, 1978.
Simon Axelrod and Rafael for property prediction and molecular generation, 2022. URL httpp::/aaxxx..oo/aaaa/////////a
Geom: Energy-annotated conformatoss
05531
Jan-Hendrik Bastek, WaiChing Sun, and Dennis M Kochmann. Physics-informed diffusion models. arXiv preprint arXiv:2403. 14404, 2024.
Alain Billionnet and Sourour Elloumi. Using mixed integer quadratic programming solver for the unconstrained quadratic 0-1 problem. Mathematical programming, 109:55-68, 2007.
Juan Carrasquilla and Roger G. Melko. Machine learning phases of matter. Nature Physics, 13(5): 431 0344, 2017. doi: 10. 10880/8mhyy44035. /ttLlatttt::/o.orgg////// wtt..o.../////pppp...oo 1038//phy44/5
Subrahmanyan Chandrasekhar. Stochastic problems in physics and astronomy. Reviews of modern physics, 15(1):1, 1943.
Samuel Cognolato, Davide Rigoni, Marco Ballarini, Luciano Serafini, Stefano Moro, Alessandro Sperduti, al. D4: Distance diffusion for truly equivariant molecular design. In ESANN 2025- Proceedings, 33rd European Symposium on Artificial Neural Networks, Computational Intelligence and Machine Learning, pages 265-270, 2025.
Miles Cranmer, Sam Greydanus, Stephan Hoyer, Peter Battaglia, David Spergel, and Shirley Ho. Lagrangian neural networks. arXiv preprint arXiv:2003.04630. 2020.
Sean Dewar. graph-rigidity-cheeer library.
https://github.co http://tthubbbccmm//wwaa22/
rrapp-rrgiidtyychcckekr 2025.
Yilun Du and Igor Mordatch. Implicit generation and modeling with energy based models. Advances in neural information processing systems, 32, 2019.
Albert Einstein. On the motion of small particles suspended in liquids rest required by the molecular-kinetic theory of heat. Annalen der physik, 17(549-560):208, 1905.
Hans Frauenfelder, Stephen G Sligar, and Peter G Wolynes. The energy landscapes and motions of proteins. Science, 254(5038):1598-1666 1991.
Yaotian Fu and Philip W Anderson. Application of statistical mechanics np-complete problems in combinatorial optimisation. Journal of Physics Mathematical and General, 1905):6005, 1986.
Crispin W Gardiner. Handbook of stochastic methods for physics, chemistry and the natural sciences. Springer series in synergetics, 1985.
Rafael Gómezz-ombarelli, Jennifer N Wei, David Duvenaud, José Miguel Hernández-Loobate Benjamín Sänchez--Lengllingg Dennis Sheberla, Jorge Aguilera-Iparraguun Iparraguirre, Timothy D Hirzel, Ryan P Adams, and Alán Aspuru-Guzik. Automatic chemical design using data-driven continuous representation of molecules. ACS central science. 4(2):268-22 276.,27762220 2018.
Will Grathwohl, Kuan-Chieh Wang, Jörn-Henrik Jacobsen, David Duvenaud, Mohammad Norouzi, and Kevin Swersky. Your classifier is secretly energy based model and you should treat it like one. arXiv preprint arXiv: 2019.
Samuel Greydanus, Misko Dzamba, and Jason Yosinski. Hamiltonian neural networks. Advances in neural information processing systems, 32, 2019.
Majdi Hassan, Nikhil Shenoy, Jungyoon Lee, Hannes Stärk, Stephan Thaler, and Dominique Beaini. Ft-flow: Equivariant flow-matching for molecular conformer generation. Advances in Neural 2024.
Information Processing Systems,

---

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In H. Larochelle, M. Ranzato, R. Hadsell. M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 6840-6851. Curran Associates, Inc., 2020. URL https::/prroeeddings iew.reeeeddiignnnneeeeieeee neurips...... cccccccppp.. ceboeeccc544400007711100000000--
pdf
Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d, 2022. URL htllhtttt//araate http://wwwwvvr.r.././///oo abs////a.2222... 17003
Han Huang, Leilei Sun, Bowen Du, and Weifeng Lv. Learning joint 2d 3d diffusion models for complete molecule generation, 2023. httpsstss://arxi org wtw./305. 12347
Edwin T Jaynes. Information theory and statistical mechanics. Physical review, 106(4):620, 1957.
Rui Jiao, Wenbing Huang. Peijia Lin, Jiaqi Han, Pin Chen, Yutong Lu, and Yang Liu. Crystal structure prediction by joint equivariant diffusion. Advances in Neural Information Processing 3y:tems,36111 336:17444-179999 2023.
Bowen Jing, Bonnie Berger, and Tommi Jaakkola. Alphafold meets flow matching for generating protein ensembles. arXiv preprint 2402:22402.048455, 2024.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, al. Highly accurate protein structure prediction with alphafold. nature, 596(7873):583-58) 2021.
Sékou-Oumar Kaba and Siamak Ravanbakhsh. Symmetry breaking and equivariant neural networks. arXiv preprint arXiv:2312.0000 2023.
Sékou-Oumar Kaba, Arnab Kumar Mondal, Yan Zhang, Yoshua Bengio, and Siamak Ravanbakhsh. Equivariance with learned canonicalization functions. In International Conference Machine Learning, pages 15546-15566. PMLR, 2023.
Wolfgang Kabsch. A solution for the best rotation to relate two sets of vectors. Foundations of Crystallogruphy, 32(5):922-923, 1976.
Diederik Kingma and Ruiqi Gao. Understanding diffusion objectives as the elbo with simple data augmentation. Advances in Neural Information Processing Systems, 36, 2024.
Lcon Klein and Frank Noé. Transferable boltzmann generators, 2025. URL abs/2406 14426
org/
Leon Klein, Andreas Krämer, and Frank Noé. Equivariant flow matching, 2023. URL https
//arxiv. org/ abs/ 15030
Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: exact likelihood generative learning for
symmetric densities. In International conference on machine learning, pages 5361-5370. PMLR, 2020.
Hendrik Anthony Kramers. Brownian motion in field of force and the diffusion model of chemical reactions. physica, 7(4):284-304, 1940.
Michael Krivelevich, Alan Lew, and Peleg Michaeli. Rigid partitions: from high connectivity to random graphs. arXiv preprint arXiv:2311.14457 2023.
Jerome M Kurtzberg. On approximation methods for the assignment problem. Journal of the ACM (JACM), 9(4):419-439, 1962.
Gerard Laman. On graphs and rigidity of plane skeletal structures. Journal of Engineering mathemat- ics, 4(4):331-340, 1970.
Paul Langevin. On the theory of brownian motion. CR Acad. Sci. Paris, 146(530-533):533) 1908.
Hannah Lawrence, Vasco Portilheiro, Yan Zhang, and Sékou-Qumar Kaba. Improving equivariant networks with probabilistic symmetry breaking. In The Thirteenth International Conference on Learning Representations, 2025. URL httphhttt://pereee///e wwt.rriiww..eeeoooorrddddd/sL...
12

---

Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, Fujie Huang, et al. A tutorial on energy-based learning, Predicting structured data, 100) 2006.
J E Lennard-Jones. Cohesion. Proceedings of the Physical Sociery, 43(5):461, 1931, doi: 10. 0959-5309/43/5/301 URLhttps: doi orggi..rg//i 10.
Alan Lew, Eran Nevo, Yuval Peled, and Orit E. Raz. Sharp threshold for rigidity of random graphs. Bulletin of the London Mathematical Society, 55(1):490-55) 5019 2023, doi: https://doi.or//10.1112
blms. 12740.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling, The Eleventh International Conference on Learning Repre- sentations, 2023. URL http::t/opperreeeiiiie// httppp/rreiiooooooooooorrrdddddd htte:/oorrmm111-Py//DD.....
Marc Mézard, Giorgio Parisi, and Miguel Angel Virasoro. Spin glass theory and beyond: An Introduction to the Replica Method and Its Applications, volume 9. World Scientific Publishing Company, 1987.
Philip M Morse. Diatomic molecules according to the wave mechanics. ii. vibrational levels. Physical review, 34(1):57, 1929.
Vitali Nesterov, Mario Wieser, and Volker Roth. 3dmolnet: a generative network for molecular structures. arXiv preprint MxXi:2000000677,, 2020.
Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(6457) 1147, 2019.
Frank Noé, Gianni De Fabritüs,, and Cecilia Clementi. Machine learning for protein folding and dynamics. Current opinion in structural biology, 60:77-84, 2020.
Seong Ho Pahng and Michael Brenner. Predicting ground state configuration of energy landscape ensemble using graph neural network. arXiv preprint arXiv:2008.08227 2020.
Raj Kumar Pathria. Statistical Mechanics: International Series of Monographs in Natural Philosophy, volume 45, Elsevier, 2017.
Yuval Peled. Sharp threshold for rigidity of random graphs. Presented at IPAM UCLA, 2024.
Xingang Peng, Jiaqi Guan, Qiang Liu, and Jianzhu Ma. Moldiff: Addressing the atom-bond inconsistency problem in 3d molecule diffusion generation. arXiv preprint wxXiv:335.007503 2023.
Kristina Preuer, Philipp Renz, Thomas Unterthiner, Sepp Hochreiter, and Günter Klambauer. Fréchet chemnet distance: metrie for generative models for molecules in drug discovery. Journal of Chemical Information and Modeling, 1839):173-1174) 2018. doi:
Maziar Raissi, Paris Perdikaris, and George Karniadakis. thysics-informed neural networks: deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-707, 2019.
Raghunathan Ramakrishnan, Pavlo O. Dral, Matthias Rupp. and O. Anatole von Lilienfeld. Quantum chemistry structures and properties 134 kilo molecules. Scientific Data, 1(1):140022, 2014. ISSN 2052-4463. doi: 10.1038/sdata.20) 2014.22. URL httph:///dooi... org/ 1033//data........ /wwta.daaaaaa//aaa/eaaa 2data.2014.2.......
Kevin Ryan, Jeff Lengyel, and Michael Shatruk. Crystal structure prediction via deep learning. Journal of the American Chemical Society, 140(3)):0115-11016,, 2018
Benjamin Sanchez-Lengeiee and Alán Aspuru-Guzik. Inverse molecular design using machine learning: Generative models for matter engineering. Science, 361(6400):66-336 2018.
Kusha Sareen, Daniel Levy, Arnab Kumar Mondal, Sékou-Oumar Kaba, Tara Schound-Sadegll and Siamak Ravanbakhsh. Symmetry- generative modeling through learned canonicalization.
2025. URLptttps: /// aarxivv org/ abs httg//ab/25000.10 07773
13

---

Victor Garcia Satorras. Emiel Hoogeboom, Fabian B. Fuchs, Ingmar Posner, and Max Welling. E(n) equivariant normalizing flows, 2022. Jttptsss://axxi org/ http:/abb/22/1...51 09016
Norbert Schuch and Frank Verstraete. Computational complexity of interacting electrons and
fundamental limitations of density functional theory. Nature physics, 5(10):732-735, 2009,
Gregor NC Simm and José Miguel Hernández-Lobato A generative model for molecular distance
geometry. arXiv preprint arxiv: 1459, 2019.
Tess E. Smidt, Mario Geiger, and Benjamin Kurt Miller. Finding symmetry breaking order parameters with euclidean neural networks. Physical Review Research, 3()::L012002. 2021.
Jascha Sohl-Dickstee Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In Francis Bach and David Blei, editors, Pro- ceedings of the 32nd International Conference on Machine Learning, volume of Proceedings of Machine Learning Research, pages 2256 2265, Lille, France, 07-09 Jul 2015. PMLR. URL htt///hhlllldeee htttss:/proceeddingg mlr http:ppse//////oooccccccccciaaaa html
Yang Song, Jascha Sohl- Sohl-Dicksteen, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastie differential equations. In International Conference on Learning Representations, 2021. URL /ttp:ttt///opaaeeeeeee.ee net/ wtt.:/wwwweeeeoummmmm/mm id=
PXTIG12RRHS
M.F. Thorpe and P.M. Duxbury. Rigidity Theory and Applications. Fundamental Materials Research. Springer US, 1999. ISBN 9780306461 156. URL https::/bbokks http//bokkkoooglee..... wwwgle.ca.books.....e.eoooooo caw.baakbooksdeeeaoooaoa......aa
3XgykKbZymKC
Soledad Villar, David W Hogg, Kate Storey-Fisher, Weichi Yao, and Ben Blum-Smith. Scalars are universal: Equivariant machine learning, structured like classical physics. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances Neural Information Processing Systems, 2021.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computa- tion, 23(7):1661-1677 2011.
David Wales, Jonathan PK Doye, Mark A Miller, Paul N Mortenson, and Tiffany R Walsh. Energy landscapes: from clusters biomolecules. Advances in Chemical Physics, 115:1-111, 2000.
Shih-Hsin Wang, Yuhao Huang, Justin M. Baker, Yuan-En Sun, Qi Tang, and Bao Wang. A theoretically- prrinccplly-piineppeeedd sparse, connected, and rigid graph representation of molecules. In The Thirteenth International Conference on Learning Representations, 2025. URL https
YuQing Xie and Tess Smidt. Equivariant symmetry breaking sets. Transactions on Machine Learning Research, 2024. ISSN 2835-8856. URL htttttt///opperevieeee wwt.:nteoooooomdddddtdteoooeoeee
Minkai Xu, Shitong Luo, Yoshua Bengio, Jian Peng, and Jian Tang. Learning neural generative
dynamics for molecular conformation generation. arXiv preprint arXiv:2102.102400 2021.
Soojung Yang and Rafael Gómezz-ombarelll. Chemically transferable generative backmapping of coarse-grained proteins. arXiv preprint arXiv:2303.0f 2023.
Claudio Zeni, Robert Pinsler, Daniel Zügner, Andrew Fowler, Matthew Horton, Xiang Fu, Zilong Wang, Aliaksandra Shysheya, Jonathan Crabbé. Shoko Ueda, et al. A generative model for inorgamic materials design. Nature, pages 1-3, 2025.
Leo Zhang, Kianoosh Ashouritaklimi, Yee Whye Teh, and Rob Cornish. Symdiff: Equivariant diffusion via stochastic symmetrisation. arXiv preprint arXiv:2410.06262, 2024.
Xuan Zhang, Limei Wang, Jacob Helwig. Youzhi Luo, Cong Fu, Yaochen Xie, Meng Liu, Yuchao Lin, Zhao Xu, Keqiang Yan, et al. Artificial intelligence for science in quantum, atomistic, and continuum systems. arXiv preprint arXiv:2307.08422. 2023.
14

---

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: As stated in the introduction, we introduce new perspective on loss functions for physical systems (Section 3 analyze their properties (Section 4.1 Section 4.2). and
demonstrate the benefits this approach via multiple experiments (Section 6
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
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 be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discuss assumptions made theoretical claims, and in the conclusion Section 7 we point out further limitations that point to future research directions.
Guidelines:
The answer NA means that the paper has 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 separate "I "Limitations" section in 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 on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or 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 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 in the paper. The authors should use their best 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.
3. Theory assumptions and prool's
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
15

---

Answer:
[Yes]
Justification: Proofs are included in Appendix A.1
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 or referenced in the statement of any theorems.
The proof's 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 a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or 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 not)?
Question:
Answer: [Yes]
Justification: We provide experimental details in the paper, and further details are provided in Appendix C
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 dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is 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.. g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) 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-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
16

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [No]
Justification:
We intend to make the code public upon acceptance.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cm cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a 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
//nips //iips.cc/pill wttp:ccppppcccceeeeeeeeee..c.ett httaa//gguudeee/eeeeebiiiisssss /Cc////oddiiisssseeeeeee ssiiiiiiiiyiyyyyy 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.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they 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
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 understand the results?
Question:
Answer: [Yes]
Justification: Basic details about training and testing are included in Section 6 with full details included in Appendix C
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to 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: [Yes]
Justification:
We explain our use of error bars in Appendix C based on multiple seeds.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
17

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one 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 errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or 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 execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Compute resource details are included in Appendix
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 as 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 it into the paper).
9. Code ethics
Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Question:
Answer: [Yes]
Justification:
The paper conforms with the Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification: Our contributions meant be general enough to apply very wide class problems within the field of AI for science, including chemistry, physics, and biology. We consider these applications be too broad discuss any particular social impacts, positive or negative.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
18

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment 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 is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are 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 to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
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)?
Question:
Answer: [NA]
Justification:
No new data pretrained models are released in this paper.
Guidelines:
The answer NA means that the paper poses no 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
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We cite sources for the datasets that we train models on, and the models that we compare against.
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 or dataset.
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.
19

---

If assets are released, the license, copyright information, and terms of use in the thcode has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssit
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: No assets are released.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
No crowdsourcing or human subjects were used.
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 in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or 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 the requirements of your country institution) were obtained?
Answer: [NA]
Justification: No crowdsourcing or human subjects were used.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending 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.
20

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere 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
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: LLMs were not used as an important, original, or standaadd component of the core methods of this research.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM for what should or should not be described.
21

---

A
Additional Theory
A.1
Proofs
A.1.1
Proof of Proposition 4.2
Proof. For any (g1,92) 92) E E(d) x (Aut (k (y)) n Aut (Ay)) where 91 E E(d) and 92 E
(Aut (k (y)) n Aut (Ay)), we have
E ((gn, 8,y) =
kiy (y) (1yrs - y5ll (1)1,92) 8; - (91,92) 85ll)² (17)
By linearity of the actions, we have
E Elg1, 92) y,y) Σ kis (y) (lly: -yyll 1911 (92 5z 92- $1))
(18)
Since the Euclidean norm of difference is E(d)) have
E (())g2)) y.y)
Σ
kig (y) (lly; Ysll - -1(022 8; 92 $1)11)
(19)
2
2
We then have
E ((91)922) V.y)
Σ
n
1
kis (y) (1y: 1351- poll-1166 -
(20)
1,j
Using the fact that 92 E (Aut (1s (y)) n Aut (Ay)).
E ((g1, 92) y.
1 '1 (x) (1x05'.-- yor', - -11550'' - -8555')1
2
(21)
2
n2
E ((91, (
Σ 92'
1
kis (y) (lly: -ysl 11(8=- -33)113
(22)
2
Since the sum is permutation invariant, have
E ((91.9) 92) y.y) Σ
Res (y) (lly,- - y,ll 11(: - };)1)²²
(23)
1
E ((g1. 92) y.y) = E (},y) For the second argument, we similarly have
(24)
E (ý,, (91 92) y) = (y,y),
(25)
due to the -ivaariive of the Euclidean norm and the automorphism of y.
This completes the proof.
A.1.2 Proof of Corollary 4.3
Proof. First,
E(y.y)= = =0= min yeRuxd
E(j,y)
(26)
Then, given
E Ey.y): Σ
kis (y) (lly: - y5ll - -18: -))1))
(27)
22

---

we see that for any Kay (y) > O the loss is minimized when each term of the sum is zero which implies that Ay = Ag when the loss minimized. This implies
arg Efy,y) -g-y yyll g B(d))
(28)
yERexa
Since the action of the group Aut (I+ (y)) O Aut (Ay) on by definition trivial, the desired result follows.
A.1.3 Proof of Proposition 4.4
Proof. We first establish some preliminaries. Since SE(d) is not compact, do not assume that the density p (xx) is normalized. This is not an issue since the results depend on the score Vx, log p (x:), which is independent of the normalization.
We also parametrize the score in terms of noise prediction align with practice. The diffusion objective for time is proportional to (multiplication by constants do not change the minimizer)
1
0t Eg
Je (0) expp)
X₁
x
(29)
0t
az
where EΦ = fo tx, X axx ofe and p(c) = N (0, I). The sample prediction is given by
&0 (x (xxxoc((xxx.t)
The expectation can be rewritten
1
Jt
(0) Ex, )x
(30)
a
at
1
Tt
(0) = Ex ((xx) Excp(x)a))
-
(31)
Og
Ott
Je (0) = =Exxp((xx)) [Te (xe, 0)]
(32)
The objective is minimized if for all we have the following noise prediction
1 Xt
MSE arg min Je (Xe, 0) arg min Excop(o(xx)) ERoad
Es
x
(33)
EyERaxd
Since expectation and minimization commute for the MSE, have
1
MSS Exxxp(xixi)
arg min G
En
(34)
O4
O4
MSB
-Exxp(xxxx) [X: - [xx]
(35)
Gt.
Using Tweedie's formula, we obtain the usual score matching relationship:
MSE a. Vx, log p (x:)
(36)
To prove our result, we consider the energy loss objective,
To (xe, 0) = Ex-p(x(x))
E
1
Tt as
20, x
(37)
a1
where
n
71
E(x, x)
Σ (||x; x; 1| x111)²
Σ
(x)); - d d3)))
(38)
23

---

and d: R2xd + R= computes the distance matrix.
We perform Taylor approximation of d (x) and dl (8) around
d (x) = d (xx) J (xx) (x xc)
(39)
where J (ý) is the Jacobian of d.
This approximation is valid when x x 1 and x - 1 which corresponds |o, <
and |oy]| < respectively. Assuming that e follows a standard normal distribution, that is bounded, for usual diffusion schedules (with monotonically increasing noise) the approximation will be valid for small values of
This yields
E(x,x) = ||J (xx) (x x)1?
(40)
Replacing in Equation
(37). we have
|
[_xx 04
Os En at
>1
To (xt, Expp(x((x))
(41)
Since Jj convex function of é, its minimization with respect éø can be performed by solving for vanishing gradient, which leads to the minimizer é" satisfying
J (xx)² J (xx)
E Fx-p(xix)) [x+ [xx] = 0
(42)
Ti
J (x,1² (xx) (ē" of Vx (xi)(x:)) =0
(43)
Because of the SE(() invariance d, J (xx) does have full-rank. Denoting the kernel space of (xt) as ker (J (x:)), we obtain
é" = P
J(x.) ((x))
Xx logp(xi) +v, v e ker (J (xx))
(44)
where P (xx)) is the orthogonal projector onto range (J (x:)).
Using the assumption that the measure p (xx) is orthogonal to any v E ker (J (x:)). The norm é* therefore minimized when it equal to the score. This completes the proof.
the score lies in range (J (xx)) and
A.1.4 Proof of Proposition 4.5
Proof. We assume the same approximation regime as Proposition 4.4
The Monte-Carke estimators associated with the two losses (for the distance loss, we assume the minimum norm minimizer) are given by
N Σ X (1)
1 N T N
Eiss = - ((x)))
1 N
EMSE
Σ X: (1)
(45)
where the N samples (0) drawn i.i.d. from p(x xt).
Bias [ = 0:
This simply follows from Equation (36) and the fact that the true score for an
density
lies in range (J (x:)), so that oe J(x)) (x,) Vx, (x.).
Var [ S Var [
We have
Cov [etal] P P(()))) Cov [EMsEl P
Cov
f(x.)
(46)

---

We can then obtain,
Cov Tr (Cov (ist)) = Tr (P (P (xx, Cov [EMe (x(x)))
(47)
Tr (Cov (Eiial) Tr (P [(xx) Cov (EMSFE))
(48)
which compared
Tr (Cov (éMEE)) Tr (P J(x.) Cov (ÉMSEl) + Tr ((I P ((x.)) Cov (cMSEl
(49)
since (I- [I-P ((x))) and ((x,) are both positive semi-definite, we conclude
Tr (Cov (Cdia) S Tr (Cov (ÉMSE))
(50)
which completes the proof.
At a high-level, the projector ((x,) removes the components associated with rigid motions from the estimated score. The minimum norm assumption amounts to the network denoinsing to the closet sample to x, in the orbit of x.
A.2 The Boltzmann Distribution is the Steady Distribution
We consider the potential energy y) around configuration y. We assume in an approximate equilibrium configuration of the system, and therefore a local minimum of the potential. We model the evolution of the system using Newton's 's equation (in Hamiltonian form), with an additional term modeling the noise. This gives the SDES
0K
DE 0y dt Pdt + VdTdW.
dý dt, OP
dP
(51)
where W is generically taken as standard Brownian noise and K the kinetic energy the system. The term proportional to P in the second equation represents the effect of friction and has the effect of bringing the system back towards the equilibrium configuration. This is the Langevin equation, introduced by Einstein| [1905]. Langevin [1908] to study Brownian motion. It describes the evolution of a system in a heat bath at temperature T; this is the uncertainty model we consider.
We suppose momentum variables are not of interest. We therefore take the kinetic energy (or masses) as negligible, obtaining the overdamped limit Kramers 1940)
0E dt /2TaW
dy
(52)
ay
The probability density y evolves according to the Fokker-Plank equation Chandrasekhar [943
Öply, t y) TV3P($, [y) + Vy [p(,, v))
(53)
ØL
For anything but the simplest potentials, this equation admits no known closed- form solution. How- ever, under some growth conditions on the potential Vyy) (see e.g. Gardiner [9855). ). the Fokker-Plank equation admits the Boltzmann distribution unique steady-state solution when f
p( y)
T exp (-)(((()))TT
(54)
where Z(y,77 T) the partition function.
A.3 Second-order Taylor Approximations of Potential Energies
Below we justify the two practical choices for the coefficients kij (y) (constant and inverse-squared decay) by performing second-order Taylor expansions of two standard pair potentials around their equilibrium bond length r. Throughout we write the equilibrium distance between atoms = l!y: y2ll and the difference between the equilibrium distance and the prediction òr f
25

---

Morse potential
The Morse potential (typically used to describe bonded pairs) is given
EMorse (f, D exp(-a(f- (r)])
(A.2)
Expanding the exponential for small ör gives exp(--ar) or 2r +0 The quadratic approximation is
Eooss (f, = Da28r2 = (MD) k8)
(55)
with
k(M: 2Da
(A.4)
Because D and a are fixed per bond type, (M)) is constant in r.
Lennard-Jones potential The 12-6 Lennard-Jone potential (commonly used for non- non-bonding interactions) can be written terms of its equilibrium separation 21/0
ru( - ( 12 -
(A.5)
Expanding for small ör and keeping terms up to second order,
Ey, r) = - k120 + O
(56)
0 EL
72 22
(
(A.6)
Op2
Fmr
Because e is fixed for a given atom-pair type, the resulting spring constant
(L)) x r-
decays with the inverse square of the equilibrium distance.
A.4 Invariance and Invariant Loss Functions
Physical systems frequently have inherent symmetry, and machine learning models built for such systems benefit from handling these symmetries. Two common symmetries are the E(3) symmetry of Cartesian coordinates, and the S(n) symmetry of exchangeable objects (such as atoms in atomistic systems). In the context of generative models, we note three approaches to designing models that respect the symmetry of the underlying distribution:
Invariant Distribution If we decide that the probability distribution we want to generate is invariant symmetry transformations, a common strategy is first sample from distribution that invariant to the group of interest, and then applying function that is equivariant to the group [Köhler et al.) 2020). For example, [Hoogeooom et al. [0022] samples from a Gaussian distribution that invariant to rotations in 5O(3), and uses an (qu)-equvva neural network to parametrize denoising diffusion model. The probability distribution of the resulting structures generated by their model is therefore invariant to rotations.
Alignment In many situations, we want predictions from a neural network to match ground truth data, up to some symmetry transformation. This is often the case in generative models for proteins and molecules, where there is symmetry to SE(3) and S(n). Recent works such as Hassan et al. [2024),
Klein et al. 2023 handle this in flow-matching context by performing an alignment procedure between samples from prior and data samples, finding an element of S(n) and SO(3) that minimizes the distance between them. While costly, this optimal alignment can be found by using the Hungarian algorithm (Kurtzberg 1962] and the Kabsch algorithm [Kabsch] 1976]. Similarly, (Sarenn et al. (2025) use an equivariant network to learn alignments that brings data samples into canonical representatives [Kaba et al. 2023), J, which are then fed into a generative model. Zhhang et al. (0244) use similar symmetrization method to achieve equivariance. This technique obviates the need to use an expensive cquivariant network for the generative model.
26

---

Invariance-based loss Another method to handle symmetry is to use loss based only on invariant features. Works such as Xu al.) (02T], [Simm and Hernandez
[2019] and Nesterov al. [2020] learn to model interatomic distance matrices rather than atomic coordinates, and convert from distance matrices into coordinates as a post-processing step. Because distances are invariant under E(3) transformations, the resulting method is invariant. Our proposed loss function Equation (5) falls into this last category.
A.5 More on Rigidity Theory
Background on theory We begin by describing the setting of rigidity theory and the necessary properties for our energy loss scale linearly the number of vertices.
Rigidity theory defines a framework as a graph = (V,E) and a map é V + Rd which can be interpreted as the physical coordinates of a given vertex. We call framework globally rigid every ø' V - Ra that yields the same distances between adjacent vertices is obtained from by an isometry. A framework is rigid if there are non-trivial continuous motions of vertices starting from & that preserves the distance between adjacent vertices. Trivial motions, in this case, correspond to group elements of E(d). The central problem of rigidity theory is determine under which conditions different families of frameworks are rigid or globally rigid [Peled 2024) and Duxbury [999].
(hhore
In our application, suffices to neglect modeling some interactions between atoms (terms in the sum in Equation 77. long the minimum loss configuration is unique and thus corresponds to the data. In other words, we require the edges describing this sum are globally rigid.
Due to the simplicity of their construction, we mention few recent results on the rigidity of random graphs. The first result is that a random k-regular graph is rigid with high probability (w.h.p.) D dimensions for A: 2 D It has been conjectured that k > 2D should be enough for rigidity, but the existing proof is limited to D = [Krivelevich et al. 2023 [Peled 2024]. Alternatively, using Erdos-Renyi random graphs, in D dimensions, one could keep adding edges at random until the minimum node degree becomes D, at which point the graph becomes rigid w.h.p. [Low et al., 2003) Note that promising construction for rigid graphs was also recently proposed in the context of machine learning Wang et al., 2025).
2D-regular graphs are globally rigid w.h.p. To construct a sparse version of the energy loss, we require that the edges over which the pairwise distances differences are summed in Equation 7 make random 2D Dereuula graphs in the sparse version of
up globally rigid graph. For simplicity, the energy loss since the number of edges seales linearly in N.
We empirically confirm the conjecture for D and N in ranges relevant for our setting. We construct 1000 random 6-regular graphs and check the fraction of the graphs that are globally rigid using rigidity checking code from Dewar [2025). This code implements rank check of the rigidity matrix and random stress test. The key result is depicted in F5
These graphs are used in the sparse energy loss for shape prediction. For molecules, we use a symmetrized version of these graphs according molecule symmetries. In both settings, 100 random graphs are pre- generated for each number of vertices and random graph from this pool is chosen when computing the loss.
Symmetrization procedure for random graphs We notice that Proposition 4.2 does not necessarily hold for Equation (7) when the edges are not complete. This means atoms symmetric under Aut (Ay) will not necessarily obtain the same gradients. To remedy this, we introduce a symmetrization procedure for given k-regular graph and molecule.
We can select random edges for a representative node in each orbit and symmetrize the adjacency matrix over the orbits via
A
Σ
A-1-
(57)
aym
gEAut(Ag)
Notice that the total number of edges we get by this procedure is
27

---

Fraction of global rigid graphs vs k
1.0
n= ==
n- 52, 3 0.0-
Fraction of global rigid graphs
n- 128,
d 3
0.6
0.4
0.2
0.0
Figure 5: Global rigidity testing of random k-regular graphs. Here, n denotes the number of vertices and d the dimension.
|syml = O
Σ NEEE
|uG||bG)
(58)
which can be quadratic in [V| if there are large orbits (ie. size O((VI)) directly connected by edges but is linear otherwise. It may be possible to ensure this is linear w.h.p. by choosing the k edges in a way that depends on G but we leave this for future work. We empirically verify these edges are sparse in our setting.
A.6 More on Spin Energy
For Ising configuration y (-1, the energy change caused by flipping spin y is AE, 2y, (y)) with h ((((y)= Σj Jıyy,- Setting ho = therefore makes
ErF(ý,y)
Σ N (y)
SEA
proportional this exact spin-flip energy around the ground state. The local energy propose is therefore sensible linear approximation to the true energy.
To obtain convex loss, make the local field weighting is always positive, For this, add
global offset fn in (y) = E, + ho)y,, Since each site has at most four neighbours with
Jigl < 1, we have -4 < htme < 4; choosing single ho > ensures weight (Jij + ho) >> 0 and penalises energetically costly
more heavily.
A.7 Extension to Flow Matching
The energy loss can be extended to Gaussian flow matching [Lipman et al. 2023). We show hereafter the correspondence for the conditional vector field. The noisy sample is given by the interpolation
x
(1 t)x +
The flow matching objective aims at regressing the vector field:
Given a vector field prediction, the corresponding sample prediction is
X; -taa(xe))
28

---

The correspondence between MSE on the vector field prediction and on sample prediction is therefore:
nn 1|2
t2 x
x 112
Therefore, the associated energy objective obtained by replacement of the regression MSE:
E(x x8).
Our theoretical results relating to score estimation properties also transfer to flow matching. This is because Gaussian flow matching also implicitly provides method for score estimation similar to diffusion models. Given the optimal vector field, u* (xx). the score is given by:
Vx log = (1 1)u* (x)) + Xt
B Related works
Several different lines of research have also incorporated concept of energy into machine learning framework. this section, we distinguish our framework from distinct but related areas of research.
Energy-based models Traditional energy-based models approach learning shaping energy landscape, where observed configurations correspond to low-energy states (LeCun et al. [2066]. Deep counterparts and their connection to discriminative training have also been extensively explored in many recent works (e.g.. Du and Mordatch [2019), Grathwohl et al. 2019)). A key distinction of the existing literature on energy-based models and our energy loss approach is that, because they minimize the forward KI. (max-likelihooe or alternatives such as contrastive and large-margin losses), they need deal with the minimization of the partition function or its surrogates i.c., the energy of arbitrary points in the domain must be high. In contrast, our treatment remains close to supervised learning losses and avoids the partition function altogether.
Physics-informed neural networks (PINNs) Raissi et al., 2019 has proposed PINNs way to learn PDEs by penalizing residuals directly the loss, enforcing solutions consistent with physical constraints. They have recently been used in the context of diffusion models [Bastek et al. 2024 In contrast to our approach, these models do not rely on training data and models are instead learned to satisfy known differential equations randomly generated points from the domain. Another family of inysice-informed losses appears in Hamiltonian Neural Networks (HNN) [Greydanus al.) 2019 and Lagrangian Neural Networks (LNN) CCrnmmer et al., 2020
Energy Sampling and Boltzmann Generators A separate line of research incorporating energies and generative modeling has been in Boltzmann Generators (Noé al.) [2019 Köhler al. 2020 Klein and Noč] 20255. These models are designed to sample physical configurations according to Boltzmann distribution stemming from a known energy function. While our framework is also based on an assumption of data belonging to Boltzmann distribution, ours is instead simply an approximation of the local landscape around each data point and does not assume the existence of callable energy function.
C Experimental Details and Additional Results
C.1 Regular shape generation
Experimental details and hyperparameters We use a 2 hidden-layer MLP with hidden dimension 64 for this task. We conducted a sweep over hidden dimension and find behaviour is relatively consistent. Models are trained in parallel on an Nvidia Quadro RTX 8000 using the Adam optimizer. The dataset size 100K randomly generated samples and we train all models 50 epochs. sweep over dataset size showed fairly consistent results. We conduct thorough sweeps for learning rate for each loss, shape degree and augmentation angle. For each setting, the model giving the highest quality is chosen.
29

---

Table 6: Best hyperparameters for GDM.
Loss
Coefficient
Learning Rate
Positional Loss Weight
Energy
Constant Inv. Dist.
9e-4
0.05
7e-4
0.05
Inv. Sq. Dist.
6e-4
0.05
Exp. Dist.
4e-4
0.05
MSE
le-3
1.5
MAE
le-3
0.8
Kabsch Align
8e-4
0.8
Table 7: Best hyperparameters for EDM.
Loss
Coefficient
Learning Rate
Positional Loss Weight
Energy
Constant Inv. Dist.
le-4
0.05
le-4
0.05
Inv. Sq. Dist.
lc-4
0.1
Exp. Dist.
le-4
0.05
MSE
3e-4
1.0
MAE
3e-4
1.0
Kabsch Align
3e-4
0.8
Evaluation To evaluate shape quality, we introduce metric that captures how regular the angular differences and radial distances are across the shape. In a well-formed, regular shape, we expect both the variation angular differences (a (A and the variation in radial distances (Gradius) to
be small. In particular, we choose Quality == = In ('uungle + Sradum)) This is, of course, design 2n
choice used to map a shape to single number. We record both Asunttc and Tradius ensure both terms are well-represented in the quality and find that visually this metric is a good reflection of the visual regularity of generated shape. For reference, above quality of 5-6, shapes look nearly visually perfect, as in Figure 3.. For quality below this, they become slightly irregular and below 2 they look very disordered in Figure 3p.
C.2
Molecule generation
C.2.1
QM9 Dataset
Experimental details and hyperparameters On QM9, we match the setup in Hoogeboom et al. (022) closely as possible. We train GDMs and EDMs with with layers and 256 node features on 100k samples from the dataset. The diffusion process has 1000 diffusion steps with polynomial noise schedule and precision 1 10- 5 An L2 denoising loss is used with mini-batch size 512 on GDM and 400 on EDM. We use the Adam optimizer. An EMA decay of 0.9999 is used. Runs were conducted on single 48G GPUs mainly on the Nvidia Quadro RTX 8000, A6000 and L40S. A full run of 3000 epochs takes 2-4 days on a single GPU.
We conduct extensive sweeps for learning rate and positional loss weight for all losses. We tune the positional loss weight to ensure there is balance between loss on positions and atom-type for all losses. Learning rates were searched for in broadly in the range [1e-5, le-2] before narrowing the range to [2e-3, 4e-4J. For the positional loss weight, we choose values in [0.05, 0.1, 0.5. 0.8, 1.0, 1.5). We find final performance are summarized in Table 6 and Table 7
very sensitive to the positional loss weight. Tuned hyperparameters
Additional results Here, we include results for all settings for GDM, GDM-aug and EDM. The results follow in Table Results are averaged across seeds.
30

---

Table Complete results on QM9.
Loss
Mol. stab. (%)
Atom stab. (%)
Valid. (%)
Unique (%)
GDM
MSE
81.7 + 3.3
98.3 + 0.3
93.3 = 1.7
99.98 + 0.04
MAE
76.3 2.0
97.7 0.3
91.1 1.2
99.96 0.05
Kabsch Align
81.7 + 2.2
98.4
93.1 = 1.2
99.93 0.13
Energy (Sparse)
86.1 2.3
99.0
96.2 1.4
100.0
Energy
86.2 + 2.1
98.9 + 0.2
96.6 1.3
100.0 + 0.0
GDM-aug
MSE
83.7 +
98.3 0.004
93.6 1.7
100.0
MAE
76.4 + 0.9
08.1
92.6 1.2
99.99 0.02
Kabsch Align
82.3
97.8 0.004
90.8
100.0
Energy (Sparse)
89.1 + 0.9
99.0 0.1
97.4 2.5
100.0 + 0.0
Energy
89.8 2.8
99.3 + 0.3
97.7 1.4
99.99 0.002
EDM
MSE
82.4 + 3.4
98.8 + 1.7
93.0 = 2.5
99.89 0.32
MAE
74.8 1.7
97.8 0.3
88.6
99.96 0.07
Kabsch Align
80.6 + 3.0
98.3 - 3.0
92.5 = 3.0
99.91 = 0.07
Energy
86.6 1.6
99.0 + 0.20
96.8
99.96 0.06
Table 9: Complete results for GDM and GDM-aug on GEOM-Drup
Loss
Mol. stab. (%)
Atom stab. (%)
Valid. (%)
Unique (%)
GDM
MSE
0.3
84.7
93.8
100
Energy
21.1
95.8
89.6
100
GDM-aug
MSE
0.8
85.6
94.8
100
Energy
24.6
96.0
89.7
100
C.2.2 GEOM-Drugs Dataset
Experimental details and hyperparameters On GEOM-Drugs, we use similar setting to QM9 but now train models with 4 layers and 256 node features, following Hoogeboom Hollowoogoooog et al. [2022]. We train the model for 13 epochs. Training is distributed across 80G Nvidia 1000 GPUs and single run takes roughly 2.5 days.
use a batch size 128 with the Adam optimizer.
We start with optimal learning rate and positional loss weight from QM9 and do a sweep over learning rates [5e-4, le-3, 2e-3] for MSE and [le-4, 4e-4, le-3] for energy loss. The hyperparameters in Table gave the best results. We use exponential coefficients for the energy loss.
Additional results aug in Table 9
We additionally report the performance of MSE and energy losses with GDM-
C.3 Spin ground state prediction
Experimental details hyperparameters The CNN we use for the spin prediction task a layer ResNet type architecture with 256 hidden layer size. All networks are trained with learning rate x 10-3 until convergence. We use the Adam optimizer with batch size 256. Temperature set to T 0.1 for the local energy loss. Training takes around 5 hours on Nvidia V100 GPUs.
31

---

C.4
More on sparse energy loss
C.4.1
Timing
Our objective in including the sparse energy loss to demonstrate our method can efficiently generalize to systems with many particles where loss calculation may contribute significantly to running time (e.g. very large point clouds). This not the case for molecules, where the neural network (GNN or Transformer) is typically fully connected and thus scales as N As Table 10 shows, the most expensive loss calculation is less than 1% the total backward and forward time.
Table 10: Wall-times for loss computation on QM9 on an NVIDIA L40S.
Component
Loss Type
Time (ms)
Loss computation
MSE
0.18 0.01
Energy
0.51 + 0.02
Sparse Energy
0.57 + 0.02
Kabsch Align
1.14
Forward pass
74 16
Backward pass Optimizer step
94 3
1,43 0.01
To better understand the scale which this becomes relevant consideration and the utility of the sparse energy loss, see the following wall clock times from the shape generation setting in Table 11
Table 11: Runtime for different loss functions as the number of nodes increases.
Nodes
Energy (ms)
Sparse Energy (ms)
MSE (ms)
Kabsch Align (ms)
30
0.240 + 0.0021
0,255 + 0.0012
0.058 + 0.0004
0.807 + 0.0027
300
0.245 + 0.0011
0.257 + 0.0024
0.059 + 0.0004
0.804 + 0.0048
3000
20.120 + 0.0055
0.275 + 0.0018
0.0779 + 0.0008
0.944 + 0,0028
30000
0.293 = 0.0029
0.0740 + 0.0011
3.242 = 0.1452
300000
2.652 + 0.0050
0.131 + 0.0004
24.381 + 0.0272
At 30000+ nodes, the energy loss requires too much memory to compute. Importantly, the sparse energy is cheaper than the Kabsch Align by factor 5-10x. Note that all losses have some constant cost that does not scale with N contributing to the wall-clock time. This explains why for QM9 (avg. 28 atoms) energy loss is marginally faster than sparse energy and why in the scaling table wall-clock times start to increase with N only after a certain point.
Interestingly, using an equivariant network with EDM takes 129.71 + 0.044 ms for the forward pass and 180.07 0.070 ms for the backward pass. Using the energy loss imparts 0.3% increase one backward pass through the model, while using an equivariant architecture imparts a 94% increase, while providing inferior benefits. The energy loss with non-equivariant architecture results in more improvement than using an equivariant architecture, at negligible computational cost, which we think a significant finding.
C.5 Sparse energy loss on larger systems
Table 12:
Sparse energy loss with GDM on GEOM- GEOM-Drugs.
Loss
Mol. stab. (%)
Atom stab. (%)
Valid. (%)
Unique (%)
MSE
0.3
84.7
93.8
100
Energy
21.1
95.8
89.6
100
Sparse Energy (Inv. Dist)
7.4
91.9
92.6
100
We include the sparse energy results on Geom -rugs in Table 12 distance decay in the coefficient worked better when the edges are sparse and random. These results highlight potential compute-performance tradeoff for this version of the loss on larger graphs. This result can likely be improved by being more intentional in the selection of sparse edges. We leave this to future work.
We found using more gradual
32


Reviews:
[{"review_id": "O1Vb7eURLz", "reviewer": "Reviewer_f4bE", "summary": "This work establishes a new method for physics-guided Machine Learning, where embedding physical principles into loss functions significantly improves performance in molecular generation and spin prediction. Its core contribution is shifting loss design from statistical error minimization to physics-driven optimization, ensuring the model contains physical constrains such as symmetry during predictions. The new model optimizes prediction errors, efficiency and reliability.", "questions": "1. Energy losses assume thermal equilibrium. How can the framework adapt to far-from-equilibrium systems (e.g., chemical reaction pathways)? Should dynamic energy landscapes be incorporated?\n2. The main focus is loss-level innovation. Could coupling with equivariant architectures enhance performance? For example: losses enforcing energy constraints + architectures handling symmetries.\n3. Will energy coefficients, the choice of potential energy equations and noise change under different temperature?\n4. How to modulate the energy loss functions when it comes to the physical systems that don’t obey Boltzmann distribution?", "limitations": "1. Atomic system coefficients (exponential decay) rely on experience without much theoretical guidance. A complex system (e.g., proteins) may need much more complicated energy approximations at prohibitive computational cost.\n2. Theoretical guarantees for score estimation valid only at low noise, high noise requires unaddressed corrections.\n3. Sparse rigid graphs reduce complexity but lack reported speedups, Fig. 3c only qualitatively claims O(N) scaling.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n\nThis paper introduces energy loss functions, a novel framework that integrates thermal equilibrium properties of physical systems into Machine Learning loss design. Unlike traditional losses such as MSE, which optimize statistical prediction errors, this work leverages the Boltzmann distribution and reverse KL divergence to derive physics-based loss functions. This ensures model predictions adhere to physical constraints such as symmetry and energy minimization.\n\nOriginality: First systematic integration of physics into loss functions, beyond Machine Learning model architectural stage.\nQuality: A comprehensive theory with formal proofs in Appendix, extensive experiments (molecule generation and spin prediction) and code release commitment.\nClarity: Showing comparison between MSE and energy loss functions with clear physics figures such as Fig. 1. Comprehensive mathematical derivations in Appendix.\n\nWeaknesses:\n1. The selection of energy coefficients lacks discussion. The choice of Morse and L-J potential seems arbitrary, adding a few sentences to prove the completeness of the choice may help. In addition, it may help if the authors give a few examples of energy coefficients in different physical systems, such as protein and crystalline systems.\n2. Lack of discussions of how to build the sparse version of the energy loss. The main text mentions the sparse model reduce computation by O(N) operations, but without clear details.\n3. Temperature remains unexplored. Temperature is treated as a hyperparameter rather than a physical variable. The authors may need to add a few sentences to discuss what will change in energy loss building under different temperature.\n4. Incomplete comparison with equivariant architectures. May need to add studies to demonstrate the complementary of benefits of energy loss functions.", "was_revised": true, "final_justification": "The authors have provided a thorough and well-reasoned rebuttal addressing all points raised in the review. They have clearly justified their choice of potentials, clarified the general applicability of their method via Taylor approximation, and convincingly demonstrated the practical benefits of the sparse energy loss, particularly in large-scale settings, through detailed wall-clock timing analysis. The comparison between equivariant architectures and the energy loss highlights a key contribution, which shows that energy-based regularization can yield substantial gains even without architectural symmetry, at minimal computational overhead. The discussion on temperature scaling, the complementarity of invariance and energy minimization, and the extension to state-of-the-art models further strengthen the paper’s technical depth and relevance.", "rebuttal": "We thank the reviewer for their valuable feedback. We are glad that they have appreciated the originality of our work, the comprehensive theory and the extensive experiments on different domains. \n\nSome of the raised questions are answered in the Appendix of the paper. We will make use of the added camera-ready page (if accepted) to move some material to the main body.\n\nWe address their concerns and questions in details below:\n\n### Selection of the energy coefficients\nThe choice of focusing on the Morse and Lennard-Jones potentials is justified by the fact that they are quite general and are the most widely used pairwise potentials. Lennard-Jones is the archetypal example for simple but realistic inter-molecular interactions across physical systems while the Morse potential is a model for bonding in a diatomic molecule. There is a very large number of potentials with explicit parametric forms that are used by practitioners, so we could not go into details of all of them. We will however include references in the paper to the relevant literature (such as [1,2]) to guide the reader towards alternatives.\n\nAn important feature of our method is that it provides a clear method to find the coefficients given any choice of potential, by second-order Taylor approximation. The Morse and LJ potentials offer important examples (see derivations in Appendix A.3 of the paper). \nWe will add a few sentences to clarify the general procedure in the manuscript following your suggestion.\n\n### Sparse loss function\nThe sparse version of the loss function as well as the algorithm to implement it is discussed in Appendix A.5. It is based on random regular graphs (which can optionally be symmetrized to improve the invariance of the loss as we discuss). We will make use of the additional camera-ready page to move some explanations to the main text. \n\nWe have also included wall-clock times benchmarking the losses on a single NVIDIA L40s following your suggestion. Our main objective in including the sparse energy loss was to demonstrate our method can efficiently generalize to systems with many particles where loss calculation may contribute significantly to running time (e.g. very large point clouds). This is not the case for molecules, where the neural network (GNN or Transformer) is typically fully connected and thus scales as N^2. In the results below on QM9 with GDM, the most expensive loss calculation is less than 1% of the total backward and forward time.\n\n| Component            | Loss Type     | Time (ms) (ms) \n| -------------------- | ------------- | ------------------- |\n| **Loss computation** | MSE           | 0.1840 ± 0.0084     |\n|                      | ENERGY        | 0.5100 ± 0.0161     |\n|                      | ENERGY-SPARSE | 0.5712 ± 0.0242     |\n|                      | Kabsch Align          | 1.1361 ± 0.0300     |\n| **Forward pass**     | –             | 73.55 ± 16.38       |\n| **Backward pass**    | –             | 94.47 ± 3.37        |\n| **Optimizer step**   | –             | 1.43 ± 0.0072       |\n\n\nTo better understand the scale at which this becomes a relevant consideration and the utility of the sparse energy loss, see the following wall clock times from the shape generation setting:\n\n| # Nodes | Energy (ms)     | Sparse Energy (ms) | MSE (ms)        | Kabsch Align (ms) |\n| ------- | --------------- | ------------------ | --------------- | ----------------- |\n| 30      | 0.240 ± 0.0021  | 0.255 ± 0.0012     | 0.058 ± 0.0004  | 0.807 ± 0.0027    |\n| 300     | 0.245 ± 0.0011  | 0.257 ± 0.0024     | 0.059 ± 0.0004  | 0.804 ± 0.0048    |\n| 3000    | 20.120 ± 0.0055 | 0.275 ± 0.0018     | 0.0779 ± 0.0008 | 0.944 ± 0.0028    |\n| 30000   | -               | 0.293 ± 0.0029     | 0.0740 ± 0.0011 | 3.242 ± 0.1452    |\n| 300000  | -               | 2.652 ± 0.0050     | 0.131 ± 0.0004  | 24.381 ± 0.0272   |\n\nAt 30000+ nodes, the energy loss requires too much memory to compute. Importantly, the sparse energy is cheaper than the Kabsch Align by a factor 5-10x. Note that all losses have some constant cost that does not scale with N contributing to the wall-clock time. This explains why for QM9 (avg. 28 atoms) energy loss is marginally faster than sparse energy and why in the scaling table wall-clock times start to increase with N only after a certain point.\n\nInterestingly, using an equivariant network with EDM takes 129.71 ± 0.044 ms for the forward pass and 180.07 ± 0.070 ms for the backward pass. Using the energy loss imparts a 0.3% increase on one backward pass through the model, while using an equivariant architecture imparts a 94% increase, while providing inferior benefits. \n\nA crucial conclusion is that our results show that the energy loss with a non-equivariant architecture results in more improvement than using an equivariant architecture, at negligible computational cost, which we think is a significant finding. We will discuss it further in the main text.\n\n### Effect of temperature\nThe temperature parameter simply scales the loss (see Equation 5). Optimizers like Adam are invariant to loss scaling. For optimizers that are not invariant (e.g. SGD), the temperature is absorbed in the learning rate.\n\n### Coupling with equivariant architectures\nWe have indeed seen that coupling the energy loss with equivariant architectures leads to improved performance. Our results in Appendix C.2.1 of the paper show that the energy loss results in significant improvements even for equivariant architectures. \nWe argue that the benefits of invariance of the loss function are complementary to equivariance of the architecture. Essentially, equivariance guarantees that for a given input output pair, the output associated with a transformed input will be the transformed output. However, it does not guarantee that the output will be the correct one, it ensures generalization to samples related by symmetry. The energy loss helps in learning the correct output, by appropriately weighting the errors and gradients during training. In the case of symmetries, instead of being forced to regress the original configuration, the network can learn to regress any configuration related by symmetry, which equivariance does not allow. This results in a reduction of variance in the gradients as we show in Proposition 4.5.\n\nWe obtained new results that confirm the increase in performance on a more modern equivariant architecture (see below).\nThe improvement offered by the energy loss is therefore orthogonal to architectural improvements and offers another avenue for development. \n\nWe also highlight that our results show that the energy loss with a non-equivariant architecture results in more improvement than using an equivariant architecture, at a much cheaper cost, which we think is a significant finding. We will discuss it further in the main text.\n\n### Extension to non-equilibrium and non-Boltzmann systems\nAs said in the paper, the energy loss is not as well motivated for systems that are not at thermal equilibrium. Our method still applies to a wide range of problems, but we have been careful in limiting its scope.\nHowever, a crucial motivation for our work is that the MSE loss is already an approximation to an energy function, simply a very crude and unphysical one. Compared to it, our approximations are much better motivated. Using more elaborate approximations is an exciting avenue for future research, but introduces a tradeoff between complexity and fidelity. What we have introduced in this work sits at an advantageous point with respect to this tradeoff.\n\n### Additional results\nFinally, following the suggestion from reviewer uwRJ, we have **significantly improved** our experimental evaluation by obtaining results on a more recent, state-of-the-art equivariant architecture for molecular generation (JODO [4]). Replacing the loss function used in this model by the energy loss function resulted in an improvement across nearly all 3D metrics. Our method therefore provides complementary gains to state-of-the-art architectures.\n\nWe hope this response addresses your concerns. If any further clarification would help you reconsider your overall assessment, we would be happy to provide it.\n\n[1] Stone, The theory of intermolecular forces, 2013\n\n[2] Frenkel and Smit, Understanding molecular simulation, 2001\n\n[3] Huang et al. Learning joint 2D and 3D diffusion models for complete molecule generation, 2023"}, {"review_id": "p1ee0yBiCU", "reviewer": "Reviewer_J6c1", "summary": "The paper introduces novel energy based loss functions for physical systems, like molecules and lattice systems. Although the authors mention that the energy of the target system could be utilized directly if available, they propose a different strategy and explain why using the true energy function can be infeasible. For many particle systems like molecules, they propose a loss based on distances between particles. In this case, the energy is given by the difference between the distance matrix of the prediction and the target. They use this to derive an alternative loss for diffusion models for many particle systems evaluate on small molecules, showing improvements over traditional diffusion models. \nMoreover, they also investigate a spin lattice system, where they use a different energy function based on the difference of the predicted spin and the target spin and its surroundings.", "questions": "- Suggestions:\n    - Many flow modes, especially Boltzmann Generators, use the target energy function during training. Although it is different to what is discussed in this paper, I think it should be mentioned in the related work sections. And potentially also the abstract and introduction should be changed accordingly, as the phrase \"previous approaches mostly focused on incorporating physical insights at the architectural level. In this paper, we propose a framework to leverage physical information directly into the loss function ...\" is not correct.\n   - 44: Concurrent work of Klein et al., 2023 should be cited as they introduce the same training objective, i.e. \"Equivariant Flow Matching with Hybrid Probability Transport\" by Yuxuan Song et al.\n\n\n- Questions:\n  - Could the energy diffusion loss also be directly applied to flow matching? \n  - Why is the number of valid molecules in table 1 lower for the energy based loss? This seems counterintuitive, but maybe there is a good explanation. Moreover, are there errors for the results reported in this table?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- The paper is well written and clearly motivated.\n- The idea of using a surrogate energy function in the loss in the proposed way is to the best of my knowledge novel. However, there are some related methods that utilize the actual target energy in the loss, which should be discussed in the paper (see suggestions)\n- They evaluate their method on very different systems, i.e., molecules and a lattice spin system, showing in both cases the advantages of their proposed loss. \n- As there are many applications of diffusion models for molecular systems, the proposed energy based loss instead of the usual diffusion loss might be useful for many applications.\n\nWeaknesses:\n- See suggestions and questions", "was_revised": true, "final_justification": "The authors were able to address all my concerns and provided insightful answer to my questions.", "rebuttal": "We thank the reviewer for their review. We appreciate that they have found the paper well-written, the idea novel, and the evaluation comprehensive.\n\nWe answer their questions hereafter:\n\n### Discussion of Boltzmann generators\nDiscussion of Boltzmann generators is included in the related works (in Appendix B). In sampling the target energy is queried and therefore has a direct role. In generative modelling and regression however, the true energy is generally not available (see section 3 of the manuscript for discussion.) Our work in considering energy loss is novel in this pure generative modelling context, for which only data is available. We will nuance the statements referred by the reviewer to address their concern, and will be happy to move some discussion on the Boltzmann generators to the main text.\n\n### Related papers\nThank you for the suggestion, we will cite Yuxuan Song et al. as another example of an invariant loss achieved using alignment.\n\n### Extension to flow matching\nThe energy loss can indeed be applied to Gaussian flow matching [1, 2]. We show hereafter the correspondence for the conditional vector field [1].\nThe noisy sample is given by the interpolation\n$$\nx_t = (1-t) x + t \\epsilon\n$$\n\nThe flow matching objective aims at regressing the vector field:\n$$\nu = \\frac{x_t - x}{t}\n$$\n\nGiven a vector field prediction, the corresponding sample prediction is\n$$\nx = x_t - t u_\\theta(x_t)\n$$\n\nThe correspondence between MSE on the vector field prediction and on sample prediction is therefore:\n$$\n\\lVert u - u_\\theta \\rVert^2 = \\frac{1}{t^2} \\lVert x - x _\\theta \\rVert^2\n$$\n\nTherefore the associated energy objective is obtained by replacement of the regression MSE:\n\n$$\n\\frac{1}{t^2} E (x - x _\\theta)\n$$\n\nOur theoretical results relating to score estimation properties also transfer to flow matching. This is because Gaussian flow matching also implicitly provides a method for score estimation similar to diffusion models [3]. Given the optimal vector field, $u^*(x_t)$ , the score is given by\n\n$$\n\\nabla_{x_t} \\log p(x_t) = -\\frac{(1-t) u^*(x_t) + x_t}{t}.\n$$\n\nWe will discuss the extension to flow matching in the revised paper.\n\n### Validity metric\nThe reasons for this decrease is not fully clear, because we only observe a decrease for the GEOM-Drugs dataset and it occurs despite a significant increase in stability, which we attribute to the physical inductive bias of our method. Validity is measured by first converting the 3D molecules into SMILES strings and then using RDKit sanitization, which is more of an opaque metric than stability. Due to the significant expense of training on GEOM-Drugs, we did not obtain error bars, following [4,5].\n\n### Additional results\nFinally, following the suggestion from reviewer uwRJ, we have **significantly improved** our experimental evaluation by obtaining results on a more recent, state-of-the-art equivariant architecture for molecular generation (JODO [4]). Replacing the loss function used in this model by the energy loss function resulted in an improvement across nearly all 3D metrics. Our method therefore provides complementary gains to state-of-the-art architectures.\n\nPlease let us know if you have any additional concerns or questions we can answer.\n\n[1] Lipman et al. Flow matching for generative modelling, 2023\n\n[2] Liu et al. Flow straight and fast: learning to generate and transfer data with rectified flow, 2023\n\n[3] Albergo et al. Stochastic interpolants: A unifying framework for flows and diffusions, 2023\n\n[4] Huang et al. Learning joint 2D and 3D diffusion models for complete molecule generation, 2023\n\n[5] Hoogeboom et al. Equivariant diffusion for molecule generation in 3D, 2022"}, {"review_id": "L71VLMO5gH", "reviewer": "Reviewer_3SJQ", "summary": "This paper proposes a method to incorporate knowledge of physical systems into the loss function, whereas previous approaches have primarily focused on model architecture. This method is realized through a reverse KL-type loss function, which allows for an analytical expression of the loss. The authors demonstrate that this approach can be applied to diffusion models and enables the incorporation of invariances, with theoretical justification. Finally, they show the effectiveness of the method through several experiments, including molecule generation and spin ground-state prediction.", "questions": "1. This energy-based physical knowledge can also be incorporated into model architectures. Could you discuss the strengths and weaknesses of incorporating it into the loss function, compared to incorporating it into the architecture?\n\n2. In Section 6.2, the results show that the energy-based loss yields stability. Could you provide a more detailed explanation of why the proposed approach leads to such stability?\n\n3. Can this approach be applied to other generative models, such as flow matching? Additionally, do the theoretical results discussed in Section 4.2 scale to those cases?", "limitations": "Yes.", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths\n\n・Overall, the paper is clearly written and easy to follow.\n\n・The idea of circumventing the intractability of the normalizing term by using the reverse KL divergence is simple yet interesting.\n\n・The proposed approach is broadly applicable to various machine learning problems, including generative tasks with diffusion models.\n\nWeaknesses\n\n・The difference between incorporating physical knowledge into model architectures and into loss functions should be more clearly articulated. At present, it remains unclear how important it is to include physical information in the loss.", "was_revised": true, "final_justification": "I have understood the benefit of incorporating constraints into the losses. While I acknowledge that this research is significant and the proposed method is convincing, the technical novelty and theoretical soundness are relatively incremental, which prevents me from assigning a higher score.", "rebuttal": "We thank the reviewer for their review. We are happy that they have found the paper clearly written, the method interesting and broadly applicable and that they only note one weakness.\n\nWe address it and their questions below:\n\n### Difference between incorporating physical biases in architecture or in loss function\nThis is a good question. One type of physically motivated inductive bias that can be incorporated in the architecture is equivariance (if the reviewer has something else in mind, please let us know). We argue that the benefits of invariance of the loss function are complementary to equivariance of the architecture. Essentially, equivariance guarantees that for a given input output pair, the output associated with a transformed input will be the transformed output. However, it does not guarantee that the output will be the correct one, it ensures generalization to samples related by symmetry. The energy loss helps in learning the correct output, by appropriately weighting the errors and gradients during training. In the case of symmetries, instead of being forced to regress the original configuration, the network can learn to regress any configuration related by symmetry, which equivariance does not allow. This results in a reduction of variance in the gradients as we show in Proposition 4.5. Thus, using the energy loss results in significant improvements even for equivariant architectures (see results in Appendix C.2.1, Table 7 and new results on more modern equivariant architecture JODO in the reply to reviewer uwRJ). \n\nA crucial point is that our results show that the energy loss with a non-equivariant architecture results in more improvement than using an equivariant architecture, at a much cheaper cost, which we think is a significant finding. We will discuss it further in the main text.\n\nThere is therefore an important difference between loss level and architecture level improvement. We will improve the discussion of this in the paper, following your feedback.\n\n### Improvement in stability\nThe stability metric is determined by valency checks on all of the atoms within a molecule. This valency is determined by the bonds drawn between atoms, and for the task of 3D molecule generation this is determined by a lookup table of what bond orders correspond to the interatomic distances predicted by the model. Since our loss encourages the model to generate molecules with accurate interatomic distances, this results in a greatly improved stability. This result adds evidence to the usefulness of having a loss that is more physically motivated than MSE.\n\n### Extension to flow matching\nThe energy loss can indeed be applied to Gaussian flow matching [1, 2]. We show hereafter the correspondence for the conditional vector field of [1].\nThe noisy sample is given by the interpolation\n$$\nx_t = (1-t) x + t \\epsilon\n$$\n\nThe flow matching objective aims at regressing the vector field:\n$$\nu = \\frac{x_t - x}{t}\n$$\n\nGiven a vector field prediction, the corresponding sample prediction is\n$$\nx = x_t - t u(x_t)\n$$\n\nThe correspondence between MSE on the vector field prediction and on sample prediction is therefore:\n$$\n\\lVert u - u_\\theta \\rVert^2 = \\frac{1}{t^2} \\lVert x - x _\\theta \\rVert^2\n$$\n\nTherefore the associated energy objective is obtained by replacement of the regression MSE:\n$$\n\\frac{1}{t^2} E (x - x _\\theta)\n$$\n\nOur theoretical results relating to score estimation properties also transfer to flow matching. This is because Gaussian flow matching also implicitly provides a method for score estimation similar to diffusion models [3]. Given the optimal vector field, $u^*(x_t)$, the score is given by\n\n$$\n\\nabla_{x_t} \\log p(x_t) = -\\frac{(1-t) u^*(x_t) + x_t}{t}.\n$$\n\nWe will discuss the extension to flow matching in the revised paper.\n\n### Additional results\nFinally, following the suggestion from reviewer uwRJ, we have **significantly improved** our experimental evaluation by obtaining results on a more recent, state-of-the-art equivariant architecture for molecular generation (JODO [4]). Replacing the loss function used in this model by the energy loss function resulted in an improvement across nearly all 3D metrics. Our method therefore provides complementary gains to state-of-the-art architectures.\n\n\nWe hope this response addresses your concerns. If any further clarification would help you reconsider your overall assessment, we would be happy to provide it.\n\n[1] Lipman et al. Flow matching for generative modelling, 2023\n\n[2] Liu et al. Flow straight and fast: learning to generate and transfer data with rectified flow, 2023\n\n[3] Albergo et al. Stochastic interpolants: A unifying framework for flows and diffusions, 2023"}, {"review_id": "nu8Z4mFqMo", "reviewer": "Reviewer_uwRJ", "summary": "The authors propose to replace generic standard training objectives like MSE or cross-entropy with energy-based loss functions that embed known physics of a system into the training signal: each datapoint is treated as an equilibrium state and the loss is simply the (approximate) energy difference between the model's prediction and that state, evaluated under a surrogate potential that respects the system's symmetries. The authors construct losses of this type for atomistic coordinates, diffusion-model training, and discrete spin systems and show that they are architecture-agnostic and symmetry-invariant without costly alignments. They demonstrate across shape-generation, molecular-generation, and spin-glass tasks that show that this physically grounded loss yields higher-quality samples and better data efficiency using standard datasets and models.", "questions": "Your derivation relies on the data being close to local minima of the (unknown) true energy. How sensitive is the loss to violations of that assumption? Can you discuss or demonstrate this?\n\nCould you provide numbers on the wall-clock time and efficiency/accuracy tradeoffs of the different loss types on the more realistic benchmarks? The small reduction in accuracy for the sparse objective on QM9 would look stronger if you could demonstrate how much compute is saved\n\nCould you please clarify exactly how molecular stability was calculated and where the numbers place your results compared to prior publications? Preferably you could include numbers from other publications in the tables you already have produced.\n\nIn my opinion, the biggest improvement would be to show that using one of the more modern architectures (e.g. JODO or EQGAT) trained with your energy-based loss similarly leads to better results on the metrics typically reported. This would very likely push results beyond current state-of-the-art if the improvements are comparable. Even one small-scale experiment (same hyper-parameters, swapped-in loss) would show whether the gains you see are orthogonal to architecture advances, and would strengthen the claim that the loss is generally useful.", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n\n- The idea is creative and a refreshing rethink of foundational concepts of ML\n- Through examples and theory, the authors show that their energy-based loss functions mirror real energy landscapes and can respect rigid motions and permutations. This is a strong property, since e.g. default DDPM optimizer gradients can pull predictions away from valid symmetry-equivalent structures.\n- The versatility of the recipe for constructing such loss functions is a strength; it works for several machine learning tasks without requiring architectural changes. \n- Losses lead to better gradients and avoids pathologies of e.g. MSE\n- On the benchmarks shown, the energy-based loss functions lead to tangible improvements. \n\nWeaknesses\n- The authors derive the loss function from a Boltzmann distribution; the derivation only makes sense here if training samples are near local minima of the true energy landscape. The authors note this assumption, but it is unclear to me how crucial it is to the framework. For example, if our data comes from MD simulations and contains e.g., non-equillibrium and transition states, can this result in loss gradients pointing the wrong way in situations where MSE stays agnostic?\n- The driving example of the paper uses a quadratic pair potential which results in a loss that scales with the number of particles N squared. While the authors mitigate this issue by selecting an O(N) subset of spring edges from a globally rigid graph, those missing long-range edges no longer penalise collective shears or electrostatic couplings, so it is unclear whether the sparse loss loses some of its informative gradient signal, even if its optima remain unchanged in theory. Table 2 shows that results are comparable to the non-sparse energy for QM9, but these are all small systems. I wonder what the tradeoff will be for systems that are larger and thus where the sparsification will be most significant for runtime.\n- Generally, I would have appreciated wall-clock speed comparing the default losses vs the dense and sparse training losses, preferably versus molecule size.\n- While the paper's derivations are mostly convincing, I find the experimental section thin where it matters most for adoption. From an application point of view, the molecule generation task with GDM and EDM is the most interesting, and the paper lacks a comparison that clearly indicates whether it can push molecule generation tasks beyond state of the art. For example, it would have been useful to see numbers from other publications reported in the papers' tables, in particular if the same architecture was used. For example, the JODO paper (Huang et al.) reports molecular stability numbers for QM9 of 93.4 % which is higher than the 89.8 reported in the present paper. It is plausible that the authors' energy-based loss objective could outperform if a similar newer architecture was used, but it is unclear whether the loss function improvement is orthogonal to the improvements from newer architectures with more embedded physics.  I also believe this would be very interesting in terms of the molecular stability on the GEOM-Drugs dataset, since this number seems to be either remarkably bad or remarkably good (I believe the latter is the case but see next point). It is, however, simply unclear where these numbers place the authors compared to prior work, and the experiments section thus does not lift this paper from theoretically interesting to broadly applicable, which I very well believe it could be. \n- Some background is missing. It would be useful to understand exactly how the stability was calculated, since the molecular stability number can be calculated in several ways (See e.g. Table 2 in the JODO paper, Huang et al., 2023). It would also be useful to discuss some of the points made in the discussion of stability calculations in GEOM-Drugs presented in Nikitin et al. (GEOM-Drugs Revisited:.. 2025), although this paper is very new and the authors haven't had a chance to include these results/discussions in their paper.  Generally, performing and discussing experiments with newer architectures and clearly reporting numbers that are fully comparable to previous publications could contribute to a better understanding of whether a modern architecture paired with an energy-based loss could achieve state-of-the-art molecular generation. If affirmative, this would very significantly impact my assessment of this paper.", "was_revised": true, "final_justification": "I choose to raise my initial score to 4, a borderline accept.\n\nMy reasons are that I appreciate the thorough rebuttal and added experiments. The JODO+Energy runs on QM9 demonstrate that the loss provides consistent gains beyond the original architectures, and the wall-clock discussion plus the sparse, globally-rigid construction convinces me runtime isn’t a blocker. Those were my main practical reservations, and they’ve been adequately addressed.\n\nMy remaining caveat is breadth. This work proposes a new loss family, so I expect wider evaluation across task types and datasets—not mostly generative settings. The current results show it works where tested, which is meaningful, but to move from a promising direction to something with general impact, I need to see the method tried on a mixed set of tasks (e.g., supervised/conditional regression, inverse problems/dynamics, larger systems) and situated with standardized comparisons. Even if the pending GEOM-Drugs result lands well, that alone doesn’t resolve the central question of generality. I also note, more lightly, that the reverse-KL derivation presumes near-equilibrium data and the diffusion analysis holds best at low noise; neither is fatal, but both strengthen the case for breadth.\n\nWeighting these points: the idea is simple, principled, and architecture-agnostic; the added evidence shows real gains with acceptable cost, and I would lean towards accept. I stop short of an outright Acceptance (i.e. a score of 5) because the experimental scope doesn’t yet match the ambition of introducing a general-purpose loss.", "rebuttal": "We thank the reviewer for the review and useful feedback. We appreciate that they have found our idea creative and sound, the theory convincing and that experiments show meaningful improvements.\n\nWe address their concerns and questions in details below:\n\n### Comparison with more recent architectures\nWe thank the reviewer for this suggestion, and agree. Our goal is to show that improvements resulting from the energy loss are orthogonal to architectural improvements and can indeed push the state-of-the-art. Following the reviewer’s suggestion, we have therefore run experiments using the recently proposed JODO architecture on QM9. \n\nGiven the time constraint, we ran the model on 500k steps, rather than the 1.5M reported in the paper. For transparency, we include both settings for JODO. We can update these results later in the discussion period.\n\nMetric-3D\n| Model                           | Atom stable ↑ | Mol stable ↑ | Validity ↑ | Completeness ↑ | FCD ↓ |\n| ------------------------------- | ------------- | ------------ | ---------- | -------------- | ----- |\n| JODO (paper)                    | 99.2%         | 93.4%        | —          | —              | 0.885 |\n| JODO (500k steps)               | 98.96%        | 91.57%       | 98.04%     | 97.93%         | 0.839 |\n| JODO + Energy Loss (500k steps) | 99.53%        | 96.00%       | 98.29%     | 98.19%         | 1.558 |\n\nMetric-Align\n| Model                           | Bond ↓ | Angle ↓ | Dihedral ↓ |\n| ------------------------------- | ------ | ------- | ---------- |\n| JODO (paper)                    | 0.1475 | 0.0121  | 6.29e-4    |\n| JODO (500k steps)               | 0.1994 | 0.0189  | 0.001029   |\n| JODO + Energy Loss (500k steps) | 0.1464 | 0.0164  | 0.006370   |\n\nOur method outperforms JODO (which uses an equivariant network and Kabsch alignment) on nearly all 3D metrics and has more accurate bond lengths and angles than JODO at matched training compute. We note that the reported run uses exponential coefficients, a choice based on our results with EDM. It is possible other forms of coefficients could be best when targeting other metrics.\n\n### Local minimum assumption\nUsing the true energy would rely on the assumption that the data is at a minimum: if the data is no longer at a local minimum, the assumption is not satisfied and the model would no longer be trained to fit the data. This is one of the reasons why we argue for not using the true energy. We construct our approximation in such a way that for any data points (equilibrium or not), the data sits at a minimum of the approximate energy, while the landscape still captures energetic considerations. So the response to the question asked by the reviewer is that the loss used in the paper will not result in gradients pointing in wrong directions.\nThe MSE loss function can be seen as an (simplest possible) energy approximation that yields an energy minimum at the data, but with an unphysical landscape around this minimum. In practice, what we really aim for is an improvement over the MSE loss.\nWe will augment and clarify the discussion of this in the paper regarding this point.\n\n### Wall-clock speed\nWe have included wall-clock times on a single NVIDIA L40s following your suggestion. Our main objective in including the sparse energy loss was to demonstrate our method can efficiently generalize to systems with many particles where loss calculation may contribute significantly to running time (e.g. very large point clouds). This is not the case for molecules, where the neural network (GNN or Transformer) is typically fully connected and thus scales as N^2. In the results below on QM9 with GDM, the most expensive loss calculation is less than 1% of the total backward and forward time.\n\n| Component            | Loss Type     | Time (ms) (ms) \n| -------------------- | ------------- | ------------------- |\n| **Loss computation** | MSE           | 0.1840 ± 0.0084     |\n|                      | Energy        | 0.5100 ± 0.0161     |\n|                      | Sparse Energy | 0.5712 ± 0.0242     |\n|                      | Kabsch Align          | 1.1361 ± 0.0300     |\n| **Forward pass**     | –             | 73.55 ± 16.38       |\n| **Backward pass**    | –             | 94.47 ± 3.37        |\n| **Optimizer step**   | –             | 1.43 ± 0.0072       |\n\n\nTo better understand the scale at which this becomes a relevant consideration and the utility of the sparse energy loss, see the following wall clock times from the shape generation setting:\n\n| # Nodes | Energy (ms)     | Sparse Energy (ms) | MSE (ms)        | Kabsch Align (ms) |\n| ------- | --------------- | ------------------ | --------------- | ----------------- |\n| 30      | 0.240 ± 0.0021  | 0.255 ± 0.0012     | 0.058 ± 0.0004  | 0.807 ± 0.0027    |\n| 300     | 0.245 ± 0.0011  | 0.257 ± 0.0024     | 0.059 ± 0.0004  | 0.804 ± 0.0048    |\n| 3000    | 20.120 ± 0.0055 | 0.275 ± 0.0018     | 0.0779 ± 0.0008 | 0.944 ± 0.0028    |\n| 30000   | -               | 0.293 ± 0.0029     | 0.0740 ± 0.0011 | 3.242 ± 0.1452    |\n| 300000  | -               | 2.652 ± 0.0050     | 0.131 ± 0.0004  | 24.381 ± 0.0272   |\n\nAt 30000+ nodes, the energy loss requires too much memory to compute. Importantly, the sparse energy is cheaper than the Kabsch Align by a factor 5-10x. \n\nNote that all losses have some constant cost that does not scale with N contributing to the wall-clock time. This explains why for QM9 (avg. 28 atoms) energy loss is marginally faster than sparse energy and why in the scaling table wall-clock times start to increase with N only after a certain point.\n\nInterestingly, using an equivariant network with EDM takes 129.71 ± 0.044 ms for the forward pass and 180.07 ± 0.070 ms for the backward pass. Using the energy loss imparts a 0.3% increase on one backward pass through the model, while using an equivariant architecture imparts a 94% increase, while providing inferior benefits. A crucial conclusion is that our results show that the energy loss with a non-equivariant architecture results in more improvement than using an equivariant architecture, at negligible computational cost, which we think is a significant finding. We will discuss it further in the main text.\n\n### Sparse loss on larger systems\nThis is a good suggestion, we plan to include results using the sparse version of the loss for GEOM-drugs, which includes bigger molecules than QM9, during the discussion period. \nWe also want to emphasize that even if it were the case that some physical information was lost, it would still be better than MSE, since the sparse energy at least includes some physical information while having the same minimizer.\nThat being said, we consider that generalizing the proposed sparse energy approximation to explicitly account for long-range effects (as well as other contributions to the energy) is a very relevant avenue for future work.\n\n### Stability calculation\n\nYou are correct that we use the 3D calculation for molecular stability as in EDM and the second part of Table 2 in JODO. A molecule is stable if all atoms have the correct valence. Valence is determined by bonds that are drawn according to a bond distance table. Our new results show this metric can be improved by 3-5% when using energy loss over MSE.\n\nWe hope this response addresses your concerns. If any further clarification would help you reconsider your overall assessment, we would be happy to provide it."}]

Author rebuttal:

Meta-review: The paper introduces energy loss functions through the reverse kullback leibler divergence as a loss function for learning on data reflecting physical systems. This idea is quite similar to established methods such as energy-based training in Boltzmann Generators, however, is different in the sense that the approach can be applied even when a ground truth potential is not available, e.g. by assuming some other potential. The approach is interesting, however, some reviewers highlighted the incremental nature of the work technical and conceptual contributions, and while the loss is touted as a general learning paradigm, experiments only cover one learning modality. Nevertheless, the paper is solid and provides an interesting direction which may have impact across a broader range of ML applications.
Decision: Accept (poster)