Attention on the Sphere

W1 - Review of Prior work We agree with the reviewer that our review of prior work was incomplete and could have been more extensive. Since the review, we have added additional references to equivariant graph transformers as suggested by reviewer WwPf ([1,2,3,4]) along side additional baselines (SphericalCNN[5] and EGFormer[6]) to the related work sections and results sections.

Our focus in the manuscript is on extending the popular attention mechanism used by transformer & segformer architectures to the spherical domain. Given this goal, we did not set out to beat the state of the art, but instead ensure that our technique is applicable across several image and simulation applications, with benefits over naive euclidean approaches. Despite this, we agree that baselines are helpful in estimating the performance of the approach and have added SphericalCNN[5] and EGFormer[6] to the segmentation and depth estimation problems.

W2 - Suitability of the 2D3DS dataset The 2D3DS dataset provides problems which require mapping spherical signals to other spherical signals, in contrast to other datasets which often deal with estimating pose or other features. This makes depth estimation and segmentation a natural choice given that the architecture is aimed at operator learning (i.e. mapping spherical signals to other spherical signals). Moreover, the 2D3DS dataset has been a popular choice for manuy related works such as [7].

For both segmentation and depth estimation, the input is an RGB image that was downsampled to 128x256 and 256x512 to reduce the computational cost and enable both non-local and non-spherical attention mechanisms as baselines, which otherwise become intractable. We agree that this could have been exposed better in the text and have edited the manuscript accordingly.

Q1 - Are all models trained with the same input/outputs. Is this appropriate to validate the method? Yes - all of the architectures are trained in the same manner to ensure comparability. Moreover the same loss formulation is used for all architectures, and architectures are kept as close to eachother in terms which makes this a the good test to evaluate the advantage of the spherical attention formulation.

One weakness of the dataset is that the segmentation and depth content closer to the poles was less variable. For the shallow water equations, a large source of visually noticeable artifacts arise at the poles, where our equivariant methods have no issues. Despite this we see advantages with our methods for all three datasets, which validates the approach.

Finally, the dataset has been used by many previous works giving a good comparison, at both a visual and quantitative (e.g., IoU) level for the applicability of our method. Metrics have been added to verify the validity of the method beyond just the validation loss.

Q2 - 2D3DS data preprocessing Yes - The dataset is downsampled from it’s original 4K representation to a more suitable resolution of 128x256 on an equiangular grid on the sphere. The target in the depth estimation case is a logarithm of the depth whereas in the segmentation case it is a class label enumerated by an integer number. The input data is z-score normalized to ensure proper training of the architectures. We agree that this should have been clarified better and have edited the manuscript accordingly

Q3 - Other application areas The depth estimation and segmentation datasets were chosen due to their relevance in robotics, whereas the Shallow Water Equation example was chosen due to the relevance of the dataset in geophysics and especially atmospheric physics. Regarding cosmology, we are not aware of a suitable image-based dataset, however there is a big body of work that deals with signal processing of spherical data in cosmology. Given the short period of time of the rebuttal phase we chose to focus on adding additional baselines and an equivariance test (see other reviews). We adapted the language in the manuscript to only mention cosmology in the context of possible application domains.

Q4 - Comparison against spherical CNNs We have implemented the SphericalCNN [5] architecture and the EGFormer [6] and added them as additional baselines to the tables provided below (Table 1 and 2). In our initial, time-limited, attempt to implement the baselines, the segmentation IoU performance is worse, which is visually notable in the segmentation results. Our technique relies on a suitable quadrature rule to enable numerical integration on the given input grid, and to use an icosohedral representation would require numerical integration rules for the different grid as well as updates to our neighborhood implementation to enable gathering of the icosohedron within the geodesic neighborhood. That said, one could imagine plugging into an adapted neighborhood attention mechanism into the SphericalCNN architecture after the input grid is projected to the icoshedral grid.

Table1: Updated table for the segmentation problem

