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.

summaries: 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.

summaries: 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.

summaries: 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.

authors:
Yaxin An
Sanket A. Deshmukh

title: Machine learning approach for accurate backmapping of coarsegrained models to all-atom models

venue: Chemical Communications,

authors:
Brandon Anderson
Truong Son Hy
Risi Kondor

title: Cormorant: Covariant molecular neural networks

venue: In Advances in Neural Information Processing Systems,

authors:
Alexander T. Baker
Alexander Greenshields-Watson
Lynda Coughlan
James A. Davies
Hanni Uusi-Kerttula
David K. Cole
Pierre J. Rizkallah
Alan L. Parker

title: Diversity within the adenovirus fiber knob hypervariable loops influences primary receptor interactions

venue: Nature Communications,

authors:
Simon Batzner
Tess E. Smidt
Lixin Sun
Jonathan P. Mailoa
Mordechai Kornbluth
Nicola Molinari
Boris

title: Kozinsky. SE(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials

authors:
E. Bayro-Corrochano
S. Buchholz
G. Sommer

title: A new self-organizing neural network using geometric algebra

venue: In Proceedings of 13th International Conference on Pattern Recognition,

authors: Jörg Behler. Four generations of high-dimensional neural network potentials. Chemical Reviews

title: ISSN 0009-2665

authors:
Helen M. Berman
John Westbrook
Zukang Feng
Gary Gilliland
T.N. Bhat
Helge Weissig
Ilya N. Shindyalov
Philip E. Bourne

title: The Protein Data Bank

venue: Nucleic Acids Research,

authors:
Emanuele Boattini
Marjolein Dijkstra
Laura Filion

title: Unsupervised learning for local structure detection in colloidal systems

venue: The Journal of Chemical Physics,

authors:
Michael M. Bronstein
Joan Bruna
Taco Cohen
Petar Veličković

title: Geometric deep learning: Grids, groups, graphs, geodesics, and gauges

authors:
Sven Buchholz
Gerald Sommer

title: On Clifford neurons and Clifford multi-layer perceptrons

venue: Neural Networks: The Official Journal of the International Neural Network Society,

authors:
Stefan Chmiela
Alexandre Tkatchenko
Huziel E. Sauceda
Igor Poltavsky
Kristof T. Schütt
Klaus-Robert Müller

title: Machine learning of accurate energy-conserving molecular force fields

venue: Science Advances,

title: Characterising soft matter using machine learning

venue: Soft Matter,

authors:
Julia Dshemuchadse
Pablo F. Damasceno
Carolyn L. Phillips
Michael Engel
Sharon C. Glotzer

title: Moving beyond the constraints of chemistry via crystal structure discovery with isotropic multiwell pair potentials

venue: Proceedings of the National Academy of Sciences,

authors:
Marc Finzi
Max Welling
Andrew Gordon Wilson

title: A practical method for constructing equivariant multilayer perceptrons for arbitrary matrix groups

authors:
Barbara Franke
Marta Veses-Garcia
Kay Diederichs
Heather Allison
Daniel J. Rigden
Olga Mayans

title: Structural annotation of the conserved carbohydrate esterase vb_24B_21 from Shiga toxin-encoding bacteriophage Φ24B

venue: Journal of Structural Biology,

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

title: Ionotropic glutamate receptor GluA2 in complex with bicyclic pyrimidinedione-based compounds: When small compound modifications have distinct effects on binding interactions

venue: ACS Chemical Neuroscience,

authors:
Fabian Fuchs
Daniel Worrall
Volker Fischer
Max Welling

title: SE(3)-transformers: 3D rototranslation equivariant attention networks

venue: Advances in Neural Information Processing Systems,

authors:
Yoshihiko Furuike
Yuka Akita
Ikuko Miyahara
Nobuo Kamiya

title: ADP-ribose pyrophosphatase reaction in crystalline state conducted by consecutive binding of two manganese(ii) ions as cofactors

authors:
Steffen Glöckner
Andreas Heine
Gerhard Klebe

title: A proof-of-concept fragment screening of a hit-validated 96-compounds library against human carbonic anhydrase

venue: ii. Biomolecules,

authors:
Yulan Guo
Hanyun Wang
Qingyong Hu
Hao Liu
Li Liu
Mohammed Bennamoun

title: Deep learning for 3D point clouds: A survey

venue: IEEE transactions on pattern analysis and machine intelligence,

authors:
Y. Hirano
S. Kimura
T. Tamada

title: High-resolution crystal structures of the solubilized domain of porcine cytochrome b5

venue: Acta Crystallographica Section D: Biological Crystallography,

authors:
H. Kamberaj
R.J. Low
M.P. Neal

title: Time reversible and symplectic integrators for molecular dynamics simulations of rigid molecules

authors:
Diederik P. Kingma
Jimmy Ba

title: Adam: A method for stochastic optimization

venue: In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations,

authors:
Johannes Klicpera
Janek Groß
Stephan Günnemann

title: Directional message passing for molecular graphs

venue: In International Conference on Learning Representations,

authors:
Johannes Klicpera
Florian Becker
Stephan Günnemann

title: GemNet: Universal directional graph neural networks for molecules

title: N-body networks: A covariant hierarchical neural network architecture for learning atomic potentials

authors:
Wei Li
Craig Burkhart
Patrycja Polińska
Vagelis Harmandaris
Manolis Doxastakis

title: Backmapping coarse-grained macromolecules: An efficient and versatile machine learning approach

title: A survey of geometric algebra and geometric calculus

venue: Advances in Applied Clifford Algebras,

authors:
Siewert J. Marrink
D. Peter Tieleman

title: Perspective on the Martini model

venue: Chemical Society Reviews,

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

title: Water-mediated recognition of simple alkyl chains by heart-type fatty-acid-binding protein

venue: Angewandte Chemie International Edition,

authors:
Michael J. Mehl
David Hicks
Cormac Toher
Ohad Levy
Robert M. Hanson
Gus Hart
Stefano Curtarolo

title: The AFLOW library of crystallographic prototypes: Part 1

venue: Computational Materials Science,

authors:
Pavlo Melnyk
M. Felsberg
Mårten Wadenbäck

title: Embed me if you can: A geometric perceptron

venue: ArXiv,

authors:
Walter Mickel
Sebastian C. Kapfer
Gerd E. Schröder-Turk
Klaus Mecke

title: Shortcomings of the bond orientational order parameters for the analysis of disordered particulate matter

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

title: Newton’s cradle" proton relay with amide–imidic acid tautomerization in inverting cellulase visualized by neutron crystallography

authors:
Hiraku Ohno
Kazuki Takeda
Satomi Niwa
Tomotaka Tsujinaka
Yuya Hanazono
Yu Hirano
Kunio Miki

title: Crystallographic characterization of the high-potential iron-sulfur protein in the oxidized state at 0.8 Å resolution

venue: PLOS ONE,

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

title: How directed evolution reshapes the energy landscape in an enzyme to boost catalysis

venue: doi: 10.1126/science

authors:
Cheol Woo Park
Mordechai Kornbluth
Jonathan Vandermause
Chris Wolverton
Boris Kozinsky
Jonathan P. Mailoa

title: Accurate and scalable graph neural network force field and molecular dynamics with direct force architecture

venue: npj Computational Materials,

authors:
Janet M. Paterson
Amy J. Shaw
Ian Burns
Alister W. Dodds
Alpana Prasad
Ken B. Reid
Trevor J. Greenhough
Annette K. Shrive

title: Atomic-resolution crystal structures of the immune protein conglutinin from cow reveal specific interactions of its binding site with N-acetylglucosamine

venue: Journal of Biological Chemistry,

authors:
J.K. Pearson
D.L. Bisset

title: Neural networks in the Clifford domain

venue: IEEE International Conference on Neural Networks (ICNN’94),

authors:
Justin Pearson
David Bisset

title: Back propagation in a Clifford algebra

venue: Proceedings of the International Conference of Artificial Neural Networks (ICANN),

authors:
Marina Plaza-Garrido
Ma Carmen Salinas-García
José C. Martínez
Ana Cámara-Artigas

title: The effect of an engineered ATCUN motif on the structure and biophysical properties of the SH3 domain of c-Src tyrosine kinase

venue: JBIC Journal of Biological Inorganic Chemistry,

authors:
Charles R. Qi
Hao Su
Mo Kaichun
Leonidas J. Guibas

title: PointNet: Deep learning on point sets for 3D classification and segmentation

venue: IEEE Conference on Computer Vision and Pattern Recognition (CVPR),

authors:
Vyas Ramasubramani
Bradley D. Dice
Eric S. Harper
Matthew P. Spellings
Joshua A. Anderson
Sharon C. Glotzer

title: Freud: A software suite for high throughput analysis of particle simulation data

venue: Computer Physics Communications,

authors:
Anna Rogers
Olga Kovaleva
Anna Rumshisky

title: A Primer in BERTology

venue: What We Know About How BERT Works. Transactions of the Association for Computational Linguistics,

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

title: Ubiquitin stimulated reversal of topoisomerase 2 DNA-protein crosslinks by TDP2

authors:
K.T. Schütt
P.-J. Kindermans
H.E. Sauceda
S. Chmiela
A. Tkatchenko
K.-R. Müller

title: SchNet: A continuous-filter convolutional neural network for modeling quantum interactions

venue: In Proceedings of the 31st International Conference on Neural Information Processing Systems,

authors:
Matthew Spellings
Sharon C Glotzer

title: Machine learning for crystal identification and discovery

venue: AIChE Journal,

authors:
Paul J. Steinhardt
David R. Nelson
Marco Ronchetti

title: Bond-orientational order in liquids and glasses

venue: Physical Review B,

authors:
K. Takaba
Y. Tai
H. Eki
H.-A. Dao
Y. Hanazono
K. Hasegawa
K. Miki
K. Takeda

title: Subatomic resolution X-ray structures of green fluorescent protein

venue: IUCrJ,

authors:
Hugues Thomas
Charles R. Qi
Jean-Emmanuel Deschaud
Beatriz Marcotegui
François Goulette
Leonidas J. Guibas

title: KPConv: Flexible and deformable convolution for point clouds

venue: Proceedings of the IEEE International Conference on Computer Vision,

authors:
Nathaniel Thomas
Tess Smidt
Steven Kearnes
Lusann Yang
Li Li
Kai Kohlhoff
Patrick Riley

title: Tensor field networks: Rotation- and translation-equivariant neural networks for 3d point clouds

authors:
Wei Tian
Peiqiang Yan
Ning Xu
Arghya Chakravorty
Robert Liefke
Qiaoran Xi
Zhanxin Wang

title: The HRP3 PWWP domain recognizes the minor groove of double-stranded DNA and recruits HRP3 to chromatin

authors:
Oliver T. Unke
Markus Meuwly

title: PhysNet: A neural network for predicting energies, forces, dipole moments, and partial charges

venue: Journal of Chemical Theory and Computation,

authors: Klaus-Robert Müller. Machine learning force fields. Chemical Reviews

title: ISSN 0009-2665

authors:
Ashish Vaswani
Noam Shazeer
Niki Parmar
Jakob Uszkoreit
Llion Jones
Aidan N Gomez
Łukasz Kaiser
Illia Polosukhin

title: Attention is all you need

venue: Advances in Neural Information Processing Systems

authors:
Alessandro Vergara
Marco Caterino
Antonello Merlino

title: Raman-markers of X-ray radiation damage of proteins

venue: International Journal of Biological Macromolecules,

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

title: Conformational heterogeneity of Savinase from NMR, HDX-MS and X-ray diffraction analysis

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.

heading: BACKGROUND

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.

heading: RELATED WORK

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.

heading: DOMAIN-SPECIFIC APPLICATIONS

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.

heading: METHODS

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.

heading: ATTENTION FROM GEOMETRIC PRODUCTS

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.

heading: MODEL ARCHITECTURES

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

heading: RESULTS

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.

heading: CRYSTAL STRUCTURE IDENTIFICATION

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.

heading: MOLECULE FORCE REGRESSION

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.

heading: PROTEIN COARSE-GRAIN BACKMAPPING

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.

heading: DISCUSSION

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).

heading: CONCLUSION

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.

heading: REPRODUCIBILITY STATEMENT

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.

heading: A GEOMETRIC ALGEBRA

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.

heading: ROTATION EQUIVARIANCE

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 + ...) .

heading: B MODEL DETAILS

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(),

heading: Activation(’relu’), Dense(32)])

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.

heading: BASELINE MODELS

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.

heading: MODEL TRAINING

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.

heading: C PDB ENTRIES USED FOR COARSE-GRAINING TASK

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)

heading: D MODEL EVALUATION THROUGHPUT MEASUREMENTS

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