Geometric Algebra Attention Networks for Small Point Clouds

Much of the success of deep learning is drawn from building architectures that properly respect underlying symmetry and structure in the data on which they operate—a set of considerations that have been united under the banner of geometric deep learning . Often problems in the physical sciences deal with relatively small sets of points in two- or three-dimensional space wherein translation , rotation , and permutation equivariance are important or even vital for models to be useful in practice . In this work , we present rotation- and permutation-equivariant architectures for deep learning on these small point clouds , composed of a set of products of terms from the geometric algebra and reductions over those products using an attention mechanism . The geometric algebra provides valuable mathematical structure by which to combine vector , scalar , and other types of geometric inputs in a systematic way to account for rotation invariance or covariance , while attention yields a powerful way to impose permutation equivariance . We demonstrate the usefulness of these architectures by training models to solve sample problems relevant to physics , chemistry , and biology . BACKGROUND . Deep learning has been immensely successful in solving a wide range of problems over the last several years , driven in large part by identifying appropriate ways to embed structure of data and symmetry of problems directly into the architecture of the network—an idea at the core of geometric deep learning ( Bronstein et al. , 2021 ) . Some applications of geometric deep learning include the use of convolutional filters in CNNs to attain translational equivariance , or graph convolutions in graph neural networks for permutation equivariance.1 Building symmetry into the architecture of a deep neural network can improve the network ’ s data efficiency and guarantee important analytical properties without having to rely on the model to learn to approximate them from training data . In this work , we derive a family of architectures that is useful in applications from physics to biology , where problems often deal with relatively small point clouds of labeled coordinates . These could be local environments of particles assembling into a crystal ( Dshemuchadse et al. , 2021 ) , atoms in a molecule interacting with other atoms ( Chmiela et al. , 2017 ) , or coarse-grained beads representing parts of a protein ( Marrink & Tieleman , 2013 ) . In many of these applications without the influence of an external field , we are interested in modeling attributes of the system—such as the identity of a particle ’ s local self-assembly environment , or the potential energy of a group of atoms—which are invariant with respect to rotation of the input coordinates , as well as permutation in the ordering of points . As shown in Figure 1 , here we attain rotation invariance by constructing functions from rotation-invariant attributes of geometric products of input vectors from geometric algebra , and permutation invariance by using an attention mechanism to reduce the set of vector products . 1In this work , we use the following terms to discuss symmetry of functions f and operations ρ : f is invariant to ρ if it does not change when ρ changes : f ◦ ρ = f . If f and ρ commute , then we say that f is covariant with respect to the operation of ρ : f ◦ ρ = ρ ◦ f ( some sources call this equivariance or same-equivariance ; the typical definition of equivariance is more general , but we will only discuss f and ρ that are endomorphisms in the context of covariance ) . Here we use equivariance to broadly mean considerations of covariance as well as invariance ( since scalars of interest are typically invariant to translation and rotation in physical applications ) for simplicity of discussion . RELATED WORK . “ Large ” point clouds . Point clouds are a ubiquitous data structure found in many domains , even outside of the physical sciences . For the purposes of this work , we focus on relatively small sets of points where each point is more information-rich—for example , carrying information about atom identities , local environments , or other attributes—in contrast to those commonly found in computer vision and robotics which may represent the geometry of a mesh or be sampled from an object , but otherwise not have as much information content . We refer to Guo et al . ( 2020 ) for a survey of this field , but a few of the recently-developed notable approaches include PointNet ( Qi et al. , 2017 ) , deep sets ( Zaheer et al. , 2017 ) , and kernel point convolutions ( Thomas et al. , 2019 ) . Geometric approaches for small point clouds . Many architectures have been proposed to incorporate rotation equivariance by augmenting graph neural networks with geometric attributes that are known to be rotation-invariant , such as bond lengths and angles . SchNet ( Schütt et al. , 2017 ) learns distance-based convolution filters which are used to propagate signals over graphs . PhysNet ( Unke & Meuwly , 2019 ) also refines node representations based on bond lengths , while incorporating a learnable attention mechanism . DimeNet ( Klicpera et al. , 2019 ) extends the information used to calculate node-level representations to include angles between bonds . GemNets ( Klicpera et al. , 2021 ) are related to DimeNets , and incorporate angle information on quadruplets of points within a graph message passing scheme using an architecture optimized for molecular datasets . GNNFF ( Park et al. , 2021 ) generates rotation-covariant results by computing a weighted sum of modulated input vectors based on a graph message passing scheme . Group representation-based approaches . These methods take advantage of group representation theory by first transforming inputs into a space in which rotation-equivariant maps are more easily expressed . This set of methods is powerful , having been used in the past to design rotationand permutation-equivariant models ( Kondor , 2018 ; Thomas et al. , 2018 ; Anderson et al. , 2019 ) , and have even been expanded recently for arbitrary groups ( Finzi et al. , 2021 ) . Attention-based models have also been utilized in this area : SE ( 3 ) Transformers ( Fuchs et al. , 2020 ) extend tensor field networks ( Thomas et al. , 2018 ) with a self-attention mechanism for increased expressivity by incorporating value- and geometry-dependent attention weights . Multivector-valued networks . The approach described herein applies geometric algebra to train deep learning models on point clouds , but using geometric algebra ( also known as Clifford algebra ) to structure the operations of neural networks is not a new concept . Prior work has introduced architectures which generate multivectors using the mathematical structure of geometric algebra as an extension of complex numbers ( Pearson & Bisset , 1992 ; 1994 ) and directly for geometric applications ( Bayro-Corrochano et al. , 1996 ; Buchholz & Sommer , 2008 ; Melnyk et al. , 2020 ) . The approach we present here is similar to many of the ideas presented above ; however , rather than specifying particular rotation-invariant quantities to utilize or learning maps that operate on irreducible representations , we leverage the structure provided by geometric algebra to calculate rotation-invariant and -covariant quantities are of interest . Finally , we use an attention mechanism to attain permutation equivariance in a flexible manner . DOMAIN-SPECIFIC APPLICATIONS . We use three applications , spanning scientific domains from basic physics to biology , to motivate the development of equivariant models . Crystal structure identification . From atoms to colloidal particles , matter often organizes itself into ordered two- or three-dimensional structures . One of the core ideas of materials science is that structure is one of the key determining factors for material behavior . With this perspective in mind , when studying computational models of self-assembling systems we often first identify what structures , if any , have formed in our simulations—a task complicated by naturally-occurring thermal noise , crystallographic defects , dataset scale , and potentially the complexity of the structures themselves . Early efforts to automatically characterize structure led to the widely-used Steinhardt order parameters ( Steinhardt et al. , 1983 ; Mickel et al. , 2013 ; Boattini et al. , 2019 ) , which are rotationallyinvariant magnitudes of sums of spherical harmonics over local particle environments . While Steinhardt order parameters can be effective when studying phase transitions or distinguishing among a small number of phases , determining appropriate hyperparameters—including neighborhood size to consider , spherical harmonic order ℓ to use , and thresholds to identify behaviors of interest—can be difficult ( Mickel et al. , 2013 ) . For this reason , data-driven approaches to analyzing structure have been the subject of great interest in recent years ( Clegg , 2021 ) . Here we identify local environments extracted from ordered structures using a rotation-invariant classifier built on our attention mechanism . While this task is fairly straightforward to perform based on geometry for single-component systems , suitable methods to apply Steinhardt order parameters to systems with multiple types of particles are less obvious ; our deep learning architecture more naturally solves this problem using signals associated with each particle type for every bond . Molecular force regression . One of the most dramatic contributions of deep learning to the field of chemistry lies in constructing fast , accurate approximations of expensive physical calculations ( Unke et al. , 2021 ; Behler , 2021 ) . Machine learning models can be many orders of magnitude faster than the methods used to generate their training data , which can bring vastly more detailed and longer-time simulations into the realm of possibility . Central to the applicability of these methods are issues of symmetry and equivariance : any imperfection in rotational invariance of a learned potential energy function could ruin the thermodynamic behavior of a system , so care must be taken in model design to ensure physical behavior . Here we train models using our attention mechanism to predict the atomic forces calculated using ab initio molecular dynamics and density functional theory ( Chmiela et al. , 2017 ) . These models are conservative , permutation-invariant , and rotationequivariant by construction—all crucial attributes when using models in simulation . Backmapping of coarse-graining operators . When simulating large molecules—such as proteins or other polymers—it is common to employ coarse graining : a process by which groups of particles are merged into ( fewer ) distinct beads , enabling faster simulations by decreasing the number of degrees of freedom of the model ( Marrink & Tieleman , 2013 ) . Although data-driven approaches have been highly successful to formulate coarse graining operations in the forward direction ( that is , from more-detailed to less-detailed systems ) , some problems are best solved using the original , fine-grained system coordinates , which are not directly available in coarse-grained simulations . To demonstrate the potential for our geometric algebra attention mechanism on this task , we train models to predict the coordinates of the heavy atoms that form an amino acid from the centers of mass of the nearest-neighbor amino acids in protein entries found within the Protein Data Bank ( Berman et al. , 2000 ) . Machine learning approaches in the past have treated this problem using simple multilayer perceptrons ( An & Deshmukh , 2020 ) or by encoding the geometry of the problem as image-type data and learning an image translation process ( Li et al. , 2020 ) . We treat this problem much like language translation models using transformers ( Vaswani et al. , 2017 ) , by “ translating ” the point cloud of the coarse-grained neighboring amino acid coordinates into a set of constituent atom coordinates for a central amino acid .

Summary: The paper proposes a method for learning functions that take small point clouds as input, such that they are both permutation and rotation equivariant. To aggregate information from a tuple of points in a rotation equivariant way, it utilizes the geometric product of the points' 3D coordinates. Permutation equivariance is guaranteed by the standard attention framework. The model is evaluated on three scientific tasks: Crystal structure identification, Molecular force regression, and backmapping of coarse-grained operators in molecule simulations.