Can We Statically Locate Knowledge in Large Language Models? Financial Domain and Toxicity Reduction Case Studies (2024)

Abstract: Current large language model (LLM) evaluations rely on benchmarks to assess model capabilities and their encoded knowledge. However, these evaluations cannot reveal where a model encodes its knowledge, and thus little is known about which weights contain specific information. We propose a method to statically (without forward or backward passes) locate topical knowledge in the weight space of an LLM, building on a prior insight that parameters can be decoded into interpretable tokens. If parameters can be mapped into the embedding space, it should be possible to directly search for knowledge via embedding similarity. We study the validity of this assumption across several LLMs for a variety of concepts in the financial domain and a toxicity detection setup. Our analysis yields an improved understanding of the promises and limitations of static knowledge location in real-world scenarios.

## introduction the impressive text generation abilities of large language models (llms) arise from complex interactions among billions of parameters. these parameters encode a vast range of knowledge, allowing models to answer closed-book fact-based questions across dozens of domains. llm evaluations focus on model abilities: what can the model do and what does it know? however, evaluations based on model outputs cannot answer where this knowledge is stored in the network. if a model correctly answers the question "what city in the united states had the first subway?", why does it matter what parameters store the answer "boston"? first, we may want to know in what domains is the model capable by simply "looking" at the knowledge-storing parameters, which can provide insights about the model's inner working. second, we may want to edit or remove a * work done during an internship at bloomberg. model's knowledge or behavior, e.g., outdated information, offensive terminology, or stereotypes. removing this information may be more effective than fine-tuning the model not to express it. third, in the search for better model architectures, we may want to enhance a model's knowledge storage ability. all of these goals and more require knowing where information is stored inside a model. previous research on locating knowledge in language models is divided into dynamic and static analyses. dynamic analyses focus on examining how model activations and outputs change with different inputs to identify where knowledge is stored and how model capabilities function @xcite . instead, we are focused on investigating the possibilities of static knowledge location, that is, without any forward or backward passes. while static knowledge location is challenging due to having access to strictly less information than dynamic methods, its potential simplicity makes it interesting for practical reasons and scientific curiosity. one static approach to examine the weights in a transformer network @xcite is to project the model parameters into the word embedding or vocabulary space for interpretation @xcite . @xcite developed the idea that all transformer parameters, including all multi-layer perceptron (mlps) and attention layers, can be interpreted by projecting them into the vocabulary space. they also use the embedding space interpretation to align parameters across models based on their vector similarity. the advantage of this approach is that it operates directly on the model parameters without requiring specific inputs or a forward pass. rather than projecting all parameters into the vocabulary space, or aligning them based on their vector similarity, we posit that if the embedding space interpretation of the parameters holds, then we should be able to directly locate specific knowl-edge given a query in the embedding space. we propose a straightforward method based on embedding similarity that identifies what knowledge is contained within a model and where that knowledge is located without any forward pass. we take two real-world case studies, information extraction tasks focusing on the financial domain and a toxicity reduction setup, and run our experiments at scale (up to 176b parameters). we study the relevance of the parameters identified by our method to the target knowledge, and investigate how specific these locations are by measuring the downstream performance on seemingly unrelated tasks when ablating the found parameters. finally, we utilize the method to gain insights into how internal model representations vary across layers and how distributed they are, and to have a better understanding of the possibilities and limits of static knowledge location. ## statically locating parameters our goal is to statically identify where knowledge is stored within the parameters of a large language model. we assume a running example of knowing the names of ceos for companies. this information is part of the financial domain, and is a specific type of information (ceo relation) that applies to many different companies. what parameters in the model store the identities of these ceos? popular approaches to locating this information use a forward pass through the model with different inputs to measure how the outputs or activations of the model vary, e.g. by inspecting attention weights @xcite or gradients @xcite . however, input-based approaches require forward and/or backward passes, which are computationally expensive and may not generalize beyond the tested inputs. instead, we are interested in locating parameters statically, that is, without input, forward or backward passes by building on recent methods of static interpretation of transformers @xcite . these methods can directly interpret model parameters in the embedding space. we represent the model parameters in embedding space, formulate a task-relevant query in the same embedding space, and use it to search over the parameters. for example, to locate parameters containing ceo identities, we take the query "ceo" and search for relevant parameters in the embedding space. we describe our method in several steps: how the model parameters are interpreted into the embedding space, how a query is represented in this same space, and efficient search of the parameter space. ## experiments we evaluate the effectiveness of our static search method in identifying parameters that contain knowledge related to the given query. unlike typical llm benchmarks, we do not know the "right answer " nor can we evaluate our search in terms of accuracy. therefore, we develop probes and metrics to measure the extent of the relevant knowledge contained in the identified parameters. we focus on domain-specific knowledge by searching for several types of financial information across several models. we measure model performance at knowledge tasks under ablation experiments, where we zero out parameters identified by our method as containing relevant knowledge. specifically, we eliminate the top-k parameter units identified as being the most similar to the given query, e.g., "ceo" for ceo knowledge questions. we increase the number of ablated units (k corresponds to 0.1%, 0.2%, etc. of the all units) to measure the performance degradation on the knowledge task. we establish baseline performance by comparing to random ablations (see appendix d.) we expect that ablating parameters will hurt a model; we seek to show that ablating specific parameters removes specific knowledge. therefore, we also report performance on control tasks unrelated to the query for which the ablations should have little effect. for example, we query the model about national capitals, unrelated to our financial queries. additionally, we select this task since it relies on general domain knowledge, for which even small llms are likely to do well. we explore the robustness of our results to this choice of control task in appendix e. ## results we summarize our main findings for select models, and include additional results for all tasks and models in the appendix. in preliminary experiments, we found that knowledge localization was more accurate for mlps; we thus focused on mlp ks and vs for most experiments. figure 1 shows the accuracy of the ceos task as we ablate an increasing number of units (mlp ks) for opt-6.7b. the accuracy drops sharply when the closest units in embedding space to "ceo" are removed (solid blue line), while the control task (dashed blue line, capitals task) remains largely unaffected. orange lines show accuracy when an equal number of randomly selected units are ablated, which has a significantly smaller effect on the performance. this shows that our method can statically identify parameters responsible for storing relevant knowledge with a certain degree of specificity. figure 2 shows the corresponding layer-wise distribution of the ablated units (close to the embedding of "ceo" in embedding space), with darker blue indicating fewer units selected. while we observe a strong effect in the first and last layers, the overall pattern is far from trivially selecting the units of the layers close to the embedding layers. a peak between layers 5 and 10 hints at some early processing of the concept of "ceo". we find similar trends for the ceo task across all models foot_2 . ## toxicity reduction our evaluation so far has been on model knowledge: can we identify where information is stored? we now turn to model behavior: can we identify what model parameters are responsible for toxic language generation? llms can generate toxic text that contains offensive language and biased beliefs and stereotypes @xcite . several strategies exist to mitigate these generations during inference @xcite , remove behavior during fine-tuning @xcite , filter pretraining data to remove biases @xcite , and neuronlevel interventions that edit the model parameters directly @xcite . we adopt a strategy similar to the neuronintervention methods and use our technique to identify and ablate model parameters associated with toxic generation. we apply our method to opt 1.3b, opt 6.7b, opt iml 1.3b, and opt iml max 1.3b and measure the reduction in toxicity. unlike prior methods that locate toxic units by projecting them into the vocabulary space @xcite or by learning ablation masks from finetuning models on the toxic dataset @xcite , we hypothesize that toxicity can be removed by ablating parameters found using our static knowledge localization method. we form a query to represent a toxic concept by averaging the embeddings of 24 toxic tokens. we locate the k nearest units to this concept embedding and ablate them as we did in our above evaluations. we measure the proportion of toxic outputs generated from each ablated model (following hanu and unitary team, 2020), as well as the perplexity of 2,000 prompts sampled from realtoxicityprompts @xcite and wikipedia (foundation). additional details on the methodology and evaluation are in appendix f. ablating mlp-v units increases the toxic language modeling perplexity with minor impairment of the generic one. figure 4 shows that, when ablating ∼1% of mlp-v units on models with size 1.3b, the perplexity of model generations on toxic prompts increases from 31.81 to 49.03 (+17.22), while the perplexity of model generation on non-toxic text (wikipedia) increases from 11.19 to 15.07 (+3.88). the mlp-v units' ablation impedes the ability of the model to model toxic language while having a smaller effect on the overall language modeling performance. however, for opt 6.7b, the degradation in perplexity for wikipedia is significant, increasing from 8.62 to 64.12 (+55.5), implying that the scale of the model is inversely correlated with the sparsity of units that encode broad concepts related to language modeling, or that this methodology of defining a toxicity embedding does not scale to larger models. ablated toxicity mlp-v units are distributed around the early layers. ## conclusions we demonstrated that by casting the parameters of an llm into embedding space and directly performing embedding similarity search with respect to a query, we can localize stored knowledge without a forward pass. we have studied the performance after ablating the selected parameters and the layer-wise distribution of these parameters in two real-world settings on a diverse range of models to gain insights on the promises and limits of static knowledge location.