DBM fit

A DAS fit asks which directions of an activation carry the answer. A DBM fit asks which units — which coordinates of the residual stream, in the basis the model already uses. The parameter is one number per unit, and what a run keeps is a mask: this unit comes from the counterfactual, that one stays.

That makes it the same sentence as a subspace swap with "units" for "directions", and the same document with three fields changed. This page runs documents/v2/dbm.json, which sweeps the sparsity pressure, so the document is a sparsity curve.

At a glance #

The gate featurizer, in three lines:

featurize(x) = (x, 0)
inverse(f, 0, x) = m ⊙ f + (1 − m) ⊙ x

m = σ(θ / T)   while the fit is updating θ      (soft, so a gradient reaches it)
m = [θ > 0]    whenever the mask is scored      (hard, so the score is of a real mask)

The features are the activation — nothing is projected — which is why a gate needs no k. The parameter is θ, one number per unit of the site, and the two readings of it are the one piece of mode in the library.

The document #

{
  "header": {
    "description": "DBM on the tiny CPU Llama: learn WHICH units of layer 0's output carry the answer, as a binary mask, rather than which directions (compare das.json \u2014 the only differences are the featurizer's kind, one objective term and the anneal). The mask is soft, \u03c3(\u03b8/T), while the fit updates it and hard, \u03b8 > 0, whenever it is scored; `mask.mask` in the objective is the mean of the soft mask, the L1 pressure toward keeping few units. The L1 weight is swept, so the document is the sparsity curve: measured 2026-09-21, 4, 2, 1 and 0 of 16 units kept."
  },
  "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"
    }
  },
  "featurizers": {
    "mask": {
      "kind": "gate"
    }
  },
  "interventions": {
    "dbm": {
      "reads": {
        "v_cf": {
          "site": "target",
          "pos": -1,
          "model": "original",
          "input": "counterfactual",
          "featurizer": "mask"
        },
        "logits": {
          "site": "head",
          "pos": -1,
          "model": "patched",
          "input": "base"
        }
      },
      "writes": {
        "patch": {
          "site": "target",
          "pos": -1,
          "featurizer": "mask",
          "mechanism": "swap",
          "operand": "v_cf"
        }
      },
      "models": {
        "patched": {
          "input": "base",
          "writes": [
            "patch"
          ]
        }
      },
      "metrics": {
        "iia": {
          "kind": "logit_diff",
          "of": "logits",
          "a": "cf_answer",
          "b": "base_answer",
          "token_form": "space_prefixed"
        },
        "ce": {
          "kind": "cross_entropy",
          "of": "logits",
          "target": "cf_answer",
          "token_form": "space_prefixed"
        }
      }
    }
  },
  "steps": {
    "fit": {
      "kind": "fit",
      "interventions": "dbm",
      "rows": {
        "base": "weekdays/data#train",
        "counterfactual": "weekdays/data#train"
      },
      "params": [
        "mask"
      ],
      "objective": [
        [
          1.0,
          "ce"
        ],
        [
          {
            "sweep": [
              0.0001,
              0.002,
              0.01,
              0.05
            ]
          },
          "mask.mask"
        ]
      ],
      "epochs": 20,
      "pairs": 16,
      "seed": 0,
      "optimizer": {
        "name": "adamw",
        "lr": 0.05,
        "weight_decay": 0.0
      },
      "early_stop": {
        "metric": "ce",
        "mode": "min",
        "patience": 20
      },
      "eval": {
        "rows": {
          "base": "weekdays/data#test",
          "counterfactual": "weekdays/data#test"
        },
        "saves": [
          {
            "value": "iia",
            "file_path": "held_out_iia.json"
          }
        ]
      },
      "anneal": {
        "mask": {
          "start": 1.0,
          "end": 0.05
        }
      }
    },
    "score": {
      "kind": "observe",
      "interventions": "dbm",
      "rows": {
        "base": "weekdays/data#train",
        "counterfactual": "weekdays/data#train"
      },
      "saves": [
        {
          "value": "iia",
          "file_path": "iia.json"
        },
        {
          "value": "ce",
          "file_path": "ce.json"
        }
      ]
    },
    "weights": {
      "kind": "weights",
      "names": [
        "mask"
      ],
      "saves": [
        {
          "value": "mask",
          "file_path": "mask.safetensors"
        }
      ]
    }
  }
}

Reading it in the order the compiler does:

• model is the tiny random-weight Llama at layer 0, whose residual stream is 16 wide — so the mask is 16 numbers and you can print it.
• featurizers.mask is {"kind": "gate"} and nothing else. There is nothing else to say.
• interventions.dbm is an interchange: read v_cf through the gate on the counterfactual prompt, write it through the same gate into the base run. Two metrics — ce, the cross entropy of the counterfactual's answer, and iia, which in this document is a logit_diff, the margin between the two answers.
• steps.fit trains mask for 20 epochs, then steps.score scores the fitted mask on the training rows and steps.weights publishes θ for the save.

The three differences from a DAS fit #

The header says the only differences from documents/v2/das.json are the featurizer's kind, one objective term and the anneal. Renamings aside, that is what diff shows, plus the optimizer settings a different parameter wants:

The two new fields #

<gate>.mask in the objective #

An objective is a list of [weight, term] pairs, summed and minimized. The terms are usually metrics. A gate publishes one more under its own name: mask.mask is the soft mask, and the objective takes its mean — so a positive weight on it is L1 pressure toward keeping few units.

