Spiral RoPE: Rotate Your Rotary Positional Embeddings in the 2D Plane

TL;DR

Spiral RoPE addresses a key limitation of Axial 2D RoPE by distributing rotary frequencies across multiple directions in the 2D plane, enabling oblique spatial reasoning in vision transformers. It preserves the relative-position encoding property of RoPE while using a grouped interleaved frequency allocation to maintain the same frequency budget. Across ImageNet-1k classification, ADE20k semantic segmentation, and diffusion-based image generation, Spiral RoPE yields consistent improvements and better extrapolation to higher resolutions. Qualitative analyses of attention maps and Fourier reconstructions show more localized activations and richer spatial representations, suggesting that multi-directional positional encoding is a general, effective architectural improvement for vision transformers.

Abstract

Rotary Position Embedding (RoPE) is the de facto positional encoding in large language models due to its ability to encode relative positions and support length extrapolation. When adapted to vision transformers, the standard axial formulation decomposes two-dimensional spatial positions into horizontal and vertical components, implicitly restricting positional encoding to axis-aligned directions. We identify this directional constraint as a fundamental limitation of the standard axial 2D RoPE, which hinders the modeling of oblique spatial relationships that naturally exist in natural images. To overcome this limitation, we propose Spiral RoPE, a simple yet effective extension that enables multi-directional positional encoding by partitioning embedding channels into multiple groups associated with uniformly distributed directions. Each group is rotated according to the projection of the patch position onto its corresponding direction, allowing spatial relationships to be encoded beyond the horizontal and vertical axes. Across a wide range of vision tasks including classification, segmentation, and generation, Spiral RoPE consistently improves performance. Qualitative analysis of attention maps further show that Spiral RoPE exhibits more concentrated activations on semantically relevant objects and better respects local object boundaries, highlighting the importance of multi-directional positional encoding in vision transformers.

Figures (16)

• Figure 1: Frequency support visualization comparing Axial 2D RoPE (left) and Spiral RoPE (right) with eight rotation groups ($K=8$). As shown, the Axial 2D RoPE places all frequencies on horizontal and vertical axes only, while our Spiral RoPE distributes frequencies across multiple directions in a spiral pattern, offering a broader directional coverage.
• Figure 2: Fourier reconstruction comparisons. Given binary input images (left), we retain only frequencies representable by each RoPE method and reconstruct via inverse-FFT. Axial 2D RoPE (middle) produces artifacts along horizontal and vertical axes. Spiral RoPE achieves more faithful reconstruction due to broader directional coverage.
• Figure 3: Query-token attention visualization. The query token (marked by a red box) is located on a foreground object. Across diverse scenes, Spiral RoPE consistently produces sharper and more localized attention around the queried region compared to APE and Axial RoPE, indicating improved spatial alignment.
• Figure 4: Multi-resolution performance on ImageNet-1k. Models are trained at $224 \times 224$ and evaluated across resolutions from $144 \times 144$ to $512 \times 512$. Spiral RoPE demonstrates superior robustness to resolution changes compared to APE and RoPE-Mixed baselines across both model sizes.
• Figure 5: Generated samples from our DiT-XL/2 model with Spiral RoPE trained for 7M steps on ImageNet $256 \times 256$. Images are generated using classifier-free guidance with scale 4.0.