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

Authors: jordi  armengol-estapé
lingyu  li
sebastian  gehrmann
achintya  gopal
david  rosenberg
gideon  mann
mark  dredze
♥ ♢ bloomberg
Venue: EMNLP_2024_Workshop_BlackBoxNLP
Published: 2024-08-15 04:19:56

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.

Main paper: 
 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 (Vig et al., 2020; Olsson et al., 2022) . 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 (Vaswani et al., 2017) is to project the model parameters into the word embedding or vocabulary space for interpretation (Elhage et al., 2021; Geva et al., 2022b) . Dar et al. (2022) 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 (Vashishth et al., 2019; Clark et al., 2019) or gradients (Dai et al., 2022) . 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 (Geva et al., 2022c; Dar et al., 2022) .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.We begin by interpreting the model parameters in the semantic embedding space. Transformers consist of two blocks: a multi-layer perceptron and a self-attention module, which are applied consecutively with residual connections. Elhage et al. (2021) note that this sum of the output of all the previous layers and the original embedding could be understood as a shared communication channel among layers, referred to as the residual stream. This property can be exploited to project intermediate outputs into the vocabulary space (nostalgebraist, 2020; Din et al., 2023) . In other words, activations across the model seem to be in the same embedding space as the input embeddings, E, which in most implementations share weights with the language modeling head (E T ).However, we need the parameters, and not only the activations, to be in a shared embedding space. Building upon Geva et al. (2022c Geva et al. ( , 2020) ) and Elhage et al. (2021) , Dar et al. (2022) proposed extending the residual stream view by projecting Transformer parameters into the original vocabulary space. We propose using this same insight to represent the Transformer's parameters in the same embedding space as text queries, allowing us to directly (semantically) compare a query with model parameters. 1 In summary, we will identify which model parameters contain names of CEOs by finding those that are most semantically similar to "CEO". We describe the procedure for both types of model parameters: MLPs and attention.Prior work has shown that MLP blocks function as key-value memories, allowing the Transformer to store knowledge (Geva et al., 2020) . The first layer of the MLP block is parameterized (omitting biases) by the weight matrix W in ∈ R D ′ ×D , the "keys" of the "memory", where D is the embedding dimension and D ′ is the hidden dimension of the MLPs. Similarly, the second layer of the MLP block is parameterized by a weight ma-trix W out ∈ R D×D ′, the "values". In Geva et al. (2022b) , the embedding space interpretation of those weights is that each W out column can be seen as an embedding vector in the same space as tokens. Geva et al. (2022c) extended this view to the parameters in the first layer; each row in W in can be seen as an embedding vector in the same space as the tokens. We refer to these parameters as MLP-K (key) and MLP-V (value).Attention blocks are associated with contextual processing rather than knowledge storage, though previous work has been able to statically interpret these parameters (Dar et al., 2022; Millidge and Black, 2022) . Attention blocks are parameterized (omitting biases) by 3 kinds of parameters for each attention head i ∈ {1, 2, ...N }: W i Q , the queries' projection; W K , the keys' projection; and W V , the values' projection. Additionally, there is a shared attention output projection, W O , which can also be split as separate W i O for each head. Dar et al. (2022) propose the subhead view, which forms embedding space interpretations for the individual units in W Q , W K , W V , and W O analogously to those we saw for MLPs. Dar et al. (2022) made assumptions to theoretically justify this embedding interpretation. For instance, they omit biases and layer normalization, and approximate the inverse of E, needed for extending the embedding space interpretation to the first layer of MLPs and attention subheads, with E T . We keep these choices since they empirically work well in Dar et al. (2022) . In our work, we don't need to explicitly use the inverse as we do not project the parameters to the vocabulary space.We now have an interpretation of the different transformer parameters in the embedding space E. We cast a given text query e in the same space to compute semantic similarity. We assume that our query effectively represents a knowledge type, e.g., the token "CEO" represents the CEO relation, and that this token is either directly present in E, or by pooling multiple tokens present in E. Given this embedding query e ∈ R D , we retrieve the k Nearest Neighbors (k-NN) over all layers for a specific parameter type (layers in MLP or attention blocks) by returning the projected parameter indices that maximize the cosine similarity for the query embedding. For example, for the first layer in MLPs:Q(e) = topk {s(e, p) | p ∈ L ⋃ i=1 rows(W i in )} ,where s is the cosine similarity and L is the number of layers. We posit that the most similar parameters p contain knowledge relevant to e.We deal with massive models with tens or hundreds of billions of parameters. Efficiently searching through this space for the parameters most similar to a query is not an easy task. Large models can be loaded into memory efficiently using model parallelism (Shoeybi et al., 2019; Rasley et al., 2020) , and we need to rearrange the Transformer parameters to match the embedding interpretation and allow for efficient search. As we saw earlier, the embedding space interpretation for the first layer of MLPs is that each row in W in can be seen as an embedding vector in the same space as tokens. Thus, the shape of the weights already matches its embedding interpretation. However, W out ∈ R D×D ′ requires transposition to correspond to the embedding shape. We refer by unit to the individual weight vectors interpreted in the shape required by the embedding space interpretation, be it a row or a column of the row of the matrix weights depending on the embedding space interpretation of each parameter kind. Similarly, the attention weights require rearrangement corresponding to the subhead view. We refer to the appendix for additional implementation details. 
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.We consider several different model sizes and families: GPT2-medium (355M) (Radford et al., 2019) , OPT-{1.3B, 6.7B, 66B} (Zhang et al., 2022) , GPT-NeoX (20B) (Black et al., 2022) , and Bloom 176B (BigScience). We also consider two instructiontuned models based on OPT-1.3B: OPT-IML-1.3B and OPT-IML-MAX-1.3B (Iyer et al., 2023) .We consider several tasks that rely on different types of financial information:• CEOs: A dataset of 500 CEO-company pairs (S&P 500).The model is asked for the name of the CEO of a given company in a zero-shot setting.• Tickers: Same as CEOs but asking for stock tickers for a company.• Ticker extraction (NER-ED): The ticker extraction task in Wu et al. (2023) , in which models must extract the tickers of the named entities (companies) appearing in the text.• Authors: A dataset of best-selling fictionauthor pairs. 2 The model is asked to reply with the author name of a given work.2 https://w.wiki/7Lhr• Directors: A dataset of Academy Awards movie-director pairs. 3 The model is asked to reply with the director of a given movie.• Arithmetic: add_sub_multiple subset in the test split of the Deepmind Mathematics dataset (Saxton, 2019) . 
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 : OPT-6.7B accuracy wrt. MLP-Ks ablations, on the CEOs task. Blue lines represent the accuracies with increasing ablated units using the embedding weight localization; orange lines represent those with random ablation. Solid lines correspond to the target task accuracies, while dashed lines correspond to the control task. We can see that the solid blue line decreases faster than both the lines corresponding to the control task and the random ones.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 3 https://w.wiki/7Li5 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 4 . Table 1 reports the differences between target and control accuracies. We confirm the difference between random and control ablations and the targeted ablation to be positive in all cases. Table 2 shows a compact version of the accuracy plots and histogram we saw before, this time for a selection of the models (OPT-{6.7B, 66B}, GPT-Neox-20B, Bloom-176B) on the CEOs task. For all models, the accuracy on the CEOs task drops sharply, while the accuracy on the control task remains less affected.How does model size influence knowledge localization? Figure 3 shows accuracy in target (solid lines) and control (dashed lines) tasks as more units close in embedding space are removed. We first consider the CEO task with MLP-K ablations. Across all model scales (from 355M to 176B) and families (GPT2, OPT, GPT-Neox, Bloom), there is an early, sharp drop in the target task performance.In contrast, the control task remains largely unaffected until significantly more weights are ablated.For the CEO task, at the extremes, Bloom-176B seems most impacted by ablation (i.e., the largest gap between dashed and solid lines), while GPT2-355M is the model least affected by the targeted ablation. Other than this observation, we see no correlation between model size and knowledge localization effectiveness; the technique works similarly across model families and sizes.How do other model features affect knowledge localization? Figure 3 also shows instructiontuned variants (OPT-IML). Interestingly, their accuracy curves have very similar shapes, and for the target task, the less fine-tuned the model, the lower the accuracy seems to drop, while a seemingly reverse pattern seems to hold for the control task. In Table 1 , we observe that in the case of targeted and random ablation on the target task (B-A), the instruction-tuned versions of OPT-1.3B present very similar numbers across ablation levels. In short, instruction-tuned models from the same base model show very similar behaviors when ablated compared to the original base model. Instruction tuning does not change where knowledge is stored or our ability to locate it. We further observe that the control task for GPT-Neox is disproportionately affected. This could be explained by the effect of post-Kaplan (Kaplan et al., 2020) scaling laws used in its pre-training compute. Longer pre-training may lead to less sparse representations, a hypothesis we seek to explore in future work.Figure 3 presents results for other tasks: Tickers and Directors. The significant delineation between target and control task shown in Figure 3 for CEOs and Directors tasks suggest that the effectiveness of localization depends on the specificity of the query token. However, the more challenging task of Ticker Extraction also led to promising results, especially with OPT-66B MLP-V (last row in Table 3 ). We refer to Appendix C for interesting results for the Arithmetic task, despite ablating the seemingly innocuous token "0". Knowledge localization worked for both OPT-66B and Bloom-176B, but less well for GPT-Neox, on which the tickers extraction task didn't show promise either. Overall, GPT-Neox layerwise distributions hint at a more significant signal in the first (MLP-Ks) and last (MLP-V) layers than in the rest of the models.What parameters store knowledge? Examining which parameters were selected confirms that we did not trivially select units from the first and last layer, where the representations are closer to the embedding layer. It also provides insights into Accuracy given 2.0% Ablation Table 1 : Results on the CEOs task with MLP-Ks. We show accuracies at different ablation levels (0.1% ablated units, 0.5% ablated units, etc). For each ablation level, for each model, we report 4 accuracies: A) the targeted ablation results (i.e., ablating units close to "CEO") on the target task (CEOs task), B) the random ablation results on the target task, C) the targeted ablation results on the control task (Capitals task), and D) the random ablation results on the control task. We also report the differences between these accuracies. We expect B-A and C-A to be positive, which means that the random ablation has less effect on the performance than the targeted ablation on the target task, and that the control task performance is less affected than the target task, respectively. how the model stores and processes knowledge. Tables 2, 3 , and 5 also present the histograms of the layer unit counts. The overall trend, similar to the findings in Vig et al. (2020) and nostalgebraist (2020), is for the unit density in either the first or last layers. More interestingly, some models show distinctive fingerprint-like patterns. For example, all OPT results with MLP-K ablation have a peak around the first third together with another peak at the last layer, especially in the larger variants. OPT results with MLP-Vs ablations follow a U-shape distribution. In other model families, GPT-Neox units concentrate on the first layers with MLP-K ablations and the last layers with MLP-V ablations.Model GPT2-355M OPT-1.3B OPT-IML-1.3B OPT-IML-Max-1.3B OPT-6.7B GPT-Neox-20B OPT-66B Bloom-176BHow well can knowledge be localized in different parameter types? As mentioned at the beginning of section 4, we found that knowledge localization was more accurate for MLPs. We still include Attention results in the Appendix, as shown in Tables 6. We successfully select Attention weights (as in, the target task performance drops faster than control) in about half the cases, but less reliably than in the case of MLPs. This is consistent with the fact that a) we evaluated knowledge-intensive tasks, and b) prior work (Geva et al., 2020 (Geva et al., , 2022b) ) suggests that MLPs are more involved in this kind of task than attention modules.  
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 (Gehman et al., 2020) . Several strategies exist to mitigate these generations during inference (Dathathri et al., 2020) , remove behavior during fine-tuning (Liu et al., 2021) , filter pretraining data to remove biases (Zhang et al., 2022) , and neuronlevel interventions that edit the model parameters directly (Geva et al., 2022c; Li et al., 2023) .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 (Geva et al., 2022c) or by learning ablation masks from finetuning models on the toxic dataset (Li et al., 2023) , 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 (Gehman et al., 2020) and Wikipedia (Foundation). Additional details on the methodology and evaluation are in Appendix F. Figure 5 shows that the ablation of MLP-V units reduces the toxicity of the models, with a drop of more than 35% when only 1% of the units are ablated. Ablating the same number of random units leads to unchanged toxicity. Curiously, ablating MLP-K units has no significant effect on toxicity, and ablating Attention units reduces the toxicity of the models, but not as effectively as MLP-V's.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. Table 4 shows the MLP-V layer distribution of the ablated units. Surprisingly, all models demonstrate that toxicity related units tend to be concentrated in the early layers, which is rarely seen in other ablation tasks. The early layers are though to be associated with shal- (Geva et al., 2020 (Geva et al., , 2022b;; Elhage et al., 2021; Dar et al., 2022) . Geva et al. (2022a) investigated keyword search over Transformer parameters, although their work is limited to the second layer of MLPs, and search over the tokens projected from the parameters, rather than directly on the parameters themselves. 
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. 
 
 Appendix: In this work, we have not included evaluations for more recent models such as Llama (Touvron et al., 2023) . We base our work on Dar et al. (2022) , which does not support SwiGLU (Shazeer, 2020) out of the box (due to the added parameters).Additionally, the evaluations are limited to domain-specific tasks, and control and target tasks are not necessarily equally easy to ablate. We explore the robustness of our results to this choice of control task in the Appendix E.In this section, we review the fundamental ideas in Vaswani et al. (2017) and Dar et al. (2022) required to better understand our proposed solution from a technical standpoint.Transformers (Vaswani et al., 2017) are mainly characterized by the following hyperparameters 5 :• D, embedding/hidden size• |V |, vocabulary size • L layers • N attention heads • D n = D N head dimension • D ′ = hidden dimension of MLPs (typically 4D)Embedding and unembedding As a first step in the Transformer, each input token from the discrete sequence is embedded into a real-valued vector through the embedding matrix E ∈ R D×|V | . Each token t is embedded as follows:x = embed(t) = E[∶, t id ]with x ∈ R D and where t id is the token id (index) corresponding to the token t. This maps tokens to the embedding space of vectors in R D . Some models also apply layer normalization after the embedding layer, but we omit it here.Similarly, an unembed (or language modeling head) layer is used as the last step, to go back to the vocabulary space of logits in R |V | . Typically, it's also a linear transformation with a weight matrixU ∈ R |V |×D : unembed(x) = U xAgain, some models apply layer normalization in unembed but we omit it here. Since U has the transposed dimensions of E, many implementations tie embedding and unembedding layers to the same layers (e.g., GPT-2). 6 In this work, we assume that the unembed layer is linear and tied to the embedding layer, which is a realistic assumption for most models, with U = E T 5 While we focus on decoder-only language models, our method is not restricted to this kind of Transformers.6 Other works such as Devlin et al. (2018) opt for a nonlinear unembedding layer.Transformer layer After the embedding layer, 7 a decoder-only model is composed of N identical layers. Let X ∈ RT ×D be a sequence of embedded tokens, with T being the sequence length, and x be an individual embedded token (real-valued vector with dimension D) in the sequence. Each layer has the following structure:Layer(X) = MLP{LN[MHA(LN(X))+X]}+Xwhere MLP stands for Multi-Layer Perceptron (MLP) and MHA stands for Multi-Head Attention.The exact order of application of LN (Layer Norm) varies across implementations.Multi-Head Attention Transformers mix information from the token embeddings in a given sequence with pairwise dot-product multi-head attention. We will first see how each attention head is (independently) defined for each attention head, and later we will see how the head outputs are aggregated. The first step is projecting the token representations in the input sequence X into Queries (Q), Keys (K), and Values (V) as follows:Q = XW Q +b q ; K = XW K +b k ; V = XW V +b v Head(X) = softmax ( QK T √ D k ⊙ M ) Vwhere b {q,k,v} are the bias terms and ⊙M is the element-wise multiplication by the masking matrix, defined as:M = [m i.j ] i,j ∈ R T ×T m i,j = 1(i ≤ j) -∞ ⋅ 1(i > j) ∀i, j ∈ [T ]with 1(a) being an indicator function returning 1if the predicate a is true and 0 otherwise. This element-wise multiplication by M has the effect of creating a causal, triangular attention mask that prevents leaking future token information.In decoder models, Q, K, and V come all from the input sequence, attending to itself. We omit unmasked self-attention or cross-attention due to our focus on decoder-only models.The outputs of each head Head n (X)∀n ∈ [N ] are then concatenated and projected back to the embedding dimension: MLPs MLPs are independently applied elementwise to each x ∈ X:[Head n (X)∀n ∈ [N ]] W o + b oMLP(x) = act(XW in + b in )W out + b outwhere act is the activation function (typically GELU (Hendrycks and Gimpel, 2016) ) and b {in,out} are the bias terms. The first layer parameters (W in ) can be referred to as keys and the second layer ones (W out ), as values, following the findings in Geva et al. (2020) that Transformer MLPs act as Key-Value memories, but they are not to be confused with the keys and values of the attention block.Layernorm Layernorm (LN) is a normalization function that is applied to hidden states of the Transformer before/after (depending on the implementation) attention and MLP blocks:LN (x) = x -µ(x) √ σ 2 (x) + ϵ ⊙ γ + βwhere γ and β ∈ R D are learnable parameters, ϵ is a constant value added to the denominator for numerical stability. and µ and σ are computed independently for each token in the sequence x ∈ R D :µ(x) = 1 D ∑ i x i ∈ R σ 2 (x) = 1 D ∑ i (x i -µ(x)) 2 ∈ RNote that each individual LN (the different ones in each layer and across different layers) has an independent parameterization of γ and β.Residual stream Both MLP and attention blocks are applied as residual connections, meaning that the output of these blocks is summed to the previous hidden states. Elhage et al. (2021) noted that this sum of the output of all the previous layers and the original embedding could be understood as a shared communication channel among layers. Taking this view to the extreme, hidden states should be able to be projected into the original vocabulary space.We omit dropout (Srivastava et al., 2014) and positional encodings since their presence and specifics vary across implementations.There is a vast literature on Transformers' interpretability from different lenses. In this work, we are interested in the view that Transformer parameters can be analyzed in the embedding space (Dar et al., 2022) . Geva et al. (2022c) showed that the values (second layer in MLPs) of Transformers can be interpreted in the embedding space. Elhage et al. (2021) showed that attention parameters could be interpreted in embedding space in small models. Dar et al. (2022) generalized those previous findings to all Transformer parameters (both attention and MLPs). In the remainder of this section, we summarize the parameter projections proposed in More specifically, the matrix M can correspond to a weight matrix. For interpreting a row vector of M , v, Geva et al. (2022c) (corresponding to the weights of an individual unit) follow two steps. The first one is the projection to project v into the vocabulary space, vE. The second step is to take the top-k argmax vE, and these top-k tokens would correspond to the tokens the most related to the unit parameterized by v. Geva et al. (2022c) posits that this interpretation is sound since the most activated vector coordinates contribute the most when added to the residual stream. Geva et al. (2022c) can only apply this method to the values of MLPs (weights of the second layer), because these are the ones directly being added to the residual stream. Dar et al. (2022) posit that inner products and matrix multiplications in a Transformer can be interpreted in the embedding space if we assume a right inverse of E, E′ , such that we can approximately reconstruct the original matrix.MLPs blocks have been shown to work as keyvalue memories where the Transformer stores knowledge (Geva et al., 2020) , so we expect to be able to locate knowledge in their parameters.The first layer of the MLP block is parameterized seen as an embedding vector in the same space as tokens. Geva et al. (2022c) extended this view to the keys (first layer of MLPs) as follows. According to the view of Transformer MLPs as Key-Value memories (Geva et al., 2020) , with x being the hidden state input to the MLP (the "queries" to the memory):xW T in = xEE ′ W T in = xEE ′ W T in = xE(W in E ′T )T Assuming the residual stream interpretation, according to which xE should be interpretable in the vocabulary pace, then W in E′T should also be interpretable in the vocabulary space since they directly interact through an inner product in the MLPs' "memory". Thus, each row in W in can be seen as an embedding vector in the same space as the tokens. Finally, note that W in E ′T can be approximated as W in E sOmitting biases, attention blocks are parameterized by 3 kinds of parameters for each head i ∈ {1, 2, ...N }: W i Q , the queries' projection; W K , the keys' projection; and W V , the values' projection. Additionally, there is a shared attention output projection, W O , that can also be split as separate W i O for each head (the part of the projection matrix interacting with the corresponding head after the concatenation). Dar et al. (2022) consider two possibilities for projecting these weight matrices into the vocabulary space, namely, the interaction matrices, and the subheads view. Elhage et al. (2021) proposed interpreting attention through the interaction matrices of queries-values, W QK , and valuesoutput projection W V O . From the dot-product attention formula, it's easy to see that If we omit biases and define Q = XW Q and K = XW K , it's easy to see that W Q and W K interact directly and in an input-independent way when computing the dot product:QK T = XW Q (XW K ) T = XW Q W T K XSimilarly, we can see that W V and W O interact directly after the concatenation of the different head outputs. All in all, for each head i we can define:W i QK = W i Q W i T K ∈ R D×D W i VO = W i V W i O ∈ R D×DSimilarly to what we saw for MLPs' keys, we can now follow how (Dar et al., 2022) interpret these matrices in the embedding space. Like MLPs' values, the output of the attention block is directly added to the residual stream and, thus, we expect it to be meaningfully projected into the embedding space. With a reasoning analogous to what we saw in the case of an MLP, making use of E ′ and interpreting inner products in the embedding space, Dar et al. (2022) showed that the other parameter kinds (W Q , W K , and W V ) can also be approximately projected into the vocabulary space in a meaningful way.Subhead view Dar et al. (2022) propose an alternative view to the attention interaction matrices that has the advantage of being able to project individual units. Using the identity AB = ∑ b j=1 A ∶,j B j,∶ :W i VO = D N ∑ j=1 W i,j V W i,j O , W i QK = D N ∑ j=1 W i,j Q W i,j T KThis allows for the definition of subheads. Subheads are the vector columns ofW i Q , W i K , W i V , that is, W i,j Q , W i,j K , W i,j V ∈ RD×1 , respectively. They can be approximately projected to the vocabulary space by multiplication by E. Additionally, the row vectorsW i,j O ∈ R 1×D of f W iO are also subheads, and they can be directly projected to the vocabulary space by multiplication by E without any approximation.Layer-norm is ignored in this approach, following Elhage et al. (2021) 's observation that it can be ignored because normalization changes only magnitudes and not the direction of the update. Biases and the effects of positional encoding are also omitted in this approach for the sake of simplicity.Finally, we note that Dar et al. (2022) propose using ET as an approximation to the right inverse E ′ due to a) being a good enough approximation, and b) yielding more interpretable results. However, in our case, since we never need to project parameters to the vocabulary space, we sidestep the need for directly using this inverse approximation.Storing rearranged weights with a list of tensor views (one list for each parameter kind, each element in the list being a tensor view corresponding to the weight matrix of a given layer and parameter kind), rather than creating a new tensor with all the parameters, allows to store the reshaped weights' data as a reference to the original one. This has the benefits of a) decreasing memory overhead, b) keeping the original device sharding in case of LLMs, and c) being able to directly modify the model weights by modifying the rearranged ones.Each model architecture requires its own weight loader, since weight storage varies across implementations (e.g., the GPT-2 family implements the MLP as a 1-D convolution layer, meaning that the weights are transposed; in some implementations, the attention keys, queries, and values projections are stored as a single linear layer). Finally, we note that optimizations typically employed in k-NN settings would be directly applicable here.In all cases, we study parameter kinds separately: W in , W out , and attention (for attention, we separately select the top K units for each among the 4 parameter -kinds, queries, keys, values and output projection -and ablate all of them at once). We use the simplest ablation method by setting the corresponding weights to zero to validate the hypothesis that the selected weights are related to the erased concept in the most extreme case. Zeroing out the weights has an indirect effect on the general layer statistics, which might explain the drop in performance of the control task after a significant amount of (presumably unrelated) weights have been ablated.In all cases, we use a fixed set of numbers for setting K when conducting an experiment on a given model, i.e. 10, 50, 100, 500, 1k, 5k, 10k. However, the total number of units that a model contains varies from one model to another. For easing the comparison of different models in plots, we transform the K values to the proportions over the total number of units that a given model has. For reporting the ablated models' accuracy in tables, we apply interpolations on each model's accuracy given their transformed proportions of ablated units to achieve the same set of proportions of ablated units, i.e. we select points 0.1%, 0.5%, 1.0%, 2.0%.Table 5 summarized the results on the Arithmetic task when ablating using "0" as query. In all cases, we conduct the random ablations as a comparison to the targeted ablation. We use a random number generator to pick the K units randomly and ablate them. For the small models with < 10B parameters, we use 5 different random seeds to run random ablations for 5 times and use their mean accuracy and standard deviation for the plots and tables. For the big models with > 10B parameters, we only use one seed to run random ablation for the plots and tables. In the main sections, we use the "Capitals" task as a control task to evaluate how specific our method can ablate the target task's knowledge embedding, that is, reducing the target task accuracy while not reducing the control task's. In this section, we use more control tasks to investigate the robustness of our results when only the "Capitals" control task is used. For any given target task, we compare the accuracy drop of the "Capitals" task as opposed to a set of control tasks(i.e., the control group). For example, in Figure 8 , we compare the results on the "CEOs" target task where we use the "Capitals" as our control task (Figure 8a ) against where we use the control group including the "Capitals", "Tickers", "Directors" and "Authors" tasks (Figure 8b ). We plot the control group using its members' mean accuracy and the standard deviation. We repeat this analysis for the "Tickers", "Directors", and "Authors" target tasks.Model GPT2-355M OPT-1.3B OPT-IML-1.3B OPT-IML-Max-1.3B OPT-6.7B GPT-Neox-20B OPT-66B Bloom-176BAs shown in Figure 8 , 9 10 and 11, we find: On the "CEOs" target task, the control group mean accuracy drops about 30% -40% on average when ablating about 2% units on models MLP-K layers, while the "Capitals" control task accuracy drops about 0% -10%. This implies that our ablation method affects control tasks differently. However, at 2% units ablation on the MLP-K layers, the control group's largest accuracy drop (about 50%) is less than the target task's accuracy drop (about 80% -90% for the most models, 60% for GPT2-335M), which shows our ablation method is still effective on locating and ablating the knowledge parameters. While figures of the "Tickers", "Directors" or "Authors" target task report different numbers for the above comparison, the above implication still applies to these target tasks.To further understand why our ablation method affects the control group's tasks differently, we compare the control tasks in the control group given a target task's ablation on the MLP-K and MLP-V layers. Similar to the Appendix E.1, we iterate this analysis over all four target tasks. As shown in Figure 12 , we find: On the "CEOs" target task at 2.0% units ablation on MLP-K layers, the least accuracy drop (about 10%) is from the "Capitals" control task while the largest accuracy drop (about 50%) is from the "Tickers" control task. On the "Tickers" target task at 2.0% units ablation on MLP-K layers, the least accuracy drop (about 15%) is still from the "Capitals" control task while the largest accuracy drop (about 70%) is from the "CEOs" control task. This implies that the "CEOs" and "Tickers" tasks are more correlated than the rest of the tasks while the "Capitals" task is more independent than other tasks. We observe the similar behavior between the "Authors" and "Directors" tasks when used as the target tasks. Inspired by the topic knowledge lookup method (Dar et al., 2022) , we demonstrate that toxicity can be represented as a topic embedding, which is utilized to locate knowledge neurons. To obtain this topic embedding of toxicity, we follow these steps:Model GPT2-355M OPT-1.3B OPT-IML-1.3B OPT-IML-Max-1.3B OPT-6.7B1. Identify the most common toxic keywords by sampling from a commonly-used list of offensive words. 9 This process yields 24 of the 9 List of Dirty, Naughty, Obscene, and Otherwise most toxic tokens.2. Retrieve the corresponding embeddings of these 24 toxic tokens from the model's embedding table.3. Compute the average of these embeddings to obtain the topic embedding of toxicity.After getting this topic embedding, we retrieve the K nearest neighbors of projected parameters using cosine similarity and zero out these parameters of selected K knowledge neurons. We perform zero-ablations on various types of layer separately: attention, MLP-K, and MLP-V.We evaluate our method on a subset of RealToxic-ityPrompts (Gehman et al., 2020) , a collection of 600 prompts designed to elicit toxic responses generated from models. We employ Detoxify (Hanu and Unitary team, 2020), a toxicity classifier, to assess the toxicity of a model's output when presented with a prompt. Detoxify provides 6 metrics, and we classify the model's output as toxic if any of the output scores from these 6 metrics are > 0.5. Given the total 600 prompts, we evaluate the toxicity of each ablated model by calculating the proportion of toxic outputs generated from each ablated model. To demonstrate the effectiveness of our method, we conduct random ablation experiments with 5 different seeds as a baseline method to compare with our method.For an additional check, we measure the perplexity of the ablated models when presented with 2000 prompts sampled from RealToxicityPrompts to assess their language generation performance. We also measure the perplexity of the ablated models when presented with another 2000 prompts from Wikipedia (Foundation) as a control task. OPT-1.3BAtt.OPT-IML-1.3BAtt.OPT-IML-Max-1.3BAtt.OPT-6.7BAtt.GPT-Neox-20BOPT-66BBloom-176BTable 6 : Results on the CEOs and Tickers task, including accuracies for the target and control task when using the embedding weight location method and the layer-wise unit distribution. A similar idea was proposed byDar et al. (2022) to align parameters across models through vector similarity.For comparing the performance of different ablated models, we use the percentage of ablated units over the total model units rather than unit counts since the total model units vary by different model size.Some models also have a positional embedding in this step, which we omit here. 


