Use case

DAS fit

Distributed alignment search: fitting a subspace, then scoring it held out.

The layer sweep found where "today" lives: the day token's residual stream, through layer 11. That is 2048 numbers. This page asks how many of them it takes — not which coordinates, but which directions.

Distributed alignment search answers it by learning the directions. Fit an orthonormal k-dimensional basis Q of the activation, swap the counterfactual's component in that subspace and leave the orthogonal complement alone; train Q so that the swap makes the model answer the counterfactual's day. Score it on pairs the fit never saw. Sweep k and the document is a rank curve.

At a glance #

Held-out interchange accuracy of a fitted subspace of the day token's residual stream at layer 8 of Llama-3.2-1B, against its rank. 30 training pairs, 12 held-out. Run on CPU, fp32, 2026-09-23.

One direction carries nothing. Four carry most of it. Sixteen carry all of it, and sixty-four carry no more. "Today" is a roughly 16-dimensional subspace of a 2048-wide stream.

The document #

documents/real/weekdays_das.json — the document this page runs, verbatim

{
  "header": {
    "description": "DAS on a real model: a k-dimensional subspace of the day token's residual stream that, swapped alone, carries 'today' from the counterfactual. Fit on 30 pairs, scored on 12 held-out pairs; k is swept, so the document is the rank curve."
  },
  "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": [
        8
      ]
    },
    "head": {
      "component": "lm_head"
    }
  },
  "featurizers": {
    "rot": {
      "kind": "subspace",
      "k": {
        "sweep": [
          1,
          4,
          16,
          64
        ]
      }
    }
  },
  "interventions": {
    "das": {
      "reads": {
        "v_cf": {
          "site": "target",
          "pos": -4,
          "model": "original",
          "input": "counterfactual",
          "featurizer": "rot"
        },
        "logits": {
          "site": "head",
          "pos": -1,
          "model": "patched",
          "input": "base"
        }
      },
      "writes": {
        "patch": {
          "site": "target",
          "pos": -4,
          "mechanism": "swap",
          "operand": "v_cf",
          "featurizer": "rot"
        }
      },
      "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"
        },
        "ce": {
          "kind": "cross_entropy",
          "of": "logits",
          "token_form": "space_prefixed",
          "target": "cf_answer"
        }
      }
    }
  },
  "steps": {
    "fit": {
      "kind": "fit",
      "interventions": "das",
      "rows": {
        "base": "weekdays_full/data#train",
        "counterfactual": "weekdays_full/data#train"
      },
      "params": [
        "rot"
      ],
      "objective": [
        [
          1.0,
          "ce"
        ]
      ],
      "epochs": 20,
      "pairs": 10,
      "seed": 0,
      "optimizer": {
        "name": "adamw",
        "lr": 0.003,
        "weight_decay": 0.0
      },
      "early_stop": {
        "metric": "iia",
        "mode": "max",
        "patience": 20
      },
      "eval": {
        "rows": {
          "base": "weekdays_full/data#test",
          "counterfactual": "weekdays_full/data#test"
        },
        "saves": [
          {
            "value": "iia",
            "file_path": "held_out_iia.json"
          }
        ]
      }
    },
    "score": {
      "kind": "observe",
      "interventions": "das",
      "rows": {
        "base": "weekdays_full/data#test",
        "counterfactual": "weekdays_full/data#test"
      },
      "saves": [
        {
          "value": "iia",
          "file_path": "iia.json"
        },
        {
          "value": "logit_diff",
          "file_path": "logit_diff.json"
        }
      ]
    },
    "weights": {
      "kind": "weights",
      "names": [
        "rot"
      ],
      "saves": [
        {
          "value": "rot",
          "file_path": "rot.safetensors"
        }
      ]
    }
  }
}

Reading it in the order the compiler does:

model pins Llama-3.2-1B by revision, fp32.
roles and sites are the layer sweep's: block_output, here fixed at layer 8, and lm_head.
featurizers.rot is the new block. kind: "subspace", and k is a sweep over 1, 4, 16, 64 — so this one document is four independent fits.
interventions.das is the layer sweep's interchange with one field added on each side: "featurizer": "rot" on the read and on the write. Three metrics now, because a fit needs a differentiable one and a report wants a legible one.
steps has three. fit trains rot; score measures it on the held-out split; weights publishes the parameter so the save can write it.

What the featurizer does #

featurize(x)      = (Qᵀx, 0)
inverse(f, 0, x)  = Q f + (x − Q Qᵀ x)

A write is inverse(do(featurize(x)), err, x) for every featurizer, so with do a swap this reads: take the subspace from the counterfactual, keep everything else from here. The complement comes from x, the pre-write activation, which is why the whole activation has to reach the featurizer. Nothing is discarded that inverse could not recover, so the error term is zero.

`k`, `d`, and the Cayley parametrization #

Q is never stored. The stored parameter is weight, an ordinary (d, k) tensor, and Q is the Cayley transform of the skew matrix it names, applied to the first k columns of the identity:

A = W Eᵀ − E Wᵀ          skew by construction
Q = (I − A)⁻¹ (I + A) E  orthonormal by construction

(I − A)⁻¹(I + A) is orthogonal whenever A is skew, so QᵀQ = I for every value of the parameter. That is the reason to parametrize rather than constrain: the optimizer is a plain AdamW over an unconstrained tensor, with no retraction step, no penalty term and no projection after the update, and the rotation cannot drift off the manifold however long the fit runs.

the start matters, and only at a real width

An unfitted subspace is a reproducible random rank-k basis drawn from seed — not the first k coordinate axes, which is what a zero parameter would give. The draw is scaled by 1/√d, and that scale is load-bearing: unscaled, the skew matrix's singular values grow like √d and the Cayley transform saturates, its derivative falling off like 1/σ², so at d = 2048 the basis stops responding to its parameter.

FINDINGS.md §19 records the measurement: at unit variance, loss 3.23 → 3.29 over 60 updates and held-out IIA 0.00 at every k; scaled, 3.22 → 0.11 and IIA 1.00. At d = 16, where the tiny fixture lives, the two are indistinguishable. A whole class of defect is invisible until a real width runs.

d is not in the document. It is the width of the site, which the engine reports for the address the site resolves to — 2048 for block_output on this model, 16 on the tiny fixture. The document says k; the model says d.

What the fit block says #

field	in this document	what it is
`rows`	`weekdays_full/data#train`	the pairs the fit sees, per role
`params`	`["rot"]`	the optimizer's parameter list, and the only trainability declaration in the format
`objective`	`[[1.0, "ce"]]`	`Σ wᵢ · termᵢ`, minimized; the sign of the weight is the direction
`epochs`	20	passes over the training rows
`pairs`	10	rows per minibatch — 30 training pairs give 3 updates an epoch
`seed`	0	the draw the parameter starts from, and the minibatch order
`optimizer`	`adamw`, lr 0.003	weight decay 0, over the one parameter
`early_stop`	`iia`, max, patience 20	stop after 20 epochs without a better eval; at 20 epochs, never
`eval`	`…#test`	the pass at the end of each epoch, and what it saves

A fit is not a second engine. An update is an ordinary observation pass over one minibatch, and the only things between updates are an optimizer step and, at the end of an epoch, an eval pass with the gradients off and on rows the fit never saw. The model is frozen: params is the whole of it.

Early stopping watches the mean of one eval metric. mode says which direction counts as better, patience how many epochs without an improvement to tolerate. Patience 20 with 20 epochs means the fit runs to the end and the field is there to be turned down.

What `iia` means here #

iia is a name this document chose; the kind is what it computes. Here it is match against the cf_answer column: 1.0 when the patched model's top token at the last position is the counterfactual's answer, 0.0 otherwise. Averaged over rows, that is interchange intervention accuracy — the fraction of pairs on which swapping the subspace moved the answer all the way.

It is not the objective. The objective is ce, the cross entropy of the counterfactual answer, because an argmax has no gradient. So the fit minimizes a soft target and is scored on a hard one, and early stopping watches the hard one.

the same name in another document

In documents/v2/das.json the metric called iia is a logit_diff. Read the kind, not the name — and explain prints both, as iia/match or iia/logit_diff.

What it compiles to #

$ uv run causalab-mini explain documents/real/weekdays_das.json
<root>: Plan
  k=1: Plan
    featurizers: Featurizers
        rot: subspace k=1 d=2048 cayley seed=0 trained=True
    fit: Fit  20 epochs x 3 update  lr=0.003 objective=((1.0, 'ce'),) params=('rot',)  early_stop='iia' patience=20
        epochs[0][0]: Observe  metrics=['iia/match', 'logit_diff/logit_diff', 'ce/cross_entropy']
            forward 'original' on 'counterfactual'  (10 rows x 8 tokens)
              read  'v_cf' at block_output[8] pos={index:-4} via 'rot'
            forward 'patched' on 'base'  (10 rows x 8 tokens)
              write 'patch' at block_output[8] pos={index:-4} swap(v_cf) via 'rot'
              read  'logits' at lm_head pos={index:-1} via 'identity'
        evaluation: Observe  metrics=['iia/match', 'logit_diff/logit_diff', 'ce/cross_entropy']  saves=['held_out_iia.json']
            forward 'original' on 'counterfactual'  (12 rows x 8 tokens)
              read  'v_cf' at block_output[8] pos={index:-4} via 'rot'
            forward 'patched' on 'base'  (12 rows x 8 tokens)
              write 'patch' at block_output[8] pos={index:-4} swap(v_cf) via 'rot'
              read  'logits' at lm_head pos={index:-1} via 'identity'
    score: Observe  metrics=['iia/match', 'logit_diff/logit_diff', 'ce/cross_entropy']  saves=['iia.json', 'logit_diff.json']
        forward 'original' on 'counterfactual'  (12 rows x 8 tokens)
          read  'v_cf' at block_output[8] pos={index:-4} via 'rot'
        forward 'patched' on 'base'  (12 rows x 8 tokens)
          write 'patch' at block_output[8] pos={index:-4} swap(v_cf) via 'rot'
          read  'logits' at lm_head pos={index:-1} via 'identity'
    weights: Weights  names=['rot']  saves=['rot.safetensors']
  k=4: Plan
    featurizers: Featurizers
        rot: subspace k=4 d=2048 cayley seed=0 trained=True
