Preprint
Sparse Readout Prism
Explaining Logit-Lens Scores in Features Instead of Tokens
Department of Computer Science and Technology, University of Cambridge
A lens reading reflects both the hidden state and the readout. SRP explains token scores as sums of sparse readout features, using only the readout's weights.
Abstract
A language model's prediction of its next token develops across layers, and lens methods track this process by decoding intermediate hidden states into tokens. But a lens reading reflects both the hidden state and the readout (the unembedding matrix) used to decode it. Many lenses are fit on a corpus, and we show that two lenses differing only in their fitting corpus can report different tokens for the same hidden states. We call this dependence corpus conditionality. To examine readout structure independently of the fitting corpus, we introduce Sparse Readout Prism (SRP), which decomposes the readout using only its weights and expresses any token logit or logit difference as a sum of contributions from sparse readout features. This reveals readout features as a new unit of analysis for lens readings, exposing structure that token identities can obscure and enabling comparisons across tokens, contexts, layers, and lenses. Replacing the original readout with SRP's sparse approximation reconstructs 8.9–17.3 percentage points more of the tested logit differences than the strongest of six baselines built on geometric relations among readout rows. Ablating features shifts logit differences in proportion to their SRP contributions. Although token readings vary with the fitting corpus, the dominant readout feature remains stable. Because SRP uses no corpus in its construction, it provides a control independent of the fitting corpus for lens analyses.
Method
Key results
SRP outperforms every baseline built on row geometry
| Readout | SRP coverage | Best alternative | Lead (pp) |
|---|---|---|---|
| Qwen3.5-0.8B | 0.754 | 0.665 | +8.9 |
| Qwen3.5-2B | 0.780 | 0.625 | +15.5 |
| Qwen3.5-9B | 0.812 | 0.639 | +17.3 |
| Ministral-3-8B | 0.777 | 0.648 | +12.9 |
| R1-Distill-Qwen-7B | 0.716 | 0.579 | +13.7 |
| R1-Distill-Llama-8B | 0.751 | 0.624 | +12.7 |
Coverage is the fraction of about 1,350 logit differences per model that the replaced readout reconstructs with the correct sign and error below half the score's magnitude. The best alternative is the strongest of six baselines built directly on the geometry of the readout rows, such as nearest-row ridge, k-means centroids and PCA.
Context decides which features carry a token
A fitted lens reports the language of its fitting corpus
Use the dictionaries
List the LM-head tokens that score highest on one readout feature, as in the paper's figures. This example downloads the 134 MB Qwen3.5-0.8B dictionary and the model's weights (about 1.9 GB in total); all dictionaries are on Hugging Face, and the code reproduces the paper.
# pip install huggingface_hub safetensors torch transformers
import torch
from huggingface_hub import hf_hub_download
from safetensors import safe_open
from transformers import AutoTokenizer
# SRP dictionary for Qwen3.5-0.8B: 16,384 readout features, k = 128.
repo, model = "hematteo/sparse-readout-prism", "Qwen/Qwen3.5-0.8B"
path = hf_hub_download(repo, "qwen3.5-0.8b/k128_16x/checkpoint.pt")
sae = torch.load(path, map_location="cpu", weights_only=True)["model_state_dict"]
W_enc = sae["encoder.weight"] # (16384, 1024)
# The model's readout. Qwen3.5-0.8B ties its LM head to the token embedding.
weights = hf_hub_download(model, "model.safetensors-00001-of-00001.safetensors")
with safe_open(weights, "pt") as f:
W_U = f.get_tensor("model.language_model.embed_tokens.weight").float()
tok = AutoTokenizer.from_pretrained(model)
# The dictionary reads LM-head rows centered on their mean and scaled to unit norm.
X = W_U - W_U.mean(0) # (248320, 1024)
X /= X.norm(dim=1, keepdim=True)
feature = 14186
top = (X[: len(tok)] @ W_enc[feature]).topk(8).indices # rows scoring highest on it
print([tok.decode(i) for i in top])
# ['的水果', '水果', ' fruit', ' Fruit', ' frutas', ' фруктов', ' fruits', '蔬果']
Related
Learning to Read Out: Unembedding Dynamics in Language Model Pretraining (NeurIPS 2026, arXiv).
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}
}