cs.LGApr 10, 2026

Efficient Matrix Implementation for Rotary Position Embedding

Authors: Chen Minqi, Zhongqi Yue, Shihao Zhang, Yun Xu, Peng Wu, kaixiang Xu, Zeyi Huang, Hanwang Zhang

Organizations: Huawei Technologies · Nanyang Technological University

Abstract

Rotary Position Embedding (RoPE) has become a core component of modern Transformer architectures across language, vision, and 3D domains. However, existing implementations rely on vector-level split and merge operations that introduce non-negligible computational overhead, often overlooked in attention optimization. The problem is further amplified in multi-dimensional settings (e.g., 2D and 3D RoPE), where additional vector operations and uneven feature partitions degrade hardware utilization. To overcome these limitations, we propose RoME (Rotary Matrix position Embedding), a mathematically equivalent yet computationally efficient reformulation of RoPE that replaces vector operations with unified matrix transformations. RoME eliminates dimension-specific operations, simplifies implementation, and enables fused parallel execution across Cube and Vector units on modern NPUs. Experiments show that RoME delivers substantial acceleration at both the operator and full-model levels. The implementation is available at https://gitcode.com/cann/ops-transformer/blob/master/experimental/posembedding/rope_matrix/README.md.

Figures & tables

Explore similar work

Aug 30, 2026

Higher-Dimensional Rotary Position Embedding

Transformers rely on position embedding mechanisms in long context modeling in most cases. Rotary Position Embedding (RoPE) embeds positional information with independent 2D rotations, forming relative position terms in self-attention. However, its pairwise, block-based, and decoupled structure limits deep mixing and robustness across channels. We propose HD-RoPE, which extends RoPE from independent 2D rotations to higher-dimensional rotations and introduces a Paley-I orthogonal basis to obtain balanced, isotropic, and dense phase mixing within each rotation subspace. This significantly enhances channel coupling and rotational degrees of freedom while maintaining orthogonal stability and the relative position closure property. Furthermore, HD-RoPE is easily optimized for engineering efficiency without introducing additional trainable parameters. We have conducted extensive evaluation results demonstrating that HD-RoPE achieves significant performance improvements over standard RoPE across various popular benchmarks and in both long and short contexts.
Jun 10, 2026

nD-RoPE: A Generalized RoPE for n-Dimensional Position Embedding

Rotary Position Embedding (RoPE) is widely adopted in Transformer models, yet its extension to high-dimensional domains lacks a unified theoretical formulation. Most existing approaches either apply rotations independently along each axis or empirically mix frequencies, which limits cross-dimensional interactions and yields direction-dependent representations. To address these limitations, we propose nD-RoPE, a decomposition-free generalization of RoPE to arbitrary dimensions. From a translation-invariant formulation in continuous Hilbert space, we derive a spectral condition for isotropy that requires treating positions and frequencies as coupled nn-dimensional vectors. We instantiate this formulation with a multi-scale regular-simplex wave-vector design, which provides non-degenerate spatial coverage and a symmetric, directionally balanced second-order response. Experiments across images, videos, and point clouds demonstrate consistent performance gains and improved generalization in high-dimensional settings.
Mar 12, 2026

Fractional Rotation, Full Potential? Investigating Performance and Convergence of Partial RoPE

Rotary Positional Embedding (RoPE) is widely used in transformers to encode relative positional information, yet the effect of applying RoPE to only a subset of hidden dimensions remains underexplored. This design choice can substantially reduce memory usage, especially at long context lengths. We show that partial RoPE achieves up to 10×\times memory savings over standard RoPE cache while maintaining comparable final loss. We systematically study its impact on training dynamics and convergence across architectures, model sizes, sequence lengths, and datasets. Our results reveal that: (1) applying RoPE to only a small fraction of dimensions (around 10%) matches the convergence of full RoPE; (2) these trends remain consistent across settings; and (3) NoPE (No Positional Encoding) can produce loss spikes, which are mitigated by minimal RoPE usage or QK-Norm application. Overall, our findings highlight partial RoPE as an effective trade-off between efficiency and stability.