Google Researchers Introduce An Open-Source Library in JAX for Deep Learning on Spherical Surfaces

Deep learning, a machine learning subset, automatically learns complex representations from the input. Its applications are used in many fields, such as image and speech recognition for language processing, object detection, and medical imaging diagnostics; finance for algorithmic trading and fraud detection; autonomous vehicles using convolutional neural networks for real-time decision-making; and recommendation systems for personalized content. 

Convolutional neural networks (CNNs) and vision transformers (ViT), two examples of deep learning models for computer vision, analyze signals by assuming planar (flat) regions. Digital photographs, for example, are presented as a grid of pixels on a flat surface. Nonetheless, this data type represents only a fraction of the diverse data encountered in scientific applications.

However, a few things could be improved by processing spherical signals using a planar approach. First, there is a sampling issue, meaning it is impossible to define uniform grids on the sphere—necessary for planar CNNs and ViTs—without significant distortion. Second, rotations frequently confuse signals and local patterns on the sphere.  To ensure that the model learns the features accurately, we need equivariance to 3D rotations. As a result, the model parameters are used more effectively, and training with less data is possible.

Intuitively, both molecular property prediction and climate forecasting problems should benefit from spherical CNNs. The intrinsic properties of molecules are invariant to rotations of the 3D structure (atom positions), so rotation equivariant representations would provide a natural way to encode this symmetry.

Consequently, the researchers have formulated an open-source library in JAX for deep learning on spherical surfaces. It outperforms state-of-the-art results on benchmarks for molecular property prediction and weather forecasting, typically handled by transformers and graph neural networks.

The researchers highlighted that these can solve both the problems of sampling and of robustness to rotation. It does by leveraging spherical convolution and cross-correlation operations. Spherical CNNs offer promising applications in two critical domains: medical research and climate analysis, holding the potential to catalyze transformative advancements for society.

Spherical CNNs present a theoretical advantage in addressing challenges related to predicting chemical properties and understanding climate states. Leveraging rotation-equivariant representations becomes particularly logical in capturing the inherent symmetries of molecular structures, where the properties remain invariant to 3D rotations (atom locations).

Since atmospheric data is naturally displayed on a sphere, spherical CNNs are well suited for this task. They can also effectively manage repeated patterns in such data at various places and orientations.

The researchers said that their models exceed or match neural weather models based on traditional CNNs on a number of weather forecasting benchmarks. The model forecasts the values of several atmospheric variables six hours in advance, and the results from a test environment are shown below. Then, the model is further evaluated up to five days in advance during training and makes predictions up to three days in advance.

Additionally, the models exhibit exceptional performance across various weather forecasting scenarios, demonstrating the effectiveness of spherical CNNs as neural weather models in a ground-breaking accomplishment. This study outlines the best strategies for scaling spherical CNNs and provides real data to support their applicability in these particular applications.

Check out the Paper, Github, and Google Blog. All Credit For This Research Goes To the Researchers on This Project. Also, don’t forget to join our 31k+ ML SubReddit, 40k+ Facebook Community, Discord Channel, and Email Newsletter, where we share the latest AI research news, cool AI projects, and more.

If you like our work, you will love our newsletter..

We are also on WhatsApp. Join our AI Channel on Whatsapp..