LieRE: Generalizing Rotary Position Encodings to Higher Dimensional Inputs

Authors: Sophie Ostmeier, Brian Axelrod, Michael Moseley, Akshay S Chaudhari, Curtis Langlotz
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: Position Encoding, Attention, Transformer, Computer Vision, Machine Learning

Abstract:
Rotary Position Embeddings (RoPE) have demonstrated efficacy and gained widespread adoption in natural language processing. However, their application to other modalities has been less prevalent. This study introduces Lie group Relative position Encodings (LieRE), which extend beyond RoPE by accommodating n-dimensional inputs. LieRE encodes positions of tokens by replacing the RoPE rotation matrix with a dense, high-dimensional, rotation matrix generated via a learned map. We conducted empirical evaluations of LieRE on 2D and 3D image classification tasks, comparing its performance against established baselines including DeiT III, RoPE-Mixed, and Vision-Llama.
Our findings reveal significant advancements across multiple metrics as compared to the DEIT III basline: LieRE leads to marked relative improvements in accuracy (10.0% for 2D and 15.1% for 3D compared to DeiT). A 3.9-fold reduction in training time for the same accuracy was observed. LieRE required 30% less training data to achieve comparable results.
These substantial improvements suggest that LieRE represents a meaningful advancement in positional encoding techniques for multi-dimensional data. The implementation details and reproducibility materials will be made openly available.

Decision: reject
Reviewer scores: 3, 10, 3

Reviews and discussion:
<thread>
creation date:
1730429927409

thread_type:
official_review

<thread_object>
creator: reviewer_brzg

<content>
summary:
The paper proposed an extension to the prior work RoPE attention. The central idea is to extend the 1-d position encoding in RoPE to handle higher-dimension position indices. In RoPE, the feature vectors are rotated before the attention inner products to take into account the distance between two tokens. For each feature vector of d-dimension, d/2 rotations in 2d space were used. The authors argue that using 1 rotation of d-dimension space, taking into account multiple token indices, is better.

Lie theory is used to map position indices to "rotations", via the exponential map.

strengths:
The use of Lie groups for position-aware attention is a novel and interesting idea.

weaknesses:
- The method is not clearly defined, with proper mathematical description. By code segment in the appendix, is also not explained, particularly not in a way that connects with the main text.
- It is not intuitive why a single rotation is better than multiple rotations. Actually, multiple rotations of RoPE have the benefit of adaptively capturing different levels of distance effects.

questions:
- Please define the A matrix in detail. Clarify the notation n and d in Section 3.2 (use them consistently please).
- Can you conduct experiments on a data set with varying image sizes and resolutions? In other words, a single dataset with different scales of spatial correlation.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730569703679

thread_type:
official_review

<thread_object>
creator: reviewer_ecnh

<content>
summary:
This paper proposes a symmetry-aware method for positional encoding (LieRE) that outperforms the state-of-the-art methods (RoPE), by considering subspace rotations of dimension higher than $2$ in the self-attention's dot-product (the orthogonal symmetry group action given by $(R,(q,k))\mapsto (Rq,Rk)$). Using a Lie group theory the resulting rotation encoding is deduced as $R=\exp(\sum_iA_ix_i)$ where $i$ is each domain dimension and $A_i$ is learnable skew-symmetric matrices. Furthermore, shuffling patches causes more accuracy drop in LieRE than RoPE, which verifies that the model relies more on positional encodings by using LieRE. It also outperforms existing methods in terms of training time and dependency on dataset size.

strengths:
- The positional encoding structure is fundamentally important in Transformer architectures. Improving over the state-of-the-art has a great impact. 
- I find it particularly interesting to treat the spatial and embedding dimensions of the hidden states equally, which I believe results in a unified understanding of the embedding space as a tensor field. Lie group generalizes rotary positional embedding from dimension $2$ to $n$ by regarding RoPE as a commutative special case.
- I find it interesting that sharing parameters across layers (LieRE-Commute over RoPE-mixed) improve learnable positional encodings.

**Update**: Thanks for adding the experiments to validate LieRE’s performance. I have raised the score from 8 to 10, the presentation score to 4, and confidence to 5. Good luck.

weaknesses:
1. Although the relative gain over existing methods is fair and remarkable, <70% accuracy on CIFAR100 and ImageNet and ~50% on UCF101 is far from optimal. For example, the referred paper (Heo et al. 2024) reports >80% accuracy. It would be more convincing to improve the baseline.
2. The 3.5x reduction in training time is compared under the wall time of 200 epochs, which means the same performance is obtained at around 57 epochs for LieRE. I wonder how these methods compare in terms of the best test loss, and the converged training loss (which means after 200 epochs). Running longer experiments may also help remedy poor baselines.
3. I find the compute efficiency less informative than the learning curve. The FLOPs analysis is of practical interest but looks trivial since positional encoding is a lightweight part of the model.

