DAPE: Data-Adaptive Positional Encoding for Length Extrapolation

Dear Reviewer QhvW,

Thank you very much for appreciating our work. We will address your concerns below.

Q1: The speed of training

A1: We will answer the question in three parts: 1) The potential way to improve the speed of CAPE; 2) The additional training ratio will gradually decrease with a larger model size.; 3) Moreover, CAPE indeed can speed up training. Please refer to the more detailed answer to CAPE computation cost in the Author Rebuttal by Authors.

The potential way to improve the speed of CAPE:

• Reduce the size of $D_{CAPE}$. The computation cost is $O(hN^2D_{CAPE})$. Therefore, by reducing the $D_{CAPE}$ to be half, the CAPE computation cost will be half.
• Algorithm that is more efficient than MLP. The MLP in CAPE can be changed to other more efficient operations, such as sparse-MLP.
• Sparsity: we can make the MLP operation sparse to speed up CAPE. As shown in Figure 6, our CAPE could work well with only $D_{CAPE}=4$, while the Table 1 count with $D_{CAPE}=32$
  • Pruning: use pruning on MLP to remove the redundant parameters.
  • Dynamic sparsity: we could use dynamic sparse training methods, such as Sparse Momentum, Dynamic Sparse Reparameterization (DSR), Dynamic Sparse Training (DST) and so on.
• Data Parallel: The CAPE consists of fully-connected layer. Therefore, better data parallel can significantly reduce the time.
• GPU with high memory width. The CAPE read/write the matrix with size $N^2$. Therefore, a GPU with a higher memory width will help improve the speed of CAPE.
• GPU with high computation compatibility. The CAPE consists of a multiple-layer perception network, which is computationally dense. Therefore, GPU has better computation compatibility can help speed up the
• Finally, with the increase in hardware, the cost of CAPE will be acceptable. For example, with the development of GPU, the Large Model gradually is accepted, while it is unimaginable to train a 175B model 10 years ago.

The additional training ratio will gradually decrease with a larger model size, compared to baseline Kerple.

The following is the time cost with a training length of 512 with micro_gpu_batch_size 1.

Apparently, when the model becomes large, the additional computational cost of CAPE gradually decreases. Therefore, the CAPE may be a potential good choice for an extremely large language model.

Moreover, CAPE indeed can speed up training, compared to current popular RoPE

With the same training token, the CAPE with a training length of 512 and batch size of 8 can even with comparable performance with a RoPE training length of 4096 and batch size of 1. Also, the CAPE with a training length of 512 and batch size only takes 192.45ms, while RoPE takes 265.48 ms. Therefore, the CAPE could be a choice for speeding up training in the future.

Q2: How CAPE is using long context information and haystack test

A2: We analyze how CAPE uses the long context information via visualization analysis, as shown in Figure 1 and Appendix. According to Figure 1 and Appendix, the CAPE not only helps the model to pay attention to the local information but also helps the model to look at the information far away, with different heads having different functions.

The haystack test, requires more training (including pretrain and alignment) so that the model can follow the instructions to finish the test, while currently, we do not have such resource to train such a model. However, in our CHE benchmark, there is one task named missing duplicate that may be similar to haystack test.

• haystack test:
  • It works by embedding specific, targeted information (the “needle”) within a larger, more complex body of text (the “haystack”).
  • The goal is to assess an LLM’s ability to identify and utilize this specific piece of information amidst a vast amount of data.
• CHE Benchmark Missing Duplicate Task:
  • The input is a binary string of the form $ww$ where w is itself a binary string. One token in this string has been hidden, and the network must find out which one it is, by looking at the value at the same position but on the other side of the string. For instance, if $x = ab$_$aba$ (i.e., w = aba), then $y = a$.
  • The goal is also to assess a model's ability to identify and utilize this specific piece of information amidst a vast amount of data.
• As shown in Table 2, Kerple only gets 79.06% accuracy, while CAPE-Kerple gets 87.56%. This suggests the ability of CAPE on the haystack test.

If there are any questions, please let us know. And if you think that we have addressed your concerns, could you please consider raising the score? Thank you very much for your support.

Dear Reviewer NzBk,

Thank you very much for appreciating our work. We will address your concerns below.

Q1: The efficiency of CAPE

A1: With the model size increase, the additional computing cost ratio will decrease, compared to baseline Kerple. Moreover, the CAPE can even speed up the training because a smaller training length of CAPE can still achieve good performance. We discuss this in CAPE computation cost in Author Rebuttal by Authors.

Q2: The choice of $D_{CAPE}$

A2: The CAPE is relatively robust to the choice of $D_{CAPE}$, as shown in Figure 6. To achieve satisfactory performances, we find that CAPE with a modest size (not large $D_{CAPE}$) is sufficient. For convenience, we will copy the results (Arxiv dataset) below.

A3: We further conduct experiments on 2.7B and 6.7B model sizes, which proves that the CAPE still works well. We further analyze the result of 2.7B and 6.7B in Result on Large Model Size 2.7B and 6.7B in Author Rebuttal by Authors. We add larger model (2.7 and 6.7BB) experiments in the following, with micr_gpu_batch_size 4 and length 512 (Books Dataset).

Q4: Can CAPE adapt to LLMs, for example, Llama3-8B, with few post-training steps?

A4: The comment is quite helpful. Although we train from scratch for all transformer models with CAPE (that is the main reason that we did not try very large transformer models due to the computation limitation), it is still possible that we merely train on the CAPE part (i.e., the introduced MLP) but freeze (or fine-tune) all other parameters of transformer models. This strategy will definitely accelerate the training and require fewer post-training steps. We will discuss the possibility of adapting CAPE to LLMs with fewer post-training steps in the paper.

