Group Theory for Machine Learning
Youth in high dimensions 2022
Mario Geiger
Postdoc at MIT with Prof. Smidt
This Talk is about Equivariant Neural Networks
input
output
Illustration of a neural network equivariant to rotations in 3D
Plan
What affects data efficiency in equivariant neural networks?
Group
\(a, b, c, e \in G\)
- \((ab)c = a(bc)\)
- \(ea=ae=a\)
- \((a^{-1})a=a (a^{-1})=e\)
Representations of Rotations
The Vectors
\(\begin{bmatrix}x_1\\x_2\\x_3 \end{bmatrix}\longrightarrow R \begin{bmatrix}x_1\\x_2\\x_3 \end{bmatrix} \)
Few examples
Representations of Rotations
The Vectors
\(\begin{bmatrix}x_1\\x_2\\x_3 \end{bmatrix}\longrightarrow R \begin{bmatrix}x_1\\x_2\\x_3 \end{bmatrix} \)
The Scalars
\(x\longrightarrow x\)
Few examples
Representations of Rotations
The Vectors
\(\begin{bmatrix}x_1\\x_2\\x_3 \end{bmatrix}\longrightarrow R \begin{bmatrix}x_1\\x_2\\x_3 \end{bmatrix} \)
The Scalars
\(x\longrightarrow x\)
Signal on the Sphere
\(f: S^2\to \mathbb{R}\)
\(f'(x)=f(R^{-1}x)\)
Few examples
Representations of Rotations
The Vectors
\(\begin{bmatrix}x_1\\x_2\\x_3 \end{bmatrix}\longrightarrow R \begin{bmatrix}x_1\\x_2\\x_3 \end{bmatrix} \)
The Scalars
\(x\longrightarrow x\)
Scalar Field
\( f: \mathbb{R}^3 \to \mathbb{R}\)
\(f'(x)=f(R^{-1}x)\)
Signal on the Sphere
\(f: S^2\to \mathbb{R}\)
\(f'(x)=f(R^{-1}x)\)
Few examples
Group Representations
\((\rho, V)\)
\(\rho:G \to (V\to V)\) \(g,g_1,g_2 \in G\) \(x, y \in V\)
- \(\rho(g)(x+\alpha y) = \rho(g)(x) + \alpha \rho(g)(y)\)
- \(\rho(g_2)(\rho(g_1)(x)) = \rho(g_2 g_1)(x) \)
The Vectors
\(\begin{bmatrix}x_1\\x_2\\x_3 \end{bmatrix}\longrightarrow R \begin{bmatrix}x_1\\x_2\\x_3 \end{bmatrix} \)
The Scalars
\(x\longrightarrow x\)
Scalar Field
\( f: \mathbb{R}^3 \to \mathbb{R}\)
\(f'(x)=f(R^{-1}x)\)
Signal on the Sphere
\(f: S^2\to \mathbb{R}\)
\(f'(x)=f(R^{-1}x)\)
Irreducible Representations
The Vectors
\(\begin{bmatrix}x_1\\x_2\\x_3 \end{bmatrix}\longrightarrow R \begin{bmatrix}x_1\\x_2\\x_3 \end{bmatrix} \)
The Scalars
\(x\longrightarrow x\)
Scalar Field
\( f: \mathbb{R}^3 \to \mathbb{R}\)
\(f'(x)=f(R^{-1}x)\)
Signal on the Sphere
\(f: S^2\to \mathbb{R}\)
\(f'(x)=f(R^{-1}x)\)
irreducible
irreducible
reducible
reducible
Irreducible Representations
Scalar Field
\( f: \mathbb{R}^3 \to \mathbb{R}\)
\(f'(x)=f(R^{-1}x)\)
reducible
Irreducible Representations
=
\(c_1 \times\)
\(c_2 \times\)
\(c_3 \times\)
\(c_4 \times\)
\(c_5 \times\)
irreducible
\(c_6 \times\)
Irreps of Rotations
| Index | Name | Examples of quantities |
|---|---|---|
| L=0 | Scalars | temperature, norm of a vector, orbital s, ... |
| L=1 | Vectors | velocity, force, orbital p, ... |
| L=2 | orbital d | |
| L=3 | orbital f | |
| L=4 | orbital g | |
| L=5 | ... | |
| L=6 | ||
| L=7 | ||
| L=8 | ||
| L=9 | ||
| L=10 | ||
| L=11 |
Irreps of Rotations
| Index | Name | Examples of quantities |
|---|---|---|
| L=0 | Scalars | temperature, norm of a vector, orbital s, ... |
| L=1 | Vectors | velocity, force, orbital p, ... |
| L=2 | orbital d | |
| L=3 | orbital f | |
| L=4 | orbital g | |
| L=5 | ... | |
| L=6 | ||
| L=7 | ||
| L=8 | ||
| L=9 | ||
| L=10 | ||
| L=11 |
Stress Tensor
(3x3 matrix)
\(\}\)
\(\sigma\longrightarrow R\sigma R^T\)
Everything can be decomposed into irreps:
Tensor Product
\(\rho_1 \otimes \rho_2\) is a representation
acting on the vector space \(V_1 \otimes V_2\)
\(X \in \mathbb{R}^{\dim V_1\times\dim V_2}\)
\(X \longrightarrow \rho_1(g) X \rho_2(g)^T \)
Tensor Product
\(\rho_1 \otimes \rho_2\) is a representation
acting on the vector space \(V_1 \otimes V_2\)
\(X \in \mathbb{R}^{\dim V_1\times\dim V_2}\)
\(X \longrightarrow \rho_1(g) X \rho_2(g)^T \)
(\(X_{ij} \longrightarrow \rho_1(g)_{ik}\rho_2(g)_{jl} X_{kl} \))
Tensor Product
reducible
=
direct sum of
irreducible
\(\rho_1 \otimes \rho_2\)
\(\rho_3 \oplus \rho_4 \oplus \rho_4\)
Tensor Product
\(G\)
\(\rho_1\)
\(\rho_2\)
\(\rho_3\)
\(\rho_4\)
\(\rho_5\)
\(\otimes\)
\(\rho_5\)
\(\rho_1\)
\(\rho_2\)
Tensor Product of Rotations
Example:
\(D_2 \otimes D_1 = D_1 \oplus D_2 \oplus D_3\)
\(D_L\) is the irreps of order L
Tensor Product of Rotations
\(D_L\) is the irreps of order L
General formula:
\(D_j \otimes D_k = D_{|j-k|} \oplus \dots \oplus D_{j+k}\)
Example:
\(D_2 \otimes D_1 = D_1 \oplus D_2 \oplus D_3\)
Equivariant Neural Network
Using the tools presented previously you can create any equivariant polynomials
Equivariant
Polynomial
\(\theta\)
\(\rho_1\)
\(\rho_2\)
\(\rho_2\)
\(\rho_3\)
\(\rho_1\)
\(\rho_1\)
\(\rho_2\)
\(\rho_3\)
\(\rho_1\)
\(\rho_2\)
\(\rho_4\)
\(\rho_4\)
\(\rho_1\)
\(\otimes\)
\(\otimes\)
\(\otimes\)
\(\oplus\)
\(\oplus\)
\(\oplus\)
\(\oplus\)
\(\otimes\)
\(\otimes\)
\(\otimes\)
\(\oplus\)
\(\oplus\)
\(\oplus\)
\(\oplus\)
Equivariant Neural Network
\(\theta\)
\(\theta\)
Equivariant Neural Networks Architectures
| Group | Name | Ref |
|---|---|---|
| Translation | Convolutional Neural Networks | |
| 90 degree rotation 2D | Group Equivariant CNN | 1602.07576 |
| 2D Rotations | Harmonic Networks | 1612.04642 |
| 2D Scale | Deep Scale-spaces | 1905.11697 |
| 3D Rotations | 3D Steerable CNN, Tensor Field Network | 1807.02547 1802.08219 |
| Lorentz | Lorentz Group Equivariant NN | 2006.04780 |
Library to make ENN for Rotations
We wrote python code to help creating Equivariant Neural Networks
$ pip install e3nn
We wrote python code to help creating Equivariant Neural Networks
$ pip install e3nn
Library to make ENN for Rotations
import e3nn e3nn.o3.spherical_harmonics(2, x, True)
Spherical Harmonics are Equivariant Polynomials
Graph Convolution
Nequip
(TFN: Nathaniel Thomas et al. 2018)
(Nequip: Simon Batzner et al. 2021)
source
dest.
\(h\)
\(\vec r\)
\(m = h \otimes Y(\vec r)\)
\(m\)
* this formula is missing the parameterized radial function
Nequip Learning Curve
(Nequip: Simon Batzner et al. 2021)
max L of the messages
The Curse of Dimensionality
and the Learning Curve
\(P =\) size of trainset
\(d =\) dimension of the data
\(\delta =\) distance to closest neighbor
Bach (2017)
The Curse of Dimensionality
and the Learning Curve
\(P =\) size of trainset
\(d =\) dimension of the data
\(\delta =\) distance to closest neighbor
\(\epsilon =\) test error
Hestness et al. (2017)
regression + Lipschitz continuous
Luxburg and Bousquet (2004)
MACE
(MACE: Ilyes Batatia et al. 2022)
source
\(1\)
dest.
\(h_1\)
\(\vec r_1\)
\(m = F_\theta(\{h_i\otimes Y(\vec r_i)\}_{i=1}^\nu)\)
\(m\)
source
\(2\)
source
\(\nu\)
\(h_2\)
\(h_\nu\)
\(\vec r_2\)
\(\vec r_\nu\)
MACE
(MACE: Ilyes Batatia et al. 2022)
L
L
3
\(m = F_\theta(\{h_i\otimes Y(\vec r_i)\}_{i=1}^\nu)\)
any L and \(\nu=1\)
\(h \otimes Y(\vec r)\)
L=0 and \(\nu=2\)
\(h_1Y(\vec r_1) \cdot h_2Y(\vec r_2)\)
Legendre polynomials
L=0 and \(\nu=3\)
\((h_1Y(\vec r_1) \otimes h_2Y(\vec r_2)) \cdot h_3Y(\vec r_3)\)
any L and \(\nu=3\)
\(h_1\otimes Y(\vec r_1) \otimes h_2\otimes Y(\vec r_2) \otimes h_3\otimes Y(\vec r_3)\)
Kind of operations in MACE
Conclusion
Equivariant Neural Networks are more data efficient if they incorporate Tensor Products of order \(L \geq 1\)
but not necessary as features (MACE)
Thanks for listening
The slides are available at
https://slides.com/mariogeiger/youth2022
Group Theory and Machine Learning
By Mario Geiger
