Circuit discovery: composing steps into one experiment

Note

Generated from notebooks/applications/circuit_discovery.ipynb, which you can run yourself. The outputs below are the ones it produced.

The other notebooks each ran one analysis. This one chains several into a single experiment that recovers a real circuit.

The logic: direct logit attribution finds the heads that write the answer into the logits (the name movers); activation patching finds heads that carry the signal but never touch the logits directly; path patching tests the connection between them.

Key questions

• Are the two groups of heads connected, or independently important?
• Does the upstream group write into the downstream heads’ query, key, or value?
• How do we know the effect is real and not an artifact of the method?

Structure

1. Set up
2. Find who writes the answer
3. Find who carries the signal
4. Test the connection with path patching
5. Package the finding as a step
6. Save

Model and data. GPT-2 small in float32. It is small enough that bfloat16 rounding visibly degrades its predictions, so the notebooks load it explicitly in float32. Eight indirect-object prompt pairs from murano.tasks.ioi.

Requirements. pip install murano-interp[plot]

1. Set up

import torch

from murano import MuranoModel, NodeSet, Pipeline, keys
from murano.nodes import Node, Side
from murano.plotting import plot_sweep, save_figure
from murano.results import Results
from murano.steps import (
    LoadPaired,
    LogitAttribution,
    LogitDiffStep,
    Logits,
    Patch,
    PathPatch,
    RecoveredMetricStep,
    Save,
    SelectComponents,
    Sweep,
)
from murano.steps.base import Step

OUTPUT_DIR = "murano_outputs/circuit_discovery"

# GPT-2 is small enough that bfloat16 rounding degrades its predictions.
model = MuranoModel("gpt2", dtype=torch.float32)
print(f"{model.model_id}: {model.n_layers} layers, {model.n_heads} heads")

gpt2: 12 layers, 12 heads

This notebook uses the indirect-object-identification task from murano.tasks. Each clean prompt names two people and then has the subject give a drink, so the next token should be the other name, the indirect object. The corrupt prompt swaps who gives, which flips the answer.

clean  : When Mary and John went to the store, John gave a drink to ___  ->  " Mary"
corrupt: When Mary and John went to the store, Mary gave a drink to ___  ->  " John"

The metric throughout is the logit difference: the correct name’s logit minus the incorrect name’s. Positive means the model prefers the right answer.

from murano.tasks import ioi

task = ioi(n=8)
print("clean  :", task.clean[0], "->", repr(task.correct[0]))
print("corrupt:", task.corrupt[0], "->", repr(task.incorrect[0]))
print(f"{len(task.clean)} prompt pairs")

clean  : When Mary and John went to the store, John gave a drink to -> ' Mary'
corrupt: When Mary and John went to the store, Mary gave a drink to -> ' John'
8 prompt pairs

baseline = Pipeline(
    [
        LoadPaired(task),
        Logits(model),
        Logits(
            model,
            prompts_key=keys.CORRUPT_PROMPTS,
            logits_key=keys.CORRUPT_LOGITS,
            mask_key=keys.CORRUPT_MASK,
            targets=None,
        ),
        LogitDiffStep(task.correct, task.incorrect, output_key="ld_clean", model=model),
        LogitDiffStep(
            task.correct,
            task.incorrect,
            logits_key=keys.CORRUPT_LOGITS,
            mask_key=keys.CORRUPT_MASK,
            output_key="ld_corrupt",
            model=model,
        ),
    ]
).run()

print(
    f"clean {baseline['ld_clean'].value:+.4f} / corrupt {baseline['ld_corrupt'].value:+.4f}"
)

clean +2.3971 / corrupt -3.1112

2. Find who writes the answer

Attribution over the clean prompts, then take the strongest heads. These write directly into the logits.

SelectComponents writes a ComponentSelection under a key of our choosing, and every causal step can read its targets from such a key rather than a hand-copied node list. We will use that twice.

attributed = Pipeline(
    [
        LoadPaired(task),
        LogitAttribution(model, correct=task.correct, incorrect=task.incorrect),
        SelectComponents(
            top_k=4, by="abs", modules="self_attn", output_key="name_movers"
        ),
    ]
).run()

name_movers = attributed["name_movers"].nodes
contributions = attributed[keys.LOGIT_ATTRIBUTION].contributions
for node in name_movers:
    print(f"  {str(node):<24} {contributions[node]:+.3f}")

  L10.self_attn.h0         +1.546
  L10.self_attn.h7         -1.491
  L9.self_attn.h9          +1.316
  L9.self_attn.h6          +1.066

3. Find who carries the signal