There is no maximize flag anywhere: the sign of the weight is the direction. [1.0, "ce"] minimizes a cross entropy; [-1.0, "logit_diff"] would maximize a margin. Here the two terms pull against each other — the cross entropy wants whatever units make the patch work, the mask term wants none of them — and the sweep over 0.0001, 0.002, 0.01, 0.05 is the trade-off, one point per weight.

anneal #

{"start": 1.0, "end": 0.05} sets T on the first update and on the last, and the updates between get a geometric interpolation. The point is that σ(θ/T) approaches the step function [θ > 0] as T falls, so the soft mask the optimizer differentiates ends up close to the hard mask the score uses. This document has 20 updates — 20 epochs of one minibatch, because the training split is smaller than pairs: 16:

The eval pass at the end of each epoch runs with training cleared, so every held-out number on this page is of a hard mask, whatever T was.

Running it #

$ uv run causalab-mini explain documents/v2/dbm.json
<root>: Plan
  objective=0.0001: Plan
    featurizers: Featurizers
        mask: gate d=16 θ=0 trained=True
    fit: Fit  20 epochs x 1 update  lr=0.05 objective=((1.0, 'ce'), (0.0001, 'mask.mask')) params=('mask',)  early_stop='ce' patience=20 anneal=(('mask', 1.0, 0.05),)
        epochs[0][0]: Observe  metrics=['iia/logit_diff', 'ce/cross_entropy']
            forward 'original' on 'counterfactual'  (2 rows x 11 tokens)
              read  'v_cf' at block_output[0] pos={index:-1} via 'mask'
            forward 'patched' on 'base'  (2 rows x 11 tokens)
              write 'patch' at block_output[0] pos={index:-1} swap(v_cf) via 'mask'
              read  'logits' at lm_head pos={index:-1} via 'identity'
        evaluation: Observe  metrics=['iia/logit_diff', 'ce/cross_entropy']  saves=['held_out_iia.json']
            forward 'original' on 'counterfactual'  (2 rows x 9 tokens)
              read  'v_cf' at block_output[0] pos={index:-1} via 'mask'
            forward 'patched' on 'base'  (2 rows x 9 tokens)
              write 'patch' at block_output[0] pos={index:-1} swap(v_cf) via 'mask'
              read  'logits' at lm_head pos={index:-1} via 'identity'
    score: Observe  metrics=['iia/logit_diff', 'ce/cross_entropy']  saves=['iia.json', 'ce.json']
        forward 'original' on 'counterfactual'  (2 rows x 11 tokens)
          read  'v_cf' at block_output[0] pos={index:-1} via 'mask'
          …
    weights: Weights  names=['mask']  saves=['mask.safetensors']
  objective=0.002: Plan
…
  objective=0.05: Plan
…

A fit is not a second engine. An update is an ordinary observation pass over one minibatch, and the only things that happen between updates are an optimizer step and — at the end of an epoch — an eval pass on rows the fit never saw. params is the whole of the trainability declaration: the model is frozen and nothing else in the run carries a gradient.

$ time uv run causalab-mini run documents/v2/dbm.json --out out/dbm
out/dbm/document.json
out/dbm/run.json
out/dbm/objective=0.0001/document.json
out/dbm/objective=0.0001/held_out_iia.json
out/dbm/objective=0.0001/iia.json
out/dbm/objective=0.0001/ce.json
out/dbm/objective=0.0001/mask.safetensors
…
out/dbm/objective=0.05/mask.safetensors

real	0m10.051s

Four fits, 80 updates and 80 eval passes, in ten seconds — because the model is 16 wide and the split is two rows.

Reading the mask #

A weights step publishes the fitted parameter as a result, and the save writes it as a safetensors bundle holding one slot, weight, with the featurizer's identity stamped in the header. For a gate that slot is θ, and the mask is θ > 0:

import glob
from safetensors import safe_open

for path in sorted(glob.glob("out/dbm/objective=*/mask.safetensors")):
    with safe_open(path, "pt") as bundle:
        theta = bundle.get_tensor("weight")
    print(path, int((theta > 0).sum()), "of", theta.numel(), "kept")

$ python mask.py
out/dbm/objective=0.0001/mask.safetensors 4 of 16 kept
out/dbm/objective=0.002/mask.safetensors 2 of 16 kept
out/dbm/objective=0.01/mask.safetensors 1 of 16 kept
out/dbm/objective=0.05/mask.safetensors 0 of 16 kept

Four, two, one, none — the same four numbers documents/v2/dbm.json's header records from 2026-09-21, from the same seed. That is the curve the sweep exists to draw: turn up the L1 weight and the fit keeps fewer units.

The bars are the fraction of the 16 units whose θ came out positive, measured from zero.

How to read the result — and this one #

A mask fit has two numbers, and both are needed. The size of the mask is what the sweep moves. The held-out score is whether a mask that small still carries the answer. A sparsity curve is only worth reading against the score beside it:

Both splits have two rows, so each cell is both of them — held_out_iia.json and ce.json, row by row.

The bundle #

Each point's mask.safetensors carries a header saying what the mask is of, and a later document that loads it is checked against that header key by key — model, revision, dtype, component, layer, kind, d. A mask fitted at one layer loaded at another would run, and would be nonsense, so it is refused naming the key that differs.

A gate nobody is fitting is always hard, including one loaded from a file. A loaded mask is a mask and nothing else: there is no temperature to set and no soft reading to fall into.