CoCA: Fusing Position Embedding with Collinear Constrained Attention in Transformers for Long Context Window Extending (2024)

Abstract: Self-attention and position embedding are two crucial modules in transformer-based Large Language Models (LLMs). However, the potential relationship between them is far from well studied, especially for long context window extending. In fact, anomalous behaviors that hinder long context extrapolation exist between Rotary Position Embedding (RoPE) and vanilla self-attention. Incorrect initial angles between Q and K can cause misestimation in modeling rotary position embedding of the closest tokens. To address this issue, we propose Collinear Constrained Attention mechanism, namely CoCA. Specifically, we enforce a collinear constraint between Q and K to seamlessly integrate RoPE and self-attention. While only adding minimal computational and spatial complexity, this integration significantly enhances long context window extrapolation ability. We provide an optimized implementation, making it a drop-in replacement for any existing transformer-based models. Extensive experiments demonstrate that CoCA excels in extending context windows. A CoCAbased GPT model, trained with a context length of 512, can extend the context window up to 32K (60×) without any fine-tuning. Additionally, incorporating CoCA into LLaMA7B achieves extrapolation up to 32K within a training length of only 2K. Our code is publicly available at: https://github.com/codefuseai/Collinear-Constrained-Attention

heading: 1 Introduction
text: In the seminal work of Transformer (Vaswani et al., 2017), it claims the ability of "extrapolating to sequence length longer than the ones encountered during training". This is an ideal hypothesis, but actually not work in practice for vanilla Transformer. Several subsequent works, collectively known as
∗Equal Contribution. This work was completed during Jing Ye’s internship at Ant Group.
†Corresponding Author
long context extrapolation, have delved into exploring the capabilities of large language models (LLMs) trained within the range of [1, N − 1] to effectively extend the testing sequence ≥ N .
Existing studies mainly focus on either the attention kernel (Beltagy et al., 2020; Ding et al., 2023; Han et al., 2023) or position embedding (Huang et al., 2023). However, they often neglect the intrinsic relationship between these two key components. Attention bias is an alternative to the explicit encoding of positional information. ALibi (Press et al., 2022) and KERPLE (Chi et al., 2022), incorporate heuristic and compositional triangle kernel-based negative causal attention bias, respectively. While these approaches effectively maintain low perplexity, they fall short in capturing long-range dependencies due to introducing local hypotheses to context tokens. Another branch of methods involve simply scaling Rotary Position Embedding (RoPE) (Su et al., 2024) to extrapolate the inference context length with minimal or
no fine-tuning. For instance, Position Interpolation (PI) (Chen et al., 2023) employs linear scaling on each position number from n to n/k, where k is the extrapolation ratio. NTK-aware Scaled RoPE (bloc97, 2023) and dynamic-NTK (Emozilla, 2023) combine high-frequency extrapolation and low-frequency interpolation. They scale the basis in RoPE upon the sequence length to adapt to the unseen position indices. Furthermore, CLEX (Chen et al., 2024) optimizes extrapolation capability by modeling the continuous dynamics over different scaling factors. However, these methods primarily alleviate the problem of modeling the rotation angles in out-of-distribution positions, without recognizing the intrinsic correlation between attention matrices and rotation angles. Therefore, these methods still suffer from a limited context window extending ratio.
Here, we present a new perspective on the relationship between position embedding (with a focus on RoPE) and the self-attention mechanism. RoPE utilizes a rotation matrix to encode absolute positions while simultaneously incorporating explicit relative position dependencies within the self-attention formulation (Su et al., 2024). It is designed based on the relative angular difference between the queries (Q) and keys (K). However, there is an initial angle between Q and K in 2-D subspace. Incorrect initial angles between Q and K introduce additional rotary borders, resulting in misestimation in modeling the rotary position embedding of the closest tokens. This rotational anomaly severely compromises the performance of long context extrapolation methods.
To address this undesirable behavior, we propose an innovative architecture called Collinear Constrained Attention (CoCA). Specifically, we enforce a collinear constraint between Q and K by initializing the angle between every two hidden dimensions in the Q and K vectors to 0. This allows for a seamless integration of RoPE and selfattention. The architecture of CoCA and its comparison with RoFormer (Su et al., 2024) is illustrated in Figure 2.
Extensive experiments show that a CoCA-based GPT model, trained within a 512 context length, can seamlessly extend the context window up to 32K (60×) without perplexity divergence. A comprehensive comparison between our method and existing methods is presented in Figure 1. Furthermore, it enhances long-context retrieval ability, achieving a passkey retrieval accuracy of 50%+
even when extrapolating to 16× longer than its training context length by applying dynamic-NTK (Emozilla, 2023). Additionally, by dropping CoCA in LLaMA-7B, we achieve extrapolation up to 32K within only a 2K training length.
Our main contributions can be summarized as follows:
• We unveil the undesirable context boundary behavior resulting from the lack of modeling the relationship between position embedding and self-attention.
• To tackle the undesirable context boundary behavior, we propose Collinear Constrained Attention (CoCA) to seamlessly integrate the position embedding and self-attention, achieving excellent long context window extrapolation performance.
• CoCA extends its context window from 512 to 32K without fine-tuning, achieving over 50% passkey accuracy even when 16 × longer than its training length. Using CoCA in LLaMA7B, we achieve extrapolation up to 32K within just 2K training length.
• CoCA introduces minimal computational and spatial complexity compared to vanilla selfattention. We provide an optimized implementation of CoCA, making it able to be a drop-in replacement for existing transformerbased models.

