Equivariance with Learned Canonicalization Functions

• The paper places its approach in the single-view-plus-transformation family, contrasting it with invariant-based and multiple-view approaches.The authors describe the transformation approach as having received less attention in deep learning.Source: Figure 1 · p. 1, §1 · p. 2
• The framework learns mappings to canonical samples and can be plugged into existing non-equivariant architectures.The authors present it as a systematic, end-to-end alternative to hand-engineered canonicalization.Source: §1 · p. 2
• In some settings, the resulting models are universal approximators of equivariant functions.Source: §1 · p. 2
• Experiments across images, N-body systems, and point clouds support the hypothesis that learned canonicalization performs better than predefined heuristics.Source: §1 · p. 2
• The framework offers efficient implementation variants for specific domains.Source: §1 · p. 2

• Spatial Transformer Networks and PointNet learn input transformations or alignments, but provide no equivariance guarantees.Source: §2 · p. 2
• Equivariant Spatial Transformer variants based on canonical coordinates do not exactly handle groups whose dimension exceeds that of the data grid.Source: §2 · p. 2
• The paper reports evidence that canonicalization networks can be made shallower and simpler without affecting performance.Source: §2 · p. 2

3.1 Problem Setting

• For a group action on inputs and outputs, equivariance maps transformed inputs to correspondingly transformed outputs, while invariance maps an orbit to one output.Source: §3.1 · pp. 2–3

3.2 General Formulation

• Canonicalization transforms an input to a canonical sample before prediction and transforms the prediction back; equivariance follows when the canonicalizer is equivariant.The formulation decouples the canonicalization and prediction components, allowing a simple canonicalizer to accompany the prediction function.Source: §3.2 · p. 3

3.3 Partial Canonicalization and Lattice of Subgroups

• Partial canonicalization leaves K-equivariance to the predictor, and smaller K yields more expressive models within the subgroup family.For a semidirect product G ≃ J ⋉ K, the canonicalizer can be J-equivariant and K-invariant.Source: §3.3 · p. 3

3.4 Universality Result

• If the canonicalizers are continuous and the predictor is a universal approximator of K-equivariant functions, the composed model is a universal approximator of G-equivariant functions.Source: §3.4 · pp. 3–4
• A continuous G-equivariant canonicalizer paired with an MLP predictor is universal for G-equivariant functions, without requiring an expressive canonicalization network.Source: §3.4 · p. 4

4 Design of Canonicalization Functions

• The direct approach uses an equivariant network to output a group transformation, while the optimization approach selects a canonical sample by minimizing an input-dependent function.The optimization objective must satisfy an equivariance condition and have a minimum unique up to input symmetry; its energy can be a non-equivariant neural network.Source: §4 · p. 4

4.1 Euclidean Group

• For images, the paper describes discretized Euclidean transformations and an optimization-based G-CNN canonicalizer, simplified to identify orientation by minimizing over the rotation dimension.A shallow network with larger filters addresses the cost and artifacts of finer rotation grids; a straight-through estimator addresses the non-differentiable argmin.Source: §4.1 · p. 5
• For point clouds, the canonicalizer outputs rotation or reflection and translation components, with equivariance constraints that can be met using existing architectures.The rotation/reflection output can be made orthogonal by orthonormalizing vectors with Gram-Schmidt; Deep Sets as a backbone gives a universal approximator of Euclidean-equivariant and permutation-equivariant functions.Source: §4.1 · p. 5

4.2 Symmetric Group

• For permutation symmetry, Sinkhorn-based optimal transport gives an equivariant direct canonicalizer, while sorting and optimal transport can also be framed as optimization.The direct construction yields convex combinations of permutations represented by doubly stochastic matrices.Source: §4.2 · pp. 5–6

5.1 Image classification

• The image experiments compare learned and optimization-based canonicalization with CNN and G-CNN baselines, using a shallow G-CNN to produce canonicalizations for randomly rotated digits.The learned variant uses a CNN prediction function; the optimization variant minimizes an energy over input rotations.Source: §5.1 · p. 6
• On Rotated MNIST, direct canonicalization outperforms the CNN baseline and is comparable to G-CNNs; the optimization variant is weaker, partly because gradient descent can get stuck in flat regions.PCA and frozen-weight canonicalization also improve performance over the CNN baseline.Source: §5.1 · p. 6

