Interpreting Arithmetic Mechanism in Large Language Models through Comparative Neuron Analysis
This work provides insights into the interpretability of neural mechanisms in LLMs, with incremental contributions to model optimization and fairness.
The study identified that arithmetic ability in large language models is localized to specific attention heads and introduced Comparative Neuron Analysis to reveal a four-stage internal logic chain, leading to applications in model pruning and bias reduction.
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.