GEOMETRIC GRAPH NEURAL DIFFUSION FOR STABLE MOLECULAR DYNAMICS SIMULATIONS

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.

1 INTRODUCTION

Molecular dynamics (MD) simulations rely on force fields to approximate the underlying potential energy surface and generate long-temporal trajectories of molecular systems. Geometric graph neural networks (Geo-GNNs) have transformed MD simulations by providing a computationally efficient 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 on the accuracy of predicting forces and overlook the performance evaluation in real MD simulations, e.g., whether the real MD could reveal detailed physical mechanisms (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 throughout 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 a 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ács et al., 2021), as the dataset contains molecular geometries sampled at 300 K, 600 K, and 1200 K. Each temperature setting naturally induces a 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 1 (a) and (b). We then trained the representative Geo-GNNs VisNet (Wang et al., 2024c) and SEGNO (Liu et al., 2024) at 300 K and tested them at 300 K, 600 K, and 1200 K. VisNet is the state-of-the-art Geo-GNN for simulating MD. While VisNet demonstrates strong within-domain performance (300 K), its accuracy degrades sharply under shifted conformation spaces (Figure 1 (c)). In contrast, SEGNO improves the generalizability via explicitly embedding physical biases. It 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 to stable MD simulations.

Figure 1: Geometric Topological Shift Analysis of the 3BPA Dataset and Extrapolation Performance 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.

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 conceptualize domain variations in conformational spaces as “geometric topological shifts.” Correspondingly, we introduce the diffusion process with two novel operators–equivariant gradient and diffusivity operators–to capture the invariance to conformational changes while maintaining equivariance. In particular, the gradient operator captures variations in node features across the topology of the geometric graph by characterizing differences between arbitrary nodes, while the diffusivity operator regulates the rate and extent of information propagation. Together, these operators drive the evolution of node representations, capturing all-pair information flows over a 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 a plug-in module, seamlessly integrating with existing EGNNs to enhance 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 conformational 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 as a sequence of geometric graphs indexed by time . Suppose we have atoms in the system, then the molecular system at each snapshot can be represented as a point cloud denoted as , where is the atom coordinate matrix and is the node feature matrix. typically contains atomic types or charge features, and it is generally time-invariant. Given the molecular structure , the objective of the machine learning force field is to predict the energy or forces with the molecular graph input .

Figure 2: 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 at 300 K) and the testing set (e.g., 3BPA at 600 K).

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 conformational changes caused by the environment, the biochemical topology generally remains invariant, as it 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; Snijders & 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 , is formed by a graphon—a continuous graph limit defined as a symmetric, measurable function —serving as an unobserved latent variable, alongside a modeling method that specifies a cutoff radius .

To elucidate the node-level structure of this graph, each node is associated with an independent and identically distributed (i.i.d.) latent variable . The vector and scalar features, and , are random variables derived from each through node-wise functions and , respectively. Next, the geometric topological adjacency matrix is a random variable determined by a pairwise function , which depends on the environment and the modeling method . Changes in , such as transitions from training to testing, lead to variations in the distribution of . Beyond the node features and adjacency structure, the label also varies due to conformational variations. We assume Y is formed by a set function , with a specific realization denoted as . We denote specific realizations of these random variables as matrices , , , and .

Extrapolation and stability. Extrapolation remains a fundamental challenge in MD simulations, particularly in the application of data-driven machine learning methods to MD. From an ML perspective, 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 unseen molecular geometries, such as samples from different temperatures or non-equilibrium states. Improving the ability to extrapolate to unseen conformations is critical for a model to ensure stable MD simulations, a principle supported by numerical MD methods (Barth & Schlick, 1998; Miao & Ortoleva, 2009). To tackle extrapolation, MatterSim (Yang et al., 2024) employs an active learning approach to address both categories, though it relies on costly high-quality data collection. Besides, SEGNO (Liu et al., 2024) integrates second-order motion laws to enhance the generalization of equivariant graph neural networks, yet it fails to address geometric topological shifts. Despite the importance of conformational space extrapolation for stable MD, enabling Geo-GNNs to generalize under geometric topological shifts remains an unresolved challenge. A brief overview of MD, generalization, and equivariance is provided in Appendix A.

