ARTICLE · 1125353
球面AI工具包来了
球面AI工具包来了
A library for differentiable signal processing and machine learning on the sphere
图1
研究背景
球面是自然界中最基本的几何形状之一。从地球表面的大气环流、海洋温度分布,到宇宙微波背景辐射、分子表面的电子密度分布,大量科学数据天然地定义在球面上。在大气科学领域,数值天气预报模式、气候模拟和遥感卫星观测数据无一不涉及球面几何。然而,传统的深度学习工具几乎全部针对欧几里得空间(如平面图像)设计,卷积神经网络、Transformer等经典架构在球面上无法直接应用——强行将球面数据投影到平面上会引入严重的几何畸变,尤其在极地地区误差急剧放大。 近年来,几何深度学习为解决这一问题提供了新思路。研究者提出了基于球面网格的图神经网络和基于球谐函数的等变网络等方法。球谐函数是球面上的正交基函数,类似于平面上的傅里叶基,能够以旋转等变的方式表示球面信号。然而,这些方法的实际推广一直受限于一个关键瓶颈:缺乏高效、可微分且可扩展的实现。现有工具要么计算效率低下,要么不支持GPU加速,要么无法嵌入现代深度学习框架,导致球面机器学习研究门槛极高。 正是在这一背景下,torch-harmonics库应运而生。该库由NVIDIA、加州理工学院等机构的研究者联合开发,旨在为球面信号处理和机器学习提供一个完整的、GPU加速的、与PyTorch生态无缝集成的开源解决方案。
核心问题
本研究的核心问题是:如何构建一个高效、可微分、可扩展的球面信号处理与机器学习库,使研究者能够像使用标准深度学习工具处理图像一样,方便地处理球面数据,同时严格保持球面的拓扑与对称性质?
研究方法
torch-harmonics的核心是球谐变换(SHT)的高效GPU实现。球谐变换将球面上的函数分解为一系列球谐系数的过程,类似于傅里叶变换将一维信号分解为不同频率的正弦波。研究团队设计了自定义CUDA内核,利用GPU的大规模并行能力加速这一计算过程,并通过分布式内存并行技术将计算扩展到此前被认为计算上不可行的高分辨率场景。 在SHT的基础上,库提供了多种球面卷积操作。离散-连续(DISCO)卷积允许在球面采样点与连续球面信号之间进行卷积运算,谱卷积则直接在球谐系数域中进行滤波操作。此外,库还实现了矢量球谐变换,能够处理球面上的矢量场(如风场、磁场等),以及全局和邻域球面注意力机制,为构建球面Transformer架构提供了基础组件。 整个库完全集成在PyTorch生态系统中,支持自动微分,这意味着所有操作都可以直接用于神经网络的端到端训练。研究团队在多个分辨率下进行了性能基准测试,验证了库的计算效率和可扩展性。
主要发现
torch-harmonics在多个维度上展现了显著优势。首先,在计算效率方面,自定义CUDA内核和分布式并行策略使该库能够处理此前计算上不可行的高分辨率球面数据,将球谐变换的计算时间从数小时缩短至分钟甚至秒级。其次,在功能完备性方面,该库覆盖了从基础球谐变换到高级球面注意力机制的完整工具链,是目前最全面的球面机器学习库之一。 在实际应用验证中,torch-harmonics已成功支撑了多个跨学科应用。在大气科学领域,基于该库开发的球面傅里叶神经算子(SFNO)在天气与气候建模中展现了优异的稳定性和精度。此外,该库还被应用于太阳风建模、全向3D视觉以及分子力场计算等领域,充分证明了其作为通用球面学习工具的广泛适用性。
创新亮点
与已有工作相比,torch-harmonics的创新主要体现在三个方面:一是首次在PyTorch生态中实现了完整的、GPU加速的球面信号处理工具链;二是通过自定义CUDA内核和分布式并行策略,将球面机器学习的计算规模推向了新的高度;三是提供了从传统球谐变换到现代球面Transformer的完整构建模块,使研究者无需从零实现底层算法。
研究意义
torch-harmonics的发布对大气科学领域具有深远影响。天气和气候模式天然定义在球面上,该库为数据驱动的天气预报和气候模拟提供了高效的计算基础设施。研究者可以更方便地构建旋转等变的神经网络模型,避免传统方法在极地地区的精度损失,从而提升全球天气预报的准确性和物理一致性。
未来展望
当前版本的torch-harmonics已在多个领域展现潜力,但仍有进一步发展的空间。未来工作可能包括支持更高效的球面采样策略、扩展到其他流形几何、以及与更多科学计算框架的集成。随着球面机器学习在气象、气候和行星科学中的需求持续增长,torch-harmonics有望成为该领域的基础设施级工具。
📎 论文链接:http://arxiv.org/abs/2609.39737v1
📅 发布时间:2026-09-30