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.
On 30k held-out OpenWebText tokens, an activation-correlation alignment
with
legible-weights/sae-mdlm-owt-l6-v0.1
shows:
Full writeup, methodology, and reproducible pipeline:
github.com/legibleweights/diffusion-vs-ar-saes1import torch
2from huggingface_hub import hf_hub_download
3from safetensors.torch import load_file
4
5class TopKSAE(torch.nn.Module):
6 def __init__(self, d_in=768, d_hidden=12288, k=32):
7 super().__init__()
8 self.k = k
9 self.encoder = torch.nn.Linear(d_in, d_hidden, bias=True)
10 self.decoder = torch.nn.Linear(d_hidden, d_in, bias=False)
11 self.pre_bias = torch.nn.Parameter(torch.zeros(d_in))
12
13 def forward(self, x):
14 pre = self.encoder(x - self.pre_bias)
15 vals, idx = pre.topk(self.k, dim=-1)
16 vals = torch.relu(vals)
17 acts = torch.zeros_like(pre)
18 acts.scatter_(-1, idx, vals)
19 return self.decoder(acts) + self.pre_bias, acts
20
21path = hf_hub_download("legible-weights/sae-gpt2-small-l6-v0.1", "sae.safetensors")
22sae = TopKSAE()
23sae.load_state_dict(load_file(path))
MIT.
