Paper deep dive
SO(3)-RoPE for Spherical Transformers
Christian Libner, Chase van de Geijn, Alexander S. Ecker, Maurice Weiler
Intelligence
Status: not_run | Model: - | Prompt: - | Confidence: 0%
Entities (0)
Relation Signals (0)
No relation signals yet.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Spherical data arise in many scientific applications. Often spherical transformers disregard the geometry of the underlying spherical domain, causing distortions and coordinate singularities near the poles. We introduce SO(3)-RoPE, a relative positional embedding that incorporates spherical geometry into transformer attention through unitary SO(3) representations. Our formulation is SO(3)-equivariant and compatible with FlashAttention, retaining efficiency of vanilla transformers. On shallow water dynamics prediction over a rotating sphere, our SO3ViT outperforms an S2Transformer baseline with lower errors and reduced runtime.
Tags
Links
- Source: https://arxiv.org/abs/2610.06229v1
- Canonical: https://arxiv.org/abs/2610.06229v1
Trouble viewing inline? Open PDF directly →
Full Text
49,193 characters extracted from source content.
Expand or collapse full text
Representations for the Physical Sciences SO(3)-RoPE for Spherical Transformers Christian Libner Chase van de Geijn Alexander S. Ecker Maurice Weiler Affiliation: Institute of Computer Science & CIDAS, University of Göttingen Affiliation: Max Planck Institute for Dynamics and Self-Organization, Göttingen Affiliation: MIT CSAIL Abstract Spherical data arise in many scientific applications. Often spherical transformers disregard the geometry of the underlying spherical domain, causing distortions and coordinate singularities near the poles. We introduce SO(3)SO(3)-RoPE, a relative positional embedding that incorporates spherical geometry into transformer attention through unitary SO(3)SO(3) representations. Our formulation is SO(3)SO(3)-equivariant and compatible with FlashAttention, retaining efficiency of vanilla transformers. On shallow water dynamics prediction over a rotating sphere, our SO3ViT outperforms an S2Transformer baseline with lower errors and reduced runtime. 1 Introduction Transformers have become a standard architecture across machine learning [46, 19, 39, 25]. Unlike convolutional networks, pure attention operations are not informed about spatial locations of their input tokens. Positional encodings provide this missing information. A widely used approach for encoding relative positions in sequences and Euclidean spaces are Rotary Positional Embeddings (RoPE) [43, 28]. We extend this idea to spherical transformers with SO(3)-RoPE, a relative positional encoding that incorporates the geometry of spherical data into attention. Our approach builds on the established line of work on spherical convolutional neural networks [11, 49]. A common operation in spherical CNNs is to lift feature maps from S2S^2 to the rotation group SO(3)SO(3), where features depend both on their spatial location and orientation. We employ such lifting convolutions as patch-embeddings of dense spherical data to tokens on SO(3)SO(3). SO(3)SO(3)-RoPE incorporates the geometry of these tokens into self-attention, analogously to how conventional RoPE incorporates Euclidean geometry. Euclidean RoPE acts via complex phase modulation (or rotations of ×22\!×\!2-blocks of features), which are unitary translation group representations. Analogously, SO(3)SO(3)-RoPE acts via unitary representations of the rotation group, specifically via Wigner D-matrices Dℓ(g)D (g), the irreducible representations of SO(3)SO(3). Because SO(3)-RoPE modifies queries and keys independently, our method can take advantage of hardware-optimised attention kernels like FlashAttention [17] out of the box. As a consequence, it preserves similar throughput, memory efficiency and scalability as vanilla transformers. We evaluate our model SO3ViT (Fig. 1) on next-timestep prediction for shallow water dynamics on a rotating sphere. Compared to an S2Transformer [5] baseline — a non-SO(3)SO(3)-equivariant spherical transformer architecture — we are able to scale to larger models and achieve significantly lower prediction errors, while drastically reducing runtime. Moreover, S2Transformer suffers from high errors due to geometry distortion around the poles. Incorporating the correct spherical geometry, our SO3ViT model shows a more uniform error distribution. The following sections introduce SO(3)SO(3)-RoPE (Sec. 2), our implementation (Sec. 3), and experiments (Sec. 4). More related work, mathematical definitions, and theorems are in the Appendix. Figure 1: SO3ViT. We patchify the sphere and process those patches with a transformer. In the attention we embed the queries and keys with our SO(3)SO(3) RoPE. 2 SO(3) Rotary Position Embedding Our transformers process feature maps f:S2→ℝcf:S^2 ^c that lie on the surface of the sphere S2S^2. For instance, our experiments predict the dynamics of the shallow water equation on S2S^2, whose state is given by spherical velocity fields and fluid heights. In practice, the input signal often comes discretised on a point cloud or grid, and is represented by a feature tensor F∈ℝN×cF ^N× c and Euclidean positions P∈ℝN×3P ^N× 3 on the unit sphere S2⊂ℝ3S^2 ^3. A key property of the spherical geometry is its rotational symmetry. Our networks are designed to be SO(3)SO(3) rotation equivariant – this ensures that they respect the spherical geometry, treating tokens at any locations equivalently, independent from their absolute positions. Intuitively, SO(3)SO(3)-equivariance means that moving a spherical input signal results in an equivalent movement of the predicted output. Formally, such movements are given by SO(3)SO(3) group actions [R⊳f](s):=f(R−1s)[R f](s):=f(R^-1s) on spherical signals, where R∈SO(3)R (3). An operator or network T is called SO(3)SO(3) equivariant if it commutes with these actions, that is, if [R⊳f]=R⊳[f]T[R f]=R [f] for any R∈SO(3)R (3). As in standard ViTs, we tokenise the sphere via convolutional patch embeddings. Following spherical CNNs, we embed spherical signals via so-called lifting convolutions (Def. 4) [11, 31, 48]. Applying a spherical kernel in any SO(3)SO(3)-pose, these convolutions result in tokens on the group SO(3)SO(3). In our implementation, we sample a discrete set of SO(3)SO(3) poses for patch tokenisation. All our tokens are therefore defined on the lifted domain SO(3)SO(3) rather than the input signal domain S2S^2. To inform the attention operation about spherical geometry we positionally embed our tokens. Following RoPE, we multiply token features f(R)f(R) at R∈SO(3)R (3) by a position-dependent unitary matrix π(R)∈ℝc×cπ(R) ^c× c. In standard RoPE one uses position-dependent ×22\!×\!2 rotation matrices, corresponding to the unitary representations of the 1D translational group. Analogously, SO(3)SO(3)-RoPE acts via unitary SO(3)SO(3) representations π. Due to unitarity and the composition property we get π(U)⊤π(V)=π(U)−1π(V)=π(U−1V)for any U,V∈SO(3).π(U) π(V)\ =\ π(U)^-1π(V)\ =\ π(U^-1V) for any U,V (3). (1) Definition 1 (SO(3)SO(3)-RoPE). Let f:SO(3)→ℝcf:SO(3) ^c be tokens on SO(3)SO(3) and π:SO(3)→ℝc×cπ:SO(3) ^c× c any unitary SO(3)SO(3)-representation (Def. 6). Our SO(3)SO(3)-RoPE embedding operator is given by [f](R):=π(R)f(R)for any R∈SO(3) [f ](R)\ :=\ π(R)\,f(R) for any \ \ R (3) (2) As in conventional RoPE, we embed our queries and keys, however, in our case on SO(3)SO(3). The attention inner product between query q(U)q(U) at U and key k(V)k(V) at V depends then only on their relative SO(3)SO(3)-position U−1VU^-1V: [q](U)⊤[k](V)=Def. 1q(U)⊤π(U)⊤π(V)k(V)=(1)q(U)⊤π(U−1V)k(V) [q ](U) \, [k ](V)\ Def.~ def:rope=\ q(U) π(U) \,π(V)\,k(V)\ ( eqn:group_rep)=\ q(U) \,π(U^-1V)\,k(V) (3) As a result, rotating the queries R⊳qR q and keys R⊳kR k leads to accordingly rotated attention scores: [R⊳q](U)⊤[R⊳k](V) [R q ](U) \, [R k ](V) =(3)[R⊳q](U)⊤π(U−1V)[R⊳k](V) ( eqn:grope_rel)=\ [R q ](U) \,π(U^-1V)\, [R k ](V) (4) =q(R−1U)⊤π(U−1V)k(R−1V) =\ q(R^-1U) \,π(U^-1V)\,k(R^-1V) =q(R−1U)⊤π(U−1RR−1V)k(R−1V) =\ q(R^-1U) \,π(U^-1R\,R^-1V)\,k(R^-1V) =(3)[q](R−1U)⊤[k](R−1V) ( eqn:grope_rel)=\ [q ](R^-1U) \, [k ](R^-1V) Therefore our SO(3)SO(3)-RoPE embedding produces a rotation-equivariant attention. Theorem 2 (SO(3)SO(3)-RoPE Attention is equivariant). Let q,k,v:SO(3)→ℝcq,k,v:SO(3) ^c be query, key and value functions, and attn[q,k,v]attn[q,k,v] be the dot-product attention operator (Def. 3). Then attn[[R⊳q],[R⊳k],R⊳v]=R⊳attn[[q],[k],v]for any R∈SO(3).attn [ [R q],\, [R k],\,R v ]\ =\ R [ [q],\, [k],\,v ] for any \ \ R (3). (5) Theorem 2 is proved in Appendix E. 3 Implementation Patch embedding & Lifting. Our spherical input signals are sampled on an icosphere, constructed by recursively subdividing the triangular faces of an icosahedron [24, 56, 14]. This produces a hierarchical, nearly uniform discretisation of S2S^2. To embed patches, we convolve triangular kernels on the icospherical grid, reducing resolution to a few tokens and lifting to SO(3)SO(3). Specifically, we use the grid’s hierarchical structure to divide it into M non-overlapping triangle patches and pick the corresponding SO(3)SO(3)-poses R1,…,R3MR_1,…,R_3M such that triangle kernels land exactly on every patch in each of three possible directions, thus resulting in 3M3M tokens. Assuming our triangle patches are roughly equilateral allows us to compute this lifting convolution without interpolation. We store the tokens’ SO(3)SO(3) positions R(α,β,γ)=Rz(α)Ry(β)Rz(γ)R(α,β,γ)=R_z(α)R_y(β)R_z(γ) via their Euler angles α,β,γα,β,γ. Unpatchification. For dense predictions on the original icospherical grid we apply a final transposed convolution. It unpatchifies tokens via resolution-increasing triangular kernels (Def. 5). SO(3) RoPE. As unitary representation π:SO(3)→ℝc×cπ:SO(3) ^c× c for token embedding we use a block-diagonal operator built from Wigner D-matrices Dℓ(R)∈ℝ(2l+1)×(2l+1)D (R) ^(2l+1)×(2l+1) [51, 48]. We pick each order l with multiplicity one, choosing as many blocks as our embedding dimension c allows: π(R(α,β,γ))=(D0(α,β,γ)0⋯00D1(α,β,γ)⋯0⋱00⋯DL(α,β,γ))π (R(α,β,γ) )\ =\ pmatrixD^0(α,β,γ)&0&·s&0\\ 0&D^1(α,β,γ)&·s&0\\ & & & \\ 0&0&·s&D^L(α,β,γ) pmatrix (6) Irreducible representations are the smallest possible building blocks, allowing for a cheap encoding. After applying SO(3)SO(3)-RoPE, we use the standard FlashAttention [17] implementation in PyTorch to compute self-attention [38]. Equivariance. Only attention layers mix information between tokens; all other layers act token-wise. With SO(3)SO(3)-RoPE, attention depends only on relative poses and is therefore equivariant to continuous rotations in SO(3)SO(3) (Theorem 2). The lifting layer, however, relies on the icosphere, so the full model is exactly equivariant only to the icosahedral subgroup of SO(3)SO(3), which is commonly used to discretize rotations [10, 14]. We verify numerically that a randomly initialised model is equivariant to all 60 rotations up to floating-point precision (Appendix C). Symmetry breaking by longitude encoding. Not all spherical problems are fully SO(3)SO(3)-equivariant – instead, the symmetry might be broken to rotations around the z-axis [1, 2, 4]. In such cases, we use an additive absolute position embedding of the longitude, summed to spherical input signals. Specifically, the longitude angle β is encoded via Legendre polynomials [P1(cosβ),…,Pc(cosβ)]⊤[P^1( β),...,P^c( β)] . 4 Experiments 4.1 Setup Dataset. We choose the shallow water equation on a rotating sphere as our benchmark, using the same on-the-fly data generation and evaluation metrics proposed in [5]. Due to the sphere’s rotation there is a Coriolis force. It breaks the full SO(3)SO(3) symmetry of the sphere down to the subgroup of SO(2)SO(2) rotations around the z-axis. Baseline. As a baseline model we pick the S2Transformer [5] with local neighborhood attention. DISCO convolutions [37] patchify the sphere to tokens that lie on an equiangular grid. Positions are embedded via additive spherical harmonics embeddings — this is an absolute position encoding, thus breaking equivariance. Training. We train for 100 epochs, using the Adam optimiser [27] and a ReduceOnPlateau scheduler, as in Bonev et al. [5]. In addition, we test the models also with a cosine scheduler. We do a grid search for the number of layers [4,8,16,32,64], the number of heads [1,2,4,8] and the embedding dimension [128,256,512]. Per layer and embedding dimension we pick the number of heads and scheduler based on the best ℓ2 _2 error. We train on a single A100 with 40GB memory and a batch size of 16. Configuration consuming too much memory are skipped. 4.2 Results Figure 2: The ℓ2 _2 error of predicting the next simulation timestep vs. the median epoch time. SO3ViT outperforms S2Transformer [5] given the same computational budget. Figure 3: SO3ViT has lower and more homogeneously distributed prediction errors than S2Transformer. Accuracy-Efficiency Trade-off. We train multiple configurations of S2Transformer and SO3ViT. Figure 3 plots the trained models’ ℓ2 _2 prediction errors versus their median epoch times. Our models achieve lower errors at comparable wall-clock times, and lower wall-clock times at each error level, forming a significantly improved Pareto frontier. This is despite the fact that the S2Transformer uses the torch-harmonics [4, 5] implementations with customised CUDA kernels. We attribute the efficiency of SO3ViT to the use of FlashAttention and the icosphere discretisation. The equiangular discretisation of the S2Transformer oversamples the poles resulting in wasted tokens, whereas our icosphere discretisation samples the sphere uniformly, allowing us to use fewer tokens. SO3ViT also requires less memory and achieves faster inference than S2Transformer. We provide a detailed analysis in Appendix B. Both models benefit from scaling parameters, although S2Transformer saturates at higher parameter counts. SO3ViT benefits more from scaling its embedding dimension rather than scaling its depth. This might be due to larger embedding dimensions allowing for a higher maximal order of Wigner D-matrices, and therefore positional resolution in our SO(3)SO(3)-RoPE embeddings. Error distribution. Next we investigate the spatial structure of prediction errors, visualized in Figure 3. For each point and model, we compute the average error over 200 simulations. S2Transformer shows an inhomogeneous error distribution, peaking around the poles. This is caused by two design choices: First, they operate on an equirectangular grid, which is inhomogeneous with diverging sampling density towards the poles. Moreover, their patch embedding utilizes a variant of DISCO convolutions which is merely SO(2)SO(2)-equivariant rather than SO(3)SO(3)-equivariant. As discussed in [48], this approach suffers inevitably from topological discontinuities around the poles. Our SO3ViT, on the other hand, shows a more symmetric error distribution. One can discern the icosahedral grid pattern, which is due to slight triangle distortions when mapping the icosahedron to an icosphere. Note that the two plots use a different color range, with the largest error magnitude of SO3ViT corresponding roughly to the smallest error magnitude of S2Transformer. These results underline the importance of building an architecture that truly respects the geometry of the data domain. 5 Conclusion We proposed SO(3)SO(3)-RoPE, a relative position encoding for spherical transformers, acting via unitary SO(3)SO(3)-irreps. It is FlashAttention-compatible and outperforms S2Transformer at equal compute. Acknowledgements. This research is funded by the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation), project number 501927736, within the DFG Priority Programme 2353: Daring More Intelligence - Design Assistants in Mechanics and Dynamics’. This project has received funding from the European Research Council (ERC) under the European Union’s Horizon Europe research and innovation programme (Grant agreement No. 101041669). The authors gratefully acknowledge the computing time granted by the Resource Allocation Board and provided on the supercomputer Emmy/Grete at NHR-Nord@Göttingen as part of the NHR infrastructure. The calculations for this research were conducted with computing resources under the project nii00194. References [1] J. Amato, Y. Xie, L. Medina-Varela, A. Aljerwi, A. McCutcheon, T. S. Rippentrop, K. Gonzalez, J. Delabrouille, M. Ishak, and N. Ruozzi (2025) CMB-ml: a cosmic microwave background dataset for the oldest possible computer vision task. 2025 IEEE/CVF International Conference on Computer Vision (ICCV), p. 9418–9430. Cited by: §3. [2] I. Armeni, S. Sax, A. R. Zamir, and S. Savarese (2017) Joint 2d-3d-semantic data for indoor scene understanding. arXiv preprint arXiv:1702.01105. Cited by: §3. [3] Y. Benny and L. Wolf (2024) SphereUFormer: a u-shaped transformer for spherical 360 perception. 2025 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), p. 940–950. Cited by: Appendix A, Appendix A, Appendix A. [4] B. Bonev, T. Kurth, C. Hundt, J. Pathak, M. Baust, K. Kashinath, and A. Anandkumar (2023) Spherical fourier neural operators: learning stable dynamics on the sphere. In International Conference on Machine Learning, Cited by: §3, §4.2. [5] B. Bonev, M. Rietmann, A. Paris, A. Carpentieri, and T. Kurth (2025) Attention on the sphere. ArXiv abs/2505.11157. Cited by: §1, Figure 3, §4.1, §4.1, §4.1, §4.2. [6] W. Boomsma and J. Frellsen (2017) Spherical convolutions and their application in molecular modelling. Advances in neural information processing systems 30. Cited by: Appendix A. [7] M. M. Bronstein, J. Bruna, T. Cohen, and P. Veličković (2021) Geometric deep learning: grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2104.13478. Cited by: Appendix A. [8] M. M. Bronstein, J. Bruna, Y. LeCun, A. Szlam, and P. Vandergheynst (2017) Geometric deep learning: going beyond euclidean data. IEEE Signal Processing Magazine 34 (4), p. 18–42. Cited by: Appendix A. [9] O. Carlsson, J. E. Gerken, H. Linander, H. Spiess, F. Ohlsson, C. Petersson, and D. Persson (2023) HEAL-swin: a vision transformer on the sphere. 2024 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), p. 6067–6077. Cited by: Appendix A, Appendix A. [10] G. Cesa, L. Lang, and M. Weiler (2022) A program to build E(N)-equivariant steerable CNNs. In International Conference on Learning Representations (ICLR), External Links: Link Cited by: §3. [11] T. Cohen, M. Geiger, J. Köhler, and M. Welling (2018) Spherical cnns. ArXiv abs/1801.10130. Cited by: Appendix A, Appendix A, §1, §2. [12] T. S. Cohen, M. Geiger, and M. Weiler (2019) A general theory of equivariant cnns on homogeneous spaces. Advances in neural information processing systems 32. Cited by: Appendix A. [13] T. S. Cohen and M. Welling (2016) Steerable cnns. arXiv preprint arXiv:1612.08498. Cited by: Appendix A. [14] T. Cohen, M. Weiler, B. Kicanaoglu, and M. Welling (2019) Gauge equivariant convolutional networks and the icosahedral cnn. ArXiv abs/1902.04615. Cited by: Appendix A, Appendix A, Appendix A, §3, §3. [15] T. Cohen and M. Welling (2016) Group equivariant convolutional networks. In International conference on machine learning, p. 2990–2999. Cited by: Appendix A. [16] B. Coors, A. P. Condurache, and A. Geiger (2018) Spherenet: learning spherical representations for detection and classification in omnidirectional images. In European conference on computer vision, p. 525–541. Cited by: Appendix A. [17] T. Dao, D. Fu, S. Ermon, A. Rudra, and C. Ré (2022) Flashattention: fast and memory-efficient exact attention with io-awareness. Advances in neural information processing systems 35, p. 16344–16359. Cited by: §1, §3. [18] M. Defferrard, M. Milani, F. Gusset, and N. Perraudin (2020) DeepSphere: a graph-based spherical cnn. ArXiv abs/2012.15000. Cited by: Appendix A. [19] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, et al. (2020) An image is worth 16x16 words: transformers for image recognition at scale. arXiv preprint arXiv:2010.11929. Cited by: §1. [20] J. R. Driscoll and D. M. Healy (1994) Computing fourier transforms and convolutions on the 2-sphere. Advances in applied mathematics 15 (2), p. 202–250. Cited by: Appendix A. [21] C. Esteves, C. Allen-Blanchette, A. Makadia, and K. Daniilidis (2018) Learning so (3) equivariant representations with spherical cnns. In Proceedings of the european conference on computer vision (ECCV), p. 52–68. Cited by: Appendix A. [22] K. M. Gorski, E. Hivon, A. J. Banday, B. D. Wandelt, F. K. Hansen, M. Reinecke, and M. Bartelmann (2005) HEALPix: a framework for high-resolution discretization and fast analysis of data distributed on the sphere. The Astrophysical Journal 622 (2), p. 759–771. Cited by: Appendix A, Appendix A. [23] M. M. Islam, R. Anand, D. R. Wessels, F. de Kruiff, T. P. Kuipers, R. Ying, C. I. S’anchez, S. P. Vadgama, G. Bökman, and E. J. Bekkers (2025) Platonic transformers: a solid choice for equivariance. ArXiv abs/2510.03511. Cited by: Appendix A. [24] C. M. Jiang, J. Huang, K. Kashinath, Prabhat, P. Marcus, and M. Niessner (2019) Spherical CNNs on unstructured grids. In International Conference on Learning Representations, External Links: Link Cited by: Appendix A, §3. [25] J. Jumper, R. Evans, A. Pritzel, T. Green, M. Figurnov, O. Ronneberger, K. Tunyasuvunakool, R. Bates, A. Žídek, A. Potapenko, et al. (2021) Highly accurate protein structure prediction with alphafold. nature 596 (7873), p. 583–589. Cited by: §1. [26] A. Kazemnejad, I. Padhi, K. Natesan Ramamurthy, P. Das, and S. Reddy (2023) The impact of positional encoding on length generalization in transformers. In Advances in Neural Information Processing Systems (NeurIPS), Vol. 36. Cited by: Appendix A. [27] D. P. Kingma and J. Ba (2014) Adam: a method for stochastic optimization. arXiv preprint arXiv:1412.6980. Cited by: §4.1. [28] K. Kogkalidis, J. Bernardy, and V. Garg (2024) Algebraic positional encodings. Advances in Neural Information Processing Systems 37, p. 34824–34845. Cited by: Appendix A, Appendix A, §1. [29] R. Kondor, Z. Lin, and S. Trivedi (2018) Clebsch–gordan nets: a fully fourier space spherical convolutional neural network. Advances in neural information processing systems 31. Cited by: Appendix A. [30] R. Kondor and S. Trivedi (2018) On the generalization of equivariance and convolution in neural networks to the action of compact groups. In International conference on machine learning, p. 2747–2755. Cited by: Appendix A. [31] R. Kondor and S. Trivedi (2018) On the generalization of equivariance and convolution in neural networks to the action of compact groups. In Proceedings of the 35th International Conference on Machine Learning, J. Dy and A. Krause (Eds.), Proceedings of Machine Learning Research, Vol. 80, p. 2747–2755. External Links: Link Cited by: §2. [32] P. J. Kostelec and D. N. Rockmore (2008) FFTs on the rotation group. Journal of Fourier analysis and applications 14 (2), p. 145–179. Cited by: Appendix A. [33] Y. Lee, J. Jeong, J. Yun, W. Cho, and K. Yoon (2019) SpherePHD: applying cnns on a spherical polyhedron representation of 360° images. In 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), p. 9173–9181. Cited by: Appendix A. [34] Z. Ling, Z. Xing, X. Zhou, M. Cao, and G. Zhou (2023) PanoSwin: a pano-style swin transformer for panorama understanding. In 2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), p. 17755–17764. Cited by: Appendix A, Appendix A. [35] H. Liu, L. Lin, J. Sun, Z. Shangguan, M. A. Alvarez, and H. Zhou (2025) Rethinking rope: a mathematical blueprint for n-dimensional positional embedding. arXiv preprint arXiv:2504.06308. Cited by: Appendix A, Appendix A. [36] T. Miyato, B. Jaeger, M. Welling, and A. Geiger (2024) Gta: a geometry-aware attention mechanism for multi-view transformers. In International Conference on Learning Representations, Vol. 2024, p. 8172–8208. Cited by: Appendix A. [37] J. Ocampo, M. A. Price, and J. D. McEwen (2022) Scalable and equivariant spherical cnns by discrete-continuous (disco) convolutions. ArXiv abs/2209.13603. Cited by: §4.1. [38] A. Paszke, S. Gross, S. Chintala, G. Chanan, E. Yang, Z. DeVito, Z. Lin, A. Desmaison, L. Antiga, and A. Lerer (2017) Automatic differentiation in pytorch. In NIPS-W, Cited by: §3. [39] A. Radford, J. W. Kim, T. Xu, G. Brockman, C. McLeavey, and I. Sutskever (2023) Robust speech recognition via large-scale weak supervision. In International conference on machine learning, p. 28492–28518. Cited by: §1. [40] S. Rasp, S. Hoyer, A. Merose, I. Langmore, P. Battaglia, T. Russell, A. Sanchez-Gonzalez, V. Yang, R. Carver, S. Agrawal, et al. (2024) WeatherBench 2: a benchmark for the next generation of data-driven global weather models. Journal of Advances in Modeling Earth Systems 16 (6), p. e2023MS004019. Cited by: Appendix A. [41] P. Shaw, J. Uszkoreit, and A. Vaswani (2018) Self-attention with relative position representations. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 2 (Short Papers), p. 464–468. Cited by: Appendix A. [42] Z. Shen, C. Lin, K. Liao, L. Nie, Z. Zheng, and Y. Zhao (2022) PanoFormer: panorama transformer for indoor 360° depth estimation. In European Conference on Computer Vision, Cited by: Appendix A, Appendix A. [43] J. Su, M. Ahmed, Y. Lu, S. Pan, W. Bo, and Y. Liu (2024) Roformer: enhanced transformer with rotary position embedding. Neurocomputing 568, p. 127063. Cited by: Appendix A, §1. [44] Y. Su and K. Grauman (2017) Learning spherical convolution for fast features from 360 imagery. Advances in neural information processing systems 30. Cited by: Appendix A. [45] C. van de Geijn, T. Lüddecke, P. Turishcheva, and A. Ecker (2026) A circular argument: does rope need to be equivariant for vision?. Advances in Neural Information Processing Systems 38, p. 20388–20424. Cited by: Appendix A, Appendix A. [46] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin (2017) Attention is all you need. Advances in neural information processing systems 30. Cited by: Appendix D, §1. [47] M. Weiler and G. Cesa (2019) General e (2)-equivariant steerable cnns. Advances in neural information processing systems 32. Cited by: Appendix A. [48] M. Weiler, P. Forré, E. Verlinde, and M. Welling (2023) Equivariant and Coordinate Independent Convolutional Networks. External Links: Link Cited by: §2, §3, §4.2. [49] M. Weiler, P. Forré, E. Verlinde, and M. Welling (2026) Equivariant and coordinate independent convolutional networks: a gauge field theory of neural networks. World Scientific. Cited by: Appendix A, §1. [50] M. Weiler, M. Geiger, M. Welling, W. Boomsma, and T. Cohen (2018) 3d steerable cnns: learning rotationally equivariant features in volumetric data. Advances in Neural information processing systems 31. Cited by: Appendix A. [51] E. Wigner (1931) Gruppentheorie und ihre anwendung auf die quantenmechanik der atomspektren. Cited by: §3. [52] B. Wu, A. Wan, X. Yue, and K. Keutzer (2018) Squeezeseg: convolutional neural nets with recurrent crf for real-time road-object segmentation from 3d lidar point cloud. In 2018 IEEE international conference on robotics and automation (ICRA), p. 1887–1893. Cited by: Appendix A. [53] Q. Yan, Q. Wang, K. Zhao, B. Li, X. Chu, and F. Deng (2022) Spheredepth: panorama depth estimation from spherical domain. In 2022 International Conference on 3D Vision (3DV), p. 1–10. Cited by: Appendix A. [54] S. Yogamani, C. Hughes, J. Horgan, G. Sistu, P. Varley, D. O’Dea, M. Uricár, S. Milz, M. Simon, K. Amende, et al. (2019) Woodscape: a multi-task, multi-camera fisheye dataset for autonomous driving. In Proceedings of the IEEE/CVF international conference on computer vision, p. 9308–9318. Cited by: Appendix A. [55] I. Yun, C. Shin, H. Lee, H. Lee, and C. E. Rhee (2023) Egformer: equirectangular geometry-biased transformer for 360 depth estimation. In 2023 IEEE/CVF International Conference on Computer Vision (ICCV), p. 6078–6089. Cited by: Appendix A, Appendix A. [56] C. Zhang, S. Liwicki, W. Smith, and R. Cipolla (2019) Orientation-aware semantic segmentation on icosahedron spheres. In 2019 IEEE/CVF International Conference on Computer Vision (ICCV), p. 3532–3540. Cited by: Appendix A, §3. [57] M. Zhdanov, A. Lucic, M. Welling, and J. van de Meent (2026) (Sparse) attention to the details: preserving spectral fidelity in ml-based weather forecasting models. ArXiv abs/2604.16429. Cited by: Appendix A. Appendix Appendix A Related Work / Background Signals that naturally live on the sphere arise throughout the sciences and engineering, from climate and weather fields defined over the globe [40] and the cosmic microwave background in cosmology [22], to omnidirectional and panoramic imagery [44, 16], molecular surfaces [6], and data from fisheye or LiDAR sensors [54, 52]. Processing such data faithfully requires the spherical geometry to be encoded into the model itself, rather than treated as an afterthought. In this section, we review three strands of work: the established theory of spherical CNNs, recent efforts to bring transformers to the sphere, and general advances in encoding geometry into transformers through positional encodings. Spherical convolutions. Geometric Deep Learning (GDL)[8, 49] studies how to incorporate data geometry into deep learning models. The guiding principle of GDL is that of the Erlangen programme[7] which characterizes geometries through group theory and symmetries. Thus, to faithfully represent the geometry, a model should respect the symmetries of its data. For a domain which is a sequence or line, such as in language, the most relevant symmetry is translation. Convolutions are the only linear operators that are translation-equivariant[30]. On the sphere, S2S^2, the relevant symmetries are 3D rotations; any rotation – or more formally any element of SO(3)SO(3), i.e. the space of 3×33× 3 special (unit determinant) orthogonal matrices – leaves S2S^2 invariant. Thus, the translations of the common convolution are replaced by rotations. The operator we seek is a linear map that is rotation-equivariant, which turns out to be a spherical convolution [11, 20]. A well-known obstacle is that a linear operator mapping a signal on S2S^2 back to a signal on S2S^2 can only be SO(3)SO(3)-equivariant when its kernel is isotropic, i.e. depends solely on the geodesic distance between points. Such zonal kernels are blind to orientation and therefore limit a model’s expressivity. Spherical CNNs overcome this by lifting features from the sphere to the rotation group[11, 21, 14]: the first layer maps a signal on S2S^2 to a signal on SO(3)SO(3), after which subsequent layers convolve on SO(3)SO(3) itself, where anisotropic and far more expressive kernels remain equivariant. Two complementary realisations of this idea are common. Group convolutions[15] enforce equivariance directly in the spatial domain by correlating over transformed copies of the kernel. Steerable convolutions[13, 50, 47, 12] instead operate in the spectral domain and are the group-theoretic counterpart of computing a convolution through the (generalized) Fourier transform, representing signals through the irreducible representations of the symmetry group. For SO(3)SO(3) these irreps are the Wigner-D matrices, with the corresponding transform expanding signals in spherical harmonics [29, 32]. We leverage both the early works’ insights on lifting and the latter works’ insights on irreps of SO(3)SO(3) to construct our positional encoding. Spherical Transformers. Just as transformers have largely displaced convolutions in vision, a growing line of work adapts them to spherical data. Before extending a transformer to the sphere, one must first decide how to discretize and represent it. While planar images allow for regular pixel grids at arbitrary resolutions, the only regular grids on the sphere correspond to Platonic solids and are of rather coarse resolution. To achieve higher resolutions, one needs to fall back to irregular sampling schemes: equirectangular (latitude–longitude) grids, icospheres obtained by repeatedly subdividing an icosahedron, or the equal-area HEALPix pixelisation [22] are all common choices. Each carries a different trade-off between sampling uniformity, hierarchical structure, and ease of implementation. In our work, we choose to operate on icosphere subdivisions. For each sampling scheme, several transformer architectures have been proposed. On the equirectangular grid, the main challenge is counteracting the projection’s polar distortion within the attention mechanism itself: PanoFormer [42] attends over tangent patches with a learnable token flow, EGformer [55] biases local attention along the panorama’s vertical and horizontal directions, and PanoSwin [34] adapts the shifted-window design of the Swin transformer to the wrap-around topology. These methods keep the convenience of a rectangular tensor, but must spend architectural effort to undo the distortion it introduces. Icosphere-based transformers avoid this distortion at the source by operating on a near-uniform mesh. SphereUFormer [3] builds a U-shaped hierarchy of local self-attention modules directly on the icosphere, using the recursive subdivision of the mesh to define the pooling and unpooling operations between resolutions. This continues a line of icosahedral networks [24, 14, 33, 56, 18, 53] that exploit the icosahedron’s high degree of symmetry, which we likewise build upon. Others adopt the HEALPix grid, an equal-area, hierarchically nested pixelization, which several transformer-based methods exploit [9, 57]. We do not follow this route: while HEALPix offers equal-area sampling and a convenient hierarchical layout, its pixelization lacks the symmetry of the icosphere, which we favor in this work to simplify lifting. Local attention A common thread across these architectures is a reliance on local attention, which restricts each query to a neighborhood of nearby tokens rather than the full sphere—both to control the quadratic cost of attention and to supply an inductive bias for locality. On the equirectangular grid this appears as deformable pixel neighborhoods in PanoFormer [42] and geometry-biased windows in EGformer [55], while PanoSwin [34], HEAL-SWIN [9]. SphereUFormer [3] likewise attends within mesh neighborhoods on the icosphere. By contrast, we control the cost by patchifying analogous to the vanilla ViT. This patchification is one of the reasons we must lift to SO(3)SO(3). The benefit is that the model captures local dependencies through the patches while the attention focuses on long range. Positional encodings. In a transformer, self-attention is permutation-equivariant and carries no intrinsic notion of position; geometry can therefore enter only through the positional encoding [28] or through geometry-dependent masking such as local or causal attention [26]. Most modern transformers rely on some form of relative positional encoding (RPE) [41], which in language corresponds to encoding a translational symmetry [36, 45, 23]. The dominant instance is the Rotary Positional Encoding (RoPE) [43], which has been shown to arise axiomatically from a small set of desirable assumptions on sequential data [28, 35]. A practical virtue of RoPE is that it is applied to the query and key vectors independently, without materializing a position-dependent attention matrix; this separability is what makes it cheap to compute and easy to extend. Extensions beyond one dimension have, however, focused almost exclusively on toric domains, yielding the widely used axial RoPE for N-dimensional inputs [28, 35, 45]. The irreps perspective that explains RoPE’s success has not been carried beyond this setting. We take up this perspective and use it to extend RoPE to the sphere—to our knowledge the first positional encoding to do so. This gap is also visible in existing spherical transformers. Because they typically resort to local attention for tractability [3], they must pair it with a local positional encoding, but these either fall back on an isotropic kernel—echoing the limited expressivity of the zonal kernels of spherical convolutions—or ignore the sphere’s gauge problem and thereby cause problems at the poles [14]. Local attention is moreover hardware-inefficient and restricted in expressivity, making such models both slower and less accurate than the patchification scheme we propose. Appendix B Memory & Inference speed Table 1: Peak GPU memory during training and inference time. For every S2Transformer configuration, some SO3ViT configuration performs better in ℓ1 _1 error, ℓ2 _2 error, memory, and inference speed. Model Depth Width Params ℓ1 _1 ℓ2 _2 Memory Inference Time S2Transformer 4 128 0.53 M 0.031 0.040 8.01 GB 141.4 ms S2Transformer 4 256 2.12 M 0.031 0.040 15.96 GB 198.8 ms S2Transformer 4 512 8.43 M 0.027 0.035 31.92 GB 282.1 ms S2Transformer 8 256 4.22 M 0.031 0.039 21.99 GB 471.0 ms SO3ViT 4 128 0.98 M 0.045 0.058 0.74 GB 6.1 ms SO3ViT 4 256 3.54 M 0.034 0.043 1.28 GB 9.3 ms SO3ViT 4 512 13.37 M 0.028 0.036 2.40 GB 16.2 ms SO3ViT 16 512 51.12 M 0.023 0.029 8.82 GB 55.2 ms SO3ViT 32 512 101.45 M 0.020 0.026 17.37 GB 109.9 ms We compare memory consumption and inference speed of the S2Transformer and SO3ViT (Tab. 1). At matched accuracy, SO3ViT uses less memory and runs faster, despite having more parameters. SO3ViT tokenises the input on an icosphere with two subdivisions (320 faces). The lifting layer then produces 3⋅320=9603· 320=960 tokens. The S2Transformer instead discretises the input on a 64×12864× 128 equirectangular grid, which gives 8,1928,192 tokens. Coarser patchification was not possible for the S2Transformer. Otherwise, it exceeded the 40 GB memory limit at a batch size of 16. The reason is the torch-harmonics implementation. It first upsamples and then applies DISCO convolutions at full output resolution. These convolutions store an interpolation matrix for every latitude, and its size grows with the kernel size. Our patchification and unpatchification avoid this cost, since they require no interpolation. This is possible because our attention mechanism allows a free choice of discretisation of SO(3)SO(3) (Sec. 3) and our assumption of equilateral triangles of the icosphere. Appendix C Equivariance Test Our SO(3)SO(3)-RoPE attention mechanism is equivariant to continuous rotations in SO(3)SO(3). We discretize our data on an icospherical grid, which breaks continuous symmetries to the subgroup of icosahedral symmetries. This preserves exact icosahedral group equivariance and approximate continuous SO(3)SO(3)-equivariance of the full model. To verify this, we sample a random input for a randomly initialised model and permute it according to each rotation of the icosahedral group. We then compare the output for the rotated input with the rotated output for the original input. For all 60 rotations, the two agree up to floating-point precision. Appendix D Definitions In order to avoid discretization artifacts and preserve symmetries, all of our definitions are formulated in the continuous setting. We use scaled dot product attention [46], replacing discrete sums over tokens by continuous integrals over SO(3)SO(3). Definition 3 (Continuous Attention). Let q,k,v:SO(3)→ℝcq,k,v:SO(3) ^c. We define scaled dot-product attention on SO(3)SO(3) as: attn[q,k,v](U)=∫SO(3)softmax[A[q,k]](U,V)v(V)Vfor any U∈SO(3),attn[q,k,v](U)\ =\ _SO(3) -24.0musoftmax\! [A[q,k] ](U,V)\;v(V)\;dV for any \ \ U (3), (7) with softmax[A](U,V)=eA(U,V)∫SO(3)eA(U,Q)Qsoftmax[A](U,V)= e^A(U,V) _SO(3)e^A(U,Q)\,dQ and the attention score between the query q(U)q(U) at U and key k(V)k(V) at V given by A[q,k](U,V)=q(U)⊤k(V)A[q,k](U,V)=q(U) \,k(V) (8) Note that we absorb the c-dependent scaling factor 1c 1 c into the queries and keys; in the implementation we scale the attention scores by it as usual. Our patch embedding from spherical signals to tokens on SO(3)SO(3) is done via usual lifting group convolutions. Definition 4 (Lifting convolution S2→SO(3)S^2 (3)). Let fin:S2→ℝcinf_in:S^2 ^c_in be a spherical signal and ψ↑:S2→ℝc×cin _ :S^2 ^c× c_in an arbitrary filter. The lifting convolution of finf_in with ψ↑ _ is given by: (ψ↑⋆fin)(R)=∫S2ψ↑(R−1s)fin(s)sfor any R∈SO(3),( _ f_in)(R)\ =\ _S^2 _ (R^-1s )\,f_in(s)\,ds for any \ \ R (3), (9) Intuitively, this operation applies a spherical kernel in every SO(3)SO(3)-rotated pose. Similarly, we use projection convolutions to map tokens on SO(3)SO(3) back to spherical signals: Definition 5 (Projection convolution SO(3)→S2SO(3)→ S^2). Let f:SO(3)→ℝcf:SO(3) ^c be a signal over the SO(3)SO(3) group and ψ↓:S2→ℝcout×c _ :S^2 ^c_out× c a filter. The projection convolution of f with ψ↓ _ is the map ψ↓⋆f:S2→ℝcout _ f:S^2 ^c_out given by (ψ↓⋆f)(s)=∫SO(3)ψ↓(R−1s)f(R)Rfor any s∈S2,( _ f)(s)\ =\ _SO(3) _ (R^-1s )\,f(R)\,dR for any \ \ s∈ S^2, (10) We define SO(3)SO(3)-RoPE as acting via unitary representations of the rotation group. For completeness, the following definition establishes what it means for a representation to be unitary: Definition 6 (Unitary SO(3)SO(3) representation). We call π:SO(3)→ℝc×cπ:SO(3) ^c× c a unitary group representation if: (I) π(U)⊤π(U)=Iπ(U) \,π(U)\,=\,I for all U∈SO(3)U (3) (unitary); (I) π(U)π(V)=π(UV)π(U)π(V)=π(UV) for all U,V∈SO(3)U,V (3) (group homomorphism). where I∈ℝc×cI ^c× c is the identity matrix. Appendix E Proofs Proof of Theorem 2. Let q,k,v:SO(3)→ℝcq,k,v:SO(3) ^c be the query, key and value functions, and let π:SO(3)→ℝc×cπ:SO(3) ^c× c be a unitary representation (Def. 6). We first record two elementary properties of unitary representations. The identity element e∈SO(3)e (3) is represented by the identity matrix I∈ℝc×cI ^c× c: π(U)π(e)=(I)π(Ue)=π(U) π(U)π(e)\ def:homomorphism=π(Ue)=\ π(U) (11) ⇒ π(U)⊤π(U)π(e)=π(U)⊤π(U) π(U) π(U)π(e)\ =\ π(U) π(U) =(I) def:unitary= π(e)=I π(e)\ =\ I Consequently, the inverse element U−1U^-1 is represented by the inverse matrix π(U)−1π(U)^-1: π(U−1)π(U)=(I)π(U−1U)=11π(e)=I π(U^-1)π(U)\ def:homomorphism=π(U^-1U)\ eqn:group_identity=\ π(e)=I (12) ⇒ π(U−1)=π(U)−1 π(U^-1)=π(U)^-1 Together, unitarity (I) and the composition property (I) give the main ingredient of the proof: the product π(U)⊤π(V)π(U) π(V) depends on U,V∈SO(3)U,V (3) only through their relative rotation U−1VU^-1V, π(U)⊤π(V)=(I)π(U)−1π(V)=12π(U−1)π(V)=(I)π(U−1V)π(U) \,π(V)\ def:unitary=\ π(U)^-1\,π(V)\ eqn:group_inverse=π(U^-1)\,π(V)\ def:homomorphism=\ π(U^-1V) (13) In particular, the product is unchanged when both U and V are rotated by the same R−1∈SO(3)R^-1 (3): π((R−1U)−1R−1V)=π(U−1RR−1V)=π(U−1V)π ((R^-1U)^-1\,R^-1V )=π(U^-1R^-1\ V)=π(U^-1V) (14) With this, the embedded attention score ARoPE[q,k]:=A[[q],[k]]A_RoPE[q,k]:=A [ [q], [k] ] is rotation-equivariant: for all R∈SO(3)R (3) and all U,V∈SO(3)U,V (3), ARoPE[R⊳q,R⊳k](U,V) A_RoPE[R q,\,R k](U,V) (15) =Def. 1[R⊳q](U)⊤[R⊳k](V) Def.~ def:rope=\ [R q](U) \, [R k](V) =q(R−1U)⊤π(U)⊤π(V)k(R−1V) =\,\,\ \ \ q(R^-1U) \,π(U) π(V)\ k(R^-1V) =(13)q(R−1U)⊤π(U−1V)k(R−1V) ( eqn:relative_pi)=\ \ \ q(R^-1U) \,π(U^-1V)\,k(R^-1V) =q(R−1U)⊤π(U−1RR−1V)k(R−1V) =\,\,\ \ \ q(R^-1U) \,π(U^-1R\,R^-1V)\ k(R^-1V) =(13)q(R−1U)⊤π(R−1U)⊤π(R−1V)k(R−1V) ( eqn:relative_pi)=\ \ \ q(R^-1U) \,π(R^-1U) \,π(R^-1V)\,k(R^-1V) =Def. 1A[[q],[k]](R−1U,R−1V) Def.~ def:rope=\ A [ [q], [k] ](R^-1U,R^-1V) =ARoPE[q,k](R−1U,R−1V). =\ \ \ \ A_RoPE[q,k](R^-1U,R^-1V). The softmax inherits this equivariance. Its denominator integrates over all of SO(3)SO(3), so we may substitute D:=R−1QD:=R^-1Q without changing the integral. softmax[ [ ARoPE[R⊳q,R⊳k]](U,V) A_RoPE[R q,\,R k] ](U,V) (16) =eARoPE[R⊳q,R⊳k](U,V)∫SO(3)eARoPE[R⊳q,R⊳k](U,Q)Q =\ \ \ \ \ \ \, e^A_RoPE[R q,\,R k](U,V) _SO(3)e^A_RoPE[R q,\,R k](U,Q)\,dQ =(15)eARoPE[q,k](R−1U,R−1V)∫SO(3)eARoPE[q,k](R−1U,R−1Q)Q ( eqn:attn_score_eqi)=\ \ \ \ \ \ e^A_RoPE[q,k](R^-1U,R^-1V) _SO(3)e^A_RoPE[q,k](R^-1U,R^-1Q)\,dQ =D:=R−1QeARoPE[q,k](R−1U,R−1V)∫SO(3)eARoPE[q,k](R−1U,D)D D:=R^-1Q=\ e^A_RoPE[q,k](R^-1U,R^-1V) _SO(3)e^A_RoPE[q,k](R^-1U,D)\,dD =softmax[ARoPE[q,k]](R−1U,R−1V) =\ \ \ \ \ \ \,softmax [A_RoPE[q,k] ](R^-1U,R^-1V) Similar, the RoPE-embedded attention is SO(3)SO(3)-equivariant in its argument U∈SO(3)U (3): attn[ [ [R⊳q],[R⊳k],R⊳v](U) [R q],\, [R k],\,R v ](U) (17) =Def. 3∫SO(3)softmax[ARoPE[R⊳q,R⊳k]](U,V)(R⊳v)(V)V Def.~ def:cont_attn=\ \ \ \ _SO(3)softmax [A_RoPE[R q,\,R k] ](U,V)\,(R v)(V)\,dV =(16)∫SO(3)softmax[ARoPE[q,k]](R−1U,R−1V)v(R−1V)V ( eqn:softmax_eqi)=\ \ \ \ \ \ _SO(3)softmax [A_RoPE[q,k] ](R^-1U,R^-1V)\,v(R^-1V)\,dV =D:=R−1V∫SO(3)softmax[ARoPE[q,k]](R−1U,D)v(D)dD D:=R^-1V=\ _SO(3)softmax [A_RoPE[q,k] ](R^-1U,D)\,v(D)\,dD =Def. 3attn[[q],[k],v](R−1U) Def.~ def:cont_attn=\ \ \ \ \ \ \ \ attn [ [q],\, [k],\,v ](R^-1U) =(R⊳attn[[q],[k],v])(U) = (R [ [q],\, [k],\,v ] )(U) ∎