Memorizing Transformers

§ 1 · Premise

The 8K ceiling and the routes around it

A 2022-era Transformer decoder at 8K context and 1024-dim heads carries $2 \cdot 8192 \cdot 1024 = 16$M scalar entries per layer in its KV cache — manageable at $\sim 32$ MB per layer fp16 but quadratic in attention compute: $T^2 d_h = 8.6 \cdot 10^{10}$ multiply-adds per layer per head. Doubling $T$ quadruples it. By the time Wu et al. wrote Memorizing Transformers (arXiv 2203.08913) in early 2022, two routes around the ceiling had been published:

• Lengthening attention itself. Transformer-XL recurrence (Dai et al. 2019) and the sparse / kernel families (Reformer, Longformer, Performer) modify the attention computation so $T^2$ becomes $T \log T$ or $T \cdot d$.
• Lossy summarization of the past. Compressive Transformer (Rae et al. 2019) pools or attention-summarizes evicted activations into a coarser long-term memory.

Both modify the model. Wu et al.’s wager is that the simpler primitive is retrieval: keep the regular self-attention mechanism untouched, but at one chosen layer also let each query look up the $k$ nearest neighbours of its own vector inside a vast external store of past keys/values produced by that same layer. Retrieved $(K, V)$ pairs are spliced into the local attention through a learned gate. The store can be billions of entries; it is not backpropagated through. The contribution is the demonstration that a single mid-stack non-differentiable kNN attention substantially improves language-model perplexity at fixed compute, generalises across domains, and scales sub-linearly with memory size.

§ 2 · Derivation

Attention extended by a kNN-retrieved set

Starting point — single-layer attention. Let $\mathbf{q}_t, \mathbf{k}_t, \mathbf{v}_t \in \mathbb{R}^{d_h}$ be the query, key, value at position $t$ for one attention head. Standard causal attention over a local window of length $W$ is

In Wu et al.’s setup $W$ matches the Transformer-XL recurrent context size (1024 or 2048 tokens, §3). Anything older than $W$ has been evicted.

Step 1 — the external memory. At one chosen layer $\ell^\star$, write every $(\mathbf{k}, \mathbf{v})$ pair produced by that layer into an external memory $\mathcal{M}$. The store is indexed by the keys: the index supports approximate-nearest-neighbour (ANN) search using inner product as the similarity. Wu et al. implement this with ScaNN (Guo et al. 2020), which gives sub-linear lookup time even at $|\mathcal{M}| = 10^9$ entries (§3.1). Memory is per attention head per layer but only at the single layer $\ell^\star$ — not the whole stack.

Step 2 — retrieval at inference and training. At every query position $t$ in layer $\ell^\star$, run a top-$k$ search on the external memory:

Wu et al. use $k = 32$ or $k = 64$ (Table 1). The returned set $\{(\mathbf{k}^{(m)}, \mathbf{v}^{(m)})\}_{m \in \mathcal{N}_k(\mathbf{q}_t)}$ is treated as constants — no gradient flows back into $\mathcal{M}$. From these, compute an external attention output exactly as standard attention:

The external softmax is independent of the local softmax — two separate normalizations.

Step 3 — the learned gate. The two outputs $\mathbf{o}^{\text{local}}_t$ and $\mathbf{o}^{\text{ext}}_t$ are combined via a per-head learnable scalar gate $g_h \in (0, 1)$ (Eq. 1, §3.2):

where $\sigma$ is the sigmoid and $g_h$ is a single scalar trained alongside everything else. Wu et al. choose the gate-then-mix form (rather than concatenating $K^{\text{local}}$ and $K^{\text{ext}}$ into one softmax) for two reasons:

1. Scale incompatibility. External keys come from arbitrary past contexts and may have very different norms from local keys; sharing a softmax would let one set dominate the other in a way that’s hard to control.
2. Training-time stability. The gate starts near $\sigma(g_h) \approx 0.05$ so external attention is off-by-default; the model learns to open the gate only when retrieval helps. Sharing a softmax would couple the gradient signal of “did retrieval help?” to “did the right key happen to be retrieved?” — Wu et al. report training instabilities under the shared form (§3.2 footnote).

