关键字:大型语言模型 、双块注意力 、上下文窗口扩展 、无需训练放缩 、长距离依赖
摘要
大型语言模型 (LLMs) 在处理和生成连贯文本时的能力显著减弱,当输入代币数量超出它们的预训练长度时。考虑到使用更长序列来微调大规模模型的昂贵开销,我们提出双块注意力 (DCA),它使得LLAMA2 70B能够支持超过100k代币的上下文窗口,而不需要持续训练。通过将长序列的注意力计算分解为基于块的模块,DCA能够有效地捕获相同块内代币之间(Intra-Chunk)和不同块之间(Inter-Chunk)的相对位置信息,并且与Flash Attention无缝整合。除了其令人印象深刻的推断能力之外,DCA在实际长上下文任务上的表现可媲美或甚至优于微调后的模型。向比其他专有模型,我们无需训练的70B模型达到了gpt-3.5-16k性能的94%,表明它是一个可行的开源替代品。本文中使用的所有代码和数据均已开源。
核心方法
核心方法包含三部分,分别是:
- Intra-Chunk Attention: 用于计算相同块内词元的注意力分数。
- Inter-Chunk Attention: 用于计算不同块之间词元的注意力分数。
- Successive-Chunk Attention: 用于计算相邻块之间词元的注意力,并保持局部性。
具体细节如下:
-
Intra-Chunk Attention:在同一块内执行注意力分数计算时,将序列分为n段,使块内的位置索引不超过块大小s。
-
Inter-Chunk Attention:为了聚合其他块的信息,引入了Inter-Chunk Attention。这需要一种新的位置索引集合,在块之间保持流动性信息。
-
Successive-Chunk Attention:特别针对相邻块之间的注意力计算,保持了模型的局部性,即模型倾向于依赖临近的词元来预测下一个词元。
关键点如下:
- 避免在RoPE中线性缩小位置指数或增加基频。
- 重用预训练模型中的原始位置指数及其嵌入,重新设计相对位置矩阵的构建。
- 使用Flash Attention在开源社区里进行长上下文放缩。
实验说明
整体来看,实验结果表明:
-
语言模型测试中,DCA在训练自由方法中表现出色,将LLMs从4k上下文窗口扩展到32k以上,同时保持了PPL的微小增长。
-
DCA与现有的流行的放缩位置编码(如PI、NTK)正交,验证了存在的长上下文LLMs可以进一步外推到192k上下文长度。
-
在长上下文理解方面,DCA在零次学习和少次学习设置中的系列长上下文理解基准的表现与或超过现有的最先进模型相当。
结论
本文中,我们介绍了双块注意力(DCA),一种新的无需训练的框架,用于扩展LLMs的上下文窗口。通过将注意力分数计算分段为块,DCA能够有效地捕捉长距离和短距离依赖关系。此外,基于块的注意力计算与Flash Attention无缝集成,对于开源社区中的长上下文放缩是关键要素。我们在各种任务中全面评估了模型,发现DCA的实际长上下文任务性能可媲美或甚至优于连续训练的模型。与专有模型相比,我们无需训练的70B模型可以达到gpt-3.5-16k性能的94%,为开源模型提供了一个有力的替代。所有代码和数据均已公开。