Attention Lens: A Tool for Mechanistically Interpreting the Attention Head Information Retrieval Mechanism
This addresses the need for mechanistic interpretability in AI research, providing a tool for researchers to decode attention mechanisms, though it is incremental as it builds on existing reverse engineering efforts.
The paper tackles the problem of understanding the specific roles of attention heads in transformer-based LLMs by proposing Attention Lens, a tool that translates attention head outputs into vocabulary tokens using learned transformations, with preliminary findings indicating that attention heads have highly specialized functions.
Transformer-based Large Language Models (LLMs) are the state-of-the-art for natural language tasks. Recent work has attempted to decode, by reverse engineering the role of linear layers, the internal mechanisms by which LLMs arrive at their final predictions for text completion tasks. Yet little is known about the specific role of attention heads in producing the final token prediction. We propose Attention Lens, a tool that enables researchers to translate the outputs of attention heads into vocabulary tokens via learned attention-head-specific transformations called lenses. Preliminary findings from our trained lenses indicate that attention heads play highly specialized roles in language models. The code for Attention Lens is available at github.com/msakarvadia/AttentionLens.