Automated Circuit Discovery: Patching, Attribution & Decomposition at Scale

Symbols. a, b, c ∈ [0, 1] are source activations of the clean record and a′, b′, c′ those of the donor cache. e_A, e_B ∈ {0, 1} mark whether edges A and B are retained in a graph H. ã = e_Aa + (1 − e_A)a′ and b̃ = e_Bb + (1 − e_B)b′ are the slots O reads. z = ã + b̃ − ãb̃ = 1 − (1 − ã)(1 − b̃), m = 2z − 1 and p = σ(m). y is a recorded label and never enters O. m₀ and p₀ are the full graph’s margin and probability on the clean record. τ > 0 is in nats.

The exact expansion. With increments δã, δb̃ at a base (ã, b̃), z changes by δã + δb̃ − (ãδb̃ + b̃δã + δãδb̃), so Δm = 2(1 − b̃)δã + 2(1 − ã)δb̃ − 2δãδb̃ with no remainder. The first two terms are the gradient at the base times the increments.

The divergence. With softplus(t) = log(1 + e^t), log p = −softplus(−m), log(1 − p) = −softplus(m) and softplus(−m) = softplus(m) − m. Substituting into the Bernoulli KL gives F(H) = softplus(m_H) − softplus(m₀) − p₀(m_H − m₀), whose derivative in m_H is σ(m_H) − p₀. Swapping the arguments gives another function; the half donor shows the difference.

The pass, beyond this record. Two other clean records, both with order A then B, checked by the helper and asserted by the code: with clean (0.5, 0.5, ·) and donor (0, 0.75, ·), Δ₁ = 0.0303 and Δ₂ = −0.0303, so a signed decrease passes and both edges go. With clean (1, 0.5, ·) and donor (0.5, 0, ·), Δ₁ = 0.0263 and Δ₂ = 0.0846 are both below τ, and the final F = 0.1109 exceeds τ: τ limits each step, not the result. (In the order B then A the second record keeps A, because its second increment is 0.1109.) In a two-edge pass that starts from the full graph, a negative second increment is always greater than −τ, so comparing |Δ| instead of Δ would make the same decisions here; in longer passes it need not.

Arithmetic and domain. Records must give exactly a, b and c as finite numbers in [0, 1]; edges are A and B only, without repeats; an order names both; τ is finite and positive. Anything else is refused, not repaired. On this domain z ∈ [0, 1] and m ∈ [−1, 1], also after rounding, so no probability needs clamping. The divergence functions in the code accept only margins in [−1, 1]: this is the divergence of this graph, not a general-purpose one, and far outside that range its softplus form would subtract huge, nearly equal terms. F is exactly 0 when m_H = m₀ exactly; within a few rounding units of m₀ it is rounding-dominated (about 1e−15) and is used as computed, with no tolerance in the comparison.

It runs as written with Python 3 and NumPy, prints the six passes and the estimates above, and asserts the strict and signed comparison, the reversed-inequality step, a signed decrease, the unused c and the direction of the divergence. It refuses the inputs the page’s TypeScript helper refuses, except JavaScript-only shapes (inherited or accessor properties, sparse arrays), which Python does not have. The pass under the printed Appendix B.1 Eq. 2 is computed only by the TypeScript helper.

import math
import numpy as np

# The graph, given here. Sources a, b, c. One receiver O reads two input slots, a~ and b~, through
# edge A (a -> O) and edge B (b -> O). There is no edge from c to O.
EDGES = ("A", "B")
CLEAN = {"a": 1.0, "b": 1.0, "c": 1.0}
DONORS = {"zero": {"a": 0.0, "b": 0.0, "c": 0.0},   # the recorded y = 0 row
          "half": {"a": 0.5, "b": 0.5, "c": 0.5},   # chosen, not an observed row
          "same": {"a": 1.0, "b": 1.0, "c": 1.0}}   # equal to the clean record: a no-op control
TAU = 0.1                                           # nats, allowed per trial deletion


def number(v):
    return not isinstance(v, bool) and isinstance(v, (int, float)) and math.isfinite(v)


def check_record(x, name):
    """Exactly a, b and c, each a finite number in [0, 1]. Nothing is coerced."""
    if type(x) is not dict or set(x) != {"a", "b", "c"}:
        raise ValueError(f"{name} must give exactly a, b and c")
    if not all(number(x[k]) and 0 <= x[k] <= 1 for k in "abc"):
        raise ValueError(f"{name}: a, b and c must be finite numbers in [0, 1]")


def check_slots(s, name):
    """A slot pair (a~, b~) of finite numbers in [0, 1]."""
    if not isinstance(s, (list, tuple)) or len(s) != 2 or not all(number(v) and 0 <= v <= 1 for v in s):
        raise ValueError(f"{name} must be two finite numbers in [0, 1]")


