Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Shiyi Zhu, Jing Ye, Wei Jiang, Siqiao Xue, Qi Zhang, Yifan Wu, Jianguo Li

TL;DR

This work tackles the challenge of long-context extrapolation in transformer models by analyzing the interaction between Rotatory Position Embedding (RoPE) and self-attention. It proposes Collinear Constrained Attention (CoCA), which enforces a collinear relation between $Q$ and $K$ to tightly couple RoPE with attention, and introduces a Slack variant for memory efficiency. Empirical results show CoCA dramatically improves long-context perplexity and retrieval performance, achieving effective extrapolation up to $32K$ tokens and strong compatibility with other RoPE-extension techniques like dynamic NTK. The approach yields a drop-in replacement with modest computational overhead and demonstrates robust long-context capabilities in both language modeling and retrieval tasks, with potential applicability across RoPE-based LLMs.

Abstract

Self-attention and position embedding are two key 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 harming long context extrapolation exist between Rotary Position Embedding (RoPE) and vanilla self-attention unveiled by our work. To address this issue, we propose a novel attention mechanism, CoCA (Collinear Constrained Attention). 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 show that CoCA performs extraordinarily well in extending context windows. A CoCA-based GPT model, trained with a context length of 512, can seamlessly extend the context window up to 32K (60$\times$), without any fine-tuning. Additionally, by dropping CoCA in LLaMA-7B, we achieve extrapolation up to 32K within only 2K training length. Our code is publicly available at: https://github.com/codefuse-ai/Collinear-Constrained-Attention

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

TL;DR

This work tackles the challenge of long-context extrapolation in transformer models by analyzing the interaction between Rotatory Position Embedding (RoPE) and self-attention. It proposes Collinear Constrained Attention (CoCA), which enforces a collinear relation between and to tightly couple RoPE with attention, and introduces a Slack variant for memory efficiency. Empirical results show CoCA dramatically improves long-context perplexity and retrieval performance, achieving effective extrapolation up to tokens and strong compatibility with other RoPE-extension techniques like dynamic NTK. The approach yields a drop-in replacement with modest computational overhead and demonstrates robust long-context capabilities in both language modeling and retrieval tasks, with potential applicability across RoPE-based LLMs.

Abstract

Self-attention and position embedding are two key 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 harming long context extrapolation exist between Rotary Position Embedding (RoPE) and vanilla self-attention unveiled by our work. To address this issue, we propose a novel attention mechanism, CoCA (Collinear Constrained Attention). Specifically, we enforce a collinear constraint between and 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 show that CoCA performs extraordinarily well in extending context windows. A CoCA-based GPT model, trained with a context length of 512, can seamlessly extend the context window up to 32K (60), without any fine-tuning. Additionally, by dropping CoCA in LLaMA-7B, we achieve extrapolation up to 32K within only 2K training length. Our code is publicly available at: https://github.com/codefuse-ai/Collinear-Constrained-Attention
Paper Structure (34 sections, 2 theorems, 33 equations, 7 figures, 6 tables)

This paper contains 34 sections, 2 theorems, 33 equations, 7 figures, 6 tables.

Table of Contents

  1. Introduction
  2. Method
  3. Rotary Position Embedding
  4. Anomalous Behavior between RoPE and Attention Matrices
  5. Collinear Constrained Attention
  6. Slacking the Constraint on Query
  7. Experimental Setting
  8. Training Data
  9. Model Variants
  10. Implementation Detials
  11. Experiment Results
  12. Long Sequence Language Modeling
  13. Long Context Retrieval
  14. Impact of Strict and Slack Constraint on Q
  15. Conclusion
  16. ...and 19 more sections

Key Result

Theorem 1

(Dual implementation of CoCA) For any attention score defined in Equation (eq:8), there exists an equivalent form as follows: with constraint:

Figures (7)

  • Figure 1: Perplexity evaluation on 100 PG-19 documents with a sliding window strategy (Stride = 512). The perplexity of RoFormer Su2021RoFormerET sharply exceeds 1000 beyond its training length, while CoCA maintains a low plateau even at 60 $\times$ its training length. ALibi Press2021TrainST encounters Out of Memory (OOM) issues for input $N_{max}>$ 8000 due to flash-attention dao2022flashattention incompatibility, we suppose it maintains perplexity for $N_{max}>$ 8000.
  • Figure 2: Architecture comparison between RoFormer and CoCA. (a) RoFormer; (b) CoCA; (c) The implementation detial of K in CoCA. Q, T, and V are produced using projection matrices identical to those employed in the vanilla self-attention. T undergoes a halving operation, with the other half being duplicated. K is then computed as the element-wise product of Q and T, adhering to a collinear constraint with Q. Note that $\mathbf{k}_n \in \mathbb{R}^{N \times d}$, where $n\in [1,N]$ is the positional index of key, $d$ is the head dimension, $N$ is the sequence length.
  • Figure 3: Anomalous behavior of RoPE in 2-D plane. The inner product of vectors $\mathbf{q}_j$ and $\mathbf{k}_j$ is contingent upon the relative angle $\theta{(\mathbf{q}_j,\mathbf{k}_j)}$, defined as $\Theta_j + (m-n)\theta_j$. Here, $\Theta_j$ represents the initial angle, and $(m-n)\theta_j$ signifies the position-dependent rotation angle. (a) $m < n$ and $\Theta_j \leq \pi$. (b) $m > n$ and $\Theta_j \leq \pi$. (c) $m < n$ and $\Theta_j > \pi$. (d) $m > n$ and $\Theta_j > \pi$.
  • Figure 4: Comparison of effective context window between RoFormer + NTK and RoFormer + NTK & CoCA.
  • Figure 5: Passkey accuracy distribution on 4 range of distances. CoCA outperforms RoFormer for all distances and scaling factors of NTK.
  • ...and 2 more figures

Theorems & Definitions (2)

  • Theorem 1
  • Theorem 2