SO(3)-RoPE for Spherical Transformers

AI summary

Overview

Research area: Geometric deep learning — specifically equivariant transformer architectures for data on the sphere (spherical transformers), with a focus on positional encoding.

Technical level: Advanced. The paper assumes familiarity with group theory (SO(3), unitary representations, Wigner D-matrices), rotary positional embeddings, and equivariant neural networks.

Scope: The paper introduces SO(3)-RoPE, a relative positional embedding for spherical transformers derived from unitary SO(3) representations, and evaluates it in a model called SO3ViT on shallow water dynamics prediction over a rotating sphere.

What This Paper Is About

Transformers have no built-in sense of where their input tokens are located, so positional encodings are needed to supply that information. Existing spherical transformers either distort the sphere (for example, by using latitude–longitude grids that oversample the poles and have coordinate singularities there) or use absolute positional encodings that break rotational equivariance. The goal of this paper is to extend Rotary Positional Embeddings (RoPE) — which encode translational symmetry in sequences — to the rotation group SO(3), so that attention on spherical data respects the geometry of the sphere.

Key Contributions

1. SO(3)-RoPE, a relative positional embedding for spherical transformers that acts on tokens via a position-dependent unitary SO(3) representation, constructed as a block-diagonal operator of Wigner D-matrices $D^\ell(R)$ — the irreducible representations of SO(3) — chosen with each order $\ell$ appearing with multiplicity one, using as many blocks as the embedding dimension $c$ allows.

2. An equivariance theorem (Theorem 2) showing that dot-product attention with SO(3)-RoPE is SO(3)-equivariant when queries and keys are both embedded, because the embedded inner product depends only on the relative SO(3)-position $U^{-1}V$. The proof is given in Appendix E.

3. A FlashAttention-compatible design: because SO(3)-RoPE modifies queries and keys independently, the method can use the standard FlashAttention implementation in PyTorch out of the box, preserving throughput, memory efficiency, and scalability comparable to vanilla transformers.

4. The SO3ViT architecture and an empirical study on shallow water dynamics over a rotating sphere, patchifying an icosphere via lifting convolutions into tokens on SO(3) and unpatchifying with a final transposed (projection) convolution for dense prediction. SO3ViT is compared against an S2Transformer baseline at matched compute.

Main Findings

• Better accuracy–efficiency frontier: Across multiple trained configurations, SO3ViT achieves lower errors at comparable wall-clock times and lower wall-clock times at each error level than S2Transformer, forming a significantly improved Pareto frontier — despite S2Transformer using torch-harmonics implementations with customised CUDA kernels. The authors attribute SO3ViT's efficiency to FlashAttention and the icosphere discretisation.

• Token efficiency from uniform sampling: The S2Transformer discretises the input on a 64 × 128 equirectangular grid, giving 8,192 tokens, whereas SO3ViT tokenises on an icosphere with two subdivisions (320 faces), and its lifting layer produces 3 · 320 = 960 tokens. The equiangular grid oversamples the poles, wasting tokens; the icosphere samples uniformly, allowing fewer tokens. Coarser patchification was not possible for the S2Transformer, which otherwise exceeded the 40 GB memory limit at a batch size of 16.

• Lower memory and faster inference at matched accuracy: In the configuration table (Appendix B), S2Transformer configurations range from 0.53 M to 8.43 M parameters, using 8.01 GB to 31.92 GB of memory with inference times of 141.4 ms to 471.0 ms, and ℓ₁ errors from 0.027 to 0.031 and ℓ₂ errors from 0.035 to 0.040. SO3ViT configurations range from 0.98 M to 101.45 M parameters, using 0.74 GB to 17.37 GB with inference times of 6.1 ms to 109.9 ms, and ℓ₁ errors from 0.020 to 0.045 and ℓ₂ errors from 0.026 to 0.058. The largest SO3ViT configuration reported (depth 32, width 512, 101.45 M parameters) reaches ℓ₁ 0.020 and ℓ₂ 0.026.

• More homogeneous error distribution: Averaging error over 200 simulations per point and model, S2Transformer shows an inhomogeneous error distribution peaking around the poles, caused by its equirectangular grid and by a patch embedding that is merely SO(2)-equivariant rather than SO(3)-equivariant, which the paper says inevitably suffers topological discontinuities around the poles. SO3ViT shows a more symmetric error distribution, with a discernible icosahedral grid pattern attributable to slight triangle distortions when mapping the icosahedron to an icosphere. The two plots use different color ranges, and the largest error magnitude of SO3ViT corresponds roughly to the smallest error magnitude of S2Transformer.

• Scaling behaviour differs: Both models benefit from scaling parameters, but S2Transformer saturates at higher parameter counts. SO3ViT benefits more from scaling its embedding dimension than its depth, which the authors suggest may be because larger embedding dimensions allow a higher maximal order of Wigner D-matrices, and therefore more positional resolution in the SO(3)-RoPE embeddings.

• Exact discrete equivariance verified: The full model is exactly equivariant only to the icosahedral subgroup of SO(3), because the lifting layer relies on the icosphere. Numerically, a randomly initialised model agrees with its rotated counterpart up to floating-point precision for all 60 rotations of the icosahedral group (Appendix C).