Q5: What's the context length boundary of CAPE? In this paper, experiments just demonstrate the maximum context length of 8192.

A5: We further validate CAPE on length 16384, which proves that CAPE still works well. For the length of 32768, the GPU reports out-of-memory. Therefore, we believe that the context boundary of CAPE is more than 16484, while the training length is 128.

Q6: What does "semantically dependent" mean?

A6: Thank you for pointing out. The "semantically dependent" indicates that our position encoding value depends on the semantics (which is the attention score in this paper). For clear presentation, we will change such words to context adaptive or contextually dependent. And we will also change explain that the context in the paper is attention score.

Q7: Why does CAPE need to incorporate such a dynamic selection mechanism, which is already learned by LLMs, into the PE

A7: The reason is that the CAPE has more expressiveness.

Suppose that the ``optimal'' attention mechanism is composed of the key-query multiplication (denoted as $A(x)=XW_Q(XW_K)^T$) and the additive positional encoding bias (denoted as $B(x)$), i.e., $A_{optimal}(x)=XW_Q(XW_K)^T+B(x)$. Here, x and X are the input sequence and the corresponding token embeddings. In previous static PE, $B(x)$ is set as a constant for all input sequence, i.e., $B(x) = B$ and the constant $B$ is optimized across all samples $\{x\}$ during training. However, our main contribution and claim is that the optimal positional encoding should vary for different sequences. Therefore, we proposed the dynamic context-adaptive PE (CAPE), where $B(x)$ depends on both the sequence and the positional information. In contrast with the static PE, the proposed CAPE can adjust dynamically with the input context, and is optimal for each input.

Speaking from higher-level, we see that the general fixed and static PE (even though the PE is learned from data) is an averaged optimal positional encoding over all training samples, while the dynamic PE is context-dependent and is optimal for each sample. That is the core motivation of using dynamic PE rather than the static PE.

If the the optimal solution is $B(x) = B$, the $f(XW_Q(XW_K)^T, B)$ can be reduced to zero as $f(*)$ is a universal approximate function (two-layer mlp with activation). Therefore, our CAPE is at least no worse than the baseline.

If there are any questions, please let us know. And if you think that we have addressed your concerns, could you please consider raising the score? Thank you very much for your support.

Ok, I appreciate the authors' professional response.

In your paper, on line 233, the authors mention, "As shown in Figure 4 and Table 5, CAPE consistently outperforms established baselines such as RoPE" (I know you mention Figure 2, please correct this in the final revision). It is surprising for me to know that Additive Position Embedding works better than RoPE-based methods.

I also read the general response by the authors, the ``Q2: Result on Large Model Size 2.7B and 6.7B (Reviewer tTna, Reviewer NzBk)'' also surprised me, as it indicates that additive RPE can achieve much better results than RoPE under the context scaling settings.

1. Can you explain why the additive RPE is better than the RoPE-based model here intuitively?

2. Are all the experiments in Figure 2 fair and consistent? Including model size, training data, etc.

3. From the results in Figure 2, it seems that the mainstream RoPE is not as effective as Additive RPE. What do you think about this phenomenon? Should the mainstream RPE now use additive RPE rather than cumulative RPE like RoPE?

I hope the authors answer the above questions, which are very important for my judgment of your work.

Dear Reviewer NzBk,

Thank you very much for your reply. We will answer your questions below.

Q1: The RoPE and its B(x)

A1: We claim "B(x) is set as a constant for all input sequence" under the discussion of additive RPE. Usually, the RoPE is not considered as an additive relative position encoding, as discussed in FIRE [3] paper Section 2.2.

We will answer the question about RoPE in two parts: 1) The definition of Additive RPE; 2) General view of the $B(x)$ and RoPE

The definition of Additive RPE, which should have formulation $A(x) = XW_{Q}(XW_{K})^{T} + B$ and B is induced by position encoding function

According to FIRE paper [3], the additive relative position encoding can be represented by the formulation $A(x) = XW_{Q}(XW_{K})^{T} + B$, while the B $\in R^{N \times N}$ is induced by the position encoding function $N^{*2} \to R$. To be more specific, $B=g(D)$, where $g(.)$ is a position encoding function and $D$ is the distance matrix. Usually, the $D$ is the following:

D= [0 0 0]
   [1 0 0]
   [2 1 0]

Therefore, when we are discussing the $A(x) = XW_{Q}(XW_{K})^{T} + B$, usually the $B$ should fulfilled the mentioned requirements. According to FIRE paper [3], the RoPE is not an additive relative position encoding so the RoPE actually does not have such $B$. Therefore, under the original definition, RoPE actually does not have either $B$ or $B(x)$. And for our $B(x)$, we mainly focus how to utilize the additive RPE matrix $B$ and attention score, so that there should be an additive relative position encoding bias matrix $B$.

Therefore, considering the previous definition, we will have the following:

• The naive query and key implementation: $XW_{Q}(XW_{K})^{T}$
• The part that utilizes both bias matrix B and $XW_{Q}(XW_{K})^{T}$: $B(x)$
• The part is that fixed after training: $B$
• For the previous naive additive position encoding method: $B(x)=B$
• For our CAPE: $B(x)=B+f(XW_{Q}(XW_{K})^{T}, B)$.

General view of the $B(x)$ and RoPE

If we do not consider the definition from FIRE[3] and make the definition of $B$ becomes more general, then we may have the following.

