Revisiting End-To-End Sparse Autoencoder Training: A Short Finetune Is All You Need
Paper • 2503.17272 • Published
End-to-end (KL-finetuned) GemmaScope-2 SAE — google/gemma-3-12b-it, layer 24 (on-policy)
A GemmaScope-2 JumpReLU residual-stream SAE made more causally faithful by a short KL+MSE ("end-to-end") fine-tune (Karvonen 2025, arXiv:2503.17272), trained on on-policy chat rollouts to match GemmaScope 2's instruction-tuned SAE training distribution. A drop-in replacement for the base GemmaScope-2 SAE with a much smaller cross-entropy gap when spliced into the model.
What it is
- Base SAE:
gemma-scope-2-12b-it-res::layer_24_width_16k_l0_medium(Gemma Scope 2; Lieberum et al. 2024, arXiv:2408.05147) - Base model:
google/gemma-3-12b-it, residual stream, layer 24 - Architecture: JumpReLU (unchanged);
d_in=3840,d_sae=16384 (~16k) - What changed:
W_enc, b_enc, W_dec, b_decfine-tuned; the JumpReLU threshold was frozen.
Method & data
Per step the frozen LM is run once with the SAE off (target logits) and once with it spliced
in; loss (KL·alpha + MSE)·0.5 with alpha = (MSE/KL).detach(). AdamW, lr 5e-5, linear
decay, decoder rows kept unit-norm, threshold frozen. ~25M tokens, ctx 1024, SAE in fp32.
Training corpus — on-policy rollouts. First-user-turn prompts from
lmsys/lmsys-chat-1m, with responses
generated on-policy by google/gemma-3-12b-it (temperature 0.7), matching GemmaScope 2's
IT-SAE recipe (real model rollouts). Only special tokens (BOS/EOS/pad) are masked from the
loss. The exact rollouts are published at
iarcuschin/gemma-3-12b-it-lmsys-onpolicy-rollouts.
Results — delta CE (increase in CE when the SAE is spliced in; lower = better)
On-policy rollout held-out split (native distribution):
| baseline (base SAE) | fine-tuned | |
|---|---|---|
| delta CE | 0.395 | 0.049 |
| mean L0 | 61 | 99 |
clean CE 1.075; 87.6% delta-CE reduction.
IT-template (pile) held-out split — vs the released GS2 base SAEs (16k, layer 24):
| SAE | mean L0 | delta CE |
|---|---|---|
| base (L0 small) | 22 | 1.106 |
| base (L0 medium) | 65 | 0.462 |
| base (L0 big) | 172 | 1.180 |
| this SAE (e2e, on-policy) | 102 | 0.024 |
clean CE 2.437; spliced ≈ clean — faithful even on this cross-distribution split, and far below every base SAE.
Caveat — sparsity drift. The threshold is frozen, so the fine-tuned encoder pushes more pre-activations above it and L0 rises (61→99). Part of the fidelity gain is higher L0, not purely better directions; for a strict same-sparsity comparison use a base SAE at the matched L0.
Usage
from huggingface_hub import snapshot_download
from sae_lens import SAE
path = snapshot_download(repo_id="iarcuschin/gemma-scope-2-12b-it-res-l24-e2e")
sae = SAE.load_from_disk(path, device="cuda")
# identical interface to the base GemmaScope-2 SAE
The raw finetuned_sae.npz (keys W_enc, b_enc, W_dec, b_dec, threshold) is included for
GemmaScope-style loading.
Provenance
- Method: Karvonen 2025, Revisiting End-To-End Sparse Autoencoder Training (arXiv:2503.17272).
- Base suite: Gemma Scope 2 (Lieberum et al. 2024, arXiv:2408.05147).
- Produced by the on-policy rollout finetune pipeline (repo commit
3985da2+ this session's rollout-corpus changes). Full config inmetrics.json/ the run'sprovenance.json.
License
Released under CC-BY-4.0, matching the GemmaScope weights it derives from; attribute Google DeepMind (GemmaScope) and cite the Karvonen method. Use of the underlying Gemma model is governed by the Gemma license. The training rollouts derive from LMSYS-Chat-1M, whose license applies to that dataset.
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support