Scalable and Equivariant Spherical CNNs by Discrete-Continuous (DISCO) Convolutions

TL;DR

The paper tackles the challenge of building spherical CNNs that are simultaneously rotationally equivariant and scalable to high-resolution data. It introduces DISCO, a discrete-continuous group convolution that keeps the filter and group actions continuous while discretizing input signals, yielding near-exact equivariance with linear scaling via sparse tensors. Key contributions include the DISCO Group Convolution, the DISCO Spherical Convolution (and its transposed variant for dense predictions), theoretical support for $\text{SO}(3)$ equivariance (and asymptotic equivariance on $\text{SO}(3)/\text{SO}(2)$), memory- and gradient-optimization strategies, and state-of-the-art results on semantic segmentation and depth estimation on sphere-based benchmarks. Practically, this enables scalable, high-resolution, equivariant spherical CNNs suitable for panoramic vision and cosmology tasks, with substantial computational and memory savings over prior methods.

Abstract

No existing spherical convolutional neural network (CNN) framework is both computationally scalable and rotationally equivariant. Continuous approaches capture rotational equivariance but are often prohibitively computationally demanding. Discrete approaches offer more favorable computational performance but at the cost of equivariance. We develop a hybrid discrete-continuous (DISCO) group convolution that is simultaneously equivariant and computationally scalable to high-resolution. While our framework can be applied to any compact group, we specialize to the sphere. Our DISCO spherical convolutions exhibit $\text{SO}(3)$ rotational equivariance, where $\text{SO}(n)$ is the special orthogonal group representing rotations in $n$-dimensions. When restricting rotations of the convolution to the quotient space $\text{SO}(3)/\text{SO}(2)$ for further computational enhancements, we recover a form of asymptotic $\text{SO}(3)$ rotational equivariance. Through a sparse tensor implementation we achieve linear scaling in number of pixels on the sphere for both computational cost and memory usage. For 4k spherical images we realize a saving of $10^9$ in computational cost and $10^4$ in memory usage when compared to the most efficient alternative equivariant spherical convolution. We apply the DISCO spherical CNN framework to a number of benchmark dense-prediction problems on the sphere, such as semantic segmentation and depth estimation, on all of which we achieve the state-of-the-art performance.

Figures (7)

• Figure 1: Spherical CNN categorization.
• Figure 2: Comparison of computational cost and memory usage of the DISCO spherical convolution and the most efficient alternative equivariant spherical convolution (an axisymmetric harmonic convolution). The DISCO spherical convolution achieves linear scaling in the number of spherical pixels, i.e. $O(L^2)$, for both computational cost and memory usage.
• Figure 3: ${\text{SO}(3)}$ and ${\text{SO}(3)} / {\text{SO}(2)}$ rotations.
• Figure 4: DISCO spherical convolutional residual U-Net architecture.
• Figure 5: Example predictions for semantic segmentation of 2D3DS data.