Step 4 — non-differentiability and what it costs. The discrete top-$k$ step in $\mathcal{N}_k(\mathbf{q}_t)$ is non-differentiable. Gradients flow:

• Into the retrieved $\mathbf{k}^{(m)}, \mathbf{v}^{(m)}$ used at this query position through the external attention softmax — these values are constants in the graph, so the gradient terminates there.
• Through the query $\mathbf{q}_t$, which means the model learns to issue better queries but cannot learn to retrieve better keys. The memory is a fixed function of the model’s past forward passes; whatever the model deposited in the past is what later queries see.
• Through the gate $g_h$ via the mixing equation, which means the model learns when to trust retrieval.

The trade is explicit: cheap memory (no autograd buffers), unbounded in size, at the price of losing the ability to learn the indexing function itself.

Step 5 — memory layer placement. Why only one layer? Two reasons. Memory cost is $O(|\mathcal{M}| \cdot 2 d_h)$ scalars; replicating across all $L$ layers multiplies that by $L$ for negligible benefit (the layer-wise ablation in Table 5 shows the 9th layer of 12 captures most of the gain). And the model only needs to learn one gate; multiple memory layers would compete for the same “open the gate when retrieval is useful” signal during early training.

Cost accounting. Per query in the memory layer, with local window $W$ and retrieval budget $k$:

With $W = 1024, k = 32$ the dense attention overhead is $\sim 3\%$ on top of the local-window attention. The ANN lookup cost depends on the index but scales as $O(\sqrt{|\mathcal{M}|})$ or better for ScaNN. Memory footprint: $|\mathcal{M}| = 65$K entries at $d_h = 128$ in fp32 is $\sim 64$ MB per head per layer — cheap; $|\mathcal{M}| = 262$K is $\sim 256$ MB, etc.

Parameter count. One learnable scalar gate per head per memory layer; negligible vs. the base model.

§ 3 · Reference implementation

kNN-augmented attention layer in pseudocode

# State:
#   memory_index: an ANN store keyed by past K^(l*), holding (K, V) at the memory layer
#   gate: learnable scalar per head, init so sigmoid(gate) ~= 0.05

def memorizing_attn(q, k_local, v_local, memory_index, gate, k_retrieve=32):
    # q, k_local, v_local: [B, T, H, d_h]
    # Local attention over the in-window window
    o_local = scaled_dot_product_attention(q, k_local, v_local)        # [B, T, H, d_h]

    # External retrieval: top-k nearest keys for each query (no gradient through index)
    with torch.no_grad():
        retrieved_k, retrieved_v = memory_index.ann_search(q, k_retrieve)
        # retrieved_*: [B, T, H, k_retrieve, d_h], treated as constants

    # External softmax over only the retrieved set, independent normalization
    scores = einsum("bthd,bthkd->bthk", q, retrieved_k) / sqrt(d_h)
    weights = softmax(scores, dim=-1)
    o_ext = einsum("bthk,bthkd->bthd", weights, retrieved_v)

    g = sigmoid(gate)[None, None, :, None]                              # [1, 1, H, 1]
    o = (1 - g) * o_local + g * o_ext

    # Add this chunk's local (K, V) to the memory for future queries
    memory_index.add(k_local.detach(), v_local.detach())
    return o

The sketch elides three production concerns: (1) the ANN index rebuild cadence — Wu et al. periodically re-fit the ScaNN partitioning every $N$ tokens written; (2) per-document memory isolation (each document’s memory is private to avoid cross-document leakage at training time, §3.3); (3) the rotary or relative-position handling for retrieved keys, which lacks a natural “distance” to the current query — Wu et al. simply omit position encoding for the retrieved set.

§ 4 · Empirical evidence

What is and isn’t known

Introducing paper (Wu et al. 2022).

