Equivariance by Contrast: Identifiable Equivariant Embeddings from Unlabeled Finite Group Actions

AI summary

Overview

Research area: Machine learning — self-supervised learning, equivariant representation learning, nonlinear Independent Component Analysis (ICA), and group theory.

Technical level: Advanced. The paper leans on group representations, nonlinear ICA identifiability results, and contrastive learning theory; readers without that background will find Sections 2 and 3 dense.

Scope: The paper introduces Equivariance by Contrast (EbC), an encoder-only method that learns an equivariant embedding and an implicit group representation from paired observations related by unknown finite-group actions, with a proof of identifiability and empirical validation on synthetic groups (SO(n), O(n), GL(n)) and the infinite dSprites dataset.

What This Paper Is About

Observations in vision, biology, and neuroscience often come in pairs related by a structured transformation (a rotation, a translation, a perturbation), but the transformation itself is usually unknown. The paper asks whether a model can recover both a latent space and a linear representation of the underlying group purely from such paired observations, without group-specific architectural design. The stated goal is to learn an encoder φ and a representation R′ satisfying φ(g·y) = R′(g) φ(y), with a formal guarantee that this solution is identifiable up to a linear transformation.

Key Contributions

1. EbC algorithm. An encoder-only method that jointly learns the feature encoder and an implicit group representation by contrastive learning on pairs (y, g·y), with no generative model and no group-specific inductive biases. The group representation is estimated non-parametrically as R̂(X, X′) = (XᵀX)⁻¹(XᵀX′), the least-squares solution to min_{R∈GL(d)} ||X′ − XRᵀ||_F².

2. Identifiability theorem. Theorem 1 shows that under diversity conditions (building on Roeder et al. [40]), the learned composition h := φ∘f recovers h(x) = Lx up to an ambiguity L ∈ GL(d), and recovers a representation R̂(h(X), h(gX)) = L R(g) L⁻¹; Corollary 1 converts this into an equivariance statement h(gx) = g h(x).

3. Content/style factorization. An optional structured parametrization constraining the minimization to diag(GL_n, I_m) ⊂ GL(m+n), producing an n-dimensional equivariant subspace and an m-dimensional invariant (content) subspace within a single contrastive loss rather than multiple losses.

4. First general-purpose demonstration. The authors describe their results as the first successful demonstration of encoder-only equivariant learning from group action observations alone, including non-trivial non-abelian groups and a product group motivated by modeling affine equivariances in computer vision.

Main Findings

• Near-perfect latent recovery on synthetic data. EbC reaches an R² of >99% in recovering the ground-truth latent space on synthetic data, substantially outperforming a linear baseline (60–70%). Reported values include R²(x) of 99.7 ± 0.22 (SO₃), 99.8 ± 0.05 (O₃), and 99.8 ± 0.03 (GL₃).

• Faithful group representations. The R²(G) metric measuring recovery of the group structure exceeds 98% on synthetic data (99.7 ± 0.25 for SO₃, 99.7 ± 0.04 for O₃, 99.7 ± 0.06 for GL₃), versus 0.0 for the InfoNCE, LDS, and SLDS baselines.

• Baselines only learn invariants. InfoNCE, +LDS, and +SLDS recover content well (typically >98% accuracy) but show essentially zero R²(x) and R²(G), confirming they capture invariance without equivariance.

• Content/style trade-off on idSprites. On the product group R_m × ℤ_n × ℤ_n, EbC reaches Acc(G, 5) of 99.91 ± 0.05 while content accuracy drops to 74.04 ± 1.91, against roughly 99.8–99.97% content accuracy for the invariant baselines.

• Learned latent geometry. The optimal embedding for R_m × ℤ_n × ℤ_n is described as a 3-torus with the three component groups along circular coordinates; the reconstructed EbC embedding is reported to be qualitatively stable under variation of content (shape and size).

• Robustness to over-parameterization. When the encoder output dimension is varied (2 dimensions for content, 4–7 for the group), group structure preserves a kNN above 99% and content classification stays above 80%, though 4× the group dimensionality is required when single transitions are observed and 8× for compound actions.

• Limit on mixing complexity. The R²(x) metric confirms quality of the mapping up to n=4 mixing layers, with degradation beyond that; forward prediction through the group representation stays high up to 4 mixing layers.

• Observable hyperparameter selection. The Acc(G) metric, which requires no ground-truth latents, shows a clear peak at the correct latent dimensionality (d=3) in the misspecification study, making it usable for model selection. Misspecifying the content dimension is less critical, except for a degenerate case at d=1.

• Symmetric loss helps on harder groups. Table 2 shows that on SO_n, O_n, and GL_n, using both y and y′ as negative samples is crucial for GL(n); on O(n) and GL(n) a symmetric loss additionally improves performance (e.g., GL_n with symmetric loss and both negatives: Acc(G,1) 99.98 ± 0.01, Acc(G,5) 99.93 ± 0.16).

• Scaling gap between theory and practice. The theory mandates d+1 examples for estimating the embedding space; empirically substantially more are needed, with roughly 6–8× giving good results through the practical criterion, and the correct dimensionality peak visible in Acc(G).

Methodology in Plain English

