Attention on the Sphere

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Bonev, Boris, Rietmann, Max, Paris, Andrea, Carpentieri, Alberto, Kurth, Thorsten
Format: Preprint
Publié: 2025
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866916739942449152
author Bonev, Boris
Rietmann, Max
Paris, Andrea
Carpentieri, Alberto
Kurth, Thorsten
author_facet Bonev, Boris
Rietmann, Max
Paris, Andrea
Carpentieri, Alberto
Kurth, Thorsten
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2505_11157
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Attention on the Sphere
Bonev, Boris
Rietmann, Max
Paris, Andrea
Carpentieri, Alberto
Kurth, Thorsten
Machine Learning
Artificial Intelligence
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.
title Attention on the Sphere
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2505.11157