Pruning redundant attention heads cuts AI memorisation by over 90%

A new study finds that deleting a large share of a language model’s attention heads (components that decide which earlier words to focus on) before it generates text can sharply cut how often the model reproduces memorised passages from its training data. The paper, titled “Fewer Heads, Fewer Memories,” was written by Aryan Chowdhury and Riza Batista-Navarro and posted to arXiv’s computation-and-language category.

The authors built on two separate lines of earlier research. One strand of theory holds that attention heads act as associative memory, storing and retrieving specific sequences, so a model’s capacity to memorise text scales with how many heads it has. A second strand, from empirical pruning studies, found that many attention heads are redundant and can be removed without hurting a model’s scores on standard benchmarks. According to the paper, no prior work had tested what happens to memorised-sequence reproduction when those two findings are combined.

To test it, the researchers used GPT-2-XL, a 1.5-billion-parameter language model, and checked it against sequences already known to be memorised verbatim from its training set. They identified 729 attention heads that a benchmark task had classified as redundant and zeroed them out, meaning each head’s output was set to contribute nothing, without retraining the model or altering its training data.

With those heads removed, the model reproduced memorised sequences 91.5% less often, the paper reports. At the same time, its perplexity, a standard measure of how well a model predicts text where lower is better, improved by 0.92 points on the WikiText-103 benchmark. The pruned model became both less likely to repeat memorised passages and slightly better at general language prediction, according to the authors’ results.

The finding points to a low-cost option for a problem that has worried researchers and companies alike. Prior work has shown large language models can leak training data verbatim, raising concerns about exposing private information or copyrighted text buried in a training set. Most proposed fixes require retraining a model or filtering its training data in advance, both costly once a model is already deployed. Pruning redundant heads, by contrast, is applied after training is finished, so the authors say it could be adopted by practitioners without rebuilding a model.

The results come from one model family tested on one set of known memorised sequences and one benchmark, and the paper has not completed peer review; arXiv lists its DOI registration as pending. It is not yet established whether the same pattern of head redundancy, or the same scale of memorisation suppression, holds for newer or larger models, or for memorised text outside the sequences the authors checked.

What to watch next is whether the authors or other labs repeat the test on more recent model architectures and publish code for identifying which heads are redundant, which would let other practitioners apply the technique beyond the paper’s own experiments on GPT-2-XL.