Complete and Efficient Covariants for 3D Point Configurations with Application to Learning Molecular Quantum Properties
Hartmut Maennel Oliver T. Unke Klaus-Robert Müller


Reply to Comment on “A slightly oblate dark matter halo revealed by a retrograde precessing Galactic disk warp"
    Haibo Yuan
    September 9, 2024
===============================================================================================================

§ ABSTRACT 
When modeling physical properties of molecules with machine learning, it is desirable to 
incorporate SO(3)-covariance. While such models based on low body order features are not complete, we formulate and prove general completeness properties for higher order methods, and show that 6k-5 of these features are enough for up to k atoms. We also find that the Clebsch–Gordan operations commonly used in these methods can be replaced by matrix multiplications without sacrificing completeness, lowering the scaling from O(l^6) to O(l^3) in the degree of the features. We apply this to quantum chemistry, but the proposed methods are generally applicable for problems involving 3D point configurations.



§ INTRODUCTION

Atomistic simulations have proven indispensable for advancing chemistry and materials science, providing insights into the behavior of matter at the atomic level.  In the past, these simulations have been computationally demanding, but the advent of Density Functional Theory (DFT) <cit.> significantly enhanced the accessibility of atomistic simulations, and recent breakthroughs in machine learning (ML) have further accelerated progress <cit.>. ML methods trained on ab initio data now enable the fast and accurate prediction of quantum properties orders of magnitude faster than traditional calculations <cit.>.
A cornerstone of these methods, whether utilizing kernel-based approaches <cit.> or deep learning <cit.>, lies in the effective representation of molecules <cit.> or materials <cit.> through carefully chosen features or descriptors. Early examples include the Coulomb Matrix representation <cit.> and SOAP <cit.>, while recent advancements extend this principle beyond rotationally invariant representations with the design of equivariant model architectures <cit.>.

However, <cit.> pointed out that commonly available descriptors are not able to uniquely identify some molecular structures <cit.>. This can lead to ambiguities (two distinct structures may be mapped to the same descriptor) that hamper the performance of ML models. Effectively, a lack of uniqueness is similar to introducing a high level of noise into the learning process and may hinder generalization. 
A second important shortcoming of some modern ML architectures was discussed by <cit.> and only becomes visible when running molecular dynamics (MD) simulations <cit.>. It was observed that ML models with excellent prediction accuracy for energies and forces can nevertheless show unphysical instabilities (e.g. spurious bond dissociation) when simulating longer MD trajectories — limiting their usefulness in practice.  Equivariant architectures, however, as broad anecdotal evidence and some theoretical analyses have shown <cit.>,  were found to enable stable MD simulations over long timescales  <cit.>. 



Both aspects lead to the interesting theoretical question of how to construct a provably unique invariant, or more generally, a “complete” (to be defined below) equivariant and computationally efficient representation of descriptors for atomistic simulations.
We will study this challenge  both by theoretical means and by performing empirical atomistic simulations.
Let us assume that  the origin of our coordinate system was fixed meaningfully and we are looking for unique 
descriptors of point sets that are equivariant under rotations in SO(3).


To get invariant features, we can use a rotationally invariant function of n points (e.g. distance from the origin for n=1, or angles between two points for n=2), and then sum over all n–tuples of points in the configuration. 
Such descriptors are called “(n+1)–body functions”. It was recently shown that
descriptors based on 2- and 3-body information (distances and angles) are unable to distinguish some non-equivalent environments <cit.>. Even 4-body information (dihedrals) is not sufficient in all cases
(see Fig. <ref>B) and it is necessary to include 
higher m-body information for some structures. Other methods that construct descriptors implicitly, e.g. by message-passing <cit.>, suffer from similar problems <cit.>.



§ RESULTS AND DISCUSSION