heading: 2 Preliminary

heading: 2.1 Rotary Position Embedding
text: Position embedding is a crucial component in transformer-based models. Here we focus on Rotary Position Embedding (RoPE) (Su et al., 2024), which is widely used by LLMs including LLaMA (Touvron et al., 2023a), LLaMA-2 (Touvron et al., 2023b), GPT-NeoX (Black et al., 2022) and Qwen (Bai et al., 2023). Suppose the positional index is an integer n ∈ [1, N ], and the corresponding input vector x = [x0, x1, ..., xd−1]T, where N is the sequence length, d is the dimension of the attention head. RoPE defines a vector-valued complex function f(x, n):
f(x, n) = [(x0 + ix1)e inθ0 , (x2 + ix3)e inθ1 ,
. . . , (xd−1 + ixd)e inθd/2−1 ]T,
where θj = B−2j/d,
(1)
here, the base B = 10, 000.
After the application of RoPE, the transformed vectors for query (q) and key (k) become f(q,m) and f(k, n), respectively. Here, m,n ∈ [0, N ] represent the positional indices of q and k. The attention operation is computed as the inner product between f(q,m) and f(k, n), defined as follows:
a(m,n) = Re(⟨f(q,m), f(k, n)⟩)
= Re d/2−1∑ j=0 (q2j + iq2j+1)(k2j − ik2j+1)ei(m−n)θj 
= d/2−1∑ j=0 [(q2jk2j + q2j+1k2j+1) cos((m− n)θj)
+ (q2jk2j+1 − q2j+1k2j) sin((m− n)θj)] (2) The attention score a(m− n) depends on the relative position (m− n).

heading: 2.2 Anomalous Behavior between RoPE and Attention Matrices
text: After applying RoPE, the attention score a(m− n) can be interpreted as the sum of d/2 inner products of complex numbers, as illustrated in Equation (2). For any pair of qj = (q2j , q2j+1) and kj = (k2j , k2j+1), which is the 2-dimensional slicing of q (or qm) and k (or kn), we introduce the
initial angle Θj between them, measured counterclockwise from kj to qj in the complex plane. Throughout our analysis, we keep the position of kj fixed, systematically rotating qj to comprehensively examine their relative positions. The final angle between qj and kj is represented as θ(qj ,kj) = Θj + (m− n)θj , where m and n are positional indices of qj and kj .
In this concept, the attention score can be formulized as:
a(m,n) = d/2−1∑ j=0 |qj ||kj | cos(θ(qj ,kj)) (3)
Refer to Figure 3 for a visual representation of this concept for any individual j ∈ [0, d/2] in the 2-D subspace. There are four distinct scenarios between qj and kj after rotation.
(1) Scenario (b) and (c): When m > n and Θj ≤ π, or m < n and Θj > π, the value of cos(θ(qj ,kj)) between qj and kj decreases with the expanding distance between m and n. In these 2 scenarios, no anomalous behavior is observed, as the attention score naturally decreases with the positional distance. This trend persists until the relative angle θ(qj ,kj) rotates beyond the boundary of π.
(2) Scenario (a) and (d): When m < n and Θj ≤ π, or m > n and Θj > π, intriguing phenomena emerge. As the distance between m and n grows, the value of cos(θ(qj ,kj)) between qj and kj paradoxically increases. This anomaly has a notable impact on attention scores, particularly affecting the τ closest tokens. In this context, τ is defined as Θj/θj for scenario (a) and (2π−Θj)/θj for scenario (d). Consequently, attention scores for these tokens are abnormally diminished.
For bidirectional language models, all four cases may occur. For causal models, only scenario (b) and (d) manifest, as m consistently exceeds n.
The attention score a(m− n) is the sum of d/2 inner-products, one of them appearing anomalous might be insignificant, however, experimental results confirm its significance.

heading: 2.3 Rotary Borders Analysis
text: Without lose of generality, in this section, we focus on case (d) to further analyse the anomalous behavior during context extrapolation. As shown in Figure 4, three distinct boundaries appear during the rotation process. These boundaries are kj , −qj , and qj , respectively. Once the relative angle between qj and kj crosses these boundaries, the monotonicity of < qj ,kj > will reverse. Specifically, the monotonicity of < qj ,kj > learns the
pattern described as follows:
< qj ,kj >=  ↑ (m− n),∀ − (2π −Θj) ≤ θ(qj ,kj) < 0 ↓ (m− n),∀0 ≤ θ(qj ,kj) < π ↑ (m− n),∀π ≤ θ(qj ,kj) < 2π ...
↑ (m− n), ∀(2k − 1)π ≤ θ(qj ,kj) < (2k)π ↓ (m− n), ∀(2k)π ≤ θ(qj ,kj) < (2k + 1)π
(4)
where θ(qj ,kj) = Θj + (m − n)θj defined in Section 2.2. This introduces confusion into the model during direct context extrapolation.
Apart from the interval −(2π − Θj) ≤ θ(qj ,kj) < 0 in Equation (4), the boundaries introduced by −qj and qj exhibit regularity with a periodicity of 2π. RoPE Scaling methods, such as PI and NTK, can effectively address these two boundary issues.
Conversely, the boundaries introduced by kj present significant challenges. These boundaries disrupt the 2π periodicity, increasing the complexity of the extrapolation. Additionally, there are a total of d/2 ∗ h ∗ L different such boundaries, where d, h, and L represent the head dimension, number of heads, and number of layers, respectively. Methods like PI and NTK are ineffective at these boundaries, even causing more positions to fall into this abnormal region. For PI, this increases by a factor of k (where k is the scaling factor), and for NTK, it increases by a factor of λ2j/d (where λ is the scaling factor). From this perspective, the positional concentration of PI results in more issues than NTK, with additional positions falling into the abnormal area during context extrapolation. This partially explains why NTK can be directly used in standard self-attention mechanisms without fine-tuning, while PI requires fine-tuning.
We resolve this third border-related challenge associated with kj by constraining kj to be collinear with qj , making it possible for extrapolation with regular periodicity.

