Attention on the Sphere

TL;DR

This work addresses learning on spherical domains by generalizing attention to the sphere, ensuring $SO(3)$-equivariance and geometric fidelity. It derives continuous spherical attention mechanisms (global and neighborhood) and discretizes them with quadrature weights, enabling neural operator behavior on arbitrary sphere discretizations. The authors implement spherical Transformers and SegFormers in PyTorch with CUDA support and validate on segmentation on the sphere, depth estimation, and shallow water equations, showing improved performance over planar baselines. The work demonstrates that incorporating spherical geometry as a prior improves stability and accuracy for physics-informed and vision tasks on spherical data, with potential impact in geoscience, cosmology, and robotics.

Abstract

We introduce a generalized attention mechanism for spherical domains, enabling Transformer architectures to natively process data defined on the two-dimensional sphere - a critical need in fields such as atmospheric physics, cosmology, and robotics, where preserving spherical symmetries and topology is essential for physical accuracy. By integrating numerical quadrature weights into the attention mechanism, we obtain a geometrically faithful spherical attention that is approximately rotationally equivariant, providing strong inductive biases and leading to better performance than Cartesian approaches. To further enhance both scalability and model performance, we propose neighborhood attention on the sphere, which confines interactions to geodesic neighborhoods. This approach reduces computational complexity and introduces the additional inductive bias for locality, while retaining the symmetry properties of our method. We provide optimized CUDA kernels and memory-efficient implementations to ensure practical applicability. The method is validated on three diverse tasks: simulating shallow water equations on the rotating sphere, spherical image segmentation, and spherical depth estimation. Across all tasks, our spherical Transformers consistently outperform their planar counterparts, highlighting the advantage of geometric priors for learning on spherical domains.

Figures (5)

• Figure 1: Neighborhood attention on the sphere: (a) Euclidean neighborhoods use planar distances, distorting receptive fields near poles. (b) Naive spherical attention uses geodesic neighborhoods but equally weights points, oversampling points and breaking $\mathrm{SO}(3)$ symmetry. (c) Our method introduces quadrature weights to account for non-uniform sampling density, preserving rotational equivariance regardless of grid alignment through proper discretization of the continuous formulation.
• Figure 2: Illustration of the spherical Transformer architecture.
• Figure 3: Samples from the segmentation dataset using the Euclidean and $S^2$ neighborhood SegFormers.
• Figure 4: Sample predictions from our transformers alongside ground truth data from the spherical depth estimation task.
• Figure 5: Comparison of predictions from $\mathbb{R}^2$ and $S^2$ Transformers to a ground truth solution of the shallow water equations.