Let us start defining an appropriate mathematical language. In applications to
chemistry, the points in the point set can belong to different atom types/elements which have to be treated
differently. We assume there is a fixed finite set  of “colors” (the atom types/elements), and each point in the
point set S is assigned a color in , i.e. S = ⋃_γ∈ S_γ.

We propose to take as potential features all polynomial point set descriptors (PPSDs), i.e. all scalar expressions that can be written down for colored point sets, using the coordinates of points, constants from , addition, multiplication, and summations over all points of a given color, such that these expressions can be evaluated for any point set independent of the number of points (See Appendix  for  formal definitions).
In practice, a variant of PPSDs is more useful, using polynomials only for the angular part (i.e. as a function on the sphere ^2) and some other function space for the radial part. With the assumptions that
these radial functions are analytic and allow approximation of continuous functions in the radius, we can 
(with some extra effort) prove almost the same theorems, see Appendix  for the definitions, and later sections for details and proofs.

General theorems:
We now describe informally a series of mathematical theorems about PPSDs that we prove in this work,
see respective Appendices 
for the precise formulations and proofs.
We first observe (see Appendix ) that the computation of any scalar PPSD
can be arranged into two steps:

  
  * Evaluate expressions involving only one summation sign: ∑_ r∈ S_γ P( r) for some color γ and polynomial P:^3 acting on point coordinates r.
  We call them fundamental features.
  
  * Evaluate polynomials in fundamental features.

This separability into two steps allows any PPSD to be evaluated in time O(n) where n is the number of points (here atoms), which is a major advantage over e.g. descriptors based on rational functions, for which this is generally not possible.

We call a PPSD that can be written such that all polynomials in fundamental features have degree d “homogeneous 
of order d” [Note that this is only the degree of the polynomial in fundamental features (step 2), it does not take into account the degrees of the polynomials used to construct the fundamental
features themselves. When we multiply out and move all summations
to the left (see Appendix ) this order corresponds to the depth of the summations, since each fundamental feature comes with one summation sign]. 
The order of such a PPSD is unique, for a proof and a refinement of this notion see Appendix
. PPSDs of order d are also said to be of “body order d+1” (this 
convention includes one atom at the origin of the coordinate system in the count).

In this language, there are infinitely many independent SO(3)–invariant PPSDs of body order 3, but the examples in <cit.> show that there are inequivalent configurations that cannot be distinguished by invariant functions of body orders ≤ 4 (see Fig. <ref>B).
Our Topological Completeness Theorem (Theorem  in Appendix ) says that this problem vanishes
when we allow arbitrary body orders, even when we restrict the functions to be polynomial invariants: Any two SO(3)–inequivalent configurations can be distinguished by SO(3)–invariant PPSDs, i.e. taking the values of all polynomial SO(3)–invariant functions gives a unique descriptor.
In general for covariant functions the values of PPSDs change when we rotate a configuration, so this uniqueness property has to be expressed differently: We prove that there are enough SO(3)–covariant PPSDs to approximate any continuous SO(3)–covariant function of colored point sets.
Without bound on the number of points in the configurations it is of course necessary to use infinitely many
independent invariant functions to distinguish all SO(3)–inequivalent configurations, as these form an infinite dimensional space. However, we can ask how many features are necessary to uniquely identify configurations of up to k points. Our Finiteness Theorem (Theorem  in Appendix ) gives a linear upper bound of 6k-5, with some guarantees for the distance of non–equivalent configurations. Its proof is based on real algebraic geometry and subanalytic geometry.

Practical construction:
We will  now  show how to produce unique features in such a way that we
never leave the space of covariant features: 

