PaperlayerView exact revision on arXiv

Source-linked AI summary

AtP*: An efficient and scalable method for localizing LLM behaviour to components

János Kramár, Tom Lieberum, Rohin Shah, Neel Nanda

arXiv:2403.00745v1cs.LGcs.CL

TL;DR

Exhaustive causal attribution can be prohibitively expensive for large language models, so this paper studies faster approximations to Activation Patching. It finds that AtP outperforms the investigated alternatives, proposes AtP* to reduce false negatives, and supplies a diagnostic for remaining errors.

  • Problem

    Exhaustively attributing behavior to model components can be prohibitively expensive because the number of components may reach billions.

    Source: §1 · p. 1

  • Method

    The paper evaluates Attribution Patching as a fast prefilter for Activation Patching and proposes AtP* with two changes targeting AtP’s false-negative modes.

    Source: §1 · pp. 1–2, §3.1 · p. 5

  • Results

    AtP significantly outperforms the investigated alternatives, with AtP* providing further significant improvement.

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

  • Takeaways & Limitations

    A diagnostic using AtP filtering and random patching can estimate residual error and bound the probability of remaining false negatives without exhaustive verification.

    Source: §1 · p. 2

    The evaluation used a small set of prompt-pair distributions, often a single pair, so results may not generalize to other distributions.

    Source: §5.1 · p. 17

Abstract

from arXiv · show

Activation Patching is a method of directly computing causal attributions of behavior to model components. However, applying it exhaustively requires a sweep with cost scaling linearly in the number of model components, which can be prohibitively expensive for SoTA Large Language Models (LLMs). We investigate Attribution Patching (AtP), a fast gradient-based approximation to Activation Patching and find two classes of failure modes of AtP which lead to significant false negatives. We propose a variant of AtP called AtP*, with two changes to address these failure modes while retaining scalability. We present the first systematic study of AtP and alternative methods for faster activation patching and show that AtP significantly outperforms all other investigated methods, with AtP* providing further significant improvement. Finally, we provide a method to bound the probability of remaining false negatives of AtP* estimates.

1. Introduction

Exhaustive attribution of model behavior can be prohibitively costly, motivating faster ways to locate causally important components. The paper evaluates AtP and alternatives, proposes AtP* to address AtP false negatives, and offers diagnostics and practical guidance.

1. Introduction

  • Exhaustive attribution of small components in state-of-the-art models can require millions or billions of forward passes.The introduction cites 2.7 · 10^9 neuron nodes for Chinchilla 70B on a 1024-token prompt.Source: §1 · p. 1
  • AtP prefilters nodes by the absolute estimate, then Activation Patching verifies candidates and filters false positives.Because the goal is often a small set of top contributors, verification can be cheaper than checking every node.Source: §1 · pp. 1–2
  • AtP has two classes of false negatives; AtP* addresses them while retaining scalability.For query and key nodes, it recomputes attention softmax before applying a gradient approximation; backward-pass dropout addresses brittle cancellations of positive and negative effects.Source: §1 · p. 2
  • AtP significantly outperforms the investigated alternatives, and AtP* provides further significant improvement.The paper also introduces alternative approximation methods that outperform brute-force Activation Patching.Source: §1 · p. 2
  • A diagnostic filters high-impact nodes with AtP, patches random subsets of the remainder, and bounds the probability of residual false negatives.The authors say this can help practitioners gauge AtP reliability without exhaustive verification.Source: §1 · p. 2
  • The paper also offers guidance on performing causal attribution and choosing useful attribution methods for different circumstances.The introduction points readers to practical guidance on when methods are likely to be useful.Source: §1 · p. 2

2. Background

The background formalizes component-level causal attribution and explains why exhaustive evaluation is costly. It then introduces Attribution Patching as a gradient-based approximation to Activation Patching.

