[論文レビュー] Scaling Spherical CNNs

Carlos Esteves, Jean-Jacques Slotine+1 more|arXiv (Cornell University)|Jun 8, 2023

Machine Learning and Data Classification被引用数 6

ひとこと要約

本論文は、ハードウェアアクセラレータに最適化されたコア演算の最適化、活性化関数や正規化などのモデル部品の再設計、アプリケーション固有の入力表現の開発を通じて、球面畳み込みニューラルネットワーク（Spherical CNNs）のスケーラブルなフレームワークを導入する。このアプローチにより、QM9分子性質予測ベンチマークで最先端の性能を達成し、気象予測においても競争力のある結果を示した。これにより、Spherical CNNsが大規模な科学的応用に効果的にスケーリング可能であることが示された。

ABSTRACT

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.

研究の動機と目的

• 既存の球面CNNのスケーラビリティの制限を解決すること。これは、計算コストが高いため、低解像度・浅いモデルに限定されてきたためである。
• 分子性質予測や気象予測などの科学的分野における球面CNNの大規模応用を可能にすること。
• スペクトル領域における球面畳み込みの計算ボトルネックを、ハードウェアに配慮した実装とアーキテクチャ的革新によって克服すること。
• 球面CNNが、グラフニューラルネットワークやトランスフォーマーなどの最先端モデルと同等またはそれを上回る性能を、現実世界の科学的ベンチマークで示せることを実証すること。

提案手法

• 表現力と効率性を向上させた新しい球面CNNのファミリーの設計。これには、新規の活性化関数、正規化層、残差ブロックが含まれる。
• TPUアーキテクチャに最適化されたスピン重み付き球面調和変換の実装。これにより、スペクトル領域畳み込みの高速化が実現された。
• 3次元原子幾何構造と回転対称性をよりよく表現する、分子の新規な球面表現の導入。
• 空間畳み込みの代わりにスペクトルプーリングを適用することで、高周波成分を保持し、モデルの汎化性能を向上させる。
• 分離可能畳み込みと制御されたチャネルスケーリングを用いることで、気象予測タスクの性能を維持しながらFLOPsを削減する。
• 各アーキテクチャ的要素の性能寄与度を検証するための体系的で包括的なアブレーションスタディの実施。

Figure 1: Previous spherical CNNs were limited to low resolutions and relatively shallow models. In this work, we scale spherical CNNs by one order of magnitude and show that they can be competitive and even outperform state-of-the-art graph neural networks and transformers on scientific application

実験結果

リサーチクエスチョン

• RQ1 計算コストが著しく増大するおそれがある中で、球面CNNは、大規模なモデル容量と高解像度にスケーリング可能だろうか？
• RQ2 活性化関数や正規化、プーリング層のアーキテクチャ的変更は、球面CNNの性能と効率性にどのように影響するだろうか？
• RQ3 球面CNNは、従来グラフニューラルネットワークやトランスフォーマーが支配的であったQM9ベンチマークで、最先端の性能を達成できるだろうか？
• RQ4 アプリケーション固有の入力表現は、分子予測および気象予測タスクにおけるモデル性能をどの程度向上させるだろうか？
• RQ5 大規模な球面学習タスクにおいて、入力解像度、特徴マップサイズ、モデルの深さは、予測精度にどのような影響を及えるだろうか？

主な発見

• 提案された球面CNNモデルは、QM9の原子化エンタルピー回帰タスクで15.25 meVのテストMAEを達成し、先行研究の球面CNNを上回り、新たな最先端性能を樹立した。
• マグニチュード活性化関数ベースラインと比較して16.1%の誤差低減、前回の最先端性能と比較して25.1%の誤差低減を達成した。
• 気象予測においては、500 hPaにおけるジオポテンシャルでRMSEが34.93 m²/s²、850 hPaにおける温度で0.62 Kを達成し、既存のベースラインと同等またはそれを上回った。
• スペクトルプーリングは空間プーリングよりもQM9で1.5%の性能向上をもたらした。また、フェーズクラッシュ活性化関数はゲーテッド非線形性よりも1.5%の向上を示した。
• 気象予測において、入力解像度を256×256から64×64に低下させると、RMSEが100%増加した。これは、高解像度の入力がモデルの精度に不可欠であることを示している。
• 32個のTPUv4アクセラレータで学習したモデルは、4日間で最先端の結果を達成した。これは、計算コストが高額であるものの、実現可能性を示している。

Figure 2: Our efficient residual block contains spin-weighted spherical Fourier transforms (FT) and inverses (IFT), multiplication with filter coefficients ( $*K$ ), activation ( $\sigma$ ) and spectral batch normalization (BN). The residual connection happens in Fourier space. Optionally, spectral