Minor issues: 
- Table 2: I find the word "stem" in Table 2 confusing and unnecessary. Clarifying it in the text rather than just in Figure 3a would help.
- Many \citet should be \citep
- Table 2 line 381: Rope should be {RoPE}. 
- Line 400: Figure ??

questions:
1. I don't understand why commutativeness in the RHS of Equation (2) is not $\exp(V)^{-1}\exp(U)$, but $\exp(U)^{-1}\exp(V)$. Is it a typo?
2. Should line 201 "We present attention with LieRE and LieRE-Commute in Algorithm 1 and 2" reverse the order? 
3. Would it be better to clarify the relation between "LieRE-Commute, learnable RoPE and RoPE-mixed" already here in the method section (now this does not appear until 5.3.3), and remind the reader in the background why the group $SO(d)$ not abelian if $d>2$, preferably with some one-line intuitions?
4. Do you have more intuition to justify the reason why the stem option outperforms layer or head heterogeneous implementations? For example, does it suffice to think of the $O(n)$ symmetry in the dot product that only needs one representative element in the quotient group, instead of striving to learn it repeatedly in every layer or head subspace?
Based on this, would you think it is possible to unify the representation by rotationally aligning the subspaces of each layer?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731121187034

thread_type:
official_review

<thread_object>
creator: reviewer_cjwd

<content>
summary:
The authors propose a position encoding technique that can be used with attention mechanisms for 2D and 3D data. In contrast to RoPE which learns a block diagonal (2x2 blocks) rotation matrix transformation of key and query matrices from relative position information, LieRE learns a general rotation matrix transformation of key and query matrices from absolute position information. While this introduces additional parameters, the authors mitigate this by sharing parameters across attention heads. The authors show that with the combination of these strategies, LieRE improves predictive performance, training efficiency and data efficiency.

strengths:
The authors propose a novel position encoding scheme with improved performance over baseline models. The improved performance is in terms of predictive performance, training efficiency, and data efficiency. Moreover the model can be used for 2D and 3D data.

weaknesses:
The writing could be improved. The document would benefit from additional proofreading. For example, in several places (Lines 169, 173) the text reads ‘equation equation’; ‘figure’ and ‘table’ should be capitalized (Line 187, 397); and the notation for updated keys and queries is inconsistent (Line 188, 189, 209). The clarity of the document would be improved. For example, it would benefit the reader if the equation for the attention mechanism were provided 3.1; some text describing the algorithm would benefit the reader; it would be nice to show that LieRE-Commute and RoPE-mixed are special cases of LieRE.

The organization could be improved. Using subsubsections in the related work doesn’t seem necessary, and takes up space that might be used to clarify the method section. Some details of the method are presented for the first time in the Results section (e.g., that the parameters of LieRE and its variants are shared across attention heads). Figures are often far from where they are referenced.

questions:
* When the authors say ‘generator space’ do they mean Lie algebra?
* In the definition of R_{LieRE} consider using A(x) as the argument to the exponential map
* (Line 400) Broken figure reference 
* Do the authors have thoughts on why LieRE improves training/data efficiency 
* Should the final statement in eq 2 be exp(V)^-1 exp(U)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732742109932

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We sincerely thank the reviewers for their thoughtful feedback and for taking the time to evaluate our submission.

In response to the reviewers’ suggestions, we have made several significant improvements to the manuscript.

**Resolution Generalization Experiments:** As per reviewer 3’s suggestion, we have added experiments to evaluate the method’s generalization across resolutions. LieRE performs strongly, with both a simple single-resolution pre-training recipe and a two resolution training recipe mirroring DEIT III.

