摘要
在各种大型模型中,对长序列进行建模至关重要;然而,将现有架构扩展到处理更长序列会带来重大的技术和资源挑战。本文提出了一种高效灵活的注意力架构,与其他优秀方法相比,它能够以更少的计算资源和微调时间来扩展大型语言模型中的上下文长度。具体来说,我们引入了相关性感知选择和合并机制,以促进高效的稀疏注意力。此外,我们还提出了一种新的数据增强技术,涉及位置编码,以增强对未见过位置的泛化能力。结果如下:首先,使用单个 A100,我们在 Llama2-7B 上实现了 32K 序列长度的微调,这比其他依赖子集进行回归的方法更有效。其次,我们提出了一种全面的方法,用于在预训练、微调和推理阶段扩展上下文长度。在预训练期间,我们的注意力机制在令牌选择过程中部分打破了翻译不变性,因此我们仅将位置编码应用于所选令牌。这种方法实现了相对较高的性能和显著的推断能力。对于微调,我们引入了循环的、随机截断的和动态增长的 NTK 位置嵌入(CRD NTK)。这种设计允许使用仅 16K 的序列长度进行微调,使 Llama2-7B 和 Mistral-7B 等模型能够以高达 1M 或甚至任意长度的上下文长度进行推理。我们的方法在 4M 上下文长度的通行证任务上实现了 100% 的准确率,并在 1M 上下文长度下保持稳定的困惑度。与传统的全注意力机制相比,这至少减少了 64 倍的资源需求,同时仍然实现了具有竞争力的性能。