causalab-mini

Use case · worked example

Layer sweep: where is "today" held?

One document, 32 experiments. Interchange the residual stream between two prompts at every layer, at two positions, and watch the answer cross from one token to the other.

Llama-3.2-1B completes "If today is Thursday, tomorrow is" with " Friday". Somewhere inside it, the identity of today has to move from the day word to the last token, where the next-token prediction is made. This page finds out where — by taking the residual stream from a run on a different day and writing it into the same place on this one, at each of the sixteen layers in turn.

That is an interchange intervention. If the patched model now answers the counterfactual day's successor, the thing you swapped was carrying the answer. The metric is iia: the fraction of rows whose top token after the patch is the counterfactual's answer.

At a glance #

Interchange accuracy per layer, at two positions Patching the day token carries the answer for layers 0 to 11 and then stops; patching the last token does nothing until layer 12 and then carries it. The two curves cross between layer 11 and layer 12. 0.000.250.500.751.00 0123456789101112131415 layer — patch the day token (pos -4) - - patch the last token (pos -1)
Interchange intervention accuracy, 42 ordered day pairs, one point per (layer, position). Run on CPU, fp32, Llama-3.2-1B — the numbers in the table below.

The crossover is between layer 11 and layer 12, and it is sharp. Up to layer 11, "today" is at the day token and the last token holds nothing; from layer 12 it is the other way round.

The data #

All 42 ordered pairs of weekdays. Every weekday name is a single token in this tokenizer, which is what makes the task scorable at one position.

$ uv run causalab-mini data weekdays_full/data
weekdays_full/data: 42 rows
  columns: base_answer, cf_answer, cf_entity, counterfactual_inputs, entity, example_id, input, split
  splits:  test, train
  sample:  {"example_id": "Sun-Mon", "input": "If today is Sunday, tomorrow is", "counterfactual_inputs": ["If today is Monday, tomorrow is"], "entity": "Sunday", "cf_entity": "Monday", "base_answer": " Monday", "cf_answer": " Tuesday", "split": "test"}

A metric column has to resolve to one token, and the CLI will tell you before you run anything:

$ uv run causalab-mini tokens meta-llama/Llama-3.2-1B " Monday" " Tuesday" "Sunday"
' Monday'                1 token(s)
' Tuesday'               1 token(s)
'Sunday'                 1 token(s)

The document #

documents/real/weekdays_layer_sweep.json — the document this page runs, verbatim
{
  "header": {
    "description": "Where is 'today' held? Interchange the residual stream between a base and a counterfactual day at every one of the 16 layers, at two positions: the day token itself (-4 of 'If today is X , tomorrow is') and the last token. 32 experiments, one document. `pos` is a NAMED axis: the read and the write move together, so it is 2 positions and not the 4 combinations a cross product would give."
  },
  "model": {
    "key": "meta-llama/Llama-3.2-1B",
    "revision": "4e20de362430cd3b72f300e6b0f18e50e7166e08",
    "dtype": "fp32"
  },
  "roles": {
    "base": {
      "field": "input"
    },
    "counterfactual": {
      "field": "counterfactual_inputs[0]"
    }
  },
  "sites": {
    "target": {
      "component": "block_output",
      "layers": [
        {
          "sweep": {
            "range": [
              0,
              16
            ]
          }
        }
      ]
    },
    "head": {
      "component": "lm_head"
    }
  },
  "interventions": {
    "patching": {
      "reads": {
        "v_cf": {
          "site": "target",
          "pos": {
            "sweep": [
              -4,
              -1
            ],
            "as": "pos"
          },
          "model": "original",
          "input": "counterfactual"
        },
        "logits": {
          "site": "head",
          "pos": -1,
          "model": "patched",
          "input": "base"
        }
      },
      "writes": {
        "patch": {
          "site": "target",
          "pos": {
            "sweep": [
              -4,
              -1
            ],
            "as": "pos"
          },
          "mechanism": "swap",
          "operand": "v_cf"
        }
      },
      "models": {
        "patched": {
          "input": "base",
          "writes": [
            "patch"
          ]
        }
      },
      "metrics": {
        "iia": {
          "kind": "match",
          "of": "logits",
          "token_form": "space_prefixed",
          "expected": "cf_answer"
        },
        "logit_diff": {
          "kind": "logit_diff",
          "of": "logits",
          "token_form": "space_prefixed",
          "a": "cf_answer",
          "b": "base_answer"
        }
      }
    }
  },
  "steps": {
    "score": {
      "kind": "observe",
      "interventions": "patching",
      "rows": {
        "base": "weekdays_full/data",
        "counterfactual": "weekdays_full/data"
      },
      "saves": [
        {
          "value": "iia",
          "file_path": "iia.json"
        },
        {
          "value": "logit_diff",
          "file_path": "logit_diff.json"
        }
      ]
    }
  }
}

