gemma-scope-9b-it-res-canonical-l20-e2e

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

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-9b-it-res-canonical :: layer_20/width_16k/canonical (Lieberum et al. 2024, Gemma Scope (arXiv:2408.05147))
• Base model: google/gemma-2-9b-it, residual stream, layer 20
• Architecture: JumpReLU (unchanged); d_in=3584, 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.

delta-CE reduction: 96.2%.

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-9b-it-res-canonical-l20-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 ea58cbf8cec21573f54c9563e73f429d24e1bec5
• 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.