Do Localization Methods Actually Localize Memorized Data in LLM s? A Tale of Two Benchmarks (2024)

Abstract: The concept of localization in LLMs is often mentioned in prior work; however, methods for localization have never been systematically and directly evaluated. We propose two complementary benchmarks that evaluate the ability of localization methods to pinpoint LLM components responsible for memorized data. In our INJ benchmark, we actively inject a piece of new information into a small subset of LLM weights, enabling us to directly evaluate whether localization methods can identify these “ground truth” weights. In our DEL benchmark, we evaluate localization by measuring how much dropping out identified neurons deletes a memorized pretrained sequence. Despite their different perspectives, our two benchmarks yield consistent rankings of five localization methods. Methods adapted from network pruning perform well on both benchmarks, and all evaluated methods show promising localization ability. On the other hand, even successful methods identify neurons that are not specific to a single memorized sequence.

## introduction large language models (llms) memorize many sequences from their pretraining corpora @xcite @xcite . for example, @xcite show that gpt2 @xcite can leak some private contact information verbatim. this paper studies whether we can localize a piece of memorized data, i.e., identify components in llms responsible for generating a sequence (near) verbatim. successful localization may inform further work in machine unlearning @xcite ; for instance, one could apply "neural surgery" to the located components to make the llm forget a piece of sensitive information. prior work on knowledge editing suggests that we can locate a small set of llm parameters that store factual knowledge @xcite . these works demonstrate localization success by showing knowledge editing success when updating only the located llm parameters. however, @xcite argue that editing success and localization are actually uncorrelated. similarly, prior methods that identify subnetworks in llms @xcite usually focus on the performance of downstream classification tasks, lacking direct evaluation on localization per se. hence, the degree of existing methods' localization success remains unclear. this paper studies the open question, "do localization methods actually localize memorized data in llms?" we first propose decoupling localization success from downstream success in our inj benchmark. our key insight is to actively create the ground-truth weights responsible for data memorization. specifically, we force llms to use a small set of pre-decided weights to memorize a piece of new information unseen during pretraining. therefore, we have the ground-truth locations where the new information is injected. we can then directly evaluate how well different localization methods recall the indices of the injected weights. we further apply the localization methods to a real-world scenario: identifying a small set of neurons in an llm responsible for memorizing a pretrained sequence. in this setting, evaluating localization success is more challenging because the ground-truth "location" of each memorized sequence is unknown. we propose the del benchmark, inspired by knockouts @xcite , a reverse-engineering approach that removes a set of nodes from the computation graph to observe their importance for specific model behavior. we first collect a set of memorized sequences, and for each sequence, we drop out the located neurons to measure their importance to memorizing that target sequence. a successful localization should cleanly erase the target sequence from an llm without hurting the memorization of the other sequences in the set after dropout. our two benchmarks complement each other: the inj benchmark provides a direct evaluation of localization methods under a well-controlled setup, while del benchmark answers if the methods can localize pretrained sequences that llms have already memorized. we systematically evaluate five methods on our two benchmarks, including existing localization methods @xcite , a brute-force method that searches for the most important neurons (zero-out), and two methods we adapt from network pruning @xcite , slimming and hard concrete. our two benchmarks rank the five methods in the same order, showing especially strong localization ability for hard concrete. for example, dropping out only 0.5% of neurons in pythia-6.9b @xcite identified by hard concrete makes the model forget 57.7% of the target memorized tokens on average. on the other hand, the del benchmark shows all methods struggle to balance between erasing the target sequence and retaining other memorized data, indicating that the identified neurons are also relevant for memorizing some other sequences. overall, both benchmarks agree all evaluated localization methods are promising, but precise localization of a single sequence remains difficult. ## background and task terminology a transformer layer @xcite consists of multi-head self-attention and a feed-forward network @xcite . prior work shows that llms use their ffns rather than self-attention as "memories" to store knowledge @xcite @xcite . here, an ffn has two fully connected layers with a non-linear activation function σ: they show that different concepts are stored in different v l i , and that we can view each activation h l i as a memory coefficient to retrieve a concept. neurons. @xcite observe the existence of knowledge neurons, a small set of neurons in ffn hidden states h l that corresponds to a relational fact, where a neuron means a component of the vector h l . for example, given the input "the capital of ireland is ", they can increase the model probability on the correct token "dublin" by amplifying the activation h l i of the identified knowledge neurons. with eq. 3, we can view increasing activation h l i as promoting the concept stored in v l i . in this work, we only search for neurons in ffns responsible for memorizing a sequence, following @xcite . in the inj benchmark, we ensure that ffns act as neural memories by only updating a set of weight vectors v l i to memorize the new information. as each v l i corresponds to a neuron in h l , locating the updated weights is equivalent to locating the corresponding neurons. in the rest of the paper, we refer to neurons as the neurons in {h l } l @xmath0, where l is the number of layers. dropout. different from @xcite , we drop out located neurons at test time to erase a memorized sequence from the llm. we can view dropping out the i-th neuron in h l as excluding the contribution of v l i from the output o l in eq. 3. memorized sequences. consider a sequence @xmath1, s) that consists of a prefix p and a suffix s. given the prefix as the prompt, if an llm is able to nearly reconstruct the suffix with greedy decoding, we say x is a memorized sequence by the llm. we discuss in §3.2 our criteria on suffix reconstruction, where we tolerate near-verbatim memorization; we also ensure every sequence has a non-trivial suffix. localization. @xcite provides a general definition of localization: identifying components of a model responsible for a certain behavior. under this definition, we consider components as a small set of neurons and behavior as the llm's generation of a memorized sequence. although some components are necessary for generation, e.g., the input and output token embeddings, we exclude them as they are not specific to a target sequence. localization methods. given an llm, a memorized sequence x, and a fixed number k, a localization method outputs k% of neurons at each layer as the predictions to localize sequence x in the llm. ## two localization benchmarks how do we know whether a method is successful in localization? we propose two benchmarking approaches: one injects a new piece of information into specific parameters in llms, while another deletes an existing memorized sequence from llms via dropout. a successful localization method should do well on both benchmarks. ## localization methods we benchmark five localization methods. each method assigns an attribution score a l (i) to each neuron n l i , the i-th neuron in the l-th layer, representing its importance in memorizing a sequence x. at test time, we select the top-k% of neurons in each layer for each method in terms of attribution scores as the located neurons for x by that method. several methods involve calculating the language modeling loss of an llm θ on the suffix of the target sequence @xmath2, s). we denote the loss as memorization loss, ℓ mem θ (x). formally, zero-out. we introduce an exhaustive method that drops out neurons one by one and uses the resulting change in memorization loss on x as the 2 we do not drop out neurons in the bottommost layer, as it hurts llms' overall memorization indiscriminately ( §5.4). we denote ℓ mem θ\n l i as the memorization loss of the llm θ after dropping out a neuron n l i . the larger the change in the loss, the more important the neuron is for memorization. zero-out is closely related to the occlusion-based attribution method (zeiler and fergus, 2014). ## related work and discussion localization identifies function-specific components, including neurons @xcite , layers @xcite , or subnetworks @xcite @xcite . for example, dai et al. (2022) find knowledge neurons for each relational fact. @xcite locate relational facts to middle ffns, specifically when llms process the last token of the subject. @xcite discover sparse knowledge subnetworks in gpt2 with a differentiable weight masking method. however, there is no standard approach to evaluate the effectiveness of localization methods. we are the first to systematically and directly compare different methods on llms of different sizes, including knowledge neurons (ig) and differentiable masking methods slimming and hard concrete. we take the view that llm memorization of a sequence is different from learning a type of knowledge. memorization is reproducing a long sequence (near) verbatim. in contrast, knowledge, often represented as @xmath3, relation, @xmath4.7% 1 layer neg-0.1% 47 layers self-4.7% 1 layer self-0.1% 47 layers triplet, occurs in variable contexts, where paraphrases are treated as equivalent expressions of the same knowledge. localization of verbatim memorization helps unlearn private or copyrighted data, for example, @xcite apply ig to localize and then erase private data from a bert finetuned on enron email dataset @xcite . our del benchmark differs from @xcite in two main ways: (1) we delete sequences that llms have naturally memorized during pretraining, (2) we locate neurons for each sequence independently, rather than finding a shared set of neurons, as our collected datasets cover diverse sequences. localization can also prevent overfitting: @xcite drop out pre-allocated neurons tied to memorizing mislabeled examples. in contrast with these works, we focus on benchmarking localization ability, since successful localization is the basis of its downstream applications. ## conclusion we propose two benchmarking approaches to define the success of llm localization, focusing on locating a small set of neurons in an llm that are responsible for memorizing a sequence. the inj benchmark enables a direct evaluation of localization methods, while the del benchmark evaluates methods on naturally memorized sequences, using dropout to measure localization success. the two benchmarks complement each other and show consistent rankings of methods. we find promising localization ability of all five methods we evaluate, especially for hard concrete. meanwhile, all methods confuse memorized sequences in the same or related categories. this finding suggests a need for better localization methods and poses the open question of whether llms use a shared set of neurons to memorize related sequences such that perfect localization is not possible.