Notes: Equivariant Features

By chance, I came across a presentation [1] by Professor Tess Smidt at MIT on equivariant neural networks; its really concise and clear introduction and I highly recommend it. I took some notes during the presentation, which I'm logging here, but practically the same content is available in her slides which can be found here. Her group also maintains a package called e3nn [2]. Below are my notes, which I'll probably revisit to update and turn into a more formal blog.

Notes

• Invariant neural networks (NN) means the outputs will be the same regardless of the transformations performed on the inputs, i.e., scalar representation.
• Equivariant NN means the transformations applied are global and get pass through to the outputs but consistently. In other words if we apply a transformation on the inputs the same transformation gets applied on the output. You can think of this as vector outputs that can be rotated.
• In higher order equivariant representations you can have very data efficient learning.
• Physical systems use coordinate systems (i.e. Euclidean geometry) to describe the physical system, we get to choose!
• Transformations between coordinate systems in Euclidean space utilize proper symmetries.
• I spoke a bit about this in my post on Graph NN potentials.
• Equivariant NN will understand the same motifs even if they are rotated or in different locations.
• How do you build a equivariant NN
• Irreducible representations to express transformations, i.e., symmetry groups.
• Spherical harmonics to featurize the geometry --> that is what are the basis functions that give use a geometry
• Tensor products to build a representation. (equivariant multiplication)

Example

Take simple, moving, two-particle system representation and indicate how it transforms.

• Define the coordinate system geometry geometry = [[x0, y0, z0], [x1, y1, z2]]
• Specify the features, which are present features = [[m0, v0x, v0y, v0z, a0x, a0y, a0z], [m1, v1x, v1y, v1z, a1x, a1y, a1z]]
• Construct the irreducible representation group, i.e. $\mathit{O}(3)$
• Using the e3nn package define the scalar and vector symmetries: scalar = e3nn.o3.Irrep('0e') , vector = e3nn.o3.Irrep('1o')
• Here the parity is encoded in the e or o , i.e., flip sign under inversion/mirroring.

Irreducible representations

• So we want to define the fundamental transformations that exist for objects/features.
• Higher-dimensions features, the data is represented as tensors in 3D space.
• The transformations of these data types is done using hydrogenic atomic spherical harmonics of differnt angular momentum.
• For example a scalar is transformed by L=0 or s-orbital, where as a vector is transformed through L=1 or p-orbital. Furthermore, a matrix would be transformed through L=2 or d-orbital.

Basis of descriptors: Spherical Harmonics

• Given a distance vector of one atom to another atom, $\vec{r}_{ij}$, find the spherical projections, i.e., convolution, we have $\sum_{lm} Y_{lm}(\hat{r}_{ij})Y_{lm}(\hat{x})$
• This gives a magnitude and a sign of the coefficients. Can think of it as an array of these. In addition, it is possible to map these coefficients on to a sphere where the radius from the center of the sphere is the magnitude.
• Now if you have $N$ atoms you can calculate these coefficients for every pairing with $i$ we can take linear sum given these harmonics form a orthogonal basis (i.e. vector space)
• This gives a nice fixed signal length and describes a distribution of vectors.
• The linear sum cancels out many components because many of them are just rotations of the same signal and therefore symmetry of the environment cancels out certain irreducible representation
• Fairly robust against noise in the signal.

Tensor Products

• So how do we create more complex features/descriptors: multiply them.
• This is similar to a kernel method but here we use tensor products, ex., $(x,y,z)\otimes(x,y,z) = (x^2,xy,xz,y^2,yz,z^2)$.
• Rotations acting on such features space are dense (reducible) but we can decompose the rotations into a irreducible representation by changing basis, where the space now is more block diagonal.

Creating equivariant signals

• Combining the spherical harmonic projections with tensor products provides invariant signals.
• Recall spherical harmonic projections are equivariant because rotation of the coordinate system changes the function.
• Power spectra, $x ⊗ x$, Bispectra, $x ⊗ x ⊗ x$, etc.
• Bicspectra will consist of scalars and pseudoscalars and smooth under distortion and break symmetry.

Applications: Clustering local environments of materials

• Can find the different symmetries in local environments by determining the Bispectra terms, i.e., can use higher order representation to enable expressive features of local atomic environments.
• Check out https://e3nn.org.

References

[1] Smidt, T., Harnessing the properties of equivariant neural networks to understand and design materials, From Molecules to Materials: ICLR 2023 Workshop on Machine learning for materials (ML4Materials), May, (2023). https://iclr.cc/virtual/2023/14141.

[2] Geiger, M., and T. Smidt. E3nn: Euclidean Neural Networks. arXiv, (2022). https://doi.org/10.48550/arXiv.2207.09453.

Reuse and Attribution