Title: Attention on the Sphere
Venue: NeurIPS 2025
Accepted: yes
Keywords: Attention

neighborhood attention

transformers

sphere

geometric machine learning

equivariance

scientific machine learning

spherical signal processing

CUDA
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors introduce a generalized attention mechanism for spherical domains, enabling Transformers to effectively process data on the 2D sphere. By incorporating numerical quadrature weights into the attention mechanism, an approximately rotationally equivariant and geometrically faithful spherical attention is obtained, offering stronger inductive biases and better performance than Cartesian methods. To improve scalability and locality, the authors additionally propose neighborhood attention that limits interactions to geodesic neighborhoods, reducing computational cost while maintaining symmetry. They provide optimized CUDA implementations and demonstrate their method’s effectiveness across three tasks: shallow water simulation, spherical image segmentation, and depth estimation. The spherical Transformers consistently outperform planar baselines, emphasizing the benefits of geometric priors for learning on spherical domains.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n1. The motivation is clear, and the problem is important.\n2. The paper is well written and clearly organized.\n3. Validation is done on three different tasks.\n\nWeaknesses:\n\n1. The literature review in the related work section is insufficient. It is mentioned in the introduction that spherical convolutional neural networks and geometric deep learning models have been used for learning on the sphere, the paper does not provide a detailed review or comparison of these prior works. Additionally, the relationship between the proposed method and neighborhood Transformers is not clear.\n2. The suitability of the Stanford 2D3DS dataset for the proposed approach is questionable. It is not clear what input representation was used during training with this dataset, raising concerns about whether the data aligns well with the model’s intended application.\n\nQUESTIONS:\n1. Are the $\\mathbb{R}^{2}$ and $S^{2}$ versions of Transformer, Transformer (local), SegFormer, SegFormer (local), and SphereUFormer trained using the same input data? If yes, is this input appropriate and optimal for validating the advantages of the proposed method?\n2. Can the authors clarify the input representation used for the Stanford 2D3DS dataset? Specifically, how is the data preprocessed or adapted to fit the spherical model’s input requirements?\n3. Given the abstract and introduction, there is an expectation to see validation on diverse domains such as atmospheric physics, cosmology, and robotics. Can the authors provide additional experimental results or validation on datasets from these or other relevant domains?\n4. How does the proposed method compare quantitatively and qualitatively to existing spherical CNNs and geometric deep learning models?\n5. How generalizable is the model to data with varying resolutions or noise characteristics?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors start with a continuous formulation of attention for the entire sphere as well as for neighborhoods defined by the geodesic distance. The continuous kernel integrals are discretized using quadrature and neighborhoods defined as above. The architecture is a modified ViT (but a Segformer can be used as well) having spherical convolutions for patch encoding, spherical multi-head attention, and a decoder with a bilinear spherical interpolation. Position encoding is spectral using the spherical harmonic basis. \nResults on the author's spherical version of ViT and Segformer outperform transformers defined on the regular grid of the sphere in segmentation, depth estimation, and the shallow water problem.\n\nSTRENGTHS AND WEAKNESSES:\nStrength: The paper handles in a correct way the discretization of the sphere by using neighborhoods defined by the geodesic thershold. Such a representation does not suffer from the singularity of the equirectangular grid or icosahedral domains.  \n\nWeaknesses: The paper has several weaknesses which I will list here and I would be eager to see the authors' responses and how the paper could be improved. \n\nW1: The authors claim as their main contribution the continuous formulation of the transformer. However, their implementation is the discrete sum of neighboorhood $k$ based on the geodesic distance. This is nothing else than a graph transformer with edges only between points where the geodesic distance is less than the $\\theta_{cutoff}$. \n\nW2: Related to W1, the authors misrepresent the literature in lines 81-84. All SO(3) equivariant transformers like the ones below have the spherical attention as a subcase, I will explain below.\n\nFuchs, F., Worrall, D., Fischer, V., & Welling, M. (2020). SE(3)-transformers: 3d roto-translation equivariant attention networks. Advances in neural information processing systems, 33, 1970-1981.\n\nand their implementation in Alphafold2.\n\nLiao, Yi-Lun, and Tess Smidt. \"Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs.\" The Eleventh International Conference on Learning Representations.\n\nW3: The authors ignore the processing of the positional encoding in the implementation.  The paper addresses only the scalar-valued version of an equivariant graph transformer. There are papers which combine scalar features (like the features in this paper) and spherical encodings like the equivariant perceiver I/O, for example, see 3.3.2 in:\n\nXu, Y., Chen, D., Liu, K., Zakharov, S., Ambrus, R. A., Daniilidis, K., & Guizilini, V. C. (2024). SE(3) Equivariant Ray Embeddings for Implicit Multi-View Depth Estimation. The Thirty-Eighth Annual Conference on Neural Information Processing Systems. \n\nW4: The notion of continuous as used in the paper it is not in the sense of a neural operator as cited in [25]. Requirement of a neural operator (see Table 1 in [25]) is to to \"query the output at any point\". That would be achieved by a decoder that has as an input a query spherical coordinate vector and conditioned on the attention output compute the segmentation or depth value at this point on the sphere. This is, for example, done in the equivariant attention decoder in \n\nChatzipantazis, E., Pertigkiozoglou, S., Dobriban, E., & Daniilidis, K. (2023). SE(3)-Equivariant Attention Networks for Shape Reconstruction in Function Space. Int. Conference Learning Representations ICLR. \n\nThe current paper though uses a bilinear interpolator in the decoder keeping a discrete structure. \n\nW5: The paper is limited to scalar valued functions failing to address applications like vector field prediction on the earth or even optical flow on images. In contrast, the SO(3) equivariant transformer papers can handle both vector  valued functions on any graph defined by the geodesic threshold as well as the processing of the spherical embedding. \n\nW7: The paper claims that the icosahedral discretization \"do not take into account the rotational symmetry of the sphere, but rather discrete permutations on the resulting graph\". If the icosahedral approaches would have used the same positional encoding as the authors here they would also account for rotation symmetry.\n\nW8: While the authors compare with ShapeUformer in the segmentation task, they do not do in the depth prediction although the ShapeUformer paper contains results (with different metrics).\n\nW9: ShapeUFormer contains depth estimation comparisons with several other approaches like Panoformer, EGformer, SFSS, HexRUnet, HealSWIN, Elite360D.\n\nIn summary, the paper lacks in novelty and technical soundness. It needs a significant enhancement to be elevated from a subcase of other papers to a paper competing in scope and properties with the state of the art.\n\nQUESTIONS:\nQ1 (and W3): The authors do not mention what they do with the positional encodings. If they add them or concatenate them to the features then the invariance formula in line 117 does not hold anymore because $q$ and $k$ do not behave like scalars (except the $L=0$ order harmonics). The rest have to be multiplied with a $W_q$ that is the product of a radial function times the spherical harmonics basis as described, in the references above and here: \nhttps://alchemybio.substack.com/p/spherical-equivariant-graph-transformer\n\nQ2 (W4): Clarify the notion of continuous.\n\nAddress all other weaknesses.\n\nLIMITATIONS:\nnone identified.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes an extension of the attention mechanism for spherical domains which is geometry aware and approximately rotationally equi-variant, enabling transformer architectures to process data natively defined on the 2D sphere. The authors derive continuous formulations for both global and neighborhood attention mechanisms on spherical domains, then discretize them using quadrature rules to ensure approximate SO(3) rotational equivariance. The method is implemented as neural operators that can handle arbitrary discretizations of the sphere. \n\nThe authors evaluate their approach on three diverse tasks: shallow water equation simulation, spherical image segmentation, and spherical depth estimation, demonstrating consistent improvements over euclidean counterparts. Popular architectures like ViT and SegFormer are systematically adapted for spherical domains by replacing standard attention mechanisms with either global spherical attention or neighborhood spherical attention. Additionally, the encoding and decoding laters are modified to handle spherical data by substituting conventional CNN operations with spherical convolutions, etc.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n- The paper provides a rigorous maths derivation by extending attention mechanisms to spherical domains through proper geometric treatment. The continuous formulation followed by quadrature-based discretization is mathematically sound.\n- The method achieves approximate SO(3) rotational equivariance, a crucial inductive bias for spherical data that standard Transformers lack. This addresses key limitations when applying planar methods to spherical domains.\n- Custom CUDA kernel implementations for both variants, full and local spherical attention will be made available, making the approach practically applicable.\n- The continuous formulation makes these architectures neural operators, enabling evaluation on arbitrary sphere discretizations while preserving geometric structure.\n- The evaluation maintains similar parameter counts and hyperparameters (embedding dimensions, layer depths) between spherical and Euclidean variants, ensuring fair comparisons that isolate the impact of geometric considerations.\n- The inclusion of sample outputs across all three tasks provides valuable visual evidence of improvements, particularly highlighting reduced distortions near poles and better preservation of spherical structure.\n- In Table 3 and Fig 4 - all S2 variants beat their R2 counterparts for shallow water equations on the rotating sphere\n\n**Weaknesses**\n- The paper lacks thorough computational complexity analysis and scalability comparisons, especially for the custom CUDA implementations versus optimized baselines. \n- In some experiments (e.g., segmentation), the improvements over Euclidean baselines are modest. The IoU improvements are often within error margins, raising questions about practical significance. The paper claims that “Across all models, we see the spherical variants outperforming their Euclidean counterparts” - this is not exactly true as S2 Transformers are worse in most metrics in Table 1.\n- The evaluation primarily compares against naive euclidean applications to spherical data. Beyond brief mentions of SphereUFormer, comprehensive comparisons with other geometric deep learning approaches would strengthen the claims.\n- The experiments are conducted on relatively small datasets (e.g., 1,621 images for segmentation with only 35-40 test samples), limiting the ability to assess generalisation and robustness. Larger-scale evaluations would provide more convincing evidence of the method's effectiveness.\n\nQUESTIONS:\n- In Table 1, why do S2 Transformers perform worse than R2 transformers for both, local and neighbourhood variants?\n- For comparison with SphereUFormer, why double the embedding dimension? - to equate for number of parameters? \n- How does the quadrature-based approximation quality degrade with grid resolution? Can you provide empirical evidence of the approximate SO(3) equivariance?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

