Equivariant Polynomials with e3nn-jax

Mario Geiger

Postdoc at

This Talk is about Equivariant Neural Networks

input

output

Illustration of a neural network equivariant to rotations in 3D

What is e3nn?

https://github.com/e3nn/e3nn

What is e3nn?

protein folding

Geometric deep learning of RNA structure R. TOWNSHEND et al. 

molecular dynamics

Nequip S. Batzner et al.    MACE I. Batatia et al.

solid state physics

Prediction of Phonon Density Z. Chen et al. 

molecular electron densities

Cracking the Quantum Scaling Limit with Machine Learned Electron Densities J. Rackers

Cosmology

Medical Images

What is e3nn?

Data efficient!

(Nequip: Simon Batzner et al. 2021)

max L of the messages

What is e3nn?

Data efficient!

(Nequip: Simon Batzner et al. 2021)

max L of the messages

What is e3nn?

Data efficient!

(MACE: Ilyes Batatia et al. 2022)

L

L

3

\(m = (\sum_i h_i\otimes Y(\vec r_i))^{\otimes\nu}\)

(

)

What is e3nn?

Also a jax library 🚀

https://github.com/e3nn/e3nn-jax

Why Equivariant Polynomials?

All linear functions

Why Equivariant Polynomials?

All linear functions

All polynomials

Why Equivariant Polynomials?

All linear functions

All polynomials

All analytical functions

Why Equivariant Polynomials?

All linear functions

All polynomials

All analytical functions

All smooth functions

Why Equivariant Polynomials?

All linear functions

All polynomials

All analytical functions

All smooth functions

All continuous functions

Why Equivariant Polynomials?

All linear functions

All polynomials

All analytical functions

All smooth functions

All continuous functions

All functions

Why Equivariant Polynomials?

All linear functions

All polynomials

All analytical functions

All smooth functions

All continuous functions

All functions

Most physics is described by smooth functions

Random Facts about these classes of functions

Why Equivariant Polynomials?

All linear functions

All polynomials

All analytical functions

All smooth functions

All continuous functions

All functions

Some Phase transitions are characterized by non continuous functions

Random Facts about these classes of functions

Why Equivariant Polynomials?

All linear functions

All polynomials

All analytical functions

All smooth functions

All continuous functions

All functions

Some Phase transitions are characterized by non continuous functions

... that are limits of analytical functions (Landau theory)

Random Facts about these classes of functions

Why Equivariant Polynomials?

All linear functions

All polynomials

All analytical functions

All smooth functions

All continuous functions

All functions

"Convergence of the training can then be related to the positive-definiteness of the
limiting NTK. We prove the positive-definiteness of the limiting NTK when the
data is supported on the sphere and the non-linearity is non-polynomial."

TNK Arthur Jacot 2020

Random Facts about these classes of functions

Why Equivariant Polynomials?

All linear functions

All polynomials

All analytical functions

All smooth functions

All continuous functions

All functions

analytical functions are the limits of polynomials

Random Facts about these classes of functions

Why Equivariant Polynomials?

All linear functions

All polynomials

All analytical functions

All smooth functions

All continuous functions

All functions

It's a good start to be able to build any polynomial 🤷🏼‍♂️

Group and Representations

Group and Representations

"what are the operations"

"how they compose"

Group and Representations

"what are the operations"

"how they compose"

"vector spaces on which the action of the group is defined"

Group and Representations

"what are the operations"

"how they compose"

rotations, parity, (translations)

scalars, vectors, pseudovectors, ...

"vector spaces on which the action of the group is defined"

Group and Representations

Group \(G\)

• \(\text{identity} \in G\)
• associativity \(g (hk) = (gh)k\)
• inverse \(g^{-1} \in G\)

"what are the operations" "how they compose"

"vector spaces on which the action of the group is defined"

Group and Representations

"what are the operations" "how they compose"

"vector spaces on which the action of the group is defined"

Group \(G\)

• \(\text{identity} \in G\)
• associativity \(g (hk) = (gh)k\)
• inverse \(g^{-1} \in G\)

Group and Representations

"what are the operations" "how they compose"

"vector spaces on which the action of the group is defined"

Group \(G\)

• \(\text{identity} \in G\)
• associativity \(g (hk) = (gh)k\)
• inverse \(g^{-1} \in G\)

