Upload folder using huggingface_hub

gemma-2-2b/layer-12/hsae_forest/hsae_forest_32k_seed0_v3/DECISIONS.md

@@ -0,0 +1,178 @@
+# HSAE-Forest (Gemma-2-2b scale) — interpretation and deviation record
+
+Module: `benchmark/methods/train/hsae_forest.py`. Paper: Luo et al., "From Atoms
+to Trees", arXiv:2602.11881 (no official code). Starting point: Kirill's
+clean-room synthetic implementation
+`synth_sae_bench_baseline/experiments/synthetic/synthsaebench/hsae_forest/{model.py,LEDGER.md}`.
+"Ledger" below = that LEDGER.md. Every item is a paper-exact port unless marked
+**kept (ledger choice)** or **deviation**.
+
+## Architecture (paper Sec. 3)
+- 3 levels, widths (8192, 16384, 32768), d_in 2304; finest level = twin's 32k. Flat
+  index space = level 0, then level 1, then level 2. (Paper used 4 levels
+  2048..16384 on Gemma-2-2b; width choice is ours, ledger M2 analogue.)
+- Each level is an independent two-layer JumpReLU SAE with per-feature threshold,
+  `b_dec` subtracted before encoding (ledger "harness convention"; paper omits
+  `b_dec`).
+- Decoder rows renormalised to unit norm after every optimizer step (paper Sec. 3.3).
+- Init: decoder rows kaiming-uniform then rescaled to norm 0.1, encoder = decoder^T,
+  biases 0, threshold 0.01 (**kept (ledger P4)**, sae-lens conventions; the paper
+  does not state init).
+
+## Losses (paper Sec. 3.1)
+- `L_SAE,l = ||recon_l - x||^2` (sum over d, mean over batch) `+ lambda_l * L0_l`
+  with the sae-lens Step STE (rectangular kernel, bandwidth 0.05) for the L0 term
+  and the JumpReLU STE for the threshold gradient of the MSE term (**kept (ledger P4)**).
+- Parent-child loss `sum_i ||f_i - sum_{j in C_i} f_j||^2` over parents with a
+  non-empty child set, computed on CLEAN (unperturbed) features, `rho = 0.01`
+  (paper Sec. 3.2 value). Evaluated via the exact algebraic expansion of the
+  ledger (bit-equivalent to the naive per-parent form; tested against a naive
+  reference in `tests/test_train_hsae_forest.py`, value and gradients).
+- **Engineering change vs Kirill**: the sibling cross term
+  `2 sum_{j<j'} a_j a_j' (d_j . d_j')` is a custom autograd function evaluated in
+  pair chunks (`pc_pair_chunk`, default 32768) instead of materialising
+  `(batch, n_pairs)` tensors under autograd. Same math; peak memory is
+  `O(batch * chunk)` even when a hub parent collects thousands of children.
+  The sibling-pair enumeration is vectorised (no Python loop over parents).
+
+## Perturbation (paper Sec. 3.2)
+- `r = 0.05`. **Kept (ledger M1)**: substitute WITH probability r
+  (`z ~ Bern(r)`, parent contribution replaced by its children's sum when z = 1);
+  the printed equation has the roles of z and 1-z swapped (erratum per ledger).
+- One Bernoulli draw per parent feature per step, shared across the batch
+  ("sampled independently across features at each training step"); only parents
+  with a non-empty child set; never the last level; applied only inside the
+  reconstruction MSE of that level (ledger P8/P10). Dedicated RNG (`seed + 1`),
+  state checkpointed.
+- Engineering: only the substituted parents' / their children's columns are
+  gathered for the two extra matmuls (Kirill masked the full-width matrices).
+
+## Hierarchy update (paper Sec. 3.3, App. F.1)
+- Every 5000 gradient steps, no-grad: each child picks the parent with the highest
+  encoder-vector cosine similarity (paper default metric); children whose best
+  similarity lies in the bottom 20% quantile of that boundary's best scores are
+  left unassigned (-1). Adjacent levels only; single parent per child.
+- All children unassigned until the first update (ledger P6); one extra update on
+  the final weights, performed inside `export()` (ledger P6). Dead children are
+  assigned like any other (paper does not special-case them).
+- Similarity computed in child chunks of 8192 (memory), fp32.
+
+## Sparsity controller (paper App. A.2)
+- Per level target L0 = 150 (twin's k). EMA of batch L0 with momentum 0.999;
+  desired change `-eta (L0_ema - target)` with `eta = 0.001`; lambda is scaled by
+  `exp(kappa * (actual - desired) / target)` clamped to [0.98, 1.02] per step,
+  `kappa = 1`, `lambda_init = 0.05`, lambda clamped to [1e-8, 1e3]
+  (**kept (ledger P3)**: the paper gives eta and the EMA momentum but not the
+  update magnitude; direction matches App. A.2: faster-than-target reduction ->
+  lambda down, denser than expected -> lambda up).
+- Batch L0 is computed in fp32 from `acts > 0` (not from the STE surrogate).
+
+## Dead latents — **deviation** from the paper's ghost gradients (Sec. 3.3)
+- A latent is dead if it has not fired (act > 0) in the last 10M tokens
+  (`tokens_since_fired` buffer over the flat index space, checkpointed).
+- Aux-k loss (Gao et al. 2024 / sae-lens TopK convention): the top `aux_k = d_in // 2 = 1152`
+  dead latents per token, ranked by ReLU'd pre-activation, reconstruct the
+  detached residual `x - recon_l` of their level's clean reconstruction through
+  `W_dec` (no `b_dec`); loss `(aux_recon - residual)^2` summed over d, mean over
+  batch, scaled by `min(n_dead / aux_k, 1)`, coefficient `aux_coeff = 1/32`.
+  Rationale: ghost grads are under-specified in the paper; aux-k is the
+  standard modern replacement and has the same purpose (give dead latents a
+  reconstruction gradient). Kirill skipped dead handling (rare on synthetic data).
+
+## Input scaling (paper App. A.1)
+- Fixed global scalar `alpha = sqrt(d_in) / E||x||_2` ("RMS norm expectation of
+  sqrt(d)" read as `E||alpha x||_2 = sqrt(d)`, ledger Amendment 1), estimated as
+  a running mean over the first `scale_estimate_tokens = 65536` raw tokens (16
+  batches; the provisional value is used for those steps), then frozen and
+  stored. Applied inside `step` and inside the exported `encode` (which takes
+  RAW activations). Exported `W_dec` is divided by alpha so it is in raw units.
+
+## Optimisation (paper App. A.1, ledger P2)
+- Adam betas (0, 0.995) (paper), peak lr 3e-4 (**kept (ledger P2)**; paper does
+  not state it), cosine schedule with linear warm-up over the first 10% of steps
+  and decay to zero (paper), computed statelessly from `(step, total_steps)`;
+  global grad-norm clip 1.0 (ledger). Batch 4096 tokens (driver; paper used 1024).
+- Compute: forward/backward under `torch.autocast("cuda", bf16)`; fp32 master
+  weights and Adam state; pre-activations, JumpReLU/Step STEs, MSE, PC loss,
+  L0 statistics and the controller in fp32 (the custom sibling-cross function
+  disables autocast and casts its inputs to fp32).
+
+## Checkpoint / export
+- `state_dict`: model params + buffers (lambdas, L0 EMA, parent indices,
+  dead-latent counters, scale stats and frozen flag, hierarchy-update count),
+  Adam state, perturbation RNG state, last step, config. `load_state_dict`
+  deep-copies the optimizer state (PyTorch aliases matching tensors otherwise).
+- Export: `weights.safetensors` (per-level W_enc/W_dec/b_enc/b_dec/threshold,
+  parent indices, dead counters), `hierarchy.json` (parent index per child per
+  boundary, -1 unassigned), `config.json` (config, scale, offsets, final
+  hierarchy stats, lambdas, L0 EMA, dead fractions). The weights file (~1 GB)
+  exceeds the driver's 200 MB copy-back limit and stays on the Modal volume.
+
+## Known caveats
+- JumpReLU thresholds are `exp(log_threshold)` with lr 3e-4, so thresholds can
+  grow at most ~e^(lr) per step; leaving the dense regime takes O(10k) steps
+  (ledger Amendment 2 observed the same on synthetic data). With 122k steps for
+  500M tokens this is expected to be fine, but the L0 trajectory should be
+  checked early in the full run.
+- The exact PC loss couples parents to dead children: a parent whose children
+  are all dead is pushed toward zero activation by `rho * a_i^2`. Paper-exact;
+  not mitigated.
+
+## Smoke run (smoke_hsae_forest, A100-80GB, 4M tokens = 976 steps, overrides
+## hierarchy_update_every=300, dead_tokens=1M to exercise every code path)
+- 11.0k tokens/s steady state (incl. Gemma forward); 24.2 GB peak CUDA memory.
+- Final per-level batch L0 6158 / 8727 / 10093, lambdas railed at the 1e3 clamp,
+  mean threshold (level 2) 0.0109 -> 0.0134: the dense regime was not left in
+  976 steps (LR schedule compressed to the run length), consistent with the
+  Adam-rate-limited log-threshold caveat above. `threshold_lr_mult` (default
+  1.0) is provided as an override should the full run's early checkpoints show
+  the same stall.
+
+## Gemma 500M run: threshold_lr_mult = 20 (decided 2026-09-14)
+
+The first 500M launch with the ledger's threshold parametrisation stayed
+dense: at step 8000 (inside the 10% LR warm-up) mean thresholds had only
+moved 0.010 -> 0.022, total L0 was ~19.5k, all three lambdas were railed at
+their 1e3 cap and MSE was rising (36 -> 126) because the L0 penalty
+dominated. Adam bounds the log-threshold step to ~lr per step, so this is a
+parametrisation speed limit, not a controller bug. The run was killed and
+relaunched with `threshold_lr_mult: 20`; everything else unchanged.
+
+## PC loss: dense path for hub parents (implementation change, 2026-09-15)
+
+Symptom in the 500M run (`forest_v3`): each hierarchy update assigned more
+children to a few hub parents, and the sibling cross term of the PC-loss
+expansion (`_SiblingCross`, cost `O(batch * n_sibling_pairs)`) grew with the
+square of the hub sizes: step time went from 0.09 s to ~4 s by step 50k.
+
+Change (pure implementation change: the loss `sum_i ||f_i - sum_{j in C_i} f_j||^2`,
+its value and all its gradients are unchanged; `state_dict`/checkpoint/export
+formats untouched, runs resume as before):
+- Parents with at least `pc_dense_min_children` (default 8) children are
+  evaluated per parent as `||[-a_i, a_{C_i}] @ [d_i; D_{C_i}]||_F^2` with a
+  padded batched matmul (`_DenseParentChild`, parents grouped by child count so
+  padding wastes < 2x, processed in memory-bounded chunks, products recomputed
+  in the backward). Cost is `O(batch * |C_i| * d)` per hub, i.e. linear in the
+  number of children (bounded overall by one decoder matmul), independent of
+  sibling pairs.
+- Parents with fewer children stay on the algebraic expansion, whose pair count
+  is now bounded by `(pc_dense_min_children - 1) / 2 * n_assigned`.
+- Why not a purely sparse (active-latents-only) formulation: the exact gradient
+  of the PC loss w.r.t. an INACTIVE child activation is `-2 (f_i - sum_j f_j) . d_j`,
+  non-zero whenever its parent or a sibling is active, and the JumpReLU
+  threshold STE consumes it (rectangle window around the threshold). A
+  gather-only implementation would silently drop that threshold gradient; the
+  dense per-hub form keeps it exact while still being linear in hub size.
+- The perturbation path (`training_losses`) was already linear: it gathers the
+  substituted parents' and their children's columns for two plain matmuls.
+  Unchanged.
+- Verified in `tests/test_train_hsae_forest.py` against the original code path
+  (kept as `ForestSAE._pc_loss_reference`): value and gradients w.r.t. both
+  activation tensors and both decoders agree to < 1e-12 in float64 for hub
+  thresholds 1 (all dense) .. 1000 (all pairs), plus the whole training
+  objective / parameter gradients with the perturbation and aux losses active.
+- CPU micro-benchmark (`d_in=32`, boundary with 3 hubs x 200 children + ~250
+  small parents, batch 64, fp32, forward+backward, 2 threads): old 176 ms/iter
+  (59.9k sibling pairs) vs new 6.5 ms/iter (27x). Scaling with hub size (3 hubs
+  x 50/100/200/400 children): old 7.5 / 26 / 170 / 377 ms, new 2.8 / 2.6 / 3.4 /
+  3.7 ms -- quadratic vs linear in hub size.

gemma-2-2b/layer-12/hsae_forest/hsae_forest_32k_seed0_v3/README.md

@@ -0,0 +1,5 @@
+# hsae_forest_32k_seed0_v3
+
+Architecture `hsae_forest` 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/hsae_forest/hsae_forest_32k_seed0_v3/config.json

@@ -0,0 +1,56 @@
+{
+ "arch": "hsae_forest",
+ "run_name": "hsae_forest_32k_seed0_v3",
+ "n_tokens": 500000000,
+ "batch_tokens": 4096,
+ "total_steps": 122070,
+ "overrides": {
+  "threshold_lr_mult": 5
+ },
+ "seed": 0,
+ "trainer_config": {
+  "d_in": 2304,
+  "widths": [
+   8192,
+   16384,
+   32768
+  ],
+  "target_l0": 150.0,
+  "rho": 0.01,
+  "perturb_r": 0.05,
+  "hierarchy_update_every": 5000,
+  "exclusion_quantile": 0.2,
+  "lambda_init": 0.05,
+  "controller_eta": 0.001,
+  "l0_ema_momentum": 0.999,
+  "controller_kappa": 1.0,
+  "controller_clamp": [
+   0.98,
+   1.02
+  ],
+  "lambda_bounds": [
+   1e-08,
+   1000.0
+  ],
+  "jumprelu_bandwidth": 0.05,
+  "jumprelu_init_threshold": 0.01,
+  "decoder_init_norm": 0.1,
+  "lr": 0.0003,
+  "adam_betas": [
+   0.0,
+   0.995
+  ],
+  "warmup_frac": 0.1,
+  "grad_clip": 1.0,
+  "threshold_lr_mult": 5,
+  "dead_tokens": 10000000,
+  "aux_coeff": 0.03125,
+  "aux_k": 1152,
+  "scale_estimate_tokens": 65536,
+  "pc_pair_chunk": 32768,
+  "pc_dense_min_children": 8,
+  "pc_dense_chunk_elems": 268435456,
+  "autocast": true,
+  "seed": 0
+ }
+}

gemma-2-2b/layer-12/hsae_forest/hsae_forest_32k_seed0_v3/export/config.json

@@ -0,0 +1,100 @@
+{
+ "arch": "hsae_forest",
+ "paper": "Luo et al., arXiv:2602.11881",
+ "config": {
+  "d_in": 2304,
+  "widths": [
+   8192,
+   16384,
+   32768
+  ],
+  "target_l0": 150.0,
+  "rho": 0.01,
+  "perturb_r": 0.05,
+  "hierarchy_update_every": 5000,
+  "exclusion_quantile": 0.2,
+  "lambda_init": 0.05,
+  "controller_eta": 0.001,
+  "l0_ema_momentum": 0.999,
+  "controller_kappa": 1.0,
+  "controller_clamp": [
+   0.98,
+   1.02
+  ],
+  "lambda_bounds": [
+   1e-08,
+   1000.0
+  ],
+  "jumprelu_bandwidth": 0.05,
+  "jumprelu_init_threshold": 0.01,
+  "decoder_init_norm": 0.1,
+  "lr": 0.0003,
+  "adam_betas": [
+   0.0,
+   0.995
+  ],
+  "warmup_frac": 0.1,
+  "grad_clip": 1.0,
+  "threshold_lr_mult": 5,
+  "dead_tokens": 10000000,
+  "aux_coeff": 0.03125,
+  "aux_k": 1152,
+  "scale_estimate_tokens": 65536,
+  "pc_pair_chunk": 32768,
+  "pc_dense_min_children": 8,
+  "pc_dense_chunk_elems": 268435456,
+  "autocast": true,
+  "seed": 0
+ },
+ "n_latents": 57344,
+ "level_offsets": [
+  0,
+  8192,
+  24576,
+  57344
+ ],
+ "scale": 0.3360775411128998,
+ "scale_frozen": true,
+ "scale_estimate_tokens_seen": 65536,
+ "last_step": 122069,
+ "n_hierarchy_updates": 25,
+ "final_hierarchy_stats": {
+  "boundary_0": {
+   "n_assigned": 13107,
+   "n_children": 16384,
+   "cutoff": 0.21655727922916412,
+   "n_changed_vs_previous": 5,
+   "assigned_best_cos_mean": 0.6290030479431152,
+   "parents_with_children": 7946,
+   "children_per_parent_max": 761,
+   "n_sibling_pairs": 295359
+  },
+  "boundary_1": {
+   "n_assigned": 26214,
+   "n_children": 32768,
+   "cutoff": 0.24024911224842072,
+   "n_changed_vs_previous": 3,
+   "assigned_best_cos_mean": 0.6191921234130859,
+   "parents_with_children": 15030,
+   "children_per_parent_max": 3465,
+   "n_sibling_pairs": 6011919
+  }
+ },
+ "final_lambdas": [
+  0.6942468881607056,
+  0.6768897175788879,
+  0.6708834171295166
+ ],
+ "final_l0_ema": [
+  150.22886657714844,
+  150.12648010253906,
+  150.096923828125
+ ],
+ "frac_dead_per_level": [
+  0.0380859375,
+  0.06402587890625,
+  0.085540771484375
+ ],
+ "weights_file": "weights.safetensors",
+ "hierarchy_file": "hierarchy.json"
+}

gemma-2-2b/layer-12/hsae_forest/hsae_forest_32k_seed0_v3/export/weights.safetensors

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

gemma-2-2b/layer-12/hsae_forest/hsae_forest_32k_seed0_v3/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/hsae_forest/hsae_forest_32k_seed0_v3/model.py

@@ -0,0 +1,820 @@
+"""HSAE-Forest trainer for Gemma-2-2b layer-12 activations (``driver.py`` arch ``hsae_forest``).
+
+Luo et al., "From Atoms to Trees" (arXiv:2602.11881): ``L`` parallel,
+weight-independent JumpReLU SAEs of increasing width over the same input,
+coupled by (i) the parent-child loss ``rho * sum_i ||f_i - sum_{j in C_i} f_j||^2``,
+(ii) train-time stochastic substitution of a parent's contribution by its
+children's (perturbation rate ``r``), and (iii) a periodic no-grad hierarchy
+update that re-assigns every child to its encoder-cosine argmax parent, with
+the bottom 20% of best-similarities left unassigned. Per-level lambda is set
+by the paper's velocity-based sparsity controller (Appendix A.2).
+
+Ported from Kirill's synthetic-scale clean-room implementation
+(``synth_sae_bench_baseline/.../hsae_forest/model.py``); every interpretation
+and deviation is listed in ``HSAE_FOREST_DECISIONS.md`` next to this file.
+
+Flattened latent index space: level 0 latents, then level 1, then level 2.
+``encode`` of the export accepts RAW residual activations (the frozen input
+scale is applied inside).
+"""
+
+from __future__ import annotations
+
+import copy
+import json
+import math
+from dataclasses import asdict, dataclass, fields
+from pathlib import Path
+from typing import Any
+
+import torch
+from torch import nn
+
+WEIGHTS_FILE = "weights.safetensors"
+CONFIG_FILE = "config.json"
+HIERARCHY_FILE = "hierarchy.json"
+
+
+# --------------------------------------------------------------------------- config
+@dataclass
+class ForestConfig:
+    d_in: int = 2304
+    widths: tuple[int, ...] = (8192, 16384, 32768)
+    target_l0: float = 150.0
+    # paper mechanisms (Sec. 3.1-3.3)
+    rho: float = 0.01
+    perturb_r: float = 0.05
+    hierarchy_update_every: int = 5000
+    exclusion_quantile: float = 0.20
+    # sparsity controller (paper App. A.2: eta, EMA; magnitudes per ledger P3)
+    lambda_init: float = 0.05
+    controller_eta: float = 0.001
+    l0_ema_momentum: float = 0.999
+    controller_kappa: float = 1.0
+    controller_clamp: tuple[float, float] = (0.98, 1.02)
+    lambda_bounds: tuple[float, float] = (1e-8, 1e3)
+    # JumpReLU internals (ledger P4: sae-lens harness conventions)
+    jumprelu_bandwidth: float = 0.05
+    jumprelu_init_threshold: float = 0.01
+    decoder_init_norm: float = 0.1
+    # optimisation (paper App. A.1 + ledger P2)
+    lr: float = 3e-4
+    adam_betas: tuple[float, float] = (0.0, 0.995)
+    warmup_frac: float = 0.1
+    grad_clip: float = 1.0
+    threshold_lr_mult: float = 1.0  # lr multiplier for log_threshold params (1.0 = ledger P4 behaviour)
+    # dead-latent handling (aux-k loss; deviation from the paper's ghost grads)
+    dead_tokens: int = 10_000_000
+    aux_coeff: float = 1.0 / 32.0
+    aux_k: int = 0  # 0 -> d_in // 2
+    # input scaling: alpha with E||alpha x||_2 = sqrt(d_in), from the first batches
+    scale_estimate_tokens: int = 65_536
+    # engineering
+    pc_pair_chunk: int = 32_768
+    # parents with >= this many children take the dense per-parent path of the
+    # PC loss (cost linear in #children); smaller parents use the sibling-pair
+    # expansion. Same loss either way; see HSAE_FOREST_DECISIONS.md.
+    pc_dense_min_children: int = 8
+    pc_dense_chunk_elems: int = 1 << 28  # fp32 elements per dense-path work chunk (~1 GB)
+    autocast: bool = True
+    seed: int = 0
+
+    def __post_init__(self) -> None:
+        self.widths = tuple(int(w) for w in self.widths)
+        self.controller_clamp = tuple(float(v) for v in self.controller_clamp)
+        self.lambda_bounds = tuple(float(v) for v in self.lambda_bounds)
+        self.adam_betas = tuple(float(v) for v in self.adam_betas)
+        if self.aux_k <= 0:
+            self.aux_k = self.d_in // 2
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+def _make_config(d_in: int, overrides: dict[str, Any]) -> ForestConfig:
+    names = {f.name for f in fields(ForestConfig)}
+    unknown = set(overrides) - names
+    if unknown:
+        raise KeyError(f"unknown hsae_forest overrides: {sorted(unknown)}")
+    return ForestConfig(d_in=d_in, **overrides)
+
+
+# --------------------------------------------------------------------------- JumpReLU STE
+# Identical to sae-lens ``JumpReLU`` / ``Step`` (rectangular kernel), inlined to
+# avoid importing sae-lens inside the training container.
+def _rectangle(x: torch.Tensor) -> torch.Tensor:
+    return ((x > -0.5) & (x < 0.5)).to(x)
+
+
+class _JumpReLU(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx: Any, x: torch.Tensor, threshold: torch.Tensor, bandwidth: float) -> torch.Tensor:
+        ctx.save_for_backward(x, threshold)
+        ctx.bandwidth = bandwidth
+        return x * (x > threshold)
+
+    @staticmethod
+    def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, None]:
+        x, threshold = ctx.saved_tensors
+        bw = ctx.bandwidth
+        x_grad = (x > threshold) * grad_output
+        threshold_grad = torch.sum(-(threshold / bw) * _rectangle((x - threshold) / bw) * grad_output, dim=0)
+        return x_grad, threshold_grad, None
+
+
+class _Step(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx: Any, x: torch.Tensor, threshold: torch.Tensor, bandwidth: float) -> torch.Tensor:
+        ctx.save_for_backward(x, threshold)
+        ctx.bandwidth = bandwidth
+        return (x > threshold).to(x)
+
+    @staticmethod
+    def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[None, torch.Tensor, None]:
+        x, threshold = ctx.saved_tensors
+        bw = ctx.bandwidth
+        threshold_grad = torch.sum(-(1.0 / bw) * _rectangle((x - threshold) / bw) * grad_output, dim=0)
+        return None, threshold_grad, None
+
+
+# --------------------------------------------------------------------------- sibling cross term
+class _SiblingCross(torch.autograd.Function):
+    """Per-sample ``sum_{(j1,j2) sibling pairs} a[:,j1] a[:,j2] (d_j1 . d_j2)``.
+
+    Same value/gradient as materialising ``(B, n_pairs)`` tensors under
+    autograd (Kirill's version), but evaluated in pair chunks so peak memory is
+    ``O(B * chunk)`` regardless of how many children a hub parent collects.
+    """
+
+    @staticmethod
+    @torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float32)
+    def forward(ctx: Any, acts: torch.Tensor, w_dec: torch.Tensor, j1: torch.Tensor, j2: torch.Tensor, chunk: int) -> torch.Tensor:
+        out = acts.new_zeros(acts.shape[0])
+        for s in range(0, j1.numel(), chunk):
+            i1, i2 = j1[s : s + chunk], j2[s : s + chunk]
+            v = (w_dec[i1] * w_dec[i2]).sum(-1)
+            out += (acts[:, i1] * acts[:, i2] * v).sum(-1)
+        ctx.save_for_backward(acts, w_dec, j1, j2)
+        ctx.chunk = chunk
+        return out
+
+    @staticmethod
+    @torch.amp.custom_bwd(device_type="cuda")
+    def backward(ctx: Any, g: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, None, None, None]:
+        acts, w_dec, j1, j2 = ctx.saved_tensors
+        chunk = ctx.chunk
+        g_acts = torch.zeros_like(acts)
+        g_w = torch.zeros_like(w_dec)
+        gcol = g[:, None].to(acts.dtype)
+        for s in range(0, j1.numel(), chunk):
+            i1, i2 = j1[s : s + chunk], j2[s : s + chunk]
+            v = (w_dec[i1] * w_dec[i2]).sum(-1)
+            a1, a2 = acts[:, i1], acts[:, i2]
+            g_acts.index_add_(1, i1, gcol * a2 * v)
+            g_acts.index_add_(1, i2, gcol * a1 * v)
+            q = (gcol * a1 * a2).sum(0)
+            g_w.index_add_(0, i1, q[:, None] * w_dec[i2])
+            g_w.index_add_(0, i2, q[:, None] * w_dec[i1])
+        return g_acts, g_w, None, None, None
+
+
+# --------------------------------------------------------------------------- dense hub-parent term
+def _dense_gather(
+    a_p: torch.Tensor, a_c: torch.Tensor, d_p: torch.Tensor, d_c: torch.Tensor,
+    parents: torch.Tensor, idx: torch.Tensor, valid: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Per parent ``p`` in ``parents``: ``A_p = [-a_p, a_{C_p}]`` (B, S+1) and
+    ``D_p = [d_p; d_{C_p}]`` (S+1, d), children padded to ``S`` with zero
+    decoder rows (``idx`` already clamped into range, ``valid`` marks real children)."""
+    B = a_p.shape[0]
+    h, S = idx.shape
+    a_kids = a_c[:, idx.reshape(-1)].reshape(B, h, S).transpose(0, 1)  # (h, B, S)
+    a_full = torch.cat([-a_p[:, parents].T[:, :, None], a_kids], dim=2)  # (h, B, S+1)
+    d_kids = d_c[idx] * valid[:, :, None].to(d_c.dtype)  # (h, S, d)
+    d_full = torch.cat([d_p[parents][:, None, :], d_kids], dim=1)  # (h, S+1, d)
+    return a_full, d_full
+
+
+class _DenseParentChild(torch.autograd.Function):
+    """``sum_t sum_{p in parents} ||a_tp d_p - sum_{j in C_p} a_tj d_j||^2`` evaluated
+    directly per parent as ``||A_p @ D_p||_F^2`` with a padded batched matmul.
+
+    Cost is ``O(B * (|C_p| + 1) * d)`` per parent -- linear in the number of
+    children, never in sibling pairs -- and bounded by one decoder matmul in
+    total. Gradients are the exact dense ones (including w.r.t. inactive
+    activations, which the JumpReLU threshold STE consumes), identical to the
+    algebraic-expansion path. Parents are processed in chunks of ``n_chunk`` so
+    peak memory is ``O(n_chunk * B * (d + S))``; the backward recomputes the
+    per-chunk products instead of storing them.
+    """
+
+    @staticmethod
+    @torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float32)
+    def forward(
+        ctx: Any, a_p: torch.Tensor, a_c: torch.Tensor, d_p: torch.Tensor, d_c: torch.Tensor,
+        parents: torch.Tensor, idx: torch.Tensor, valid: torch.Tensor, n_chunk: int,
+    ) -> torch.Tensor:
+        ctx.save_for_backward(a_p, a_c, d_p, d_c, parents, idx, valid)
+        ctx.n_chunk = n_chunk
+        total = a_p.new_zeros(())
+        for s in range(0, parents.numel(), n_chunk):
+            a_full, d_full = _dense_gather(a_p, a_c, d_p, d_c, parents[s : s + n_chunk], idx[s : s + n_chunk], valid[s : s + n_chunk])
+            total = total + torch.bmm(a_full, d_full).pow(2).sum()
+        return total
+
+    @staticmethod
+    @torch.amp.custom_bwd(device_type="cuda")
+    def backward(ctx: Any, g: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, None, None, None, None]:
+        a_p, a_c, d_p, d_c, parents, idx, valid = ctx.saved_tensors
+        n_chunk = ctx.n_chunk
+        B, n_c = a_c.shape
+        d = d_c.shape[1]
+        g_ap = torch.zeros_like(a_p)
+        g_ac = torch.zeros_like(a_c)
+        g_dp = torch.zeros_like(d_p)
+        g_dc = torch.zeros_like(d_c)
+        two_g = 2.0 * g.to(a_p.dtype)
+        for s in range(0, parents.numel(), n_chunk):
+            P, I, V = parents[s : s + n_chunk], idx[s : s + n_chunk], valid[s : s + n_chunk]
+            a_full, d_full = _dense_gather(a_p, a_c, d_p, d_c, P, I, V)
+            u = torch.bmm(a_full, d_full) * two_g  # (h, B, d) = 2 g * residual
+            ga = torch.bmm(u, d_full.transpose(1, 2))  # (h, B, S+1); pad columns are 0 (zero rows)
+            gd = torch.bmm(a_full.transpose(1, 2), u)  # (h, S+1, d)
+            g_ap.index_add_(1, P, -ga[:, :, 0].T)
+            g_ac.index_add_(1, I.reshape(-1), ga[:, :, 1:].permute(1, 0, 2).reshape(B, -1))
+            g_dp.index_add_(0, P, gd[:, 0, :])
+            gd_kids = gd[:, 1:, :] * V[:, :, None].to(gd.dtype)  # drop the padding rows' (spurious) grads
+            g_dc.index_add_(0, I.reshape(-1), gd_kids.reshape(-1, d))
+        return g_ap, g_ac, g_dp, g_dc, None, None, None, None
+
+
+# --------------------------------------------------------------------------- model
+class ForestLevel(nn.Module):
+    """One JumpReLU SAE level (per-feature threshold, b_dec pre-subtracted)."""
+
+    def __init__(self, d_in: int, width: int, cfg: ForestConfig) -> None:
+        super().__init__()
+        w_dec = torch.empty(width, d_in)
+        nn.init.kaiming_uniform_(w_dec)
+        w_dec = w_dec / w_dec.norm(dim=-1, keepdim=True) * cfg.decoder_init_norm
+        self.W_dec = nn.Parameter(w_dec)
+        self.W_enc = nn.Parameter(w_dec.T.clone().contiguous())
+        self.b_enc = nn.Parameter(torch.zeros(width))
+        self.b_dec = nn.Parameter(torch.zeros(d_in))
+        self.log_threshold = nn.Parameter(torch.full((width,), math.log(cfg.jumprelu_init_threshold)))
+        self.bandwidth = cfg.jumprelu_bandwidth
+
+    @property
+    def threshold(self) -> torch.Tensor:
+        return self.log_threshold.exp()
+
+    def pre_acts(self, x: torch.Tensor) -> torch.Tensor:
+        return ((x - self.b_dec) @ self.W_enc + self.b_enc).float()
+
+    def acts(self, pre: torch.Tensor) -> torch.Tensor:
+        return _JumpReLU.apply(pre, self.threshold, self.bandwidth)
+
+    def l0_step(self, pre: torch.Tensor) -> torch.Tensor:
+        return _Step.apply(pre, self.threshold, self.bandwidth).sum(dim=-1)
+
+    def decode(self, acts: torch.Tensor) -> torch.Tensor:
+        return (acts @ self.W_dec).float() + self.b_dec
+
+    @torch.no_grad()
+    def renormalize_decoder(self) -> None:
+        self.W_dec.data /= self.W_dec.data.norm(dim=-1, keepdim=True).clamp(min=1e-8)
+
+
+class ForestSAE(nn.Module):
+    def __init__(self, cfg: ForestConfig) -> None:
+        super().__init__()
+        self.cfg = cfg
+        self.levels = nn.ModuleList([ForestLevel(cfg.d_in, w, cfg) for w in cfg.widths])
+        n_levels = len(cfg.widths)
+        for b, child_w in enumerate(cfg.widths[1:]):
+            self.register_buffer(f"parent_idx_{b}", torch.full((child_w,), -1, dtype=torch.int64))
+        self.register_buffer("lambdas", torch.full((n_levels,), cfg.lambda_init))
+        self.register_buffer("l0_ema", torch.full((n_levels,), cfg.target_l0))
+        self.register_buffer("n_hierarchy_updates", torch.zeros((), dtype=torch.int64))
+        # dead-latent tracking over the flattened index space
+        self.register_buffer("tokens_since_fired", torch.zeros(sum(cfg.widths), dtype=torch.int64))
+        # input scale alpha = sqrt(d) / E||x||_2, estimated from the first batches
+        self.register_buffer("norm_sum", torch.zeros((), dtype=torch.float64))
+        self.register_buffer("norm_count", torch.zeros((), dtype=torch.int64))
+        self.register_buffer("scale", torch.ones(()))
+        self.register_buffer("scale_frozen", torch.zeros((), dtype=torch.bool))
+        self._pc_cache: dict[int, dict[str, Any]] = {}
+
+    # ---- bookkeeping -------------------------------------------------------
+    @property
+    def n_levels(self) -> int:
+        return len(self.cfg.widths)
+
+    @property
+    def n_boundaries(self) -> int:
+        return len(self.cfg.widths) - 1
+
+    @property
+    def offsets(self) -> list[int]:
+        out = [0]
+        for w in self.cfg.widths:
+            out.append(out[-1] + w)
+        return out
+
+    def parent_idx(self, boundary: int) -> torch.Tensor:
+        return getattr(self, f"parent_idx_{boundary}")
+
+    def _apply(self, fn, recurse=True):  # type: ignore[override]
+        self._pc_cache = {}
+        return super()._apply(fn, recurse)
+
+    def load_state_dict(self, *args: Any, **kwargs: Any):  # type: ignore[override]
+        result = super().load_state_dict(*args, **kwargs)
+        self._pc_cache = {}
+        return result
+
+    def dead_masks(self) -> list[torch.Tensor]:
+        dead = self.tokens_since_fired >= self.cfg.dead_tokens
+        off = self.offsets
+        return [dead[off[l] : off[l + 1]] for l in range(self.n_levels)]
+
+    # ---- input scale ---------------------------------------------------------
+    @torch.no_grad()
+    def update_scale(self, x_raw: torch.Tensor) -> None:
+        if bool(self.scale_frozen):
+            return
+        self.norm_sum += x_raw.float().norm(dim=-1).double().sum()
+        self.norm_count += x_raw.shape[0]
+        mean_norm = self.norm_sum / self.norm_count.double()
+        self.scale.fill_(float(math.sqrt(self.cfg.d_in) / mean_norm))
+        if int(self.norm_count) >= self.cfg.scale_estimate_tokens:
+            self.scale_frozen.fill_(True)
+
+    # ---- hierarchy update (no-grad; paper Sec. 3.3 "Similarity-Based Hierarchy Update")
+    @torch.no_grad()
+    def hierarchy_update(self, child_chunk: int = 8192) -> dict[str, Any]:
+        stats: dict[str, Any] = {}
+        for b in range(self.n_boundaries):
+            parent, child = self.levels[b], self.levels[b + 1]
+            e_parent = nn.functional.normalize(parent.W_enc.T.float(), dim=-1)  # (n_p, d)
+            e_child = nn.functional.normalize(child.W_enc.T.float(), dim=-1)  # (n_c, d)
+            bests, argmaxes = [], []
+            for s in range(0, e_child.shape[0], child_chunk):
+                sim = e_child[s : s + child_chunk] @ e_parent.T
+                best, argmax = sim.max(dim=-1)
+                bests.append(best)
+                argmaxes.append(argmax)
+            best = torch.cat(bests)
+            argmax = torch.cat(argmaxes)
+            cutoff = torch.quantile(best, self.cfg.exclusion_quantile)
+            assigned = best >= cutoff  # bottom-quantile best-scores stay unassigned
+            idx = torch.where(assigned, argmax, torch.full_like(argmax, -1))
+            previous = self.parent_idx(b).clone()
+            self.parent_idx(b).copy_(idx)
+            branching = torch.bincount(idx[assigned], minlength=self.cfg.widths[b])
+            stats[f"boundary_{b}"] = {
+                "n_assigned": int(assigned.sum()),
+                "n_children": int(idx.numel()),
+                "cutoff": float(cutoff),
+                "n_changed_vs_previous": int((idx != previous).sum()),
+                "assigned_best_cos_mean": float(best[assigned].mean()) if bool(assigned.any()) else None,
+                "parents_with_children": int((branching > 0).sum()),
+                "children_per_parent_max": int(branching.max()),
+                "n_sibling_pairs": int((branching.double() * (branching.double() - 1) / 2).sum()),
+            }
+        self.n_hierarchy_updates += 1
+        self._pc_cache = {}
+        return stats
+
+    # ---- parent-child loss ---------------------------------------------------
+    @staticmethod
+    def _sibling_pairs(js: torch.Tensor, ps: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor] | tuple[None, None]:
+        """All sibling pairs ``(j, j')``, ``j`` before ``j'``, of children ``js`` sorted by parent ``ps``."""
+        n = ps.numel()
+        if n == 0:
+            return None, None
+        dev = ps.device
+        change = torch.ones(n, dtype=torch.bool, device=dev)
+        change[1:] = ps[1:] != ps[:-1]
+        run_id = torch.cumsum(change.long(), 0) - 1
+        run_sizes = torch.bincount(run_id)
+        run_end = torch.cumsum(run_sizes, 0)
+        pos = torch.arange(n, device=dev)
+        counts = run_end[run_id] - pos - 1  # pairs (pos, pos') with pos' > pos in the same run
+        total = int(counts.sum())
+        if total == 0:
+            return None, None
+        j1_pos = torch.repeat_interleave(pos, counts)
+        pair_start = torch.cumsum(counts, 0) - counts
+        within = torch.arange(total, device=dev) - torch.repeat_interleave(pair_start, counts)
+        j2_pos = j1_pos + 1 + within
+        return js[j1_pos], js[j2_pos]
+
+    def _build_pc_structure(self, boundary: int, dense_min_children: int | None) -> dict[str, Any]:
+        """Assignment-dependent index structure of one boundary.
+
+        Parents with fewer than ``dense_min_children`` children ("small") are
+        evaluated by the algebraic expansion (``has_child``/``j``/``pj`` masks and
+        the sibling pairs ``j1``/``j2`` cover ONLY those); parents with at least
+        that many ("hubs") are listed in ``dense_groups`` for the per-parent
+        dense path. ``dense_min_children=None`` puts every parent on the
+        expansion path (the original implementation, kept for the reference).
+        """
+        p_idx = self.parent_idx(boundary)
+        n_p, n_c = self.cfg.widths[boundary], self.cfg.widths[boundary + 1]
+        dev = p_idx.device
+        assigned = p_idx >= 0
+        has_child = torch.zeros(n_p, dtype=torch.bool, device=dev)
+        struct: dict[str, Any] = {"any_assigned": bool(assigned.any()), "j1": None, "j2": None, "has_child": has_child, "dense_groups": []}
+        if not struct["any_assigned"]:
+            return struct
+        j = torch.nonzero(assigned, as_tuple=True)[0]
+        pj = p_idx[j]
+        order = torch.argsort(pj, stable=True)
+        js, ps = j[order], pj[order]  # children sorted by parent
+        counts = torch.bincount(pj, minlength=n_p)  # children per parent
+        if dense_min_children is None:
+            small = torch.ones_like(js, dtype=torch.bool)
+        else:
+            small = counts[ps] < dense_min_children
+        js_s, ps_s = js[small], ps[small]
+        has_child[ps_s] = True
+        struct["j"], struct["pj"] = js_s, ps_s
+        struct["j1"], struct["j2"] = self._sibling_pairs(js_s, ps_s)
+        if dense_min_children is not None:
+            hubs = torch.nonzero(counts >= dense_min_children, as_tuple=True)[0]
+            if hubs.numel() > 0:
+                starts = torch.cumsum(counts, 0) - counts  # first position of each parent's run in js
+                sizes = counts[hubs]
+                order_h = torch.argsort(sizes, descending=True, stable=True)
+                hubs, sizes = hubs[order_h], sizes[order_h]
+                sizes_l = sizes.tolist()
+                # group parents of similar child count so padding wastes < 2x
+                groups: list[tuple[int, int]] = []
+                start = 0
+                for k in range(1, len(sizes_l) + 1):
+                    if k == len(sizes_l) or sizes_l[k] * 2 < sizes_l[start]:
+                        groups.append((start, k))
+                        start = k
+                for g0, g1 in groups:
+                    P, sz = hubs[g0:g1], sizes[g0:g1]
+                    h, S = P.numel(), int(sz[0])
+                    idx = torch.full((h, S), n_c, dtype=torch.int64, device=dev)
+                    rows = torch.repeat_interleave(torch.arange(h, device=dev), sz)
+                    cols = torch.arange(int(sz.sum()), device=dev) - torch.repeat_interleave(torch.cumsum(sz, 0) - sz, sz)
+                    idx[rows, cols] = js[starts[P][rows] + cols]
+                    valid = idx < n_c
+                    struct["dense_groups"].append({"parents": P, "idx": idx.clamp(max=n_c - 1), "valid": valid, "S": S})
+        return struct
+
+    def _pc_structure(self, boundary: int) -> dict[str, Any]:
+        """Structure for the configured hub threshold, cached between hierarchy updates."""
+        cached = self._pc_cache.get(boundary)
+        if cached is None:
+            cached = self._build_pc_structure(boundary, self.cfg.pc_dense_min_children)
+            self._pc_cache[boundary] = cached
+        return cached
+
+    def _pc_expansion(self, boundary: int, struct: dict[str, Any], acts_parent: torch.Tensor, acts_child: torch.Tensor) -> torch.Tensor:
+        """Per-sample ``sum_i ||a_i d_i - sum_{j in C_i} a_j d_j||^2`` over the parents covered by
+        ``struct``'s masks, via the exact expansion
+        ``a_i^2||d_i||^2 - 2 sum_j a_i a_j (d_i.d_j) + sum_j a_j^2||d_j||^2 + 2 sum_{j<j'} a_j a_j' (d_j.d_j')``."""
+        parent, child = self.levels[boundary], self.levels[boundary + 1]
+        d_parent_sq = parent.W_dec.pow(2).sum(-1)
+        term1 = (acts_parent.pow(2) * (d_parent_sq * struct["has_child"])).sum(-1)
+        if not struct["any_assigned"]:
+            return term1
+        j, pj = struct["j"], struct["pj"]
+        dots_pc = (child.W_dec[j] * parent.W_dec[pj]).sum(-1)
+        a_c = acts_child[:, j]
+        term2 = (acts_parent[:, pj] * a_c * dots_pc).sum(-1)
+        term3 = (a_c.pow(2) * child.W_dec[j].pow(2).sum(-1)).sum(-1)
+        total = term1 - 2.0 * term2 + term3
+        if struct["j1"] is not None:
+            total = total + 2.0 * _SiblingCross.apply(acts_child, child.W_dec, struct["j1"], struct["j2"], self.cfg.pc_pair_chunk)
+        return total
+
+    def pc_loss_boundary(self, boundary: int, acts_parent: torch.Tensor, acts_child: torch.Tensor) -> torch.Tensor:
+        """``mean_b sum_i ||a_i d_i - sum_{j in C_i} a_j d_j||^2`` over parents with children.
+
+        Parents with fewer than ``cfg.pc_dense_min_children`` children use the
+        sibling-pair expansion (``_pc_expansion``); hub parents use the dense
+        per-parent form (``_DenseParentChild``, cost linear in #children). Both
+        give the same value and gradients (see ``_pc_loss_reference``).
+        """
+        struct = self._pc_structure(boundary)
+        total = self._pc_expansion(boundary, struct, acts_parent, acts_child).mean()
+        if struct["dense_groups"]:
+            parent, child = self.levels[boundary], self.levels[boundary + 1]
+            B, d = acts_parent.shape[0], self.cfg.d_in
+            dense = acts_parent.new_zeros(())
+            for grp in struct["dense_groups"]:
+                n_chunk = max(1, self.cfg.pc_dense_chunk_elems // (B * (d + grp["S"] + 1)))
+                dense = dense + _DenseParentChild.apply(
+                    acts_parent, acts_child, parent.W_dec, child.W_dec, grp["parents"], grp["idx"], grp["valid"], n_chunk
+                )
+            total = total + dense / B
+        return total
+
+    def _pc_loss_reference(self, boundary: int, acts_parent: torch.Tensor, acts_child: torch.Tensor) -> torch.Tensor:
+        """Original implementation (every parent on the sibling-pair expansion; cost
+        ``O(batch * n_sibling_pairs)``). Kept, uncached, as the test reference."""
+        struct = self._build_pc_structure(boundary, None)
+        return self._pc_expansion(boundary, struct, acts_parent, acts_child).mean()
+
+    # ---- training losses -----------------------------------------------------
+    def training_losses(self, x: torch.Tensor, generator: torch.Generator | None, dead_masks: list[torch.Tensor], n_dead: list[int]) -> dict[str, Any]:
+        cfg = self.cfg
+        pres = [level.pre_acts(x) for level in self.levels]
+        acts = [level.acts(pre) for level, pre in zip(self.levels, pres, strict=True)]
+        recons = [level.decode(a) for level, a in zip(self.levels, acts, strict=True)]
+
+        # Perturbation (ledger M1: substitute WITH probability r), one Bernoulli
+        # draw per parent per step shared across the batch; only parents with
+        # children; never the last level; only inside the reconstruction MSE.
+        mse_terms = []
+        for l, level in enumerate(self.levels):
+            recon = recons[l]
+            if cfg.perturb_r > 0 and l < self.n_boundaries:
+                p_idx = self.parent_idx(l)
+                assigned = p_idx >= 0
+                if bool(assigned.any()):
+                    n_p = cfg.widths[l]
+                    z = torch.rand(n_p, device=x.device, generator=generator) < cfg.perturb_r
+                    has_child = torch.zeros(n_p, dtype=torch.bool, device=x.device)
+                    has_child[p_idx[assigned]] = True
+                    sub = torch.nonzero(z & has_child, as_tuple=True)[0]
+                    if sub.numel() > 0:
+                        sub_mask = torch.zeros(n_p, dtype=torch.bool, device=x.device)
+                        sub_mask[sub] = True
+                        cidx = torch.nonzero(assigned & sub_mask[p_idx.clamp(min=0)], as_tuple=True)[0]
+                        child_level = self.levels[l + 1]
+                        parent_part = (acts[l][:, sub] @ level.W_dec[sub]).float()
+                        child_part = (acts[l + 1][:, cidx] @ child_level.W_dec[cidx]).float()
+                        recon = recon - parent_part + child_part
+            mse_terms.append((recon - x).pow(2).sum(-1).mean())
+
+        l0_losses = [level.l0_step(pre).mean() for level, pre in zip(self.levels, pres, strict=True)]
+        sae_loss = sum(mse + self.lambdas[l] * l0 for l, (mse, l0) in enumerate(zip(mse_terms, l0_losses, strict=True)))
+        pc = sum(self.pc_loss_boundary(b, acts[b], acts[b + 1]) for b in range(self.n_boundaries))
+        aux = self._aux_loss(x, pres, recons, dead_masks, n_dead)
+        total = sae_loss + cfg.rho * pc + cfg.aux_coeff * aux
+        with torch.no_grad():
+            batch_l0 = torch.stack([(a > 0).float().sum(-1).mean() for a in acts])
+            fired = torch.cat([(a > 0).any(dim=0) for a in acts])
+        return {
+            "loss": total,
+            "mse_terms": [t.detach() for t in mse_terms],
+            "pc_loss": pc.detach() if torch.is_tensor(pc) else torch.tensor(float(pc)),
+            "aux_loss": aux.detach() if torch.is_tensor(aux) else torch.tensor(float(aux)),
+            "batch_l0": batch_l0,
+            "fired": fired,
+            "acts": acts,
+        }
+
+    def _aux_loss(self, x: torch.Tensor, pres: list[torch.Tensor], recons: list[torch.Tensor], dead_masks: list[torch.Tensor], n_dead: list[int]) -> torch.Tensor | float:
+        """Aux-k residual reconstruction on dead latents (Gao et al. 2024 style;
+        deviation from the paper's ghost grads). Top ``aux_k`` dead latents by
+        ReLU'd pre-activation reconstruct the detached residual of the clean
+        reconstruction; no ``b_dec``; scaled by ``min(n_dead / aux_k, 1)``."""
+        cfg = self.cfg
+        total: torch.Tensor | float = 0.0
+        for l, level in enumerate(self.levels):
+            if n_dead[l] == 0:
+                continue
+            k = min(cfg.aux_k, n_dead[l])
+            residual = (x - recons[l]).detach()
+            masked = torch.where(dead_masks[l][None, :], torch.relu(pres[l]), torch.full_like(pres[l], -torch.inf))
+            top = masked.topk(k, dim=-1, sorted=False)
+            aux_acts = torch.zeros_like(pres[l]).scatter_(-1, top.indices, top.values)
+            aux_recon = (aux_acts @ level.W_dec).float()
+            scale = min(n_dead[l] / cfg.aux_k, 1.0)
+            total = total + scale * (aux_recon - residual).pow(2).sum(-1).mean()
+        return total
+
+    # ---- sparsity controller (paper App. A.2; magnitudes per ledger P3) --------
+    @torch.no_grad()
+    def controller_step(self, batch_l0: torch.Tensor) -> None:
+        cfg = self.cfg
+        m = cfg.l0_ema_momentum
+        prev = self.l0_ema.clone()
+        self.l0_ema.mul_(m).add_((1 - m) * batch_l0.float())
+        desired_delta = -cfg.controller_eta * (prev - cfg.target_l0)
+        actual_delta = self.l0_ema - prev
+        surplus = (actual_delta - desired_delta) / max(cfg.target_l0, 1.0)
+        factor = torch.exp(cfg.controller_kappa * surplus).clamp(*cfg.controller_clamp)
+        self.lambdas.mul_(factor).clamp_(*cfg.lambda_bounds)
+
+    @torch.no_grad()
+    def renormalize_decoders(self) -> None:
+        for level in self.levels:
+            level.renormalize_decoder()
+
+    @torch.no_grad()
+    def update_dead_tracking(self, fired: torch.Tensor, n_tokens: int) -> None:
+        self.tokens_since_fired += n_tokens
+        self.tokens_since_fired[fired] = 0
+
+    def parents_per_boundary(self) -> list[list[int]]:
+        return [self.parent_idx(b).tolist() for b in range(self.n_boundaries)]
+
+
+# --------------------------------------------------------------------------- trainer
+def lr_multiplier(step: int, total_steps: int, warmup_frac: float) -> float:
+    warmup = max(1, int(total_steps * warmup_frac))
+    if step < warmup:
+        return (step + 1) / warmup
+    progress = (step - warmup) / max(1, total_steps - warmup)
+    return 0.5 * (1.0 + math.cos(math.pi * min(progress, 1.0)))
+
+
+class HSAEForestTrainer:
+    def __init__(self, cfg: ForestConfig, device: str = "cuda") -> None:
+        self.cfg = cfg
+        self.device = torch.device(device)
+        self.d_in = cfg.d_in
+        self.n_latents = sum(cfg.widths)
+        self.model = ForestSAE(cfg).to(self.device)
+        thr = [p for n, p in self.model.named_parameters() if n.endswith("log_threshold")]
+        rest = [p for n, p in self.model.named_parameters() if not n.endswith("log_threshold")]
+        self.optimizer = torch.optim.Adam(
+            [{"params": rest, "lr_mult": 1.0}, {"params": thr, "lr_mult": cfg.threshold_lr_mult}],
+            lr=cfg.lr, betas=cfg.adam_betas,
+        )
+        self.perturb_gen = torch.Generator(device=self.device)
+        self.perturb_gen.manual_seed(cfg.seed + 1)
+        self.last_step = -1
+        self.last_hierarchy_stats: dict[str, Any] = {}
+        self.use_autocast = cfg.autocast and self.device.type == "cuda"
+
+    def config_dict(self) -> dict[str, Any]:
+        return self.cfg.to_dict()
+
+    def step(self, x: torch.Tensor, step: int, total_steps: int) -> dict[str, float]:
+        cfg, model = self.cfg, self.model
+        x = x.to(self.device, dtype=torch.float32, non_blocking=True)
+        model.update_scale(x)
+        x = x * model.scale
+        lr = cfg.lr * lr_multiplier(step, total_steps, cfg.warmup_frac)
+        for group in self.optimizer.param_groups:
+            group["lr"] = lr * group["lr_mult"]
+
+        dead_masks = model.dead_masks()
+        n_dead = torch.stack([m.sum() for m in dead_masks]).tolist()
+        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=self.use_autocast):
+            out = model.training_losses(x, self.perturb_gen, dead_masks, n_dead)
+        out["loss"].backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)
+        self.optimizer.step()
+        self.optimizer.zero_grad(set_to_none=True)
+        model.renormalize_decoders()
+        model.controller_step(out["batch_l0"])
+        model.update_dead_tracking(out["fired"], int(x.shape[0]))
+        if (step + 1) % cfg.hierarchy_update_every == 0:
+            self.last_hierarchy_stats = model.hierarchy_update()
+            print(json.dumps({"hierarchy_update": step + 1, **self.last_hierarchy_stats}), flush=True)
+        self.last_step = step
+
+        batch_l0 = out["batch_l0"].tolist()
+        metrics: dict[str, float] = {
+            "loss": float(out["loss"].detach()),
+            "l0": float(sum(batch_l0)),
+            "pc_loss": float(out["pc_loss"]),
+            "aux_loss": float(out["aux_loss"]),
+            "frac_dead": float(sum(n_dead)) / self.n_latents,
+            "lr": lr,
+            "scale": float(model.scale),
+        }
+        if self.device.type == "cuda":
+            metrics["cuda_max_mem_gb"] = torch.cuda.max_memory_allocated(self.device) / 1e9
+        lambdas, emas = model.lambdas.tolist(), model.l0_ema.tolist()
+        mses = [float(t) for t in out["mse_terms"]]
+        for l in range(model.n_levels):
+            metrics[f"l0/level{l}"] = batch_l0[l]
+            metrics[f"l0_ema/level{l}"] = emas[l]
+            metrics[f"lambda/level{l}"] = lambdas[l]
+            metrics[f"mse/level{l}"] = mses[l]
+            metrics[f"frac_dead/level{l}"] = n_dead[l] / cfg.widths[l]
+            metrics[f"threshold_mean/level{l}"] = float(model.levels[l].threshold.mean())
+        return metrics
+
+    # ---- checkpointing -------------------------------------------------------
+    def state_dict(self) -> dict[str, Any]:
+        return {
+            "model": self.model.state_dict(),
+            "optimizer": self.optimizer.state_dict(),
+            "perturb_gen": self.perturb_gen.get_state(),
+            "last_step": self.last_step,
+            "config": self.cfg.to_dict(),
+        }
+
+    def load_state_dict(self, state: dict[str, Any]) -> None:
+        self.model.load_state_dict(state["model"])
+        # Optimizer.load_state_dict aliases state tensors that already match the
+        # param dtype/device (no copy), so deep-copy to keep the source intact.
+        self.optimizer.load_state_dict(copy.deepcopy(state["optimizer"]))
+        self.perturb_gen.set_state(state["perturb_gen"].cpu())
+        self.last_step = int(state.get("last_step", -1))
+
+    # ---- export --------------------------------------------------------------
+    @torch.no_grad()
+    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)
+        # Ledger P6 / paper Sec. 3.3: one extra hierarchy update on the final weights.
+        final_stats = self.model.hierarchy_update()
+        tensors: dict[str, torch.Tensor] = {}
+        for l, level in enumerate(self.model.levels):
+            tensors[f"level{l}.W_enc"] = level.W_enc.detach().float().cpu().contiguous()
+            tensors[f"level{l}.W_dec"] = level.W_dec.detach().float().cpu().contiguous()
+            tensors[f"level{l}.b_enc"] = level.b_enc.detach().float().cpu().contiguous()
+            tensors[f"level{l}.b_dec"] = level.b_dec.detach().float().cpu().contiguous()
+            tensors[f"level{l}.threshold"] = level.threshold.detach().float().cpu().contiguous()
+        for b in range(self.model.n_boundaries):
+            tensors[f"parent_idx_{b}"] = self.model.parent_idx(b).cpu().contiguous()
+        tensors["tokens_since_fired"] = self.model.tokens_since_fired.cpu().contiguous()
+        save_file(tensors, str(out_dir / WEIGHTS_FILE))
+        parents = self.model.parents_per_boundary()
+        (out_dir / HIERARCHY_FILE).write_text(json.dumps({"widths": list(self.cfg.widths), "parents": parents}))
+        dead = self.model.dead_masks()
+        (out_dir / CONFIG_FILE).write_text(json.dumps({
+            "arch": "hsae_forest",
+            "paper": "Luo et al., arXiv:2602.11881",
+            "config": self.cfg.to_dict(),
+            "n_latents": self.n_latents,
+            "level_offsets": self.model.offsets,
+            "scale": float(self.model.scale),
+            "scale_frozen": bool(self.model.scale_frozen),
+            "scale_estimate_tokens_seen": int(self.model.norm_count),
+            "last_step": self.last_step,
+            "n_hierarchy_updates": int(self.model.n_hierarchy_updates),
+            "final_hierarchy_stats": final_stats,
+            "final_lambdas": self.model.lambdas.tolist(),
+            "final_l0_ema": self.model.l0_ema.tolist(),
+            "frac_dead_per_level": [float(m.float().mean()) for m in dead],
+            "weights_file": WEIGHTS_FILE,
+            "hierarchy_file": HIERARCHY_FILE,
+        }, indent=1))
+
+
+def build_trainer(d_in: int, overrides: dict[str, Any], device: str = "cuda") -> HSAEForestTrainer:
+    cfg = _make_config(d_in, dict(overrides or {}))
+    return HSAEForestTrainer(cfg, device=device)
+
+
+# --------------------------------------------------------------------------- export consumer
+class ExportedForest:
+    """``ExportedSAE`` over the flattened (level 0, level 1, level 2) index space."""
+
+    def __init__(self, cfg: ForestConfig, tensors: dict[str, torch.Tensor], parents: list[list[int]], scale: float, device: str) -> None:
+        self.cfg = cfg
+        self.widths = list(cfg.widths)
+        self.scale = float(scale)
+        self.device = torch.device(device)
+        n_levels = len(self.widths)
+        self.W_enc = [tensors[f"level{l}.W_enc"].to(self.device) for l in range(n_levels)]
+        self.b_enc = [tensors[f"level{l}.b_enc"].to(self.device) for l in range(n_levels)]
+        self.b_dec = [tensors[f"level{l}.b_dec"].to(self.device) for l in range(n_levels)]
+        self.threshold = [tensors[f"level{l}.threshold"].to(self.device) for l in range(n_levels)]
+        w_dec_scaled = torch.cat([tensors[f"level{l}.W_dec"] for l in range(n_levels)]).to(self.device)
+        # Decoder directions in RAW activation units: x ~= sum_j a_j (d_j / alpha) + b_dec / alpha.
+        self.W_dec = w_dec_scaled / self.scale
+        self.parents = parents
+        self.n_latents = sum(self.widths)
+        self.offsets = [0]
+        for w in self.widths:
+            self.offsets.append(self.offsets[-1] + w)
+
+    @torch.no_grad()
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        xs = x.to(self.device, dtype=torch.float32) * self.scale
+        outs = []
+        for l in range(len(self.widths)):
+            pre = (xs - self.b_dec[l]) @ self.W_enc[l] + self.b_enc[l]
+            outs.append(pre * (pre > self.threshold[l]))
+        return torch.cat(outs, dim=-1)
+
+    def native_ancestors(self) -> dict[int, set[int]]:
+        out: dict[int, set[int]] = {}
+        for b, plist in enumerate(self.parents):
+            child_off, parent_off = self.offsets[b + 1], self.offsets[b]
+            for j, p in enumerate(plist):
+                if p >= 0:
+                    out[child_off + j] = {parent_off + p}
+        return out
+
+    def latent_level(self) -> list[int]:
+        return [l for l, w in enumerate(self.widths) for _ in range(w)]
+
+
+def load_exported(export_dir: Path, device: str = "cuda") -> ExportedForest:
+    from safetensors.torch import load_file
+
+    export_dir = Path(export_dir)
+    meta = json.loads((export_dir / CONFIG_FILE).read_text())
+    cfg = ForestConfig(**meta["config"])
+    tensors = load_file(str(export_dir / meta.get("weights_file", WEIGHTS_FILE)))
+    parents = json.loads((export_dir / meta.get("hierarchy_file", HIERARCHY_FILE)).read_text())["parents"]
+    return ExportedForest(cfg, tensors, parents, meta["scale"], device)

gemma-2-2b/layer-12/hsae_forest/hsae_forest_32k_seed0_v3/summary.json

@@ -0,0 +1,8 @@
+{
+ "run_name": "hsae_forest_32k_seed0_v3",
+ "arch": "hsae_forest",
+ "steps": 122070,
+ "tokens": 499998720,
+ "seconds": 15422.808210134506,
+ "export": "/vol/train/hsae_forest_32k_seed0_v3/export"
+}