Table 2 - Depth Estimation We have updated the depth estimation data with new baselines (EGFormer & SphericalCNN). Please find the updated values below:

Q5 - Varying input resolutions and noise As noted in Table 1 (segmentation results) we tried two resolutions with the local S2 segformer against the local R2 segformer, and both resolutions had similar evaluation scores to their lower resolution counterparts. The continuous formulation of the architecture makes the method applicable to any reasonable resolution and sampling scheme which allows accurate approximation of the integrals in the attention formulation. Moreover, our implementation of sparse attention (geodesic circle) on the sphere is equivalent with standard cartesian blocked neighborhood attention from a memory and computational complexity standpoint and thus scales favorably compared to full attention as resolution grows. Regarding performance with noise: The 2D3DS dataset is composed of RGB images of the real world and shows that our models are robust against noise found in non-synthetic images. That said, explicitly evaluating noise performance, was beyond the scope of our work, although we have no reason to expect it to perform differently than a traditional Transformer or Segformer model on euclidean data.

[1] Fuchs, 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.

[2] Liao, Yi-Lun, and Tess Smidt. "Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs." The Eleventh International Conference on Learning Representations.

[3] Xu, 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.

[4] Chatzipantazis, E., Pertigkiozoglou, S., Dobriban, E., & Daniilidis, K. (2023). SE(3)-Equivariant Attention Networks for Shape Reconstruction in Function Space. Int. Conference Learning Representations ICLR.

[5] 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.

[6] Zelin Zhang, Tao Zhang, KediLI, Xu Zheng. EGFormer. arXiv preprint arXiv:2505.14014, 2025

[7] Jeremy Ocampo, Matthew A. Price, Jason D. McEwen. Scalable and Equivariant Spherical CNNs by Discrete-Continuous (DISCO) Convolutions. arXiv preprint arXiv:2209.13603, 2022•arxiv.org

We thank the reviewer for the clarification.

Our work is motivated by the popularity of vision transformers in computer vision problems and Scientific ML problems such as weather prediction. Our aim is to present a generalization of the local and global attention mechanisms to the spherical domain, not to present a novel SOTA architecture for spherical domains. As such, we believe that the most appropriate baselines are comparable vision transformers, which we took as starting point for our transformers.

Whether the method can be expected to universally work better for certain types of spherical data than e.g. Spherical CNNs - this is impossible to say. We expect the method to work well in situations where vision transformers [1] and especially neighborhood attention transformers [2] are effective for euclidean data. This includes depth estimation, segmentation, classification etc. If a spherical vision transformer is desired, this paper describes how to achieve it and provides efficient implementations of the respective modules.

The components in our architecture are simply spherical counterparts to the operations encountered in the Vision Transformer [1] and Segformer. For instance, patch embeddings can be understood as a special case of a convolution, which is why we chose to replace them with spherical convolutional embeddings. As such, our architectures are kept as close to their euclidean counterparts as possible. Beyond this we are not aware of other ways to motivate or justify the individual components. We would also like to point out that works such as [1] do not really justify components beyond this, either.

[1] Dosovitskiy, Alexey, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani et al. "An image is worth 16x16 words: Transformers for image recognition at scale." arXiv preprint arXiv:2010.11929 (2020).

[2] Hassani, Ali, Steven Walton, Jiachen Li, Shen Li, and Humphrey Shi. "Neighborhood attention transformer." In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 6185-6194. 2023.

First of all, we would like to thank the reviewer for taking their time and providing us with this review. In the following we address the reviewers concerns:

W1 - Lack of novelty, only a graph transformer, no equivariance for vector-valued

The reviewer argues that this is only a graph transformer where edges are restricted to nodes within a certain geodesic distance from each other. This is not the case, as a graph transformer does not impose a quadrature rule on the summation imposed from a continuous formulation. A consequence of this is that a graph transformer will have permutation equivariance within the nodes inside this cutoff radius. This is not the case here - we break this permutational equivariance and get SO(3) equivariance - a notable difference and a strong inductive bias. We have emphasized this point in the manuscript to make it clearer.

