Information Flow Routes: Automatically Interpreting Language Models at Scale (2024)

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 computations. 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. Unlike with 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 analyze model behavior in general, for specific types of predictions, or different domains. We experiment with Llama 2 and show that some attention head roles are 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.

## introduction current state-of-the-art language models (lms) are built on top of the transformer architecture @xcite @xcite . inside the model, each representation evolves from the current input token embedding to the final representation used to predict the next token. this evolution happens through additive updates coming from attention and feed-forward blocks. the resulting stack of same-token representations is usually referred to as * work done during an internship at meta ai. figure 1 : the important information flow routes for a token (mary) prediction. gpt2-small, τ = 0.04. while during a forward pass all the edges are present, computations important for each prediction are likely to form a small portion of the original graph @xcite @xcite . we extract this important subgraph in a top-down manner by tracing information back through the network and, at each step, leaving only edges that were relevant (figure 1 ). to understand which edges are important, we rely on an attribution method @xcite and refuse from activation patching typical for the existing mechanistic interpretability workflows @xcite @xcite . firstly, patching requires human efforts to create templates and contrastive examples, thus is only applicable to a few pre-defined templates. secondly, explaining a single prediction demands a substantial number of interventions @xcite . given that each intervention needs a forward pass, studying large models becomes increasingly impractical. in contrast, our method is about 100 times faster. in the experiments, we first show that our information flow routes rely on the same task-specific attention heads found in patching circuits @xcite . however, our method is more versatile and informative than patching: it can evaluate the importance of model components both (i) overall for a prediction, and (ii) compared to a contrastive example (i.e., patching setting). aditionally, we argue that patching is fragile: its results can vary depending on the choice of the contrastive template (i.e., human preference). next, we come to the settings unreachable to patching, i.e. broad set of predictions and overall importance of model components. for llama 2, we show that some attention heads' functions are overall important, e.g. previous token heads and subword merging heads. then, we find that information inside llama 2 flows similarly when handling tokens of the same part of speech. finally, some model components are specialized on domains such as coding or multilingual texts. overall, our contributions are as follows: extracting the important subgraph. while during a forward pass all edges in figure 2 are present, computations relevant for each prediction are likely to form a small portion of the original graph @xcite . we extract this important subgraph, the information flow routes, in a top-down manner by tracing information back through the network (figure 3 ). we start from a single node -the representation on top of the residual stream. then, we go over immediately connected lower nodes and, if important at this step, we add them to the subgraph along with the corresponding edges. the algorithm requires setting a threshold τ for the minimum edge importance. ## information flow vs patching circuits first, we compare our information routes to the circuits found in previous work for the indirect object identification (ioi) @xcite and greater-than @xcite tasks. ## general experiments information flows inside the network via mechanisms implemented in the model. in this section, we try to understand whether the important mechanisms depend on the properties of processed by the model tokens and we look at general patterns in component importances. ## domain-specific model components we saw that some attention heads' roles are generally important across predictions. now, let us see whether not universally important components might be specialized for specific domains. setting. we evaluate the average importance of llama 2-7b's components on 1000 sentences from different domains. specifically, we consider (i) c4 @xcite , (ii) different languages: english, russian, italian, spanish flores-200 devtest sets @xcite , (iii) code data from codeparrot, (iv) addition/subtraction data @xcite . ## additional related work earlier works evaluating the importance of model components include @xcite who identified relevant neurons in @xcite who looked at the relevance of entire attention heads in the transformer. instead of overall importance, recent research largely focuses on specific tasks and aims to find the important lm subparts via activation patching @xcite . patching has been also used to locate factual knowledge in lms @xcite and to discover task vectors behind in-context learning capabilities @xcite . our work can be related to the concurrent line of research developing methods to approximate activation patching @xcite @xcite . ## conclusions we view computations inside the transformer as information flowing between token representations through model components. using this view, we propose to interpret language model predictions by extracting the important part of the overall information flow. our method for extracting these impor-tant information flow routes is automatic, highly efficient, applicable to any prediction, more versatile and informative compared to existing pipelines.