Metrics: turning a change into a number you can defend

Note

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

Every intervention in Murano produces logits. A metric turns those logits into a scalar, and the scalar is what you argue with. Choosing it badly is the easiest way to reach a confident wrong conclusion, because different metrics disagree about the same intervention.

This notebook runs every metric step on one intervention and shows them disagreeing.

Key questions

• What does each metric step actually measure?
• Can a metric say nothing changed while the behavior flips?
• Can a metric say nothing changed while the logit difference moves?

Structure

1. Set up
2. Logit difference: did the answer flip?
3. KL divergence: did the distribution move?
4. Answer log-probability: how sure is the model?
5. Recovered fraction: normalizing against two baselines
6. Comparing two tensors directly
7. Answer rank: the metric that saturates
8. 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. No extras: pip install murano-interp

1. Set up

import torch

from murano import MuranoModel, Pipeline, keys
from murano.nodes import Node
from murano.steps import (
    Ablate,
    AnswerLogProbStep,
    AnswerRankStep,
    ComparisonComputationStep,
    KLDivergenceStep,
    LoadPaired,
    LoadPrompts,
    LogitDiffStep,
    Logits,
    RecoveredMetricStep,
    Save,
)

OUTPUT_DIR = "murano_outputs/metrics"

# 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

One intervention runs through the whole notebook: zero-ablate head L9H9 on the clean prompts. Logits writes the clean logits, Ablate writes the damaged ones, and every metric below reads that same pair.

HEAD = Node(9, "self_attn", head=9)

results = Pipeline(
    [
        LoadPrompts(task.clean),
        Logits(model),
        Ablate(model, targets=[HEAD], method="zero"),
    ]
).run()

print("clean logits :", tuple(results[keys.FINAL_LOGITS].shape))
print("ablated logits:", tuple(results[keys.ABLATED_LOGITS].shape))

clean logits : (8, 14, 50257)
ablated logits: (8, 14, 50257)

2. Logit difference: did the answer flip?

LogitDiffStep subtracts the incorrect answer’s logit from the correct one, at the answer position. It is the standard metric for circuit work because it is task aware: it asks about the two tokens the task is about, and ignores the other 50,255.

results = LogitDiffStep(
    task.correct, task.incorrect, output_key="ld_clean", model=model
)(results)
results = LogitDiffStep(
    task.correct,
    task.incorrect,
    logits_key=keys.ABLATED_LOGITS,
    output_key="ld_ablated",
    model=model,
)(results)

clean = results["ld_clean"].value
ablated = results["ld_ablated"].value
print(f"clean   : {clean:+.4f}")
print(f"ablated : {ablated:+.4f}   (change {ablated - clean:+.4f})")

clean   : +2.3971
ablated : +2.3897   (change -0.0073)

3. KL divergence: did the distribution move?

KLDivergenceStep compares the two output distributions at the answer position. It is task agnostic: it does not know which tokens matter, only how much probability mass moved anywhere.

Sanity first. A distribution against itself must score zero.

results = KLDivergenceStep(
    p_key=keys.FINAL_LOGITS, q_key=keys.FINAL_LOGITS, output_key="kl_self"
)(results)
results = KLDivergenceStep(
    p_key=keys.FINAL_LOGITS, q_key=keys.ABLATED_LOGITS, output_key="kl_ablated"
)(results)

print(f"KL(clean || clean)   = {results['kl_self'].value:.6f}   (must be 0)")
print(f"KL(clean || ablated) = {results['kl_ablated'].value:.6f}")

KL(clean || clean)   = 0.000000   (must be 0)
KL(clean || ablated) = 0.028441

Hold these two numbers side by side. The KL divergence is small: the output distribution barely moved. The logit difference also barely moved, because zeroing this head happens to be a weak intervention.

The trap appears when they disagree. A task-aware metric can collapse while KL stays tiny, because the two answer tokens carry a vanishing share of the total probability mass. A small KL is not evidence that a component does not matter. ablation.ipynb shows exactly this: resampling L9H9 destroys the logit difference while the KL stays around 0.03.

4. Answer log-probability: how sure is the model?

AnswerLogProbStep reports the log-probability of the correct token, or with as_loss=True, its negative, which is the per-example cross-entropy of the answer.

Unlike the logit difference, this is sensitive to the whole distribution: making some third token more likely lowers it, even if the correct token still beats the incorrect one.

results = AnswerLogProbStep(correct=task.correct, output_key="alp", model=model)(
    results
)
results = AnswerLogProbStep(
    correct=task.correct, as_loss=True, output_key="alp_loss", model=model
)(results)

