Moment kernels: a simple and scalable approach for equivariance to rotations and reflections in deep convolutional networks

TL;DR

This paper introduces moment kernels, a simple and scalable method for achieving rotation and reflection equivariance in convolutional neural networks, simplifying previous complex approaches and demonstrating effectiveness in biomedical image analysis tasks.

Contribution

The authors propose moment kernels as a straightforward alternative to complex representation theory-based methods for equivariance, and prove all equivariant kernels must be of this form.

Findings

Moment kernels enable equivariance using standard convolution modules.

Implemented neural networks achieve invariance and equivariance in biomedical tasks.

Effective in classification, registration, and segmentation tasks.

Abstract

The principle of translation equivariance (if an input image is translated an output image should be translated by the same amount), led to the development of convolutional neural networks that revolutionized machine vision. Other symmetries, like rotations and reflections, play a similarly critical role, especially in biomedical image analysis, but exploiting these symmetries has not seen wide adoption. We hypothesize that this is partially due to the mathematical complexity of methods used to exploit these symmetries, which often rely on representation theory, a bespoke concept in differential geometry and group theory. In this work, we show that the same equivariance can be achieved using a simple form of convolution kernels that we call ``moment kernels,'' and prove that all equivariant kernels must take this form. These are a set of radially symmetric functions of a spatial position $x$, multiplied by powers of the components of $x$ or the identity matrix. We implement equivariant neural networks using standard convolution modules, and provide architectures to execute several biomedical image analysis tasks that depend on equivariance principles: classification (outputs are invariant under orthogonal transforms), 3D image registration (outputs transform like a vector), and cell segmentation (quadratic forms defining ellipses transform like a matrix).