5.2 N-body dynamics prediction

• On N-body prediction, the method achieves state-of-the-art results and outperforms intrinsically equivariant baselines in accuracy and efficiency.The task predicts future positions of five charged particles from initial positions and velocities; the prediction GNN was not optimized, and the canonicalization network has around 20 times fewer parameters.Source: Table 3 · p. 7, §5.2 · p. 7

5.1 Image classification

• As rotation order increases, the method’s inference speed stays roughly constant, while its performance improves significantly relative to G-CNN; G-CNN is slightly better for p4.The comparison reduces G-CNN channels to match parameter counts.Source: Figure 3 · p. 6, §5.1 · p. 7

5.2 N-body dynamics prediction

• For N-body dynamics, centroid-based translation canonicalization is only marginally worse than fully learned canonicalization, while learned orthogonal canonicalization significantly outperforms PCA-based Frame Averaging.Learned canonicalization improves performance by close to 50% over frozen canonicalization, which itself improves by more than 20% over the GNN alone.Source: §5.2 · p. 7

5.3 Point cloud classification and segmentation

• For point clouds, the canonicalization network predicts three vectors, orthonormalizes them into a rotation matrix, and uses that matrix to canonicalize the input.The study tests linear and nonlinear Vector Neuron canonicalizers, including a frozen-weight variant, with PointNet and DGCNN predictors.Source: §5.3 · p. 8
• Point-cloud results are competitive across rotation settings, with learning slightly improving performance except when test clouds are already aligned; canonicalization overhead is negligible and inference is faster than Vector Neurons.Shallow linear canonicalization also achieves good results, and the prediction network can bottleneck performance.Source: §5.3 · p. 8, Table 5 · p. 8, Table 6 · p. 9

5.4 Discussion

• Across images, N-body systems, and point clouds, the results support learning canonicalization over predefined heuristics.The authors suggest learned canonicalizations provide consistency for prediction and implicit data augmentation, combining some augmentation advantages with exact equivariance.Source: §5.4 · p. 9

6 Conclusion

• Learned canonicalization functions can be plugged into existing architectures, with described strategies for Euclidean and symmetric groups.Source: §6 · p. 9
• Across image, dynamical-systems, and point-cloud studies, the approach matches or exceeds baselines on invariant tasks, and learned canonicalization outperforms fixed mappings.The conclusion also reports that a shallow equivariant canonical network does not hinder performance.Source: §6 · p. 9
• The method does not guarantee smooth canonicalization, so small input changes may cause large variations in the prediction function’s input.Source: §6 · p. 9
• Canonicalization may also be difficult in domains lacking semantic content to identify a meaningful canonical form, such as some astronomical or biological images.Source: §6 · p. 9

A Symmetric inputs and relaxed equivariance

• Relaxed equivariance is introduced to address the difficulty that symmetric inputs pose for standard equivariance of the canonicalization function.Source: A Symmetric inputs and relaxed… · p. 13
• For relaxed equivariance, a transformed input may correspond to an element in a coset determined by the input’s stabilizer, rather than requiring a single standard-equivariant match.Source: A Symmetric inputs and relaxed… · p. 13
• Relaxed equivariance permits asymmetric outputs from symmetric inputs and can model symmetry breaking or assign different embeddings to nodes in the same orbit.Source: A Symmetric inputs and relaxed… · p. 14
• A translation-equivariant function cannot map a uniform image to a non-uniform image, and a permutation-equivariant function cannot assign different outputs to nodes in the same orbit.These examples illustrate why standard equivariance can exclude desired symmetry-breaking outputs.Source: A Symmetric inputs and relaxed… · p. 14

B Optimization approach to canonicalization

