Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding

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

Q1: Could the authors provide additional results on length extrapolation (training on short sequences, testing on longer ones) and short-text capabilities, such as MMLU?

Thank you for bringing this up. We have included evaluation results, including MMLU, in Table 7, Appendix B. The results demonstrate that TAPE performs competitively on these tasks, even when the base model is not pre-trained with the same architecture. For details on length extrapolation, please refer to our response to Q2.

Q2: If an LLM is trained with TAPE, is directly extending the model's context window as challenging as with RoPE? If extending the context window is difficult, are there methods similar to NTK or PI to achieve effective and efficient context window expansion?

Thank you for raising this question. In Section 4.1, we evaluate TAPE’s zero-shot length extrapolation ability using an arithmetic learning task, where the test length is twice the training length. TAPE demonstrates competitive performance in this setup. Inspired by your excellent suggestion, we have explored YarN [1] to further enhance TAPE’s length extrapolation capability in Appendix B. As shown in Figure 5, YaRN effectively enhances the length extrapolation performance of TAPE with RoPE initialization, achieving a relative improvement of 3.4%.

Q3: Could the authors conduct ablation studies on the impact of the Attention and MLP blocks in updating TAPE?

Yes, we have included ablation studies in the Appendix B to validate our design choices for incorporating position contextualization in both the attention and MLP layers. We evaluate models’ perplexity results on the test set use the same pretraining setup described in Sec. 4.2. A summarized version of Table 6 in Appendix B is provided below, showing that removing any component increases perplexity.

References:

[1] Peng, Bowen, et al. "YaRN: Efficient Context Window Extension of Large Language Models." The Twelfth International Conference on Learning Representations.

We greatly thank Reviewer 8dnP for reviewing our work and provide insightful suggestions on further strengthening the manuscript. We address the concerns as follows.

