S PATIAL A TTENTION K INETIC N ETWORK WITH E( N )-E QUIVARIANCE
Anonymous authors
Paper under double-blind review
A BSTRACT
Neural networks that are equivariant to rotations, translations, reflections, and permutations on n -dimensional geometric space have shown promise in physical modeling for tasks such as accurately but inexpensively modeling complex potential energy surfaces to guiding the sampling of complex dynamical systems or forecasting their time evolution.
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 and chemical 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 it comes to 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 displayed utility in predicting quantum chemical potential energies, and considerable effort has been made in the finetuning 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 completely describe node environments, 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 repeatedly and keeps track of a fictitious velocity variable, which we call a spatial attention kinetic network with E(n)-Equivariance , or SAKE (pronounced sah-keh , like the Japanese rice wine).
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 state-of-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 5) while requiring only a fraction of their training and inference time.
Furthermore, we incorporate
SAKE into an equivariant normalizing flow framework and show its utility in exact likelihood sampling tasks.
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]
∀ 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 D 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]
(3) where edges incident to a node v pool their embeddings to form aggregated neighbor embedding a v , and finally a node update ,
[Equation]
[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 D 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; 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.
They nonetheless 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 5) 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 5).
In ablation study (Appendix Section 11), 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.
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 invariant continuous function g ( h e vui , x v , x u i ) on E(n) spatial attention ϕ
SA can approximate g with arbitrarily small error ϵ with some { λ } , f , and µ .
This universal approximative power on E ( n ) -invariant functions can led to parametrization of
E ( n ) -equivariant functions, which we furthermore studied in Remark 2.
The distinctness condition corresponds to the full-rank requirement in GMN (Huang et al., 2022).
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 Appendix
Section 11.
Edge embedding.
To embed pairwise interactions, combine the SchNet (Schütt et al., 2017)-style continuous filter convolution and the simple concatenation of the scalar-valued distance as in 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.
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 ( k ) e uv ∗ α
X×H uv ▷ Euclidean and semantic attention, Sec 5 Eq 8 for u ∈ N ( v ) do h ( k +1) SA v = ϕ
SA ( h ( k ) e uv ,⃗ e uv )
▷ Spatial attention, Sec 5 Eq 6 end for a ( k ) 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 ) λ i ( h ( k ) e uv ) f
( ⃗ e kuv )
▷ Vel.
update, Sec 5 Eqx ( K ) v ← x ( k ) v + v ( K ) v ▷ Position update, Sec 5 Eqh ( k +1) v ← ϕ v ( h ( k ) v , a ( k ) v , h ( k ) SA v )
▷ Node update, Sec 2.2 Eqreturn { h ( k +1) v } , { x ( k +1) v } , { v ( k +1) v } end for end function
Euclidean and semantic attention To promote anisotropy (which Dwivedi et al.
(2020) finds useful) in the pooling operation, apart from spatial attention introduced in Section 4, we used two more types of attention mechanisms, to weigh interactions (edges) based on Euclidean and semantic distances.
For a node v and a set of neighbors { u ∈ N ( v ) } , positioned at x v and x u , connected by directed edges ⃗ e uv respectively, the Euclidean attention weights α X
u
i v are calculated in an additive manner (Bahdanau et al., 2014) with an adaptive scaling factor γ before Softmin .
And semantic attention weights α H u
i v are calculated similar to Graph Attention Networks (Veliˇckovi´c et al., 2018).
To produce models for the purpose of molecular simulation, To ensure 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 via yet another Softmax operation (Matveev et al., 2020):
[Equation]
It is trivial to expand to multi-head attention with k sets of ( a T 1 ,...,k , γ 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 10 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 , 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 10) 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 10) corresponds to 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.
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 generate a (distribution of) coördinates conditioned on initial coördinates and embedding.
See Appendix Section 9 for experimental details and settings.
6. I NVARIANT TASKS : MACHINE LEARNING POTENTIAL CONSTRUCTION MD
(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 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.
In terms of training cost, most state-of-the-art models requires days of training, whereas the MD17 experiment was completed within 6 hours.
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 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 [Table] 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
[Table]
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
[Table] [Table]
[Table]
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.
It 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.
Walking motion capture.
(CMU, 2003) is used by us to demonstrate SAKE’s general capacity to forecast dynamic systems beyond microscopic scale.
In this task, closely following the experiment setting of Huang et al.
; Kipf et al.
, we predict the position of a walking person (subject 35 in CMU motion capture database (CMU, 2003))
based on their initial position.
Again, we observe that SAKE outperform other models by a large margin (Table 5) and is significantly faster.
[Table]
E QUIVARIANT EXACT LIKELIHOOD SAMPLING WITH SAKE
[Table] Current equivariant normalizing flow models (Satorras et al., 2022; Köhler et al., 2020) (brielfy reviewed in Appendix Section 10), 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 23 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 22 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]
which is, up to a constant, an evidence lower bound for x :
[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]
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
[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.
To experimentally benchmark this innovation, we adopted the benchmark systems of DW4 and LJ in Köhler et al.
(2020), with 4 particles on 2-dimensional space and 13 particles on 3-dimensional space, whose energy landscapes are respectively determined by a pairwise double-well potential and a Lennard-Jones potential.
As shown in Table 6, SAKE consistently outperforms other models with the exception of DW4 system with 10 3 training examples.
The slightly worse performance on this particular case might be attributed to the fact that SAKEFlow discretizes (with stochasticity) the integration of Equation 23 and requires one set of parameters for each layer rather than shared parameters as in Köhler et al.
Satorras et al.
The overall superior performance
SAKE demonstrated in this task is indicative that SAKE can be used in the direct sampling of highdimensional many-body conformations, namely Boltzmann type distributions p ( x ) ∝ exp( − u ( x )) as discussed in Noé et al.
(2019)—we leave this for future study.
[Table]
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 more expensive 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 and, when equipped with an augmented invertible flow scheme, sampling.
Limitations In this work, 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—an ablation study might be needed to clearly distinguish which components of the network are contributing to the high performance the most.
For example, thorough experiments are needed to study the exact likelihood sampling performance of
SAKEModel with ODE integration or vise versa —other equivariant model with noise-augmented generative process.
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
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]
20) 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 .
[-As such ϕ DPS ( v ; λ k , λ q ) = 12 ϕ SA ( v ; λ + ) 
-]
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://anonymous.4open.science/
r/sake-EF 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 5) 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 10 (velocity update), and µ in Equation 6 (spatial attention).
SiLU is used everywhere as activation, except in Equation 10 (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 8 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 7.
All random seeds are fixed as 2666, in an homage to Bolaño (2004).
[Table]
(Table 1)
[Table]
(Table 2)
[Table]
Table 3)
[Table]
Table 4)
[Table] [Table]
9. D ATA A VAILABILITY
The source and details of benchmark datasets are summarized in Table
[Table]
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]
and conversely from a sample on the target space to the latent space z
= f xz ( x ; θ ) whose likelihood is given by
[Equation]
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 21, if an unnormalized target distribution is given, or Equation 22, 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] [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.
A BLATION STUDY
To elucidate the contribution of each attention modules on the performance of SAKE, we delete certain parts of it and retrain the MD17 force prediction task.
Interestingly, depriving the model of either semantic or Euclidean attention alone doesn’t drastically change it performance, while deleting both is detrimental, possibly suggesting these two attention mechanisms complement each other.
[Table]