Moreover, the reviewer cites a number of works that construct SE(3) equivariant transformers based on the equivariant SE(3) Graph attention. While the equivariant Graph attention does indeed ensure SE(3) and therefore SO(3) equivariance, a fundamental difference is that the Equiformer and related approaches mainly consider features/point clouds as inputs whereas we consider entire functions that are densely sampled on the sphere. Moreover the argument that this is a special case of aforementioned architectures does not hold. For instance, their depth-wise product (DTP) ensures equivariance by introducing a weight parameterized by the pairwise distance of features. In contrast , we make use of the proper discretization of the invariant Haar measure to obtain an equivariant formulation. This results in a pre-determined integration weight dependent on the absolute location. As such this formulation cannot be obtained as a special casse from the equiformer.

Finally, we would like to emphasize that there is significant value in realizing special cases of architectures and motivating them from a different theoretical perspective. Moreover, we provide efficient implementations and experimental evaluations that this approach is valid. As such this work presents a valuable addition to the literature, regardless of whether it can be derived from another architecture.

W2 - misrepresented related works, S2 attention as subcase of SE(3)

We thank the reviewer for bringing these works to our attention. We believe these works are relevant and should have been added to the manuscript. We have added [1,2,3,4] to the related work section of our manuscript.

W3 - Ignoring positional embedding, scalar-valued equivariant graph transformer

In contrast to equivariant graph transformers, we consider functions on the sphere and not features. A natural way of processing the positional embeddings in an equivariant manner is to treat them as inputs, therefore rotating them alongside the input if a rotation is applied. We have clarified this in the text. Regarding vector-valued functions see W5.

W4 - Clarify notion of continuous

The notion of continuous refers to all blocks being defined in the continuous domain, such that discretizations of the architecture can be derived at any resolution/grid that do not alter the parameterization. As such, the interpolation is not problematic as it is well-defined for any resolution, given that the discretizations are reasonable. There exist other Neural Operator architectures such as the Convolutional Neural Operator [5] which make heavy use of interpolation. Moreover, the decoder does not require this operation - it is added to avoid aliasing as is common in many modern convolutional architectures. We have clarified this in the manuscript.

W5 - Equivariance of vector-valued inputs

We do consider signals defined on the surface of the sphere, rather than features which denote positions or directions in euclidean space. As such rotations are applied as passive rotations to the input signals rather than to the outputs of the functions as well (as would be required). Upon popular request, we have added an equivariance test which demonstrates the approximate equivariance (see Table 3). This test considers vector-valued signals on the sphere which are rotated before the attention is applied and then rotated back. The L2 errors remain within a factor 2 of the interpolation incurred by rotating the signals and interpolating them onto the new grid and then reversing this operation. We have added this equivariance test alongside plots to visually demonstrate the equivariance of the attention mechanism (including vector-valued signals).

Table 3: Upon suggestion from multiple reviewers , we have added an equivariance test. This test takes a Gaussian random linear combination of spherical Harmonics to create a vector-valued noisy signal on the sphere, which is cut off at l=12, m=12.Equivariance errors are reported alongside interpolation errors that arise from applying the rotations discretely. All errors are reported in terms of the L2 loss on the sphere.

W6 - Positional embeddings could fix the icosahedral approaches

The icosahedral approaches that are referenced in the paper do use graph-based attention and rely on its properties to achieve equivariance[6]. However, this is limited to discrete rotations which map nodes of the icosahedron back to nodes of the icosahedron, irrespective of the positional embeddings. This is stated in their work [6]. In contrast, our method only achieves approximate equivariance (see Table 3), however for arbitrary rotations in SO(3).

W7 - SphereUFormer not applied to depth

We took the results for SphereUFormer directly from their paper and used only metrics which are comparable. However, we have added two additional baselines to the depth example.

W9 - Lack of baselines compared to SphereUFormer

