Gradient contract¶

Every edit either keeps the autograd graph through the values it writes or cuts it. An edit that keeps the graph lets the gradient of a loss on the output reach the parameters upstream of the edited site; an edit that cuts it gives those parameters no gradient through the edited elements. Cutting is correct for analysis (a patched value is evidence from another run, not a function of this one) and silent for training, so a caller who needs the gradient says so with torchlens.intervention.keep_grad(edit).

The rules¶

• Verbs whose result is a function of the site value keep the graph: scale, clamp, add_noise, steer, add, project_onto, project_out, mean_ablate(dim=...) without source= (the gradient flows through the mean), scramble_elements(seed=...) without source=, splice_module, permute_batch, set_direction_mean (through the orthogonal part), the subspace_* edits (through the complement), compose and at_positions of keeping verbs, and hook callables that compute from their input.
• Verbs that write a value computed elsewhere cut it: zero_ablate, replace_with, swap_with, patch_from (the donor is detached by design), resample_rows_from, and mean_ablate or scramble_elements with source= (a tensor or a population).
• A selection door edits only the selected elements, so the rest of the site keeps the graph: under do(tl.units(...)) a cutting verb still lets the gradient through the unselected elements, and under do(tl.subspace(...)) through the orthogonal complement.
• The capture doors, fork.do (replay), do(model=, inputs=) (rerun), do with InterventionOptions(engine="set_only") + run and spec.bind(model) keep whatever the verb keeps. do(site, edit, intervention=InterventionOptions(differentiable=True)) re-leafs the replayed values by design, so no gradient reaches the parameters upstream of the edit under any verb; the gradient stops at the replay frontier. To train parameters upstream of an edit, use a live door (tl.when(site, edit).bind(model)(x) or tl.trace(model, x, intervene=...)).

Asking for the gradient: keep_grad¶

torchlens.intervention.keep_grad(edit) returns the same edit with a fire-time check. When the values the edit writes carry no autograd graph while the site value had one, it raises InvalidArgumentError with code edit_detached_from_graph, the site label and a spelling that keeps the graph. It also raises when the site value has no graph to keep (reason site_has_no_graph) or when gradients are disabled (grad_disabled). InterventionDecision(keep_grad=True) runs the same check; keep_grad=False, the default, checks nothing. The check is on the edit, not the door: under differentiable replay a keeping verb passes it and still delivers no gradient upstream (see the matrix), so there keep_grad warns once per session with code keep_grad_replay_frontier: the edit stays connected inside the replayed computation, parameters upstream of the replay frontier get no gradient, and the live door tl.when(site, keep_grad(edit)).bind(model)(x) reaches them.

The graph-keeping spellings:

• a zero: keep_grad(tl.scale(0.0)) writes the same zeros as tl.zero_ablate() with a zero gradient;
• a fixed value: keep_grad(tl.intervention.compose(tl.scale(0.0), tl.add(value))) writes value with a zero gradient;
• an interchange that trains: do(tl.subspace(site, basis, origin=...), tl.patch_from(source)) keeps the gradient through the orthogonal complement, while the written component is the source's (detached).

# docs-test: expected=[True, True, True]
import torch
from torch import nn
import torchlens as tl
from torchlens.intervention import keep_grad

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(4, 4), nn.Linear(4, 2)).eval()
x = torch.randn(3, 4)
site = tl.module("0")

try:
    tl.trace(model, x, intervene=tl.when(site, keep_grad(tl.zero_ablate())))
    refused = False
except ValueError as error:
    refused = error.fields["code"] == "edit_detached_from_graph"

trace = tl.trace(model, x, intervene=tl.when(site, keep_grad(tl.scale(0.0))))
out = trace[trace.output_layers[0]].out
out.sum().backward()

RESULT = [
    refused,
    out.requires_grad,
    bool((model[0].weight.grad == 0).all()),
]

The matrix¶

One row per door and verb; the two columns are the outcome without and with keep_grad. Each cell is checked by tests/test_gradient_contract_matrix.py on a float64 two-layer model with the edit at the first layer's output: reaches means output.sum().backward() gives the first layer's weight a gradient, equal to the gradient of the same edit applied by hand in a plain torch forward hook; cut means it gives none. raises:<code> is a typed refusal. Under do(tl.subspace(...)) the verbs act on the subspace coordinates or write a value's component (see subspace edits), and the projection edits refuse subspace_edit_unsupported there. set_affine_coordinates is a source verb: it keeps the graph under do(tl.subspace(...)) (its new coordinates W x + b are a function of the site value) and refuses affine_source_needs_subspace at every other door. Every do(differentiable=True) cell reads cut because the replay frontier re-leafs the edited value; each of its keep_grad cells that does not raise also carries the keep_grad_replay_frontier warning, and the test checks it.