iarcuschin/gemma-scope-2b-pt-res-canonical-l12-e2e

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

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.

As far as we could determine at release time, this is the first publicly available end-to-end / KL-finetuned SAE built on the GemmaScope (JumpReLU) suite; prior released e2e SAEs cover GPT-2 small / TinyStories (Apollo) or SAEBench TopK/ReLU SAEs (Karvonen), not GemmaScope.

What it is

• Base SAE: gemma-scope-2b-pt-res-canonical :: layer_12/width_16k/canonical (Lieberum et al. 2024, Gemma Scope (arXiv:2408.05147))
• Base model: google/gemma-2-2b, residual stream, layer 12
• Architecture: JumpReLU (unchanged); d_in=2304, 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: monology/pile-uncopyrighted, ~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 monology/pile-uncopyrighted.

Metric	Baseline (GemmaScope)	Fine-tuned
delta CE	0.4567	0.0543
clean CE	1.9149	1.9149
spliced CE	2.3716	1.9692
mean L0	89.5970	103.3570

delta-CE reduction: 88.1%.

Caveat — sparsity drift. The JumpReLU threshold is frozen, but the fine-tuned encoder shifts more pre-activations above it, so L0 rises (see table). Part of the fidelity gain is the SAE firing more features, not purely more functionally faithful directions. For a strict same-sparsity comparison, compare against a GemmaScope SAE at the matched L0. Evaluated only on the training distribution and a single layer.

Usage

from huggingface_hub import snapshot_download
from sae_lens import SAE

path = snapshot_download(repo_id="iarcuschin/gemma-scope-2b-pt-res-canonical-l12-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 a00635afe5ed8587493dc6ec2d516934c7bf992d
• 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.