• The naive query and key implementation: $XW_{Q}(XW_{K})^{T}$
• The part that removes naive query and key implementation: $B(x)$
• The part is that fixed after training for additive RPE: $B$
• For the previous naive additive position encoding method: $B(x)=B$
• For our CAPE: $B(x)=B+f(XW_{Q}(XW_{K})^{T}, B)$.
• For RoPE: $B(x)=Qf(K)^{T} + f(Q)K^{T} + f(Q)f(K)^{T}$. From the general additive RPE perspective, this may be why RoPE is better than its baseline Sinusoidal encodings.

Q2: Why CAPE is better than other RPE methods e.g.

A2: Compared to other RPEs e.g. YaRN (YaRN is also a great work), we further improve the performance of the current great method Kerple and FIRE[3] via dynamic position encoding to get better performance. Also, YaRN improves performance via sequence-length-related adjustment.

• Our baseline is powerful. We may check the FIRE paper[3] Figure 1 (on page 2). Figure 1 of the FIRE paper compares the performance between Kerple, and FIRE and other RPEs (including YaRN). The results prove that Kerple and FIRE achieve relatively good performance.
• CAPE solves the baseline Kerple and FIRE's limitation, whose position encoding is fixed after training. As we claim in our paper, though additive RPE achieves good performance, its position encoding is fixed after training. Hence, we propose to further improve the additive relative position encoding performance by dynamically adjusting the position encoding via attention score.
• Therefore, our method could achieve relatively better performance than other methods.

Q3: The definition of "better extrapolation than other RPE methods".

A3: The better extrapolation means: with the same experiment setting (training length, training tokens, and so on), our proposed method could achieve better performance on evaluation length $T_{eval}$, while $T_{eval}$ is larger than training length $T_{train}$.

For example, we train CAPE-Kerple with the maximum training length of 1024, and we could say that CAPE-Kerple could have better extrapolation performance if CAPE-Kerple achieves better performance on evaluation length 8192. We follow the definition of length extrapolation from previous works [1][2][3].

Thank you very much for your reply. If there is any other question, please let us know.

Reference:

[1] Press, O., Smith, N., & Lewis, M. Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation. ICLR, 2021.

[2] Chi, T. C., Fan, T. H., Ramadge, P. J., & Rudnicky, A. (2022). Kerple: Kernelized relative positional embedding for length extrapolation. NIPS, 2022

[3] Li, S., You, C., Guruganesh, G., Ainslie, J., Ontanon, S., Zaheer, M., ... & Bhojanapalli, S. Functional Interpolation for Relative Positions improves Long Context Transformers. ICLR, 2024.

Q3: From the results in Figure 2, it seems that the mainstream RoPE is not as effective as Additive RPE. What do you think about this phenomenon?

A3: The following is our personal opinion. The mainstream method is RoPE because of three reasons: 1) The timing of RoPE and additive RPE ; 2) Within training length, RoPE has comparable performance with additive RPE; 3) The development of LLaMA.

The timing of RoPE and additive RPE The RoPE paper is on arxiv from 20 April 2021, while there was no Alibi (27 Aug, 2021), Kerple (20 May, 2022), or FIRE (3 Oct, 2023).

Within training length, RoPE has comparable performance with additive RPE. As shown in Figure 2 and Appendix 5, for the performance within training length, the proposed RoPE achieves close or similar performance with Kerple and FIRE. For example, with a training length 512, RoPE achieves 4.5755 ppl, Kerple achieves 4.5817 and FIRE achieves 4.5741. Therefore, for the performance within training length, the different is not very large. The additive RPE presents its superiority for length extrapolation, and the length extrapolation problem has become important recently.

Therefore, as previously we mainly considered performance within training length, the RoPE is enough.

The development of LLaMA

• LLaMA uses RoPE. Therefore, if anyone is working on open-source LLM, then they will use LLaMA so that they will use RoPE.
• Therefore, we could find that the YaRN, CLEX, ChunkLLaMA or other works that focus on length extrapolation or long-context all focus on RoPE.
• Also, currently data is relatively important so the architecture from LLaMA to LLaMA 3 is not changed a lot. Hence, the current mainstream position encoding method is RoPE.

Q4: Should the mainstream RPE now use additive RPE rather than cumulative RPE like RoPE?

A4: It is an interesting question, and we are also curious about it. Both additive PE and RoPE-based PE try to incorporate the key-query similarity with the positional information of tokens. However, they adopt different operations, using addition and multiplication operations respectively. In this paper, we developed an additive context-adaptive PE. It is hard to say which kind of PE (additive or RoPE-based) will be the mainstream finally. [RoPE is widely recognized and used in Llamma models.]

Based on the current evidence and the experiment results, additive RPE may be a better choice.

• The FIRE has prove the it could achieve better performance than RoPE, whatever within training length or beyond training length.
• CAPE further helps achieve better performance within training. As proved in our paper (Figure 2 and Table 5), the CAPE could help additive RPE achieves better performance within training length.
• CAPE further helps achieve better length extrapolation performance. Also, as proved in our paper, the CAPE could help additive RPE achieve better performance beyond training length.
• Therefore, based on the current evidence and experiment results, the additive RPE may be better.

Thank you very much for your constructive comments. If there is any further question, please let us know.

Dear Reviewer CfZF,

Thank you very much for appreciating our work. We will address your concerns below.

Q1: The performance of transformer-xl relative encoding

A1: We have shown the performance of transformer-xl below, which also presents great length extrapolation performance. Also, the Relative in Table 2 is just the transformer-xl relative encoding. The experiment is conduct on 125M model with training length 512 and batch size 32.

