Source-linked AI summary
A library for differentiable signal processing and machine learning on the sphere
Thorsten Kurth, Max Rietmann, Mauro Bisson, Andrea Paris, Alberto Carpentieri, Jean Kossaifi, Anima Anandkumar, Christian Hundt, Boris Bonev
TL;DR
Spherical machine learning needs tools that respect the domain’s topology and symmetry, while practical adoption of existing geometric methods has been hindered by a lack of efficient, differentiable, scalable implementations. torch-harmonics provides differentiable spherical signal-processing and machine-learning operations, with CUDA kernels and distributed computing to support high-resolution models. The library’s operations also underpin spherical architectures and applications across weather, vision, and scientific computing.
Problem
Spherical domains are important across scientific fields, but existing geometric deep-learning approaches have lacked efficient, differentiable, and scalable implementations for practical adoption.
Source: §1 · p. 1Method
torch-harmonics provides differentiable PyTorch implementations of spherical transforms, convolutions, and attention, with CUDA kernels and distributed computing support.
Source: §1 · pp. 1–2, §2.1 · p. 3, §3 · p. 4Results
Custom CUDA kernels and distributed computing enable training high-resolution models that were previously computationally prohibitive.
Source: §1 · p. 2Takeaways & Limitations
The library supports spherical architectures and has been used in weather forecasting, solar wind modeling, omnidirectional 3D vision, and molecular force fields.
Source: §5 · p. 4
Abstract
from arXiv · showhide
The two-dimensional sphere embedded in three-dimensional Euclidean space S2, plays a central role in a variety of scientific and engineering domains, including geophysics, planetary science, geodesy, atmospheric physics, quantum chemistry, cosmology, and virtual reality, among many others. As machine learning increasingly permeates these fields, the demand grows for robust tools that process and model functions on the sphere, while respecting the inherent topological and symmetry properties of the domain. We present torch-harmonics, a comprehensive library that offers efficient, differentiable implementations of advanced signal processing and machine learning (ML) methods for spherical data. These include the spherical harmonic transform (SHT), the spherical analogue of the Fourier transform, vector spherical harmonics, discrete-continuous and spectral convolutions, as well as both global and neighborhood spherical attention mechanisms. Beyond traditional representations, torch-harmonics provides the building blocks for state-of-the-art spherical ML architectures such as spherical transformers in order to enable scalable, rotationally-aware learning and inference in modern scientific and engineering applications.
1 Introduction
Spherical applications need tools that respect the sphere’s topology and symmetry, but existing geometric deep learning methods have lacked practical implementations. torch-harmonics provides differentiable spherical signal-processing operations and uses CUDA kernels and distributed computing to support high-resolution models.
1 Introduction
- Standard Euclidean deep-learning operations do not naturally extend to the sphere without distortions or singularities.Existing geometric deep-learning approaches have also faced limited practical adoption because efficient, differentiable, and scalable implementations were lacking.Source: §1 · p. 1
- torch-harmonics is a PyTorch library for differentiable signal processing on the sphere, including spherical harmonic and vector spherical harmonic transforms and spherical convolutions.The introduction also identifies spherical attention among its operations.Source: §1 · pp. 1–2
- Custom CUDA kernels and distributed computing support training high-resolution models that were previously computationally prohibitive.The library serves as a foundation for the Spherical Fourier Neural Operator (SFNO) (Bonev et al., 2023) and FourCastNet3 (FCN3) (Bonev et al., 2025a).Source: §1 · p. 2
2 Library Design and Functionality
torch-harmonics combines differentiable PyTorch operations with CUDA acceleration, PyTorch fallbacks, and distributed computing. Its spherical signal-processing components cover transforms, interpolation, quadrature, convolutions, and attention.
2 Library Design and Functionality
- All operations support PyTorch backpropagation, with custom CUDA kernels, pure-PyTorch fallbacks, and distributed memory support for high-resolution data.These are the library’s stated design goals of differentiability, efficiency, compatibility, and scalability.Source: §2 · p. 2
2.1 Core Signal Processing Operations
- The spherical harmonic transform decomposes grid values into spherical harmonic basis functions and is implemented through Fourier and Legendre transforms.The Legendre step uses batched matrix multiplications, while the Fourier step uses an FFT.Source: §2.1 · pp. 2–3
- Vector spherical harmonic transforms decompose vector fields into divergence-free and curl-free components, supporting direct computation of divergence and curl in the spectral domain.The passage highlights their usefulness for solving PDEs on the sphere.Source: §2.1 · p. 3
- The library supports spectral interpolation for alias-free resolution changes, bilinear interpolation using Haversine distance, and quadrature for spherical integrals and norms.Quadrature rules are available for all supported grid types.Source: §2.1 · p. 3
- Spherical convolution is defined through the rotation group SO(3), with spectral and discrete-continuous approaches supporting global and local anisotropic filters.DISCO filters can use piecewise linear hat functions, Zernike polynomials, or wavelet-like functions.Source: §2.1 · p. 3
Spherical Attention
- Global and neighborhood attention use quadrature weights to account for non-uniform sampling density on spherical grids.Distributed implementations use halo exchanges for attention mechanisms.Source: Spherical Attention · p. 3
3 Machine Learning Architectures
The paper describes spherical architectures using spectral, localized, attention-based, and hybrid physics–machine-learning approaches. These include models that combine global and local operators, continuous attention, and differentiable spectral solvers.
Spherical Fourier Neural Operators
- Spherical Fourier Neural Operators (Bonev et al., 2023) apply learnable spectral filters between spherical harmonic transforms to provide global receptive fields on the sphere.Their blocks apply a pointwise nonlinearity after the inverse transform.Source: §3 · p. 3
Local Spherical Neural Operators
- Local Spherical Neural Operators (Liu-Schiaffini et al., 2024) use compact-support DISCO convolutions for localized anisotropic features, while FourCastNet 3 combines LSNO and SFNO layers.LSNOs maintain approximate rotation equivariance and suit processes with preferential directions.Source: §3 · pp. 3–4
Spherical Transformers
- Spherical Transformers (Bonev et al., 2025b) use continuous attention on spherical geometry, with geodesic distances and quadrature weights in the library’s global and neighborhood attention implementations.The passage identifies data assimilation and irregularly sampled observations as useful settings.Source: §3 · p. 4
Hybrid Physics-ML Architectures
- Differentiable spectral transforms support hybrid methods that incorporate machine learning into numerical solvers, including a shallow-water-equation solver.The passage says such solvers form the backbone of global circulation models and enable training of neural GCMs.Source: §3 · p. 4
4 Applications and Examples
The repository provides examples spanning differentiable spherical solvers, neural operators, image processing, and reusable model code, alongside reported deployments in weather forecasting and scientific computing.
4 Applications and Examples
- Examples cover a differentiable shallow-water solver, SFNO and LSNO training on its generated data, and spherical-transformer image processing for 360° depth estimation and segmentation against Euclidean baselines.The repository also includes complete examples of spherical layers, models, solvers, and their integration into standard training pipelines; the library is reported in production weather forecasting and scientific computing.Source: §4 · p. 4
5 Conclusion
The conclusion presents torch-harmonics as an open-source library for differentiable spherical signal processing and machine learning, while Appendix A reviews the mathematical foundations of its operations.
5 Conclusion
- torch-harmonics provides GPU-accelerated spherical transforms, convolutions, and attention integrated with PyTorch, with CUDA kernels and distributed parallelism supporting high-resolution work.The conclusion notes applications in weather and climate, solar wind modeling, omnidirectional 3D vision, and molecular force fields.Source: §5 · p. 4
Appendix A. Signal processing on the sphere
- Appendix A introduces the spherical domain, coordinate systems, functions, rotations, spherical harmonics, convolutions, and attention used by the library.It frames the mathematical background for the signal-processing operations implemented in torch-harmonics.Source: Appendix A · p. 8
A.1 Coordinate systems and functions on the sphere
- Spherical coordinates use colatitude ϑ and longitude φ; functions on S2 are treated in L2(S2) with the rotation-invariant measure dµ(x) = sin ϑ dϑ dφ.The North Pole is at ϑ = 0, and φ ranges from 0 to 2π.Source: §A.1 · p. 8
A.2 The rotation group SO(3)
- SO(3) consists of three-dimensional rotations and acts transitively on S2 by mapping each point x to Rx.A function on the sphere can likewise be transformed by a rotation operator associated with R.Source: §A.2 · p. 8
A.3 Grids and Quadrature
- Spherical signals are discretized on grids of points with quadrature weights; Gaussian grids integrate spherical harmonics exactly through degree 2nlat − 1.Equiangular grids use equally spaced spherical coordinates, while Gaussian grids use Gauss–Legendre latitude nodes.Source: §A.3 · pp. 8–9
A.4 Spherical Harmonics
- Spherical harmonics form a complete orthonormal basis for L2(S2), and the SHT expands spherical functions in that basis.The transform can be factorized by exploiting the tensor-product structure of the harmonics.Source: §A.4 · p. 9, Implementation Details · p. 9
Implementation Details
The implementation separates the spherical harmonic transform into longitude and latitude operations, uses spectral coefficients for vector-field derivatives, and defines spherical convolution through rotations.
Implementation Details
- The SHT first applies a real-valued FFT along longitude, truncated to M modes, then contracts over latitude with Legendre coefficients and quadrature weights.The inverse transform reverses the order, using Legendre synthesis followed by an inverse FFT, without quadrature weights in the inverse Legendre step.Source: Implementation Details · p. 9
- The angular power spectrum measures signal energy by spatial scale and is rotation-invariant, providing a compact summary of spectral content.It quantifies the contribution of scale ℓ to total signal energy.Source: Implementation Details · p. 10
A.5 Vector Spherical Harmonics
- Vector spherical harmonics uniquely decompose tangent fields into toroidal, divergence-free and poloidal, curl-free components, whose spectral coefficients support differential-operator calculations.Divergence is computed from the poloidal potential and curl from the toroidal potential via the spherical Laplacian.Source: §A.5 · p. 10
A.6 Group Convolutions
- Spherical convolution generalizes translation-equivariant convolution through a rotated-filter inner product on SO(3); restricting rotations to SO(3)/SO(2) yields an output on S2.The rotation used for the spherical output maps the North Pole to the point x.Source: §A.6 · p. 10
A.7 The Convolution Theorem
The section contrasts efficient global isotropic spectral convolution with DISCO, which directly approximates continuous convolution for localized, anisotropic filters.
A.7 The Convolution Theorem
- For zonal filters, the spherical convolution theorem reduces convolution to pointwise multiplication of spherical-harmonic coefficients, enabling efficient global isotropic operations.This spectral-domain computation underlies the Spherical Fourier Neural Operator.Source: §A.7 · pp. 10–11
A.8 Discrete-Continuous (DISCO) Convolutions
- DISCO approximates the continuous convolution integral by quadrature to support local, anisotropic filters.Its filter can be represented by basis functions on a tangent plane or disk centered at the North Pole; rotating the filter centers it at each output point and provides approximate rotation equivariance.Source: §A.8 · p. 11
- Alternative DISCO bases include Morlet-like wavelets on a compact disk and Zernike polynomials, which are orthogonal on the disk.The Morlet-like basis uses a Hann window to produce smooth, compactly supported filters while keeping the convolution tensor sparse.Source: §A.8 · p. 11
- DISCO can parameterize filters with piecewise-linear radial and angular hat functions on a polar grid, forming a flexible local basis suited to polar geometry.The radial and angular basis functions are combined as a tensor product.Source: §A.8 · p. 12
A.8.1 Implementation Details
The implementation precomputes and sparsifies DISCO geometry-dependent operations, while spherical attention uses quadrature-aware global integration or a local geodesic mask.
A.8.1 Implementation Details
- The DISCO convolution tensor depends only on the grid geometry and resolution, so it can be precomputed and stored sparsely; longitudinal shifts permit reuse across output locations.For coarser output grids, the shift uses a stride determined by the input and output longitude resolutions.Source: §A.8.1 · p. 12
- DISCO convolution is implemented as sparse-times-dense multiplication with a shift, using a custom CUDA kernel for its complex memory access patterns.The transpose convolution reverses the operation by summing over output indices.Source: §A.8.1 · p. 12
A.9 Spherical Attention
- Spherical attention discretizes its continuous integral with quadrature weights, accounting for nonuniform grid sampling and approximately respecting spherical geometry and SO(3) equivariance.The weights address sampling-density variation such as point clustering near the poles on equiangular grids.Source: §A.9 · pp. 12–13
- Neighborhood attention masks scores outside a local geodesic neighborhood, sparsifying the attention matrix while preserving local spherical symmetries.The mask uses the neighborhood radius r to restrict each query's attention computation.Source: §A.9 · p. 13
Appendix B. Implementation Examples
Appendix B illustrates how torch-harmonics components can be assembled into spherical neural architectures, including an SFNO spectral layer. Its transforms handle coordinate changes, while spectral multiplication supplies global convolution.
B.1 Spherical Fourier Neural Operator
- An SFNO layer transforms inputs to spectral coefficients, applies learnable spectral weights, and uses the inverse SHT to return outputs to the spatial domain.The example uses RealSHT and InverseRealSHT and defines learnable weights over channels and spectral modes.Source: §B.1 · p. 13
- The forward and inverse SHT handle coordinate transformations, while spectral multiplication implements a global convolution.Source: §B.1 · p. 13
- Driscoll-Healy spectral convolutions are isotropic and spherically equivariant, while a generalized spectral bias term supports anisotropy.The Driscoll-Healy weights are real-valued and depend only on m.Source: §B.1 · p. 14
B.2 DISCO Convolutions for Local Processing
The implementation examples span local DISCO filters, spherical attention, hybrid blocks, distributed spectral computation, a shallow-water solver, and GPU memory-layout optimization. Together they show how library operators and implementation choices support these use cases.
B.2 DISCO Convolutions for Local Processing
- DISCO convolutions provide local, anisotropic filters with filter rotation and quadrature handled internally, yielding approximate rotation equivariance and localized receptive fields.Source: §B.2 · p. 14
B.3 Spherical Attention
- Spherical attention uses quadrature weights so its integral accounts for varying grid-cell areas and respects spherical geometry.Source: §B.3 · p. 15
B.4 Hybrid Model: Combining SFNO and DISCO
- The hybrid block combines SFNO for large-scale wave dynamics with DISCO for small-scale local features, targeting complex physical systems such as atmospheric flows.Source: §B.4 · pp. 15–16
B.5 Distributed Spectral Convolution
- Distributed spectral convolution decomposes the sphere across latitude and longitude process groups, with custom autograd handling layer collectives and additional reductions for shared weight gradients.The example also describes a separate data-parallel group; spectral-convolution weights are shared along longitude and require gradient reduction there.Source: §B.5 · pp. 16, 18
B.6 Shallow Water Equations Solver
- The shallow-water solver uses vector spherical harmonic transforms to compute divergence and curl in spectral space, avoiding numerical instabilities common in finite-difference schemes.The equations describe velocity and fluid-height evolution on a rotating sphere.Source: §B.6 · pp. 19, 21
B.7.1 Memory Layout Optimizations for SHT
- Stride-1 contraction layouts let the Legendre step run as a high-performance batched GEMM that exploits GPU memory hierarchy and tensor cores.The layout places the contraction index in the fastest-varying position in both the input and precomputed weight tensors.Source: §B.7.1 · pp. 21–22
B.8 Implementation Details for Distributed Operations
The appendix details distributed spherical harmonic and DISCO operations, then describes custom CUDA operators for sparse aggregation workloads. Its distributed SHT uses all-to-all communication to preserve serial-transform outputs exactly.
B.8.1 Distributed SHT
- The distributed SHT partitions latitude and longitude across a two-dimensional process grid and uses NCCL all-to-all collectives to transpose data for local FFT and Legendre-transform steps.The forward transform redistributes spatial dimensions, computes a local FFT, and then performs a local weighted Legendre projection.Source: §B.8.1 · pp. 22–23
- The distributed SHT's all-to-all approach produces outputs bit-wise identical to the corresponding serial transform, unlike a distributed matrix multiplication that can differ because of operation ordering.Source: §B.8.1 · p. 23
B.8.2 Distributed DISCO Convolution
- Distributed DISCO convolution makes longitude local, applies per-rank sparse kernels to latitude slabs, and combines partial outputs with a polar all-reduce and scatter.The latitude dimension can be split, while the full longitude extent must be locally available because the longitudinal shift couples all longitudes.Source: §B.8.2 · pp. 23–24
B.9 Custom CUDA Operators
- Custom CUDA forward and backward kernels implement DISCO and spherical neighborhood attention as sparse graph aggregations with quadrature weighting.The passage notes that PyTorch lacks high-level tools for implementing these operators efficiently.Source: §B.9 · p. 25
B.9.1 Spherical Neighborhood Attention
The section describes custom kernels for spherical neighborhood attention and contrasts their data layouts with DISCO convolution. It reports faster tensor permutation and a load-balancing strategy that improves runtime in experiments.
B.9.1 Spherical Neighborhood Attention
- Attention kernels use channel-last layout for coalesced feature-vector access, with register-resident and shared-memory variants to accommodate different channel widths.The kernels can also use vectorized 128-bit loads and stores when alignment permits.Source: §B.9.1 · pp. 25–26
- ∼6.2 TB/s versus 0.5 TB/s: custom BCHW↔BHWC transpose kernels on a GB200 make permutation cost negligible relative to aggregation, compared with PyTorch permute().The comparison used ∼2.0 GB of data; the custom-kernel throughput is described as an order-of-magnitude improvement.Source: §B.9.1 · pp. 25–26
- Presorting CSR rows by descending length yielded up to 20% runtime improvement on a GB200 by scheduling dense neighborhoods first and improving GPU utilization.The sorting cost is described as negligible because reordering depends only on CSR row length.Source: §B.9.1 · p. 26
B.9.2 DISCO convolution
- Unlike attention's feature-vector dot products, DISCO reduces input vectors element-wise per channel and therefore processes tensors in native BCHW layout without permutations.Its thread mapping also reduces redundant reads and backward-pass atomic operations.Source: §B.9.2 · p. 27