symm-learning 0.2.2

Torch modules and utilities of equivariant/invariant learning

Symmetric Learning

Lightweight python package for doing geometric deep learning using ESCNN. This package simply holds:

• Generic equivariant torch models and modules that are not present in ESCNN.
• Linear algebra utilities when working with symmetric vector spaces.
• Statistics utilities for symmetric random variables.

Installation

pip install symm-learning
# or
git clone https://github.com/Danfoa/symmetric_learning
cd symmetric_learning
pip install -e .

Structure

---

Linear Algebra

• lstsq: Symmetry-aware computation of the least-squares solution to a linear system of equations with symmetric input-output data.
• invariant_orthogonal_projector: Computes the orthogonal projection to the invariant subspace of a symmetric vector space.

---

Statistics

• var_mean: Symmetry-aware computation of the variance and mean of a symmetric random variable.
• cov: Symmetry-aware computation of the covariance / cross-covariance of two symmetric random variables.

---

Models

• iMLP: Invariant MLP for learning invariant functions.
• eMLP: Equivariant MLP for learning equivariant functions.

---

Torch Modules

Change2DisentangledBasis

Module for changing the basis of a tensor to a disentangled / isotypic basis.

IrrepSubspaceNormPooling

Module for extracting invariant features from a geometric tensor, giving one feature per irreducible subspace/representation.%

EquivMultivariateNormal

Utility layer to parameterize a G-equivariant multivariate Gaussian/Normal distribution:

\begin{aligned}
y &\sim \mathcal{N} \bigl(\mu(x), \Sigma(x)\bigr)& \\
\text{s.t.}
&\rho_Y(g)\mu(x) = \mu \bigl(\rho_X(g)\cdot x\bigr) \\
&\rho_Y(g)\Sigma(x)\rho_Y(g)^{\top} = \Sigma\bigl(\rho_X(g)\cdot x\bigr),
\quad \forall\, g \in G.
\end{aligned}

Such that the conditional probability distribution of y given x is $\mathbb{G}$-invariant to the simultaneous group action on $\mathcal{X}$ and $\mathcal{Y}$:

$$ P(y \mid x) = P(\rho_Y(g) y \mid \rho_X(g) x) \quad \forall g \in \mathbb{G}. $$

This means that if you want to parameterize a $\mathbb{G}$-equivariant stochastic function $y = f(x)$ using neural networks, you can use any backbone architecture whose output are the input parameters of a EquivMultivariateNormal distribution, as shown below:

from escnn.group import CyclicGroup
from escnn.nn import FieldType
from symm_learning.models.emlp import EMLP
from symm_learning.nn.equiv_multivariate_gaussian import EquivMultivariateNormal

G = CyclicGroup(3)
x_type = FieldType(escnn.gspaces.no_base_space(G), representations=[G.regular_representation])
y_type = FieldType(escnn.gspaces.no_base_space(G), representations=[G.regular_representation] * 1)
# Instanciate the output equivariant multivariate normal distribution in order to get the NN output type
e_normal = EquivMultivariateNormal(y_type, diagonal=True)
# Instanciate your NN model to output the parameters of the distribution
nn = EMLP(in_type=x_type, out_type=e_normal.in_type)
# Sample from the distribution
x = x_type(torch.randn(10, x_type.size))
z = nn(x) # (B, dim_y + n_dof_cov)
dist = e_normal.get_distribution(z) # instance of  torch.distributions.MultivariateNormal
y = dist.sample()  # (B, n)

Here, $z$ is a (batch_size, dim_y + n_dof_cov) input tensor with the first dim_y entries defining the mean of the distribution $\mu(x)$ and the next n_dof_cov entries define the free degrees of freedom from the symmetry constrained covariance matrix. See below