Information Flow Routes: Automatically Interpreting Language Models at Scale

Abstract: Information flows by routes inside the network via mechanisms implemented in the model. These routes can be represented as graphs where nodes correspond to token representations and edges to operations inside the network. We automatically build these graphs in a top-down manner, for each prediction leaving only the most important nodes and edges. In contrast to the existing workflows relying on activation patching, we do this through attribution: this allows us to efficiently uncover existing circuits with just a single forward pass. Additionally, the applicability of our method is far beyond patching: we do not need a human to carefully design prediction templates, and we can extract information flow routes for any prediction (not just the ones among the allowed templates). As a result, we can talk about model behavior in general, for specific types of predictions, or different domains. We experiment with Llama 2 and show that the role of some attention heads is overall important, e.g. previous token heads and subword merging heads. Next, we find similarities in Llama 2 behavior when handling tokens of the same part of speech. Finally, we show that some model components can be specialized on domains such as coding or multilingual texts.

• The paper introduces an efficient attribution method that reduces computational time by 100x to trace key information routes in Transformer models.
• It employs a graph-based subgraph extraction approach to identify influential model components driving specific predictions.
• Empirical findings reveal pivotal roles of specific attention heads, which improves model debugging and enhances interpretability across varied domains.

Interpreting LLMs Through Information Flow

Overview

In the paper titled "Information Flow Routes: Automatically Interpreting LLMs at Scale," authors Javier Ferrando and Elena Voita offer a novel method for understanding how information flows inside LLMs, particularly those built on the Transformer architecture. Their approach creates visualizable "information flow routes" that help trace which parts of a model are most relevant for making specific predictions. This approach aims to provide a clearer interpretation of the model's internal mechanisms, moving away from traditional activation patching techniques, towards a more efficient and automatic method using attribution.

Understanding Information Flow in Transformers

Transformers process input through layers of computations, where each layer includes operations such as attention and feed-forward mechanisms. Traditionally, interpreting these operations requires a detailed dissection of which nodes (token representations) and edges (operations) in the network graph are most active during prediction. The authors of this paper introduce a methodology that tracks the flow of information through these networks more systematically.

• Graph Representation: The model's operations are represented as a graph where nodes are token embeddings and edges reflect the operations like attention links or feed-forward actions.
• Subgraph Extraction: For every output prediction, the algorithm identifies a subgraph of this larger graph, focusing on nodes and edges that significantly influence the final outcome, known as "important subgraphs."

A Novel Approach to Extract Important Subgraphs

Unlike traditional methods that involve intensive computation to understand model behavior (often requiring multiple forward passes through the model), this paper proposes using an attribution method to simplify the process.

• Efficient Attribution Over Activation Patching: The proposed approach takes roughly 100 times less computational time than older methods. It works by tracking back from an output layer to determine which parts of the network contributed most significantly to the decision.
• Versatility and General Applicability: This tool is not limited by preset conditions and can be applied to a broader range of predictions, spanning various domains and language tasks.

Empirical Insights and Implications

The application on a version of the Llama model yields some insightful observations:

• Relevance of Specific Attention Heads: Certain attention heads, like those tracking previous tokens or merging subwords, consistently play pivotal roles across different contexts and tasks.
• Domain-Specific Components: When testing across various domains (code, multilingual texts), some model components are frequently more active, suggesting specialization towards particular types of data or tasks.

Looking Forward

The ability to map out and understand these information flow routes has important implications for both theoretical and applied AI research:

• Model Debugging and Improvement: By pinpointing which areas of a model are most active for particular tasks, developers can better understand unexpected model behavior and improve model architectures.
• Enhanced Interpretability: This method provides a more granular look at how decisions are made within LLMs, which is crucial for applications requiring transparency and accountability.

Conclusions

The development of information flow routes offers a promising direction towards demystifying the often opaque internal workings of large-scale LLMs. This approach not only enhances our understanding of these complex models but also opens the door to more targeted and efficient model optimization and debugging strategies.

By providing a clear, efficient, and scalable method to visualize and interpret the roles of various components in LLMs, this research takes a significant step forward in the field of machine learning interpretability. Future research may expand these techniques to other model architectures or integrate this understanding into the development of next-generation AI systems.