PaliGemma2-3B SAE Checkpoints

Architecture

• Base model: google/paligemma2-3b-pt-224 (Gemma 2B backbone, 26 layers)
• SAE width: 16,384 features, d_in=2,304
• Training data: 50,000 VQAv2 validation samples

TopK SAE (topk/)

• Activation: top-k selection, k=50
• LR: 2e-4 / sqrt(d_sae / 16384)
• Init: Gemma Scope 2B (pretrained/) or random (random/)
• Note: Architecture mismatch — Gemma Scope natively uses JumpReLU

JumpReLU SAE (jumprelu/) — Recommended

• Activation: JumpReLU with learnable per-feature threshold
• LR: 7e-5, Adam betas=(0.0, 0.999)
• Target L0: 50, bandwidth=0.001, sparsity_coeff=1.0
• Warmup: 1000-step LR warmup, 2000-step sparsity warmup
• Init: Gemma Scope 2B (pretrained/) — architecture-matched, includes threshold
• Based on: github.com/saprmarks/dictionary_learning JumpReLU trainer

Files

• {type}/{method}/pretrained_layer_{i}.pt: SAE state_dict for layer i
• {type}/intermediate/{25,50,75}pct/: Training dynamics checkpoints
• {type}/logs/metrics_{method}_layer_{i}.csv: Training metrics