Butanium/sparse-readout-prism-qwen3.6-27b-k128-16x

AI usage disclaimer

Trained by silico powered by fable 5.1. No human reviewed the code but this is using the paper's original code so I trust it. The rest of the readme is AI generated and not reviewed

Sparse Readout Prism dictionary for Qwen3.6-27B (TopK, 16x, k = 128)

A readout SAE: a TopK sparse dictionary trained on the rows of the unembedding matrix W_U of Qwen/Qwen3.6-27B (revision 6a9e13bd6fc8f0983b9b99948120bc37f49c13e9). It lets a logit-lens score s = h · w_t for any decoded state h and token t be written exactly as a base term plus one contribution per active dictionary feature plus a residual, so a top-token list can be explained in features instead of tokens.

Status: below the paper's fidelity thresholds. On the paper's evaluation basis (first 20,000 kept rows) it passes row explained variance (0.6547 ≥ 0.60) and top-1 agreement (0.8295 ≥ 0.80) but misses readout KL (0.3556 bits > 0.35). On the full 248,068-row vocabulary all three thresholds fail (rowEV 0.4948, top-1 0.8135, KL 0.5058 bits). Read the Metrics and Limitations sections before using the feature assignments for high-id (rare, non-English) tokens.

Method and sources

• Method: Sparse Readout Prism: Explaining Logit-Lens Scores in Features Instead of Tokens, Matteo He, William F. Shen, Xinchi Qiu, Nicholas D. Lane. arXiv:2609.01936.
• Code: hematteo/sparse-readout-prism, installed unchanged at commit 68b9b22da499758959382209e68960fbe5fa9318 (v0.2.2). reproduce/sparse_readout_prism_edits.diff is empty: no library file was modified. The only additions are the experiment-side scripts in reproduce/ (extraction, segmented training driver, job script, config).

Deviations from the released code and from the paper's recipe