W1: In Section 3.2 Authors introduce a number of hyperparameters, specific to TAPE. However in 4.1 and 4.2, where models are trained from scratch, authors do not specify how TAPE was initialized. For ex, what are M, B, L, D, B', E, etc? How those where selected and why? In finetuning things are more clear, where TAPE was initialized to align with RoPE , but what's the size of W1/2 (B`) and how it was selected?

Thank you for pointing this out. The values M = 12, B = 64, B’ = 48, and L = D = 2 are used consistently across all experiments, as clarified in the updated Appendix. Specifically, L = D = 2 was chosen to align with RoPE, while M = 12 and B = 64 are consistent with prior work. B’ = 48 selected to introduce a minor (~1%) increase in parameters.

W2: Authors perform evaluation on selected tasks claiming "superior performance in language modeling, arithmetic reasoning, and long-context retrieval tasks":

1. Arithmetic learning task: superior performance belongs to Abacus Embeddings (McLeish et al., 2024), that can solve problem with up to 100 operands. Regarding embeddings including in the comparison, I find information in caption and in text confusing: on pictures, FIRE and TAPE look almost identical, your report "The average accuracy across the heatmap is 26.12%, 26.12%, 49.44% and 41.42% respectively for RoPE, RandPE, FIRE and TAPE." - i.e. TAPE performs worse than FIRE. But then in the text of Section 4.1 authors claim "Compared to FIRE, which achieves 25.24% ..., TAPE shows a remarkable 51.9% relative improvement." - which seem to contradict the picture.
2. Evaluation on SCROLLS task is further confusing: this benchmark contains 6 datasets, with test partitions kept in private. Authors report performance only on 4 selected datasets out of 6. Not clear what data where actually used for evaluation, and what was the evaluation setup. It's even more surprising to see low scores on QuALITY benchmark, which is multi-choice QA dataset, where random guess achieves 25% accuracy, while the best performance reported is 11.6%. If authors used another evaluation pipeline, it should be thoroughly described in the paper.
3. Fine-tuning: authors report perplexity on language modeling task, and accuracy on Passkey retrieval, but did not investigate performance on standard reasoning evaluation benchmarks, or even long-context QA reported in section 4.2. With all that, author's claim about "TAPE’s superiority in both arithmetic reasoning and long context language modeling task" is questionable.

We sincerely appreciate your detailed feedback and aim to address the points raised as follows:

1. Arithmetic Learning: Sorry for the confusion caused by the caption mismatch. We have corrected the error and aligned the caption description with the main text. For a fair comparison, we do not consider abacus embeddings, although they can be seamlessly integrated with TAPE. Additionally, we included an extra baseline for this experiment, and our method still achieves SOTA performance in the revised manuscript.
2. Comprehensive Evaluation on SCROLLS: The SCROLLS benchmark includes three additional summarization tasks (GovReport, QMSum, and SummScreenFD), which were initially omitted due to the higher inference time required compared to QA tasks. To provide comprehensive results, we have now completed evaluations across all 7 tasks and included the results in Table 1. A summarized version of the results below shows that TAPE achieves the best overall performance. Regarding the evaluation setup, we used the official evaluation code on the validation dataset, which will be detailed in our open-sourced code. Regarding QuALITY, please refer to response to Q1 for a clarification on its validity. | Method | QAS | CNLI | NQA | QuAL | QMS | SumS | GovR | | --- | --- | --- | --- | --- | --- | --- | --- | | RoPE | 8.39 | 65.00 | 1.77 | 0.04 | 6.34 | 5.63 | 9.71 | | ALiBi | 8.25 | 69.62 | 4.11 | 0.0 | 9.92 | 9.78 | 18.81 | | RandPE | 13.44 | 62.01 | 4.63 | 0.38 | 8.43 | 8.31 | 8.93 | | xPos | 9.02 | 71.75 | 4.83 | 0.24 | 10.73 | 9.38 | 16.38 | | TAPE (ours) | 11.52 | 72.80 | 6.79 | 11.60 | 12.42 | 10.34 | 15.18 |
3. More Evaluations on Fine-tuned Models: We have added evaluations on standard reasoning and language understanding benchmarks, ARC and MMLU, in Table 7, Appendix B. The results demonstrate that TAPE performs competitively on these tasks, even when the base model is not pre-trained with the same architecture. Regarding SCROLLS, which requires a fine-tuned model for evaluation [1,2,3], we are currently conducting these resource-intensive experiments and plan to include the results in the camera-ready version, time permitting.

W2.3 >> in Table 7, Appendix B. Do I read it correctly, that improvements in perplexity (Table 2) didn't result in improvements in accuracy? Isn't the later more important? If so, worth moving to the main text.

You are correct that improvements in perplexity (Table 2) do not necessarily translate to improvements in accuracy on standard benchmarks. This is because accuracy in such benchmarks primarily depends on the pre-trained base model and can be further enhanced by techniques like RLHF. The aim of the experiment in Table 2 is not to improve benchmark accuracy directly but to evaluate which techniques enable a model pre-trained on short-length text to adapt effectively to long-length text through fine-tuning, without the need for pre-training on long text from scratch. Perplexity is used here as a metric to assess the model's adaptation to long text, as demonstrated in prior work like LongLoRA [6].

W4: I don't think my question was answered. There are no ablations.

Please refer to our response to W1 for the specification of other parameters and ablations for B′.

Regarding SCROLLS, which requires a fine-tuned model for evaluation [1,2,3] - not sure I understand, why would it require finetuning? You have reported pre-training evaluations in Table 1

We would like to confirm that SCROLLS does require fine-tuning. Pre-training a language model on a corpus and then fine-tuning it on SCROLLS is the standard approach in previous works [1, 2, 3], as the benchmark is specifically designed for fine-tuning.

Our experiments in Section 4.2 follow a similar setup: we pre-train the model from scratch and then fine-tune and test it on downstream SCROLLS. To avoid any confusion, we’ve added further clarification in the revised paper.

For additional context, we also explored how zero-shot generation looks for Llama-7b fine-tuned with LoRA. The outputs in QuALITY were largely nonsensical, with repeated patterns like “The\n\u0409\n\u0409\n\u0409...” until the token limit was reached. This suggests that the QuALITY benchmark is not suited for zero-shot or few-shot generation and instead requires fine-tuning for meaningful evaluation.

Q1: Why do you think TAPE outperforms other baselines on long-context tasks, especially on QuALITY?

To understand why TAPE outperforms other baselines on long-context tasks like QuALITY, it is important to consider the challenges faced by Transformers with existing positional embeddings in such scenarios. Traditional positional embeddings often assume a distance-decay effect, where information from words farther apart is given less significance. While this assumption works well for most language tasks, it may fail in long-context scenarios where critical information could appear at the beginning or at positions far from the output.

TAPE addresses this limitation by incorporating positional contextualization, allowing positional embeddings to adaptively adjust the importance of different positions without relying on a fixed decay pattern. This is supported by our visualization of attention patterns in Figure 7, which shows that TAPE can effectively attend to longer contexts.

For further illustration, we compared TAPE’s output with other baselines on two specific questions from QuALITY. As shown in the following table, TAPE consistently generates either the correct answer or a response close to the correct one, unlike other methods. Related analyses and the specific questions are detailed in Table 10 and Table 11 of Appendix C.

Dear Reviewer 8dnP,

We want to thank you for the constructive comments in your review. As a follow-up on our responses, we would like to kindly remind you that the discussion period is ending soon. We hope to use this open response period to discuss the paper to solve the concerns and improve the quality of our paper. Have you gotten a chance to read our further responses and revision, which attempt to address all of your concerns? In the revised manuscript, we have included additional experiments and illustrations, such as hyperparameter ablation studies, and example QA from QuALITY. For your convenience, we have also provided detailed references in our response. We sincerely hope to have further discussions with you to see if our response solves the concerns. We would be more than happy to provide more information or clarification, should it be necessary. We hope our paper, as the first work to propose the integration of position contextualization into an enhanced Transformer, could be valued and receive a positive and fair assessment.

Best,

Paper 1199 Authors

Dear authors,

Thank you for your response. I increased my score to 8 and lean toward accepting this paper.

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

W1: Illustration of TAPE Method: It would be highly beneficial if the authors could include an illustration of the TAPE method. Visualizing it with a block diagonal matrix and context-dependent data flow should be feasible and would greatly enhance readers' understanding of TAPE’s operations.

Thank you for this excellent suggestion. We have added an illustration figure (Figure 6) in Appendix C to better visualize TAPE’s operations, including the context-dependent data flow in both attention layer and MLP layer. Please feel free to share any further ideas for improving the visualization.

W2: Missing Hyperparameters: Some important hyperparameters are not specified. For instance, is ( B = L = D = 2 ) used across all experiments? Additionally, details on hyperparameters like learning rate, batch size, and optimizer are missing and should be included.

Thank you for pointing this out. L = D = 2 is used across all experiments, as clarified in the updated Appendix A. Additional training hyperparameters, including learning rate, batch size, and optimizer details, have also been provided.

W3: Potential Error in Figure 2 Caption: The caption for Figure 2 may contain an error. It states, "The average accuracy across the heatmap is 26.12%, 26.12%, 49.44%, and 41.42% respectively for RoPE, RandPE, FIRE, and TAPE." This appears to contradict line 364, which mentions that "TAPE shows a remarkable 51.9% relative improvement."

Sorry for the confusion caused by the caption mismatch. We have corrected the error and aligned the caption description with the main text. Additionally, we included an extra baseline for this experiment, and our method still achieves SOTA performance in the revised manuscript.

W4: Explanation for Improved Arithmetic Task Performance: More explanation is needed on why TAPE improves performance on arithmetic tasks. Are the authors suggesting that TAPE has inherent advantages in this area?

Thank you for bringing this up. Here is a more detailed explanation:

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.

We have also included this explanation in Sec. 4.1 to further clarify TAPE’s advantages in arithmetic tasks.

W5: Clarification on TAPE’s Task Performance: It’s unclear how TAPE improves performance across all reported tasks. Does TAPE enable more flexible attention patterns? Could you, for instance, visualize the (\phi) function block by block in equation (6) to provide further insights?

Yes, thank you for raising this excellent point. We have added a visualization of the last layer’s attention patterns in Figure 7, Appendix C, showing that TAPE effectively captures and attends to more contextual information over longer distances compared to RoPE.

W6: Compatibility of T5 with Flash Attention: I’m unsure if it's accurate to state that T5 is incompatible with flash attention. While support is currently lacking, compatibility could be achieved with implementation.

Thank you for pointing this out. We agree that T5 could theoretically be made compatible with Flash Attention, but it has not yet been implemented. We have softened our tone in the main text to reflect this. To clarify, the main engineering challenge lies in the difficulty of implementing backward computation with gradient calculations for relative position embeddings, which is why Alibi is currently the only supported relative position embedding. In contrast, our TAPE is natively compatible with Flash Attention and could achieve even greater efficiency through our customized fused kernel, which we consider a notable practical contribution.

Will the code be released?

Yes, we will open-source the code immediately and make the model available in the near future.

Q1: Is the description of Figure 2 incorrect because it doesn't match the description provided in lines 361--365.

Yes, this has been addressed as noted in W4.

Q2: Where are the ablation results for design decisions? For example, what happens if you treat PE as a vector instead of tensor, what if the PE is the same for all blocks?

We have included thorough ablation results for design decisions Appendix B. These results cover scenarios such as using vector PE and keeping the same PE across all blocks, i.e., not incorporating positional contextualization in all layers. The models’ perplexity on the test set was evaluated using the same pretraining setup described in Sec. 4.2. A summarized version of Table 6 in Appendix B is shown below, demonstrating that any ablation increases perplexity.

Q3: Can you visualize attention patterns or some qualitative analysis on where the proposed method shines?

Yes, thank you for raising this excellent point. We have added a visualization of the last layer’s attention patterns in Figure 7, Appendix C. The results show that TAPE effectively captures and attends to more contextual information over longer distances.

Q4: The paper mainly motivates that TAPE ensures "stability." What do you mean by that and how do you measure it? Where is the empirical evidence for that?

By "stability," we mean that the representation of a sequence remains consistent under positional translation, as defined in prior work [3]. Specifically, this refers to maintaining the same sentence embedding even when padding is added at the beginning or end of the sequence. For instance, when a shifted sentence is input into a Transformer with rotary positional embedding, the positional embedding rotates in its embedding space. Our method ensures rotational equivariance, enabling the same output after encoding. Empirical evidence is presented in our ablation study in Appendix B, where removing rotational equivariance from layer updates results in a decline in performance (increasing perplexity) compared to TAPE. A quick overview of this result is provided in the summary table below.

References:

[1] Biderman, Stella, et al. "Pythia: A suite for analyzing large language models across training and scaling." International Conference on Machine Learning. PMLR, 2023.

[2] Gu, Albert, and Tri Dao. "Mamba: Linear-time sequence modeling with selective state spaces." First Conference on Language Modeling. 2024.

[3] Sun, Yutao, et al. "A Length-Extrapolatable Transformer." Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). 2023.

We greatly thank Reviewer SBfN for reviewing our work and providing insightful suggestions on further strengthening the manuscript. We address the concerns as follows.

W1: The complexity of the method is not justified. For example, it is not clear why there should be PEs corresponding to blocks of input? And why should the PE be a tensor?

Thank you for pointing this out. Here, we will justify our motivation for this design as follows. Ablation results validating the effectiveness of updating PEs and using tensorial PEs can be found in our response to Q2.

• To clarify our motivation for incorporating positional contextualization, it is important to understand the limitations of Transformers in addressing arithmetic tasks. Traditional positional embeddings often assume a distance-decay effect, where words farther apart have less significance in the output. While this assumption holds for most language tasks, it does not apply to arithmetic tasks, where every digit has equal importance regardless of its distance. Positional contextualization allows us to dynamically reweight positional importance based on the task context, maintaining effective distance decay for language tasks while handling arithmetic contexts appropriately.
• Regarding the design of tensorial embeddings versus vector embeddings, a simple way to understand their relationship is to view a tensor as a multi-channel vector. A tensor with a single channel reduces to a vector. This technique of increasing dimensionality to enhance expressive power has been adopted in several works, including RoPE [1], which extends sinusoidal PE to a 2-channel format (one for sine and the other for cosine), and the multi-head attention mechanism [2] in the Transformer.

W2: The method is motivated in terms of increasing stability of PEs, but there is no empirical evidence presented. The paper lacks many ablations that justify the design decisions of the new architecture.

We have included ablation studies in the Appendix B to validate our design choices for the attention and MLP layers, as well as for the tensorial PEs and equivariant PEs to ensure stability. For a detailed analysis of PE stability, please refer to our response to Q4.

W3: The method increases the number of parameters. A fair comparison should be provided for baselines by making sure the number of parameters is same.

While our method does increase the number of parameters, the difference is very minor—only about 1.1%—as shown in Table 3 of Sec. 4.4 (155.3M compared to 154.9M). In prior research on large language models [1,2], models with parameter counts of 125M, 130M, and 160M are often considered to be of equivalent scale. Moreover, this minor parameter increase leads to significant performance gains, such as a 22% relative improvement over the previous SOTA method in the arithmetic learning task (Sec. 4.1).

W4: It is unclear if the performance of the proposed method is stronger than existing proposals. For example, based on Figure 2 (caption), it seems to me that the method performs poorly than FIRE.

Sorry for the confusion caused by the caption mismatch. Our method actually achieves SOTA performance in this experiment. We have fixed the caption description and added an extra baseline for this experiment in the revised manuscript.

W5: Baselines such as NoPE (no position encoding) and FIRE are missing for experiments (NoPE for Sec 4.1 and 4.2, and FIRE for Sec 4.2).

Thanks for bringing this up. We have included NoPE in Sec 4.1 and FIRE in Sec 4.2, and our method continues to demonstrate the best overall performance in these experiments. We will also complete running NoPE results and add them to Sec 4.2 in the camera-ready version if time permits.

W6: The method is mainly motivated for permutaion invariance but all evaluation is performed with causal models which are permutation variant.

Thank you for pointing this out. While our experiments focus on causal language tasks, it is important to note that Transformers are inherently permutation-invariant architectures. They are commonly applied to causal tasks using causal masks and objectives during pre-training. Since our method aims to improve the positional addressing capabilities of Transformers, ensuring compatibility with their intrinsic permutation invariance is essential. We selected language tasks for evaluation because of their wide applicability and impact, even though order-irrelevant tasks could also be explored in future work.

Thank you for your response and revisions. Some of these details are important which were absent in the original submission.

Q4: "stability." What do you mean by that and how do you measure it? Where is the empirical evidence for that?

The padding experiment in the rebuttal is useful. What is the result with no padding, i.e., do you end up with the same number with and without padding? This is important too to make your point that TAPE is orthogonal variant. I am surprised that there is no definition of stability in the original submission. Your ablation in the experiment reminded me of this paper [1].

Q2: Where are the ablation results for design decisions?

In your response, why is there no row where both of them are absent? I could not find it in the revision too. Is that equivalent to ROPE? Some of this discussion has to be in the main paper, not appendix.

W6: The method is mainly motivated for permutaion invariance

I don't find this response satisfying. Motivating the entire paper with permutation invariance but choosing an evaluation setting with Transformer decoder is contradicting. Is there a way you can prove this point empirically?

Many of these details are not present in the original submission, and the revision provides some of these details in the appendix. I am increasing the score to 5 as I still believe the paper can be improved a lot with respect to its motivations and linking them to experimental results/ablations. I will not be able to engage further due to bandwidth.

[1] The curious case of absolute position embeddings. Sinha et al. 2022

To provide a clearer understanding of our work, it is crucial to differentiate between two key concepts that should not be conflated as “our method”: our design principles and the enhanced Transformer. To avoid such confusion, we commit to including these clarifications in a future version of our paper.

1. Our design principles for Transformers with positional contextualization include permutation invariance. Many previous positional embedding methods, such as RoPE and other relative embeddings, also adhere to this principle. The essence of these methods is that positional relationships matter, but absolute positional order does not, meaning positional information should be permutation-invariant. We adopt this principle in our enhanced Transformer, not only because it aligns with the permutation-invariant nature of Transformers, but also because it reflects the permutation-invariant nature of positional information itself. Empirical evidence supporting the use of this principle in positional embeddings can be found in the ineffectiveness of absolute position embeddings, which are permutation-variant, as demonstrated in [1, 2].
2. For our enhanced Transformer, the primary motivation is positional contextualization—updating positional embeddings through the Transformer architecture. This motivation arises from the limitations of Transformers in handling complex language tasks, such as arithmetic and long-context reasoning. Accordingly, we choose to evaluate our architecture on these language tasks. While our model is general and can handle other permutation-invariant tasks more effectively, this is beyond the scope of our current motivation and remains a promising direction for future work.

References

[1] Sinha, Koustuv, et al. "The Curious Case of Absolute Position Embeddings." Findings of the Association for Computational Linguistics: EMNLP 2022. 2022.

[2] Wang, Yu-An, and Yun-Nung Chen. "What Do Position Embeddings Learn? An Empirical Study of Pre-Trained Language Model Positional Encoding." Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP). 2020.

Thank you for your questions. We hope the following answers address your concerns, and we are willing to provide further clarifications if necessary.

Q2: Where are the ablation results for design decisions? → In your response, why is there no row where both of them are absent? I could not find it in the revision too. Is that equivalent to ROPE? Some of this discussion has to be in the main paper, not appendix.

There are no such a setting that both are set as “w/o” because you can find in our ablation setting that the L = 2 in equivariance ablation while L = 1 in tensor ablation, so these two configurations are not orthogonal and not directly combined. Note that we carefully need to find a representative way to ablate each design aspect—equivariance and tensorial embedding—without introducing excessive variation, as these were originally integrated into our dedicated design. As suggested, we have found that setting L = 4 and D = 1  ablates both equivariance and tensor simultaneously. The results are provided below.

Q4: What do you mean by "stability" and how do you measure it? Where is the empirical evidence for that? → The padding experiment in the rebuttal is useful. What is the result with no padding, i.e., do you end up with the same number with and without padding? This is important too to make your point that TAPE is orthogonal variant. I am surprised that there is no definition of stability in the original submission. Your ablation in the experiment reminded me of this paper [1].

In the revised manuscript, we have added the explanation (definition) for "stability" in the main text and provided empirical evidence, including comprehensive padding results, to show TAPE’s stability in Appendix B. Thank you for pointing out this work [1]. We find it a good example for emphasizing the importance of permutation invariance in positional embeddings, as absolute position embeddings are unable to model the permutation-invariant relationships between positional information. This insight aids in understanding our response to W6.

W6: The method is mainly motivated for permutaion invariance → I don't find this response satisfying. Motivating the entire paper with permutation invariance but choosing an evaluation setting with Transformer decoder is contradicting. Is there a way you can prove this point empirically?

We would like to clarify that permutation invariance is not the primary motivation; our method is primarily driven by positional contextualization. Permutation invariance, in turn, serves to enhance positional contextualization by ensuring stable updates. This is because it aligns with the permutation-invariant nature of Transformers and is embedded in several previous embedding methods, making it a natural principle for us to adopt without unnecessarily disrupting it.

Moreover, we would like to further explain the rationale behind prioritizing language tasks over permutation-invariant ones. While we acknowledge that permutation invariance is crucial for tasks where it is required, our model is not restricted to permutation-invariant tasks. Transformers, when combined with a simple causal mask, can handle language tasks, and our enhanced model can be applied in the same way. Therefore, experiments on a variety of tasks are useful for demonstrating the model's effectiveness, and it is not necessary to restrict them to permutation-invariant tasks. We prioritize language tasks to showcase the effectiveness of positional contextualization, as this technique is primarily motivated by the need to address complex language tasks, such as arithmetic and long-context reasoning.