Source-linked AI summary

Rotary Position Embedding for Vision Transformer

Byeongho Heo, Song Park, Dongyoon Han, Sangdoo Yun

arXiv:2403.13298v2cs.CVcs.LG

TL;DR

Position embeddings are crucial for Transformers, yet existing approaches can struggle when Vision Transformers change image resolution, while RoPE's vision effects remain underexplored. This paper develops and evaluates practical 2D RoPE variants, including learnable mixed-axis RoPE-Mixed, and finds improved high-resolution and downstream vision performance with negligible computation.

  • Problem

    Existing position embeddings struggle with resolution changes, and RoPE's effectiveness in vision tasks has received limited investigation.

  • Method

    The paper investigates 1D-to-2D RoPE extensions for ViTs and proposes RoPE-Mixed, which uses learnable mixed-axis frequencies for diagonal directions.

  • Results

    2D RoPE improves multi-resolution classification and downstream object detection and semantic segmentation, while RoPE-Mixed further outperforms conventional 2D RoPE across tasks.

  • Takeaways & Limitations

    RoPE provides an effective position-embedding option for vision transformers, with inference overhead of only 1.8M FLOPs, or 0.01% of ViT-B's 17.6G FLOPs.

  • Takeaways & Limitations

    Axial frequencies cannot handle diagonal directions because their frequencies depend on only one axis.

Abstract

from arXiv · show

Rotary Position Embedding (RoPE) performs remarkably on language models, especially for length extrapolation of Transformers. However, the impacts of RoPE on computer vision domains have been underexplored, even though RoPE appears capable of enhancing Vision Transformer (ViT) performance in a way similar to the language domain. This study provides a comprehensive analysis of RoPE when applied to ViTs, utilizing practical implementations of RoPE for 2D vision data. The analysis reveals that RoPE demonstrates impressive extrapolation performance, i.e., maintaining precision while increasing image resolution at inference. It eventually leads to performance improvement for ImageNet-1k, COCO detection, and ADE-20k segmentation. We believe this study provides thorough guidelines to apply RoPE into ViT, promising improved backbone performance with minimal extra computational overhead. Our code and pre-trained models are available at https://github.com/naver-ai/rope-vit

1 Introduction

Vision Transformers depend on position embeddings because self-attention is permutation-invariant, but existing absolute and relative schemes struggle when image resolution changes. This paper investigates 2D RoPE for vision and reports benefits across classification, detection, and segmentation.

  • Self-attention treats tokens without regard to index, so position embeddings provide the spatial information needed to distinguish token locations.
  • Absolute Positional Embedding uses absolute token locations, whereas Relative Position Bias adds relative biases to self-attention matrices.
  • Existing position embeddings work at fixed resolution but struggle when inference resolution changes, which can degrade dense prediction performance.
  • The paper expands RoPE from 1D to 2D vision data and proposes RoPE-Mixed with learnable mixed-axis frequencies to handle diagonal directions.
  • 2D RoPE is evaluated with ViT and Swin Transformer on ImageNet-1k classification, MS-COCO detection, and ADE20k semantic segmentation, improving high-resolution performance.

2 Related Works

Related work covers positional embeddings and earlier applications of RoPE in vision architectures. The paper distinguishes its general-purpose evaluation from prior methods requiring specialized training or targeting language-related and generative tasks.

  • ViT commonly uses learnable Absolute Positional Embedding, while hierarchical ViTs such as Swin Transformer generally use Relative Position Bias.
  • LaPE improves positional embeddings through simple scaling with adaptive layer normalization across various networks.
  • Earlier RoPE studies used 1D or axial 2D variants in ViT-related, multimodal, diffusion, and language-aligned models, with limited alignment to basic vision benchmarks.
  • CAPE and ResFormer investigate multi-resolution ViT inference using position-embedding augmentation or depth-wise convolution-based relative position embedding.
  • RoPE targets multi-resolution improvement without special training methods, preserving compatibility with existing training recipes.