Figure 3: The Illustration of Geometric Graph Neural Diffusion: (a) Our method serves as a plug-in module that integrates with local equivariant message passing. (b) The GGND uses equivariant diffusion operators (gradient and diffusivity) on a 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.

3 METHOD

Our design enhances robustness to geometric and topological variations while preserving -equivariance, thereby enabling stable MD simulations despite limited coverage of molecular conformations in training data. The model integrates two complementary components: (i) a novel geometric graph neural diffusion module with global attention and (ii) a conventional local equivariant message passing neural network (Satorras et al., 2021; Wang et al., 2024c). The GGND employs a diffusion process modeled as a partial differential equation (PDE) on the graph, incorporating global attention to capture long-range dependencies across all nodes. This global perspective mitigates 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 module, seamlessly integrable with most existing EGNN frameworks to enhance their performance in stable MD simulations. We provide an overview of our 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 to 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 a geometric graph with nodes, where each node has a scalar feature vector and a position , and is the set of edges determined by the adjacency matrix assuming full connectivity in the graph. A concise background on diffusion processes on graphs is provided in Appendix B.

Each node has initial features , which include chemical element features (invariant scalars) and positions (for equivariance). The features consist of spherical tensors labeled by irreducible representations of O(3), denoted as , where indexes channels (learnable features), is the degree (e.g., for invariants, for vectors, and higher for tensors), and indexes components. Diffusion models on geometric graphs replace discrete GNN layers with continuous time-evolving node embeddings , where and evolves according to the diffusion equation:

where are equivariant node features, with initial conditions , and is the embedding layer through by radial basis functions (RBF).

The term denotes the diffusivity over the graph, defined as an matrix-valued function dependent on , which measures the rate of information flow between node pairs. The gradient maps node fields to edge fields, while the divergence operator div is its adjoint, mapping edge fields back to nodes. This diffusion process is modeled as a 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 an equivariant diffusivity operator, which facilitate global information flow while maintaining equivariance.

Equivariant gradient operator. The gradient operator generalizes scalar differences to higher-order tensors, incorporating directional information to preserve -equivariance. It is defined as:

where are learnable weights for mixing channels. This equivariant gradient operator on the graph generalizes the scalar gradient to higher , ensuring equivariance, with the difference modulated by directional information to preserve 3D structure. Notably, ranges over all nodes in , 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.

Equivariant diffusivity operator. The diffusivity is made equivariant by defining it as a tensor-valued attention matrix. We extend scalar attention to tensors as follows:

where are Clebsch-Gordan coefficients ensuring proper equivariance of . Here, , is the unit vector, are spherical harmonics (for directional equivariance), is a learnable radial basis function derived from Bessel functions and an MLP (ensuring invariance to distance), and is an equivariant pairwise interaction (e.g., a gated tensor product). This formulation ensures that transforms correctly under , serving as an equivariant filter that captures global dependencies. The attention matrix is right-stochastic, allowing Equation (1) to be rewritten as:

Equation (4) governs the dynamics of the system from to a specified stopping time , producing geometric and topological node representations . The equation is generally nonlinear due to the dependence of the diffusivity matrix on . A linear variant emerges when attention weights are fixed; however, as static attention is impractical, we focus on the nonlinear GGND model.

Output. For energy prediction, we utilize the invariant components of , specifically , combined with local equivariant features learned by Equivariant Graph Neural Networks (EGNNs). This ensures that site-specific energy contributions remain invariant, computed as , where represents features fused from local EGNN outputs and geometric-topological invariant features via concatenation and a linear transformation: = W[; ].

Equivariance. The diffusion process on the geometric graph, as described in Equation (4), enables the learning of domain-invariant features. With equivariant gradient and diffusivity operators, it also ensures the equivariance of the learned global features. We provide a proof of the equivariance of our GGND in Appendix C.

3.2 ALLEVIATING GEOMETRIC TOPOLOGICAL SHIFTS

