Sliding Window Attention Training for Efficient Large Language Models

摘要

基于Transformer的大语言模型（LLMs）最近取得的进展显示出其在各种任务上的卓越能力。然而，它们在处理长文档时，计算复杂度与序列长度呈二次方关系，这仍然是一个重大瓶颈。因此，人们提出了许多方法，如稀疏注意力机制和状态空间模型，以提高大语言模型在处理长序列时的效率。虽然这些方法提高了效率，但它们往往需要复杂的架构和并行训练技术。这就需要一种简单而有效的模型，同时保留基本的Transformer架构。为此，我们引入了SWAT（Sliding Window Attention Training，滑动窗口注意力训练），它通过滑动窗口注意力训练实现了对长上下文的高效处理。具体来说，SWAT使用sigmoid函数代替softmax函数，以实现高效的信息压缩和保留。然后，它利用平衡的ALiBi（Attention with Linear Biases，带线性偏差的注意力）和旋转位置嵌入（Rotary Position Embedding，RoPE）来稳定训练过程。在推理时，SWAT通过滑动窗口注意力保持线性计算复杂度，同时保持模型性能。与主流的线性循环架构相比，SWAT在八个常识推理基准测试中取得了最先进的（SOTA）结果。代码可通过