Skip to content

Steps — Ablate

Ablate step: zero, mean, or resample a component and read the logits.

Turns a component off and measures what breaks. The step rewrites one or more target activations during a single forward pass and stores the resulting logits, so the existing metric steps (logit-diff, KL, recovered) can score the clean run against the ablated run and quantify the component’s causal contribution.

Three replacement methods share one mechanism and differ only in the value that overwrites the target:

  • "zero" writes zeros (the component contributes nothing).
  • "mean" writes the component’s mean activation (it contributes only its average, removing the input-specific signal).
  • "resample" writes another input’s activation at the same site (the within-batch permutation, or a provided corrupted batch), the noising baseline used by activation patching.

Targets are :class:~murano.nodes.Node addresses. A whole-component target (head is None) ablates the module output; an attention head target (Node(layer, "self_attn", head=h)) ablates only that head’s slice of the attention output projection input. One :class:Ablate call is a single mode: all whole-component or all per-head, not a mix.

Ablate target components and write the resulting forward-pass logits.

Runs one intervened forward pass over the batched prompts: each target component’s activation is replaced by zeros, its mean, or a resampled activation, and the model’s output logits are stored under logits_key. Comparing those logits against a clean :class:~murano.steps.logits.Logits run with a metric step quantifies the component’s contribution.

Reads from results:

  • results[prompts_key]: PromptBatch (default prompts), the batch to run.
  • results[source_key]: PromptBatch, only when source_key is set: the
  • batch whose activations are resampled in (the cross-run patch source).

Writes to results:

  • results[logits_key]: Tensor [B, S, vocab] of ablated logits.
  • results[mask_key]: Tensor [B, S] marking real (1) vs padding (0), so a
  • downstream metric step can locate the answer position.

Args:

  • model: Model backend to run.
  • targets: Components to ablate: a :class:NodeSet, a single address, or
  • an iterable of addresses. A whole-component target ablates the module
  • output; a head target (Node(layer, "self_attn", head=h)) ablates
  • that head. One call is a single mode (all whole-component or all
  • per-head). Pass this or targets_key, not both.
  • targets_key: Results key holding a
  • : class:~murano.artifacts.ComponentSelection a discovery step wrote,
  • read at run time so attribute-then-ablate composes in one pipeline.
  • Pass this or targets, not both.
  • method: "zero", "mean", or "resample".
  • mean_over: For method="mean", whether the mean pools over batch and
  • positions ("all") or is taken per position ("position").
  • means: For method="mean", an optional precomputed “{address:
  • tensor}“ of mean vectors (e.g. a corpus mean) used instead of the
  • current batch’s mean. Whole-component targets take one d_model
  • vector per site (keyed by Node(layer, module)); per-head targets
  • take one head_dim vector per head (keyed by “Node(layer,
  • “self_attn”, head=h)“), so a head can be mean-ablated to its mean over
  • a reference distribution.
  • source: For method="resample", an optional sequence of corrupted
  • prompts paired one-to-one with the inputs; each must be
  • token-length-matched to its clean prompt so positions align. Without
  • it (and without source_key), resample permutes within the batch.
  • source_key: For method="resample", an optional Results key naming a
  • second: class:~murano.artifacts.PromptBatch (already in results)
  • whose activations are resampled in, the cross-run patch source. Same
  • token-length-alignment requirement as source; mutually exclusive
  • with it. The: class:~murano.steps.patch.Patch step builds on this.
  • permutation: For within-batch resample, an explicit rearrangement of
  • range(batch) (for reproducibility); otherwise a random one.
  • seed: Seed for the random within-batch permutation when permutation
  • is not given.
  • positions: Token position(s) to ablate, in the
  • : func:~murano.steps.metrics._answer_positions form (an int or
  • per-example sequence, negatives allowed); None ablates every
  • position.
  • prompts_key: Results key to read the batch to run from (default
  • prompts); set it to run on, e.g., the corrupt side of a pair.
  • logits_key: Results key to write the ablated logits under.
  • mask_key: Results key to write the attention mask under.

Raises:

  • ValueError: If neither or both of targets and targets_key are
  • given, method/mean_over is unknown, targets is empty or
  • mixes modes, a method-specific argument is misused, or a precomputed
  • means table does not cover every target with a right-sized vector.
  • NotImplementedError: If per-head capture is requested for a module whose
  • architecture’s attention output projection is not recognized.
def __init__(self, model: ModelBackend, targets: NodeSet | AddressLike | Iterable[AddressLike] | None = None, method: Literal['zero', 'mean', 'resample'] = 'zero', targets_key: str | None = None, mean_over: Literal['all', 'position'] = 'all', means: dict[AddressLike, torch.Tensor] | None = None, source: Sequence[str] | None = None, source_key: str | None = None, permutation: Sequence[int] | None = None, seed: int | None = None, positions: int | Sequence[int] | torch.Tensor | None = None, prompts_key: str = keys.PROMPTS, logits_key: str = keys.ABLATED_LOGITS, mask_key: str = keys.ATTENTION_MASK):