heading: 3 Method

heading: 3.1 Collinear Constrained Attention
text: To tackle the anomalous behavior between RoPE and attention matrices, we propose a new attention mechanism, Collinear Constrained Attention, namely CoCA. Specifically, by applying a collinear constraint to any pair of qj = (q2j , q2j+1) and kj = (k2j , k2j+1), we seamlessly integrate RoPE into self-attention mechanism, achieving long context extrapolation.
To formalize this, considering a sequence of N input tokens SN = {wn}Nn=1, with corresponding word embeddings EN = {xn}Nn=1, where xn ∈ Rd is the d-dimensional word embedding vector of token wn without position information. First, the queries qm are obtained:
qm = WQxm, ∀m ∈ [1, N ] (5)
Next, we derive the keys kn with collinear constraints. This begins with the introducing of the constraint coefficient tn for each token position n, as depicted in Equation (6).
tn = WTxn,∀n ∈ [1, N ] (6)
Next, Equation (7) imposes the collinearity condition on the coefficients t2j and t2j+1, where tn = [t0, t1, ..., td−1]
T, ensuring that each pair is identical. This step effectively duplicates each 2- dimensional segment of the tensor.
t2j = t2j+1,∀j ∈ [0, d/2− 1] tn = Relu(tn)
(7)
Subsequently, the keys are calculated as shown in Equation (8), where kn are represented by the element-wise multiplication of Q = (q1, ..., qN ) and tn. This results in an expansion of dimensionality, as kn ∈ RN×d now includes an additional sequence length dimension. We address potential memory pressure by optimizing tensor contractions, ensuring no net increase in memory consumption. For an in-depth analysis, please refer to Appendix C.
kn = Q⊙ tn = (q1 ◦ tn, ..., qN ◦ tn) (8)
After that, we apply RoPE on Q and K, with the function f detailed in Equation (1).
f(qm) = f(qm,m) f(kn) = f(Q⊙ tn, n) = f(Q, n)⊙ tn (9)
Finally, the attention score of CoCA would be:
a(m,n) = Re(⟨f(qm,m), f(qm, n) ◦ tn⟩) (10)
Equation (10) illustrates the additional dimension of the keys in our CoCA mechanism. Specifically, it maps the index of each query to the additional dimension, establishing a collinear relationship between the n-th key and the m-th query. This is a critical aspect of our method.

heading: 3.2 Slacking the Constraint on Query
text: In Section 3.1, we present a theoretically precise solution for CoCA. However, practical implementation faces challenges due to the complexity of O(N2d) when storing f(Q, n). To address this issue, we provide a dual implementation with O(Nd) complexity in this section and prove their equivalence.
Theorem 1. (Dual implementation of CoCA) For any attention score defined in Equation (10), there exists an equivalent form as follows:
a(m,n) = Re(⟨f(qm,m), qm ◦ f(tn, n)⟩) (11)
with constraint:
q2j = q2j+1, ∀j ∈ [0, d/2− 1] (12)
Proof : The proof consists of two steps. Step 1. We prove that, by imposing the constraint q2j = q2j+1, ∀j ∈ [0, d/2 − 1], Re(⟨f(qm,m), qm ◦ f(tn, n)⟩) is equivalent to Re(⟨f(qm,m), f(qm, n) ◦ tn⟩).
To see this, we calculate the difference between f(qm, n) ◦ tn and qm ◦ f(tn, n):
f(qm, n) ◦ tn − qm ◦ f(tn, n)
=  t0(q0 cosnθ0 − q1 sinnθ0) t1(q0 sinnθ0 + q1 cosnθ0)
. . . td−2(qd−2 cosnθd/2−1 − qd−1 sinnθd/2−1) td−1(qd−2 sinnθd/2−1 + qd−1 cosnθd/2−1)

−  q0(t0 cosnθ0 − t1 sinnθ0) q1(t0 sinnθ0 + t1 cosnθ0)
. . . qd−2(td−2 cosnθd/2−1 − td−1 sinnθd/2−1) qd−1(td−2 sinnθd/2−1 + td−1 cosnθd/2−1)  (13)
Recall that t2j = t2j+1,∀j ∈ [0, d/2 − 1] (see Equation (7)), Equation (13) is equivalent to:
f(qm, n) ◦ tn − qm ◦ f(tn, n)
=  t0(q0 − q1) sinnθ0 t1(q0 − q1) sinnθ0
. . . td−2(qd−2 − qd−1) sinnθd/2−1 td−1(qd−2 − qd−1) sinnθd/2−1
 (14)
