Scaling Spherical CNNs (2306.05420v1)

Abstract: Spherical CNNs generalize CNNs to functions on the sphere, by using spherical convolutions as the main linear operation. The most accurate and efficient way to compute spherical convolutions is in the spectral domain (via the convolution theorem), which is still costlier than the usual planar convolutions. For this reason, applications of spherical CNNs have so far been limited to small problems that can be approached with low model capacity. In this work, we show how spherical CNNs can be scaled for much larger problems. To achieve this, we make critical improvements including novel variants of common model components, an implementation of core operations to exploit hardware accelerator characteristics, and application-specific input representations that exploit the properties of our model. Experiments show our larger spherical CNNs reach state-of-the-art on several targets of the QM9 molecular benchmark, which was previously dominated by equivariant graph neural networks, and achieve competitive performance on multiple weather forecasting tasks. Our code is available at https://github.com/google-research/spherical-cnn.

• The paper demonstrates that scaling spherical CNNs significantly improves computational efficiency by leveraging optimized TPUs with spin-weighted spherical harmonic transforms and DFT matrices.
• It introduces architectural enhancements, including novel phase collapse activations, spectral batch normalization, and efficient residual blocks to boost expressivity and accuracy.
• The study validates spherical CNNs on molecular property prediction and weather forecasting benchmarks, outperforming traditional CNNs and graph neural networks.

Scaling Spherical CNNs: A Rigorous Exploration

This paper addresses the crucial challenge of scaling spherical convolutional neural networks (CNNs), which are designed to operate on data that is inherently spherical, optimizing their applicability to large-scale problems. Spherical CNNs extend the capabilities of traditional planar CNNs to functions defined on the sphere, maintaining desirable properties like equivariance to rotations, spatial weight sharing, and localized filtering. By leveraging these properties, spherical CNNs are poised to unlock new capabilities in domains where data is naturally mapped to spherical representations, such as climate modeling and molecular structure analysis.

Key Contributions

The research provides several advancements to the foundational architecture and implementation of spherical CNNs:

1. Efficient Implementation: Through the development of a spin-weighted spherical harmonic transform fundementally optimized for TPU execution, the authors significantly improve computational efficiency over previous models. The shift from reliance on Fast Fourier Transform (FFT) algorithms to utilizing Discrete Fourier Transform (DFT) matrices exploits the hardware characteristics of TPUs, enhancing execution speed even further.
2. Architectural Enhancements: The paper introduces critical modifications to standard components within the network, such as novel activations using phase collapse mechanisms, spectral batch normalization, and an efficient residual block design that synergistically enhance the expressivity, efficiency, and accuracy of spherical CNNs.
3. Application-specific Modeling: For molecular property prediction with QM9 datasets, the authors present a new representational framework that captures Coulombic and van der Waals interactions through Gaussian smoothed power laws. In the context of weather forecasting, the capability to handle high-dimensional data over Earth's spherical surface is demonstrated through iterative forecasting tasks and benchmarks like WeatherBench, focusing on effective representation of atmospheric dynamics.

Numerical Results and Implications

The paper reports state-of-the-art performance on QM9, achieving superior accuracy over existing models in predicting twelve distinct molecular properties. Furthermore, it shows competitive to superior results against graph neural networks and transformers, traditionally the stronger performers in this area. In weather forecasting tasks, spherical CNNs effectively model atmospheric variables, outperforming traditional CNN approaches on several metrics across multiple datasets, including challenging longer-range forecasts.

Theoretical and Practical Implications

The implications of this research for the machine learning and broader scientific community are substantial:

• Optimization: This work complements the theoretical understanding of spherical CNNs while providing practical methods for their implementation, thereby broadening their applicability. It indicates that a thoughtful design of network components, coupled with an understanding of the underlying computation platform, can yield substantial gains in efficiency without the need to compromise on model capacity.
• Versatility of Spherical Data: By matching or surpassing performance of models specifically tuned to rotational symmetries and complex spatial structures, spherical CNNs validate their role as versatile tools for handling spherical data in scientific and industrial domains.
• Scalability: The proven scalability of the approach implies that future research and applications in fields requiring spherical data processing, like real-time global weather systems, astrophysics, and complex 3D molecular modeling, will benefit from enhanced model fidelity and reduced computational constraints.

Future Directions

The research paves the way for further exploration into the applications of spherical CNNs. Future developments could harness these models in unexplored domains exhibiting spherical symmetries or requiring intrinsic understanding of 3D structures. Additionally, the methods for efficient computation on parallel architectures like TPUs suggest avenues for refining deep learning models tailored to modern hardware capabilities.

Ultimately, the work elucidates the path toward embedding spherical CNNs as a mainstay in machine learning toolkits, ensuring that this class of models finds widespread utility and impacts multiple transformational applications.