Now patch every head from the clean run into the corrupt run, and keep the ones that restore behavior. Only layers strictly before the earliest name mover can feed into it, so that is where we look.

Sweep runs the patch-and-score chain once per head, forking baseline each time. SelectComponents then ranks the scores it collected. A component sweep publishes the same {Node: float} map an attribution does, so the two steps compose in one pipeline with no node list carried by hand between them.

first_mover_layer = min(node.layer for node in name_movers)
print(f"searching layers 0..{first_mover_layer - 1} for upstream carriers\n")

found = Pipeline(
    [
        Sweep(
            over=NodeSet.product(
                range(first_mover_layer), ["self_attn"], heads=range(model.n_heads)
            ),
            steps=lambda head: [
                Patch(model, targets=[head]),
                LogitDiffStep(
                    task.correct,
                    task.incorrect,
                    logits_key=keys.PATCHED_LOGITS,
                    mask_key=keys.PATCHED_MASK,
                    output_key="ld_patched",
                    model=model,
                ),
                RecoveredMetricStep("ld_clean", "ld_corrupt", "ld_patched"),
            ],
            read=keys.RECOVERED,
        ),
        SelectComponents(
            source_key=keys.SWEEP, top_k=4, by="abs", output_key="upstream"
        ),
    ]
).run(baseline.copy())

carried = found[keys.SWEEP]
upstream = found["upstream"].nodes

ranked = sorted(carried.scores.items(), key=lambda item: -abs(item[1]))
print(f"{'head':<24} {'recovered':>10} {'direct':>10}")
print("-" * 46)
for node, score in ranked[:6]:
    print(f"{str(node):<24} {score:>+10.3f} {contributions.get(node, 0.0):>+10.3f}")

searching layers 0..8 for upstream carriers

head                      recovered     direct
----------------------------------------------
L5.self_attn.h5              +0.268     -0.007
L8.self_attn.h10             +0.263     +0.155
L8.self_attn.h6              +0.240     -0.075
L7.self_attn.h9              +0.224     +0.221
L7.self_attn.h3              +0.155     +0.090
L3.self_attn.h0              +0.129     +0.005

# The recovered fractions straddle zero, so plot_sweep centers a diverging scale
# on it and a head with no effect lands in the neutral color.
fig = plot_sweep(carried, title="Recovered logit difference, upstream heads only")
save_figure(fig, f"{OUTPUT_DIR}/plots/upstream_sweep.png")
fig

Look at the two columns together. Some of these heads recover a substantial fraction of the behavior while contributing essentially nothing directly to the logits: their direct score is near zero, or even negative. Whatever they do, they do it through something else.

Others, in layers 7 and 8, score on both. A head is allowed to write to the logits and feed a downstream head; the two columns measure different things and neither subsumes the other. The heads worth chasing are the ones where the gap is large.

4. Test the connection with path patching

PathPatch isolates one edge of the computation graph. It patches a sender’s output, but only along the direct path into a named receiver, freezing everything else so no other route can carry the change.

A receiver is a specific slot: Node(layer, "self_attn", head=h, side=Side.Q) is that head’s query input. Query, key, and value are separate slots, and which one an upstream head writes into tells you what it is doing.

One direction matters: senders must be upstream of the receiver. A layer-10 head cannot feed a layer-9 head. We check that below.

Note also that Patch and PathPatch take their base from opposite sides. Patch bases on the corrupt run and splices clean activations in. PathPatch bases on the clean run and splices corrupt activations in, so a drop in the metric means the path was carrying the signal.

receiver_head = max(name_movers, key=lambda node: contributions[node])
print("receiver (strongest name mover):", receiver_head)
print("senders  (upstream carriers)  :", [str(n) for n in upstream])

receiver (strongest name mover): L10.self_attn.h0
senders  (upstream carriers)  : ['L5.self_attn.h5', 'L8.self_attn.h10', 'L8.self_attn.h6', 'L7.self_attn.h9']

Which slot do the upstream heads write into?

One more sweep, this time over the three slots rather than over components. PathPatch reads its senders from the upstream selection the last sweep wrote, so nothing is passed by hand. ld_clean already sits in the forked results, which is what each slot’s patched score is compared against.

slots = Pipeline(
    [
        Sweep(
            over=[Side.Q, Side.K, Side.V],
            steps=lambda side: [
                PathPatch(
                    model,
                    senders_key="upstream",
                    receiver=Node(
                        receiver_head.layer,
                        "self_attn",
                        head=receiver_head.head,
                        side=side,
                    ),
                ),
                LogitDiffStep(
                    task.correct,
                    task.incorrect,
                    logits_key=keys.PATH_PATCHED_LOGITS,
                    mask_key=keys.PATH_PATCHED_MASK,
                    output_key="ld_path",
                    model=model,
                ),
            ],
            read="ld_path",
        )
    ]
).run(found.copy())

