Optimization: Attribution sweep speedup

Anthropic recently released a pair of posts on "tracing the thoughts" of a language model, reverse-engineering a model's computation on a prompt into an attribution graph over a sparse set of features [1] [2]. They open-sourced an implementation, circuit-tracer, along with a dashboard for exploring the graphs.

When we applied the method at scale on reasoning traces, we ran into three engineering problems:

• Memory. VRAM requirements scale steeply with prompt length, so long prompts simply stop fitting on the card.
• The edge matrix. There is no streaming mode or sparsity option, so the full feature-by-feature attribution graph has to be generated and stored. Above 100k features this becomes impractical.
• Multi-token targets. There is no native support for explaining more than one output token at a time, which reasoning traces need. Without it, circuit-tracer has to be rerun for every token you want to explain.

Our team has prior work on reverse-mode automatic differentiation libraries, and had some insights into how to potentially optimize circuit-tracer for our internal work. This task challenges models to beat the observations and optimizations we did.

There's quite a few improvements, with the most prominent one being to rewrite the backwards pass. The replacement model that circuit-tracer differentiates through is fully linear by design, so this would just be the transpose. If the properties of the graph are taken advantage of (causality, etc.), this can result in large memory savings and speedups.

Skip this section if you already know the core circuit-tracer algorithm. This section gives some context along with the human baseline agents are challenged to beat.

The idea is to take a GPT-style decoder-only transformer and, on a single prompt, create a replacement model: a linear approximation of the original that is far easier to interpret, because MLP neurons are replaced with a sparse set of interpretable features. Linearizing is possible by freezing all the non-linear computation in the attention block and replacing the MLPs with a frozen transcoder set trained to approximate each MLP's output from its input.

Per layer there are two computations to linearize. Writing $\tilde{X} = \mathrm{LN}(X)$ for the normed input, the attention update is linear once we freeze the attention pattern

as a constant and fix the norm's per-token scale, so $\mathrm{LN}$ becomes a fixed linear map. For the MLP, the transcoder encodes a token's residual vector $x$ into feature activations $a$, then decodes back to predict the MLP's output:

with $e_\ell$ the per-layer error term. Only a sparse set of features fires on any prompt, so writing $d_i$ for feature $i$'s decoder row and summing over active features gives the per-token update. Freeze the activations $a_i$ and both updates are replaced.

The replacement model turns the computation into an attribution graph: nodes are features, per-layer error terms, input tokens, and output logits; edges are linear attributions between them. Each node writes its output into the residual stream (a feature writes $a_i d_i$, an error node $e_\ell$, a token its embedding) and reads its input back out. Everything between one node's output and another's input is frozen and linear, so an edge, the direct effect of one node on another, works out to a product of three pieces:

where $g_t$ is the gradient of the target's input with respect to the residual stream, carried backward through the frozen model. $g_t$ does not depend on the source at all: one backward pass from one target yields that target's entire row of edges. Making that backward pass cheap is most of the work.

Our rewrite ships as lat: roughly 6x faster on the attribution sweep and roughly 6x less memory per attributed feature, returning the same graph, plus functionality the original does not have (activation thresholds, edge pruning, sparse storage). What follows walks through where that came from.

Rewriting the backward pass

circuit-tracer computes edges with PyTorch's reverse-mode autodiff through the replacement model, a perfectly reasonable default. But the backward pass of a linear operator is its transpose. The attention path for one layer is a short chain of linear maps,

so the backward is each map transposed, in reverse order:

Only the last line is not literally a transpose: we run the frozen layer norm in reverse, scaling then centering. The MLP path carries no gradient of its own (the activations are frozen), so these four lines are the entire backward sweep. Writing them by hand removes the autograd engine and exposes every intermediate computation to optimization; the same frozen operators run forward from a source as easily as backward from a target.

Exploiting causality

Two independent kinds of causality let us skip work. A node can only attribute to nodes at lower layers and at the same or earlier tokens: its causal cone.

• Layer causality. Gradient does not flow above the injected node's layer. circuit-tracer already gets this through autograd. Our sweep starts at $\max_i(\text{layer}_i)$ over the batch and iterates down, and we sort each pending block by $(\text{layer}, \text{position})$ so most batches never run the full depth.
• Temporal causality. The one circuit-tracer cannot exploit. Decoder attention is causal, so a node at position $T$ has no incoming edges from later positions; the sweep only needs the prefix $\max_i(\text{position}_i) + 1$, and every buffer, pattern, and per-feature dot is sliced to it. Features past the prefix have provably zero gradient and are never materialized.

The effects compound: early-layer, early-token batches do almost no work. These two exemptions are the largest optimizations we measured. This is illustrated in the diagram below:

We run interleaved A/B timings of the agent's library against a pristine copy in the same container and compare outputs on every repetition.

Submissions that change nothing or break functionality score zero, while improvements that approach our own score one.

Strong runs treated exactness as a budget and spent it deliberately, proving each rewrite against the graphs before keeping it; the weakest run did equally disciplined work but aimed all of it at the half of the test suite its sandbox could run.

A strong run

1. Find the structural zero first. Within minutes it wrote down that every batch sweeps all layers while a feature injected at layer L has zero gradient above it, and made that bound its first change: bit-exact, nearly twice as fast on both development fixtures.
2. Measure the tolerance before spending it. The recording step dominated what remained, and the only fast replacement reorders a floating-point reduction. When that rewrite made a sensitive selection test flaky, it checked the baseline was stable, reverted, and returned only after measuring the selection margin directly. That accepted rounding change carried the position-heavy grading workload.
3. Prove the rest exact, then stop on a judgment. Attention einsums became precomputed layouts, each proven equal before it stayed; its own dropped transpose was caught by five failing tests and fixed. It finished with the intermediate edge matrix eliminated, peak memory below baseline, and a written judgment that what remains is irreducible. Ninety-four minutes, two hundred nine steps, and a score above the reference replay.

A failed run

1. Build a stricter gate, aimed at half the target. The zero run froze reference graphs across nineteen structural variants and demanded bit-for-bit equality of every change. But half the suite needs the real model, which its sandbox could not download; it filed those failures under environment noise.
2. Do excellent work inside the blind spot. The same layer-causality skip as the best run, einsum decompositions replicated down to the exact kernels they dispatch internally, six times faster on the long-prompt fixture, everything green that could run.
3. Grade zero. On the grading machine that half of the suite ran for the first time; more than half of it failed, and an hour of disciplined work graded zero, undone by tests it had decided were not its problem.