While Rotary Position Embeddings (RoPE) for natural language performs well and has become widely adopted, its adoption for other modalities has been slower. Here, we introduce Lie group Relative position Encodings (LieRE) that goes beyond RoPE in supporting higher dimensional inputs. We evaluate the performance of LieRE on 2D and 3D image classification tasks and observe that LieRE leads to marked improvements in performance (up to 6%), training efficiency (3.5x reduction), data efficiency (30%) compared to the baselines of RoFormer, DeiT III, RoPE-Mixed and Vision-Llama.
This paper examines how to encode spatial information for transformers. LieRE proposes a method of encoding positions in \(\mathbb{R}^d\) as high dimensional rotations that outperforms state of the art techniques on both 2D and 3D data.
LieRE replaces the RoPE block-diagonal rotation matrix with a dense, high dimensional rotation computed as \(R=\exp(Ax)\) where \(x\) denotes the position of the token and \(A\) is a learnable linear map to skew symmetric matrices and \(\exp\) denotes the matrix exponential. The Lie group structure of high-dimensional rotations ensures R is a rotation, and thus orthogonal, and inverted in the course of the inner product of the attention mechanism. We also show that the small amount of learnable capacity introduced by \(A\) improves performance, showing that marginal increases in the number of learnable parameters of \(A\) result in improved performance, allowing the model to learn how to better use positional information.