# nous-sae

Sparse autoencoders for mechanistic interpretability (Phase 177.2), in PyTorch,
CPU-trainable end to end. Consumes activations captured by `@nous/interp-core`.

## SAE families

- **JumpReLU** (`jumprelu.py`, 177.2.1.1) — Rajamanoharan et al. 2024. Per-feature
  learnable threshold gating the pre-activations through `z·H(z−θ)`, true-L0
  sparsity penalty, and the paper's straight-through estimator (rectangle kernel
  of bandwidth ε) for the threshold gradient. Primary family.
- **TopK** (`topk.py`, 177.2.1.2) — Gao et al./OpenAI 2024. Exactly `k` active
  features per token (L0 = k by construction) with an `aux_k` dead-feature revival
  loss that reconstructs the residual from the top dead features (default 4×
  expansion).
- **Gated** (`gated.py`, 177.2.1.3) — Rajamanoharan et al. 2024. Decouples gating
  from magnitude via a shared `W_gate` rescaled by `exp(r_mag)`, with a
  frozen-decoder auxiliary reconstruction term that removes L1 shrinkage bias.

All share `SAEBase`: unit-norm decoder columns (with parallel-gradient removal),
`b_dec` pre-encoder bias, and FVU / L0 diagnostics.

## Training pipeline (`training.py`, 177.2.1.4)

Streams activations through a bounded reservoir shuffle, trains with Adam
(β₂=0.999), linear LR warmup, λ annealing, per-step decoder normalisation,
Anthropic-style dead-feature resampling toward worst-reconstructed inputs (with
Adam-state reset), and EMA parameter checkpoints.

## Feature analysis

- **Autointerp** (`autointerp.py`, 177.2.2.1) — top-activating context extraction,
  a real LLM-free `TokenStatisticsLabeller` plus a fail-loud `LLMFeatureLabeller`
  seam, and per-token simulate-then-score fidelity (Bills et al.) with a
  stratified held-out split.
- **Feature splitting** (`feature_splitting.py`, 177.2.2.2) — cross-scale decoder
  cosine matching that flags a parent feature splitting into ≥2 children and
  classifies canonical vs derivative features across an SAE width ladder.
- **Ablation steering** (`steering.py`, 177.2.2.3) — error-preserving feature
  edits (`r' = decode(edit(f)) + (r − decode(f))`), a `run_with_hooks`-compatible
  steering hook, and before/after effect measurement.
- **Concept dictionary** (`concept_dictionary.py`, 177.2.2.4) — feature → label →
  safety category export (toxicity, deception, CSAM-adjacent, self-harm,
  jailbreak-adjacent) with a keyword/LLM categoriser and per-category auxiliary
  supervision signals for Kuan Yin harm classifiers.

## Tests

`python3 -m pytest tests/ -q` — 29 tests: reconstruction on synthetic sparse
dictionaries (FVU < 0.05 for JumpReLU/Gated; exact L0=k for TopK), JumpReLU STE
gradient localisation, aux-loss activation on dead features, reservoir coverage,
LR/λ warmup, dead-feature resampling, driver-token recovery + held-out fidelity,
split detection + lineage, exact-contribution ablation, and concept-category
mapping + auxiliary supervision.

## Not locally actionable

177.2.1.5 (train the full Gemma-4-9B / Qwen3-8B SAE suite at Gemma-Scope scale)
and 177.2.1.6 (train SAEs on the `@iris/code` / `@kuanyin/safety` heads) require
GPU training against real model weights and are left unchecked.