夜雨聆风学习资料网

ARTICLE · 1125353

球面AI工具包来了

球面AI工具包来了
arXiv

球面AI工具包来了

A library for differentiable signal processing and machine learning on the sphere

球面数据无处不在,从天气预报到宇宙学,但传统深度学习工具在球面上"水土不服"。NVIDIA等机构联合发布torch-harmonics——一个GPU加速、可微分的球面信号处理与机器学习库,让球面AI像处理图像一样简单。该库已成功应用于天气气候建模、太阳风模拟等领域,有望成为球面科学计算的基础设施。
👥 作者:Thorsten Kurth、Max Rietmann、Mauro Bisson、Andrea Paris、Alberto Carpentieri、Jean Kossaifi、Anima Anandkumar、Christian Hundt、Boris Bonev 📅 发布日期:2026-09-30 🆔 arxiv ID:2609.39737v1
📝 论文摘要
嵌入三维欧几里得空间中的二维球面S2在地球物理、行星科学、大地测量学、大气物理、量子化学、宇宙学和虚拟现实等众多科学与工程领域中扮演着核心角色。随着机器学习日益渗透这些领域,人们对能够处理和建模球面函数、同时尊重该域固有拓扑与对称性质的工具需求不断增长。本文介绍了torch-harmonics,一个综合性库,提供球面数据的高效、可微分先进信号处理与机器学习方法的实现,包括球谐变换(SHT,即傅里叶变换的球面类比)、矢量球谐、离散-连续卷积与谱卷积,以及全局和邻域球面注意力机制。除传统表示外,torch-harmonics还为球面Transformer等最先进的球面机器学习架构提供了构建模块,以在现代科学与工程应用中实现可扩展的、旋转感知的学习与推理。

图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有望成为该领域的基础设施级工具。

🔑 关键词
球谐变换 球面机器学习 可微分信号处理 GPU加速 球面Transformer 天气气候建模 几何深度学习 PyTorch

📎 论文链接:http://arxiv.org/abs/2609.39737v1

📅 发布时间:2026-09-30

查看 arxiv 原文 →
论文来源 arxiv.org | 由 AI 自动翻译整理

相关学习资料