3D Rotary Position Encoding for Transformers

• 3D-RPE is a positional encoding method that extends standard rotary embeddings to three spatial or spatiotemporal dimensions, preserving geometric relationships.
• It applies axis-specific and group-theoretic rotations to token representations, achieving translation invariance and enhanced inductive bias for complex data.
• 3D-RPE improves performance in medical imaging, 3D object detection, and long-context language modeling while keeping computational overhead minimal.

3D Rotary Position Encoding (3D-RPE) refers to a family of positional encoding mechanisms for Transformers in which the standard rotary embedding paradigm is extended to three spatial (or spatiotemporal) dimensions. In 3D-RPE, the core principle is to apply axis-specific or group-structured rotations to each channel-pair in token representations, using the full 3D position of each token (e.g., grid coordinates in volumes, (x,y,z) in medical scans, (w,h,t) in video, or spherical coordinates in geospatial modeling). This approach directly enables Transformers to capture relative relationships and geometric structure of 3D data, providing translation invariance, improved long-range context, and enhanced inductive bias for volumetric, video, and other high-dimensional input domains.

1. Mathematical Formulations of 3D-RPE

Three broad classes of 3D-RPE have emerged: axis-separable 2D/3D rotary schemes, group-theoretic (Lie algebraic) encodings, and spherical/coordinate-aware generalizations.

Axis-Separable 3D Rotary Embedding

The majority of practical 3D-RPEs generalize the original 1D RoPE by separately rotating channel-pairs for each of the three Cartesian axes. For input tokens at position $p_n = (x_n, y_n, z_n)$, partition the feature $q_n \in \mathbb{R}^d$ into thirds, apply frequency-parameterized 2D rotations along each axis, and compose the result multiplicatively (or additively in angle space):

• For axis $\alpha \in \{x,y,z\}$ and frequency $i$, the rotation angle is $\theta_{n,\alpha,i} = p_{n,\alpha} \cdot \omega_i$ with

$\omega_i = 1 / 10000^{2i/d_m}, \quad d_m = d/3$

• Each channel-pair is rotated by the $2 \times 2$ matrix

$$R(\theta) = \begin{bmatrix} \cos \theta & -\sin \theta \ \sin \theta & \cos \theta \end{bmatrix}$$

• The full rotated embedding is constructed as

$$q_n^{rot} = R^z(p_{n,z})[R^y(p_{n,y})(R^x(p_{n,x})q_n)]$$

and similarly for $k_n$ (Li et al., 18 Mar 2025, Feng et al., 24 Mar 2025).

Temporal and Multi-Scale Extensions

3D-RPE is often extended to spatiotemporal cases by defining a fourth rotation for time, with angles summed or partitioned by channel group:

$$\theta_i = \theta_x[i] + \theta_y[i] + \theta_z[i] + \theta_t[i]$$

Each pair $(q_{2i}, q_{2i+1})$ is rotated by $R(\theta_i)$ (Ji et al., 17 Apr 2025). Multi-scale 3D-RPE can further scale frequencies per resolution layer.

Lie Group/Algebraic Encodings

Advancing beyond channel-separable RoPE, LieRE and STRING generalize position encodings to the exponential map of skew-symmetric matrix generators:

$$R(r) = \exp(L_x x + L_y y + L_z z),\quad L_k^\top = -L_k,\ [L_i, L_j]=0$$

This encoding is unified, translation-invariant, and can be efficiently instantiated using fast basis change, circulant, or Cayley parametrization (Ostmeier et al., 14 Jun 2024, Schenck et al., 4 Feb 2025). The action on queries/keys is $q_i \mapsto R(r_i) q_i$, and the dot product reduces to a function of $r_j - r_i$.

Bloch Sphere and Spherical Encodings

For domains with inherent spherical geometry, as in geospatial tasks, the 3D-RPE is formulated via Euler-angle composite rotations (e.g., $R_z(\omega_i \lambda) R_x(\varphi)$), acting block-diagonally over feature channels. This preserves geodesic distances and ensures the group-composition property for SO(3) (Unlu, 2023, Ma et al., 14 Jun 2024).

2. Integration into Transformer Architectures

In all variants, the central integration point is the multi-head self-attention mechanism. The process is:

• Project tokens to $Q$, $K$, $V$ as in standard attention.
• Apply axis-aligned, group-theoretic, or spherical 3D-RPE transformations to $Q$, $K$ (but not $V$).
• Compute attention weights as $$A = \operatorname{Softmax}\left(\frac{Q^{rot} (K^{rot})^T}{\sqrt{d}}\right)$$.
• Aggregate outputs by $O = AV$.

In axis-separable methods, rotations are fused with channel grouping for GPU efficiency. In Lie group methods, rotations are computed via matrix exponentials. In spherical cases, rotations are mapped to each 3-channel block using the relevant spherical harmonics or Euler angles.

A summary of architecture-specific integration techniques appears below.

Approach	Key Integration	Parameter Overhead
Separable RoPE	Axis-based channel rotations (per-head, per-layer)	None
LieRE/STRING	Learnable so(3)/SO(d) generator for each coordinate	≈ 9–3d params/head
Spherical RoPE	Block-diagonal SO(3) rotations from lon/lat	None

3. Theoretical Properties: Translation Invariance and Group Structure

A distinguishing feature of 3D-RPE is that it yields translation invariance and relative position encoding natively in three (or more) dimensions.

• With $R(r)$ a group homomorphism (as in STRING),

$R(r_i)^\top R(r_j) = R(r_j - r_i).$

