RECON: Robust symmetry discovery via Explicit Canonical Orientation Normalization

TL;DR

Recon is introduced, a class-pose agnostic, test-time canonicalization layer that corrects arbitrary canonicals via a simple right translation, yielding natural, data-aligned canonicalizations.

Abstract

Real world data often exhibits unknown, instance-specific symmetries that rarely exactly match a transformation group $G$ fixed a priori. Class-pose decompositions aim to create disentangled representations by factoring inputs into invariant features and a pose $g\in G$ defined relative to a training-dependent, \emph{arbitrary} canonical representation. We introduce RECON, a class-pose agnostic \emph{canonical orientation normalization} that corrects arbitrary canonicals via a simple right translation, yielding \emph{natural}, data-aligned canonicalizations. This enables (i) unsupervised discovery of instance-specific pose distributions, (ii) detection of out-of-distribution poses and (iii) a plug-and-play \emph{test-time canonicalization layer}. This layer can be attached on top of any pre-trained model to infuse group invariance, improving its performance without retraining. We validate on images and molecular ensembles, demonstrating accurate symmetry discovery, and matching or outperforming other canonicalizations in downstream classification.

Figures (18)

• Figure 1: (a) Class-pose methods assign arbitrary (often out-of-distribution) canonicals per class. (b) This leads to distinct relative-pose distributions $\nu_{[x]}$, obscuring the shared $\pm 30^\circ$ symmetries. recon corrects these offsets, mapping inputs under the same symmetries to the same distribution $\mu_{[x]}$ and extracting their natural pose$\gamma$. (c) Our data-aligned canonicalization removes symmetry-induced distribution shifts, improving downstream performance of pre-trained backbones without architectural restrictions or retraining.
• Figure 2: SGM
• Figure 3: IE-AE
• Figure 4: recon (Ours)
• Figure 6: Problem setup. Left: A class $[x]$ is defined by inputs clustering together in the invariant space $\mathcal{Z}$. Right: We model instances $s \in [x]$ as $\rho_{\mathcal{X}}(g)\gamma_{[x]} + \varepsilon_s$, where $\gamma_{[x]}$ is a reference frame and $g$ is drawn from a distribution over rotation angles $\mu_{[x]}$. The objective is to recover $\mu_{[x]}$ from the unlabeled data.

Theorems & Definitions (5)

• Proposition 3.1: Approximation of $\mu_{[x]}$ via Normalization
• Proposition B.1: Approximation of $\mu_{[x]}$ via Normalization
• proof
• Proposition B.2: Right-translation invariance of the log density
• proof