A TopK sparse autoencoder for the post-block residual stream of layer 6 in kuleshov-group/mdlm-owt — a 130 M-parameter discrete masked-diffusion language model from Sahoo et al. NeurIPS 2024.

Trained as one half of a controlled cross-paradigm comparison against legible-weights/sae-gpt2-small-l6-v0.1 — a matched-scale SAE on GPT-2 small trained on the same corpus. The two SAEs differ only in the base model's training objective (causal autoregressive vs. discrete masked diffusion).

Architecture

field	value
arch	TopK SAE (Gao et al. 2024)
d_in	768 (MDLM hidden_dim)
d_hidden	12 288 (16× expansion)
k	32
params	18.9 M

Training

field	value
dataset	`Skylion007/openwebtext` (streamed)
n_tokens	10,000,000
seq_len	512
diffusion_noise_t	0 (clean text — this MDLM has time_conditioning=False)
exclude_first_n	4
batch_size	4,096
optimizer	Adam, lr 5e-4
n_epochs	4
seed	0
hardware	1× RTX 4090
wall time	~3 min end-to-end

Metrics

metric	value
Final MSE	5.58
Final EV	0.90
Final L0	32

MSE differs from the GPT-2 counterpart (0.85) because MDLM activations have ~2× the magnitude. EV is scale-normalized and comparable.

Cross-paradigm finding

55 % of these features have a GPT-2 SAE counterpart at activation correlation r > 0.30 (median 0.33, max 0.99), despite median decoder cosine of −0.002. Top-correlated pairs are closed-class function-word features ( which, the, but, …) that fire on identical contexts across both models. Full writeup, methodology, and reproducible pipeline: github.com/legibleweights/diffusion-vs-ar-saes

Usage

python

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-mdlm-owt-l6-v0.1", "sae.safetensors")
sae = TopKSAE()
sae.load_state_dict(load_file(path))

Hook point: output of model.backbone.blocks[6]. The MDLM modeling file in the upstream repo depends on flash-attn. The diffusion-vs-ar-saes repo ships a patched version that uses standard PyTorch SDPA instead, so anyone can load and use this SAE without that build dependency.

License

MIT.