summaries: 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. The paper proposes a geometric algebra attention network for small point clouds. The attention based on geometric algebra’s multivector is rotation equivariant and permutation equivariant. Specifically, attention is composed of four functions that operate on tuples and respect the desired equivariance. Moreover, the paper validates its proposed geometric algebra attention on three domain-specific applications including crystal structure identification, molecular force regression, and back mapping of coarse-graining operators.  This paper proposes a rotation and permutation equivariant geometric deep learning model for problems where the data is represented as small points clouds. The equivariance properties are achieved by leveraging geometric algebra formulations. More specifically, rotation-equivariance is accomplished by geometric products of multivectors and permutation-equivariance by using an attention mechanism over invariant terms of these products. This model is evaluated in three different applications showing better or comparable results than existing approaches. These models also offer additional features like the analysis of the attention maps produced. This paper presents a deep neural network for processing small sets of points in three-dimensional space with a rotation-invariance and a permutation equivariance.  The proposed method is twofold. One utilizes a geometric product of geometric algebra on input vectors to achieve the rotation-invariant attributes. The other is a permutation-equivariant reduction over the geometric products using an attention mechanism. The authors demonstrate applications in physical science such as crystal structure identification, molecular force regression, and backmapping of coarse-graining operators. 
abstractText: 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 twoor 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 rotationand 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.
references.authors: Yaxin An Sanket A. Deshmukh
references.title: Machine learning approach for accurate backmapping of coarsegrained models to all-atom models
references.venue: Chemical Communications,
references.authors: Brandon Anderson Truong Son Hy Risi Kondor
references.title: Cormorant: Covariant molecular neural networks
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Alexander T. Baker Alexander Greenshields-Watson Lynda Coughlan James A. Davies Hanni Uusi-Kerttula David K. Cole Pierre J. Rizkallah Alan L. Parker
references.title: Diversity within the adenovirus fiber knob hypervariable loops influences primary receptor interactions
references.venue: Nature Communications,
references.authors: Simon Batzner Tess E. Smidt Lixin Sun Jonathan P. Mailoa Mordechai Kornbluth Nicola Molinari Boris
references.title: Kozinsky. SE(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials
references.authors: E. Bayro-Corrochano S. Buchholz G. Sommer
references.title: A new self-organizing neural network using geometric algebra
references.venue: In Proceedings of 13th International Conference on Pattern Recognition,
references.authors: Jörg Behler. Four generations of high-dimensional neural network potentials. Chemical Reviews March
references.title: ISSN 0009-2665
references.venue: doi: 10.1021/acs.chemrev.0c00868.
references.authors: Helen M. Berman John Westbrook Zukang Feng Gary Gilliland T.N. Bhat Helge Weissig Ilya N. Shindyalov Philip E. Bourne
references.title: The Protein Data Bank
references.venue: Nucleic Acids Research,
references.authors: Emanuele Boattini Marjolein Dijkstra Laura Filion
references.title: Unsupervised learning for local structure detection in colloidal systems
references.venue: The Journal of Chemical Physics,
references.authors: Michael M. Bronstein Joan Bruna Taco Cohen Petar Veličković
references.title: Geometric deep learning: Grids, groups, graphs, geodesics, and gauges
references.authors: Sven Buchholz Gerald Sommer
references.title: On Clifford neurons and Clifford multi-layer perceptrons
references.venue: Neural Networks: The Official Journal of the International Neural Network Society,
references.authors: Stefan Chmiela Alexandre Tkatchenko Huziel E. Sauceda Igor Poltavsky Kristof T. Schütt Klaus-Robert Müller
references.title: Machine learning of accurate energy-conserving molecular force fields
references.venue: Science Advances,
references.authors: Paul S. Clegg
references.title: Characterising soft matter using machine learning
references.venue: Soft Matter,
references.authors: Julia Dshemuchadse Pablo F. Damasceno Carolyn L. Phillips Michael Engel Sharon C. Glotzer
references.title: Moving beyond the constraints of chemistry via crystal structure discovery with isotropic multiwell pair potentials
references.venue: Proceedings of the National Academy of Sciences,
references.authors: Marc Finzi Max Welling Andrew Gordon Wilson
references.title: A practical method for constructing equivariant multilayer perceptrons for arbitrary matrix groups
references.authors: Barbara Franke Marta Veses-Garcia Kay Diederichs Heather Allison Daniel J. Rigden Olga Mayans
references.title: Structural annotation of the conserved carbohydrate esterase vb_24B_21 from Shiga toxin-encoding bacteriophage Φ24B
references.venue: Journal of Structural Biology,
references.authors: Karla Frydenvang Darryl S. Pickering Giridhar U. Kshirsagar Giulia Chemi Sandra Gemma Desiree Sprogøe Anne Mette Kærn Simone Brogi Giuseppe Campiani Stefania Butini Jette Sandholm Kastrup
references.title: Ionotropic glutamate receptor GluA2 in complex with bicyclic pyrimidinedione-based compounds: When small compound modifications have distinct effects on binding interactions
references.venue: ACS Chemical Neuroscience,
references.authors: Fabian Fuchs Daniel Worrall Volker Fischer Max Welling
references.title: SE(3)-transformers: 3D rototranslation equivariant attention networks
references.venue: Advances in Neural Information Processing Systems,
references.authors: Yoshihiko Furuike Yuka Akita Ikuko Miyahara Nobuo Kamiya
references.title: ADP-ribose pyrophosphatase reaction in crystalline state conducted by consecutive binding of two manganese(ii) ions as cofactors
references.authors: Steffen Glöckner Andreas Heine Gerhard Klebe
references.title: A proof-of-concept fragment screening of a hit-validated 96-compounds library against human carbonic anhydrase
references.venue: ii. Biomolecules,
references.authors: Yulan Guo Hanyun Wang Qingyong Hu Hao Liu Li Liu Mohammed Bennamoun
references.title: Deep learning for 3D point clouds: A survey
references.venue: IEEE transactions on pattern analysis and machine intelligence,
references.authors: Y. Hirano S. Kimura T. Tamada
references.title: High-resolution crystal structures of the solubilized domain of porcine cytochrome b5
references.venue: Acta Crystallographica Section D: Biological Crystallography,
references.authors: H. Kamberaj R.J. Low M.P. Neal
references.title: Time reversible and symplectic integrators for molecular dynamics simulations of rigid molecules
references.venue: The Journal of Chemical Physics,
references.authors: Diederik P. Kingma Jimmy Ba
references.title: Adam: A method for stochastic optimization
references.venue: In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations,
references.authors: Johannes Klicpera Janek Groß Stephan Günnemann
references.title: Directional message passing for molecular graphs
references.venue: In International Conference on Learning Representations,
references.authors: Johannes Klicpera Florian Becker Stephan Günnemann
references.title: GemNet: Universal directional graph neural networks for molecules
references.authors: Risi Kondor
references.title: N-body networks: A covariant hierarchical neural network architecture for learning atomic potentials
references.venue: [cs],
references.authors: Wei Li Craig Burkhart Patrycja Polińska Vagelis Harmandaris Manolis Doxastakis
references.title: Backmapping coarse-grained macromolecules: An efficient and versatile machine learning approach
references.venue: The Journal of Chemical Physics,
references.authors: Alan Macdonald
references.title: A survey of geometric algebra and geometric calculus
references.venue: Advances in Applied Clifford Algebras,
references.authors: Siewert J. Marrink D. Peter Tieleman
references.title: Perspective on the Martini model
references.venue: Chemical Society Reviews,
references.authors: Shigeru Matsuoka Shigeru Sugiyama Daisuke Matsuoka Mika Hirose Sébastien Lethu Hikaru Ano Toshiaki Hara Osamu Ichihara S. Roy Kimura Satoshi Murakami Hanako Ishida Eiichi Mizohata Tsuyoshi Inoue Michio Murata
references.title: Water-mediated recognition of simple alkyl chains by heart-type fatty-acid-binding protein
references.venue: Angewandte Chemie International Edition,
references.authors: Michael J. Mehl David Hicks Cormac Toher Ohad Levy Robert M. Hanson Gus Hart Stefano Curtarolo
references.title: The AFLOW library of crystallographic prototypes: Part 1
references.venue: Computational Materials Science,
references.authors: Pavlo Melnyk M. Felsberg Mårten Wadenbäck
references.title: Embed me if you can: A geometric perceptron
references.venue: ArXiv,
references.authors: Walter Mickel Sebastian C. Kapfer Gerd E. Schröder-Turk Klaus Mecke
references.title: Shortcomings of the bond orientational order parameters for the analysis of disordered particulate matter
references.venue: The Journal of Chemical Physics,
references.authors: Akihiko Nakamura Takuya Ishida Katsuhiro Kusaka Taro Yamada Shinya Fushinobu Ichiro Tanaka Satoshi Kaneko Kazunori Ohta Hiroaki Tanaka Koji Inaka Yoshiki Higuchi Nobuo Niimura Masahiro Samejima Kiyohiko Igarashi
references.title: Newton’s cradle" proton relay with amide–imidic acid tautomerization in inverting cellulase visualized by neutron crystallography
references.venue: Science Advances,
references.authors: Hiraku Ohno Kazuki Takeda Satomi Niwa Tomotaka Tsujinaka Yuya Hanazono Yu Hirano Kunio Miki
references.title: Crystallographic characterization of the high-potential iron-sulfur protein in the oxidized state at 0.8 Å resolution
references.venue: PLOS ONE,
references.authors: Renee Otten Ricardo A.P. Pádua H. Adrian Bunzel Vy Nguyen Warintra Pitsawong MacKenzie Patterson Shuo Sui Sarah L. Perry Aina E. Cohen Donald Hilvert Dorothee Kern
references.title: How directed evolution reshapes the energy landscape in an enzyme to boost catalysis
references.venue: doi: 10.1126/science
references.authors: Cheol Woo Park Mordechai Kornbluth Jonathan Vandermause Chris Wolverton Boris Kozinsky Jonathan P. Mailoa
references.title: Accurate and scalable graph neural network force field and molecular dynamics with direct force architecture
references.venue: npj Computational Materials,
references.authors: Janet M. Paterson Amy J. Shaw Ian Burns Alister W. Dodds Alpana Prasad Ken B. Reid Trevor J. Greenhough Annette K. Shrive
references.title: Atomic-resolution crystal structures of the immune protein conglutinin from cow reveal specific interactions of its binding site with N-acetylglucosamine
references.venue: Journal of Biological Chemistry,
references.authors: J.K. Pearson D.L. Bisset
references.title: Neural networks in the Clifford domain
references.venue: IEEE International Conference on Neural Networks (ICNN’94),
references.authors: Justin Pearson David Bisset
references.title: Back propagation in a Clifford algebra
references.venue: Proceedings of the International Conference of Artificial Neural Networks (ICANN),
references.authors: Marina Plaza-Garrido Ma Carmen Salinas-García José C. Martínez Ana Cámara-Artigas
references.title: The effect of an engineered ATCUN motif on the structure and biophysical properties of the SH3 domain of c-Src tyrosine kinase
references.venue: JBIC Journal of Biological Inorganic Chemistry,
references.authors: Charles R. Qi Hao Su Mo Kaichun Leonidas J. Guibas
references.title: PointNet: Deep learning on point sets for 3D classification and segmentation
references.venue: IEEE Conference on Computer Vision and Pattern Recognition (CVPR),
references.authors: Vyas Ramasubramani Bradley D. Dice Eric S. Harper Matthew P. Spellings Joshua A. Anderson Sharon C. Glotzer
references.title: Freud: A software suite for high throughput analysis of particle simulation data
references.venue: Computer Physics Communications,
references.authors: Anna Rogers Olga Kovaleva Anna Rumshisky
references.title: A Primer in BERTology
references.venue: What We Know About How BERT Works. Transactions of the Association for Computational Linguistics,
references.authors: Matthew J Schellenberg C Denise Appel Amanda A Riccio Logan R Butler Juno M Krahn Jenna A Liebermann Felipe Cortés-Ledesma R Scott Williams
references.title: Ubiquitin stimulated reversal of topoisomerase 2 DNA-protein crosslinks by TDP2
references.venue: Nucleic Acids Research,
references.authors: K.T. Schütt P.-J. Kindermans H.E. Sauceda S. Chmiela A. Tkatchenko K.-R. Müller
references.title: SchNet: A continuous-filter convolutional neural network for modeling quantum interactions
references.venue: In Proceedings of the 31st International Conference on Neural Information Processing Systems,
references.authors: Matthew Spellings Sharon C Glotzer
references.title: Machine learning for crystal identification and discovery
references.venue: AIChE Journal,
references.authors: Paul J. Steinhardt David R. Nelson Marco Ronchetti
references.title: Bond-orientational order in liquids and glasses
references.venue: Physical Review B,
references.authors: K. Takaba Y. Tai H. Eki H.-A. Dao Y. Hanazono K. Hasegawa K. Miki K. Takeda
references.title: Subatomic resolution X-ray structures of green fluorescent protein
references.venue: IUCrJ,
references.authors: Hugues Thomas Charles R. Qi Jean-Emmanuel Deschaud Beatriz Marcotegui François Goulette Leonidas J. Guibas
references.title: KPConv: Flexible and deformable convolution for point clouds
references.venue: Proceedings of the IEEE International Conference on Computer Vision,
references.authors: Nathaniel Thomas Tess Smidt Steven Kearnes Lusann Yang Li Li Kai Kohlhoff Patrick Riley
references.title: Tensor field networks: Rotation- and translation-equivariant neural networks for 3d point clouds
references.venue: [cs],
references.authors: Wei Tian Peiqiang Yan Ning Xu Arghya Chakravorty Robert Liefke Qiaoran Xi Zhanxin Wang
references.title: The HRP3 PWWP domain recognizes the minor groove of double-stranded DNA and recruits HRP3 to chromatin
references.venue: Nucleic Acids Research,
references.authors: Oliver T. Unke Markus Meuwly
references.title: PhysNet: A neural network for predicting energies, forces, dipole moments, and partial charges
references.venue: Journal of Chemical Theory and Computation,
references.authors: Oliver T. Unke Stefan Chmiela Huziel E. Sauceda Michael Gastegger Igor Poltavsky Kristof T. Schütt Alexandre Tkatchenko Klaus-Robert Müller. Machine learning force fields. Chemical Reviews March
references.title: ISSN 0009-2665
references.venue: doi: 10.1021/acs.chemrev.0c01111.
references.authors: Ashish Vaswani Noam Shazeer Niki Parmar Jakob Uszkoreit Llion Jones Aidan N Gomez Łukasz Kaiser Illia Polosukhin
references.title: Attention is all you need
references.venue: Advances in Neural Information Processing Systems
references.authors: Alessandro Vergara Marco Caterino Antonello Merlino
references.title: Raman-markers of X-ray radiation damage of proteins
references.venue: International Journal of Biological Macromolecules,
references.authors: Shanshan Wu Tam T.T.N. Nguyen Olga V. Moroz Johan P. Turkenburg Jens E. Nielsen Keith S. Wilson Kasper D. Rand Kaare Teilum
references.title: Conformational heterogeneity of Savinase from NMR, HDX-MS and X-ray diffraction analysis
references.venue: PeerJ, 8:e9408,
sections.text: 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.
sections.heading: BACKGROUND
sections.text: 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.
sections.heading: RELATED WORK
sections.text: “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.
sections.heading: DOMAIN-SPECIFIC APPLICATIONS
sections.text: 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.
sections.heading: METHODS
sections.text: In this work we formulate deep neural networks using learnable functions consisting of two main parts: (1) a set of geometric products (from the geometric algebra in three spatial dimensions) of input vectors; and (2) a permutation-equivariant reduction over these products using an attention mechanism. We describe each of these aspects below.
sections.heading: ATTENTION FROM GEOMETRIC PRODUCTS
sections.text: We begin with a description of the bare essentials of geometric algebra used in this work. For more details, we refer the reader to Appendix A. Briefly, geometric algebra provides mathematical structure to deal with geometric objects—such as points and planes—using a common language for arbitrary numbers of spatial dimensions. The primary objects dealt with using geometric algebra are multivectors, which consist of linear combinations of basis elements; in three dimensions, these basis elements are one scalar component, three vector components, three bivector components, and one trivector component. We can calculate the so-called geometric product of two multivectors to yield a new multivector; for example, the geometric product of two vectors yields a scalar (that is the dot product of the vectors) plus a bivector (related to the cross product of the vectors). The geometric product is not commutative; for two general multivectors A and B, we denote the product C as C = AB. In this work, we use the geometric product to combine groups of input vectors and systematically extract rotation-invariant quantities of interest—such as distances, bond angles, and volumes—for use in network layers.
  We specify a computational complexity (i.e. whether to compose calculations from pairs, triplets, quadruplets, or larger groups of points) and, for input point clouds with N points, we construct all N2 possible pairs, N3 triplets, and so on;2 we call these groups of points a tuple here. In addition to a coordinate r⃗i, we associate a set of values vi to each point indexed by i in some space with a given working dimension (we avoid calling these vectors to decrease the confusion with geometric vectors; values correspond to the non-geometric attributes of the point, such as type embeddings for atomic systems). To create permutation-covariant functions (producing a value for each input point of a point cloud) or permutation-invariant functions (producing a single output value for the entire cloud), we make use of a simple attention mechanism based on the rotation-invariant attributes of each tuple. Attention has been used widely in applying deep learning to a range of problem domains over the last few years, with particular success in the field of natural language processing (Vaswani et al., 2017). Instead of the dot product self-attention commonly used in language models—which operates on pairs of elements—we use a geometrically-informed attention mechanism that operates on each tuple. We specify four functions: a value-generating function V , a tuple value-merging function M, a joining function that summarizes the geometry and tuple value representations J , and a score-generating function S . The functions have the following uses within the network:
  • V produces features in the working dimension of the model from the rotation-invariant geometric quantities associated with each tuple—specifically, rotation-invariant attributes of each geometric product pijk... = r⃗ir⃗j r⃗k..., which include bond distances, angles, volumes, and other geometric attributes. Geometric algebra provides a framework to systematically extract these attributes: for a general multivector, the rotation-invariant quantities are the scalar component, trivector component, the norm of the vector component, and the norm of the bivector component. In the networks presented here, we use MLPs with a single hidden layer for V .
  2To avoid polynomial scaling, the set of products could be reduced according to the edges of a specified graph, the Voronoi diagram of the point cloud, or by randomly sampling tuples of points; however, we leave these extensions to future work.
  • M merges the 1, 2, 3, or more values associated with a tuple of input points into the working dimension of the model. The form of M could be a complex learned function, a linear projection for each tuple position, or simply the summation of input tuple values.
  • J joins the rotation-invariant representations from V and the tuple value representations from M. Like M, it could vary in complexity from a simple sum to a complex learned nonlinear function.
  • S generates score logits from the representation of each tuple, incorporating both geometry and value signals associated with points in the tuple. The results from S , passed through a softmax function, will yield the weights for the attention mechanism. We use MLPs with a single hidden layer for S in the networks presented here.
  We first calculate the multivector geometric products pijk... of all combinations of input vectors indexed by i, j, k, and so on, up to the specified rank (for example, pairwise attention would produce a two-dimensional matrix of products pij). We then use V , M, J , and S—together with a function extracting the rotation-invariant attributes of a geometric product (which are the scalar component, trivector component, and the norms of the vector and bivector components, depending on how many input vectors are joined via the geometric product) into qijk...—as follows for a network producing permutation-covariant outputs yi for each input point:
  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 over all tuples of the point cloud simultaneously. While J , V , and M could in principle be used to change the working dimension of the network by stacking permutation-covariant layers, in this work we keep the dimension constant for the sake of easily adding residual connections.
  If rotation-covariant (rather than rotation-invariant) behavior is needed for the output of the network, the same attention structure can be used with slight modifications; here, we extract a vector from a learned linear combination of the input vectors r⃗i, r⃗j , r⃗k... and the geometric product pijk... (which consists of directly taking the vector component from products of odd numbers of input vectors, or multiplying a bivector by the unit trivector to produce a vector in the case of even numbers of input vectors). These vectors can be combined with learned vector weights αn, a scalar rescaling each vector (generated by a learned function R), and the attention mechanism to yield
  r⃗i ′ = ∑ jk... wijk...R(vijk...) (α0vector(pijk...) + α1r⃗i + α2r⃗j + α3r⃗k + ...) . (2)
  We use R to provide a value- and geometry-dependent rescaling of the output vector, making it easier to generate vectors that lie outside of the convex hull of the input point cloud. Because wijk..., vijk..., and αn are rotation-invariant and the geometric algebra is coordinate-free, the resulting output r⃗i′ is also rotation-covariant; more details of this are shown in Appendix A.
sections.heading: MODEL ARCHITECTURES
sections.text: We demonstrate the utility of our geometric algebra attention scheme by training deep networks to solve three problems appearing in physics, chemistry, and biology. For simplicity, all of the geometric algebra attention models presented here utilize pairwise attention with a working dimension of 32 units. Value functions V , score functions S , and rescaling functions R are simple multilayer perceptrons with a hidden width of 64 units, with layer normalization applied to the hidden layer of V . The network for crystal structure identification uses the mean function for merge functions
  M and join functions J , while the other two applications use learned linear projections. More thorough descriptions of the model architectures—including Python code under the MIT license—are detailed in Appendix B and the Supplementary Information.
  Crystal structure identification. We train networks using the 12 nearest neighboring particles j around a central particle i extracted from a structure. We use the neighbor bond coordinate differences r⃗ij and a symmetrized one-hot embedding vector [ti − tj , ti + tj ] to encode particle type information for each bond. After projecting these per-bond signals up to 32 dimensions, we perform two steps of permutation-covariant attention (using Equation 1) before a permutation-invariant reduction over the point cloud and final projection down to class logits.
  Molecular force regression. To attain translation invariance, we first convert the set of atomic coordinates r⃗i and one-hot type representations ti into the atom-centered pairwise coordinate difference matrices r⃗j − r⃗i and symmetrized one-hot type vectors [ti − tj , ti + tj ]. After projecting the value signals up to 32 dimensions, we perform six steps of permutation-covariant attention (using Equation 1) before a final permutation-invariant reduction and projection down to a per-atom energy. These energies are summed and we calculate the gradient of the per-molecule energy with respect to the input coordinates to ensure that a conservative force field is learned. We train models by matching the predicted forces to those found in the training data via a mean squared error loss.
  Backmapping of coarse-graining operators. We first apply two layers of permutation-covariant attention (using Equation 1) to the center-of-mass coordinates and identities of amino acids surrounding a central residue, to incorporate local geometry information into the coarse-grained bead representations. We then apply a translation layer to convert coarse-grained amino acid coordinates to atomic coordinates, which produces vectors according to Equation 2 by augmenting the tuple representation vijk... of Equation 1 with labels corresponding to the identity of the atom that should be produced, so that the value is calculated as
  vatom,ijk... = J (vatom,V(qijk...),M(vi, vj , vk, ...)). (3)
  Following this layer, two rotation-covariant layers (using Equation 2) are applied to the atomcentered, fine-grained atomic coordinates to further refine them.3
sections.heading: RESULTS
sections.text: Numerical results are reported using the mean and standard error of the mean over 5 independent samples. Baseline models are described in Appendix B.
sections.heading: CRYSTAL STRUCTURE IDENTIFICATION
sections.text: As shown in Figure 2, we select 8 prototypes of single- and two-component crystals from the AFLOW Encyclopedia of Crystallographic Prototypes (Mehl et al., 2017; Hicks et al., 2019). These structures are chosen to demonstrate that models can learn not only geometric information (cF4Cu and hP2-Mg are similar structures but with a different stacking of their close-packed layers; the clathrates cP46-Si and cF136-Si are also similar to each other, with a different arrangement of many common motifs), but also the information encoded within each point (cP2-CsCl and cF8-ZnS have identical geometry to cI2-W and cF8-C, respectively, differing only by the type assignments of their particles). For each structure, we rescale the unit cell so that the shortest nearest-neighbor distance over the structure is 1 before replicating the unit cell to consist of at least 2048 particles. We then create three samples of each structure by adding Gaussian noise with a standard deviation of 10−3, 5 · 10−2, and 0.1 separately to the particle coordinates, in order to emulate thermal noise. We find the nearest neighbors of each particle using the freud (Ramasubramani et al., 2020) python library and use these point clouds to train models to identify the source structure type for each particle.
  These networks rapidly learn to identify structures after a few epochs, with a final overall accuracy of 98.98% ± 0.08% after training for roughly 50 minutes on an NVIDIA Titan Xp GPU. We note that we do not expect this to be a difficult task in general: a baseline method using a rotation-invariant spherical harmonic featurization of varying-sized local neighborhoods of particles (Spellings & Glotzer, 2018) is also able to quickly reach 96.1% ± 0.1% accuracy, although we find encoding particle type information to be more natural in the graph neural network-like setting of the geometric algebra attention networks than when using spherical harmonic featurizations.
sections.heading: MOLECULE FORCE REGRESSION
sections.text: In a method similar to Batzner et al. (2021), we train models to predict the atomic forces calculated using ab initio molecular dynamics and density functional theory available in the MD17 dataset (Chmiela et al., 2017) for eight different molecules. Consistent with previous benchmarks on this dataset, we train networks using the mean squared error loss for each molecule using 1,000 snapshots of forces each as training, validation, and test data sets. Training a model on an individual molecule’s data takes between roughly 20 hours (ethanol and malonaldehyde, with nine atoms each) to 90 hours (naphthalene, with eighteen atoms) on an NVIDIA Titan Xp GPU. Test set losses, expressed as the mean absolute error over each force component for each sample, are presented in Table 1.
  Overall, these models perform favorably compared to NequIP and SchNet, while being slightly worse than GemNets. The better performance of GemNets could be due to the operation of GemNetQ on quadruplets of atoms (the geometric algebra attention networks learned here only operate on pairs of bonds from a central atom), the incorporation of energy in addition to force labels in the GemNet training scheme (in this work, we choose to only train on force data, rather than using a
  3For applications of this method to systems at nonzero temperature, we would expect to be better-served by using an architecture that generates output distributions instead of only point values, but we disregard this here for simplicity; in other words, here we are teaching models to memorize training data directly, rather than model a thermodynamic ensemble.
  weighted loss incorporating both quantities), or simply a more suitable architecture or hyperparameters for the GemNet models.
  Another useful attribute of attention-based models lies in analyzing learned attention maps (Rogers et al., 2020); to that end, we present sample attention maps of molecules in the training set in Figure 3. Notably, the attention weights tend to be more broadly distributed when calculating the forces on atoms participating in more complex bonding environments—such as aromatic rings—which we may expect due to physical effects, such as electron delocalization.
sections.heading: PROTEIN COARSE-GRAIN BACKMAPPING
sections.text: We use 19 protein structures, listed in Appendix C, which have high-resolution structural refinements (with resolution error less than or equal to 1.0Å) and were published on the Protein Data Bank (Berman et al., 2000) between 2015 and 2020 for training data. Because the variety of structural refinement workflows could lead to a large variance in structure fidelity compared to the learning capacity of the models, we use the training set error to characterize model performance instead of measuring generalization capability with a typical split of training, validation, and test data4. We compare geometric algebra attention models to baselines using geometrically-naïve transformer models in order to probe the impact of rotation equivariance on performance. This baseline model does not have rotation equivariance built in, and must learn to approximate it through data augmentation. We test a series of models using representations of varying width to compare the data efficiency of our geometric algebra attention scheme. In addition to comparing the overall accuracy of model predictions, we record the number of epochs required to reach 75% of the total improvement in training error. As shown in Figure 4, equivariant models are typically faster to train and more dataefficient, producing higher-quality predictions using fewer parameters than non-equivariant models. Furthermore, we are able to train larger equivariant models than geometrically-naïve models: basic transformers begin to show difficulty converging to a low training error above 64 neurons, while equivariant models do not saturate in performance until roughly 96 neurons in width.
sections.heading: DISCUSSION
sections.text: We find the architectures formulated here to be useful for a variety of tasks. Rather than being limited to operating on bond distances and angles as in SchNet (Schütt et al., 2017), PhysNet (Unke & Meuwly, 2019), and DimeNet (Klicpera et al., 2019), geometric algebra provides a systematic way to build functions with the desired rotation- and permutation-equivariance, with the flexibility to incorporate other types of geometric objects (such as the orientation quaternion commonly used for anisotropic particles in molecular dynamics methods (Kamberaj et al., 2005)) into the framework. The attention mechanism presented here provides a simple, powerful method to account for both geometric and node-level signals. The primitives of our geometric algebra attention scheme—distances,
  4In other applications of this method, we would expect data to be more homogeneous so that typical training, validation, and test data splitting methods can be used; here, we use PDB data as a proxy for simulation outputs using typical simulation methods.
  areas, angles, and volumes—and the calculated attention weights naturally lend themselves to interpretability, which we believe will prove useful in distilling insights from trained models.
  Although the architectures presented here work well for the problems we have selected, creating geometric products of vectors is only a subset of the valid combinations that could be generated. In these cases we have carefully chosen sums and differences of input vectors to respect symmetries we would like to impose on the system—such as using the pairwise distance of all input coordinates for the molecular force regression task to impose translation invariance—but it is possible that incorporating learned linear combinations of inputs or even generating intermediate multivector quantities could yield more powerful models.
  An obvious limitation to using higher-degree correlations lies in the computational complexity and memory scaling of generating tuples, which are both proportional to Nr for neighborhoods of N coordinates and tuples of length r. Polynomial scaling behavior can be ameliorated by restricting which combinations of input points are considered, essentially treating the attention weights of all other combinations as 0. These combinations could be randomly sampled from all valid indices ijk... or use more physically-relevant restrictions, such as utilizing the molecular connectivity graph for molecular force regression or edges derived from the Voronoi tessellation for other applications. If smoothness of model predictions is a concern—as may be the case for learning general N-body interaction potentials, for example—the architectures presented here could be augmented by incorporating weights that decay to 0 as bonds are broken in the Voronoi diagram graph (Mickel et al., 2013).
sections.heading: CONCLUSION
sections.text: In this work, we have presented a strategy for developing rotation- and permutation-equivariant neural network architectures by combining geometric algebra and attention mechanisms. These architectures operate directly on vector, scalar, and other geometric quantities of interest to produce outputs which respect desirable symmetries by construction. We believe that the mathematical simplicity and the insights derived from inspecting attention maps are particularly appealing aspects of the algorithms presented here. We hope that these architectures will help a wider range of scientific disciplines reap the benefits of geometric deep learning.
sections.heading: REPRODUCIBILITY STATEMENT
sections.text: In addition to the mathematical description of the architectures used in the Methods section, schematic diagrams and pseudocode to generate models for each task in the style of the tensorflow.keras API are listed in Appendix B. Descriptions of baseline methods for the structure identification and coarse-grain backmapping tasks are also available in Appendix B. The full code used to train and analyze the new model architectures presented here as well as baselines are available in the Supplementary Information.
sections.heading: A GEOMETRIC ALGEBRA
sections.text: The geometric algebra was developed in the 19th century and provides a consistent framework for dealing with scalars and other geometric quantities—such as vectors, areas, and volumes found in three-dimensional space—in arbitrary dimensions (Clifford, 1878). Here, we will describe the essential parts of geometric algebra as related to our proposed attention mechanism, and defer to other works, such as Macdonald (2017), for a more thorough description. The geometric algebra specifies a binary operator, the geometric product, that works on multivectors. Multivectors can be expressed as linear combinations of terms from a fixed basis set for a given space, such as R2 or R3; in threedimensional space, this yields scalars, vectors, bivectors (which specify signed areas within a plane and have 3 components), and trivectors (which specify signed volumes and have 1 component)—a total of 8 linearly independent terms for each multivector.5 When rotation invariance is desired, as in Equation 1, we can utilize the rotation-invariant components of a multivector: scalars, trivectors, the norms of vectors, and the norms of bivectors are rotation-invariant. Geometric algebra provides a general framework that can be used to build up expressive functions as linear combinations and geometric products of multivector inputs; in this work, we derive rotation-equivariant quantities from geometric products to suit the application and symmetry of our problems of interest.
  Terms for arbitrary geometric products can be calculated using basic mathematical attributes of the geometric product: for two orthogonal vectors e⃗1 and e⃗2, e⃗1e⃗1 = e⃗2e⃗2 = 1 and e⃗1e⃗2 = −e⃗2e⃗1. As an example, multiplying a vector A = αe⃗1+ γe⃗3 and vector-bivector combination B = δe⃗1+ ζe⃗2+ µe⃗1e⃗3 + νe⃗2e⃗3 proceeds as follows:
  AB = (αe⃗1 + γe⃗3)(δe⃗1 + ζe⃗2 + µe⃗1e⃗3 + νe⃗2e⃗3)
  = α(δe⃗1e⃗1 + ζe⃗1e⃗2 + µe⃗1e⃗1e⃗3 + νe⃗1e⃗2e⃗3) + γ(δe⃗3e⃗1 + ζe⃗3e⃗2 + µe⃗3e⃗1e⃗3 + νe⃗3e⃗2e⃗3)
  = α(δ + ζe⃗1e⃗2 + µe⃗3 + νe⃗1e⃗2e⃗3) + γ(−δe⃗1e⃗3 − ζe⃗2e⃗3 − µe⃗1e⃗3e⃗3 − νe⃗2e⃗3e⃗3) = α(δ + ζe⃗1e⃗2 + µe⃗3 + νe⃗1e⃗2e⃗3) + γ(−δe⃗1e⃗3 − ζe⃗2e⃗3 − µe⃗1 − νe⃗2) = αδ︸︷︷︸
  scalar αδ −γµe⃗1 − γνe⃗2 + αµe⃗3︸ ︷︷ ︸ vector (−γµ,−γν, αµ) +αζe⃗1e⃗2 − γδe⃗1e⃗3 − γζe⃗2e⃗3︸ ︷︷ ︸ bivector (αζ,−γδ,−γζ) +ανe⃗1e⃗2e⃗3︸ ︷︷ ︸ trivector αν
  (4)
  More generally, the types of elements produced by a geometric product of two multivectors in R3 with the given components are summarized in Table 2 below.
  From Table 2, we can see that successive products of vectors—such as pijk... in Equation 1— alternate between producing two types of multivectors: products of even numbers of vectors yield a scalar and bivector ((v + t)v = vv + tv = (s + b) + b → s + b), while products of odd numbers of vectors produce a vector and trivector ((s+ b)v = sv + bv = v + (v + t) → v + t). Generating rotation-invariant quantities from these products as described in the main text is the primary application of geometric algebra in this work, although in general the method could be used to incorporate different types of scalar, vector, bivector, and trivector quantities; for example, rotations could be input as quaternions, which are isomorphic to the scalar-and-bivector product of even numbers of vectors.
  5Multivectors form a vector space: the individual components of multivectors (bivectors, for example) can be directly summed elementwise, but multivector components of different types stay separate and are multiplied using the distributive property of geometric products when needed, producing output types according to Table 2. The geometric product has an identity of the scalar 1 and is associative; in other words, it forms a monoid over multivectors.
sections.heading: ROTATION EQUIVARIANCE
sections.text: Because the scalars used in Equation 2 are rotation-invariant and geometric algebra is coordinatefree—that is, we can perform a change of basis to replace the basis vectors e⃗1, e⃗2, e⃗3 with rotated versions Qe⃗1, Qe⃗2, Qe⃗3 to transform the point cloud using a rotation matrix Q and the vector component of the product pijk... will transform in the same way—the output of Equation 2 is also rotation-covariant: for a rotation matrix Q,
  Qr⃗i ′ = Q ∑ jk... wijk...R(vijk...) (α0vector(pijk...) + α1r⃗i + α2r⃗j + α3r⃗k + ...)
  = ∑ jk... wijk...R(vijk...) (α0Qvector(pijk...) + α1Qr⃗i + α2Qr⃗j + α3Qr⃗k + ...)
  = ∑ jk... wijk...R(vijk...) (α0vector((Qr⃗i)(Qr⃗j)(Qr⃗k)...) + α1Qr⃗i + α2Qr⃗j + α3Qr⃗k + ...) .
sections.heading: B MODEL DETAILS
sections.text: We list python-style pseudocode—using the semantics of the Keras API within Tensorflow (TensorFlow Developers, 2021)—for the models used in this work below. We also use the following geometric algebra attention layers, which are built with MLPs given by the listed make_scorefun (S) and make_valuefun (V) functions described in Equation 1:
  • VectorAttention implements the value-generating geometric algebra attention calculation described in Equation 1. It accepts a point cloud and set of values associated with each point, and produces new values.
  • Vector2VectorAttention implements the coordinate-generating (i.e., rotationcovariant) geometric algebra attention calculation described in Equation 2. It accepts a point cloud and set of values associated with each point, and produces new coordinates.
  • LabeledVectorAttention Implements the label-augmented point cloud translation layer described in Equation 3. It takes three arguments (a set of label values, a set of reference coordinates, and a set of reference values) and produces a set of new coordinates.
  Listing 1: Sample model-building code for rotation-invariant crystal structure identification. def make_scorefun():
  return Sequential([Dense(64, activation=’relu’), Dropout(0.5), Dense(1)])
  def make_valuefun(): return Sequential([Dense(64), LayerNormalization(),
  Activation(’relu’), Dropout(0.5), Dense(32)])
  # Inputs are given as N-length sets of coordinates # or values for each nearest-neighbor bond r_in = Input((None, 3)) v_in = Input((None, D_in)) # Project values to the desired working dimension last_v = Dense(D_working)(v_in)
  # Calculate attention in a series of blocks for _ in range(2):
  residual_in = last_v last_v = VectorAttention()(r_in, last_v) last_v = Dense(2*D_working, activation=’relu’)(last_v)
  last_v = Dense(D_working)(last_v) + residual_in
  # Permutation-invariant attention summation (NxD_working) last_v = VectorAttention(permutation_invariant=True)(r_ij, last_v) # Final MLP last_v = Dense(2*D_working, activation=’relu’)(last_v) logits = Dense(num_classes, activation=’softmax’)(last_v) model = Model([r_in, v_in], logits)
  Listing 2: Sample model-building code for conservative, rotation-equivariant neural networks for molecular force regression. def make_scorefun():
  return Sequential([Dense(64, activation=’swish’), Dense(1)])
  def make_valuefun(): return Sequential([Dense(64), LayerNormalization(),
sections.heading: Activation(’relu’), Dense(32)])
sections.text: # Inputs are given as N-length sets of coordinates # or values for each atom in a molecule r_in = Input((None, 3)) v_in = Input((None, D_in)) # Convert to atom-centered coordinates (NxNx3) and values (NxNx[2*D_in]) r_ij = PairwiseDifference()(r_in) v_ij = PairwiseSumDifference()(v_in) # Project values to the desired working dimension last_v = Dense(D_working)(v_ij)
  # Calculate attention in a series of blocks for _ in range(6):
  residual_in = last_v last_v = VectorAttention()(r_ij, last_v) last_v = Dense(2*D_working, activation=’swish’)(last_v) last_v = Dense(D_working)(last_v) + residual_in
  # Permutation-invariant attention summation (NxD_working) last_v = VectorAttention(permutation_invariant=True)(r_ij, last_v) # Final MLP and summation last_v = Dense(2*D_working, activation=’swish’)(last_v) atomic_energy = Dense(1, use_bias=False)(last_v) # The energy of the molecule is the sum of each atom’s energy molecule_energy = NeighborhoodSum()(atomic_energy) # The force is the negative gradient of the energy # with respect to input coordinates force_output = NegativeGradient()(molecule_energy, r_in) model = Model([r_in, v_in], force_output)
  Listing 3: Sample model-building code for rotation-covariant coarse-grain operator backmapping. def make_scorefun():
  return Sequential([Dense(64, activation=’relu’), Dense(1)])
  def make_valuefun(): return Sequential([Dense(64), LayerNormalization(),
  Activation(’relu’), Dense(32)])
  # Inputs are given as N-length sets of coordinates # or values for each coarse-grained neighbor, and a set # of atomic labels for the central coarse-grained bead r_in = Input((None, 3)) v_in = Input((None, D_in)) child_types_in = Input((None,)) # Project values to the desired working dimension child_labels = Embedding(num_child_types)(child_types_in) last_v = Dense(D_working)(v_in)
  # Calculate attention in a series of blocks for _ in range(2):
  residual_in = last_v last_v = VectorAttention()(r_in, last_v) last_v = Dense(2*D_working, activation=’relu’)(last_v) last_v = Dense(D_working)(last_v) + residual_in
  # Translate coarse-grained to fine-grained coordinates vecs = LabeledVectorAttention()(r_in, last_v, child_labels)
  # Convert to atom-centered representations for refinement delta_v = PairwiseDifferenceSum()(child_labels) delta_v = Dense(D_working)(delta_v)
  # Refine atomic coordinates for _ in range(2):
  residual_in = vecs vecs = PairwiseDifference()(vecs) vecs = Vector2VectorAttention(permutation_invariant=True)(
  vecs, delta_v) model = Model([r_in, v_in, child_types_in], vecs)
  We also present diagrams summarizing the geometric algebra attention model architectures used in this work in Figure 5.
sections.heading: BASELINE MODELS
sections.text: Crystal structure identification. We train MLPs on spherical harmonic descriptions of varyingsized local environments of particles as described in Spellings & Glotzer (2018). To deal with particle types using this method, we concatenate two sets of local environment spherical harmonics: one considering only bonds of like-typed particles (i.e. bonds to particles of type A if the central particle is of type A), and another set considering only bonds to differently-typed particles. We consider environments of up to the 12 nearest neighbors (disregarding particle types) for each set of spherical harmonics. The MLPs used consist of a batch normalization layer followed by two blocks of a dense layer 64 neurons in width using a ReLU activation function feeding into a dropout layer with a dropout rate of 0.5. The final result is projected down to class logits.
  Backmapping of coarse-graining operators. We train models using dot product attention: vertex representations are concatenated directly with their geometric coordinates and fed into a similar architecture to that shown in Figure 5, with geometric algebra attention layers replaced by standard dot-product attention layers. For the intermediate translation layer, query embeddings associated with each atom type are used. Because this baseline model is not rotation-covariant by construction, it must learn to generate rotation-covariant predictions; to impart this information to the model, we apply a random rotation to each pair of point clouds during training.
sections.heading: MODEL TRAINING
sections.text: Networks for the crystal structure identification and protein coarse-grain backmapping tasks are trained for up to 800 epochs using the adam optimizer (Kingma & Ba, 2015); the learning rate is decreased by a factor of 0.75 after the validation set loss does not decrease for 20 epochs, and training is ended early if the validation set loss does not decrease for 50 epochs. Models trained on the molecular force regression task use the longer training schedule presented in Batzner et al., which decreases the learning rate by a factor of 0.8 after 1000 epochs of no improvement in the validation loss and ends training after 2500 epochs of no improvement, or 50000 epochs in total.
sections.heading: C PDB ENTRIES USED FOR COARSE-GRAINING TASK
sections.text: • 3X2L (Nakamura et al., 2015) • 3X32 (Hirano et al., 2015) • 4TKJ (Matsuoka et al., 2015) • 3X0J (Furuike et al., 2016)
  • 5WQR (Ohno et al., 2017) • 6EQE (Austin et al., 2018) • 6ETK (Vergara et al., 2018) • 6FJN (Baker et al., 2019) • 6JGJ (Takaba et al., 2019) • 6RYG (Paterson et al., 2019) • 6IIP (Tian et al., 2019) • 6Q01 (Schellenberg et al., 2020) • 6XVM (Plaza-Garrido et al., 2020) • 6Y5S (Wu et al., 2020) • 6YP6 (Franke et al., 2020) • 7A5M (Barone et al., 2020) • 7K4T (Otten et al., 2020) • 6YK4 (Frydenvang et al., 2020) • 6SAY (Glöckner et al., 2020)
sections.heading: D MODEL EVALUATION THROUGHPUT MEASUREMENTS
sections.text: Typical times required to evaluate models on a single point cloud using an NVIDIA GeForce RTX 2060 GPU, amortized over a batch, are reported in Table 3 below.
title: SMALL POINT CLOUDS