Can Fewer Attention Heads Make Language Models Forget More?

Large language models can learn broad linguistic patterns, but they can also reproduce passages from their training data word for word. When those passages contain personal information, copyrighted material, or other sensitive content, memorization becomes a privacy and intellectual-property concern. Existing responses often focus on training-data filtering, retraining, or output-level safeguards. A new arXiv paper examines a lighter intervention: disabling selected attention heads before inference.

What the study found

• Attention heads may act as memory components. Prior theoretical work has proposed that attention heads can function as associative memory mechanisms, with potential memorization capacity increasing with the number of heads.
• The pruning is targeted, not random. The authors identify heads that are redundant for a language-modeling task and zero them out at inference time rather than removing parameters through retraining.
• Verbatim reproduction drops sharply. In GPT-2-XL, which has 1.5 billion parameters, pruning 729 task-redundant heads reduced the count of known memorized sequences reproduced by the model by 91.5%.
• The benchmark score did not deteriorate. On WikiText-103, the benchmark used to identify redundant heads, perplexity improved by 0.92. The supplied material does not establish whether the same effect holds across other tasks or datasets.

The important claim is not that removing attention heads universally makes a model safer. Rather, the result suggests that some exact-recall behavior may depend on a subset of heads that are not essential for the evaluated language-modeling task. If those heads help retrieve training sequences while contributing little to ordinary prediction, selectively suppressing them could reduce one form of memorization while preserving much of the model’s normal operation.

Why it matters—and what remains open

For deployment teams, inference-time pruning is attractive because it does not require changing the training corpus or retraining the model. It also broadens the role of pruning: beyond compression and efficiency, structural interventions could become a way to modulate undesirable behaviors after training.

The findings nevertheless need careful qualification. The experiment centers on GPT-2-XL, known memorized sequences, and WikiText-103. The abstract does not provide results for larger models, other architectures, unseen memorization cases, or a broad set of downstream capabilities. It also does not show whether disabling the heads affects factual recall, long-context behavior, or other functions. Accordingly, the method should be viewed as a promising inference-time mitigation technique, not a complete solution to model-memory governance. Broader evaluations will be needed to determine whether the reduction in verbatim reproduction is robust and whether it introduces meaningful trade-offs.

Source: arXiv