PaperlayerView exact revision on arXiv

Source-linked AI summary

Information Flow Routes: Automatically Interpreting Language Models at Scale

Javier Ferrando, Elena Voita

arXiv:2403.00824v2cs.CLcs.AI

TL;DR

Activation-patching approaches rely on human-designed templates and repeated interventions, limiting their reach across predictions and making large-model studies impractical. This paper uses attribution to automatically extract prediction-specific information-flow routes. In Llama 2, the analyses identify broadly important heads, shared patterns across same-part-of-speech tokens, and domain-specialized components.

  • Problem

    Activation patching requires human-designed templates and many interventions, restricting it to a few predefined templates and making large-model studies increasingly impractical.

    Source: §1 · p. 2

  • Method

    The paper automatically extracts important information-flow routes for predictions using attribution, rather than activation patching.

    Source: §1 · p. 2, §7 · p. 9

  • Results

    Llama 2 routes identify generally important previous-token and subword-merging heads, similar information flow for tokens with the same part of speech, and domain-specialized components.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The routes support analysis of model behavior across predictions, prediction types, and domains, beyond template-defined patching settings.

    Source: §1 · p. 2, §7 · p. 9

    The part-of-speech analysis uses the dataset's next token rather than a model-generated token, adding noise while assuming their parts of speech are often similar.

    Source: §4.1 · p. 6

Abstract

from arXiv · show

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.

1. Introduction

The paper represents Transformer computations as information-flow graphs and extracts prediction-relevant routes through attribution rather than activation patching. Its Llama 2 analyses identify broadly important attention-head functions, similarities across tokens of the same part of speech, and components specialized to domains.

1. Introduction

  • The method models token representations as graph nodes and model operations as edges, then traces backward from a prediction to retain important information-flow routes.The authors use attribution to assess edge importance rather than activation patching.Source: §1 · p. 1
  • Patching depends on human-designed templates and repeated interventions, whereas attribution-based route extraction is about 100 times faster.Because each intervention requires a forward pass, patching becomes increasingly impractical for large models.Source: §1 · p. 2
  • Information-flow routes recover task-specific attention heads found by patching, while also evaluating overall component importance and comparison with contrastive examples.The authors argue that patching results can vary with the chosen contrastive template.Source: §1 · p. 2
  • Llama 2 routes reveal generally important previous-token and subword-merging heads, similar information flow for tokens with the same part of speech, and components specialized to coding or multilingual text.The specialized components are active for those domains and not otherwise.Source: §1 · p. 2

2. Extracting Information Flow Routes

The paper represents transformer computation as a graph of token representations connected by model operations, then traces prediction-relevant information backward to extract a smaller subgraph. It scores edges by attribution and retains those above a threshold, decomposing attention updates into per-input contributions.

2. Extracting Information Flow Routes

  • The extracted information flow route is a top-down subgraph of the computation graph, retaining edges above a minimum importance threshold τ.Extraction starts from the prediction representation and traces through connected lower nodes, adding important edges and nodes.Source: §2 · p. 2

2.1.1. IMPORTANCE IN GENERAL CASE

  • Edge importance follows ALTI: an incoming vector matters in proportion to its proximity to the resulting sum, with contributions beyond the sum’s L1 length ignored.Proximity uses negative distance; vectors farther from the sum receive less importance.Source: §2.1.1 · p. 3

2.1.2. DEFINING FFN EDGES

  • FFN edge vectors represent the residual connection and feed-forward output as separate contributions to the next representation.The method defines these edge vectors from the corresponding updates.Source: §2.1.2 · p. 3

2.1.3. DEFINING ATTENTION EDGES

  • For attention heads, each input representation is transformed into a raw output and weighted by prediction-specific attention, forming an independent sub-edge into the next residual state.The head-specific output transformation treats attention weights as constants for the prediction.Source: §2.1.3 · p. 3
  • Attention-layer updates combine per-input sub-edges across heads, while also accounting for the current token’s residual connection.Sub-edge importances are aggregated across heads for connections between the same pair of nodes.Source: §2.1.3 · pp. 3–4, Figure 5 · p. 4