We analyze the extrapolation capability of our geometric graph neural diffusion model with respect to geometric topological shifts, as defined in Section 2. Our focus is on the extrapolation error of the parametric function , instantiated as the continuous equivariant diffusion model in Equation (4), when transferring from training data generated under environment (and modeling method ) to testing data under (and ). Such shifts may arise from variations in adjacency matrices due to different cutoff radii or environmental conditions affecting inter-node distances in molecular systems.

Denote the training dataset of size as , drawn from , and let be a bounded loss function. The training error is

Our objective is to minimize the expected loss on testing data from :

When , this reduces to the standard in-distribution setting, where the extrapolation gap is bounded by

With , the Rademacher complexity of the function class is induced by , and the upper bound is determined by dataset size and model complexity.

In the out-of-distribution regime where , geometric topological shifts complicate the analysis. Changes in geometric topologies alter node representations in the equivariant graph diffusion equations (4), expressible as . The extrapolation gap could be decomposed into three terms (Wu et al., 2025). Assume and are Lipschitz continuous. For geometric graph data generated per Section 2, with probability at least , the extrapolation gap satisfies:

We denote the first as OOD model error and the second as OOD label error . Since is independent of testing data under , the impact of geometric topological shifts on extrapolation hinges on and : the former captures variation in due to shifting topologies (e.g., adjacency changes from varying cutoff radii or conformations), while the latter reflects label differences across environments or methods. is dictated by the data-generating process, whereas depends on , specifically the sensitivity of representations to shifts. We next examine as in Equation (4), adapted for equivariance.

Theorem 3.1 For geometric graph data per Section 2, if and are injective, the geometric graph neural diffusion model in Equation (4) reduces the representation variation to any order , where is an arbitrary polynomial, , and (with incorporating geometric distances via cutoff radius and is the diagonal degree matrix of ).

This indicates that the geometric graph neural diffusion model controls representation changes at arbitrary rates relative to , maintaining robust force prediction for conformation variations in molecular dynamics. The injectivity of and are mild assumptions, mapping from compact latent spaces to high-dimensional vector and scalar features. Applying Equation (8) yields the following.

This bound ensures controllable extrapolation error at any rate relative to . The model achieves desired extrapolation ability under shifts, such as in machine learning force fields or simulations with conformational changes or varying cutoffs. In contrast, the change rate of features produced by the local message passing model has an exponential upper bound. We presented the proof for the Corollary 3.2 in the Appendix D.

4 EXPERIMENTS

4.1 EXPERIMENTAL SETUP

Datasets. We used the 3BPA and SAMD23 datasets to evaluate our model’s performance, particularly in the presence of geometric and topological shifts. The 3BPA dataset consists of 500 training structures of the flexible, drug-like molecule 3-(benzyloxy)pyridin-2-amine at 300 K, with test data provided at 300 K, 600 K, 1200 K, and different dihedral angles (Kovács et 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).

Baselines. Our proposed equivariant graph neural diffusion can be integrated with any local equivariant message-passing-based method. To evaluate performance improvements, we selected four representative methods—NequIP (Batzner et al., 2022), MACE (Batatia et al., 2022), SEGNO (Liu et al., 2024), and VisNet (Wang et al., 2024c)—as baselines and compared our approach when combined with them against these baselines alone. Additionally, we conducted a comprehensive comparison 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 P3M (Wang et al., 2024b), LSRM (Li et al., 2024b), and FreeCG (Shao et al., 2025),

Metrics. Accuracy: we evaluate the predictive performance of our model using the mean absolute error (MAE) for energy and force predictions. For the SAMD23 dataset, which includes SiN molecules with atom counts ranging from 16 to 510, we report the energy per atom to ensure comparability across molecular sizes. Stability: following the methodology in (Fu et al., 2023), we assess the stability of flexible molecules by monitoring bond length deviations. A real MD simulation is classified as unstable at time if the maximum deviation of any bond length from its equilibrium value exceeds a threshold, formally defined as: , where denotes the set of all bonds, and are bond endpoints, is the equilibrium bond length, and is the stability threshold. For systems with periodic boundary conditions, stability is evaluated using the radial distribution function (RDF). A simulation is deemed unstable at time when: , where represents the time-averaging operator, is a 1 ps time window, and is set to 1.0. We perform constant-energy (NVE) molecular dynamics simulations at the specified temperature, employing Velocity Verlet integration over 100 ps with a 1 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 dynamics 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.

