[论文解读] Scalable and Equivariant Spherical CNNs by Discrete-Continuous (DISCO) Convolutions

Jeremy Ocampo, Matt A. Price+1 more|arXiv (Cornell University)|Sep 27, 2022

Medical Imaging and Analysis被引用 5

一句话总结

本文提出了一种用于球面卷积神经网络的 DISCO（离散-连续）组卷积，实现了计算可扩展性与 SO(3) 旋转等变性，解决了现有方法中的关键局限。通过将离散采样与通过稀疏张量操作实现的连续对称性保持相结合，DISCO 实现了计算和内存的线性缩放，在高分辨率球面分割和深度估计基准测试中优于先前方法。

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 $ ext{SO}(3)$ rotational equivariance, where $ ext{SO}(n)$ is the special orthogonal group representing rotations in $n$-dimensions. When restricting rotations of the convolution to the quotient space $ ext{SO}(3)/ ext{SO}(2)$ for further computational enhancements, we recover a form of asymptotic $ ext{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.

研究动机与目标

• 解决缺乏同时具备计算可扩展性与旋转等变性的球面卷积神经网络框架的问题。
• 克服连续球面卷积神经网络的计算不可行性以及离散方法的等变性失效问题。
• 在球面数据上支持高分辨率密集预测任务（例如语义分割、深度估计）。
• 开发基于稀疏张量的实现，使其计算成本和内存使用与球面像素数量呈线性关系。
• 在包括 2D3DS、Omni-SYNTHIA 和 Pano3D 在内的基准测试中达到最先进性能。

提出的方法

• 提出一种混合离散-连续（DISCO）组卷积，既保持 SO(3) 旋转等变性，又实现高效计算。
• 通过在商空间 SO(3)/SO(2) 上定义 DISCO 卷积，降低计算成本，同时保持渐近的 SO(3) 等变性。
• 使用稀疏张量表示实现框架，实现计算成本和内存使用量的线性缩放。
• 利用快速球面谐波变换和采样定理，在离散网格框架内保持对连续对称性的感知。
• 为密集预测任务设计转置 DISCO 卷积，支持高分辨率输入与输出。
• 在 TensorFlow 中使用自定义梯度，确保端到端可微性与硬件加速兼容性（GPU/TPU）。

实验结果

研究问题

• RQ1 能否设计一种球面卷积神经网络框架，同时实现计算可扩展性与精确的 SO(3) 旋转等变性？
• RQ2 如何将球面上的离散采样与连续对称性结合，以在不产生禁止性计算成本的前提下保持等变性？
• RQ3 当将卷积限制在商空间 SO(3)/SO(2) 时，DISCO 框架在多大程度上仍能保持等变性？
• RQ4 与现有球面卷积神经网络相比，DISCO 框架在推理成本、内存使用和密集预测任务性能方面表现如何？
• RQ5 DISCO 框架能否在高分辨率球面基准（如 2D3DS、Omni-SYNTHIA 和 Pano3D）上实现最先进性能？

主要发现

• DISCO 球面卷积神经网络在保持 SO(3) 旋转等变性的同时，计算成本与内存使用量均与球面像素数量呈线性关系，相比最高效的先前等变方法，计算成本降低 10^9 倍，内存使用降低 10^4 倍。
• 在球面 MNIST 基准测试中，DISCO-Axisymmetric 达到 98.6% 的测试准确率，优于平面 CNN（98.3%）和非等变离散方法。
• 在 2D3DS 室内 360° 数据集上，DISCO-Directional-Aug 达到 45.7% 的 mIoU 和 62.7% 的 mAcc，超越所有先前方法，包括 CubeNet（45.0% mIoU）。
• 在 Omni-SYNTHIA 户外数据集上，DISCO-Separable 和 DISCO-Separable-Aug 分别达到 48.3% 和 49.2% 的 mIoU，创下新的最先进性能记录。
• 在 Pano3D 深度估计基准测试中，DISCO-Directional 在 9 项指标中有 6 项位列顶尖，仅使用 658k 参数（比先前最先进模型少 40 倍），且未使用表面法线监督。
• DISCO-Directional 架构优于 DISCO-Axisymmetric，因其具备方向感知的定向滤波器与渐近的 SO(3) 等变性。