Table of Contents
Fetching ...

Adaptive Canonicalization with Application to Invariant Anisotropic Geometric Networks

Ya-Wei Eileen Lin, Ron Levie

TL;DR

This work introduces adaptive canonicalization, a framework where the canonical form of an input is chosen in a network- and input-dependent way to respect symmetries while ensuring continuity and universal approximation. Through prior maximization, the method selects transformations that maximize predictive confidence, yielding end-to-end models that are both symmetry-preserving and continuous. The authors develop theory establishing equicontinuity and UAT for adaptive canonicalized functions, and demonstrate two concrete applications: resolving eigenbasis ambiguities in spectral graph neural networks and handling rotational symmetries in point clouds. Empirically, adaptive canonicalization with anisotropic nonlinear spectral filters and anisotropic point-cloud networks outperforms data augmentation, fixed canonicalization, and equivariant architectures across toy tasks, graph benchmarks, and large-scale molecular datasets. The approach offers a flexible, efficient alternative to heavy group-specific designs with broad applicability beyond the tested domains.

Abstract

Canonicalization is a widely used strategy in equivariant machine learning, enforcing symmetry in neural networks by mapping each input to a standard form. Yet, it often introduces discontinuities that can affect stability during training, limit generalization, and complicate universal approximation theorems. In this paper, we address this by introducing adaptive canonicalization, a general framework in which the canonicalization depends both on the input and the network. Specifically, we present the adaptive canonicalization based on prior maximization, where the standard form of the input is chosen to maximize the predictive confidence of the network. We prove that this construction yields continuous and symmetry-respecting models that admit universal approximation properties. We propose two applications of our setting: (i) resolving eigenbasis ambiguities in spectral graph neural networks, and (ii) handling rotational symmetries in point clouds. We empirically validate our methods on molecular and protein classification, as well as point cloud classification tasks. Our adaptive canonicalization outperforms the three other common solutions to equivariant machine learning: data augmentation, standard canonicalization, and equivariant architectures.

Adaptive Canonicalization with Application to Invariant Anisotropic Geometric Networks

TL;DR

This work introduces adaptive canonicalization, a framework where the canonical form of an input is chosen in a network- and input-dependent way to respect symmetries while ensuring continuity and universal approximation. Through prior maximization, the method selects transformations that maximize predictive confidence, yielding end-to-end models that are both symmetry-preserving and continuous. The authors develop theory establishing equicontinuity and UAT for adaptive canonicalized functions, and demonstrate two concrete applications: resolving eigenbasis ambiguities in spectral graph neural networks and handling rotational symmetries in point clouds. Empirically, adaptive canonicalization with anisotropic nonlinear spectral filters and anisotropic point-cloud networks outperforms data augmentation, fixed canonicalization, and equivariant architectures across toy tasks, graph benchmarks, and large-scale molecular datasets. The approach offers a flexible, efficient alternative to heavy group-specific designs with broad applicability beyond the tested domains.

Abstract

Canonicalization is a widely used strategy in equivariant machine learning, enforcing symmetry in neural networks by mapping each input to a standard form. Yet, it often introduces discontinuities that can affect stability during training, limit generalization, and complicate universal approximation theorems. In this paper, we address this by introducing adaptive canonicalization, a general framework in which the canonicalization depends both on the input and the network. Specifically, we present the adaptive canonicalization based on prior maximization, where the standard form of the input is chosen to maximize the predictive confidence of the network. We prove that this construction yields continuous and symmetry-respecting models that admit universal approximation properties. We propose two applications of our setting: (i) resolving eigenbasis ambiguities in spectral graph neural networks, and (ii) handling rotational symmetries in point clouds. We empirically validate our methods on molecular and protein classification, as well as point cloud classification tasks. Our adaptive canonicalization outperforms the three other common solutions to equivariant machine learning: data augmentation, standard canonicalization, and equivariant architectures.

Paper Structure

This paper contains 87 sections, 10 theorems, 48 equations, 4 figures, 13 tables, 1 algorithm.

Key Result

Theorem 6

Let $\mathcal{N}(\mathcal{K},\mathbb{R}^D)$ be a universal approximator of $C_0(\mathcal{K},\mathbb{R}^D)$, and $f\circ \rho_f$ an adaptive canonicalized continuous function. Then, for every $\epsilon>0$, there exists a network $\theta\in \mathcal{N}(\mathcal{K},\mathbb{R}^D)$ such that for every $g

Figures (4)

  • Figure 1: Illustration of prior maximization adaptive canonicalization in classification. The adaptive canonicalization optimizes the transformations $\beta_{x,\Psi_j}$ of the inputs $x$ to the classifiers $\Psi_j$, while, during training, $\Psi_j$ are simultaneously trained w.r.t. the adaptively canonicalized inputs $\pi(\beta_{x,\Psi_j})x$.
  • Figure 2: Hyperparameter sensitivity with respect to grid size, noise level, and hidden dimension.
  • Figure 3: Mean geodesic distance on $\mathcal{SO}(3)$ between the canonicalizations between consecutive epochs.
  • Figure 4: The canonicalized point clouds for the chair class.

Theorems & Definitions (26)

  • Definition 1
  • Definition 2
  • Definition 3
  • Definition 4
  • Definition 5
  • Theorem 6: Universal approximation of adaptive canonicalized functions
  • proof
  • Definition 7
  • Theorem 8
  • proof
  • ...and 16 more