comments: SUMMARY:
This paper revisits attention but on spherical data. Attention mechanisms are derived clearly and implemented successfully on spherical data. The method outperforms planar counter parts in clear evaluation of the method.

STRENGTHS AND WEAKNESSES:
+ The paper has clear theoretical contribution to attention with spheres
+ An implementation is made available and outperforms planar transformers in evaluation
+ Computational runtime is evaluated

- while the evaluation is fairly complete, the paper would benefit from ablation study which tests equivariant transformations of the sphere
- other spherical convolution method could enhance the experiments. For example 2D3DS is frequently reported in spherical convolution works.
- I am wondering if an icosahedron mesh could help with efficiency issues. See

[1] Chao Zhang, Stephan Liwicki, Sen He, William Smith, Roberto Cipolla:
HexNet: An Orientation-Aware Deep Learning Framework for Omni-Directional Input. IEEE Trans. Pattern Anal. Mach. Intell. 45(12): 14665-14681 (2023)
[2]  Y. Lee, J. Jeong, J. Yun, W. Cho, and K.-J. Yoon, “Spherephd: Applying cnns on a spherical polyhedron representation of 360deg
images,” in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2019, pp. 9181–9189.

QUESTIONS:
Minor additional ablation study could improve the paper.
Future work may investigate alternative ways to compute convolutions of the transformer.

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 4
IMPACT: 4
ORIGINALITY: 4
DATE: 2025-07-16