DAPE V2: Process Attention Score as Feature Map for Length Extrapolation

## Abstract

Abstract The attention mechanism is a fundamental component of the Transformer model, contributing to interactions among distinct tokens, in contrast to earlier feed-forward neural networks. In general, the attention scores are determined simply by the key-query products. However, this work’s occasional trial (combining DAPE and NoPE) of including additional MLPs on attention scores without position encoding indicates that the classical key-query multiplication may limit the performance of Transformers.
In this work, we conceptualize attention as a feature map and apply the convolution operator (for neighboring attention scores across different heads) to mimic the processing methods in computer vision. Specifically, the main contribution of this paper is identifying and interpreting the Transformer length extrapolation problem as a result of the limited expressiveness of the naive query and key dot product, and we successfully translate the length extrapolation issue into a well-understood feature map processing problem. The novel insight, which can be adapted to various attention-related models, reveals that the current Transformer architecture has the potential for further evolution. Extensive experiments demonstrate that treating attention as a feature map and applying convolution as a processing method significantly enhances Transformer performance.

## 1 Introduction

Transformer-based models (Vaswani et al., ) have delivered exceptional performances across widespread applications, including language processing (Zhang et al., ; Guo et al., ; Ainslie et al., ), computer vision (Chen et al., ; Peebles & Xie, ), quantitative research (Zhou et al., ; Liu et al., ; Wu et al., ), and scientific machine learning (Taylor et al., ; Geneva & Zabaras, ).
However, the quadratic cost of the key-query multiplication for processing a sequence raised much concern about the modern architecture of Transformers especially for long context inputs. To address the issue of storage and computation efficiency, recent research delves into developing more efficient architectures, such as sparse structural attention (Xiao et al., ; Zhu et al., ), adaptive key selection (Xiao et al., ; Fountas et al., ), and hybrid models (Lieber et al., ).
While these adaptations enhance efficiency, they often involve tradeoffs with model effectiveness.

At the same time, there is another voice advocating for refining the model design for tackling complex tasks, rather than prioritizing efficiency. Positional encoding is one of the key components of the attention mechanism.
Although the widely recognized decoder-based Transformer can implicitly incorporate the positional information of tokens, growing evidence both theoretically and empirically shows that the well-designed explicit positional encoding significantly enhances the model performances, especially in long-context tasks (Su et al., ; Press et al., ; Zhao et al., ).
In practice, Transformers depend on positional encoding to explicitly incorporate positional information, enabling the model to make meaningful token predictions. Without these encodings, token generation would lack the necessary contextual order, rendering the outputs nonsensical.
The well-recognized RoPE (Su et al., ), which is adopted in LLaMA (Touvron et al., ), distinguishes the token order by rotating with different angles depending on the token position. However, it demonstrated a notable performance degradation, failing entirely when the input length is double that of the training length (Peng et al., ; Chen et al., ; Ding et al., ). The undesirable performance degradation is also observed for other positional encoding methods, e.g., ALiBi (Press et al., ) and Kerple (Chi et al., ) . FIRE (Li et al., ) alleviates the long-context extrapolation by learnable positional encodings, trying to capture the suitable positional representation by MLPs.
However, a common characteristic among these positional encodings is their predefined and static nature.
Specifically, they are fixed across various tasks and models, which may lead to their inability to adapt to varying input lengths and contexts effectively.
Recently, the data-adaptive positional encoding method, namely DAPE (Zheng et al., ), which adjusts dynamically with context, enhances the length generalization by incorporating the attention scores and positional information with a more complex mechanism.

