serino28/probesae-checkpoints

ProbeSAE — Probe-Guided Sparse Autoencoder Checkpoints

Fine-tuned SAE checkpoints from "Addressing the Concept Faithfulness Gap through Probe-Guided Sparse Autoencoders" by Lorenzo Malandri, Fabio Mercorio and Antonio Serino, accepted to the Findings of the Association for Computational Linguistics: EMNLP 2026.

• Paper: https://openreview.net/forum?id=mthcFoe0g6
• Code: https://github.com/serino28/probeSAE

ProbeSAE fine-tunes a pre-trained Sparse Autoencoder (SAE) with a joint loss:

L = MSE(x, x̂) + λ · ||z||₀ + β · BCE(σ(wᵀz + b), y)

where z are the SAE sparse activations, λ controls sparsity (enforced via JumpReLU threshold or TopK selection), and w is a learned sparse probe head.
This directly reduces the concept faithfulness gap — the performance delta between a dense linear probe and a sparse probe built on SAE features.

---

Repository Contents

toxicity/
  gemma/    # Gemma-2-2B JumpReLU SAE, layers 13/16/24/25, toxicity domain
  llama/    # LLaMA-3.1-8B TopK SAE, layers 16/17/18/19, toxicity domain
sentiment/
  gemma/    # Gemma-2-2B JumpReLU SAE, layers 22/23/24/25, sentiment domain
  llama/    # LLaMA-3.1-8B TopK SAE, layers 16/17/18/19, sentiment domain

48 checkpoints total (2 families × 2 domains × 4 layers × 3 conditions: ProbeSAE, the β=0 ablation, and the G-SAE baseline).
Each file is ~290 MB (Gemma) or ~1 GB (LLaMA).

---

Naming Convention

{family}_L{layer}_B{beta}_N{samples}[_E{epochs}].pt   # ProbeSAE and the β=0 ablation
{family}_L{layer}_GSAE.pt                             # G-SAE baseline

Field	Meaning
family	gemma or llama
L{layer}	Residual-stream layer index
B{beta}	Classification loss weight β
N{samples}	Fine-tuning samples used (always 100000)
E{epochs}	Number of training epochs (omitted in older files = 5 epochs)

Condition mapping

β value	Condition	Description
B0.0	β=0 (ablation)	MSE only — fine-tuned without classification signal
B1.0	ProbeSAE	Optimal β for toxicity (Gemma and LLaMA) and sentiment (LLaMA)
B10.0	ProbeSAE	Optimal β for LLaMA sentiment (L16)
B20.0	ProbeSAE	Optimal β for Gemma sentiment (L22, L24, L25)
B10.0	ProbeSAE	Optimal β for Gemma sentiment (L23)
(no β field)	G-SAE	Concept-conditioned baseline (Härle et al., 2026), retrained here on the same base SAEs, data and layers

β values were selected via grid search (β ∈ {1, 5, 10, 15, 20, 50}) on a held-out validation split.

---

Checkpoint Format

The layout differs by family. Every file carries args (training hyperparameters) and metrics (in-distribution test-set evaluation); the weights are stored as:

Family	Config key	Weights key	Weight names
Gemma (JumpReLU)	sae_cfg	model_state	sae.W_enc, sae.b_enc, sae.W_dec, sae.b_dec, sae.threshold
LLaMA (TopK)	cfg	state_dict	W_enc, b_enc, W_dec, b_dec

Shapes: W_enc [d_in, d_sae], b_enc [d_sae], W_dec [d_sae, d_in], b_dec [d_in], threshold [d_sae]. Gemma: d_in=2304, d_sae=16384. LLaMA: d_in=4096, d_sae=32768.

The trained sparse probe head (sparse_probe.weight [1, d_sae], sparse_probe.bias [1]) is stored only in the Gemma ProbeSAE files. The LLaMA checkpoints keep SAE weights only, and the G-SAE files have no probe head by construction. This costs nothing in practice: every evaluation in the paper refits the probe on the SAE features rather than reusing the training head.

probesae/models/load_sae.py in the code repository handles all of these formats.

metrics keys

Key	Description
auroc	Sparse probe AUROC on fine-tuned SAE features (in-distribution test set)
auroc_base_sparse	Sparse probe AUROC on base (pre-fine-tuning) SAE features
auroc_dense	Dense linear probe AUROC (upper bound)
r2	Reconstruction R² (measures how well the SAE reconstructs activations)
mse	Reconstruction MSE
l0	Mean number of active features per token

---

