Interpreting Arithmetic Mechanism in Large Language Models through Comparative Neuron Analysis

TL;DR

This paper introduces the Comparative Neuron Analysis method to interpret how large language models perform arithmetic, revealing specific neuron roles and internal logic stages, and explores LoRA's enhancement mechanism, with applications in pruning and bias reduction.

Contribution

It presents a novel interpretability method that uncovers the internal logic and neuron roles in arithmetic reasoning within large language models, and analyzes LoRA's effect on predictions.

Findings

Arithmetic ability is localized in specific attention heads.

Four stages of internal logic from input to prediction are identified.

LoRA amplifies the contribution of relevant FFN neurons.

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.