适用于科学应用的可扩展球形 CNN

典型的计算机视觉深度学习模型，如卷积神经网络(CNN) 和视觉变换器(ViT)，在假设平面空间的情况下处理信号。例如，数字图像表示为平面上的像素网格。然而，这种类型的数据只占我们在科学应用中遇到的数据的一小部分。从地球大气中采样的变量，如温度和湿度，自然地表示在球面上。某些类型的宇宙学数据和全景照片也是球面信号，因此最好这样处理。

使用为平面图像设计的方法处理球面信号存在一些问题，原因如下。首先，存在采样问题，即无法在球面上定义均匀的网格，而平面 CNN 和 ViT 需要这些网格，并且不产生严重失真。

将球体投影到平面时，红色圆圈所代表的块在极点附近严重扭曲。此采样问题会损害传统 CNN 和 ViT 对球面输入的准确性。

其次，球体上的信号和局部模式通常因旋转而变得复杂，因此模型需要一种方法来解决这个问题。我们希望与 3D 旋转等方差，以确保学习到的特征遵循输入的旋转。这可以更好地利用模型参数，并允许使用更少的数据进行训练。在大多数输入没有首选方向的设置中，例如 3D 形状和分子，与 3D 旋转等方差也很有用。

使用全景摄像机进行无人机比赛。这里的急转弯导致球面图像发生较大的 3D 旋转。我们希望我们的模型能够承受这种旋转。

在大气中，常见的情况是看到相似的图案出现在不同的位置和方向。我们希望我们的模型能够共享参数来识别这些图案。

考虑到上述挑战，我们在ICML 2023上展示的“扩展球形 CNN ”中引入了JAX中的开源库，用于球形表面的深度学习。我们展示了该库的应用如何在天气预报和分子特性预测基准上达到或超越最先进的性能，这些任务通常使用 Transformer 和图神经网络来解决。

球形 CNN 的背景

球面 CNN利用球面卷积和互相关运算解决了采样问题和旋转鲁棒性问题，这些运算通常通过广义傅里叶变换计算得出。然而，对于平面，使用小滤波器的卷积速度更快，因为它可以在规则网格上执行而无需使用傅里叶变换。球面输入的计算成本较高，迄今为止，这限制了球面 CNN 在小型模型和数据集以及低分辨率数据集中的应用。

我们的贡献

我们在JAX中实现了自旋加权球形 CNN 的球形卷积，重点关注速度，并使用数据并行性在大量 TPU 上实现分布式训练。我们还引入了新的相位坍缩激活和光谱批量归一化层，以及一个可提高准确度和效率的新残差块，从而可以训练比以前大 100 倍的更准确的模型。我们将这些新模型应用于分子特性回归和天气预报。

与文献相比，我们在特征大小和模型容量方面将球形 CNN 提高了两个数量级：  Cohen'18、  Esteves'18、  Esteves'20和 Cobb'21。VGG -19被列为传统 CNN 参考。  我们最大的天气预报模型有 256 x 256 x 78 个输入和输出，在训练期间运行 96 个卷积层，最低内部分辨率为 128 x 128 x 256。

分子性质回归

预测分子特性可用于药物研发，其目标是快速筛选大量分子以寻找具有所需特性的分子。类似的模型也可能适用于针对蛋白质间相互作用的药物设计。计算或实验量子化学中的当前方法成本高昂，这促使人们使用机器学习。

分子可以用一组原子及其在三维空间中的位置来表示；分子的旋转会改变位置，但不会改变分子特性。这促使球形 CNN 因其旋转等方差而得到应用。然而，分子并非定义为球面上的信号，因此第一步是将它们映射到一组球形函数。我们通过利用分子原子之间的物理相互作用来实现这一点。

每个原子由一组球形信号表示，这些信号累积了与每种类型的其他原子的物理相互作用（显示在右侧的三个面板中）。例如，氧原子（O；顶部面板）有一个氧通道（左侧标记为“O”的球体表示）和氢通道（“H”，右侧）。氧原子相对于两个氢原子的累积 库仑 力由标记为“H”的球体底部的红色阴影区域表示。由于氧原子不向自身施加任何力，“O”球体是均匀的。我们为范德华力添加了额外的 通道 。

将球形 CNN 应用于每个原子的特征，然后将结果组合起来以产生属性预测。这在大多数属性中产生了最先进的性能，正如QM9基准中通常评估的那样：

与 QM9 的 12 个属性上的最新技术进行错误比较（  详情 请参阅数据集论文）。我们展示了TorchMD-Net 和 PaiNN 的 结果，将 TorchMD-Net 错误标准化为 1.0（越低越好）。我们的模型（以绿色显示）在大多数目标上都优于基线。

天气预报

准确的气候预报是及时预警极端天气事件、实现有效的水资源管理以及指导明智的基础设施规划的宝贵工具。在一个日益受到气候灾害威胁的世界里，迫切需要比一般环流模型更快、更准确地在更长的时间范围内提供预报。预测模型对于预测旨在应对气候变化的努力（例如气候干预）的安全性和有效性也很重要。目前最先进的技术使用基于流体动力学和热力学的昂贵数值模型，这些模型往往会在几天后出现偏差。

鉴于这些挑战，机器学习研究人员迫切需要解决气候预测问题，因为数据驱动技术既可以降低计算成本，又可以提高长距离精度。球形 CNN 适合这项任务，因为大气数据本身就呈现在球体上。它们还可以有效处理此类数据中常见的不同位置和方向的重复模式。

我们将我们的模型应用于多个天气预报基准，其表现 优于或匹敌基于传统 CNN 的神经天气模型（具体来说，1、2和3）。下面我们展示了测试设置中的结果，其中模型将多个大气变量作为输入，并预测其值提前六小时。然后，该模型迭代应用于其自身预测，以生成更长的预测。在训练期间，该模型最多可预测三天，最多可评估五天。Keisler为这项任务提出了一个图神经网络，但我们表明，球形 CNN 可以在相同设置下匹配 GNN 精度。

使用球形 CNN 进行长达五天（120 小时）的迭代天气预报。动画显示了给定压力下的特定湿度预测及其误差。

使用球形 CNN 预测风速和温度。

其他资源

我们的高效球形 CNN JAX 库现已可用。我们已经展示了分子特性回归和天气预报的应用，我们相信该库将在其他科学应用以及计算机视觉和 3D 视觉方面有所帮助。

天气预报是 Google 的一个活跃研究领域，其目标是构建更准确、更强大的模型（例如Graphcast，一种最近的基于 ML 的中程预测模型），并构建能够推动整个研究界进一步进步的工具，例如最近发布的WeatherBench 2。

致谢

这项工作是与 Jean-Jacques Slotine 合作完成的，基于之前与 Kostas Daniilidis 和 Christine Allen-Blanchette 的合作。我们感谢 Stephan Hoyer、Stephan Rasp 和 Ignacio Lopez-Gomez 帮助处理和评估数据，感谢 Fei Sha、Vivian Yang、Anudhyan Boral、Leonardo Zepeda-Núñez 和 Avram Hershko 提供的建议和讨论。我们感谢 Michael Riley 和 Corinna Cortes 对这个项目的支持和鼓励。