nathandelisle/massive-saes

Massive SAEs

SAEs from Massive Activations Confound Sparse Autoencoders (Delisle and Tan).

TopK SAEs (k = 64, 4x width), one layer per model, five seeds (42, 137, 0, 1, 2), in two versions:

• standard: a normal SAE.
• bypass: trained with PC1 removed from its input. PC1 is added back to the output.

model	layer
gpt2	6
pythia-70m-deduped	3
pythia-160m	6
pythia-1.4b	12
pythia-6.9b	16
gemma-2-2b	13
gemma-2-9b	21
llama-3.1-8b	16
qwen-2.5-7b	14

Layer L means the output of block L, which is hidden_states[L + 1] in transformers.

Files

• <model>/<standard|bypass>/seed<k>/sae.safetensors: weights.
• <model>/<standard|bypass>/seed<k>/sae.pt: the same weights, in the original file the paper's numbers were computed from.
• <model>/<standard|bypass>/seed<k>/config.json: training settings and hashes.
• <model>/preprocessing/mean.npy and pca.npz: the training mean and PC1 (v1). Every SAE needs the mean. Bypass SAEs also need v1.

Usage

import numpy as np, torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

repo, model, sae = "nathandelisle/massive-saes", "gpt2", "gpt2/bypass/seed42"
get = lambda f: hf_hub_download(repo, f)
W = load_file(get(f"{sae}/sae.safetensors"))
mu = torch.tensor(np.load(get(f"{model}/preprocessing/mean.npy")), dtype=torch.float32)
v1 = torch.tensor(np.load(get(f"{model}/preprocessing/pca.npz"))["v1"], dtype=torch.float32)
v1 = v1 / v1.norm()

def reconstruct(h, bypass=True, k=64):
    x = h.float() - mu
    if bypass:
        c = x @ v1
        x = x - c[..., None] * v1
    pre = torch.relu((x - W["b_dec"]) @ W["W_enc"].T + W["b_enc"])
    z = torch.zeros_like(pre).scatter(-1, pre.topk(k).indices, pre.topk(k).values)
    out = z @ W["W_dec"] + W["b_dec"] + mu
    return out + c[..., None] * v1 if bypass else out

For a standard SAE, use bypass=False.

License

CC BY 4.0. The base models' own licenses also apply (Llama 3.1 Community License, Gemma Terms of Use).