How to get started with equivariant deep learning
Evangelos Chatzipantazis, Stefanos Pertigkiozoglou
Tutorial Overview
Case Study� Building an Equivariant Point Encoder
Equivariance in Different Applications
Where to go next?
Symmetry Groups
Invariance/Equivariance
Equivariant Layers
Tutorial Overview
Case Study� Building an Equivariant Point Encoder
Equivariance in Different Applications
Where to go next?
Symmetry Groups
Invariance/Equivariance
Equivariant Layers
Equivariant Deep Learning in Computer Perception
Classical Equivision.
SIFT Features (D.Lowe 2004)
=
+
Steerable Filters (W. Freeman 1991)
Early Deep Learning
LeNet (Y. LeCun 1998)
Neocognitron(K. Fukushima 1980)
Symmetry Groups
generators
Frieze
Wallpaper
Crystallographic
elements
Transformation Groups
A group is a set and an operation that satisfies:
A transformation group is a group of transformations
where the operation is composition.
The Math behind:
Group Actions
The same group can act on different objects.
Same group, “same” object, different actions. Redundant representations give rise to more symmetries.
A (left) group action on a set is a function that satisfies:
(linear) Representations
A linear representation of a group on a vector space is a map with the property:
It is a group action on a vector space that preserves the vector space property too. (G- space) If n-dimensional then n-by-n invertible matrices. “Geometric vectors!”
Why linear? Even if base space not linear e.g. sphere, functions on this sphere induce linear actions. The (left) regular representation
Equivalence
Finite dimensional: single for all
Theorem (Maschke, Peter-Weyl): Any finite dimensional unitary representation of a compact group (on finite dimensional vector space) over a field with characteristic 0 is completely reducible into a direct sum of irreducible representations
SO(3) and its Irreducible Representations
SO(3) can rotate:
A vector
An inertia matrix
A scalar function (color)
A vector function:
SO(3) is compact. All these representations can be decomposed into irreps.
Irreps are like primes for natural numbers. Building blocks for larger representations!
Symmetry group? Yes, for preserves: norm and handedness, vector space
3
1
2
3
1
2
SO(3) and its Irreducible Representations
Spherical
Harmonics
The irreducibles of SO(3) are called Wigner-D matrices of size
Example: Rotation of Inertia Matrix :
Tutorial Overview
Case Study� Building an Equivariant Point Encoder
Equivariance in Different Applications
Where to go next?
Symmetry Groups
Invariance/Equivariance
Equivariant Layers
Invariance and Equivariance
Input space:
Action:
Output space:
Action:
Equivariance
car
Invariance
Equivariance
Group / Space / Action : The story so far
Does my task have symmetries?
Specify group, input space, input action, output space, output action.
Next step: build layers that are parametric and process the input+action to produce output+action aka Equivariant layers.
Point Clouds
Occupancy Function
Benefits of Equivariance
Equivariance vs Canonicalization
Canonicalize instead! (i.e. choose an invariant input representation)
PCA
PCA
Intrinsic / Picasso Problem
Non-smoothness
Not robust to noise
Even selecting robust and general canonical object involves learning which in turn utilizes equivariant deep networks.
Equivariance vs Data Augmentation
Tutorial Overview
Case Study� Building an Equivariant Point Encoder
Equivariance in Different Applications
Where to go next?
Symmetry Groups
Invariance/Equivariance
Equivariant Layers
Building equivariant layers: Group / Space / Action
Equivalently: (EMLP by Finzi et al. 2021)
Input Representation: Output Representation:
Equivariant Linear Layer:
Schur’s Lemma: If irreducibles (on ) then:
If else,
Matrix group
Example: Irreducible Representations.
(VNN C.Deng et al. 2021)
(2D-VNN)
Example 2: Lie Group infinitesimal generators
Example 3: Discrete, generators
Equivariance:
Composing Lifting Convolution, Regular Group Convolution and a Projection operator (like Group Max pooling) constitutes the Regular G-CNNs (T. Cohen et. al 2016)
from function on (homogeneous) space
to function on group.
Equivariance:
Input/Output Action
(Tensor field)
Steerable Equivariant Layers (Weiler et. al 2018)
All (continuous) equivariant linear operators
Harmonic Networks (D.Worall et al. 2016)
Steerability Constraint
1. Lifting Convolution: expanded domain
2.Steerable Convolution:
expanded co-domain
Connection via Fourier Transform!
Equivariant Layers on Lie Groups
Equivariant Layers on Lie Groups
Riemannian Manifolds
Clifford Algebras
Ray Space
More general spaces
Tutorial Overview
Case Study� Building an Equivariant Point Encoder
Equivariance in Different Applications
Where to go next?
Symmetry Groups
Invariance/Equivariance
Equivariant Layers
Tutorial Overview
Case Study� Building an Equivariant Point Encoder
Equivariance in Different Applications
Where to go next?
Symmetry Groups
Invariance/Equivariance
Equivariant Layers
Case Study: Equivariant Point Encoder
0
2
3
1
0
2
3
1
Case Study: Equivariant Point Encoder
0
2
3
1
0
2
3
1
Translation Equivariance
0
2
3
1
0
2
3
1
Case Study: Equivariant Point Encoder
0
2
3
1
0
2
3
1
1
3
0
2
Permutation+Translation Equivariance
1
3
0
2
Case Study: Equivariant Point Encoder
3
2
1
0
Message Passing
Case Study: Equivariant Point Encoder
3
2
1
0
Message Passing
Case Study: Equivariant Point Encoder
3
2
1
0
Message Passing
Case Study: Equivariant Point Encoder
3
2
1
0
Message Passing
Case Study: Equivariant Point Encoder
0
3
2
1
Message Passing
Case Study: Equivariant Point Encoder
0
2
3
1
0
2
3
1
0
2
3
1
0
2
3
1
Rotational�Equivariance
Case Study: Equivariant Point Encoder
Rotational�Equivariance
0
2
3
1
0
2
3
1
0
2
3
1
0
2
3
1
Case Study: Equivariant Point Encoder
0
2
3
1
0
2
3
1
0
2
3
1
0
2
3
1
Case Study: Equivariant Point Encoder
0
2
3
1
0
2
3
1
0
2
3
1
0
2
3
1
Case Study: Equivariant Point Encoder
3
2
1
0
Message Passing
Case Study: Equivariant Point Encoder
Message Passing
3
2
1
0
Case Study: Equivariant Point Encoder
Message Passing
3
2
1
0
Case Study: Equivariant Point Encoder
�Equivariant Message Passing
3
2
1
0
Case Study: Equivariant Point Encoder
Equivariant Message Passing
Case Study: Equivariant Point Encoder
Equivariant Message Passing
Steerable Convolutional Kernel
Case Study: Equivariant Point Encoder
Equivariant Message Passing
Steerable Convolutional Kernel
0
2
3
1
0
2
3
1
Examples of Equivariant Linear Layers
Equivariant Linear Layer Example
We show an example of using Vector Neurons
class VNLinear(nn.Module):
def __init__(self, in_channels, out_channels):
super(VNLinear, self).__init__()
self.map_to_feat = nn.Linear(in_channels, out_channels, bias=False)
def forward(self, x):
'''
x: point features of shape [B, N_feat, 3, N_samples, ...]
'''
x_out = self.map_to_feat(x.transpose(1,-1)).transpose(1,-1)
return x_out
Examples of Equivariant Linear Layers
Equivariant Linear Layer Example
We show an example of using Vector Neurons
class VNLinear(nn.Module):
def __init__(self, in_channels, out_channels):
super(VNLinear, self).__init__()
self.map_to_feat = nn.Linear(in_channels, out_channels, bias=False)
def forward(self, x):
'''
x: point features of shape [B, N_feat, 3, N_samples, ...]
'''
x_out = self.map_to_feat(x.transpose(1,-1)).transpose(1,-1)
return x_out
SO(3) Steerable Convolutional Kernel
We show an example of using Tensor Field Network, by utilizing the E3NN library�������
Examples of Equivariant Linear Layers
Equivariant Linear Layer Example
We show an example of using Vector Neurons
class VNLinear(nn.Module):
def __init__(self, in_channels, out_channels):
super(VNLinear, self).__init__()
self.map_to_feat = nn.Linear(in_channels, out_channels, bias=False)
def forward(self, x):
'''
x: point features of shape [B, N_feat, 3, N_samples, ...]
'''
x_out = self.map_to_feat(x.transpose(1,-1)).transpose(1,-1)
return x_out
SO(3) Steerable Convolutional Kernel
We show an example of using Tensor Field Network, by utilizing the E3NN library�������
Examples of Equivariant Linear Layers
Equivariant Linear Layer Example
We show an example of using Vector Neurons [citations]
class VNLinear(nn.Module):
def __init__(self, in_channels, out_channels):
super(VNLinear, self).__init__()
self.map_to_feat = nn.Linear(in_channels, out_channels, bias=False)
def forward(self, x):
'''
x: point features of shape [B, N_feat, 3, N_samples, ...]
'''
x_out = self.map_to_feat(x.transpose(1,-1)).transpose(1,-1)
return x_out
SO(3) Steerable Convolutional Kernel
We show an example of using Tensor Field Network [citations], by utilizing E3NN library�������
SO(3) Steerable Convolutional Kernel
from e3nn import o3
from e3nn.nn import FullyConnectedTensorProduct
class CGTensorProduct(torch.nn.Module):
def __init__(irreps_input="5x0e+5x1o", irreps_output="5x0e+1x1o",
max_harmonics=2):
super().__init__()
self.irreps_input=o3.Irreps(irreps_input), o3.Irreps(irreps_output)
self.irreps_output=o3.Irreps(irreps_output)
self.edge_irreps=o3.spherical_harmonics(max_harmonics)
self.tp=FullyConnectedTensorProduct(self.irreps_input,
self.edge_irreps,
self.irreps_output,
internal_weights=False)
self.phi=MLP(in_channels=1,out_channels=self.tp.weight_numel)
def forward(self,f,edge_diff):
harm=o3.spherical_harmonics(self.edge_irreps,edge_diff)
weights=self.phi(torch.norm(edge_diff))
return self.tp(f,harm,weight)
Examples of Equivariant Linear Layers
Equivariant Linear Layer Example
We show an example of using Vector Neurons [citations]
class VNLinear(nn.Module):
def __init__(self, in_channels, out_channels):
super(VNLinear, self).__init__()
self.map_to_feat = nn.Linear(in_channels, out_channels, bias=False)
def forward(self, x):
'''
x: point features of shape [B, N_feat, 3, N_samples, ...]
'''
x_out = self.map_to_feat(x.transpose(1,-1)).transpose(1,-1)
return x_out
SO(3) Steerable Convolutional Kernel
We show an example of using Tensor Field Network [citations], by utilizing E3NN library�������
SO(3) Steerable Convolutional Kernel
from e3nn import o3
from e3nn.nn import FullyConnectedTensorProduct
class CGTensorProduct(torch.nn.Module):
def __init__(irreps_input="5x0e+5x1o", irreps_output="5x0e+1x1o",
max_harmonics=2):
super().__init__()
self.irreps_input=o3.Irreps(irreps_input), o3.Irreps(irreps_output)
self.irreps_output=o3.Irreps(irreps_output)
self.edge_irreps=o3.spherical_harmonics(max_harmonics)
self.tp=FullyConnectedTensorProduct(self.irreps_input,
self.edge_irreps,
self.irreps_output,
internal_weights=False)
self.phi=MLP(in_channels=1,out_channels=self.tp.weight_numel)
def forward(self,f,edge_diff):
harm=o3.spherical_harmonics(self.edge_irreps,edge_diff)
weights=self.phi(torch.norm(edge_diff))
return self.tp(f,harm,weight)
Examples of Equivariant Linear Layers
Equivariant Linear Layer Example
We show an example of using Vector Neurons [citations]
class VNLinear(nn.Module):
def __init__(self, in_channels, out_channels):
super(VNLinear, self).__init__()
self.map_to_feat = nn.Linear(in_channels, out_channels, bias=False)
def forward(self, x):
'''
x: point features of shape [B, N_feat, 3, N_samples, ...]
'''
x_out = self.map_to_feat(x.transpose(1,-1)).transpose(1,-1)
return x_out
SO(3) Steerable Convolutional Kernel
We show an example of using Tensor Field Network [citations], by utilizing E3NN library�������
SO(3) Steerable Convolutional Kernel
from e3nn import o3
from e3nn.nn import FullyConnectedTensorProduct
class CGTensorProduct(torch.nn.Module):
def __init__(irreps_input="5x0e+5x1o", irreps_output="5x0e+1x1o",
max_harmonics=2):
super().__init__()
self.irreps_input=o3.Irreps(irreps_input), o3.Irreps(irreps_output)
self.irreps_output=o3.Irreps(irreps_output)
self.edge_irreps=o3.spherical_harmonics(max_harmonics)
self.tp=FullyConnectedTensorProduct(self.irreps_input,
self.edge_irreps,
self.irreps_output,
internal_weights=False)
self.phi=MLP(in_channels=1,out_channels=self.tp.weight_numel)
def forward(self,f,edge_diff):
harm=o3.spherical_harmonics(self.edge_irreps,edge_diff)
weights=self.phi(torch.norm(edge_diff))
return self.tp(f,harm,weight)
Examples of Equivariant Linear Layers
Equivariant Linear Layer Example
We show an example of using Vector Neurons [citations]
class VNLinear(nn.Module):
def __init__(self, in_channels, out_channels):
super(VNLinear, self).__init__()
self.map_to_feat = nn.Linear(in_channels, out_channels, bias=False)
def forward(self, x):
'''
x: point features of shape [B, N_feat, 3, N_samples, ...]
'''
x_out = self.map_to_feat(x.transpose(1,-1)).transpose(1,-1)
return x_out
SO(3) Steerable Convolutional Kernel
We show an example of using Tensor Field Network [citations], by utilizing E3NN library�������
SO(3) Steerable Convolutional Kernel
from e3nn import o3
from e3nn.nn import FullyConnectedTensorProduct
class CGTensorProduct(torch.nn.Module):
def __init__(irreps_input="5x0e+5x1o", irreps_output="5x0e+1x1o",
max_harmonics=2):
super().__init__()
self.irreps_input=o3.Irreps(irreps_input), o3.Irreps(irreps_output)
self.irreps_output=o3.Irreps(irreps_output)
self.edge_irreps=o3.spherical_harmonics(max_harmonics)
self.tp=FullyConnectedTensorProduct(self.irreps_input,
self.edge_irreps,
self.irreps_output,
internal_weights=False)
self.phi=MLP(in_channels=1,out_channels=self.tp.weight_numel)
def forward(self,f,edge_diff):
harm=o3.spherical_harmonics(self.edge_irreps,edge_diff)
weights=self.phi(torch.norm(edge_diff))
return self.tp(f,harm,weight)
Examples of Equivariant Linear Layers
Equivariant Linear Layer Example
We show an example of using Vector Neurons [citations]
class VNLinear(nn.Module):
def __init__(self, in_channels, out_channels):
super(VNLinear, self).__init__()
self.map_to_feat = nn.Linear(in_channels, out_channels, bias=False)
def forward(self, x):
'''
x: point features of shape [B, N_feat, 3, N_samples, ...]
'''
x_out = self.map_to_feat(x.transpose(1,-1)).transpose(1,-1)
return x_out
SO(3) Steerable Convolutional Kernel
We show an example of using Tensor Field Network [citations], by utilizing E3NN library�������
SO(3) Steerable Convolutional Kernel
from e3nn import o3
from e3nn.nn import FullyConnectedTensorProduct
class CGTensorProduct(torch.nn.Module):
def __init__(irreps_input="5x0e+5x1o", irreps_output="5x0e+1x1o",
max_harmonics=2):
super().__init__()
self.irreps_input=o3.Irreps(irreps_input), o3.Irreps(irreps_output)
self.irreps_output=o3.Irreps(irreps_output)
self.edge_irreps=o3.spherical_harmonics(max_harmonics)
self.tp=FullyConnectedTensorProduct(self.irreps_input,
self.edge_irreps,
self.irreps_output,
internal_weights=False)
self.phi=MLP(in_channels=1,out_channels=self.tp.weight_numel)
def forward(self,f,edge_diff):
harm=o3.spherical_harmonics(self.edge_irreps,edge_diff)
weights=self.phi(torch.norm(edge_diff))
return self.tp(f,harm,weight)
Examples of Equivariant Linear Layers
Equivariant Linear Layer Example
We show an example of using Vector Neurons
class VNLinear(nn.Module):
def __init__(self, in_channels, out_channels):
super(VNLinear, self).__init__()
self.map_to_feat = nn.Linear(in_channels, out_channels, bias=False)
def forward(self, x):
'''
x: point features of shape [B, N_feat, 3, N_samples, ...]
'''
x_out = self.map_to_feat(x.transpose(1,-1)).transpose(1,-1)
return x_out
SO(3) Steerable Convolutional Kernel
We show an example of using Tensor Field Network, by utilizing E3NN library�������
from e3nn import o3
from e3nn.nn import FullyConnectedTensorProduct
class CGTensorProduct(torch.nn.Module):
def __init__(irreps_input="5x0e+5x1o", irreps_output="5x0e+1x1o",
max_harmonics=2):
super().__init__()
self.irreps_input=o3.Irreps(irreps_input), o3.Irreps(irreps_output)
self.irreps_output=o3.Irreps(irreps_output)
self.edge_irreps=o3.spherical_harmonics(max_harmonics)
self.tp=FullyConnectedTensorProduct(self.irreps_input,
self.edge_irreps,
self.irreps_output,
internal_weights=False)
self.phi=MLP(in_channels=1,out_channels=self.tp.weight_numel)
def forward(self,x,edge_diff):
harm=o3.spherical_harmonics(self.edge_irreps,edge_diff)
weights=self.phi(torch.norm(edge_diff))
return self.tp(x,harm,weight)
Equivariant Non-Linearities
Equivariant Constraint
General Recipe for Equivariant Pointwise NonLinearity
Non-Linearity
Invariant Features
Equivariant Features
Input Features
Equivariant Non-Linearities
Equivariant Constraint
Tensor Field Network Non-Linearity
Non-Linearity
Invariant Features
Equivariant Features
Input Features
Equivariant Non-Linearities
Equivariant Constraint
Vector Neurons Non-Linearity
Dot product
ReLU
Incorporating Attention Layers
3
2
1
0
Incorporating Attention Layers
3
2
1
0
Incorporating Attention Layers
3
2
1
0
softmax
Incorporating Attention Layers
3
2
1
0
softmax
Incorporating Attention Layers
3
2
1
0
softmax
Incorporating Attention Layers
3
2
1
0
3
2
1
0
Incorporating Attention Layers
3
2
1
0
softmax
3
2
1
0
Incorporating Attention Layers
3
2
1
0
softmax
3
2
1
0
Group Convolution
Lifting to Group
Group of discrete rotations C4
Group Convolution
Lifting to Group
Group Convolution
Lifting to Group
Group Convolution
Lifting to Group
Lifting to Group
Group Convolution
Lifting to Group
Lifting to Group
Tutorial Overview
Case Study� Building an Equivariant Point Encoder
Equivariance in Different Applications
Where to go next?
Symmetry Groups
Invariance/Equivariance
Equivariant Layers
Equivariant Point Encoder Applications
Point Cloud Segmentation/Classification
Point Cloud Reconstruction
Robotic Manipulation
Molecular Property Prediction
“Equivariant Descriptor Fields: SE(3)-Equivariant Energy-Based Models for End-to-End Visual Robotic Manipulation Learning”, Ryu et al. (ICLR 2023)
Figure from “SE(3)-Equivariant Attention Netowrks for Shape Reconstruction in Function Space” Chatzipantazis et al. (ICLR 2023)
2D Keypoint Extraction
The pre-specified 2D grid of images allows for precomputation of the steerable kernel
The ESCNN library allows for easy implementation �“A Program to Build E(N)-Equivariant Steerable CNNs” Cesa et al. (ICLR 2022)
“S-Trek: Sequential Translation and Rotation Equivariant Keypoints for local feature extraction”, Santellani et al. 2023
From Equivariance to Bi-Equivariance
Assume we are interested on the task of finding the transform that connects two puzzle pieces
T
From Equivariance to Bi-Equivariance
Assume we are interested on the task of finding the transform that connects two puzzle pieces
T
T
From Equivariance to Bi-Equivariance
From Equivariance to Bi-Equivariance
From Equivariance to Bi-Equivariance
From Equivariance to Bi-Equivariance
From Equivariance to Bi-Equivariance
Bi-Equivariance Example: Relative rotation estimation
Equivariant Encoder
Equivariant Encoder
Bi-Equivariant Feature
Bi-Equivariance Example: Relative rotation estimation
Equivariant Encoder
Equivariant Encoder
Bi-Equivariant Feature
SVD
Procrustes
Bi-Equivariance Example: Relative rotation estimation
Equivariant Encoder
Equivariant Encoder
Bi-Equivariant Feature
SVD
Procrustes
Bi-Equivariant Applications
Point Cloud Registration
Bi-Equivariant Applications
Point Cloud Registration
Pick and Place
"Diffusion-EDFs: Bi-equivariant Denoising Generative Modeling on SE(3) for Visual Robotic Manipulation", Ryu et al. (CVPR 2024)
Bi-Equivariant Applications
Point Cloud Registration
Pick and Place
"Diffusion-EDFs: Bi-equivariant Denoising Generative Modeling on SE(3) for Visual Robotic Manipulation", Ryu et al. (CVPR 2024)
Protein Docking
Protein Z-dependent inhibitor
Protein Z
“INDEPENDENT SE(3)-EQUIVARIANT MODELS FOR END-TO-END RIGID PROTEIN DOCKING”, Ganea et al. 2021
Permutation Equivariance
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
Permutation Equivariance
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| | | |
| | | |
| | | |
| | | |
| | | |
| | | |
| | | |
| | | |
| | | |
| | | |
| | | |
| | | |
Row and Column Permutation Equivariance
| | | |
| | | |
| | | |
| | | |
Permutation Equivariance
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| | | |
| | | |
| | | |
| | | |
| | | |
| | | |
| | | |
| | | |
| | | |
| | | |
| | | |
| | | |
| | | |
| | | |
| | | |
| | | |
Row and Column Permutation Equivariance
“Deep Models of Interactions Across Sets” Hartford et. al. 2018
Permutation Equivariant Applications
Permutation Equivariant SfM
[“Deep Permutation Equivariant Structure from Motion”, Moran et al. 2021]
Permutation Equivariant Applications
Permutation Equivariant SfM
“Deep Permutation Equivariant Structure from Motion”, Moran et al. 2021
Symmetries in Deep Weight Space
“Equivariant Architectures for Learning in Deep Weight Spaces”� [Navon et al. 2021]
Tutorial Overview
Case Study� Building an Equivariant Point Encoder
Equivariance in Different Applications
Where to go next?
Symmetry Groups
Invariance/Equivariance
Equivariant Layers
Limitations and Open Questions
There are still various open problems and limitation on applying equivariant representations
�
Limitations and Open Questions
There are still various open problems and limitation on applying equivariant representations
�
Limitations and Open Questions
There are still various open problems and limitation on applying equivariant representations
�
Limitations and Open Questions
There are still various open problems and limitation on applying equivariant representations
�
Limitations and Open Questions
There are still various open problems and limitation on applying equivariant representations
�
Limitations and Open Questions
There are still various open problems and limitation on applying equivariant representations
�
Limitations and Open Questions
There are still various open problems and limitation on applying equivariant representations
�
Limitations and Open Questions
There are still various open problems and limitation on applying equivariant representations
Limitations and Open Questions
There are still various open problems and limitation on applying equivariant representations
Limitations and Open Questions
There are still various open problems and limitation on applying equivariant representations
Limitations and Open Questions
Regardless the benefits there are still open problems and limitation on apply equivariant representations
Limitations and Open Questions
Regardless the benefits there are still open problems and limitation on apply equivariant representations
Limitations and Open Questions
There are still various open problems and limitation on applying equivariant representations
�
Useful Resources
Libraries
�Books/Courses