Towards Diffeomorphism-Equivariant Neural Networks via Canonicalization

Incorporating group symmetries via equivariance into neural networks has emerged as a robust approach for improving efficiency and overcoming the data requirements 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 canonicalization. Formulating equivariance as an optimization 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 generalizes to unseen transformations without relying on extensive data augmentation or retraining.

In the real world, data is rarely homogeneous: objects can appear stretched, bent, or smoothly deformed due to changes in viewpoint, motion, or natural variation. While humans can usually recognize the same object and corresponding symmetries despite these changes, neural networks often struggle when images are distorted in ways they did not encounter during training. A common solution is to retrain models on much additional data, often synthetically altered, but this can require large amounts of augmented data and computational resources. We investigate whether pre-trained neural networks can be made robust to such changes without retraining them from scratch. Our core idea is to introduce a canonicalization step that transforms an input into a standardized (canonical) form before applying the network. We focus on the transformation group of diffeomorphisms, naturally occurring smooth transformations that bend or deform an image without tearing or breaking it, and formulate the canonicalization step as an optimization problem. Our results show that this approach improves the consistency of image classification and segmentation under previously unseen distortions. This could make neural networks more reliable and less dependent on the training datasets.