You can deploy this model and start earning money today!

hematteo/
sparse-readout-prism

Name: sparse-readout-prism
Availability: InStock

sparse-readout-prism

sparse-autoencoder

mechanistic-interpretability

Views

No views yet

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

python

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

bibtex

@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.