2.2. Extracting the Important Subgraph

  • The route-building procedure thresholds edge importances; FFN blocks generally appear more important than attention heads because FFNs are not further decomposed.The authors note that separate thresholds for attention and FFN edges are possible, but they did not test this.Source: §2.2 · p. 4

3. Information Flow vs Patching Circuits

Across IOI and Greater-than, information flow routes recover task-relevant heads while also exposing components and differences that patching’s contrastive comparisons can miss. The method’s comparisons show template sensitivity in patching, and the reported IOI circuit-discovery run takes 5 seconds rather than 8 minutes with ACDC.

3. Information Flow vs Patching Circuits

  • Information flow routes identify all components contributing to a prediction, whereas patching identifies components important relative to a contrastive baseline.The routes can also be compared against contrastive templates.Source: §3 · p. 4, §3.2 · p. 5

3.1. Indirect Object Identification

  • In IOI, routes largely contain previously discovered circuit heads, while also identifying generic heads that disappear in contrastive differences; previous-token heads matter in both prediction types.The route comparison also distinguishes generic components from task-specific heads.Source: §3.1 · p. 4, §3 · p. 4, Figure 7 · p. 5

3.2. Greater-than

  • Greater-than heads remain important with the YY = 01 contrastive template but become less important with YY = 00, illustrating patching’s sensitivity to template choice.For YY = 01, the original task’s important heads also remain important for the contrastive template.Source: §3.2 · p. 5

3.3. OPT-125m vs GPT2-Small

  • Information flow routes for IOI and Greater-than are similar between OPT-125M and GPT2-small.Source: §3.3 · p. 5

3.4. Hundred Times Faster than Patching

  • 5 seconds versus 8 minutes: the method discovers the IOI circuit about 100x faster than ACDC on a batch of 50 examples.The comparison is reported for a single-GPU ACDC run.Source: §3.4 · p. 5

4. General Experiments

Using component-importance patterns in Llama 2-7B, the section finds token- and position-sensitive processing, generally important attention heads, and first-period residual streams that can act like BOS tokens.

4. General Experiments

  • The analysis uses 1,000 C4 sentences and projects per-prediction component-importance vectors with t-SNE, coloring points by input or next-token part of speech.Contribution patterns depend more on input-token than next-token part of speech; next-token tags come from the dataset rather than model-generated tokens.Source: §4 · p. 6, §4.1 · p. 6

4.1. Component Importance for POS

  • For function-word inputs, component contributions cluster by part of speech; content-word patterns are more mixed, and first versus later subwords have distinct contribution patterns.Verbs and nouns are somewhat separated, but content words overall form a large mixed cluster; the first-versus-later distinction also separates number tokens into first and later digits.Source: §4.1 · p. 6

4.2. Bottom-to-Top Patterns

  • Attention heads and FFN blocks are more active in lower layers, while the last-layer FFN is highly important, substantially changing the representation before prediction.The authors interpret lower-layer activity as processing inputs and extracting general information, with later operations more fine-grained and specialized.Source: §4.2 · p. 6, Figure 9 · p. 7

4.3. Positional and Subword Merging Heads

  • Several previous-token heads are among the most important attention heads in their layers, while bottom-layer subword-merging heads are important specifically for later subwords.Previous-token heads put more than half of their influence on the preceding token in at least 70% of cases; subword-merging heads draw information from earlier subwords of the same word.Source: §4.3 · p. 7

4.4. Peculiar Information Flow Patterns, or Periods Acting as BOS

  • For first-period examples, contributions from most attention and FFN blocks are near zero, with processing largely carried by the residual stream, which can act as BOS through the last three layers.This also occurs when the period is not sentence-final; the authors suggest future work examine whether this behavior might cause incorrect generation.Source: §4.4 · pp. 7–8

