arXiv:2410.04211cs.CLcs.AI2024-10

提出新型稀疏注意力机制,低成本实现大模型长文本处理。

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

  • 通过相关性感知选择与合并,实现高效稀疏注意力。
  • 单张A100可微调32K长度的Llama2-7B,资源需求降低64倍。
  • 支持预训练、微调、推理全阶段长序列扩展,适合长文本任务研究者。

建模长序列对大规模模型至关重要,但扩展现有架构以处理更长序列面临显著技术和资源挑战。本文提出一种高效灵活的注意力架构,可在减少计算资源和微调时间的前提下,实现大语言模型上下文长度扩展。具体而言,引入相关性感知选择与合并机制,促进高效稀疏注意力;同时提出一种新型位置编码数据增强技术,提升对未见位置的泛化能力。实验结果表明:仅用一张A100,即可在32K序列长度上微调Llama2-7B,优于依赖子集回归的其他方法;提出一套完整的预训练、微调与推理阶段长序列扩展方案。预训练阶段,注意力机制部分破坏了平移不变性,仅对选中标记应用位置编码,实现高性能与强外推能力。微调阶段引入循环、随机截断与动态增长的NTK位置嵌入(CRD NTK),仅需16K微调长度,即可使Llama2-7B和Mistral-7B在长达100万甚至任意长度下推理。在400万上下文长度的passkey任务中达到100%准确率,100万上下文长度下困惑度稳定,资源消耗相比传统全注意力机制降低至少64倍,性能仍具竞争力。

原文摘要 · Abstract (English)

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.

长序列稀疏注意力微调优化上下文扩展

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。