chanind-goodfire/baseline-saes

Baseline SAEs

BatchTopK and Matryoshka BatchTopK sparse autoencoders trained on the output of Gemma-2-2b decoder block model.layers.12, using streaming monology/pile-uncopyrighted and SAELens 6.51.1.

Exports use the following directory structure. Only completed models are uploaded; available token budgets can be found in the repository's file tree.

gemma-2-2b/layer-12/
  {50M,500M}/
    {batchtopk,matryoshka-batchtopk}/
      width-16384/k100/
        cfg.json
        sae_weights.safetensors
        sparsity.safetensors
        training_config.json

Both training architectures have 16,384 latents and BatchTopK k=100. Matryoshka uses nested widths [4096, 16384]. The activation dimension is 2304. The inference threshold is learned during training, so inference L0 can differ from the training k. SAELens exports these BatchTopK-trained models as JumpReLU inference SAEs; their training architectures are recorded in cfg.json metadata. This repository currently publishes only the BatchTopK and Matryoshka training variants.

Training uses seed 0, context length 1024, batches of 4096 activation tokens, Hugging Face model execution with bfloat16 autocasting, and float32 SAE parameters. The learning rate is 3e-4. The 500M runs decay it over the last 20% of training; the 50M runs use a constant learning rate. Actual token counts can exceed the named budget by fewer than 4096 tokens because updates use whole batches. Each export's training_config.json records its actual count.

Loading

from pathlib import Path

from huggingface_hub import snapshot_download
from sae_lens import SAE

subdir = "gemma-2-2b/layer-12/500M/batchtopk/width-16384/k100"
root = snapshot_download(
    "chanind-goodfire/baseline-saes",
    allow_patterns=f"{subdir}/*",
)
sae = SAE.load_from_disk(Path(root) / subdir, device="cpu")

Use matryoshka-batchtopk for the corresponding Matryoshka export, or replace 500M with 50M when that completed budget is available.