Local-global attention

Overview

Local-global attention gives each query two reaches. The local part lets it read keys inside a fixed window near its own position, so the cost of that part per query does not grow with the full length. The global part gives it a route to text far away. The papers differ on where that route comes from.

Gemma 2 keeps the route in separate layers. Some layers use a 4096-token sliding window and the others span 8192 tokens, which is the context length 2408.00118 §2 tab. 1. LongT5 keeps every encoder layer sparse. Each layer adds one summary token per 16-token block of input, built from that block's embeddings and discarded after the layer 2112.07916 §3.1.2. HiLo keeps the route inside every attention layer. A share of the heads reads only non-overlapping windows, and the rest read a pooled version of the whole grid 2205.13213 §4.1.

Variants

Alternating sliding-window and global layers (Gemma 2)

Layers alternate between a local sliding-window layer and a global layer. In both kinds the mask is causal, so a query never reads a later position. Local layers also drop keys that lie more than 4096 tokens back [arxiv:2408.00118 §2]. The same attention layers use logit soft-capping at 50.0, a separate choice that this concept does not cover 2408.00118 §2. Card: Gemma 2: Improving Open Language Models at a Practical Size.

Local window with transient global tokens (LongT5 TGlobal)

Every encoder layer has the same two parts. The local part is a window of radius r = 127, so each token attends to 127 tokens on each side and the window adds no parameters 2112.07916 §3.1.1. The global part cuts the input into blocks of k = 16 tokens. Each block gives one global token: the sum of its embeddings, then normalization 2112.07916 §3.1.2. Each token attends to its window and to every global token of its own example. The global tokens are rebuilt in every layer and are not stored. The paper counts about 10k extra parameters for base, 25k for large and 50k for xl 2112.07916 §3.1.2, fn. 1. Card: LongT5: Efficient Text-To-Text Transformer for Long Sequences.

Local window only (LongT5 Local)

LongT5 also reports the local part alone, with no global tokens. Its cost is O(l × r) for input length l 2112.07916 §3.1.1. The TGlobal cost is O(l × (r + l/k)) 2112.07916 §3.1.2, so the global term is the extra work. The same card gives the measured trade-off. On a 4x8 TPUv3 slice at input length 2048, base models ran at 860 sequences per second for Local, 765 for TGlobal and 479 for T5.1.1 base. The longest input before out-of-memory was 36k for base Local, 12k for base TGlobal and 6k for T5.1.1 base; for large models it was 10k, 6k and 3k 2112.07916 §5.1, fig. 3. On Natural Questions the paper reports Local clearly below TGlobal and T5.1.1 [arxiv:2112.07916 §5.2, fig. 4]. At input 512 with large models, the reported F1 is 58.00 for Local and 61.53 for TGlobal 2112.07916 tab. 4. These are author-reported single runs.

Split heads: windowed and pooled-global heads (HiLo)

Each attention layer splits its heads between two branches, with no separate global layers. A fraction alpha of the heads, 0.9 in the ImageNet setting, does self-attention inside non-overlapping s x s windows, with s = 2 and no shifting 2205.13213 §4.1 §5.1. The other heads take queries from every token, but their keys and values come from the grid after s x s average pooling, which leaves N / s² entries. The two outputs are concatenated 2205.13213 §4.1 eq. 6. The pooled branch still has a term quadratic in N, which grows at high resolution unless s grows. The paper therefore suggests a larger window there and tests window 4 for detection 2205.13213 §4.1 §5.2. Card: Fast Vision Transformers with HiLo Attention.

The paper's ImageNet-1K result for LITv2-S is 82.0% top-1 at 1,471 images per second on an RTX 3090. Plain multi-head self-attention in the same backbone gets 82.3% at 1,293 images per second 2205.13213 tab. 4. The paper's sweep finds that alpha 0 (windows only) performs badly, and that pure pooled attention (alpha 1.0) is competitive on ImageNet but 0.6 mIoU behind alpha 0.9 on ADE20K 2205.13213 §5.4 fig. 4. These are author-reported single runs.

When to use which

Implementation sketch

The pieces are a per-layer mask, a layer schedule, and, for TGlobal, a block summary. HiLo replaces the schedule with a per-layer head split. The papers fix the sizes, but not every ordering or boundary choice.

# Attention mask for one layer  [arxiv:2408.00118 §2]
# T: query positions, S: key positions, window: 4096 tokens
def attention_mask(T, S, is_local, window=4096):
    # Decoder-only: a query never attends to a later position  [§2]
    mask = causal_mask(T, S)
    # Local layers also drop keys more than `window` tokens back  [§2]
    if is_local:
        mask = mask and within_window(T, S, window)
    return mask

