Long-Short-Range Message-Passing: A Physics-Informed Framework to Capture Non-Local Interaction for Scalable Molecular Dynamics Simulation
Authors: Yunyang Li
Yusong Wang
Lin Huang
Han Yang
Xinran Wei
Jia Zhang
Tong Wang
Zun Wang
Bin Shao
Tie-Yan Liu
Venue: ICLR 2024 poster (2024)
Keywords: Molecular Modeling
Quantum Chemistry
Fragmentation
Non-Local Interactions
EGNN
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We proposed a long-range short-range message-passing framework to capture non-local interactions and demonstrate the state-of-the-art results with up to 40% error reduction for molecules in MD22 and Chignolin datasets.

Abstract: Computational simulation of chemical and biological systems using *ab initio* molecular dynamics has been a challenge over decades. Researchers have attempted to address the problem with machine learning and fragmentation-based methods. However, the two approaches fail to give a satisfactory description of long-range and many-body interactions, respectively. Inspired by fragmentation-based methods, we propose the Long-Short-Range Message-Passing (LSR-MP) framework as a generalization of the existing equivariant graph neural networks (EGNNs) with the intent to incorporate long-range interactions efficiently and effectively. We apply the LSR-MP framework to the recently proposed ViSNet and demonstrate the state-of-the-art results with up to 40% MAE reduction for molecules in MD22 and Chignolin datasets. Consistent improvements to various EGNNs will also be discussed to illustrate the general applicability and robustness of our LSR-MP framework. The code for our experiments and trained model weights could be found at https://github.com/liyy2/LSR-MP.

Published as a conference paper at ICLR 2024
LONG-SHORT-RANGE MESSAGE-PASSING: A PHYSICS-INFORMED FRAMEWORK TO CAPTURE NON-LOCAL INTERACTION FOR SCALABLE MOLEC- ULAR DYNAMICS SIMULATION
Yunyang Li *
Yusong Wang V
Lin Huang =0 Han Yang*
Xinran Wei* Bin Shao Tin-Yan Liu
Jia Zhang A
Tong Wang &
Zun Wang*
Yale University
*MSR A14Science
Xi'an an Jiattong University
thuagg.ing
jia.zhang, watong @emicoosoft. com.
ABSTRACT
Computational simulation of chemical and biological systems using ab initio molec- ular dynamics has been a challenge over decades. Researchers have attempted methods.
to address the problem with machine learning and fragmentation-baste However, the two approaches fail to give satisfactory description of long-range and many-body interactions, respectively. Inspired by fragmentation- Message-Passe Paassgg-Psssingggg (LSR-MP) framework
meth-
ods, we propose the Long-Short-Raae a generalization of the existing equivariant graph neural networks (EGNNs) with the intent to incorporate long-range interactions efficiently and effectively. We apply the LSR-MP framework to the recently proposed ViSNet and demonstrate the state-of-the-art results with up 40% MAE reduction for molecules in MDD2 and Chignolin datasets. Consistent improvements to various EGNNS will also be discussed illustrate the general applicability and robustness of our LSR-MP framework. The code for our experiments and trained model weights could be found / wwt. wtt.h///..o//1121111---1
INTRODUCTION
Ab initio molecular dynamics (AIMD) Car & Parrinello (985) has been an indispensable tool in the fields of chemistry, biology, and material science. By virtue of its effective description of the kinetic and thermodynamic properties in molecular and condensed systems, AIMD capable of elucidating numerous phenomena of interest, such as chemical reactions Hwang al. (0015), protein folding Cellmer et al. (2011) and electron-phonon interactions Karsai et al. (2018); Kundu et al. 2202T 2022). However, AIMD simulations driven by conventional quantum chemical methods
such as density functional theory (DFT) Hohenberg & Kohn (994)); Kohn & Sham (9655 become
prohibitively expensive as the size of the system increases, due the high computational cost SZabo
& Ostlund (201))
Computational chemists have sought to address the challenges associated with understanding molecu- lar interactions by adopting a divide-and-comquee strategy. This approach has significantly spurred the evolution of fragmentation-
methods in recent years Yang (199I): Exner & Mezey (2003); He & Zhang (2055 2006); Li et al. (2007); He et al. 2014).. Central to these methods is the principle of chemical locality He et al. (2014), which posits that chemical subsystems often exhibit minimal or weak interactions with one another. This principle is grounded in the broader hypothesis that molecular behaviors and interactions largely arise from localized phenomena. Empirical evidence further supports this notion Collins & Bettens (2015), demonstrating that molecular reactivity is primarily dictated by the presence, composition, and spatial configuration of specific functional groups. Consequently, by subdividing large molecules into smaller, more manageable subsystems, one can facilitate parallelized studies and gain deeper insights into their intricate behaviors. Following
Work done during internship at Microsoft.
1 Equal contribution.

---

