Memorizing Transformers

In plain words

A normal transformer language model only sees a fixed window of recent text. This paper gives one of its attention layers a side store that holds internal representations of text the model has already processed. For each token, the layer looks up the stored entries that most resemble the current query, reads their values, and mixes that result with ordinary attention over the recent text using a learned gate.

The store is filled with what the model computed earlier and is not trained by gradient descent. Because the lookup is a nearest-neighbour search, the store can be large, and the model can reach back to text far outside its window, such as a definition given much earlier in a document. The authors report lower prediction error on long documents of several kinds, and say the same idea can be added to an already trained model and then fine-tuned.

TL;DR

• Adds a large external memory of past keys and values to one layer of a decoder-only transformer §3.1.
• Each query fetches its top-k stored pairs by approximate kNN, attends over them, and mixes the result with local attention through a learned per-head gate §3.1 eq. 1 eq. 2.
• Author-reported perplexity gains on five long-text datasets, growing with memory size up to 65K tokens §4.3.
• Relevant if you want long-range recall without a longer attention window; the memory is not trained by gradients, and the official repository also contains a dense (non-kNN) variant.

Problem

Attention cost grows with the context length, so models are usually trained on short windows. Information beyond the window can only reach the model through weights trained over many steps, and the paper argues that storing facts as weights is slow to learn §1. Earlier long-context methods compress, pool, sparsify or cache past tokens, and some average distant tokens instead of keeping them exact §2. The authors' premise is that exact retrieval of distant keys and values by kNN can scale further than these alternatives §1 §2.

Method

Architecture

The model is a vanilla decoder-only transformer with causal masking §3. Documents are cut into subsequences of 512 tokens. Each subsequence is one training step, and the subsequences of a document are fed in order, not shuffled, as in Transformer-XL §3.

A Transformer-XL style cache keeps keys and values of the previous step and prepends them to the current keys and values. A sliding-window causal mask gives each token a local context of the previous 512 tokens §3.

kNN-augmented attention layer

One layer near the top of the stack (the 9th of 12 in the main experiments §4.2) gets the extra memory. Its local branch is ordinary dense self-attention over the current subsequence. Its memory branch does the following §3.1:

1. Keep a per-head store of the most recent M (key, value) pairs.
2. For every query token, retrieve the top-k stored pairs by approximate nearest-neighbour search. The retrieved set differs for each query.
3. Compute dot products between the query and the retrieved keys, apply softmax over those k scores only, and return the weighted sum of the retrieved values.
4. Combine the memory branch with the local branch through a learned gate eq. 1 eq. 2.

The gate is a scalar bias per head, passed through a sigmoid. Each head can therefore decide how much to rely on memory and how much on local context. The authors observed that most heads ended up attending almost only to memory §3.1. The gate does not depend on token content, though the authors note that a content-dependent gate would be easy to add §3.1.

# kNN-augmented attention layer  [§3.1] [eq. 1] [eq. 2]
# Input:  queries, keys, values of the current subsequence (B, L, H, d)
#         memory: stored (key, value) pairs per head, M entries, no gradients
#         gate_bias: one learned scalar per head (H,)
# Output: layer output (B, L, H, d)

def knn_augmented_attention(queries, keys, values, memory, gate_bias):
    # 1. Dense causal self-attention over the local subsequence  [§3.1]
    local_out = causal_attention(queries, keys, values)   # (B, L, H, d)

    # 2. Per query and per head, fetch the top-k stored pairs  [§3.1] [§3.3]
    #    k = 32 in the main experiments  [§4.2]
    #    The search is approximate, about 90% recall  [§3.3]
    retrieved_keys, retrieved_values = approx_top_k(
        queries, memory.keys, memory.values, k)           # (B, L, H, k, d)

    # 3. Softmax over the k retrieved scores only  [§3.1]
    #    The scale factor on the dot product is not stated in the text
    scores = einsum(queries, retrieved_keys)              # (B, L, H, k)
    weights = softmax(scores, axis=-1)
    memory_out = einsum(weights, retrieved_values)        # (B, L, H, d)

    # 4. Gate memory against local attention, one gate per head  [eq. 1]
    gate = sigmoid(gate_bias)                             # (H,)
    return memory_out * gate + local_out * (1 - gate)     # [eq. 2]

The two gate equations as written in the paper: the gate is the sigmoid of a per-head bias eq. 1, and the combined attention is a gated sum of the memory output and the local output eq. 2 §3.1.

Memory update and batching

After each training step, the keys and values of the current subsequence are appended to the memory. When the memory is full, the oldest entries are dropped §3.1. The memory is not updated by gradients, so the keys and values it returns are whatever the model computed on earlier steps.

