Rethinking RoPE: A Mathematical Blueprint for N-dimensional Positional Embedding
This work provides a foundational mathematical blueprint for researchers and practitioners in machine learning, particularly those working on positional embeddings for LLMs and higher-dimensional data like images, though it is incremental as it builds on existing RoPE concepts.
The authors tackled the lack of a unified theoretical framework for Rotary Position Embedding (RoPE) in higher-dimensional domains by proposing a systematic mathematical framework based on Lie group and Lie algebra theory, which unifies existing designs and enables principled extensions to N-dimensional modalities.
Rotary Position Embedding (RoPE) is widely adopted in large language models (LLMs) due to its efficient encoding of relative positions with strong extrapolation capabilities. However, while its application in higher-dimensional input domains, such as 2D images, have been explored in several attempts, a unified theoretical framework is still lacking. To address this, we propose a systematic mathematical framework for RoPE grounded in Lie group and Lie algebra theory. We derive the necessary and sufficient conditions for any valid $N$-dimensional RoPE based on two core properties of RoPE - relativity and reversibility. We demonstrate that RoPE can be characterized as a basis of a maximal abelian subalgebra (MASA) in the special orthogonal Lie algebra, and that the commonly used axis-aligned block-diagonal RoPE, where each input axis is encoded by an independent 2x2 rotation block, corresponds to the maximal toral subalgebra. Furthermore, we reduce spatial inter-dimensional interactions to a change of basis, resolved by learning an orthogonal transformation. Our experiment results suggest that inter-dimensional interactions should be balanced with local structure preservation. Overall, our framework unifies and explains existing RoPE designs while enabling principled extensions to higher-dimensional modalities and tasks.