The data arrive in batches of n+1 paired samples {(yᵢ, y′ᵢ)}, where every pair in a batch is assumed to be related by the same, unknown group element g. An encoder φ maps each observation into a latent space. From the batch, the method estimates the matrix that best linearly maps the encoded first members onto the encoded second members — this matrix is the estimate of the group action's representation. That matrix is then applied to a query sample y, and the model must pick the correct transformed sample y′ out of a set S that mixes the positive with randomly drawn negatives.

This objective is written as a softmax likelihood over negative-sample distances (Eq. 3), giving a loss closely related to InfoNCE but with the extra structure needed for group learning. Training minimizes the negative log-likelihood over all pairs and uniformly sampled negatives. Because the group representation itself is defined entirely through the encoder, the whole model reduces to learning a single feature map.

For content/style separation, the least-squares fit is constrained to block-diagonal matrices with GL_n in the equivariant block and an identity in the invariant block, so the embedding splits into a part the group acts on and a part it does not.

Experiments use a three-layer MLP with 512 hidden units for φ, the gels least-squares solver in PyTorch, 84 sample pairs for idSprites and 12 for synthetic data, 1024 positive and 16k negative samples per batch, Adam at learning rate 10⁻³ for 20k steps, and 80/10/10 train/valid/test splits. Results are reported across three dataset seeds and three model seeds, with metrics computed out-of-sample. Synthetic data use a random 3-layer MLP mixer followed by a random linear map into 50 dimensions, up to 1000 sampled matrices R(g) and 1M pairs, with n=3, d=3 and |C|=100 unless stated otherwise. Baselines use the identical training protocol.

Why This Matters

The paper shifts equivariant representation learning away from hand-designed, group-specific architectures and toward a general contrastive objective with an identifiability guarantee. Where prior contrastive methods such as CARE restrict to orthogonal representations on the hypersphere, STL allows nonlinear latent relations, and the neural Fourier transform requires a generative model, EbC is encoder-only and learns general linear representations. It also demonstrates the approach on non-abelian groups (O(n), GL(n)) and on a product group representing combined rotation and translation, which is the kind of structure that appears in real vision problems.

Potential real-world applications:

• Computer vision. Models that must track a scene under rotations and translations, where affine equivariance is the natural structure, as suggested by the R_m × ℤ_n × ℤ_n setup on infinite dSprites.
• Single-cell biology. The introduction cites single-cell transcriptomic datasets with observations of cells before and after perturbations such as gene knockouts or pharmacological intervention, where the perturbation is the unknown group action.
• Neuroscience. Neural activity recorded across changing brain states under sensory input or behavioral output, where paired observations relate through a shared latent transformation.
• Time-series data. The paper notes that pairs can be derived from time series under a "slowness prior" assuming nearby points share a common action.

Industry relevance: The method requires no labeling of transformation parameters and no architecture engineering per symmetry, which lowers the cost of building transformation-aware representations for domains where the symmetry group is unknown or composite. The Acc(G) metric is available without ground-truth latents, making practical model selection possible in deployment settings. The authors' own caveat is that broad evaluation on real-world data remains future work.

Future Directions

• Extending evaluation to full-scale image datasets such as 3DIdent, and conducting more baseline comparisons, which the authors say is limited by the lack of extensive benchmarks; additional real-world experiments are reported in Appendix C.5.
• Improving the content/style subspace prior so that a clear separation can be expected beyond the case of minimal group dimensionality, possibly with theoretical guarantees on optimal separation.
• Extending the theory to bound behavior under limited samples, since the current guarantee mandates only d+1 examples while practice requires roughly 6–8× more; replacing the out-of-the-box least squares estimator with more stable or better-adapted algorithms is suggested to close this gap.
• Addressing the accuracy trade-off for large numbers of content classes, where EbC continues to estimate the group representation but content classification degrades once class dimensions reach multiple hundreds or more, possibly because it must implicitly learn a prototype per class.

Target Audience

Researchers and graduate students working on self-supervised learning, equivariant and group-equivariant neural networks, identifiable representation learning, or nonlinear ICA. It is most useful to readers who already understand group representations and contrastive objectives and want a method that avoids group-specific architectural priors; practitioners in vision, computational biology, or neuroscience interested in learning symmetries from paired data without labeled transformations will also find the experimental protocol and the Acc(G) model-selection criterion directly relevant.

Authors’ abstract

We propose Equivariance by Contrast (EbC) to learn equivariant embeddings from observation pairs $(\mathbf{y}, g \cdot \mathbf{y})$, where $g$ is drawn from a finite group acting on the data. Our method jointly learns a latent space and a group representation in which group actions correspond to invertible linear maps -- without relying on group-specific inductive biases. We validate our approach on the infinite dSprites dataset with structured transformations defined by the finite group $G:= (R_m \times \mathbb{Z}_n \times \mathbb{Z}_n)$, combining discrete rotations and periodic translations. The resulting embeddings exhibit high-fidelity equivariance, with group operations faithfully reproduced in latent space. On synthetic data, we further validate the approach on the non-abelian orthogonal group $O(n)$ and the general linear group $GL(n)$. We also provide a theoretical proof for identifiability. While broad evaluation across diverse group types on real-world data remains future work, our results constitute the first successful demonstration of general-purpose encoder-only equivariant learning from group action observations alone, including non-trivial non-abelian groups and a product group motivated by modeling affine equivariances in computer vision.