def check_edges(edges, name):
    """Distinct edges of O. c has no edge, so "C" is refused, not ignored."""
    if not isinstance(edges, (list, tuple)) or len(edges) > 2 or any(e not in EDGES for e in edges) \
            or len(set(edges)) != len(edges):
        raise ValueError(f"{name} must list distinct edges from A and B")


def check_pair(g, name):
    """Two finite numbers of any sign: a gradient or an increment."""
    if not isinstance(g, (list, tuple)) or len(g) != 2 or not all(number(v) for v in g):
        raise ValueError(f"{name} must be two finite numbers")


def slots(x, donor, kept):
    """(a~, b~) = e * clean + (1 - e) * donor, with retained-edge bits e. A removed edge's slot
    receives the cached donor value; the sources themselves keep their values."""
    check_record(x, "clean")
    check_record(donor, "donor")
    check_edges(kept, "retained edges")
    e = np.array([1.0 if "A" in kept else 0.0, 1.0 if "B" in kept else 0.0])
    s = e * np.array([float(x["a"]), float(x["b"])]) + (1.0 - e) * np.array([float(donor["a"]), float(donor["b"])])
    return float(s[0]), float(s[1])


def sigmoid(m):
    if m >= 0:
        return 1.0 / (1.0 + math.exp(-m))
    e = math.exp(m)
    return e / (1.0 + e)


def output(at, bt):
    """What O computes from its slots: z, the margin m and the probability p. c is not an input."""
    z = at + bt - at * bt
    m = 2.0 * z - 1.0
    return z, m, sigmoid(m)


def softplus(m):
    """log(1 + e^m) without overflow."""
    return m + math.log1p(math.exp(-m)) if m > 0 else math.log1p(math.exp(m))


def forward_kl(m0, m):
    """KL(Bern(sigmoid(m0)) || Bern(sigmoid(m))) in nats. m0 is the full graph's margin on the clean
    record: a fixed reference, never the current graph's output. Margins lie in [-1, 1], the range
    this graph produces; far outside it this form is not used and not trusted."""
    if not (number(m0) and number(m) and -1 <= m0 <= 1 and -1 <= m <= 1):
        raise ValueError("margins must be finite numbers in [-1, 1], the margins this graph produces")
    return softplus(m) - softplus(m0) - sigmoid(m0) * (m - m0)


def removes(delta, tau):
    """Strict and signed: a negative increment passes whatever its size; equality does not."""
    if not number(delta) or not number(tau) or not tau > 0:
        raise ValueError("delta must be finite and tau a finite positive number of nats")
    return delta < tau


def prune(clean, donor, order, tau):
    """One pass: test each edge against the graph currently kept, and remove it iff
    F(trial) - F(current) < tau, strictly and with sign."""
    check_record(clean, "clean")
    check_record(donor, "donor")
    if not isinstance(order, (list, tuple)) or len(order) != 2:
        raise ValueError("a pass order names A and B once each")
    check_edges(order, "pass order")
    if not number(tau) or not tau > 0:
        raise ValueError("tau must be a finite positive number of nats")
    m0 = output(*slots(clean, donor, EDGES))[1]          # full graph: the reference for every step
    kept, steps = list(EDGES), []
    for edge in order:
        trial = [e for e in kept if e != edge]
        f_now = forward_kl(m0, output(*slots(clean, donor, kept))[1])   # from H itself, not carried over
        f_try = forward_kl(m0, output(*slots(clean, donor, trial))[1])
        delta = f_try - f_now                             # signed
        removed = removes(delta, tau)
        steps.append((edge, list(kept), trial, f_now, f_try, delta, removed))
        if removed:
            kept = trial                                  # an accepted removal persists
        # otherwise nothing is assigned: H stays the graph it was before the trial
    return m0, kept, steps


def margin_gradient(base):
    """dm/da~ = 2(1 - b~) and dm/db~ = 2(1 - a~) at a base (a~, b~)."""
    check_slots(base, "base")
    return 2.0 * (1.0 - base[1]), 2.0 * (1.0 - base[0])


def margin_change(base, new, grad=None):
    """Exact m(new) - m(base) against the first-order estimate. Increments are replacement minus base;
    the gradient is the base's own unless another is supplied (that is how a stale one is reused)."""
    check_slots(base, "base")
    check_slots(new, "replacement")
    da, db = new[0] - base[0], new[1] - base[1]
    g = margin_gradient(base) if grad is None else grad
    check_pair(g, "gradient")
    return output(*new)[1] - output(*base)[1], g[0] * da + g[1] * db, -2.0 * da * db