2.1. Problem Statement

  • The choice of graph nodes sets attribution granularity: layer-level nodes are coarser, while finer-grained components cost more to evaluate.The paper primarily studies fine-grained attributions.Source: §2.1 · p. 2
  • The setup models a decoder-only transformer as a computational graph and defines behavior over clean and reference prompts, with component contribution measured by counterfactual impact.The formalism is also described as straightforwardly applicable to other model classes, though experiments use decoder-only transformers.Source: §2.1 · pp. 2, 4

2.2. Attribution Patching

  • Computing each node’s contribution across all prompt pairs can be prohibitively expensive on state-of-the-art models with billions of nodes.Precise evaluation scales with the support of the prompt-pair distribution.Source: §2.2 · p. 4
  • Attribution Patching uses a first-order Taylor approximation and takes the absolute value inside the expectation to reduce cancellation across prompt pairs.It requires two forward passes and one backward pass to estimate all nodes for a prompt pair.Source: §2.2 · pp. 4–5

3. Methods

The section diagnoses two classes of AtP false negatives, introduces query/key corrections and GradDrop, and describes subset-sampling diagnostics and alternative baselines.

3.1.1. False negatives from attention saturation

  • On Pythia-12B CITY-PP, both AtP improvements reduce false negatives, with most gains from the QK fix; many missed keys and queries show large attention-probability changes.For query and key nodes, the correction recomputes attention weights before applying a gradient approximation to the rest of the model.Source: §3.1.1 · p. 6, Figure 4 · p. 8

3.1.2. False negatives from cancellation

  • GradDrop disrupts direct/indirect-effect cancellation by zeroing downstream-layer gradients iteratively, improving estimates for the largest neuron false negatives.The corrected neurons had direct-effect ratios of 5.35, 12.2, and 0, compared with a median of 0 across all nodes.Source: §3.1.2 · pp. 7–8, Figure 5 · p. 9

3.2. Diagnostics

  • Subset sampling can bound the effect magnitudes of nodes outside AtP*’s top 1024; bounds can tighten with more samples, subject to Welch-test assumptions.The diagnostic relies on populations being approximately unbiased, normally distributed, and not skewed; a nonparametric bootstrap may be more reliable but costs more compute.Source: §3.2 · pp. 9–10, Figure 6 · p. 10

3.3. Baselines

  • Iterative patching tests nodes in random order and is inefficient; Blocks and Hierarchical instead prioritize high-contribution groups, trading traversal cost against the work to finish all groups.Grouping-based methods also scale with the size of the prompt-pair distribution’s support, as well as the number of nodes.Source: §3.3 · pp. 10–12
  • For prompt-pair distributions, alternating node measurements with verification ranks unverified nodes by average observed effect, helping find large nodes sooner when their effects appear across many pairs.The approach balances compute between measuring nodes on prompt pairs and verifying candidates.Source: §3.3 · p. 10
  • Under approximate node additivity, subset sampling estimates node effects using Bernoulli-sampled node sets and prompt pairs; it is unbiased without interactions and has bias proportional to p under a simple interaction model.The method’s summary statistics also support bounding false-negative magnitudes.Source: §3.3 · pp. 10–11

4. Experiments

The experiments assess attribution methods by the compute needed to verify high-impact nodes, across model sizes, node granularities, prompt pairs, and distributions. AtP performs well on distributions, especially with the QK fix, while averaging across prompts diminishes GradDrops’ benefit.

4.1. Setup

  • Experiments use Pythia models from 410M to 12B parameters and compare neuron nodes with attention nodes, defined at token-position and site granularity.The reported loss is negative log probability on the clean prompt’s target; the authors briefly explore other metrics in Appendix B.4.Source: §4.1 · p. 12

4.2. Measuring Effectiveness and Efficiency

  • Efficiency is measured by total forward-pass cost to verify the top K nodes, combining estimation cost with the cost of verifying nodes in estimated-magnitude order.The study focuses on finding high-effect nodes; false positives can be screened during verification, whereas false negatives are harder to remedy without checking all nodes.Source: §4.2 · pp. 12–13

