Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Learning the RoPEs: Better 2D and 3D Position Encodings with STRING

Connor Schenck, Isaac Reid, Mithun George Jacob, Alex Bewley, Joshua Ainslie, David Rendleman, Deepali Jain, Mohit Sharma, Avinava Dubey, Ayzaan Wahid, Sumeet Singh, René Wagner, Tianli Ding, Chuyuan Fu, Arunkumar Byravan, Jake Varley, Alexey Gritsenko, Matthias Minderer, Dmitry Kalashnikov, Jonathan Tompson, Vikas Sindhwani, Krzysztof Choromanski

TL;DR

STRING generalizes rotary position encodings through a Lie-group framework to produce separable translationally invariant position encodings applicable to multi-dimensional token coordinates. It offers computationally efficient instantiations (Cayley-STRING, Circulant-STRING) and learnable variants, while subsuming RoPE as a special case. Empirically, STRING yields consistent improvements over RoPE and absolute encodings across ImageNet/Places365, open-vocabulary 2D/3D object detection, and both simulated and real-world robotic manipulation, including 3D depth-enabled tasks on KUKA robots. The work demonstrates that 3D-aware STRING encodings can be integrated with Vision Transformers to enhance generalization and robustness with favorable compute characteristics.

Abstract

We introduce STRING: Separable Translationally Invariant Position Encodings. STRING extends Rotary Position Encodings, a recently proposed and widely used algorithm in large language models, via a unifying theoretical framework. Importantly, STRING still provides exact translation invariance, including token coordinates of arbitrary dimensionality, whilst maintaining a low computational footprint. These properties are especially important in robotics, where efficient 3D token representation is key. We integrate STRING into Vision Transformers with RGB(-D) inputs (color plus optional depth), showing substantial gains, e.g. in open-vocabulary object detection and for robotics controllers. We complement our experiments with a rigorous mathematical analysis, proving the universality of our methods.

Learning the RoPEs: Better 2D and 3D Position Encodings with STRING

TL;DR

STRING generalizes rotary position encodings through a Lie-group framework to produce separable translationally invariant position encodings applicable to multi-dimensional token coordinates. It offers computationally efficient instantiations (Cayley-STRING, Circulant-STRING) and learnable variants, while subsuming RoPE as a special case. Empirically, STRING yields consistent improvements over RoPE and absolute encodings across ImageNet/Places365, open-vocabulary 2D/3D object detection, and both simulated and real-world robotic manipulation, including 3D depth-enabled tasks on KUKA robots. The work demonstrates that 3D-aware STRING encodings can be integrated with Vision Transformers to enhance generalization and robustness with favorable compute characteristics.

Abstract

We introduce STRING: Separable Translationally Invariant Position Encodings. STRING extends Rotary Position Encodings, a recently proposed and widely used algorithm in large language models, via a unifying theoretical framework. Importantly, STRING still provides exact translation invariance, including token coordinates of arbitrary dimensionality, whilst maintaining a low computational footprint. These properties are especially important in robotics, where efficient 3D token representation is key. We integrate STRING into Vision Transformers with RGB(-D) inputs (color plus optional depth), showing substantial gains, e.g. in open-vocabulary object detection and for robotics controllers. We complement our experiments with a rigorous mathematical analysis, proving the universality of our methods.

Paper Structure

This paper contains 43 sections, 4 theorems, 37 equations, 13 figures, 6 tables.

Table of Contents

  1. Introduction and Related Work
  2. Preliminaries
  3. STRING: Separable Translationally Invariant Position Encodings
  4. Introducing STRING
  5. Computationally efficient STRING
  6. Learnable frequencies with STRING.
  7. Loose ends
  8. Experiments
  9. General Experiments: Classification and Retrieval
  10. Improving Open-Vocabulary Object Detection
  11. Open-Vocabulary Object Detection in 2D
  12. Open-Vocabulary Object Detection in 3D
  13. Simulation-Based Robot Manipulation: ALOHA
  14. Real-World 3D Robot Manipulation: KUKA
  15. Setting
  16. ...and 28 more sections

Key Result

Theorem 3.2

Consider the set of mappings $\mathbf{R}(\cdot):\mathbb{R}^{d_c} \to \mathbb{R}^{d \times d}$ that satisfy the group-like translational invariance property $\mathbf{R}(\mathbf{r}_i)^\top \mathbf{R}(\mathbf{r}_j) = \mathbf{R}(\mathbf{r}_j - \mathbf{r}_i) \space \forall \space \mathbf{r}_i, \mathbf{r}

Figures (13)

  • Figure 1: Top: Successful diffusion policy conditioned on a STRING-enhanced Transformer vision encoder, attempting the double-insertion task on Aloha-sim. Bottom: Same experiment, but with a regular vision encoder for which the policy fails. STRING provides strong improvements for training dexterous robotics policies, outperforming previous position encoding algorithms such as RoPE.
  • Figure 2: Example outputs for the 3D detection task for baseline, RoPE, and Cayley-S. Green boxes: groundtruth. Blue boxes: predictions.
  • Figure 3: HandOverBanana task for the ALOHA 2 robot: real (top) and the corresponding simulated (bottom) evaluation.
  • Figure 4: Mean success rate across all tasks (i.e. MultiTask) evaluated 10 times every 10K train steps over 1M train steps.
  • Figure 5: Performance of STRING with 3D input vs. baselines on real-robot tasks (with 2 seeds). 2D baseline performance without depth input is $\approx65\%$. Incorporating depth through surface normal maps (nmap) reduces performance to $42\%$. Using 3D STRING for incorporating depth improves the performance in both scenarios - with and without normal maps to $53\%$ and $74\%$ respectively. Mean/stdev shown above were calculated from $35$ evaluation runs.
  • ...and 8 more figures

Theorems & Definitions (6)

  • Definition 3.1
  • Theorem 3.2: STRING is general
  • Theorem 3.3: RoPE is a type of STRING #1
  • Theorem 3.4: RoPE is a type of STRING #2
  • Theorem 3.5: Circulant-STRING is fast
  • proof