3 Method

The method reviews conventional positional embeddings and extends Rotary Position Embedding from one-dimensional sequences to two-dimensional vision inputs. It uses mixed learnable frequencies to represent diagonal relative directions while retaining RoPE’s relative-position formulation.

  • Conventional positional embeddings: APE adds positional embeddings to patch tokens, whereas RPB adds learned relative-position biases directly to each self-attention matrix.APE is commonly used in ViT, while RPB is widely used in hierarchical ViTs.
  • Rotary Position Embedding: RoPE applies position-dependent rotations to query and key vectors, injecting relative token positions directly into attention similarities.The resulting attention term contains the relative displacement n − m in rotational form.
  • Rotary Position Embedding: RoPE uses multiple frequencies across query and key channel dimensions to encode relative positions through periodic rotations.The frequency schedule is θ_t = 10000^(-t / (d_head/2)).
  • 2D RoPE: Axial frequencies provide a simple 2D expansion but cannot represent diagonal directions because each frequency depends on only one spatial axis.The limitation arises because RoPE already uses query-key multiplication to encode relative distance, leaving no way to mix axial frequencies for diagonals.
  • Mixed learnable frequency: Mixed frequencies combine x- and y-axis components in the rotation matrix and attention formulation, enabling diagonal-direction handling.The mixed learnable-frequency implementation uses separate frequency sets for each head and self-attention layer.
  • Mixed learnable frequency: The mixed learnable-frequency parameters add only approximately 0.01% of network parameters in ViT-B.This implementation makes RoPE learnable in the vision domain while keeping the parameter increase negligible.

2D Fourier analysis.

The analysis compares RoPE frequency designs through 2D Fourier reconstruction and examines how RoPE changes attention behavior across resolutions. Mixed frequencies yield sharper positional representations, while RoPE produces longer-range and more varied attention interactions, with negligible computation cost.

  • Fourier reconstruction: 2D FFT followed by inverse FFT with limited RoPE frequencies evaluates the representational differences between RoPE-Axial and RoPE-Mixed.Because only dhead/2 frequencies are available, reconstructions are imperfect and reveal each design’s representation pattern.
  • Fourier reconstruction: Axial frequencies produce artifacts along axial lines, whereas Mixed frequencies use diverse 2D frequencies to produce sharper positional representations.The comparison uses RoPE-Mixed frequencies from ViT-B trained on ImageNet-1k.
  • Attention analysis: RoPE increases attention distance and entropy at middle layers, especially when inference resolution differs from the training resolution.Higher entropy indicates that more tokens participate in attention, while attention distance measures average spatial interaction distance.
  • Computation costs: RoPE requires only 1.8M FLOPs for ViT-B, or 0.01% of its 17.6G FLOPs, because the rotation matrix is pre-computed.The Hadamard product is the only inference-time computation required.

4 Experiments

Experiments evaluate 2D RoPE variants in ViT and Swin Transformer architectures across multi-resolution classification, object detection, and semantic segmentation. RoPE generally improves extrapolation and downstream performance, with variant and task differences in the value of combining conventional position embeddings.

  • Experimental setup: 2D RoPE variants are evaluated in ViT and Swin Transformer architectures across ImageNet-1k classification, MS-COCO detection, and ADE20k segmentation.ViT uses APE while Swin Transformer uses RPB, allowing RoPE to be assessed as a replacement for either conventional embedding.
  • Multi-resolution classification: For ViT classification, both RoPE-Axial and RoPE-Mixed outperform APE at resolutions larger than 224, while RoPE-Mixed performs better at all input resolutions.The models are trained at 224 × 224, so larger resolutions constitute extrapolation cases.
  • Conventional position embeddings: With ViT, APE improves interpolation and nearly doubles RoPE-Mixed’s improvement there, but reduces extrapolation gains; with Swin, RPB provides no advantage over RoPE-Mixed.The paper therefore presents the preferable combination as dependent on the target resolution and task.
  • Object detection: DINO-ViTDet gains more than +1.0pp box AP from replacing positional embeddings with RoPE, with RoPE-Mixed reaching +1.8pp for ViT-B and ViT-L.DINO-Swin also outperforms RPB with RoPE, but gains are smaller because its pretrained window size provides no extrapolation.
  • Semantic segmentation: RoPE improves semantic segmentation across ViT-UperNet and Swin-Mask2Former, while Mixed+APE achieves +2.3 and +2.5 mIoU in ViT-UperNet.Mixed+APE is best across the reported ViT-UperNet cases, whereas +RPB benefits Swin-based segmentation only in limited cases.

