Schaalvergroting van Sferische CNN's

Scaling Spherical CNNs

June 8, 2023

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

cs.AI

Samenvatting

Sferische CNN's generaliseren CNN's naar functies op de bol, door sferische convoluties als de belangrijkste lineaire bewerking te gebruiken. De meest nauwkeurige en efficiënte manier om sferische convoluties te berekenen is in het spectrale domein (via het convolutietheorema), wat nog steeds duurder is dan de gebruikelijke vlakke convoluties. Om deze reden zijn toepassingen van sferische CNN's tot nu toe beperkt gebleven tot kleine problemen die met een lage modelcapaciteit kunnen worden aangepakt. In dit werk laten we zien hoe sferische CNN's kunnen worden opgeschaald voor veel grotere problemen. Om dit te bereiken, maken we cruciale verbeteringen, waaronder nieuwe varianten van veelgebruikte modelcomponenten, een implementatie van kernbewerkingen om de kenmerken van hardwareversnellers te benutten, en toepassingsspecifieke invoerrepresentaties die de eigenschappen van ons model uitbuiten. Experimenten tonen aan dat onze grotere sferische CNN's state-of-the-art bereiken op verschillende doelen van de QM9-moleculaire benchmark, die voorheen werd gedomineerd door equivariante grafische neurale netwerken, en competitieve prestaties leveren op meerdere weersvoorspellingstaken. Onze code is beschikbaar op 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.