SAE features

An sae featurizer is an encoder/decoder pair loaded from a safetensors bundle. It is the one featurizer in mini with an error term: what the dictionary does not explain is carried around the write and put back, so a write changes the features and nothing else.

Nothing here trains one. sae, linear and pca are loaded-only kinds: a document that declares one without a file_path is refused, and a fit may not name one. The dictionary comes from somewhere else — sae_lens, dictionary_learning, your own training run — and mini's job is to put it in the forward pass and check that it belongs there.

At a glance #

$ uv run causalab-mini explain documents/v2/sae_feature_ablation.json
<root>: Plan
  features=0: Plan
    featurizers: Featurizers
        dictionary: sae k=32 d=16 loaded from documents/artifacts/sae.safetensors trained=False
    ablate: Observe  metrics=['logit_diff/logit_diff', 'p_answer/token_prob']  outputs=['active=latents']  saves=['zeroed.json', 'p_answer.json']
        forward 'original' on 'base'  (4 rows x 11 tokens)
          read  'latents' at block_output[0] pos={index:-1} via 'dictionary'
        forward 'zeroed' on 'base'  (4 rows x 11 tokens)
          write 'ablate' at block_output[0] pos={index:-1} swap(0.0) via 'dictionary' on its features [0]
          read  'logits' at lm_head pos={index:-1} via 'identity'
  …
  features=3: Plan
    featurizers: Featurizers
        dictionary: sae k=32 d=16 loaded from documents/artifacts/sae.safetensors trained=False
    ablate: Observe  metrics=['logit_diff/logit_diff', 'p_answer/token_prob']  outputs=['active=latents']  saves=['zeroed.json', 'p_answer.json']
        forward 'original' on 'base'  (4 rows x 11 tokens)
          read  'latents' at block_output[0] pos={index:-1} via 'dictionary'
        forward 'zeroed' on 'base'  (4 rows x 11 tokens)
          write 'ablate' at block_output[0] pos={index:-1} swap(0.0) via 'dictionary' on its features [3]
          read  'logits' at lm_head pos={index:-1} via 'identity'

Four points, one per ablated latent. Two things on those lines are what this page is about: via 'dictionary', which says the read and the write happen in the dictionary's space rather than the residual stream's, and on its features [0], which says the write touches one coordinate of that space.

The document #

{
  "header": {
    "description": "SAE feature ablation. The residual stream at layer 0 is read through a loaded sparse autoencoder \u2014 an encoder/decoder pair with an ERROR TERM, so what the dictionary does not explain is put back exactly \u2014 and one latent at a time is set to zero (`features: [j]`, swept). An ablated latent that was already zero on these prompts changes nothing, which is the point of doing this in feature space: the sweep's own points are the comparison. The bundle here is a random 16->32 dictionary, a fixture: the mechanics are real, the latents mean nothing."
  },
  "model": {
    "key": "hf-internal-testing/tiny-random-LlamaForCausalLM",
    "revision": "9fb191250dd56d0ba7ec9785a025ed29c03d5998",
    "dtype": "fp32"
  },
  "roles": {
    "base": {
      "field": "input"
    }
  },
  "sites": {
    "target": {
      "component": "block_output",
      "layers": [
        0
      ]
    },
    "head": {
      "component": "lm_head"
    }
  },
  "featurizers": {
    "dictionary": {
      "kind": "sae",
      "file_path": "documents/artifacts/sae.safetensors"
    }
  },
  "interventions": {
    "ablate": {
      "reads": {
        "logits": {
          "site": "head",
          "pos": -1,
          "model": "zeroed",
          "input": "base"
        },
        "latents": {
          "site": "target",
          "pos": -1,
          "model": "original",
          "input": "base",
          "featurizer": "dictionary"
        }
      },
      "writes": {
        "ablate": {
          "site": "target",
          "pos": -1,
          "mechanism": "swap",
          "operand": 0.0,
          "featurizer": "dictionary",
          "features": {
            "sweep": [
              [
                0
              ],
              [
                1
              ],
              [
                2
              ],
              [
                3
              ]
            ]
          }
        }
      },
      "models": {
        "zeroed": {
          "input": "base",
          "writes": [
            "ablate"
          ]
        }
      },
      "metrics": {
        "logit_diff": {
          "kind": "logit_diff",
          "of": "logits",
          "a": "cf_answer",
          "b": "base_answer"
        },
        "p_answer": {
          "kind": "token_prob",
          "of": "logits",
          "token": "cf_answer"
        }
      }
    }
  },
  "steps": {
    "ablate": {
      "kind": "observe",
      "interventions": "ablate",
      "rows": {
        "base": "weekdays/train"
      },
      "saves": [
        {
          "value": "logit_diff",
          "file_path": "zeroed.json"
        },
        {
          "value": "p_answer",
          "file_path": "p_answer.json"
        }
      ],
      "outputs": {
        "active": {
          "read": "latents"
        }
      }
    }
  }
}

The parts that are not in an ordinary patching document:

• featurizers.dictionary declares the parameter set: kind: "sae" and the path to its bundle. No k — for an encoder the rank comes from the bundle's own shapes.
• the read latents carries "featurizer": "dictionary", so what it returns is the feature vector, not the activation.
• the write ablate carries the same featurizer and "features": {"sweep": [[0], [1], [2], [3]]} — which coordinates of that space the mechanism acts on. The mechanism is swap with operand 0.0: set those latents to zero.
• steps.ablate.outputs.active publishes the read for any later step in the same plan. This document has one step, so nothing consumes it.

