Long Context · May 2023

Landmark Attention

intermediate

Introduced by Amirkeivan Mohtashami, Martin Jaggi (EPFL) in Landmark Attention: Random-Access Infinite Context Length for Transformers (arXiv:2305.16300)

long-context

Give attention 'random access' to a long context by summarizing each chunk with a learned 'landmark' token. Queries attend to landmarks first to decide which chunks to fully expand, then attend within only the selected chunks.

§ 1 · Premise

Where length and resolution collide

At training length $T = 2{,}048$ , the LLaMA-7B used as Mohtashami & Jaggi’s testbed already materializes a full attention matrix of $T^2 = 4.2 \times 10^6$ entries per head per layer (Mohtashami & Jaggi 2023, §1). Doubling context to 4K quadruples the FLOPs and KV-cache footprint; pushing to 32K — the inference regime the paper actually targets — would cost $256\times$ the attention compute at fixed depth. By 2023 several families of workaround existed, and each gave up something specific:

Sliding-window attention (Beltagy et al. 2020, arXiv 2004.05150; StreamingLLM keeps an attention-sink prefix on top, Xiao et al. 2023) keeps cost linear in $T$ but discards keys older than the window. A token at distance $> W$ is unreachable in one layer; deep stacks recover it slowly and lossily.
Compressive memory (Transformer-XL, Dai et al. 2019; Compressive Transformer, Rae et al. 2019) keeps everything but at coarsened resolution: the old tokens survive as pooled or attention-summarized vectors. Resolution degrades with distance.
kNN memory (Memorizing Transformers, Wu et al. 2022) retrieves a small slice of old K/V at full resolution, but the retrieval index is non-differentiable; the model cannot learn to index better.

Landmark Attention’s wager is that the right primitive is random access rather than uniform decay or sliding eviction: most queries need only a few specific chunks of the distant past at full resolution, and the model itself can learn which chunks those are. The contribution is a single special token — the landmark — that summarizes a chunk well enough to act as its addressable handle inside the attention softmax.

§ 2 · Derivation

From dense attention to landmark-gated retrieval

Starting point. Standard causal attention at position $t$ over a key/value cache of length $T$ computes

\mathrm{Attn}(\mathbf{q}_t, K_{1:t}, V_{1:t}) = \mathrm{softmax}\!\left(\frac{\mathbf{q}_t K_{1:t}^\top}{\sqrt{d_h}}\right) V_{1:t},

with $\mathbf{q}_t \in \mathbb{R}^{d_h}$ , $K_{1:t}, V_{1:t} \in \mathbb{R}^{t \times d_h}$ . The $O(T)$ keys all participate in the same softmax; both memory and FLOPs scale as $\Theta(T \cdot d_h)$ per query position.

Step 1 — chunking with landmark slots. Partition the sequence into contiguous blocks of $W$ tokens (Mohtashami & Jaggi use $W = 50$ in their LLaMA-7B fine-tune, §4.1). For each block $b$ , append a single learnable landmark token $\boldsymbol{\ell}_b$ that participates in self-attention exactly like a real token:

\underbrace{[x_1, \ldots, x_W, \boldsymbol{\ell}_1]}_{\text{block 1}},\; \underbrace{[x_{W+1}, \ldots, x_{2W}, \boldsymbol{\ell}_2]}_{\text{block 2}},\;\ldots

Each $\boldsymbol{\ell}_b$ is a sequence position; it has its own query, key, value via the shared $W_Q, W_K, W_V$ , and through normal causal attention it absorbs information from the $W$ tokens of its block. There is no auxiliary loss; the landmark’s representation is shaped entirely by next-token prediction on the chunked sequence.

Step 2 — grouped softmax at training time. Define $\mathcal{B}(t)$ as the index of the block containing token $t$ . Partition the keys before $t$ into blocks $b = 1, \ldots, \mathcal{B}(t)$ . A landmark $\boldsymbol{\ell}_b$ stands for its whole block. Mohtashami & Jaggi modify the softmax so that, for a query $\mathbf{q}_t$ and a non-landmark key $\mathbf{k}_j$ in block $b$ , the normalized attention weight is the product of two factors — (§3.2, Eq. 1):