# Memory update after one training step  [§3.1]
# Input: memory with capacity M per head, keys and values of the segment

def update_memory(memory, keys, values):
    # Append the segment's (key, value) pairs  [§3.1]
    memory.append(keys, values)
    # Drop the oldest pairs beyond capacity M  [§3.1]
    memory.keep_most_recent(M)

Each sequence in a batch comes from a different document and has its own memory. That memory is cleared when a new document starts §3.

Positions

Local attention uses T5 relative position bias. The retrieved memories get no position bias. The authors cite earlier PG-19 analysis that relative position matters little at long range, and note that T5's bias puts all long-range tokens in one bucket anyway §3.

Staleness and normalisation

The keys in memory were produced by older parameters. The authors call this distribution shift and note that very old entries may be stale §3.2. Keys and queries are normalised to keep old and new keys comparable in magnitude. This does not remove staleness, but the authors found it helped stability with the Transformer-XL cache §3.2.

Training from scratch with a large memory was sometimes worse than pretraining with a small memory (8192) and then fine-tuning with a larger one. The authors suspect staleness, though models could partly adapt their queries to it §3.2.

Approximate kNN

The authors use an approximate search built for TPUs that reaches about 90% recall: about 90% of the true top-k are returned §3.3. They mention that Faiss and ScaNN are alternatives on CPU and GPU §3.3. The paper reports that the approximation quality did not noticeably hurt model quality §4.2.

Hyperparameters

Scaling runs §4.4:

Both scaling runs use no XL cache. Their comparisons are on arXiv only §4.4.

Fine-tuning runs §4.5: for memories of 131K tokens or more, the model is pretrained with memory 8192 or 65K for 500K steps, then fine-tuned for a further 20K steps at the larger memory. Converting a pretrained 1B vanilla model to use a 65K memory is also reported §4.5.

Results

Author-reported. Test-set token-level perplexity, lower is better, after 500K training steps §4.3. Setup: decoder-only models of the size above, trained on each dataset separately §4.1 §4.2.

Source: main results table in §4.3. The 2048-context, no-memory row is the Transformer-XL baseline at that context.

Fine-tuning on arXiv §4.5, test perplexity:

Headline claims from the text:

• Adding memory helps both vanilla and Transformer-XL models on all five datasets §4.3. Larger memory helped more, and 65K was best throughout §4.3.
• A vanilla model with a small 1536-token memory gains almost as much as one with a 2048-token window and no memory, though the memory is not differentiable and sits in one layer only §4.3.
• Scaling: an 8B model is reported only in plots. The text says the 8K-memory Memorizing Transformer matches the perplexity of a vanilla model with five times more parameters §4.4. Figure values are not given in the text.
• Fine-tuning a pretrained 1B vanilla model to a 65K memory closes 85% of the gap to the memory-trained 1B model within 20K steps, and the full gap after 100K steps §4.5.
• Qualitative: the gains concentrate on a small share of tokens, mostly rare names, references and function names §4.6. On Isabelle, the model retrieved the definition of a lemma's referenced object in 8 of 10 manually checked lemma-name predictions §4.6.

Ablations app. A.1, at context 512, XL cache 512, memory 8192:

Two kNN layers gave no further benefit app. A.1.

Paper vs code

