[论文解读] 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分子性质预测基准上实现了最先进性能，在天气预报任务中也取得了具有竞争力的结果，表明球面CNN能够有效扩展至大规模科学应用。

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可在真实世界科学基准上超越或匹配图神经网络与Transformer等最先进模型的性能。

提出的方法

• 设计新型球面CNN家族，提升表达能力与效率，包括新型激活函数、归一化层和残差块。
• 为TPU架构量身定制高效的旋量加权球谐变换，加速频域卷积运算。
• 提出一种新型分子球面表征方式，更优地编码三维原子几何结构与旋转对称性。
• 采用谱池化替代空间池化，以保留高频信息并提升模型泛化能力。
• 应用可分离卷积与受控通道缩放策略，在保持天气预报任务性能的同时降低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基准上实现最先进性能，该基准长期由图神经网络与Transformer主导？
• RQ4 特定应用的输入表征在分子与天气预报任务中能在多大程度上提升模型性能？
• RQ5 输入分辨率、特征图尺寸与模型深度对大规模球面学习任务中预测准确率的影响如何？

主要发现

• 所提出的球面CNN模型在QM9原子化焓回归任务上的测试平均绝对误差（MAE）达到15.25 meV，优于以往球面CNN，并创下新SOTA记录。
• 与幅度激活基线相比，误差降低16.1%；与此前SOTA相比，误差降低25.1%。
• 在天气预报任务中，模型在500 hPa高度位势能的均方根误差（RMSE）为34.93 m²/s²，850 hPa温度预测的RMSE为0.62 K，表现与现有基线相当或更优。
• 在QM9任务中，谱池化相比空间池化性能提升1.5%；相位崩溃激活函数相比门控非线性函数提升1.5%。
• 将输入分辨率从256×256降低至64×64，天气预报任务的RMSE增加100%，凸显高分辨率输入对模型准确率的关键作用。
• 该模型在32个TPUv4加速器上训练4天即达到SOTA性能，尽管计算成本较高，但仍证明了其可行性。

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