LieRE vs RoPE-Mixed generalization across resolutions (Please see the figure at https://imgur.com/a/bnPmERw)

| Resolution (imsize) | 192           	| 224           	| 256           	| 320           	| 384           	| 448           	|
|---------------------|-------------------|-------------------|-------------------|-------------------|-------------------|-------------------|
| liere           	| 64.6 (64.1, 65.0) | 69.3 (68.9, 69.7) | 71.4 (71.0, 71.8) | 73.2 (72.8, 73.6) | 73.1 (72.7, 73.5) | 71.7 (71.3, 72.1) |
| liere_finetune  	| 64.0 (63.5, 64.4) | 68.7 (68.3, 69.1) | 71.5 (71.1, 71.9) | 73.8 (73.4, 74.2) | 74.0 (73.6, 74.3) | 73.3 (72.9, 73.6) |
| naver           	| 63.7 (63.3, 64.2) | 68.8 (68.4, 69.2) | 71.1 (70.7, 71.5) | 72.5 (72.1, 72.9) | 72.1 (71.7, 72.5) | 69.7 (69.3, 70.1) |
| naver_finetune  	| 63.1 (62.7, 63.5) | 68.2 (67.8, 68.6) | 71.2 (70.8, 71.6) | 73.3 (72.9, 73.6) | 73.3 (72.9, 73.7) | 71.7 (71.3, 72.1) |


**Clarity and stylistic changes, particularly of the methods section (section 3):** We have rewritten the methods section to provide a clearer and more detailed explanation of LieRE, making it easier for readers to follow. We have also fixed numerous more minor issues throughout the paper.

We also identified a memory-related issue involving the layout of one of our parameter tensors after a reshape operation, which impacted certain experimental results, in particular learning across heads and layers for RoPE-Mixed and LieRE. We have since resolved this issue and include the top line results before and after the fix below. X’s indicate trainers that are in progress and we will update in the comments and the paper once the last trainer as finished. Additionally, we have thoroughly reviewed our codebase.  We also include the performance of LieRE with the constraint of using rotation matrices that are block diagonal with block size 64 and 8 instead shared across heads and layers based on the performance (Table 2). Due to compute restrains with changed the 3D depth to 8 (default in the Video MAE paper) instead of 32.

| Method | CIFAR-100 | ImageNet-1k | UCF101 |
|------------------------|--------------------------|-------------------------|------------------------|
| Abs. Pos. E.$^{*,**}$ | 63.9 (62.9-65.8) | 66.1 (65.7-66.5) | 40.9 (40.5-41.3) |
 | VisionLlama RoPE | 65.5 (64.6-66.5) | 65.4 (65.0-65.8) | 45.0 (44.6-45.4) |
| RoPE-Mixed | 68.8 (67.9-69.7) | 68.8 (68.4-69.2) | 46.3 (45.9-46.7) |
| LieRE$_{8}$ | **70.3 (69.4-71.2)** | **69.6 (69.2-70.0)** | **47.0 (46.6-47.4)** |
 | LieRE$_{64}$ | 70.0 (69.1-70.9) | 69.3 (68.9-69.7) | 44.7 (44.3-45.1) |

| Method            	| CIFAR-100       	| ImageNet-1k    	| UCF101           	|
|-----------------------|---------------------|--------------------|----------------------|
| Abs. Pos. E.*.**  	| 63.9 (62.9-65.8)   | 66.1 (65.7-66.5)   | 42.7 (42.3-43.1) 	|
| VisionLlama RoPE  	| 65.5 (64.6-66.5)   | 65.4 (65.0-65.8)   | 50.5 (50.1-50.9) 	|
| RoPE-Mixed        	| 63.7 (61.8-65.6)   | 65.9 (65.5-66.3)   | 45.7 (45.2-46.0) 	|
| LieRE-Commute_shared     	| 66.9 (66.0-67.9)   | 68.4 (67.9-68.9)   | 52.1 (51.7-52.5) 	|
| **LieRE_64_shared**         	| **69.9 (68.9-70.8)** | **69.5 (69.1-69.9)** | **53.6 (53.2-54.0)** |

We will upload an updated version of the paper once the final trainers complete so that all the differences are clearly visible in pdfdiff.

We believe these revisions not only address the reviewers' concerns but also strengthen the overall contribution of our work. We are excited about LieRE’s potential and its ability to advance positional encoding techniques for high-dimensional data and look forward to sharing with the field. Thank you again for your constructive feedback, and we look forward to any further comments or suggestions.
</content>
</thread_object>

</thread>


Metareview:
metareview:
Lie group Relative position Encodings (LieRE) is an alternative to or in a sense extension of Rotary Position Embeddings (RoPE). While RoPE learns a structured rotation matrices for keys and queries, LieRE learns general rotation matrices, and in effect has more degrees of freedom for the learned transformation. These degrees of freedom allow for the encoding of absolute and relative position information. Experiments on 2D image and 3D video classification show improved accuracy, training time, and data efficiency on the standard datasets of CIFAR-10/100, ImageNet and common but outdated dataset of UCF101. The empirical improvements and applicability to different dimensionalities are the key strengths, while lack of clarity and insufficient analysis are the key weaknesses. This work is missing a fully consistent analysis of results vs. relative encodings like RoPE, as some results are reported against absolute position baselines instead, while both are relevant to LieRE as an extension of both. This work is also missing sufficient clarity to be understood and informative to its audience, as evidenced by the two confirmed ratings for rejection and comments about level of detail and writing. The meta-reviewer highlights that the revision has made progress on this front, but at the same time the scale of its edits merits another round of review. The meta-reviewer sides with rejection on the balance, due to the issues of clarity and comparisons, although there is definite empirical promise for the proposed method. The authors are encouraged to further incorporate the reviewer feedback to resubmit a clearer and more thoroughly-analyzed edition of this work.

Note: the authors are advised to use parenthetical citations (`\citep`). This had no bearing on the decision, but is offered as miscellaneous typesetting advice.