Ship the revision's experiment code and refresh release metadata (0.2.0)

CITATION.cff

@@ -23,12 +23,12 @@ authors:
     family-names: Lane
     affiliation: University of Cambridge
 license: MIT
-version: "0.1.0"
-date-released: "2026-06-12"
+version: "0.2.0"
+date-released: "2026-08-31"
 repository-code: "https://github.com/hematteo/sparse-readout-prism"
 preferred-citation:
   type: unpublished
-  title: "Sparse Readout Prism: A Sparse LM-Head Basis for Logit-Lens Readouts"
+  title: "Sparse Readout Prism: Explaining Logit-Lens Scores in Features Rather Than Tokens"
   authors:
     - given-names: Matteo
       family-names: He

README.md

@@ -6,7 +6,7 @@
 [![ruff](https://img.shields.io/badge/lint-ruff-261230)](https://docs.astral.sh/ruff/)
 
 Code and configs for the paper
-**"Sparse Readout Prism: A Sparse LM-Head Basis for Logit-Lens Readouts"**
+**"Sparse Readout Prism: Explaining Logit-Lens Scores in Features Rather Than Tokens"**
 (preprint forthcoming).
 
 Sparse Readout Prism factorizes language-model unembedding rows into a
@@ -37,8 +37,8 @@ after replacing `W_U` with its reconstruction, `KL` is the readout KL in bits:
 
 | Model | d_features | rowEV | top1 | KL (bits) |
 |---|---:|---:|---:|---:|
-| Qwen3.5-0.8B | 32768 | 0.877 | 0.891 | — |
-| Qwen3.5-2B | 65536 | 0.847 | 0.887 | — |
+| Qwen3.5-0.8B | 32768 | 0.877 | 0.891 | 0.135 |
+| Qwen3.5-2B | 65536 | 0.847 | 0.887 | 0.136 |
 | Qwen3.5-9B | 131072 | 0.857 | 0.900 | 0.105 |
 
 Full per-model / per-width / per-`k` numbers (including the fixed-budget k = 128
@@ -48,6 +48,24 @@ see Appendix K of the paper.
 
 ---
 
+## What 0.2.0 adds
+
+Code behind the experiments added in the paper's revision, ported from the
+research workspace with behaviour preserved (every artifact is mapped to its
+script and command line in [`docs/REPRODUCE.md`](docs/REPRODUCE.md) §2):
+
+- the cross-lens study: Jacobian-lens and ridge-translator fitting, per-layer
+  readout dumps, EN–ZH and EN–DE aggregators, prompt banks under `data/cross_lens/`;
+- causal validation of feature contributions
+  (`scripts/eval/run_causal_contribution_validation.py`);
+- the redesigned direct-geometry baselines (weighted row-kNN, k-means centroid
+  dictionaries) and the pooled error-tail analysis;
+- matched-KL frontier controls in the lexical-edit stress test and the paired
+  cluster bootstrap behind them;
+- cross-seed stability, feature-group matching and leave-one-out core recovery,
+  with the seed-variation and low-k training configs;
+- the sense-labelled evaluation on CoarseWSD-20.
+
 ## Install
 
 Requires Python 3.11 or 3.12.
@@ -58,6 +76,16 @@ uv sync                 # or: pip install -e .
 
 `uv sync` is recommended — it installs from the pinned `uv.lock` and picks
 the default PyTorch wheel for your platform (CPU/MPS on macOS, CUDA on Linux).
+
+The cross-lens study scripts (`scripts/run/fit_jlens.py` and the scripts that
+consume its lenses) additionally need the Jacobian-lens reference
+implementation, which is not on PyPI:
+
+```bash
+uv sync --extra lens    # adds jlens from github.com/anthropics/jacobian-lens (Apache-2.0)
+```
+
+Everything else installs and tests without it; those scripts import it lazily.
 If your GPU needs a specific CUDA build, pin `torch` to a `download.pytorch.org`
 index in `[tool.uv.sources]` locally and re-lock.
 
@@ -121,7 +149,7 @@ separately in the paper's LaTeX source from the committed metrics.
 The same call works on any HuggingFace causal LM — you just need a dictionary
 trained on **that model's** `W_U`. The selected paper dictionaries are published
 on the Hugging Face Hub at
-[`matteohe/sparse-readout-prism`](https://huggingface.co/matteohe/sparse-readout-prism)
+[`hematteo/sparse-readout-prism`](https://huggingface.co/hematteo/sparse-readout-prism)
 (`<model>/<operating_point>/checkpoint.pt`); download one with
 `huggingface_hub.hf_hub_download`, or train your own
 (`scripts/data/extract_model_readout.py` + the trainer — Pythia-160M is
@@ -150,7 +178,7 @@ row_mean, row_norms, rows_normalized = preprocess_rows(W_U)
 # 3. Load the matching trained dictionary and decompose.
 #    For a published model, pull the dictionary straight from the Hub, e.g.:
 #      from huggingface_hub import hf_hub_download
-#      ckpt = hf_hub_download("matteohe/sparse-readout-prism", "qwen3.5-2b/k256_32x/checkpoint.pt")
+#      ckpt = hf_hub_download("hematteo/sparse-readout-prism", "qwen3.5-2b/k256_32x/checkpoint.pt")
 #    For pythia-160m here, train one locally first (see docs/REPRODUCE.md).
 sae = load_factorizer("checkpoint.pt", freeze=True)       # trained on this model's W_U
 d = decompose_token_logit(
@@ -182,22 +210,38 @@ src/sparse_readout_prism/         core library
   token_display.py, utils.py      token/text display strings + shared IO / model-loader helpers
   research/                       flat helpers shared by several scripts/ entry
                                   points (Qwen readout + query-decomposition
-                                  toolkits, prompt banks, registry, run IO);
+                                  toolkits, prompt banks, registry, run IO,
+                                  seed-stability contrast pipeline);
                                   logic used by a single script stays inline in
                                   that script
 
 scripts/                          self-contained CLI entry points
-  train/, run/, eval/             SAE training, experiment suites, evaluation
-  analyze/, data/                 component inspection, data prep
+  train/                          SAE training from a run config
+  run/                            experiment suites: query-fidelity bank, direct-
+                                  geometry baselines, lexical-edit stress test,
+                                  Jacobian / ridge lens fitting + cross-lens
+                                  readout dumps, CoarseWSD-20 bundles
+  eval/                           task-fidelity evaluation, dense-control
+                                  diagnostic, causal validation, error tails,
+                                  matched-KL bootstrap, cross-lens aggregators,
+                                  cross-seed stability / group matching / LOO
+  analyze/                        feature directions, row-norm tails, label
+                                  audit, nearest-rows reading, antonym-layer
+                                  and three-lens tables, sense-group alignment
+  data/                           readout extraction, query banks, polysemy
+                                  pairs, the EN-DE cross-lens bank
   figures/                        compute + persist the metrics behind the paper
                                   figures (CSV/JSON example panels, lens grids,
-                                  decompositions); no figures are rendered here
+                                  decompositions, the shared-feature panel);
+                                  no figures are rendered here
 
 configs/
   models/, sweeps/, registries/   model configs, sweep grids, paper-selected SAEs
 
 data/
   query_banks/                    curated prompt banks (paper inputs)
+  cross_lens/                     EN-ZH / EN-DE cross-lens prompt banks + exclusion report
+  wsd/                            optional sense-gloss anchors for the CoarseWSD-20 runs
   appendix/                       audited literals behind the Appendix K tables
   audit/                          feature-label audit annotations + counts (Appendix L)
 
@@ -215,10 +259,12 @@ docs/                             REPRODUCE.md (figure → script + config map),
 | [`docs/REPRODUCE.md`](docs/REPRODUCE.md) | Figure / table → script + config map. Hardware budget. |
 | [`docs/DATA.md`](docs/DATA.md) | Where artefacts live, schema, how to adapt to a different model. |
 | [`data/query_banks/README.md`](data/query_banks/README.md) | Paper-input JSONL schema + per-file purpose. |
+| [`data/cross_lens/README.md`](data/cross_lens/README.md) | Cross-lens prompt-bank schema, families, controls, provenance. |
+| [`configs/sweeps/README.md`](configs/sweeps/README.md) | Training configs, including the seed-variation and low-k cells. |
 
 Pretrained readout-feature dictionaries (the selected SAEs cited in the
 paper) are published on the Hugging Face Hub at
-[`matteohe/sparse-readout-prism`](https://huggingface.co/matteohe/sparse-readout-prism)
+[`hematteo/sparse-readout-prism`](https://huggingface.co/hematteo/sparse-readout-prism)
 — 17 dictionaries across 8 base models, laid out as
 `<model>/<operating_point>/checkpoint.pt`.
 [`configs/registries/exp2_selected_sae_checkpoints.yaml`](configs/registries/exp2_selected_sae_checkpoints.yaml)
@@ -248,8 +294,11 @@ The suite is fast and CI-safe: identity correctness of the decomposition,
 schema invariants of the task-fidelity evaluator, checkpoint-registry
 resolution, the data-loading branches (token-mask filtering, val-split
 fallbacks), the trainer's L0 controller and resume path, the public-API
-surface, and layout guardrails for the `research/` sub-package. No GPU, no
-model loads, no figure rendering.
+surface, and layout guardrails for the `research/` sub-package. The
+experiment code added in 0.2.0 ships its own CPU tests (baseline identity for
+the row-kNN and k-means methods, the causal-validation self-test, the paired
+matched-KL bootstrap, the cross-lens aggregators, sense-group selection, and
+the seed-stability pipeline). No GPU, no model loads, no figure rendering.
 
 ## Related repository
 
@@ -263,7 +312,7 @@ analysis.
 
 ```bibtex
 @misc{he2026sparsereadoutprism,
-  title  = {Sparse Readout Prism: A Sparse LM-Head Basis for Logit-Lens Readouts},
+  title  = {Sparse Readout Prism: Explaining Logit-Lens Scores in Features Rather Than Tokens},
   author = {He, Matteo and Shen, William F. and Qiu, Xinchi and Lane, Nicholas D.},
   year   = {2026},
   note   = {Preprint forthcoming; see the repository for the up-to-date reference},

configs/registries/exp2_selected_sae_checkpoints.yaml

@@ -14,7 +14,7 @@
 # The Pythia-160M continuity anchor uses a legacy state-dict schema and is NOT
 # published on the Hub.
 
-huggingface_repo: matteohe/sparse-readout-prism
+huggingface_repo: hematteo/sparse-readout-prism
 archive_root: ${SRP_ARCHIVE_ROOT}
 
 headline:

configs/registries/result1_query_fidelity_cluster.yaml

@@ -18,7 +18,7 @@
 # (fidelity -> <model>/k256_32x, strict_budget -> <model>/k128_16x; qwen9b
 # strict_budget -> qwen3.5-9b/k128_8x). The `checkpoint:` paths below remain the
 # canonical on-cluster training locations.
-huggingface_repo: matteohe/sparse-readout-prism
+huggingface_repo: hematteo/sparse-readout-prism
 archive_root: ${SRP_ARCHIVE_ROOT}
 
 models:

configs/sweeps/README.md

@@ -37,6 +37,9 @@ Two kinds of file live here:
 | `deepseek_r1_distill_qwen_7b_paper.yaml`, `…_16x_k128.yaml` | run config | R1-Distill-Qwen-7B finalists (32×/k256 and strict-budget 16×/k128). |
 | `deepseek_r1_distill_llama_8b_paper.yaml`, `…_16x_k128.yaml` | run config | R1-Distill-Llama-8B finalists. |
 | `ministral3_8b_paper.yaml`, `…_16x_k128.yaml` | run config | Ministral-3-8B-Base finalists (32×/k256 and strict-budget 16×/k128). |
+| `qwen35_2b_seedvar_base.yaml` | run config | Qwen3.5-2B seed-variation base (32×/k256, seed 0). The stability appendix's three-seed families and the multi-seed rows of the low-k table are this file with `--set` overrides of `run.seed` / `run.init_seed` / `data.data_seed`, `factorizer.d_features` / `factorizer.k` / `evaluation.k`, and `run.name` / `run.output_dir` (cells listed in the file header). |
+| `qwen35_0p8b_paper_topk_32x_k256_s1.yaml`, `…_s2.yaml` | run config | Qwen3.5-0.8B 32×/k256 seed-window cells (seeds 1 and 2; seed 0 is the archived finalist). |
+| `lowk_qwen35_2b_16x_base.yaml`, `lowk_qwen35_0p8b_16x_base.yaml` | run config | Low activation-budget cells at 16× width (`k` in {32, 64} via `--set factorizer.k=… --set evaluation.k=…`), behind the sparsity-budget table. |
 
 The provisional Gemma operating points in the registries were trained with the
 same converge recipe as the Qwen models (`paper_phase2_base.yaml` protocol);
@@ -50,6 +53,14 @@ The run configs set `data.fallback: error` and a cluster `data.path`
 `FileNotFoundError` naming the path. (Grid templates are rejected earlier still,
 by the run-config guard, before any data is loaded.)
 
+The run files behind the paper's seed-variation and low-k cells carried a few
+extra keys (`data.row_preprocessing`, `training.init_mode`, `lr_schedule`,
+`lambda_mode`, `final_lr_frac`, `auxk_*`, `artifacts.make_figures`) that no
+code path in the shipped trainer reads, and did not read in the trainer that
+produced those cells either; they are omitted here rather than shipped as if
+they were live. Every value the trainer does read is identical to the paper
+files.
+
 The cross-model finalist configs (R1-Distill-Qwen-7B, R1-Distill-Llama-8B,
 Ministral) are launched with the generic trainer —
 `train_readout_sae_from_config.py --config <file>` — and have no dedicated

configs/sweeps/lowk_qwen35_0p8b_16x_base.yaml

@@ -0,0 +1,71 @@
+# Qwen3.5-0.8B readout SAE, 16x width, low activation-budget cells
+# (tab:app-low-k-dead-features, the single-seed Q-0.8B rows k=32 and k=64).
+# Same recipe as the seed-variation configs (TopK, 20000 steps, batch 4096,
+# hybrid frequency/uniform row sampling); d_features fixed at 16384
+# (16 x d_model 1024), seed 0.
+#
+# A launch overrides per cell via --set
+#   factorizer.k / evaluation.k        the budget (evaluation.k = factorizer.k)
+#   run.name / run.output_dir          qwen35_0p8b_d16384_k<k>_s0
+# and nothing else.
+#
+# Paper cells (d_features, k, seed): (16384, 32, 0), (16384, 64, 0).
+# The row-centred explained variance (rowEV) was not recorded for these two
+# cells; the table enters it as ---.
+#
+# Example cell (k=32):
+#   uv run python scripts/train/train_readout_sae_from_config.py \
+#     --config configs/sweeps/lowk_qwen35_0p8b_16x_base.yaml \
+#     --set factorizer.k=32 --set evaluation.k=32 \
+#     --set run.output_dir=results/lowk_qwen35_0p8b_16x/qwen35_0p8b_d16384_k32_s0 \
+#     --set run.name=qwen35_0p8b_d16384_k32_s0
+#
+# data.path is the {W_U_orig, h_LN} extraction of Qwen3.5-0.8B
+# (scripts/data/extract_model_readout.py), the artifact
+# configs/registries/result1_query_fidelity_cluster.yaml resolves for this model.
+
+run:
+  name: lowk_qwen35_0p8b_16x
+  seed: 0
+  init_seed: 0
+  output_dir: results/lowk_qwen35_0p8b_16x/qwen35_0p8b_d16384_k64_s0
+  device: auto
+
+data:
+  path: ${SRP_ARCHIVE_ROOT}/data/qwen35-0p8b/qwen35_0p8b.pt
+  fallback: error
+  max_rows: null
+  max_hidden: 2048
+  val_hidden: 512
+  data_seed: 0
+
+factorizer:
+  architecture: topk
+  d_features: 16384
+  k: 64
+  encoder_init_scale: 0.02
+
+training:
+  steps: 20000
+  batch_size: 4096
+  lr: 0.001
+  weight_decay: 0.0001
+  code_l2: 0.000001
+  prism_top_n: 1
+  lambda_prism: 0.001
+  row_sampling: hybrid_50freq_50uniform
+
+evaluation:
+  k: 64
+  max_eval_rows: 20000
+  max_eval_hidden: 1024
+  topk_overlap: 5
+  contribution_top_k: 8
+
+artifacts:
+  save_checkpoint: true
+
+autoresearch:
+  score_thresholds:
+    max_dead_feature_rate: 0.5
+    max_val_logit_kl_bits_mean: 20.0

configs/sweeps/lowk_qwen35_2b_16x_base.yaml

@@ -0,0 +1,74 @@
+# Qwen3.5-2B readout SAE, 16x width, low activation-budget cells
+# (tab:app-low-k-dead-features, the single-seed Q-2B k=32 row and seed 0 of the
+# k=64 row). Same recipe as qwen35_2b_seedvar_base.yaml (TopK, 20000 steps,
+# batch 4096, hybrid frequency/uniform row sampling); d_features fixed at 32768
+# (16 x d_model 2048), seed 0.
+#
+# A launch overrides per cell via --set
+#   factorizer.k / evaluation.k        the budget (evaluation.k = factorizer.k)
+#   run.name / run.output_dir          qwen35_2b_d32768_k<k>_s0
+# and nothing else.
+#
+# Paper cells (d_features, k, seed): (32768, 32, 0), (32768, 64, 0).
+# Seeds 1 and 2 of the k=64 row come from qwen35_2b_seedvar_base.yaml with
+# --set factorizer.d_features=32768 --set factorizer.k=64. The row-centred
+# explained variance (rowEV) was not recorded for the two seed-0 cells here;
+# the table's k=64 rowEV range is over the other two seeds and the k=32 row
+# enters it as ---.
+#
+# Example cell (k=32):
+#   uv run python scripts/train/train_readout_sae_from_config.py \
+#     --config configs/sweeps/lowk_qwen35_2b_16x_base.yaml \
+#     --set factorizer.k=32 --set evaluation.k=32 \
+#     --set run.output_dir=results/lowk_qwen35_2b_16x/qwen35_2b_d32768_k32_s0 \
+#     --set run.name=qwen35_2b_d32768_k32_s0
+#
+# data.path is the {W_U_orig, h_LN} extraction of Qwen3.5-2B
+# (scripts/data/extract_model_readout.py), the artifact
+# configs/registries/result1_query_fidelity_cluster.yaml resolves for this model.
+
+run:
+  name: lowk_qwen35_2b_16x
+  seed: 0
+  init_seed: 0
+  output_dir: results/lowk_qwen35_2b_16x/qwen35_2b_d32768_k64_s0
+  device: auto
+
+data:
+  path: ${SRP_ARCHIVE_ROOT}/data/qwen35-2b/qwen35_2b.pt
+  fallback: error
+  max_rows: null
+  max_hidden: 2048
+  val_hidden: 512
+  data_seed: 0
+
+factorizer:
+  architecture: topk
+  d_features: 32768
+  k: 64
+  encoder_init_scale: 0.02
+
+training:
+  steps: 20000
+  batch_size: 4096
+  lr: 0.001
+  weight_decay: 0.0001
+  code_l2: 0.000001
+  prism_top_n: 1
+  lambda_prism: 0.001
+  row_sampling: hybrid_50freq_50uniform
+
+evaluation:
+  k: 64
+  max_eval_rows: 20000
+  max_eval_hidden: 1024
+  topk_overlap: 5
+  contribution_top_k: 8
+
+artifacts:
+  save_checkpoint: true
+
+autoresearch:
+  score_thresholds:
+    max_dead_feature_rate: 0.5
+    max_val_logit_kl_bits_mean: 20.0