4.3. Single Prompt Pairs versus Distributions

  • The clean-pair experiments include IOI-PP and CITY-PP, chosen to represent attention-head-heavy behavior and factual recall, respectively.The study reports 100% verified-recall costs for these settings and also examines 90% recall; Figure 7 notes that AtP* false negatives are a small minority of nodes.Source: §4.3 · p. 13, Figure 7 · p. 14
  • Distributional tests use a 120-pair reduced IOI setting with AttentionNodes and an “a”/“an” task with NeuronNodes.AtP, AtP*, and Subsampling scale to distributions; AtP methods are inexpensive to run per pair, while Subsampling averages across the distribution.Source: §4.3 · p. 14
  • Across prompt-pair distributions, AtP performs well, especially with the QK fix, while averaging across prompts diminishes GradDrops’ benefit.The authors also report that Subsampling may make false-negative diagnostics reasonably confident with less overhead than in the single-pair setting.Source: §4.3 · pp. 14, 17

5. Discussion

The discussion qualifies the findings by the tested models, prompt distributions, node choices, and evaluation target. It also considers extensions and applications, then recommends method choices according to attribution granularity and the confidence–compute tradeoff.

5.1. Limitations

  • Results may not generalize beyond the small set of prompt-pair distributions or the Pythia family, especially to SoTA-scale or non-standard decoder-only models.The authors also note that the study evaluates fine-grained nodes; AtP may be less suitable for coarser components such as layers or residual streams.Source: §5.1 · pp. 17–18
  • The activation-patching target can disagree with direct effect, and distributional averaging can underweight nodes that matter only for a particular prompt.The study evaluates relative rank rather than effect-size accuracy, filters false positives through verification, and does not compare its selected nodes with human-discovered ones.Source: §5.1 · p. 18

5.2. Extensions/Variants

  • AtP can extend to edge attribution, and AtP* may improve it, but GradDrops is omitted from recommended edge-AtP* because its benefit is lessened while its cost remains substantial.The authors leave investigation of coarser nodes, layer-normalization variants, and denoising to future work; they describe other ablation forms as possible applications of their tweaks.Source: §5.2 · p. 19

5.3. Applications

  • AtP* may aid circuit discovery and targeted steering, while SAE sparsity creates a tradeoff between active-node counts and reconstruction error.Some cited SAEs have 10–20 active features per 500 neurons at a token position, reducing node counts 20–50× relative to the MLP setting; feasible sparsity at relevant or SoTA scale remains open.Source: §5.3 · pp. 19–20

5.4. Recommendation

  • Method choice depends on localization granularity and confidence versus compute: the authors favor ordinary patching for coarser components and recommend Blocks for single pairs or Subsampling for distributions when forward-pass costs remain prohibitive.When using AtP, they recommend AtP* for single pairs, AtP+QKFix for AttentionNodes on distributions, and AtP for NeuronNodes on distributions.Source: §5.4 · p. 20
  • A practitioner can add a diagnostic after AtP when greater confidence is needed, but the useful conclusions remain bounded by whether the prompt distribution captures the behavior of interest.The authors say their recommendations are best substantiated for focused prompt pairs or distributions, attention or neuron sites, nodewise attribution, and clean-prompt next-token cross-entropy.Source: §5.4 · p. 20

6. Related work

The related work situates activation patching within causal mediation and distinguishes low-level node localization from work on higher-level causal abstractions. It also contrasts this study’s analysis of AtP and its alternatives with concurrent circuit-finding work.

Activation Patching

  • Activation patching intervenes on a component’s output and has been widely applied to identify causally important nodes in neural networks.The paper places it in the broader causal-mediation tradition, where intervening on intermediate nodes reveals their effects.Source: §6 · p. 21

Localization and Mediation Analysis

  • Unlike work on high-level causal abstractions, this paper targets important low-level variables in the computational graph rather than their semantics or groupings.The cited causal-abstraction studies include Geiger et al. (2020, 2021, 2022, 2023).Source: §6 · p. 21

Attribution Patching / Gradient-based Masking

  • Compared with concurrent work using AtP for circuit finding, this study analyzes AtP failure modes, proposes AtP*, and evaluates methods against a ground truth independent of human judgment.The paper also notes that concurrent work combined edge attribution patching with ACDC.Source: §6 · p. 21

