Localizing Paragraph Memorization in Language Models
This work addresses the issue of data memorization in language models for AI safety and interpretability, though it is incremental in localizing specific mechanisms.
The paper tackled the problem of localizing memorization mechanisms in language models, finding that gradients of memorized paragraphs are larger in lower layers and identifying a specific attention head focused on rare tokens, with memorized continuations being harder to unlearn or corrupt.
Can we localize the weights and mechanisms used by a language model to memorize and recite entire paragraphs of its training data? In this paper, we show that while memorization is spread across multiple layers and model components, gradients of memorized paragraphs have a distinguishable spatial pattern, being larger in lower model layers than gradients of non-memorized examples. Moreover, the memorized examples can be unlearned by fine-tuning only the high-gradient weights. We localize a low-layer attention head that appears to be especially involved in paragraph memorization. This head is predominantly focusing its attention on distinctive, rare tokens that are least frequent in a corpus-level unigram distribution. Next, we study how localized memorization is across the tokens in the prefix by perturbing tokens and measuring the caused change in the decoding. A few distinctive tokens early in a prefix can often corrupt the entire continuation. Overall, memorized continuations are not only harder to unlearn, but also to corrupt than non-memorized ones.