Sliding-window attention

Overview

In full causal attention, a token at position i reads every earlier token, so the work per token grows with the length of the text and the cache grows with it 2310.06825 §1 fig. 1. Sliding-window attention restricts each position to a fixed window of positions. In a causal decoder the window is the most recent positions. With window size W, a position reads at most about W keys in each layer 2310.06825 §2 fig. 1. The cost per token is then set by W and not by the prompt length 2310.06825 §2.

The window does not have to look backwards only. In the video encoder of ActionFormer, each time step attends to steps around it, so the window reaches both earlier and later steps 2202.07925 §3.1. That use is not causal, and it works on features extracted in advance rather than on generated text.

A single layer sees only W positions, but the layers are stacked. A position in layer k can carry information from about k \times W positions back, so the theoretical reach of the last layer grows with depth 2310.06825 §2 fig. 1. This is a span through the layers, not direct attention to distant tokens. The Mistral paper does not measure accuracy at that length 2310.06825 §2 §3.

The Mistral paper pairs the window with a rolling-buffer cache, which keeps only the last W keys and values per layer 2310.06825 §2. Without that cache, the window reduces compute but not the memory that grows with the text. This is a different axis from shifted-window attention, which splits an image or video into fixed blocks and shifts the grid between blocks. The alternating local layers of local-global attention are covered in the Gemma 2 variant below.

Variants

Window with stacked layers (causal, Mistral 7B)

Each layer reads only the previous layer's states inside the window, and the model is built from many such layers 2310.06825 §2 tab. 1. Mistral 7B uses W = 4096 with 32 layers 2310.06825 tab. 1 §2. The paper reports the reach of this stack as about 131K positions, which is 32 × 4096 2310.06825 §2 fig. 1. Card: Mistral 7B.

The paper states the window as the positions between i - W and i, while its figure describes at most W tokens 2310.06825 §2 fig. 1. These two readings differ by one key per position, and the paper does not decide between them 2310.06825 §2.

Alternating local and global layers (Gemma 2)

Gemma 2 alternates one local layer with one global layer. The local layers use a sliding window of 4096 tokens and the global layers span 8192 tokens 2408.00118 §2 tab. 1. Local layers therefore read at most 4096 earlier tokens directly, while global layers read the full context. The paper says the layers alternate without naming the first one; the code starts with a local layer, as the card's Paper vs code section records. Card: Gemma 2.

The Gemma 2 card lists no ablation that isolates this alternation, so its effect on quality is not stated. Its only window test is a shrink at inference on the 9B model, which changed perplexity little 2408.00118 tab. 10.

Two-sided window over a sequence (ActionFormer)

Each time step attends to the steps around it, within a window of size W, in every layer of the encoder 2202.07925 §3.1. The same W is used on every pyramid level, so W = 19 on the 16× level covers 304 steps 2202.07925 §3.1. The paper uses W = 19 on THUMOS14, 9 on EPIC-Kitchens 100 and 11 on ActivityNet-1.3 2202.07925 §4.1 §4.2 §4.3. Card: ActionFormer.

On THUMOS14 with I3D features, windows of 9, 19, 25 and 37 give 66.5, 66.8, 66.4 and 66.7 average mAP. Full attention gives 66.8. At W = 19 the model costs 45.3 GMACs against 57.8 for full attention 2202.07925 tab. 3b. The authors also report that the Transformer runs about 2× slower than their convolutional baseline at similar MACs, because self-attention is hard to parallelise on GPUs 2202.07925 §4.4 tab. 3b.

Rolling-buffer cache

The keys and values for absolute position i are written to slot i \bmod W of a buffer with W slots per layer, so new entries overwrite the oldest ones 2310.06825 §2 fig. 2. The cache size is therefore fixed, whatever the length of the text. The paper says this cuts cache memory by a factor of 8 at 32K tokens without hurting quality 2310.06825 §2. This is the paper's own analysis. The quality claim has no table or ablation behind it 2310.06825 §2 §3. The Gemma 2 and ActionFormer cards do not describe a rolling cache.