Item	Paper's Qwen/Gemma finalist recipe	This dictionary
Encoder/decoder init	row-seeded init	random-normal (encoder_init_scale 0.02, the shipped trainer's default)
Learning-rate schedule	warmup then cosine 1e-3 → 1e-4	constant 1e-3 (the shipped trainer has no schedule)
Prism penalty	delayed ramp	λ = 1e-3 from step 1 (no ramp in the shipped trainer)
Precision	not stated	fp32, no TF32 / autocast / compile
Training in segments	single run	four 5,000-step segments resumed from the library's own train_ckpt.pt (weights, optimizer state and sampling RNG restored), which the trainer treats as one 20,000-step run; each rung was evaluated with evaluate_model (see metrics/ladder_metrics.json)
Held-out states	wikitext-2 prompts in the authors' extraction script	NeelNanda/pile-10k documents, last position only (see below)
Model adapter	authors' extract_model_readout.py	reproduce/extract_readout.py, same payload schema {W_U_orig, h_LN, token_mask} and the same adapter check (lm_head input × W_U = model logits, bf16 relative error 2.1e-3, structural error 0)

The three recipe items in the first rows are not implemented in the released trainer at this commit and were not re-implemented; they are the untried levers most likely to close the KL gap.

Operating point

Architecture	TopK SAE over centered, per-row L2-normalized unembedding rows
Width	81,920 features (16 × d_model 5,120)
k	128 active features per row
Steps / batch	20,000 steps × 4,096 rows
Optimizer	AdamW, lr 1e-3 constant, weight decay 1e-4, code L2 1e-6
Prism penalty	λ_prism = 1e-3 on the top-1 target (prism_top_n 1), from step 1
Row sampling	hybrid 50% frequency-weighted / 50% uniform
Rows	all 248,068 kept token ids of the 248,320-row W_U (token_mask from the tokenizer drops specials and byte-fallback rows)
Centering	live: full-vocabulary row mean of W_U (stored in the checkpoint as row_mean)
Training states (prism penalty)	2,048 held-out final-norm states
Seeds	data_seed 0, init_seed 0
Hardware / wall time	1 × H100 80 GB; 10,916 s of training (0.56 s/step) plus about 45 s per evaluation rung; extraction 500 s; total job 3 h 22 min
Software	torch 2.11.0+cu130, sparse-readout-prism 0.2.2 @ 68b9b22

Held-out hidden states

4,200 documents from NeelNanda/pile-10k (train split, shuffled with seed 42, documents with at least 200 characters, tokenized with max_len 256, mean 227 tokens). For each document the input to lm_head (the post-final-RMSNorm state) at the last non-pad position was captured in fp32 (norms: mean 135.4, min 76.7, max 146.8). The library then draws 2,048 + 2,000 of these states with torch.randperm (seed data_seed + 11); the first 2,048 are the training states for the prism penalty, the next 2,000 are the held-out evaluation states. val_overlaps_train is false.

Metrics

Both bases use the same 2,000 held-out states. The paper's basis restricts W_U to the first 20,000 kept rows and takes the softmax over that sub-vocabulary (this matches the authors' sweep configs, evaluation.max_eval_rows 20000); the full-vocabulary basis uses all 248,068 kept rows. KL is the mean over states of KL(original ‖ reconstructed) in bits (the paper's glossary describes a median; the library computes a mean, and the 0.35 threshold was applied to the mean).

Metric	Paper basis (20,000 rows)	Full vocabulary (248,068 rows)	Paper threshold
Row explained variance (centered)	0.6547	0.4948	≥ 0.60
Row cosine (centered)	0.8029	0.6976
Top-1 agreement (val states)	0.8295	0.8135	≥ 0.80
Top-5 overlap	0.7919	0.7683
Readout KL, mean (bits)	0.3556	0.5058	≤ 0.35
Top-1 logit residual fraction, mean	0.0546	0.0651
Top-1 logit residual, mean abs	1.100	1.346
Top-8 abs contribution coverage	0.8589	0.8497
Top-8 positive contribution coverage	0.8876	0.8779
Dead feature rate	0.2123	0.1192
Rare feature rate	0.7562	0.7773
Always-on feature rate	7.3e-5	0.0
Feature usage entropy	0.8296	0.8397
Mean L0 (p10 / p50 / p90)	127.94 (128 / 128 / 128)	128.00 (128 / 128 / 128)
Identity error, max / mean	5.3e-5 / 1.0e-5	5.3e-5 / 1.0e-5
Duplicate decoder fraction	0.0	0.0
Library selection score	2.895	2.583

Sign agreement between the feature sum and the exact score held in all 320 top-5 token decompositions of the downstream steering-vector readouts this dictionary was built for; the relative residual ρ0 = |ε| / |s| over those 160 raw decompositions had median 0.276 (IQR 0.134–0.464), near 0.05 for frequent rows and 0.43–0.54 for content words.

Convergence ladder (paper basis): rowEV 0.4308 / 0.5588 / 0.6165 / 0.6547, top-1 0.6795 / 0.7710 / 0.8100 / 0.8295, KL 1.5606 / 0.6518 / 0.4557 / 0.3556 bits, dead rate 0.4084 / 0.3184 / 0.2527 / 0.2123 at 5k / 10k / 15k / 20k steps.

Paper reference points at the same k: Qwen3.5-9B 8×/k128 rowEV 0.621, top-1 0.846, KL 0.296; Qwen3.5-2B 16×/k128 0.712 / 0.858 / 0.261.

Files

File	Content
checkpoint.pt	the trained dictionary, the library's checkpoint layout (below); sha256 in checkpoint.sha256
checkpoint.sha256	sha256sum of checkpoint.pt
config.yaml	the training config as written by the library run (reproduce/qwen36_27b_topk_16x_k128.yaml is the source)
metrics/metrics.json	paper-basis metrics (20,000 rows)
metrics/metrics_all_rows.json	full-vocabulary metrics
metrics/ladder_metrics.json, metrics/fidelity_rungs.json	per-rung metrics at 5k/10k/15k/20k steps
metrics/history.json, metrics/compute_meta.json	training history and compute metadata written by the library
feature_labels_library.json	the library's own labels for its top 256 features
feature_top_rows.jsonl	one line per feature (all 81,920): top 12 unembedding rows by encoder activation, verbatim tokenizer.decode([id]), with token ids and activations; rows with zero activation dropped, so dead features have an empty list
reproduce/	everything needed to retrain: COMMANDS.md, job_dictionary.sh, extract_readout.py, train_ladder.py, qwen36_27b_topk_16x_k128.yaml, steering_analysis/ (model loader used by the extraction), pyproject.toml, uv.lock, sparse_readout_prism_edits.diff (empty)

Checkpoint layout

torch.load(path, weights_only=True) returns a dict with

• model_state_dict: encoder.weight (81920 × 5120), encoder.bias (81920), decoder (81920 × 5120); nothing else;
• factorizer: {architecture: topk, d_features: 81920, k: 128, encoder_init_scale: 0.02};
• evaluation: the evaluation config;
• row_mean (5120): the centering mean used in training;
• row_norms (248068): per-row norms of the centered kept rows;
• row_token_ids (248068): the token ids of those rows;
• metrics: the paper-basis metrics.

Loading

import torch
from huggingface_hub import hf_hub_download
from sparse_readout_prism.factorizers import load_factorizer          # pip install "sparse-readout-prism @ git+https://github.com/hematteo/sparse-readout-prism@68b9b22da499758959382209e68960fbe5fa9318"
from sparse_readout_prism.data import preprocess_rows, center_normalize_rows

path = hf_hub_download("Butanium/srp-qwen3.6-27b-k128-16x", "checkpoint.pt")
ckpt = torch.load(path, map_location="cpu", weights_only=True)
sae = load_factorizer(ckpt, freeze=True)          # uses ckpt["factorizer"]; note: `load_sae` does not accept this layout

# W_U of Qwen3.6-27B at the pinned revision: lm_head.weight (248320 x 5120). Only the lm_head shard is needed;
# reproduce/steering_analysis/jlens/head.py::load_head_from_hub downloads just that shard.
W_U = ...                                          # torch.float32, (248320, 5120)

# Rows are centered on the full-vocabulary mean and L2-normalized, exactly as in training.
row_mean, row_norms, rows_normalized = preprocess_rows(W_U)
# the live mean and the stored training mean differ by 0.02% of the mean centered-row norm
assert (row_mean - ckpt["row_mean"]).norm() / (W_U - row_mean).norm(dim=1).mean() < 1e-3

with torch.no_grad():
    codes = sae.encode(rows_normalized[:4096])     # (n, 81920), 128 nonzeros per row
    recon = sae.decode(codes)                      # unit-norm decoder columns

# Decomposing a logit s = h . w_t for a decoded state h (post final-norm) and token t:
#   s = h . row_mean  +  sum_f  h . (row_norms[t] * codes[t, f] * decoder[f])  +  residual
# The library's research.cross_lens.decompose_token / research.qwen_readout.encode_topk implement this exactly.

Intended use and limitations

Intended for interpretability research on Qwen3.6-27B logit-lens or J-lens / R-lens readouts: explaining which readout features a decoded state drives, comparing lenses at the feature level, and checking whether tokens in different scripts share features. It was trained for one such study (readouts of two steering vectors) and is released so that study can be reproduced and the basis reused.

• Fidelity below the paper's thresholds. On the full vocabulary, where rare and non-English rows live, row explained variance is 0.49 and readout KL 0.51 bits; dominant-feature assignments for content words typically leave 40–55% of the logit in the residual. Frequent rows and punctuation are reconstructed much better (ρ0 near 0.05).
• Dead and rare features. 21% of features are dead on the 20,000-row basis (12% on the full vocabulary) and about 76% are rare; feature_top_rows.jsonl has empty lists for dead features.
• Labels are rows, not descriptions. Feature labels are the feature's top unembedding rows by encoder activation, verbatim. No model-generated or human descriptions are included.
• Generic trainer. The paper's row-seeded init, LR decay and delayed prism ramp were not used; a dictionary trained with them would likely be more faithful.
• Not a model of hidden-state activations: this dictionary factorizes W_U rows, not residual-stream activations.

Reproduction

See reproduce/COMMANDS.md. The two scripts are run from the root of the codebase with the pinned library on PYTHONPATH; extract_readout.py produces the {W_U_orig, h_LN, token_mask} payload from the HF model and train_ladder.py runs the library's trainer and evaluations. pyproject.toml / uv.lock are the dependency declarations as used (the job image supplied torch, transformers and numpy; the library was installed with --no-deps).

License

MIT, matching the sparse-readout-prism code. The base model Qwen/Qwen3.6-27B keeps its own license; this repository contains no model weights other than a factorization of its unembedding matrix.

Citation

@article{he2026sparsereadoutprism,
  title   = {Sparse Readout Prism: Explaining Logit-Lens Scores in Features Instead of Tokens},
  author  = {He, Matteo and Shen, William F. and Qiu, Xinchi and Lane, Nicholas D.},
  journal = {arXiv preprint arXiv:2609.01936},
  year    = {2026}
}