Sliding Window Attention Training for Efficient Large Language Models

TL;DR

This work tackles the quadratic bottleneck of attention in long-context LLMs by introducing Sliding Window Attention Training (SWAT). SWAT replaces softmax with sigmoid attention and augments it with balanced ALiBi and RoPE to preserve information across sliding windows, enabling effective long-context processing within the standard Transformer architecture. Through extensive experiments on eight benchmarks and two model scales, SWAT achieves competitive or superior performance with linear-time attention, and ablations show the importance of the sigmoid activation, AliRope, and bidirectional slope design. The approach offers a practical, scalable path to efficient long-context language modeling with broad potential impact on real-world processing of long documents.

Abstract

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.

Figures (5)

• Figure 1: The demonstration of the SWA mechanism in Transformers.
• Figure 2: The $\log_{10}$ perplexity of four LLMs (Llama-2-7b, Llama-3.1-8B, Qwen2-7B and Mistral-7B-v0.1) on the third book of PG-19 test set using SWA inference. The window sizes are set not to exceed their respective training sequence lengths. The x-axis represents the sliding window size, and the y-axis represents the evaluation sequence length. For a fixed window size, perplexity increases (color shifts to blue) as the evaluation length grows.
• Figure 3: Heatmaps of attention scores (top four squares) and token embedding variance (bottom four lines) across different layers of Qwen2-7B. Higher token variance corresponds to stronger attention, highlighting their correlation. The two color bars indicate respective scales.
• Figure 4: The demonstration of the SWA mechanism in Transformers, where the model's information coverage includes residual and active tokens, depending on the model depth and window size.
• Figure 5: The training loss of models with different modules including Sigmoid, RoPE, and ALiBi, with the balanced slopes.