Decision: Interesting static analysis of models on multiple tasks

Reviews: review: paper_summary: paper_summary: For a given query (/question), this paper studies techniques to localize the knowledge (/answer) to specific parameters in a transformer. Using techniques from prior work, they project the parameters (mostly MLP) into the embedding space and then perform kNN search between the query embedding and parameter embeddings. They evaluate the impact of these identified parameters by zeroing them out and comparing it against a control setting. On multiple zero-shot question-answering tasks, author(s) show the effectiveness of knowledge localization on variants of GPT-2, GPT-Neo, OPT and Bloom models.

Results show interesting trends for tasks such as company name --> CEO name prediction, but results are mixed for other tasks such as tickers, authors, directors etc.,
strength_weakness: strength_weakness: - The current paper overtly focuses on the CEO task in the main paper text. As the author(s) hint in section 4.3, the generalization to other tasks is not guaranteed. On tasks such tickers, and directors the differences between targeted and controlled ablations are less pronounced (Table 3, Figure 3).
- Some parts of the analysis are quite hard to follow, specifically around layer-wise distributions, ablations on tickers and directors tasks. The paper can benefit from editing parts of section 4 to make it easier to understand the contributions. In the current form, author(s) refer back-and-forth between many figures and tables. See my questions below.
- As author(s) acknowledge, MLP seems more useful than attention layers (lines 437-450). I think this should be mentioned at the start of section 4 to set expectations for the readers. This is important because sections 2 and 3 do not make this distinction.
questions: questions: - I think the argument in lines 36-39 is problematic. Authors cite "we don't know what to ask" as an issue with existing probing methods, but this paper also uses a 'query' to look for knowledge in the parameters. Am I missing something here?
- Is there any difference with the approach of Dar et al., 2022? Section 2.1 seem to indicate its the same.
- The legend for control task in Figure 1 is unclear.
- What are $n$ and $N$ in Figure 2 caption? Also, are $k$ and $K$ different in lines 271-275? $k$ is not used anywhere else in the text, so you could skip defining a variable for this.
- OPT-66B presents interesting behavior on the control task in Table 1. Why is there a significant drop in performance with 0.1% ablation on the control task? Did you happen to investigate this?
- From Table 2, I couldn’t find a consistent pattern amongst the layers. Does this imply different models store ‘CEO’ knowledge in different layers? Some discussion on this can be helpful.
- To be consistent with Figure 3, I would recommend using % ablated units as x-axis in Figure 1.
- I couldn’t follow the reasoning in first paragraph of section 4.2, specifically about Bloom-176B and GPT2-355M. The explanation about model size extremities only holds for ‘CEOs’ category. From Figure 3 (’Tickers’), all OPT models seem to show similar drops as Bloom-176B. So, the argument about model size extremities doesn’t seem to hold. Also, any reason for not report Bloom-176B for ‘Directors’?
paper_summary: paper_summary: The paper focuses on identifying which parameters contribute to specific knowledge/ behaviors in the LLM. They use static analysis to project weights into token space and computing embedding similarity between tokens and weights to isolate weight specific behaviors.
strength_weakness: strength_weakness: Experimentation could be done across domains/tasks to analyze the generalizability of the proposed approach,
questions: questions: Any reasoning behind picking financial domain and toxicity specifically for experimentation?
paper_summary: paper_summary: This paper introduces a novel method for statically identifying where specific knowledge is stored within the parameters of a large language model (LLM). The method leverages the idea that model parameters can be interpreted as embeddings in the same embedding space. By searching for parameters most similar (in embedding space) to a query representing the desired knowledge, the authors are able to locate relevant parameters without requiring forward or backward passes.
strength_weakness: strength_weakness: Given the promising results in the financial domain and toxicity cases, it would be exciting to see how well the method performs in other domains and tasks to further explore its generalizability.
questions: questions: Overall, this research offers a promising approach to statically locating knowledge in LLMs, paving the way for further exploration of how these models encode and store information. Additionally, I appreciate the authors' clear presentation and the easy-to-follow explanations of their experimental setup and methodology.