Group and Representations

"what are the operations" "how they compose"

"vector spaces on which the action of the group is defined"

Group \(G\)

• \(\text{identity} \in G\)
• associativity \(g (hk) = (gh)k\)
• inverse \(g^{-1} \in G\)

Representation \(D(g, x)\)

• \(g\in G\), \(x \in V\)
• Linear \(D(g, x+y) = D(g,x) + D(g,y) \)
• Follow the structure of the group
  \(D(gh,x) = D(g, D(h,x))\)

Group and Representations

"what are the operations" "how they compose"

"vector spaces on which the action of the group is defined"

Group \(G\)

• \(\text{identity} \in G\)
• associativity \(g (hk) = (gh)k\)
• inverse \(g^{-1} \in G\)

Representation \(D(g, x)\)

• \(g\in G\), \(x \in V\)
• Linear \(D(g, x+y) = D(g,x) + D(g,y) \)
• Follow the structure of the group
  \(D(gh,x) = D(g, D(h,x))\)

Group and Representations

"what are the operations" "how they compose"

"vector spaces on which the action of the group is defined"

Group \(G\)

• \(\text{identity} \in G\)
• associativity \(g (hk) = (gh)k\)
• inverse \(g^{-1} \in G\)

Representation \(D(g, x)\)

• \(g\in G\), \(x \in V\)
• Linear \(D(g, x+y) = D(g,x) + D(g,y) \)
• Follow the structure of the group
  \(D(gh,x) = D(g, D(h,x))\)

Group and Representations

"what are the operations" "how they compose"

"vector spaces on which the action of the group is defined"

Group \(G\)

• \(\text{identity} \in G\)
• associativity \(g (hk) = (gh)k\)
• inverse \(g^{-1} \in G\)

Representation \(D(g, x)\)

• \(g\in G\), \(x \in V\)
• Linear \(D(g, x+y) = D(g,x) + D(g,y) \)
• Follow the structure of the group
  \(D(gh,x) = D(g, D(h,x))\)

Group and Representations

"what are the operations" "how they compose"

"vector spaces on which the action of the group is defined"

Group \(G\)

• \(\text{identity} \in G\)
• associativity \(g (hk) = (gh)k\)
• inverse \(g^{-1} \in G\)

Representation \(D(g, x)\)

• \(g\in G\), \(x \in V\)
• Linear \(D(g, x+y) = D(g,x) + D(g,y) \)
• Follow the structure of the group
  \(D(gh,x) = D(g, D(h,x))\)

Equivalent notation \(D(g) x\)

• \(D(g) : V\to V\)
• \(D(g) \in \mathbb{R}^{d\times d}\)
• \(D(gh) = D(g) D(h)\)

Group and Representations

"what are the operations" "how they compose"

"vector spaces on which the action of the group is defined"

Group \(G\)

• \(\text{identity} \in G\)
• associativity \(g (hk) = (gh)k\)
• inverse \(g^{-1} \in G\)

Representation \(D(g, x)\)

• \(g\in G\), \(x \in V\)
• Linear \(D(g, x+y) = D(g,x) + D(g,y) \)
• Follow the structure of the group
  \(D(gh,x) = D(g, D(h,x))\)

Group and Representations

"what are the operations" "how they compose"

"vector spaces on which the action of the group is defined"

Equivalent notation \(D(g) x\)

• \(D(g) : V\to V\)
• \(D(g) \in \mathbb{R}^{d\times d}\)
• \(D(gh) = D(g) D(h)\)

Group \(G\)

• \(\text{identity} \in G\)
• associativity \(g (hk) = (gh)k\)
• inverse \(g^{-1} \in G\)

Representation \(D(g, x)\)

• \(g\in G\), \(x \in V\)
• Linear \(D(g, x+y) = D(g,x) + D(g,y) \)
• Follow the structure of the group
  \(D(gh,x) = D(g, D(h,x))\)

Group and Representations

"what are the operations" "how they compose"

"vector spaces on which the action of the group is defined"

Equivalent notation \(D(g) x\)

• \(D(g) : V\to V\)
• \(D(g) \in \mathbb{R}^{d\times d}\)
• \(D(gh) = D(g) D(h)\)