Reading it in the order the compiler does:

  • model pins the checkpoint by revision and the dtype. The revision is part of the experiment, not a convenience.
  • roles maps each input role to a column: base is input, counterfactual is counterfactual_inputs[0]. Every step supplies rows for both.
  • sites names two places. target is block_output — the residual stream leaving a layer — at a layer the sweep fills in. head is lm_head, which takes no layer because it is not per-layer.
  • interventions.patching is the experiment. It reads v_cf from the unmodified model on the counterfactual prompt; it writes patch, a swap whose operand is that read; and models.patched is the base prompt run with that write applied. The second read, logits, is taken from the patched model.
  • metrics scores the logits read. iia is a match against the cf_answer column; logit_diff is the counterfactual answer's logit minus the base answer's.
  • steps.score is one observation pass over weekdays_full/data for both roles, saving the two metric tables.
the sweep axis is named

Two fields carry {"sweep": [-4, -1], "as": "pos"} — the read's position and the write's. Because they share the name pos, they move together: two points, not the four a cross product would give. The layer sweep is unnamed, so it is an axis of its own. 16 layers × 2 positions = 32 experiments, one document.

What it compiles to #

A sweep is lowered on the raw JSON before anything is compiled: each point becomes an ordinary document and is compiled on its own, with its own addresses, its own tokenization and its own digest. explain prints the whole tree, and it needs no weights:

$ uv run causalab-mini explain documents/real/weekdays_layer_sweep.json
<root>: Plan
  layers=0,pos=-4: Plan
    score: Observe  metrics=['iia/match', 'logit_diff/logit_diff']  saves=['iia.json', 'logit_diff.json']
        forward 'original' on 'counterfactual'  (42 rows x 8 tokens)
          read  'v_cf' at block_output[0] pos={index:-4} via 'identity'
        forward 'patched' on 'base'  (42 rows x 8 tokens)
          write 'patch' at block_output[0] pos={index:-4} swap(v_cf) via 'identity'
          read  'logits' at lm_head pos={index:-1} via 'identity'
  layers=0,pos=-1: Plan
    score: Observe  metrics=['iia/match', 'logit_diff/logit_diff']  saves=['iia.json', 'logit_diff.json']
        forward 'original' on 'counterfactual'  (42 rows x 8 tokens)
          read  'v_cf' at block_output[0] pos={index:-1} via 'identity'
        forward 'patched' on 'base'  (42 rows x 8 tokens)
          write 'patch' at block_output[0] pos={index:-1} swap(v_cf) via 'identity'
          read  'logits' at lm_head pos={index:-1} via 'identity'
…
  layers=15,pos=-1: Plan
    score: Observe  metrics=['iia/match', 'logit_diff/logit_diff']  saves=['iia.json', 'logit_diff.json']
        forward 'original' on 'counterfactual'  (42 rows x 8 tokens)
          read  'v_cf' at block_output[15] pos={index:-1} via 'identity'
        forward 'patched' on 'base'  (42 rows x 8 tokens)
          write 'patch' at block_output[15] pos={index:-1} swap(v_cf) via 'identity'
          read  'logits' at lm_head pos={index:-1} via 'identity'

Each point is two forwards over the same 42 rows: the counterfactual prompts through the unmodified model to get the value, then the base prompts with that value written in. pos={index:-4} is the day word — the prompts are If today is X , tomorrow is, so it is four tokens from the end — and pos={index:-1} is the token the prediction is read from.

Running it #

$ uv run causalab-mini validate documents/real/weekdays_layer_sweep.json
ok: documents/real/weekdays_layer_sweep.json is a valid plan-shaped document
$ time uv run causalab-mini run documents/real/weekdays_layer_sweep.json \
      --device-map cpu --out out/sweep
out/sweep/document.json
out/sweep/run.json
out/sweep/layers=0,pos=-4/document.json
out/sweep/layers=0,pos=-4/iia.json
out/sweep/layers=0,pos=-4/logit_diff.json
out/sweep/layers=0,pos=-1/document.json
…
out/sweep/layers=15,pos=-1/logit_diff.json

real	1m33.389s
what you need to reproduce this

Llama-3.2-1B is a gated checkpoint, so the first run needs huggingface-cli login. No GPU: the run above is CPU-only, fp32, 93 seconds wall clock including the model load, on a 12-core Ryzen 9 5900X. With --device-map cuda it is about 12 seconds (documents/real/README.md).

What lands on disk #

A swept document's root plan holds one child plan per point, and a plan's path in the tree is its path on disk — which is where layers=12,pos=-1/ comes from. The root carries the document and the run record; each point carries its own compiled document and its two metric tables.

out/sweep/
├── document.json          the document as written, sweeps and all
├── run.json               engine, library versions, package digest, served dtype
├── layers=0,pos=-4/
│   ├── document.json      this point's document, the sweep lowered out
│   ├── iia.json           42 rows
│   └── logit_diff.json    42 rows
├── layers=0,pos=-1/
…
└── layers=15,pos=-1/
$ cat out/sweep/run.json
{
 "engine": "NNterpEngine",
 "remote": false,
 "batch_size": null,
 "versions": {
  "torch": "2.14.0",
  "transformers": "5.17.0",
  "nnsight": "0.8.1.dev125+ga8ee93782",
  "nnterp": "1.3.1.dev64+g0e4b401af",
  "pydantic": "2.13.5"
 },
 "causalab_mini": "76c8c9d0c0d1f3e0b8e2eaeb09564107d1c0a70cb9e8b7d6cecab09fdc92d8a1",
 "served_dtype": "torch.float32"
}

