Geometric Graph Neural Diffusion for Stable Molecular Dynamics Simulations
Authors: Haokai Hong
Wanyu Lin
Zhang Chusong
KC Tan
Venue: ICLR 2026 Poster (2026)
Keywords: Machine learning force field
graph neural network
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: Geometric graph neural networks (Geo-GNNs) have revolutionized molecular dynamics (MD) simulations by providing accurate and fast energy and force predictions. However, minor prediction errors could still destabilize MD trajectories in real MD simulations due to the limited coverage of molecular conformations in training datasets. Existing methods that focus on in-distribution predictions often fail to address extrapolation to unseen conformations, undermining the simulation stability. To tackle this, we propose Geometric Graph Neural Diffusion (GGND), a novel framework that can capture geometrically invariant topological features, thereby alleviating error accumulation and ensuring stable MD simulations. The core of our framework is that it iteratively refines atomic representations, enabling instantaneous information flow between arbitrary atomic pairs while maintaining equivariance. Our proposed GGND is a plug-and-play module that can seamlessly integrate with existing local equivariant message-passing frameworks, enhancing their predictive performance and simulation stability. We conducted sets of experiments on the 3BPA and SAMD23 benchmark datasets, which encompass diverse molecular conformations across varied temperatures. We also ran real MD simulations to evaluate the stability. GGND outperforms baseline models in both accuracy and stability under significant topological shifts, advancing stable molecular modeling for real-world applications.

Published as a conference paper at ICLR 2026
GEOMETRIC GRAPH NEURAL DIFFUSION FOR STABLE MOLECULAR DYNAMICS SIMULATIONS
Haokai Hong' Wanyu Lin 2,2; Chusong Zhang 1,3, Kay Chen Tan'
'Department of Data Science and Artificial Intelligence, DDepartment of Computing, The Hong Kong Polytechnic University, Hong Kong SAR, China. "Zhejiang University. haokai hongeconeeet.......o.oooooooo... www.ooonnnnttttt.....eooo.eommee {wan-yu. 1 lin, ketan (rtttan@@poly....mm edu. hk
ABSTRACT
Geometric graph neural networks (Geo-GNNs) have revolutionized molecular dy- namics (MD) simulations by providing accurate and fast energy and force predic- tions. However, minor prediction errors could still destabilize MD trajectories in real MD simulations due to the limited coverage of molecular conformations in training datasets. Existing methods that focus on in-distribution predictions often fail to address extrapolation to unseen conformations, undermining the simulation stability. To tackle this, we propose Geometric Graph Neural Diffusion (GGND), a novel framework that can capture geometrically invariant topological features, thereby alleviating error accumulation and ensuring stable MD simulations. The core of our framework is that it iteratively refines atomic representations, enabling instantaneous information flow between arbitrary atomic pairs while maintaining equivariance. Our proposed GGND is plug-and and-play module that can seamlessly integrate with existing local equivariant message- passing frameworks, enhancing their predictive performance and simulation stability. We conducted sets of exper- iments on the 3BPA and SAMD23 benchmark datasets, which encompass diverse molecular conformations across varied temperatures. We also ran real MD simu- lations to evaluate the stability. GGND outperforms baseline models in both accu- racy and stability under significant topological shifts, advancing stable molecular modeling for real-world applications.
INTRODUCTION
Molecular dynamics (MD) simulations rely on force fields approximate the underlying poten- tial energy surface and generate long-temporal trajectories of molecular systems. Geometric graph neural networks (Geo-GNNs) have transformed MD simulations by providing a computationally ef- ficient alternative to quantum mechanical methods, while maintaining high accuracy in predicting energies and forces (Wang et al., 2024a; Batatia et al., 2022: Wang et al., 2024c). Existing Geo- GNN evaluations mainly focus the accuracy of predicting forces and overlook the performance evaluation in real MD simulations, c.g., whether the real MD could reveal detailed physical mecha- nisms (Lane et al., 2011). Recent studies have shown that even small errors in predicting forces could lead to catastrophic failure in real long-time simulations (Fu et al., 2023). This is because through- the long-temporal trajectory, it can exhibit molecular conformations that are out of the training distribution. More specifically, due to the lack of extrapolation capability, most Geo-GNNs cannot produce accurate force prediction for unseen conformation, introducing pathological behaviors, i.e., unphysical chemical bonding. in real MD simulation. Such a phenomenon can be quantified via chemical bonding connectivity in a real MD simulation, and is termed stability (Fu et al., 2023).
To examine the influence of conformation shifts on current Geo-GNNs, we use the 3BPA dataset (Kováes et al., 2021), as the dataset contains molecular geometries sampled 300 K, 600 K, and 1200 K. Each temperature setting naturally induces distinct conformation domain. We quantify conformation variations using edge-frequency distributions of atom pairs and prove that discrepancies grow systematically with increasing temperature gaps (Figure (a) and (b). We then
trained the representative Geo-GNNs VisNet (Wang et al., 2024c) and SEGNO (Liu et al., 2024) at
"Coreeponniigg author.

---

Published as a conference paper at ICLR 2026
(b)
25
(c)
Training Set 300 K
300 K
600 K
1200 K
Metries
Energy
Foroz
Stubility
Energy
Force
Stability
Energy
Force
Stuiliit
VisNet
1.4053
3.4642
1.4043
0.004
SEGNO
99.812
0.9082
0,8925
1.2375
0.009
Figure 1: Geometric Topological Shift Analysis of the 3BPA Dataset and Extrapolation Per- formance Across Conformational Domains in the 3BPA Dataset: (a) Distribution of adjacency matrix of 3BPA in training data (300 K); (b) distributional difference of Adjacency matrix of 3BPA in testing set (k=300 K, 600 K, and 1200 K) and training set (300 K); (c) extrapolation performance evaluation across conformational domains in the 3BPA Dataset.
300 K and tested them at 300 K, 600 K, and 1200 K. VisNet is the state-of-the-ard Geo-GNN for sim- ulating MD. While VisNet demonstrates strong within-domant performance (300 K), its accuracy degrades sharply under shifted conformation spaces (Figure (c)). In contrast, SEGNO improves generalizability via explicitly embedding physical biases. indeed improves the extrapolation ability, but suffers from in-domain performance degradation. These findings confirm the urgent need for a new Geo-GNN that can remain robust and extrapolate effectively across various conformation domains, leading stable MD simulations.
To fill the gap, we propose a new framework, dubbed geometric graph neural diffusion (GGND), inspired by the graph heat equation- a generalization of the diffusion equation rooted in spectral graph theory (Chung, 1997). Specifically, to facilitate the theoretical analysis, we first conceptual- ize domain variations in conformational spaces "geometric topological shifts." Correspondingly, we introduce the diffusion process with two novel operators-eequivaiinn gradient and diffusivity operators-to capture the invariance conformational changes while maintaining equivariance. In particular, the gradient operator captures variations in node features across the topology of the ge- ometric graph by characterizing differences between arbitrary nodes, while the diffusivity operator regulates the rate and extent of information propagation. Together, these operators drive the evolu- tion of node representations, capturing all plir information flows over complete molecular graph, thereby remaining invariant to conformational changes. Our main contributions are outlined below:
First, we propose geometric graph neural diffusion (GGND) that can extrapolate effectively across various conformation domains, leading to stable MD simulations.
Second, we provide a theoretical analysis of GGND, establishing a regret bound under geometric topological shifts and proving the equivariance of the model. This regret-bound guarantees improved performance in extrapolating to unseen molecular conformations and enhances the stability of MD simulations.
Third, GGND functions as plug-in module, scamlessly integrating with existing EGNNs to en- hance their extrapolation capabilities. We evaluate GGND's performance on the 3BPA (Kovács et al., 2021) and SAMD23 (Kim et al., 2023) datasets, focusing on stability metrics for unseen molecular conformations. Our results demonstrate robust generalization across diverse conforma- tional spaces and superior stability in real-world MD simulations compared to all baselines.
2 PRELIMINARIES AND RELATED WORKS
Molecular graph. In this paper, we explore the dynamics simulation of large-scale molecular systems, represented sequence of geometric graphs & indexed by time t. Suppose have N atoms the system, then the molecular system G at each snapshot can be represented as a point cloud denoted as (X, H), where X :i;...;:Nl R is the atom coordinate matrix
2

---

Published as a conference paper at ICLR 2026
Unobserved variables
Observed graph variables
Messurable
Atomm: nodees
3BPA
function
500
Favironment
Geoenotric Topological Shift
Training
variations
Gruph
Graph model with predefined cutoff
labels
3BPA
600
Graph geneation
Favironmental
Model
cause
couse
The Illustration of Geometric Topological Shifts Caused by Environment E and Model with Predefined Cutoff C: The unobserved measurable function W and environment E (temperature or pressure), along with the modeling method M, serve as causes in the graph topology formation process, influencing unobserved variables and leading to variations in observed graph variables. This mechanism demonstrates how environmental causes (e.g., temperature changes) and model causes (predefined cutoff C) result in geometric topological shifts between the training set (e.g., 3BPA 300 K) and the testing (e.g., 3BPA at 600 K).
Figure 2:
and H = I.. hy] E NR is the node feature matrix. H typically contains atomic types or charge features, and it is generally time-invariant. Given the molecular structure G, the objective of the machine learning force field is to predict the energy or forces with the molecular graph input G.
Topology of 3D graph. In this study, we focus on the geometric topology of 3D molecular graphs. For geometric topology, nodes are atoms, and edges are established based on a predefined model- related radius cutoff distance threshold, such that pairs of atoms within this cutoff are considered neighbors. The term "topology" may also refer to the biochemical topology (or 2D molecular graph); however, unless specified otherwise, this paper focuses on geometric topology. Under conforma- tional changes caused by the environment, the biochemical topology generally remains invariant, as is defined by the fixed chemical connectivity of the molecule. Conversely, the geometric topology is dynamic, varying with the predetermined cutoff distance and the spatial coordinates of atoms, which may shift due to conformational changes.
Geometric topological shift. We propose a causal mechanism for geometric topology formation within a molecular system, as illustrated in Figure 2. building on prior work (Medvedev, 2014; Sni- jders & Nowicki, 1997). Unlike a 2D graph, our approach generalizes the data-forming mechanism to incorporate both geometric topological adjacency and node features. Specifically, a 3D graph with geometric topology, denoted as = (X, H, A). is formed by a graphon-4 continuous graph limit defined as symmetric, measurable function W [o, 1J2 + [0, 1]-serving an unobserved latent variable, alongside a modeling method that specifies a cutoff radius C.
To elucidate the node-level structure of this graph, each node E V is associated with an inde- pendent and identically distributed (i.i.d) latent variable Ug ~ 24(0, 1]. The vector and sealar features, X [X.]] and H [Hu]. are random variables derived from each U,, through node- wise functions X, = f(U,; W) and Ha = h(U,; W), respectively. Next, the geometric topo- logical adjacency matrix A [A(u, ui]] is random variable determined by pairwise function A(u, v) hUU,, U.: W, E, M). which depends on the environment E and the modeling method M. Changes in E, such transitions from training testing, lead variations in the distribution of A. Beyond the node features and adjacency structure, the label Y also varies due to conforma- tional variations. We assume Y is formed by function Y r(U, V, A: W), with a specific realization denoted as Y. We denote specific realizations of these random variables as matrices X, H, A, and Y.
Extrapolation and stability. Extrapolation remains fundamental challenge in MD simulations, particularly the application data-driven machine learning methods to MD. From an ML per- spective, extrapolation in MD can be categorized into two types: 1) extrapolation to chemical space, and 2) extrapolation to conformational space. The former entails predicting properties or dynamics for molecules absent from the training set, while the latter involves forecasting dynamics for un-
3

