Unsupervised Learning of Group Invariant and Equivariant Representations

TL;DR

This work presents an unsupervised framework for learning representations that are simultaneously group-invariant and group-equivariant. By factorizing the latent space into a $G$-invariant code and a learnable group-action component, the model reconstructs inputs via a predicted transformation, formalized through $x = \rho_X(\psi(x))\delta(\eta(x))$ and a learnable $\psi$. The authors provide a general construction $\psi = \xi \circ \mu$ valid for any group $G$, with explicit instantiations for $SO(2)$, $SO(3)$, $S_n$, $T_n$, and $SE(n)$, and demonstrate the approach across Rotated MNIST, sets, point clouds, molecules, and ShapeNet. Experiments show rotation/translation/permutation invariances in the latent space, improved downstream performance, and meaningful structure in the embeddings. The framework offers a general, scalable path toward unsupervised, symmetry-aware representation learning with broad applicability to 3D geometry, graphs, and sets, and it opens avenues for probabilistic group-action modeling and generative tasks.

Abstract

Equivariant neural networks, whose hidden features transform according to representations of a group G acting on the data, exhibit training efficiency and an improved generalisation performance. In this work, we extend group invariant and equivariant representation learning to the field of unsupervised deep learning. We propose a general learning strategy based on an encoder-decoder framework in which the latent representation is separated in an invariant term and an equivariant group action component. The key idea is that the network learns to encode and decode data to and from a group-invariant representation by additionally learning to predict the appropriate group action to align input and output pose to solve the reconstruction task. We derive the necessary conditions on the equivariant encoder, and we present a construction valid for any G, both discrete and continuous. We describe explicitly our construction for rotations, translations and permutations. We test the validity and the robustness of our approach in a variety of experiments with diverse data types employing different network architectures.

Figures (11)

• Figure 1: a) Schematic of the learning task this work is concerned with. Data points $x\in X$ are encoded to and decoded from latent space $Z$. Points in the same orbit in $X$ are mapped to the same point (orbit) $z\in Z = X/G$. Latent points $z$ are mapped to canonical elements $\hat{x}\in \{\rho_X(g)x|\forall g \in G\}$. b) Schematic of our proposed framework with data points $x$, encoding function $\eta$, decoding function $\delta$, canonical elements $\hat{x}$, group function $\psi$ and group action $g$
• Figure 2: TSNE embedding of the encoded test dataset for a classical and our proposed SO(2) invariant autoencoder.
• Figure 3: Input and predicted output for rotated versions of three MNIST images. Top row shows the input image successively rotated by $45^{\circ}$. Middle row shows the decoded (canonical) image and bottom row shows the decoded image after applying the predicted rotation.
• Figure 4: a) Element-wise reconstruction accuracy of our proposed permutation invariant autoencoder (cross) and a classical non-permutation invariant autoencoder (diamond) for different embedding and set sizes. b) Example set $x$ with 100 elements with its canonical reconstruction $\hat{x}$ and the predicted permutation matrix $P$ (resulting into a perfect reconstruction). One can confirm for oneself that, e.g., $x[38] = x'[0]$, matching $P[38, 0]=1$. c) Best viewed in colour. Visualization of the two-dimensional embedding of a permutation-invariant autoencoder for all 5151 sets of 100 elements with 3 different element classes. Each point represents one set, colours represent set compositions (proportion of each element class, independent of the order).
• Figure 5: a) Five different Tetris shapes represented by points at the center of the four blocks respectively. Input points, output points and rotated (predicted group action) output points as reconstructed by our proposed SE($3$)- and S($N$)-invariant autoencoder are visualized. b) Two-dimensional latent space for all Tetris shapes augmented with Gaussian noise ($\sigma=0.01$). Colors of points match colors of shapes on the right. c) Two molecular conformations and their reconstructions represented as point cloud and ball-and-stick model (left true, right predicted).

Theorems & Definitions (12)

• Proposition 2.3
• Proposition 2.4
• Lemma 2.5
• Proposition 2.7
• Proposition 2.8
• Proposition 2.9
• Proposition A.1
• Proposition A.2
• Lemma A.3
• Proposition A.4