causalab-mini

A small, readable intervention engine

A causal-intervention experiment is a JSON document and nothing else.

A document is pure data: which model, which rows, which places inside it to read and to write, and what to score. The compiler resolves everything that needs a model but not a row — the widths, a module path, an interior's operation, the prompts' token ids — on the client, without weights, and hands back a plan of strings, integers and tuples. An engine walks that plan in one session and turns it into tensors. The results come home by filling the plan in: every number sits on the node that produced it, and a node's path in the tree is its path on disk.

It is a clean-room reimplementation of causalab's intervention engine, deliberately small — about 7,400 lines — so that you can read the whole thing. It imports nothing from causalab.

The plan that goes out to the engine is the plan that comes back. Nothing is re-parsed on the way home: a result is set on the node that produced it, and a node's path in the tree is its path on disk.

Install #

git clone https://github.com/JadenFiotto-Kaufman/causalab-mini
cd causalab-mini
uv sync

That is the whole setup. The read-only verbs of the command line need no GPU, no weights and no network beyond a model's config, and run in seconds:

uv run causalab-mini vocab                  # components, mechanisms, kinds, position forms
uv run causalab-mini data weekdays/train    # columns, rows, a sample
uv run causalab-mini validate documents/v2/patching.json
uv run causalab-mini explain  documents/v2/patching.json

Run one document #

The smallest real experiment in the repository is an interchange intervention: read the residual stream at the last token of a counterfactual prompt, swap it into the same place on the base prompt, and score what the model then says. It runs on a tiny random-weight Llama, on CPU, in a few seconds.

documents/v2/patching.json — the whole document, nothing elided

{
  "header": {
    "description": "The smallest real interchange \u2014 one swap at the last position of layer 0 \u2014 in causalab-mini's plan-shaped format. Compare documents/minimal_cpu.json, the same experiment in the protocol's. Here there is one step, because the run is one pass."
  },
  "model": {
    "key": "hf-internal-testing/tiny-random-LlamaForCausalLM",
    "revision": "9fb191250dd56d0ba7ec9785a025ed29c03d5998",
    "dtype": "fp32"
  },
  "roles": {
    "base": {
      "field": "input"
    },
    "counterfactual": {
      "field": "counterfactual_inputs[0]"
    }
  },
  "sites": {
    "target": {
      "component": "block_output",
      "layers": [
        0
      ]
    },
    "head": {
      "component": "lm_head"
    }
  },
  "interventions": {
    "patching": {
      "reads": {
        "v_cf": {
          "site": "target",
          "pos": -1,
          "model": "original",
          "input": "counterfactual"
        },
        "logits": {
          "site": "head",
          "pos": -1,
          "model": "patched",
          "input": "base"
        }
      },
      "writes": {
        "patch": {
          "site": "target",
          "pos": -1,
          "mechanism": "swap",
          "operand": "v_cf"
        }
      },
      "models": {
        "patched": {
          "input": "base",
          "writes": [
            "patch"
          ]
        }
      },
      "metrics": {
        "iia": {
          "kind": "match",
          "of": "logits",
          "expected": "cf_answer",
          "token_form": "space_prefixed"
        },
        "logit_diff": {
          "kind": "logit_diff",
          "of": "logits",
          "a": "cf_answer",
          "b": "base_answer",
          "token_form": "space_prefixed"
        }
      }
    }
  },
  "steps": {
    "score": {
      "kind": "observe",
      "interventions": "patching",
      "rows": {
        "base": "weekdays/train",
        "counterfactual": "weekdays/train"
      },
      "saves": [
        {
          "value": "iia",
          "file_path": "iia.json"
        },
        {
          "value": "logit_diff",
          "file_path": "logit_diff.json"
        }
      ]
    }
  }
}

Four blocks do the work. roles says which column of a row each input reads. sites names the places — a component, and a layer where the component is per-layer. interventions is the experiment: the reads, the writes, the models a write applies to, and the metrics. steps says which rows to run over and which results to write out.

explain compiles it against the model without loading any weights, and prints what would run:

$ uv run causalab-mini explain documents/v2/patching.json
<root>: Plan
  score: Observe  metrics=['iia/match', 'logit_diff/logit_diff']  saves=['iia.json', 'logit_diff.json']
      forward 'original' on 'counterfactual'  (4 rows x 11 tokens)
        read  'v_cf' at block_output[0] pos={index:-1} via 'identity'
      forward 'patched' on 'base'  (4 rows x 11 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'

and run does it and writes the files:

$ uv run causalab-mini run documents/v2/patching.json --device-map cpu --out out
out/document.json
out/run.json
out/iia.json
out/logit_diff.json

Read the output #

A metric table is a JSON array of row objects, one file per metric, with the labels repeated on every row so jq and a person can both read it. One row of out/iia.json:

{
 "example_id": "0",
 "metric": "iia",
 "value": 0.0,
 "eligible": true,
 "positions": null,
 "reason": "",
 "unit": "fraction",
 "estimand_version": "match/v1",
 "produced_by": "ffd162b5011bc74abf2f400d71595f63849a393a1fcf78aa234d26bbc77e791e"
}

what the fields are for

eligible is false when the measurement was excluded — the row's data had no answer to score, or the run could not place the position — and then value is null and reason says which. An excluded row is still a row, so it can never be read as a zero or silently shorten a denominator. produced_by is the digest of the experiment that made it: everything in the document except header.

Beside the metrics, run writes document.json (the experiment, verbatim) and run.json (the engine, the library versions, a digest of this package, and the dtype the model was actually served in). A result is never a file with no way back to what produced it.

These numbers are zero because the checkpoint has random weights — it is a fixture that proves the machinery, not a model that knows anything. For a real result, see the worked example below.

A map of the site #

The document #

Use cases #

Ten worked experiments, each one a real document run end to end. The layer sweep is the one to read first: it is a complete result on a real model, from the document to the curve.