\alpha_{t,j} \;=\; \underbrace{\frac{\exp(\mathbf{q}_t \cdot \boldsymbol{\ell}_b / \sqrt{d_h})}{\sum_{b'=1}^{\mathcal{B}(t)} \exp(\mathbf{q}_t \cdot \boldsymbol{\ell}_{b'} / \sqrt{d_h})}}_{\text{block-level (over landmarks)}} \;\cdot\; \underbrace{\frac{\exp(\mathbf{q}_t \cdot \mathbf{k}_j / \sqrt{d_h})}{\sum_{j' \in \text{block } b} \exp(\mathbf{q}_t \cdot \mathbf{k}_{j'} / \sqrt{d_h})}}_{\text{within-block (excluding landmark)}}.

The first factor is a softmax over landmarks only — exactly $\mathcal{B}(t)$ logits, one per block. The second factor is a per-block softmax over the $W$ ordinary keys, with the landmark removed. Reading this equation backwards is what makes the rest of the design work: the landmark $\boldsymbol{\ell}_b$ is forced to carry whatever signal the model needs to compute its block’s outer-softmax weight, because the gradient of $\alpha_{t,j}$ with respect to $\boldsymbol{\ell}_b$ flows only through the outer factor.

Why two softmaxes rather than one. A single softmax over $\{\mathbf{k}_j\} \cup \{\boldsymbol{\ell}_b\}$ would let landmarks and tokens compete on the same scale; the landmark would have to win each token’s mass directly. The grouped form decouples the selection problem (which block?) from the content problem (which token within the block?). The model can learn to make landmark logits very different in scale from token logits without that choice corrupting within-block attention.

Step 3 — sparse inference via landmark gating. At inference, materializing the grouped softmax above costs the same as dense attention. The point of the construction is that, with the landmarks already trained, we can replace the inner softmax over every block with a sparse approximation:

\widehat{\mathrm{Attn}}(\mathbf{q}_t) = \sum_{b \in \mathrm{TopK}(\mathbf{q}_t \cdot L)} \;\;\;\; \sum_{j \in \text{block } b} \alpha_{t,j} \, \mathbf{v}_j,

where $L = [\boldsymbol{\ell}_1, \ldots, \boldsymbol{\ell}_C] \in \mathbb{R}^{C \times d_h}$ stacks the landmarks for the $C$ past blocks and $\mathrm{TopK}$ picks the highest-scoring landmark indices. Mohtashami & Jaggi report $K = 2$ retrieved blocks works for their LLaMA-7B fine-tune (§4.3). For blocks outside the top- $K$ set the weight is treated as zero — the inner softmax never runs.

Step 4 — cost accounting. Let $C = T / W$ be the number of past blocks at position $t$ , $W$ the block size, $K$ the retrieval budget, $d_h$ the head dimension. Per query, per head:

\mathrm{FLOPs} \approx \underbrace{C \cdot d_h}_{\text{landmark scoring}} + \underbrace{K \cdot W \cdot d_h}_{\text{in-block attention}} \;=\; d_h \cdot \left( \tfrac{T}{W} + K W \right).

This is minimized at $W \approx \sqrt{T/K}$ — the classic two-level lookup balance. For $T = 32{,}768$ , $W = 50$ , $K = 2$ that gives $\approx 755 d_h$ FLOPs per query vs. $T d_h = 32{,}768 d_h$ for dense attention, a $\sim 43\times$ reduction at the FLOP level. Memory follows the same shape: only the landmark matrix $L \in \mathbb{R}^{C \times d_h}$ and the top- $K$ blocks need to live in fast memory at any given step.

Step 5 — boundary handling. Two details Mohtashami & Jaggi treat carefully:

The current block (containing position $t$ ) and a small recent window of $r$ prior tokens are exempted from the gating — they are always attended to in full. This keeps short-range attention exact and avoids the failure mode where the top- $K$ blocks all sit deep in the past.
The training-time grouped softmax uses teacher-forced chunk boundaries that match what the inference loop will see; the model never learns under a different partition than it deploys under (§3.4).

Parameter count. The only added parameters are the embedding of the landmark token — $d_h \cdot H$ scalars in total, where $H$ is the number of heads, since the same landmark embedding is used in every position. Across a typical 32-layer 4096-dim model that is on the order of $\sim 4{,}000$ extra parameters, not per layer: the landmark is a special input token, not a per-layer projection.

Chunk size B16Top-K chunks to expand2

Each chunk ends with a learnable landmark token (dark blue tick). At inference, the query first attends only to the C landmarks — a cheap selection step that gates which chunks are interesting. Then, full attention runs only within the top-K chunks (the highlighted regions). Per-query cost drops from O(L) to O(C + K·B). The landmark tokens are trained jointly with the rest of the network, so they learn to encode each chunk's content well enough to support the selection.

§ 3 · Reference implementation

Grouped softmax in PyTorch-style pseudocode

# x: [B, T, d_model]  block-chunked with one landmark token appended per W real tokens
# block_id[t] in [0, C) tells which block position t belongs to
# is_landmark[t] is True iff position t is a landmark slot
def landmark_attn_train(q, k, v, block_id, is_landmark):
    # q, k, v: [B, T, H, d_h]
    logits = einsum("bthd,bshd->bhts", q, k) / sqrt(d_h)
    causal_mask(logits)

    # Outer softmax: only landmark keys participate, normalized over past blocks
    land_logits = logits.masked_fill(~is_landmark[None, None, None, :], -inf)
    outer = softmax(land_logits, dim=-1)            # [B, H, T, T], nonzero only at landmark cols

    # Inner softmax: only non-landmark keys, normalized within each block
    tok_logits = logits.masked_fill(is_landmark[None, None, None, :], -inf)
    # group_softmax normalizes within each block_id group along the key axis
    inner = group_softmax(tok_logits, group=block_id, dim=-1)

    # Combine: weight = outer-block-prob * inner-within-block-prob
    block_prob = scatter_to_token(outer, block_id)  # broadcast block weight to its tokens
    weights = block_prob * inner                    # [B, H, T, T]
    return einsum("bhts,bshd->bthd", weights, v)

def landmark_attn_infer(q_t, k_landmarks, v_landmarks, blocks_kv, K, recent_window):
    # k_landmarks: [C, H, d_h]   one per past block
    # blocks_kv: list of (K_b, V_b), each [W, H, d_h]
    scores = einsum("hd,chd->hc", q_t, k_landmarks) / sqrt(d_h)
    topk = scores.topk(K, dim=-1).indices            # [H, K]
    out = recent_attn(q_t, recent_window)            # always-on local window
    for h in range(H):
        for b in topk[h]:
            out[h] += attend(q_t[h], blocks_kv[b])   # full softmax inside selected block
    return out

The sketch elides three production concerns: (1) fused-kernel realization of the grouped softmax — Mohtashami & Jaggi rely on a custom Triton kernel (repo epfml/landmark-attention); (2) position encoding for landmarks, which they handle by giving the landmark a position equal to the last token of its block; (3) head-merged top- $K$ — heads can either vote individually or share a single retrieved block set.

§ 4 · Empirical evidence

What is and isn’t known

Introducing paper (Mohtashami & Jaggi 2023). Two main results:

Perplexity at training length. On English Wikipedia, a Transformer trained from scratch with landmark attention at block size $W = 50$ matches a dense-attention baseline within $\sim 0.05$ perplexity at sequence length 512, the training length (Table 1). The grouped softmax is not a quality regression at the trained length.
Extrapolation to longer contexts. Fine-tuning a LLaMA-7B checkpoint with landmark attention on 15B tokens of RedPajama at chunk size $W = 50$ produces a model that correctly retrieves a hidden passphrase from contexts up to $32{,}768$ tokens — well beyond LLaMA’s 2K pretraining length — at $> 90\%$ accuracy on the “Passkey Retrieval” task they introduce (Figure 3, §4.3). A dense LLaMA-7B baseline drops to chance at $\sim 4{,}000$ tokens.

