Nonlinearities in Steerable SO(2)-Equivariant CNNs

TL;DR

This paper analyzes how nonlinearities affect steerable SO(2)-equivariant CNNs and introduces an FFT-based method to maintain equivariance, enabling improved symmetry-preserving neural network architectures for 2D and 3D data.

Contribution

It develops a novel FFT-based algorithm to compute Fourier representations of nonlinear activations while preserving equivariance, enhancing steerable CNN design.

Findings

Exact equivariance for polynomial nonlinearities.

Approximate equivariance with tunable accuracy for general functions.

Competitive accuracy on 2D and 3D symmetry-preserving tasks.

Abstract

Invariance under symmetry is an important problem in machine learning. Our paper looks specifically at equivariant neural networks where transformations of inputs yield homomorphic transformations of outputs. Here, steerable CNNs have emerged as the standard solution. An inherent problem of steerable representations is that general nonlinear layers break equivariance, thus restricting architectural choices. Our paper applies harmonic distortion analysis to illuminate the effect of nonlinearities on Fourier representations of SO(2). We develop a novel FFT-based algorithm for computing representations of non-linearly transformed activations while maintaining band-limitation. It yields exact equivariance for polynomial (approximations of) nonlinearities, as well as approximate solutions with tunable accuracy for general functions. We apply the approach to build a fully E(3)-equivariant network for sampled 3D surface data. In experiments with 2D and 3D data, we obtain results that compare favorably to the state-of-the-art in terms of accuracy while permitting continuous symmetry and exact equivariance.