Spin Spherical Fourier Neural Operators: Efficient Equivariant Learning for Vector and Spin Fields on the Sphere
Abstract
Many fields of physical interest on the sphere, such as wind, the polarisation of the cosmic microwave background and the shear of weak gravitational lensing, are spin fields, whose values change phase when the local frame rotates. We extend spherical Fourier neural operators (SFNOs) to fields of any spin, through lifted spherical convolutions that are exactly rotationally equivariant. Lifting spin fields to Mackey functions on the rotation group provides the conceptual framework, while computation is performed exactly and efficiently through their harmonic representation and by appealing to sampling theorems on the sphere and rotation group. We construct spin SFNOs (S$^2$FNOs) from these building blocks. We demonstrate empirically that rotational equivariance holds to machine precision for these layers. We achieve state-of-the-art performance on existing benchmarks, and introduce new benchmarks on which spin-aware models outperform their scalar counterparts.
The following analysis is generated by AI and may contain errors. It is separate from the authors’ abstract and is not peer review.
Research summary
The paper presents a rigorous and efficient extension of Fourier Neural Operators to spin fields on the sphere, achieving exact rotational equivariance and state-of-the-art performance on new benchmarks in weather and cosmology.
Methodology
The paper introduces Spin Spherical Fourier Neural Operators (S$^2$FNOs), extending the SFNO architecture to handle spin-weighted fields (vectors, tensors) on the sphere. The core innovation is the use of lifted spherical convolutions based on Mackey functions on the rotation group SO(3). This allows for exact rotational equivariance for linear layers. The authors provide a rigorous theoretical framework connecting spin fields to functions on the rotation group, proving that their harmonic-space implementation is both efficient (avoiding full group convolutions where possible) and exactly equivariant. They address the challenge of non-linearities by distinguishing between equivariant non-linearities on the sphere (modulus-based) and those on the rotation group (arbitrary pointwise maps), providing a flexible block design. The parameterization via interpolated profiles in harmonic space ensures resolution independence, a key feature of neural operators.
Results & evaluation
The authors introduce three new benchmarks: Spin Spherical MNIST (classification/dense prediction), Global Wind Forecasting (spin-1), and Weak Gravitational Lensing (spin-2 to spin-0). These are significant contributions as public benchmarks for spin fields were previously scarce. Empirically, the spin-aware models outperform scalar counterparts (which treat vector components as independent channels) by significant margins (e.g., 44.9% MSE reduction in wind forecasting). The paper also verifies equivariance to machine precision ($10^{-15}$) for linear layers, validating the theoretical claims. The comparison with existing methods like Spin-CNNs and group convolutions highlights the efficiency advantage (fewer parameters) while maintaining or improving accuracy.
Reproducibility
The paper states that code and benchmarks will be released with the `s2ai` codebase. Detailed architectural descriptions, initialization schemes, and hyperparameters are provided in the appendices. The use of standard libraries (JAX) and well-defined mathematical operations (SHT, Wigner transforms) enhances reproducibility. The explicit mention of the `s2fft` code for transforms further aids implementation.
Limitations
The primary limitation is the computational cost of the rotation group convolutions, although the authors mitigate this by restricting kernels to diagonal spins in the group construction. The non-linearities on the sphere are restricted to modulus-based functions to maintain equivariance, which may limit expressivity compared to arbitrary non-linearities on the group. The benchmarks, while novel, are specific to certain physical domains; generalization to other spin-field applications needs further validation. The paper relies on bandlimited signals, which may not capture high-frequency details in some real-world data without careful truncation.
Broader impact
This work has significant implications for scientific machine learning, particularly in cosmology (CMB polarization, weak lensing) and geophysics (weather forecasting, ocean currents). By providing a principled way to handle spin fields with exact equivariance, it enables more accurate and data-efficient models for these domains. The introduction of new benchmarks will facilitate future research in equivariant learning on the sphere. The methods could also be extended to other manifolds or groups, broadening the scope of geometric deep learning. The paper presents a rigorous and efficient extension of Fourier Neural Operators to spin fields on the sphere, achieving exact rotational equivariance and state-of-the-art performance on new benchmarks in weather and cosmology.
Author of this paper? Send a correction.