Let l be the irreducible (2l+1)–dimensional (real) representation
of SO(3), and Y_l:^3l be a SO(3)–covariant polynomial (which is unique on the sphere ^2 up to a scalar constant factor, see Appendix ).
These Y_l are given by (real valued) spherical harmonics of degree l.
We now proceed again in two stages:

  
  * Evaluate spherical harmonics [To be precise, we multiply the spherical harmonics with radial basis functions. We use exponent 2k here to have only
    polynomial functions. In practice, we would rather use different, decaying functions,
    this is treated as “case ii” in general as one of the variations, see Appendix 
    .]
    :

    ∑_ r∈ S_γ | r|^2k Y_l( r)
    for all colors γ and k=0,1,2,... and l=0,1,2,... 

    These are covariant fundamental features (i.e. of order 1) with values in l.
  
  * Iterate for d=1,2,...: Compute Clebsch–Gordan operations [These project
    the tensor product of two representations to an irreducible component, see e.g. <cit.>.]
    ^l_1⊗^l_2^l_3 for |l_1-l_2| ≤ l_3 ≤ l_1+l_2, where the feature in ^l_1 is a fundamental feature, and the feature in ^l_2 is of order d. This gives covariant features of order d+1.

Clearly, this construction appears to be somewhat special, so we may ask whether it actually gives “enough” invariants (i.e. achieves completeness).
This is in fact true in a very strong sense: Our Algebraic Completeness Theorem (Theorem ) says that all  invariant / covariant
PPSDs can be obtained as a linear combination of them; in a sense this is just the isotypical decomposition of 
the space of all PPSDs (see Appendix ).
While the above strategy to construct invariant / covariant functions
has been used e.g. in <cit.>, our novel completeness theorems 
show that this avoids the potential incompleteness problem pointed out in
<cit.>.  In fact, by our algebraic completeness theorem we get all polynomial covariant functions,
and by the topological completeness theorem those are sufficient to approximate any continuous covariant function. We also get an algebraic completeness theorem for features constructed from
tensor products and contractions as in <cit.>, see 
Theorem ; this is based on classical invariant theory.

Computational bottleneck:
We now turn to a particular efficient variant of our construction. 
Since invariant PPSDs of order <4 are 
not sufficient for distinguishing all SO(3)–equivalence classes, we need to construct covariants of higher body orders, i.e. in the above procedure we need to use the Clebsch–Gordan 
products. Note that their computational cost
is independent of the number of points, and is linear in the number of products, but 
scales as O(l^6) when we take the tensor product of two representations of the form
0⊕...⊕l. But with unrestricted number of points in our configuration, we cannot bound the l, even if we are just considering 
configurations on ^2⊂^3: Using Y_l:^2l only for l=0,1,...,L-1 yields a ||· L^2–dimensional vector space of fundamental features ∑_ r∈ S_γ Y_l( r)
(and all PPSDs are polynomials in the fundamental features). So this could only describe a configuration space of 
a dimension ≤ ||· L^2, not the ∞–dimensional space of configurations on ^2
with an unbounded number of points.
Consequently, the bottleneck for a larger number of points (necessitating using larger l for constructing the fundamental features) can be determined as the O(l^6) Clebsch–Gordan operation.
We will now propose how to construct a subset of local descriptors that alternatively to 
 Clebsch–Gordan operations relies only on 
matrix-matrix multiplication. This procedure scales as only O(l^3) for 
bilinear operations on two representations of the form
0⊕...⊕l. 
A similar speedup was published recently in <cit.>, replacing the Clebsch–Gordan operation 
by the multiplication of functions. However, since multiplying functions (instead of matrices) is commutative, this does not reproduce the anti–commutative part of the Clebsch–Gordan operations. Therefore 
the construction  in <cit.> does not have the full expressivity desired and would not satisfy our Algebraic Completeness Theorem or Theorem <ref> below. In particular, since commutative products cannot produce
pseudo–tensors, its invariants could not distinguish configurations from their mirror images.

Matrix Construction: 
Our key idea for removing the computational bottleneck is to apply the Clebsch–Gordan relation

    a⊗b≃|a-b|⊕|a-b|+1⊕... ⊕a+b

