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

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_dec fine-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):

clean CE 1.075; 87.6% delta-CE reduction.

IT-template (pile) held-out split — vs the released GS2 base SAEs (16k, layer 24):

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 in metrics.json / the run's provenance.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.