Mise à l'échelle des réseaux de neurones sphériques

Scaling Spherical CNNs

June 8, 2023

Auteurs: Carlos Esteves, Jean-Jacques Slotine, Ameesh Makadia

cs.AI

Résumé

Les réseaux de neurones convolutifs sphériques généralisent les CNN classiques aux fonctions définies sur la sphère, en utilisant des convolutions sphériques comme opération linéaire principale. La méthode la plus précise et efficace pour calculer ces convolutions sphériques se fait dans le domaine spectral (via le théorème de convolution), ce qui reste plus coûteux que les convolutions planaires usuelles. Pour cette raison, les applications des CNN sphériques ont jusqu'à présent été limitées à des problèmes de petite taille, abordables avec une faible capacité de modèle. Dans ce travail, nous montrons comment les CNN sphériques peuvent être mis à l'échelle pour des problèmes bien plus vastes. Pour y parvenir, nous apportons des améliorations cruciales, notamment des variantes innovantes de composants de modèle courants, une implémentation des opérations de base exploitant les caractéristiques des accélérateurs matériels, et des représentations d'entrée spécifiques à l'application qui tirent parti des propriétés de notre modèle. Les expériences montrent que nos CNN sphériques à plus grande échelle atteignent l'état de l'art sur plusieurs cibles du benchmark moléculaire QM9, précédemment dominé par les réseaux de neurones graphiques équivariants, et obtiennent des performances compétitives sur plusieurs tâches de prévision météorologique. Notre code est disponible à l'adresse https://github.com/google-research/spherical-cnn.

English

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.