…

Four points, each a complete experiment with its own parameter, its own optimizer and its own digest. d=2048 is printed because the compiler has already asked the model; the document never said it.

The fit, the score and the weights step of one point all run inside the same session as every other point — one model load, one traced context, 240 optimizer steps. That is what makes a swept fit a single document rather than a shell loop.

Running it #

$ uv run causalab-mini validate documents/real/weekdays_das.json
ok: documents/real/weekdays_das.json is a valid plan-shaped document

$ time uv run causalab-mini run documents/real/weekdays_das.json \
      --device-map cpu --out out/das_real
out/das_real/document.json
out/das_real/run.json
out/das_real/k=1/document.json
out/das_real/k=1/held_out_iia.json
out/das_real/k=1/iia.json
out/das_real/k=1/logit_diff.json
out/das_real/k=1/rot.safetensors
…
out/das_real/k=64/rot.safetensors

real	10m0.438s

Four fits, 240 updates and 80 eval passes on a 1B-parameter model in fp32, on twelve CPU cores, in ten minutes. The recorded A100 run of the same document took 61 seconds, and 36 seconds through a self-hosted NDIF (documents/real/README.md).

two tables, two names

A fit's eval pass is a step rather than a plan, so it writes into its enclosing plan's directory rather than one of its own. It and the score step both save a table of iia, so they need two file names: held_out_iia.json and iia.json. Two saves with one name in one plan do not collide loudly — the later one is what is on disk.

The rank curve #

The mean of score's table over the 12 held-out pairs, per point:

k	held-out IIA	rows moved	held-out logit diff
1	0.00	0 / 12	−0.81
4	0.83	10 / 12	+1.52
16	1.00	12 / 12	+3.40
64	1.00	12 / 12	+3.48

The IIA column is the same 0.00 / 0.83 / 1.00 / 1.00 that documents/real/README.md records from the A100 run — a fit, on a different device, in a different month, landing on the same four numbers from the same seed.

The logit-difference column is what the accuracy column hides. At k = 16 the answer has not merely flipped; the counterfactual's day beats the base's by 3.4 logits. At k = 64 it is 3.48 — four times the rank buys 0.08. Whatever "today" is, sixteen directions hold it.

The bundle #

steps.weights publishes the parameter and the save writes it as a safetensors bundle with one slot, weight, and the featurizer's identity in the header:

$ python -c "…"   # safe_open('out/das_real/k=16/rot.safetensors').metadata()
{
 "engine": "causalab-mini",
 "site": "target",
 "kind": "subspace",
 "produced_by": "67ce7ff4200b70dfe8e3947c68a01c17f84240be49e699548fea293930203231",
 "layer": "8",
 "featurizer_dtype": "fp32",
 "k": "16",
 "model_revision": "4e20de362430cd3b72f300e6b0f18e50e7166e08",
 "component": "block_output",
 "model_key": "meta-llama/Llama-3.2-1B",
 "d": "2048",
 "model_dtype": "fp32",
 "parametrization": "cayley"
}
weight (2048, 16)

A rotation is a grid of numbers and nothing in the numbers says which model or which layer it came from. The header does, and a document that loads one with file_path is checked against it key by key before anything runs. A mismatch is refused naming the key:

$ uv run causalab-mini explain das_apply_at_layer_1.json
…
causalab_mini.plan.plan.PlanError: featurizer 'rot': 'documents/artifacts/rot.safetensors' is not this featurizer — layer: bundle says '0', document says '1'

A rotation fitted at layer 0 and loaded at layer 1 would run, and would be nonsense. documents/v2/das_apply.json is the honest form of that load: the same experiment with file_path on the featurizer, no fit step and no weights step — a subspace scored, not trained.

What the rank curve does not say #

read against a control

A high held-out IIA at k = 64 is less impressive than the same number at k = 4, and at large enough k it stops being evidence at all: swap enough of a 2048-wide activation and you have swapped the activation. The comparison that makes a rank meaningful is an untrained subspace of the same rank — documents/random_subspace_cpu.json is that control — and a principal-component basis of the same rank, which documents/v2/pca_harvest.json and pca_control.json are.

And a subspace that carries the answer is still a place, not a mechanism. The fit shows that swapping those 16 directions at layer 8 of the day token suffices to change the answer. It does not show that the model computes the answer there, nor that the 16 directions mean anything individually — a rotation is fitted to a task, and another task would fit another rotation.

twelve held-out pairs

The split is 30 train, 12 test out of 42 ordered pairs of weekdays. One row is 8 percentage points of the IIA column, which is why k = 4 reads 0.83 — ten rows moved and two did not. A curve this shape on twelve rows is worth believing about the gap between 1 and 16; it is not worth reading between 16 and 64.