Correlation-Aware Select and Merge Attention for Efficient Fine-Tuning and Context Length Extension

• URL: http://arxiv.org/abs/2410.04211v1
• Date: Sat, 5 Oct 2024 15:59:32 GMT
• Title: Correlation-Aware Select and Merge Attention for Efficient Fine-Tuning and Context Length Extension
• Authors: Ning Wang, Zekun Li, Tongxin Bai, Guoqi Li,
• Abstract summary: We introduce correlation-aware selection and merging mechanisms to facilitate efficient sparse attention. We also propose a novel data augmentation technique involving positional encodings to enhance generalization to unseen positions. Our method achieves 100% accuracy on the passkey task with a context length of 4M and maintains stable perplexity at a 1M context length.
• Score: 21.729875191721984
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Modeling long sequences is crucial for various large-scale models; however, extending existing architectures to handle longer sequences presents significant technical and resource challenges. In this paper, we propose an efficient and flexible attention architecture that enables the extension of context lengths in large language models with reduced computational resources and fine-tuning time compared to other excellent methods. Specifically, we introduce correlation-aware selection and merging mechanisms to facilitate efficient sparse attention. In addition, we also propose a novel data augmentation technique involving positional encodings to enhance generalization to unseen positions. The results are as follows: First, using a single A100, we achieve fine-tuning on Llama2-7B with a sequence length of 32K, which is more efficient than other methods that rely on subsets for regression. Second, we present a comprehensive method for extending context lengths across the pre-training, fine-tuning, and inference phases. During pre-training, our attention mechanism partially breaks translation invariance during token selection, so we apply positional encodings only to the selected tokens. This approach achieves relatively high performance and significant extrapolation capabilities. For fine-tuning, we introduce Cyclic, Randomly Truncated, and Dynamically Growing NTK Positional Embedding (CRD NTK). This design allows fine-tuning with a sequence length of only 16K, enabling models such as Llama2-7B and Mistral-7B to perform inference with context lengths of up to 1M or even arbitrary lengths. Our method achieves 100\% accuracy on the passkey task with a context length of 4M and maintains stable perplexity at a 1M context length. This represents at least a 64-fold reduction in resource requirements compared to traditional full-attention mechanisms, while still achieving competitive performance.

Related papers

• UIO-LLMs: Unbiased Incremental Optimization for Long-Context LLMs [111.12010207132204]
  UIO-LLMs is an incremental optimization approach for memory-enhanced transformers under long-context settings. We refine the training process using the Truncated Backpropagation Through Time (TBPTT) algorithm. UIO-LLMs successfully handle long context, such as extending the context window of Llama2-7b-chat from 4K to 100K tokens with minimal 2% additional parameters.
  arXiv  Detail & Related papers  (2024-06-26T08:44:36Z)
• Sparser is Faster and Less is More: Efficient Sparse Attention for Long-Range Transformers [58.5711048151424]
  We introduce SPARSEK Attention, a novel sparse attention mechanism designed to overcome computational and memory obstacles. Our approach integrates a scoring network and a differentiable top-k mask operator, SPARSEK, to select a constant number of KV pairs for each query. Experimental results reveal that SPARSEK Attention outperforms previous sparse attention methods.
  arXiv  Detail & Related papers  (2024-06-24T15:55:59Z)
• SinkLoRA: Enhanced Efficiency and Chat Capabilities for Long-Context Large Language Models [4.497551890206997]
  Self-attention mechanism scales quadratically with sequence length. LongLoRA proposed shifted sparse attention (S(2)-Attn), effectively enabling context extension. SinkLoRA is still not as efficient as vanilla attention, reaching only 39% of the perplexity improvement compared to full attention.
  arXiv  Detail & Related papers  (2024-06-09T07:23:34Z)
• LongVQ: Long Sequence Modeling with Vector Quantization on Structured Memory [63.41820940103348]
  Self-attention mechanism's computational cost limits its practicality for long sequences. We propose a new method called LongVQ to compress the global abstraction as a length-fixed codebook. LongVQ effectively maintains dynamic global and local patterns, which helps to complement the lack of long-range dependency issues.
  arXiv  Detail & Related papers  (2024-04-17T08:26:34Z)
• Bidirectional Long-Range Parser for Sequential Data Understanding [3.76054468268713]
  We introduce BLRP (Bidirectional Long-Range), a novel and versatile attention mechanism designed to increase performance and efficiency on long-sequence tasks. We show the benefits and versatility of our approach on vision and language domains by demonstrating competitive results against state-of-the-art methods.
  arXiv  Detail & Related papers  (2024-04-08T05:45:03Z)
• BurstAttention: An Efficient Distributed Attention Framework for Extremely Long Sequences [96.74779792715819]
  We propose a distributed attention framework named BurstAttention'' to optimize memory access and communication operations. The experimental results under different length settings demonstrate that BurstAttention offers significant advantages for processing long sequences.
  arXiv  Detail & Related papers  (2024-03-14T12:51:58Z)
• Training-Free Long-Context Scaling of Large Language Models [114.53296002607993]
  We propose Dual Chunk Attention, which enables Llama2 70B to support context windows of more than 100k tokens without continual training. By decomposing the attention for long sequences into chunk-based modules, DCA manages to effectively capture the relative positional information of tokens.
  arXiv  Detail & Related papers  (2024-02-27T12:39:23Z)
• LongLoRA: Efficient Fine-tuning of Long-Context Large Language Models [67.58275666573496]
  LongLoRA is an efficient fine-tuning approach that extends the context sizes of pre-trained large language models. We demonstrate strong empirical results on various tasks on Llama2 models from 7B/13B to 70B.
  arXiv  Detail & Related papers  (2023-09-21T17:59:11Z)
• Efficiently Modeling Long Sequences with Structured State Spaces [15.456254157293836]
  We propose a new sequence model based on a new parameterization for the fundamental state space model. S4 achieves strong empirical results across a diverse range of established benchmarks, including (i) 91% accuracy on sequential CIFAR-10 with no data augmentation or auxiliary losses, on par with a larger 2-D ResNet.
  arXiv  Detail & Related papers  (2021-10-31T03:32:18Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.