sae-gpt2-small-l6-v0.1

SAE — GPT-2 small, layer 6 residual stream (v0.1)

A TopK sparse autoencoder for the post-block residual stream of layer 6 in openai-community/gpt2. Trained as one half of a controlled cross- paradigm comparison against legible-weights/sae-mdlm-owt-l6-v0.1 — a matched-scale SAE on the MDLM masked-diffusion LM trained on the same corpus.

The two SAEs were trained with identical hyperparameters on activations from OpenWebText. The only axis of variation is the base model's training objective (causal next-token vs. discrete masked diffusion).

Architecture

Training

Metrics

Cross-paradigm finding

On 30k held-out OpenWebText tokens, an activation-correlation alignment with legible-weights/sae-mdlm-owt-l6-v0.1 shows:

• 55 % of MDLM SAE features have a GPT-2 SAE counterpart at Pearson r > 0.30
• max correlation 0.989 (closed-class function words like which, the)
• median decoder cosine between paired features: −0.002

Full writeup, methodology, and reproducible pipeline: github.com/legibleweights/diffusion-vs-ar-saes

Usage

import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

class TopKSAE(torch.nn.Module):
    def __init__(self, d_in=768, d_hidden=12288, k=32):
        super().__init__()
        self.k = k
        self.encoder = torch.nn.Linear(d_in, d_hidden, bias=True)
        self.decoder = torch.nn.Linear(d_hidden, d_in, bias=False)
        self.pre_bias = torch.nn.Parameter(torch.zeros(d_in))

    def forward(self, x):
        pre = self.encoder(x - self.pre_bias)
        vals, idx = pre.topk(self.k, dim=-1)
        vals = torch.relu(vals)
        acts = torch.zeros_like(pre)
        acts.scatter_(-1, idx, vals)
        return self.decoder(acts) + self.pre_bias, acts

path = hf_hub_download("legible-weights/sae-gpt2-small-l6-v0.1", "sae.safetensors")
sae = TopKSAE()
sae.load_state_dict(load_file(path))

Hook point: output of model.transformer.h[6].

License

MIT.