Clearly, if we impose the constraint q2j = q2j+1, ∀j ∈ [0, d/2 − 1], the vector in Equation (14) becomes null and we deduce that:
f(qm, n) ◦ tn − qm ◦ f(tn, n) = 0 (15)
Consequently, with the constraint q2j = q2j+1, ∀j ∈ [0, d/2− 1], we have:
Re(⟨f(qm,m), qm ◦ f(tn, n)⟩) = Re(⟨f(qm,m), f(qm, n) ◦ tn⟩)
(16)
Step 2. We further demonstrate that, q2j = q2j+1, ∀j ∈ [0, d/2− 1] is in fact a redundant constraint when calculating Re(⟨f(qm,m), f(qm, n)◦ tn⟩). To verify this, we expand the inner product:
Re(⟨f(qm,m), f(qm, n) ◦ tn⟩)
= d/2−1∑ j=0 [(q22jt2j + q 2 2j+1t2j+1) cos((m− n)θj)
+ (q2jq2j+1t2j − q2j+1q2jt2j+1) sin((m− n)θj)] (17)
Recall again t2j = t2j+1, ∀j ∈ [0, d/2 − 1], we have
Re(⟨f(qm,m), f(qm, n) ◦ tn⟩)
= d/2−1∑ j=0 t2j [(q 2 2j + q 2 2j+1) cos((m− n)θj)]
= d/2−1∑ j=0 t2j |qj |2 cos((m− n)θj)
(18)
This implies that Re(⟨f(qm,m), f(qm, n) ◦ tn⟩) depends solely on the magnitude of qj = (q2j , q2j+1) in 2-D subspace, demonstrating the independence of the relationship between q2j and q2j+1. Refer to Appendix D.2 for the rigorous proof.
Now we conclude that, with the constraint q2j = q2j+1,∀j ∈ [0, d/2 − 1], Re(⟨f(qm,m), qm ◦ f(tn, n)⟩) is equivalent to Re(⟨f(qm,m), f(qm, n) ◦ tn⟩) with no constraint on query.
By removing q2j = q2j+1 constraint, we designate this modified version as CoCA-Slack. The mathematical definition is provided in Appendix D.3.

heading: 4 Experimental Setting
text: This section provides an overview of the experimental setup, including details regarding the training data utilized and the baseline models employed to evaluate the effectiveness of the proposed method.

heading: 4.1 Training Data
text: Our model undergoes training on a combination of datasets, including the Pile training dataset (Gao et al., 2020), BookCorpus (Zhu et al., 2015), and the Wikipedia Corpus (Foundation, 2021). Additionally, we integrate manually collected opensource code from GitHub repositories with at least 1 star. From these datasets, we derive a sample of approximately 50B tokens, maintaining a composition of 75% text and 25% code.

heading: 4.2 Model Variants
text: To evaluate the effectiveness of our proposed approach, we train 3 models from scratch under identical experimental settings, including ALibi (Press et al., 2022), RoFomer (Su et al., 2024), and RoFormer+CoCA. All models share common specifications, featuring a size of 350M, 24 layers, a hidden dimension of 1024, 16 attention heads, and a maximum sequence length of 512. The key distinctions among them lie in variations in self-attention mechanisms and position embeddings. The implementation is optimized based on EleutherAI GPTNeoX1. Training a model from scratch demands substantial computational resources. Therefore, we also conduct experiments involving fine-tuning existing LLMs integrated with CoCA module. For this purpose, we utilize the LLaMA-7B model (Touvron et al., 2023a), which was trained with a context length of 2,048. Additionally, we employ dynamic-NTK for all the above models.
In summary, our comparison models are categorized as follows:
• Training from scratch: ALibi, RoFormer, RoFormer+CoCA, RoFormer+dynamic NTK, and RoFormer+dynamic NTK & CoCA.
• Fine-tuning LLM: LLaMA-7B, LLaMA7B+CoCA, LLaMA-7B+dynamic NTK, and LLaMA-7B+dynamic NTK & CoCA.

heading: 4.3 Implementation Detials
text: Pre-training Procedure We train all models using the next token prediction objective. We use AdamW (Loshchilov and Hutter, 2017) with β1 = 0.9 and β2 = 0.95. The learning rate follows a linear warm-up of 1% of total steps, starting from 1e-7. Subsequently, the learning rate is adjusted to 1e-4 with linear decay, eventually reaching 1e-5. The training utilizes 8 A100 GPUs, with a global
1https://github.com/EleutherAI/gpt-neox/tree/v2.0
batch size of 256 and 2 gradient steps accumulation, taking approximately 96 hours for 2 epochs. Fine-tuning Procedure To integrate CoCA in LLaMA, we employ a three-stage fine-tuning strategy: (1) only updating the K projection (7% of parameters). This stage aims to reconstruct the K projection in CoCA. By freezing the other parameters, we maintain attention scores as closely as possible to those of vanilla self-attention. (2) updating the QKV projection (21% of parameters). This stage aims to address intrinsic over-fitting in vanilla self-attention caused by undesired behaviors between RoPE and attention matrices. (3) fine-tuning all parameters. Each stage involves 15K steps, totaling 7.5B tokens (22B tokens overall), using the next token prediction objective. The training length of LLaMA-7B + CoCA remains at 2,048 as in the original model. All experiments are conducted with 32 A100 GPUs, setting a per-device batch size to 8 without gradient accumulation.

heading: 5 Experiment Results
text: We conducted experiments to shed light on the following reasonable doubts:
• Can CoCA improve the long context extrapolation performance of existing models?
• Can combining CoCA with other extending methods for RoPE effectively solve the three types of rotational boundary problems?
• Can CoCA maintain its performance in short context windows? (Detailed in Appendix B.4.)