According to the experiment, the transformer-xl achieves good performance on length extrapolation. Therefore, this also suggests that the position encoding should interact with attention/query/key to further improve the performance.

Q2: The computation expense

A2: Thank you very much for your precious comment on the compuation cost. Yes, we concatenate the attention and bias on the head dimension, and then we use an mlp to process them to dynamically adjust the position encoding values. We have further analyze the CAPE cost in CAPE computation cost in Author Rebuttal by Authors.

Q3: Mention appendix implementation section G when discussing function f and the multi-head CAPE in the main paper.

A3: Thank you very much for your suggestion. We will revise the paper with the following sentence:

Original: It then outputs $h$-dimensional vectors, where each element corresponds to the CAPE for the respective head.

Revised: It then outputs $h$-dimensional vectors, where each element corresponds to the CAPE for the respective head. We have shown the code implementation in Appendix G.

Q4: Present the pseudocode with einops-style syntax

A4: Thank you very much for your suggestion. We will add the einops-style syntax in Appendix G.

import torch
import torch.nn as nn
from einops import rearrange, repeat
class CAPE(nn.Module):
    def __init__(self, head_number=12, mlp_width=12):
        """
        CAPE attention bias module.

        Args:
          head_number: number of attention heads.
          mlp_width: Width of MLP.
        """
        super(CAPE, self).__init__()

        self.mlp = nn.Sequential(
            nn.Linear(2 * head_number, mlp_width),
            nn.LeakyReLU(),
            nn.Linear(mlp_width, head_number)
        )

    def forward(self, attention: torch.Tensor, bias: torch.Tensor):
        """
        Args:
          attention: input sequence, which is q^T * k,
             shape [bsz, num_heads, seq_len, seq_len]
          bias: bias matrix, which can be generated by Alibi, Kerple 
          FIRE or other additive position encodings
             shape [1, num_heads, seq_len, seq_len]

        Returns:
          attention with CAPE,
          shape [bsz, num_heads, seq_len, seq_len]
        """
        # Repeat the bias for batch size
        bias_tile = repeat(bias, '1 h l1 l2 -> b h l1 l2', b=attention.shape[0])
        
        # Concatenate attention and bias
        attention_bias_concat = torch.cat((attention, bias_tile), dim=1)
        
        # Rearrange the dimensions for MLP processing
        attention_bias_concat = rearrange(attention_bias_concat, 'b h l1 l2 -> b l1 l2 h')
        
        # Apply the MLP
        attention_bias_concat = self.mlp(attention_bias_concat)
        
        # Rearrange back to original dimensions
        attention_bias_concat = rearrange(attention_bias_concat, 'b l1 l2 h -> b h l1 l2')
        
        return attention + bias + attention_bias_concat

If there are any questions, please let us know. And if you think that we have addressed your concerns, could you please consider raising the score? Thank you very much for your support.

Dear Reviewer CfZF,

Thank you very much for appreciating our work. We will follow your suggestion to add the new experiment results and discuss the transformer-xl in our paper, and we will update it immediately when we are allowed to revise the paper. Also, we will answer the two questions mentioned in the following:

Q1: The Experiment of Alibi

A1: The experiment of Alibi is the following.

For model size 2.7B:

For model size 6.7B:

The reason for missing Alibi experiment result: with the model size increase, we face some engineering challenges:

• Our implementation is based on the framework GPT-NeoX, which already implements the Alibi position encodings.
• You may notice that our maximum evaluation length is 8192 for 125M and 350M, while the maximum evaluation length for 2.7B is 4096 and the maximum evaluation length for 6.7B is 2048. The reason is that the Alibi position encoding faces out-of-memory challenges for the evaluation length 8192 for the 2.7B model size and 4096 for the 6.7B model size.
• Therefore, considering two aspects ( 1. evaluation as long as possible; 2. the most popular position encoding method is RoPE), we mainly present the results of RoPE, T5's bias, Kerple and CAPE-Kerple in rebuttal so that the evaluation length can be longer (evaluation length 4096) for both 2.7B model size and 6.7B model size.
• If you would like to see more experiment results, please let us know. We will try our best to finish the experiment as soon as possible.

Q2: Discussion with XPos

A2: The CAPE is designed for Additive Relative Position Encoding, while XPos is an improved version of RoPE (Note: RoPE is Relative Position Encoding, but NOT Additive Relative Position Encoding). Therefore, CAPE and XPos focus on different directions of position encodings.

The Difference between XPos and CAPE

The XPos:

• The implementation of RoPE: $f_q(q,n)=qe^{i\theta n}$
• The implementation of XPos: $f_q(q,n)=qe^{\xi n+i\theta n}$
• Apparently, the XPos is an improved version of RoPE, while XPos degrades to RoPE when $\xi$ becomes zero.

The CAPE:

• The implementation of Additive Relative Position Encoding: $A(x)=XW_Q(XW_K)^T+B$. The $B$ is the bias matrix, which could be Alibi, Kerple, FIRE, or other potential additive relative position encoding methods.
• The implementation of CAPE: $A(x)=XW_Q(XW_K)^T+B + f(XW_Q(XW_K)^T, B)$.
• Apparently, the CAPE can be applied to any potential additive relative position encoding methods.

The performance of XPos and CAPE-Kerple

The previous paper BiPE [1] conducted experiments using XPos on the Arxiv dataset. As shown in the BiPE [1] paper's Figure 4, with the training length 1024, the perplexity of XPos increases quickly from about 6 ppl (at evaluation length 1024) to about 16 ppl (at evaluation length 6144), while our proposed method (CAPE-Kerple) decreases the ppl from 5.21 (evaluation length 1024) to 5.00 ppl (at evaluation length 8192) with training length 128. This suggests that our method should be better than the mentioned XPos.