Official code: google-research/meliad (Apache-2.0), commit cc1aef50bdcda007e7ea15be9b9a89b95996d17a. Its README says it is "not an officially supported Google product" and provided as-is [code: raw/code/2203_08913/README.md#L3]. The paper's authors work at Google, and the repository is under the google-research organisation; that is the basis for treating it as the official code.

1. Dense attention, not kNN, in the memory2 path. The paper's method is top-k kNN over memory §3.1. The memory2 implementation that base_memory2.gin selects attends densely over the full memory [code: transformer/memory2/dense_kv_memory.py#L199-L229] with a memory of 65536 entries base_memory2.gin [code: transformer/configs/base_memory2.gin#L114]. The README calls it a re-implementation "with dense rather than Top-K attention" [code: README.md#L127-L128]. The top-k path is transformer/memory/ [code: transformer/memory/memory_layer.py#L268-L317]. Do not expect the paper's numbers from the dense path.
2. Neighbours. The paper uses k = 32 in all experiments §4.2. The 8K memory config sets memory_num_neighbors = 128 [code: transformer/configs/memory/external_memory_8k.gin#L13]. The paper's ablation shows 32 and 128 give similar perplexity app. A.1.
3. Memory layer index. The paper uses the 9th of 12 layers §4.2. The 8K config uses memory_layer_indices = (-4,) [code: transformer/configs/memory/external_memory_8k.gin#L10], which is the 9th layer counting from 1. base_memory2.gin sets index 6 (the 7th layer) base_memory2.gin. The defaults disagree, so set the index explicitly.
4. Dimensions. base_memory2.gin matches the paper's 1024 embedding, 8 heads of 128, FFN 4096 base_memory2.gin. base_memory.gin uses 512 embedding, FFN 2048 and 512-token segments [code: transformer/configs/base_memory.gin#L21-L25].
5. Learning-rate schedule. The paper says linear warmup then square-root decay §4.2. The shared trainer config selects a cosine schedule with max_lr = 0.01 and min_lr = 0.001 [code: transformer/configs/trainer_configuration.gin#L36-L38] [code: transformer/configs/trainer_configuration.gin#L51-L52], with Adafactor learning_rate = 1.0 trainer_configuration.gin. A square-root option exists optimizer_config.py but is not the default. Which schedule produced the paper's numbers is not confirmed by the code.
6. Gate. Matches the paper: one scalar per head, passed through a sigmoid, combined with local output as a gated sum [code: transformer/memory/transformer_layer.py#L253-L254] [code: transformer/memory/transformer_layer.py#L562-L567]. The code initialises the bias to zero, so the gate starts at 0.5. The paper does not state an initial value. The default combine mode is TRAINABLE_WEIGHTED_MEAN [code: transformer/memory/transformer_layer.py#L83].
7. Memory reset per document. Matches the paper: the memory is cleared at the start of a new document dense_kv_memory.py and memory_reset_on_new_doc defaults to true [code: transformer/memory/transformer_layer.py#L82].
8. Write-back. The paper appends the segment after each step and drops the oldest entries §3.1. The dense code writes into a ring buffer indexed modulo the memory size [code: transformer/memory2/dense_kv_memory.py#L158-L197].
9. Approximate top-k. The paper gives about 90% recall §3.3. The code's TPU search is a bucketed approximation with a soft selection step [code: transformer/memory/memory_layer.py#L268-L288]. It does not state a recall figure.
10. Window and memory. The README warns that memory should not be used with a sliding window, so the window and segment should match [code: README.md#L58-L59]. base_memory2.gin sets a 512-token window inside a 4096-token segment base_memory2.gin. The paper uses a 512-token local subsequence §3. Check this when reproducing.
11. Datasets not released. The code has no data. C4 and PG-19 are public. Github, Isabelle and arXiv Math were built by the authors from public sources §4.1 [Reproducibility statement].

Pitfalls

• Use the kNN path. The dense memory2 path in the same repo is not the paper's method dense_kv_memory.py.
• Memory size must be a multiple of the segment length. The dense code rejects other sizes dense_kv_memory.py.
• Stored keys go stale. Entries come from earlier parameters, and large memories trained from scratch can be unstable. The paper's workaround is to pretrain small and fine-tune large §3.2 §4.5. Key normalisation helps but does not remove it.
• Not differentiable. Gradients reach the memory only indirectly, through the current segment's attention dense_kv_memory.py. Keys and values are reused from prior steps, which is what keeps the cost low §1.
• Unstated values. Attention scale factor, dropout, weight decay, the gate's initial value and fine-tune warmup are not stated in the paper. The code's choices are listed in the table and in Paper vs code.
• Cost. At about 200M parameters on TPUv3, the step time triples with a 65K memory (0.2 s to 0.6 s) §4.2. The paper calls the cost manageable on one device for 131K to 262K tokens, but gives no multi-device figures §1.
• Retrieval depends on approximation. The paper reports no quality loss from lower-quality kNN, but gives no sweep §4.2.
• Comparisons confound memory with XL cache. Most memory rows also have an XL cache, so compare like with like §4.3.
• Author-reported. Perplexities come from the authors' own runs. The gains are large relative to the seed spread reported in app. A.1. The Isabelle definition-retrieval result is a manual check of 10 examples §4.6.
• Floats not verified against the PDF. Table and figure numbers were not checked against the compiled PDF, so the card cites sections instead.

Lineage

• Builds on: Transformer-XL, which caches the previous step's keys and values and prepends them to the current segment §2. This paper keeps that cache and adds the kNN memory on top.
• Related: Memory-efficient Transformers, which replace dense attention with a kNN lookup for speed and memory §2; Expire-span, which prunes far-away tokens §2; Feedback and compressive transformers §2; a non-differentiable (key, token) table that replaces the output softmax §2.
• Superseded by: not stated in the paper or code.