C o CA : 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 \textbf{Co} llinear \textbf{C} onstrained \textbf{A} ttention 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 CoCA-based GPT model, trained with a context length of 512, can extend the context window up to 32K (60 \times ) without any fine-tuning.Additionally, incorporating CoCA into LLaMA-7B achieves extrapolation up to 32K within a training length of only 2K.Our code is publicly available at: https://github.com/codefuse-ai/Collinear-Constrained-Attention

## introduction in the seminal work of transformer @xcite , 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 figure 1 : perplexity evaluation on 100 pg-19 documents with a sliding window strategy (@xmath0. the perplexity of roformer @xcite sharply exceeds 1000 beyond its training length, while coca maintains a low plateau even at 60 × its training length. alibi @xcite encounters out of memory (oom) issues for input @xmath1, we suppose it maintains perplexity for @xmath2. existing studies mainly focus on either the attention kernel @xcite @xcite or position embedding @xcite . 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 @xcite and kerple @xcite , 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) @xcite to extrapolate the inference context length with minimal or no fine-tuning. for instance, position interpolation (pi) @xcite employs linear scaling on each position number from n to n/k, where k is the extrapolation ratio. ntk-aware scaled rope @xcite and dynamic-ntk @xcite 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 @xcite 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 @xcite . 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 @xcite . 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 @xcite 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 @xcite . 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: ## experimental setting 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. ## experiment results we conducted experiments to shed light on the following reasonable doubts: ## conclusion 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.