“backwards” to efficiently encode a collection of features in
^(|a-b|)⊕ ... ⊕^(a+b) as a (2a+1)× (2b+1) matrix in

    Lin(a,b)
      ≃a^* ⊗b≃a⊗b

(see Appendix ) and then the matrix multiplication is a covariant map of representations

    Lin(a,b) × Lin(b,c)  Lin(a,c).

With Schur's Lemma one can show that it can be expressed as a linear combination of Clebsch–Gordan operations, so unless some coefficients are zero, we can expect this operation to be as useful as the Clebsch–Gordan operations for constructing covariant features 
of higher body order. This is indeed the case and to formulate the
corresponding theorem, we define the involved features:

Let ι_a,b,l: l Mat_2b+1,2a+1 be the embedding given by (<ref>)


and define the “matrix moments”

    M_a,b,l() := ι_a,b,l∑_ r ∈ S_ Y_l( r)

which are (2b+1)×(2a+1) matrices 
(see Appendix  for some examples for explicit formulas).
Then the result of the multiplication

    M_a_m-1, a_m, l_m(γ_m)· ... 
       · M_a_1, a_2, l_2(γ_2) · M_0,a_1, l_1(γ_1)

23mm

    < g r a p h i c s >

























































with l_1=a_1 and 
|a_1-a_2| ≤ l_2 ≤ a_1+a_2,...,|a_m-1-a_m| ≤ l_m ≤ a_m-1+a_m
are covariant a_m× 1 matrices, i.e. vectors in a_m, given by polynomials of degree l_1+...+l_m, and computing them takes O(m· a^3) steps for an upper bound a ≥ a_i.
 
.

  Any SO(3)–covariant feature with values in a l can be written as a linear combination
  of the SO(3)–covariants (<ref>) with a_m=l.

  For O(3)–covariants it is enough to use those features given by (<ref>) 
  with the appropriate parity of l_1+...+l_m.

For the proof see Appendix .
Learning a linear combination:
While computing one given invariant of the form (<ref>) would not be more efficient than 
with Clebsch–Gordan operations (as it would waste whole matrices for encoding only one feature), for applications in
Machine Learning we always compute with linear combinations of features (with learnable coefficients), and both the
Clebsch–Gordan operation and Matrix Multiplication define maps

    (0⊕...⊕l) ⊗(0⊕...⊕l)   0⊕...⊕l

which are used to build up different linear combinations of covariants of higher body
order. In the Clebsch–Gordan case we also can add to the learnable coefficients of the input features
further learnable parameters that give different weights to the individual parts 
l_1⊗l_2l
that contribute to the same l in the output, whereas in the Matrix Multiplication case these
mixture coefficients are fixed (but depend on the shape of the matrices involved).
However, Theorem <ref> shows that using different shapes of matrices is already sufficient to generate all possible covariants, so both methods can
in principle learn the same functions.
Matrix of Matrices construction for efficiency: For practical applications it is important how to organize the matrix multiplications
efficiently. In particular when using GPUs / TPUs with hardware support
for matrix multiplication, it is much more favorable to compute with a few large matrices
than with many small matrices.
Therefore we will use linear combinations of ι_a,b,l for l=|a-b|,...,a+b
to fill a (2b+1) × (2a+1) matrix, and pack r × r small matrices for a,b in 
{l_1, l_2,...,l_r} into a large square matrix of side length (2l_1 +1) + ... + (2l_r+1).

Then k-1 such matrices are multiplied to get a matrix built out of covariants
of body order k.

    < g r a p h i c s >





















  







































This matrix can then be applied to n_1 column vectors
from l_1⊕l_2⊕ ... ⊕l_r 
to get covariant vectors of body order k+1.
8mm

    < g r a p h i c s >







































If the end result should be scalars, we can take scalar
products of the n_1 column vectors in l with n_2 new
covariants in l to obtain
r· n_1· n_2 invariants of body order k+2;
also the traces of the 
square submatrices of the matrix product give invariants of body order k (marked in
color in the above example diagram).
Full architectures:
The proposed matrix products approach can be readily used to replace Clebsch–Gordan operations across all possible learning architectures giving rise to significant efficiency gains. 


