Elamur Research⌕
☰
home / papers / 2024 / AtP*: An efficient and scalable method for localizing LLM behaviour to components

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

Paper · 2024 · Janos Kramar, Tom Lieberum, Rohin Shah et al.

Summary

Mechanistic interpretability often relies on Activation Patching. This method replaces one component's activation with its value from a corrupted prompt and measures how a behaviour metric changes. Run exhaustively, it needs a separate forward pass for every component, which is impractical for large LLMs with billions of neurons. The Google DeepMind authors study Attribution Patching (AtP). AtP is a first-order gradient approximation that scores all nodes at once. The highest-scoring nodes are then checked with real patching.

They identify two sources of false negatives in AtP:

  • Saturated attention softmax, where the local gradient misrepresents the true effect.
  • Cancellation between direct and indirect effects.

Their variant, AtP*, addresses these with two changes. It recomputes the softmax exactly when patching queries and keys, and it applies dropout in the backward pass. In the first systematic comparison against several faster patching baselines on Pythia models, AtP outperforms all of them, and AtP* improves further. A subsampling-based diagnostic statistically bounds any remaining missed effects. Together, these make fine-grained causal localization feasible at scale.

Key points

  • Exhaustive patching is costly: a 1024-token prompt in Chinchilla 70B has about 2.7·10⁹ neuron nodes. AtP scores every node with only two forward passes and one backward pass per prompt pair.
  • AtP* combines two fixes, QK-fix (exact softmax recomputation) and GradDrop (backward-pass dropout). GradDrop has a noticeable upfront cost but reduces false negatives.
  • The evaluation covers MLP neurons and attention nodes on prompt-pair tasks such as CITY-PP and IOI-PP, using Pythia models including Pythia-12B. Cost is measured in forward passes relative to an oracle verification order.
  • The authors give practical recommendations depending on the setting:
    • AtP* for single prompt pairs.
    • AtP+QKFix for attention nodes over distributions.
    • Plain AtP for neuron nodes over distributions.
  • The diagnostic first filters out high-impact nodes with AtP, then patches random subsets of the remaining nodes to bound undetected effects.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🏢Google DeepMindLondon · United Kingdom82 works · 84,829 citations

 Fields, topics and concepts

📚 Shelves 3The field of science this work belongs to (OpenAlex fields).Computer Science · 19,503 worksArtificial Intelligence · 18,805 worksSoftware · 150 works

🧭 Topics 1Research subjects and languages the work is about.Language models & pretraining · 4,145 works

Cite

Cite as Janos Kramar, Tom Lieberum, Rohin Shah and Neel Nanda. 2024. AtP*: An efficient and scalable method for localizing LLM behaviour to components. arXiv:2403.00745.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2403.00745 · arXiv 2403.00745

@misc{kramar2024atp,
  title = {{AtP*: An efficient and scalable method for localizing LLM behaviour to components}},
  author = {Janos Kramar, Tom Lieberum, Rohin Shah and Neel Nanda},
  year = {2024},
  doi = {10.48550/arxiv.2403.00745},
  eprint = {2403.00745},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2403.00745}
}

References

