Sparse Transformer
Weights for the paper "Sparse Transformer: a Controlled Testbed for Mechanistic Interpretability".
Code: github.com/AntonKorznikov/SparseTransformer
The Sparse Transformer (ST) is a 5-layer decoder-only language model trained on TinyStories whose MLP layers are wide and exactly K-sparse:
MLP(x) = V_out · TopK(ReLU(V_in x))
Only K out of N = M · 256 neurons fire per token, so every MLP output is a sum of K known value vectors. The neurons are highly monosemantic (AutoInterp 0.866 for M=32, K=8 vs 0.605 for a loss-matched dense model), which makes the ST a known reference basis for evaluating SAEs, Transcoders and circuit discovery on natural language.
Models
This repository contains all 35 models of the grid
- width multiplier M ∈ {1, 2, 4, 8, 16, 32, 64} (N = 256·M MLP neurons per layer)
- sparsity K ∈ {8, 16, 32, 64, None} (
None= no TopK, dense ReLU MLP baseline)
each with 11 checkpoints (initialisation and every 10% of training):
main_mlp_mult_{M}_K_{K}/ckpt_iter{0,30000,60000,...,300000}.pt
The primary model used throughout the paper is main_mlp_mult_32_K_8 (8192 neurons per layer, 8 active).
| Layers / heads / d_model | 5 / 8 / 256 |
| Context length | 256 |
| Tokenizer | GPT-2 (tiktoken, vocab 50257) |
| Positional encoding / norm | RoPE / RMSNorm |
| Embeddings | untied |
| Training | 300k steps, batch 64 (~4.9B tokens), AdamW, cosine LR 3e-4 → 1e-5, dead-neuron resampling every 1k steps |
Final validation loss
| M \ K | 8 | 16 | 32 | 64 | None (dense) |
|---|---|---|---|---|---|
| 1 | 1.627 | 1.611 | 1.617 | 1.618 | 1.623 |
| 2 | 1.591 | 1.572 | 1.573 | 1.564 | 1.570 |
| 4 | 1.540 | 1.527 | 1.522 | 1.513 | 1.517 |
| 8 | 1.508 | 1.490 | 1.474 | 1.463 | 1.454 |
| 16 | 1.475 | 1.445 | 1.420 | 1.403 | 1.398 |
| 32 | 1.466 | 1.419 | 1.390 | 1.376 | 1.340 |
| 64 | 1.470 | 1.426 | 1.384 | 1.352 | 1.305 |
Usage
The model class lives in the code repository
(model_transformer.py).
import torch
from huggingface_hub import hf_hub_download
from model_transformer import SparseTransformerLM, TransformerConfig
path = hf_hub_download("AntonKorznikov/SparseTransformer", "main_mlp_mult_32_K_8/ckpt_iter300000.pt")
ckpt = torch.load(path, map_location="cpu", weights_only=False)
model = SparseTransformerLM(TransformerConfig(**ckpt["model_cfg"]))
model.load_state_dict(ckpt["model"])
model.eval()
# MLP neuron i in layer l:
# key k_i = model.blocks[l].mlp.fc1.weight[i] (input pattern)
# value v_i = model.blocks[l].mlp.fc2.weight[:, i] (vector written to the residual stream)
Each checkpoint is a dict with model (state dict), model_cfg (architecture), train_cfg,
iter and best_val.
To download a whole model folder:
huggingface-cli download AntonKorznikov/SparseTransformer --include "main_mlp_mult_32_K_8/*" --local-dir out
Citation
@article{korznikov2026sparsetransformer,
title = {Sparse Transformer: a Controlled Testbed for Mechanistic Interpretability},
author = {Korznikov, Anton and others},
year = {2026}
}