gemma-scope-2-12b-it-res-l31-e2e

End-to-end (KL-finetuned) GemmaScope SAE — google/gemma-3-12b-it, layer 31

A GemmaScope JumpReLU sparse autoencoder made more causally faithful by a short KL+MSE ("end-to-end") fine-tune, following Karvonen 2025, Revisiting End-To-End Sparse Autoencoder Training: A Short Finetune Is All You Need (arXiv:2503.17272). It is a drop-in replacement for the corresponding GemmaScope SAE, with a lower cross-entropy gap when spliced into the model.

What it is

• Base SAE: gemma-scope-2-12b-it-res :: layer_31_width_16k_l0_medium (Lieberum et al. 2024, Gemma Scope (arXiv:2408.05147))
• Base model: google/gemma-3-12b-it, residual stream, layer 31
• Architecture: JumpReLU (unchanged); d_in=3840, d_sae=16384 (~16k)
• What changed: W_enc, b_enc, W_dec, b_dec fine-tuned; the JumpReLU threshold was frozen.

Method

Per-step, the frozen LM is run once with the SAE off (target logits) and once with the SAE spliced in; the loss is (KL * alpha + MSE) * 0.5 where alpha = (MSE / KL).detach() rescales the KL term to the MSE's magnitude. Optimizer AdamW, lr 5e-05, linear decay; decoder rows kept unit-norm.

• Data: lmsys/lmsys-chat-1m, ~25,000,000 tokens, ctx 1024
• Trained in: fp32 (SAE), bf16 (model), eager attention (for correct Gemma-2 logit soft-capping)

Results

Fidelity metric is delta CE: the increase in cross-entropy when the SAE reconstruction is spliced into the model (lower is better), measured on held-out lmsys/lmsys-chat-1m.

delta-CE: 0.1566 baseline, -0.0406 fine-tuned (125.9% reduction toward zero).

Caveat — evaluation scope. delta CE is measured on held-out samples from the training distribution (on-policy chat rollouts) at a single layer, so it can overstate fidelity relative to a broad-distribution or all-layer evaluation. L0 does not rise under this fine-tune (see table), so the gain is not sparsity drift; for a strict same-sparsity comparison, still compare against a GemmaScope 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-l31-e2e")
sae = SAE.load_from_disk(path, device="cuda")
# identical interface to the original GemmaScope SAE

The raw finetuned_sae.npz (keys: W_enc, b_enc, W_dec, b_dec, threshold) is also included for GemmaScope-style loading.

Provenance

• Method: Karvonen 2025, Revisiting End-To-End Sparse Autoencoder Training: A Short Finetune Is All You Need (arXiv:2503.17272)
• Base SAE suite: Lieberum et al. 2024, Gemma Scope (arXiv:2408.05147)
• Produced by commit f5e18356f67735898a30459761d99b563668eb06
• Full training config in metrics.json / the run's provenance.json

License

Released under CC-BY-4.0, matching the GemmaScope weights it derives from; please attribute Google DeepMind (GemmaScope) and cite the Karvonen method. Use of the underlying Gemma model is additionally governed by the Gemma license.