---

Published as a conference paper at ICLR 2026
Local Equivariant (b)
Geometric Graph Diffusion
Message Passing
Geometric
ropologecll dhilts
Equivariance difftasion operations
Generalizable energy and force prediction
Diffusion on fully-- connected graph
Geometric sopelogical invariant features
Staiiity fhreshold
Local environment
Stable MD cimmiitions
Figure 3: The Illustration of Geometric Graph Neural Diffusion: (a) Our method serves plug-in module that integrates with local equivariant message passing. (b) The GGND uses equiv- ariant diffusion operators (gradient and diffusivity) on fully connected graph to capture domain- invariant geometric topological features. (c) The local message passing and the equivariant diffusion operators are combined to address geometric topological shifts, enabling generalizable energy and force predictions for stable molecular dynamics simulations.
seen molecular geometries, such as samples from different temperatures or non-equilibrium states. Improving the ability to extrapolate to unseen conformations is critical for model to ensure stable MD simulations, a principle supported by numerical MD methods (Barth & Schlick, 1998; Miao & Ortoleva, 2009). To tackle extrapolation, MatterSim (Yang al., 2024) employs an active learning approach to address both categories, though relies on costly high-quality data collection. Besides, SEGNO (Liu al., 2024) integrates second-order motion laws to enhance the generalization of cquivariant graph neural networks, yet it fails to address geometric topological shifts. Despite the importance of conformational space extrapolation for stable MD, enabling Geo-GNNs general- ize under geometric topological shifts remains an unresolved challenge. A brief overview of MD, generalization, and equivariance is provided in Appendix A.
METHOD
Our design enhances robustness to geometric and topological variations while preserving SB(3)- equivariance, thereby enabling stable MD simulations despite limited coverage of molecular con- formations in training data. The model integrates two complementary components: (i) novel geometric graph neural diffusion module with global attention and (ii) conventional local equivari- ant message passing neural network (Satorras al., 2021; Wang et al., 2024c). The GGND employs diffusion process modeled as partial differential equation (PDE) on the graph, incorporating global attention to capture long-range dependencies across all nodes. This global perspective miti- gates challenges posed by geometric and topological shifts by facilitating information propagation beyond local neighborhoods. In contrast, the EGNN focuses on local interactions, updating node features and positions through message passing within local neighborhoods. The GGND serves as a plug-in in module, seamlessly integrable with most existing EGNN frameworks to enhance their performance in stable MD) simulations. We provide overview of method in Figure 3.
3.1 GEOMETRIC GRAPH NEURAL DIFFUSION
The geometric graph neural diffusion model is designed to learn equivariant features that are robust shifts in geometric topology. To enable diffusion on geometric graphs, we incorporate higher- order equivariant message passing, which facilitates accurate modeling of such graphs. Given geometric graph G (V,E) with n [V] nodes, where each node E V has a scalar feature vector h, d and position x; E R3, and È is the set of edges determined by the adjacency matrix A, assuming full connectivity in the graph. A concise background on diffusion processes on graphs is provided in Appendix B.
Each node has initial features z, (0), which include chemical element features h; (invariant scalars) and positions x; (for equivariance). The features (t) consist of spherical tensors labeled by inre- ducible representations of O(3), denoted as (t). where & indexes channels (learnable fea- tures), T. is the degree (e.g. L for invariants, L. 1 for vectors, and higher L. for ten-

---

Published as a conference paper at ICLR 2026
sors), and M ....... L indexes components. Diffusion models geometric graphs replace discrete GNN layers with continuous time-evolving node embeddings Z(t) = {z, (t) where (t) [0, 00) - Rd and evolves according to the diffusion equation:
az(t) div (s(, VZ(t) 0,
Ōt
where Z(t) (+) are equivariant node features, with initial conditions Z(0) de (X, H), and de is the embedding layer through by radial basis functions (RBF).
The term (tZ(t), X, t) denotes the diffusivity over the graph, defined as an n n matrix-values valued function dependent on Ag. which measures the rate of information flow between node pairs. The gradient VZ maps node fields to edge fields, while the divergence operator div is its adjoint, map- ping edge fields back to nodes. This diffusion process is modeled as partial differential equation (PDE) on the graph, adapted to handle equivariant features via higher-order messages.
The GGND module is designed to learn features invariant to geometric topological shifts, enabling extrapolation to unseen molecular conformations. To achieve this, we introduce two novel operators: an equivariant gradient operator and equivariant diffusivity operator, which facilitate global information flow while maintaining equivariance.
Equivariant gradient operator. order tensors, incorporating directional information to preserve 3E(3)-qquivariance It is defined as:
The gradient operator V generalizes scalar differences higher-
(Vz) (Vz)ij.klmm
Σ W kklt (2 5,2mm 22,2mm2
where W are learnable weights for mixing channels. This equivariant gradient operator on the graph generalizes the scalar gradient to higher L,, ensuring equivariance, with the difference 2;- Z; modulated by directional information to preserve 3D structure. Notably, J ranges over all nodes in V, aligning with latent interactions among nodes determined by the underlying data manifold. This induces all-pair information flows over a complete graph and remains invariant to changes in & due to conformational variations.
The diffusivity S(t) is made equivariant by defining it as a tensor-valued attention matrix. We extend scalar attention to tensors follows:
Equivariant diffusivity operator.
((t) klymms
Σ
Clame Camiltama tats [|xx)|])))) (t),
(3)
11.22.m1..m2
where Clams damm are Clebsch-Gordan coefficients ensuring proper equivariance of S(t). Here, ||xx<|| = 11xy-x;li Ets the unit vector, Y!, are spherical harmonics (for directional equivariance), Rat, is a learnable radial basis function derived from Bessel functions and an MI.P (ensuring invariance to distance), and an equivariant pairwise interaction (e.g., a gated tensor product). This formulation ensures that (t)) transforms correctly under SE(3), serving as equivariant filter that captures global dependencies. The attention matrix S = (s(x,, x,)) is right-stochastic, allowing Equation (1) to be rewritten as:
OZ(t) Øt (S(Z(t), (,) ()) OstsT.
(4)
Equation (4) governs the dynamics of the system from = 0 to a specified stopping time T, produc- ing geometric and topological node representations Z(T). The equation is generally nonlinear due to the dependence of the diffusivity matrix on Z. linear variant emerges when attention weights fixed; however, as static attention is impractical, we focus on the nonlinear GGND model.
Output. For energy prediction, we utilize the invariant components of Z(T), specifically 24, (T), combined with local equivariant features learned by Equivariant Graph Neural Networks (EGNNs). This ensures that site-specific energy contributions E, remain invariant, computed as E = dD (f) = Ex ΣE Wgf. where f, represents features fused from local EGNN outputs 1., and geometric-topole topologiaal invariant features 20,600 via concatenation and a linear transfor- mation: f W (l 20, (T)).
5

---

Published as a conference paper at ICLR 2026
Equivariance. The diffusion process the geometric graph, as described in Equation (4), enables the learning of domain- domain-invariam invariant features. With equivariant gradient and diffusivity operators, also ensures the equivariance of the learned global features. We provide proof of the equivariance of our GGND in Appendix C.
3.2 ALLEVIATING GEOMETRIC TOPOLOGICAL SHIFTS
We analyze the extrapolation capability our geometric graph neural diffusion model with respect to geometric topological shifts, defined in Section 2. Our focus is the extrapolation error of the parametric function o, instantiated as the continuous equivariant diffusion model in Equations (4), when transferring from training data generated under environment Etr (and modeling method Mcr) to testing data under Ere (and Mur). Such shifts may arise from variations in adjacency matrices due different cutoff radii environmental conditions affecting inter-node distances in molecular systems.
Denote the training dataset of size Nur as {(X(, E Etr, M Mr.) ) and let -(--)) be a bounded loss function. The training error
HO,
Y drawn from p(X, A,Y
No.
(T Fer, Mor)
Σ
(Too(
H A ( Y()).
(5)
Nur
Our objective is to minimize the expected loss on testing data from p(X, H. A,Y E = Ete, Mc Mir):
O(T:; Exc, Mur A B. | H. Y (6)
When Ete = Eur, this reduces to the standard in-distribution setting, where the extrapolation gap is bounded by
C(T:: Eu, Mu)) (T's Eur Mx) S (Tee Eir, Mus Nu) 2H(TTe +0 (1
log(1/0)
Nir
(7)
With H(Ta), the Rademacher complexity of the function class is induced by To, and the upper bound is determined by dataset size and model complexity.
In the diitii--iitiiiitionn regime where Free * Ftr, geometric topological shifts complicate the analysis. Changes in geometric topologies alter node representations Z(T) in the equivariant graph diffusion equations (4), expressible as ZTT A) (()( A). The extrapolation gap could be decomposed into three terms (Wu et al., 2025). Assume é and do are Lipschitz continuous. For geometric graph data generated per Section 2, with probability least 1 - 6, the extrapolation gap satisfies:
I(T:: Ere, Md) (To Eur. Mc)) S Din (To, Eur, Mu) M.)
+ O E AA M.) Z[T A') A(T;A)]la) A)
(8)
Y -Yllz))
We denote the first O(-) as OOD model error (T Eur, Mus Ete) and the second O(-) as OOD label error Dr. (FEr, Mar, F.c)). Since D., is independent of testing data under Exr * Fir. the impact geometric topological shifts on extrapolation hinges on DM and Dr.: the former captures variation in Z(T; A) due shifting topologies (e.g., adjacency changes from varying cutoff radii or conformations), while the latter reflects label differences across environments or methods. Dr, is dictated by the data-generating process, whereas DM depends on I's, specifically the sensitivity of representations to shifts. We next examine Ie as in Equation (4), adapted for equivariance.
Theorem 3.1 For geometric graph data per Section if f and h are injective, the geometric graph neural diffusion model in Equation (4) reduces the representation variation [Z A') Z(T; l to any order (*[[)A (|])), where vp is an arbitrary polynomial, AÅ À' A, and
D 1/2 1/2 1/1 10 1/2 (with A incorporating geometric distances via cutoff radius and D is the diagonal degree matrix of A).

---

Published as a conference paper at ICLR 2026
This indicates that the geometric graph neural diffusion model controls representation changes arbitrary rates relative to | maintaining robust force prediction for conformation variations in molccular dynamics. The injectivity of f and h are mild assumptions, mapping from compact latent spaces to high-dimensional vector and scalar features. Applying Equation (8) yields the following.
Corollary 3.2 Under the condition of Theorem 3.1, the model-dependent extrapolation bound in Equation (8) reduces to arbitrary polynomial orders with respect to geometric topological shifts: DTa(To, o, Eur, Ere, Mur) O (RAmp(AlBrr.MMa Mer).. A A'mppA/ (.p(A(Er.,Me)) (+0 (+())
This bound ensures controllable extrapolation enror any rate relative to | The model achieves desired extrapolation ability under shifts, such as in machine learning force fields sim- ulations with conformational changes or varying cutoffs. contrast, the change rate of features produced by the local message passing model has exponential upper bound. We presented the proof for the Corollary 3.2 in the Appendix D.
EXPERIMENTS
4.1 EXPERIMENTAL SETUP
Datasets. We used the 3BPA and SAMD23 datasets to evaluate our model's performance, par- ticularly in the presence of geometric and topological shifts. The 3BPA dataset consists of 500 training structures of the flexible, drug-like molecule test data provided at 300 K, 600 K, 1200 K, and different dihedral angles (Kovács al., 2021). The SAMD23 dataset comprises simulations of the semiconductor materials SiN and HfO under various conditions, including variations in initial structures, stoichiometry, temperature, strain, and defects, with unit cells containing up to 510 atoms (Kim et al., 2023).
at 300 K, with
Baselines. Our proposed equivariant graph neural diffusion can be integrated with any local equiv- ariant paassgge-aasing-basstte method. To evaluate performance improvements, we selected four representative
(Batzner et al., 2022), MACE (Batatia et al., 2022), SEGNO (Liu et al., 2024), and VisNet (Wang et al., 2024c) as baselines and compared approach when com- bined with them against these baselines alone. Additionally, we conducted a comprehensive com- parison with several SOTA models on the SAMD23 dataset, including Allegro (Musaelian et al., 2023), Equiformer V2 (Liao et al., 2024), QuinNet (Wang et al., 2023), Neural P'M (Wang et al., 2024b), LSRM (Li al., 2024b), and FreeCG (Shao et al., 2025),
Metrics. Accuracy: we evaluate the predictive performance of our model using the mean abso- lute error (MAE) for energy and force predictions. For the SAMD23 dataset, which includes SiN molecules with atom counts ranging from to 510, we report the energy per atom to ensure compa- rability across molecular sizes. Stability: following the methodology (Fu et al., 2023), we assess the stability of ficxible molccules by monitoring bond length deviations. A real MD simulation is classified as unstable at time T if the maximum deviation of any bond length from its equilibrium value exceeds a threshold, formally defined as: maxi,jeB [i( - xj(T)] bıyl > A, where B denotes the set of all bonds, and j are bond endpoints, b,j is the equilibrium bond length, and A is the stability threshold. For systems with periodic boundary conditions, stability is evaluated using the radial distribution function (RDF). A simulation deemed unstable at time T when: (r)Ft(r))t = +T|rr > A, where (-) represents the time-averaging operator, is 1 ps time window, and A is set to 1.0. We perform constant-energ. (NVE) molecular dynam- ics simulations at the specified temperature, employing Velocity Verlet integration over 100 ps with fs timestep. The stability metric is defined as the first timestep (in ps, ranging from 0 to 100) at which an unstable molecular configuration occurs. We conduct five independent molecular dy- namies simulations and report the average stability metric as the final result. Higher stability values indicate better performance in maintaining long-term stable molecular dynamics simulations.
Jo (RDF(r))
4.2 RESULTS AND ANALYSIS
7

---

Published as a conference paper at ICLR 2026
Table 1: Accuracy and Stability on the 3BPA Dataset. MAE for energy (E, eV), force (F, eV/A), and stability (S. ps) of three baseline models and our proposed model (+GGND), trained on config- urations of the flexible drug-like molecule 3BPA at 300 K and evaluated on 300 K, 600 K. 1200 K. and varied dihedral angles. Best results are in bold; tied results underlined.
Conformation
Metries
MACE
+GGND
NequlP
+GGND
SEGNO
+GGND
VisNet
+GGND
B (1)
0.113
0.010
0.165
0.094
0.593
0.293
0.002
0.002
300K
F (()
0.165
0.022
0.113
0.104
0.359
0.183
0.006
0.006
S (t)
100
100
100
100
99.812
100
100
100
E (L)
0.161
0.023
0.335
0.122
0.908
0.295
1.405
0.022
600K
(L)
0.335
0.044
0.161
0.153
0.893
0.193
0.997
0.041
S (t)
100
100
98.271
100
59.892
100
25.358
100
(()
0.271
0.109
0.770
0,477
2.836
0.503
3,464
0,583
1200K
F())
0.770
0.111
0,271
0.269
1.238
0.285
1.404
0.304
S (T)
1.965
29.218
0,018
17.052
0.009
16,201
0.004
11.209
E E()
0.169
0.012
0.387
0.375
0.923
0.267
0.789
0.050
Dihedral Slices
F())
0.289
0.017
0.242
0.189
0.795
0.192
0.697
0.039
S (T)
100
100
89.119
100
72.282
100
47.785
100
Table 2: Accuracy and Stability on the 3BPA Dataset. MAE for energy per atom (E/A, cV), force (F, eVIA), and stability (S. ps) obtained by SOTA models and our proposed model (GGND), trained on SiN and HfO semiconductor molecular system. Best results are in bold.
Molecule
Splits
Metrics
NequlP
MACE
Allegro
Neural P3M
QuinNet
Equiffrmmr V2
LSRM
FreeCG
GGND
E/A
0.013
0.012
0.015
0,030
0
0,010
0,010
0.011
0.009
Test
[()
0.598
0.526
0.673
0.485
0.490
0.451
0.490
0.494
0.443
S (1)
69.009
78.845
63.583
88,280
83,286
98,284
81.000
84,500
100
SIN
(J)
0.022
0.018
0.017
0.021
OOD
F())
1.018
0.912
1.185
0.837
0.836
0.972
0.832
0.844
0.754
(1)
63.733
65.710
55.824
85,888
86.512
82,031
74,217
76.631
99,892
E/A (J)
0.007
0.006
0.007
0.006
0.006
0.005
0.006
0.006
0.005
Test
F())
0.377
0.335
0.385
0.311
0.304
0.298
0.312
0.315
0.179
(()
65.377
78.054
64.282
89.034
97.184
87.353
85.040
HIO
E/A (J)
0.011
0.010
0.012
0.009
0.009
0.010
0.010
0.009
0.068
0OD
F())
0.430
0.570
0.457
0.683
0.593
(()
61.621
65.689
60.982
84,
85.453
79.762
86.373
75.916
97.928
Performance on 3BPA datasets. Experimental results on the 3BPA dataset in Table show that integrating the proposed GGND module with baseline Geo-GNN models (MACE, NequlP, SEGNO, and VisNet) significantly enhances performance, particularly in extrapolating to conformational do- mains with improved stability. On the in-domain 300 K test set, GGND improves energy and force prediction accuracy for most baselines while achieving perfect stability at 100 ps. For instance, SEGNO's energy MAE decreases from 0.593 eV to 0.293 eV and force MAE from 0.359 eV/Å to 0.183 eV/Á, with stability rising from 99.812 ps to 100 ps. At 600 K. where domain shifts occur, GGND's advantages are more pronounced; it reduces VisNet's energy MAE from 1.405 eV to 0.022 eV and force MAE from 0.997 eV/À to 0.041 eV/À. boosting stability from 25.358 ps 100 ps. Comparable improvements are observed for SEGNO, with energy MAE dropping from 0.908 eV to 0.295 eV and stability from 59.892 ps to 100 ps. These findings indicate that SEGNO addresses universal generalization but not geometric topological shifts. Additionally, GGND outperforms all baselines on dihedral slices.
Under severe geometric topological shifts at 1200 K, baselines suffer catastrophic degradation in stability (e.g.. VisNet at 0.004 ps. MACE 1.965ps. SEGNO 0.009 ps), whereas GGND restores robustness, increasing MACE's stability to 29.218 ps (15-fold), NequlP to 17.052 ps (947-fold), SEGNO to 16.201 ps (1800-fold), and VisNet to 11.209 ps (2802-fold). Concurrent accuracy gains include VisNet's s energy MAE reduction from 3.464 eV to O. 108 eV. These results highlight GGND's efficacy mitigating geometric topological shifts via all-pair information diffusion, facilitating sta- ble long-term MD simulations in unseen conformations without additional DFT data.
Performance on SAMD23 dataset. The GGND model outperforms baselines across SiN and HfO datasets in both Test and OOD splits, as shown in Table 2. For SiN, GGND achieves lower en-
8

---

Published as a conference paper at ICLR 2026
Table 3: Ablation Analysis on the 3BPA Dataset. MAE for energy (E, eV), force (F, eVIA), and stability (S. ps) of baseline model, GGND, and two variants of GGND. Best results are in bold; tied results are underlined.
Conformation
300 K
600
1200
Variations
E (4)
F (L)
S (f)
E (L)
(()
S (1)
E (1)
F (4)
S (1)
Baseline
0.002
0.006
100
1.405
0.997
25.358
3,464
1.404
0.004
GGND
0.013
0.058
100
0.982
0.998
39.075
3.049
1.406
0.291
GGND 0.015
0.072
98.827
0.643
0.661
69.292
1.908
0.882
2.892
GGND 0.002
0.006
100
0.022
0.041
100
0.583
0.304
11.209
GGND with local diffusion graph.
$ Local message passing baseline plus fully-connected message passing.
ergy per atom (F/A) error of 0.009 eV and a force MAF of 0.443 eV/Á in the Test split, improving force predictions by approximately 9% over Neural P'M and LSRM, with a perfect stability score of 100 ps. In the OOD split, GGND maintains robust performance with stability score of 99.89 ps, significantly surpassing QuinNet and Neural P'M. For HIO, GGND records an E/A of 0.005 eV and a force MAE of 0.179 eV/λ in the Test split, reducing force errors by over 40% compared to Neural P'M and QuinNet, and achieving perfect stability score of 100 ps. In the OOD split, its stability score of 97.93 ps notably exceeds baselines. GGND's equivariant diffusion process effec- tively captures all-pair interactions, ensuring insensitivity to conformational changes and enhancing stability in molecular dynamics simulations. The experimental results highlight the remarkable abil- ity of GGND to address geometric topological shifts, as evidenced by its consistent outperformance across the SiN and HfO datasets. The model's outstanding stability scores of 100 ps in both Test splits and near pear-perfect scores in OOD splits (99.892 ps for SiN and 97.928 ps for HfO) suggest that the equivariant diffusion process effectively captures all-pair information flows, making GGND insensitive to conformational changes.
Stability Visualization and Analysis. The stability metrie indicates the first time step at which the MD simulation becomes unstable. To better characterize the stability throughout the entire MD process, we visualize the max- imum bond length deviation in Figure In 100 ps MD simulations on the 3BPA dataset, GGND outperforms the MI.-based baselines by maintaining stability. Notably, although GGND exhibits instability around 30 ps, these unsta- ble states occur randomly, whereas both VisNet and MACE show persistent instability after a certain time step.
Stability on 1200 K. MD
Classical
MACE
Max bund longtth deviation (4)
VilNe
GOND
Averaged Max Deviation of MACE (12,3042..
Averaged Deviation VisNet (109 0293)
Averagal Mas Deviation CGGND (0:0700)
Stability Thechild
10
20000
40000
Ablation Study. To evaluate the impact of fully-connected diffusion in our proposed GGND model, GGND which uses local diffusion on the graph, and GGND1, which combines the base- line with fully-connected message passing. The ablation study on the 3BPA dataset (Table 3) demon- strates the superior generalization of GGND to unseen conformational domains 600 and 1200 K, while matching the baseline's 's optimal performance at 300 K (energy MAE: 0.002 eV, force MAE: 0.006 eV/Á, stability: 100 ps). In contrast, GGND;, limited by local diffusion, fails to generalize effectively, with performance close to the baseline (e.g., stability of 0.291 at 1200 K), as it cannot capture all-pair interactions. GGND; shows some generalization potential (e.g., stability of 2.892 ps 1200 K) but underperforms GGND due training challenges, highlighting the advantage of fully-com colle-conneeted diffusion in enabling robust, equivariant information flow for stable and accurate molecular dynamics simulations across diverse conformations.
Simulation time (fa)
curated two variants: Figure 4: Stability of MD Simulations on 3BPA.

---

Published as a conference paper at ICLR 2026
5 CONCLUSION
In this study, we investigate the stability of MD simulations and identify extrapolation to unseen con- formations key challenge. To address this, propose GGND, novel framework that improves the stability and generalizability of MD simulations by capturing geometrically invariant topological features through equivariant diffusion process. By mitigating geometric topological shifts arising from conformational variations, GGND reduces error accumulation, ensures robust energy and force predictions for unseen molecular conformations, leading stable molecular dynamics simulations. Our theoretical analysís establishes a regret bound under such shifts, providing formal guarantees of stability. Designed as a plug-and-play module, GGND integrates seamlessly with existing lo- cal equivariant message-passing networks, boosting out-of-domain performance while preserving in-domain accuracy. Comprehensive experiments on the 3BPA and SAMD23 datasets show that GGND surpasses baseline models in both accuracy and simulation stability,
ACKNOWLEDGMENTS
This work was supported in part by the Hong Kong Research Grants Council General Research Fund Under Ref. No 15208725, the Hong Kong Polytechnic University Internal Research Fund Under PO057774, the Research Grants Council of the Hong Kong SAR (Grant No. PolyU15215623, PolyU15229824, C5052-23G, and SRFS2526-5504) and the Hong Kong Polytechnic University (P0058445).
ETHICS STATEMENT
This study adheres to the ICLR Code of Ethics, with careful consideration of the ethical implications of our work, particularly its societal impacts, which are comprehensively addressed in Appendix Our methodology does not involve human participants, sensitive data, applications with signif- icant misuse potential. We have prioritized fairness and transparency in development of our models and findings, addressing potential biases in the dataset and model design in the referenced appendix. No conflicts of interest or funding concerns compromise the integrity of this research. The use of large language models is detailed in Appendix K.
REPRODUCIBILITY STATEMENT
Complete proofs for all theoretical claims are included in Appendices D and C. This study uti- lizes the publicly available datasets 3BPA and SAMD23, accessible at https: /pubs.u.. acs. org/doi/10 11102//ascs..o pttp//cccccc......oooo/e/ooococc 1:00647 and https: //ggttutubb///gg com/ SASSSATPPP (PPublic MLFF-Framewor
respectively. We adhere to the data splits specified in the publications as- sociated with these datasets, with all relevant parameters documented in Appendix E. The code is available at ht tps:: //tithub. com/ Maakkiilnng/eee GGND. These resources collectively enable full replication of our experiments and results.
REFERENCES
Yassine Abbahaddou, Fragkiskos D. Malliaros, Johannes F. Lutzeyer, Amine M. Aboussalah, and Michalis Vazirgiannis. Graph neural network generalization with gaussian mixture model based augmentation. Forty-second Intern International Conference on Machine Learning, 2025.
Eric Barth and Tamar Schlick. Overcoming Stability Limitations in Biomolecular Dynamies. Combining Force Splitting via Extrapolation with Langevin Dynamics in LN. The Journal of Chemical Physics, 109(5):6177-1622, 1998.
Ilyes Batatia, David P Kovacs, Gregor Simm, Christoph Ortner, and Gábor Csányi. MACE: Higher order equivariant message passing neural networks for fast and accurate force fields. Advances in Neural Information Processing Systems, 35:11423-11436, 2022.
Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P Mailoa, Mordechai Ko- mbluth, Nicola Molinari, Tess E Smidt, and Boris Kozinsky. E (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature Communications, 13(1): 2453, 2022.
10

---

Published as a conference paper at ICLR 2026
Chi Chen and Shyue Ping Ong. A Universal Graph Deep Learning Interatomic Potential for the Periodic Table. Nature Computational Science, 2(11):718-722 2022.
Fan RK Chung. Spectral Graph Theory, volume 92. American Mathematical Soc., 1997.
Taoyong Cui, Chenyu Tang, Dongzhan Zhou, Yuqiang Li, Xingao Gong, Wanli Ouyang, Mao Su, and Shufei Zhang. Online Test-Time Adaptation for Better Generalization of Interatomic Poten- tials to Out-of-Distribution Data. Nature Communications, 16(1))1891, 2025.
Bowen Deng, Peichen Zhong, KyuJung Jun, Janosh Riebesell, Kevin Han, Christopher J Bartel, and Gerbrand Ceder. CHGNet as Pretrained Universal Neural Network Potential for Charge- Informed Atomistic Modelling. Nature Machine Intelligence, 5(9):1031-1041, 2023.
Shaohua Fan, Xiao Wang, Chuan Shi, Peng Cui, and Bai Wang. Generalizing graph neural networks on iitt-f--diitrbbution graphs. IEEE Transactions on Pattern Analysis and Machine Intelligence, 46(1):322-331) 2024. doi: 10. 109/PPAMI.20233322100
Xiang Fu, Zhenghao Wu, Wujie Wang. Tian Xie, Sinan Keten, Rafael Tommi Jaakkola. Forces are not Enough: Benchmark and Critical Evaluation for Machine Learn- ing Force Fields with Molecular Simulations. Trunsactions on Machine Leurning Research, 2023. ISSN 2835-8856. Survey Certification.
and
Jiaqi Han, Wenbing Huang, Tingyang Xu, and Yu Rong. Equivariant graph hierarchy- -based neural networks. Advances in Neural Information Processing Systems, 35:9176-99 7177699177-1 2022.
Haokai Hong, Wanyu Lin, and Kay Chen Tan. Accelerating 3D Molecule Generation via Jointly Geometric Optimal Transport. In The Thirteenth International Conference on Learning Repre- sentations, 2025a.
Haokai Hong, Wanyu Lin, Ming Yang, and Kay Chen Tan. Distributional priors guided diffusion for generating 3d molecules in low data regimes. In The Fortieth AAAI Conference on Artificial Intelligence, 2025b.
Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant Dif- fusion for Molecule Generation in 3D. In International Conference Machine Learning, pp- 8867-8887. PMLR, 2022.
Geonu Kim, Byunggook Na, Gunhee Kim, Hyuntae Cho, Seungjin Kang, Hee Sun Lee, Saerom Choi, Heejae Kim, Seungwon Lee, and Yongdeok Kim. Benchmark of Machine Learning Force Fields for Semiconductor Simulations: Datasets, Metrics, and Comparative Analysis. Advances in Neural Information Processing Systems,
2023.
Diederik Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. In International Conference on Learning Representations (ICLR), San Diega, CA, USA, 2015.
Dávid Péter Kováes, Cas van der Oord, Jiri Kucera, Alice EA Allen, Daniel Cole, Christoph Ortner, and Gábor Csányi. Linear Atomic Cluster Expansion Force Fields for Organic Molecules: Beyond RMSE. Journal of Chemical Theory and Computation, 17(12):769-7771 2021.
Thomas J Lane, Gregory R Bowman, Kyle Beauchamp, Vincent A Voelz, and Vijay S Pande. Markov State Model Reveals Folding and Functional Dynamies in Ultra-Long MD Trajectories. Journal of the American Chemical Society, 133(45): 14411-18499, 2011.
Hongkang Li, Meng Wang, Tengfei Ma, Sijia Liu, ZAIXI ZHANG. and Pin-Yu Chen. What im- proves the generalization of graph transformers? a theoretical dive into the self-attention and positional encoding. In Forty-first International Conference on Machine Learning, 2024a.
Yunyang Li, Yusong Wang, Lin Huang, Han Yang, Xinran Wei, Jia Zhang, Tong Wang, Zun Wang, Bin Shao, and Tie-Yan Liu. Long-Short-Raud Range Message-Passingg A Physies-Informed Frame- work to Capture Non-Local Interaction for Scalable Molecular Dynamics Simulation. In The Twelfth International Conference on Learning Representations, 2024b.
Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant Graph Attention Transformer for 3D Atom- istic Graphs. In The Eleventh International Conference on Learning Representations, 2023.
11

---

Published as a conference paper at ICLR 2026
Yi-Lun Liao, Brandon M Wood, Abhishek Das, and Tess Smidt. EquiformerV2: Improved Equiv- ariant Transformer for Scaling to Higher-Degre, Representations. In The Twelfth International Conference on Learning Representations, 2024.
Yang Liu, Jiashun Cheng, Haihong Zhao, Tingyang Xu, Peilin Zhao, Fugee Tsung, Jia Li, and Yu Rong. SEGNO: Generalizing equivariant graph neural networks with physical inductive bi- ases. In The Twelfth International Conference on Learning Representations, 2024.
Georgi Medvedev. The Nonlinear Heat Equation on Dense Graphs and Graph Limits, SIAM
Journal on Mathematical Analysis, 46(4):2743-2766, 2014.
Yinglong Miao and Peter Ortoleva. Molecular Dynamics/Order Parameter Extrapolation for Bio- nanosystem Simulations. Journal of Computational Chemistry, 30(3):423-437, 2009.
Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron J Owen, Mordechai Ko- rnbluth, and Boris Kozinsky. Learning local equivariant representations for large-scale atomistic dynamics. Nature Communications, 14(1):579, 2023.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural net- works. In International Conference on Machine Learning, pp.
PMLR, 2021.
Shihao Shao, Haoran Geng, Zun Wang, and Qinghua Cui. FreeCG: Free the Design Space of Clebsch-Gordan Transform for Machine Learning Force Fields. In International Conference on Learning Representations, 2025.
Tom AB Snijders and Krzysztof Nowicki. Estimation and Prediction for Stochastic Blockmodels for Graphs with Latent Block Structure. Journal of Classification, 14(1):75-100) 1997.
Huayi Tang and Yong Liu. Towards Understanding Generalization of Graph Neural Networks. In International Conference on Machine Learning, pp. 33674-33719. PMLR, 2023.
Tong Wang, Xinheng He, Mingyu Li, Yatao Li, Ran Bi, Yusong Wang, Chaoran Cheng, Xiangzhen Shen, Jiawei Meng, He Zhang, et al. Ab initio characterization of protein molecular dynamics with ai2bmd. Nature, pp. 1-9, 2024a.
Yusong Wang, Chaoran Cheng, Shaoning Li, Yuxuan Ren, Bin Shao, Ge Liu, Pheng-Ann Heng, and Nanning Zheng. Neural 35m3:m: A long-range interaction modeling enhancer for geomet- ric GNNs. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024b. URL httpt::t///........ httpst:/
net / forum? id=ncqauwS11..
Yusong Wang, Tong Wang, Shaoning Li, Xinheng He, Mingyu Li, Zun Wang, Nanning Zheng, Bin Shao, and Tie- Yan Liu. Enhancing geometric representations for molecules with equivariant vector-scalar interactive message passing. Nature Communicattions, 15(1):313, 2024c.
Zun Wang, Guoqing Liu, Yichi Zhou, Tong Wang, and Bin Shao. Efficiently incorporating quintuple interactions into geometric deep learning force fields. Advances in Neural Information Processing Systems, 36:77043-77055. 2023.
Qitian Wu, Chenxiao Yang, Kaipeng Zeng, and Michael M. Bronstein. Supercharging graph trans- formers with advective diffusion. In Forty-second International Conference on Machine Learn ing, 2025.
Minkai Xu, Jiaqi Han, Aaron Lou, Jean Kossaifi, Arvind Ramanathan, Kamyar Azizzadenesheli, Jure Leskovec, Stefano Ermon, and Anima Anandkumar. Equivariant Graph Neural Operator for Modeling 3D Dynamics. Learning, ICML'24. JMLR.org. 2024.
Proceedings of the 41st International Conference on Machine
Chenxiao Yang, Qitian Wu, Jiahua Wang, and Junchi Yan. Graph neural networks are inherently good generalizers: Insights by bridging GNNs and MLPs. In The Eleventh International Confer- ence on Learning Representations, 2023.
Han Yang, Chenxi Hu, Yichi Zhou, Xixian Liu, Yu Shi, Jielan Li, Guanzhi Li, Zekun Chen, Shuizhou Chen, Claudio Zeni, et al. Mattersim: A Deep Learning Atomistic Model Across Elements, Temperatures and Pressures. arXiv preprint arXiv:2405.049m 2024.
12

---

Published as a conference paper at ICLR 2026
APPENDIX
A RELATED WORKS
A.1 MACHINE LEARNING MOLECULAR DYNAMICS SIMULATION AND ITS GENERALIZATION
Machine-learning for molecular dynamics simulation, including ML interatomic potentials and ML force fields, has emerged as an accurate and computationally efficient surrogate for quantum me- chanical caleulations in MD simulations and related atomistic modeling tasks. However, the gener- alization of MI. methods for MD simulations remains unaddressed.
A major advance in improving the generalization of these models has been the development of "uni- versal" or pretrained interatomic potentials trained on chemically diverse and large-scale datasets. Models such M3GNet (Chen & Ong, 2022) and CHGNet (Deng al., 2023) show that broad coverage of elements, bonding motifs, and structures can substantially improve transferability and reduce retraining costs. Nevertheless, these approaches often falter when faced with distribution shifts outside the training domain, such as conformations sampled under different thermodynamic conditions.
Another complementary line of work focuses on architectural inductive biases, especially equivari- ant geometric graph neural networks that explicitly encode physical symmetries. Equivariant mes- sage passing architectures, including recent efficient implementations (Batatia et al., 2022; Wang et al., 2024c), typically deliver higher accuracy and better sample efficiency than non-equivariant baselines. While these models interpolate well within the training regime, symmetry constraints alone are insufficient to guarantee robustness under changes the conformational distribution.
Alongside architectural innovations, there is growing recognition that standard energy and force test errors can be misleading as proxies for MD stability. Benchmarking studies such as (Fu et al., 2023) emphasize trajectory-leve evaluations, including long-term stability, conservation laws, and reproduction of thermodynamic observables. These evaluations often reveal substantial degradation under nut-of-distributin (OOD) conditions even when pointwise prediction errors remain low.
To address such failures, some works integrate active learning and on-the-fly adaptation into MLFF/MLIP workflows. Active learning pipelines and uncertainty-awant simulation controllers selectively query new data in high-uncertainty regions (Yang et al., 2024), while recent methods such as TAIP (Cui et al., 2025) perform test-time adaptation to reduce the impact of train-test distri- bution gaps. These strategies mitigate specific failure modes but can be computationally demanding and lack systematic mechanisms detect and react to distribution shifts in a theoretical manner.
Our work builds on these threads by proposing geometric GNN that detects and responds to distri- bution shifts in conformational space, aiming improve stability and accuracy when generalizing across thermodynamie regimes.
A.2 GENERALIZABLEE GRAPH NEURAL NETWORK
Recent advancements in graph neural network (GNN) generalization have progressed along four complementary dimensions: theoretical foundations, architectural innovations, training methodolo- gies, and data-centric strategies. Theoretical studies have established sample-complexity and stabil- ity bounds, elucidating structural factors- such as propagation depth and graph connectivity- that influence generalization and inform design decisions (Tang & Liu, 2023; Yang et al., 2023). Ar- chitectural advancements demonstrate that incorporating inductive biases or attention mechanisms enhances generalization performance (Liu et al., 2024; Wu et al., 2025; Li et al., 2024a). Practical training and data-centric approaches translate theoretical insights into practice through causal learn- ing and data augmentation techniques, effectively mitigating empirical generalization gaps (Abba- haddou et al., 2025; Fan et al., 2024). Collectively, these efforts outline coherent research agenda: leverage theoretical insights to identify generalization bottlenecks, design architectures with induc- tive biases to reduce sample complexity, and employ training and data strategies, alongside out-of- distribution (OOD)-awart mechanisms, ensure robust generalization.
13

---

Published as a conference paper at ICLR 2026
However, existing research GNN generalization has largely overlooked the challenge of confor- mational distribution drift in MLFF or MD. Specifically, environmental factors, such as temperature, induce drifts in the geometric structure distribution of molecular graphs, which subsequently al- ter their topological configurations.
A.3 EQUIVARIANT GRAPII NEURAL NETWORK
Physical systems obey symmetry principles that constrain valid model behavior. For molecular modeling, the governing laws are invariant or equivariant under the Euclidean group SE(3). in- cluding rotations, translations, and reflections. Neural networks that respect these symmetries have demonstrated substantial improvements in data efficiency, physical consistency, and generalization for atomistic modeling.
Equivariant graph neural networks enforce transformation consistency by ensuring that node rep- resentations transform according to group representations when input coordinates are transformed. These approaches can be divided into incorporating symmetry through invariant features derived from interatomic distances (Satorras al., 2021; Xu al., 2024) and propagating higher-order ge- ometric tensors and directional information through equivariant message passing (Liao Smidt, 2023; Batatia et al., 2022; Batzner et al., 2022). These architectures enable accurate modeling of energies and forces while preserving physically meaningful coordinate dependencies. Equivariant graph neural networks has been validated effective for range molecular modeling tasks, includ- ing force field prediction (Batatia et al., 2022; Batzner et al., 2022; Wang et al., 2024c), molecular dynamics simulation (Xu al., 2024; Han et al., 2022), and molecular design (Hoogeboom al., 2022: Hong et al., 2025a;b).
In molecular dynamics simulation, equivariance is particularly critical because force predictions must transform consistently with spatial rotations to ensure stable trajectory integration. Conse- quently, modern geometric deep learning approaches for interatomic potentials and force fields widely adopt equivariant representations as fundamental inductive bias. However, symmetry preservation alone does not guarantee robustness under distributional shifts in geometric topology. Our work builds upon equivariant message passing by introducing global diffusion mechanism that maintains equivariance while improving robustness to conformational variation.
B DIFFUSION PROCESSES ON GRAPHS
This section provides concise background diffusion processes defined graphs. The presen- tation clarifies how Equation can be interpreted as a graph partial differential equation.
B.1 DIFFUSION AS FEATURE PROPAGATION
Let G = (V,E) be graph with = [V] nodes. A time-dependent node feature field is denoted by
Z(t) [2] (t)........ ... z, (t)] Rxd
where (t) represents the feature of node i at time 2 0.
Graph diffusion models describe the evolution of node features through a continuous-time process
az(t) C(Z()),
ai
where L is graph differential operator that governs information propagation across nodes. This formulation generalizes discrete message passing by replacing layer-wise updates with a continuous dynamical system.
B.2 GRADIENT AND DIVERGENCE GRAPHS
Diffusion on graphs can be derived from discrete analogues of gradient and divergence operators.
14

---

Published as a conference paper at ICLR 2026
Graph gradient. Given node features Z Rexd the gradient operator maps node features to edge features:
(VZ)
Zj Zi+
for node pairs (i.j). This operator measures feature variation across the graph topology.
Graph divergence. The divergence operator aggregates edge features back to nodes:
(div F); -E F..
where F,j are edge features. It serves as the adjoint of the gradient operator and controls how edge-level information influences node states.
B.3 DIFFUSION EQUATION ON GRAPHS
Combining gradient and divergence yields the graph diffusion equation
@z(() div (S(t) O VZ(()),
where S()) R=x diffusivity operator that modulates information flow between node pairs and @ denotes element-wise interaction between edge fields.
When S(t) symmetric and row-stochastic, the diffusion dynamics can be equivalently written
OZ(t) (S(() ())
Ōt
which corresponds to continuous-tin random walk on the graph. This formulation shows that diffusion performs feature smoothing while preserving global consistency across nodes.
C PROOF EQUIVARIANCE GGND
We prove that the GGND process, defined by the PDE in Equation ( equution and its rewritten form in Equation ( equation 4), is 3E(3)-qquivariant Specifically, the learned features Z(T) trans- form correctly under SE(3) transformations (rotations and translations) applied to the input positions X and invariant scalar features H.
SE(3)--quivariance means that if we apply a transformation g SE(3) the inputs, the output features transform accordingly:
Z'(T) = (Tg)Z(T),
where ZT) is the solution of the PDE for the transformed inputs X' = gX and H' (since H invariant scalars), and (gg) denotes the group representation acting the spherical tensor features (irreducible representations of O(3), labeled by L and M).
Translations handled trivially because the model depends only relative positions |x,a|| (in- variant) and unit vectors XŞ; (equivariant under rotations but invariant under translations). The initial embedding de (X, H) uses radial basis functions (RBFs) distances, which are translation- invariant, Thus, the dynamics preserve translation invariance.
We focus on rotation equivariance under g E SO(3). The representation D4(g) acts on each irrep component as:
Σ =E
Z,,LMM
where the action the same for all nodes
Assume the initial condition is equivariant: Z'( (0)) We need to show that if Z'(l)
D(g)Z(I) at time I, then the time derivative preserves this property:
OZ'(t) = D(g)
(t
Öt
0t
15

---

Published as a conference paper at ICLR 2026
This requires showing that the Right-hand side operator f(Z, div (S( (,), VZ(t)) (S [t) is equivariant:
(Z) (())
We prove this by showing that each component-thhe gradient VZ, the diffusivity S, and their combination-ss equivariant.
Equivariance of the Gradient Operator.
The equivariant gradient defined
(VZ)ij,klam
Σ
kita (Z 3,l2ma "1,kl2m22/
where W are learnable scalar weights (invariant under rotations), and the gradient operates per irrep L la and channel, without changing L.
For the transformed features and positions:
2
D ED'a D' (2) k2gkkmm kr,
3. 3lkaama
m1
Thus,
(Vz)) )ij.hamm?
Σ W kkiz
Σ D migm (2)tam, 2t,ktam's)
Σ
D'a myms (g)
(V)
mi
since W is invariant. The gradient transforms the same irrep, so V is equivariant: VZ (VZ).
Equivariance of the Diffusivity.
The diffusivity (attention matrix) is:
(i.) Jkames
Σ
Clama Comidam Rad,tats (I)
11.42.m1.mng
where C are Clebsch-Gordan coefficients (invariant), R is a learnable radial function (depends invariant distance [|x,;[],, Yth are spherical harmonics, and @ is an equivariant pairwise interaction (e.g., gated tensor product).
-|x/fell = ||,را|| (invariant), Xxx = 9Xحو. and assuming the conven- tion where r' -p Ddr (z) (Rjx) (equivariant filter in NequlP and e3nn), 6124, 2.)m = Ep, D' (g)Ø(2, Z3 Maps (by assumption that á is equivariant).
Under transformation:
Substituting:
Sii, S Jkkgmm
Σ
Clama lamidama
R((xxjil)
> D (0)x1 (85ri)
Σ Día (g)0taps
This is:
Σ
Σ lylgmima
Clams Oamiltam RY dieps Di D'i (') (g).
Since the tensor product representation is D @ D and the CG decomposition to la commutes with the group action (CG coefficients invariant and define equivariant basis change), the entire expression transforms as the output irrep:
S' S'ii.jkkgmm
Σ m/
D'a mamg (o)
S(Z', Thus, S(Z', X',)) = -
in the sense of the adjoint action on linear maps, but since treat S as producing equivariant filters, the composition preserves equivariance.
16

---

Published as a conference paper at ICLR 2026
Equivariance of the PDE Operator. The operator is f (S I)Z, which corresponds to Σ, Sij @ (z) z;) (assuming @ is an equivariant bilinear operation, such as channel-wist multi- plication or tensor contraction, consistent with the right-stochastic property).
Since VZ' VZ and S' D(g)SD(g)~ (adjoint for maps), the product S' VZ
(S o VZ), because:
D(g)SD(9)~11 o ((g)(VZ) (S + VZ),
assuming @ commutes with (g) (as does for tensor products or contractions in irrep bases).
The divergence div is sum over j. which is permutation-invarian and commutes with (g):
div VZ Σ (Ss) (zj z:)) P(g) div(S @ VZ].
Thus, f(Z',X') (Z,)
Since the operator is equivariant, the solution the PDE (e.g., via numerical integration like Euler steps) preserves equivariance: Z' (t) = ((g)Z(t for all t, including = T.
Finally, the energy prediction uses the invariant components k,00 (T) (L 0, scalar invariants), which are unchanged under (g)), ensuring the total energy EGGND is SE(3)-invariant.
D PROOFS OF REGRET BOUND
Proof of Equation (8). Assume that the loss function é and the decoder óasc Lipschitz continu-
ous with constants L2 and L2, respectively. According to the topology formation hypothesis in Sec- tion 2, decompose the joint distribution p(X, H, A, M) p(X, H E, M M)A, X, H, E, M). Since X f(U; W) and H = h(U; W') are independent of E and M, it follows that p(X, H Eu,, Mr) p(X, Ere, Mu)..
The extrapolation gap is:
IT Exe, Mcr) (To:o; Eur, Mu) S [C(T:; Ere, Ma) - L(To: Eir, Mt)
[[Te; Eur, Me) Cu oT:: Eur Mu)]
The second term is bounded by Din Erc, Mrr, No.) with probability at least by standard extrapolation bounds (Equation (7)).
For the first term:
L(To: Ere, Mir) L(To; Eir, M..) M. ((To H H'))), Y')]
(10)
Since the marginals over X, H are identical:
= Ex.
(II)
By the triangle inequality and Lipschitz continuity of
< (EA
(12)
+ +E,, (e(Tol) A), -(T,,X H),,
where A,Y Y p(- Eu, Ma)) and A', A',Y p(- E Mu.). Applying Lipschitz constants:
< L1 - ZT; A') Z(T; L2 E 1|22.
yielding the decomposition after rescaling constants into ((-).
17

---

Published as a conference paper at ICLR 2026
Proof of Theorem 3.1. The equivariant graph neural diffusion model assumes full connectivity, with diffusivity S(t) defined based on positions X and features (tt) (Equation (3)), independent of the adjacency A. Thus, the solution (TT) to Equation (4) does not depend on A, implying [2(T; A') Z(T; A) O.
Since 0 (((())A))22) for any arbitrary polynomial function v. the variation magnitude is reduced to any order. The injectivity / and ensures that latent variables U be recovered from X and H. enabling the equivariant attention in S(t) to capture latent interactions from the graphon W, independent of shifts in A.
Proof of Corollary 3.2. into the DM-
The conclusion follows directly by substituting the result of Equation (8)
E PARAMETER
The training details are outlined below, with sataset-specifie parameters provided in Table 4.
Table 4: Dataset Information and Dataset-specific Parameters
Dataset
3BPA
SAMD23
Molecule
3BPA
SiN
HfO
Atoms
27
16-510
96
Batch size
2
Epochs
1,000
200
200
Training Time (h)
30
24
Training
1. Optimizer: Adam (Kingma & Ba, 2015) optimizer is used with constant learning rate of 10-4 as our default training configuration.
2. GPU: NVIDIA GeForce RTX 3090
3. CPU: Intel(R) Xeon(R) Platinum 8338C CPU
4. Memory: 512 GB
For baselines, we adopt the recommended parameters from their original publications. Specifically, for the message-passinngg passing layer in implementations with GGND, include additional layer for geometric graph neural diffusion. To ensure a fair comparison with similar model capacity, we reduce message-passing layer in the baselines during integration.
F COMPUTATIONAL OVERHEAD ANALYSIS AND COMPARISON
Table 5 summarizes the computational and memory overhead incurred when integrating GGND into ViSNet on the 3BPA dataset. All models are trained using identical hyperparameters: ViSNet in- cludes an additional 500 training epochs to provide more comprehensive point of comparison. Ex- periments were conducted on an NVIDIA GeForce RTX 3090 GPU paired with an Intel(R) Xeon(R) Platinum 8338C CPU.
Overall, incorporating GGND increases training time by 26.54% and MD simulation time by 15.57%. GPU memory consumption increases by 15.28% during training and 14.87% during in- ference. Considering the substantial gains in energy/force accuracy and stability, these additional computational and memory costs are acceptable in practice.
G DATASETS
The 3BPA dataset consists of configurations the flexible, diug-like molecule 3-(benzyloxy) pyridin- -riidin---amine. Initial configurations were generated from short (0.5 ps) molecular dynamics sim-
18

---

Published as a conference paper at ICLR 2026
Table Computational and Memory Overhead Introduced by GGND
Metrics
VisNet
VisNer
+GGND
Training Time (h)
6.822
10.326
8.633
Time
Inference Time (s)
13.872
13.672
16.032
MD Time for 100 ps (h)
1.958
1.949
2.282
Training Memory
20.457
20.455
23.582
Memory
Inference Memory (GiB)
14.125
14.248
16.225
ViSNet* includes an additional 500 training epochs.
300K
600K
1200K
Figure 5: Representative 3BPA Molecular Configurations Sampled from MD at 300, 600, and 1200 K.
ulations using the ANI-Ix force field bias sampling toward lower-energy regions of the potential energy surface. In addition, longer 25 ps MD simulations were performed at three temperatures 300, 600, and 1200 K-using a Langevin thermostat with fs time step. A selection of these configurations is visualized in Figure 5.
H MD SIMULATION UNDER NVE AND NVT ENSEMBLES
We perform constant-energy (NVE) and constant- tenttaatttemppratuure (NVT) molecular dynamics simula- tions at 1200 K on the 3BPA dataset. We use Selocity-Verlee integration for 200 ps with 1 fs timestep for the NVE ensembles. We perform Langevin dynamics at a temperature of 1200K, timestep of 1.0 fs, and friction coefficient of 0.01 fs for 200,000 steps, corresponding to 200 ps for NVT ensemble.
In both the NVE and NVT molecular dynamics simulations, GGND demonstrates markedly superior stability compared to MACE and VisNet (Figures and 7). GGND achieves the lowest Averaged Max Bond Length Deviation across the full 200 ps trajectories and maintains stable geometries, with only a few isolated instances where the Max Bond Length Deviation slightly exceeds the threshold. In contrast, MACE begins violate the bond-length threshold at approximately 13 ps in the NVE
19

---

Published as a conference paper at ICLR 2026
Stability 1200 K MD under NVE Ensemble
MACE
Max bund lengih deviation (A)
10
GOND
Avergee
Averaged Max Deviation ViaNa (297.8551)
Avcraged Devistin (00530)
Suaiiit Threshold
10
167
25000
50000
75000
100000
125000
150008
175000
200000
Simalation time (B)
Figure 6: Stability of MD Simulations under NVE Ensemble on 3BPA.
Stubility on 1200 MD under NVT ensemble
MACE
ViANe
Max bondd lenghs devistion (A)
GAND
Avergeed Devistion MACE (214 7088)
Avenged Deviation VisNer (218. (218.1154)
Averaged Max Deviaton of GGND (000701)
Stabiiit T'reshold
10
10
25000
30000
100000
150000
175000
Simulation time (B)
Figure 7: Stability of MD Simulations under NVT Ensemble on 3BPA.
ensemble and around 27 ps in the NVT ensemble, after which the deviations steadily grow. Vis- Net performs even worse, showing threshold-breaking behavior almost immediately and displaying rapidly increasing deviations throughout the simulation. Overall, these results highlight the robust- ness and stability of GGND in unseen temperature MD settings.
I LIMITATIONS
Although the GGND framework demonstrates strong performance in enhancing stability and accu- racy for small to medium-sized molecular systems, such as those in the 3BPA and SAMD23 datasets with up to 510 atoms, its scalability larger biomolecular systems like proteins comprising millions of atoms remains a critical limitation. The method's reliance a fully-connected graph for the diffusion process, which facilitates all-pair information flows, introduces quadratic computational complexity in both time and memory with respect to the number of atoms, rendering it impractical for real-world applications involving extensive simulations. Future developments could incorpo- rate sparse approximations or hierarchical diffusion mechanisms to mitigate these issues and extend GGND's s utility to large-scale protein dynamics.
J IMPACT STATEMENTS
This paper presents research aimed at advancing Artificial Intelligence (AI) applications in scientific domains, including materials science, chemistry, and biology. The insights and expertise gained will significantly enhance AI technologies, accelerating the process of scientific discovery.
20

---

Published as a conference paper at ICLR 2026
Machine learning for molecular dynamics enables rapid molecular analysis. However, the potential for misuse and unintended consequences underscores the need for stringent ethical guidelines, robust regulations, and responsible deployment to safeguard individuals and society from harm.
K THE USE LARGE LANGUAGE MODELS
The core method development and research ideation in this paper conducted independently of LLMs, and LLMs did not contribute to any original or non-standard components of the work. The authors utilized LLMs solely general-purpont assist tool for checking grammar and improving the clarity of the manuscript, as well as for aiding in the comprehension of existing literature. All content in this submission, including any text refined with LLM assistance, has been thoroughly reviewed by the authors, who take full responsibility for its accuracy, integrity, and compliance with ethical standards. No LLMs considered contributors eligible for authorship.
21


Reviews:
[{"review_id": "GIipZxZUL7", "reviewer": "Reviewer_sTeQ", "summary": "The manuscript proposes Geometric Graph Neural Diffusion (GGND), a plug-in SE(3)-equivariant module that performs fully connected, diffusion-style updates on atom features and is then fused with a local equivariant GNN backbone. The central claim is that this global diffusion makes node representations insensitive to \"geometric topological shifts\" caused by temperature-dependent conformations and cutoff-dependent neighbor graphs, thereby improving out-of-distribution (OOD) stability in molecular dynamics. Formally, the model evolves features by a PDE-like update, $\\partial_t Z(t)=(S(Z, X, t)-I) Z(t)$, and the paper presents a regret-style bound to argue reduced OOD error under adjacency shifts. Experiments on 3BPA (train at 300K, test at 600/1200K and dihedral slices) and on the SAMD23 SiN/HfO benchmark report notable gains in both force/energy MAE and a stability metric based on long-rollout NVE simulations up to 100 ps . The manuscript positions GGND as a general \"plug-and-play\" layer to retrofit existing local EGNNs.", "questions": "I would like to see precise clarifications to the following questions. First, is $S$ ever computed from a sparse $A$ in training or inference, or is the operator strictly complete-graph at all times? Please fix the contradiction between §3.1 and the proof of Theorem 3.1, and state the exact dependence on $A$ in the method section. Second, please correct units in every table and paragraph and re-evaluate claims like \"below chemical accuracy\"; the present mix of $\\mathrm{eV} / \\mathrm{meV}$ is not acceptable. Third, what happens beyond 100 ps and under thermostats? Showing stability curves up to $\\geq 1 \\mathrm{~ns}$ and including NVT/NPT would reduce concerns about integrator-specific artefacts. Fourth, why is there no comparison to Allegro and EquiformerV2 as strong global baselines? A single paragraph in Related Work is not sufficient; the comparison needs to be empirical on the same splits, with matched cutoffs, time steps, and neighbor policies. Fifth, can the authors release scripts for the stability metric (bond- and RDF-based) and the exact MD settings (neighbor list rebuilds, cutoffs, barostat/thermostat where applicable), since small choices there can create large apparent gaps? Finally, is there a way to incorporate environment conditioning, something like GG-ODE, on top of GGND to see whether invariance and conditioning are complementary rather than alternatives?", "limitations": "", "rating": 4, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "From the angle of originality, the paper makes a concrete architectural move: it injects an equivariant, all‑pairs diffusion operator into otherwise local, cutoff‑dependent message passing, with spherical harmonics and Clebsch–Gordan structure to preserve symmetry. This is not only technically consistent with modern E(3)/SE(3) practice but is also aligned with a well‑motivated failure mode: changing neighbor graphs across temperatures or cutoffs degrade MLFFs even when in‑distribution errors look small. The formalization of “geometric topological shift” and the decomposition of the OOD gap give the work a clearer problem statement than usual in this area. The empirical section goes beyond pointwise MAE to include trajectory stability under NVE, which, in the opinion of the reviewer, is the right stress test to separate superficially accurate models from those that do not explode in practice.", "weaknesses": "The theoretical core reads fragile. In §3.1 the diffusivity $S(Z, X, t)$ is described as dependent on the adjacency $A$, which matches the paper's narrative about topology shift; however, the proof of Theorem 3.1 later asserts the model \"assumes full connectivity\" and that $S$ depends only on positions/features and thus is independent of $A$, hence $Z(T)$ does not depend on $A$ at all. Under that assumption, the main bound collapses to a trivial statement. This contradiction must be resolved unambiguously, because the claimed robustness hinges on the exact role of $A$. \n\nThe reporting has unit inconsistencies that undermine trust. Table 1 explicitly states energy MAE in eV, yet the surrounding text speaks of meV for the same numbers and conclusions (e.g., improvements claimed “below chemical accuracy”). This is a three‑orders‑of‑magnitude discrepancy that changes the physical interpretation and must be corrected before one can judge the effect size.\n\nThe experimental design does not yet establish that the new diffusion mechanism is the decisive factor. The ablation includes a “fully‑connected message passing” variant, but there is no head‑to‑head against strong global equivariant baselines that already propagate information across all pairs, such as Allegro (strictly local but designed to capture many‑body correlations without message passing) and modern equivariant Transformers like EquiformerV2. Without those comparisons, one cannot attribute the OOD gains specifically to the PDE‑style update rather than to the mere presence of global interactions. A fair suite here would at least include NequIP and MACE (already present) plus Allegro and EquiformerV2, all tuned for the same cutoff and time step.\n\nThe evaluation scope feels narrow for the strength of the claims. Stability is reported up to 100 ps in NVE with Velocity‑Verlet and a 1 fs step; that is useful but short. The paper should probe longer horizons, and also include NVT (Nosé‑Hoover and Langevin) to show the method is not sensitive to ensemble or integrator/thermostat coupling. The “Forces are not Enough” paper makes exactly this point about trajectory‑level validation; aligning with that protocol would improve credibility.\n\nRelated work is incomplete in a way that directly affects positioning. GG‑ODE (KDD’23) explicitly addresses cross‑environment dynamics by learning a shared graph‑neural ODE with environment‑specific latent exogenous factors (e.g., temperature), a complementary strategy to the paper’s “invariance to topology change.” It is not discussed or compared, and the reader cannot see whether conditioning on environment would reduce or even remove the need for diffusion on a fully connected graph.\n\nFinally, scalability is acknowledged as $O\\left(N^2\\right)$ but not quantified. If the core is all-pairs diffusion, memory and wall-clock scaling should be reported on realistic MD supercells, not only 27-atom 3BPA or 96-atom HfO cells. A convincing story would include sparsification or hierarchical diffusion and a clear regime where GGND remains practical.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YSRI30ZJxx", "reviewer": "Reviewer_eXFQ", "summary": "The paper proposes Geometric Graph Neural Diffusion (GGND), a module augments local SE(3)-equivariant GNNs with a global graph diffusion process to propagate information across all atom pairs while preserving equivariance.  GGND introduces equivariant gradient and equivariant diffusivity (attention) operators, aiming to learn features that are invariant to geometric–topological shifts to improve long-horizon MD stability beyond in-distribution accuracy. Experiments on 3BPA and SAMD23 report gains in energy/force MAE and higher stability in OOD scenarios.", "questions": "What is the compute and memory increase when adding GGND to GNN backbone?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "*  GGND is a plug-in-play diffusion module that can integrate with common equivariant GNN backbones, preserves equivariance, and facilitates global information mixing.\n\n* On 3BPA across 600–1200 K, GGND increases stable rollout length by orders of magnitude versus baselines while simultaneously reducing energy/force MAE. On SAMD23 (SiN/HfO), indicating robustness well beyond in-distribution regimes.", "weaknesses": "Disclaimer: The theory analysis of the paper is beyond my expertise, so I’m open to the authors’ clarifications and to other reviewers’ perspectives.\n\nThe theorem 3.1 justifies how geometric graph neural diffusion is able to reduce the representation variation of the graph neural network, but I'm generally curious why (and when) graph neural diffusion is superior over  equivariant attention models (e.g., Equiformer, TorchMD-Net) once those baselines are given sufficient effective depth/receptive field. A discussion or an ablation study would help clarify the conditions under which GGND is actually superior.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "12FTCkdMgY", "reviewer": "Reviewer_WP7J", "summary": "This paper addresses a critical challenge for Geometric Graph Neural Networks (Geo-GNNs) in molecular dynamics (MD) simulations: minor prediction errors can accumulate and destabilize long-term MD trajectories. The authors attribute this problem to the model's poor ability to extrapolate to unseen conformations, particularly those arising from environmental changes like different temperatures, a phenomenon they formalize as \"Geometric Topological Shifts\".\n\nTo tackle this, the paper proposes a novel framework called Geometric Graph Neural Diffusion (GGND). The core idea of GGND is to serve as a \"plug-and-play\" module that integrates with existing local equivariant message-passing frameworks. It achieves this by performing an equivariant diffusion process on a fully-connected graph, enabling instantaneous information flow between arbitrary atomic pairs. This process utilizes novel \"equivariant gradient\" and \"equivariant diffusivity\" operators, designed to capture geometrically invariant topological features, thereby mitigating error accumulation and ensuring stable MD simulations.", "questions": "1. Scalability of GGND:\n\nWhile the proposed method shows significant improvements in smaller molecular systems, can GGND scale efficiently to larger molecules or systems with a higher number of atoms? What are the computational bottlenecks, and how might they be addressed?\n\n2. Comparison to Other Approaches:\n\nHow does GGND compare with other domain adaptation methods, such as active learning or test-time adaptation, in terms of improving stability and extrapolation? Have these methods been considered as part of your baseline comparisons?\n\n3. Model Limitations:\n\nThe theoretical analysis shows GGND's ability to manage geometric topological shifts. Are there any specific types of conformational shifts or molecular behaviors that GGND might still struggle with, even with its diffusion mechanism?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "### Novel and Effective Method: \nThe GGND framework is a novel contribution. Its core design—applying equivariant diffusion on a fully-connected graph —is intuitively sound. It overcomes the limitations of local message passing during topological changes by capturing \"all-pair information flows\". The ablation study  effectively demonstrates the superiority of \"fully-connected diffusion\" over \"local diffusion\" and \"fully-connected message passing,\" validating the design's efficacy.\n\n### Sufficient Experimental Validation:\n- The experiment on the 3BPA dataset is well-designed, using different temperatures (300K, 600K, 1200K) to simulate conformational shifts. The results are impressive; for example, at 1200K, VisNet's stability is boosted from 0.004 ps to 11.209 ps, and MACE's stability also sees a 15-fold improvement.\n\n- The comparison against multiple SOTA models on the SAMD23 dataset shows that GGND not only enhances stability but is also highly competitive in force and energy prediction accuracy on OOD splits.\n\n- Solid Theoretical Support: The authors provide a theoretical analysis for GGND, including a proof of SE(3) equivariance and a regret bound under geometric topological shifts , which adds to the work's rigor.\n\n### Strong Practicality: \nGGND is designed as a plug-and-play module that can be seamlessly integrated into various existing equivariant GNN architectures. This significantly increases the method's potential applicability and impact.", "weaknesses": "### Scalability: \nThis is the paper's most significant limitation. The method's reliance on diffusion over a fully-connected graph incurs $O(N^2)$ computational and memory complexity (where N is the number of atoms). The authors acknowledge this makes the method \"impractical\" for large biomolecular systems. While the SAMD23 dataset reaches up to 510 atoms, this is still orders of magnitude smaller than many common systems in MD simulations\n\n### Justification for the Diffusion Framework: \nThe ablation study shows that fully-connected diffusion (GGND) outperforms a fully-connected message passing baseline (GGND‡). The paper attributes the baseline's poorer performance to \"training challenges\". This explanation is somewhat brief and does not fully explore why the PDE-based diffusion process is inherently superior to a standard (but also fully-connected) message-passing layer for this task.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UqFldXb885", "reviewer": "Reviewer_Xkok", "summary": "The paper addresses a critical instability issue in Geometric Graph Neural Networks (Geo-GNNs) used for molecular dynamics (MD) simulations. It identifies the cause as \"geometric topological shifts,\" where models trained on certain molecular conformations (e.g., at 300 K) fail to extrapolate to unseen conformations (e.g., at 1200 K).\n\nThe authors propose Geometric Graph Neural Diffusion (GGND), a plug-and-play module that runs an equivariant diffusion process on a fully-connected graph. This allows the model to capture global, all-pair atomic information, making it robust to the local topological changes that cause standard models to fail. The paper provides a theoretical regret bound to justify this improved robustness and shows strong empirical results, with GGND massively improving simulation stability on the 3BPA and SAMD23 datasets.", "questions": "**Scalability**: Given the $O(N^2)$ bottleneck, have you explored approximations? For example, what is the performance trade-off if you use a large-radius cutoff (e.g., 20-30 Å) instead of a fully-connected graph?\n\n**Computational Cost**: What is the practical wall-clock time overhead for a simulation step when adding the GGND module to a baseline like VisNet?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**Significance**: It tackles a major practical bottleneck for ML force fields: long-term simulation stability. Solving this OOD extrapolation problem is critical for reliable MD simulations.\n\n**Originality**: The solution is novel. Augmenting local message passing with a parallel, global equivariant diffusion process is a creative and technically sound approach to capturing physics robustly.\n\n**Quality**: The paper is high-quality. The theoretical analysis provides a formal justification for the method's robustness. The experiments are excellent, using temperature variations (300 K vs 1200 K) to create a perfect test for the hypothesis, and the results are dramatic and convincing.\n\n**Clarity**: The paper is well-written, clearly motivating the problem as a \"geometric topological shift\" and explaining the GGND solution .", "weaknesses": "**Scalability**: This is the most significant weakness. The method's reliance on a fully-connected graph introduces $O(N^2)$ computational and memory complexity, which the authors admit is \"impractical\" for the large-scale systems (e.g., proteins) where this stability is most needed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YHu7mmfvpG", "reviewer": "Reviewer_Ev5w", "summary": "The study aim to address a critical limitation of existing Geometric Graph Neural Networks (Geo-GNNs) in molecular dynamics (MD) simulations: minor force prediction errors can destabilize long-temporal trajectories, especially when extrapolating to unseen molecular conformations (e.g., different temperatures) due to limited training data coverage. To solve this, the authors propose GGND, a plug-and-play framework that captures geometrically invariant topological features via an equivariant diffusion process on fully connected graphs, mitigating error accumulation and ensuring stable MD simulations. Experiments were conducted on the 3BPA and SAMD23 datasetsshow that GGND outperforms baselines in both accuracy and stability.", "questions": "The authors should seriously address the concerns shown in Weakness point by point to improve the quality of the paper.", "limitations": "", "rating": 2, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The paper provides solid theoretical support: it formalizes the causal mechanism of geometric topology formation (linking environment, cutoff radius, and latent graphons) and derives a regret bound for extrapolation error under topological shifts. Importantly, it proves GGND’s SE(3) equivariance via detailed analysis of gradient/diffusivity operators\n2.  It compares with 4 representative baselines (NequIP, MACE, SEGNO, VisNet) and shows its usefulness in different models.", "weaknesses": "1. Scalability Limitation for Large Biomolecular Systems​\nGGND’s fully connected graph diffusion introduces quadratic computational complexity (O(N²) for N atoms), making it impractical for large systems like proteins (millions of atoms) or complex biomolecules. The paper only tests systems up to 510 atoms (SAMD23 HfO), and its memory/time costs will grow exponentially with larger N. This restricts its real-world applicability to small/medium molecules, as large-scale MD simulations (e.g., protein folding) are common in biology and drug discovery.\n2. Incomplete stability analysis: While bond length and RDF are used to measure stability, the paper does not evaluate other critical MD metrics (e.g., energy conservation, pressure fluctuations, or thermodynamic observables like melting points), which are essential for validating real-world utility.\n3. Unclear relationship between diffusion operation and geometric topology. It lacks illustration how the diffusion operations map diffusion structural topology.\n4. Lack strcuture analysis. Simulating a biomolecule at 1200K is WEIRD and meanningless. You can get nothing from the unrealistic experiment. The author should carefully check the structure changes during the simulation or perform another one under a normal environment.\n5. lack interpretability: while it mentions attention weights reflect interaction strength, no visualizations are provided to illustrate how it captures unobserved topological features.\n6. Time consumption\nAgain, the diffusion module seems ultra time-consuming. The authors should seriously analyze and demonstrate the extra huge computational cost during model training compared to the baseline model, e.g., ViSNet. Will it be applied to large molecular systems? \n7. Misleading concept of Generalization\nThe defination in this article seems a little misleading. When taking about generalization, the generalizable ability among different kinds of molecules, i.e., model trained in a molecules and tested on other kinds of molecules is much more important and convincing than that for different conformations of the same kind of molecule. Unfortunately, through the whole article, the authors neglected the important point for some reason.\n8. Unphysical structures\nAgain, it's so werid and totally meaningless to sample protein structures at such high temerature (1200K). First, the author must show the structes they sampled at that temperature. I guess most of structures are totally collapsed. Second, the authors should perform classical MD simulations at such meaningless temperature to compare the RMSD/system energy or other observables with the models shown in the article.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)