S PATIAL A TTENTION K INETIC N ETWORK WITH E( N )-E QUIVARIANCE
A BSTRACT
Neural networks that are equivariant to rotations, translations, reflections, and permutations on n -dimensional geometric space have shown promise in physical modeling—from modeling potential energy surfaces to forecasting the time evolution of dynamical systems.
Current state-of-the-art methods employ spherical harmonics to encode higher-order interactions among particles, which are computationally expensive.
In this paper, we propose a simple alternative functional form that uses neurally parametrized linear combinations of edge vectors to achieve equivariance while still universally approximating node environments.
Incorporating this insight, we design spatial attention kinetic networks with E(n)-equivariance, or SAKE, which are competitive in many-body system modeling tasks while being significantly faster.
I NTRODUCTION
Encoding the relevant symmetries of systems of interest into the inductive biases of deep learning architectures has been shown to be crucial in physical modeling.
Graph neural networks (GNNs) (Kipf and Welling, 2016;
Xu et al., 2018; Gilmer et al., 2017; Hamilton et al., 2017; Battaglia et al., 2018), for instance, preserve permutation equivariance by applying indexing-invariant pooling functions among nodes (particles) and edges (pair-wise interactions) and have emerged to become a powerful workhorse in a wide range of modeling tasks for many-body system (Satorras et al., 2021).
When describing not only the topology of the system but also the geometry of the state, relevant symmetry groups for three-dimensional systems are SO(3) (rotational equivariance), SE(3) (rotational and translational equivariance), and E(3) (additionally with reflectional equivariance).
A ubiquitous and naturally invariant first attempt to encode the geometry of such systems is to employ only radial information, i.e., interparticle distances.
This alone has empirically shown utility in predicting quantum chemical potential energies, and considerable effort has been made in the fine-tuning of radial filters to achieve quantum chemical accuracy—1 kcal/mol, the empirical threshold to qualitatively reliably predict the behavior of a quantum mechanical system—and beyond (Schütt et al., 2017).
Nonetheless, radial information alone is not sufficient to fully describe node environments—the spatial distribution of neighbors around individual particles.
The relative locations of particles around a central node could drastically change despite maintaining distances to these neighbors unaltered.
To describe node environments with completeness, one needs to address these remaining degrees of freedom.
Current state-of-the-art approaches encode angular distributions by employing a truncated series of spherical harmonics to generate higher-order feature representations; while these models have been shown to be data efficient for learning properties of physical systems, these features are expensive to compute, with the expense growing rapidly with the order of harmonics included (Thomas et al., 2018; Klicpera et al., 2021a;
Fuchs et al., 2020; Batzner et al., 2021;
Anderson et al., 2019).
The prohibitive cost would prevent this otherwise performant class of models from being employed in materials and drug design, where rapid simulations of large systems are crucial to provide quantitative insights.
Here, we design a simple functional form, which we call spatial attention , that uses the norm of a set of neurally parametrized linear combinations of edge vectors to describe the node environment.
Though simple in form, easy to engineer, and ultra-fast to compute, spatial attention is capable of universally approximating any functions defined on local node environment while preserving
E(n)-invariance/equivariance in arbitrary n -dimensional space.
After demonstrating the approximation universality and invariance of spatial attention, we incorporate it into a novel neural network architecture that uses spatial attention to parametrize fictitious velocity and positions equivariantly, which we call a spatial attention kinetic network with E(n)-Equivariance , or SAKE (pronounced saké ( sah-keh ), like the Japanese rice wine) 1 .
To demonstrate the robustness and versatility of SAKE, we benchmark its performance on potential energy approximation and dynamical system forecasting and sampling tasks.
For all popular benchmarks, compared to stateof-the-art models, SAKE achieves competitive performance on a wide range of invariant (MD17:
Table 1, QM9: Table 3, ISO17:
Table 2) and equivariant (N-body charged particle Table 4, walking motion: Table 6) while requiring only a fraction of their training and inference time.
B ACKGROUND
In this section, we provide some theoretical background on physical modeling, equivariance, and graph neural networks to lay the groundwork for the exposition of spatial attention networks.
2. E QUIVARIANCE : PERMUTATIONAL , ROTATIONAL , TRANSLATIONAL , AND REFLECTIONAL
A function f : X → Y is said to be equivariant to a symmetry group G
if
[Equation] [Equation]
) ∀ g ∈ G and some equivalent transformations on the two spaces respectively T g : X → X and
S g : Y → Y .
If on a n -dimensional space X = Y = R n , we have permutation P and T g ( x ) =
P x , S g ( y ) =
P y satisfying Equation 1, we say f is permutationally equivariant; if T g ( x ) =
x R where R ∈ R n × n is a rotation matrix RR T
=
I , and S g ( y )
=
y R we say f is rotationally equivariant; if T g ( x ) = x + ∆ x and S g ( y ) =
y + ∆ x , where x ∈ R n we say f is translationally equivariant; finally, if
T g ( x ) =
Ref θ ( x ) and S g ( y ) =
Ref θ ( y ) , and Ref θ is a reflection on n -dimensional space, we say f is reflectionally equivariant.
2. G RAPH NEURAL NETWORKS
Modern GNNs, which exchanges and summarizes information among nodes and edges, are better analyzed through the spatial rather than spectral lens, according to Wu et al.
(2019)’s classification.
Following the framework from Gilmer et al.
; Xu et al.
; Battaglia et al.
(2018), for a node v with neighbors u ∈ N ( v ) , in a graph G , with h ( k ) v denoting the feature of node v at the k -th layer (or k -th round of message-passing) and h 0 v ∈ R C the initial node feature on the embedding space, the k -th message-passing step of a GNN can be written as three steps:
First, an edge update ,
[Equation]
where the feature embeddings h u of two connected nodes u and v update their edge feature embedding h e uv , followed by neighborhood aggregation ,
[Equation]
where edges incident to a node v pool their embeddings to form aggregated neighbor embedding a v , and finally a node update ,
[Equation]
where N ( · ) denotes the operation to return the multiset of neighbors of a node and ϕ e and ϕ v are implemented as feedforward neural networks.
Since the neighborhood aggregation functions ρ e → v are always chosen to be indexing-invariant functions, namely a SUM or a MEAN operator,
Equation 3, and thereby the entire scheme, is permutationally invariant.
Problem Statement.
We are interested in designing a class of parametrized functions f θ : X ×H →
X × H that map from and to the joint spaces of ( n -dimensional) geometry X ∈ R n and semantic embedding H ∈ R C such that it is permutationally, rotationally, translationally, and reflectionally equivariant on X and invariant on H .
That is, for a given coördinate x ∈ X , embedding h ∈ H and any transformation mentioned in Section 2.1 T : X → X (rotation, translation, and reflection), we have:
[Equation]
R ELATED W ORK : I NVARIANT AND E QUIVARIANT GNN S Invariant GNN.
Although the notion of graph might not appear in their original publications, a plethora of neural architectures could be viewed as de facto graph neural networks with invariance on geometric space (for a survey, see (Han et al., 2022)).
SchNet (Schütt et al., 2017), for example, approximates the potential energy of molecules by applying radially-symmetric convolutions that only operate on the distances between atoms to ensure the model is E(3)-invariant, effectively using a graph neural network architecture where nodes denote atoms and edges denote the distance separation between pairs of atoms within a cutoff radius.
In addition to interatomic distances, invariant models could also incorporate information about the angular distribution of atom neighbors by computing atomic features that incorporate angular information of neighbor-atom-neighbor triplets (Smith et al., 2017; Behler, 2011; Klicpera et al., 2020a; Wang et al., 2022; 2023; Liu et al., 2022).
Compared to SAKE, these models are not capable of equivariant modeling; even in an invariant setting, they are empirically less performant (See invariant tasks performance Section 6.1).
Equivariant GNN with spherical harmonics.
Architectures achieving SE(3)-
equivariance by leveraging Bessel functions and spherical harmonics to encode higher-order interactions (Thomas et al., 2018; Klicpera et al., 2021a; Fuchs et al., 2020; Anderson et al., 2019; Klicpera et al., 2021b;
Brandstetter et al., 2021; Liu et al., 2022) shows outstanding performance and data efficiency, some of which is on par with SAKE (Table 1, 4).
Villar et al.
(2021) discusses that these higher-order representation is not necessary to construct invariantly and equivariantly universal models.
Meanwhile, they tend to be difficult to engineer and expensive to train and run.
(See runtime benchmarks in aforementioned tables.)
Equivariant GNN with dot product scalarization.
Recent efforts (Schütt et al., 2021;
Thölke and
Fabritiis, 2022; Huang et al., 2022) relying on dot product of edge tensors as equivariant nonlinearity have achieved competitive results on machine learning potential construction, among which none has been validated to perform well on both equivariant and invariant tasks .
Experimentally, they are consistently outperformed by SAKE in invariant (Section 6.1: Table 1) and equivariant (Section 6.2:
Table 6) tasks.
Message passing between emebedding and geometric spaces.
Satorras et al.
(2021) has formalized the link between graph neural networks and and equivariance and provided a generalizing framework encompassing iterative geometric-to-embedding and embedding-to-geometric type update.
Like our proposed architecture, E(N) Equivariant Graph Neural Networks (EGNN), proposed in
Satorras et al.
(2021), also uniquely describes the geometry of a n -body system, although their argument was based on a global scale given sufficient steps of message passing, while we can sufficiently describe the local geometric environment.
This advantage is evidenced by extensive experiments (invariant: Table 3; equivariant: Table 4 and 6).
In ablation study (Section 6.3), we show that without spatial attention , EGNN is not competitive in potential energy modeling, even when all other tricks used in this paper were added.
Similar to EGNN (Satorras et al., 2021), our model is equivariant w.r.t.
a general E(n) group and are not restricted to E(3).
Also worth noting is that the Geometric
Vector Perceptrons algorithm proposed in Jing et al.
(2021) could be regarded as a special case of our framework where the attention weights are learned globally, whereas we learn them in an amortized manner and thus can be transductively generalized across systems.
T HEORY : S PATIAL A TTENTION
Given a node v with embedding h v ∈ H = R C (where C denotes the embedding dimension) and position x v ∈ X = R n (where n denotes the geometry dimension) in a graph G , its neighbors u ∈ N ( v ) with connecting edges { e uv } , with displacement vector ⃗ e uv = x v − x u and embedding h e uv = ρ v → e ( h v , h u ) with some aggregation function ρ v → e , we define spatial attention ϕ : X ×H →
H as
[Equation]
where λ i : H → R 1 , i = 1 , ..., N λ is a set of arbitrary attention weights-generating function that operates on the edge embedding, f : X → X is an equivariant function that operate on the edge vector, µ : N λ
→ H is an arbitrary function that takes the norms of N λ linear combinations, and ⊕ denotes concatenation.
We drop the explicit dependence of ϕ SA ( v ) on both geometric and embedding properties of v and u for simplicity.
We name this functional form attention because it characterizes the alignment between edge vectors in various projections.
When implemented, λ and µ take the form of feed-forward neural networks.
In other words, for each node v , we first parametrize a N λ = |{ λ i }| (is analogous to the number heads in multi-head attention) sets of attention weights based on the edge embeddings
λ i ( h e uv ) .
At the same time, we also equivariantly (on E(n) group) transform the edge vector into f ( ⃗ e uv ) .
Next, we use those N λ sets of attention weights to linearly combine these vectors among edges and take the Euclidean norm of each of these linear combinations, resulting in N λ norms.
Lastly, we concatenate these norms and put them into a feed-forward neural network µ to compute the node embedding.
It naturally follows that:
Remark 1.
Spatial attention is permutationally, rotationally, transalationally, and reflectionally invariant on E(n), since h e uv is invariant w.r.t.
indexing and the Euclidean norm function is invariant to rotation, translation, and reflection.
With all local degrees of freedom incorporated in spatial attention, it is perhaps intuitive to see that ϕ SA ( v ) can uniquely define the local environment of nodes up to E(n) symmetry.
We formalize this finding (Proof in Appendix Section 8.1):
Theorem 1.
For a node v in a graph G with neighbors u ∈ N ( v ) connected to u by edges with positions x v and x u distinct embeddings
h e vui ̸ = h e vuj , ∀ 1 ≤ i, j ≤ |N ( v ) | , and any E(n)-invariant continuous function g ( x v , x u i |
h e vui ) spatial attention ϕ SA can approximate g with arbitrarily small error ϵ with some { λ } , f , and µ .
It is worth noting here that we only consider the universal approximation power of a mapping from the geometry space X to a scalar space, without considering that on the space of node (and edge) embedding; regardless of geometry, GNNs that operate on neighborhood node embeddings are generally not universal (Xu et al., 2018;
Corso et al., 2020).
This universal approximative power on
E ( n ) -invariant functions can lead to universally approximative parametrization of E ( n ) -equivariant functions, which we show in Remark 2.
Practically, the inequality condition is not difficult to satisfy: we implement h e as dependent upon edge length so even if the semantic embedding of edges are identical, the inequality h e vui ̸ = h e vuj holds as long as the system is not strictly symemtrical at all times (namely the mirror symmetry presented in hydrogen molecule) despite distortions resulted from vibrations.
A RCHITECTURE : S PATIAL A TTENTION K INETIC N ETWORK (SAKE)
Leveraging the simple functional form introduced in Section 4, we design a fast, efficient, and easy-to-implement architecture termed a spatial attention kinetic network with E(n)-equivariance , or
SAKE.
Here, we introduce the remaining components of SAKE before assembling them in a modular way.
The necessity of these designs are experimentally justified with ablation study in Section 6.3.
Edge embedding.
To embed pairwise interactions, we combine the SchNet (Schütt et al., 2017)- style continuous filter convolution and the simple concatenation of the scalar-valued distance as in function SAKEL AYER ( { h ( k ) v } , { x ( k ) v } , { v ( k ) v } , G )
▷ Input embedding, position, and velocity for v ∈ V do for u ∈ N ( v ) do h ( k ) e uv ← ϕ e
( h kv , h ku , || ⃗ e uv || )
▷ Edge update, Sec 5 Eq 7 end for h ( k +1)
e uv ← h ( ) e uv ∗ α X×H uv
Semantic and cutoff, Sec 10 for u
N ( v ) do h (
1) SA v =
SA uv , e uv ) ▷ Spatial
, Sec
6 end for a (
) v ← (cid:80) u ∈N ( v ) h
( k ) e uv ▷ Neighborhood aggregation, Sec 2.2 Eqv ( k +1) v ← ϕ v →V ( h ( k ) v ) v ( k ) + W v (cid:80) i (cid:80) u ∈N ( v )
[Table]
i ( h (
[Table]
) e uv ) f
[Table]
e kuv )
▷ Vel.
update, Sec 5 Eqx
←
▷ Position update, Sec 5 Eqh
) v ← ϕ v
▷ Node update, Sec 2.2 Eqreturn {
) v } ,
} end for end function Satorras et al.
(2021) to achieve the balance between high radial resolution and large receptive field.
The resulting edge embedding is thus
[Equation]
where ⊙ denotes Hadamard product and f r is a (filter-generating) feed-forward network as in Schütt et al.
Semantic attention and distance cutoff.
To promote anisotropy (which Dwivedi et al.
(2020) finds useful in GNNs) in the pooling operation, apart from spatial attention introduced in Section 4, we compute the attention score on semantic and geometry space to weight interactions among particles based on embedding similarity and distances on n -dimensional space.
The distance weights are calculated using the cutoff function proposed in Unke and Meuwly (2019):
[Equation]
to filter out interactions outside a certain range ( d 0 ).
And semantic attention weights α H u i v are calculated similar to Graph Attention Networks (Veliˇckovi´c et al., 2018),
[Equation] [Equation]
To produce models for the purpose of molecular simulation by ensuring continuous forces and gradients, one would need to choose as sigma σ as at least C2-continuous activation functions, namely CeLU (Barron, 2017).
These weights are combined and normalzied:
[Equation]
[Equation]
It is trivial to expand to multi-head attention with k sets of ( a T 1 ,...,k ) and the resulting combined attention weights concatenated.
Since the edge embedding after mixed Euclidean and semantic attention h e uv ∗ α
X×H uv already encodes the desired inductive bias that nodes farther away from each other would have less impact on each other’s embedding, we directly use this representation and simply set λ i as a linear projection (with weights W λ ) and f as identity in Equation 6:
[Equation]
We leave more elaborate choices of λ and f functions for future study.
Fictitious velocity integration.
Similar to Satorras et al.
(2021), we keep track of a fictitious velocity variable v v for each node and linearly combine it (with weights W λ ) update it and positions x v in turns like a Euler-discretized Hamiltonian integration.
[Equation]
During velocity update, we scale the current velocity and, for the sake of parameter saving, reuse the same set of linear combinations of edge vectors used in Equation 6 as the additive term.
Remark 2.
If h e uv is distinct among edges, there exist sets of λ
i (even when f = I is the identity) such that Equation 12 can approximate any vector on the subspace spanned by edge vectors and is thus universal up to equivariance thereon.
Equivariance analysis.
We inlay these modular components into the framework of graph neural networks (described in Section 2.2) to produce the SAKE architecture, which we show in Algorithm 1.
(A SAKEModel is defined by applying SAKELayer iteratively from k = 0 to k = K − 1 , the depth of the network.)
We have previously discussed that spatial attention is E(n)-invariant in Remark 4.
With the E(n) equivariance of velocity update and position update (Equation 12) being proved in
Satorras et al.
(2021) and the rest of the model only takes the norm of edges and is E(n) invariant,
SAKE is E(n)-equivariant.
Runtime analysis.
The runtime complexity of spatial attention (Equation 6), which is the bottleneck of this algorithm, is O
( |E| N λ CD ) , i.e., linear in the number of graph edges |E| , number of attention weights N λ , embedding dimension C , and geometric dimension D .
When implemented, the FOR loops in Algorithm 1 can be packed into multi-dimensional tensors that could benefit from GPU acceleration.
Relation to spherical harmonics-based models.
Under the framework of TFN (Thomas et al., 2018), the embedding and position input in Algorithm 1 { h ( k ) v } , { x ( k ) v } corresponds to the l = 0 and l = 1 type tensors.
The concatenation followed by neural network in Equation 6 is loosely analogous to the direct sum operation in Clebsch-Gordon decomposition.
While Smith et al.
Schütt et al.
Satorras et al.
(2021) operates on l = 0 tensors only, the spatial attention mechanism we propose (Section 4 Equation 6) and velocity/position update (Section 5 Equation 12) corresponds to 1 ⊕ 1 → 0 and 1 ⊕
→ 1 type networks, respectively.
Klicpera et al.
(2021a);
Villar et al.
(2021) have discovered that l = 0 , 1 are the complete levels of tensor to universally describe geometry on
E (3) , while higher-order tensors are not necessary to completely describe the node environment.
Kovács et al.
(2021) also discusses the concept of density projection where body-order functional forms can be recovered by lower-order terms.
E XPERIMENTS
As discussed in Section 2, SAKE provides a mapping from and to the joint space of geometry and embedding X × H , while being equivariant on geometric space and invariant on embedding space.
We are therefore interested to characterize the performance of SAKE on two types of tasks: invariant modeling (Section 6.1), where we model some scalar property of a physical system, namely potential energy; and equivariant modeling (Section 6.2), where we predict coördinates conditioned on initial position, velocity, and embedding.
On both classes of tasks, SAKE displays competitive performance while requiring significantly less inference time compared to current state-of-the-art models.
See Appendix Section 9 for experimental details and settings.
6. I NVARIANT TASKS : MACHINE LEARNING POTENTIAL CONSTRUCTION MD17 potential energy
(Chmiela et al., 2017) tests the capacity of the model in the extreme small data regime.
Its training set contains merely 1000 configurations and quantum chemical energies and forces of single small molecules in vacuum computed using density functional theory (DFT).
As
[Table]
summarized in Table 1, SAKE consistently outperforms all benchmarked models with the exception of NequIP (Batzner et al., 2021), which has comparable performance but is noticeably slower.
(Note that, to avoid inaccurate report on the suboptimal software configuration, we only report runtime data directly quoted from their original publications and replicate the hardware environment by ourselves.
The same applies hereafter.)
In terms of training cost, most state-of-the-art models requires days of training, whereas the MD17 experiment was completed within 6 hours.
This significant advantage in speed would allow SAKE to be more rapidly trained and deployed on realistic applications involving molecular dynamics simulations.
ISO (Schütt et al., 2017) goes beyond the single-molecule regime.
It involves a slightly more diverse chemical space containing 5000-step ab initio molecular dynamics simulation trajectories (with energies and forces) of 129 molecules with the same formula C 7 H 10 O 2 .
The test set of ISO is split into known (which Kovács et al.
(2021) argues to be very close to training set) and unknown molecules, based on the chemical identity (topology) at the beginning of the simulation, which could be regarded as interpolative and extrapolative tasks, respectively.
As shown in Table 2, SAKE significantly outperforms other models on the unknown molecules in the test set, indicating that
SAKE is capable of extrapolating and generalizing onto unseen chemical spaces when trained on limited data.
QM (Ramakrishnan et al., 2014) tests the transductive generalizability across distinct small chemical graphs.
It entails a very diverse chemical space of 134k molecules with annotated phsyical properties calculated with B3LYP/6-31G(2df,p) level of quantum chemistry, albeit all with atequilibrium (low energy) conformations.
In Table 3, SAKE achieves state-of-the-art performance at predicting HOMO, LUMO, and ∆ ϵ properties—a class of most crucial molecular properties closely related to reactivity.
Interestingly, SAKE performs competitively on predicting extensive physical properties but not intensive ones (also see discussions in Pronobis et al.
We hypothesize that this will be mitigated by choosing size-invariant pooling functions—we leave this for future study.
6. E QUIVARIANT TASKS Charged N-body dynamical system forecasting
(Kipf et al., 2018; Fuchs et al., 2020) tests if a model can predict the evolution of a physical system sufficiently long after initial conditions.
This simple system consists of 5 charge-carrying particles with initial positions and velocities, and the position at a given moment is predicted.
As shown in Table 4, although the interactions (Coulomb forces) are entirely pairwise, we see here that the additional expressiveness SAKE affords lead to competitive performance on this demonstrative task.
[Table]
[Table] [Table]
MD17 forecast
(Chmiela et al., 2017; Huang et al., 2022) involves a simulation forecast task with slightly more complicated systems compared to Table 4.
It uses the same dataset in
Table
[Table]
but predicts the time evolution of the molecular dynamics simulation directly, rather than predicting the mapping from the geometry to potential energy.
Following the protocol in Huang et al.
, we predict the atom position based on the velocity and coördinate 3000 steps prior.
As shown in Table 5, SAKE achieves superior performance on 6 out of 8 systems, without leveraging spherical harmonics (as in TFN (Thomas et al., 2018) or SE(3)-TF (Fuchs et al., 2020)) or hand-coded edges (as in
GMN (Huang et al., 2022)).
Aspirin Benzene Ethanol Malonaldehyde Naphthalene Salicylic Toluene Uracil
[Table]
Walking motion capture
(CMU, 2003) has a higher system complexity and noise and is adopted to demonstrate SAKE’s general capacity to forecast dynamic systems beyond microscopic scale.
In this task, again closely following the experiment setting of Huang et al.
; Kipf et al.
[Table]
we predict the position of a walking person (subject 35 in CMU motion capture database (CMU, 2003))
based on their initial position.
Again
[Table]
we observe that SAKE outperform other models by a large margin (Table 6) and is significantly faster.
6. A BALATION STUDY
[Table] [Table]
To elucidate each component’s contribution towards the final performance, we perform an ablation study on one of the most popular tasks studied so far—MD17 potential energy modeling (Section 6.1, Table 1).
More specifically, we focus on the most complicated molecular system in the dataset, aspirin.
We inspect three components proposed in the paper— spatial attention (Equation 6), which we argue is the main novelty of the paper, as well as semantic attention (Equation 10) and velocity and position update (Equation 12).
As is shown in Table 7, spatial attention improves the performance regardless of the rest of the configuration.
To clearly demonstrate the utility of these auxiliary modules, we also add these components (except spatial attention) to the backbone of EGNN (Satorras et al., 2022), and summarize the results in Appendix Table 10—qualitatively, with these modules, EGNN achieves decent performance on par with SchNet (Schütt et al., 2017), with the performance gap comes primarily from its inability to conceive the node angular environments.
D ISCUSSION
Conclusions In this paper we have introduced an invariant/equivariant functional form termed spatial attention that uses neurally parametrized linear combinations of edge vectors to equivariantly yet universally characterize node environments at a fraction of the cost of state-of-the-art equivariant approaches that makes use of spherical harmonics.
Equipped with this spatial attention module, we use it to build a spatial attention kinetic network (SAKE) architecture which is permutationally, translationally, rotationally, and reflectionally equivariant.
We have demonstrated the utility of this model in n -body physical modeling tasks ranging from potential energy prediction to dynamical system forecasting.
Limitations Theoretical: The universality condition is only discussed w.r.t.
node geometry, without considering node embeddings; moreoever, the inequality condition in Theorem 1, although rarely violated in physical modeling (even when system is highly symmetrical), can be potentially overrestricting.
We plan to generalize this framework to consider the expressive power of functions on the joint space of node embedding and geometry in future works.
Experimental:
Herein, apart from the novel functional form spatial attention (Equation 6), the rest of the architecture has not been thoroughly optimized and analyzed in the context of the growing design space of equivariant neural networks.
Future directions We plan to conduct more thorough experiments on (bio)molecular systems to explore the potential of SAKE in building general protein/small molecule force fields and enhanced sampling methods that could facilitate large-scale simulations useful in therapeutics and material discovery.
Social impact and ethics statement
This work provide an accurate and extremely efficient way to approximate properties and dynamics of molecular systems and physical states.
It may advances research in a wide range of disciplines, including physics, chemistry, biochemistry, biophysics, and drug and material discovery.
As with all molecular machine learning methods, negative implications may be possible if used in the design of explosives, toxins, chemical weapons, and overly addictive recreational narcotics.
Reproducibility statement The software package containing the algorithm proposed here is distributed open source under MIT license.
All necessary code, data, and details to reproduce the experiments can be found in Appendix Section 9.
Funding Research reported in this publication was supported by the National Institute for General
Medical Sciences of the National Institutes of Health under award numbers R01GM132386 and
R01GM140090.
YW acknowledges funding from NIH grant R01GM132386 and the Sloan Kettering
Institute.
JDC acknowledges funding from NIH grants R01GM132386 and R01GM140090.
Disclaimer The content is solely the responsibility of the authors and does not necessarily represent the official views of the National Institutes of Health.
Disclosures JDC is a current member of the Scientific Advisory Board of OpenEye Scientific
Software, Redesign Science, Ventus Therapeutics, and Interline Therapeutics, and has equity interests in Redesign Science and Interline Therapeutics.
The Chodera laboratory receives or has received funding from multiple sources, including the National Institutes of Health, the National Science
Foundation, the Parker Institute for Cancer Immunotherapy, Relay Therapeutics, Entasis Therapeutics,
Silicon Therapeutics, EMD Serono (Merck KGaA), AstraZeneca, Vir Biotechnology, Bayer, XtalPi,
Interline Therapeutics, the Molecular Sciences Software Institute, the Starr Cancer Consortium, the
Open Force Field Consortium, Cycle for Survival, a Louis V. Gerstner Young Investigator Award, and the Sloan Kettering Institute.
A complete funding history for the Chodera lab can be found at http://choderalab.org/funding .
Acknowledgments The authors would like to thank the ICLR reviewers for providing constructive feedback that substantially improved the quality and clarity of this manuscript.
YW thanks Theofanis
Karaletsos, Insitro and Andrew D. White, University of Rochester, for useful discussions and Leo
Klein, Freie Universität Berlin, for catching an embarrassing bug.
P ROOFS 8. P ROOF FOR T HEOREM
Proof.
(Informal)
An appropriate set of choices to satisfy the requirement is: f
= I , and λ takes the set of |N ( v ) | + |N ( v ) | 2 elements:
[Equation]
14) where [ · ] is the Iverson bracket.
The first |N ( v ) | -set of one-indexed functions give the node-toneighbor distances whereas the second |N ( v )
| 2 -set of two-indexed functions give the distances among neighbors.
This recovers the distance matrix A X ( v ) among the node and its neighbors
A X ( v ) ii =
|| x
v − x u i || and A X ( v ) ij =
||
x
u j − x u
i || , i ̸
= j, 1 ≤ i, j ≤ |N ( v ) | .
As such, the relative positions of the node and its neighbors are uniquely defined up to E(n) symmetry (Havel, 1998).
In other words, x v and x u i can be uniquely embedded on an arbitrary coördinate system on which g is evaluated.
Since { λ } is neural, by the universal approximation theorem (Hornik et al., 1989), A X ( v ) can still be recovered, at least through recovering the set of { λ } discussed above.
Considering again the universal approximation of µ , any function of the geometric node environment can be approximated by ϕ SA .
D ETAILED M ETHODS 9. C ODE A VAILABILITY
The corresponding software package and all scripts used to conduct the experiments in this paper are distributed open source under MIT license at: https://github.com/choderalab/sake
This package can be installed via: pip install sake-gnn .
9. H ARDWARE C ONFIGURATION
All models are trained on NVIDIA Tesla V100 GPUs.
Following the settings reported in the publications of baseline models, the inference time benchmark experiments (Section 6.2, 6.1) are done on NVIDIA GeForce GTX 1080 Ti GPU (For Table 1 and Table 6) and NVIDIA GeForce GTX
Ti GPU (For Table 4).
9. A RCHITECTURE AND O PTIMIZATION D ETAILS
One-layer feed-forward neural networks are used as f r in Equation 2 (edge update); two-layer feed-forward neural networks are used as ϕ e in Equation 2 (edge update), ϕ v in Equation 4 (node update), ϕ v →V in Equation 12 (velocity update), and µ in Equation 6 (spatial attention).
SiLU is used everywhere as activation, except in Equation 12 (velocity update) where the last activation function is chosen as y = 2 ∗ Sigmoid( x ) to constraint the velocity scaling to between 0 and 2 and in Equationwhere CeLU is used before attention; additionally, tanh is applied on the additive part of Equationto constraint it to between -1 and 1.
4 attention heads are used with γ in Equation 10 spaced evenly between 0 and 5 Å. 50 RBF basis are used, spacedly evenly between 0 and 5 Å. All models are optimized with Adam optimizer.
We summarize the hyperparameters used in these experiments in
Table 8.
All random seeds are fixed as 2666, as an homage to Bolaño (2004).
[Table]
Table 1)
[Table]
Table 1)
[Table]
(Table 1)
[Table]
(Table 1) [Table]
[Table]
(Table 1)
[Table]
MD17 (Table 1)
[Table]
Trajectory Forcast (Table 5)
[Table]
[Table]
9. D ATA A VAILABILITY
The source and details of benchmark datasets are summarized in Table
[Table]
A BALTION STUDY ON EGNN B RIEF INTRODUCTION OF EQUIVARIANT NORMALIZING FLOWS
Normalizing flows (Rezende and Mohamed, 2016; Papamakarios et al., 2021) are a family of learnable bijections
f zx : Z → X that transform a tractable distribution on latent space
q z
( z ) to another on the target space x = f zx ( z ; θ ) (dropping dependency on parameter thereafter)
whose density can be analytically written as
[Equation]
[Table] and conversely from a sample on the target space to the latent space z = f xz ( x ; θ ) whose likelihood is given by
[Equation]
(16) where f zx = f − 1 xz .
To close the gap between the intractable p x and the tractable q x , so that one can sample on the target space efficiently, the flow is then trained by maximizing either Equation 15, if an unnormalized target distribution is given, or Equation 16, if samples are given.
The concept of equivariant normalizing flow was first introduced in Köhler et al.
Rezende et al.
; Satorras et al.
(2022), where they adopted the framework of continuous normalizing flow (Chen et al., 2018) to define the bijection f zx as
[Equation]
and restrict f ′ zx as E(n)-equivariant w.r.t.
z .
Integration, or the sum over infinitely many equivariant functions, does not alter the equivariance.
To numerically approximate this integration Chen et al.
(2019) involves evaluating f ′ zx multiple times and is therefore expensive.
E QUIVARIANT EXACT LIKELIHOOD SAMPLING WITH SAKE F LOW .
Current equivariant normalizing flow models (Satorras et al., 2022; Köhler et al., 2020) (brielfy reviewed in Appendix Section 11), relies on ODE-based numerical integration(Chen
et al., 2019), are computationally expensive.
We propose a much simpler invertible flow model that uses our SAKE model, termed SAKE Flow.
First, following a scheme introduced in Huang et al.
(2020) (in which it is argued that with flexible enough kernels Equation 17 could be approximated arbitrarily well), we extend the space X with an auxiliary space A .
Correspondingly, we extend the tractable distribution to be q ( z , a ) =
q z ( z ) q a ( a ) and the target distribution to be p ( x , a ) =
p x ( z ) q a ( a ) .
We then change the problem statement of Equation 16 to: find a parametrized function f zx ( z , a ; θ ) =
f − 1 xz ( x , a ; θ ) that is a bijection on the space
A × X to maximize the joint likelihood:
[Equation]
18) which is, up to a constant, an evidence lower bound for x :
[Equation]
[Equation]
[Equation]
The equality holds when the (non-negative)
Kullback–Leibler divergence between the tractable distribution q a
and the conditional distribution given samples p ( a | x ) is zero.
As such, an unbiased estimate of the marginal likelihood can be given by
[Equation] [Equation]
Similar to Huang et al.
Dinh et al.
(2017), we define f zx ( z , a , θ ) as a series of alternating affine coupling:
[Equation]
with analytic inverse.
Composing these transformations g z
→ a ◦ g a → z ◦ ...g z
→ a ◦ g a → z , we get our bijection f zx
( z , a ; θ ) .
Now, it only remains to define the structure of the translation functions { T } and scaling functions { S } .
As is outlined in Satorras et al.
(2022), if f zx is translation-invariant , it is impossible to have (cid:82) p x ( x ) d x = 1 , so we drop the translation invariance requirement and design f zx , as well as the composing { T, S } , to be only rotation and reflection equivariant.
Correspondingly, we require all tractable distributions to be confined on a ( |V| − 1) n -dimensional subspace with 0 gravity center.
(Note that this reduces the symmetry group we work on from E(n) to SO(n).)
A valid choice both for q z ( z ) and q a ( a ) with is a centered Gaussian distribution (Satorras et al., 2022):
[Equation]
with (cid:80) v ∈V x v = 0 .
Moreover, to keep the gravity center at 0 , we require that { T, S } shall not change the gravity center.
A general recipe to construct such { T, S } is to use the equivariant and invariant outputs of SAKE to parametrize T and S respectively.
[Equation]
To preserve the center of gravity, we center the translation to have zero center of gravity and enforce the same scalar to be used across particles as the scaling factor.