hematteo/sparse-readout-prism

Sparse Readout Prism — pretrained dictionaries of readout features

Pretrained dictionaries for Sparse Readout Prism (SRP), which decomposes a language model's readout (the unembedding matrix W_U) using only its weights. SRP factorizes the unembedding rows into a dictionary of reusable readout features and expresses any token logit or logit difference as a sum of signed feature contributions plus an explicit residual

h · W_U[token]  ≈  base  +  Σ_i z_i (h · d_i)  +  residual

Fidelity diagnostics reported with every decomposition measure how well the fitted rows stand in for the originals. The dictionaries here are trained on the rows of the final readout W_U, which distinguishes them from residual-stream SAEs trained on activations.

• Code: https://github.com/hematteo/sparse-readout-prism
• Paper: Sparse Readout Prism: Explaining Logit-Lens Scores in Features Instead of Tokens, arXiv:2609.01936, https://arxiv.org/abs/2609.01936

Operating points

Most base models ship two dictionaries, a high-fidelity point (k256, 32× width, k = 256) and a strict-budget point (k128, 16× width — 8× for Qwen3.5-9B — k = 128). Qwen3.5-9B additionally ships a 16×/k256 capacity point, so it has three. The exact width of each is in the width column below.

Checkpoints

Layout: <model>/<operating_point>/checkpoint.pt. A machine-readable copy of the table below, with file sizes and each file's layout, is manifest.json.

Path	base model	width	d_features	k	rowEV	top-1	KL (bits)
qwen3.5-0.8b/k128_16x	Qwen/Qwen3.5-0.8B	16×	16384	128	0.760	0.844	0.277
qwen3.5-0.8b/k256_32x	Qwen/Qwen3.5-0.8B	32×	32768	256	0.877	0.891	0.135
qwen3.5-2b/k128_16x	Qwen/Qwen3.5-2B	16×	32768	128	0.712	0.858	0.261
qwen3.5-2b/k256_32x	Qwen/Qwen3.5-2B	32×	65536	256	0.847	0.887	0.136
qwen3.5-9b/k128_8x	Qwen/Qwen3.5-9B	8×	32768	128	0.621	0.846	0.296
qwen3.5-9b/k256_16x	Qwen/Qwen3.5-9B	16×	65536	256	0.761	0.874	0.167
qwen3.5-9b/k256_32x	Qwen/Qwen3.5-9B	32×	131072	256	0.857	0.900	0.105
gemma-4-e2b/k128_16x	google/gemma-4-E2B-it	16×	24576	128	0.714	0.623	1.94
gemma-4-e2b/k256_32x	google/gemma-4-E2B-it	32×	49152	256	0.834	0.333	6.37
gemma-4-e4b/k128_16x	google/gemma-4-E4B-it	16×	40960	128	0.693	0.669	1.82
gemma-4-e4b/k256_32x	google/gemma-4-E4B-it	32×	81920	256	0.827	0.736	1.22
ministral-3-8b/k128_16x	mistralai/Ministral-3-8B-Base-2512	16×	65536	128	0.806	0.885	0.130
ministral-3-8b/k256_32x	mistralai/Ministral-3-8B-Base-2512	32×	131072	256	0.888	0.904	0.087
r1-distill-qwen-7b/k128_16x	deepseek-ai/DeepSeek-R1-Distill-Qwen-7B	16×	57344	128	0.709	0.695	0.777
r1-distill-qwen-7b/k256_32x	deepseek-ai/DeepSeek-R1-Distill-Qwen-7B	32×	114688	256	0.844	0.760	0.489
r1-distill-llama-8b/k128_16x	deepseek-ai/DeepSeek-R1-Distill-Llama-8B	16×	65536	128	0.796	0.725	0.536
r1-distill-llama-8b/k256_32x	deepseek-ai/DeepSeek-R1-Distill-Llama-8B	32×	131072	256	0.888	0.754	0.434

rowEV is row-centered explained variance for the reconstructed unembedding rows, top-1 is agreement between the original and reconstructed vocabulary argmaxes on held-out hidden states, and KL is the readout KL in bits. The Qwen and Gemma numbers come from the paper's selection tables (Appendices C and D), with the Gemma top-1/KL computed through the final-logit softcap (softcap-correct, see the paper). The Ministral and R1-Distill numbers are the checkpoints' held-out eval, stored in each file under metrics and reported in the paper's Appendix E.

How they were trained

TopK factorizer on the centered and row-normalized W_U rows, with the shared converged finalist recipe — 20k steps, batch 4096, AdamW lr 1e-3 (warmup then cosine), prism penalty lambda_prism = 1e-3 with a delayed linear ramp, hybrid (50% frequency / 50% uniform) row sampling, and row-seeded init. The operating point's k is the audit k used for decomposition.

Usage

from huggingface_hub import hf_hub_download
from sparse_readout_prism import load_factorizer  # uv sync (or pip install -e .) from the GitHub repo

path = hf_hub_download("hematteo/sparse-readout-prism", "qwen3.5-2b/k256_32x/checkpoint.pt")
sae = load_factorizer(path, freeze=True)   # rebuild + load_state_dict + eval, one call

Each checkpoint.pt is a weights_only=True-loadable dict. Two layouts ship, and load_factorizer handles both:

• Qwen and Gemma files hold model_state_dict and config, the full training config, with the factorizer block (architecture, k, d_features) under config.factorizer. No row statistics are embedded.
• Ministral and R1-Distill files hold model_state_dict, a top-level factorizer block, evaluation, the held-out metrics, and the preprocessing pinned at training time as row_mean, row_norms and row_token_ids.

To decompose you also need the centered and row-normalized preprocessing. Recompute it from the model's W_U with preprocess_rows(W_U), which centers on the mean of all of W_U. That is what the paper's analysis runs use, while training centered on the text-token rows only; on Qwen3.5-2B the two means differ by about 0.04% of a centered row norm. The code's --centering {live,trained} flag selects between them, and live is the paper's default. Decomposing against a different preprocessing breaks the identity. See the GitHub README quickstart for the full decomposition snippet.

Intended use & limitations

Research artifact for mechanistic interpretability of the final readout. The fidelity diagnostics reported with each decomposition (residual size, sign agreement) measure how well the fitted rows stand in for the originals on that query, and high rowEV alone does not establish that. These dictionaries describe the readout side of a logit and say nothing about why a hidden state arose (no residual-stream or circuit attribution).

Citation

@misc{he2026sparsereadoutprismexplaining,
  title         = {Sparse Readout Prism: Explaining Logit-Lens Scores in Features Instead of Tokens},
  author        = {Matteo He and William F. Shen and Xinchi Qiu and Nicholas D. Lane},
  year          = {2026},
  eprint        = {2609.01936},
  archivePrefix = {arXiv},
  primaryClass  = {cs.CL},
  url           = {https://arxiv.org/abs/2609.01936},
}

License: MIT.