Upon request from the reviewers we have added SphericalCNN [2] alongside EGFormer [3] to the results section (See Table 2 and Table 1 in he rebuttal for reviewer nUh3). We would like to emphasize that the goal of the paper is not to introduce a novel architecture but rather demonstrate how the attention mechanism is correctly adapted to spherical data. This is why in our experiments we use the same transformer architectures formulated on $R^2$ and demonstrate improvements that can be achieved by formulating the attention mechanism correctly. We expect that this “building block” may find its way into other architectures dealing with spherical data.

Table 2 - Depth Estimation We have updated the depth estimation data with new baselines (EGFormer & SphericalCNN). Please find the updated values below:

References

[1] Fuchs, 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.

[2] Liao, Yi-Lun, and Tess Smidt. "Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs." The Eleventh International Conference on Learning Representations.

[3] Xu, 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.

[4] Chatzipantazis, E., Pertigkiozoglou, S., Dobriban, E., & Daniilidis, K. (2023). SE(3)-Equivariant Attention Networks for Shape Reconstruction in Function Space. Int. Conference Learning Representations ICLR.

[5] Bogdan Raonić, Roberto Molinaro, Tim De Ryck, Tobias Rohner, Francesca Bartolucci, Rima Alaifari, Siddhartha Mishra, Emmanuel de Bézenac. Convolutional Neural Operators for robust and accurate learning of PDEs. Advances in Neural Information Processing Systems, 2023

[6] Sungmin Cho, Raehyuk Jung, Junseok Kwon. Spherical Transformer . arXiv preprint arXiv:2202.04942, 2022•arxiv.org

[7] Taco S. Cohen, Mario Geiger, Jonas Koehler, Max Welling. Spherical CNNs. arXiv preprint arXiv:1801.10130, 2018

[8] Zelin Zhang, Tao Zhang, KediLI, Xu Zheng.

First of all we would like to thank the reviewer for taking their time to review our manuscript and providing us with feedback.

W1 - Complexity & performance analysis At time of submission, the CUDA implementation was still too immature for serious complexity and performance experiments, particularly against optimized baselines like NATTEN or PyTorch’s full attention implementation. Since the submission, we have reworked the parallelization strategy to enable inline softmax and much improved coalesced memory access patterns such that the complexity benefit compared to full attention can be realized. Experiments done for this review show that updated local vs full S2 segformer is about 1.5x faster in our tests. Our implementation is made publicly available and we aim to provide timing information in the revised version of the manuscript.

W2 - Discuss segmentation and depth similarity across image models We have softened the language for the performance of the S2 models on the segmentation and depth datasets, which are within the experimental bounds of the R2 models. Moreover, we have discovered a bug in the implementation of the standard S2 transformer, which has made the results more consistent. Given the out-performance of the S2 models on the shallow water equations, we assume this similar performance is a result of the lack of detail at the poles of the spherical images, which doesn’t allow the S2 models to outperform the R2 models on depth and segmentation particularly at the poles.

W3 - More comparisons against other spherical techniques While the focus of the paper is to prove the effectiveness of the spherical attention formulation, we agree that more baselines could be added. Upon suggestion we have added an implementation of Spherical CNNs and EGFormer to our segmentation and depth estimation datasets.

Table1: Updated table for the segmentation problem

Table 2 - Depth Estimation We have updated the depth estimation data with new baselines (EGFormer & SphericalCNN). Please find the updated values below:

W4 - 2D3DS dataset size The lack of images is definitely a weakness of the 2D3DS dataset. However, we find that numerous works in the literature use this dataset to evaluate the effectiveness of spherical architectures. We chose to keep our architectures small as the goal of the paper was to evaluate the effectiveness of the spherical attention formulation over the euclidean counterparts, and not to develop a SOTA architecture. We believe that for this limited setting, this is sufficient.

Q1 - why do S2 transformers perform worse than R2 transformers Since the submission of the manuscript, we have discovered a bug in the implementation of the S2 transformer, which has been fixed. Updated results are shown in Table 1. Both S2 and R2 perform very similarly, especially considering the standard deviation.