Therefore, the XPOS is an improved version of RoPE (note that RoPE is not additive relative position encoding), while CAPE generally improves the performance of the additive relative position encodings, including Alibi, Kerple, FIRE and so on.

Finally, if you have any questions or would like to discuss anything, please let us know. We will try our best to share our opinions or conduct experiments.

Reference:

[1] He, Z., Feng, G., Luo, S., Yang, K., He, D., Xu, J., ... & Wang, L. (2024). Two Stones Hit One Bird: Bilevel Positional Encoding for Better Length Extrapolation. arXiv preprint arXiv:2401.16421.

Dear Reviewer CfZF,

Thank you very much for your reply, and thank you very much for your support. We promise that we will include the additional details/experiments in the paper/appendix.

Dear Reviewer tTna,

Thank you for the detailed review. We will address your concerns below.

Q1: The major concern is in the experimental evaluation part..

A1: We have presented the experimental results besides PPL, as shown in Page 9 Section 4.7 and Appendix D Experiments on Chomsky Hierarchy Evaluation Benchmark (the benchmark and its variants are used in previous works), which uses accuracy as evaluation metrics.

Four convenience, we directly copy the result of CHE Benchmark here. We train with 200K steps and length 40, while we test on length 500. The random accuracy is 50%, except for modular arithmetic simple, cycle navigation, bucket sort, solve equation, and modular arithmetic brackets, where it is 20%. \dagger$$\dagger$$\dagger denotes permutation-invariant tasks, which are expected to be solved without positional information.

Comparative performance improvements. CAPE consistently enhanced performance across various tasks, especially on permutation-variant tasks. Specifically, CAPE improved upon Alibi and FIRE's results in all 11 tested permutation-invariant tasks. Similarly, it outperformed Kerple in 10 of these tasks.

Q2: Experiments on large model size

A2: The following is the result of the experiment on 2.7B and 6.7B, with training length 512 and micro_gpu_batch_size 4. We further discuss the experiment on large model size in Result on Large Model Size 2.7B and 6.7B in Author Rebuttal by Authors

Q3: Explanation of Equation 2 and Equation 3

A3: The next step of our operation is $softmax$. Therefore, $h(A_{CAPE}(x))=softmax(XW_Q(XW_K)^T+f(XW_Q(XW_K)^T, B))!=softmax(XW_Q(XW_K)^T+B)=h(A(x))$, while $A(x)$ is the naive transformer attention score calculation without CAPE.

Let us explain the attention module in Transformer step by step:

• Step 1 (Calculate attention score via query $XW_Q$ and $XW_K$):
  • Previous implementation of $A_{score}=XW_Q(XW_K)^T+B$
  • CAPE implementation of $A_{score}=XW_Q(XW_K)^T+B+f(XW_Q(XW_K)^T, B)$
• Step 2 (use softmax on attention-score, row-wise. Therefore, the next step is $**softmax**$):
  • $A_{scoreSoftmax}=softmax(A_{score})$
• Step 3 $A_{scoreSoftmax}$ and value $XW_K$ to get embedding of each token:
  • output=$A_{scoreSoftmax}XW_V$

If we understand correctly, the neural network h(∗) mentioned in the comment corresponds to the feedforward network layer (FFN), that follows the attention layer. We would like to clarify that the MLPs introduced in the CAPE model are not duplicated. In CAPE, these MLPs dynamically adjust the positional encodings based on context information (CAPE: Step 1). Subsequently, softmax operations are applied across attention score, row-wise. It is important to note that the CAPE's adjustments, applied during the attention phase, do not directly alter the token values but act on the attention computation. However, the FFN layer that follows attention modifies (FFN: after Step 3) each token through a nonlinear transformation.

If there are any questions, please let us know. And if you think that we have addressed your concerns, could you please consider raising the score? Thank you very much for your support.

Dear Reviewer tTna,

We would like to thank you again for your detailed reviews. We have updated the experiment results and the explanation of our method in the above response.

As the discussion period will be closed soon, we would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions. We would be happy to do any follow-up discussion or address any additional comments.

Again, thank you very much for your attention to our work.

Dear Authors,

Similar to question#2 raised by AC, I also have this question (see Question #1 and #3 posed in Re-Response to Authors by NzBk).

The authors can combine these questions to respond, thx.

Reviewer NzBk

Q3: Could you clarify why for CAPE in contrast to all other baselines you observe improvement with the longer context, while other degrade always?

A3: The CAPE performance is highly related to the baseline method, such as Kerple or FIRE. Though CAPE could usually improve additive RPE method performance, CAPE may achieve worse performance if the baseline additive RPE method collapses.

The following are the results of FIRE and CAPE-FIRE on the Books3 dataset. With a training length of 512, CAPE can consistently improve performance.

And if we change the training length from 512 to 128, we will have the following results.

According to the experiments, we can see that the CAPE-FIRE performance could be even worse when the FIRE performance almost collapsed (which is 802.76 ppl at evaluation 8192).

Therefore, the CAPE performance is highly related to the baseline additive RPE performance. CAPE could improve performance when the baseline method could achieve acceptable performance under the current situation, while the performance may degrade if the baseline method collapses under the current situation.

Q4: Could it be that with even longer context the trend will continue?

A4: As we discussed above, the performance of CAPE is highly related to the baseline method performance, such as Kerple and FIRE. For Kerple, our method still works well when the training length 128 and the evaluation length is 16384.

