---
license: apache-2.0
tags:
  - sparse-autoencoder
  - mechanistic-interpretability
  - probing
  - gemma
  - llama
language:
  - en
base_model:
  - google/gemma-2-2b
  - meta-llama/Llama-3.1-8B
---

# 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

```bash
pip install torch sae-lens>=5.0.0
```

### Loading a Gemma (JumpReLU) checkpoint

```python
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

```python
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)

```python
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](https://huggingface.co/google/gemma-scope-2b-pt-res) |
| LLaMA-3.1-8B | LlamaScope | TopK k=50, 32 768 features | [fnlp/LlamaScope](https://huggingface.co/fnlp/LlamaScope) |

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

---

## Training Data

| Domain | Training set | Size |
|--------|-------------|------|
| Toxicity | [ToxiGen](https://huggingface.co/datasets/skg/toxigen-data) (balanced) | 100k |
| Sentiment | [Sentiment140](https://huggingface.co/datasets/stanfordnlp/sentiment140) (balanced) | 100k |

---

## Citation

```bibtex
@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}
}
```