Sliding Window Attention Training for Efficient Large Language Models

LLM Long Context ML TAVR
2025年02月26日
最近在基于Transformer的大型语言模型(LLMs)方面的进展已经在各种任务中展示了卓越的能力。然而,这些模型在处理长文档时,由于其与序列长度呈二次关系的计算复杂度,仍然存在显著的瓶颈。因此,许多努力,如稀疏注意力和状态空间模型等方法已经被提出,以提高LLMs在长序列上的效率。尽管这些方法有效,但它们要么牺牲了性能,要么引入了结构复杂性。这促使我们寻求一种既简单又高效的模型,同时保留基本的Transformer架构。为此,我们引入了SWAT,通过滑动窗口注意力训练(Sliding Window Attention Training)实现高效处理长上下文。本文首先将Transformer的低效归因于由softmax操作高方差导致的注意力汇聚现象。然后,我们用sigmoid函数替代softmax,并利用平衡的ALiBi和旋转位置嵌入(Rotary Position Embedding),以实现高效的信息压缩和保留。实验表明,SWAT在八个基准测试中,相较于最先进的线性递归架构,实现了最佳性能。代码可在https://anonymous.4open.science/r/SWAT-attention获取。
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://anonymous.4open.science/r/SWAT-attention.
许愿