4.2 RESULTS AND ANALYSIS

Performance on 3BPA datasets. Experimental results on the 3BPA dataset in Table 1 show that integrating the proposed GGND module with baseline Geo-GNN models (MACE, NequIP, SEGNO, and VisNet) significantly enhances performance, particularly in extrapolating to conformational domains 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 to 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 at 1.965ps, SEGNO at 0.009 ps), whereas GGND restores robustness, increasing MACE’s stability to 29.218 ps (15-fold), NequIP 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 energy MAE reduction from 3.464 eV to 0.108 eV. These results highlight GGND’s efficacy in mitigating geometric topological shifts via all-pair information diffusion, facilitating stable 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 a lower energy per atom (E/A) error of 0.009 eV and a force MAE of 0.443 eV/ Å in the Test split, improving force predictions by approximately 9% over Neural P3M and LSRM, with a perfect stability score of 100 ps. In the OOD split, GGND maintains robust performance with a stability score of 99.89 ps, significantly surpassing QuinNet and Neural P3M. For HfO, 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 P3M and QuinNet, and achieving a 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 effectively captures all-pair interactions, ensuring insensitivity to conformational changes and enhancing stability in molecular dynamics simulations. The experimental results highlight the remarkable ability 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-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 metric 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 maximum bond length deviation in Figure 4. In 100 ps MD simulations on the 3BPA dataset, GGND outperforms the ML-based baselines by maintaining stability. Notably, although GGND exhibits instability around 30 ps, these unstable states occur randomly, whereas both VisNet and MACE show persistent instability after a certain time step.

Figure 4: Stability of MD Simulations on 3BPA.

Ablation Study. To evaluate the impact of fully-connected diffusion in our proposed GGND model, we curated two variants: GGND†, which uses local diffusion on the graph, and GGND‡, which combines the baseline with fully-connected message passing. The ablation study on the 3BPA dataset (Table 3) demonstrates the superior generalization of GGND to unseen conformational domains at 600 K and 1200 K, while matching the baseline’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 ps at 1200 K), as it cannot capture all-pair interactions. GGND‡ shows some generalization potential (e.g., stability of 2.892 ps at 1200 K) but underperforms GGND due to training challenges, highlighting the advantage of fully-connected diffusion in enabling robust, equivariant information flow for stable and accurate molecular dynamics simulations across diverse conformations.

5 CONCLUSION

In this study, we investigate the stability of MD simulations and identify extrapolation to unseen conformations as a key challenge. To address this, we propose GGND, a novel framework that improves the stability and generalizability of MD simulations by capturing geometrically invariant topological features through an 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 to stable molecular dynamics simulations. Our theoretical analysis establishes a regret bound under such shifts, providing formal guarantees of stability. Designed as a plug-and-play module, GGND integrates seamlessly with existing local 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 P0057774, the Research Grants Council of the Hong Kong SAR (Grant No. PolyU15215623, PolyU15229824, C5052-23G, and SRFS2526-5S04), 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 J. Our methodology does not involve human participants, sensitive data, or applications with significant misuse potential. We have prioritized fairness and transparency in the 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 utilizes the publicly available datasets 3BPA and SAMD23, accessible at https://pubs.acs. org/doi/10.1021/acs.jctc.1c00647 and https://github.com/SAITPublic/MLFF-Framework, respectively. We adhere to the data splits specified in the publications associated with these datasets, with all relevant parameters documented in Appendix E. The code is available at https://github.com/HaokaiHong/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. In Forty-second International Conference on Machine Learning, 2025.

Eric Barth and Tamar Schlick. Overcoming Stability Limitations in Biomolecular Dynamics. I. Combining Force Splitting via Extrapolation with Langevin Dynamics in LN. The Journal of Chemical Physics, 109(5):1617–1632, 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 Kornbluth, 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.

