Attention reveals key model circuits in one pass
Attention as a Routing Graph: Live Circuit Extraction from a Single Forward Pass
Artificial IntelligenceMachine Learning
Summary
Language models solve problems using internal circuits, which are usually found by testing many changes step-by-step. The authors tried a simpler way by looking at attention in just one run to find important routes that lead to the answers. They showed that removing these routes harms the model more than removing random ones, meaning these routes really matter. Their method is much faster than previous approaches but doesn’t find all circuits perfectly.
What this means in practice
- •For ai model engineers: Quickly identify critical attention routes to debug or improve language models on tasks where circuits are known.
- •For machine learning operations teams: Reduce computational cost in causal analysis of language models by extracting important edges from a single forward pass instead of expensive patch tests.
Authors
Ash Manvi, Samreena Tajreen
Abstract
Finding circuits in language models usually means running many careful interventions. We try something simpler: treat attention as a routing map from one forward pass, keep a small set of routes that point toward the answer, and ask whether those routes actually matter. They often do. On induction and IOI (tasks where the "right" circuit is already known), ablating our extracted edges hurts the model much more than ablating a random set of the same size. We evaluate n=100 prompts per cell on GPT-2 Small, GPT-2 Medium, and Pythia-410M, with paired gap tests and bootstrap confidence intervals. The extract step costs one forward; a head-by-head patch sweep costs about two orders of magnitude more. We are not claiming a complete circuit atlas. We are claiming a cheap sketch that carries real causal signal on known tasks, with clear failure modes when it does not. Code and evaluation artifacts are at https://github.com/Aquinf03/live-circuit-routing.