clean = found["ld_clean"].value
print(f"{'slot':<8} {'change in logit difference':>28}")
print("-" * 38)
for side, patched in slots[keys.SWEEP].scores.items():
    print(f"{side.name:<8} {patched - clean:>+28.4f}")

slot       change in logit difference
--------------------------------------
Q                             -0.5784
K                             -0.0066
V                             -0.0007

The upstream heads write into the query, not the key or the value. The key moves the metric by about one percent as much, and the value by a tenth of that again.

That is a mechanistic claim about what they compute. The name mover’s value carries what to copy; its query decides where to look. So these upstream heads are not supplying the name. They are telling the name mover which name to skip, by marking the subject as already-mentioned.

The layer 7 and 8 members of this set are the S-inhibition heads. The layer 5 head is a duplicate-token head: it spots that one name occurs twice, which is the signal the S-inhibition heads need in the first place.

A control: the effect must vanish for downstream senders

Path patching from heads that sit after the receiver cannot possibly affect it. If the method is sound, the measured change is exactly zero. Not small, zero.

# Any head strictly after the receiver will do. If the receiver sat in the last
# layer there would be no downstream sender at all, and a "control" that quietly
# tested a same-layer head would still pass while proving nothing.
next_layer = receiver_head.layer + 1
assert next_layer < model.n_layers, (
    "receiver is in the last layer: no sender is downstream"
)
downstream = [Node(next_layer, "self_attn", head=0)]

control = Pipeline(
    [
        PathPatch(
            model,
            senders=downstream,
            receiver=Node(
                receiver_head.layer, "self_attn", head=receiver_head.head, side=Side.Q
            ),
        ),
        LogitDiffStep(
            task.correct,
            task.incorrect,
            logits_key=keys.PATH_PATCHED_LOGITS,
            mask_key=keys.PATH_PATCHED_MASK,
            output_key="ld_path",
            model=model,
        ),
    ]
).run(found.copy())

effect = control["ld_path"].value - clean
print(f"downstream senders {[str(n) for n in downstream]}")
print(f"change in logit difference: {effect:+.6f}")

assert abs(effect) < 1e-9, "a downstream sender must not affect an upstream receiver"
print("\nzero to floating-point precision, as it must be")

downstream senders ['L11.self_attn.h0']
change in logit difference: +0.000000

zero to floating-point precision, as it must be

5. Package the finding as a step

The pipeline abstraction pays off when a bespoke analysis becomes a reusable component. A step needs only reads, writes, and __call__.

class CircuitSummary(Step):
    """Collect the discovered circuit into one artifact."""

    reads = ["name_movers"]
    writes = ["circuit"]

    def __init__(self, upstream: list[Node], receiver: Node, slot: Side):
        self.upstream = upstream
        self.receiver = receiver
        self.slot = slot

    def __call__(self, results: Results) -> Results:
        results["circuit"] = {
            "name_movers": [str(node) for node in results["name_movers"].nodes],
            "upstream": [str(node) for node in self.upstream],
            "writes_into": self.slot.name,
            "receiver": str(self.receiver),
        }
        return results


attributed = CircuitSummary(upstream, receiver_head, Side.Q)(attributed)
for key, value in attributed["circuit"].items():
    print(f"{key:>12}: {value}")

 name_movers: ['L10.self_attn.h0', 'L10.self_attn.h7', 'L9.self_attn.h9', 'L9.self_attn.h6']
    upstream: ['L5.self_attn.h5', 'L8.self_attn.h10', 'L8.self_attn.h6', 'L7.self_attn.h9']
 writes_into: Q
    receiver: L10.self_attn.h0

What we found

Reading the results back as one mechanism:

• A group of heads in layers 7 and 8 identify the repeated subject name.
• They write that information into the query of the name movers in layers 9 and 10.
• The name movers then attend to the other name, the indirect object, and copy it into the logits, exactly the copying behavior attention.ipynb read off the OV circuit.
• Some name movers are negative and suppress the answer instead.

This is the indirect-object-identification circuit.

6. Save

run_dir = Save(
    output_dir="murano_outputs", run_name="circuit_discovery", model_id=model.model_id
)(attributed)[keys.OUTPUT_DIR]

print("saved to:", run_dir)

saved to: murano_outputs/circuit_discovery

What next

• ../reproductions/wang2023_ioi.ipynb reproduces the full circuit, with more prompt templates and every head class.
• custom_pipeline.ipynb builds a different custom experiment, around probing and steering.