• The optimization approach defines a canonicalization function by minimizing a score over group transformations, with sufficient conditions making the result relaxed-equivariant.One condition constrains the minimizers to a coset associated with the input stabilizer.Source: B Optimization approach to… · p. 15
• The score can be designed using an equivariant function, and for continuous groups convex neural networks can ease optimization.The authors conjecture that suitable neural-network functions satisfy the uniqueness condition almost surely under weak assumptions.Source: B Optimization approach to… · p. 15

C Proof of Theorem 3.1

• Theorem 3.1 establishes equivariance for a general subgroup K under the stated condition, using K-equivariance of the prediction function.The proof invokes a k ∈ K supplied by equation 3.Source: C Proof of Theorem 3.1 · p. 16

D Proof of Theorem 3.2

• For G ≃ K ⋊ J with K normal, each orbit’s canonical sample is obtained from a member x as h(x)^-1 x.Source: D Proof of Theorem 3.2 · p. 16
• When canonicalization is J-equivariant and K-invariant, the proof shows its image can be chosen as ρ(J).It fixes the canonical sample’s K-component by requiring k = e.Source: D Proof of Theorem 3.2 · p. 16

• The proof is inspired by the symmetrization approach of Yarotsky (2022) and Puny et al. (2022).Source: E Proof of Theorem 3.3 · p. 17
• For a compact set K, the proof uses continuity of the inverse action and evaluation maps to establish the needed compact-set construction.The inverse map is continuous in the Banach algebra of linear operators, and the evaluation map is continuous because X is locally compact and Hausdorff.Source: E Proof of Theorem 3.3 · p. 17
• The approximation argument uses universality for K-equivariant functions to approximate G-equivariant functions, choosing δ = ϵ/c for the desired bound.The constant c is the supremum of the norm of h′ on K.Source: E Proof of Theorem 3.3 · p. 17

F Proof of Theorem 5

• Theorem F.1 states that the Gram-Schmidt process is O(n)-equivariant.Source: F Proof of Theorem 5 · p. 17
• For linearly independent input vectors, Gram–Schmidt first constructs orthogonal vectors and then normalizes them to form an orthonormal basis.Source: F Proof of Theorem 5 · p. 17
• Strong induction establishes equivariance of the orthogonal vectors; invariance of the dot product and Euclidean norm extends it to the orthonormal basis.Equivariance of the full process follows because compositions of equivariant functions are equivariant.Source: F Proof of Theorem 5 · p. 18

G.1 Image classification experiments

• Image experiments train a seven-layer CNN and G-CNN for 100 epochs, while canonicalization optimizes an image-derived point-cloud energy with five gradient-descent steps.The CNN uses batch normalization and ReLU, with dropout at layers 4 and 7; the energy network uses Deep Sets and implicit differentiation.Source: §G.1 · p. 18

G.2 N-body dynamics prediction experiments

• For N-body prediction, canonicalization uses a two-layer Vector Neurons Deep Sets network, paired with an Satorras et al. (2021) GNN prediction network.Training uses MSE, Adam, 10,000 epochs, early stopping, weight decay 10^-8, and dropout p = 0.5.Source: §G.2 · pp. 18–19

G.3 Point Cloud Classification and Segmentation experiments

• For point-cloud tasks, the canonicalization network pools three 3D vectors, orthonormalizes them with Gram–Schmidt, and uses PointNet or DGCNN prediction networks.The experiments train with cross entropy and SGD for 200 epochs, using cosine annealing from 0.1 to 0.001.Source: §G.3 · p. 19

H.1 Image classification

• Higher-order rotation equivariance in the canonicalization function improves image-classification performance over adding layers; a shallow network suffices at sufficiently high discrete rotation order.For digit 7, learned p64 canonicalizations align similar-looking samples, while PCA and frozen canonicalization parameters struggle to produce stable mappings; Table 7 compares layer count and rotation order.Source: §H.1 · p. 19, Table 7 · p. 19, Figure 5 · p. 20

I Algorithm for Image Inputs

• The image-input algorithm derives group-element features from a shallow G-CNN, then uses them to select reflection and rotation transformations for the input.Features are pooled across channels and spatial dimensions; the resulting transformation reflects the image when selected and rotates it by the negative predicted angle.Source: I Algorithm for Image Inputs · p. 21