Q2 - Why double the embedding dimension in the SphereUformer comparison The results for SphereUFormer were taken directly from their paper and they were carried out on double the resolution. To bring our architecture into a similar parameter count, we doubled the embedding dimension.

Q3 - How does quadrature scale with mesh resolution, and empirical test of SO(3) approximate equivariance

We thank the reviewer for this excellent suggestion. On equiangular grids such as for 2D3DS, we use the Clenshaw-Curtis quadrature rule, whereas on the Gauss grids (Shallow Water equations), we use Gauss quadrature. Theory suggests that both of these methods have an error that bounded by $\mathcal{O}((2 N )^{-k} / k)$ for a k-times differentiable integrand. We have added a test to estimate the approximate equivariance in Table 3. In this test, equivariance is tested by rotating smooth signals (degree up to l=12, m=12) discretely and applying the attention mechanism to the rotated signal before rotating it back. The result is compared to a direct application of the attention mechanism. We find that the error is dominated by the interpolation error which is incurred when the signal is rotated discretely and observe that the total error is within 1.5x of the interpolation error incurred by the rotation. We have added a more detailed explanation of this test alongside the results to the manuscript.

First of all, we would like to thank the reviewer for taking their time and providing us with this review. In the following we address the reviewers concerns:

W1 - equivariant experiments

As suggested by the reviewer, we have added an equivariance test to the manuscript. The convergence study is reported in Table 1, which shows equivariance and interpolation convergence for both the full and neighborhood transformer as resolution is increased. We see that the error is dominated by the interpolation error incurred when rotating the signal numerically.

Table 1 - Equivariance test Upon suggestion from multiple reviewers , we have added an equivariance test. This test takes a Gaussian random linear combination of spherical Harmonics to create a vector-valued noisy signal on the sphere, which is cut off at l=12, m=12.Equivariance errors are reported alongside interpolation errors that arise from applying the rotations discretely. All errors are reported in terms of the L2 loss on the sphere.

W2 - spherical convolution & other baselines

We were able to run the suggested SphericalCNN[1] (icosohedral projection) model on our depth and segmentation datasets, with the results presented in the revised Tables 2 & 3. We additionally implemented the EGFormer model for the segmentation dataset with results also presented in Tables 2 & 3. That said, with the limited time for review, it is likely that both models could be improved. For example, we noted that our EGFormer[2] model experienced significant overfitting when training.

Table 2 - Segmentation For segmentation, we added two new baselines (EGFormer & Spherical CNN) and reran several S2 and R2 attention-based models to ensure comparability across the models.

Table 3 - Depth Estimation We have updated the depth estimation data with new baselines (EGFormer & SphericalCNN). Please find the updated values below:

W3 - icosohedral mesh

The choice of an equiangular discretization has several advantages, including allowing direct comparison against standard R2 attention methods, ease of visualization, and direct use of the data, which is frequently given in an equiangular format. Since submission, we have worked on our CUDA implementation with an improved parallelization strategy that enabled online softmax and coalesced channel memory accesses, which have dramatically improved the forward and backward kernel performance. That said, nothing rules out the use of an icosahedron mesh except for finding an appropriate quadrature rule to enable the discretized spatial integration, and adapting the implementation to use indirect addressing to gather the mesh elements within the geodesic neighborhood. The lower density of points at the poles of the sphere could offset the performance losses due to indirect addressing lookups.

Q - Alternate ways of convolutions of the transformer

If the reviewer is referring to the choice of the encoding / decoding approaches in the spherical transformer, they were chosen to ensure equivariance throughout the architecture. In particular, because spherical convolutional patch embeddings can be seen as sparse convolution on the sphere.

References

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

[2] Zelin Zhang, Tao Zhang, KediLI, Xu Zheng. EGFormer. arXiv preprint arXiv:2505.14014, 2025