Q5: Any intuition from the learned representation why this happens? Maybe compare Fig 1 for position 2k, 4k, 8k to see the trend of change in how the bias is learnt?

A5: We have presented more examples of the Figure 1 visualization in Appendix F, with length 512, 2048 and 8192. And the training length 512. Intuitively, our CAPE works better because it has both local and anti-local position patterns, which is achieved by different heads.

• The CAPE could really adjust the bias matrix value for different samples. In Appendix F, for each validation length (512, 2048 or 8196), we present the visualization of two different samples. Let us first focus on the evaluation length 512.
  • For sample 1 with evaluation length 512, its layer 2's $4^th$ head output bias value is mostly smaller than 5.
  • For sample 2 with evaluation length 512, its layer 2's $4^th$ head output bias value could be larger than 5.
  • This suggests that the CAPE can really adjust the bias matrix for different samples.
• Different heads have different functions (trends), whatever the evaluation length is 512, 2048, or 8192 .
  • Whatever the evaluation length 512, 2048 or 8192, it seems that the different layers' different heads have their own functions after training, which may help improve the model performance.
  • For the original attention score, it seems that different heads may not have different functions. This may be why NoPE does not work very well.
  • For the Kerple bias, all heads pay attention to local information.
  • For our CAPE, the different heads have different functions
    • The layer 1's head 8 is used to pay attention to long-distance information.
    • The layer12's head 7 is used for the local information.

Dear Area Chair a2or,

Thank you very much for your reply. We will answer your questions below.

Q1: Wanna to exclude inductive bias on $B$ here

A1: We present the CAPE $QK^T+f(QK^T, B)$ in the following tables (Figure 5).

According to the question, the CAPE variant $QK^T+f(QK^T, B)$ may be what we need. If we need to try other variants, please let us know.

Q2: Looking at the Appendix F if I correctly interpret the leftmost plot - this is the final attention.

A2: The leftmost plot is $QK^T$. We have described the figure Attention, Figure Kerple Bias and CAPE bias in Figure 1. We directly copy it here.

• (1) The Attention (the leftmost plot) is $QK^T$, which is the naive attention implementation
• (2) The Kerple bias (the middle plot) is $B$;
• (3) The CAPE (with Kerple) bias (rightmost plot) is $f(QK^T,B)$.
• Therefore, the final attention is (1)+(2)+(3). We promise that we will provide the final attention visualization in the final revision.

Q3: Baseline of NoPE

A3: We have run the baseline of NoPE, as shown in Figure 1. For convenience, we directly copy our results here.

With sufficient training, the NoPE may achieve relatively good performance

• When the training length is 128, the NoPE is worse than RoPE, from the evaluation length 128 to 8192
• When the training length is 512, the NoPE could get better performance than RoPE on evaluation length 1024.
• Therefore, NoPE actually could potentially have acceptable length extrapolation performance, while it needs sufficient training.

For the NoPE attention visualization

According to your suggestion, we also visualize NoPE.

On Arxiv Dataset

• With evaluation length 512:

  • For the early layer (layer 1 to 2): its attention distribution is similar to some leftmost plots, where all heads look in average per position. This may because early layer does not receive sufficient implicit position information.
  • For the middle layer (layer 3 to 7)
    • Most attention heads present long-term decay. This decay may follow either a linear relationship with distance or a logarithmic pattern.
  • For the latter layer (layer 8 to 12)
    • The attention distribution gradually becomes relatively similar to the early layer.

• With the increase in the evaluation length, such long-term decay becomes more obvious.

On Books3 Dataset

• The Books3's attention distribution in different layers is relatively similar to the Arxiv dataset's head pattern.
• Also, we observed the non-local attention head pattern, such as layer 2's head 9. For example, with the distance increase, the attention score of layer 2's head 9 does not decrease.
• This may explain why additive RPE and CAPE are better: explicitly help provide the needed local and anti-local attention pattern.

Finally, we promise that we will provide the NoPE attention visualization and discussion in the final revision.

Thank you very much for your constructive comments. If there is any further question, please let us know.

Dear Area Chair a2or,

Thank you very much for your reply. We will answer your question below.

Q1: The attention pattern of Kerple and CAPE-Kerple

A1: To purely analyze the function $f$, we further analyze the visualization with $QK^T+f(QK^T,B)$ (The Equation 2 for CAPE in our paper) so that the total bias matrix is $f(QK^T,B)$, without being affected by the residual $B$.

For CAPE-Kerple bias with CAPE formulation $QK^T+f(QK^T,B)$.

According to the visualization, the total bias matrix $f(QK^T,B)$ presents clearly anti-local attention pattern. For example, the layer 1's head 8 and head 2 present clearly anti-local attention pattern, which is relatively similar to the anti-local pattern shown in the Appendix F layer 1's $f(QK^T,B)$ in formulation $QK^T+B+f(QK^T,B)$. Hence, we could have confidence to say that the proposed CAPE bias could present anti-local attention patterns. The function $f$ not only changes the slope of decay but also provides the anti-local attention pattern if needed.

The final attention pattern comparison between Kerple and CAPE-Kerple

We also further visualize the Kerple final attention(which is $QK^T+B$) and CAPE-Kerple final attention (which is $QK^T+f(QK^T,B)$) . As the CAPE total bias matrix $f(QK^T,B)$ could clearly provide the anti-local attention pattern, the CAPE-Kerple final attention ($QK^T+f(QK^T,B)$) successfully clearly presents the anti-local attention pattern on some attention heads so that Kerple final attention pattern and CAPE final attention pattern are different.