Independent follow-up. Passkey retrieval became the canonical lightweight long-context probe after this paper; later work directly compares against landmark attention:

Activation Beacon (Zhang et al. 2024, arXiv 2401.03462, Table 5) reproduces landmark attention on LLaMA-2-7B and reports comparable passkey accuracy at 32K context but a $\sim 2$ point PG-19 perplexity gap versus the beacon-based summarizer at equal context length — i.e., landmark attention preserves retrieval well but loses fluency on broad-context tasks.
LongLoRA (Chen et al. 2023, arXiv 2309.12307, §5) uses passkey as one of its evaluation tasks and confirms landmark attention’s near-perfect recall at 32K, while reporting that LongLoRA’s S²-attention reaches comparable recall with $\sim 30\%$ less fine-tuning compute.
Lost in the Middle (Liu et al. 2023, arXiv 2307.03172) and the broader “needle-in-a-haystack” literature inherited the chunked-passkey evaluation protocol that landmark attention’s paper popularized.

Sensitivity studies — what is not publicly known. The introducing paper sweeps block size $W \in \{25, 50, 100\}$ and retrieval budget $K \in \{2, 4\}$ (Table 4), reporting that $W = 50, K = 2$ is the sweet spot for their LLaMA-7B setup. There is no public study of: how landmark attention behaves on tasks that require aggregating information across many distant blocks (e.g., document-level summarization), how the design interacts with RoPE position scaling at extrapolation lengths beyond 32K, or how it composes with GQA / MLA-style KV compression. I don’t know of an independent reproduction at $> 13$ B parameters; the landmark-attention codebase pinned at LLaMA-7B is what the field has compared against.

Production adoption. None recorded in this knowledge base. The technique sits in the research lineage that fed into Activation Beacon, LongLoRA, and the broader 2024 long-context retrieval-via-attention thread, but no frontier dense or MoE model ships landmark attention as its long-context primitive.

Lineage

Predecessors: Compressive TransformerCompressive

Cite

BibTeX entry for the original paper

@article{arxiv2305_16300,
  title  = {Landmark Attention: Random-Access Infinite Context Length for Transformers},
  author = {Amirkeivan Mohtashami, Martin Jaggi (EPFL)},
  year   = {2023},
  eprint = {2305.16300},
  archivePrefix = {arXiv},
  url    = {https://arxiv.org/abs/2305.16300}
}

Or cite the paper directly: arXiv:2305.16300.

Export

BibTeX

@article{arxiv_2305_16300,
  title         = {Landmark Attention: Random-Access Infinite Context Length for Transformers},
  author        = {Amirkeivan Mohtashami and Martin Jaggi (EPFL)},
  year          = {2023},
  eprint        = {2305.16300},
  archivePrefix = {arXiv},
  url           = {https://arxiv.org/abs/2305.16300}
}

CSL JSON

{
  "id": "arxiv_2305_16300",
  "type": "article-journal",
  "title": "Landmark Attention: Random-Access Infinite Context Length for Transformers",
  "author": [
    {
      "literal": "Amirkeivan Mohtashami"
    },
    {
      "literal": "Martin Jaggi (EPFL)"
    }
  ],
  "issued": {
    "date-parts": [
      [
        2023
      ]
    ]
  },
  "URL": "https://arxiv.org/abs/2305.16300",
  "number": "2305.16300",
  "source": "arXiv"
}

RIS

TY  - JOUR
TI  - Landmark Attention: Random-Access Infinite Context Length for Transformers
AU  - Amirkeivan Mohtashami
AU  - Martin Jaggi (EPFL)
PY  - 2023
JO  - arXiv
AN  - arXiv:2305.16300
UR  - https://arxiv.org/abs/2305.16300
ER  -