37 references, 11 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. N. Belrose, D. Schneider-Joseph, S. Ravfogel, R. Cotterell, E. Raff, and S. Biderman. Leace: Perfect linear concept erasure in closed form. arXiv preprint arXiv:2306.03819, 2023. — [arXiv:2306.03819](https://arxiv.org/abs/2306.03819) find
  2. S. Biderman, H. Schoelkopf, Q. G. Anthony, H. Bradley, K. O’Brien, E. Hallahan, M. A. Khan, S. Purohit, U. S. Prashanth, E. Raff, A. Skowron, L. Sutawika, and O. van der Wal. Pythia: A suite for analyzing large language models across training and scaling. In A. Krause, E. Brunskill, K. Cho, B. Engel… find
  3. PMLR, 2023. URL https://proceedings.mlr.press/v202/ biderman23a.html . ↗
  4. T. Bricken, A. Templeton, J. Batson, B. Chen, A. Jermyn, T. Conerly, N. Turner, C. Anil, C. Denison, A. Askell, R. Lasenby, Y. Wu, S. Kravec, N. Schiefer, T. Maxwell, N. Joseph, Z. Hatfield-Dodds, A. Tamkin, K. Nguyen, B. McLean, J. E. Burke, T. Hume, S. Carter, T. Henighan, and C. Olah. Towards mon… find
  5. N. Cammarata, S. Carter, G. Goh, C. Olah, M. Petrov, L. Schubert, C. Voss, B. Egan, and S. K. Lim. Thread: Circuits. Distill, 2020. doi: 10.23915/distill.00024. https://distill.pub/2020/circuits. — [doi:10.23915/distill.00024](https://doi.org/10.23915/distill.00024) ↗
  6. N. D. Cao, L. Schmid, D. Hupkes, and I. Titov. Sparse interventions in language models with differentiable masking, 2021. find
  7. L. Chan, A. Garriga-Alonso, N. Goldwosky-Dill, R. Greenblatt, J. Nitishinskaya, A. Radhakrishnan, B. Shlegeris, and N. Thomas. Causal scrubbing, a method for rigorously testing interpretability hypotheses. AI Alignment Forum, 2022. https://www.alignmentforum.org/posts/ JvZhhzycHu2Yd57RN/causal-scrub… ↗
  8. A. Conmy, A. N. Mavor-Parker, A. Lynch, S. Heimersheim, and A. Garriga-Alonso. Towards automated circuit discovery for mechanistic interpretability, 2023. find
  9. H. Cunningham, A. Ewart, L. Riggs, R. Huben, and L. Sharkey. Sparse autoencoders find highly interpretable features in language models, 2023. find
  10. J. Feng and J. Steinhardt. How do language models bind entities in context?, 2023. find
  11. M. Figurnov, A. Ibraimova, D. P. Vetrov, and P. Kohli. Perforatedcnns: Acceleration through elimination of redundant convolutions. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016. URL h… find
  12. M. Finlayson, A. Mueller, S. Gehrmann, S. Shieber, T. Linzen, and Y. Belinkov. Causal analysis of syntactic agreement mechanisms in neural language models. In C. Zong, F. Xia, W. Li, and R. Navigli, editors, Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and… find
  13. L. Gao, S. Biderman, S. Black, L. Golding, T. Hoppe, C. Foster, J. Phang, H. He, A. Thite, N. Nabeshima, S. Presser, and C. Leahy. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027, 2020. find
  14. A. Geiger, K. Richardson, and C. Potts. Neural natural language inference models partially embed theories of lexical entailment and negation, 2020.
  15. A. Geiger, H. Lu, T. Icard, and C. Potts. Causal abstractions of neural networks, 2021. find
  16. A. Geiger, Z. Wu, H. Lu, J. Rozner, E. Kreiss, T. Icard, N. D. Goodman, and C. Potts. Inducing causal structure for interpretable neural networks, 2022. find
  17. A. Geiger, C. Potts, and T. Icard. Causal abstraction for faithful model interpretation, 2023. find
  18. M. Geva, J. Bastings, K. Filippova, and A. Globerson. Dissecting recall of factual associations in auto-regressive language models, 2023. find
  19. N. Goldowsky-Dill, C. MacLeod, L. Sato, and A. Arora. Localizing model behavior with path patching, 2023.
  20. W. Gurnee, N. Nanda, M. Pauly, K. Harvey, D. Troitskii, and D. Bertsimas. Finding neurons in a haystack: Case studies with sparse probing, 2023.
  21. M. Hanna, O. Liu, and A. Variengien. How does gpt-2 compute greater-than?: Interpreting mathematical abilities in a pre-trained language model, 2023.
  22. P. Hase, M. Bansal, B. Kim, and A. Ghandeharioun. Does localization inform editing? surprising differences in causality-based localization vs. knowledge editing in language models, 2023. find
  23. R. Hendel, M. Geva, and A. Globerson. In-context learning creates task vectors, 2023.
  24. J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. de Las Casas, L. A. Hendricks, J. Welbl, A. Clark, T. Hennigan, E. Noland, K. Millican, G. van den Driessche, B. Damoc, A. Guy, S. Osindero, K. Simonyan, E. Elsen, O. Vinyals, J. Rae, and L. Sifre. An empirical analysis o… find
  25. J. Huang, A. Geiger, K. D’Oosterlinck, Z. Wu, and C. Potts. Rigorously assessing natural language explanations of neurons, 2023.
  26. E. Jenner, A. Garriga-Alonso, and E. Zverev. A comparison of causal scrubbing, causal abstractions, and related methods. AI Alignment Forum, 2022. https://www.alignmentforum.org/posts/uLMWMeBG3ruoBRhMW/ a-comparison-of-causal-scrubbing-causal-abstractions-and . ↗
  27. O. Jorgensen, D. Cope, N. Schoots, and M. Shanahan. Improving activation steering in language models with mean-centring, 2023. find
  28. K. Li, O. Patel, F. Viégas, H. Pfister, and M. Wattenberg. Inference-time intervention: Eliciting truthful answers from a language model, 2023.
  29. T. Lieberum, M. Rahtz, J. Kramár, N. Nanda, G. Irving, R. Shah, and V. Mikulik. Does circuit analysis interpretability scale? evidence from multiple choice capabilities in chinchilla, 2023.
  30. C. Louizos, M. Welling, and D. P. Kingma. Learning sparse neural networks through 𝑙 0 regularization, 2018. find
  31. C. McDougall, A. Conmy, C. Rushing, T. McGrath, and N. Nanda. Copy suppression: Comprehensively understanding an attention head, 2023.
  32. T. McGrath, M. Rahtz, J. Kramár, V. Mikulik, and S. Legg. The hydra effect: Emergent self-repair in language model computations, 2023.
  33. K. Meng, D. Bau, A. Andonian, and Y. Belinkov. Locating and editing factual associations in gpt, 2023. find
  34. J. Merullo, C. Eickhoff, and E. Pavlick. Circuit component reuse across tasks in transformer language models, 2023. find
  35. P. Michel, O. Levy, and G. Neubig. Are sixteen heads really better than one? In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https://proceedings.neurips.cc/… ↗
  36. P. Molchanov, S. Tyree, T. Karras, T. Aila, and J. Kautz. Pruning convolutional neural networks for resource efficient inference. In International Conference on Learning Representations, 2017. URL https://openreview.net/forum?id=SJGCiw5gl . ↗
  37. N. Nanda. Attribution patching: Activation patching at industrial scale. 2022. URL https://www.neelnanda.io/mechanistic-interpretability/attribution-patching . ↗

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/kramar-2024-atp-efficient-scalable-method-localizing/