We promise that we will provide the visualization of Kerple and CAPE-Kerple ($QK^T+f(QK^T,B)$) in the final revision

Q2: Why you don't have perfect generalization to the longer sequences? Any explanation/intuition?

A2: The previous work already proves that the NoPE performance may not be uniform and independent from the sequence length, such as NoPE Paper[1] and FIRE [2].

• According to the NoPE paper[1], the performance of NoPE is not uniform and independent of the sequence length. For example, for the NoPE Paper [1] Figure 3, the performance of NoPE will be worse if the evaluation length is larger than the training length. The NoPE paper claims that there is potential for NoPE to have better extrapolation performance than other position encodings, but it does not claim that the NoPE performance is uniform and independent from the sequence length.
• The FIRE paper also compares the performance of NoPE in FIRE paper's Figure 1. According to the figure, we can find that the NoPE performance becomes worse (perplexity also increases) when the evaluation length is larger than the training length.
• Our personal opinion.
  • As claimed by NoPE Paper [1], the NoPE implicitly learns the position encoding information, which may not be easy to learn and is inefficient.
  • Therefore, we may need to sufficiently train the model if we use NoPE.
• Therefore, according to the previous works (such as [1][2]), the NoPE performance may not be uniform and independent of the sequence length.

Thank you very much for your constructive comments. If there is any further question, please let us know.

Reference:

[1] Kazemnejad, A., Padhi, I., Natesan Ramamurthy, K., Das, P., & Reddy, S. (2024). The impact of positional encoding on length generalization in transformers. Advances in Neural Information Processing Systems, 36.

[2] Li, S., You, C., Guruganesh, G., Ainslie, J., Ontanon, S., Zaheer, M., ... & Bhojanapalli, S. Functional Interpolation for Relative Positions improves Long Context Transformers. In The Twelfth International Conference on Learning Representations.

After the reply submission, it seems that the openreview system does not send email and notification. Therefore, we resubmit the reply below

Q4: The performance of Kerple

A4: This may be caused by the evaluation protocol or training strategy.

For our evaluation protocol, we select the last 256 tokens in the sequence to evaluation so that the evaluation is more critical and could better reflect the model performance with the current sequence length.

According to the Kerple paper, you may find that its reported results achieves relatively better performance than our paper, FIRE and BiPE paper[1]. For example, the RoPE result in Kerple's paper achieves 5.94 ppl with an evaluation length of 2048 on the Github dataset (training length is 512), while BiPE reports that RoPE achieves 15 ppl with an evaluation length of 2048 on the Github dataset (training length is 1024). This may be caused by different evaluation protocols or training strategies.

• Moreover, even if we directly use previous works' results, our proposed method (with a more critical evaluation protocol and the baseline is lower) could still achieve better performance.
  • For example, on the arxiv dataset Kerple's paper reports that it could achieve 5.01 ppl on evaluation length 8192 (with training length 512), and our proposed method achieves 3.70 ppl on evaluation length 8192 (with training length 512).
  • This suggests that even with a more critical evaluation protocol, and proposed method could still achieve good performance.

Q5: how all results on 128 length are helpful for overall conclusions and claims of the paper

A5: It suggests that our proposed method has the potential to extrapolate 64 times, compared to the training length

We would like to know the context boundary. Because of the limitation of computational resources, we would like to know the largest many times that our proposed method could extrapolate. With a training length of 128, we successfully extrapolate to 8192, suggesting the potential that our proposed method could extrapolate 64 times the training length.

Q6: The question for prior papers and Figure 5

A6: We will answer the question in two parts: 1) The question for the prior paper; 2) The Figure 5 baseline is Kerple, and thank you very much for your notice.

The question for the prior paper As we mentioned above, the prior paper may use different evaluation protocols and experiment settings so that their performance may be better. For example, the Alibi uses the sliding-window evaluation protocol.

For the previous paper that claims they could extrapolate to 16K. We may pay attention to the FIRE paper. According to FIRE figure 6, we can find that it is hard for the prior paper method to keep performance on evaluation length 8192 (when the training length is 512).

For Figure 5 caption

Thank you very much for your notice. The Figure 5 caption should be CAPE-Kerple but not CAPE-Alibi. We will revise it in the final revision.

Q7: Question about NoPE

A7: The training setting may be different. As we said above, the NoPE may need more training tokens and length. For our experiments, we train with 50K with length 512 (following Kerple setting) and 8 GPUs on the Arxiv or Books dataset, while FIRE runs 600K steps with length 2048 and 128 GPUs.

Also, as we said above, the NoPE implicitly learns the position information so that we may sufficiently train the model so that the NoPE can present its ability to outperform other methods.

Q8: Not clear why the function will be able to do better generalization in general.

A8: We further conduct an ablation study on the $f$, proving that $f$ help enhance the bias matrix. CAPE improves the bias matrix for $A_{final}$, while the $A_{final}$ is used to calculate $A_{final}K$. For the unseen position, the $B$ partially could handle it (FIRE changes the problem to interpolation), but is not accurate enough so that CAPE helps enhance the bias matrix $B$ via attention score

• The CAPE $f(QK^T, B)$ is better than naive $f(B)$, suggesting that the context-adaptive is important.
• The $QK^T+B+f(B)$ is better than $QK^T+B$, suggesting that benefiting from improving the expressiveness of the bias matrix.

We promise that we will release the code. We are sure that the experiment setting is fair and the experiment results are reproducible

Thank you very much for your constructive comments. If there is any further question, please let us know.

Reference