One row of a metric table, from the point where the crossover has just happened:

$ head -13 out/sweep/layers=12,pos=-1/iia.json
[
 {
  "example_id": "Sun-Mon",
  "metric": "iia",
  "value": 1.0,
  "eligible": true,
  "positions": null,
  "reason": "",
  "unit": "fraction",
  "estimand_version": "match/v1",
  "produced_by": "d62e457f5a2e46afc0d95bc208ea5112f589c5bbf5a555421be6a16f421c5724"
 },
…

value is per row, not per point: 1.0 means this pair's patched model did answer the counterfactual's day. positions is null and reason empty because every position in this document is a fixed integer — there was nothing the run had to resolve against the text, and nothing that could fail to be placed. produced_by is this point's digest, so two points of one sweep are never confusable.

The curve #

The mean over the 42 rows, per point. Both metrics, both positions:

layerIIA, day tokenIIA, last tokenlogit diff, daylogit diff, last
01.000.00+2.40-2.40
11.000.00+2.40-2.41
21.000.00+2.38-2.41
31.000.00+2.39-2.40
41.000.00+2.43-2.40
51.000.00+2.42-2.39
61.000.00+2.40-2.37
71.000.00+2.37-2.33
81.000.00+2.28-2.26
91.000.00+2.23-2.21
101.000.00+2.14-2.10
110.980.00+1.24-1.21
120.100.88-1.18+1.15
130.100.88-1.18+1.16
140.001.00-2.32+2.31
150.001.00-2.40+2.40

What the numbers mean #

Three things are visible, and each of them is a claim about the model.

The day token holds "today" through layer 11. Swapping block_output[0..10] at position −4 moves every one of the 42 rows to the counterfactual's answer — IIA 1.00, logit difference about +2.4 in favour of the counterfactual day. Early in the stack, the day word's own residual stream is the fact.

It crosses to the last token at layer 12. Between layer 11 and layer 12 the day token's IIA falls from 0.98 to 0.10 and the last token's rises from 0.00 to 0.88; by layer 14 it is 1.00 at the last token and 0.00 at the day token. The two curves are close to mirror images, which is what you want to see: the information has not been destroyed, it has moved. Layers 12 and 13 are the transition — 0.88, not 1.00 — so the move is not quite instantaneous.

Before layer 12, the last token holds nothing. IIA 0.00 and a logit difference of −2.4: patching it leaves the base answer exactly where it was. That is the control this sweep gets for free, and it is why sweeping the position alongside the layer is worth the second axis.

what this does not say

An interchange result localizes a place where the counterfactual value is sufficient to change the answer. It does not say the model computes the answer there, it does not say the representation is a direction rather than a whole activation, and at layer 15 — the last block output, one norm away from the logits — a perfect IIA is close to tautological. The neighbouring documents in documents/real/ are what narrow it: head patching finds which head does the moving, DAS finds how many dimensions carry it.

Recorded results for the same document #

The numbers above were produced by running the document on CPU while writing this page. documents/real/README.md records an earlier run of the same file on an A100 and through a self-hosted NDIF, and reports the surrounding experiments:

documentwhat it asksrecorded result
weekdays_behaviordoes the model do the task at all? 42/42 correct, p(answer) 0.45 mean
weekdays_layer_sweepthis page day token 1.00 to layer 10, 0.98 at 11, then 0.10, 0.10, 0.00, 0.00; last token 0.00 to layer 11, then 0.88, 0.88, 1.00, 1.00
weekdays_head_patchingone head's result at the last token, 32 heads × layers 10–12 one head of 96 stands out: L12 H28, mean logit difference −2.40 (median head) to +0.05
weekdays_dasa k-dim subspace of the day token at layer 8, fit on 30 pairs held-out IIA 0.00 / 0.83 / 1.00 / 1.00 at k = 1 / 4 / 16 / 64
weekdays_pattern_patchingswap only a head's attention pattern at the last token nothing moves: L12 H28 gives −2.396 against a −2.403 median

Read together: "today" sits at the day token through layer 11, crosses to the last token at layer 12 largely through one head, and is a roughly 16-dimensional subspace of a 2048-wide stream. And it is what that head moves, not where it looks — its pattern from another prompt changes nothing, while its values from another prompt change the answer.

The recorded run also notes that a self-hosted NDIF served bf16 whatever the document asked, so the remote numbers are close but not equal: the last-token sweep reads 0.86 there where fp32 reads 0.88. The two local engines — nnsight and plain forward hooks — agree exactly on all 672 numbers of a layer sweep.

See also #

  • Sweeps — named axes, the range form, and what a sweep may not touch
  • Positions — why -4 is a fixed integer here and what to write when it cannot be
  • Locations — block_output, lm_head, and the rest of the component vocabulary
  • Attention pattern — the follow-up that separates what a head moves from where it looks
  • DAS fit — the follow-up that asks how many dimensions carry it