heading: 5.1 Long Sequence Language Modeling
text: We evaluate the long sequence language modeling performance of both our model and baseline models on the test splits of the PG-19 dataset (Rae et al., 2020). For this evaluation, we randomly select a subsample comprising 100 documents, each containing at least 32,768 SentencePiece (Kudo and Richardson, 2018) tokens. We then truncate each test document to its initial 32,768 tokens. The evaluation involves calculating perplexity across different context window sizes using a sliding window approach, as described by (Press et al., 2022), with a stride of 512. The perplexity results for both our models and baselines are presented in Table 1 and Figure 1.
Based on our experiments, the evaluation results indicate that models integrated with CoCA exhibit significantly improved perplexity with longer
inference sequence length. For pre-trained models, by increasing the context window size from 512 (training context window size) to 32k, the perplexity of CoCA only increases from 20.11 to 171.63, whereas the perplexity of RoFormer becomes inf . Additionally, by increasing the context window size from 2K to 32K, the perplexity of finetuned LLaMA-7B+CoCA only increases 21.68, while LLaMA-7B with other extending methods increases more than 100. In general, we observe a consistent trend of CoCA achieving better perplexity with longer context windows. This suggests that CoCA has a more robust position embedding, enabling it to handle long context more effectively.
In contrast, we observe that models extended through the direct application of dynamic NTKaware Scaled RoPE exhibit a larger increase in perplexity at longer sequences. The perplexity of both RoFormer+dynamic NTK and LLaMA7B+dynamic NTK remains significantly higher than that combining CoCA. This difference becomes more pronounced as the sequence length increases. When the inference sequence length reaches 32k, the perplexity of RoFormer+dynamic NTK increases to 380.75, while the result for RoFormer+CoCA is only 171.63. Similarly, the perplexity of LLaMA-7B+dynamic NTK reaches 133.87, whereas LLaMA-7B+CoCA is only 29.95.
It is worth noting that the model achieves the best performance when both dynamic NTK and CoCA are combined. Particularly, LLaMA-7B+dynamic NTK & CoCA consistently maintains a very low perplexity. Even when the inference sequence length has reached 32k (16 × longer than the training length), the perplexity is only 13.89. This indicates that combining CoCA with other extending methods for RoPE can effectively address the three types of rotational boundary problems, achieving robust long-text extrapolation modeling capabilities.

heading: 5.2 Long Context Retrieval
text: Perplexity evaluates the performance of language model in predicting the next token. However, it is insufficient for a comprehensive assessment of the effective context window size. To address this, we conducted experiments using a passkey retrieval task (Mohtashami and Jaggi, 2023) to evaluate our method and baselines. The task involves identifying and retrieving a randomly hidden passkey within a lengthy document. More details of task definition and test sample generation settings can
Method Evaluation Context Window Size (Perplexity ↓)
Method Evaluation Context Window Size (Accuracy↑)
be found in Appendix B.1. Table 2 illustrates the accuracy of all tested models and their variants.
It is evident that ALibi exhibits failures when tested on sequences that were 1× longer than its training length, attributed to its local hypothesis. In contrast, our model consistently demonstrated superior accuracy. RoFormer+dynamic NTK & CoCA maintained a 50% accuracy, even with the test sequence length expanded to 16× its training length. Similarly, LLaMA-7B+dynamic NTK & CoCA still maintained a 30% accuracy when the test length was up to 32K.

heading: 5.3 Impact of Strict and Slack Constraint on
text: Q
As mentioned in Section 3.2, we implement a slack version of CoCA, referred to as CoCA-Slack. In this section, under the same experimental settings, we implement two versions of CoCA based on RoFormer-350M, labeled as CoCA-Slack and
CoCA-Strict. The comparison results between them are shown in Table 3.
We observe that the CoCA-Strict and CoCASlack exhibit similar performance in long sequence language modeling, as evidenced by comparable perplexity results. However, in the passkey retrieval task, contrary to our initial expectations, the CoCA-Strict model produces significantly lower results. This unexpected outcome suggests that models with a slack constraint may offer additional performance advantages, such as a larger effective context window size.
Understanding the reasons behind the superiority of slack constraints will be a key focus of our future work. In this regard, we provide some theoretical insights in Appendices D.2 and D.3. These insights aim to shed light on the underlying mechanisms that contribute to the observed differences and lay the groundwork for a more comprehensive analysis in subsequent research.

heading: 6 Conclusion
text: In this paper, we introduce Collinear Constrained Attention (CoCA), a novel approach that integrates position embedding with the self-attention mechanism. This innovation addresses undesired behaviors occurring around the context window boundary, which stem from discrepancies between RoPE and attention matrices. To the best of our knowledge, we are the first to analyze the initial angles between queries and keys in the self-attention mechanism, which give rise to anomalous phenomena in RoPE. Furthermore, we deduce a slack constraint for our implementation of CoCA. Extensive experiments demonstrate that incorporating CoCA into existing models significantly enhances performance in both long sequence language modeling and long context retrieval tasks. Additionally, the simultaneous integration of CoCA with other extended RoPE methods (e.g., dynamic-NTK) effectively mitigates three types of rotational boundary issues, resulting in remarkably improved capabilities for long context extrapolation.