Published as a conference paper at ICLR 2024
Table 1: Comparison of the methods. N number of atoms, m number of basis.
Muthod
Efficiency
Many-hody
Lang-raage
Accuracy
Methods
O[N*s ja 9,m
Accunte
Pragmentation Methods
OIN) + ONN (n 2 2 D(A*][a 1-
1-5 kcal/mel
EGNN
X
Optinaily lccalmol
EGNN- LSR MP
0[8 1-
Optinally < kcaltmol
the computation on these subsystems, there manually designed techniques available to reconstitute the holistic properties of the entire molecule. Notable examples of these assembly methods are rooted in energy- and density matrix-based approaches [Li al. (2007 2021). Regardless of the method selected, the many-body interactions among subsystems must be carefully calibrated, as this is the primary source of error in the fragmentation- the two- and/or three-body interactions. However, such approximation compromises the accuracy of the fragmentation- baaee methods and thus has significantly limited their applicability Li et al. (021).
method. A common practice is to take into account
Recent advances in the application of machine learning in chemistry Zhang et al. (2018): Behler (20TT: Unke et al. (202T): Deringer al. (202T); Zhang et al.) (2022) have facilitated AIMD simulations of larger molecules without significant loss of accuracy, in contrast to the prohibitive computational costs of conventional DFT calculations. In particular, the recent equivariant graph neural networks (EGNNs) introduce the inductive bias of symmetry into the model Han et al. (2022) which further improves the model's capacity to represent the geometric structures the molecular systems. Existing EGNNs model interactions in local environments by incorporating structural information such as bond lengths, bond angles, and dihedral angles within a predefined radius cutoff. Although being effective for datasets consisting of small molecular systems, e.g. MD17 Chmiela eet al. (2018); |Schütt et al. (2017): Chmiela et al. (017)) and QM9 |Ramakrishnan et al. (2014), this approach results in substantial loss of information for larger molecules in which long-range interactions, including electrostatic and van der Waals forces, are non-negligible. Increasing the radius cutoff and stacking more layers are common methods to remedy the information loss, but they inevitably bring about new issues, e.g. loss of data efficiency, convergence difficulty, and information over-squahhin Alon & Yahav (2020).
In light of these challenges, we introduce the tragmentation-basee message passing on top of the current EGNNs, called the Long-Shorr-BR Short-Range Message- (LSR-MP) framework, with the objective to (1) explicitly incorporate long-range interactions, (2) compensate for the loss of higher-order many-body interaction in existing fragmentation-
approaches and (3) maintain computational efficiency. As illustrated in Fig. [ the input molecule is processed at two different levels. In the first level, we introduce the short-range module to encode the many-body interactions under the local neighborhood. In the second level, formulate the fragmentation module and the long-range module, which are dedicated to modeling long- range interactions between fragments and atoms. We compare the key differences of these methodologies in Table I Our contribution could be listed as follows:
1. We introduce LSR-MP, a novel message-passing framework for neural network-based poten- tial modeling. Leveraging BRICS fragmentation techniques and long-range message- message-passing. LSR-MP effectively captures long-range information, showcasing its robust performance in modeling large molecular systems.
2. We present an exemplary implementation of the LSR-MP framework, dubbed FiSNet-LSRM Our approach leverages the interaction between vectorial and scalar embeddings maintain equivariance Wang et al. (2022a b). Notably, ViSNet-LSRM achieves competitive perfor- mance large molecular datasets such as MD22 (em222Chmllaaaa al. (2023) and Chignolin Wang
et al. (2022a), while utilizing fewer parameters and offering up to 2.3x faster efficiency.
3. We illustrate the general applicability of the LSR-MP framework using other EGNN back- bones, such Equiformer, PaiNN, and ET, which results in substantial performance im- provements over the original models. This highlights the framework's 's superior adaptability and effectiveness.
2

---

Published as a conference paper at ICLR 2024
2 RELATED WORK
2.1
FRAGMENTATION -BSSE METHOOS
Fragmentation techniques, historically employed facilitate quantum mechanical (QM) computations for expansive molecular systems, present solutions to intricate scaling issues Ganesh et al. (2006); Li al. (2007). In recent research trajectories, there has been discernible bifurcation into two prominent fragmentation methodologies; density matrix-based and energy-based Within the density matrix paradigm, the overarching strategy entails the construction of an aggregate density matrix, sourced from the individual matrices of subsystems. Subsequent property extractions hinge upon this consolidated matrix. Noteworthy implementations in this domain encompass the divide-and-conquer (DC) technique Yang [99I), the adjustable density matrix assembler (ADMA) framework [Exner Mezey 2003), the elongation approach Gu et al. (2004), and the molecular fractionation with conjugated caps (MFCC) mechanism [H & Zhang (2005). Conversely, energy-based fragmentation methodologies pivot on direct assimilation of subsystem properties to deduce the global property. A seminal contribution by He al. involved the adaptation of the MFCC model into more streamlined formulation He & Zhang (2006), which was subsequently enhanced with electrostatic embedding schemes, capturing long-range interactions (He et al. (2014). Parallel advancements in the energy-based arena include the molecular tailoring approach Gadre & Ganesh (2006), the molecules- in-molecules (MIM) paradigm Mayhall & Raghavachari (2011), and the systematic fragmentation
methodology Deev & Collins 2005).
2.2 EQUIVARIANT GRAPH NEURAL NETWORKS
Equivariant Graph Neural Networks (EGNNs) inherently weave symmetry's 's inductive bias into model architectures. Traditional EGNNs predominantly lean on group theory, invoking irredueible representations Thomas et al. 2018)) : Fuchs et al. (2020); Anderson et al. (2019) and leveraging the Clebsch- Gordan (CG) product guarantee equivariance, Contemporary strides in this domain, such as those presented in Batzner et al. (2022); Frank et al. (2022b); Batatia et al. (2022). have incorporated higher-order spherical harmonic tensors, resulting in notable performance enhancements. Parallel to these, alternative advancements have been steered by methods like PaiNN [Schitt et al. (202T) and TorchMD-NET and vectorial features. ViSNet Wang et al. (2022a) builds upon PaiNN's foundation, integrating runtime geometric computation (RGC) and vector-sealar interactive message passing (ViS-MP).
& De Fabritiis (2022), which prioritize iterative updates of scalar
NONG-SHORT-RANGE MESSAGE- PASSAGE-PASSING
Notations: We use h and ø to represent short-range scalar and vectorial embeddings, respectively, and x and j7 for their long-range counterparts. Capitalized fragment embeddings are signified by H and V. Subscripts align with atom or fragment indexes, while superscripts denote layers in a multi-layered network. Additionally, product, and ( for vector scalar products. 1d is a dimension d column vector of ones. Lishort the number of short-range layer, Loong the number of long-range layer. Functions include: DENSE(.) for an activated linear layer with bias, LINEAR(.) (-) for a biased linear layer, and (-) for an unbiased linear layer. Network parameters are typically unshared unless specified.
dist(-) for the Euclidean distance, @ for the Hadamard
3.1 SHORT-RANGE MODULE
Considering short- range radius graph Gshon with node set V for atoms and edge set Eshon which is defined as Exhont = {etz ddist(ā;, FT) S T Tsbort, Vi,j E V}- The short-range module performs message passing on Gshont by taking the atomic numbers
N and positions ş e = (n is the number of atoms in the system) as input, and is targeted model the geometric information (bonds, angle, dihedral, improper) on Gatont-
For illustration purposes, suppose the short-range module operates scalar embeddings and vectorial embeddings, while generalizations to higher-order spherical harmonics embeddings could be easily derived with the Clebsch-Gordan tensor product, and is included in Appendix L
3

---

Published as conference paper at ICLR 2024
(e) Short-Rangg Module
LSR-MP Framework
Short-Range Modulia
Bond
Chedral
Sehart
Band Angie
Improper
Shart Module 3.1)
[b]
Pragmentation Module
NHz
Split Layer
HO
OH
(b) Fragmentation
(c) Long-Range Modull
OH OH
king
Fragmentations (Sec 3.2)
(oj Long-Raange Module
Bectmstatics
Output Layer
Gredient
Energy
Force
Long-Ranger Module (Sec 3.3)
Figure LSR-MP Framework: (a) center atom (in red) uses short-range equívariant message passing radius graph for its embedding. A linear projection and split process in the split layer channel some info to the fragmentation and long-range modules. (b) The fragmentation module dissects the system and derives fragment representations. bipartite equivariant message passing between atoms and fragments. The output layer merges both long and short-range embeddings, then sums the graph for graph-level embeddings.
(c) The
module facilitates
In general, the short-range module adopts an iterative short-range message-passing scheme:
16, o7 clstort 44
Σ mi
Σ
(I)
JEN(S
JEN(
The ('hhor-)) defines a message-passing framework, and (-)- denotes the first-order neighbor set of particular node. mij and miy denote the scalar message and vectorial message between node and its first order neighbor j.
3.2 FRAGMENTATION MODULE
To incorporate long-range components, look towards the fragmentation-bast based methods prevalent in quantum chemistry. While fragmentation-basel methods might not adept at depicting many- body interactions, they can potentially capture long-range interactions with good accuracy. For fragmentation- based methods, the preservation of intrinsic chemical structures during fragmentation paramount, this would minimize the artifacts induced by fragmentation based methods. The BRICS method, as implemented in RDKit challenge. Formally, the BRICS fragmentation procedure defined as: BRICS(Z, where S {0, 1mmxm n is the number of atoms, is the number of fragments.
B Landrum et al. (2013), offers robust solution to this
BRICS utilizes 16 predefined chemical motifs, selected based on organic chemistry principles, to direct the fragmentation process towards meaningful substructures. It aims to reduce energy discrepancies arising from bond breakages, ensuring chemically appropriate local environment. Additionally, BRICS applies post- tost-fragmentation techniques: removing small fragments. handling duplicates, and correcting overlapping fragment issues. Detailed implementation BRICS in the fragmentation module is provided in Appendix G
4

---

Published as a conference paper at ICLR 2024
Fragment Representation Learning: To obtain effective representations, the fragment learning function summarized the atom's embeddings within fragments. We denote an arbitrary atom sealar embedding h, vectorial embeddings as v Analogous to atoms, the fragment's scalar embedding invariant under SO(3) transformations, and its vectorial representation is equivariant to SO(3) transformations. Characterization of vectorial fragment emheddings enables the network to model cquivariant long-range interactions. The general learning scheme is defined as follows:
H = Σ a. O h²: V
Σ B P, = Σ YPP + G V.
(2)
iES(f)
BES())
SES())
in which Hj. V and Pj denotes the scalar embedding, vectorial embedding, and position of fragment j. respectively. j is the index for fragments, and ( is the set induced by the assignments of fragment j. a 14. E Rd are weight vectors for each atom within the fragments and should be SO(3) invariant. Additionally, to guarantee translational equivariance, are parameterized so that EiESt) 7: = 1,75
3.3 LONG-RANGE MODULE: GEOMETRIC BIPARTITE MESSAGE PASSING
Considering bipartite radius graph Giong with node set (V,u). V is the atoms set and U is the fragments set, and the edge set defined as:
Exee = {eus dist (p., S Clong Vi N, Vj vjv..
(3)
Mann is the long-range cutoff and is normally chosen be much larger than the radius of the short- range neighbor graph Tshort. The bipartite geometric message-passingg is performed to characterize long-range interactions:
2/ n. 1/long
1-1
Σ JEN(()
M Σ M
(4)
JEN())
(i())oog is the general bipartite message passing framework, a! the long-range scalar embedding. 14, is the long-range vectorial embedding, i, are index for atom, fragment respectively, and N()) is the neighborhood of atom i on the atom-fragment bipartite graph. Mij and Mij denote the bipartite scalar message and vectorial message between atom i and its incident fragment j.
3.4 PROPERTIES PREDICTION
The short-range embeddings, denoted as and ग, likely capture local interactions, while the long- range embeddings, represented by x and f7, encompass long-range interactions. These two types of embeddings supplement each other, offering comprehensive modeling of the systems. To effectively combine these embeddings, LSR-MP employs late fusion strategy:
hout
DENSE ([h [[=
(5)
Uout U ([zhhmm j
(6)
Lyhon and Liong are the number of layers for the short-range module and long-range module respec- tively. The output layer predicts scalar properties using the sealar embedding hou or predict tensorial properties using Tont and hoat-
4 VISNET-LSRM
Based on LSR-MP framework, we provided an exemplary implementation called ViSNet-LSRM. It uses ViSNet Wang al. (2022a) the backbone for the short-range module, thus ht, of VISNET( Z, p). For the long-range module, we will give a detailed architecture as follows. A visual representation of the architecture be found in Appendix[H] H. The design principles of this long-range module adopt the following proposition:
Proposition 4.1. The Hadamard product of a scalar representation by vectorial representation results in a vectorial representation. The inner product of two vectorial representations results in a scalar representation.
h, v could be short-range embeddings h vi or long-range emheddings v, fi see 3.3
5

---

Published as a conference paper at ICLR 2024
4.1 LONG-RANGE MODULE IN VISNET-LSRM
Layer Normalization of Fragment Representation: For each fragment, we commence by applying a layer norm to its scalar embedding and norm-based min-max normalization to the vectorial embedding:
V
| min (()
Vaorm.d
(7)
||Vall
(max([] min (IVE)
min(-) and max(-) are applied to the channel dimension, || is 12-norm applied the spatial dimension. Empirically, we observed that this normalization is succinct and effective technique to bolster model convergence.
Histance-Dependee Bipartite Geometric Transformer: Considering central atom, certain fragments may establish robust long-range interactions due to factors such charge, polarity, distance, and so forth, while others may not. It is intuitive that fragments in close proximity generally exert more substantial influence than those farther away. light of these observations, we propose distance-based equivariant multi-headed attention mechanism for atom-fragment message passing. For simplicity, we will omit the notation for heads the subsequent derivations. We commence by encoding the distance between fragments and atoms utilizing continuous filters:
(8) (-) is the radial basis function. The distance between the fragment and the atom depends on the size of both the atom and the fragment, as well as the position of the atom relative to the fragment, we employed a size-dependent distance encoding which is parametrized by a linear function:
S = DENSE (e (F.)))
رU w (2.. H, dist((pi Pý) + b(2, H,).
Subsequently, we employed the attention mechanism Vaswani et al. 2017); Thölke & De Fabritiis (2022) for atom-fragment interaction:
q: -W ky W,H,, vy W.
(10)
where WG, WN, W,, are projection maps. The query is computed using atom scalar embeddings; the key and value are computed using fragment scalar embeddings. The attention weights Aij between atom-i and fragment-j are obtained through element-wise wise multiplication query ges key ky, and encoded distance stz:
Aij = SiL.U(sum(q: O kj O (iy));
(11)
The output the attention mechanism results in a value vector weighted by attention weights, from which a scalar message is derived using a dense layer. Concurrently, the relative direction of the fragment to each atom and the vector embeddings of each fragment form the vectorial message, both gated by a dense layer acting on the attention output:
Mij DENNE
(12)
P.
mig = DENSE( uy)
DENSH(Aye;) o Vj.
(13)
11 P,
The aggregation function of the scalar and the vectorial message is summation bipartite fragment neighbors:
m
Σ mlgi
m
Σ my
(14)
JEN()
JEN())
Define 0, = ( (,+--),, (xt-1)) that is the scalar product of vectorial embeddings under two different projections. The updated scalar embeddings and vectorial embedding obtained as:
o + LINEAR (m!) @ Ú, + LINEAR (m!): p = m LI)) o pp-1 (15)
5 EXPERIMENT RESULTS
5.1 RESULTS MD22 DATASET AND CHIGNOLIN DATASET
MD22 is novel dataset presenting molecular dynamics trajectories, encompassing examples from four primary classes of biomolecules: proteins, lipids, carbohydrates, nucleic acid The data split of
BRICS fragmentation methods fail to fragment the supramolecules in MD22, detail their implementation MRR-MP in Appendix R.9

---

Published as a conference paper at ICLR 2024
MD22 is chosen to be consistent with Chmiela al. (2022) ) Chignolin, the most elementary artificial protein, boasts 166 atoms. The dataset from Wang et al. (2022a) features 9,543 conformations, determined at the DFT level. It poses significant challenges given its origin from replica exchange MD (REMD). The REMD method, integrating MD simulation with the Monte Carlo algorithm, adeptly navigates high energy barriers and adequately samples the conformational space of proteins. This results in diverse modes from the conformational space being represented in both the training and test sets. We partitioned the dataset into training, validation, and test sets in an 7:1 :2 ratio. For benchmarks, we considered sGDML 5GDML/Chmiela et al. (202)), ViSNet Wang et al. (2022a), PaiNN Schütt al. (202T), Thöleeee & De Fabritiis (2022). and the recently introduced GemNetOC Gasteiger et al. 2022). MACE Batatia et al. (2022). SO3krates Frank et al. (2022a). Allegro Musaelian et al. (2023) and Equiformer Liao & Smidt (2022). Comprehensive settings and implementation details are provided in Appendix K.
Table 2: Mean absolute enrors (MAE) of energy (kcal/mol) and force (kcal/mol/A) for five large biomolecules on MD22 and Chignolin. The best one in each category is highlighted bold.
Fall
50leaa Paace
Aepe
HVC
Egelbree
Vidmo
(ag Ser LS04
0300
0.121
0100
03013
60831
04718
0.05
DMA
1311
1.10)
9.138
0101
0120
0.153
03079
C181
Dats
12.9
4.21
088
6000
18890
0862
0356
eibn
0665
0996
0090
0.0
086
1381
8308
1303
6126
E34
0053
6150
0,68
0.35
6112
Clipeii
3209
1330
Table 3: Comparison of the number of parameters and training speed of various methods when forces
MAE is comparable on the molecule AT-AT-CG-CG. Detail settings can be found in Appendix K.2
Methods (MAF)
ViSNet (0.16)
ViSNer-L (0.13)
PaiNN (0.35)
ET (0.29)
Allegro (0.13)
Equiformer (0.13)
# of Parameters
2.21M
1.70M
3.20M
3.46M
15.11M
3.02M
Training Time Epoch (s)
44
26
818
155
the mean absolute errors (MAE) of both energy and forces have significantly decreased in comparison to the original ViSNet, with a more pronounced improvement in performance as the system size expands. This enhancement highlights the importance of explicitly characterizing long-range interactions for large biomolecules. In comparison previous state-of-the-art models such as Equiformer, Allegro, and sGDML, ViSNet-LSRM demonstrates competitive performance while delivering superior efficiency. Furthermore, we conducted comparative study validate the efficacy of our approach in terms of both performance and parameter counts. Our results, presented in Table 3 show that when performance is comparable, our framework uses fewer parameters and offers faster speed, indicating that our approach is more efficient and effective at capturing long-range interactions. This comprehensive analysis underscores the advantages of ViSNet-LSRM terms of both performance and efficiency, making it promising solution for large biomolecular systems.
As illustrated in Table 2
5.2 STUDY LONG-RANGE INTERACTIONS
In this subsection, we conduct an analysis of the long-range interactions based on the molecule AT-AT-CG-CG in MD22, aiming to address the following questions:
Question 1. Can existing EGNNs capture long-range interactions by increasing the radius cutoffs?
Contrary to expectations, the answer negative. As depicted in Fig 2() and (b), all three models exhibit optimal performance when the short-range cutoff is set to 4 or 5. A further increase in cutoff could significantly deteriorate the performance. This finding suggests that existing EGNNs may be bottlnecked by their limited capacity in information representation when handling a large number of neighbors, leading to information squasinn......
Question 2. Are existing EGNNs capable of capturing long-range interactions by increasing the number of layers?
Intuitively, deeper models possess a larger receptive field to capture long-range interactions. The results in Figures 2 d) and (e) across all three models with varying layer numbers confirm this intuition. For -SNettLSRMM layers at 2. Performance improves as layer number increases for all models. However, a 3-layr ViSNet-I SRM (marked in red in Fig. ootperrorrss 8-layer ViSNet (marked in yellow in Fig. 2 and
vary the number of short-range layers while fixing the longerange
7

---

Published as a conference paper at ICLR 2024
21.3
Exeeyy (cskioo
Fonee MAKE prokmoll
1.4
PuNNNN
Infen0e Time (180
Puhhh
15.1
11.1
022
5.1
Culom(Ai
(a)
(b)
(c)
VSShet
--- -15hae--S/MM
17.1
Vishin-LRRM
Forst MAE (sckiockh)
0.4
Infeenss Trme
12.1
21.1
022
biwgg
01
711
0.16
0.1
31
Numrber of Beywen
Namber of bywer
Namber of layaer
(d)
Comparative Studies on the (short-range) cutoff and the number of (short-range) layers for three methods, including ViSNet-LSRM, ViSNet, and PaiNN. For ViSNet-LSRM, we consistently used 2 long-range layers and 9 Å long-range cutoff. The mean absolute errors (MAE) of energy and forces for AT-AT-CG-CG are shown in each sub-graph. From top to bottom, the performance of three methods regarding the cutoff and the number of interactions increasing the cutoff radius could lead to information over squashing which is prevalent across all EGNNs, including the short-range component of MONNet-LSRM d, demonstrate increasing the depth (shown horizontally) of the short-range models not as effective as introducing the long-range layer (shown vertically).
Figure 2:
shown. a, b, illustrate
PaiNN, suggesting incorporating explicit long-ram rangrragge modules more effective than simply deepening EGNNs. Furthermore, the shallow ViSNet-LSRM is more efficient than the deeper baselines (Fig. f). Di Giovanni et (2023) has pointed out deepened models would inevitably induce vanishing gradients.
Question How to demonstrate that the proposed method improves the performance due to successful long-range interactions? How to justify the BRICS- CrCSSfframmentttoon method?
To address these questions, we conducted an ablation study Table modifying key components of our framework. Instead of the BRICS fragmentation method, we employed distance- based k- means clustering, keeping the rest of the model aligned with LSR-MP. Results indicate that even without chemical insights, the k-means-based LSR-MP surpasses the baseline that lacks LSR-MP, underscoring the importance of incorporating long-range components. However, the k-means method fragments chemical systems non-canonically, potentially masking critical long-range interactions. As previously noted, the BRICS method reduces energy loss from breaking chemical bonds, offering more informative fragmentation for molecular modeling. The improvements of BRICS upon K- means have corroborated these findings. We further benchmarked LSR-MP against another prevalent framework for long-range interaction, incorporating single global node counteract information oversquashing. Our improvements on this framework further accentuate the efficacy and robustness of L.SRMM..
Based on the experiments, we hypothesize two-fold contribution of BRICS methods: (I) Incorporating Chemical Insights: Leveraging domain-specific motif matching, BRICS refines representation learning to better capture intrinsic chemical environments. With foundation in chem- ical locality and appropriate fragment representation learning, the resulting fragment representations could effectively highlight essential compound properties while omitting superfluous information. (2) Prioritizing Negative Curvature Edges: This formulation builds upon the Balanced Forman curvature introduced in Topping al. (2021). At its core, the Balanced Forman curvature is negative
Discussion:
8

---

Published as a conference paper at ICLR 2024
Datovaa
Dignoli
ATLAT 1__00_p0
XTAT
Energy
30
0.7
0.0
Energy AoeE jos saak)
Fores MIE 1kgp poak)
0.
0.
0.28
PaNe
Equiloeeet
visiaal
Paki
Equitamer
Visklo
Figure 3: General applicability of LSR-MP framework with PaiNN, ET, Equiformer, ViSNet. Dark color schemes indicate an EGNN without LSR-MP, and light color schemes indicate with LSR-MP.
when the edge behaves bridge between its adjacent first-order neighborhood, while it is positive when the neighborhood stay connected after removing the edge. The phenomenon of over-squashing correlates with edges exhibiting high negative curvature Topping al. (202I) A key observation we present (See Appendix [E.5) is that bond-breaking induced by BRICS normally have high negative Balanced Forman curvature in the molecular graph, serving as the bridge the bottleneck propa- gate information across different functional groups. To alleviate this, the proposed long-range module rewires the original graph to alleviate graph's 'ss strongly negatively curved edges. Furthermore, have proved that LSR-MP could improve the upper bound of the Jacobian obstruction proposed Di Giovanni al. (0023). The proof included Appendix B
Table 4: Ablation studies of LSR-MP AT-AT-CG-GG.
Method (Average Fragment Size)
Energy MAE
Force MAE
BRICS (8.43)
0.1064
0.1135
K-Means (8.43)
0.1276
0.1246
Single Global Node
0,1482
0.1696
w/o Long- range Component
0.1563
0.1995
5.3 GENERAL APPLICABILITY OF LSR-MP
We evaluated the general applicability of L.R-MP AT-AT, AT-AT-CG-CG, Chignolin by using PaiNN, ET, Equiformer, and ViSNet as the short-range module. To ensure rigorous comparative standards, 6-layer EGNN model was juxtaposed with an LSR-MP model integrating 4 short- range layers and long-range layers. As shown in Fig. 3 all models show significant performance improvements under the LSR-MP framework. The force MAE of Equiformer, ViSNet, PaiNN, and ET is reduced by 15.00% 24.88% 58.8%, and 55.0% in AT-AT-CG-CG while similar improvement shown on energy MAE. This suggests that the LSR-MP framework can extend various EGNNs to characterize short-range and long-range interactions, and thus be generally applied to larger molecules.
CONCLUSIONS
In this work, inspired by the fragmentation- method in quantum chemistry, we proposed a novel message-passing framework, LSR-MP, to capture long-range and short-range interactions. With an implementation based on ViSNet. it showed superior performance on large molecules with efficiency and effectiveness. It also showed general applicability with application other EGNNs.
Limitations and Future Work: supramolecules, and canonical clustering methods were used alternatiess. In our future work, differentiable fragmentation method may be further developed to learn optimal fragmentation in an end-to-end manner.
As shown in the MD22 dataset, BRICS has failed to deal with

---

Published as a conference paper at ICLR 2024
REFERENCES
Uri Alon and Eran Yahav. On the bottleneck of graph neural networks and its practical implications. arXiv preprint arXiv:2006.05200 2020.
Brandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. Advances in neural information processing systems, 32, 2019.
llyes Batatia, David Peter Kovacs, Gregor N. C. Simm, Christoph Ortner, and Gabor Csanyi. MACE: Higher order equivariant message passing neural networks for fast and accurate force fields. In Alice H. Oh, Alckh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL httpttt:///ppeeeeeteeee wwwp://wwwwww...eeeeeeeeeee
id=YPpSngE-24 20
Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P Mailoa, Mordechai Korblluhh Nicola Molinari, Tess E Smidt, and Boris Kozinsky. E (3)-eqiivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature communications, 13(1):1-11, 2022
Jörg Behler. Four generations of high-dimensional neural network potentials. Chemical Reviews, 121 (16):10077-1007,, 2021.
Richard Car and Mark Parrinello. Unified approach for molecular dynamics and density- ffnntiall theory. Physical review letters, 55(22):2471, 1985.
Troy Cellmer, Marco Buscaglia, Eric R Henry, James Hofrichter, and William A Eaton. Mak- ing connections between ultrafast protein folding kinetics and molecular dynamics simulations. Proceedings of the National Academy of Sciences, 101(15):103--6108 2011.
Stefan Chmiela, Alexandre Tkatchenko, Huziel E Sauceda, Igor Poltavsky, Kristof Schütt, and Klaus-Robert Müller. Machine learning of accurate energy-conserving molecular force fields. Science advances, ((5)e160301s 2017.
Stefan Chmiela, Huziel E Sauceda, Klaus-Robert Müller, and Alexandre Tkatchenko. Towards exact molecular dynamics simulations with machine-leared force fields. Nature communications, 9(1): 1-10, 2018
Stefan Chmiela, Valentin Vassilev-Galinde Oliver T. Unke, Adil Kabylda, Huziel E. Sauceda, Alexandre Tkatchenko, and Klaus-Rober Müller. Accurate global machine learning force fields for molecules with hundreds of atoms, 2022. URL /attts://arxtaa http/aaxvvvvooo org abs 2209.02.0 14865)
Stefan Chmicla, Valentin
Oliver T. Unke, Adil Kabylda, Huziel E. Sauceda, Alexandre Tkatchenko, and Robert Müller. Accurate global machine learning force fields for molecules with hundreds of atoms. Science Advances, e()eaad88873, 2023. doi: 10. 106scaav www. wwwssciieccc..ccc....oo......e.s org/ doi / haaii/abs/000 1126 sciadv adf0873
adf0873. URL
Michael Collins and Ryan PA Bettens. Energy-based molecular fragmentation methods. Chemical reviews, 1 115(12)5607-5662, 2015.
Vitali Deev and Michael A Collins. Approximate ab initio energies by systematic molecular fragmen- tation. The Journal of chemical physics, 122(15):154102, 2005.
Volker Deringer, Albert P Bartók, Noam Bernstein, David M Wilkins, Michele Ceriotti, and Gábor Csányi. Gaussian process regression for materials and molecules. Chemical Reviews, 121(16):
10073-10141, 2021.
Francesco Di Giovanni, Lorenzo Giusti, Federico Barbero, Giulia Luise, Pietro Lio, and Michael M. Bronstein. On over-squashing in message passing neural networks: The impact width, depth, and topology. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 7865 7885. PMLR, 23-29 Jul 2023, URL https: htttt:p://rroceeeee http:////ooceeeee/see/ee mlr. press w wttp:////2//iiiiooooooo222222222 html
10

---

Published as a conference paper at ICLR 2024
Vijay Prakash Dwivedi and Xavier Bresson. A generalization of transformer networks to graphs. arXiv preprint arXiv:2012.09699, 2020.
Thomas Exner and Paul Mezey. Ab initio quality properties for macromolecules using the adma approach. Journal of computational chemistry, 24(16):1980-199) 2003.
J Thorben Frank, Oliver T Unke, and Klaus-Robert Müller. Soōkrates-self-attenti for higher-order geometric interactions on arbitrary length-scales. arXiv preprint arXiv:2205. /4276, 2022a
J Thorben Frank, Oliver T Unke, Klaus-Robert Müller, and Stefan Chmiela. From peptides to nanostructures: A euclidean transformer for fast and stable machine learned force fields. arXiv preprint arXiv:2309. 15126, 2023.
Thorben Frank, Oliver Thorsten Unke, and Klaus Robert Muller. So3krates: Equivariant attention for interactions on arbitrary length-scales in molecular systems. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022b. URL Tttt:::://pperrviiee net /wt/teeuun/a wowumm t 10nxtAmcJa
Xinyu Fu, Jiani Zhang, Ziqiao Meng, and Irwin King. Magnn: Metapath aggregated graph neural network for heterogeneous graph embedding. In Proceedings of The Web Conference 2020, pp. 2331-2341. 2020.
Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. Se (3)-transformers: 3d roto- translation equivariant attention networks. Advances in Neural Information Processing Systems, 13:1970-1981, 2020.
Shridhar R Gadre and V Ganesh. Molecular tailoring approach: towards pc-based ab initio treatment of large molecules. Journal Theoretical and Compulational Chemistry, 5(04))835-855 2006.
Ganesh, Rameshwar K Dongare, P Balanarayan, and Shridhar R Gadre. Molecular tailoring approach for geometry optimization of large molecules: Energy evaluation and parallelization strategies. The Journal of chemical physics, 125(10):104109, 2006.
Johannes Gasteiger, Muhammed Shuaibi, Anuroop Sriram, Stephan Günnemann, Zachary Ulissi, C Lawrence Zitnick, and Abhishek Das. Gemnet-oc: developing graph neural networks for large and diverse molecular simulation datasets. arXiv preprint arXiv:2204.0277 027822 2022.
Andrea Grisaf and Michele Ceriotti. Incorporating long-range physics in atomic-scale machine learning. The Journal of Chemical Physics, 121,,0), nov 2019. doi: 10.1063/1.5128375 URL
https:://do http///oorgg//o///////////oo 11w...016611111.. 512837775
Feng Long Gu, Yuriko Aoki, Jacek Korchowiec, Akira Imamura, and Bernard Kirtman, A new localization scheme for the elongation method. The Journal of chemical physics, 120(21) 10391, 2004.
Jiaqi Han, Yu Rong, Tingyang Xu, and Wenbing Huang. Geometrically equivariant graph neural networks: survey, arXiv preprint arXiv:2202.07230, 2022.
Xiao He and John ZH Zhang. new method for direct calculation of total energy of protein. The Journal of chemical physics, 122(3):031 1003, 2005.
Xiao He and John ZH Zhang. The generalized molecular fractionation with conjugate caps/molecclar mechanics method for direct calculation of protein energy. The Journal of chemical physies, 124 118):1 2006.
Xiao He, Tong Zhu, Xianwei Wang. Jinfeng Liu, and John ZH Zhang. Fragment quantum mechanical calculation of proteins and its applications. Accounts of chemical research, 479))7778-27577 2014.
Pierre Hohenberg and Walter Kohn. Inhomogeneous electron gas. Physical review, 136(3B):B864, 1964.
Huiting Hong, Hantao Guo, Yucheng Lin, Xiaoqing Yang, Zang Li, and Jieping Ye. An attention- based graph neural network for heterogeneous structural learning. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 4132-4139, 2020.
11

---

Published as a conference paper at ICLR 2024
Ziniu Hu, Yuxiao Dong, Kuansan Wang, and Yizhou Sun. Heterogeneous graph transformer. In Proceedings of the web conference 2020, pp. 7704-2710, 2020.
Gyeong Hwang. Haley M Stowe, Eunsu Pack, and Dhivya Manogaran. Reaction mechanisms of aqueous monoethanolamine with carbon dioxide: combined quantum chemical and molecular dynamics study. Physical Chemistry Chemical Physics, 17(2):831-83) 2015.
Ferenc Karsai, Manuel Engel, Espen Flage-Larsen, and Georg Kresse. Electron Dlectroo-phhnon coupling in semiconductors within the gw approximation. New Journal Physics, 20(12): (2)1133000 2018.
Walter Kohn and Lu Jeu Sham. Self-consistent equations including exchange and correlation effects. Physical review, 140(4A):A1133, 1965.
Arthur Kosmala, Johannes Gasteiger, Nicholas Gao, and Stephan Günnemann. Ewald-based long- range message passing for molecular graphs. arXiv preprint arXiv:2303.04799, 2023.
David Peter Kovacs, Ilyes Batatia, Eszter Sara Arany, and Gabor Csanyi. Evaluation of the mace force field architecture: from medicinal chemistry to materials science. arXiv preprint 1Xiv:2305.14247. 2023.
Arpan Kundu, Marco Govoni, Han Yang. Michele Ceriotti, Francois Gygi, and Giulia Galli. Quantum vibronic effects on the electronic properties of solid and molecular carbon. Physical Review Materials, 5(7): 2021.
Arpan Kundu, Yunxiang Song, and Giulia Galli. Influence of nuclear quantum effects the electronic properties of amorphous carbon. Proceedings of the National Academy of Sciences, 119 (31):2220308111 19, 2022.
Greg Landrum et al. Rdkit: A software suite for cheminformatiis, computaiioa chemistry, and predictive modeling. Greg Landrum, 2013. URL BtLtttts://www. wwwit. .org/
Wei Li, Shuhua Li, and Yuansheng Jiang. Generalized energy-based fragmentation approach for computing the ground-state energies and properties of large molecules. The Journal of Physical Chemistry A, 11 111(11):193-2119 2007.
Wei Li, Haibo Ma, Shuhua Li, and Jing Ma. Computational and data driven molecular material design assisted by low scaling quantum mechanics calculations and machine learning. Chemical Science, 2021.
Yujia Li, Daniel Tarlow, Mare Brockschmidt, and Richard Zemel. Gated graph sequence neural networks. arXiv preprint
2015.
Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. arXiv preprint arXiv:2206. 19900, 2022.
Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic
graphs. In International Conference Learning Representations, 2023. URL https:// openreview net/ www.eeee/oooeedooeeeee. FARgOTD
Nicholas J Mayhall and Krishnan Raghavachari. Molecules in-molecusss: An extrapolated fragment- based approach for accurate calculations on large molecules and materials. Journal of chemical theory and computation, 7(5):1336-1343, 201
Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron Owen, Mordechai Kornbluth, and Boris Kozinsky. Learning local equivariant representations for large-scale atomistic dynamics. Nature Communications, 14(1 (4(1)9579 2023.
Maho Nakata and Toshiyuki Maeda. Pubchemqe 531yp/6-31g"//mmd dataset: the electronic structures of 86 million molecules using b31yp/6-31g* calculations. arXiv preprint arXiv:2305.1884 2023.
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data 1(1):17, 2014.
Michael Schlichtkrull, Thomas N. Kipf, Peter Bloem, Rianne van den Berg, Ivan Titov, and Max Welling. Modeling relational data with graph convolutional networks, 2017.
12

---

Published as a conference paper at ICLR 2024
Kristof Schütt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In International Conference on Machine Learning, pp. 9377-93888. PMLR, 2021.
Kristof T Schütt, Farhad Arbabzadah, Stefan Chmiela, Klaus R Müller, and Alexandre Tkatchenko. Quantum-chemical insights from deep tensor neural networks. Nature communications, 8(1):1-8, 2017.
Nima Shoghi, Adeesh Kolluru, John R Kitchin, Zachary W Ulissi, Lawrence Zitnick, and Bran- don M Wood. From molecules to materials: Pre-training large generalizable models for atomic property prediction. arXiv preprint (..ii:2211..68000 2023.
Attila Szabo and Neil S Ostlund. Modern quantum chemistry: introduction to advanced electronic structure theory. Courier Corporation, 2012.
Philipp Thölke and Gianni De Fabritiis. Torchmd-net: Equivariant transformers for neural network based molecular potentials. The International Conference on Learning Representations, 2022.
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang. Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and translation-equivarian neural networks for 3d point clouds. arXiv preprint 18X.v 19..8.2...11,, 2018.
Jake Topping. Francesco Di Giovanni, Benjamin Paul Chamberlain, Xiaowen Dong, and Michael M Bronstein. Understanding over-squashing and bottlenecks on graphs via curvature. In International Conference on Learning Representations, 2021.
Oliver T Unke, Stefan Chmiela, Huziel E Sauceda, Michael Gastegger, Igor Poltavsky, Kristof Schütt, Alexandre Tkatchenko, and Klaus-Robert Muüller. Machine learning force fields. Chemical Reviews, 121(16):01422-0118 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. CoRR, abs/1 2017. URL
http:// /aapiiv./aiiv.////a org/ www./bag/////oooooe/.../e..eeee/ / 03762
Hongwei Wang, Fuzheng Zhang, Mengdi Zhang, Jure Leskovec, Miao Zhao, Wenjie Li, and Zhongyuan Wang. Knowledge graph convolutional networks for recommender systems with label smoothness regularization. CoRR, abs/1905.04413. 2019a.
abs/1905 04413
Shike Wang, Fan Xu, Yunyang Li, Jie Wang, Ke Zhang, Yong Liu, Min Wu, and Jie Zheng. Kg4sl: knowledge graph neural network for synthetic lethality prediction in human cancers. Bioinformatics, 2021.
37(Supplement 1atteeett___1111181
Xiao Wang, Houye Ji, Chuan Shi, Bai Wang, Yanfang Ye, Peng Cui, and Philip S Yu. Heterogeneous graph attention network. In The world wide web conference, pp. 2022-2032, 2019b.
Yusong Wang, Shaoning Li, Xinheng He, Mingyu Li, Zun Wang. Nanning Zheng, Bin Shao, Tong Wang, and Tie- Yan Liu. Visnet: scalable and accurate geometric deep learning potential for molecular dynamics simulation. arXiv preprint arXiv:2210.16511 2022a.
Yusong Wang, Shaoning Li, Tong Wang, Zun Wang, Xinheng He, Bin Shao, and Tie-Yan Liu. An ensemble of visnet, transformer-m, and pretraining models for molecular property prediction in ogb lcrge-scale challenge@ neurips 2022. arXiv preprint wtXiv:2211.1299) 2022b.
Zun Wang, Chong Wang, Sibo Zhao, Yong Xu, Shaogang Hao, Chang Yu Hsieh, Bing-Lin Gu, and Wenhui Duan. Heterogeneous relational message passing networks for molecular dynamics simulations. npj Computational Materials, 8(1):33, 2022c.
Zun Wang. Hongfei Wu, Lixin Sun, Xinheng He, Zhirong Liu, Bin Shao, Tong Wang, and Tic- Yan Liu. Improving machine learning force fields for molecular dynamics simulations with fine-grained force metrics. The Journal of Chemical Physics, 159(3), 2023.
Weitao Yang. Direct calculation of electron density in letters, 66(1 14:188, 1991.
theory. Physical review
13

---

Published as a conference paper at ICLR 2024
Linfeng Zhang, Jiequn Han, Han Wang, Roberto Car, and EJPRL Weinan. Deep potential molecular dynamics: a scalable model with the accuracy of quantum mechanics. Physical review letters, 120 (14):143001, 2018.
Linfeng Zhang, Han Wang, Maria Carolina Muniz, Athanassios Panagiotopoulos, Roberto Car, and Weinan E. A deep potential model with long-range electrostatic interactions. The Journal of Chemical Physics, 156(12) 2022.
14

---

Published as a conference paper at ICLR 2024
Contents
Additional Related Work
17
LSR-MP can alleviate Over-squashing
17
Justification of the Method
18
Pitfalls of Fragmentation- Methods
20
Additional Experiments
20
PubChem
20
Electrostatics Binding Energy
E.2.1
Decay of Interactions
21
Study on Transferability
E,4
Simulation
22
E,5
Examine Curvature of BRICS Prioritized Edges
22
E.6
Comprehensive Force Evaluation
23
E.7
Comparison Different Fragmentation Schemes
25
E.8
Attention Weight Analysis
26
Applicability to Standard Graph Clustering Method
27
10
Ablation study of the long-range cutoff
27
Impact of Fragment Size
28
Geometric Optimization
29
F.1
Geometric Optimization Acceleration with Hybrid Models
29
G BRICS and Modification of BRICS
30
Long-Range Module Visualization
32
Proof of Equivariance
32
J I Loss Function
36
Hyperparameter setting
36
K.1 Hyperparameter in Table
36
K.2 Hyperparameter in Table
36
Higher-order LSRM
36
Preliminary
37
Notation
38
Rhort-Range Module
38
15

---

Published as a conference paper at ICLR 2024
L.4
Fragmentation Module
39
Long-Range Module
39
L.6
Properties Prediction
39
Equiformer-LSRM
39
16

---

Published as a conference paper at ICLR 2024
A ADDITIONAL RELATED WORK
Learning Heterogeneous Graphs: Heterogeneous graphs are vital in recommendation systems Wang et al. (0) and scientific fields like molecular dynamics Wang et al. (022cc) and systems biology Wang al. (002T) Heterogeneous Graph Neural Networks (HGNNs) and metapath-free methods. Metapath-
divided into two HGNNs first gather features from similar semantic neighbors and then combine different semantic types. For example. Relational GCN Schlichtkrull et al. (2017) uses unique weights for each relationship type. HAN Wang et al. (2019b) leverages metapaths differentiate semanties, merging structural details from each metapath in the aggregation process, and then combining these details for each node. MAGNN Fu et al.] (2020) expands this idea by considering all nodes in metapath instance. On the other hand, metapath free HGNNs, like traditional GNNs, collect information from all neighboring nodes simultaneously. Among metapath-free HGNNs, HetSANN Hong et al. (2020) uses a multi-layer network to create specific attentions for different relationships, and HGT Hu et al. (2020) introduces the Transformer model to handle different types of nodes and edges.
types: metapath-
Learning Long-Range Dependency: Recent advancements in graph neural networks (GNNs) have shown promising developments in addressing long-range dependencies within graph structures, a challenge historically difficult due to the inherent limitations of local neighborhood aggregation. Pioneering approaches, like those presented by Li et al. 2015)) introduced Gated Graph Sequence Neural Networks, leveraging gated recurrent units to better capture long-range interactions. The emergence of Graph Transformers, as proposed by Dwivedi et al. (2020) Dwivedd & Bresson (2020), marks significant leap forward, employing self-attention mechanisms directly model relationships between distant nodes. Additionally, works such as those by bylonn & Yahav (2020) on spectral-basee GNNs have demonstrated innovative ways to encapsulate long-range dependencies by extending traditional convolutional methods. Di Giovanni et al. (2023) introduced curvature- -ased rewiring methods to handle long-range dependency. the realm of equivariant graph neural networks, [Frank et al. (2022a) introduced extra neighbors in the hidden space to capture long-range interactions. Kosmala et al. (202)) introduced Ewald summation, which imposes a Fourier transformation with frequency cutoff to capture long-range interactions.
B LSR-MP CAN ALLEVIATE SVER-SQUASHING
The definition of information over squashing is defined as symmetric Jacobian obstruction Di Gio- vanni et al. (2023). The mathematical representation of this concept is given by:
1 Ohm d, ah
Oh
J (v.
do Dht
A quantity to characterize oversquashing
Sensitivity to myself
Sensitivity nodes
where m > & both index layer numbers.
Proposition B.1 (EGNN module over-squashing is dependent on commute time Di Giovanni et al. (2023)). Given a short-range module, with Sy == c,I + A as the graph shift operator adopted by the short-range module. Let O (v, u) be the symmetric Jacobian obstruction of nodes v, u after m layers. Assume that each path in the computational graph is activated with equal probability p. Let and be the maximal spectral norm and minimal singular value of the weight matrices of the short-range module, Let 2.. If plc, + Ca) S 1, we have:
((GShhrt, AT) u) S (lm)) (v. u) & AT(1, w) < 2
wherein T(C, u) referred to as the commute time, quantifies the expected number of steps a random wall originatiin at node v would take to reach node and return. In certain molecular structures, this commute time can be considerably large. Consider, for instance, an elongated carbon chain or unfolded amino acid chain with hundreds of amino acids, the commute time spanning from end to the other could be ((EE)).
Random walks refers uniform sampling of the adjacency matrix each step with backtracking enabled.
17

---

Published as a conference paper at ICLR 2024
Proposition B.2 (Long-range module alleviate over-sqashing). Given long- the same assumption holds as in Proposition [.I] if floug could cover the whole molecule, the Long-range obstruction function is upper bounded by a function independent of dist(pi. P)))
long-range range module, with
2p
(v, w) S
J.Ca
Proof. The commute time could be expanded as:
TVv u) = 2
where R(U, u) here refers to the voltage difference if unit current flows from v to u, with each edge assigned with unit resistance.
In long-range module:
Ve, u E V, R(v, u)
R1
R2
where R R2 refers to the length of the paths from e to u, the total number of paths.
We proceed to find a path from v, u in LSR-MP. Recall the fragmentation module introduces virtual nodes contained within a fragment. Hence, in the long-range graph, each node is connected to the virtual nodes of its corresponding fragments. Furthermore, if rlong could cover the entire molecule, this translates to a complete bipartite graph where each node connected to each virtual node. Hence, this gives:
nodes connecting
Ve, u, 3i E [n]. such that R² = 2.
This gives an upper bound on R(e, u):
R(v, u)
Rmin S R₁ =2,
R1
R2
R.
T(U, u) S 4/44,,
pT (T v)
2p
(mm) (v, u) S
2/cc. E
p.C.
C JUSTIFICATION OF THE METHOD
The total energy molecule consisting of atoms writes,
E (N) E(p .P.)
(16)
where pú represents the position of the i-th atom. Due to the complexity, the exact form of the above function is unknown with limited exceptions of very small systems. Our development starts from the ansatz that the total energy of the N N-body system can be decomposed as follows,
LIN
(p.) ΣΣ
(
Σ
(3) (9.
(17)
i<jck
where (1)) is the energy associated to a single atom, and (2) and ((3) are the two- and three-body interactions, respectively. We omit higher orders of interactions in our derivation for simplicity.
To model the total energy function in Eq. 17 existing EGNNs typically adopt the strategy that neglects the interactions beyond a truncation radius ssaart
2 pT) 0, if dist(p., (75)
Tshort.
(3) (p. Pk)) = 0, if dist(p.. p,) > Fahort or dist(p,,pul pk) > Fshort or dist p.) > Fshort:
(18)
18

---

Published as a conference paper at ICLR 2024
where dist(:,)) )) is metric function. As result, Eq. 17 is approximated as,
N
(p.)
N
e (2) (1s, 22)00rshonn - dist(pi [3j]
N
N
N
e Pa) PE) - dist - [j)]][[short - dist( - -.k) [r))a(rsharr dist - [.]]
KjjK
(19)
where 0||] is the Heaviside step function.
The neglect of the long-range interactions results in loss of accuracy, and detailed discussion can be found in the main text. To accommodate long-range interactions, naïve method is to further increase the truncation radius Ishort in Eq. [9] However, this idea is known to suffer from the loss of data efficiency and prohibitive computational cost system size increases.
To resolve this issue, we adopt a long-short strategy that introduces a second truncation radius Flong- which much larger than Tshorts and the interactions between atoms modeled according to their distances relative to rshort and Tlong- In particular, the two- and three-body interactions can be approximated as follows.
(2 py), pp).
dist(p:.Pp) S Tshort
(p..
Tshort < dist(p., p,)) S Tlong
0.
otherwise
cs Pk),
dist(pi. P,) S Tshort and dist(p,, Pk) S rshort and dist(pk, Pi) S Tshort dist(p:.P3) Plong and dist(pj,Pa) > Plong and dist(pe, p.) Plong
(p. Pk.
~
0, Cs pk),
otherwise
(20)
where and e' are approximated forms of the interactions to be learned with neural networks. Our method is generalization of existing EGNNs and can be reduced to EGNN in the limit of Tlong = shoort- Apparently, the introduction of second truncation radius Flong covers significant amount of long-range interactions that do not present in Eq. 19
Now we turn to the implementation the method with neural networks. We first assume that the local environment of each atom and the interactions for atoms within the short truncation radius Tshort can be learned through short-range model,
h; SHORT RANGE DESCRIPTOR (pt)
E2 E2 57) SHORT RANGE INTERACTION(1;; (hi; ht) 6g (7, Pk) = SHORT RANGE INTERACTION (h; hi, hx)
(21)
where h; is the descriptor of the local environment for i-th atom. The long-range interactions are assumed to be learned also from the local descriptors,
(p. ATOMWISE LONG RANGE(R,, hT) (P. = ATOMWISE LONG RANGE(h;, h, h2r)
(22)
The interactions are written directly in terms of atomic descriptors, thus we name them atomwise long-range models. We note that the atomwise range models must not be the same the short- range ones, as it makes our method an EGNN with larger cutoff. Without doubt, the computational cost of atomwise long-range model quickly becomes prohibitive with respect the system size.
19

---

Published as a conference paper at ICLR 2024
Thus, we divide the system into multiple fragments, whose descriptors be learned from atomwise features,
fa FRAGMENT DESCRIPTOR((h,... Fa.
(23)
and approximate atomwise long-rang range interactions with Klom-fragment ones: <2 (5) + eg Pk. +...
ATOMWISE LONG RANGE(h,, h1) ATOMWISE LoNG RANGE(h;, ltg, ha)
ATOM FRAGMENT LONG RANGE(H,, FRAGMENT DESCRIPTOR( (( hys. }}) ATOM FRAGMENT LONG RANGE(h,, fa).
(24)
It is worthy noting that the atom-fragment long-range model only explicitly approximates two- fragment interactions. While long-range fragment-fragman interactions higher orders are assumed be negligible. short-range interactions that involve more than fragments are implicitly encoded in the short-range model (Eq. 21).
D PITFALLS OF FRAGMENTATION-BASER METHODS
Fragmentation-base based methods offer scalable solutions for quantum mechanical problems by breaking down large systems into computationally manageable pieces. However, capturing many-body effects remains challenge. This appendix elucidates the reasons behind these limitations, focusing on the many-body expansion (MBE) as representative example.
The MBE represents the energy system divided into N fragments as:
Evoal Σ E,
EEg +
Σ
AEijk +...
(25)
iejek
Where:
E, denotes the energy the ith fragment computed in isolation.
Eaj represents the two-body interaction energy between fragments i and
A Fajk signifies the three-body interaction energy among fragments i, j. and k, and so on.
Truncation Errors: Including higher-order terms often leads to truncation after a certain order (e.g.. three-body terms):
Eranee ΣE+ AE - Σ AEigh
(26)
KS
Kjjc
For systems where higher-order terms are essential, this truncation fails depict the many-body nature of interactions.
Non-additivity of Many-Body Interactions: Many-body interactions are inherently non-al non-additive, making their accurate representation in fragmentation methods challenging.
Basis Set Inconsistency: Basis Set Superposition Error (BSSE) arises due to overlapping basis functions when fragments are combined, distorting the true many-body interactions.
Mutual Polarization: Many-body effects often emerge from mutual polarization of multiple frag- ments. When fragments are treated in isolation or only via pairwise interactions, this mutual polarization missed.
In conclusiton, while MBE offers systematic framework for accounting for fragment interactions, its practical implementation in fragmentation methods fraught with challenges. The inherent approximations, truncation of the series, and non-additive effeects introduce difficulties in capturing the true many-body nature molecular interactions.
E
ADDITIONAL EXPERIMENTS
E. PUBCHEM
PubChemQC B3LYP/631G+//PM6 database Nakata & Maeda (202),, which contains electronic properties calculated using density functional theory for 85.9 million small molecules from PubChem.
20

---

Published as a conference paper at ICLR 2024
To testify gaggraage interactions, extracted molecules with molecular weight larger than 500 to perform training. The atom number of the dataset ranges from 40 to 100. We used 26000 molecules as the training set and 8000 molecules the test set. The experiment results are included in Table 5
E.2 ELECTROSTATICS BINDING ENERGY
The binding energy of the charged dimers dataset & Ceriotti (2019) consists of 661 diverse organic molecular dimers containing H. C. N, and atoms, with at least one monomer each dimer carrying net charge, extracted from the BioFragment Database. For each dimer, the dataset provides 13 configurations with varying inter-monomer distances from 3-8 A, and the corresponding binding curves (interaction energies versus distance) are calculated using DFT. 600 dimers are used for training machine learning models and 61 for testing. Isolated monomers are also included to provide the dissociation limit. This dataset offers a realistic challenge for assessing model performance in predicting binding curves dominated by long- mange electrostatic interactions across wide range of chemical environments. The test results are attached in Table
Table 5: Energy MAE (kcal/mol)
Electrostatics Binding Energy PubChem
Vi& VISNeI-LSRM
2.978
0.0654
2.012
E.2.1 DECAY OF INTERACTIONS
We studied the decay of interactions by separating two molecules in dimer configuration and plotting the energy function of distance. The results attached Fig- 4 These experiments demonstrate that, compared to local model, our model exhibits a more appropriate decaying behavior. This finding crucial as suggests that our model captures the long-range interactions more effectively.
RoGa rocal
OPPREOO
0.15
(
0.10
0.05
0.00
[A)
R(A]
Figure 4: Visualization of the decay interaction for electrostatics binding energy dataset. a) The decay curve of one CH₂COO and 4-Meetylimidazole. b) The decay curve of two CH₂COO
E.3 STUDY ON TRANSFERABILITY
In this section, we targeted examine the transferability extrapolation capacity of our model to larger unseen molecules. In particular, we performed three experiments to illustrate this:
Zero-Shot Experiment: study transferability, our first experiment was a zero-shot setup. We trained molecules including ATAT, Stachyose, DHA, and Ac-Ala3-NhMe, and then tested directly on larger molecule, ATATCGCG. The zero-shot results are shown in Table 6 This experiment revealed that direct transferability without demonstration challenging for MD22 trajectories.
Few-Shot Learning Experiment: To further explore transferability, we conducted a few- shot learning experiment, shown in Table 6 By adding small set of 50 ATATCGCG
21

---

Published as a conference paper at ICLR 2024
training samples the original zero-shot training model demonstrated significant improvement over the baseline model. This suggests that with minimal additional training data, our model can adapt to larger molecular systems more effectively than traditional local models.
PubChem In our study, we further assessed our model's capabilities using the PubChem dataset, as elaborated in Appendix E.I The dataset features heterogeneous molecules of size ranging from 40 to 100. We recalculated the dataset using 1-zvp as the basis set to improve accuracy. Notably, included molecular force, which remains informative signals given that the molecules were relaxed only through semi- cemi-empirical approach. For dataset division, we used molecules with fewer than 60 atoms (30,545 samples) for training and those with more than 60 atoms (3,455 samples) for testing. Our results in shown in Table 6 Compared to the baseline ViSNet model, our model showed enhanced performance on larger molecules, underlining its robust transferability and wide applicability in diverse molecular contexts.
Table 6: Transferability Experiments
Experiment
Metric
ViSNet
ViSNet-LSRM
Energy 184.44
150.23
Zero Shot
Force
10.93
10.21
Energy
2.575
2.167
Few Shot
Force 0.7448
0.6556
Energy 4.458
3.339
PubChem
Force 0.3303
0.2395
E.4 MD SIMULATION
We performed an MD simulation for relatively large molecule, ATAT, for 20ps, matching the duration of the AT-AT simulation in the MD22 dataset. This was done a constant energy ensemble (NVE). These simulations were driven by our ViSNet-L.SR and DFT with time step of T = fs. allowing us to analyze the vibrational spectra of the AT-AT molecule. As depicted in Fig. 5 both the trajectory in MD22 and the trajectory simulated by ViSNet-LSRM show similar vibrational spectra, albeit with minor differences in peak intensities compared to DFT. This suggests that our simulations accurately mimic the actual vibrational modes of the molecules over relatively long time periods.
To test the smoothness of ViSNet-LSRM, we ran longer 200 ps NVE simulation with a time step of r Ifs for this molecule. The total energy profile is displayed in Fig. 6 The total energy is conserved within a reasonable range (+ - 0,0001%) of fluctuation, validating the capability of the proposed ViSNet-I.SRM under long simulations.
E.5 EXAMINE CURVATURE OF BRICS PRIORITIZED EDGES
To investigate the capacity of BRICS methods to prioritize edges characterized by high negative curvatures, we analyzed the Balanced Forman curvature on two datasets: MD22 and chignolin. In particular, we classified an edge spanning two fragments as BRICS-prioritizer edge. Conversely, edges not meeting this criterion were labeled as Von-BRICS- SereCCS-rirritzeedd edges. In Figure we offer a visualization the fragmentation outcomes in relation to curvature. Here, the node color scheme represents distinct fragments resulting from BRICS, while the edge coloring reflects the curvature value. This visualization underscores the tendency of BRICS fragmentation to give precedence to edges with negative curvature. To substantiate this observation, on curvatures across the six systems under scrutiny (See Figure/B3 Consistently, our findings affirm the propensity of BRICS prioritize edges with negative curvature.
applied the Mann-Whitnes Whitnhhhhyyyyy U-Test
22

---

Published as a conference paper at ICLR 2024
Velocity auto-correlation function of AT-AT
DFT
MD22
ViSNet-LSRM
Density of states (arb. units)
Wwmmmm
250
500
750
1000
1250
1500
1750
2000
Frequency (cm-1)
Figure 5: Velocity autocorrelation function of ATAT.
Total Energy (Etot) vs. Steps
50083395
-50086.09
Elo. (40
50084-22
30004-36
Figure 6: Visulization of Total Energy. Each step represents Ifs. The fluctuation is within 0.0001% % percentage of the total energy.
E.6 COMPREHENSIVE FORCE EVALUATION
We have incorporated additional force evaluation metrics, drawing from the methodologies presented by Wang et al. (2023), which encompass:
Global Metrics
Mean absolute error (Fmac)
Max absolute (Fmax-err)
Mean normalized error (FNMmae)
Max normalized error (FNMmax-err)
Clement-Based Metrics:
Fmae, Fmax-err, FNMmae, FNMmax-err for each element type. The results shown in Figure
23

---

Published as a conference paper at ICLR 2024
A
0
-4
-6
-4
22
y
0
4
Sachys
DIK
08-
0.
9.-
22-
-95
-16
-10-
Oliguin
60.5
-2.4
-14
-6.6
-10
-15
-0.5
-1.0
BRIS Prrilaed NoN-RUCS
BIICS Parflaed
Merdil Nerlool
1RSSSPortoed Vowell Pretued
Figure 7: A. A visualization of BRICS fragmentation and Balanced Forman curvature. The blue edge indicates negative curvature. The red edge indicates positive curvature. B. Statistical testing of BRICS prioortzie edges and Non- Non-BRICS prioritized edges. *** means p-value < 0.001
24

---

Published as a conference paper at ICLR 2024
Carbon C_f mae
Nitrogen- N.m.mee
Ooygen -O_f mee
3.75
2.23
2.125
5100
2475
0.4
0.2
5200
2200
Stadhyoe
Stadyyse
Statfyoe
Carbon C_f_mae
N_fr_mae
Oxygen 0_ffr_mee
0.80
0.35
0.6
0.50
0.25
0..0
0.15
0.0
02
A.Al2 10-90
DHA
Stadhyood
ATAT ATATCOOO
DHA Stadyooo
ATAT ATATOUOO
Acd DHA Staclosow
ATAT
ATATCOCU
Carbon max
Nitrogen
Orygen O
08
6.5
6.
01
010
AnAbb3-4
ATAT ATATOGOG
80
Acsko1-4
Stachyooe
ATAT ATATOGOG
60
Abdia308 16e
DIA
Stachyooe
ATAT
ATATOGOG
Carbon- C__frmaac
Nitrogen -N. Nafr.max
Oxygen-O O0.frm0mee.
100
200
17.5
12
200
58
25
00
Stachyose
Stadhyose
Staclyoee
Figure 8: Element-based force evaluation for ViSNet-LSRM and ViSNet. The dark color indicates ViSNet-LSRM and the light color refers to ViSNet. It could be seen that our models outperform the baseline ViSNet in terms of Mean absolute error (Fmae), Max absolute error (Fmax-err), Mean normalized error (FNMmae), and Max normalized error (FNMmax-err) in most cases. Also, is worth notifying that ViS has greatly reduced the prediction error carbon atoms when compared to ViSNet.
E.7 COMPARISON OF DIFFERENT FRAGMENTATION SCHEMES
The sensitivity of our LSR-MP framework to different fragmentation schemes is a crucial aspect of research. Fragmentation schemes be broadly categorized into two groups: knowledge-
25

---

Published as a conference paper at ICLR 2024
based fragmentation methods, such as divide-and wwide-and-conquer (DC) and molecular fractionation with conjugated caps (MFCC), and conventional graph clustering techniques like k-means clustering.
We have conducted comprchensive analysis of several knowledge-basee based fragmentation methods, including DC and MFCC, which is shown in Table 7 Our findings indicate that the LSR-MP framework is generally robust to various fragmentation schemes. However, some schemes may exhibit superior performance depending on the specific system under investigation.
In addition knowledge-basee based methods, we have explored the use of conventional graph clustering techniques by incorporating k-means clustering as an alternative to fragmentation-basee methodolo- gies within our L.RRMP framework. Our experiments show that knowledge- approaches, which draw upon chemical domain expertise, generally outperform k-means clustering methods for most molecules. Nevertheless, the LSR-MP framework, when combined with k-means clustering, consis- tently surpasses comparable models that do not utilize the LSR-MP approach. Moreover, our results for two supramolecules, employing k-means and distance- spectral clustering. significantly exceed the performance of equivalent baseline methods.
In conclusion, knowledge-based fragmentation approaches generally outperform k-means clustering methods for the majority of molecules, as k-means is distance-based clustering method that does not consider chemical properties like atom types and bond types, potentially resulting in chemically insignificant fragments. Furthermore, the LSR-MP framework, when combined with various fragmentation schemes, demonstrates better performance than baseline methods, highlighting the versatility and broad applicability of our method.
Table 7: MAE for different Fragmentation schemes Biomolecules of MD22, the best- perst-erformin methods are highlighted bold.
Molecule
Metries
LSRM MFCC
LSRM Divide and Conquer
LSRM Kmeans
LSRM Brics
Ac-AAla-NhMMe
Energy
0.0637
0.0824
0.0662
0.0654
Force
0.0928
0.1064
0,0956
0.0942
DHA
Energy
0.0815
0.1374
0.0966
0.0873
Force
0.0562
0.0742
0.0620
0.0598
Stachyose
Energy
0.1295
0.1259
0.1199
0.1055
Force
0.1016
0.0904
0,0821
0.0767
AT-AT
Energy
0.0772
0.1081
0.1033
0.0772
Force
0.0790
0.0929
0.091
0.0781
AT-AT-CG-CG
Energy 0.1135
0.1438
0,1446
0.1135
Force
0.1064
0.1421
0.1476
0.1064
E.8 ATTENTION WEIGHT ANALYSIS
We analyzed the attention coefficients obtained in our model to establish the connection between model predictions and interpretability.. In particular, we extracted the attention weights in the long- range modules to study the atom- kkom-framment interactions in AT-AT-CG-CG. For each attention head, we visualized the atom-fragment interactions with the largest attention weights, as shown in Figure 9 The long-range modules attend some short-range interactions such as hydrogen bondings (N-NH2) as well as long-range interactions (C-C5H30N4, C-C5H3N4). This is compatible with the physical intuition that hydrogen bonds are essential components in the nucleic acids base- mase-pairing system. This also suggests a two-fold contribution of long-rang models: (1) they explicitly characterize the long- range interactions and (2) they partly restore the information lost in short-range further studied the attention weights averaged over atoms and fragments, which are shown in Figure 10 (a.b). We find that the interactions between N-CH3 pair have the smallest attention coefficients, and this can be interpreted group. We then visualized the two atom-fraae framfframmett interactions with the largest weights in Figure I0 (c.d). Figure (c) suggests the interactions between nitrogenous bases and distant oxygen atoms, counterintuitively, play significant roles model predictions, which further corroborates our model's capacity to capture interactions beyond local environments.
cesssage-passing We
the polarity difference between the nitrogen atom and the methyl
26

---

Published as a conference paper at ICLR 2024
N NH2
H NH2
C5H3ON4
C- NH2
C4H5ON3
C- CSH3N4
Figure 9: Visulization of atom-fraan -rom-framment interactions with large attention coefficients in AT-AT-CG-
CG. Central atoms are denoted in red, and the corresponding fragments are shown in yellow.
E.9 APPLICABILITY TO STANDARD GRAPH CLUSTERING METHOD
For two supramolecules in MD22, Buckyball catcher and Double-walled nanotube, their conforma- tions do not fall into the chemical prototypes specified in BRICS, leading to a failure fragmentation. To address this issue, we employ K-Means clustering for Buckyball cutcher and distance-based spec- clustering for Double-walled nanotube. As demonstrated in Table ?? ViSNet-LSRM achieves competitive performance compared to other EGNN methods for these result indicates that our framework is compatible with standard graph clustering methods, which could make the method more universally applicable. Nevertheless, the development of general fragmentation algorithm for such supramolecules warrants further investigation. In addition, we have conducted comparative analysis of fragmentation schemes focusing the BRICS method and standard graph clustering applied to biomolecules in MD22. The results are included in Appendix F
supramolecules. This
Table 8: Mean absolute errors (MAE) of energy (kcal/mol) and force (kcal/mol/A) for two supramolecules on MD22 compared with state-of-the-art models. The best one in each category is highlighted in bold.
Molecule
atoms
<GMML
PaiNN
TorchMD NET
Allegro
Equiformer
ViSNet
ViSNet-
Buckyball catcher
energy forces
0.4563
0.5258
0.6820
0.4098
0.3318
0.0887
0.1114
0.1335
0.1026
Double-walled nanotube
370
forces
energy
4.0122
0.5231
0.9168
1.1124
1.4732
1.0031
0.3428
2,2007
0.2747
1,1945
0.3959
0.3391
E. ABLATION STUDY THE LONG-RANGE CUTOFF
As shown in Table. the performance of ViSNet-LSRM improves when the long-range cutoff changes from 6 to and fluctuates slightly as it continues to increase. This is likely because all relevant fragments have already been included within 9A, and further increase does not introduce extra information. In addition, large long-range cutoff does not significantly increase the computational cost or lead to the information over-squashing. since the number of fragments is small, When dealing with larger molecules, increasing the long-range cutoff may be useful and still efficient. When compared with original ViSNet with 5 layers, ViSNet-LSRM with layers has similar inference time and better performance. In conclusion, all studies suggest that our LSR-MP framework is extremely
27

---

Published as a conference paper at ICLR 2024
(a)
0.575 (b)
0.15
0.13
0.475
0.450
0.11
-0.09
0.377
CEH302N2
CAHSONS
CSH3N
CSHION
H2W
SON
Che
043
16
(c)
(d)
10.0
75
300
0.0 N
5.0
5.0
18.0
-15
10.0
7400
-7.5 .50_0_3 nn
0
10.0
-2.5
-5.0 -7.
7.5
100
C5H33N4
C5H3ON4
Figure 10: Visulization of attention coefficients in AT-AT-CG-CG, (a) attention coefficients are averaged by atoms and fragments. (b) attention coefficients are only averaged by atoms. (c. d) visualization of the two most salient interactions suggested by (a).
efficient and effective for modeling the long-range interactions rather than either deepening the model or increasing the short-range cutoff.
Table 9: Study of long-range cutoff in LSR-MP framework AT-AT-CG-CG in MD22 dataset. The best results are shown in bold.
Long-range cutoff'
Energy MAR (kcal/mol)
Forces MAR
Inference time
(A)
(kcal/mol/AA)
(ms)
6
0.1234
0.1100
12.02
9
0.1
0.1064
12.14
12
0,1117
0.1074
12.26
15
066
0.1116
12.26
E. IMPACT OF FRAGMENT SIZE
In Figure we investigate the relationship between the Average Fragment Size AT-AT-CG-CG and two evaluation metrics: Force MAE and Energy MAE. As the fragment size increase, both metrics exhibit first decrease and then increase trajectory. Notably, both the Force MAE and Energy MAE share a similar trend, emphasizing the significance of fragment size in influencing these outcomes. This observation underscores the significance of choosing optimal fragment size.
28

---

Published as a conference paper at ICLR 2024
Force MAE and Energy MAE vs. Average Fragment Size
Force MAE
0.16
Energy MAE
0.15
0.14
MAE
0.12
0.11
0
10
15
20
Average Fragment Size on AT-AT-CG-CG
Figure 11: Visulization of the relationship between the Average Fragment Size on AT-AT-CG-CG and Model Errors
F GEOMETRIC OPTIMIZATION
F.1 GEOMETRIC OPTIMIZATION ACCELERATION WITH HYBRID MODELS
Geometric optimization, often referred to energy minimization molecular mechanics optimiza- tion, computational technique used in molecular modeling. The primary goal of geometric optimization is to find the stable or "lowest-energy" structure of molecule. This is accomplished by iteratively adjusting the atomic positions to minimize the potential energy of the system, usually using gradient descent or related algorithms. Once the molecule reaches a state where the force on each atom is essentially zero, the molecule is said to be in its optimized geometry or at local energy minimum. In our evaluation, we randomly sampled 5 initial configurations from the MD22 AT-AT-CG-CG molecule's 's test set. Using m06-2x DFT, we defined a 'reference geometry' through geometric optimization. To assess the neural potential model, initially leveraging the neural potential for optimization, and then refining with DFT until conver- gence. Our key metries the DFT iteration counts required for refinement and the rottmean--suare deviation (RMSD) between the neural potential potential-convergel geometry and the reference. Fewer DFT iterations suggest the model's practical utility in GO acceleration. Concurrently, the RMSD provides a direct measure of the model's reliability in mirroring DFT PES. As depicted in Figure [2 ViSNet- LSRM, in comparison ViSNet, more accurately replicates the DFT-converged geometry. Moreover, ViSNet-LSRM substantially diminishes the number of DFT steps required in hybrid models when compared to ViSNet, achieving acceleration rates of up to 50%.
adopted two-stage approach:
Geometric Optimization (GO) and Molecular Dynamics (MD) are pivotal computational method- ologies designed to probe molecular systems and kinetics. While GO primarily scouts the energy landscape for minimal energy configurations. MD provides insights into the temporal changes of molecular structures, accommodating specific thermodynamic ensembles like NVT and NPT. The core objective of both methodologies is to elucidate the dynamic attributes of molecular systems, positioning them as critical tools for understanding kinetics in computational studies.
Our study leveraged Density Functional Theory (DFT) not supervised learning component but as an evaluation measure for optimized geometry. We utilized force fields trained on the MD22 dataset, inferring kinetics from potential energy surface gradients. The GO process integrated approximately
29

---

Published as a conference paper at ICLR 2024
Comparison
b
Comparison
D0
city prutersiin
20
Sheee
1,
Geometric Optimization. x-axis represents five different initial configurations. (a) Comparison of Neural Optimized Geometry when juxtaposed with the DFT optimized Geometry (b) Comparison of the number of DFT steps required in Hybrid Models.
Figure 12:
100 trajectories, all distinct from the training set, emphasizing the task's ' relevance in showcasing the force field's generalization and kinetic representation.
Considering five structures in our analysis, each GO trajectory required about 20 minutes on 24-core CPU, culminating in week of computation on 128-core CPU using the PySCF software for DFT tasks. Although computationally intensive, we deem this effort critical for robust insights. We also ensured that the initial configurations for GO were excluded from the Machine Learning Force Field (MLFF) training to avoid overlaps.
In essence, GO operates as a time-series mechanism, transitioning from specific conformation 81 to 8, until convergence. Importantly, each step in this process is based on the conformational position of the previous step and the corresponding forces. This is fundamentally similar to the working principle of MD (Molecular Dynamics), whereas MD under specific thermodynamic ensemble conditions, such as NVT (constant temperature and volume) or NPT (constant temperature and pressure). For better understanding, enclose a pseudocode for GO as follows:
For each initial structure (not included the training set), we commence by performing Geomet cCpt imi zationCG until convergence using DFT FF/ViSNet FF/V FF. Notice that ViSNet FF and ViSNet LSRM FF were trained on MD22 DFT labels.
We extract the last item in the crajectoryliss which refer as the converged geometry.
We proceed by comparing the rmsd of the DFT Converged Geometry and initial geometry, ViSNet Converged Geometry and the DFT Converged Geometry, ViSNet-L.RRM Converged Geometry and the DFT Converged Geometry.
G BRICS AND MODIFICATION OF BRICS
The BRICS method is fragmentation technique designed identify local chemical environments indicated by link atoms of different types. By breaking active/weak bonds, a series of small, active fragments is produced. BRICS takes the chemical environment of each bond type and the surrounding substructures into consideration, resulting in fragment assignments that are more in line with chemistry and reducing the energy loss caused by bond breaking. However, this method can produce too-smail fragments, just one two atoms. To address this issue, a minimum fragment size and maximum fragment size are set, and any fragment smaller than the minimum merged with the smallest neighboring fragment if their is less than the maximum. This greatly reduces the number of small fragments. A visual representation of BRICS fragmentation results is shown Figure [4]
The Breaking of Retrosynthetically Interesting Chemical Substructures (BRICS) method is one of the most widely employed strategies in the communities of quantum chemistry, chemical retrosynthesis, and drug discovery. We summarize the key points of BRICS as follows:
30

---

Published as a conference paper at ICLR 2024
Algorithm 1: Geometric Optimization with Conjugate Gradient
Data: molecule, forceFieles energy/ olerance, gradientTolerancee, manGradient) olerance, max I terations
Result: currend.Structure. trajectory List. begin
trajectoryList Initialize:
currentStructure t poo.eeull.getnnntiiSSStnee
currentEnergy f ce ield wwwppueeeeegggyuureeeeeeeeeeeeee
wow.puuueecnnnennomuuuuuuuttturr serechDireetion + currentrradientt
currentGradient + force Field. www.e.e..oo.....
iteration +- 0:
converged + False;
while not converged and iteration < maxiterations do
10
a + LineSenchccurennSruuurr,, searchDirection, forceField);
newStructure currentStructure
searchDireectin;
12
newEnergy + force ield, wwwp://oeeeeeeeeeeeeeeestttteee
newGradient wwwoeeeeeee.....ooeere /one sowercaeesaessiiieeeeesestrttrrr
gradientRMS + V mean(newGradient
14
15
gradientMax max(abs(newGradden))))
+ (newGradient new Gradient)/(currenGGeen currentrradiett);
16
17
searchDireetion new Gradient + x searchDirection;
energyDiffference < abs(newFnergy currentFnergy):
(energyDifference < energyTolerance) or (gradientRMS < gradien(Tolerance) or
19
(gradientMax maxGradientTTleranne)) then
29
converged True;
21
else
currentStructure < newSttucctur;
23
currentEnergy + newEnergy;
currentGradient +- new Gradient;
snjjecoryysssaapenncccreeeessss
25
iteration+ 1;
return currentStructure, LrajectoryList;
31

---

Published as a conference paper at ICLR 2024
1. A compound is first dissected into multiple substructures at predefined 16 types of bonds that are selected by organic chemists. In addition, BRICS also takes into account the chemical environment near the bonds, e.g. the types of atoms, to make sure that the size of each fragment is reasonable and the characteristics of the compounds are kept as much as possible.
2. BRICS method then applies substructure filters to remove extremely small fragments (for example single atoms), duplicate fragments, and fragments with overlaps.
3. Finally, BRICS concludes the fragmentation procedure by adding supplementary atoms (mostly hydrogen atoms) to the fragments at the bond-breaking points and makes them chemically stable.
NH2
Quantum Chemikrr
Raw Coorcinatus
Calcultiom
NH2
HO
O-H
OH
H
OH OH
OH OH
Row Molecules
Stabalized Fragments (Not Performed LSR MP)
R
HO
O-R
R
OH
OH
Mabira Frepmants
juddrngf)
NH2
Wotf Matcing
R
NH
NH2
Cheteil Hoods
HO
OH
OH
OH OH
OH OH
Proponal Frogments
Fitered Fregments (Uhad in 150-MP)
Figure 13: Illustration of the BRICS algorithm and its integration with LSR-MP.
Additionally, to enhance the method's accessibility broader audience, we provide pseudo-code for BRICS in Algorithm 2 For deeper understanding of the method, we encourage readers to consult the original BRICS paper for further details.
H LONG-RANGE MODULE VISUALIZATION
On top of the derivations in the main text, we provide visualization of the long-range module which is implemented in ViSNet-LSRM, which is shown in Figure
I PROOF OF EQUIVARIANCE
Definition 1. (Rotation Invariance). f. X y is rotation-invariant if AR SO(3), X
X x(X R.) = =[(X).
Definition 2. (Rotation Equivariance). f X + y is rotation-equivariann Vf E S0(3), X e X, T E T, (Xx (X R) (X)()).
32

---

Published as a conference paper at ICLR 2024
Algorithm 2: BRICS Algorithm for Fragmentation
Input :Molecule
Oufput:Set of final fragments
BRICS_Algox thm (molecule)
honds_to_brean +- F nd __dddds (molecule)
fragments + BeBBeeeeaaa_aBBe Bond (predefined_boneee
Eltered_fragmeed Applllppppsaslllall wtt_SSuusrutttr___ieeee (fragments)
stabilized stabilizee_ffragments- + Stabi Baaeoliiiz_FFrammmni (filtered_fragmente
return caabbllized_frggmeeee
Band_Ronde (molecule)
Data: molecule
Result: List of bonds break
bonds to break + empty list for each bond the molecule do
10
I1
if the bond and its chemical environment match one of the 16 predefined bond types then Add the list of bonds break
12
end
14
end
15
return bonds to break
17 Break_ Break_Bond Bond (bonds to br_ak)
Data: bonds_to. foadd_tt_bbreeee to_brek
Result: fragments
fragments empty list for each bond the bonds to break do break the bond and add resulting fragment to the fragment list
19
21
end
22
return fragments
23
Apply_Sute wtw_Suusstuctuuuu wtw.:usstuucuuue__eeeeeeee (fragments)
Data: list of fragments Result: Filtered list of fragments filtered_fraet fraggmenfssgmntssscom
empty list for each frugment frugments do
26
if fragment size is reasonable and not duplicate overlapping with other fragments
then
Add it the list of filtered fragments end end
29
30
return fraggeenffaggents
32
Stabilize Saabilize_Fragments (fragments) Data: list of fragments Result: List of stabilized fragments
stabilized stabilizeddfrrgmentt + empty list for each fragment in fiitered_framents do
34
36
Add supplementary atoms (e.g.. hydrogen atoms) to make the fragment chemically stable Add the stabilized fragment the list of stabilized fragments
37
end
39
return guubllized_fragmen
33

---

Published as a conference paper at ICLR 2024
Ac-ia3--1
Stacwyose
DHA
ASAT
chigrelts
Figure 14: Visulization of BRICS fragmentation results
RBF
R83
Radial funetion
8 --
AGC
Aggregation
W/o bizd
w/o bias
prodlct
AGG
Add
Rademard product
AN
Figure 15: Structure of long-range module.
Proposition L.. The rotation invariance of the LSR-MP is preserved if the rotation invariance of short-range message-passen passsgg and long-range bipartite peessag-ppassing is preserved.
Proof. The goal is to show that if a rotation transformation R E SO(3) is applied to p. the output how remains unchanged.
First, consider the short-range peessage-passing which is rotationally invariant. This implies (Equa- tion 1) that when a rotation transformation R is applied p. the output h' remains the same as the original output h:
h' = h SHORTR
(27)
The fragmentation learning module combines the resulting embeddings linearly. The output H remains unchanged after applying the rotation transformation, as shown by the equality between H! and Hy:
H3
Σ OO (ESÚ)
Σ a, 2 H,
(28)
iESU)
Next, consider the long-range bipartite message-passing, which is also rotationally invariant. When rotation transformation is applied to ft, V. and P. the output & remains unchanged:
34

---

Published as a conference paper at ICLR 2024
LONGRANGEMODULEG(, H, Ríi RT. RV, RP).
(29)
By combining Eq and Eq 29 the transformed long-range scalar embeddings can be expressed as MC
μ.
(30)
The final output Nom remains unchanged after applying the rotation transformation to the input features. The Dense function combines the outputs from the short-range and long-range message- message-passing modules:
You DENSE (IV' DENSE [h,,)) hout.
(31)
L
Proposition 1.2. The rolation equivariance of the LSR-MP is preserved the rotation equivariance of short-range message saaauggg-pssigggggoe and long-rang, bipartite message-passing is preserved.
Proof. This proposition is equivalent to: If rotation transformation & (0( is applied to p. one could find a predictable transformation T T to the vectorial embeddings Tout
Since short-range paassgg is rotationally equivariant, we can obtain the transformed short- range vectorial embeddings F' by applying the rotation transformation R to the input positions p and then passing them through the ShortRang:Module:
T (T) THTORTRANCEEMODUEE.... Rt).
(32)
In the context of vectorial embeddings, the transformed short-range vectorial embeddings d are equal the rotation transformation applied to the original vectorial embeddings o:
R SIIOERRANGEMODULEE2,
(33)
The fragmentation learning module composed linear combinations of the resulting embeddings. This yields the transformed fragmentation learning module embeddings V and Pj, which can be obtained by applying the rotation transformation R to the original embeddings V, and Pz:
V =
Σ
BOT= =
Σ
G RJ RR Σ
O = RV,
(34)
(ES()
iES())
iES()
P/
Σ (ESU)
XP. Σ
RR = R. Σ 7e RP,,
(35)
(ESU)
(ES0)
As the long-range bipartite message-passing is also equivariant, this indicates that the transformed long-range vectorial embeddings could be obtained via RJT:
R LONGRANGEMODULE(h..P R RD. RV RP).
(36)
Combing Eq. 34 and Eq.
36 have:
n'
Ry LONGRANGEMODULLCC. H1,,, Ru, Rp, RV RPP..
(37)
This equation shows that the transformed long- range vectorial embeddings jl' equal to the rotation transformation R applied to the original vectorial embeddings Finally, the output vectorial embeddings are given by:
35

---

Published as a conference paper at ICLR 2024
E.. U(
= Ri
(38)
This equation demonstrates that the transformed output vectorial embeddings Tou can be obtained by applying the rotation transformation R to the concatenated original vectorial embeddings T and ji, and then passing them through the linear combination function U. The result is equal to the rotation transformation R applied to the original output vectorial embeddings Vout-
In summary. the proof shows that if the rotation cquivariance of short-range message-passing passing and longerange bipartite pesssag-paassing preserved, the rotation equivariance of the LSR-MP is also preserved.
J Loss FUNCTION
In our study, we used a combination of mean squared error (MSE) loss functions for energy and force to train models. Specifically, minimized the weighted sum of MSE between the predicted and true energy values and force values during training. The weight for energy was set to 1 p, while weight for force
p.
L= (1
E2
E2)
F; - F,
(39)
n2
3n
where the number of samples, E; is the true energy value for sample i, É, is the predicted energy value for sample F; is the true force value for sample and F, is the predicted force value for sample i. The trade-off parameter controls the relative importance of minimizing energy versus force errors during training. The ||-I| symbol represents the 7.2 norm, which calculates the Euclidean distance between the true and predicted force vectors, providing measure of the overall force prediction error.
K
HYPERPARAMETER SETTING
K.
HYPERPARAMETER IN TABLE 2
For baseline methods that provided an official implementation (Allegro, PaiNN, ET), we used their code directly, while for methods that not open source (ViSNet), we reimplemented them based on their papers. All models were trained using NVIDIA A6000 GPUs. For training on MD22, we used kcal/mol as the default unit for regression. For training on Chignolin, we used eV.
Table 10: Comparison of Hyperparameters used for MD22 and Chignolin. NA indicates the hyperpa-
rameter is not used.
VENG
Equiforme
Allearo 1:1,
Ecualomerr
128
138
128
FTK
17800:t0 6411:1 3310:2)
leuning use
3,00
3.00 04
LIKE-BG
1,00E.0 .3
S.0IE-04
2000 08
5,00E-04
raagei
NA
NA
Warm
1030
1000
1000
1000
1000
1060
eaty slop parence
500
500
500
500
500
1000
500
schodaler
Coene Auscalin
Annoaing
Redusct B.Daaai
Cleizz Armesinng
Annenling
1/900
1280
199
17750
1/30
1000
1780
(dvin) creff
4
4
17.41
Jose easoff
NA
NA
NA
NA
NA
K.2 HYPERPARAMETER IN TABLE 3
L HIGHER-ORDER LSRM
In this section, we detail the application of LSR-MP to another series of EGNNs that employ the Clebsh-Gorder (CG) tensor product in its core architecture. In particular, chose the state-of- the-art Equiformer as the short-range model. For simplicity, we would ignore parity in the ensuing discussion.
36

---

Published as a conference paper at ICLR 2024
Table 11: Comparison the number of parameters and training speed of various methods when forces MAE is comparable the molecule AT AT-AT-CG-CG.
Methods (MAK)
ViSNet (0.16)
ViSNaI-LSRM (0.13)
PaiNN (0.35)
ET (0,29)
Allegro (0.13)
Equilormer (0.13)
Parameters
2.21M
1.70M
3.20M
3.46M
15.11M
3.02M
Layer Number
8
4+2
3
hidden channels
128
64(1
08
(1
128
128
128
1280/ 0), 64 13/ 0)
Training Time Epoch
155
L.1
PRELIMINARY
group representation characterizes how group elements, such as rotations and translations, act on vector space. In the 3D Euclidean group E(3) context, we consider scalars and Euclidean vectors in R Scalars remain unchanged under rotation, while Euclidean vectors transform accordingly. To address translation symmetry, we work with relative positions.
Irreducible representations (irreps) are the minimal building blocks of group representations. They consist of transformation matrices acting specific vector spaces. For each y E SO(3) group, representing 3D rotations, we have irreps matrices Dr(g) R(2L+1)x(2L+)), known as Wigner-I) matrices. These matrices act on vector spaces of dimension (2L + 1), where L is a non-negative integer. Vectors transformed by (g) are type-L vectors. In EGNNs, each of these vectors could serve as hidden neuron, and neurons with L Z are normally referred to as vector neurons. We concatenate multiple type-L vectors construct (E(3) (e3u-vrrraair irreps features. For example, I xd the main text are composed of type-0 vectors, with I order, and do channels. Vectorial features 7 E 3xdd are composed of type-1 vectors, with 3 distinct orders, and di channels. In general. type-L features fL E (2)++,))dd, have 2L 1 distinct orders and dz channels. The irreducible representation could written as a data structure that aggregates irreps of different types:
scalar features IR
f (Jo Rxxdd .1 3x ,..., Sr. E (2
(40)
Spherical harmonics (SH) are functions capable of projecting Euclidean vectors R& into type-L vectors:
SH() FER St. 22 y
(41)
which exhibit E(3) equiivarrae,,,,,,, preserving the group structure during vector transformations:
Vg E S0(3), L ())
(D= = (2))
(42)
By employing SH of relative positions between node-i and node-j, we could generate the initial set of irreps features:
(fL) = Y
(43)
Equivariant information propagates through irreps features via operations such tensor-product- based directional message passing.
Clebsh-Gorden Tensor Product (CG product) is used form interactions between two different vectors of type-L, and type-L2, and output vector of type- -Lg:
fta = ft. & fta:
(44)
where @ denotes CG product. The tensor product utilizes Clebsch-Gordan coefficients, to combine vectors of different types and different orders:
La
La
f(La.ma) = (fr, 8 fta)my
Σ
Σ
cllasms) (L2. Clemm
(45)
mL La m2=-La L2
where indexes to the order of type-L tensor, and (La.ms)
coefficient, CG coefficient non-zeros only when: (L1mm))( is the Clebsch-Gordan
T1 T22 S La < Th T.2h
(46)
37

---

Published as a conference paper at ICLR 2024
which restricts the output type of the operations. For example, type-1 vectors and type-1 vectors could not produce type-3 vectors, since the output types are restricted in {0, 1.
Parameterized CG product: The tensor product could be formulated as parameterized operation where learnable parameter Wh1,, Ig is assigned to each combination of (L1, L2 Lz). Each type-L
features are consist of d channels of type- vectors, thus the learnable parameters could be generalized to the following forms:
(L,, d1, L2, d2, La,, d3) + W:lag.ds)
(47)
where d1, d2 and d3 corresponds to the channel index of the input fi, f2 and the output fa. This leads to the parameterized CG product which could be written as:
fLa = fL. W fiz.
(48)
L.2 NOTATION
Notations: To distinguish short-range shortragg as h, with h being:
and long-range embeddings, we denote short-range embeddings
= ((h) E 1xxd, (h1) E F. Raxd, ..... (hL) E 1221-1)xdd.
(49)
while long-range embeddings as The fragments embeddings are capitalized, which are denoted as H. The short-range message is denoted as my and the long-range message is denoted as May- H, x. majs My are all irreducible representations, taking the same form as Eq. 49 In (hz)), indexes to the type of the irreps, / indexes to the layer number in multilayer framework, and indexes to different atoms. Meanwhile, we also used a short-hand notation hi, where the index of the type of the irreps is ignored.
L.3 SHORRT-ANGE MODULE
Similar to the short-range module introduced in the main text, the short-range module performs message passing Gsheet by taking the atomic numbers Z N and positions 5 R (n is the number atoms in the system) as input, and model the geometric information on Gabon-
For type-O vectors, they are initialized with atom number embeddings:
(ho) = (2))
(50)
where 2; Z is the atomic number, EMBED(-) N + Rd is learnable embedding map, with do being the number of hidden channels for type-0 vectors.
For Lype-L vectors with L 2 1, they are initialized be zeros:
(
0 R (2L+1)xx, if L 2
(51)
where a type-L L irreps are composed of d, channels of type-L vectors.
In general, the short-range module adopts an iterative short-range message-passing passing scheme:
Oshort
1 1-1
Σ
maj
1-1
(52)
JEN(E)
The @sson() (-) defines message-passieee passsing framework, and N-) denotes the first-order neighbor set of particular node, mij denote the message message between node i and its first order neighbor j. mij is composed of irreps of different types and is computed using the following message functions:
myj = o(h, hy, p.)).
(53)
Commonly, spherical harmonics used turn P.j into irreducible representations to interact with the node irreps:
my ((,, hy, SH(P)))
(54)
where (- .) is the spherical harmonics, and rij is the distance from node-i to node j. This interaction could be performed via parameterized CG product:
mij h4 8 W(re) (H
(55)
38

---

Published as a conference paper at ICLR 2024
L.4 FRAGMENTATION MODULE
To obtain the irress for given fragment, the type-L irreps of the contained node representation summed:
HHL) = Σ
(oL) の (hz)),
(56)
(ES())
in which (HL) denotes the type-L irreps of fragment j is the index for fragments, and S()) is the set induced by the assignments of fragment j. af are weight vectors for each atom within the fragments.
L.5 LONG-RANGE MODULE
The long-range module is targeted to capture possible atom-fragment interactions on Goong Generally, the long-range embeddings are initialized based the short-range embeddings of layer- Shont Type-0 vectors are initialized with DENSE(.) acting on type-0 irreps of the short-range representation. For 1, type-L vectors are initialized with (-) (linear layer without bias) acting on type-L L irreps of the short-range representation:
DENSE ((ho) U (1h2)) Lomen
if 0,
(xr))
(57)
if 2 1,
The geometric cessage-passing is performed to characterize long-range interactions:
t'long
x
Σ M JEN()
(58)
qb())on is the general bipartite message passing framework, x is the long-range irreps, i,j are index for atom, and (-) is the neighborhood of atom on the atom-fragmmet bipartite graph. Mij. denoting the message between atom i and its incident fragment j. is comprised of irreps of different types and are obtained via the message functions (-)):
Mij vb (hg Hy. Hy.
(59)
Similarly, spherical harmonics commonly used to form the interaction between h;, Hy:
Mi = 4b (hi, SH; PJ), l1 P,11) -
(60)
L.6 PROPERTIES PREDICTION
A fusion strategy can be applied to integrate the long-range irreps and short-range irreps:
{ DENSE ([(ho)Lamn Loaan (xo (sne) (U [[(hr,)/oor (xx,) (
if T-= 0, i > 1;
( out
(61)
Lshon and Loong are the number of layers for the short-range module and long-range module respec- tively.
L.7 EQCIFORMER-LSRM.
Based on LSR-MP framework and irreducible representation, we provided another exemplary im- plementation of LSR-MP called Equiformer-LSRM. It uses Equiformer Liao & Smidi (023) as the backbone for the short-range module, thus h' EQUIFORMER Z, p). In the fragment learning module, we choose 7/ to be Er 2, ie. setting the center of any fragment to be the center of
mass of all contained atoms. a; are chosen be 1d For the long-range module, reused the
39

---

Published as a conference paper at ICLR 2024
LAYERNORM and DPTRANSBLOCK implemented in Equiformer Liao & Smidt (2023) pass messages in the atom -ragmeeteent bipartite graph. In particular, we first perform layer normalization on the fragment irreducible representation and concatenated the atom representations and fragment representations to form a new set of node representations:
= [h, Hf).
(62)
The atom positions and fragment positions were also concatenated:
p'
P
(63)
Message passing were performed on the atom-fragment bipartite graph, thus:
-DPTRANSBLOCK( Glong 17) DT)
(64)
40


Reviews:
[{"review_id": "uP7zvCta1w", "reviewer": "Reviewer_cAHU", "summary": "This paper proposes a new framework for machine learning of molecular dynamics, called Long-Short-Range Message-Passing (LSR-MP). LSR-MP combines short-range and long-range message passing on graphs to capture both local and non-local interactions in chemical and biological systems. LSR-MP uses a fragmentation-based method inspired by quantum chemistry to divide large molecules into smaller subsystems and model their long-range interactions efficiently and effectively. LSR-MP is implemented on top of an existing equivariant graph neural network (EGNN) called ViSNet, and achieves state-of-the-art results on large molecular datasets with fewer parameters and faster speed. LSR-MP is also applied to other EGNN models and shows consistent improvements, demonstrating its general applicability and robustness.", "questions": "Q1. How do you justify the choice of the LSR-MP framework as a generalization of the existing EGNNs? What are the advantages and limitations of this framework compared to other possible ways of incorporating long-range interactions, such as attention mechanisms, continuous filters, or Fourier features?\n\nQ2. How do you ensure the stability and accuracy of the BRICS fragmentation method for different types of molecules and systems? How sensitive is the performance of the LSR-MP framework to the choice of the fragmentation method and the number and size of the fragments?\n\nQ3. How do you evaluate the scalability and efficiency of the LSR-MP framework for larger and more complex molecular systems? What are the computational and memory costs of the LSR-MP framework, and how do they compare with the conventional quantum chemical methods and other machine learning methods? How do you handle the trade-off between accuracy and efficiency in the LSR-MP framework?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "The paper presents a novel and elegant framework for long-short-range message passing on graphs, which can capture both local and non-local interactions in chemical and biological systems. \nThe paper draws inspiration from quantum chemistry and adopts a fragmentation-based method to divide large molecules into smaller subsystems and model their long-range interactions efficiently and effectively. This is a clever and creative way to overcome the computational and memory challenges of existing methods. The paper implements the proposed framework on top of an existing equivariant graph neural network (EGNN) called ViSNet, and demonstrates its superior performance on two large molecular datasets, MD22 and Chignolin. The paper shows that the proposed method achieves state-of-the-art results with fewer parameters and faster speed than the baselines, which is impressive and convincing.\nThe paper also applies the proposed framework to other EGNN models, such as PaiNN, ET, and Equiformer, and shows consistent improvements across different architectures and datasets. This demonstrates the general applicability and robustness of the proposed framework, and suggests that it can be easily integrated with other existing methods.", "weaknesses": "The paper does not provide a clear analysis of the stability, and error bounds and how sensitive the performance of the method is to the choice of these modules and parameters.", "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": "h1tPYcwuou", "reviewer": "Reviewer_4VzQ", "summary": "This paper proposes a novel framework for molecular dynamics simulations using machine learning. The framework, called Long-Short-Range Message-Passing (LSR-MP), combines equivariant graph neural networks (EGNNs) with fragmentation-based methods to capture both short-range and long-range interactions among atoms. The authors demonstrate that LSR-MP can achieve state-of-the-art results on large molecular datasets, while being more efficient and effective than existing methods. The authors also conduct ablation studies and analysis to validate the importance of incorporating long-range components and the advantages of using BRICS fragmentation.", "questions": "- As far as I know, there are many long-range message-passing modules designed for graph-structured data. Can you compare the proposed method with other long-range message-passing modules?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- **Problem Definition**: This paper addresses a challenging and important problem of modeling large molecular systems with high accuracy and low computational cost.\n\n- **Methodology**: This paper introduces a novel message-passing framework that leverages domain knowledge from quantum chemistry to incorporate long-range interactions efficiently and effectively.\n\n- **Performance**: This paper shows significant performance improvements over existing methods on various benchmarks, while using fewer parameters and offering faster speed.\n\n- **Generalizability**: This paper illustrates the general applicability and robustness of the LSR-MP framework by applying it to different EGNN backbones and showing consistent improvements.\n\n- **Implementation**: This paper provide sufficient details on experimental setups and how the method is implemented.", "weaknesses": "- **Novelty**: I could not find any distinct weaknesses in this paper, but I might have missed one since I am not an expert in Molecular Modeling. One major concern is regarding the novelty of the proposed long-range message-passing module. As far as I know, long-range message-passing is one of the highlighted research topics in GNN literature. It would be better to discuss this line of work.", "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": "U5E9z1VDan", "reviewer": "Reviewer_GkgW", "summary": "This paper proposes new message-passing neural networks that capture long-range interactions by generalizing equivariant graph neural networks inspired by fragmentation-based approaches. For the implementation, BRICS fragmentation was leveraged. The authors demonstrated the effectiveness of the proposed method with a recently proposed architecture ViSNet and achieved considerable improvement in large molecule benchmarks: MD and Chignolin datasets. To evaluate the proposed method’s applicability, the authors provided results with other EGNNs such as Equiformer, PaiNN, and ET. They show consistent improvement.", "questions": "1. How about inference time? It was not clear how ViSNet has more parameters than ViSNet-LSRM. Also, the training time was reported, but inference time was not available. In real-world applications, inference time is more important for deployment. I believe that shorter training time would imply shorter inference time, but it should be explicitly discussed to be more comprehensive. Fig. 2, (c)(f) partially show the inference time for the subset of baselines\n2. Figure 3 is confusing. The legend should be updated. \n3. Table 4 is not explicitly referred to in the text, although the paragraph of the text of Q3 in Section 5.2 discusses the result. It will be a quick fix. \n4. Typo (?) in Proposition 4.1 Hamdard -> Hadamard product (?)", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. **Strong empirical results** The proposed method achieved competitive performance in MD22 and Chignolin datasets. In Table 2 for MD22, the proposed method with ViSNet shows achieved the best performance in various settings. More importantly, the proposed method, LSRM, shows consistent improvement compared to the vanilla model without LSRM.\n2. **General applicability** The proposed method with various EGNNs base networks shows consistent performance improvement. \n3. **Computational efficiency** The proposed method shows great performance without significant computational overhead. Rather, the proposed method has the smallest model size and the shortest training time.\n4. **Comprehensive experimental results** The paper provides many details and additional experiments.", "weaknesses": "1. **Limited impact.** The technical contributions of this paper has limited impact. Although the method show overall comparable performance, it is a model that is manually-designed by domain knowledge.\n2. **Narrow perspective.** Basically, the proposed method uses two different types of graphs. Then the problem can be viewed as learning on heterogeneous graphs. In recent years, learning graph neural networks on heterogeneous graphs by manually/automatically transforming graphs has been actively studied. The authors may want to include the related work and potentially compare with them. Beyond long-range dependency, non-local/semantic relations also have been utilized. \n3. **Far-fetched claim.** In Figure 2., I do not see anything but overall performance gap. I do not think that the graphs support the claim that LSRM helps models to capture long-range dependency. All three models exhibit similar behaviors.", "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": "PVakXtAHKv", "reviewer": "Reviewer_gmV7", "summary": "This paper proposes a fragment-based approach to propagate long-range information in Graph Neural Networks (GNNs). A set of fragments is constructed using the BRICS fragmentation method, which leverages chemical structures to define well-behaved fragments. The framework operates at two levels: first, it performs a message-passing step on the short-range graph, and then uses these results to define fragment-level features that are also message-passed at the fragment-level graph. The method demonstrates competitive accuracy on the MD22 benchmark, which contains large structures, and shows some improvements over short-range baselines.", "questions": "- How well do you expect your model to transfer to new, unseen systems, particularly those of larger sizes?\n\n- Could you plot the typical decay learned by your interactions, assuming the fragmentation approach allows for it? You could try separating two molecules and plotting the energy as a function of distance. Without sensible decay, the model stands little chance of extrapolating effectively.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper is clearly written and pedagogical.\n- Various ablation studies were conducted on both the long-range modules and the fragmentation methods.\n- Results show improvements over short-range baselines for large molecules.\n- Some limitations of the fragmentation methods are discussed.", "weaknesses": "- The main weakness of the method, as shared in the paper, is the definition of fragments. First, as mentioned, it is not clear how to define these for most systems, including materials. My biggest concern is the issue of smoothness. In molecular dynamics (MD) simulations, it is crucial to ensure that the predictions are smooth. I can envision many MD scenarios where such partitioning might cause problems, and I would be very interested in seeing the behavior of this model over long simulations.\n\n- The Equiformer and VisNet models have 4 layers with a 4Å cutoff, resulting in a receptive field of 32Å in diameter. Most of the MD22 molecules fit well within their receptive fields. While this does not detract from the improvement offered by the method, it should be clearly highlighted.\n\n- The importance of long-range effects beyond a 12Å radius is subtle, as large effects are usually screened in most systems. One would expect to see little difference in errors between a short-range and a long-range model. However, observables computed from MD simulations might vary significantly, as these long-range effects do not average out over long timescales. To capture these observables accurately, the most crucial factor is the decay of interactions, rather than raw accuracy. There is no reason to believe that your approach would correctly capture this decay, enabling accurate observables in these simulations. I want to stress that long-range effects in large biomolecular systems are mostly relevant for observables, and justifying the method solely through raw accuracy has limited scientific relevance.\n\n- One of the main challenges of long-range modeling is transferability, especially for models without typical decay behaviors. I would be very interested in seeing how this model extrapolates to longer, unseen molecules, and whether it performs better than a local model in this context. This is the only relevant setting for practical applications, particularly for modeling systems where ab initio computations are not feasible.", "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)