Upload folder using huggingface_hub

gemma-2-2b/layer-12/tree_sae/tree_sae_32k_k5_seed0/DECISIONS.md

@@ -0,0 +1,131 @@
+# Tree SAE (Gemma-2-2b, 32k) — interpretation & deviation ledger
+
+Module: `benchmark/methods/train/tree_sae.py`. Paper: Cao et al., arXiv:2605.07922
+(v2). Starting point: Kirill's synthetic-scale clean-room port
+(`synth_sae_bench_baseline/experiments/synthetic/synthsaebench/tree_sae/`, ledger
+`PAPER_DERIVED_DECISIONS.md`, cited below as "ledger #N"). Markers: **PE** paper-exact,
+**PI** paper-inferred, **CR** clean-room choice, **KIRILL** = unchanged from his ledger,
+**CHANGED** = differs from his implementation.
+
+## Architecture / forward pass
+- One shared encoder/decoder + one bias; features partitioned layer-major into 4 privilege
+  layers (§4 "Tree Structure"; ledger #1). **PE / KIRILL**
+- Widths `[2048, 4096, 12288, 14336]` = paper's `[1536:3072:9216:10752]` (App. C) = exact
+  1:2:6:7 ratio scaled to 32768 (the twin's width). **PE ratio / PI scale**
+- Per-layer TopK `[105, 21, 13, 11]` (sum 150 = twin k): paper's 4-layer L0=80 row
+  `[56, 11, 7, 6]` (Table 1) × 150/80 = `[105, 20.6, 13.1, 11.3]`, rounded. The paper's own
+  rule (round the per-prefix L0 of a matched Matryoshka) was not run. **PI**
+- **TopK before gate** (Eq. 6: `f*` "is subsequently utilized as the final activation"); the
+  realized L0 can fall below 150 (ledger #4). **PI / KIRILL**. Override `gate_before_topk:
+  true` runs the per-layer TopK on gate-allowed candidates instead (realized L0 == k when
+  enough allowed features are positive). Evidence for that reading: Table 1 lists the Tree
+  SAE per-layer L0s as exact integers, and the 4M-token smoke run under the default reading
+  had layers 2-3 almost silent (L0 0.012 / 0.001 of 13 / 11) because children of randomly
+  assigned sparse parents rarely survive the gate. Not "clearly" stated, so the default is
+  unchanged. **OPEN**
+- Recursive gate `f*_i = f_i · 1[f*_parent > 0]`, root-attached features ungated (Eq. 6);
+  layers gated in ascending order so the recursion is exact. **PE / KIRILL**
+- Bias added exactly once per cumulative reconstruction `xhat_{<=l} = b + Σ_{t<=l} xhat_t`
+  (ledger #6). **PI / KIRILL**
+- Reconstruction loss `Σ_l ||xhat_{<=l} − x||²`, batch mean (Eq. 7; ledger #7). **PE / KIRILL**
+
+## Auxiliary loss
+- Per-layer TopK-style aux loss (Eq. 8), aux-k 256 (App. C). **PE**
+- `α_1..4 = 1/128` — the paper's stated Gemma-2-2b setting (App. C: "setting auxiliary
+  coefficient as α_{1,2,3,4} = 1/128 ... can be beneficial on ... Gemma-2-2b"). Kirill used the
+  GPT-2 recipe `[1/32, 0, 0, 0]`. **PE (Gemma) / CHANGED**
+- Aux features = top-256 DEAD features of the layer by **ungated** ReLU pre-activation
+  (ledger #9). **CR / KIRILL**
+- Aux target is the **detached** residual `(x − xhat_{<=l}).detach()`: the paper says it uses
+  "the auxiliary loss proposed by Gao et al. (2024)", whose reference implementations (and
+  SAELens/Matryoshka) reconstruct a detached residual so the aux term does not push the
+  primary weights. Kirill did not detach. **PI / CHANGED**
+- Nothing is dead before 10M tokens have been seen (ledger #16), so the aux loss is inactive
+  for the first ~2441 steps of a 4096-token-batch run. **PI / KIRILL**
+
+## Dynamic allocation
+- Capacity `C_p` += per-example primary cumulative loss for every example on which `p` fires
+  (final gated act > 0); aux excluded; running sum from the start, never reset (ledger #10,
+  #11). **PI / KIRILL**
+- Accumulator dtype: batch increment `loss @ fires` in fp32 on GPU, added into a **float64**
+  accumulator. The task spec said fp32; a 500M-token fp32 running sum of O(1e5) per-token
+  losses would exceed fp32's 24-bit mantissa and silently drop increments. **CR / CHANGED
+  (Kirill: fp64 on CPU)**
+- Firing statistics (`fire_count`, batch-granularity `last_fire` token stamps) are int64 on
+  GPU; all per-step bookkeeping is vectorised (`fires.sum(0)`, one mat-vec, one `where`).
+- Parent eligibility: firing rate ≥ 1 per 50,000 tokens, running rate from the start of
+  training (App. G; ledger #12). **PE rule / PI window**
+- Dead = no positive final activation in the most recent 10M tokens (App. G; ledger #16).
+  **PE**
+- Greedy Algorithm 1 (heap, max-min payoff), tie-break smaller feature id (ledger #14).
+  Candidate parents for layer l = ALL features of earlier layers (`m_l`). **PE / CR tie**
+- Only DEAD children are reassigned; live edges never touched; first-fit in ascending parent
+  id into eligible parents whose LIVE child count is below `k*`; unplaceable dead children
+  keep their parent (App. G "first fit", "less than the optimal number of non-dead child
+  features"; ledger #15). Implemented as a vectorised `repeat_interleave` slot list — same
+  result as Kirill's pointer loop. **PE / KIRILL**
+- **Root child quota (Gemma option, App. G)**: "Optionally, we can enforce a prerequisited
+  number of child features for the root node at l=0 ... helps reducing number of dead
+  features in Gemma-2-2b". The paper gives no value. Implementation: at every reallocation
+  event, layer l reserves `q_l = floor(root_child_quota_frac · s_l)` ROOT slots (default
+  1/16 → 256/768/896); the greedy distributes the remaining `s_l − q_l` among eligible
+  learned parents; ROOT deficit slots (`q_l − live root-attached children`) are filled
+  FIRST by dead children, then learned parents. Not applied at initialisation. `0` restores
+  Kirill's no-quota behaviour (ledger #13). **CR value / CHANGED (Kirill: none)**
+- Reallocation schedule: first event after 3000 steps, interval ×2 after each event, capped
+  at 10,000 (App. C "increase the interval by two"; ledger #17) → events at steps 3000, 9000,
+  19000, 29000, … (13 events in 122,070 steps). The schedule state (`realloc_interval`,
+  `next_realloc_step`) is checkpointed. **PI / KIRILL**
+- Dead-to-root event once the 1-based step count reaches `0.5 · total_steps` (step 61,035;
+  App. C "at half of the training"; ledger #18). **PI**
+- Per-step order: optimizer update → statistics update → reallocation (if due) →
+  dead-to-root (if due), as in Kirill's loop.
+- Initial allocation: every non-first-layer feature gets a uniformly random parent among all
+  earlier-layer features (dedicated CPU `torch.Generator`, seed `cfg.seed`); no eligibility
+  filter at init; first layer → ROOT (ledger #19). **PI / KIRILL**
+
+## Optimisation / numerics
+- Adam, lr 1e-4, betas (0.9, 0.999), eps 1e-8, constant LR, no warmup/decay/clipping
+  (App. C; ledger #22). **PE lr / CR defaults**
+- Decoder rows unit-norm: per-row tangent-space gradient projection before the update,
+  renormalisation after; the whole-model gradient is NOT normalised (App. C "normalise the
+  gradient as well as the decoder vector"; ledger #20). **PI / KIRILL**
+- Init: kaiming-uniform decoder rows normalised to unit norm, encoder tied (`W_dec.T`),
+  zero bias (ledger #21). **CR / KIRILL**
+- Batch 4096 tokens and 500M tokens are fixed by the driver (twin recipe); the paper used
+  batch 5120. **CR (driver)**
+- Forward/backward under `torch.autocast("cuda", bfloat16)`, fp32 master weights and Adam
+  state; losses computed in fp32; per-layer TopK selection runs on bf16 pre-activations.
+  Kirill trained in fp32. Exported-model `encode` runs in fp32 without autocast. **CR / CHANGED**
+
+## Metrics, export, hierarchy
+- `step` returns: `loss`, `recon_loss` (Eq. 7), `aux_loss`, `mse` (per-element MSE of the final
+  cumulative reconstruction), `l0` (mean gated-active features per token, all layers),
+  `l0_layer{0..3}`, `n_dead`/`frac_dead` (dead mask at the start of the step),
+  `n_root_attached` (non-first-layer features whose parent is ROOT), `n_realloc_events`,
+  `dead_to_root_done`.
+- Export: `weights.safetensors` (`W_enc` (d_in, n), `W_dec` (n, d_in) unit rows, `b`),
+  `allocation.json` (parent per feature, ROOT = −1; layer per feature; dead mask),
+  `config.json` (config, step, tokens, event log). Flattened latent index = feature index.
+- `ExportedTreeSAE.encode(x)` = final gated activations of raw residuals;
+  `native_ancestors()` = `{child: {parent}}` from the final allocation, root-attached
+  features contribute no edge, no graded score (ledger #25); `latent_level()` = privilege
+  layer. **PE**
+- Checkpoint (`state_dict`): weights, Adam state, allocation vector, capacities, fire counts,
+  last-fire stamps, tokens seen, schedule state, dead-to-root flag, event log, generator
+  state, step → exact resume (test `test_checkpoint_roundtrip_reproduces_next_step`).
+
+## Open uncertainties (author questions in Kirill's `AUTHOR_QUESTIONS.md` still apply)
+- Root quota value and fill order (CR here); TopK-before-gate; capacity semantics/reset;
+  aux residual detach; whether "increase the interval by two" is ×2; Adam schedule.
+
+## Gemma 500M run: gate BEFORE TopK (decided 2026-09-14)
+
+The 500M-token run uses `gate_before_topk: true`. Evidence: the paper's
+Table 1 reports Tree SAE per-layer L0s as exact integers equal to the
+configured per-layer K ([21, 5, 4, 2], [26, 6], ...), which is only possible
+if the recursive gate is applied first and TopK then selects K survivors per
+layer; with TopK first the realised L0 is below K and non-integer (our
+smoke run: [105, 2.3, 2.7, 2.2] vs [105, 21, 13, 11] with gate-first).
+Kirill's frozen synthetic reading (TopK first) is kept available via the
+override for comparison.

gemma-2-2b/layer-12/tree_sae/tree_sae_32k_k5_seed0/README.md

@@ -0,0 +1,5 @@
+# tree_sae_32k_k5_seed0
+
+Architecture `tree_sae` trained on google/gemma-2-2b layer 12 (`model.layers.12`, resid_post) with the twin-SAE recipe: 499,998,720 tokens of chanind/pile-uncopyrighted-gemma-1024-abbrv-2B, batch 4096 tokens, context 1024.
+
+Load with `model.py::load_exported(export_dir)` from this folder (see `interface.py`); `encode(x)` takes raw residual activations and returns the flattened latents; `native_ancestors()` gives the architecture's own parent claims. Interpretation decisions vs the paper are in `DECISIONS.md`; `config.json`/`log.jsonl` are the training config and metrics.

gemma-2-2b/layer-12/tree_sae/tree_sae_32k_k5_seed0/config.json

@@ -0,0 +1,68 @@
+{
+ "arch": "tree_sae",
+ "run_name": "tree_sae_32k_k5_seed0",
+ "n_tokens": 500000000,
+ "batch_tokens": 4096,
+ "total_steps": 122070,
+ "overrides": {
+  "gate_before_topk": true,
+  "layer_ks": [
+   2,
+   1,
+   1,
+   1
+  ]
+ },
+ "seed": 0,
+ "init_seed": 0,
+ "group": [
+  "hsae_32k_k5_seed0",
+  "tree_sae_32k_k5_seed0",
+  "hsae_forest_32k_k5_seed0"
+ ],
+ "trainer_config": {
+  "d_in": 2304,
+  "layer_widths": [
+   2048,
+   4096,
+   12288,
+   14336
+  ],
+  "layer_ks": [
+   2,
+   1,
+   1,
+   1
+  ],
+  "aux_k": 256,
+  "aux_alphas": [
+   0.0078125,
+   0.0078125,
+   0.0078125,
+   0.0078125
+  ],
+  "root_child_quota_frac": 0.0625,
+  "lr": 0.0001,
+  "adam_betas": [
+   0.9,
+   0.999
+  ],
+  "adam_eps": 1e-08,
+  "dead_token_window": 10000000,
+  "parent_eligibility_rate": 2e-05,
+  "realloc_first_interval": 3000,
+  "realloc_growth": 2,
+  "realloc_cap": 10000,
+  "dead_to_root_frac": 0.5,
+  "seed": 0,
+  "autocast_bf16": true,
+  "gate_before_topk": true,
+  "n_latents": 32768,
+  "layer_starts": [
+   0,
+   2048,
+   6144,
+   18432
+  ]
+ }
+}

gemma-2-2b/layer-12/tree_sae/tree_sae_32k_k5_seed0/export/config.json

@@ -0,0 +1,502 @@
+{
+ "arch": "tree_sae",
+ "config": {
+  "d_in": 2304,
+  "layer_widths": [
+   2048,
+   4096,
+   12288,
+   14336
+  ],
+  "layer_ks": [
+   2,
+   1,
+   1,
+   1
+  ],
+  "aux_k": 256,
+  "aux_alphas": [
+   0.0078125,
+   0.0078125,
+   0.0078125,
+   0.0078125
+  ],
+  "root_child_quota_frac": 0.0625,
+  "lr": 0.0001,
+  "adam_betas": [
+   0.9,
+   0.999
+  ],
+  "adam_eps": 1e-08,
+  "dead_token_window": 10000000,
+  "parent_eligibility_rate": 2e-05,
+  "realloc_first_interval": 3000,
+  "realloc_growth": 2,
+  "realloc_cap": 10000,
+  "dead_to_root_frac": 0.5,
+  "seed": 0,
+  "autocast_bf16": true,
+  "gate_before_topk": true
+ },
+ "n_latents": 32768,
+ "layer_starts": [
+  0,
+  2048,
+  6144,
+  18432
+ ],
+ "step": 122070,
+ "tokens_seen": 499998720,
+ "n_dead": 2226,
+ "n_root_attached": 10671,
+ "events": [
+  {
+   "kind": "realloc",
+   "step": 3000,
+   "tokens_seen": 12288000,
+   "layers": [
+    {
+     "layer": 1,
+     "n_dead": 3584,
+     "n_moved": 3584,
+     "n_eligible_parents": 467,
+     "k_star_sum": 3840,
+     "root_quota": 256,
+     "n_to_root": 256
+    },
+    {
+     "layer": 2,
+     "n_dead": 11010,
+     "n_moved": 11010,
+     "n_eligible_parents": 987,
+     "k_star_sum": 11520,
+     "root_quota": 768,
+     "n_to_root": 768
+    },
+    {
+     "layer": 3,
+     "n_dead": 13015,
+     "n_moved": 13015,
+     "n_eligible_parents": 2006,
+     "k_star_sum": 13440,
+     "root_quota": 896,
+     "n_to_root": 896
+    }
+   ]
+  },
+  {
+   "kind": "realloc",
+   "step": 9000,
+   "tokens_seen": 36864000,
+   "layers": [
+    {
+     "layer": 1,
+     "n_dead": 1579,
+     "n_moved": 1579,
+     "n_eligible_parents": 467,
+     "k_star_sum": 3840,
+     "root_quota": 256,
+     "n_to_root": 1
+    },
+    {
+     "layer": 2,
+     "n_dead": 9682,
+     "n_moved": 9682,
+     "n_eligible_parents": 1780,
+     "k_star_sum": 11520,
+     "root_quota": 768,
+     "n_to_root": 240
+    },
+    {
+     "layer": 3,
+     "n_dead": 11879,
+     "n_moved": 11879,
+     "n_eligible_parents": 3691,
+     "k_star_sum": 13440,
+     "root_quota": 896,
+     "n_to_root": 189
+    }
+   ]
+  },
+  {
+   "kind": "realloc",
+   "step": 19000,
+   "tokens_seen": 77824000,
+   "layers": [
+    {
+     "layer": 1,
+     "n_dead": 35,
+     "n_moved": 35,
+     "n_eligible_parents": 683,
+     "k_star_sum": 3840,
+     "root_quota": 256,
+     "n_to_root": 0
+    },
+    {
+     "layer": 2,
+     "n_dead": 7938,
+     "n_moved": 7938,
+     "n_eligible_parents": 2864,
+     "k_star_sum": 11520,
+     "root_quota": 768,
+     "n_to_root": 61
+    },
+    {
+     "layer": 3,
+     "n_dead": 10447,
+     "n_moved": 10447,
+     "n_eligible_parents": 5786,
+     "k_star_sum": 13440,
+     "root_quota": 896,
+     "n_to_root": 39
+    }
+   ]
+  },
+  {
+   "kind": "realloc",
+   "step": 29000,
+   "tokens_seen": 118784000,
+   "layers": [
+    {
+     "layer": 1,
+     "n_dead": 13,
+     "n_moved": 13,
+     "n_eligible_parents": 907,
+     "k_star_sum": 3840,
+     "root_quota": 256,
+     "n_to_root": 0
+    },
+    {
+     "layer": 2,
+     "n_dead": 6640,
+     "n_moved": 6640,
+     "n_eligible_parents": 3698,
+     "k_star_sum": 11520,
+     "root_quota": 768,
+     "n_to_root": 15
+    },
+    {
+     "layer": 3,
+     "n_dead": 9206,
+     "n_moved": 9206,
+     "n_eligible_parents": 7401,
+     "k_star_sum": 13440,
+     "root_quota": 896,
+     "n_to_root": 8
+    }
+   ]
+  },
+  {
+   "kind": "realloc",
+   "step": 39000,
+   "tokens_seen": 159744000,
+   "layers": [
+    {
+     "layer": 1,
+     "n_dead": 3,
+     "n_moved": 3,
+     "n_eligible_parents": 1085,
+     "k_star_sum": 3840,
+     "root_quota": 256,
+     "n_to_root": 0
+    },
+    {
+     "layer": 2,
+     "n_dead": 5618,
+     "n_moved": 5618,
+     "n_eligible_parents": 4315,
+     "k_star_sum": 11520,
+     "root_quota": 768,
+     "n_to_root": 1
+    },
+    {
+     "layer": 3,
+     "n_dead": 8138,
+     "n_moved": 8138,
+     "n_eligible_parents": 8605,
+     "k_star_sum": 13440,
+     "root_quota": 896,
+     "n_to_root": 2
+    }
+   ]
+  },
+  {
+   "kind": "realloc",
+   "step": 49000,
+   "tokens_seen": 200704000,
+   "layers": [
+    {
+     "layer": 1,
+     "n_dead": 3,
+     "n_moved": 3,
+     "n_eligible_parents": 1260,
+     "k_star_sum": 3840,
+     "root_quota": 256,
+     "n_to_root": 0
+    },
+    {
+     "layer": 2,
+     "n_dead": 4782,
+     "n_moved": 4782,
+     "n_eligible_parents": 4736,
+     "k_star_sum": 11520,
+     "root_quota": 768,
+     "n_to_root": 0
+    },
+    {
+     "layer": 3,
+     "n_dead": 7137,
+     "n_moved": 7137,
+     "n_eligible_parents": 9502,
+     "k_star_sum": 13440,
+     "root_quota": 896,
+     "n_to_root": 1
+    }
+   ]
+  },
+  {
+   "kind": "realloc",
+   "step": 59000,
+   "tokens_seen": 241664000,
+   "layers": [
+    {
+     "layer": 1,
+     "n_dead": 0,
+     "n_moved": 0,
+     "n_eligible_parents": 1416,
+     "k_star_sum": 3840,
+     "root_quota": 256,
+     "n_to_root": 0
+    },
+    {
+     "layer": 2,
+     "n_dead": 4065,
+     "n_moved": 4065,
+     "n_eligible_parents": 5069,
+     "k_star_sum": 11520,
+     "root_quota": 768,
+     "n_to_root": 0
+    },
+    {
+     "layer": 3,
+     "n_dead": 6273,
+     "n_moved": 6273,
+     "n_eligible_parents": 10241,
+     "k_star_sum": 13440,
+     "root_quota": 896,
+     "n_to_root": 0
+    }
+   ]
+  },
+  {
+   "kind": "dead_to_root",
+   "step": 61035,
+   "tokens_seen": 249999360,
+   "n_moved": 9270
+  },
+  {
+   "kind": "realloc",
+   "step": 69000,
+   "tokens_seen": 282624000,
+   "layers": [
+    {
+     "layer": 1,
+     "n_dead": 0,
+     "n_moved": 0,
+     "n_eligible_parents": 1570,
+     "k_star_sum": 3840,
+     "root_quota": 256,
+     "n_to_root": 0
+    },
+    {
+     "layer": 2,
+     "n_dead": 253,
+     "n_moved": 253,
+     "n_eligible_parents": 5319,
+     "k_star_sum": 11520,
+     "root_quota": 768,
+     "n_to_root": 0
+    },
+    {
+     "layer": 3,
+     "n_dead": 1077,
+     "n_moved": 1077,
+     "n_eligible_parents": 11514,
+     "k_star_sum": 13440,
+     "root_quota": 896,
+     "n_to_root": 0
+    }
+   ]
+  },
+  {
+   "kind": "realloc",
+   "step": 79000,
+   "tokens_seen": 323584000,
+   "layers": [
+    {
+     "layer": 1,
+     "n_dead": 1,
+     "n_moved": 1,
+     "n_eligible_parents": 1710,
+     "k_star_sum": 3840,
+     "root_quota": 256,
+     "n_to_root": 0
+    },
+    {
+     "layer": 2,
+     "n_dead": 346,
+     "n_moved": 346,
+     "n_eligible_parents": 5486,
+     "k_star_sum": 11520,
+     "root_quota": 768,
+     "n_to_root": 0
+    },
+    {
+     "layer": 3,
+     "n_dead": 1275,
+     "n_moved": 1275,
+     "n_eligible_parents": 12782,
+     "k_star_sum": 13440,
+     "root_quota": 896,
+     "n_to_root": 0
+    }
+   ]
+  },
+  {
+   "kind": "realloc",
+   "step": 89000,
+   "tokens_seen": 364544000,
+   "layers": [
+    {
+     "layer": 1,
+     "n_dead": 0,
+     "n_moved": 0,
+     "n_eligible_parents": 1826,
+     "k_star_sum": 3840,
+     "root_quota": 256,
+     "n_to_root": 0
+    },
+    {
+     "layer": 2,
+     "n_dead": 805,
+     "n_moved": 805,
+     "n_eligible_parents": 5630,
+     "k_star_sum": 11520,
+     "root_quota": 768,
+     "n_to_root": 0
+    },
+    {
+     "layer": 3,
+     "n_dead": 1351,
+     "n_moved": 1351,
+     "n_eligible_parents": 13388,
+     "k_star_sum": 13440,
+     "root_quota": 896,
+     "n_to_root": 0
+    }
+   ]
+  },
+  {
+   "kind": "realloc",
+   "step": 99000,
+   "tokens_seen": 405504000,
+   "layers": [
+    {
+     "layer": 1,
+     "n_dead": 1,
+     "n_moved": 1,
+     "n_eligible_parents": 1912,
+     "k_star_sum": 3840,
+     "root_quota": 256,
+     "n_to_root": 0
+    },
+    {
+     "layer": 2,
+     "n_dead": 826,
+     "n_moved": 826,
+     "n_eligible_parents": 5731,
+     "k_star_sum": 11520,
+     "root_quota": 768,
+     "n_to_root": 0
+    },
+    {
+     "layer": 3,
+     "n_dead": 1432,
+     "n_moved": 1432,
+     "n_eligible_parents": 13768,
+     "k_star_sum": 13440,
+     "root_quota": 896,
+     "n_to_root": 0
+    }
+   ]
+  },
+  {
+   "kind": "realloc",
+   "step": 109000,
+   "tokens_seen": 446464000,
+   "layers": [
+    {
+     "layer": 1,
+     "n_dead": 1,
+     "n_moved": 1,
+     "n_eligible_parents": 1962,
+     "k_star_sum": 3840,
+     "root_quota": 256,
+     "n_to_root": 0
+    },
+    {
+     "layer": 2,
+     "n_dead": 873,
+     "n_moved": 873,
+     "n_eligible_parents": 5796,
+     "k_star_sum": 11520,
+     "root_quota": 768,
+     "n_to_root": 0
+    },
+    {
+     "layer": 3,
+     "n_dead": 1432,
+     "n_moved": 1432,
+     "n_eligible_parents": 14020,
+     "k_star_sum": 13440,
+     "root_quota": 896,
+     "n_to_root": 0
+    }
+   ]
+  },
+  {
+   "kind": "realloc",
+   "step": 119000,
+   "tokens_seen": 487424000,
+   "layers": [
+    {
+     "layer": 1,
+     "n_dead": 0,
+     "n_moved": 0,
+     "n_eligible_parents": 1991,
+     "k_star_sum": 3840,
+     "root_quota": 256,
+     "n_to_root": 0
+    },
+    {
+     "layer": 2,
+     "n_dead": 894,
+     "n_moved": 894,
+     "n_eligible_parents": 5825,
+     "k_star_sum": 11520,
+     "root_quota": 768,
+     "n_to_root": 0
+    },
+    {
+     "layer": 3,
+     "n_dead": 1413,
+     "n_moved": 1413,
+     "n_eligible_parents": 14136,
+     "k_star_sum": 13440,
+     "root_quota": 896,
+     "n_to_root": 0
+    }
+   ]
+  }
+ ]
+}

gemma-2-2b/layer-12/tree_sae/tree_sae_32k_k5_seed0/export/weights.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:be05f1bd5973d2169788d604f6d28839706fad227c78e2a01cd68828bc88caac
+size 603989232

gemma-2-2b/layer-12/tree_sae/tree_sae_32k_k5_seed0/interface.py

@@ -0,0 +1,59 @@
+"""Interface every trainable hierarchy SAE implements for ``driver.py``.
+
+The driver owns: the activation stream (twin recipe: 500M tokens of the
+twin's corpus, batches of 4096 tokens, Gemma-2-2b layer 12), step counting,
+checkpoint/resume on the Modal volume, metric logging, and the final export.
+Everything architecture-specific — losses, optimizer, LR schedule, decoder
+renormalisation, hierarchy updates, reallocation, dead-feature handling —
+lives inside ``step``.
+
+Exports are consumed by the benchmark through :class:`ExportedSAE`: a single
+concatenated latent index space (all levels / layers / experts flattened),
+``encode(x) -> (batch, n_latents)`` non-negative activations, and the
+architecture's own ancestor claims ``native_ancestors()`` as
+``{child_index: {ancestor_indices}}`` (direct parents; the benchmark applies
+transitive closure).
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Protocol
+
+import torch
+
+
+class HierarchyTrainer(Protocol):
+    d_in: int
+    n_latents: int  # size of the flattened latent index space
+
+    def step(self, x: torch.Tensor, step: int, total_steps: int) -> dict[str, float]:
+        """One optimizer step on a ``(batch, d_in)`` float32 CUDA batch.
+
+        Returns scalar metrics (must include ``"loss"`` and ``"l0"``, where
+        ``l0`` is the mean number of active latents per token across the
+        whole flattened index space).
+        """
+        ...
+
+    def state_dict(self) -> dict[str, Any]: ...
+
+    def load_state_dict(self, state: dict[str, Any]) -> None: ...
+
+    def export(self, out_dir: Path) -> None:
+        """Write everything :class:`ExportedSAE` needs to ``out_dir``."""
+        ...
+
+
+class ExportedSAE(Protocol):
+    n_latents: int
+    W_dec: torch.Tensor  # (n_latents, d_in) decoder directions (for cosine baselines)
+
+    @torch.no_grad()
+    def encode(self, x: torch.Tensor) -> torch.Tensor: ...
+
+    def native_ancestors(self) -> dict[int, set[int]]: ...
+
+    def latent_level(self) -> list[int]:
+        """Level / layer / expert-vs-child tag per flattened latent (for reporting)."""
+        ...

gemma-2-2b/layer-12/tree_sae/tree_sae_32k_k5_seed0/model.py

@@ -0,0 +1,549 @@
+"""Tree SAE trainer for Gemma-2-2b layer 12 (``driver.py --arch tree_sae``).
+
+Clean-room port of Cao et al., "Tree SAE: Learning Hierarchical Feature
+Structures in Sparse Autoencoders" (arXiv:2605.07922) to the benchmark's
+``HierarchyTrainer`` / ``ExportedSAE`` interface (``base.py``). Starting point
+was Kirill's synthetic-scale implementation
+(``synth_sae_bench_baseline/.../tree_sae/``); every interpretation of the
+paper and every deviation from that implementation is listed in
+``TREE_SAE_DECISIONS.md`` next to this file.
+
+Architecture (one shared dictionary, layer-major feature ids):
+
+- ``pre = (x - b) @ W_enc``; ``relu = ReLU(pre)``; independent TopK inside each
+  privilege layer's slice (``layer_ks``); recursive parent gate (paper Eq. 6)
+  applied AFTER TopK: ``f*_i = f_i * 1[f*_parent(i) > 0]`` (root-attached
+  features are ungated);
+- cumulative reconstructions ``xhat_{<=l} = b + sum_{t<=l} gated_t @ W_dec_t``
+  (bias added once); loss ``sum_l ||xhat_{<=l} - x||^2`` (Eq. 7, batch mean);
+- per-layer TopK-style auxiliary dead-feature loss (Eq. 8) on the DETACHED
+  residual, aux-k 256, alpha_l = 1/128 for all layers (paper's Gemma setting);
+- dynamic allocation (Alg. 1-2): capacity = running sum of the per-example
+  primary loss over examples the feature fires on; parent eligibility >= 1
+  firing per 50k tokens; dead = no final activation in 10M tokens;
+  reallocation events after 3000 steps then interval x2 capped at 10000;
+  optional root child quota (paper's Gemma option); all remaining dead
+  features moved to ROOT at half the token budget.
+
+Flattened latent index = feature index; ``ROOT = -1`` in the allocation
+vector means "no learned parent". Invariant: ``allocation[i] == ROOT`` or
+``allocation[i] < layer_start(layer_of(i))`` (parent strictly earlier layer).
+"""
+
+from __future__ import annotations
+
+import contextlib
+import heapq
+import json
+from dataclasses import asdict, dataclass, fields
+from pathlib import Path
+from typing import Any
+
+import torch
+
+ROOT = -1
+_TUPLE_FIELDS = {"layer_widths", "layer_ks", "aux_alphas", "adam_betas"}
+
+
+@dataclass(frozen=True)
+class TreeSAEConfig:
+    d_in: int = 2304
+    # Paper 4-layer ratio [1536:3072:9216:10752] scaled to 32768 (twin width).
+    layer_widths: tuple[int, ...] = (2048, 4096, 12288, 14336)
+    # Paper 4-layer L0=80 row [56, 11, 7, 6] scaled proportionally to 150.
+    layer_ks: tuple[int, ...] = (105, 21, 13, 11)
+    aux_k: int = 256
+    # Paper App. C: alpha_{1,2,3,4} = 1/128 "on other models such as Gemma-2-2b".
+    aux_alphas: tuple[float, ...] = (1.0 / 128.0,) * 4
+    # Paper App. G Gemma option: a required number of ROOT children per layer,
+    # here a fraction of the layer width at every reallocation event; 0 disables.
+    root_child_quota_frac: float = 1.0 / 16.0
+    lr: float = 1e-4
+    adam_betas: tuple[float, float] = (0.9, 0.999)
+    adam_eps: float = 1e-8
+    dead_token_window: int = 10_000_000
+    parent_eligibility_rate: float = 1.0 / 50_000.0
+    realloc_first_interval: int = 3000
+    realloc_growth: int = 2
+    realloc_cap: int = 10_000
+    dead_to_root_frac: float = 0.5  # fraction of total_steps for the dead-to-root event
+    seed: int = 0
+    autocast_bf16: bool = True  # bf16 autocast on CUDA (fp32 master weights / Adam state)
+    # False (default, ledger #4): per-layer TopK on raw ReLU pre-activations, THEN the parent
+    # gate (realized per-layer L0 <= k). True: gate first, TopK among gate-allowed features
+    # (realized L0 == k whenever enough allowed features are positive).
+    gate_before_topk: bool = False
+
+    def __post_init__(self) -> None:
+        if not (len(self.layer_widths) == len(self.layer_ks) == len(self.aux_alphas)):
+            raise ValueError("layer_widths, layer_ks and aux_alphas must have equal length")
+        if any(k > w or k < 1 for k, w in zip(self.layer_ks, self.layer_widths, strict=True)):
+            raise ValueError("per-layer k must be in [1, layer width]")
+        if not 0.0 <= self.root_child_quota_frac < 1.0:
+            raise ValueError("root_child_quota_frac must be in [0, 1)")
+
+    @property
+    def n_layers(self) -> int:
+        return len(self.layer_widths)
+
+    @property
+    def n_features(self) -> int:
+        return sum(self.layer_widths)
+
+    def layer_starts(self) -> tuple[int, ...]:
+        starts = [0]
+        for w in self.layer_widths[:-1]:
+            starts.append(starts[-1] + w)
+        return tuple(starts)
+
+    def layer_slices(self) -> tuple[slice, ...]:
+        return tuple(slice(s, s + w) for s, w in zip(self.layer_starts(), self.layer_widths, strict=True))
+
+    def root_quota(self, layer: int) -> int:
+        if layer == 0:
+            return self.layer_widths[0]
+        return int(self.root_child_quota_frac * self.layer_widths[layer])
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> TreeSAEConfig:
+        known = {f.name for f in fields(cls)}
+        unknown = set(d) - known
+        if unknown:
+            raise ValueError(f"unknown TreeSAEConfig keys: {sorted(unknown)}")
+        fixed = {k: (tuple(v) if k in _TUPLE_FIELDS else v) for k, v in d.items()}
+        return cls(**fixed)
+
+
+# ---------------------------------------------------------------------------
+# Functional core (shared by trainer and exported model)
+# ---------------------------------------------------------------------------
+
+
+def layer_index_vector(cfg: TreeSAEConfig) -> torch.Tensor:
+    """(n,) int64: privilege layer of every feature."""
+    return torch.repeat_interleave(
+        torch.arange(cfg.n_layers), torch.tensor(cfg.layer_widths, dtype=torch.long)
+    )
+
+
+def verify_allocation(cfg: TreeSAEConfig, allocation: torch.Tensor) -> None:
+    starts = torch.tensor(cfg.layer_starts(), dtype=torch.long, device=allocation.device)
+    start_of = starts[layer_index_vector(cfg).to(allocation.device)]
+    ok = (allocation == ROOT) | ((allocation >= 0) & (allocation < start_of))
+    if not bool(ok.all()):
+        raise AssertionError("allocation invariant violated: parent not in a strictly earlier layer")
+    if not bool((allocation[cfg.layer_slices()[0]] == ROOT).all()):
+        raise AssertionError("first-layer features must be ROOT-attached")
+
+
+def firing_mask(cfg: TreeSAEConfig, relu: torch.Tensor, allocation: torch.Tensor) -> torch.Tensor:
+    """(B, n) bool: per-layer TopK selection AND recursive parent gate (Eq. 6).
+
+    Default (ledger reading): TopK on the raw pre-activations first, gate second.
+    With ``cfg.gate_before_topk`` the TopK runs on the gated candidates instead.
+    Layers are processed in ascending order, so every parent's final firing state
+    is known when its children are gated. No gradient flows through the mask.
+    """
+    with torch.no_grad():
+        fires = torch.zeros_like(relu, dtype=torch.bool)
+        positive = relu > 0
+        for layer, (sl, k) in enumerate(zip(cfg.layer_slices(), cfg.layer_ks, strict=True)):
+            allowed = None
+            if layer > 0:
+                parents = allocation[sl]
+                is_root = parents == ROOT
+                if not bool(is_root.all()):
+                    allowed = is_root.unsqueeze(0) | fires[:, parents.clamp(min=0)]
+            cand = relu[:, sl]
+            if allowed is not None and cfg.gate_before_topk:
+                cand = cand * allowed
+            _, idx = cand.topk(k, dim=1)
+            fires[:, sl].scatter_(1, idx, True)
+            fires[:, sl] &= positive[:, sl]
+            if allowed is not None:
+                fires[:, sl] &= allowed
+    return fires
+
+
+def cumulative_reconstructions(
+    cfg: TreeSAEConfig, gated: torch.Tensor, W_dec: torch.Tensor, b: torch.Tensor
+) -> list[torch.Tensor]:
+    """``xhat_{<=l} = b + sum_{t<=l} gated_t @ W_dec_t`` for every layer l."""
+    out: list[torch.Tensor] = []
+    running = b.unsqueeze(0)
+    for sl in cfg.layer_slices():
+        running = running + gated[:, sl] @ W_dec[sl]
+        out.append(running)
+    return out
+
+
+def greedy_optimal_counts(capacities: list[float], eligible: list[bool], n_children: int) -> list[int]:
+    """Paper Algorithm 1 (heap greedy max-min payoff). Ties -> smaller id."""
+    m = len(capacities)
+    k_star = [0] * m
+    heap = [(-capacities[p], p) for p in range(m) if eligible[p]]
+    if not heap or n_children <= 0:
+        return k_star
+    heapq.heapify(heap)
+    assigned = 0
+    while heap and assigned < n_children:
+        _, p = heapq.heappop(heap)
+        k_star[p] += 1
+        heapq.heappush(heap, (-capacities[p] / (k_star[p] + 1), p))
+        assigned += 1
+    return k_star
+
+
+# ---------------------------------------------------------------------------
+# Trainer
+# ---------------------------------------------------------------------------
+
+
+class TreeSAETrainer:
+    def __init__(self, cfg: TreeSAEConfig, device: str | torch.device = "cuda") -> None:
+        self.cfg = cfg
+        self.device = torch.device(device)
+        self.d_in = cfg.d_in
+        self.n_latents = cfg.n_features
+        self.gen = torch.Generator(device="cpu").manual_seed(cfg.seed)
+        n, d = cfg.n_features, cfg.d_in
+
+        # kaiming-uniform bound for fan_in=d (a=0): sqrt(6/d); rows unit-norm; tied encoder.
+        w_dec = torch.empty(n, d, dtype=torch.float32).uniform_(-((6.0 / d) ** 0.5), (6.0 / d) ** 0.5, generator=self.gen)
+        w_dec /= w_dec.norm(dim=1, keepdim=True)
+        self.W_dec = torch.nn.Parameter(w_dec.to(self.device))
+        self.W_enc = torch.nn.Parameter(w_dec.T.contiguous().to(self.device))
+        self.b = torch.nn.Parameter(torch.zeros(d, dtype=torch.float32, device=self.device))
+        self.optimizer = torch.optim.Adam(
+            [self.W_enc, self.W_dec, self.b], lr=cfg.lr, betas=cfg.adam_betas, eps=cfg.adam_eps
+        )
+
+        # Allocation state (all on device, vectorised per-step bookkeeping).
+        alloc = torch.full((n,), ROOT, dtype=torch.long)
+        for sl in cfg.layer_slices()[1:]:
+            alloc[sl] = torch.randint(0, sl.start, (sl.stop - sl.start,), generator=self.gen)
+        self.allocation = alloc.to(self.device)
+        self.capacity = torch.zeros(n, dtype=torch.float64, device=self.device)
+        self.fire_count = torch.zeros(n, dtype=torch.long, device=self.device)
+        self.last_fire = torch.full((n,), -1, dtype=torch.long, device=self.device)
+        self.tokens_seen = 0
+        self.step_count = 0
+        self.realloc_interval = cfg.realloc_first_interval
+        self.next_realloc_step = cfg.realloc_first_interval
+        self.dead_to_root_done = False
+        self.events: list[dict[str, Any]] = []
+
+        self._layer_index = layer_index_vector(cfg).to(self.device)
+        self._first_layer = self._layer_index == 0
+        self._use_autocast = bool(cfg.autocast_bf16) and self.device.type == "cuda"
+        verify_allocation(cfg, self.allocation)
+
+    # ------------------------------------------------------------------ misc
+    def config_dict(self) -> dict[str, Any]:
+        return {**self.cfg.to_dict(), "n_latents": self.n_latents, "layer_starts": list(self.cfg.layer_starts())}
+
+    def _autocast(self):
+        if self._use_autocast:
+            return torch.autocast("cuda", dtype=torch.bfloat16)
+        return contextlib.nullcontext()
+
+    # ------------------------------------------------------------ statistics
+    def dead_mask(self) -> torch.Tensor:
+        """(n,) bool: no positive final activation in the last ``dead_token_window`` tokens."""
+        if self.tokens_seen < self.cfg.dead_token_window:
+            return torch.zeros(self.n_latents, dtype=torch.bool, device=self.device)
+        return (self.tokens_seen - self.last_fire) >= self.cfg.dead_token_window
+
+    def eligible_parent_mask(self) -> torch.Tensor:
+        """(n,) bool: firing rate >= parent_eligibility_rate (running, from start)."""
+        if self.tokens_seen == 0:
+            return torch.zeros(self.n_latents, dtype=torch.bool, device=self.device)
+        return self.fire_count.double() / float(self.tokens_seen) >= self.cfg.parent_eligibility_rate
+
+    # --------------------------------------------------------------- forward
+    def _forward(self, x: torch.Tensor):
+        """Returns (relu, fires, cumulative, per_example_loss)."""
+        with self._autocast():
+            relu = torch.relu((x - self.b) @ self.W_enc)
+            fires = firing_mask(self.cfg, relu, self.allocation)
+            gated = relu * fires
+            cumulative = cumulative_reconstructions(self.cfg, gated, self.W_dec, self.b)
+        per_example = torch.stack([((c.float() - x) ** 2).sum(dim=1) for c in cumulative], dim=0).sum(dim=0)
+        return relu, fires, cumulative, per_example
+
+    def _aux_loss(self, x: torch.Tensor, relu: torch.Tensor, cumulative: list[torch.Tensor], dead: torch.Tensor) -> torch.Tensor:
+        """Per-layer TopK-style aux loss (Eq. 8): top aux_k dead features of layer l (by
+        ungated ReLU pre-activation) reconstruct the detached residual x - xhat_{<=l}."""
+        total = x.new_zeros(())
+        for l, (sl, alpha) in enumerate(zip(self.cfg.layer_slices(), self.cfg.aux_alphas, strict=True)):
+            if alpha == 0.0:
+                continue
+            dead_ids = torch.nonzero(dead[sl], as_tuple=True)[0] + sl.start
+            n_dead = int(dead_ids.numel())
+            if n_dead == 0:
+                continue
+            with self._autocast():
+                dead_acts = relu[:, dead_ids]
+                vals, idx = dead_acts.topk(min(self.cfg.aux_k, n_dead), dim=1)
+                picked = torch.zeros_like(dead_acts).scatter_(1, idx, vals)
+                e_hat = picked @ self.W_dec[dead_ids]
+            resid = (x - cumulative[l].float()).detach()
+            total = total + alpha * ((e_hat.float() - resid) ** 2).sum(dim=1).mean()
+        return total
+
+    @torch.no_grad()
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """Final gated activations (fp32, no autocast) for raw residual activations."""
+        x = x.to(self.device, torch.float32)
+        relu = torch.relu((x - self.b) @ self.W_enc)
+        return relu * firing_mask(self.cfg, relu, self.allocation)
+
+    # ------------------------------------------------------------------ step
+    def step(self, x: torch.Tensor, step: int, total_steps: int) -> dict[str, float]:
+        cfg = self.cfg
+        x = x.to(self.device, torch.float32)
+        B = x.shape[0]
+        dead = self.dead_mask()
+
+        relu, fires, cumulative, per_example = self._forward(x)
+        recon_loss = per_example.mean()
+        aux_loss = self._aux_loss(x, relu, cumulative, dead)
+        loss = recon_loss + aux_loss
+
+        self.optimizer.zero_grad(set_to_none=True)
+        loss.backward()
+        self._project_decoder_grads()
+        self.optimizer.step()
+        self._renormalize_decoder()
+
+        with torch.no_grad():
+            per_feature = fires.sum(dim=0)  # (n,) int64
+            inc = per_example.detach() @ fires.to(torch.float32)  # (n,) fp32 batch increment
+            self.capacity += inc.double()
+            self.fire_count += per_feature
+            fired = per_feature > 0
+            self.last_fire = torch.where(fired, torch.full_like(self.last_fire, self.tokens_seen + B), self.last_fire)
+            self.tokens_seen += B
+            self.step_count = step + 1
+
+            n_dead = dead.sum()
+            n_root = (self.allocation == ROOT).sum() - int(cfg.layer_widths[0])
+            layer_l0 = torch.stack([per_feature[sl].sum() for sl in cfg.layer_slices()]).double() / B
+            mse = ((cumulative[-1].float() - x) ** 2).mean()
+            scalars = torch.stack([
+                loss.detach().double(), recon_loss.detach().double(), aux_loss.detach().double(),
+                mse.double(), n_dead.double(), n_root.double(), per_feature.sum().double() / B,
+            ]).tolist()
+            layer_l0 = layer_l0.tolist()
+
+        # --- allocation events (post-step, 1-based step count) ---
+        count = step + 1
+        if count == self.next_realloc_step:
+            self.reallocate()
+            self.realloc_interval = min(self.realloc_interval * cfg.realloc_growth, cfg.realloc_cap)
+            self.next_realloc_step = count + self.realloc_interval
+        if not self.dead_to_root_done and count >= int(cfg.dead_to_root_frac * total_steps):
+            self.move_dead_to_root()
+
+        metrics = {
+            "loss": scalars[0], "recon_loss": scalars[1], "aux_loss": scalars[2], "mse": scalars[3],
+            "n_dead": scalars[4], "frac_dead": scalars[4] / self.n_latents,
+            "n_root_attached": scalars[5], "l0": scalars[6],
+            "n_realloc_events": float(sum(e["kind"] == "realloc" for e in self.events)),
+            "dead_to_root_done": float(self.dead_to_root_done),
+        }
+        for l, v in enumerate(layer_l0):
+            metrics[f"l0_layer{l}"] = v
+        return metrics
+
+    @torch.no_grad()
+    def _project_decoder_grads(self) -> None:
+        g = self.W_dec.grad
+        if g is None:
+            return
+        rows = self.W_dec / self.W_dec.norm(dim=1, keepdim=True)
+        g -= (g * rows).sum(dim=1, keepdim=True) * rows
+
+    @torch.no_grad()
+    def _renormalize_decoder(self) -> None:
+        self.W_dec.data /= self.W_dec.data.norm(dim=1, keepdim=True)
+
+    # ------------------------------------------------------------ allocation
+    @torch.no_grad()
+    def reallocate(self) -> dict[str, Any]:
+        """One reallocation event (Alg. 2): per layer compute k* (Alg. 1) over all earlier
+        features, then first-fit the layer's DEAD children into parents whose live-child
+        count is below k* (ROOT quota slots first, then eligible parents in ascending id).
+        Live edges are never modified; unplaceable dead children keep their parent."""
+        cfg = self.cfg
+        dead = self.dead_mask().cpu()
+        eligible = self.eligible_parent_mask().cpu()
+        cap = self.capacity.cpu()
+        alloc = self.allocation.cpu().clone()
+        event: dict[str, Any] = {"kind": "realloc", "step": self.step_count, "tokens_seen": self.tokens_seen, "layers": []}
+        for layer, sl in enumerate(cfg.layer_slices()):
+            if layer == 0:
+                continue
+            m, width = sl.start, sl.stop - sl.start
+            quota = cfg.root_quota(layer)
+            k_star = torch.tensor(
+                greedy_optimal_counts(cap[:m].tolist(), eligible[:m].tolist(), width - quota), dtype=torch.long
+            )
+            children, child_dead = alloc[sl], dead[sl]
+            live_learned = children[(children != ROOT) & ~child_dead]
+            live_counts = torch.bincount(live_learned, minlength=m)
+            live_root = int(((children == ROOT) & ~child_dead).sum())
+            deficit = (k_star - live_counts).clamp(min=0) * eligible[:m].long()
+            parent_seq = torch.cat([
+                torch.full((max(quota - live_root, 0),), ROOT, dtype=torch.long),
+                torch.repeat_interleave(torch.arange(m), deficit),
+            ])
+            dead_children = torch.nonzero(child_dead, as_tuple=True)[0] + sl.start  # ascending ids
+            n_assign = min(int(dead_children.numel()), int(parent_seq.numel()))
+            if n_assign > 0:
+                alloc[dead_children[:n_assign]] = parent_seq[:n_assign]
+            event["layers"].append({
+                "layer": layer, "n_dead": int(dead_children.numel()), "n_moved": n_assign,
+                "n_eligible_parents": int(eligible[:m].sum()), "k_star_sum": int(k_star.sum()),
+                "root_quota": quota, "n_to_root": int((parent_seq[:n_assign] == ROOT).sum()),
+            })
+        verify_allocation(cfg, alloc)
+        self.allocation = alloc.to(self.device)
+        self.events.append(event)
+        return event
+
+    @torch.no_grad()
+    def move_dead_to_root(self) -> int:
+        movable = self.dead_mask() & ~self._first_layer
+        self.allocation = torch.where(movable, torch.full_like(self.allocation, ROOT), self.allocation)
+        self.dead_to_root_done = True
+        verify_allocation(self.cfg, self.allocation)
+        moved = int(movable.sum())
+        self.events.append({"kind": "dead_to_root", "step": self.step_count, "tokens_seen": self.tokens_seen, "n_moved": moved})
+        return moved
+
+    # ------------------------------------------------------------ checkpoint
+    def state_dict(self) -> dict[str, Any]:
+        return {
+            "config": self.cfg.to_dict(),
+            "model": {"W_enc": self.W_enc.detach().clone(), "W_dec": self.W_dec.detach().clone(), "b": self.b.detach().clone()},
+            "optimizer": _clone_tree(self.optimizer.state_dict()),
+            "allocation": {
+                "allocation": self.allocation.clone(), "capacity": self.capacity.clone(),
+                "fire_count": self.fire_count.clone(), "last_fire": self.last_fire.clone(),
+                "tokens_seen": self.tokens_seen, "realloc_interval": self.realloc_interval,
+                "next_realloc_step": self.next_realloc_step, "dead_to_root_done": self.dead_to_root_done,
+                "events": list(self.events),
+            },
+            "generator": self.gen.get_state().clone(),
+            "step": self.step_count,
+        }
+
+    def load_state_dict(self, state: dict[str, Any]) -> None:
+        with torch.no_grad():
+            self.W_enc.copy_(state["model"]["W_enc"])
+            self.W_dec.copy_(state["model"]["W_dec"])
+            self.b.copy_(state["model"]["b"])
+        self.optimizer.load_state_dict(_clone_tree(state["optimizer"]))
+        a = state["allocation"]
+        self.allocation = a["allocation"].to(self.device, torch.long)
+        self.capacity = a["capacity"].to(self.device, torch.float64)
+        self.fire_count = a["fire_count"].to(self.device, torch.long)
+        self.last_fire = a["last_fire"].to(self.device, torch.long)
+        self.tokens_seen = int(a["tokens_seen"])
+        self.realloc_interval = int(a["realloc_interval"])
+        self.next_realloc_step = int(a["next_realloc_step"])
+        self.dead_to_root_done = bool(a["dead_to_root_done"])
+        self.events = list(a["events"])
+        self.gen.set_state(state["generator"].cpu())
+        self.step_count = int(state["step"])
+        verify_allocation(self.cfg, self.allocation)
+
+    # ---------------------------------------------------------------- export
+    def export(self, out_dir: Path) -> None:
+        from safetensors.torch import save_file
+
+        out_dir = Path(out_dir)
+        out_dir.mkdir(parents=True, exist_ok=True)
+        save_file(
+            {"W_enc": self.W_enc.detach().cpu().contiguous(), "W_dec": self.W_dec.detach().cpu().contiguous(), "b": self.b.detach().cpu()},
+            str(out_dir / "weights.safetensors"),
+        )
+        dead = self.dead_mask()
+        (out_dir / "allocation.json").write_text(json.dumps({
+            "parent": self.allocation.cpu().tolist(),
+            "layer": self._layer_index.cpu().tolist(),
+            "dead": dead.cpu().tolist(),
+        }))
+        (out_dir / "config.json").write_text(json.dumps({
+            "arch": "tree_sae", "config": self.cfg.to_dict(), "n_latents": self.n_latents,
+            "layer_starts": list(self.cfg.layer_starts()), "step": self.step_count,
+            "tokens_seen": self.tokens_seen, "n_dead": int(dead.sum()),
+            "n_root_attached": int((self.allocation == ROOT).sum()) - int(self.cfg.layer_widths[0]),
+            "events": self.events,
+        }, indent=1))
+
+
+def _clone_tree(obj: Any) -> Any:
+    if isinstance(obj, torch.Tensor):
+        return obj.clone()
+    if isinstance(obj, dict):
+        return {k: _clone_tree(v) for k, v in obj.items()}
+    if isinstance(obj, list):
+        return [_clone_tree(v) for v in obj]
+    if isinstance(obj, tuple):
+        return tuple(_clone_tree(v) for v in obj)
+    return obj
+
+
+def build_trainer(d_in: int, overrides: dict, device: str = "cuda") -> TreeSAETrainer:
+    cfg = TreeSAEConfig.from_dict({**dict(overrides), "d_in": d_in})
+    return TreeSAETrainer(cfg, device=device)
+
+
+# ---------------------------------------------------------------------------
+# Exported model
+# ---------------------------------------------------------------------------
+
+
+class ExportedTreeSAE:
+    def __init__(self, cfg: TreeSAEConfig, W_enc: torch.Tensor, W_dec: torch.Tensor, b: torch.Tensor, allocation: torch.Tensor, device) -> None:
+        self.cfg = cfg
+        self.device = torch.device(device)
+        self.n_latents = cfg.n_features
+        self.W_enc = W_enc.to(self.device, torch.float32)
+        self.W_dec = W_dec.to(self.device, torch.float32)
+        self.b = b.to(self.device, torch.float32)
+        self.allocation = allocation.to(self.device, torch.long)
+        verify_allocation(cfg, self.allocation)
+
+    @torch.no_grad()
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        x = x.to(self.device, torch.float32)
+        relu = torch.relu((x - self.b) @ self.W_enc)
+        return relu * firing_mask(self.cfg, relu, self.allocation)
+
+    @torch.no_grad()
+    def decode(self, acts: torch.Tensor) -> torch.Tensor:
+        return acts.to(self.device, torch.float32) @ self.W_dec + self.b
+
+    def native_ancestors(self) -> dict[int, set[int]]:
+        alloc = self.allocation.cpu().tolist()
+        return {child: {parent} for child, parent in enumerate(alloc) if parent != ROOT}
+
+    def latent_level(self) -> list[int]:
+        return layer_index_vector(self.cfg).tolist()
+
+
+def load_exported(export_dir: Path, device: str = "cuda") -> ExportedTreeSAE:
+    from safetensors.torch import load_file
+
+    export_dir = Path(export_dir)
+    meta = json.loads((export_dir / "config.json").read_text())
+    cfg = TreeSAEConfig.from_dict(meta["config"])
+    weights = load_file(str(export_dir / "weights.safetensors"))
+    alloc = torch.tensor(json.loads((export_dir / "allocation.json").read_text())["parent"], dtype=torch.long)
+    return ExportedTreeSAE(cfg, weights["W_enc"], weights["W_dec"], weights["b"], alloc, device)

gemma-2-2b/layer-12/tree_sae/tree_sae_32k_k5_seed0/summary.json

@@ -0,0 +1,13 @@
+{
+ "run_name": "tree_sae_32k_k5_seed0",
+ "arch": "tree_sae",
+ "steps": 122070,
+ "tokens": 499998720,
+ "seconds": 33316.1808950901,
+ "export": "/vol/train/tree_sae_32k_k5_seed0/export",
+ "group": [
+  "hsae_32k_k5_seed0",
+  "tree_sae_32k_k5_seed0",
+  "hsae_forest_32k_k5_seed0"
+ ]
+}