Sliding Window Attention Training for Efficient Large Language Models

📄 arXiv: 2502.18845v2 📥 PDF

作者: Zichuan Fu, Wentao Song, Yejing Wang, Xian Wu, Yefeng Zheng, Yingying Zhang, Derong Xu, Xuetao Wei, Tong Xu, Xiangyu Zhao

分类: cs.CL, cs.AI, cs.LG

发布日期: 2025-02-26 (更新: 2025-06-04)

备注: 14 pages, 5 figures

🔗 代码/项目: GITHUB

💡 一句话要点

提出SWAT：通过滑动窗口注意力训练实现高效的大型语言模型

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 长文本处理 滑动窗口注意力 Sigmoid注意力 ALiBi Rotary位置嵌入 注意力机制 Transformer

📋 核心要点

  1. Transformer处理长序列时，计算复杂度呈平方级增长，成为其应用瓶颈，现有稀疏注意力等方法往往牺牲性能或引入结构复杂性。
  2. SWAT通过滑动窗口注意力训练，用Sigmoid函数替代Softmax，并结合ALiBi和Rotary位置嵌入，实现高效长文本处理。
  3. 实验结果表明，SWAT在多个基准测试中，性能超越了当前最优的线性循环架构，验证了其有效性。

📝 摘要（中文）

基于Transformer的大型语言模型（LLMs）在各种任务中展现了卓越的能力。然而，其计算复杂度随序列长度呈二次方增长，这仍然是处理长文档的重大瓶颈。为了提高LLMs在长序列上的效率，涌现了诸如稀疏注意力机制和状态空间模型等方法。尽管这些方法有效，但它们牺牲了性能或引入了结构复杂性。因此，需要一种简单而高效的模型，同时保留Transformer的基本架构。为此，我们提出了SWAT，它通过滑动窗口注意力训练实现高效的长上下文处理。本文首先将Transformer的低效归因于softmax操作的高方差导致的注意力沉没现象。然后，我们用sigmoid函数代替softmax，并利用平衡的ALiBi和Rotary位置嵌入来实现高效的信息压缩和保留。实验表明，在八个基准测试中，与最先进的线性循环架构相比，SWAT实现了SOTA性能。

🔬 方法详解

问题定义：论文旨在解决大型语言模型处理长序列时计算复杂度过高的问题。现有方法，如稀疏注意力机制和状态空间模型，虽然能降低计算复杂度，但往往会牺牲模型性能或引入额外的结构复杂性，使得模型训练和部署更加困难。因此，需要一种既能保持Transformer基本架构，又能高效处理长序列的方法。

核心思路：论文的核心思路是通过滑动窗口注意力训练（Sliding Window Attention Training, SWAT）来降低计算复杂度。SWAT的关键在于两点：一是使用Sigmoid函数替代Softmax函数，以缓解注意力沉没现象；二是结合ALiBi和Rotary位置嵌入，实现高效的信息压缩和保留。

技术框架：SWAT模型基于标准的Transformer架构，主要改进在于注意力机制部分。具体来说，SWAT使用滑动窗口注意力，即每个token只关注其周围固定窗口内的token。在滑动窗口注意力中，使用Sigmoid函数替代Softmax函数来计算注意力权重。同时，结合ALiBi（Attention with Linear Biases）和Rotary Position Embedding来编码位置信息。ALiBi通过线性偏差来表示位置信息，Rotary Position Embedding则通过旋转矩阵来编码位置信息。

关键创新：SWAT的关键创新在于将滑动窗口注意力、Sigmoid注意力、ALiBi和Rotary位置嵌入结合起来，形成一个高效且易于实现的解决方案。Sigmoid函数的使用缓解了Softmax函数带来的注意力沉没问题，使得模型能够更好地关注重要的token。ALiBi和Rotary位置嵌入则提供了高效的位置编码方式，有助于模型理解长序列中的位置关系。

关键设计：SWAT的关键设计包括：滑动窗口的大小、Sigmoid函数的参数、ALiBi的线性偏差以及Rotary Position Embedding的旋转角度。论文中可能给出了这些参数的具体设置，以及损失函数和训练策略的细节。例如，滑动窗口的大小决定了每个token能够关注的上下文范围，Sigmoid函数的参数影响注意力权重的分布，ALiBi和Rotary Position Embedding的参数影响位置编码的精度。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，SWAT在八个基准测试中取得了SOTA性能，超越了现有的线性循环架构。具体性能数据未知，但论文强调SWAT在保持Transformer基本架构的同时，实现了显著的性能提升，证明了其有效性。

🎯 应用场景

SWAT具有广泛的应用前景，尤其是在需要处理长文本的场景中，例如：长文档摘要、机器翻译、问答系统、代码生成等。通过降低计算复杂度，SWAT使得大型语言模型能够更高效地处理这些任务，从而提高模型的性能和用户体验。此外，SWAT的简单性和易于实现的特点，也使其更容易被集成到现有的系统中。

📄 摘要（原文）

Recent advances in transformer-based Large Language Models (LLMs) have demonstrated remarkable capabilities across various tasks. However, their quadratic computational complexity concerning sequence length remains a significant bottleneck for processing long documents. As a result, many efforts like sparse attention and state space models have been proposed to improve the efficiency of LLMs over long sequences. Though effective, these approaches compromise the performance or introduce structural complexity. This calls for a simple yet efficient model that preserves the fundamental Transformer architecture. To this end, we introduce SWAT, which enables efficient long-context handling via Sliding Window Attention Training. This paper first attributes the inefficiency of Transformers to the attention sink phenomenon resulting from the high variance of softmax operation. Then, we replace softmax with the sigmoid function and utilize a balanced ALiBi and Rotary Position Embedding for efficient information compression and retention. Experiments demonstrate that SWAT achieves SOTA performance compared with state-of-the-art linear recurrent architectures on eight benchmarks. Code is available at https://github.com/Fzkuji/swat-attention.