---
library_name: pytorch
license: mit
base_model: kuleshov-group/mdlm-owt
tags:
  - sparse-autoencoder
  - interpretability
  - mechanistic-interpretability
  - sae
  - mdlm
  - diffusion-language-model
  - cross-paradigm-comparison
datasets:
  - Skylion007/openwebtext
language:
  - en
---

# 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](https://huggingface.co/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](https://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](https://github.com/legibleweights/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.