Group \(G\)

• \(\text{identity} \in G\)
• associativity \(g (hk) = (gh)k\)
• inverse \(g^{-1} \in G\)

Representation \(D(g, x)\)

• \(g\in G\), \(x \in V\)
• Linear \(D(g, x+y) = D(g,x) + D(g,y) \)
• Follow the structure of the group
  \(D(gh,x) = D(g, D(h,x))\)

Group and Representations

"what are the operations" "how they compose"

"vector spaces on which the action of the group is defined"

Equivalent notation \(D(g) x\)

• \(D(g) : V\to V\)
• \(D(g) \in \mathbb{R}^{d\times d}\)
• \(D(gh) = D(g) D(h)\)

Group \(G\)

• \(\text{identity} \in G\)
• associativity \(g (hk) = (gh)k\)
• inverse \(g^{-1} \in G\)

Examples of representations

\(\begin{bmatrix} a^1\\a^2\\a^3\\a^4\\a^5\\a^6\\a^7\\a^8\\a^9\end{bmatrix}\in \mathbb{R}^9\)

Examples of representations

Representations are like data types

It tells you how to interpret the data with respect to the group action

\(\begin{bmatrix} a^1\\a^2\\a^3\\a^4\\a^5\\a^6\\a^7\\a^8\\a^9\end{bmatrix}\in \mathbb{R}^9\)

Examples of representations

\(\begin{bmatrix} a^1\\a^2\\a^3\\a^4\\a^5\\a^6\\a^7\\a^8\\a^9\end{bmatrix}\in \mathbb{R}^9\)

Knowing that \(a_1, a_2, a_3\) are scalars tells you that they are not affected by a rotation of your system

Representations are like data types

It tells you how to interpret the data with respect to the group action

Examples of representations

\(\begin{bmatrix} a^1\\a^2\\a^3\\a^4\\a^5\\a^6\\a^7\\a^8\\a^9\end{bmatrix}\in \mathbb{R}^9\)

Representations are like data types

It tells you how to interpret the data with respect to the group action

Examples of representations

\(\begin{bmatrix} a^1\\a^2\\a^3\\a^4\\a^5\\a^6\\a^7\\a^8\\a^9\end{bmatrix}\in \mathbb{R}^9\)

If the system is rotated, the 3 components of the vector change together!

Representations are like data types

It tells you how to interpret the data with respect to the group action

Examples of representations

\(\begin{bmatrix} a^1\\a^2\\a^3\\a^4\\a^5\\a^6\\a^7\\a^8\\a^9\end{bmatrix}\in \mathbb{R}^9\)

The two vectors transforms indepentently

Representations are like data types

It tells you how to interpret the data with respect to the group action

Examples of representations

\(\begin{bmatrix} a^1\\a^2\\a^3\\a^4\\a^5\\a^6\\a^7\\a^8\\a^9\end{bmatrix}\in \mathbb{R}^9\)

system rotated by \(g\)

