Equivariance with Learned Canonicalization Functions

The paper introduces a framework for building an invariant or equivariant model/function approximator by learning a canonicalization function. Instead of forcing the whole prediction network to be equivariant by architectural design, the model first maps the input to a canonical representative, applies a prediction network there, and then transforms the output back if needed.

Problem Setting

The authors consider inputs $x \in \mathcal{X}$ and outputs $y \in \mathcal{Y}$, where both spaces are finite-dimensional vector spaces. A symmetry group $G$ acts on the input space through a representation $\rho : G \to \mathrm{GL}(\mathcal{X})$ and on the output space through another representation $\rho' : G \to \mathrm{GL}(\mathcal{Y})$.

The goal is to learn a function $\phi : \mathcal{X} \to \mathcal{Y}$, that is equivariant. A function is $G$-equivariant if transforming the input by $g$ causes the output to transform consistently:

$$\phi(\rho(g)x)=\rho'(g)\phi(x).$$

If $\rho'(g)=I$, then the output does not transform, and the function is invariant.
The orbit of an input $x$ is the set of all transformed versions of $x$:

$$\rho(G)x = \{\rho(g)x \mid g \in G\}$$

Canonicalization means choosing one representative from this orbit and mapping all equivalent inputs to that representative.

General Canonicalization Framework

The main construction writes an equivariant function as

$$\phi(x)=h'(x)f(h(x)^{-1}x).$$

where

• $h(x)$ is the canonicalization function.
• $h(x)^{-1}x$ is the canonicalized input.
• $f$ is the prediction function applied in canonical space.
• $h'(x)$ transforms the output back to the correct pose/frame.

Conceptually, this means that the input $x$ is canonicalized, then the predictor function is applied, then the output is transformed back from the canonical space.

For invariant tasks, such as classification, the final transformation $h'(x)$ is unnecessary because the output should not change under the group action. In that case,

$$h'(x)=I$$

and the model reduces to

$$\phi(x)=f(h(x)^{-1}x).$$

Partial Canonicalization

The paper then generalizes the framework by allowing only part of the symmetry to be handled by canonicalization. Let $K \leq G$ be a subgroup. The prediction function $f$ is assumed to be $K$-equivariant, while the canonicalizer handles the remaining symmetry.
In other words, if the predictor that is used already handles part of the symmetry naturally, the canonicalization only needs to handle the remaining symmetry.

For example, if we want to design an architecture that is equivariant to rotations and translations, we can use a canonicalizer with MLP. However, a CNN is already translation-equivariant, so the model can let the CNN handle translations and use the canonicalizer only for rotations/reflections.

Design of Canonicalization Functions

The paper presents two main approaches.

Direct Approach

In the direct approach, the canonicalization function is an equivariant neural network that directly outputs a group element:

$$h(x) \in \rho(G).$$

For example, if the symmetry group is rotations, the network outputs a rotation. If the group is permutations, it outputs a permutation or a relaxed permutation. If the group is Euclidean transformations, it outputs a rotation/reflection and translation.

The direct approach is simple conceptually, but it requires the canonicalization network itself to be equivariant.

Optimization Approach

In the optimization approach, the canonicalization function does not directly predict the transformation. Instead, it searches for the transformation that makes the input look most canonical.

The model defines an energy function $E$ that measures how “canonical” a transformed input looks. Then it chooses the transformation that gives the lowest energy:

$$h(x)=\arg\min_{\rho(g)} E(\rho(g)^{-1}x)$$

Conceptually, this is like mentally rotating an object until it appears upright or standard. The advantage is that the energy function itself can be ordinary; the structure comes from evaluating it on transformed versions of the input.