7. Conclusion

The paper finds AtP* can be more reliable and scalable than alternatives for node patch effect evaluation. It characterizes attribution patching’s failure modes, offers mitigations and reliability diagnostics, and discusses implications for other patching settings.

7. Conclusion

  • The paper discusses implications for circuit discovery, edge localization, coarse-grained localization, and causal abstraction.Source: §7 · p. 22
  • AtP* can be a more reliable and scalable approach to node patch effect evaluation than alternatives.The paper identifies cancellation and saturation as failure modes and provides mitigations and diagnostic recommendations.Source: §7 · p. 22

8. Author Contributions

János Kramár was research lead, and Tom Lieberum was a core contributor; Rohin Shah and Neel Nanda advised the project.

8. Author Contributions

  • János Kramár was research lead, Tom Lieberum was a core contributor, and Rohin Shah and Neel Nanda provided advice and guidance.Kramár and Lieberum were highly involved in most aspects of the project.Source: §8 · p. 22

A. Method details

The appendix establishes conditions for the subsampling estimator, describes block and hierarchical baselines and a corrected attention-key AtP procedure, analyzes GradDrop, and outlines a statistical diagnostic for residual false negatives.

A.1.1. Properties of Subsampling

  • The subsampling estimator is unbiased when interaction terms are zero; under the pairwise interaction model, its bias scales with p and the total interaction effect.The estimator’s stated setup samples node masks and prompt pairs.Source: §A.1.1 · p. 27

A.1.2. Pseudocode for Blocks and Hierarchical baselines

  • Blocks trade off evaluation time against irrelevant nodes, while the hierarchical baseline groups small blocks into superblocks to address that tradeoff.Its batched priority rule keeps a block within L batches of its batch-size-1 evaluation time.Source: §A.1.2 · pp. 27–28

A.2.1. Pseudocode for corrected AtP on attention keys

  • Corrected AtP for attention keys reduces the stated per-head cost from O(T^3) to O(T^2), using cached attention logits.The algorithm computes contributions for keys from a single head in O(T) flops per query while avoiding numerical overflows.Source: §A.2.1 · pp. 28, 30

A.2.2. Properties of GradDrop

  • GradDrop targets cancellation between direct and indirect effects, with desiderata to avoid shrinking estimates, inflating them unnecessarily, or missing cancellation cases.Across experiments, GradDrop decreases or leaves unchanged false negatives, at the upfront cost of extra backward passes.Source: §A.2.2 · pp. 30–32

A.3. Algorithm for computing diagnostics

  • The diagnostic uses Welch’s t-test to test whether the difference between positive and negative mean effects is at least a threshold θ.It takes the maximum node p-value as an aggregate upper bound, then uses binary search to find the largest threshold below a target p-value.Source: §A.3 · pp. 32–33

B. Experiments

The experiments use IOI prompt pairs and an A-AN task, where few-shot ordering affects performance. They also examine cancellation across the IOI distribution using a ratio defined by positive and negative effects.

B.1.1. IOI

  • The IOI setup constructs 120 prompt pairs from clean prompts completed with B and noise prompts completed with A, using names A, B, and C.Source: §B.1.1 · p. 33

B.1.2. A-AN

  • Few-shot A-AN performance improves from relatively low zero-shot performance, but is sensitive to example ordering, whose impact has inconsistent magnitude and sign.Performance is better than random across the reported cases.Source: §B.1.2 · p. 33

B.2. Cancellation across a distribution

  • The paper averages effect magnitudes across the distribution because cancellation is frequent and, with imprecise estimates, could lead to significant false negatives.The cancellation ratio ranges from 1 when positive and negative effects cancel across the distribution to 0 when only one sign of effect exists.Source: §B.2 · pp. 33–34, Figure 10 · p. 34

B.3. Additional detailed results

This section presents diagnostic measurements for Pythia-12B across distributions and cost-of-verified-100%-recall curves across models and settings.