Chi Chen and Shyue Ping Ong. A Universal Graph Deep Learning Interatomic Potential for the Periodic Table. Nature Computational Science, 2(11):718–728, 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 Potentials 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 A 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 out-of-distribution graphs. IEEE Transactions on Pattern Analysis and Machine Intelligence, 46(1):322–337, 2024. doi: 10.1109/TPAMI.2023.3321097.

Xiang Fu, Zhenghao Wu, Wujie Wang, Tian Xie, Sinan Keten, Rafael Gomez-Bombarelli, and Tommi Jaakkola. Forces are not Enough: Benchmark and Critical Evaluation for Machine Learning Force Fields with Molecular Simulations. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. Survey Certification.

Jiaqi Han, Wenbing Huang, Tingyang Xu, and Yu Rong. Equivariant graph hierarchy-based neural networks. Advances in Neural Information Processing Systems, 35:9176–9187, 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 Representations, 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, Vıctor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant Diffusion for Molecule Generation in 3D. In International Conference on 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, 36:51434–51476, 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ács, Cas van der Oord, Jiri Kucera, Alice EA Allen, Daniel J 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):7696–7711, 2021.

Thomas J Lane, Gregory R Bowman, Kyle Beauchamp, Vincent A Voelz, and Vijay S Pande. Markov State Model Reveals Folding and Functional Dynamics in Ultra-Long MD Trajectories. Journal of the American Chemical Society, 133(45):18413–18419, 2011.

Hongkang Li, Meng Wang, Tengfei Ma, Sijia Liu, ZAIXI ZHANG, and Pin-Yu Chen. What improves 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-Range Message-Passing: A Physics-Informed Framework 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 Atomistic Graphs. In The Eleventh International Conference on Learning Representations, 2023.

Yi-Lun Liao, Brandon M Wood, Abhishek Das, and Tess Smidt. EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree 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 biases. In The Twelfth International Conference on Learning Representations, 2024.

Georgi S 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 J Ortoleva. Molecular Dynamics/Order Parameter Extrapolation for Bionanosystem Simulations. Journal of Computational Chemistry, 30(3):423–437, 2009.

Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron J Owen, Mordechai Kornbluth, and Boris Kozinsky. Learning local equivariant representations for large-scale atomistic dynamics. Nature Communications, 14(1):579, 2023.

Vıctor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural networks. In International Conference on Machine Learning, pp. 9323–9332. 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 p$ˆ3$m: A long-range interaction modeling enhancer for geometric GNNs. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024b. URL https://openreview.net/forum?id=ncqauwSyl5.

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 Communications, 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 transformers with advective diffusion. In Forty-second International Conference on Machine Learning, 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. In Proceedings of the 41st International Conference on Machine Learning, ICML’24. JMLR.org, 2024.

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 Conference 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.04967, 2024.

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 mechanical calculations in MD simulations and related atomistic modeling tasks. However, the generalization of ML methods for MD simulations remains unaddressed.