print(f"{'log-prob':<26} {results['alp'].value:+.4f}")
print(f"{'negative log-prob (loss)':<26} {results['alp_loss'].value:+.4f}")

log-prob                   -0.9642
negative log-prob (loss)   +0.9642

5. Recovered fraction: normalizing against two baselines

Raw scores are hard to compare across tasks and models. RecoveredMetricStep rescales a score against a clean and a corrupt baseline:

recovered = (patched - corrupt) / (clean - corrupt)

0 means the intervention achieved nothing, 1 means it fully restored clean behavior. Values outside [0, 1] are meaningful: negative means it made things worse than corrupt.

paired = 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()

# Score the clean run against itself: by construction it recovers exactly 1.0.
paired["ld_patched"] = paired["ld_clean"]
paired = RecoveredMetricStep("ld_clean", "ld_corrupt", "ld_patched")(paired)

print(f"clean     : {paired['ld_clean'].value:+.4f}")
print(f"corrupt   : {paired['ld_corrupt'].value:+.4f}")
print(f"recovered : {paired[keys.RECOVERED].value:.4f}   (a clean run recovers 1.0)")

clean     : +2.3971
corrupt   : -3.1112
recovered : 1.0000   (a clean run recovers 1.0)

6. Comparing two tensors directly

ComparisonComputationStep is the escape hatch: give it two tensors already in Results and it returns their element-wise difference or their row-wise cosine similarity. It is not a metric step, it writes a tensor, not a MetricScore.

The cosine result below is a warning worth internalizing.

for kind in ["difference", "cosine_similarity"]:
    results = ComparisonComputationStep(
        key_a=keys.FINAL_LOGITS,
        key_b=keys.ABLATED_LOGITS,
        output_key=f"cmp_{kind}",
        comparison_type=kind,
    )(results)
    out = results[f"cmp_{kind}"]
    print(f"{kind:<18} shape={tuple(out.shape)}  mean={float(out.float().mean()):+.6f}")

difference         shape=(8, 14, 50257)  mean=+0.470193
cosine_similarity  shape=(8, 14)  mean=+0.999999

The cosine similarity between the clean and ablated logits is essentially exactly 1.0, even though the ablation demonstrably changed the model’s output.

That is not a bug. A logit vector over 50,257 tokens is dominated by the same enormous shared structure in every forward pass: almost every token is very unlikely, and the vector points in nearly the same direction regardless. Cosine similarity on logits is close to useless.

Use the difference when you want to know which tokens moved, KL when you want the distributions compared honestly, and the logit difference when you have a task.

7. Answer rank: the metric that saturates

AnswerRankStep counts how many tokens outscore the correct answer at the answer position. Rank 0 means the model would emit the answer greedily.

It reads the answer position from the attention mask rather than taking the last column, which is the same resolution every metric step here uses. That matters: under right-side padding the last column is a padding position, and a metric that scored it would report a confident number about a token the model never predicted.

results = AnswerRankStep(task.correct, model=model)(results)
results = AnswerRankStep(
    task.correct,
    logits_key=keys.ABLATED_LOGITS,
    mask_key=keys.ATTENTION_MASK,
    output_key="answer_rank_ablated",
    model=model,
)(results)

print(f"mean rank of the correct name, clean   : {results['answer_rank'].value:.2f}")
print(
    f"mean rank of the correct name, ablated : {results['answer_rank_ablated'].value:.2f}"
)
print(f"per-example (clean): {results['answer_rank'].per_example}")

mean rank of the correct name, clean   : 0.00
mean rank of the correct name, ablated : 0.00
per-example (clean): [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]

Both numbers are 0: the correct name is the model’s top prediction before and after the ablation. The metric sees nothing.

That is not a failure of the code, it is the notebook’s thesis arriving one last time. Rank is quantized. It cannot distinguish “barely winning” from “winning by a mile”, so a weak intervention is invisible to it while the logit difference registers a change and the KL divergence registers another. Every metric is a lens, and each one is blind to something. Pick the one that can see the effect you are claiming, and say which one you picked.

When no shipped metric can see your effect, write a Step: read the keys, write a MetricScore, and every other step composes with it unchanged. custom_pipeline.ipynb does exactly that.

8. Save

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

print("saved to:", run_dir)

saved to: murano_outputs/metrics

What next

• ablation.ipynb uses these metrics to show that the choice of ablation decides the answer.
• activation_patching.ipynb builds a causal map out of the recovered fraction.
• custom_pipeline.ipynb writes the metric step the library does not ship.