Chunked pre-fill

A long prompt can be pushed into the cache in chunks. The paper suggests a chunk size of W, and for each chunk the attention covers the cache through the window and the chunk itself with a causal mask 2310.06825 §2 fig. 3. The chunk size is a suggestion, not a fixed value 2310.06825 §2. The reference implementation's default is to treat the whole prompt as one chunk; see the card's Paper vs code section.

When to use which

Implementation sketch

Three parts: the causal sliding mask for a full-length pass, the rolling-buffer write for incremental decoding, and the two-sided window used for offline sequences. Batch, heads and rotary embeddings are omitted. The Mistral paper does not settle whether the causal window holds W or W + 1 keys, so the mask below uses the text's reading and should be tested. The ActionFormer paper does not give the offsets of its window, so the two-sided version assumes a centred window and should be checked against its code.

# Sliding-window attention for one layer  [arxiv:2310.06825 §2] [fig. 1]
# Input:  hidden (T, d) from the previous layer; window W
# Output: new_hidden (T, d)

def sliding_window_attention(hidden, window):
    # Position i reads previous-layer states from i - W up to i  [§2]
    for position i in range(T):
        # Boundary: W or W + 1 keys; the paper leaves this open  [§2] [fig. 1]
        visible = hidden[max(0, i - window) : i + 1]
        new_hidden[i] = attend(hidden[i], visible, visible)
    return new_hidden

# Rolling-buffer cache write  [arxiv:2310.06825 §2] [fig. 2]
# Input:  cache_k, cache_v of shape (W, n_kv_heads, head_dim)
#         key, value for absolute position i
# Output: updated cache

def write_to_cache(cache_k, cache_v, i, key, value, window):
    # Position i goes to slot i mod W; older entries are overwritten  [§2]
    slot = i % window
    cache_k[slot] = key
    cache_v[slot] = value
    return cache_k, cache_v

# Two-sided local attention for one layer  [arxiv:2202.07925 §3.1]
# Input:  feats (T, D) at one pyramid level; window W
# Output: new_feats (T, D)

def local_attention_two_sided(feats, window):
    # Each step attends to steps around it, including later ones  [§3.1]
    # The paper gives no offsets; a centred window is an assumption to check
    half = window // 2
    for step t in range(T):
        visible = feats[max(0, t - half) : min(T, t + half + 1)]
        new_feats[t] = attend(feats[t], visible, visible)
    return new_feats

Contradictions

None across the three papers on a shared measurement, because each measures something different: Mistral reports only a 16K speed figure and the cache saving, ActionFormer compares window sizes with full attention on video, and Gemma 2 tests a shrunk window only at inference on 9B 2310.06825 §2 2202.07925 tab. 3b 2408.00118 tab. 10.

Within the Mistral paper, the text and the figure describe the window boundary differently (see Variants) 2310.06825 §2 fig. 1. The abstract and the results text also differ from the table: the paper says Mistral 7B beats Llama 2 13B on every benchmark, while its own table shows Llama 2 13B ahead on one row. That row is not about attention, but it bears on how far the paper's claims can be read 2310.06825 §1 §3 tab. 2.

Papers

• Mistral 7B: introduces the model with window W = 4096, a rolling-buffer cache and chunked pre-fill, and reports the 16K speed figure and the 8× cache saving at 32K. The paper names sparse Transformers and Longformer as earlier uses of the window (arXiv:1904.10509, arXiv:2004.05150) 2310.06825 §1; this page uses none of their content.
• ActionFormer: uses a two-sided local self-attention window in a video Transformer encoder, with window-size ablations on THUMOS14 2202.07925 §3.1 tab. 3b.
• Gemma 2: alternates local sliding-window (4096) and global (8192) attention layers in a decoder-only model 2408.00118 §2 tab. 1.