Skalierung von sphärischen CNNs

Scaling Spherical CNNs

June 8, 2023

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

cs.AI

Zusammenfassung

Sphärische CNNs verallgemeinern CNNs auf Funktionen auf der Kugel, indem sie sphärische Faltungen als Hauptoperation verwenden. Die genaueste und effizienteste Methode zur Berechnung sphärischer Faltungen erfolgt im Spektralbereich (über das Faltungstheorem), was jedoch immer noch aufwändiger ist als die üblichen planaren Faltungen. Aus diesem Grund waren Anwendungen sphärischer CNNs bisher auf kleine Probleme beschränkt, die mit geringer Modellkapazität angegangen werden können. In dieser Arbeit zeigen wir, wie sphärische CNNs für deutlich größere Probleme skaliert werden können. Um dies zu erreichen, führen wir entscheidende Verbesserungen ein, darunter neuartige Varianten gängiger Modellkomponenten, eine Implementierung der Kernoperationen zur Nutzung der Eigenschaften von Hardwarebeschleunigern und anwendungsspezifische Eingabedarstellungen, die die Eigenschaften unseres Modells ausnutzen. Experimente zeigen, dass unsere größeren sphärischen CNNs auf mehreren Zielen des QM9-Molekülbenchmarks den State-of-the-Art erreichen, der bisher von äquivarianten Graph-Neural-Networks dominiert wurde, und wettbewerbsfähige Leistungen bei mehreren Wettervorhersageaufgaben erzielen. Unser Code ist verfügbar unter 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.