What the bundle has to be #

The tensor names and orientations are sae_lens's. A missing W_dec means tied weights, W_enc transposed.

d is not read from the file: it is the width of the site the featurizer acts at, worked out on the client from the model, and the bundle's W_enc must match it. k is W_enc.shape[-1], and it may exceed d — a dictionary is overcomplete. The one in documents/artifacts/ is 16 → 32 for a 16-wide tiny Llama.

The compiler reads the file, checks it, and re-serializes the tensors into the plan as bytes — safetensors.torch.save(tensors). That is the only tensor-shaped thing a fresh plan carries, and it is an input of the experiment in the same way the token ids are: it pickles, it compares, it ships to a server, and at four bytes a number it is a format a real dictionary of tens of millions of parameters fits in. The run turns the bytes back into tensors once, when the featurizers for a step are built.

The header is checked against the document #

A dictionary is a grid of numbers and nothing in the numbers says which model or which layer it came from. A bundle mini wrote carries that in its safetensors header — model key, revision, dtype, component, layer, kind, k, d — and a load is refused by the key that disagrees:

$ uv run causalab-mini explain out/sae_stamped_mismatch.json
Traceback (most recent call last):
  …
causalab_mini.plan.plan.PlanError: featurizer 'dictionary': 'out/sae_layer1.safetensors' is not this featurizer — layer: bundle says '1', document says '0'

Reading feature activations #

A read with a featurizer returns featurize(x), the feature vector. The document publishes it as active but does not save it. With one save added — {"value": "active", "file_path": "active.safetensors"} — it lands as a [4, 1, 32] tensor: four rows, one position, thirty-two latents.

latent    0       1       2       3       4      …   26      …   31
row 0     0.0512  0.0000  0.0432  0.0587  0.0000     0.4411     0.1295

Sixteen of the thirty-two are exactly zero on all four prompts, latent 1 among them — a relu encoder, and half its latents never fire here. The four the document sweeps are latents 0 to 3, so one of the four ablations is of a latent that was not on.

Writing a feature #

A write with features is the mechanism applied to those coordinates of the feature vector and to nothing else. The vector is cloned, the named columns are replaced by mechanism(f[..., index], operand), and the rest pass through unchanged. An operand in the same feature space is cut the same way; a bare number — 0.0 here — is used as it is.

So swap with operand 0.0 on features: [j] is: set latent j to zero, keep the other thirty-one, decode, and add the error term back.

The error term #

featurize(x)          = (f, x - decode(f))     f = act((x - b_dec) @ W_enc + b_enc)
inverse(f', err, x)   = decode(f') + err

The second element of featurize is the reconstruction residual, and inverse adds it back. What that buys is that the write's effect is the change to f and only that. Without it every SAE intervention would also be "replace the activation with its reconstruction", and the two effects could not be told apart.

How much that matters depends on how good the dictionary is, and this one is a random fixture. Measured on the four activations this document actually reads, ‖x − decode(f)‖ / ‖x‖ is 7.12: the "reconstruction" is seven times the size of the thing it reconstructs. And yet:

>>> f, err = enc.featurize(x)          # x: the four activations, ‖x‖ = 0.1505
>>> float(err.norm() / x.norm())
7.123898029327393
>>> float((enc.inverse(f, err, x) - x).abs().max())
1.4901161193847656e-08

With f' = f the write is the identity to 1.5e-08 on an activation whose own norm is 0.15 — fp32 rounding in decode(f) + err, not the dictionary's error. A worthless dictionary and an exact no-op, which is the property the error term exists for.

What the run produced #

$ uv run causalab-mini run documents/v2/sae_feature_ablation.json --device-map cpu --out out/sae
out/b-sae/document.json
out/b-sae/run.json
out/b-sae/features=0/document.json
out/b-sae/features=0/zeroed.json
out/b-sae/features=0/p_answer.json
out/b-sae/features=1/document.json
out/b-sae/features=1/zeroed.json
out/b-sae/features=1/p_answer.json
out/b-sae/features=2/document.json
out/b-sae/features=2/zeroed.json
out/b-sae/features=2/p_answer.json
out/b-sae/features=3/document.json
out/b-sae/features=3/zeroed.json
out/b-sae/features=3/p_answer.json

One directory per swept feature, named after the field the wrapper sits on. The logit_diff column of each, and — from the same document with the writes removed — the same metric with nothing written:

Latent 1 was already zero on every row, and zeroing it changes nothing — the four numbers agree with the un-written run to six decimals. Measured exactly, the largest difference is 6.0e-08, which is the same fp32 rounding as the round-trip above: the write still runs decode(f) + err where the un-written model runs neither.

That is what doing this in feature space buys. In the residual stream there is no such thing as zeroing a coordinate that was already off; in the dictionary's space, an ablation of a latent that did not fire is a no-op you can verify.

The linear kind #

linear is the same Encoder with no nonlinearity, so f may be negative and the error term is still there. W_dec, b_enc and b_dec are all optional; with none of them it is x @ W_enc with tied decoding. Use it for a dictionary that is not sparse — a fixed probe basis, a transcoder's input map.

It is not the same as pca. A pca featurizer is an orthonormal (d, k) basis, its error term is always zero, and its inverse keeps the orthogonal complement of the pre-write activation. An encoder's error term is the reconstruction residual, which is a different quantity and is not zero.

What is refused #

All six are refused before any weights are loaded — explain reports them, and so does the compile at the top of a run.