Performance Summary (in-distribution test sets)

Toxicity domain (ToxiGen test set)

Model	Layer	Condition	AUROC ↑	AUROC base ↑	AUROC dense ↑	L0	R² ↑
Gemma-2-2B	L13	ProbeSAE (β=1)	0.867	0.813	0.806	92.2	0.894
Gemma-2-2B	L16	ProbeSAE (β=1)	0.847	0.809	0.803	75.8	0.906
LLaMA-3.1-8B	L17	ProbeSAE (β=1)	0.916	0.830	0.913	50.0	0.930

Sentiment domain (Sentiment140 test set)

Model	Layer	Condition	AUROC ↑	AUROC base ↑	AUROC dense ↑	L0	R² ↑
Gemma-2-2B	L24	ProbeSAE (β=20)	0.827	0.790	0.823	38.9	0.934
LLaMA-3.1-8B	L17	ProbeSAE (β=1)	0.879	0.712	0.888	50.0	0.956

---

How to Load

Dependencies

pip install torch sae-lens>=5.0.0

Loading a Gemma (JumpReLU) checkpoint

import torch
from sae_lens import SAEConfig
from sae_lens.saes.jumprelu_sae import JumpReLUSAE

def load_gemma_probesae(path: str, device: str = "cpu") -> JumpReLUSAE:
    ckpt = torch.load(path, map_location=device, weights_only=False)

    sae_cfg = {**ckpt["sae_cfg"], "architecture": "jumprelu"}
    sae = JumpReLUSAE(SAEConfig.from_dict(sae_cfg))

    # Load only the SAE weights (exclude sparse_probe.*)
    sae_state = {
        k.replace("sae.", ""): v
        for k, v in ckpt["model_state"].items()
        if k.startswith("sae.")
    }
    sae.load_state_dict(sae_state, strict=False)
    sae.to(device=device, dtype=torch.float32)
    sae.eval()
    return sae

sae = load_gemma_probesae("toxicity/gemma/gemma_L16_B1.0_N100000.pt")
print("Stored metrics:", torch.load("toxicity/gemma/gemma_L16_B1.0_N100000.pt",
                                    map_location="cpu", weights_only=False)["metrics"])

Loading a LLaMA (TopK) checkpoint

import torch
from copy import deepcopy

def load_llama_probesae(path: str, base_sae, device: str = "cpu"):
    """
    base_sae: a loaded LlamaScope TopKSAEWrapper (see probesae/models/topk_sae.py).
    The fine-tuned checkpoint shares the same cfg as the base SAE.
    """
    ckpt = torch.load(path, map_location=device, weights_only=False)
    ft_sae = deepcopy(base_sae)
    ft_sae.load_state_dict(ckpt["state_dict"], strict=True)
    ft_sae.to(dtype=torch.float32)
    ft_sae.eval()
    return ft_sae

Accessing the sparse probe head (Gemma ProbeSAE files only)

ckpt = torch.load("toxicity/gemma/gemma_L16_B1.0_N100000.pt",
                  map_location="cpu", weights_only=False)

probe_weight = ckpt["model_state"]["sparse_probe.weight"]  # [1, d_sae]
probe_bias   = ckpt["model_state"]["sparse_probe.bias"]    # [1]

# Classify a batch of residual-stream activations x: [B, d_in]
z = sae.encode(x)                               # [B, d_sae]
logit = (z @ probe_weight.T) + probe_bias       # [B, 1]
prob  = torch.sigmoid(logit).squeeze(-1)        # [B]

---

Base SAEs

These checkpoints fine-tune the following publicly available base SAEs:

Model	SAE	Architecture	Source
Gemma-2-2B	GemmaScope	JumpReLU, 16 384 features	google/gemma-scope-2b-pt-res
LLaMA-3.1-8B	LlamaScope	TopK k=50, 32 768 features	fnlp/LlamaScope

The base SAE weights are not included in this repository; download them separately.

---

Training Data

Domain	Training set	Size
Toxicity	ToxiGen (balanced)	100k
Sentiment	Sentiment140 (balanced)	100k

---

Citation

@inproceedings{probesae2026,
  title     = {Addressing the Concept Faithfulness Gap through Probe-Guided Sparse Autoencoders},
  author    = {Malandri, Lorenzo and Mercorio, Fabio and Serino, Antonio},
  booktitle = {Findings of the Association for Computational Linguistics: EMNLP 2026},
  year      = {2026},
  url       = {https://openreview.net/forum?id=mthcFoe0g6}
}