Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding

We greatly thank Reviewer udZG for appreciating our contributions, providing valuable suggestions on improving the work, and supporting the acceptance of this work. We address the questions as follows.

W1: Existing positional encoding methods have introduced improvements to RoPE to better adapt it for long-sequence tasks, such as NTK-aware RoPE [4]. Could you provide comparative experimental results of this method on long-context tasks?

Thank you for your suggestion. We have implemented YaRN as an additional baseline in our SCROLLS benchmark experiments. The training is currently in progress, and we will update the results once they are available.

Q1: Computing a positional embedding at each layer increases the computational cost for larger-scale language models. Is the added time acceptable, and does the proposed method have scalability? Please provide a time complexity analysis.

As shown in Table 3, TAPE introduces <1% additional parameters and approximately 12% increased computational cost (FLOPs). Since scaling typically involves stacking these layers, the overhead remains proportional. We believe the modest increase in memory usage and ~10% additional compute time is a reasonable trade-off for the benefits provided.

We greatly thank Reviewer poAq for appreciating our contributions, providing valuable suggestions on improving the work, and supporting the acceptance of this work. We address the questions as follows.

W1: The authors provide the running time of attention layers as experimental results in Table 4. However, since the proposed method updates positional features layer-wise through interactions and joint training with token representations in every Transformer layer, what is the experimental runtime during training and inference for large-scale language models?

Thank you for your suggestion. We have included the experimental runtime for full model inference in the updated table below.

W2: Since the proposed method updates positional features layer-wise, how does this affect gradient in the model? Could the authors provide an analysis or empirical results on gradient changes to better understand the impact of this method on training process?

We appreciate this insightful question. During our experiments, we monitored gradient norms and did not observe significant differences compared to RoPE. However, as this aspect falls outside our primary focus (motivation, methodology, and significance), we did not conduct an in-depth analysis of gradient behavior.

Q1: In Table 2 and Table 4, the authors present experimental results showing that the proposed method performs well on long-context tasks. I am curious whether this strong performance in long-context scenarios has theoretical support.

TAPE's superior performance in long-context scenarios can be attributed to two key factors: First, its learnable nature enables dynamic adaptation to varying context lengths. Second, as formally established in Proposition 3.1, our relative position encoding scheme possesses inherent generalization capabilities to unseen sequence lengths. In contrast, conventional positional encoding methods exhibit fundamental limitations in long-context settings, as they either rely on predetermined distance-decay patterns, or lack this crucial relativity property.

We sincerely appreciate Reviewer G28F's positive assessment of our contributions and strong endorsement for acceptance. The reviewer provided one suggestion, to which we respond below:

I do not see any major weaknesses. One suggestion is about the study of hyper-parameters. I think the authors choose B, L, R to align with RoPE for comparison. Since the proposed embeddings are learnable, it would be interesting to see how those hyper-parameters could affect the performance.

Thank you for your valuable suggestion. We have conducted hyperparameter exploration (detailed in the Appendix D), and our experiments demonstrate that the RoPE-like initialization yields the best performance among the configurations we tested.

We greatly thank Reviewer qpJK for appreciating our contributions. We address the concerns as follows.

W1: The motivation of the technique is a bit confusing. The authors claim that relative positional encodings are crucial for “stability and generalization to varying sequence lengths” (L223-224), and use relative positional encodings in their formulation (equation 6) and prove that the transformer is invariant to token index shift (Prop 3.1), but then highlight an arithmetic task where absolute positions are necessary.

We apologize for the inaccuracy in our original statement (Line 316) regarding absolute positions that caused confusion. Please refer to our response to Q1 to see whether this clarification addresses the point.

---

W2: Also, the use of tensors is not very well-motivated. Overall, although the experimental performance is seemingly very good, the design of TAPE seems a bit ad-hoc.

We acknowledge that certain implementation choices in TAPE, such as the tensorial embedding design, were empirically driven to optimize performance. As noted in Appendix D, our ablation studies validate the effectiveness of these architectural decisions.

---

W3: The paper also cites geometric learning as inspiration — “This approach is inspired from the studies for geometric deep learning which processes graphs and point clouds by integrating token features with their geometric properties while preserving inherent physical symmetries” (L83) — but isn’t this property (positional encodings depending on relative distances) already satisfied by RoPE? Some additional ablations that could help clarify the design of TAPE include ablating the tensorial nature (reducing dimensions), ablating the dependence on context in the positional embeddings, etc.

Yes, RoPE also satisfies this property, which is precisely why we employ it as one of our instantiation methods. We do not claim this property as our novelty, but rather use it as a principle to motivate our design choices regarding contextualized position embeddings. As shown in Appendix D, we provide ablation studies that: (1) validate the design of equivariant and tensorial embeddings, (2) analyze the effects of Attention and MLP layers, and (3) investigate the impact of hyperparameter choices. Note that ablating the context-dependence in our positional embeddings (our core design) reduces to RoPE, which is consequently included as a baseline in our experiments.

---

Q1: The arithmetic task is one where absolute positions are necessary (L316), but TAPE uses relative positions (Prop 3.1) — how then does TAPE do well? (Clarification of TAPE’s advantages and understanding of experimental results)

We appreciate this opportunity to clarify. There was an inaccuracy in the original L316 statement regarding absolute positions - we have now corrected this, as absolute positions are not necessarily required for arithmetic. The critical factor is learning the relative importance of different positions within the sequence. As detailed in L328 onwards, TAPE is able to learn these position-dependent importance relationships within the task context. A detailed explanation is also attached:

In arithmetic tasks, every digit has equal importance to the equation, regardless of its distance from the output. Traditional positional embeddings often assume a distance-decay effect, where words farther apart are less significant in the output. While this assumption is valid for most language tasks, it does not hold for arithmetic tasks. Positional contextualization enables dynamic reweighting of positional importance based on the task context, preserving effective distance decay for language tasks while addressing arithmetic contexts appropriately. This highlights TAPE’s potential advantages in arithmetic tasks.

---

Q2: Is there an indexing problem with equation 7? Don’t the two sums cancel each other out? (Clarification for assessing paper’s accuracy)

No, the two sums in Equation 7 form a linear combination of vectors with weights summing to 1, making they cannot cancel out.

---

Q3: Do the same hyperparameter choices (dimensions, etc) work across different tasks? (Affects my evaluation of advantages/disadvantages of the method)

Yes, we maintain consistent hyperparameters across all tasks (with L=R=2 in our main experiments). Additionally, Appendix D provides detailed ablation studies analyzing the impact of different hyperparameter choices.

---

Response to Suggestions:

1. Prop 3.1: The orthogonality of R is indeed crucial for Proposition 3.1 to hold, as non-orthogonal transformations would violate the invariance properties demonstrated in Appendix B. While generalizing the assumptions is an interesting direction, we believe this extension merits dedicated future research.
2. Figure 2: The x- and y-axes represent the sequence lengths of the two operands respectively.
3. Typos: Thank you for your careful review. We have fixed all of them.