[1] He, Z., Feng, , ... & He, D. Two Stones Hit One Bird: Bilevel Positional Encoding for Better Length Extrapolation. ICML, 2024.

[2] Ruoss, A., , ... & Veness, J. (2023). Randomized positional encodings boost length generalization of transformers.

Dear Area Chair a2or,

Thank you very much for your reply. We will answer your question below.

For the question related to experiments:

• Our Table 2 CHE results (with 200K training steps) are almost aligned with the baseline[2].
• Previous work (such as Alibi) uses the sliding-window evaluation protocol, which divides sentences with L tokens and then evaluates. This will make the performance look better.
• Previous works (such as Kerple) also use non-overlapping evaluation, while we choose the last 256 tokens so that the evaluation is more critical. Our Kerple implementation comes from its released code.
  • Kerple's paper uses non-overlapping evaluation: To train on or evaluate a sequence longer than $L$ tokens (which is usually the training length), it is typical to segment the sequence into L-length subsequences and train on or evaluate them independently. This will make the performance look better, compared to our evaluation protocol. For details, please refer to Alibi paper
• With the different evaluation protocols and training settings, the reported results may be different, such as the Kerple performance difference between Kerple paper and FIRE paper. In Kerple's paper, its ppl gradually decrease with the evaluation length increases. However, the FIRE paper Figure 6 (train length 512) reports that the Kerple method ppl gradually increases with the evaluation length increases.
• Moreover, even if we directly use previous works' results, our proposed method (with a more critical evaluation protocol so that the baseline is lower) could still achieve better performance.
  • For example, on the arxiv dataset Kerple paper reports that it could achieve 5.01 ppl on evaluation length 8192 (with training length 512), and our proposed method achieves 3.70 ppl on evaluation length 8192 (with training length 512).
  • This suggests that even with a more critical evaluation protocol, and proposed method could still achieve good performance.
• All our experiments are fair and consistent. We directly use the popular GPT-NeoX framework for our experiments (it has Alibi implementation), and we promise that we will release the code. We are sure that the experiment setting is fair and the experiment results are reproducible.

Q1: Why then in your comparison in Figures 2, 3, 4 and above Tables FIRE is the worst in the generalization? ? Even worse than RoPE?

A1: This may be caused by the different training strategies and settings.

For our experiments, we train with 50K with length 512 (following Kerple setting) and 8 GPUs, while FIRE runs 600K steps with length 2048 and 128 GPUs. The FIRE uses MLP to predict the position encoding so it may need more training tokens and other potential training strategies.

Also, FIRE performance is actually better than RoPE, with a larger training length

Our FIRE implementation: we directly utilize the FIRE implementation code provided in FIRE's paper.

Q2: Why did you decide to select 128 context lengths for training?

A2: We do present the results of training length 512 (Figire 2) and 1024 (Figure 3) in our paper, which follows the setting of Alibi paper and Kerple paper.

Why choose 128? The reason is that we focus on length extrapolation so we would like to know based on the training length $T_{train}$, how long our proposed method could still achieve good performance on evaluation length $T_{eval}$. Basically, there are two ways:

• Train on the longer length and then evaluate on the longer length. However, this will take a lot of computing cost, as we know that the transformer cost is $N^2$, while the N is the sequence length.
• The second way is that we train on a relatively short length (such as 128) and evaluate on a relatively longer length (such as 8192). This could reduce the computational cost.
• Based on our hardware, we additional train on 128.

Q3: Could you explain the choice of Arxiv and Books3 data as e.g. prior works used WikiText, C4, OpenWebText2, GitHub?

A3: The Arxiv dataset and Books dataset are also commonly used in previous works, and their Mean Document Size is relatively bigger.

• The mean document size: Wikipedia (en) (1.11 KiB), OpenWebText2 (3.85 KiB), GitHub (5.25 KiB)
• The Arxiv dataset is commonly used in previous works, such as Kerple, FIRE and BiPE. The Mean Document Size of Arxiv is 46.61 KiB.
• The Books dataset is also commonly used in previous works, such as LongRoPE. Also, The Mean Document Size of Books is 538.36 KiB, which is almost the largest one in Pile dataset.
• As we focus on length extrapolation, we prefer to validate our proposed method on the long-context situations so that Arxiv and Books3 are preferred.

Dear Mingyu Xu,

Thank you very much for your interest in our work. We will answer your question below.

The cost could be divided into two parts: 1) computation cost: 2) real time cost

The Computation Cost

• The cost of Feed-Forward Network is: $O(Nd_{head}^2d_{hidden}^2)$=$aNd_{head}^2d_{hidden}^2$, where a is a constant, N is the sequence length, $d_{head}$ is the attention head number and $d_{hidden}$ is the dimension for attention calculation.
• The cost of Attention: $O(N^2d_{head}d_{hidden})$=$bN^2d_{head}d_{hidden}$, where b is a constant.
• The additional cost of DAPE: $O(N^2d_{head}d_{DAPE})$=$cN^2d_{head}d_{DAPE}$, where c is a constant.

The Real Time Cost: Bontecck is I/O(Read In/ Read Out) when the length is very large, if without flash attention.

If without flash attention, we have to read the matrix with size $T \cdot T$ for more several times. Hence, the time cost bottleneck becomes the I/O time cost.

Moreover, our poster will be presented next week at the NeurIPS conference. We could have a face-to-face discussion, if necessary.

• Poster Location: East Exhibit Hall A-C #3009
• Poster Time: Wed 11 Dec 11 a.m. PST — 2 p.m. PST

Again, thank you very much for your attention, and hope to see you soon.

Best regards,

Chuanyang