Interpreting Arithmetic Mechanism in Large Language Models through Comparative Neuron Analysis (2024)

Abstract: We find arithmetic ability resides within a limited number of attention heads, with each head specializing in distinct operations. To delve into the reason, we introduce the Comparative Neuron Analysis (CNA) method, which identifies an internal logic chain consisting of four distinct stages from input to prediction: feature enhancing with shallow FFN neurons, feature transferring by shallow attention layers, feature predicting by arithmetic heads, and prediction enhancing among deep FFN neurons. Moreover, we identify the human-interpretable FFN neurons within both feature-enhancing and feature-predicting stages. These findings lead us to investigate the mechanism of LoRA, revealing that it enhances prediction probabilities by amplifying the coefficient scores of FFN neurons related to predictions. Finally, we apply our method in model pruning for arithmetic tasks and model editing for reducing gender bias. Code is on https://github.com/zepingyu0512/arithmetic-mechanism.

## introduction arithmetic ability is a crucial foundational skill of large language models (llms) @xcite @xcite , contributing significantly to reasoning @xcite and mathematical tasks @xcite . while existing studies @xcite @xcite have made significant breakthroughs in understanding arithmetic tasks, the exact mechanism still remains elusive. @xcite find that only a few attention heads significantly impact arithmetic performance, but they do not elaborate on the mechanisms of these heads or how they influence ffn layers. @xcite intervene the hidden states and find the information flow from number and operation positions to the last position. however, they do not locate the important attention heads (proved to store different abilities @xcite ) and ffn neurons (proved to store knowledge @xcite ). despite the challenge of pinpointing important ffn neurons among tens of thousands of nodes, many studies @xcite @xcite emphasize that considering ffn neurons as fundamental units is crucial for better understanding ffn layers. furthermore, as model editing typically occurs at the neuron level @xcite , it remains unclear how to effectively leverage the explanations due to the uncertainty surrounding the precise locations of important parameters. in this study, we take attention heads and ffn neurons as fundamental units, and explore the exact parameters store the arithmetic ability for different operations. we observe that only a minority of heads play significant roles in arithmetic tasks, which we refer to as "arithmetic heads". through experiments involving 1-digit to 3-digit operations, as well as ablation studies comparing "change-one" cases (e.g., 15+37=52) with "memorize" cases (e.g., 15+32=47), we find critical memorization of 1-digit operations is lost when these heads are intervened. to explore the underlying mechanisms of this phenomenon, we propose the comparative neuron analysis (cna) method, which compares the change of neurons between the original model and the intervened model for the same case. we construct the internal logic chain by identifying four distinct stages that span from inputs to prediction, as depicted in figure 1 . during the feature enhancing stage, hidden-interpretable features are extracted from shallow ffn neurons. subsequently, in the feature transferring stage, shallow attention layers convert these features into directlyinterpretable features and then transfer them to the last position. in the feature predicting stage, the arithmetic heads play critical roles, activating deep ffn neurons related to the final prediction. finally, a prediction enhancing stage exists among deep ffn neurons. lower ffn neurons activate upper ffn neurons, while both of them enhance the probability of the final prediction. based on this analysis, we investigate the mechanism of @xcite . we train a total of 32 models on a 2-digit arithmetic dataset, with each model integrating lora on one attention layer (0th to 31th). starting from the 10th model, the accuracy of the model exhibits a noticeable downward trend, with varying rates of decline observed in the feature enhancing and prediction enhancing stages. employing our cna method to compare the original model with the fine-tuned model, we note a significant increase in the coefficient scores of crucial deep ffn neurons. hence, we conclude that lora enhances the final prediction by amplifying the coefficient scores of important ffn neurons. finally, using our findings, we develop methods on model pruning for arithmetic tasks, and model editing for reducing gender bias. to summarize, our contributions are as follows: ## arithmetic heads in llms we aim to examine the localization of arithmetic ability in llama-7b @xcite , a large language model consisting of 32 layers. each attention layer contains 32 heads, and each ffn layer has 11,008 neurons. we observe the same results and mechanisms in gpt-j @xcite , detailed in appendix c. ## comparative neuron analysis for mechanistic interpretability in this section, we investigate how head 17 22 influence 1d+ and 1d-operations. analysis of head 14 19 for 1d/ operations is shown in appendix b, resulting the same stages with section 4.2-4.4. ## understanding the mechanism of lora lora @xcite is a commonly used parameter-efficient fine-tuning method @xcite @xcite . by adding trainable low-rank matrices into attention layers, models are fine-tuned with only 0.5% additional parameters, yielding favorable outcomes. intuitively, lora is similar to a head. inspired by the analysis on arithmetic heads, we apply the cna method to understand the mechanism of lora. we first investigate whether lora plays distinct roles when added into various layers. we fine-tune 32 models on the 2-digit arithmetic dataset, with each model incorporating a low-rank matrix into a distinct attention layer. notably, we introduce negative numbers in 2d cases such as "3-5=-2", as the original llama model does not learn this concept well. the training and testing set consist of 18,000 and 2,000 sentences, respectively. we determine the optimal learning rate from choices of 0.001, 0.0005, and 0.0001. the maximum epoch is set to 4. the results are depicted in figure 2 . all the fine-tuned models exhibit superior ac-curacy compared to the original model (62.96%). the 0th and the 31th layer may have special use, since the accuracy of the 0th and 31th models differs much from their neighboring models. the accuracy of the 1st -9th models is around 90%. starting from the 10th model, the accuracy keeps decreasing. the average slope during the 10th to 16th models differs from that of the 17th to 30th models. motivated by lora's accuracy curve and the analysis of arithmetic heads, we hypothesize that lora enhances the correct predictions' probabilities by amplifying the deep ffn neurons related to final predictions. we apply our cna method on the original model and five lora models analyzing the case "3+5=", detailed in across all five fine-tuned models, the coefficient scores of 25 7164 and 19 5769 surpass those of the original model. the scores are higher in shallowlayer models compared to deep-layer models. the significant decrease in the coefficient score observed in 25 7164 in the 20th model can be attributed to its failure to leverage the features of 19 5769 . for all cases, we compute the average coefficient score increase of 1st -9th, 10th -16th, and 17th-30th models on the most important neurons, detailed in ## applications in this section, we utilize our method for model pruning on arithmetic tasks and for model editing aimed at mitigating gender bias. ## discussion and conclusion we aim to discuss the mechanisms behind causal mediation analysis and static interpretation methods. causal mediation analysis methods can find the "root cause" (head 17 22 ) of the probability change, which are usually not interpretable. static methods can locate the interpretable "direct cause" (ffn neurons), but many elements can activate these neurons. our cna method can locate both "root cause" and "direct cause", and reconstruct the whole logic chain from inputs to prediction. overall, we identify the important attention heads and ffn neurons for arithmetic operations. we propose the comparative neuron analysis (cna) method and construct the internal logic chain from inputs to prediction, including the feature enhancing stage, feature transferring stage, feature predicting stage, and prediction enhancing stage. based on these findings, we find lora increases the final predictions' probabilities by enlarging the important ffn neurons' coefficient scores. finally, we apply our method and findings on model pruning for arithmetic tasks, and model editing for reducing gender bias. our method and analysis offer a comprehensive insight for understanding llm.