Sliding Window Attention
When processing long-form text with large language models (LLMs), computational complexity and GPU memory consumption remain fundamental bottlenecks. The standard Transformer architecture relies on full attention, where every token interacts with every other token in the sequence. This results in quadratic complexity, O(N²), relative to sequence length N. Once inputs scale to tens or hundreds of thousands of tokens, compute resources are quickly exhausted.
To break through this limitation, Sliding Window Attention (SWA) has emerged as a critical technique for efficient long-context inference. It is widely adopted in modern open-source models such as Mistral 7B. The core intuition is straightforward: sacrifice global perception in favor of local associations.
From Full Attention to Local Windows
In standard self-attention, each token can attend to all other tokens in the sequence. Sliding window attention restricts this interaction: each token is only allowed to attend to neighboring tokens within a fixed distance W along the sequence dimension.
For the i-th token under a unidirectional causal mask, visibility is limited to the interval [i - W, i]. As the sequence progresses, this window "slides" forward, giving the mechanism its name. Tokens outside the window are masked out, typically by setting their attention logits to negative infinity before softmax.
Key Advantages
Adopting sliding window attention delivers two primary benefits:
- Linearized computational complexity: Because each token computes attention only against a fixed window of size W, complexity drops from O(N²) to O(N × W). When W is much smaller than N, this approaches linear complexity.
- Dramatically reduced memory footprint: During inference, the KV cache traditionally stores key and value matrices for every historical token, causing memory to grow linearly with sequence length. SWA only needs to maintain a KV cache of size W, substantially saving VRAM. This enables longer contexts and larger batch sizes.
Mathematical Formulation and Attention Masking
Implementation typically involves modifying the attention mask. Given a window size W, the j-th token is visible to the i-th token if:
max(0, i - W) ≤ j ≤ i
This produces the following behavior:
- When i < W, the token can attend to all tokens from the beginning of the sequence up to the current position.
- When i ≥ W, the token can attend only to itself and the preceding W tokens.
Visually, the attention matrix becomes a banded matrix with bandwidth W around the main diagonal, while all other positions are masked.
Stacking Layers to Recover Global Context
A common question is whether restricting each token to a local window sacrifices global understanding. The answer is no. This mirrors the receptive field principle in convolutional neural networks (CNNs):
- In the first layer, each token sees only its local W neighbors.
- In the second layer, each token aggregates information from the first layer's local windows, effectively covering a span of up to 2W tokens.
- As layers stack—32 layers or more in typical LLMs—the receptive field expands, allowing local information to propagate and fuse into global semantics through the depth of the network.
Experiments show that even with per-layer local attention, deep models retain strong long-text understanding and modeling capabilities.
Conceptual Implementation
Below is simplified PyTorch-style pseudocode demonstrating how to implement sliding window attention via masking:
import torch
import torch.nn.functional as F
def sliding_window_attention(q, k, v, window_size):
"""
q, k, v: [batch_size, num_heads, seq_len, head_dim]
window_size: int, the sliding window size W
"""
batch_size, num_heads, seq_len, head_dim = q.shape
# Standard attention scores [batch_size, num_heads, seq_len, seq_len]
scores = torch.matmul(q, k.transpose(-1, -2)) / (head_dim ** 0.5)
# Causal mask
causal_mask = torch.triu(
torch.full((seq_len, seq_len), float('-inf')), diagonal=1
).to(q.device)
# Sliding window mask
row_idx = torch.arange(seq_len, device=q.device).unsqueeze(1)
col_idx = torch.arange(seq_len, device=q.device).unsqueeze(0)
window_mask = (col_idx < row_idx - window_size) | (col_idx > row_idx)
# Apply masks
scores.masked_fill_(window_mask, float('-inf'))
scores.masked_fill_(causal_mask, float('-inf'))
# Softmax normalization
attn_weights = F.softmax(scores, dim=-1)
# Weighted sum of values
output = torch.matmul(attn_weights, v)
return output
Summary and Outlook
Sliding window attention achieves a clever balance: it preserves the expressive power of large language models while sharply reducing computational and memory overhead. In practice, cutting-edge architectures often combine SWA with additional mechanisms. For example, Mistral's Rolling Buffer Cache pairs sliding windows with global anchor tokens to reconcile local efficiency with specific global task requirements.
Understanding sliding window attention is an essential step toward grasping modern efficient LLM architectures.