Interpreting Context Look-ups in Transformers: Investigating Attention- MLP Interactions (2024)

Abstract: Understanding the inner workings of large language models (LLMs) is crucial for advancing their theoretical foundations and real-world applications. While the attention mechanism and multi-layer perceptrons (MLPs) have been studied independently, their interactions remain largely unexplored. This study investigates how attention heads and next-token neurons interact in LLMs to predict new words. We propose a methodology to identify next-token neurons, find prompts that highly activate them, and determine the upstream attention heads responsible. We then generate and evaluate explanations for the activity of these attention heads in an automated manner. Our findings reveal that some attention heads recognize specific contexts relevant to predicting a token and activate a downstream token-predicting neuron accordingly. This mechanism provides a deeper understanding of how attention heads work with MLP neurons to perform next-token prediction. Our approach offers a foundation for further research into the intricate workings of LLMs and their impact on text generation and understanding.

## introduction transformer-based models have made significant advancements in natural language understanding and generation. however, the inner workings of large language models (llms) remain largely unexplored, particularly when their behavior deviates from common beliefs @xcite or when modifying their factual knowledge proves challenging @xcite . a comprehensive understanding of these models is essential for expanding their theoretical foundations, improving real-world applications, mitigating biases, and reducing potential risks. this study aims to provide insight into the mechanistic interpretability of gpt-2 (radford et al., * equal contribution. ## related work we survey previous work about interpreting language models, with primary focus on attention heads and mlp layers as separate entities. ## results and discussion we first identify meaningful attention head activity ( §5.1), and perform baseline comparisons ( §5.2). we then show that this can also be found in smaller model variants ( §5.3), we verify the relationship between the attention head and the neuron through ablation in §5.4. we also discuss the interpretability illusion in this context ( §5.5) before discussing the performance of gpt-4 as an auto-labeler in these experiments ( §5.6). ## conclusion we investigated attention heads and their role in activating next-token neurons. we found that some attention heads would only activate their associated next-token neurons in a subset of max-activating prompts. we then use gpt-4 to generate explanations for this activity, and made gpt-4 use that explanation for a classification task on a test-set of prompts to evaluate explanation quality. our findings reveal that attention heads process context and use next-token neurons to help predict a word associated with, but not previously seen in that context. this provides insight into how attention heads can work with mlp neurons to predict new words. future work could investigate whether this phenomenon happens to such neurons throughout their range of activations, which could give insights to the role of next-token neurons when the predicted token is unlikely to be the neuron's associated token. further research could also look into whether multiple attention heads could work with a single next-token neuron in a similar manner, or whether a single attention head could work with multiple next-token neurons in a similar manner. this could provide insights to whether such attention-head neuron interactions extend to more complex contexts.