1. Perplexity gains across domains. On four long-document datasets — arXiv math, GitHub code, PG-19 books, C4 web — a 200M-parameter Memorizing Transformer with $|\mathcal{M}| = 65{,}536$ entries lowers test perplexity by 1.5–2.7 nats vs. a matched non-memory baseline (Table 1). The largest gains come on arXiv math, where citation and notation reuse make retrieval particularly useful.
2. Memory-size scaling. The same architecture’s perplexity continues to drop as $|\mathcal{M}|$ grows from $1{,}536$ to $262{,}144$ entries (Figure 1) — i.e., the model genuinely uses the bigger memory and is not just absorbing a fixed amount of recent context.
3. Parameter-count tradeoff. At fixed compute, adding kNN memory matches the gain from roughly $5\times$ more parameters; e.g., a 200M memory model rivals an 800M no-memory baseline on arXiv (Table 2).
4. Layer-placement ablation. The 9th layer of 12 gives the largest gain; adding memory to every layer adds at most 0.1 nat on top (Table 5).
5. Generalization across model sizes. The technique was retested at 1B parameters with a $\sim 1$ nat perplexity gain that holds at scale (§4.2).

Independent follow-up.

• Unlimiformer (Bertsch et al. 2023, arXiv 2305.01625) generalizes the memory layer’s kNN attention to encoder-decoder summarization, reporting consistent ROUGE improvements at the cost of zero added parameters — direct architectural descendant.
• Landmark Attention (Mohtashami & Jaggi 2023, arXiv 2305.16300, §2) explicitly positions itself against Memorizing Transformers’ non-differentiability: “the retrieval is not learned, so the model cannot improve indexing through training.”
• RPT / Retrieval-Pretrained Transformer (Rubin & Berant 2023, arXiv 2306.13421) closes the loop the Wu et al. paper leaves open — they make the retrieval differentiable via a Gumbel-softmax over a small candidate set and report further perplexity gains. The headline finding: Wu et al.’s results are a lower bound on what kNN memory can do once the indexing is also trained.
• Survey coverage. Routinely cited as the canonical “external-memory” baseline in long-context surveys (Pawar et al. 2024, arXiv 2402.02244, §3.2).

Sensitivity studies — what is not publicly known. The introducing paper sweeps memory size and layer placement but not (a) the gate’s initialization or the gate-vs.-shared-softmax tradeoff in detail, (b) the choice of inner product vs. cosine vs. learned similarity inside the ANN, or (c) how retrieval quality interacts with RoPE / ALiBi position encodings (the paper uses T5-style relative bias). I don’t know of an independent reproduction at scales $> 1$B parameters; the published results stop there.

Production adoption. None recorded in this knowledge base. Memorizing Transformers’ practical legacy is the broader retrieval-augmented family (RAG-style retrieval into the prompt, KV-cache compression with retrieval gates) rather than the specific “kNN-from-attention-keys at one layer” mechanism, which was largely superseded by RAG architectures that retrieve documents into the input rather than $(K, V)$ pairs into the attention.

@article{arxiv2203_08913,
  title  = {Memorizing Transformers},
  author = {Yuhuai Wu, Markus N. Rabe, DeLesley Hutchins, Christian Szegedy},
  year   = {2022},
  eprint = {2203.08913},
  archivePrefix = {arXiv},
  url    = {https://arxiv.org/abs/2203.08913}
}

@article{arxiv_2203_08913,
  title         = {Memorizing Transformers},
  author        = {Yuhuai Wu and Markus N. Rabe and DeLesley Hutchins and Christian Szegedy},
  year          = {2022},
  eprint        = {2203.08913},
  archivePrefix = {arXiv},
  url           = {https://arxiv.org/abs/2203.08913}
}

{
  "id": "arxiv_2203_08913",
  "type": "article-journal",
  "title": "Memorizing Transformers",
  "author": [
    {
      "literal": "Yuhuai Wu"
    },
    {
      "literal": "Markus N. Rabe"
    },
    {
      "literal": "DeLesley Hutchins"
    },
    {
      "literal": "Christian Szegedy"
    }
  ],
  "issued": {
    "date-parts": [
      [
        2022
      ]
    ]
  },
  "URL": "https://arxiv.org/abs/2203.08913",
  "number": "2203.08913",
  "source": "arXiv"
}

TY  - JOUR
TI  - Memorizing Transformers
AU  - Yuhuai Wu
AU  - Markus N. Rabe
AU  - DeLesley Hutchins
AU  - Christian Szegedy
PY  - 2022
JO  - arXiv
AN  - arXiv:2203.08913
UR  - https://arxiv.org/abs/2203.08913
ER  -