LieRE: Generalizing Rotary Position Encodings to Higher Dimensional Inputs

Explanation of method: We updated the method section based on your feedback. We replicate a compressed version below for your convenience.

A matrix is skew symmetric if $A^T=-A$, and any matrix $X$ can be made skew symmetric using torch.triu(x, diagonal=1) - torch.triu(x, diagonal=1)^T. In order to encoding positions in a $d$-dimensional space LieRE learns sets of $d$ skew-symmetric matrices which we refer to as the skew-symmetric basis. Recall that in a typical ViT, there is a 1:1 mapping between tokens and patches, leading to each patch having a unique position. A token’s position vector $p$ is then encoded as a rotation matrix $R(p) = \exp{\left(\sum\limits_{i=0}^d A_i p_i\right)}$, where $\exp$ refers to the matrix exponential. Then the individual attention head’s keys and queries are modified with a tensorized batch matrix multiplication that rotates each token’s keys and queries by the rotation matrix for that token.

Outside of the section where we experiment with sharing parameters, this is done with separate skew-symmetric bases for each attention head in every layer.

We are also updating the algorithm box to reflect this greater level of detail.

Why a single high n-dimensional rotation does better than $n/2$ 2D rotations: Using a single high dimensional rotation actually increases the capacity of the mode compared to using multiple 2D rotations. Q$n/2$ 2D rotations only have $n/2$ degrees of freedom, whereas an $n$-dimensional rotation has $n(n-1)/2$ degrees of freedom. These additional degrees of freedom allow the model to learn a richer set of positional encodings. See the “Motions and dimensions” section here for a reference. This could also be seen as a corollary of skew-symmetric matrices being the Lie algebra of the high dimensional rotation matrix Lie group. It is important to keep in mind that applying a 2D rotation to coordinates 1,2, another 2D rotation to coordinates 3,4 etc, is a special case of a single high dimensional rotation meaning the LieRE can represent RoPE-style methods exactly in addition to being able to represent high dimensional rotations which cannot be broken down into a set of 2D rotation matrices on the block diagonal.

Experiments with varying image sizes and resolutions: This is an excellent suggestion. We evaluated both LieRE and RoPE-Mixed on different resolutions of the imagenet validation set below.

Please see the results at https://imgur.com/a/bnPmERw

We evaluate the methods under two training recipes. The first is the original 200 epoch imagenet training recipe used throughout this paper that always trains at a resolution of 224x224. The second one mirrors the DEIT-III recipe and adds an additional 30 epochs at a 240x240. All models perform better with the two-stage training recipe. The performance benefit of LieRE increases as the resolution increases.

This performance increase is particularly exciting in the context of RoPE-Mixed having already shown large improvements over absolute position encodings.

Please let us know if there is anything else we can do to improve your rating. Your feedback has already been very helpful and led to the exciting multi-resolution results.

Thank you for the careful review and constructive feedback of the manuscript.

Stronger baseline: The shorter training recipe was motivated by more limited computational resources. We are also cautiously optimistic that we will have results under the DEIT III training recipe used by Heo et al by the deadline, but that depends on how much GPU time we get between now and then. We will share the results once it completes. In the meantime, we fine tuned the models for an additional 30 epochs as part of the multi-resolution classification task, closing the gap somewhat with the much longer training recipe. Update on 12/2: We are on epoch 346/400 of the longer recipe pretraining and are hoping to be able resume the job soon.

Detailed FLOPs information: Thanks for the feedback. We moved the detailed FLOPs table to the appendix in order to make room for the learning curves and multi-resolution classification results. We also added a table to the appendix with the best test losses for all results.

Clarity: For table 2 we updated the column names to “All Shared”, “Shared Across Heads”, “Shared Across Layers” and corrected Rope to RoPE. We corrected the commutativity typo in equation 2, corrected the usage of citet and citep and fixed ??? references. We also reversed the order of algorithms 1 and 2 in line 208. We moved the relationship between RoPE-Mixed and Liere into the method section (section 3).

Performance of sharing LieRE parameters across heads and layers: We attribute this to the memory issue with further details in the official comment.

Test losses and learning curve: We added the validation losses and learning curves to the appendix. The losses and accuracy curves are well converged. We also trained LieRE on imagenet for 400 epochs and could not see substantial gains.

Thank you again for the review! Please let us know if there is anything else we can do to improve the paper as we await the longer training run.

Thank you for taking the time to review our paper and providing feedback that we were able to use to improve the manuscript.

Style and Proofreading: Corrected issues such as "equation equation" (Lines 169, 173), inconsistent notation (Lines 188, 189, 209), and improper capitalization of "Figure" and "Table." Used A in argument to exponential to clarify that the input is skew symmetric. Fixed commutativity typo in equation 2. Additional proofreading improved overall clarity and eliminated typos.

Clarity: Added background section for attention that includes equations. Extensive improvements to the writing in section 3 (method), expanded algorithm descriptions, and clarified how RoPE-Mixed are special cases of LieRE.

Organization: Reduced usage of subsubsections, moved details on parameter sharing to the method section, and repositioned figures, tables and algorithms closer to their references.

Why LieRE improves efficiency: LieRE enables the model to depend much more on the token positions. We refer to the patch shuffling experiment as evidence of this. Furthermore, dropping the commutativity allows the model to encode both absolute and relative positional information unlike prior methods that often focused on one or the other. The RoPE-Mixed paper also showed some improvement from using RoPE-Mixed combined with absolute position encodings. LieRE lets one combine both absolute and relative positional information in a consistent way.

Thank you again for your constructive feedback that helped improve the paper. Please let us know if there is anything else we can do to improve the paper and score.

We wanted to thank the reviewers for your constructive feedback. We are wondering if our newest comments address your concerns regarding:

• Writing style and clarity of the paper: We improved the writing of the paper. In particular, we improved on the method section describing the LieRE.
• Resolutions Generalization: We have included experiments to assess the method's generalization across resolutions compared to RoPE-Mixed. LieRE significantly outperforms RoPE-Mixed using both a straightforward single-resolution pre-training approach and a two-resolution training strategy mirroring the DEIT III training recipe.

The discussion period is coming to an end tomorrow, and we hope the comments and new manuscript address the comments. If there are any remaining gaps, would it be possible for the reviewers to share them so we have a chance to address them?

Thanks in advance!