\(\begin{bmatrix} a'^1\\a'^2\\a'^3\\a'^4\\a'^5\\a'^6\\a'^7\\a'^8\\a'^9\end{bmatrix}=D(g)\begin{bmatrix} a^1\\a^2\\a^3\\a^4\\a^5\\a^6\\a^7\\a^8\\a^9\end{bmatrix}\)

Representations are like data types

It tells you how to interpret the data with respect to the group action

Examples of representations

\(\begin{bmatrix} a^1\\a^2\\a^3\\a^4\\a^5\\a^6\\a^7\\a^8\\a^9\end{bmatrix}\in \mathbb{R}^9\)

system rotated by \(g\)

\(\begin{bmatrix} a^1\\a^2\\a^3\\a^4\\a^5\\a^6\\a^7\\a^8\\a^9\end{bmatrix}\)

Representations are like data types

It tells you how to interpret the data with respect to the group action

\(\begin{bmatrix} a'^1\\a'^2\\a'^3\\a'^4\\a'^5\\a'^6\\a'^7\\a'^8\\a'^9\end{bmatrix}=\)

Examples of representations

\(\begin{bmatrix} a^1\\a^2\\a^3\\a^4\\a^5\\a^6\\a^7\\a^8\\a^9\end{bmatrix}\in \mathbb{R}^9\)

system rotated by \(g\)

Representations are like data types

It tells you how to interpret the data with respect to the group action

\(\begin{bmatrix} a^1\\a^2\\a^3\\a^4\\a^5\\a^6\\a^7\\a^8\\a^9\end{bmatrix}\)

\(\begin{bmatrix} a'^1\\a'^2\\a'^3\\a'^4\\a'^5\\a'^6\\a'^7\\a'^8\\a'^9\end{bmatrix}=\)

Examples of representations

\(\begin{bmatrix} a^1\\a^2\\a^3\\a^4\\a^5\\a^6\\a^7\\a^8\\a^9\end{bmatrix}\in \mathbb{R}^9\)

system rotated by \(g\)

Representations are like data types

It tells you how to interpret the data with respect to the group action

\(\begin{bmatrix} a^1\\a^2\\a^3\\a^4\\a^5\\a^6\\a^7\\a^8\\a^9\end{bmatrix}\)

\(\begin{bmatrix} a'^1\\a'^2\\a'^3\\a'^4\\a'^5\\a'^6\\a'^7\\a'^8\\a'^9\end{bmatrix}=\)

Equivariance

\(V\)

\(V'\)

Equivariance

\(V\)

\(V'\)

\(D(g)\)

\(D'(g)\)

\(V\)

\(V'\)

Equivariance

\(V\)

\(V'\)

\(D(g)\)

\(D'(g)\)

\(V\)

\(V'\)

\(f\)

Equivariance

\(V\)

\(V'\)

\(D(g)\)

\(D'(g)\)

\(V\)

\(V'\)

\(f\)

\(f\)

Equivariance

\(V\)

\(V'\)

\(D(g)\)

\(D'(g)\)

\(V\)

\(V'\)

\(f\)

\(f\)

\(f(D(g) x)\)

Equivariance

\(V\)

\(V'\)

\(D(g)\)

\(D'(g)\)

\(V\)

\(V'\)

\(f\)

\(f\)

\(f(D(g) x)\)

\(D'(g) f(x)\)

Equivariance

\(V\)

\(V'\)

\(D(g)\)

\(D'(g)\)

\(V\)

\(V'\)

\(f\)

\(f\)

\(f(D(g) x)\)

\(D'(g) f(x)\)

\(=\)

Polynomials

\(x \mapsto x^2 + 2(x-4)\)

Polynomials

\(x \mapsto x^2 + 2(x-4)\)

\(\begin{bmatrix} x\\ y\\ z \end{bmatrix} \mapsto x^2 + 2(y-z)x\)

Polynomials

\(x \mapsto x^2 + 2(x-4)\)

\(\begin{bmatrix} x\\ y\\ z \end{bmatrix} \mapsto x^2 + 2(y-z)x\)

\(\begin{bmatrix} x\\ y\\ z \end{bmatrix} \mapsto \begin{bmatrix} x^2 + 2(y-z) \\ z^4 + 100 x y z\end{bmatrix}\)

Equivariant Polynomials

\(P(D(g) x) = D'(g) P(x)\)

Equivariant Polynomials

\(P(D(g) x) = D'(g) P(x)\)

Equivariant Polynomials

\(P(D(g) x) = D'(g) P(x)\)

\(R\)

not linear, probably not even invertible

\(?\)

Equivariant Polynomials

\(P(D(g) x) = D'(g) P(x)\)

\(R\)

\(?\)

(This one is actually invariant)

5 Tools to Build Equivariant Polynomials

• 🔨 Composition \(\circ\)
• 🔧 Addition \(+\)
• 🔩 Multiplication \(\otimes\)
• 💡 Linear Mixing

5 Tools to Build Equivariant Polynomials

• 🔨 Composition \(\circ\)
• 🔧 Addition \(+\)
• 🔩 Multiplication \(\otimes\)
• 💡 Linear Mixing

Bottom-Up approach !!  Combine simple polynomials into more complex ones

IrrepsArray

IrrepsArray

3 scalars

1 vector

IrrepsArray

IrrepsArray

3 scalars

1 vector

IrrepsArray

3 scalars

1 vector

🔨 Composition

two equivariant functions

\(f: V_1 \rightarrow V_2\)

\(h: V_2 \rightarrow V_3\)

\(h\circ f\) is equivariant!! 😊

\( h(f(D_1(g) x)) = h(D_2(g) f(x)) = D_3(g) h(f(x)) \)

🔨 Composition

two equivariant functions

\(f: V_1 \rightarrow V_2\)

\(h: V_2 \rightarrow V_3\)

\(h\circ f\) is equivariant!! 😊

🔧 Addition

two equivariant functions

\(f: V_1 \rightarrow V_3\)

\(h: V_2 \rightarrow V_3\)

\(h + f\) is equivariant!! 😊

\( f(D_1(g) x) + h(D_2(g)x) = D_3(g) (f(x) + h(x)) \)

🔧 Addition

two equivariant functions

\(f: V_1 \rightarrow V_3\)

\(h: V_2 \rightarrow V_3\)

\(h + f\) is equivariant!! 😊

🔩 Multiplication

\(\begin{bmatrix} {\color{red} x_1}\\{\color{red} x_2}\\{\color{red} x_3}\end{bmatrix}\)

\(\begin{bmatrix} {\color{blue} y_1}\\{\color{blue} y_2}\\{\color{blue} y_3}\\{\color{blue} y_4}\\{\color{blue} y_5} \end{bmatrix}\)

transforming with \(D(g)\)

transforming with \(D'(g)\)

\(= \begin{bmatrix} x_1y_1 & x_1y_2 & x_1y_3 & x_1y_4 & x_1y_5 \\ x_2y_1 & x_2y_2 & x_2y_3 & x_2y_4 & x_2y_5 \\ x_3y_1 & x_3y_2 & x_3y_3 & x_3y_4 & x_3y_5 \end{bmatrix}\)

🔩 Multiplication

\(\otimes\)

\(= \begin{bmatrix}{\color{red} x_1} {\color{blue} y_1}&{\color{red} x_1} {\color{blue} y_2}&{\color{red} x_1} {\color{blue} y_3}&{\color{red} x_1} {\color{blue} y_4}&{\color{red} x_1} {\color{blue} y_5}\\{\color{red} x_2} {\color{blue} y_1}&{\color{red} x_2} {\color{blue} y_2}&{\color{red} x_2} {\color{blue} y_3}&{\color{red} x_2} {\color{blue} y_4}&{\color{red} x_2} {\color{blue} y_5}\\{\color{red} x_3} {\color{blue} y_1}&{\color{red} x_3} {\color{blue} y_2}&{\color{red} x_3} {\color{blue} y_3}&{\color{red} x_3} {\color{blue} y_4}&{\color{red} x_3} {\color{blue} y_5}\end{bmatrix}\)

Tensor Product

\(\begin{bmatrix} {\color{red} x_1}\\{\color{red} x_2}\\{\color{red} x_3}\end{bmatrix}\)

\(\begin{bmatrix} {\color{blue} y_1}\\{\color{blue} y_2}\\{\color{blue} y_3}\\{\color{blue} y_4}\\{\color{blue} y_5} \end{bmatrix}\)

Reducible representations

\(D\) defined on \(V\)

is reducible if

\(\exists W \subset V\)     \(W\neq0, V\)

such that

\(D|_W\) is a representation

Reducible representations

Famous Example

\(\begin{bmatrix}{\color{red} x_1} {\color{blue} x_2}&{\color{red} x_1} {\color{blue} y_2}&{\color{red} x_1} {\color{blue} z_2}\\{\color{red} y_1} {\color{blue} x_2}&{\color{red} y_1} {\color{blue} y_2}&{\color{red} y_1} {\color{blue} z_2}\\{\color{red} z_1} {\color{blue} x_2}&{\color{red} z_1} {\color{blue} y_2}&{\color{red} z_1} {\color{blue} z_2}\end{bmatrix}\)

\(\begin{bmatrix} {\color{red} x_1}\\{\color{red} y_1}\\{\color{red} z_1}\end{bmatrix}\otimes\begin{bmatrix} {\color{blue} x_2}\\{\color{blue} y_2}\\{\color{blue} z_2}\end{bmatrix} = \)

\(D\) defined on \(V\)

is reducible if

\(\exists W \subset V\)     \(W\neq0, V\)

such that

\(D|_W\) is a representation

Reducible representations

\(\begin{bmatrix}{\color{red} x_1} {\color{blue} x_2}&{\color{red} x_1} {\color{blue} y_2}&{\color{red} x_1} {\color{blue} z_2}\\{\color{red} y_1} {\color{blue} x_2}&{\color{red} y_1} {\color{blue} y_2}&{\color{red} y_1} {\color{blue} z_2}\\{\color{red} z_1} {\color{blue} x_2}&{\color{red} z_1} {\color{blue} y_2}&{\color{red} z_1} {\color{blue} z_2}\end{bmatrix}\)

\(D\) is reducible if

\(\exists W \subset V\)

such that

\(D|_W\) is a representation

Irreducible representations

For the group of rotations (\(SO(3)\))

They are index by \(L=0, 1, 2, \dots\)

Of dimension \(2L+1\)

Irreducible representations

For the group of rotations (\(SO(3)\))

They are index by \(L=0, 1, 2, \dots\)

Of dimension \(2L+1\)

to use to achieve better data efficiency

Irreducible representations

For the group of rotations + parity (\(O(3)\))

They are index by \(L=0, 1, 2, \dots\)

and \(p=\pm 1\)

Of dimension \(2L+1\)

Even

Odd

Irreducible representations

For the group of rotations + parity (\(O(3)\))

They are index by \(L=0, 1, 2, \dots\)

and \(p=\pm 1\)

Of dimension \(2L+1\)

Even

Odd

🔩 Multiplication

\(L_1 \otimes L_2 = |L_1-L_2| \oplus \dots \oplus (L_1+L_2)\)

🔩 Multiplication

\(L_1 \otimes L_2 = |L_1-L_2| \oplus \dots \oplus (L_1+L_2)\)

generalization of \(3\times3=1+3+5\)

• \( 2 \otimes 1 = 1 \oplus 2 \oplus 3 \)
• \( 2 \otimes 2 = 0\oplus 1 \oplus 2 \oplus 3 \oplus 4 \)

🔩 Multiplication

\(L_1 \otimes L_2 = |L_1-L_2| \oplus \dots \oplus (L_1+L_2)\)

• \( 2 \otimes 1 = 1 \oplus 2 \oplus 3 \)
   
• \( 2 \otimes 2 = 0\oplus 1 \oplus 2 \oplus 3 \oplus 4 \)

🔩 Multiplication

🔩 Multiplication

🔩 Multiplication

💡 Linear Mixing

\(\begin{bmatrix} a^1\\a^2\\a^3\\a^4\\a^5\\a^6\\a^7\\a^8\\a^9\end{bmatrix}\)

\(\begin{bmatrix} b^1\\b^2\\b^3\\b^4\\b^5\\b^6\\b^7\\b^8\\b^9\end{bmatrix}\)

Linear map

💡 Linear Mixing

\(\begin{bmatrix} a^1\\a^2\\a^3\\a^4\\a^5\\a^6\\a^7\\a^8\\a^9\end{bmatrix}\)

\(\begin{bmatrix} b^1\\b^2\\b^3\\b^4\\b^5\\b^6\\b^7\\b^8\\b^9\end{bmatrix}\)

\(w_1\)

\(w_2\)

\(w_3\)

by Schur's lemma

💡 Linear Mixing

\(\begin{bmatrix} a^1\\a^2\\a^3\\a^4\\a^5\\a^6\\a^7\\a^8\\a^9\end{bmatrix}\)

\(\begin{bmatrix} b^1\\b^2\\b^3\\b^4\\b^5\\b^6\\b^7\\b^8\\b^9\end{bmatrix}\)

\(w_1\)

\(w_2\)

\(w_3\)

🔨 Composition

🔧 Addition

🔩 Multiplication

💡 Linear Mixing

Spherical Harmonics

Conclusion

• Group and representations
• Equivariance \(f({\color{purple}D(g)} x) = {\color{darkgreen}D'(g)} f(x)\)
  • 🔨 Composition
  • 🔧 Addition
• Polynomials (Bottom-Up)
  •  🔩 Tensor Product  \(L_1 \otimes L_2 = |L_1-L_2| \oplus \dots \oplus (L_1+L_2)\)
  • 💡 Linear Mixing
• ​Example
  • ​​Spherical Harmonics

Thank you for listening!