• Handling broken symmetry: Because the Coriolis force of the rotating sphere breaks full SO(3) symmetry down to SO(2) rotations around the z-axis, the authors add an absolute longitude encoding, summing Legendre polynomials $[P^1(\cos\beta), \dots, P^c(\cos\beta)]$ to the spherical input signal.

• Not reported: The dataset size is not reported, as the shallow water data is generated on the fly using the same generation procedure and evaluation metrics as the S2Transformer paper. The numerical ℓ₁/ℓ₂ errors of the configurations plotted in the accuracy-versus-epoch-time figure are not given as numbers in the text, only those in the Appendix B table.

Methodology in Plain English

The researchers start from an analogy. Standard RoPE encodes position in a sequence by multiplying feature vectors by position-dependent rotation matrices — these are unitary representations of the 1D translation group, and they make the attention score between two tokens depend only on their relative offset. The sphere's relevant symmetry is not translation but 3D rotation, so the analogous objects are unitary representations of SO(3), which are the Wigner D-matrices.

The practical recipe: represent the spherical input on an icosphere, built by recursively subdividing an icosahedron's triangular faces. Divide the grid into M non-overlapping triangle patches and choose SO(3) poses such that triangular kernels land on every patch in each of three possible directions, yielding 3M tokens on SO(3) rather than on the sphere. Store each token's SO(3) position as Euler angles via $R(\alpha,\beta,\gamma) = R_z(\alpha) R_y(\beta) R_z(\gamma)$. Embed patches with lifting convolutions and unpatchify, for dense prediction, with a final transposed convolution using resolution-increasing triangular kernels. Because the patches are assumed roughly equilateral, this lifting can be done without interpolation.

In attention, multiply queries and keys by a block-diagonal matrix built from Wigner D-matrices, then run standard FlashAttention. Because only the attention layers mix information between tokens and all other layers act token-wise, and because the embedded attention score depends only on the relative pose $U^{-1}V$, the attention is equivariant. The experiments train many configurations of SO3ViT and S2Transformer for 100 epochs with the Adam optimiser and a ReduceOnPlateau scheduler (as in the S2Transformer paper) plus a cosine scheduler variant, using a grid search over depth [4, 8, 16, 32, 64], heads [1, 2, 4, 8], and embedding dimension [128, 256, 512], selecting the number of heads and scheduler per layer and embedding dimension based on the best ℓ₂ error, on a single A100 with 40 GB memory and batch size 16. Configurations consuming too much memory were skipped.

Why This Matters

Impact on research: The paper argues that it is the first positional encoding to extend RoPE beyond toric domains to the sphere, carrying an irreps-based explanation of RoPE's success into a new setting. It connects two previously separate threads — the lifting-and-irreps machinery of spherical CNNs and the positional-encoding design of transformers — and shows that an equivariant, geometry-aware design can simultaneously improve accuracy and wall-clock efficiency, rather than trading one for the other.

Real-world applications (drawn from the domains the paper lists as naturally spherical data):

• Climate and weather fields defined over the globe, which is exactly the shallow water prediction setting used here.
• The cosmic microwave background in cosmology.
• Omnidirectional and panoramic imagery, plus data from fisheye or LiDAR sensors.
• Molecular surfaces.

Industry relevance: Any pipeline that must model global geophysical or sensor data on a sphere benefits directly: weather and climate forecasting services, robotics and autonomous systems processing omnidirectional or LiDAR input, and scientific computing groups running large simulations. The FlashAttention compatibility matters commercially because it means an equivariant spherical transformer can be dropped into existing hardware-optimised attention stacks and scaled up, rather than requiring bespoke kernels.

Future Directions

• Ensuring exact continuous SO(3) equivariance of the full model rather than only icosahedral equivariance, since the lifting layer's reliance on the icosphere currently breaks continuous symmetry to a 60-element discrete subgroup.
• Multi-step rollout evaluation: the reported benchmark is next-timestep prediction, so the stability and error accumulation of SO3-RoPE models over longer autoregressive forecasts is not addressed.
• Investigating the observed scaling behaviour further, specifically why SO3ViT benefits more from embedding dimension than depth and whether this is indeed tied to the maximal Wigner D-matrix order fitting into the embedding dimension.
• Comparing against the alternative discretisation choices the paper discusses but does not use — HEALPix pixelisation and equirectangular approaches such as PanoFormer, EGformer, PanoSwin, SphereUFormer and HEAL-SWIN — and exploring whether the SO(3)-irrep perspective on RoPE extends to other symmetry groups and homogeneous spaces beyond the sphere.

Target Audience

Researchers and graduate students in geometric deep learning and equivariant neural networks; machine-learning practitioners working on scientific computing, Earth system modelling, and physics simulation; transformer architecture researchers interested in positional encoding design beyond sequences and images; and engineers building models for spherical or rotationally symmetric sensor data who need both accuracy and FlashAttention-level throughput. Readers without a background in group theory or representation theory will need the Appendix definitions (continuous attention, lifting and projection convolutions, unitary representations) before the core argument is fully accessible.

Authors’ 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.