• The self-attention logit $q_i^T R(r_j - r_i) k_j$ then depends only on the 3D offset, matching the theoretical desiderata for translation-invariant, relative geometric bias (Schenck et al., 4 Feb 2025).
• In spherical encodings, the block-diagonal rotation scheme ensures that the attention kernel preserves (locally linear) spherical geodesic distances between token positions (Unlu, 2023).

In contrast, axis-separable 3D-RPEs approximate this property by composing independent rotations per axis, which works well in grid-like volumetric domains, though it slightly limits cross-axis expressivity compared to Lie/STRING encodings.

4. Practical Implementations and Computational Considerations

3D-RPEs introduce minimal computational and memory overhead, especially in axis-separable and block-diagonal incarnations.

• The per-token cost is $O(d)$ for standard rotary rotations (one $2\times2$ multiply per channel pair per axis), or $O(d^2)$ in less-optimized global rotation schemes.
• Frequencies are precomputed and rotations are vectorized for GPU throughput; Circulant-STRING achieves $O(d\log d)$ per token (Schenck et al., 4 Feb 2025).
• Patchification, coordinate normalization ([0,1] scaling), and channel interleaving or parallelization are standard for efficiency (Li et al., 18 Mar 2025, Ji et al., 17 Apr 2025).
• In frameworks like PyTorch, 3D-RPE is fused as a custom CUDA kernel, keeping FLOPs and memory increase typically under 2–5% for state-of-the-art volumetric models.

5. Empirical Impact Across Modalities

Extensive experimental results have established the advantages of 3D-RPE in various domains.

• Medical Volumetric Segmentation: RoMedFormer (MRI/CT segmentation) achieves $+3\%$ Dice improvements on small 3D structures and $~1.5$\% better than relative bias methods, with strong out-of-distribution shape priors and removal of boundary artifacts (Li et al., 18 Mar 2025).
• 3D Object and Video Detection: RoPETR's 3D-RPE (with spatiotemporal rotations) yields a 31% reduction in mean absolute velocity error (mAVE) in NuScenes detection, with an NDS increase from $67.6\%$ to $69.0\%$ (Ji et al., 17 Apr 2025).
• 3D Texture Synthesis: RomanTex demonstrates clear improvements in multi-view coherence, lowering Local Alignment Distance (LAD) and FID scores for texture consistency on 3D geometries (Feng et al., 24 Mar 2025).
• Long-context Language Modeling: 3D-RPE extends the context length of LLaMA2-7B models from $4$k to $100$k tokens while maintaining low perplexity ($\sim 7 - 8$), outperforming 2D RoPE-based baselines on NLU, summarization, few-shot, and code tasks (Ma et al., 14 Jun 2024).
• Spatiotemporal Reasoning in Video-LLMs: Naïve 3D partitioned RoPE-3D achieves minor baseline improvements; VRoPE further mitigates spatial bias for highly uniform attention in video+text, boosting retrieval accuracy from $54.8\%$ (RoPE) to $87.0\%$ (VRoPE) at 1024+ frames (Liu et al., 17 Feb 2025).
• General Vision/Robotics: STRING and LieRE consistently surpass both absolute and naive RoPE 3D baselines, yielding $\sim 1-2\%$ top-1 gains on 3D classification and significant robustness/transfer in open-vocabulary object detection and robotic manipulation (Ostmeier et al., 14 Jun 2024, Schenck et al., 4 Feb 2025).

6. Limitations, Design Choices, and Future Extensions

While 3D-RPE is highly general, several avenues and caveats are notable:

• Axis-separability vs. Expressivity: Standard 3D-RPEs are axis-parallel; more expressive group-guided encodings (LieRE/STRING) can offer cross-axis representation at marginal compute cost.
• Memory and Compute: LieRE and full-string variants introduce limited but nonzero matrix exponential cost per token, but are tractable with current hardware.
• Domain Adaptivity: Learned frequency schedules (RoPE-M), change-of-basis parametrizations, and per-head parameter sharing facilitate adaptation to diverse input scales and geometric structures (Schenck et al., 4 Feb 2025).
• Modal Integration and Continuity: For video-text fusion and geospatial applications, smoother index continuity and bias balancing (as in VRoPE or spherical RoPE) are essential for mitigating discontinuities and token bias (Liu et al., 17 Feb 2025, Unlu, 2023).
• Extension to Higher-Dimensional Coordinates: All approaches described generalize in principle to $D$-dimensional domains (e.g., spatiotemporal plus semantic axes), though computational structure and inductive bias should be revisited for $D > 3$.
• Empirical Open Questions: Optimal chunk size selection, adaptive or learned angular schedules, and deeper integration with sparse/local attention remain areas for future research (Ma et al., 14 Jun 2024).

7. Comparative Summary of Approaches

The following table organizes representative 3D-RPE mechanisms, highlighting their core principles, mathematical structure, and target domains:

Approach	Mathematical Form	Target Domain	Cited Work
Axis-sep. 3D RoPE	Per-axis $2 \times 2$ rot.	3D vision, video	(Li et al., 18 Mar 2025, Feng et al., 24 Mar 2025, Ji et al., 17 Apr 2025)
LieRE, STRING	$\exp(\sum L_k r^k)$	Vision, robotics	(Ostmeier et al., 14 Jun 2024, Schenck et al., 4 Feb 2025)
Spherical RoPE	$R_z(\omega \lambda)R_x(\varphi)$	Geo, geotokens	(Unlu, 2023)
VRoPE	Symmetric index splits, diagonal spatial rotation	Video-text LLM	(Liu et al., 17 Feb 2025)
Bloch Sphere 3D-RPE	Tilted circ., chunked	Long-context LM	(Ma et al., 14 Jun 2024)

Each approach preserves the essential property that attention depends only on relative position—now in three spatial or spatiotemporal axes—and brings geometric priors directly into the Transformer architecture.

---