heading: Limitations
text: Our current approach, CoCA, has so far been validated exclusively on RoPE. Experimental results demonstrate that CoCA enhances the long-context extrapolation performance of LLMs and further augments other extension methods by addressing rotational boundary issues. However, questions arise regarding its applicability to more general methods. While the effectiveness of slack posi-
tion embedding (SPE) is evident, a deeper understanding of the underlying reasons for its superior performance necessitates further investigation.

heading: Acknowledgement
text: Thanks to Jianlin Su from Moonshot AI Ltd and Zhongqiang Huang from Sangfor Technology for their valuable suggestions during the revision of the paper. They pointed out that CoCA is equivalent to relaxing the constraint on the query.

heading: A Related Work
text: Existing researches are mainly focused on the submodule of attention kernel or position embedding (Huang et al., 2023). In the following sections, we will separately introduce works on these two aspects: Section A.1 primarily addresses the former, while Section A.2 delves into the latter.

heading: A.1 Efficient Attention Mechanisms
text: Several works aim to implement efficient attention mechanisms with reduced computational demands, even achieving linear complexity. This enables extending the effective context length boundary of LLMs during inference by directly increasing Lmax in the pre-training stage (Ding et al., 2023; Mohtashami and Jaggi, 2023). Noteworthy approaches include Longformer (Beltagy et al., 2020), utilizing slide window attention, and models such as StreamingLLM (Xiao et al., 2023) and LM-Infinite (Han et al., 2023), which leverage a global-local attention mechanism. These variants have achieved success to a certain extent, but still face issues we unveiled in this work when using RoPE as their positional encoding method.

heading: A.2 Extrapolative Position Embedding Methods
text: Extrapolative position embedding methods aim to enhance the length generalization capability of LLMs.

heading: A.2.1 Attention Bias
text: In seeking alternatives to the explicit encoding of positional information, researchers have explored the integration of attention bias to capture the sequential and temporal nuances inherent in natural language. Early approaches, such as T5 (Ruder
et al., 2019), incorporate learnable attention bias. However, these methods do not explicitly address the challenge of length extrapolation. ALibi (Press et al., 2022) introduces a negative causal attention bias in a heuristic manner. Extending the ALiBi-style attention bias, KERPLE (Chi et al., 2022) treats it as a composition triangle kernel for self-attention and modifies style Xpos (Sun et al., 2023) by integrating it with RoPE. While these approaches effectively manage to maintain low perplexity levels, they fall short in capturing long-range dependencies due to introducing local hypotheses to context tokens.

heading: A.2.2 Extend RoPE
text: Besides, various strategies have been explored to extend RoPE (Su et al., 2024), a commonly employed positional encoding method in popular LLMs. Recent approaches involve simply scaling it to extrapolate the inference context length with minimal or no fine-tuning. For instance, Position Interpolation (PI) (Chen et al., 2023) applies linear scaling on each position number from n to n/k, densifying the representation space to extend the farthest length boundary by k times. Other approaches, such as NTK-aware Scaled RoPE (bloc97, 2023) and dynamic-NTK (Emozilla, 2023), combine high-frequency extrapolation and low-frequency interpolation. These training-free methods require limited code changes during inference (Peng et al., 2023). Furthermore, CLEX (Chen et al., 2024) optimizes extrapolation capability by modeling the continuous dynamics over different scaling factors. However, these methods aim solely at alleviating the problem of modeling the rotation angles in out-of-distribution (OOD) positions without recognizing the intrinsic correlation between attention matrices and rotation angles. Therefore, these methods still suffer from a limited context window extending ratio.
Previous methods independently investigate selfattention and position embedding without considering their intrinsic relationship, especially for the widely used RoPE method.

heading: B Additional Experiment

heading: B.1 Passkey Retrieval Task Definition
text: The passkey retrieval task, as proposed by Mohtashami and Jaggi (2023), involves the model recovering a randomly generated passkey hidden in a long document (see Listing 1 for the task prompt
format). Given a language model, we can determine the effective context window by assessing the upper and lower bounds. We assume a random passkey is k tokens away from the end of the input. If a model consistently fails to recover the passkey in multiple attempts, it suggests a context window size smaller than k. Conversely, successful retrievals indicate an effective context window size of at least k tokens (Chen et al., 2023).
There is an important info hidden inside a lot of irrelevant text. Find it and memorize them. I will quiz you about the important information there.
The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again. ... // Repeat x times.
// Passkey is 5 randomly generated numbers. The passkey is 12345. Remeber it. 12345 is the passkey.
The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again. ... // Repeat y times.
What is the passkey?
Listing 1: Prompt format for passkey retrieval (Mohtashami and Jaggi, 2023). The passkey is randomly generated from 10,000 to 99,999.
In our experiments, we generate test samples based on the prompt template in Listing 1, with lengths ranging from 512 to 32k. There are 100 test cases for each length. Given a language model, we input the passkey task prompt, examine the model’s output for the new 64 tokens, and calculate the accuracy.

heading: B.2 Analysis I : Consistency of Optimization in Position Embedding
text: The passkey retrieval results are presented in Section 5.2. Our model demonstrates superior passkey retrieval accuracy compared to baseline models under various conditions. However, we remain intrigued about its optimization, specifically whether it occurs within or beyond the confines of the training context window. To probe this further, we categorize the experimental data into two segments: passkey distance shorter and farther than the training context window length.
Figure 5 (a) illustrates the comparison results when the passkey is inserted less than 512 tokens away from the end token, while Figure 5 (b) illus-
trates that outside this range. When the passkey is inserted outside the 512 window, RoFormer+NTK & CoCA consistently outperforms Roformer+NTK across various lengths of inference sequences. This superiority persists when the passkey is inserted inside the 512 window. Notably, with an increase in the length of the inference sequence, RoFormer + NTK & CoCA demonstrates increasingly superior performance compared to RoFormer + NTK. These experiments suggest that our model can consistently optimize the position embedding and extend the effective context window.