# Layer schedule  [arxiv:2408.00118 §2]
# The paper says the kinds alternate, not which one comes first
def layer_types(num_layers):
    # Assumed here: even-indexed layers local, odd-indexed global
    return ["local" if i % 2 == 0 else "global" for i in range(num_layers)]

# One global token per block, rebuilt in each layer  [arxiv:2112.07916 §3.1.2]
# Input: token embeddings (l, d); block size k = 16
def block_globals(embeddings, k=16):
    # Cut the sequence into fixed blocks of k tokens  [§3.1.2]
    blocks = split_into_blocks(embeddings, k)
    # Sum each block, then normalize; not kept after the layer  [§3.1.2]
    return [layer_norm(sum(block)) for block in blocks]

# HiLo head split inside one attention layer  [arxiv:2205.13213 §4.1] [eq. 6]
# Assumes H and W are divisible by s  [see HiLo card, Pitfalls]
def hilo_layer(tokens, H, W, s, alpha, num_heads):
    # Round the pooled-branch count down, as the released code does  [card]
    lo_heads = floor(alpha * num_heads)
    hi_heads = num_heads - lo_heads
    # Windowed heads: self-attention inside each s x s block  [§4.1]
    hi_out = window_attention(tokens, H, W, s, hi_heads)
    # Pooled heads: queries at full resolution, keys and values pooled  [§4.1]
    lo_out = pooled_attention(tokens, H, W, s, lo_heads)
    return concat([hi_out, lo_out], axis=-1)

Notes on where implementations may differ:

• Starting kind (Gemma 2). The released code starts with a local layer, but the paper does not say so 2408.00118 §2. Weights are indexed by layer, so the ordering decides which parameters apply to which layer. Check it against the Gemma 2 card, which compares paper and code.
• Window bound (Gemma 2). The paper gives only the window size. The Gemma 2 card records that the released sliding mask uses a two-sided bound on position, and that under a causal mask the upper bound removes nothing. That reading was not checked by running the code.
• Window radius (LongT5). The paper defines the window as r tokens on each side 2112.07916 §3.1.1. The LongT5 card warns that a naive block implementation can reach further than that, so check the mask before relying on the exact window.
• Packed examples (LongT5). Global tokens must be masked by example. Without that mask, examples packed into one sequence attend to each other 2112.07916 §3.1.2.
• Global key and value projections (LongT5). The LongT5 card reports that the Flaxformer code reuses the local key and value projections for global tokens. The paper does not say this, and the card flags the finding as unverified.
• Head count (HiLo). The paper writes alpha times the head count for the pooled heads, which need not be an integer 2205.13213 §4.1. The HiLo card records that the released code rounds down, so 12 heads at alpha 0.9 give 10 pooled heads and 2 windowed heads.
• Divisibility (HiLo). The window reshape needs H and W divisible by s. The paper pads feature maps for dense tasks 2205.13213 app. C. The HiLo card records that the released classification path does not pad and needs divisible sizes.
• Window size 1 (HiLo). The paper says stage 4 of LITv2 uses window size 1, which makes HiLo standard self-attention 2205.13213 app. A. Keep that setting if you copy the backbone.
• Position information (HiLo). The attention layers use no position embedding or relative bias. Position comes from a depthwise convolution in the feed-forward block, as the paper argues 2205.13213 §4.2.

Contradictions

None known among the cards for this slug. The three cards measure different things, so their numbers are not comparable. Gemma 2 reports perplexity after shrinking a window of an already trained model 2408.00118 tab. 10. LongT5 compares local-only and TGlobal encoders as separate models on summarization and QA 2112.07916 tab. 4, §5.2. HiLo reports ImageNet accuracy, throughput and detection and segmentation results for a vision backbone against other vision backbones 2205.13213 tab. 1 tab. 4.

Papers

• Gemma 2: Improving Open Language Models at a Practical Size: alternates a 4096-token sliding-window layer with a global layer in a family of 2B, 9B and 27B decoder-only models, and reports the window can be shrunk at inference on the 9B model.
• LongT5: Efficient Text-To-Text Transformer for Long Sequences: an encoder-decoder with a radius-127 local window and per-block transient global tokens, with a Local-only variant and speed and memory results against T5.1.1.
• Fast Vision Transformers with HiLo Attention: splits each attention layer's heads between windowed self-attention and attention to window-pooled keys and values, in the LITv2 vision backbone, with author-reported ImageNet, COCO and ADE20K results.