sae-mdlm-owt-l6-v0.1

SAE — MDLM-OWT, layer 6 residual stream (v0.1)

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

Training

Metrics

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

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.