Localizing Paragraph Memorization in Language Models (2403.19851v1)

Abstract: Can we localize the weights and mechanisms used by a LLM 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.

• The paper reveals that specific tokens in paragraph prefixes act as memorization triggers, leading to exact retrievals from training data.
• It employs gradient-based parameter attribution on GPT-Neo 125M to localize key model components, notably identifying attention head 2 in layer 1.
• The study demonstrates that targeted sparse unlearning can modify memorized outputs, offering actionable insights into model control and privacy enhancements.

Localizing and Understanding Paragraph Memorization in LLMs

Introduction to Paragraph Memorization

In the research domain of LLMs, understanding the phenomenon of paragraph memorization is essential for advancing our knowledge on how these models store and retrieve long sequences of text from their training data. A significant challenge is localizing the model components responsible for memorizing entire paragraphs, as such memorization may have implications for both model performance and privacy. This blog post summarizes recent findings on identifying and characterizing the model internals that contribute to paragraph memorization, focusing on an in-depth analysis of the GPT-Neo 125M model trained on the publicly available Pile dataset.

Paragraph Memorization: Definition and Metrics

The paper defines paragraph memorization as the ability of a LLM to produce an exact continuation of a given prefix from its training set. Two key metrics are employed to evaluate memorization:

• Exact Match (EM): The count of tokens in the model-generated continuation that exactly match the true continuation, up to a maximum of 50 tokens.
• Negative Log-Likelihood (NLL): Measures how likely the model considers the true continuation, with lower values indicating higher likelihoods.

Based on these metrics, paragraphs are categorized into memorized and non-memorized sets, facilitating a comparative analysis of the model's behavior across these two categories.

Discovering Memorization Triggers in Prefixes

An intriguing finding is that specific tokens within the prefix of memorized paragraphs can significantly affect the model's generation, acting as "memorization triggers." By perturbing these tokens, the model's output diverges from the memorized continuation, often resulting in equally plausible but non-memorized alternatives. This suggests that memorization may be linked to distinctive, rare tokens within the prefix that act as key anchors for retrieving the stored information.

Gradient-based Localization of Memorization Components

To further understand the internal mechanisms behind memorization, the paper employs gradient-based parameter attribution. By analyzing the gradients of memorized and non-memorized paragraphs, it was observed that gradients tend to be larger in lower layers for memorized paragraphs and in higher layers for non-memorized paragraphs. Specifically, attention to details revealed that the attention head 2 in layer 1 of GPT-Neo 125M shows distinct gradient patterns associated with memorization, suggesting its significant role in processing rare or unique tokens within paragraph prefixes.

Sparse Unlearning and Editing of Memorization

The localized understanding of memorization components allows for targeted interventions, such as sparse unlearning and editing of memorized paragraphs. By fine-tuning only the most relevant weights identified through gradient attribution, the model can be effectively "unlearned" or edited to produce different continuations for previously memorized paragraphs. This supports the hypothesis that localized model components significantly contribute to paragraph memorization.

Implications and Future Directions

The identification of specific model components, particularly the attention head in layer 1, as crucial to paragraph memorization has several implications. It opens up avenues for further research into model interpretability, as understanding the role of individual model components can lead to more explainable AI. Additionally, the methods developed for localizing memorization components and manipulating memorized content offer promising approaches for addressing privacy concerns related to unintended memorization in LLMs.

In conclusion, the paper provides valuable insights into the mechanisms of paragraph memorization in LLMs, highlighting the potential for targeted modifications to control memorization behavior. As the field of LLM research advances, such findings will be critical in developing models that are both powerful and aligned with ethical considerations regarding data privacy and usage.