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 (‖W_dec − W_dec_init‖) and pairwise decoder cosine similarity were not logged for this run.

Training configuration

Hyperparameter	Value
Base model	HuggingFaceTB/SmolLM2-135M (30 layers, d_model = 576)
Hook point	residual stream, per layer L0–L29
Active latents (k)	50
Dictionary size	18,432 features (32× expansion)
AuxK	k_aux = 25, active throughout
Optimizer LR	2e-4, warmup → decay
Batch size	32,768 tokens/step (microbatch = 32,768, accum = 1)
Max steps	15,000 (EV-plateau early stopping enabled)
Seed	0

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

Layer	EV	Mean L0	Dead %	Stop step
L00	0.953	45.78	0	4751
L01	0.955	45.22	0	5000
L02	0.964	44.39	0	4501
L03	0.949	48.54	0	4001
L04	0.936	49.37	0	4001
L05	0.936	50.10	0	5000
L06	0.930	45.74	0	3751
L07	0.934	45.35	0	3751
L08	0.943	45.56	0	4251
L09	0.939	45.61	0	3751
L10	0.940	43.09	0	4751
L11	0.989	48.03	0	2251
L12	0.986	47.90	0	2001
L13	0.985	47.49	0	2001
L14	0.984	45.63	0	2001
L15	0.983	46.30	0	2001
L16	0.980	48.27	0	2251
L17	0.976	45.93	0	2001
L18	0.974	45.08	0	2001
L19	0.971	46.30	0	2251
L20	0.967	46.15	0	2251
L21	0.960	45.90	0	2501
L22	0.951	46.12	0	2501
L23	0.932	46.27	0	2751
L24	0.936	46.01	0	3001
L25	0.928	47.46	0	3501
L26	0.911	46.45	0	3501
L27	0.911	46.89	0	3001
L28	0.889	46.79	0	3251
L29	0.873	44.81	0	3251

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