5 Conclusion

The study concludes that 2D RoPE effectively supports multi-resolution classification in ViT and Swin Transformers, improves downstream detection and segmentation, and that RoPE-Mixed further enhances performance.

  • 2D RoPE is effective for multi-resolution classification in both ViT and Swin Transformers.
  • 2D RoPE improves object detection and semantic segmentation performance with a significant margin.
  • RoPE-Mixed outperforms conventional 2D RoPE across various tasks.RoPE-Mixed uses mixed axis frequency with learnable parameters.

A Experiments (cont’d)

The appendix reports additional ablations and complete multi-resolution classification results across position-embedding configurations for ViTs and Swin Transformers.

  • The appendix evaluates default position embeddings, 2D RoPE variants, and combinations of 2D RoPE with default embeddings.The evaluated variants include APE or RPB, RoPE-Axial, RoPE-Mixed, and their combined forms.
  • The experiments cover both ViTs and Swin Transformers under multi-resolution classification settings.
  • Additional ablation studies and complete performance numbers are provided for the multi-resolution experiments.

A.1 Impacts of learnable frequencies

Learnable frequencies improve RoPE-Axial at high resolutions but are ineffective at other tested resolutions, while the appendix compares these variants for ViT-S and ViT-B.

  • Learnable frequencies improve RoPE-Axial performance at high resolutions 384 and 512.
  • Learnable frequencies are ineffective for RoPE-Axial at other tested resolutions.
  • Tables A.1 and A.2 compare learnable RoPE-Axial, fixed RoPE-Axial, and RoPE-Mixed for ViT-S and ViT-B.

A.2 Multi-resolution classification – ViT

The appendix reports complete multi-resolution classification results for ViT-S, ViT-B, and ViT-L, including comparisons among 2D RoPE variants and APE.

  • 2D RoPE variants outperform APE at the smallest tested resolution, 96 × 96.The appendix reports a significant performance gap at this resolution.
  • Tables A.3, A.4, and A.5 report multi-resolution performance for ViT-S, ViT-B, and ViT-L, respectively.These tables correspond to the first, second, and third graphs in Figure 4.

A.3 Multi-resolution classification – Swin Transformer

Across Swin Transformer sizes, 2D RoPE variants significantly outperform relative position bias at small resolutions, while the appendix tables provide the corresponding multi-resolution results and comparisons against base embeddings.

  • Multi-resolution performance: 2D RoPE variants significantly outperform RPB at 96 × 96 and 128 × 128 resolutions for Swin Transformer classification.This pattern is reported across Swin-T, Swin-S, and Swin-B results corresponding to Figure 5.
  • Multi-resolution performance: Tables A.6, A.7, and A.8 report multi-resolution performance for Swin-T, Swin-S, and Swin-B, respectively.Each table corresponds to one of the three graphs in Figure 5.
  • Comparisons with base embeddings: Figure 6 reports RoPE-Mixed performance improvement over APE or RPB, with the values expressed as percentage-point improvements.Tables A.9 and A.10 provide the numerical values for the two comparisons.
  • Comparisons with base embeddings: Table A.9 gives ViT-B improvement compared to APE, while Table A.10 gives Swin-B improvement compared to RPB.These tables correspond to the left and right graphs of Figure 6, respectively.
Loading 2403.13298v2…