heading: B.3 Analysis II : Impact of Dynamic-NTK in CoCA
text: We utilize the dynamic NTK method (Emozilla, 2023) during the inference process, applying it separately to both our model and the baseline model. To comprehensively assess the robustness of these models, we conduct a thorough validation by varying scaling factors (2, 4, and 8).
The results in Figures 1 and 6 demonstrate that, with the integration of the dynamic NTK method, our model achieves higher passkey accuracy and lower perplexity. Additionally, when the scaling factor varies between 2, 4, and 8, the vanilla RoFormer model fails to maintain stable performance. In contrast, CoCA consistently outperforms RoFormer at different scaling rates. This consistent trend indicates that our model is more robust, showing minimal performance fluctuations with changes in the scaling factor.
Furthermore, it suggests that by implementing collinear constraints, we can cleverly address anomalous behavior in RoPE, allowing RoPE to better leverage other extrapolation techniques.

heading: B.4 Analysis III : Compatibility of CoCA with
text: PI

heading: B.4.1 Experiment Setup
text: We conduct experiments utilizing the LLaMA-7B model (Touvron et al., 2023a) and its LLaMA-7B + CoCA variant, as detialed in Section 4.2. To apply PI, we follow the same settings as Chen et al. (2023). All parameters are fine-tuned, and the finetuning sequence length is set to 32k. The learning rate is adjusted to 2e − 5 with no decay. The experiments are conducted on 32 A100 GPUs, with a per-device batch size of 1 and no gradient accumulation. Each experiment is carried out for 6,000 steps.

heading: B.4.2 Long Context Validation
text: The fine-tuning results with PI are shown in Table 4. For long sequence modeling, both LLaMA-7B+PI and LLaMA-7B+CoCA & PI exhibit competitive performance across sequence lengths from 512 to 32768. Both methods achieve near-perfect accuracy in long context retrieval across all sequence lengths, with scores close to 1.0.
Overall, these findings suggest that the integration of PI and the CoCA module with the LLaMA7B model yields robust performance in both long sequence modeling and long context retrieval tasks. Additionally, the CoCA module demonstrates the ability to maintain performance levels comparable to PI, particularly evident at longer sequence lengths.

heading: B.4.3 Short Context Validation
text: In addition to enhancing long-context extrapolation, it is imperative to consider the practicality and scalability of CoCA in short contexts. Hence, we
evaluate our model on OpenCompass (Contributors, 2023), which comprises various dimensions, including reasoning, understanding, knowledge, language, and examination. The results are presented in Table 5.
The results indicate that LLaMA-7B with CoCA integration achieves performance comparable to the original LLaMA-7B across all evaluated dimensions. Specifically, CoCA integration does not significantly decrease performance in reasoning, understanding, and language tasks, and even surpasses the original LLaMA in some aspects. The observed decline in the knowledge dimension may be due to the distribution bias between the finetuning data and the original pre-training data (Jiang et al., 2023). The abnormal score in the CHID task could be attributed to a bias in our training data, which includes some Chinese documents. Overall, the integrated CoCA outperforms LLaMA-7B in comprehensive scores, demonstrating its effectiveness in both long-context and short-context scenarios, and highlighting its practical value.

heading: C Computational and Spatial Complexity Analysis
text: In this section, we analyze the computational and spatial complexities of CoCA. Table 6 provides a detailed comparison between the vanilla selfattention mechanism and CoCA.
When using the operation Krot = Q ◦ Trot, the computational complexity of CoCA does not exceed twice that of the vanilla self-attention. In practice, the training and inference speed of CoCA are comparable to the vanilla self-attention mech-
anism, with only a slight increase of about 5% to 10% , as depicted in Figure 7. However, there is a significant increase in spatial complexity when expanding Krot = Q◦Trot, becoming d times that of the vanilla self-attention. This level of spatial complexity is not practical for applications.
To address this problem, we can draw inspiration from the computation of QrotKTrot, which involves two steps: element-wise multiplication between Qrot and Krot followed by summation along the hidden dimension. Optimization is attainable by condensing the hidden dimension before fully expanding the sequence length dimension. Consequently, the spatial complexity is effectively reduced from N2d to N2. This optimization strategy is equally applicable to Krot = Q ◦ Trot. These
two components can be unified as articulated in Equation (19):
QrotK T rot = Qrot(Q ◦ Trot)T (19)
The commendable work accomplished by opt_einsum (a. Smith and Gray, 2018) facilitates the optimization of Equation (19). Experimental results indicate that Roformer+CoCA only demands approximately 60GB of GPU memory during inference with a sequence length of 32k, aligning closely with the memory consumption of the vanilla self-attention mechanism.

heading: D Theoretical Proof

