Geometric Attention Networks for Small Point Clouds

Much of the success of deep learning is drawn from building architectures that1 properly respect underlying symmetry and structure in the data on which they2 operate—a set of considerations that have been united under the banner of geo-3 metric deep learning . Often problems in the physical sciences deal with relatively4 small sets of points in two- or three-dimensional space wherein translation , rota-5 tion , and permutation equivariance are important or even vital for models to be6 useful in practice . In this work , we present an architecture for deep learning on7 these small point clouds with rotation and permutation equivariance , composed of8 a set of products of terms from the geometric algebra and reductions over those9 products using an attention mechanism . The geometric algebra provides valuable10 mathematical structure by which to combine vector , scalar , and other types of11 geometric inputs in a systematic way to account for rotation invariance or covari-12 ance , while attention yields a powerful way to impose permutation equivariance.13 We demonstrate the usefulness of these architectures by training models to solve14 sample problems relevant to physics , chemistry , and biology.15 Introduction16 Deep learning has been immensely successful in solving a wide range of problems over the last17 several years , driven in large part by identifying appropriate ways to embed structure of data and18 symmetry of problems directly into the architecture of the network—an idea at the core of geometric19 deep learning [ 1 ] . Some applications of geometric deep learning include the use of convolutional20 filters in CNNs to attain translational equivariance , or graph convolutions in graph neural networks21 for permutation equivariance.1 Building symmetry into the architecture of a deep neural network22 can improve the data efficiency of the network and guarantee important analytical properties without23 having to rely on the network to learn to approximate them from training data.24 In this work , we derive a family of architectures that is useful in applications from physics to biology,25 where problems often deal with relatively small point clouds of labeled coordinates . These could26 be local environments of particles assembling into a crystal [ 2 ] , atoms in a molecule interacting27 with other atoms [ 3 ] , or coarse-grained beads representing parts of a protein [ 4 ] . In many of these28 applications without the influence of an external field , we are interested in modeling attributes of the29 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 . Submitted to 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Do not distribute . system—such as the identity of a particle ’ s local self-assembly environment , or the potential energy30 of a group of atoms—which are invariant with respect to rotation of the input coordinates , as well as31 permutation in the ordering of points . Here we attain rotation invariance by constructing functions32 from rotation-invariant components of geometric products of input vectors from geometric algebra,33 and permutation invariance by using an attention mechanism to intelligently reduce representations34 over the set of vector products.35 Related Work36 “ Large ” point clouds . Point clouds are a ubiquitous data structure and are often found in domains37 outside of the physical sciences . For the purposes of this work , we focus on comparatively small sets38 of points where the points are relatively information-rich—for example , carrying information about39 atom identities , local environments , or other information—in contrast to point clouds commonly40 found in computer vision and robotics which may represent the geometry of a mesh or be sampled41 from an object , but otherwise not have as much information associated with each point . We refer42 to Guo et al . for a survey of this field [ 5 ] , but a few of the recently-developed notable approaches43 include PointNet [ 6 ] , deep sets [ 7 ] , and kernel point convolutions [ 8 ] .44 Geometric approaches for small point clouds . Many architectures have been proposed to incor-45 porate rotation equivariance by augmenting graph neural networks with geometric attributes that46 are known to be rotation-invariant , such as bond lengths and angles . SchNet [ 9 ] learns distance-47 based convolution filters which are used to propagate signals over graphs . PhysNet [ 10 ] also refines48 node representations based on bond lengths , while incorporating a learnable attention mechanism.49 DimeNet [ 11 ] extends the information used to calculate node-level representations to include angles50 between bonds . GNNFF [ 12 ] generates rotation-covariant results by computing a weighted sum of51 modulated input vectors based on a graph message passing scheme.52 Group representation-based approaches . These methods take advantage of group representation53 theory by first transforming inputs into a space in which rotation-equivariant maps are more eas-54 ily expressed . This set of methods is powerful , having been used in the past to design rotation-55 and permutation-equivariant models [ 13–15 ] , and have even been expanded recently for arbitrary56 groups [ 16 ] . Attention-based models have also been utilized in this area : SE ( 3 ) Transformers [ 17 ] 57 extend tensor field networks [ 14 ] with a self-attention mechanism for increased expressivity by in-58 corporating value- and geometry-dependent attention weights.59 The approach we present here is similar to several of the ideas presented above ; however , rather60 than specifying particular rotation-invariant quantities to utilize or learning maps that operate on61 irreducible representations , we leverage the structure provided by geometric algebra to determine62 which rotation-invariant and -covariant quantities are of interest.63 Geometric Attention Networks64 In this work we formulate deep neural networks using learnable functions consisting of two parts:65 ( 1 ) a set of geometric products of input vectors ; and ( 2 ) a permutation-equivariant reduction over66 these products using an attention mechanism . We describe each of these aspects below.67 Geometric Algebra68 The geometric algebra was developed in the 19th century and provides a consistent framework for69 dealing with scalars and other geometric quantities—such as vectors , areas , and volumes in three-70 dimensional space—in arbitrary dimensions [ 18 ] . Here , we will describe the essential parts of geo-71 metric algebra as related to our proposed attention mechanism , and defer to other works for a more72 thorough description [ 19 ] . The geometric algebra specifies a binary operator , the geometric product,73 that works on multivectors . Multivectors can be expressed as linear combinations of terms from a74 fixed basis set for a given space , such as R2 or R3 ; in three-dimensional space , this yields scalars,75 vectors , bivectors ( which specify signed areas within a plane and have 3 components ) , and trivec-76 tors ( which specify signed volumes and have 1 component ) —a total of 8 linearly independent terms77 for each multivector2 . When rotation invariance is desired , we can utilize the rotation-invariant78 components of a multivector : scalars , trivectors , the norms of vectors , and the norms of bivectors79 are rotation-invariant . As an example , the geometric product of two vectors yields a scalar plus a80 bivector ; the scalar component is the dot product , and the bivector component is related to the cross81 product of the two.82 Geometric algebra provides a general framework that can be used to build up expressive functions83 as linear combinations and geometric products of multivector inputs ; rotation-equivariant quantities84 can then be derived from the products , depending on the application and symmetry of the problem85 of interest . The types of elements produced by a geometric product of two multivectors in R3 with86 the given components are listed in Table 1 below.87 could be used to incorporate different types of scalar , vector , bivector , and trivector quantities ; for93 example , rotations could be input as quaternions , which are isomorphic to the scalar-and-bivector94 product of even numbers of vectors.95 Attention from Geometric Products96 For input point clouds with N points , we can construct a series of successively higher-order geo-97 metric products for all N2 possible pairs , N3 triplets , and so on ; these individual points , pairs , or98 triplets we will call a tuple in this context . In addition to a coordinate r⃗i , we associate a set of values99 vi to each point indexed by i in some space with a given working dimension ( we avoid calling these100 vectors to decrease the confusion with geometric vectors ; these correspond to the non-geometric101 attributes of the point , such as type embeddings ) . To create permutation-covariant functions ( pro-102 ducing a value for each input point ) or permutation-invariant functions ( producing a single output103 value ) , we make use of a simple attention mechanism based on the rotation-invariant attributes of104 each tuple . Attention has been used widely in applying deep learning to a range of problem do-105 mains over the last few years , with particular success in the field of natural language processing [ 20 ] .106 Since we are already generating tuple-wise quantities , we choose to utilize a simpler mechanism107 than the typical dot product self-attention . We specify four functions : a value-generating function108 V , a tuple value-merging function M , a joining function that summarizes the invariant and tuple109 representations J , and a score-generating function S. The functions have the following uses within110 the network:111 • V produces features in the working dimension of the model from the invariants associated112 with each tuple.113 • M merges the 1 , 2 , 3 , or more values associated with a tuple of input points into the114 working dimension of the model . The form of M could be a complex function , a learned115 linear projection for each tuple position , or simply taking the sum of the tuple values.116 • J joins the invariant representations from V and the tuple representations from M. Like117 M , it could be a learned projection or a simple sum function.118 • S generates score logits from the representation of each tuple , which incorporates invariants119 associated with the tuple and the values being associated with each point that is part of the120 tuple . The results from S , passed through a softmax function , will yield the weights for the121 attention mechanism.122 We first calculate the multivector geometric products pijk ... of all combinations of input vectors123 i , j , k , and so on , up to a specified rank ( pairwise attention would produce a two-dimensional124 matrix of products pij ) . We then use V , M , J , and S—together with a function extracting the125 rotation-invariant attributes of a geometric product ( the scalar component , trivector component , and126 the norms of the vector and bivector components ) —as follows for a network producing permutation-127 covariant outputs yi:128 pijk ... = r⃗ir⃗j r⃗k ... qijk ... = invariants ( pijk ... ) vijk ... = J ( V ( qijk ... ) , M ( vi , vj , vk , ... ) ) wijk ... = softmax jk ... ( S ( vijk ... ) ) yi = ∑ jk ... wijk ... vijk ... ( 1 ) If a permutation-invariant reduction is desired , then the softmax and final sum can be performed over129 all tuples simultaneously , rather than for each input point individually . While J , V , and M could in130 principle be used to change the working dimension as permutation-covariant layers are stacked on131 top of each other , in this work we keep the working dimension constant for the sake of easily adding132 residual connections.133 If rotation-covariant , rather than rotation-invariant , behavior is needed for the output of the network,134 the same attention structure can be used with slight modifications ; here , we coerce a vector from135 the product pijk ... ( which consists of directly taking the vector component from products of odd136 numbers of input vectors , or multiplying a bivector by the unit trivector to produce a vector—as137 shown in the last column of Table 1—in the case of even numbers of input vectors ) . These vectors138 can be combined with a scalar rescaling each vector—generated by a learned function R—and the139 attention mechanism to yield140 r⃗i ′ = ∑ jk ... wijk ... R ( qijk ... ) vector ( pijk ... ) . ( 2 ) Results141 We demonstrate the utility of our geometric algebra attention scheme by training deep networks to142 solve three problems appearing in physics , chemistry , and biology . For simplicity , all the models143 presented here utilize pairwise attention with a working depth of 32 units . Value functions V , score144 functions S , and rescaling functions R are simple multilayer perceptrons with a hidden width of145 64 units , with layer normalization applied to the output of V . The network for crystal structure146 identification uses the mean function for merge functions M and join functions J , while the other147 two applications use learned linear projections . Networks are trained for up to 800 epochs using148 the adam optimizer [ 21 ] ; the learning rate is decreased by a factor of 0.75 after the validation set149 loss does not decrease for 20 epochs , and training is ended early if the validation set loss does not150 decrease for 50 epochs . Numerical results are reported as the mean and standard error of the mean151 over 5 samples . Python code under the MIT license implementing each experimental workflow is152 included in the supplementary information.153 Crystal Structure Identification154 On length scales ranging from those of atoms to colloidal particles , matter often organizes itself into155 ordered two- or three-dimensional structures . One of the core ideas of materials science is that struc-156 ture is one of the major determining factors for material behavior . With this perspective in mind,157 when studying computational models of self-assembling systems we often first identify what struc-158 tures , if any , have formed in our simulations—a task complicated by naturally-occurring thermal159 noise , crystallographic defects , and potentially the complexity of the structures themselves . Early160 efforts to automatically characterize structure led to the widely-used Steinhardt order parameters [ 22–161 24 ] , which are rotationally-invariant sums of spherical harmonic magnitudes over local particle162 environments . While the Steinhardt order parameters can be useful when studying phase transi-163 tions or distinguishing among a small number of phases , determining appropriate hyperparameters—164 including neighborhood size to consider , spherical harmonic order ℓ to use , and thresholds to iden-165 tify behaviors of interest—can be difficult [ 23 ] . For this reason , data-driven approaches to analyzing166 structure have been the subject of great interest in recent years [ 25 ] .167 We use geometric attention networks to identify the source structure type of small neighborhoods168 of particles extracted from bulk crystals . We select 8 prototypes of single- and two-component169 crystals from the AFLOW Encyclopedia of Crystallographic Prototypes [ 26 , 27 ] . These structures170 are chosen to demonstrate that models can learn not only geometric information ( cF4-Cu and hP2-171 Mg are similar structures but with a different stacking of their close-packed layers ; the clathrates172 cP46-Si and cF136-Si are also similar , with a different arrangement of many common motifs ) , but173 also the information encoded within each point ( cP2-CsCl and cF8-ZnS differ from cI2-W and174 cF8-C only by their particle type assignments ) . For each structure , we rescale the unit cell such that175 the shortest nearest-neighbor distance over the structure is 1 before replicating the unit cell to consist176 of at least 2048 particles . We then create three samples of each structure by adding Gaussian noise177 with a standard deviation of 10−3 , 5 ·10−2 , and 0.1 separately to the particle coordinates , in order to178 emulate thermal noise . For each particle in the structure , we find the 12 nearest neighbors and their179 associated types using the freud [ 28 ] python library , which we feed into the network as the pairwise180 distance r⃗ij = r⃗j − r⃗i and one-hot-encoded symmetrized type vector ∆tij = [ Iti − Itj , Iti + Itj ] ,181 where I is the identity matrix of dimension corresponding to the maximum number of types.182 We train classifiers with 2 permutation-covariant attention blocks before a final reduction over the en-183 tire particle neighborhood—as shown in Figure 2—in order to categorize local particle environments184 according to their source crystal structure type . These networks rapidly learn to identify structures185 after a few epochs , with a final overall accuracy of 98.7 % ± 0.2 % after training for roughly 45186 minutes on an NVIDIA Titan Xp GPU.187 Molecule Force Regression188 One of the most dramatic contributions of deep learning to the field of chemistry lies in constructing189 fast , accurate approximations of expensive physical calculations [ 29 , 30 ] . Machine learning models190 can be many orders of magnitude faster than the methods used to generate their training data , which191 can bring vastly more detailed and longer-time simulations into the realm of possibility . Central192 to the applicability of these methods are issues of symmetry and equivariance : any imperfection193 in rotational invariance of a learned potential energy function could ruin the proper thermodynamic194 behavior of a model , for example , so models must be carefully designed to ensure physical behavior.195 In a method similar to Batzner et al . [ 31 ] , we train models to predict the per-atom forces calculated196 using ab initio molecular dynamics and density functional theory available in the MD17 dataset [ 3 ] .197 As shown in Figure 3 , we first transform the raw coordinates and types of each atom in a given198 molecule into the pairwise difference and symmetric sum and difference of the coordinates and one-199 hot type encoding for each atom with respect to each other atom , respectively , to fix translation200 invariance and assign type representations to the pairwise particle bonds . We then perform a series201 of geometric attention calculations , calculating new values per atom , which are finally summed to202 produce a scalar energy . The gradient of this energy with respect to the input coordinates is used to203 produce the force output of the network , which ensures that a conservative force field is learned.204 Consistent with previous benchmarks on this dataset , we train networks using the mean squared205 distance loss for each molecule using 1,000 snapshots of forces each as training , validation , and test206 data sets . We also report results for models trained on all molecules ’ data simultaneously , or 8,000207 snapshots each for training , validation , and testing . Training a model on an individual molecule ’ s208 data takes between 30 minutes ( ethanol and malonaldehyde , with nine atoms each ) to two hours209 ( aspirin , with twenty-one atoms ) on an NVIDIA Titan Xp GPU , while the all-molecule dataset210 requires roughly 16 hours to train . Test set losses , expressed as the mean absolute error over each211 force component for each sample , are presented in Table 2.212 Our geometric algebra attention networks produce results competitive with SchNet [ 9 ] , an architec-213 ture using learned radial distance convolution filters . Although the models generated here do not out-214 perform the Neural Equivariant Interatomic Potentials by Batzner et al . [ 31 ] , we note that our models215 are trained for a fraction of the time ( 2 GPU hours and 800 epochs for our method , compared to on216 the order of 8 GPU days and 2500 epochs for NequIP ) and without drastic hyperparameter tuning217 aside from optimizing the number of residual blocks to use in the network architecture . Notably , the218 models trained on all 8,000 molecular snapshots perform significantly better than almost all of the219 specialized models , indicating that additional data could likely improve the results presented here220 even without careful hyperparameter optimization.221 Backmapping Coarse-Graining in Proteins222 When simulating large molecules—such as proteins or other polymers—it is common to employ223 coarse graining : a process by which groups of particles are merged into ( fewer ) distinct beads , en-224 abling faster simulations by decreasing the number of degrees of freedom of the model [ 4 ] . Although225 data-driven approaches have been highly successful to formulate coarse graining operations in the226 forward direction ( that is , from more-detailed to less-detailed systems ) , some problems are best227 solved using the original , fine-grained system coordinates , which are not directly available in coarse-228 grained simulations . To demonstrate the potential for our geometric algebra attention mechanism229 on this task , we train models to predict the coordinates of the heavy atoms that form an amino acid230 from the centers of mass of the nearest-neighbor amino acids . We take 19 protein structures [ 32–50 ] 231 that have high-resolution structural refinements ( with resolution error less than or equal to 1.0Å ) 232 and were published between 2015 and 2020 from the Protein Data Bank [ 51 ] . For applications of233 this method to systems at nonzero temperature , we would expect to be better-served by using an234 architecture that produces distributions instead of only point values , but we disregard this here for235 simplicity ; in other words , here we are teaching models to memorize the results of structure refine-236 ment algorithms , which may be different for each PDB entry . For every amino acid in each entry,237 we create a point cloud of its 12 nearest neighbor amino acid centers of mass , as well as a point238 cloud of the primary amino acid ’ s atomic coordinates relative to its center of mass . Two layers of239 permutation-covariant geometric product attention are applied to the coarse-grained amino acid co-240 ordinates before being passed to a layer which produces a vector output according to Equation 2 by241 augmenting the tuple representation vijk ... of Equation 1 with labels corresponding to the identity242 of the atom that should be produced , so that the value is calculated as243 vatom , ijk ... = J ( vatom , V ( qijk ... ) , M ( vi , vj , vk , ... ) ) . Following this layer—which maps coarse-grained coordinates of amino acids to fine-grained coordi-244 nates of atoms—two rotation-covariant layers are applied to the atomic coordinates to further refine245 them , as shown in Figure 4.246 Because the resolution of the structural refinement algorithms is on the order of 0.5Å or greater , we247 use the training set error as a measure of the learning progress of the models instead of performing248 a standard split of training , validation , and test set data . After training for roughly 3 hours on249 an NVIDIA Titan Xp GPU , models achieve a mean absolute error of 0.128Å ± 0.002Å ( down250 from approximately 0.5Å initially ) , indicating that they are able to learn to reconstruct atomic-scale251 coordinates from coarse-grained positions.252 Discussion253 Overall , we find the architectures formulated here to be useful for a variety of tasks . Rather254 than being limited to operating on bond distances and angles as in SchNet [ 9 ] , PhysNet [ 10 ] , and255 DimeNet [ 11 ] , geometric algebra provides a systematic way to build functions with the desired256 rotation- and permutation-equivariance , with the flexibility to incorporate other types of geomet-257 ric objects ( such as the orientation quaternion commonly used for anisotropic particles in molecular258 dynamics methods [ 52 ] ) . The attention mechanism presented here provides a simple yet powerful259 method to incorporate both geometric and node-level signals . The primitives of our geometric alge-260 bra attention scheme—distances , areas , angles , and volumes—and the calculated attention weights261 naturally lend themselves to interpretability , which we believe will prove useful in distilling insights262 from trained models.263 Limitations264 Combination of terms . Although the architectures presented here work well for the problems we265 have selected , creating geometric products of vectors is only a subset of the valid combinations that266 could be generated . In these cases we have carefully chosen sums and differences of input vectors to267 respect symmetries we would like to impose on the system—such as using the pairwise distance of268 all input coordinates for the molecular force regression task to impose translation invariance—but it269 is possible that more powerful models could be formed by incorporating learned linear combinations270 of inputs or intermediate multivector quantities . We leave this as a topic of future work.271 Computational scaling and neighborhood definition . An obvious limitation to using higher-272 degree correlations lies in the computational complexity and memory scaling of generating tuples,273 which are both proportional to Nr for neighborhoods of N coordinates and tuples of length r. Poly-274 nomial scaling behavior can be ameliorated by restricting which combinations of input points are275 considered , essentially treating the attention weights of all other combinations as 0 . These combi-276 nations could be randomly sampled from all valid indices ijk ... or use more physically-relevant re-277 strictions , such as utilizing the molecular connectivity graph for molecular force regression or edges278 derived from the Voronoi tessellation for other applications . If smoothness of model predictions is279 a concern—as may be the case for learning general N-body interaction potentials , for example—the280 architectures presented here could be augmented by incorporating weights that decay to 0 as bonds281 are broken in the Voronoi diagram graph [ 23 ] .282 Conclusion283 In this work , we have presented a strategy for developing rotation- and permutation-equivariant284 neural network architectures by combining geometric algebra and attention mechanisms . These285 architectures operate directly on the vector , scalar , and other geometric quantities of interest to pro-286 duce outputs which respect desirable symmetries by construction . We believe that the mathematical287 simplicity and the insights derived from attention maps are particularly appealing aspects of the288 algorithms presented here . We hope that these architectures will help a wider range of scientific289 disciplines reap the benefits of geometric deep learning.290 References291 [ 1 ] Michael M. Bronstein , Joan Bruna , Taco Cohen , and Petar Veličković . Geometric deep learn-292 ing : Grids , groups , graphs , geodesics , and gauges . arXiv:2104.13478 [ cs , stat ] , April 2021.293 [ 2 ] Julia Dshemuchadse , Pablo F. Damasceno , Carolyn L. Phillips , Michael Engel , and Sharon C.294 Glotzer . Moving beyond the constraints of chemistry via crystal structure discovery with295 isotropic multiwell pair potentials . Proceedings of the National Academy of Sciences , 118 ( 21 ) ,296 May 2021.297 [ 3 ] Stefan Chmiela , Alexandre Tkatchenko , Huziel E. Sauceda , Igor Poltavsky , Kristof T. Schütt,298 and Klaus-Robert Müller . Machine learning of accurate energy-conserving molecular force299 fields . Science Advances , 3 ( 5 ) : e1603015 , May 2017.300 [ 4 ] Siewert J. Marrink and D. Peter Tieleman . Perspective on the Martini model . Chemical Society301 Reviews , 42 ( 16 ) :6801–6822 , July 2013.302 [ 5 ] Yulan Guo , Hanyun Wang , Qingyong Hu , Hao Liu , Li Liu , and Mohammed Bennamoun . Deep303 learning for 3d point clouds : A survey . IEEE transactions on pattern analysis and machine304 intelligence , 2020.305 [ 6 ] Charles R. Qi , Hao Su , Mo Kaichun , and Leonidas J. Guibas . PointNet : Deep learning on point306 sets for 3D classification and segmentation . In 2017 IEEE Conference on Computer Vision and307 Pattern Recognition ( CVPR ) , pages 77–85 , July 2017.308 [ 7 ] Manzil Zaheer , Satwik Kottur , Siamak Ravanbakhsh , Barnabas Poczos , Russ R. Salakhutdinov,309 and Alexander J. Smola . Deep sets . Advances in Neural Information Processing Systems , 30,310 2017.311 [ 8 ] Hugues Thomas , Charles R. Qi , Jean-Emmanuel Deschaud , Beatriz Marcotegui , François312 Goulette , and Leonidas J. Guibas . KPConv : Flexible and deformable convolution for point313 clouds . Proceedings of the IEEE International Conference on Computer Vision , 2019.314 [ 9 ] K. T. Schütt , P.-J . Kindermans , H. E. Sauceda , S. Chmiela , A. Tkatchenko , and K.-R. Müller.315 SchNet : A continuous-filter convolutional neural network for modeling quantum interactions.316 In Proceedings of the 31st International Conference on Neural Information Processing Sys-317 tems , NIPS ’ 17 , pages 992–1002 , Red Hook , NY , USA , December 2017 . Curran Associates318 Inc.319 [ 10 ] Oliver T. Unke and Markus Meuwly . PhysNet : A neural network for predicting energies,320 forces , dipole moments , and partial charges . Journal of Chemical Theory and Computation,321 15 ( 6 ) :3678–3693 , June 2019.322 [ 11 ] Johannes Klicpera , Janek Groß , and Stephan Günnemann . Directional message passing for323 molecular graphs . In International Conference on Learning Representations , September 2019.324 [ 12 ] Cheol Woo Park , Mordechai Kornbluth , Jonathan Vandermause , Chris Wolverton , Boris Kozin-325 sky , and Jonathan P. Mailoa . Accurate and scalable graph neural network force field and326 molecular dynamics with direct force architecture . npj Computational Materials , 7 ( 1 ) :1–9,327 May 2021.328 [ 13 ] Risi Kondor . N-body networks : A covariant hierarchical neural network architecture for learn-329 ing atomic potentials . arXiv:1803.01588 [ cs ] , March 2018.330 [ 14 ] Nathaniel Thomas , Tess Smidt , Steven Kearnes , Lusann Yang , Li Li , Kai Kohlhoff , and Patrick331 Riley . Tensor field networks : Rotation- and translation-equivariant neural networks for 3d point332 clouds . arXiv:1802.08219 [ cs ] , February 2018.333 [ 15 ] Brandon Anderson , Truong Son Hy , and Risi Kondor . Cormorant : Covariant molecular neural334 networks . In Advances in Neural Information Processing Systems , volume 32 , 2019.335 [ 16 ] Marc Finzi , Max Welling , and Andrew Gordon Wilson . A practical method for constructing336 equivariant multilayer perceptrons for arbitrary matrix groups . arXiv:2104.09459 [ cs , math,337 stat ] , April 2021.338 [ 17 ] Fabian Fuchs , Daniel Worrall , Volker Fischer , and Max Welling . SE ( 3 ) -transformers : 3D339 roto-translation equivariant attention networks . Advances in Neural Information Processing340 Systems , 33:1970–1981 , 2020.341 [ 18 ] William Clifford . Applications of Grassmann ’ s extensive algebra . American Journal of Math-342 ematics , 1 ( 4 ) :350–358 , 1878.343 [ 19 ] Alan Macdonald . A survey of geometric algebra and geometric calculus . Advances in Applied344 Clifford Algebras , 27 ( 1 ) :853–891 , March 2017.345 [ 20 ] Ashish Vaswani , Noam Shazeer , Niki Parmar , Jakob Uszkoreit , Llion Jones , Aidan N Gomez,346 Łukasz Kaiser , and Illia Polosukhin . Attention is all you need . In I. Guyon , U. V. Luxburg,347 S. Bengio , H. Wallach , R. Fergus , S. Vishwanathan , and R. Garnett , editors , Advances in348 Neural Information Processing Systems 30 , pages 5998–6008 . Curran Associates , Inc. , 2017.349 [ 21 ] Diederik P. Kingma and Jimmy Ba . Adam : A method for stochastic optimization . In Yoshua350 Bengio and Yann LeCun , editors , 3rd International Conference on Learning Representations,351 ICLR 2015 , San Diego , CA , USA , May 7-9 , 2015 , Conference Track Proceedings , 2015.352 [ 22 ] Paul J. Steinhardt , David R. Nelson , and Marco Ronchetti . Bond-orientational order in liquids353 and glasses . Physical Review B , 28 ( 2 ) :784–805 , July 1983.354 [ 23 ] Walter Mickel , Sebastian C. Kapfer , Gerd E. Schröder-Turk , and Klaus Mecke . Shortcomings355 of the bond orientational order parameters for the analysis of disordered particulate matter . The356 Journal of Chemical Physics , 138 ( 4 ) :044501 , January 2013.357 [ 24 ] Emanuele Boattini , Marjolein Dijkstra , and Laura Filion . Unsupervised learning for local358 structure detection in colloidal systems . The Journal of Chemical Physics , 151 ( 15 ) :154901,359 October 2019.360 [ 25 ] Paul S. Clegg . Characterising soft matter using machine learning . Soft Matter , 17 ( 15 ) :3991–361 4005 , April 2021.362 [ 26 ] Michael J. Mehl , David Hicks , Cormac Toher , Ohad Levy , Robert M. Hanson , Gus Hart , and363 Stefano Curtarolo . The AFLOW library of crystallographic prototypes : Part 1 . Computational364 Materials Science , 136 : S1–S828 , August 2017.365 [ 27 ] David Hicks , Michael J. Mehl , Eric Gossett , Cormac Toher , Ohad Levy , Robert M. Hanson,366 Gus Hart , and Stefano Curtarolo . The AFLOW library of crystallographic prototypes : Part 2.367 Computational Materials Science , 161 : S1–S1011 , April 2019.368 [ 28 ] Vyas Ramasubramani , Bradley D. Dice , Eric S. Harper , Matthew P. Spellings , Joshua A. An-369 derson , and Sharon C. Glotzer . Freud : A software suite for high throughput analysis of particle370 simulation data . Computer Physics Communications , 254:107275 , September 2020.371 [ 29 ] Oliver T. Unke , Stefan Chmiela , Huziel E. Sauceda , Michael Gastegger , Igor Poltavsky,372 Kristof T. Schütt , Alexandre Tkatchenko , and Klaus-Robert Müller . Machine learning force373 fields . Chemical Reviews , March 2021.374 [ 30 ] Jörg Behler . Four generations of high-dimensional neural network potentials . Chemical Re-375 views , March 2021.376 [ 31 ] Simon Batzner , Tess E. Smidt , Lixin Sun , Jonathan P. Mailoa , Mordechai Kornbluth , Nicola377 Molinari , and Boris Kozinsky . SE ( 3 ) -equivariant graph neural networks for data-efficient and378 accurate interatomic potentials . arXiv:2101.03164 [ cond-mat , physics : physics ] , January 2021.379 [ 32 ] Akihiko Nakamura , Takuya Ishida , Katsuhiro Kusaka , Taro Yamada , Shinya Fushinobu , Ichiro380 Tanaka , Satoshi Kaneko , Kazunori Ohta , Hiroaki Tanaka , Koji Inaka , Yoshiki Higuchi , Nobuo381 Niimura , Masahiro Samejima , and Kiyohiko Igarashi . `` Newton ’ s cradle '' proton relay with382 amide–imidic acid tautomerization in inverting cellulase visualized by neutron crystallography.383 Science Advances , 1 ( 7 ) : e1500263 , August 2015.384 [ 33 ] Y. Hirano , S. Kimura , and T. Tamada . High-resolution crystal structures of the solubilized385 domain of porcine cytochrome b5 . Acta Crystallographica Section D : Biological Crystallog-386 raphy , 71 ( 7 ) :1572–1581 , July 2015.387 [ 34 ] Shigeru Matsuoka , Shigeru Sugiyama , Daisuke Matsuoka , Mika Hirose , Sébastien Lethu,388 Hikaru Ano , Toshiaki Hara , Osamu Ichihara , S. Roy Kimura , Satoshi Murakami , Hanako389 Ishida , Eiichi Mizohata , Tsuyoshi Inoue , and Michio Murata . Water-mediated recognition390 of simple alkyl chains by heart-type fatty-acid-binding protein . Angewandte Chemie Interna-391 tional Edition , 54 ( 5 ) :1508–1511 , 2015.392 [ 35 ] Yoshihiko Furuike , Yuka Akita , Ikuko Miyahara , and Nobuo Kamiya . ADP-ribose pyrophos-393 phatase reaction in crystalline state conducted by consecutive binding of two manganese ( ii ) 394 ions as cofactors . Biochemistry , 55 ( 12 ) :1801–1812 , March 2016.395 [ 36 ] Hiraku Ohno , Kazuki Takeda , Satomi Niwa , Tomotaka Tsujinaka , Yuya Hanazono , Yu Hirano,396 and Kunio Miki . Crystallographic characterization of the high-potential iron-sulfur protein in397 the oxidized state at 0.8 Å resolution . PLOS ONE , 12 ( 5 ) : e0178183 , May 2017.398 [ 37 ] Harry P. Austin , Mark D. Allen , Bryon S. Donohoe , Nicholas A. Rorrer , Fiona L. Kearns , Ro-399 drigo L. Silveira , Benjamin C. Pollard , Graham Dominick , Ramona Duman , Kamel El Omari,400 Vitaliy Mykhaylyk , Armin Wagner , William E. Michener , Antonella Amore , Munir S. Skaf,401 Michael F. Crowley , Alan W. Thorne , Christopher W. Johnson , H. Lee Woodcock , John E.402 McGeehan , and Gregg T. Beckham . Characterization and engineering of a plastic-degrading403 aromatic polyesterase . Proceedings of the National Academy of Sciences , 115 ( 19 ) : E4350–404 E4357 , May 2018.405 [ 38 ] Alessandro Vergara , Marco Caterino , and Antonello Merlino . Raman-markers of X-ray radia-406 tion damage of proteins . International Journal of Biological Macromolecules , 111:1194–1205,407 May 2018.408 [ 39 ] Alexander T. Baker , Alexander Greenshields-Watson , Lynda Coughlan , James A. Davies,409 Hanni Uusi-Kerttula , David K. Cole , Pierre J. Rizkallah , and Alan L. Parker . Diversity within410 the adenovirus fiber knob hypervariable loops influences primary receptor interactions . Nature411 Communications , 10 ( 1 ) :741 , February 2019.412 [ 40 ] K. Takaba , Y. Tai , H. Eki , H.-A . Dao , Y. Hanazono , K. Hasegawa , K. Miki , and K. Takeda.413 Subatomic resolution X-ray structures of green fluorescent protein . IUCrJ , 6 ( 3 ) :387–400 , May414 2019.415 [ 41 ] Janet M. Paterson , Amy J. Shaw , Ian Burns , Alister W. Dodds , Alpana Prasad , Ken B. Reid,416 Trevor J. Greenhough , and Annette K. Shrive . Atomic-resolution crystal structures of the417 immune protein conglutinin from cow reveal specific interactions of its binding site with N-418 acetylglucosamine . Journal of Biological Chemistry , 294 ( 45 ) :17155–17165 , November 2019.419 [ 42 ] Wei Tian , Peiqiang Yan , Ning Xu , Arghya Chakravorty , Robert Liefke , Qiaoran Xi , and420 Zhanxin Wang . The HRP3 PWWP domain recognizes the minor groove of double-stranded421 DNA and recruits HRP3 to chromatin . Nucleic Acids Research , 47 ( 10 ) :5436–5448 , June 2019.422 [ 43 ] Matthew J Schellenberg , C Denise Appel , Amanda A Riccio , Logan R Butler , Juno M Krahn,423 Jenna A Liebermann , Felipe Cortés-Ledesma , and R Scott Williams . Ubiquitin stimulated424 reversal of topoisomerase 2 DNA-protein crosslinks by TDP2 . Nucleic Acids Research,425 48 ( 11 ) :6310–6325 , June 2020.426 [ 44 ] Marina Plaza-Garrido , Mª Carmen Salinas-García , José C. Martínez , and Ana Cámara-Artigas.427 The effect of an engineered ATCUN motif on the structure and biophysical properties of428 the SH3 domain of c-Src tyrosine kinase . JBIC Journal of Biological Inorganic Chemistry,429 25 ( 4 ) :621–634 , June 2020.430 [ 45 ] Shanshan Wu , Tam T. T. N. Nguyen , Olga V. Moroz , Johan P. Turkenburg , Jens E. Nielsen,431 Keith S. Wilson , Kasper D. Rand , and Kaare Teilum . Conformational heterogeneity of Sav-432 inase from NMR , HDX-MS and X-ray diffraction analysis . PeerJ , 8 : e9408 , June 2020.433 [ 46 ] Barbara Franke , Marta Veses-Garcia , Kay Diederichs , Heather Allison , Daniel J. Rigden , and434 Olga Mayans . Structural annotation of the conserved carbohydrate esterase vb_24B_21 from435 Shiga toxin-encoding bacteriophage Φ24B . Journal of Structural Biology , 212 ( 1 ) :107596,436 October 2020.437 [ 47 ] Matthias Barone , Matthias Müller , Slim Chiha , Jiang Ren , Dominik Albat , Arne Soicke,438 Stephan Dohmen , Marco Klein , Judith Bruns , Maarten van Dinther , Robert Opitz , Peter Linde-439 mann , Monika Beerbaum , Kathrin Motzny , Yvette Roske , Peter Schmieder , Rudolf Volkmer,440 Marc Nazaré , Udo Heinemann , Hartmut Oschkinat , Peter ten Dijke , Hans-Günther Schmalz,441 and Ronald Kühne . Designed nanomolar small-molecule inhibitors of Ena/VASP EVH1 inter-442 action impair invasion and extravasation of breast cancer cells . Proceedings of the National443 Academy of Sciences , 117 ( 47 ) :29684–29690 , November 2020.444 [ 48 ] Renee Otten , Ricardo A. P. Pádua , H. Adrian Bunzel , Vy Nguyen , Warintra Pitsawong,445 MacKenzie Patterson , Shuo Sui , Sarah L. Perry , Aina E. Cohen , Donald Hilvert , and Dorothee446 Kern . How directed evolution reshapes the energy landscape in an enzyme to boost catalysis.447 Science , 370 ( 6523 ) :1442–1446 , December 2020.448 [ 49 ] Karla Frydenvang , Darryl S. Pickering , Giridhar U. Kshirsagar , Giulia Chemi , Sandra Gemma,449 Desiree Sprogøe , Anne Mette Kærn , Simone Brogi , Giuseppe Campiani , Stefania Butini,450 and Jette Sandholm Kastrup . Ionotropic glutamate receptor GluA2 in complex with bicyclic451 pyrimidinedione-based compounds : When small compound modifications have distinct effects452 on binding interactions . ACS Chemical Neuroscience , 11 ( 12 ) :1791–1800 , June 2020.453 [ 50 ] Steffen Glöckner , Andreas Heine , and Gerhard Klebe . A proof-of-concept fragment screening454 of a hit-validated 96-compounds library against human carbonic anhydrase ii . Biomolecules,455 10 ( 4 ) :518 , April 2020.456 [ 51 ] Helen M. Berman , John Westbrook , Zukang Feng , Gary Gilliland , T. N. Bhat , Helge Weissig,457 Ilya N. Shindyalov , and Philip E. Bourne . The Protein Data Bank . Nucleic Acids Research,458 28 ( 1 ) :235–242 , January 2000.459 [ 52 ] H. Kamberaj , R. J . Low , and M. P. Neal . Time reversible and symplectic integrators460 for molecular dynamics simulations of rigid molecules . The Journal of Chemical Physics,461 122 ( 22 ) :224114 , June 2005.462 Checklist463 1 . For all authors ... 464 ( a ) Do the main claims made in the abstract and introduction accurately reflect the paper ’ s465 contributions and scope ? [ Yes ] 466 ( b ) Did you describe the limitations of your work ? [ Yes ] 467 ( c ) Did you discuss any potential negative societal impacts of your work ? [ N/A ] We feel468 that the methods presented in this paper are sufficiently agnostic to applications that it469 is hard to directly address any negative or positive societal impacts of the work.470 ( d ) Have you read the ethics review guidelines and ensured that your paper conforms to471 them ? [ Yes ] 472 2 . If you are including theoretical results ... 473 ( a ) Did you state the full set of assumptions of all theoretical results ? [ N/A ] 474 ( b ) Did you include complete proofs of all theoretical results ? [ N/A ] 475 3 . If you ran experiments ... 476 ( a ) Did you include the code , data , and instructions needed to reproduce the main exper-477 imental results ( either in the supplemental material or as a URL ) ? [ Yes ] Code and478 scripts to reproduce results to be included as SI before the appropriate deadline.479 ( b ) Did you specify all the training details ( e.g. , data splits , hyperparameters , how they480 were chosen ) ? [ Yes ] Details of training and architectures are provided in the Results481 section , as well as the code that will be included in the SI.482 ( c ) Did you report error bars ( e.g. , with respect to the random seed after running experi-483 ments multiple times ) ? [ Yes ] 484 ( d ) Did you include the total amount of compute and the type of resources used ( e.g. , type485 of GPUs , internal cluster , or cloud provider ) ? [ Yes ] 486 4 . If you are using existing assets ( e.g. , code , data , models ) or curating/releasing new assets ... 487 ( a ) If your work uses existing assets , did you cite the creators ? [ N/A ] No existing assets488 were evaluated.489 ( b ) Did you mention the license of the assets ? [ Yes ] 490 ( c ) Did you include any new assets either in the supplemental material or as a URL ? [ Yes ] 491 Trained models and workflows will be included in the SI.492 ( d ) Did you discuss whether and how consent was obtained from people whose data493 you ’ re using/curating ? [ N/A ] 494 ( e ) Did you discuss whether the data you are using/curating contains personally identifi-495 able information or offensive content ? [ N/A ] 496 5 . If you used crowdsourcing or conducted research with human subjects ... 497 ( a ) Did you include the full text of instructions given to participants and screenshots , if498 applicable ? [ N/A ] 499 ( b ) Did you describe any potential participant risks , with links to Institutional Review500 Board ( IRB ) approvals , if applicable ? [ N/A ] 501 ( c ) Did you include the estimated hourly wage paid to participants and the total amount502 spent on participant compensation ? [ N/A ] 503

Summary: This paper studies attention from a geometric algebra perspective. Geometric algebra provides a useful tool to combine vector, scalar, and other types of geometric inputs to account for rotation invariance/equivariance and permutation invariance/equivariance. In this paper, the attention mechanism is factored into four functions: a value-generating function, a tuple value-merging function, a joining function that summarizes the invariant and tuple representations, and a score-generating function. This new perspective on attention is novel and can motivate more follow-up works on studying high-order attention. On the other hand, however, this paper doesn't introduce new module design beyond interpreting attention and the actual implementation is the same with self-attention.