juiceb0xc0de/smollm2-135m-SAE

SmolLM2-135M Layer-Complete Top-K SAE Suite

30 sparse autoencoders — one per residual-stream layer of `HuggingFaceTB/SmolLM2-135M`.

Each SAE uses Top-K hard sparsity (k=50), a dictionary of 18,432 features (32× expansion over d_model=576), and an always-on AuxK auxiliary term (k_aux=25). Identical hyperparameters across all 30 layers, seed 0.

What to expect

• —Mean explained variance: 0.949 across layers
• —Best reconstruction at mid-layers (L11, EV 0.989)
• —Weakest at the final layer (L29, EV 0.873)
• —Mean L0 ≈ 46.4 (target k = 50)
• —Zero dead features across all 30 SAEs
• —Total training: ~95.8k optimizer steps / ~3.14B tokens

What this is

A training-healthy, uniform SAE atlas for SmolLM2-135M. A reasonable starting point for layer comparison, feature-steering experiments, and cross-layer analysis. Every layer was trained under identical hyperparameters, so per-layer SAEs are directly comparable.

What this is NOT

• —Not a proven set of interpretable features. No downstream validation (CE/KL recovery under activation patching, automated interpretability scoring, or steering evaluation) has been run.
• —Not fully converged at every layer. The deepest layers (L27–L29) early-stopped while EV was still climbing and may benefit from longer training — treat their reported EV as a lower bound, not a converged value.
• —Not an unassisted zero-dead-feature result. The AuxK auxiliary term was active throughout training (effective_aux_k = 25); the discrete resampler never fired, but that is a consequence of AuxK keeping latents alive, not independent evidence that no rescue mechanism was needed. Top-K hard sparsity, the sparsity-warmup schedule, and AuxK were all present on every run — this suite does not isolate which one (or what combination) is responsible.
• —Not assessed for decoder specialization vs. redundancy. EV and dead% do not distinguish a well-specialized dictionary from one containing duplicate or near-initialization features. Decoder drift (‖Wdec − Wdec_init‖) and pairwise decoder cosine similarity were not logged for this run.

Training configuration

Hardware

L00–L04 trained on NVIDIA L40S; L05–L29 trained on NVIDIA H100 80GB HBM3 (RunPod). Throughput ranged 244K–812K tokens/sec; the step-up partway through the sweep is consistent with this hardware tier change but is inferred from throughput data, not independently confirmed in logs.

Per-layer results

Reproduction

model_id   = HuggingFaceTB/SmolLM2-135M
layers     = 0..29   (one SAE each)
k          = 50
n_features = 18432   (32x expansion over d_model=576)
k_aux      = 25
lr         = 2e-4    (warmup -> decay)
batch      = 32768 tokens/step, accum=1
n_steps    = 15000   (EV-plateau early stop)
seed       = 0

Logs

Full per-step metrics and run logs: ricks-holmberg-juiceb0xc0de/smollm2-sae on W&B.

License

MIT