Flash Interpretability: Decoding Specialised Feature Neurons in Large Language Models with the LM-Head
This provides a fast, low-compute method for interpretability in LLMs, though it is incremental as it builds on existing neuron analysis techniques.
The paper tackles the challenge of interpreting large language models by decoding neuron weights into token probabilities via the LM-head, enabling identification of specialized feature neurons like 'dog' and 'California' and validating this by clamping neurons to affect output probabilities, with results showing over 75% consistency in top associated tokens between pre-trained and instruct models and mapping all up-projection neurons in under 10 seconds.
Large Language Models (LLMs) typically have billions of parameters and are thus often difficult to interpret in their operation. In this work, we demonstrate that it is possible to decode neuron weights directly into token probabilities through the final projection layer of the model (the LM-head). This is illustrated in Llama 3.1 8B where we use the LM-head to find examples of specialised feature neurons such as a "dog" neuron and a "California" neuron, and we validate this by clamping these neurons to affect the probability of the concept in the output. We evaluate this method on both the pre-trained and Instruct models, finding that over 75% of neurons in the up-projection layers in the instruct model have the same top associated token compared to the pretrained model. Finally, we demonstrate that clamping the "dog" neuron leads the instruct model to always discuss dogs when asked about its favourite animal. Through our method, it is possible to map the top features of the entirety of Llama 3.1 8B's up-projection neurons in less than 10 seconds, with minimal compute.