In this paper, we propose that precise attention scores are crucial for improving Transformer length extrapolation, and we introduce a new perspective on the attention mechanisms. Traditionally, attention scores are computed through the dot product of the query and key vectors. As illustrated in Figure [1](#S3.F1), further processing these attention scores using a neural network—a general case of DAPE (Zheng et al., )—can significantly enhance the length generalization of Transformers, even in the absence of positional encoding (NoPE). Therefore, we suggest treating attention scores as feature maps. By conceptualizing attention as an image feature map (with dimensions $[B,C,W,H]$ for batch size, channel size, width, and height), we can achieve more accurate attention scores by applying techniques used in image processing.
In this work, we employ different kernel sizes (such as 1×3) to process attention, finding that the perplexity (ppl) of attention decreases significantly—from over 600 to just above 100—when trained on a sequence length of 128 and evaluated on a length of 8192.

In summary, our contributions are as follows:

- 1.
We highlight that the coarse attention mechanism, which is the direct result of the query and key dot product, limits the Transformer’s ability to extrapolate to longer sequences. However, Transformers can achieve good length extrapolation performance with careful processing of attention scores.
- 2.
By treating attention scores as feature maps and refining them using image processing techniques like convolution, we can enhance the Transformer’s extrapolation capabilities.
- 3.
We conducted extensive experiments on language tasks to support our claims and believe that these insights can significantly improve the Transformer’s performance in length extrapolation.

## 2 Related Works

#### Absolute Positional Encoding

Absolute positional encoding (APE), introduced by Vaswani et al. (), enables Transformers to incorporate positional information. Specifically, at the first layer, each position $i$ is assigned a real-valued encoding ${\bm{e}}_{i}\in\mathbb{R}^{d}$, which can be either learnable or a fixed sinusoidal encoding (Vaswani et al., ; Kiyono et al., ; Likhomanenko et al., ; Wang et al., ; Liu et al., ), and this encoding is then added to the input sequence. Although this approach is straightforward, Transformers relying on APE tend to struggle with generalizing to longer sequences (Press et al., ).

#### Relative Positional Encoding

Relative positional encoding (RPE) offers an alternative for embedding positional information (Shaw et al., ; Raffel et al., ; Press et al., ). A widely used RPE method in large language models is rotary positional encoding (RoPE)(Su et al., ; Chowdhery et al., ; Touvron et al., ). To address length extrapolation challenges(Press et al., ; Kazemnejad et al., ), positional interpolation (PI) has been introduced (Chen et al., ) to extend the context window. Building on this approach, models like LongLora (Chen et al., ), LongRope (Ding et al., ), YaRN (Peng et al., ), and CLEX (Chen et al., ) have emerged. Another notable direction involves additive positional encoding. For most additive RPE techniques, the computation of pre-softmax attention logits can be expressed using the formula:
${\bm{A}}_{\mathrm{RPE}}({\bm{X}})={\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top}+{\bm{B}},$
where the bias matrix ${\bm{B}}\in\mathbb{R}^{n\times n}$ is derived from the positional encoding function $b:{\mathbb{N}}^{2}\to\mathbb{R}$, with the $(i,j)$-th entry of ${\bm{B}}$ defined as $b(i,j)$. Different parameterizations of $b$ give rise to various RPE variants. Methods supporting arbitrary sequence lengths include T5’s RPE (Raffel et al., ), ALiBi (Press et al., ), Kerple (Chi et al., ), Sandwich (Chi et al., ), and FIRE (Li et al., ). Recently, DAPE (Zheng et al., ) has been introduced, employing MLPs to dynamically adjust bias values based on the input data.

#### Data-Adaptive Related Positional Encoding.

Transformer-XL (Dai et al., ) introduced the use of learnable query and key biases for adaptive positional encodings. Data-Adaptive Positional Encoding (DAPE)(Zheng et al., ) extends this idea by leveraging MLPs to adjust positional encodings based on attention over the head dimension for length extrapolation, ensuring different input data receive unique positional encodings. Contextual Positional Encoding(Golovneva et al., ) further refines this by conditioning position increments on specific tokens, as determined by the model, allowing positions to adapt based on context.”

## 3 Method

In this section, we first review the previously developed Data-Adaptive Positional Encoding method (DAPE), which incorporates attention scores and positional information through MLPs. As a proof-of-concept, our occasional trial on DAPE without the positional information (as shown in Figure [1](#S3.F1)) suggests that regarding attention as a feature map and processing it with classical operators (e.g., convolution) can enhance the Transformers’ behavior. As discussed in some previous works the perplexity scores come mostly from the associative recall (i.e., copy) tasks. In addition, we theoretically show by construction that the proposed method can explicitly realize the associative recall task, in contrast to the implicit conduct through positional encoding in standard Transformers. The two key differences between DAPE (Zheng et al., ) and this work are: 1) Insight: DAPE attributes length extrapolation performance gains to adaptive position encoding, while this work finds DAPE could still improve performance without position encoding so that we take a broader view, explaining that the Transformer’s length extrapolation ability is limited by the expressiveness of the naive query-key dot product, which can be enhanced using image processing techniques; 2) Performance: As shown in Figure [1](#S3.F1), DAPE is designed for additive RPE and may underperform with non-additive RPE (e.g., RoPE), whereas this work suggests that increasing kernel size (e.g., with $\textrm{DAPE}_{1\times 3}$) may improve RoPE’s performance. The $\textrm{DAPE}_{1\times 3}$ implementation is shown in Appendix [J](#A10).

### 3.1 Additive Relative Positional Encoding

For most additive relative positional encoding (ARPE) methods, the computation of pre-softmax attention logits can be unified under the following formula:

$$ ${\bm{A}}_{\mathrm{ARPE}}({\bm{X}})={\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top}+{\bm{B}},$ (1) $$

where the bias matrix ${\bm{B}}\in\mathbb{R}^{n\times n}$ is induced by the position encoding function $b:{\mathbb{N}}^{2}\to\mathbb{R}$ and the $(i,j)$-th entry of ${\bm{B}}$ is defined as $b(i,j)$. Various formulations and parameterizations of $b$ give rise to different variants of RPE. Examples of additive RPE include:
(1) ALiBi: $b(i,j)=-r|i-j|$, with the scaler $r>0$ as a hyper-parameter;
(2) Kerple: $b(i,j)=-r_{1}log(1+r_{2}|i-j|)$ with $r_{1}$ and $r_{2}$ are two learnable parameters;
(3) FIRE: $b(i,j)=f_{\theta}\left(\frac{\psi(i-j)}{\psi(\max\{L,i\})}\right)$, where the positional encoding function $f_{\theta}$ parameterized by $\theta$ is learned from data and $\psi$ is a transformation function aimed at assigning more model capacity to local positions.

#### Data-Adaptive Position Encoding (DAPE)

The DAPE rewrite the Equation [1](#S3.E1) as the following:

$$ ${\bm{A}}_{\mathrm{DAPE}}({\bm{X}})={\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top}+f({\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top},{\bm{B}}).$ (2) $$

Here, $f:\mathbb{R}^{T\times T}\times\mathbb{R}^{T\times T}\to\mathbb{R}^{T\times T}$ is an element-wise function and $T$ is the sequence length. Another variant of DAPE is with residual, which is the following:

$$ ${\bm{A}}_{\mathrm{DAPE}}({\bm{X}})={\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top}+{\bm{B}}+f({\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top},{\bm{B}}).$ (3) $$

In practice, DAPE (Zheng et al., ) utilizes a two-layer LeakyReLU MLP with hidden dimension $D_{\mathrm{DAPE}}$ (default value is 32) to parameterize $f(\cdot)$ due to its universal approximability (Leshno et al., ).
All parameters are learned directly from the data during the training process. This architecture allows $f(\cdot)$ to dynamically adjust positional embeddings based on the input sequence data, ensuring that the encoding method is both adaptive and dependent on the input data.

### 3.2 Special Case of DAPE: Bias is Zero

Figure: Figure 1: The result of DAPE (Zheng et al., ) (equivalent to kernel $1\times 1$ in our explanation) and $\textrm{DAPE}_{1\times 3}$ (kernel $1\times 3$ by this work), with baseline NoPE and RoPE. The model is trained with length 128 and length 512 respectively. The $\textrm{DAPE}_{1\times 3}$ denotes that we use $H\times 1\times 3$ convolutions kernel size on the attention score with shape $[B,H,T,T]$. We find that DAPE can even improve the performance of NoPE (without biased position encoding), suggesting that the explanation in Zheng et al. (), which attributes the improvement to adaptive position encoding, may have a more general underlying cause.
Refer to caption: /html/2410.04798/assets/x1.png

DAPE was originally designed to dynamically adjust the positional encoding by incorporating input data information. Generally, any additive positional encoding method that includes positional information can be represented as the matrix $\bm{B}$ in the DAPE model, as outlined in Equation [2](#S3.E2). Notably, No Positional Encoding (NoPE) (Kazemnejad et al., ) is a special case of additive RPE that assigns zero value to the matrix $\bm{B}$. The mathematical formulation of DAPE equipped with NoPE is given by:

$$ ${\bm{A}}_{\mathrm{DAPE}}({\bm{X}})={\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top}+f({\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top}).$ (4) $$

#### The result of DAPE-NoPE (the Zheng et al. ( 2024 ) only combine DAPE with ALiBi, Kerple and FIRE but not with NoPE or RoPE).

Compared with the standard Transformer architecture, DAPE-NoPE introduces additional MLPs post the key-query multiplication and prior to the softmax operator.
As shown in Figure [1](#S3.F1), experimental evidence suggests that DAPE with NoPE significantly outperforms the basic NoPE, prompting a reconsideration of the behaviors of standard Transformers.
The additional MLPs (i.e., denoted as $f(\cdot)$ in Equation [4](#S3.E4)) facilitate information sharing across attention heads and complicate the attention calculation with nonlinear transformation beyond the simple key-query multiplication.
This leads to a critical question: Is the current Transformer architecture, particularly the attention mechanism, sufficiently expressive for real-world language tasks?
Although numerous studies aim to enhance efficiency by reducing computation and storage in standard Transformers, these often come at the cost of effectiveness, potentially hindering the evolution of next-generation Transformer models.
Motivated by these insights and observations, we enhance the Transformer’s expressiveness and behavior by regarding attention as a feature map and applying convolutional operations, akin to those used in computer vision.

#### The result of DAPE-RoPE.

Building on the hypothesis that DAPE enhances Transformer performance by processing pre-softmax scores with MLPs, we explore its applicability to non-additive positional encoding methods, specifically RoPE (Su et al., ). In the DAPE-RoPE configuration, RoPE first computes the classic attention scores of key-query multiplication with RoPE, which are then refined using the MLPs described in Equation [4](#S3.E4). The visualized results of the validation perplexity for DAPE-RoPE and other positional encoding methods are presented in Figure [1](#S3.F1). The results indicate that DAPE-RoPE may degrade the performance, while $\textrm{DAPE}_{1\times 3}$-RoPE (with kernel size $1\times 3$, propsoed by this work) not only improves overall performance but also excels in length extrapolation tasks, particularly at larger sequence lengths. This finding substantiates the effectiveness of $\textrm{DAPE}_{1\times 3}$-RoPE, confirming its superior performance compared to standard RoPE, attributing to the additionally introduced convolution operations to the attention scores.

### 3.3 DAPE V2: Process Attention Scores as Feature Maps

As discussed above, improving Transformer performance necessitates refining the processing of attention score computation beyond the conventional key-query multiplication.
We propose regarding the pre-softmax attention scores as feature maps (4-dimensional tensors) and applying convolutional operators.
This approach facilitates enhanced communication across neighboring tokens and heads, drawing parallels to popular techniques used in computer vision.
This novel method aims to leverage the spatial relationships within tokens, potentially unlocking new aspects of model capabilities.

#### Rethink the DAPE formulation.

In DAPE (Zheng et al., ), MLPs are utilized to process and integrate attention and biases. Notably, these MLP operations can be equated to convolution operations with $1\times 1$ kernel  (Krizhevsky et al., ; Simonyan & Zisserman, ; He et al., ), a stride of one, and no padding. Consequently, we can reformulate the DAPE in Equation [3](#S3.E3) as the following:

$$ ${\bm{A}}_{\mathrm{DAPE}}({\bm{X}})={\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top}+{\bm{B}}+Conv(tril(({\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top},{\bm{B}})).$ (5) $$

Under such formulation, DAPE employs convolution operation to process the pre-softmax attention scores of key-query multiplication. The $\texttt{tril}(\cdot)$ returns the lower triangular part of the matrix and the other elements of the result tensor out are set to 0. The resulting attention tensor has a shape of $[B,H,T,T]$, where the four dimensions correspond to the batch size, number of heads, and the context length for both the query and key. This mirrors the structure of an image feature tensor with shape $[B,C,H,W]$, where the dimensions represent the batch size, number of channels, image height, and image width, respectively. This structural similarity underscores the feasibility of considering attention scores as a tensor of feature mappings, where popular and effective convolution operations can be leveraged for refined processing.

#### Process attention with more powerful convolution operation.

In computer vision, the limitations of $1\times 1$ kernels for processing image features are well-recognized. To improve upon the attention scores processed by these kernels (e.g., DAPE), we introduce $1\times k$ kernels with a stride of 1 and padding of $k-1$. This approach allows for wider and deeper convolution across key dimensions and heads without information leakage, as we ensure the attention scores remain lower-triangular. This mechanism is visualized in Appendix [I](#A9).
The use of $1\times k$ kernels suggests a targeted convolution along the key dimensions across heads. In general, while extending this to include the query dimensions as a standard kernel is theoretically possible, it would significantly increase computational demands. Our forthcoming analysis demonstrates that Transformers modified with $1\times k$ convolution are adept at associative recall tasks (i.e., the copy task), validating the benefits of integrating convolution in attention calculation. We left as a future work investigating the performances and the soundness of general convolution kernels, such as square sizes. The key contribution of this work is providing a novel insight that suggests applying convolution operations and processing attention as feature maps to improve Transformers’ performances.

#### Realizing associate recall tasks through convolution.

As pointed out in some previous works (Arora et al., ), the perplexity scores of Transformers mostly result from the performances on associate recall tasks (i.e., the copy tasks). Numerous studies have explored the mechanism of associative recall within Transformers, both from theoretical perspectives and experimental validations (Arora et al., ; Bietti et al., ; Golovneva et al., ). Here, we theoretically prove that the proposed model can realize the associative recall tasks. Notably, this capability is achieved independently of positional encodings, marking a significant advancement in the flexibility and applicability of the proposed architecture. By integrating convolutional operations, we enable the model to handle associative tasks more effectively, leveraging spatial relationships inherent in the data, similar to methods used in image processing.
To explain the associative recall mechanism, (Bietti et al., ) proved that the first layer of the Transformer is responsible for the previous token mechanism through the positional encoding. More specifically, given a sequence of input tokens $\bm{X}=[\bm{x}_{1},\bm{x}_{2},\cdots,\bm{x}_{N}]$ with corresponding orthogonal positional encoding vectors $[\bm{p}_{1},\bm{p}_{2},\cdots,\bm{p}_{N}]$, the first layer primarily facilitates the copying of the previous token to the current token (e.g., $\bm{x}_{i}+\bm{W}_{V}^{1}\bm{x}_{i-1}$, where $\bm{W}_{V}^{1}$ is the value matrix at the first layer of the Transformer). The input tokens are combined with positional encodings $\bm{x}_{i}+\bm{p}_{i}$ and the key-query weight matrix is defined as $\bm{W}_{K}^{1\top}\bm{W}_{Q}^{1}=\sum_{i=1}^{N}\bm{p}_{i-1}\bm{p}_{i}^{\top}$. The orthogonality of positional encoding vectors and the special choices of the key-query matrix ensure that attention scores predominantly focus on the previous token.
In contrast to this implicit mechanism in standard Transformers, our proposed method leverages a convolution operation to explicitly realize associative recall. This approach not only simplifies the process but also enhances its effectiveness by directly manipulating the spatial relationships within tokens and attention scores.
Consider a scenario where the word “Hakuna” is consistently followed by “Matata” within a lengthy paragraph. Without the loss of generality, we assume that $\bm{x}_{1}$ and $\bm{x}_{2}$ represent the tokens of “Hakuna” and “Matata” respectively, and $\bm{x}_{N}=\bm{x}_{1}$ implies that the N-th token in the sequence is “Hakuna”. Then we expect that the Transformer can predict and output the next token $\bm{x}_{N+1}$ as “Matata”. For simplicity, we consider a one-head Transformer without positional encoding. We employ a convolution operation with a kernel size of $1\times 2$ and weights $[-1,1]$.
Note that the convolution is linear and processing the attention scores along the key dimensions is effectively equivalent to applying convolutions directly to the key vectors themselves. Consequently, the key vector of $\bm{x}_{2}$ can be expressed as $\bm{W}_{K}^{1}\left(\bm{x}_{2}-\bm{x}_{1}\right)$ and the query vector for $\bm{x}_{N}$ admits $\bm{W}_{Q}^{1}\bm{x}_{N}$. By configuring the matrix $\bm{W}_{K}^{1\top}\bm{W}_{Q}^{1}$ to be $-\bm{I}$, the attention mechanism after the convolution predominantly allocates the attention values of $\bm{x}_{N}$ to the token $\bm{x}_{2}$. This ensures that the token values of $\bm{x}_{2}$ are effectively copied to $\bm{x}_{N}$, resulting in the model outputting “Matata” following “Hakuna”.

###### Proposition 1 .

Transformers incorporating convolution operations can perform associative recall tasks without the need for positional encoding.

###### Proposition 1 .

#### Comparisons with hybrid models of convolution and Transformers.

Recent developments in hybrid architectures have seen the integration of convolutional and Transformer models to capitalize on the strengths of both. For instance,
Fu et al. () introduced the FlashConv layer, which combines the efficiency of State Space Models (SSMs) with the capabilities of attention-based models. Similarly, Arora et al. () developed a gated convolution layer, noted for its effectiveness in addressing associative recall tasks. These models typically stack convolution layers directly with standard Transformer layers, resulting in modifications to the token values through convolution. In contrast, our model adopts a distinctive approach by applying convolution along the key dimension during the computation of attention scores. This method preserves the original token values while still leveraging the convolution’s benefits for processing attention.

## 4 Experiment

#### Baselines.

We evaluate the proposed $\textrm{DAPE}_{1\times 3}$ against several well-established baselines, including NoPE (Kazemnejad et al., ), RoPE (Su et al., ), T5’s Bias (Raffel et al., ), ALiBi (Press et al., ), Kerple (Chi et al., ), FIRE (Li et al., ), CoPE (Golovneva et al., ), and DAPE (Zheng et al., ). As our kernels are applied across all heads, we simplify by omitting the kernel size description at the head dimension. For example, $\textrm{DAPE}_{1\times 3}$ indicates the use of a $H\times 1\times 3$ convolution kernel size on the attention scores, with a shape of $[B,H,T,T]$.

#### Datasets.

Our analysis is based on training language models using the Arxiv and Books3 datasets, commonly employed benchmarks for assessing model performance (Press et al., ; Chi et al., ; Li et al., ; Ding et al., ). We begin our evaluation by processing entire sequences and comparing the zero-shot perplexity of the last 256 tokens across various input lengths. In addition to perplexity, we also leverage downstream datasets with randomized positional encoding (Ruoss et al., ) to further assess $\textrm{DAPE}_{1\times 3}$.

#### Experiment settings.

Initially, we compare $\textrm{DAPE}_{1\times 3}$ with other baselines at training lengths of 128, 512, and 1024, using 125M decoder-only Transformers (Brown et al., ), with model configurations detailed in Appendix [G](#A7). Subsequently, we evaluate the performance of different training lengths using the same number of training tokens but with larger model sizes (350M and 2.7B). We also explore the impact of the convolutional hidden dimension $D_{\text{DAPE}}$, the effect of information leakage, and the influence of varying kernel sizes. Additionally, we examine the computational efficiency of $\textrm{DAPE}_{1\times 3}$, focusing on processing times. Lastly, we evaluate $\textrm{DAPE}_{1\times 3}$ on algorithmic reasoning datasets using accuracy metrics. Compared to DAPE (Zheng et al., ), $\textrm{DAPE}_{1\times 3}$ demonstrates a more pronounced attention sink (Xiao et al., ), as visualized in Appendix [I](#A9).

### 4.1 Compare with Baselines

Figure: Figure 2: Comparisons with baselines: performance with training lengths 128 and 512 on Arxiv and Books3 datasets.
Refer to caption: /html/2410.04798/assets/x3.png

#### DAPE 1 × 3 subscript DAPE 1 3 \textrm{DAPE}_{1\times 3} -Kerple improves performance within training length, proving its ability to process the entire sequence.

According to Figure [2](#S4.F2), the proposed $\textrm{DAPE}_{1\times 3}$-Kerple demonstrates superior performance across various training and evaluation lengths. Specifically, $\textrm{DAPE}_{1\times 3}$-Kerple achieves the best performance where the training length is 128 or 512 and the evaluation length ranges from 128 to 8192. This performance consistency is observed across both the arXiv and Books datasets.
For instance, on the arXiv dataset with a training length of 512, $\textrm{DAPE}_{1\times 3}$-Kerple achieves a perplexity score of 4.44. This score surpasses those of other methods, such as DAPE-Kerple with a perplexity of 4.49, CoPE with 4.51, Kerple with 4.57, and RoPE with 4.57. These results indicate that $\textrm{DAPE}_{1\times 3}$-Kerple has a more robust modeling capability within the training length compared to the other methods evaluated. The Appendix [A](#A1) also presents the performance of different methods with training length 1024. The improvements are not only significant but also consistent, reinforcing the efficacy of the $\textrm{DAPE}_{1\times 3}$-Kerple approach in handling various training lengths effectively.

#### DAPE 1 × 3 subscript DAPE 1 3 \textrm{DAPE}_{1\times 3} -Kerple improves performance beyond training length.

The advantages of $\textrm{DAPE}_{1\times 3}$-Kerple extend beyond the training length. When the training length is set to 128 and the evaluation length is extended to 8192, $\textrm{DAPE}_{1\times 3}$-Kerple achieves a perplexity score of 4.60 on the arXiv dataset and 23.52 on the Books3 dataset. These scores are significantly better than those achieved by DAPE-Kerple, which records perplexity scores of 4.97 and 25.01 on the arXiv and Books3 datasets, respectively. Similarly, CoPE performs poorly with perplexity scores of 29.86 on the arXiv dataset and 90.66 on the Books3 dataset under the same conditions.
Furthermore, when the training duration is increased to 512, $\textrm{DAPE}_{1\times 3}$-Kerple continues to deliver the best performance, further validating its superior generalization capabilities. These findings highlight the scalability and robustness of $\textrm{DAPE}_{1\times 3}$-Kerple, which is attributed to the introduced convolution operator, making it a promising approach for diverse data scenarios and lengths.

### 4.2 Performance with Same Training tokens and Different Training Length

Figure: Figure 3: The performance with same training tokens and different training length. With the same training tokens, $\textrm{DAPE}_{1\times 3}$ with training length 512 could even achieve better performance than RoPE with training length 4096.
Refer to caption: /html/2410.04798/assets/x7.png

#### Compared to RoPE, with the same training tokens, DAPE 1 × 3 subscript DAPE 1 3 \textrm{DAPE}_{1\times 3} -Kerple with a training length of 128 achieves performance comparable to RoPE with a training length of 4096, for varying evaluation length.

As shown in Figure [3](#S4.F3), for $\textrm{DAPE}_{1\times 3}$-Kerple trained with a length of 128, it achieves a perplexity (ppl) of 8.15 at an evaluation length of 128 and 4.95 at an evaluation length of 4096 on the arXiv dataset. In comparison, RoPE trained with a length of 4096 achieves a ppl of 9.59 at an evaluation length of 128 and 4.92 at an evaluation length of 4096. Similarly, on the Books3 dataset, $\textrm{DAPE}_{1\times 3}$-Kerple trained with a length of 128 achieves a ppl of 31.07 at an evaluation length of 128 and 23.19 at an evaluation length of 4096, while RoPE trained with a length of 4096 achieves 38.36 and 24.58, respectively. This suggests the superiority of the proposed $\textrm{DAPE}_{1\times 3}$ with the introduced convolution operators among heads and neighboring tokens.

#### With the same training tokens, compared to DAPE 1 × 3 subscript DAPE 1 3 \textrm{DAPE}_{1\times 3} with longer training lengths, DAPE 1 × 3 subscript DAPE 1 3 \textrm{DAPE}_{1\times 3} with shorter training lengths can achieve comparable performance, indicating that DAPE 1 × 3 subscript DAPE 1 3 \textrm{DAPE}_{1\times 3} enhances the model’s understanding of text structure.

On the arXiv dataset, $\textrm{DAPE}_{1\times 3}$-Kerple with training lengths of 512 demonstrates performance close to that of training with a length of 4096 when the evaluation length is 4096. Moreover, the performance curves for training lengths of 1024, and 2048 are almost identical. This trend is also observed with the Books3 dataset. These results indicate that $\textrm{DAPE}_{1\times 3}$-Kerple effectively helps the model comprehend text structure, enabling it to extend to longer lengths.

#### Transformers may overfit their training length: training on longer sequences may decrease performance when testing on shorter sequences.

On the arXiv dataset, $\textrm{DAPE}_{1\times 3}$-Kerple with a training length of 128 achieves the best performance when the evaluation length is 128. Similarly, $\textrm{DAPE}_{1\times 3}$-Kerple with training lengths of 256, 512, 1024, and 2048 achieves the best performance at evaluation lengths of 256, 512, 1024, and 2048, respectively. Also, on evaluation 128, the RoPE with training length 4096 and batch size 1 also achieves worse performance than the RoPE with training length 128 and batch size 32. This suggests that training on longer sequences may worsen a Transformer’s performance at shorter sequence lengths.

#### DAPE 1 × 3 subscript DAPE 1 3 \textrm{DAPE}_{1\times 3} can reduce the training time cost via larger batch size and shorter training length, achieving comparable performance compared to trained on longer length.

The cost of $\textrm{DAPE}_{1\times 3}$ is $\mathcal{O}(B\cdot(h\cdot d\cdot T^{2}+h\cdot D_{\text{DAPE}}\cdot T^{2}))$, where $B$, $h$, $d$, $T$ and $D_{\text{DAPE}}$ are the batch size, attention hidden dimension, attention head number, sequence length and DAPE hidden dimension. By reducing the training length from $T$ to $\frac{T}{K}$ and increasing the batch size from $B$ to $B\cdot K$ with the same training tokens, the cost becomes $\mathcal{O}(B\cdot K\cdot(h\cdot d\cdot(\frac{T}{K})^{2}+h\cdot D_{\text{DAPE}}\cdot(\frac{T}{K})^{2}))$, which simplifies to $\mathcal{O}(\frac{B\cdot(h\cdot d\cdot T^{2}+h\cdot D_{\text{DAPE}}\cdot T^{2})}{K})$. For example, when the training length is 128 and the batch size is 32, the time cost of one step is 40.30ms. The time cost of length 256 (batch 16), length 512 (batch 8), length 1024 (batch 4), and length 2048 (batch 2) are 42.61ms, 50.38ms, 79.36ms, and 120.14ms. This reduction demonstrates the potential for significant training time savings.

### 4.3 The Effect of Larger Model Size

Figure: Figure 4: The Effect of Larger Model Size 350M. We show the results with training length 128 and training length 512 on Arxiv dataset.
Refer to caption: /html/2410.04798/assets/x9.png

#### DAPE 1 × 3 subscript DAPE 1 3 \textrm{DAPE}_{1\times 3} performs well with larger model sizes, such as 350M and 2.7B.

As illustrated in Figure [4](#S4.F4), the proposed $\textrm{DAPE}_{1\times 3}$ shows superior performance at varying evaluation lengths with a model size of 350M. For a training length of 128, $\textrm{DAPE}_{1\times 3}$-Kerple achieves a perplexity (ppl) of 7.63 at an evaluation length of 128 and 4.43 at an evaluation length of 8192, compared to DAPE’s 7.69 and 4.69, respectively. Similarly, for a training length of 512, $\textrm{DAPE}_{1\times 3}$-Kerple achieves a ppl of 4.10 at an evaluation length of 128 and 3.35 at an evaluation length of 8192, whereas DAPE achieves 4.14 and 3.44, respectively. We also present the 2.7B model size result in Appendix [B](#A2). Therefore, the proposed $\textrm{DAPE}_{1\times 3}$ demonstrates excellent performance with larger model sizes, showing the potential of including the proposed processing techniques in existing large language models.

### 4.4 The Effect of DAPE 1 × 3 subscript DAPE 1 3 \textrm{DAPE}_{1\times 3}

Figure: Figure 5: The effect of $\textrm{DAPE}_{1\times 3}$. Whatever the baseline is ALiBi, Kerple or FIRE, the proposed $\textrm{DAPE}_{1\times 3}$ can all improve their performance. The Figure [1](#S3.F1) also proves that the proposed $\textrm{DAPE}_{1\times 3}$ is effective for NoPE and RoPE.
Refer to caption: /html/2410.04798/assets/x11.png

#### For Additive Positional Encoding, DAPE 1 × 3 subscript DAPE 1 3 \textrm{DAPE}_{1\times 3} enhances performance within and beyond the training length.

As demonstrated in Figure [5](#S4.F5), for varying additive positional encoding such as ALiBi, Kerple, and FIRE, their incorporations with $\textrm{DAPE}_{1\times 3}$ (i.e., $\textrm{DAPE}_{1\times 3}$-ALiBi, $\textrm{DAPE}_{1\times 3}$-Kerple, and $\textrm{DAPE}_{1\times 3}$-FIRE) consistently improve performance. Furthermore, regardless of the specific additive positional encoding used, the proposed $\textrm{DAPE}_{1\times 3}$ (configured with a kernel size of $1\times 3$) outperforms the standard DAPE method (which employs a kernel size of $1\times 1$). Also, as shown in Figure [1](#S3.F1), $\textrm{DAPE}_{1\times 3}$ imrpoves the performance of NoPE, both within and beyond the training length
These results highlight the robustness and scalability of $\textrm{DAPE}_{1\times 3}$, suggesting its broad applicability in enhancing additive positional encoding frameworks.

#### For Non-Additive Positional Encoding, DAPE 1 × 3 subscript DAPE 1 3 \textrm{DAPE}_{1\times 3} also improves performance within and beyond the training length.

As illustrated in Figure [1](#S3.F1), $\textrm{DAPE}_{1\times 3}$ enhances the performance of RoPE, both within and beyond the training length. In contrast, naive DAPE reduces the performance of RoPE, with training lengths of 128 and 512. This indicates that the proposed $\textrm{DAPE}_{1\times 3}$ is a versatile and widely applicable method with the potential to be applied to various position encoding techniques on the language modeling task.

### 4.5 The Performance of DAPE 1 × 3 subscript DAPE 1 3 \textrm{DAPE}_{1\times 3} with Information Leakage

Figure: Figure 6: Result with information leakage.
Refer to caption: /html/2410.04798/assets/x13.png

#### The DAPE 1 × 3 subscript DAPE 1 3 \textrm{DAPE}_{1\times 3} can utilize attention data, which is supported by almost zero loss (perplexity is 1) under information leakage.

To prevent the information leakage, we use the $torch.tril$ before $\textrm{DAPE}_{1\times 3}$ to make the attention score lower-triangular matrix. For the cheating version, we do not use the torch.tril. As shown in Figure [6](#S4.F6), whatever $\textrm{DAPE}_{1\times 3}$-ALiBi, $\textrm{DAPE}_{1\times 3}$-Kerple or $\textrm{DAPE}_{1\times 3}$-FIRE, their cheating version can all achieve about zero loss within evaluation length 1024. Furthermore, the $\textrm{DAPE}_{1\times 3}$-Kerple can even aachievezero loss when the evaluation length is extended to 8096. This suggest that the proposed $\textrm{DAPE}_{1\times 3}$ can really realize and utilize the information of attention score.

### 4.6 Compare DAPE and DAPE 1 × 3 subscript DAPE 1 3 \textrm{DAPE}_{1\times 3} with Approximate Computational Cost

#### DAPE 1 × 3 subscript DAPE 1 3 \textrm{DAPE}_{1\times 3} achieves even better performance at a lower computational cost.

As shown in Appendix [C](#A3), when the training length is set to 128, $\textrm{DAPE}_{1\times 3}$-Kerple with $D_{\text{DAPE}}$ as 10 achieves a perplexity (ppl) of 8.16 at an evaluation length of 128 and 4.74 at an evaluation length of 8192. This performance is notably better than that of DAPE-Kerple with $D_{\text{DAPE}}$ as 64, which achieves perplexities of 8.21 and 4.87, respectively. Moreover, when the training length is extended to 512 and the evaluation length is smaller or equal to 4096, $\textrm{DAPE}_{1\times 3}$-Kerple with $D_{\text{DAPE}}$ as 10 continues to surpass the performance of DAPE-Kerple with $D_{\text{DAPE}}$ as 64. Also, $\textrm{DAPE}_{1\times 3}$-Kerple with $D_{\text{DAPE}}$ as 21 always achieves better performance than DAPE-Kerple with $D_{\text{DAPE}}$ as 64. This demonstrates that $\textrm{DAPE}_{1\times 3}$ not only maintains its performance advantage across different training lengths but also requires a lower computational cost.

### 4.7 The Performance with Different Kernel Sizes

#### Different experiment settings may have different optimal kernel sizes.

Appendix [D](#A4) shows the performance of DAPE with various kernel sizes, including DAPE (equivalent to a $1\times 1$ kernel size), $\textrm{\text{DAPE}}_{1\times 3}$, $\textrm{DAPE}_{1\times 5}$, and $\textrm{DAPE}_{1\times 7}$. For the Arxiv dataset, larger kernel sizes consistently achieve better performance, evaluating with training lengths of 128 or 512. However, for the Books3 dataset, $\textrm{DAPE}_{1\times 3}$ performs best when the training length is 128 and evaluated at 8192, whereas $\textrm{DAPE}_{1\times 5}$ performs best at the same evaluation level when the training length is 512. These results suggest that the optimal kernel size may vary depending on the experimental setting, ranging from $1\times 1$ to larger kernel sizes. Although larger kernel sizes contribute to stronger expressiveness from intuition, we conjecture that the performance degradation for overly large kernel sizes results from optimization challenges.

### 4.8 The Performance on CHE Benchmark with Accuracy Evaluation Metrics

#### Different tasks have different optimal kernel sizes, as shown in Appendix E and Appendix D .

For example, on Missing Duplicate task, the $\textrm{DAPE}_{1\times 3}$-Kerple improves the 87.57 of DAPE-Kerple to 99.65. However, on the Stack Manipulationtask, the $\textrm{DAPE}_{1\times 3}$-Kerple decreases the 72.04 of DAPE-Kerple to 68.18. Also, as shown in Appendix [D](#A4), the larger kernel size does not always lead to better performance. Overall, larger kernel size provides a potential way to improve the Transformer length extrapolation performance, and we usually could find a suitable kernel size (ranging from 1×1 to larger
kernel sizes) to achieve better performance than without further processing attention score.

#### The large kernel size performance improvement is related to the baseline bias matrix.

As shown in Appendix [E](#A5), the best performance is usually achieved by further processing attention scores via kernel size 1 or 3. Moreover, on 11 permutation-variant tasks, the $\textrm{DAPE}_{1\times 3}$-Kerple achieves better performance on 8 of 11 tasks compared to Kerple. And the $\textrm{DAPE}_{1\times 3}$-FIRE achieves better performance on 6 of 11 tasks compared to FIRE. This suggests that the large kernel size performance improvement is related to the baseline bias matrix.

### 4.9 The Time Cost

#### As the model size increases, the additional computational cost ratio gradually decreases.

As shown in Appendix [F](#A6), when the model size is 350M, the time cost for Kerple is 189.91 ms, while DAPE-Kerple takes 224.22 ms, and $\textrm{DAPE}_{1\times 3}$-Kerple requires 252.84 ms. Compared to $\textrm{DAPE}_{1\times 3}$-Kerple, the time cost ratios for Kerple and DAPE-Kerple are 0.7511 and 0.8868, respectively. As the model size increases from 350M to 2.7B and 6.7B, the time cost ratio for Kerple rises from 0.7511 to 0.8205 and 0.8918, respectively. Similarly, the time cost ratio for DAPE-Kerple increases from 0.8868 to 0.9361 and 0.9677. Therefore, as the model size increases, the time cost ratio also increases, indicating that the additional computational cost decreases progressively.

## 5 Conclusion

In this paper, we point out that the key of Transformer length extrapolation is the better and more accurate attention score. Therefore, we develop and analyze $\textrm{DAPE}_{1\times 3}$ by processing the attention score as feature maps via convolution operation.
Theoretically, we show that the associative recall tasks, which account for the most perplexity scores, can be realized by the proposed Transformer with convolution, in contrast to the vanilla Transformer.
We conducted comprehensive experiments on Arxiv, Books3, and CHE to validate the effectiveness of the proposed method,
where the proposed method exhibits significant superiority.

## Appendix A Compare with Baseline on Arxiv Dataset with Training Length 1024

**Table 1: The performance (ppl) on Arxiv dataset with training length 1024, compared to baselines.**
| Method | 1024 | 2048 | 4096 | 8192 |
| --- | --- | --- | --- | --- |
| NoPE (Kazemnejad et al., ) | 4.16 | 42.27 | 1854.73 | 17167.32 |
| RoPE (Su et al., ) | 4.07 | 86.20 | 237.67 | 256.12 |
| T5’s bias (Raffel et al., ) | 4.03 | 4.28 | 13.07 | 79.55 |
| ALiBi (Press et al., ) | 4.09 | 4.53 | 4.45 | 4.22 |
| Kerple (Chi et al., ) | 4.06 | 4.09 | 4.68 | 6.951 |
| FIRE (Li et al., ) | 4.06 | 9.21 | 236.18 | 440.60 |
| DAPE-Kerple (Zheng et al., ) | 3.98 | 3.91 | 3.68 | 3.41 |
| $\textrm{DAPE}_{1\times 3}$-Kerple | 3.93 | 3.86 | 3.61 | 3.37 |

## Appendix B Large Model Size

**Table 2: The performance (ppl) under large model size 2.7B on Books3 dataset.**
| Method | 512 | 1024 | 2048 | 4096 |
| --- | --- | --- | --- | --- |
| RoPE | 21.01 | 25.00 | 48.13 | 160.59 |
| T5’s bias | 21.10 | 21.88 | 23.59 | 33.23 |
| Kerple | 21.14 | 22.08 | 23.38 | 27.21 |
| DAPE-Kerple | 20.52 | 21.01 | 20.23 | 19.67 |
| $\textrm{DAPE}_{1\times 3}$-Kerple (kernel size 1x3) | 20.16 | 20.54 | 19.80 | 19.02 |

## Appendix C Compare DAPE and DAPE 1 × 3 subscript DAPE 1 3 \textrm{DAPE}_{1\times 3} with Approximate Computational Cost

Figure: Figure 7: Compare $\textrm{DAPE}_{1\times 3}$ and DAPE with the approximately same cost on Arxiv Dataset. We compare the $\textrm{DAPE}_{1\times 3}$ and DAPE with approximate cost and different $D_{\text{DAPE}}$. As the kernel size of is $\textrm{DAPE}_{1\times 3}$ $1\times 3$, the proposed $\textrm{DAPE}_{1\times 3}$ is the triple computation cost of DAPE, with the same $D_{\text{DAPE}}$.
Refer to caption: /html/2410.04798/assets/x14.png

## Appendix D The Performance with Different Kernel Size

**Table 3: The performance with different kernel sizes, with training length 128 and evaluation from length 128 to 8192. For different datasets and training length, the optimal kernel size may not always be the largest one, especially when the evaluation length is larger.**
| Dataset | Method | 128 | 256 | 512 | 1024 | 2048 | 4096 | 8192 |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Arxiv | Kerple | 8.30 | 7.10 | 5.85 | 6.91 | 9.17 | 11.48 | 12.59 |
| DAPE-Kerple (Kernel Size 1x1) | 8.21 | 6.98 | 5.38 | 5.20 | 5.33 | 5.26 | 4.97 |  |
| $\textrm{DAPE}_{1\times 3}$-Kerple (Kernel Size 1x3) | 8.15 | 6.92 | 5.29 | 5.05 | 5.11 | 4.95 | 4.60 |  |
| $\textrm{DAPE}_{1\times 5}$-Kerple (Kernel Size 1x5) | 8.13 | 6.91 | 5.27 | 5.04 | 5.10 | 4.91 | 4.57 |  |
| $\textrm{DAPE}_{1\times 7}$-Kerple (Kernel Size 1x7) | 8.12 | 6.89 | 5.26 | 5.02 | 5.09 | 4.91 | 4.57 |  |
| Books3 | Kerple | 32.10 | 29.09 | 28.10 | 35.75 | 44.68 | 56.39 | 66.23 |
| DAPE-Kerple (Kernel Size 1x1) | 31.49 | 28.27 | 24.93 | 24.31 | 23.34 | 24.38 | 25.01 |  |
| $\textrm{DAPE}_{1\times 3}$-Kerple (Kernel Size 1x3) | 31.07 | 27.81 | 24.38 | 23.57 | 22.40 | 23.19 | 23.52 |  |
| $\textrm{DAPE}_{1\times 5}$-Kerple (Kernel Size 1x5) | 31.02 | 27.79 | 24.36 | 23.57 | 22.41 | 23.32 | 23.71 |  |
| $\textrm{DAPE}_{1\times 7}$-Kerple (Kernel Size 1x7) | 30.98 | 27.76 | 24.31 | 23.47 | 22.30 | 23.00 | 23.57 |  |

**Table 4: The performance with different kernel size, with training length 512 and evaluation from length 512 to 8192. For different datasets and training length, the optimal kernel size may not always be the largest one, especially when the evaluation length is larger.**
| Dataset | Method | 512 | 1024 | 2048 | 4096 | 8192 |
| --- | --- | --- | --- | --- | --- | --- |
| Arxiv | Kerple | 4.57 | 4.37 | 5.09 | 6.80 | 9.08 |
| DAPE-Kerple (Kernel Size 1x1) | 4.49 | 4.20 | 4.17 | 3.95 | 3.70 |  |
| $\textrm{DAPE}_{1\times 3}$-Kerple (Kernel Size 1x3) | 4.44 | 4.14 | 4.09 | 3.87 | 3.58 |  |
| $\textrm{DAPE}_{1\times 5}$-Kerple (Kernel Size 1x5) | 4.44 | 4.14 | 4.10 | 3.85 | 3.59 |  |
| $\textrm{DAPE}_{1\times 7}$-Kerple (Kernel Size 1x7) | 4.43 | 4.13 | 4.08 | 3.85 | 3.57 |  |
| Books3 | Kerple | 19.83 | 19.19 | 20.48 | 28.33 | 40.94 |
| DAPE-Kerple (Kernel Size 1x1) | 19.25 | 18.28 | 17.20 | 17.58 | 17.85 |  |
| $\textrm{DAPE}_{1\times 3}$-Kerple (Kernel Size 1x3) | 18.95 | 17.92 | 16.79 | 17.05 | 17.20 |  |
| $\textrm{DAPE}_{1\times 5}$-Kerple (Kernel Size 1x5) | 18.89 | 17.87 | 16.76 | 17.09 | 17.10 |  |
| $\textrm{DAPE}_{1\times 7}$-Kerple (Kernel Size 1x7) | 18.86 | 17.82 | 16.70 | 17.01 | 17.16 |  |

## Appendix E The Performance of DAPE 1 × 3 subscript DAPE 1 3 \textrm{DAPE}_{1\times 3} on CHE Benchmark

**Table 5: Train on length 40 with 200k steps, and test from lengths 41 to 500. The random accuracy is 50%, except for Modular Arithmetic (Simple), Cycle Navigation, Bucket Sort, Solve Equation and Modular Arithmetic, where it is 20%. $\dagger$$\dagger$$\dagger$ denotes permutation-invariant tasks, which are expected to be solved without positional information. The dataset comes from Choromanski et al. (), with experiment setting from Randomized PE(Ruoss et al., ).**
|  |  | Baseline |  | DAPE (Kernel Size 1) |  | DAPE (Kernel Size 3) |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Level | Task | RoPE | Relative | ALiBi | Kerple | FIRE |  | ALiBi | Kerple | FIRE |  | ALiBi | Kerple | FIRE |
| R | Even Pairs | 99.98 | 96.60 | 73.52 | 57.50 | 73.86 |  | 99.99 | 99.58 | 100 |  | 99.99 | 100 | 100 |
| Modular Arithmetic (Simple) | 21.35 | 20.84 | 20.02 | 21.79 | 21.09 |  | 23.58 | 24.47 | 24.46 |  | 21.48 | 23.90 | 23.43 |  |
| Parity Check$\dagger$$\dagger$$\dagger$ | 50.05 | 50.09 | 50.09 | 50.07 | 50.97 |  | 50.30 | 50.07 | 50.04 |  | 50.13 | 52.51 | 50.11 |  |
| Cycle Navigation$\dagger$$\dagger$$\dagger$ | 27.63 | 26.95 | 24.64 | 29.47 | 28.41 |  | 22.99 | 34.53 | 27.54 |  | 24.43 | 24.32 | 24.34 |  |
| DCF | Stack Manipulation | 61.49 | 64.73 | 66.42 | 66.93 | 69.33 |  | 68.18 | 72.04 | 70.90 |  | 58.90 | 68.18 | 60.90 |
| Reverse String | 65.23 | 65.59 | 71.09 | 71.54 | 65.89 |  | 73.37 | 70.74 | 76.40 |  | 56.61 | 81.84 | 70.11 |  |
| Modular Arithmetic | 31.25 | 31.74 | 30.56 | 24.79 | 30.92 |  | 31.34 | 32.37 | 31.50 |  | 29.46 | 26.13 | 27.00 |  |
| Solve Equation | 21.85 | 22.93 | 19.92 | 21.15 | 22.06 |  | 20.03 | 22.49 | 22.42 |  | 20.26 | 23.95 | 23.62 |  |
| CS | Duplicate String | 64.97 | 67.66 | 65.13 | 66.72 | 69.03 |  | 70.84 | 72.95 | 72.71 |  | 52.96 | 57.03 | 66.01 |
| Missing Duplicate | 63.37 | 72.34 | 74.21 | 79.06 | 79.27 |  | 83.41 | 87.57 | 89.17 |  | 59.33 | 99.65 | 74.83 |  |
| Odds First | 61.00 | 61.57 | 59.88 | 62.59 | 63.28 |  | 63.78 | 67.08 | 66.34 |  | 57.35 | 56.87 | 56.57 |  |
| Binary Addition | 55.59 | 56.96 | 54.72 | 56.35 | 55.70 |  | 59.71 | 60.88 | 56.62 |  | 57.49 | 55.32 | 57.86 |  |
| Compute Sqrt | 51.88 | 51.63 | 50.63 | 51.11 | 50.80 |  | 51.64 | 51.33 | 52.46 |  | 52.08 | 51.76 | 51.93 |  |
| Bucket Sort$\dagger$$\dagger$$\dagger$ | 98.12 | 99.31 | 98.45 | 99.38 | 99.57 |  | 99.38 | 98.81 | 99.37 |  | 96.61 | 99.06 | 98.56 |  |

## Appendix F DAPE 1 × 3 subscript DAPE 1 3 \textrm{DAPE}_{1\times 3} Time Cost

**Table 6: The time cost (millisecond) under different testing lengths, with $D_{\textbf{\text{DAPE}}}$ as 32 and default batch size 1, with training length 512.**
| Method | 350M Total | Ratio | 2.7B Total | Ratio | 6.7B Total | Ratio |
| --- | --- | --- | --- | --- | --- | --- |
| RoPE (Su et al., ) | 210.01 | 0.8306 | 472.63 | 1.0472 | 635.57 | 0.8564 |
| T5’s bias (Raffel et al., ) | 355.16 | 1.4046 | 537.62 | 1.1912 | 808.85 | 1.0899 |
| ALiBi (Press et al., ) | 172.60 | 0.6826 | 325.95 | 0.7222 | 596.77 | 0.8041 |
| Kerple (Chi et al., ) | 189.91 | 0.7511 | 370.32 | 0.8205 | 661.82 | 0.8918 |
| FIRE (Li et al., ) | 248.13 | 0.9813 | 432.63 | 0.9586 | 797.68 | 1.0748 |
| DAPE-Kerple (Zheng et al., ) | 224.22 | 0.8868 | 422.48 | 0.9361 | 717.46 | 0.9667 |
| $\text{DAPE}_{1\times 3}$-Kerple | 252.84 | 1.0000 | 451.29 | 1.0000 | 742.10 | 1.0000 |

## Appendix G Model Configuration

All experiments are conducted on 8 GPUs. The 125M and 350M model configuration is the following.

**Table 7: Model Configurations.**
|  |  | 125M |  | 350M |
| --- | --- | --- | --- | --- |
| Training sequence length |  | $512$ |  | $512$ |
| Batch size |  | 32 $\times$ 8 |  | 32 $\times$ 8 |
| Numer of iterations |  | $50$k |  | $50$k |
| Dropout prob. |  | $0.0$ |  | $0.0$ |
| Attention dropout prob. |  | $0.0$ |  | $0.0$ |
| Attention head |  | 12 |  | 16 |
| Feature dimension |  | 768 |  | 1024 |
| Layer number |  | 12 |  | 24 |
| Optimizer |  | Adam |  | Adam |
| Optimizer parameter betas |  | [0.9, 0.95] |  | [0.9, 0.95] |
| Learning rate |  | $6\mathrm{e}-4$ |  | $3\mathrm{e}-4$ |
| Precision |  | float16 |  | float16 |

## Appendix H Data-Adaptive Related Position Encoding Performance Comparison

**Table 8: The performance comparison between data-related position encoding, with dataset Books3 and training length 128.**
| Method | 128 | 256 | 512 | 1024 | 2048 | 4096 | 8192 |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Transformer-XL | 31.57 | 28.49 | 26.07 | 26.98 | 27.90 | 32.76 | 41.12 |
| CoPE | 31.61 | 28.41 | 25.79 | 27.96 | 33.80 | 54.08 | 90.66 |
| DAPE-Kerple (Kernel Size 1x1) | 31.49 | 28.27 | 24.93 | 24.31 | 23.34 | 24.38 | 25.01 |
| $\textrm{DAPE}_{1\times 3}$-Kerple (Kernel Size 1x3) | 31.07 | 27.81 | 24.38 | 23.57 | 22.40 | 23.19 | 23.52 |

## Appendix I DAPE 1 × 3 subscript DAPE 1 3 \textrm{DAPE}_{1\times 3} Visualization

The model is trained with $\textrm{DAPE}_{1\times 3}$-Kerple on length 512. Compared to DAPE (Zheng et al., ), it seems that the $\textrm{DAPE}_{1\times 3}$ presents a more obvious attention sink (Xiao et al., ).

### I.1 Visualization on length 512

Figure: Figure 8: Evaluation Length 512 Example 1: Part 1. From Left to Right: (1) The Attention is ${\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top}$; (2) The Kerple bias is ${\bm{B}}$; (3) The $\textrm{DAPE}_{1\times 3}$ (with Kerple) bias is $f({\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top},{\bm{B}})$.
Refer to caption: /html/2410.04798/assets/x16.png

Figure: Figure 9: Evaluation Length 512 Example 1: Part 2. From Left to Right: (1) The Attention is ${\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top}$; (2) The Kerple bias is ${\bm{B}}$; (3) The $\textrm{DAPE}_{1\times 3}$ (with Kerple) bias is $f({\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top},{\bm{B}})$.
Refer to caption: /html/2410.04798/assets/x28.png

Figure: Figure 10: Evaluation Length 512 Example 1: Part 3. From Left to Right: (1) The Attention is ${\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top}$; (2) The Kerple bias is ${\bm{B}}$; (3) The $\textrm{DAPE}_{1\times 3}$ (with Kerple) bias is $f({\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top},{\bm{B}})$.
Refer to caption: /html/2410.04798/assets/x43.png

Figure: Figure 11: Evaluation Length 512 Example 2: Part 1. From Left to Right: (1) The Attention is ${\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top}$; (2) The Kerple bias is ${\bm{B}}$; (3) The $\textrm{DAPE}_{1\times 3}$ (with Kerple) bias is $f({\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top},{\bm{B}})$.
Refer to caption: /html/2410.04798/assets/x52.png

Figure: Figure 12: Evaluation Length 512 Example 2: Part 2. From Left to Right: (1) The Attention is ${\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top}$; (2) The Kerple bias is ${\bm{B}}$; (3) The $\textrm{DAPE}_{1\times 3}$ (with Kerple) bias is $f({\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top},{\bm{B}})$.
Refer to caption: /html/2410.04798/assets/x70.png

### I.2 Visualization on length 2048

Figure: Figure 13: Evaluation Length 2048 Example 1: Part 1. From Left to Right: (1) The Attention is ${\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top}$; (2) The Kerple bias is ${\bm{B}}$; (3) The $\textrm{DAPE}_{1\times 3}$ (with Kerple) bias is $f({\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top},{\bm{B}})$.
Refer to caption: /html/2410.04798/assets/x88.png

Figure: Figure 14: Evaluation Length 2048 Example 1: Part 2. From Left to Right: (1) The Attention is ${\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top}$; (2) The Kerple bias is ${\bm{B}}$; (3) The $\textrm{DAPE}_{1\times 3}$ (with Kerple) bias is $f({\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top},{\bm{B}})$.
Refer to caption: /html/2410.04798/assets/x103.png

Figure: Figure 15: Evaluation Length 2048 Example 1: Part 3. From Left to Right: (1) The Attention is ${\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top}$; (2) The Kerple bias is ${\bm{B}}$; (3) The $\textrm{DAPE}_{1\times 3}$ (with Kerple) bias is $f({\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top},{\bm{B}})$.
Refer to caption: /html/2410.04798/assets/x118.png

Figure: Figure 16: Evaluation Length 2048 Example 2: Part 1. From Left to Right: (1) The Attention is ${\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top}$; (2) The Kerple bias is ${\bm{B}}$; (3) The $\textrm{DAPE}_{1\times 3}$ (with Kerple) bias is $f({\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top},{\bm{B}})$.
Refer to caption: /html/2410.04798/assets/x124.png

Figure: Figure 17: Evaluation Length 2048 Example 2: Part 2. From Left to Right: (1) The Attention is ${\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top}$; (2) The Kerple bias is ${\bm{B}}$; (3) The $\textrm{DAPE}_{1\times 3}$ (with Kerple) bias is $f({\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top},{\bm{B}})$.
Refer to caption: /html/2410.04798/assets/x142.png

### I.3 Visualization on length 8192

Figure: Figure 18: Evaluation Length 8192 Example 1: Part 1. From Left to Right: (1) The Attention is ${\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top}$; (2) The Kerple bias is ${\bm{B}}$; (3) The $\textrm{DAPE}_{1\times 3}$ (with Kerple) bias is $f({\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top},{\bm{B}})$.
Refer to caption: /html/2410.04798/assets/x160.png

Figure: Figure 19: Evaluation Length 8192 Example 1: Part 2. From Left to Right: (1) The Attention is ${\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top}$; (2) The Kerple bias is ${\bm{B}}$; (3) The $\textrm{DAPE}_{1\times 3}$ (with Kerple) bias is $f({\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top},{\bm{B}})$.
Refer to caption: /html/2410.04798/assets/x175.png

Figure: Figure 20: Evaluation Length 8192 Example 1: Part 3. From Left to Right: (1) The Attention is ${\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top}$; (2) The Kerple bias is ${\bm{B}}$; (3) The $\textrm{DAPE}_{1\times 3}$ (with Kerple) bias is $f({\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top},{\bm{B}})$.
Refer to caption: /html/2410.04798/assets/x190.png

Figure: Figure 21: Evaluation Length 8192 Example 2: Part 1. From Left to Right: (1) The Attention is ${\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top}$; (2) The Kerple bias is ${\bm{B}}$; (3) The $\textrm{DAPE}_{1\times 3}$ (with Kerple) bias is $f({\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top},{\bm{B}})$.
Refer to caption: /html/2410.04798/assets/x196.png

Figure: Figure 22: Evaluation Length 8192 Example 2: Part 2. From Left to Right: (1) The Attention is ${\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top}$; (2) The Kerple bias is ${\bm{B}}$; (3) The $\textrm{DAPE}_{1\times 3}$ (with Kerple) bias is $f({\bm{X}}{\bm{W}}_{Q}({\bm{X}}{\bm{W}}_{K})^{\top},{\bm{B}})$.
Refer to caption: /html/2410.04798/assets/x214.png

## Appendix J Implementation

In this section, we present the implementation of the proposed $\textrm{DAPE}_{1\times 3}$ module in PyTorch (Paszke et al., ).

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: **Main concerns, questions & potential improvements:**  
- Most results (in fact, almost all) are reported with for ‘Kerple’, which seems to work well (e.g. Figure 2, Figure 3, and Figure 3) in combination with DAPEv2;  
However, when looking at the ‘broader’ applicability in Figure 5, it quickly becomes clear that results across the board are much more inconsistent!   
-> e.g. ALiBi: ALiBi performs well on its own for training seq-len 128, is improved by DAPE-ALiBi – but significantly worse for DAPEv2; 
The manuscript however states that DAPE-1x3 ‘consistently improves performance’, which is incorrect and should be discussed (including insights)
- Appendix E / Section 4.8 shows results for DAPE with kernel-size 1 and 3 – I assume ‘1’ is the classic DAPE, and ‘3’ the v2?   
If so, again – results vary a lot in terms of which one is better for which task and combined with which ‘pe-method’, and I don’t see this discussed in the manuscript appropriately. 
- General: A wider discussion of the limitations would significantly help any reader/user, and I’d suggest the authors consider being upfront about these and provide the reader with helpful guidance (Similarly when using DAPEv2 with FIRE, while there is some improvement, it still ‘diverges’ quickly)

- I’d like the authors to include actual insights based on their experiences and the background knowledge of working with these different approaches (FIRE, ALiBi, Kerple, etc.) – e.g. is one generally preferable? If not, what are the situations you would recommend combining DAPEv2 with any particular one of these?
- In Figure 6, although the model can cheat, I’d be curious why the authors think that the DAPEv2-ALiBi becomes significantly less stable (than both non-cheating and original-non-cheating) 

Additional comments:
- I’d suggest the authors replace some of the references through the seminal works in their introduction in terms of how Transformers have made an impact (e.g. noting CV but not citing ViT/DeiT isn’t good research practice, as these authors should be acknowledged)
- I’d like to suggest the authors to check and potentially slightly rework the manuscript in terms of preciseness of their wording; While I am aware this might be due to language barrier, there are multiple instances where statements are misleading/confusing/too general, e.g. 
  - Abstract: ‘[…] contributing to interactions among distinct tokens, in contrast to earlier feed-forward NNs’ -> This is not really true/correct, as any FFN can establish interactions between elements of data – e.g. a CNN establishes the same over a local window in a sequence, etc.; 
  - L 49: ‘rendering the outputs non-sensical’ -> In the context of NLP, the output will still be a valid word and hence ‘sensical’, the architecture simply loses its ability to learn relationships over a sequence and reverts back to sets/bag-of-words; Also note: The authors discuss “Transformers” in general, and there actually are multiple use cases where Transformers are used on set-based problems
  - …


---
## Update post-rebuttal: 
Some of my concerns have been addressed, and I am therefore increasing my score slightly from 5 to 6 -- but it still remains a borderline case to me.
soundness: 2
strengths: **Originality & Significance:**   
- The authors slightly expand on the insights of the original DAPE paper and provide new results on the original and their improved variant 

**Quality:**   
- Experiments conducted across two datasets in comparison with multiple popular ‘positional embedding’ methods, including NoPE, RoPE, CoPE, ALiBi, Kerple and FiRE
- Insights into how the computational complexity is affected are provided, as well as results for three model sizes

**Clarity:**  
- The paper is mostly easy to read; 
- Graphs and tables are clearly labeled and easy to interpret
summary: The authors propose to improve the extrapolation abilities of Transformer models beyond their training sequence length by building upon the previously introduced method of data-adaptive positional encodings (DAPE). The authors find that replacing DAPE’s standard MLP through a convolutional MLP further improves performance.
weaknesses: _TL;DR: While I appreciate the work the authors have put into the manuscript and their experiments, the main ‘methodological’ novelty facilitating the approach has already been presented in the original DAPE paper. The authors’ addition of using a convolution instead of an MLP (i.e. replacing a 1x1 conv with a 1x3 conv, combined with inconsistent improvements) combined with the manuscript in its current state is in my opinion not enough to pass the bar for ICLR;_

-	Minor ‘methodological’ addition to existing DAPE, with results varying from ‘improvement’ to ‘decrease in performance’ – see questions. 
-	Insufficient (no) discussion of limitations, although inconsistencies can be easily seen from the presented results – see questions.
-	Interpretation on an in-sight level of results obtained with different method (FIRE, ALiBi, etc.) could be significantly extended
-	Minor: Quality of Manuscript in terms of wording/preciseness of statements

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: Please see weaknesses (especially the third one).
soundness: 2
strengths: - The paper is clearly written.
- Extending the MLP in DAPE’s attention model to 1x3 convolution achieves improvement over multiple experiments.
summary: This paper introduces an incremental change over the prior work DAPE (Zheng et al., 2024), by extending the MLP used in the attention to 1x3 convolution. This small change achieves improvement over multiple experiment settings.
weaknesses: - The major concern is this paper only introduces an incremental change over DAPE, I.e. extending the MLP in attention model to 1x3 convolution. In addition, compared to the gap between DAPE and other baselines, the gap between this paper and DAPE is relatively small.

- This paper could be written in a more straightforward way, by directly showing the difference between it and DAPE, and highlighting why it is crucial. Readers may have confusion about the contribution of this paper and DAPE.

- Line 126: It is hard to buy the insight: *Transformer’s length extrapolation ability is limited by the expressiveness of the naive query-key dot product.* This conclusion is drawn by showing DAPE without position encoding still achieves improvement. But there exists another explanation as follows. Transformer’s length extrapolation ability is limited due to the lack of accurate position encoding. MLP in DAPE implicitly learns the spatial information from the dot product of query and key, thus improving the performance. And extending MLP to 1x3 convolution can further improve encoding the spatial information.

- Discussion about an important reference is missing. “On Translation Invariance in CNNs: Convolutional Layers can Exploit Absolute Spatial Location” (in CVPR 2020), by Osman Semih Kayhan and Jan C. van Gemert. It found that the boundary effects operate even far from the image boundary, allowing the network to exploit absolute spatial location all over the image. This may help explain why convolution introduces more gains.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 1
questions: Could the authors elaborate on "Proposition 1: Transformers incorporating convolution operations can perform associative recall tasks without the need for positional encoding"? The rationale behind this proposition is unclear and requires further explanation.
soundness: 3
strengths: + The application of convolution on attention maps to improve relative position encoding in large language models is both novel and inspiring. The use of convolution, a fast and efficient operation, allows for seamless integration into existing frameworks.

+ The proposed method demonstrates strong performance on the length extrapolation task, outperforming established techniques such as RoPE and NoPE, which underscores its effectiveness.
summary: This paper extends DAPE and proposes a new approach to address the Transformer long context extrapolation problem by treating attention scores as feature maps. The authors conceptualize attention mechanisms akin to image processing techniques, utilizing convolutional operations on attention scores across different heads. This methodology, inspired by methods in computer vision, enhances Transformer performance on extrapolation tasks across multiple lengths, both in theoretical underpinnings and through empirical validation. It has outperformed some popular position embedding methods such as RoPE and NoPE.
weaknesses: - The paper suffers from poor writing and organizational structure. Basic variables such as X, W_Q, and W_K​ are not adequately explained as the context, despite that Transformers are quite popular.

- Confusing Arguments: 
1) Line 181-182 states: "The result of DAPE-NoPE (the Zheng et al. (2024) only combine DAPE with ALiBi, Kerple and FIRE but not with NoPE or RoPE)." This sentence is confusing and seems disconnected from the preceding context. 
2) Line 191-192 mentions: "potentially hindering the evolution of next-generation Transformer models," which lacks clarity and context.
3) Line 198 states: "RoPE first computes the classic attention scores of key-query multiplication with RoPE." This description is unclear and requires further elaboration.

- The authors fail to adequately explain the rationale and motivation for applying convolution to embed position information, abruptly transitioning to technical details without sufficient context.

- The current popular solution for long-context extrapolation is to fine-tune RoPE-based LLMs on long-context data, which is not addressed in the baseline comparisons.

- There is no discussion of computational efficiency metrics such as FLOPS, which would be valuable for assessing the proposed method's practicality.

- The benchmarks employed in the study are limited, reducing the generalizability of the findings.

Meta-review: The proposed method replaces the MLP layer of DAPE with a 1x3 convolutional layer. While the author rebuttal successfully addressed most clarification requests from the reviewers, a critical concern regarding the limited technical novelty (raised by Reviewers i54z and UyHR) remains unresolved. Overall, the contributions of the paper fall short of the ICLR standard, and the AC, therefore, recommends rejecting the paper.

Additionally, the paper must differentiate itself from a significantly related study on correlation-structure-based attention:
[A] Kim et al., Learning Correlation Structures for Vision Transformers, CVPR 2024.

Decision: Reject