def probability_change(base, new):
    """The same replacement read through p = sigmoid(m): exact change and first-order estimate."""
    check_slots(base, "base")
    check_slots(new, "replacement")
    p = output(*base)[2]
    g = margin_gradient(base)                             # dp/dslot = p(1 - p) * dm/dslot
    da, db = new[0] - base[0], new[1] - base[1]
    return output(*new)[2] - p, p * (1.0 - p) * g[0] * da + p * (1.0 - p) * g[1] * db


def show(v):
    return f"{v + 0.0:.6f}"                               # + 0.0 prints -0.0 as 0.000000


for name, donor in DONORS.items():
    for order in (("A", "B"), ("B", "A")):
        m0, kept, steps = prune(CLEAN, donor, order, TAU)
        for edge, now, trial, f_now, f_try, delta, removed in steps:
            print(f"{name} {''.join(order)} | test {edge} | H {now} -> {trial} | F {show(f_now)} -> "
                  f"{show(f_try)} | delta {show(delta)} | {'remove' if removed else 'keep'}")
        print(f"{name} {''.join(order)} | kept {kept}")

intact, after_a, empty = (1.0, 1.0), (0.0, 1.0), (0.0, 0.0)   # slot pairs under the zero donor
exact, linear, interaction = margin_change(intact, empty)
print("m, both slots at (1, 1) | exact", show(exact), "| first order", show(linear), "| interaction", show(interaction))
exact, linear, _ = margin_change(after_a, empty)
stale = margin_change(after_a, empty, margin_gradient(intact))[1]
print("m, b~ -> 0 at (0, 1) | exact", show(exact), "| first order", show(linear), "| (1, 1) gradient reused", show(stale))
exact, linear = probability_change(after_a, empty)
print("p, b~ -> 0 at (0, 1) | exact", show(exact), "| first order", show(linear))

# Signed and strict: a negative increment passes, equality does not. At tau equal to the second
# increment the edge is still kept.
assert removes(-1.0, TAU) and not removes(TAU, TAU)
second = prune(CLEAN, DONORS["zero"], ("A", "B"), TAU)[2][1]
assert prune(CLEAN, DONORS["zero"], ("A", "B"), second[5])[1] == ["B"]
# Appendix B.1 Eq. 2 as printed compares F(H) - F(trial) = -delta, which would remove B as well.
assert -second[5] < TAU and not second[6]
# Two clean records this page does not use, order A then B. A negative increment passes:
_, kept, steps = prune({"a": 0.5, "b": 0.5, "c": 0.0}, {"a": 0.0, "b": 0.75, "c": 0.0}, ("A", "B"), TAU)
assert steps[1][5] < 0 and kept == []
# and tau limits each step, not the result: both increments pass and the final F exceeds tau.
_, kept, steps = prune({"a": 1.0, "b": 0.5, "c": 0.0}, {"a": 0.5, "b": 0.0, "c": 0.0}, ("A", "B"), TAU)
assert kept == [] and all(step[5] < TAU for step in steps) and steps[1][4] > TAU
# Replacing node c changes c for its consumers; O is not one, so nothing O computes changes.
assert slots(dict(CLEAN, c=0.0), DONORS["zero"], EDGES) == slots(CLEAN, DONORS["zero"], EDGES)
# The direction of the divergence: forward, from the fixed reference. The reverse gives 0.027955.
assert abs(forward_kl(1.0, 0.5) - 0.0263445859768863) < 1e-12

Five short checks on the same graph. A hint comes before the answer. Beside each item this page notes your first check and any help opened; a later correct check never changes the first. Nothing is saved or sent, and reloading clears it. A shown answer is assistance, not an unassisted success. Opening practice or help does not change the figures above.

1. Zero donor, order B then A. What is the signed increment Δ at the second trial, to four decimal places?

   Δ =
   CheckHintShow the answer
2. Appendix B.1 Eq. 2, as printed in arXiv v2, removes an edge when F(H) − F(H_new) < τ. What would the zero-donor pass with order A then B return under that rule?

   The same graph {B}.The empty graph, with final F = 0.4621 nats.The full graph {A, B}.
   CheckHintShow the answer

3. At base (ã, b̃) = (0, 1), what change in m does the first-order estimate at that base predict for replacing b̃ by 0?

   first order =
   CheckHintShow the answer
4. With the half donor the pass removes both edges. What does that support?

   O does not need a or b.The full graph ignores its inputs.With this donor, replacing both slots leaves the positive class and F below τ; the result says more about the donor than about the edges.
   CheckHintShow the answer
5. c recovers y on both recorded rows. Which comparison tests whether O uses c?

   Replace c with a different value while a and b stay fixed, and recompute O.Collect more rows and check that c still matches y.Check that c agrees with a on every row.
   CheckHintShow the answer