A major advance in improving the generalization of these models has been the development of “universal” or pretrained interatomic potentials trained on chemically diverse and large-scale datasets. Models such as M3GNet (Chen & Ong, 2022) and CHGNet (Deng et 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 equivariant geometric graph neural networks that explicitly encode physical symmetries. Equivariant message 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 in 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-level evaluations, including long-term stability, conservation laws, and reproduction of thermodynamic observables. These evaluations often reveal substantial degradation under out-of-distribution (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-aware 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 distribution gaps. These strategies mitigate specific failure modes but can be computationally demanding and lack systematic mechanisms to detect and react to distribution shifts in a theoretical manner.

Our work builds on these threads by proposing a geometric GNN that detects and responds to distribution shifts in conformational space, aiming to improve stability and accuracy when generalizing across thermodynamic regimes.

A.2 GENERALIZABLE GRAPH NEURAL NETWORK

Recent advancements in graph neural network (GNN) generalization have progressed along four complementary dimensions: theoretical foundations, architectural innovations, training methodologies, and data-centric strategies. Theoretical studies have established sample-complexity and stability 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). Architectural 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 learning and data augmentation techniques, effectively mitigating empirical generalization gaps (Abbahaddou et al., 2025; Fan et al., 2024). Collectively, these efforts outline a coherent research agenda: leverage theoretical insights to identify generalization bottlenecks, design architectures with inductive biases to reduce sample complexity, and employ training and data strategies, alongside out-of-distribution (OOD)-aware mechanisms, to ensure robust generalization.

However, existing research on GNN generalization has largely overlooked the challenge of conformational distribution drift in MLFF or MD. Specifically, environmental factors, such as temperature, can induce drifts in the geometric structure distribution of molecular graphs, which subsequently alter their topological configurations.

A.3 EQUIVARIANT GRAPH 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 , including 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 representations 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 et al., 2021; Xu et al., 2024) and propagating higher-order geometric 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 a range of molecular modeling tasks, including force field prediction (Batatia et al., 2022; Batzner et al., 2022; Wang et al., 2024c), molecular dynamics simulation (Xu et al., 2024; Han et al., 2022), and molecular design (Hoogeboom et 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. Consequently, modern geometric deep learning approaches for interatomic potentials and force fields widely adopt equivariant representations as a 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 a global diffusion mechanism that maintains equivariance while improving robustness to conformational variation.

B DIFFUSION PROCESSES ON GRAPHS

This section provides a concise background on diffusion processes defined on graphs. The presentation clarifies how Equation 1 can be interpreted as a graph partial differential equation.

B.1 DIFFUSION AS FEATURE PROPAGATION

Let be a graph with nodes. A time-dependent node feature field is denoted by

where represents the feature of node at time .

Graph diffusion models describe the evolution of node features through a continuous-time process

where is a 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 ON GRAPHS

Diffusion on graphs can be derived from discrete analogues of gradient and divergence operators.

Graph gradient. Given node features , the gradient operator maps node features to edge features: ,

for node pairs . This operator measures feature variation across the graph topology.

Graph divergence. The divergence operator aggregates edge features back to nodes:

where 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

where is a diffusivity operator that modulates information flow between node pairs and denotes element-wise interaction between edge fields.

When is symmetric and row-stochastic, the diffusion dynamics can be equivalently written as

which corresponds to a continuous-time random walk on the graph. This formulation shows that diffusion performs feature smoothing while preserving global consistency across nodes.

C PROOF OF EQUIVARIANCE OF GGND

We prove that the GGND process, as defined by the PDE in Equation ( equation 1) and its rewritten form in Equation ( equation 4), is -equivariant. Specifically, the learned features transform correctly under transformations (rotations and translations) applied to the input positions and invariant scalar features .

-equivariance means that if we apply a transformation to the inputs, the output features transform accordingly: ,

where is the solution of the PDE for the transformed inputs and (since are invariant scalars), and denotes the group representation acting on the spherical tensor features (irreducible representations of , labeled by and ).

Translations are handled trivially because the model depends only on relative positions (invariant) and unit vectors (equivariant under rotations but invariant under translations). The initial embedding uses radial basis functions (RBFs) on distances, which are translation-invariant. Thus, the dynamics preserve translation invariance.

We focus on rotation equivariance under . The representation acts on each irrep component as:

where the action is the same for all nodes .

Assume the initial condition is equivariant: . We need to show that if at time , then the time derivative preserves this property:

This requires showing that the right-hand side operator is equivariant: .

We prove this by showing that each component—the gradient , the diffusivity , and their combination—is equivariant.

Equivariance of the Gradient Operator. The equivariant gradient is defined as:

where are learnable scalar weights (invariant under rotations), and the gradient operates per irrep and channel, without changing .

For the transformed features and positions:

Thus,

since is invariant. The gradient transforms as the same irrep, so is equivariant: .

Equivariance of the Diffusivity. The diffusivity (attention matrix) is:

where are Clebsch-Gordan coefficients (invariant), is a learnable radial function (depends on invariant distance ), are spherical harmonics, and is an equivariant pairwise interaction (e.g., gated tensor product).

Substituting:

This is: ∑