Upload folder using huggingface_hub

gemma-2-2b/layer-12/hsae/hsae_32k_k15_seed0/DECISIONS.md

@@ -0,0 +1,153 @@
+# H-SAE torch port — decisions and deviations
+
+`benchmark/methods/train/hsae.py` is a PyTorch port of the authors' JAX/Equinox
+H-SAE (Muchane, Pal & Veitch, arXiv:2506.01197), vendored byte-for-byte at
+`synth_sae_bench_baseline/third_party/hsae_muchane/moe_eqx.py`. Parity against
+that file is tested in `tests/test_train_hsae.py` (forward activations, routing
+order, loss, raw gradients, and four full optimizer steps agree to ~1e-7 on CPU
+fp32; the tests assert 1e-5). Everything below is a place where the port or the
+run configuration knowingly differs from `moe_eqx.py` or from the authors' run.
+
+## Numerics / implementation (same maths, different kernels)
+
+- **Dense instead of gathered expert matmuls.** `moe_eqx.py` gathers
+  `W_down[top_k_indices]` (`(k, s, d)` per example) and `W_up[top_k_indices]`.
+  At batch 4096, k = 75, d = 2304 that gather is 11 GB per step, so the port
+  computes `x @ W_down.reshape(m*s, d).T` for *all* experts and gathers the
+  `(B, k, s)` routed subspace codes, and decodes with a scattered `(B, m*s)`
+  code matrix times `W_up`. The top-level decode likewise scatters the routed
+  values into a `(B, m)` matrix and multiplies by `D.T` instead of gathering
+  `D[:, idx]`. Forward values and gradients are identical up to reduction order
+  (verified by the parity test); the wasted FLOPs are cheap relative to the
+  Gemma forward.
+- **`naive_top_k` via a stable descending sort.** The upstream iterative
+  argmax scan returns the k largest values in descending order with ties
+  resolved to the lowest index. `torch.sort(descending=True, stable=True)` has
+  exactly that order; the test checks it against an iterative argmax on
+  tie-heavy input and against JAX (exact index-and-order equality, including a
+  zero input row where all k routed slots are ties at 0).
+- **bf16 autocast (CUDA only) for the four big matmuls**: top encoder,
+  `W_down` projection, top decoder, `W_up` lift. Offset-ReLU, TopK, the
+  expert-level (`(B, k, a, s)`) computations, masking, losses and the
+  orthogonality penalty (a `(m, m)` matmul on unit-normalised rows) stay fp32
+  (TF32 as set by the driver). Master weights and Adam state are fp32. The
+  authors trained in fp32 on TPU/A100 (JAX default matmul precision); the
+  parity test runs with autocast disabled. `autocast_bf16=False` restores
+  exact fp32.
+- **`clip_by_global_norm`** is implemented directly as `g *= min(1, c/||g||)`
+  (optax semantics) rather than `torch.nn.utils.clip_grad_norm_`, whose
+  `1e-6` denominator fudge is absent in optax.
+- **Adam** is `torch.optim.Adam(betas=(0.9, 0.999), eps=1e-8)`, which matches
+  `optax.adam(b1=0.9)` defaults (`b2=0.999, eps=1e-8, eps_root=0`) including
+  bias correction; the LR is written into the param group every step from the
+  pure-function schedule `lr_at(step, total_steps)`. Verified by the 4-step
+  parity test at lr_peak = 1e-2 (weights agree to 4e-7 after moving by 0.2).
+- **LR schedule at step 0.** optax evaluates the linear warm-up in float32,
+  where `lr_init = 1e-11` is below the ULP of the peak, so the JAX schedule
+  yields exactly 0.0 on the first step; the port yields 1e-11. The parity test
+  allows 1e-10 absolute on the LR. (The warm-up is effectively "from 0".)
+- **Regulariser warm-up and LR schedule are indexed by the driver's 0-based
+  `step`**, matching optax's optimizer count (first update uses the step-0
+  values) and Kirill's `regularizer_warmup_frac(step)`. `decay_steps =
+  total_steps - warmup` as in the upstream `get_lr_schedule`; if the run is
+  shorter than the warm-up (e.g. the 4M-token smoke run: 976 steps) the whole
+  run is warm-up and the cosine part is never reached.
+- **Init.** `he_uniform` (uniform with variance `2/fan_in`) with fan_in = d for
+  `E` and `W_down`, s for the expert encoders; `W_down` rows and expert
+  encoder rows unit-normalised; decoders are the transposes at init. Like
+  upstream, the top decoder is *not* unit-normalised at init — it becomes so
+  after the first `normalize_decoders()` (so the loss drops at step 1 for that
+  reason alone). Seeds are torch seeds; no claim of bitwise equality with
+  `jax.random.PRNGKey`.
+- **The upstream `cross_orthogonality_penalty` quirk is reproduced.**
+  `jnp.diag(ED)` there is the 1-D diagonal, so `ED - jnp.diag(ED)` broadcasts
+  `ED[i, j] - ED[j, j]` rather than zeroing the diagonal as in the paper's
+  eq. 3.4. The port keeps the code's behaviour (the parity test would fail
+  otherwise).
+- **L1 term** is on *all* top codes (post offset-ReLU) and on all `k` routed
+  slots' raw leaky codes (unmasked), exactly as in `loss_fn` (the paper says
+  "outside the top k"; the code does not).
+- **`use_bias=False`** only (the authors' setting; the bias branch in
+  `moe_eqx.py` loads a file from their filesystem and is dead code).
+- Not ported: `tensorstore_data_generator` and the upstream `run_moe_eqx.py`
+  loop (data pipeline, wandb, FSDP, the "2 epochs" quirk). The driver supplies
+  data; the schedule is defined over the driver's `total_steps`.
+
+## Readout / export conventions (not in `moe_eqx.py`)
+
+- **Flattened ids**: `0..m-1` experts, `m + j*a + l` child `l` of expert `j`
+  (identical to Kirill's `common.py` and the authors' dashboard).
+  `native_ancestors() = {child: {expert}}`, `latent_level()` = 0 / 1.
+- **`encode(x)`** applies `x / (||x||_2 + 1e-6)` per example (the authors'
+  `_prefetch` normalisation, no centring), then returns expert activations
+  (routed top value if > 0, else 0) and child activations
+  `relu(masked low code)`. A routed-but-negative leaky child code (slope 0.01)
+  does contribute to the reconstruction during training but is reported as
+  *not firing* (Kirill's frozen firing semantics, `eval_adapter.py`). `l0` is
+  the count of these non-negative flattened activations that are > 0; it is
+  structurally at most `2k`.
+- **`W_dec`**: expert rows `D[:, j]` (unit norm); child rows
+  `W_up[j] @ D_j[:, l]` (the authors' dashboard convention, Kirill's
+  "primary" child direction; the "parent + child" sensitivity variant is not
+  exported but is `W_dec[j] + W_dec[child]`).
+- The exported model runs in fp32 (`autocast_bf16=False` on load).
+- Export = `top.safetensors` (E, D), `experts.safetensors` (W_down, W_up,
+  expert enc/dec) and `config.json`; two files so each stays under the
+  driver's 200 MB copy limit (experts ≈ 151 MB at the target size).
+
+## Configuration vs the authors' run
+
+- **Width**: m = 2048 experts x a = 15 children = 32768 flattened latents to
+  match the twin's width. The authors used 8k–32k experts with 16 or 64
+  sublatents; 15 (not 16) is forced by the 32768 total. Subspace dim s = 4 is
+  the authors' value for the 16-sublatent setting ("when the number of
+  sublatents per expert is 16, the subspace dimension is 4"); Kirill used 4 too.
+- **k = 75** (authors: 32) so that active nodes per token = up to k experts +
+  k children ≈ 150, the twin's L0. The realised L0 is reported by `step`
+  (`l0`, `l0_top`, `l0_child`) — an expert routed with a top value of 0 (fewer
+  than k codes above the offset `1/sqrt(d)`) and negative leaky child codes
+  reduce it below 150.
+- **Batch 4096 tokens** (driver; authors 32512) and **500M tokens** (authors:
+  1B vectors x 4 epochs). The LR schedule (peak 5e-4, 1000-step linear
+  warm-up, cosine to 10 %, clip 0.75) and the regulariser warm-up (1000 steps)
+  are kept at the authors' *step* counts even though a step now sees 8x fewer
+  tokens; the peak LR is not rescaled. Rationale: with Adam the per-step update
+  size is set by the LR, not the batch; the authors' values were not tuned for
+  this batch, and a sqrt-scaling argument would lower the peak to ~1.8e-4, but
+  there is no evidence for either choice here, so the defaults stay the
+  authors' and `overrides` can change `lr_peak` / `lr_warmup_steps`.
+- **Data**: Gemma-2-2b layer 12 resid_post, twin corpus (authors: layer 20,
+  Wikipedia, first 256 tokens per article). Inputs are L2-normalised per
+  example as in the authors' pipeline, so the reconstruction MSE (`mse`, sum
+  over d, mean over batch — the authors' definition) is on unit-norm vectors.
+- **Dead-feature metric** (`dead_frac`, `dead_child_frac`): fraction of
+  experts / child slots that have not fired in the last
+  `dead_window_steps = 250` steps (~1M tokens at batch 4096, the paper's dead
+  criterion). Counters start at the window, so never-fired latents count as
+  dead from step 0. There is no dead-feature intervention (the authors rely on
+  the orthogonality penalty; none in the code).
+- **Checkpoint** = model weights, Adam state, dead counters, step count; the
+  schedule is a pure function of `(step, total_steps)` so no extra state is
+  needed. Resume is bitwise on CPU; on CUDA the index backward uses atomic
+  adds, so a resumed run is deterministic only up to floating-point ordering.
+
+## Driver change
+
+- `driver.py` image: added `scikit-learn` — `benchmark/__init__.py` imports
+  `sklearn` (via `grading.py`/`locate.py`), so importing any
+  `benchmark.methods.train.*` module inside the container failed without it.
+
+## Smoke run (`smoke_hsae`, A100-80GB, 4M tokens = 976 steps, all inside the LR warm-up)
+
+- Throughput: 3,997,696 tokens in 112 s of training loop (~35.7k tokens/s
+  end-to-end including the Gemma-2-2b forward; 138 s wall including model
+  load) -> a 500M-token run is ~3.9 h on an A100-80GB.
+- Realised L0 (non-zero flattened latents per token): 150.0 at step 200,
+  146.8 at 600, 128.5 at step 976 (64.3 experts + 64.2 children). L0 is
+  capped at 2k = 150 and *decreases* during training as fewer than k top codes
+  clear the `1/sqrt(d)` offset; where it settles in a full run is unknown.
+- Loss 0.285, reconstruction MSE 0.247 on unit-norm inputs (FVU 0.247) at
+  step 976 with lr still ramping (5e-4 reached at the last step); no dead
+  experts, 0.35 % dead child slots. The orthogonality term is ~1e-6 (the
+  upstream penalty is a Frobenius norm divided by m^2 - m, so at m = 2048 it
+  is tiny by construction; reported as `ortho_raw` too).

gemma-2-2b/layer-12/hsae/hsae_32k_k15_seed0/README.md

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

@@ -0,0 +1,42 @@
+{
+ "arch": "hsae",
+ "run_name": "hsae_32k_k15_seed0",
+ "n_tokens": 500000000,
+ "batch_tokens": 4096,
+ "total_steps": 122070,
+ "overrides": {
+  "k": 8
+ },
+ "seed": 0,
+ "init_seed": 0,
+ "group": [
+  "hsae_32k_k15_seed0",
+  "tree_sae_32k_k15_seed0",
+  "hsae_forest_32k_k15_seed0"
+ ],
+ "trainer_config": {
+  "d_in": 2304,
+  "num_experts": 2048,
+  "atoms_per_subspace": 15,
+  "subspace_dim": 4,
+  "k": 8,
+  "l1_penalty": 0.001,
+  "ortho_penalty": 0.1,
+  "top_recon_coeff": 0.1,
+  "reg_warmup_steps": 1000,
+  "lr_init": 1e-11,
+  "lr_peak": 0.0005,
+  "lr_warmup_steps": 1000,
+  "lr_cosine_alpha": 0.1,
+  "norm_clip": 0.75,
+  "adam_b1": 0.9,
+  "adam_b2": 0.999,
+  "adam_eps": 1e-08,
+  "eps": 1e-06,
+  "normalize_inputs": true,
+  "autocast_bf16": true,
+  "dead_window_steps": 250,
+  "seed": 0,
+  "n_latents": 32768
+ }
+}

gemma-2-2b/layer-12/hsae/hsae_32k_k15_seed0/export/config.json

@@ -0,0 +1,42 @@
+{
+ "arch": "hsae",
+ "config": {
+  "d_in": 2304,
+  "num_experts": 2048,
+  "atoms_per_subspace": 15,
+  "subspace_dim": 4,
+  "k": 8,
+  "l1_penalty": 0.001,
+  "ortho_penalty": 0.1,
+  "top_recon_coeff": 0.1,
+  "reg_warmup_steps": 1000,
+  "lr_init": 1e-11,
+  "lr_peak": 0.0005,
+  "lr_warmup_steps": 1000,
+  "lr_cosine_alpha": 0.1,
+  "norm_clip": 0.75,
+  "adam_b1": 0.9,
+  "adam_b2": 0.999,
+  "adam_eps": 1e-08,
+  "eps": 1e-06,
+  "normalize_inputs": true,
+  "autocast_bf16": true,
+  "dead_window_steps": 250,
+  "seed": 0
+ },
+ "n_latents": 32768,
+ "n_steps_done": 122070,
+ "latent_layout": "ids 0..m-1 experts; id m + j*a + l = child l of expert j",
+ "weights": {
+  "top.safetensors": [
+   "top_encoder (m,d)",
+   "top_decoder (d,m)"
+  ],
+  "experts.safetensors": [
+   "W_down (m,s,d)",
+   "W_up (m,d,s)",
+   "expert_encoder (m,a,s)",
+   "expert_decoder (m,s,a)"
+  ]
+ }
+}

gemma-2-2b/layer-12/hsae/hsae_32k_k15_seed0/export/experts.safetensors

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

gemma-2-2b/layer-12/hsae/hsae_32k_k15_seed0/export/top.safetensors

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

gemma-2-2b/layer-12/hsae/hsae_32k_k15_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/hsae/hsae_32k_k15_seed0/model.py

@@ -0,0 +1,462 @@
+"""H-SAE (Muchane, Pal & Veitch, arXiv:2506.01197) — PyTorch trainer for ``driver.py``.
+
+A faithful torch port of the authors' JAX/Equinox implementation vendored at
+``synth_sae_bench_baseline/third_party/hsae_muchane/moe_eqx.py`` (commit
+``3fa346cb``): ``MixtureOfExperts_v2`` (top-level offset-ReLU + iterative
+TopK routing, per-expert subspace projection, TopK-1 child selection with
+parent gating), ``loss_fn`` (reconstruction + top-only reconstruction + L1 +
+cross-orthogonality with linear regulariser warm-up), the authors'
+per-example L2 input normalisation, and ``update_model`` (decoder-tangent
+gradient projection, global-norm clip 0.75, Adam(b1=0.9) with linear warm-up
+from 1e-11 and cosine decay to 10% of peak, unit-norm decoder columns after
+every step). Every deliberate deviation is listed in ``HSAE_DECISIONS.md``.
+
+Flattened latent index space (Kirill's ``common.py`` / the authors' dashboard):
+ids ``0..m-1`` are experts, id ``m + j*a + l`` is child ``l`` of expert ``j``.
+
+Parameter layout (identical to the JAX fields)::
+
+    top_encoder     E       (m, d)      top codes  = offset_relu(E x)
+    top_decoder     D       (d, m)      unit-norm columns (after each step)
+    W_down                  (m, s, d)   x_sub_j    = W_down[j] x
+    W_up                    (m, d, s)   lift back:   W_up[j] x_sub_j
+    expert_encoder  E_j     (m, a, s)   low codes  = leaky_offset_relu(E_j x_sub_j)
+    expert_decoder  D_j     (m, s, a)   unit-norm columns (after each step)
+"""
+
+from __future__ import annotations
+
+import json
+import math
+from dataclasses import asdict, dataclass, replace
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+
+UPSTREAM_EPS = 1e-6  # moe_eqx.py: eps (normalisation and orthogonality penalty)
+WEIGHT_KEYS = ("top_encoder", "top_decoder", "W_down", "W_up", "expert_encoder", "expert_decoder")
+
+
+@dataclass(frozen=True)
+class HSAEConfig:
+    d_in: int = 2304
+    num_experts: int = 2048  # m
+    atoms_per_subspace: int = 15  # a  -> m + m*a = 32768 flattened latents
+    subspace_dim: int = 4  # s  (authors: 4 for ~16 sublatents/expert, else 8)
+    k: int = 75  # top-level TopK: <= k experts + <= k children active per token
+    l1_penalty: float = 1e-3
+    ortho_penalty: float = 1e-1
+    top_recon_coeff: float = 1e-1  # the hard-coded 1E-1 * reconstruction_loss_top
+    reg_warmup_steps: int = 1000  # regularisers ramp linearly from 0
+    lr_init: float = 1e-11
+    lr_peak: float = 5e-4
+    lr_warmup_steps: int = 1000
+    lr_cosine_alpha: float = 0.1  # cosine decays to alpha * peak
+    norm_clip: float = 0.75
+    adam_b1: float = 0.9
+    adam_b2: float = 0.999
+    adam_eps: float = 1e-8
+    eps: float = UPSTREAM_EPS
+    normalize_inputs: bool = True
+    autocast_bf16: bool = True  # bf16 for the four big matmuls on CUDA only
+    dead_window_steps: int = 250  # ~1M tokens at batch 4096 (paper's dead criterion)
+    seed: int = 0
+
+    @property
+    def n_latents(self) -> int:
+        return self.num_experts * (1 + self.atoms_per_subspace)
+
+    @property
+    def offset(self) -> float:
+        return 1.0 / math.sqrt(self.d_in)
+
+
+# ---------------------------------------------------------------------------
+# Pure functions ported from moe_eqx.py
+# ---------------------------------------------------------------------------
+def normalize_inputs(x: Tensor, eps: float = UPSTREAM_EPS) -> Tensor:
+    """Authors' per-example L2 normalisation ``x / (||x||_2 + eps)`` (no centring)."""
+    return x / (x.norm(dim=-1, keepdim=True) + eps)
+
+
+def leaky_offset_relu(x: Tensor, negative_slope: float, offset: float) -> Tensor:
+    return torch.where(x >= offset, x, negative_slope * x)
+
+
+def naive_top_k(data: Tensor, k: int) -> tuple[Tensor, Tensor]:
+    """Upstream ``naive_top_k`` semantics for a ``(B, m)`` batch: the k largest
+    values in descending order, ties resolved to the lowest index (iterative
+    ``argmax``). A stable descending sort yields exactly that order."""
+    _, order = torch.sort(data, dim=1, descending=True, stable=True)
+    idx = order[:, :k]
+    return torch.gather(data, 1, idx), idx
+
+
+def cross_orthogonality_penalty(encoder: Tensor, decoder: Tensor, eps: float = UPSTREAM_EPS) -> Tensor:
+    """Byte-faithful port, including the upstream quirk that ``jnp.diag(ED)`` is
+    the 1-D diagonal broadcast over rows (``ED[i, j] - ED[j, j]``), not the
+    diagonal matrix of the paper's eq. 3.4."""
+    e = encoder.T  # (d, m)
+    d = decoder.T  # (m, d)
+    e = e / (e.norm(dim=0, keepdim=True) + eps)
+    d = d / (d.norm(dim=1, keepdim=True) + eps)
+    ed = d @ e  # (m, m)
+    dim = ed.shape[0]
+    return torch.linalg.norm(ed - torch.diagonal(ed)) / (dim**2 - dim)
+
+
+def lr_at(step: int, total_steps: int, cfg: HSAEConfig) -> float:
+    """optax ``join_schedules([linear_schedule, cosine_decay_schedule], [warmup])``
+    evaluated at the optimizer count ``step`` (0-based: first update uses lr_init).
+    ``decay_steps = total_steps - warmup`` as in upstream ``get_lr_schedule``."""
+    warmup = cfg.lr_warmup_steps
+    if warmup > 0 and step < warmup:
+        return cfg.lr_init + (cfg.lr_peak - cfg.lr_init) * (step / warmup)
+    decay_steps = max(total_steps - warmup, 1)
+    count = min(max(step - warmup, 0), decay_steps)
+    cosine = 0.5 * (1.0 + math.cos(math.pi * count / decay_steps))
+    return cfg.lr_peak * ((1.0 - cfg.lr_cosine_alpha) * cosine + cfg.lr_cosine_alpha)
+
+
+def reg_warmup_frac(step: int, cfg: HSAEConfig) -> float:
+    if cfg.reg_warmup_steps <= 0:
+        return 1.0
+    return min(1.0, step / cfg.reg_warmup_steps)
+
+
+# ---------------------------------------------------------------------------
+# Model
+# ---------------------------------------------------------------------------
+class HSAE(nn.Module):
+    def __init__(self, cfg: HSAEConfig, device: str | torch.device = "cpu") -> None:
+        super().__init__()
+        self.cfg = cfg
+        m, d, s, a = cfg.num_experts, cfg.d_in, cfg.subspace_dim, cfg.atoms_per_subspace
+        gen = torch.Generator(device="cpu").manual_seed(cfg.seed)
+
+        def he_uniform(shape: tuple[int, ...], fan_in: int) -> Tensor:
+            limit = math.sqrt(6.0 / fan_in)  # jax he_uniform: variance 2/fan_in, uniform
+            return (torch.rand(shape, generator=gen, dtype=torch.float32) * 2 - 1) * limit
+
+        top_enc = he_uniform((m, d), d)
+        w_down = he_uniform((m, s, d), d)
+        w_down = w_down / w_down.norm(dim=-1, keepdim=True)
+        exp_enc = he_uniform((m, a, s), s)
+        exp_enc = exp_enc / exp_enc.norm(dim=-1, keepdim=True)
+        # Tied transposes at init (upstream); trained as independent parameters.
+        # NOTE: like upstream, the top decoder is NOT unit-normalised at init —
+        # normalize_decoders() runs after the first optimizer step.
+        self.top_encoder = nn.Parameter(top_enc.to(device))
+        self.top_decoder = nn.Parameter(top_enc.T.contiguous().to(device))
+        self.W_down = nn.Parameter(w_down.to(device))
+        self.W_up = nn.Parameter(w_down.transpose(1, 2).contiguous().to(device))
+        self.expert_encoder = nn.Parameter(exp_enc.to(device))
+        self.expert_decoder = nn.Parameter(exp_enc.transpose(1, 2).contiguous().to(device))
+
+    # -- weights I/O -------------------------------------------------------
+    @torch.no_grad()
+    def load_arrays(self, arrays: dict[str, Any]) -> None:
+        for key in WEIGHT_KEYS:
+            getattr(self, key).copy_(torch.as_tensor(arrays[key], dtype=torch.float32))
+
+    def arrays(self) -> dict[str, Tensor]:
+        return {key: getattr(self, key).detach().float().cpu().contiguous() for key in WEIGHT_KEYS}
+
+    # -- forward -----------------------------------------------------------
+    def _autocast(self):
+        use = self.cfg.autocast_bf16 and self.top_encoder.is_cuda
+        return torch.autocast("cuda", dtype=torch.bfloat16, enabled=use)
+
+    def forward(self, x: Tensor) -> dict[str, Tensor]:
+        """``x`` is the already-normalised ``(B, d)`` float32 batch. Mirrors
+        ``vmap(encode) -> mask_codes -> vmap(decode)``; low-level quantities are
+        in routed-slot order ``(B, k, ...)`` exactly like upstream."""
+        cfg = self.cfg
+        m, s, k = cfg.num_experts, cfg.subspace_dim, cfg.k
+        B = x.shape[0]
+        offset = cfg.offset
+
+        with self._autocast():
+            pre = x @ self.top_encoder.T  # (B, m)
+            subspace_all = x @ self.W_down.reshape(m * s, -1).T  # (B, m*s)
+        pre = pre.float()
+        top_codes = leaky_offset_relu(pre, negative_slope=0.0, offset=offset)
+        top_k_values, top_k_indices = naive_top_k(top_codes, k)  # (B, k)
+
+        # expert path in fp32 (tiny tensors)
+        sub_idx = (top_k_indices[:, :, None] * s + torch.arange(s, device=x.device)).reshape(B, k * s)
+        subspace = torch.gather(subspace_all, 1, sub_idx).float().reshape(B, k, s)
+        sel_enc = self.expert_encoder[top_k_indices]  # (B, k, a, s)
+        low_pre = torch.einsum("bkas,bks->bka", sel_enc, subspace)
+        low_codes = leaky_offset_relu(low_pre, negative_slope=1e-2, offset=offset)
+
+        # mask_codes: parent gate (top value > 0) and TopK-1 child by raw argmax
+        submask = (top_k_values > 0).to(low_codes.dtype)  # (B, k)
+        child_argmax = torch.argmax(low_codes, dim=-1)  # (B, k), first index on ties
+        low_mask = F.one_hot(child_argmax, cfg.atoms_per_subspace).to(low_codes.dtype)
+        masked_low = low_codes * (low_mask * submask[:, :, None])
+
+        # decode
+        z_top = torch.zeros(B, m, device=x.device, dtype=top_codes.dtype).scatter(1, top_k_indices, top_k_values)
+        sel_dec = self.expert_decoder[top_k_indices]  # (B, k, s, a)
+        decoded_sub = torch.einsum("bksa,bka->bks", sel_dec, masked_low)  # (B, k, s)
+        z_sub = torch.zeros(B, m * s, device=x.device, dtype=decoded_sub.dtype).scatter(
+            1, sub_idx, decoded_sub.reshape(B, k * s)
+        )
+        with self._autocast():
+            x_hat_top = z_top @ self.top_decoder.T  # (B, d)
+            low_recon = z_sub @ self.W_up.permute(0, 2, 1).reshape(m * s, -1)  # (B, d)
+        x_hat_top = x_hat_top.float()
+        x_hat = x_hat_top + low_recon.float()
+        return {
+            "top_codes": top_codes,
+            "top_k_values": top_k_values,
+            "top_k_indices": top_k_indices,
+            "low_codes": low_codes,
+            "masked_low": masked_low,
+            "child_argmax": child_argmax,
+            "submask": submask,
+            "x_hat": x_hat,
+            "x_hat_top": x_hat_top,
+        }
+
+    def loss(self, x: Tensor, l1_penalty: float, ortho_penalty: float) -> tuple[Tensor, dict[str, Tensor]]:
+        """Upstream ``loss_fn`` on a normalised batch (penalties already warmed up)."""
+        cfg = self.cfg
+        out = self(x)
+        B = x.shape[0]
+        recon = ((x - out["x_hat"]) ** 2).sum(-1).mean()
+        recon_top = ((x - out["x_hat_top"]) ** 2).sum(-1).mean()
+        l1 = l1_penalty * (out["top_codes"].abs().sum() + out["low_codes"].abs().sum()) / B
+        ortho = cross_orthogonality_penalty(self.top_encoder, self.top_decoder, cfg.eps)
+        total = recon + l1 + ortho_penalty * ortho + cfg.top_recon_coeff * recon_top
+        dec_norms = self.expert_decoder.detach().norm(dim=1)
+        stats = {
+            "loss": total,
+            "reconstruction_loss": recon,
+            "reconstruction_loss_top": recon_top,
+            "total_l1_loss": l1,
+            "total_ortho_loss": ortho_penalty * ortho,
+            "ortho_raw": ortho,
+            "decoder_max_norm": dec_norms.max(),
+            "decoder_min_norm": dec_norms.min(),
+        }
+        return total, {**stats, **out}
+
+    # -- flattened latents ---------------------------------------------------
+    def flatten(self, out: dict[str, Tensor]) -> Tensor:
+        """``(B, m + m*a)`` non-negative activations: expert ids first, then
+        child ``m + j*a + l``. Children use ``relu(masked_low)`` (a selected but
+        negative leaky code does not fire; it still enters the reconstruction)."""
+        cfg = self.cfg
+        m, a = cfg.num_experts, cfg.atoms_per_subspace
+        idx = out["top_k_indices"]
+        B, k = idx.shape
+        top = torch.zeros(B, m, device=idx.device, dtype=torch.float32).scatter(
+            1, idx, out["top_k_values"] * out["submask"]
+        )
+        child = torch.zeros(B, m * a, device=idx.device, dtype=torch.float32)
+        child_idx = (idx[:, :, None] * a + torch.arange(a, device=idx.device)).reshape(B, k * a)
+        child = child.scatter(1, child_idx, out["masked_low"].clamp_min(0.0).reshape(B, k * a))
+        return torch.cat([top, child], dim=1)
+
+    # -- decoder maintenance (upstream update_model) --------------------------
+    @torch.no_grad()
+    def project_away_grads(self) -> None:
+        """Remove the component of each decoder-column gradient parallel to the
+        (normalised) column itself, so the update is tangent to the unit sphere."""
+        D, g = self.top_decoder, self.top_decoder.grad
+        if g is not None:
+            dn = D / D.norm(dim=0, keepdim=True)
+            g.sub_((g * dn).sum(0, keepdim=True) * dn)
+        Dj, gj = self.expert_decoder, self.expert_decoder.grad
+        if gj is not None:
+            dn = Dj / Dj.norm(dim=1, keepdim=True)
+            gj.sub_((gj * dn).sum(1, keepdim=True) * dn)
+
+    @torch.no_grad()
+    def normalize_decoders(self) -> None:
+        self.top_decoder.div_(self.top_decoder.norm(dim=0, keepdim=True))
+        self.expert_decoder.div_(self.expert_decoder.norm(dim=1, keepdim=True))
+
+    # -- decoder directions ---------------------------------------------------
+    @torch.no_grad()
+    def decoder_directions(self) -> Tensor:
+        """``(m + m*a, d)``: ``D[:, j]`` for experts, ``W_up[j] @ D_j[:, l]`` for
+        children (the authors' dashboard convention, mapped to activation space)."""
+        m, a = self.cfg.num_experts, self.cfg.atoms_per_subspace
+        top = self.top_decoder.T
+        child = torch.einsum("mds,msa->mad", self.W_up, self.expert_decoder).reshape(m * a, -1)
+        return torch.cat([top, child], dim=0).float().contiguous()
+
+
+# ---------------------------------------------------------------------------
+# Trainer
+# ---------------------------------------------------------------------------
+class HSAETrainer:
+    def __init__(self, cfg: HSAEConfig, device: str = "cuda") -> None:
+        self.cfg = cfg
+        self.device = torch.device(device)
+        self.d_in = cfg.d_in
+        self.n_latents = cfg.n_latents
+        self.model = HSAE(cfg, device=self.device)
+        self.opt = torch.optim.Adam(
+            self.model.parameters(), lr=cfg.lr_init, betas=(cfg.adam_b1, cfg.adam_b2), eps=cfg.adam_eps
+        )
+        w = cfg.dead_window_steps
+        self.steps_since_fired = torch.full((cfg.num_experts,), w, dtype=torch.long, device=self.device)
+        self.child_steps_since_fired = torch.full(
+            (cfg.num_experts, cfg.atoms_per_subspace), w, dtype=torch.long, device=self.device
+        )
+        self.n_steps_done = 0
+
+    def config_dict(self) -> dict[str, Any]:
+        return {**asdict(self.cfg), "n_latents": self.n_latents}
+
+    @torch.no_grad()
+    def _clip_global_norm(self) -> float:
+        """optax ``clip_by_global_norm``: scale all grads by ``min(1, c / ||g||)``."""
+        grads = [p.grad for p in self.model.parameters() if p.grad is not None]
+        norm = torch.linalg.vector_norm(torch.stack([g.norm() for g in grads]))
+        coef = torch.clamp(self.cfg.norm_clip / norm, max=1.0)
+        for g in grads:
+            g.mul_(coef)
+        return float(norm)
+
+    def step(self, x: Tensor, step: int, total_steps: int) -> dict[str, float]:
+        cfg = self.cfg
+        x = x.to(self.device, torch.float32)
+        if cfg.normalize_inputs:
+            x = normalize_inputs(x, cfg.eps)
+        frac = reg_warmup_frac(step, cfg)
+        lr = lr_at(step, total_steps, cfg)
+        for group in self.opt.param_groups:
+            group["lr"] = lr
+
+        self.model.train()
+        loss, aux = self.model.loss(x, cfg.l1_penalty * frac, cfg.ortho_penalty * frac)
+        self.opt.zero_grad(set_to_none=True)
+        loss.backward()
+        self.model.project_away_grads()
+        grad_norm = self._clip_global_norm()
+        self.opt.step()
+        self.model.normalize_decoders()
+        self.n_steps_done += 1
+
+        with torch.no_grad():
+            flat = self.model.flatten(aux)
+            m = cfg.num_experts
+            active = flat > 0
+            fired_top = active[:, :m].any(0)
+            fired_child = active[:, m:].reshape(-1, m, cfg.atoms_per_subspace).any(0)
+            self.steps_since_fired += 1
+            self.steps_since_fired[fired_top] = 0
+            self.child_steps_since_fired += 1
+            self.child_steps_since_fired[fired_child] = 0
+            metrics = {
+                "loss": float(loss),
+                "l0": float(active.sum(1).float().mean()),
+                "l0_top": float(active[:, :m].sum(1).float().mean()),
+                "l0_child": float(active[:, m:].sum(1).float().mean()),
+                "mse": float(aux["reconstruction_loss"]),  # sum over d, mean over batch (authors')
+                "mse_top": float(aux["reconstruction_loss_top"]),
+                "fvu": float(aux["reconstruction_loss"] / (x**2).sum(-1).mean()),
+                "l1": float(aux["total_l1_loss"]),
+                "ortho": float(aux["total_ortho_loss"]),
+                "ortho_raw": float(aux["ortho_raw"]),
+                "dead_frac": float((self.steps_since_fired >= cfg.dead_window_steps).float().mean()),
+                "dead_child_frac": float((self.child_steps_since_fired >= cfg.dead_window_steps).float().mean()),
+                "grad_norm": grad_norm,
+                "lr": lr,
+                "reg_frac": frac,
+            }
+        return metrics
+
+    def state_dict(self) -> dict[str, Any]:
+        return {
+            "config": asdict(self.cfg),
+            "model": self.model.state_dict(),
+            "opt": self.opt.state_dict(),
+            "steps_since_fired": self.steps_since_fired,
+            "child_steps_since_fired": self.child_steps_since_fired,
+            "n_steps_done": self.n_steps_done,
+        }
+
+    def load_state_dict(self, state: dict[str, Any]) -> None:
+        self.model.load_state_dict(state["model"])
+        self.opt.load_state_dict(state["opt"])
+        self.steps_since_fired.copy_(state["steps_since_fired"].to(self.device))
+        self.child_steps_since_fired.copy_(state["child_steps_since_fired"].to(self.device))
+        self.n_steps_done = int(state["n_steps_done"])
+
+    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)
+        arrays = self.model.arrays()
+        # two files so each stays well under the driver's 200MB copy limit
+        save_file({k: arrays[k] for k in ("top_encoder", "top_decoder")}, str(out_dir / "top.safetensors"))
+        save_file(
+            {k: arrays[k] for k in ("W_down", "W_up", "expert_encoder", "expert_decoder")},
+            str(out_dir / "experts.safetensors"),
+        )
+        (out_dir / "config.json").write_text(json.dumps({
+            "arch": "hsae",
+            "config": asdict(self.cfg),
+            "n_latents": self.n_latents,
+            "n_steps_done": self.n_steps_done,
+            "latent_layout": "ids 0..m-1 experts; id m + j*a + l = child l of expert j",
+            "weights": {
+                "top.safetensors": ["top_encoder (m,d)", "top_decoder (d,m)"],
+                "experts.safetensors": ["W_down (m,s,d)", "W_up (m,d,s)", "expert_encoder (m,a,s)", "expert_decoder (m,s,a)"],
+            },
+        }, indent=1))
+
+
+def build_trainer(d_in: int, overrides: dict, device: str = "cuda") -> HSAETrainer:
+    cfg = replace(HSAEConfig(d_in=d_in), **overrides)
+    return HSAETrainer(cfg, device=device)
+
+
+# ---------------------------------------------------------------------------
+# Exported model
+# ---------------------------------------------------------------------------
+class ExportedHSAE:
+    def __init__(self, model: HSAE) -> None:
+        self.model = model.eval()
+        self.cfg = model.cfg
+        self.n_latents = model.cfg.n_latents
+        self.W_dec = model.decoder_directions()
+
+    @torch.no_grad()
+    def encode(self, x: Tensor) -> Tensor:
+        """Raw residual activations ``(B, d)`` -> ``(B, n_latents)`` non-negative
+        flattened activations (normalisation applied exactly as in training)."""
+        x = x.to(self.W_dec.device, torch.float32)
+        if self.cfg.normalize_inputs:
+            x = normalize_inputs(x, self.cfg.eps)
+        return self.model.flatten(self.model(x))
+
+    def native_ancestors(self) -> dict[int, set[int]]:
+        m, a = self.cfg.num_experts, self.cfg.atoms_per_subspace
+        return {m + j * a + l: {j} for j in range(m) for l in range(a)}
+
+    def latent_level(self) -> list[int]:
+        m, a = self.cfg.num_experts, self.cfg.atoms_per_subspace
+        return [0] * m + [1] * (m * a)
+
+
+def load_exported(export_dir: Path, device: str = "cuda") -> ExportedHSAE:
+    from safetensors.torch import load_file
+
+    export_dir = Path(export_dir)
+    meta = json.loads((export_dir / "config.json").read_text())
+    cfg = replace(HSAEConfig(**meta["config"]), autocast_bf16=False)  # evaluate in fp32
+    model = HSAE(cfg, device=device)
+    arrays = {**load_file(str(export_dir / "top.safetensors")), **load_file(str(export_dir / "experts.safetensors"))}
+    model.load_arrays(arrays)
+    return ExportedHSAE(model)

gemma-2-2b/layer-12/hsae/hsae_32k_k15_seed0/summary.json

@@ -0,0 +1,13 @@
+{
+ "run_name": "hsae_32k_k15_seed0",
+ "arch": "hsae",
+ "steps": 122070,
+ "tokens": 499998720,
+ "seconds": 31551.99144744873,
+ "export": "/vol/train/hsae_32k_k15_seed0/export",
+ "group": [
+  "hsae_32k_k15_seed0",
+  "tree_sae_32k_k15_seed0",
+  "hsae_forest_32k_k15_seed0"
+ ]
+}