Neuron-Level Knowledge Attribution in Large Language Models (2024)

Abstract: Identifying important neurons for final predictions is essential for understanding the mechanisms of large language models. Due to computational constraints, current attribution techniques struggle to operate at neuron level. In this paper, we propose a static method for pinpointing significant neurons. Compared to seven other methods, our approach demonstrates superior performance across three metrics. Additionally, since most static methods typically only identify “value neurons” directly contributing to the final prediction, we propose a method for identifying “query neurons” which activate these “value neurons”. Finally, we apply our methods to analyze six types of knowledge across both attention and feed-forward network (FFN) layers. Our method and analysis are helpful for understanding the mechanisms of knowledge storage and set the stage for future research in knowledge editing. The code is available on https://github.com/zepingyu0512/neuron-attribution.

## introduction transformer-based large language models (llms) @xcite @xcite possess remarkable capabilities for storing factual knowledge, which is important for downstream tasks including question answering @xcite and reasoning @xcite . while recent studies @xcite @xcite @xcite have made significant progress in understanding knowledge localization and the information flow from inputs to predictions, it is still hard to identify exact parameters for knowledge storage in llms due to several reasons. firstly, existing studies often depend on causal tracing @xcite and integrated gradients @xcite for knowledge attribution. however, many studies @xcite @xcite point out that the computational complexity of forward and backward operations in these methods restricts their applicability to millions of neurons in llms, which are proved as fundamental units for storing knowledge @xcite @xcite . secondly, while a few studies @xcite have devised methods for analyzing neurons, they often lack comparisons with other methods. therefore, how to identify important neurons in llms is still unclear. lastly, existing methods typically concentrate on either attention or feed-forward network (ffn) module, often lacking evaluation of the other module. it is crucial to quantitatively compare the importance of both attention and ffn layers. in this paper, we focus on neuron-level attribution methods. we analyze the distribution change caused by each neuron and discover that both the neuron's coefficient score and the final prediction's ranking, when projecting this neuron's subvalue into vocabulary space, play significant roles. based on this finding, we employ log probability increase as importance score, enabling the identification of neurons that contribute significantly to final predictions. compared with seven other static methods, our proposed method achieves the best performance on three metrics. furthermore, since the identified neurons directly contribute to the final predictions' probability, we also develop a static method to identify "query neurons" that aid in activating these "value neurons". specifically, we calculate the inner products between the query neurons and value neurons as importance scores. based on our proposed methods, we analyze six types of knowledge in both attention and ffn layers, yielding numerous valuable insights (figure 1 ): 1) both attention and ffn layers can store knowledge, and all important neurons directly contribute to knowledge prediction are in deep layers. 2) in attention layers, knowledge with similar semantics (e.g. language, country, city) tends to be stored in the same heads. knowledge with distinct semantics (e.g. country, color) is stored in different heads. 3) while numerous neurons contribute to the final prediction, intervening on a few value neurons (300) or query neurons (1000) can significantly influence the final prediction. 4) ffn value neurons are mainly activated by medium-deep attention value neurons, while these attention neurons are mainly activated by shallow/medium ffn query neurons. overall, our contributions are as follows: a) we design a static method for neuron-level knowledge attribution in large language models. compared with seven static methods, our method achieves the best performance under three metrics. b) as the identified neurons usually directly contribute to the final predictions, we design a static method to identify the "query neurons" activating these "value neurons". c) we analyze the localization of six types of knowledge in both attention and ffn layers. our analysis is helpful for understanding the mechanisms of knowledge storage in language models. ## methodology in this section, we aim to locate important neurons for specific predictions. we introduce the background in section 3.1, and analyze the distribution change caused by neurons in section 3.2. based on the analysis, we introduce our proposed method for locating the "value neurons" that contribute to the final predictions directly in section 3.3, and propose a static method to locate the "query neurons" that activate these "value neurons" in section 3.4. ## experiments in this section, we compare our neuron-level attribution method with seven other methods in section 4.1. then we analyze the localization of six types of knowledge using our method in section 4.2. ## conclusion in this study, we propose a method based on log probability increase to identify the important "value neurons". we also develop a method based on inner products to locate the "query neurons" activating these "value neurons". our method and analysis on six types of knowledge are helpful for exploring and understanding the mechanism of llms.