Diffeomorphism-Equivariant Neural Networks

TL;DR

This work tackles the challenge of extending neural networks to be diffeomorphism-equivariant under the infinite-dimensional group of diffeomorphisms $\mathcal{D}(\mathcal{X})$. It introduces DiffeoNN, which uses energy-based canonicalisation with stationary velocity fields (SVFs) to map inputs to a canonical form before applying a pretrained inner model and then reversing the transformation to restore equivariance, all without retraining. The canonicalisation energy combines a VAE loss and an adversarial discriminator with a deformation regulariser, and gradient-based optimisation with a SIREN parameterisation and the Scaling-and-Squaring method ensures practical, approximate invariance. Empirical results across synthetic segmentation, lung segmentation, and MNIST genus classification show improved robustness to unseen deformations and competitive performance relative to data-augmented baselines, demonstrating data efficiency and modularity. These findings point to promising theoretical questions about the stability, limits, and generalisation guarantees of energy-based canonicalisation for infinite-dimensional symmetry groups.

Abstract

Incorporating group symmetries via equivariance into neural networks has emerged as a robust approach for overcoming the efficiency and data demands of modern deep learning. While most existing approaches, such as group convolutions and averaging-based methods, focus on compact, finite, or low-dimensional groups with linear actions, this work explores how equivariance can be extended to infinite-dimensional groups. We propose a strategy designed to induce diffeomorphism equivariance in pre-trained neural networks via energy-based canonicalisation. Formulating equivariance as an optimisation problem allows us to access the rich toolbox of already established differentiable image registration methods. Empirical results on segmentation and classification tasks confirm that our approach achieves approximate equivariance and generalises to unseen transformations without relying on extensive data augmentation or retraining.

Figures (12)

• Figure 1: Examples of (a) diffeomorphism equivariance in lung image realData segmentation and (b) diffeomorphism invariance in predictions of topological invariants in MNIST mnist. Applying a diffeomorphism $g \in \mathcal{D}(\mathcal{X})$ to the input should result in a correspondingly transformed segmentation output, i.e., $f_\theta(g \cdot x) = g \cdot f_\theta(x)$, and preserve the classification, i.e., $f_\theta(g \cdot x) = f_\theta(x)$.
• Figure 2: Canonicalisation (DiffeoNN) for lung segmentation. The network $f_\theta$ is pre-trained on a simple chest X-ray/ground-truth segmentation dataset. Applying the trained model $f_\theta$ to a diffeomorphically transformed image without canonicalisation results in an Intersection-over-Union (IoU) of $0.8769$(naïve) with clearly visible errors in the segmentation map. Using canonicalisation, i.e., transforming the input closer to the training dataset before applying the network $f_\theta$ and reversing the transformation, achieves an almost perfect segmentation with an IoU of $0.9560$(ours).
• Figure 3: In the canonicalisation step, an input $x \in \mathcal{X}$ is canonicalised to $x_c = g_x \cdot x$, where $g_x$ is a canonicalising element that is determined by solving a task-specific optimisation problem. Then the inner model $f_\theta$ is applied to $x_c$. Its output $y_{x_c} := f_\theta(x_c)$ is transformed by the reverse canonicalisation$g_x^{-1}$ to obtain the final output $y_x = g_x^{-1} \cdot y_{x_c}$.
• Figure 4: Lung segmentation of diffeomorphically transformed chest X-ray images from realData. Shown are the steps of DiffeoNN (column two and three), as well as the final outputs of DiffeoNN, an augmented U-Net, and the inner U-Net of DiffeoNN without augmentation (naïve) (column four, five, and six). DiffeoNN produces more accurate lung segmentations than the naïve approach. While segmentations of DiffeoNN and augmented U-Net are close to the ground truth, the augmented do show some small artefacts.
• Figure 5: Examples of images $x \in X_E$ and their canonicalised form $x_c$(row one) with corresponding transformed images $g^\prime \cdot x \in X_{TE}$ and their canonicalised form $( g^\prime \cdot x)_c$(row two) with their canonicalisation energies. The experiment is described in Section \ref{['sec:exp_syn']}. The energies before the canonicalisation steps are very different. In contrast, after canonicalising the inputs, the energies are approximately the same. Furthermore, the canonicalised images look very similar to images from $X_E$. This verifies the effectiveness of the canonicalisation step and its invariance empirically.

Theorems & Definitions (3)

• Definition 1.1: Diffeomorphism
• Definition 3.1: Stationary Velocity Field (SVF)
• Definition 3.2: Flow