Spatial Attention Kinetic Network
with E(n)-Equivariance

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 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 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 sah-keh, like the Japanese rice wine)¹¹1Implementation: https://github.com/choderalab/sake. 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 6) while requiring only a fraction of their training and inference time.

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.

Given a node $v$ with embedding $h_{v}\in\mathcal{H}=\mathbb{R}^{C}$ (where $C$ denotes the embedding dimension) and position $\textbf{x}_{v}\in\mathcal{X}=\mathbb{R}^{n}$ (where $n$ denotes the geometry dimension) in a graph $\mathcal{G}$, its neighbors $u\in\mathcal{N}(v)$ with connecting edges $\{e_{uv}\}$ , with displacement vector $\vec{\mathbf{e}}_{uv}=\mathbf{x}_{v}-\mathbf{x}_{u}$ and embedding $h_{e_{uv}}=\rho^{v\rightarrow e}(h_{v},h_{u})$ with some aggregation function $\rho^{v\rightarrow e}$, we define spatial attention $\phi:\mathcal{X}\times\mathcal{H}\rightarrow\mathcal{H}$ as

where $\lambda_{i}:\mathcal{H}\rightarrow\mathbb{R}^{1},i=1,...,N_{\lambda}$ is a set of arbitrary attention weights-generating function that operates on the edge embedding, $f:\mathcal{X}\rightarrow\mathcal{X}$ is an equivariant function that operate on the edge vector, $\mu:N_{\lambda}\rightarrow\mathcal{H}$ is an arbitrary function that takes the norms of $N_{\lambda}$ linear combinations, and $\oplus$ denotes concatenation. We drop the explicit dependence of $\phi^{\operatorname{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.

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 $\phi^{\operatorname{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):

It is worth noting here that we only consider the universal approximation power of a mapping from the geometry space $\mathcal{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_{vu_{i}}}\neq h_{e_{vu_{j}}}$ 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.

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.

As discussed in Section 2, SAKE provides a mapping from and to the joint space of geometry and embedding $\mathcal{X}\times\mathcal{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.

Normalizing flows [Rezende and Mohamed, 2016, Papamakarios et al., 2021] are a family of learnable bijections $f_{zx}:\mathcal{Z}\rightarrow\mathcal{X}$ that transform a tractable distribution on latent space $q_{z}(\mathbf{z})$ to another on the target space $\mathbf{x}=f_{zx}(\mathbf{z};\theta)$ (dropping dependency on parameter thereafter) whose density can be analytically written as

and conversely from a sample on the target space to the latent space $\mathbf{z}=f_{xz}(\mathbf{x};\theta)$ whose likelihood is given by

where $f_{zx}=f_{xz}^{-1}$. 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. [2020], Rezende et al. [2019], Satorras et al. [2022], where they adopted the framework of continuous normalizing flow [Chen et al., 2018] to define the bijection $f_{zx}$ as

and restrict $f_{zx}^{\prime}$ as E(n)-equivariant w.r.t. $\mathbf{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}^{\prime}$ multiple times and is therefore expensive.

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 $\mathcal{X}$ with an auxiliary space $\mathcal{A}$. Correspondingly, we extend the tractable distribution to be $q(\mathbf{z},\mathbf{a})=q_{z}(\mathbf{z})q_{a}(\mathbf{a})$ and the target distribution to be $p(\mathbf{x},\mathbf{a})=p_{x}(\mathbf{z})q_{a}(\mathbf{a})$. We then change the problem statement of Equation 16 to: find a parametrized function $f_{zx}(\mathbf{z},\mathbf{a};\theta)=f_{xz}^{-1}(\mathbf{x},\mathbf{a};\theta)$ that is a bijection on the space $\mathcal{A}\times\mathcal{X}$ to maximize the joint likelihood:

which is, up to a constant, an evidence lower bound for $\mathbf{x}$:

The equality holds when the (non-negative) Kullback–Leibler divergence between the tractable distribution $q_{a}$ and the conditional distribution given samples $p(\mathbf{a}|\mathbf{x})$ is zero. As such, an unbiased estimate of the marginal likelihood can be given by

Similar to Huang et al. [2020], Dinh et al. [2017], we define $f_{zx}(\mathbf{z},\mathbf{a},\theta)$ as a series of alternating affine coupling:

with analytic inverse. Composing these transformations $g^{z\rightarrow a}\circ g^{a\rightarrow z}\circ...g^{z\rightarrow a}\circ g^{a\rightarrow z}$, we get our bijection $f_{zx}(\mathbf{z},\mathbf{a};\theta)$.

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 $\int p_{x}(\mathbf{x})\operatorname{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 $(|\mathcal{V}|-1)n$-dimensional subspace with $\mathbf{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}(\mathbf{z})$ and $q_{a}(\mathbf{a})$ with is a centered Gaussian distribution [Satorras et al., 2022]:

with $\sum\limits_{v\in\mathcal{V}}\mathbf{x}_{v}=\mathbf{0}$. Moreover, to keep the gravity center at $\mathbf{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.

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.