heading: D.1 Strong Form of Long-term Decay with CoCA
text: We have introduced the basic theory of Rotary Position Embedding in Section 2.1. In fact, (Su et al., 2024) shows that RoPE has the characteristic of long-term decay:
|a(s)| = ∣∣∣∣∣Re d/2−1∑
j=0
hje isθj  ∣∣∣∣∣ ≤ (max
i |hi+1 − hi|) d/2−1∑ j=0 |Sj+1|
(20)
where hj := (q2j + iq2j+1)(k2j − ik2j+1) and Sj := ∑j−1 k=0 e
isθk , s = (m− n), m for the index of query, n for the index of key. Since the value of∑d/2−1
j=0 |Sj+1| decays with the relative distance s, the attention score decays either.
This characteristic ensures the stability of RoPE during extrapolation to some extent by preventing outliers. For CoCA, a stronger deduction can be formulated as follows:
|a(s)| ≤ (max i |li+1 − li|) d/2−1∑ j=0 |Cj+1| (21)
where lj := |q2j+iq2j+1||k2j+ik2j+1|, and Cj :=∑j−1 k=0 cos(sθk). Furthermore, it holds that:
|li+1 − li| ≤ |hi+1 − hi| (22)
Proof : Notice that when the initial angle Θj between qj and kj is 0, from Equation (18), the attention score can be simplified as:
a(s) = Re d/2−1∑ j=0 hje isθj  =
d/2−1∑ j=0 lj cos(sθj)
(23)
By following the study of (Su et al., 2024), we can easily derive the estimation in Equation (21).
For Equation (22), applying the triangle inequality, we get:
|hi+1 − hi| ≥ ||hi+1| − |hi|| (24)
Reviewing the definition of hi = (q2j + iq2j+1)(k2j − ik2j+1), we will find:
|hi+1 − hi| ≥ ||hi+1| − |hi|| = ||qi+1k∗i+1| − |qik∗i || = ||qi+1ki+1| − |qiki|| = |li+1 − li|
(25)

heading: D.2 Homeomorphism of Representation Space
text: Theorem 2. (Homeomorphism of representation space) For any attention score defined as follows:
a(m,n) = Re(⟨f(qm,m), f(qm, n) ◦ tn⟩) (26)
where qm is the query, m is the index number of query, tn is the collinear coefficient of CoCA, n is the index number of key, f is the rotation operator.
Denote its representation space with respect to qm as:
F (Q) = {a(m,n)|∀qm ∈ Q ⊂ Rd} (27)
where qm = WQxm, xm ∈ EN , m ∈ [1, N ] and EN is the word embedding space, WQ is the projection matrix.
Then we have the following homeomorphism:
F (Q) ∼= F (Qhalf ) (28)
where Qhalf = Q|q2j=q2j+1,∀j∈[0,d/2−1]. Proof : We prove it by demonstrating the homeomorphism mapping G:
G : F (Q) → F (Qhalf ) F ((q0, ..., qd−1) 7→ F (( √ q20 + q 2 1
2 , ...,
√ q2d−2 + q 2 d−1
2 )
(29)
It consists of three parts: Part I (G is a bijection): recall Equation (18), we have:
G(X) = X, ∀X ∈ F (Q) (30)
which implies that G is an identity mapping, naturally injective.
Next, we prove that G is also surjective: for any Y = F ((q0, ..., qd−1)|q2j=q2j+1) ∈ F (Qhalf ), there exists Ỹ ∈ F (Q) such that G(Ỹ ) = Y . Let
Ỹ = F ((q0, ..., qd−1)|q2j=q2j+1) ∈ F (Q) (31)
obviously we have G(Ỹ ) = Y . Part II (G is continuous): For any X0 ∈ F (Q), ϵ > 0, there exists δ, such that if |X − X0| < δ, then |G(X)− G(X0)| < ϵ.
From Part I, G is an identity mapping, let δ = ϵ, then the continuity of G holds. Part III (G−1 is continuous): G is an identity mapping, so is G−1. Following Part II, we immediately deduce that G−1 is continuous.

heading: D.3 Slack Position Embedding
text: Let H be a Hilbert space, and {T (n)|n ≥ 0} ⊂ L(H) is a family of bounded linear operator on H. A is the inner-product defined on H.
If it satisfies the following property, then we call {T (n)|n ≥ 0} is a relative (bounded linear) operator on H:
∃ {S(m)|m ∈ Z} : H×H → C (X,Y ) 7→ S(m)(X,Y )
is a family of semi-bilinear operator on H
s.t. S(p− q)(X,Y ) = A(T (p)(X), T (q)(Y )) ∀ p, q ∈ [0, N ], X, Y ∈ H,
(32)
Additionally, if it satisfies the following property, then we call {T (n)|n ≥ 0} is a slack relative
(bounded linear) operator on H:
∃ {S(m)|m ∈ Z} : H×H → C (X,Y ) 7→ S(m)(X,Y )
is a family of semi-bilinear operator on H
and H′ ⊂ H,H′ ̸= ∅ s.t. S(p− q)(X,Y ) = A(T (p)(X), T (q)(Y )) ∀ p, q ∈ [0, N ], X, Y ∈ H′,
(33)
Specifically, when H represents our projection space in self-attention, and {T (n)|n ≥ 0} is a position embedding on it, such as the Rotary Position Embedding (RoPE), we refer to it as a Slack Position Embedding (SPE) if it satisfies the property described in Equation (33).

heading: Future Work
text: Understanding the reasons behind the superiority of slack constraints will be a key focus of our future work. In this regard, we provide some theoretical insights in Appendices D. 2 and D. 3. These insights aim to shed light on the underlying mechanisms that contribute to the observed differences and lay the groundwork for a more comprehensive analysis in subsequent research