SO(3)-RoPE for Spherical Transformers

Many scientific domains produce data that lives on the sphere. Most spherical transformer architectures violate the spherical geometry, introducing distortions and coordinate singularities at the poles. We develop a RoPE-like position embedding, depending only on relative spherical locations. They act via unitary SO3 representations and lead to SO3-equivariant attention. Our SO3-RoPE is compatible with FlashAttention and thus retains the throughput and memory efficiency of vanilla transformers. We evaluate our method on a shallow water dynamics prediction task on a rotating sphere. It outperforms an S2Transformer baseline while having lower wall-clock time.