As a proof of concept, in the following experiments we will focus on the simplest such architecture which only
computes a linear combination of many such invariants, see Appendix  for code and more details (e.g. in 
practice we may want to shift the matrices by the identity to obtain a similar effect to skip connections in ResNets.)


Extensions of this minimal architecture could use a deep neural network instead of a linear combination
of invariants, or can use nonlinear activation functions to modify the matrices obtained in intermediate
steps. In architectures using several layers of Clebsch–Gordan operations, such  activation
functions are restricted to functions of the scalar channel, since “you cannot apply a
transcendental function to a vector”. Maybe surprisingly, in our matrix formulation this actually
becomes possible: Applying any analytic function to our (2l+1)×(2l+1)
matrices (not element wise, but e.g. implemented as Taylor series for matrices) is also a covariant operation! Notably, Matrix exponentiation has been suggested as an efficient 
and useful operation in Neural Networks in <cit.>.



§ EXPERIMENTAL RESULTS


Our methods yield complete representations and can thus indeed distinguish (molecular) configurations that require higher order features (see <cit.>). This is demonstrated experimentally in Fig. <ref>B/D.
In Fig. <ref>C we used the library E3x (<cit.>), which allows switching between full tensor layers using the
Clebsch–Gordan operation and “Fused Tensor Layers” for which we implemented matrix multiplication instead of the Clebsch–Gordan operations. The plot shows the inference 
run time measured on CPUs for computing a function defined by two Tensor layers, depending
on the the setting of “max degree” and whether full or fused layers were used.
In another synthetic experiment, we learn an invariant
polynomial of degree 10 with
Clebsch–Gordan operations and with our matrix multiplication framework, and plot the training 
curves averaged over 10 data sets.

 
    < g r a p h i c s >

-2mm

    < g r a p h i c s >

On CPUs, almost all the time is spent in the Clebsch–Gordan operations, and replacing them
by the matrix multiplication method makes the training faster by a factor over 100.
When using GPUs, the speedup is not quite as dramatic, but still a factor of 8.4 on V100. (Details in Appendix )
As a first demonstration of our framework for atomistic simulations we show that with the simple architecture that linearly combines the resulting polynomial invariants, we 
can learn forces with local features alone that, interestingly, can match the accuracy of other more complex  methods which use
several message passing / self attention steps with nonlinear networks (So3krates, <cit.>) or global kernel methods (sGDML, <cit.>). Specifically, our experiments (see Table <ref>) show that accuracies align, notably, independent of the molecule sizes, see Appendix  for details. 
Since our model is just a linear combination of features of known body order and L,
in future studies, one could use such models to investigate body order expansions or study the influence of larger Ls in detail.
Appropriately representing chemical structure and atomic environments in molecules and materials is an important prerequisite for accurate machine learning models in chemistry.
Ideal descriptors are unique, computationally efficient, and covariant.
In this work we have established an algebraic framework that enables a practical construction of provably complete system(s) of features with these desired properties that holds for any 3D point configurations. Apart from the abstract theoretical contribution of this work, we show that our construction can be readily implemented as matrix-matrix multiplication – reducing computational complexity from O(l^6) to O(l^3) compared to Clebsch-Gordan operations. This yields large efficiency gains while maintaining the performance level of standard machine learning models for atomistic simulation.
In summary, our theoretically well founded unique, covariant, and efficient

descriptors provide a versatile basis for future atomistic modeling and potentially other applications of machine learning on point configurations. 




§ ACKNOWLEDGEMENT

The authors acknowledge valuable discussions with Romuald Elie and Zhengdao Chen.

Correspondence to HM (hartmutm@google.com) and KRM (klausrobert@google.com).



