Масштабирование сферических сверточных нейронных сетей

Scaling Spherical CNNs

June 8, 2023

Авторы: Carlos Esteves, Jean-Jacques Slotine, Ameesh Makadia

cs.AI

Аннотация

Сферические сверточные нейронные сети (CNNs) обобщают традиционные CNNs для работы с функциями на сфере, используя сферические свертки в качестве основной линейной операции. Наиболее точный и эффективный способ вычисления сферических сверток — в спектральной области (через теорему о свертке), что всё же требует больше ресурсов, чем обычные плоские свертки. По этой причине применение сферических CNNs до сих пор ограничивалось небольшими задачами, которые можно решать с использованием моделей низкой сложности. В данной работе мы показываем, как сферические CNNs можно масштабировать для решения значительно более крупных задач. Для этого мы вносим ключевые улучшения, включая новые варианты стандартных компонентов моделей, реализацию основных операций с учетом характеристик аппаратных ускорителей, а также специализированные представления входных данных, которые используют свойства нашей модели. Эксперименты показывают, что наши более крупные сферические CNN достигают наилучших результатов на нескольких целевых задачах молекулярного бенчмарка QM9, который ранее доминировали эквивариантные графовые нейронные сети, а также демонстрируют конкурентоспособную производительность на множестве задач прогнозирования погоды. Наш код доступен по адресу 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.