Rotary Position Embeddings on the Sphere for Global Weather Forecasting

Modern transformers rely on positional encodings in the attention mechanism to account for the positions of input tokens. Rotary Position Embeddings (RoPE) are among the most popular approaches to date for this. For transformers applied to sequences defined on the sphere, for example in machine learning-based weather forecasting, extending RoPE to $S^2$ is non-trivial. We therefore develop S2RoPE, a spherical rotary position encoding for queries and keys defined on the sphere. We propose two variants: channel-wise modulation by normalized complex spherical harmonics (S2RoPE-SH) and rotations within complete harmonic degree spaces (S2RoPE-Wigner). The first approach modulates both amplitude and phase, whereas the second one preserves feature norms and encodes the relative position of two tokens through a rotation between their local frames in one spherical harmonics band. We apply both forms to two global medium-range forecasting models, the WeatherGenerator and Stormer. We find improvements over baseline methods in standard metrics, particularly at short and medium lead times. In WeatherGenerator, we further find improvements in the preservation of fine-scale detail and less drift in the balance between wind and geopotential during extended rollouts.