Disabling 729 Attention Heads in GPT-2-XL Cut Verbatim Memorization by 91.5%

Language models can reproduce chunks of their training data verbatim, which risks leaking personal information and triggering copyright claims. Researchers have shown that part of this problem can be fixed without fine-tuning: it is enough to zero out "redundant" attention heads before inference. In GPT-2-XL (1.5B parameters), disabling 729 such heads cut the number of verbatim-reproducible sequences by 91.5%, and perplexity on WikiText-103 improved by 0.92.

Which heads were removed and how they were found

The idea rests on two observations. Theoretical work suggests that multi-head attention acts as an associative memory whose capacity grows with the number of heads. Empirical work has long shown that a significant share of heads can be switched off with almost no loss in quality. Nobody had tried to connect the two.

The authors freeze all weights and train only binary masks on the head outputs. The mask is a step function, so gradients pass through a sigmoid approximation (a straight-through estimator). A penalty in the loss function keeps perplexity from rising above a set threshold and at the same time rewards fewer active heads.

The setup uses the AdamW optimizer, batches of 8 sequences and 1000 steps. The final configuration is not the one taken; instead the authors pick the one with the lowest ratio of perplexity to the original. The masks are fitted on a separate part of WikiText-103 or LAMBADA (a dataset where the model guesses the last word of a paragraph) that does not overlap with the evaluation set.

Memorization was measured through k-extractability: if the model is given the beginning of a sequence, it produces the continuation word for word under greedy decoding. The tests ran on GPT-2-XL with the OpenWebText corpus, and on Pythia-2.8B and Pythia-6.9B (both Deduped) with the Pile-Deduped corpus.

The results in numbers

On GPT-2-XL, masks fitted on WikiText-103 removed about 61% of heads. Perplexity fell from 10.08 to 9.16, and in one configuration the number of memorized sequences dropped from 5590 to 315. Both Pythia models show the same picture: perplexity went from 7.56 to 7.32 and from 6.90 to 6.73, and memorization fell from 20 000 to 2875 and 3500 sequences. Across all models and both benchmarks, the reduction ranged from 44 to 94%.

Masks fitted on LAMBADA disable fewer heads, and memorization falls less, but last-word prediction accuracy holds steady or rises slightly. The sets of removed heads differ across the two benchmarks and three models, yet the effect appears everywhere. The authors conclude that memorization is distributed across the whole attention mechanism rather than sitting in a narrow group of heads.

A control experiment shows that it matters not only how many heads are disabled but also which ones. When the same number of heads was disabled at random, memorization sometimes dropped even further (330 versus 473 sequences for GPT-2-XL in one configuration), but the model fell apart. GPT-2-XL's perplexity on WikiText-103 rose from 10.08 to 146.82, and Pythia's rose from 7.56 to 18.30 and from 6.90 to 17.20. GPT-2-XL's accuracy on LAMBADA fell from 0.6128 to 0.2014. A broken model simply reproduces nothing, so the memorization count alone is not enough to evaluate the method.

Why it works and where the weak spots are

The mechanistic analysis was done on Pythia-2.8B and GPT-2-XL. On memorized sequences, heads more often "saturate" the softmax: nearly all attention goes to a single position. In addition, different heads more often point to the same position. After ablation, this gap between memorized and non-memorized text shrinks. The logit gap between the top token and the runner-up on memorized sequences was 8.79 and 7.68, roughly four times higher than on ordinary text. After ablation it fell to 3.60 and 3.86, while on non-memorized text it barely changed. This suggests the intervention hits the mechanism of excess confidence rather than the model's confidence in general.

The authors name the limitations themselves. The mechanistic part is correlational: the data show that saturation, head agreement and confidence change together, but causation is not proven. The agreement comparison is made over the pool of saturated heads, which shrinks after ablation, so part of the effect may be an artifact. Why the remaining heads look "wider" has not been tested either. The analysis does not cover Pythia-6.9B because of a lack of compute.

There is also a methodological limitation. Quality was measured on the same benchmark used to find the redundant heads, albeit on a separate sample. How the pruned model behaves on other tasks has not been tested in the paper. The experiments are limited to three small models from 2020–2023, and the result may differ on modern LLMs with grouped attention. The authors promise to release the code after the paper is accepted.

The takeaway for practitioners: teams that run a model in production and worry about training-data leaks now have a cheap tool that needs no retraining. But it lowers the probability of verbatim reproduction rather than guaranteeing privacy, and the model should be tested on your own tasks before you apply it.