5. Domain-Specific Model Components

Across Llama 2-7B domains, component importance varies by domain and task. Projections of specialized heads’ outputs show topic-related concepts consistent with their activation patterns.

5.1. Components are Specialized

  • Important heads differ across addition, code, and non-English inputs, and heads important for addition can score among the lowest on C4.Addition and subtraction share many important heads but also have heads active for only one task; the authors suggest this fine-grained specialization might relate to reasoning, pending further validation.Source: §5.1 · p. 8, §5 · p. 8
  • Feed-forward block importance also varies by domain: the last FFN is less relevant for non-English inputs, and FFN importance falls to zero in some addition and subtraction layers.The analysis averages component importance over 1,000 sentences from C4, four FLORES-200 language sets, CodeParrot code data, and addition/subtraction data.Source: §5.1 · p. 8, §5 · p. 8

5.2. Specialized Heads Output Topic-Related Concepts

  • Code-specific heads promote coding- and technology-related tokens, while heads most active for non-English inputs promote multiple languages and avoid English tokens.The non-English heads also promote tokens related to cities, countries, and currencies; the authors report that activation patterns and residual-stream updates give consistent views of head functions.Source: §5.2 · p. 9, Figure 12 · p. 9

6. Related Work

Prior work identified important neurons and attention heads, then increasingly studied task-specific language-model components using activation patching. This paper relates to work approximating patching.

6. Related Work

  • Earlier studies identified important neurons shared across NMT models and assessed Transformer attention-head relevance through attribution and pruning.The passage cites Bau et al. (2019) for neurons and Voita et al. (2019) for attention heads.Source: §6 · p. 9
  • Later work focused on task-specific LM components, using activation patching to study GPT-2 IOI and greater-than circuits, factual knowledge, and in-context-learning task vectors.The passage also notes activation patching’s introduction for analyzing LM biases by Vig et al. (2020).Source: §6 · p. 9
  • The paper is related to concurrent methods that approximate activation patching.The cited approaches are Syed et al. (2023), Kramár et al. (2024), and Hanna et al. (2024).Source: §6 · p. 9

7. Conclusions

The paper represents Transformer computation as information flowing between token representations through model components. It automatically extracts important routes for predictions.

7. Conclusions

  • The approach models Transformer computation as information flow between token representations through model components.It extracts the important part of that flow to interpret predictions.Source: §7 · p. 9
  • The method automatically extracts important information-flow routes and is described as highly efficient.Source: §7 · p. 9
  • The method applies to any prediction and is described as more versatile and informative than existing pipelines.Source: §7 · p. 9

A. Background on Activation Patching

Activation patching identifies circuits by replacing internal activations with those from a contrastive input, but template dependence and repeated interventions constrain its use. The section also describes layer normalization as a constant affine map when its standard deviation is treated as constant.

A. Background on Activation Patching

B.1. Linearizing the Layer Normalization

  • Treating σ(x) as constant makes layer normalization a constant affine transformation, with L combining centering, normalization, and scaling.The linear map subtracts the input mean, and the left matrix applies the layer-normalization weights element-wise.Source: §B.1 · p. 14

B.2. Subword Merging Head definition

A subword merging head is identified by directing influence from later subwords toward earlier subwords of the same word, while having little influence on the first subword.

B.2. Subword Merging Head definition

  • In at least 70% of cases, later subwords direct more than half their influence to earlier subwords of the same word, while first-subword influence is at most 0.005%.Source: §B.2 · p. 14

C. Examples of Routes

The appendix shows extracted information flow routes for two tasks, IOI and greater-than, using GPT2-small and OPT-125m.

C. Examples of Routes

D. Period Acting as BOS

An example Llama 2-7b attention map shows attention after the first period spread between the BOS token and that period.

D. Period Acting as BOS

Related papers

Loading 2403.00824v2…