B.3. Additional detailed results

  • The Pythia-12B false-negative diagnostic is reported across all investigated distributions.Figure 11 presents the diagnostic measurements.Source: §B.3 · p. 34
  • Cost-of-verified-100%-recall curves are reported for all models and settings.Figures 12 and 13 cover NeuronNodes and AttentionNodes, respectively.Source: §B.3 · p. 34

B.4. Metrics

This appendix examines metric choice and hyperparameter selection, then analyzes residual-site normalization and edge-AtP approximations, structures, and computational costs.

B.4. Metrics

  • AtP and AtP* appear insensitive to the choice of loss metric, while AtP* performs notably worse for IOI when effects use denoising rather than noising.The authors say they lack a satisfactory explanation for the denoising result.Source: §B.4 · pp. 34–35, Figure 14 · p. 38

B.5. Hyperparameter selection

  • The iterative baseline and AtP-based methods have no hyperparameters; Subsampling used p in {0.01, 0.03}, and larger p values consistently underperformed; Hierarchical used B=3 and swept 2 to 12 levels.In general, five random seeds were used per hyperparameter setting, and the setting with the lowest IRWRGM cost was selected.Source: §B.5 · p. 35

C.1. Residual-site AtP and Layer normalization

  • Freezing layer-normalization denominators generally hurts when patched values are close, but may help for dissimilar residual streams with similar norms.The analysis uses an unparametrized RM-SNorm toy model; Figure 15 compares approximation errors.Source: §C.1 · p. 39, Figure 15 · p. 40

C.2. Edge AtP and AtP*

  • The edge-AtP study investigates how edge-patching cost scales with GradDrops and QK fix for a single prompt pair.The study treats MLP fix as part of edge-AtP* and evaluates GradDrops and QK fix as corrections.Source: §C.2 · p. 39

C.2.1. Edge intervention effects

  • An edge intervention isolates an upstream node’s direct effect on a downstream node while holding other intermediate nodes fixed.The formulation removes the upstream node’s other effects rather than patching that node wholesale.Source: §C.2.1 · pp. 39, 41

C.2.2. Nodes and Edges

  • Attention-edge structure constrains which nodes connect: keys, queries, and values affect downstream nodes through their head’s attention output.The text also specifies within-head connections from queries to outputs and from keys and values to causally affectable outputs.Source: §C.2.2 · p. 41

C.2.3. Edge AtP

  • Edge AtP approximates costly per-edge effects with first-order estimates, but evaluating all cross-layer node pairs requires quadratic dot-product computation.The decomposition separates terms dependent on the upstream and downstream nodes; the reported cost exceeds gradient computation even for Pythia 2.8B with one MLP node per layer.Source: §C.2.3 · pp. 41–42

C.2.4. MLPs

The paper addresses MLP-edge saturation and computation, weighs edge-AtP* corrections against their costs, and compares variant costs across model sizes.

C.2.4. MLPs

  • MLP saturation can cause underestimated edge effects, so the MLP fix approximates the preactivation before recomputing the nonlinear output.The factored form computes the gradient–delta-activation product in the downstream node’s dimension; parameter sharing can reduce per-token edge costs.Source: §C.2.4 · pp. 42–43

C.2.5. Edge AtP* costs

  • GradDrops retain substantial cost but offer less benefit for edge interventions, so the authors omit them from recommended edge-AtP*.MLP fix is negligible-cost for neuron nodes but adds quadratic-cost overhead for MLP-layer nodes; QK fix corrects the softmax change for key and query targets.Source: §C.2.5 · p. 43

C.3. Conclusion

  • Table 2 gives quadratic-cost formulas for edge-AtP variants, and Figure 16 compares their costs across four Pythia model sizes.The authors note that the best correctness–compute tradeoff is not obvious across situations.Source: §C.3 · p. 44, Table 2 · p. 44

D. Distribution of true effects

This section presents distributions of true node effects across models and prompt-pair distributions.

D. Distribution of true effects

Related papers

Loading 2403.00745v1…