diff --git a/CITATION.cff b/CITATION.cff index f745a2c..4a37c68 100644 --- a/CITATION.cff +++ b/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 diff --git a/README.md b/README.md index 880c814..723a87a 100644 --- a/README.md +++ b/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) (`//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 `//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}, diff --git a/configs/registries/exp2_selected_sae_checkpoints.yaml b/configs/registries/exp2_selected_sae_checkpoints.yaml index 42ed1d5..0297083 100644 --- a/configs/registries/exp2_selected_sae_checkpoints.yaml +++ b/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: diff --git a/configs/registries/result1_query_fidelity_cluster.yaml b/configs/registries/result1_query_fidelity_cluster.yaml index 605c2bb..c801abb 100644 --- a/configs/registries/result1_query_fidelity_cluster.yaml +++ b/configs/registries/result1_query_fidelity_cluster.yaml @@ -18,7 +18,7 @@ # (fidelity -> /k256_32x, strict_budget -> /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: diff --git a/configs/sweeps/README.md b/configs/sweeps/README.md index 60228c4..43fa52d 100644 --- a/configs/sweeps/README.md +++ b/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 ` — and have no dedicated diff --git a/configs/sweeps/lowk_qwen35_0p8b_16x_base.yaml b/configs/sweeps/lowk_qwen35_0p8b_16x_base.yaml new file mode 100644 index 0000000..71e1657 --- /dev/null +++ b/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_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 diff --git a/configs/sweeps/lowk_qwen35_2b_16x_base.yaml b/configs/sweeps/lowk_qwen35_2b_16x_base.yaml new file mode 100644 index 0000000..ccbb782 --- /dev/null +++ b/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_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 diff --git a/configs/sweeps/qwen35_0p8b_paper_topk_32x_k256_s1.yaml b/configs/sweeps/qwen35_0p8b_paper_topk_32x_k256_s1.yaml new file mode 100644 index 0000000..7346d69 --- /dev/null +++ b/configs/sweeps/qwen35_0p8b_paper_topk_32x_k256_s1.yaml @@ -0,0 +1,59 @@ +# Qwen3.5-0.8B readout SAE, 32x/k=256 recipe, seed 1 (tab:app-k-qwen08b-seed-window). +# Repeats the archived seed-0 dictionary setting (qwen0p8b_k256_32x in +# configs/registries/exp2_selected_sae_checkpoints.yaml) with only +# run.seed / run.init_seed / data.data_seed changed. Sister file: +# qwen35_0p8b_paper_topk_32x_k256_s2.yaml (seed 2). +# +# Launched as-is, no --set overrides: +# uv run python scripts/train/train_readout_sae_from_config.py \ +# --config configs/sweeps/qwen35_0p8b_paper_topk_32x_k256_s1.yaml +# +# d_features 32768 = 32 x d_model (1024). 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: qwen35_0p8b_paper_topk_32x_k256_s1 + seed: 1 + init_seed: 1 + output_dir: results/experiments/proto_token_lens/qwen35_0p8b_sae/seedvar/topk_d32768_32x_k256_rowseeded_hybrid_lamramp_s1 + 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: 1 + +factorizer: + architecture: topk + d_features: 32768 + k: 256 + 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: 256 + 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 diff --git a/configs/sweeps/qwen35_0p8b_paper_topk_32x_k256_s2.yaml b/configs/sweeps/qwen35_0p8b_paper_topk_32x_k256_s2.yaml new file mode 100644 index 0000000..444ac20 --- /dev/null +++ b/configs/sweeps/qwen35_0p8b_paper_topk_32x_k256_s2.yaml @@ -0,0 +1,59 @@ +# Qwen3.5-0.8B readout SAE, 32x/k=256 recipe, seed 2 (tab:app-k-qwen08b-seed-window). +# Sister file of qwen35_0p8b_paper_topk_32x_k256_s1.yaml: the archived seed-0 +# dictionary setting (qwen0p8b_k256_32x in +# configs/registries/exp2_selected_sae_checkpoints.yaml) with only +# run.seed / run.init_seed / data.data_seed changed. +# +# Launched as-is, no --set overrides: +# uv run python scripts/train/train_readout_sae_from_config.py \ +# --config configs/sweeps/qwen35_0p8b_paper_topk_32x_k256_s2.yaml +# +# d_features 32768 = 32 x d_model (1024). 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: qwen35_0p8b_paper_topk_32x_k256_s2 + seed: 2 + init_seed: 2 + output_dir: results/experiments/proto_token_lens/qwen35_0p8b_sae/seedvar/topk_d32768_32x_k256_rowseeded_hybrid_lamramp_s2 + 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: 2 + +factorizer: + architecture: topk + d_features: 32768 + k: 256 + 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: 256 + 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 diff --git a/configs/sweeps/qwen35_2b_seedvar_base.yaml b/configs/sweeps/qwen35_2b_seedvar_base.yaml new file mode 100644 index 0000000..2e52a3a --- /dev/null +++ b/configs/sweeps/qwen35_2b_seedvar_base.yaml @@ -0,0 +1,76 @@ +# Qwen3.5-2B readout SAE, seed-variation base recipe. Backs the Appendix I +# stability analyses (tab:app-cross-seed-stability, tab:app-feature-group-matching, +# tab:app-loo-core-recovery) and the multi-seed rows of tab:app-low-k-dead-features. +# TopK, 20000 steps, batch 4096, hybrid frequency/uniform row sampling; the same +# optimizer recipe as paper_phase2_base.yaml. +# +# One file, many cells. A launch overrides per cell via --set +# run.seed / run.init_seed / data.data_seed all three = the cell's seed +# factorizer.d_features / factorizer.k / evaluation.k width x budget (evaluation.k = factorizer.k) +# run.name / run.output_dir qwen2b_d_k_s +# and leaves everything else untouched. +# +# Paper cells (d_features, k, seed) launched from this file: +# seed-variation family (Appendix I): (65536, 256, 0|1|2) and (32768, 128, 0|1|2) +# width x budget grid (tab:app-low-k-dead-features, rows with 2-3 seeds): +# (32768, 256, 0|1|2) [seed 1 did not finish; the table's row has two seeds] +# (65536, 128, 0|1|2) +# (32768, 64, 1|2) [seed 0 of this row and the k=32 row come from lowk_qwen35_2b_16x_base.yaml] +# +# Example cell (seed 1 of the 32x/k256 family): +# uv run python scripts/train/train_readout_sae_from_config.py \ +# --config configs/sweeps/qwen35_2b_seedvar_base.yaml \ +# --set factorizer.d_features=65536 --set factorizer.k=256 --set evaluation.k=256 \ +# --set run.seed=1 --set run.init_seed=1 --set data.data_seed=1 \ +# --set run.output_dir=results/qwen35_2b_seedvar/qwen2b_d65536_k256_s1 \ +# --set run.name=qwen2b_d65536_k256_s1 +# +# d_features 65536 = 32 x d_model (2048). 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: qwen35_2b_seedvar_base + seed: 0 + init_seed: 0 + output_dir: results/qwen35_2b_seedvar/qwen2b_d65536_k256_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: 65536 + k: 256 + 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: 256 + 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 diff --git a/data/cross_lens/README.md b/data/cross_lens/README.md new file mode 100644 index 0000000..c5b8196 --- /dev/null +++ b/data/cross_lens/README.md @@ -0,0 +1,100 @@ +# Cross-lens prompt banks + +Prompt banks for the cross-lens study (paper section "A Fitted Lens Reports the +Language of Its Corpus" and the "Cross-Lens Study Protocol" appendix). Each +record is one **prompt** plus the vocabulary tokens whose transported logits +the readout runner decomposes on the same frozen Qwen3.5-9B hidden states under +two or more fitted lenses. The banks are model-specific in one respect only: +every decomposed surface is a single token under the Qwen3.5 tokenizer. + +## Schema + +Every record has these keys (`lang_a` / `lang_b` only in the EN-DE bank): + +| Key | Type | Description | +|---|---|---| +| `id` | str | Unique prompt id (`ant_01`, `fac_03`, `antonym_de_01`, ...). | +| `group` | str | Prompt family (below). Families starting with `ctrl_` are controls and never pooled into the headline. | +| `concept` | str | The concept the prompt elicits; used to exclude same-concept pairs from the shuffled-pairing null. | +| `prompt` | str | The full prompt fed to the model; the final position is scored. | +| `form_a` | str | The in-context answer form (the language or script the prompt asks for). | +| `form_b` | str | The same concept's canonical form in the other language. | +| `targets` | list[str] | Every surface the runner decomposes: `form_a`, `form_b`, any extra forms, then the nulls. | +| `null_targets` | list[str] | Unrelated tokens (script- or language-matched) for the unrelated-token null floor. | +| `answer` | str | The expected answer (`form_a` without its leading space). | +| `lang_a`, `lang_b` | str | EN-DE bank only: the language of `form_a` / `form_b` (`de` or `en`). | + +Each target string is probed by the id of its first token, and the aggregators +compare the dominant readout feature of `form_a`'s decomposition across lenses +on the mid-band layers {21, 24, 26, 29}. Two null comparisons price in loose +matching: the unrelated-token null (`form_a` vs each `null_targets` entry within +one lens) and the shuffled-pairing null (prompt i under lens A vs prompt j of a +different concept under lens B, three seeded draws per prompt). + +## Files + +| File | Records | What it is | +|---|---|---| +| `cross_lens_prompts_en_zh.json` | 92 | The English-Chinese bank: 80 cross-lens prompts in seven families plus 12 controls. `form_a` is the answer-script form, `form_b` the other-script form; nulls are script-matched. Read by the English- and Chinese-fitted Jacobian lenses (100 and 300 prompts), the two ridge translators, and the Jacobian-vs-ridge comparison. | +| `cross_lens_prompts_en_de.json` | 102 | The English-German bank: 90 cross-lens prompts in seven families plus 12 controls. Both languages are Latin-script, so the surface-language call in the aggregator is lexical and cognates are excluded by construction. Built by `scripts/data/build_cross_lens_de_bank.py`. | +| `cross_lens_prompts_en_de_exclusion_report.json` | - | The builder's report for the EN-DE bank: tokenizer, candidate / kept counts per family, and the 30 of 132 candidates excluded with the reason (`cognate (lev=n)` at normalised edit distance <= 2 after diacritic folding, or `not single-token`). | +| `cross_lens_three_lens_prompts.json` | 2 | One-prompt-per-family bank for the worked example read by three lenses at once: the German antonym prompt `antonym_de_01` (answer `groß`) and its Chinese analogue `ant_01` (answer 大), each with an extended target list covering the English, Chinese and German answer forms plus two nulls. | + +### Families + +EN-ZH bank (`cross_lens_prompts_en_zh.json`): `antonym_zh` (14), `cloze_en` +(10), `cloze_zh` (18), `exemplar_zh` (8), `fact_zh` (10), `trans_en2zh` (10), +`trans_zh2en` (10); controls `ctrl_digit` (6), `ctrl_propn` (6). The antonym +family includes the prompt used in the Jacobian-lens paper's multilingual +illustration ("the opposite of 小"). + +EN-DE bank (`cross_lens_prompts_en_de.json`): `antonym_de` (9), `cloze_de` +(14), `cloze_en` (15), `exemplar_de` (10), `fact_de` (8), `trans_de2en` (17), +`trans_en2de` (17); controls `ctrl_digit` (6), `ctrl_propn` (6). + +Controls (`ctrl_digit`, `ctrl_propn`) have an expected surface form that does +not depend on language, so no lens-dependent flip is expected; they are +reported separately. + +## Provenance + +The EN-ZH bank was hand-curated: seven families of short prompts whose answer +is a common concept with a single-token realisation in both scripts, nulls +drawn from unrelated everyday nouns. The EN-DE bank was generated from a +candidate list by `scripts/data/build_cross_lens_de_bank.py` under three rules +fixed before any lens was fitted (single-token surfaces under the Qwen +tokenizer with context-appropriate spacing, cognate exclusion at normalised +Levenshtein <= 2, round-robin language-matched nulls); the exclusion report is +shipped alongside. The three-lens bank reuses `ant_01` and `antonym_de_01` +verbatim with the extended target lists used for the worked example. + +The C4 fitting corpora for the lenses (seeded English, Chinese and German +prompt dumps from `scripts/run/fit_jlens.py prompts`) are not committed, since +raw C4 web text can contain third-party PII; regenerate them with the commands +in that script's docstring. + +## Used by + +- **`scripts/run/run_cross_lens_readouts.py`** reads a bank under one fitted + lens and writes the per-layer readout dump (top-5 tokens, per-target + decompositions) that every aggregator consumes. +- **`scripts/eval/aggregate_cross_lens_en_zh.py`** and + **`scripts/eval/aggregate_cross_lens_en_de.py`** compute agreement rates, null + floors and the surface-language split from two dumps and the bank. +- **`scripts/analyze/cross_lens_three_lens_prompt.py`** tabulates one prompt + under several lenses from the three-lens dumps. +- **`scripts/figures/compute_cross_lens_shared_feature.py`** extracts the + shared-feature numbers behind the two-lens figure for one prompt and layer. + +See [`docs/REPRODUCE.md`](../../docs/REPRODUCE.md) for the artifact-to-script +mapping and the exact command lines. + +## Adapting the banks + +To test a new language pair or lens construction, write a JSON with the same +schema (`{"note": ..., "prompts": [...]}`), check that every entry of `targets` +is a single token under your tokenizer, and keep `concept` filled so the +shuffled-pairing null can exclude same-concept pairs. The EN-ZH aggregator's +surface call is by script (Latin vs CJK) and the EN-DE aggregator's is lexical; +a pair that shares a script needs the lexical variant with its word lists +adapted. diff --git a/data/cross_lens/cross_lens_prompts_en_de.json b/data/cross_lens/cross_lens_prompts_en_de.json new file mode 100644 index 0000000..719679a --- /dev/null +++ b/data/cross_lens/cross_lens_prompts_en_de.json @@ -0,0 +1,2147 @@ +{ + "note": "Cross-lens prompt bank EN-DE. Same schema and metrics as cross_lens_prompts_en_zh.json; both languages Latin-script, so the surface-language call is lexical and cognates (normalized Levenshtein <= 2) are excluded by construction. form_a = in-context answer form, form_b = other-language form; null_targets = [de, en] language-matched unrelated tokens. Built by scripts/data/build_cross_lens_de_bank.py; exclusions in the companion report cross_lens_prompts_en_de_exclusion_report.json.", + "prompts": [ + { + "id": "antonym_de_01", + "group": "antonym_de", + "concept": "big", + "prompt": "Das Gegenteil von \"hoch\" ist \"tief\". Das Gegenteil von \"klein\" ist \"", + "form_a": " groß", + "form_b": " big", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " groß", + " big", + " Salz", + " salt" + ], + "null_targets": [ + " Salz", + " salt" + ], + "answer": "groß" + }, + { + "id": "antonym_de_02", + "group": "antonym_de", + "concept": "small", + "prompt": "Das Gegenteil von \"hoch\" ist \"tief\". Das Gegenteil von \"groß\" ist \"", + "form_a": " klein", + "form_b": " small", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " klein", + " small", + " Wand", + " wall" + ], + "null_targets": [ + " Wand", + " wall" + ], + "answer": "klein" + }, + { + "id": "antonym_de_03", + "group": "antonym_de", + "concept": "slow", + "prompt": "Das Gegenteil von \"hoch\" ist \"tief\". Das Gegenteil von \"schnell\" ist \"", + "form_a": " langsam", + "form_b": " slow", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " langsam", + " slow", + " Schuh", + " shoe" + ], + "null_targets": [ + " Schuh", + " shoe" + ], + "answer": "langsam" + }, + { + "id": "antonym_de_04", + "group": "antonym_de", + "concept": "bright", + "prompt": "Das Gegenteil von \"hoch\" ist \"tief\". Das Gegenteil von \"dunkel\" ist \"", + "form_a": "hell", + "form_b": " bright", + "lang_a": "de", + "lang_b": "en", + "targets": [ + "hell", + " bright", + " Karte", + " cup" + ], + "null_targets": [ + " Karte", + " cup" + ], + "answer": "hell" + }, + { + "id": "antonym_de_06", + "group": "antonym_de", + "concept": "bad", + "prompt": "Das Gegenteil von \"hoch\" ist \"tief\". Das Gegenteil von \"gut\" ist \"", + "form_a": " schlecht", + "form_b": " bad", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " schlecht", + " bad", + " Karte", + " needle" + ], + "null_targets": [ + " Karte", + " needle" + ], + "answer": "schlecht" + }, + { + "id": "antonym_de_07", + "group": "antonym_de", + "concept": "empty", + "prompt": "Das Gegenteil von \"hoch\" ist \"tief\". Das Gegenteil von \"voll\" ist \"", + "form_a": " leer", + "form_b": " empty", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " leer", + " empty", + " Karte", + " ticket" + ], + "null_targets": [ + " Karte", + " ticket" + ], + "answer": "leer" + }, + { + "id": "antonym_de_08", + "group": "antonym_de", + "concept": "dry", + "prompt": "Das Gegenteil von \"hoch\" ist \"tief\". Das Gegenteil von \"nass\" ist \"", + "form_a": " trocken", + "form_b": " dry", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " trocken", + " dry", + " Salz", + " soap" + ], + "null_targets": [ + " Salz", + " soap" + ], + "answer": "trocken" + }, + { + "id": "antonym_de_11", + "group": "antonym_de", + "concept": "late", + "prompt": "Das Gegenteil von \"hoch\" ist \"tief\". Das Gegenteil von \"früh\" ist \"", + "form_a": " spät", + "form_b": " late", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " spät", + " late", + " Salz", + " fork" + ], + "null_targets": [ + " Salz", + " fork" + ], + "answer": "spät" + }, + { + "id": "antonym_de_12", + "group": "antonym_de", + "concept": "cheap", + "prompt": "Das Gegenteil von \"hoch\" ist \"tief\". Das Gegenteil von \"teuer\" ist \"", + "form_a": " billig", + "form_b": " cheap", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " billig", + " cheap", + " Salz", + " salt" + ], + "null_targets": [ + " Salz", + " salt" + ], + "answer": "billig" + }, + { + "id": "trans_de2en_01", + "group": "trans_de2en", + "concept": "dog", + "prompt": "Das englische Wort für \"Himmel\" ist \"sky\". Das englische Wort für \"Hund\" ist \"", + "form_a": "dog", + "form_b": " Hund", + "lang_a": "en", + "lang_b": "de", + "targets": [ + "dog", + " Hund", + " Wand", + " wall" + ], + "null_targets": [ + " Wand", + " wall" + ], + "answer": "dog" + }, + { + "id": "trans_de2en_02", + "group": "trans_de2en", + "concept": "tree", + "prompt": "Das englische Wort für \"Himmel\" ist \"sky\". Das englische Wort für \"Baum\" ist \"", + "form_a": "tree", + "form_b": " Baum", + "lang_a": "en", + "lang_b": "de", + "targets": [ + "tree", + " Baum", + " Schuh", + " shoe" + ], + "null_targets": [ + " Schuh", + " shoe" + ], + "answer": "tree" + }, + { + "id": "trans_de2en_03", + "group": "trans_de2en", + "concept": "table", + "prompt": "Das englische Wort für \"Himmel\" ist \"sky\". Das englische Wort für \"Tisch\" ist \"", + "form_a": "table", + "form_b": " Tisch", + "lang_a": "en", + "lang_b": "de", + "targets": [ + "table", + " Tisch", + " Karte", + " cup" + ], + "null_targets": [ + " Karte", + " cup" + ], + "answer": "table" + }, + { + "id": "trans_de2en_05", + "group": "trans_de2en", + "concept": "window", + "prompt": "Das englische Wort für \"Himmel\" ist \"sky\". Das englische Wort für \"Fenster\" ist \"", + "form_a": "window", + "form_b": " Fenster", + "lang_a": "en", + "lang_b": "de", + "targets": [ + "window", + " Fenster", + " Karte", + " needle" + ], + "null_targets": [ + " Karte", + " needle" + ], + "answer": "window" + }, + { + "id": "trans_de2en_06", + "group": "trans_de2en", + "concept": "horse", + "prompt": "Das englische Wort für \"Himmel\" ist \"sky\". Das englische Wort für \"Pferd\" ist \"", + "form_a": "horse", + "form_b": " Pferd", + "lang_a": "en", + "lang_b": "de", + "targets": [ + "horse", + " Pferd", + " Karte", + " ticket" + ], + "null_targets": [ + " Karte", + " ticket" + ], + "answer": "horse" + }, + { + "id": "trans_de2en_07", + "group": "trans_de2en", + "concept": "bird", + "prompt": "Das englische Wort für \"Himmel\" ist \"sky\". Das englische Wort für \"Vogel\" ist \"", + "form_a": "bird", + "form_b": " Vogel", + "lang_a": "en", + "lang_b": "de", + "targets": [ + "bird", + " Vogel", + " Salz", + " soap" + ], + "null_targets": [ + " Salz", + " soap" + ], + "answer": "bird" + }, + { + "id": "trans_de2en_08", + "group": "trans_de2en", + "concept": "mountain", + "prompt": "Das englische Wort für \"Himmel\" ist \"sky\". Das englische Wort für \"Berg\" ist \"", + "form_a": " mountain", + "form_b": " Berg", + "lang_a": "en", + "lang_b": "de", + "targets": [ + " mountain", + " Berg", + " Salz", + " fork" + ], + "null_targets": [ + " Salz", + " fork" + ], + "answer": "mountain" + }, + { + "id": "trans_de2en_09", + "group": "trans_de2en", + "concept": "rain", + "prompt": "Das englische Wort für \"Himmel\" ist \"sky\". Das englische Wort für \"Regen\" ist \"", + "form_a": "rain", + "form_b": " Regen", + "lang_a": "en", + "lang_b": "de", + "targets": [ + "rain", + " Regen", + " Salz", + " salt" + ], + "null_targets": [ + " Salz", + " salt" + ], + "answer": "rain" + }, + { + "id": "trans_de2en_12", + "group": "trans_de2en", + "concept": "door", + "prompt": "Das englische Wort für \"Himmel\" ist \"sky\". Das englische Wort für \"Tür\" ist \"", + "form_a": "door", + "form_b": " Tür", + "lang_a": "en", + "lang_b": "de", + "targets": [ + "door", + " Tür", + " Wand", + " wall" + ], + "null_targets": [ + " Wand", + " wall" + ], + "answer": "door" + }, + { + "id": "trans_en2de_01", + "group": "trans_en2de", + "concept": "dog", + "prompt": "The German word for \"sky\" is \"Himmel\". The German word for \"dog\" is \"", + "form_a": " Hund", + "form_b": " dog", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Hund", + " dog", + " Schuh", + " shoe" + ], + "null_targets": [ + " Schuh", + " shoe" + ], + "answer": "Hund" + }, + { + "id": "trans_en2de_02", + "group": "trans_en2de", + "concept": "tree", + "prompt": "The German word for \"sky\" is \"Himmel\". The German word for \"tree\" is \"", + "form_a": " Baum", + "form_b": " tree", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Baum", + " tree", + " Karte", + " cup" + ], + "null_targets": [ + " Karte", + " cup" + ], + "answer": "Baum" + }, + { + "id": "trans_en2de_03", + "group": "trans_en2de", + "concept": "table", + "prompt": "The German word for \"sky\" is \"Himmel\". The German word for \"table\" is \"", + "form_a": " Tisch", + "form_b": " table", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Tisch", + " table", + " Karte", + " needle" + ], + "null_targets": [ + " Karte", + " needle" + ], + "answer": "Tisch" + }, + { + "id": "trans_en2de_05", + "group": "trans_en2de", + "concept": "window", + "prompt": "The German word for \"sky\" is \"Himmel\". The German word for \"window\" is \"", + "form_a": " Fenster", + "form_b": " window", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Fenster", + " window", + " Karte", + " ticket" + ], + "null_targets": [ + " Karte", + " ticket" + ], + "answer": "Fenster" + }, + { + "id": "trans_en2de_06", + "group": "trans_en2de", + "concept": "horse", + "prompt": "The German word for \"sky\" is \"Himmel\". The German word for \"horse\" is \"", + "form_a": " Pferd", + "form_b": " horse", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Pferd", + " horse", + " Salz", + " soap" + ], + "null_targets": [ + " Salz", + " soap" + ], + "answer": "Pferd" + }, + { + "id": "trans_en2de_07", + "group": "trans_en2de", + "concept": "bird", + "prompt": "The German word for \"sky\" is \"Himmel\". The German word for \"bird\" is \"", + "form_a": " Vogel", + "form_b": " bird", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Vogel", + " bird", + " Salz", + " fork" + ], + "null_targets": [ + " Salz", + " fork" + ], + "answer": "Vogel" + }, + { + "id": "trans_en2de_08", + "group": "trans_en2de", + "concept": "mountain", + "prompt": "The German word for \"sky\" is \"Himmel\". The German word for \"mountain\" is \"", + "form_a": " Berg", + "form_b": " mountain", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Berg", + " mountain", + " Salz", + " salt" + ], + "null_targets": [ + " Salz", + " salt" + ], + "answer": "Berg" + }, + { + "id": "trans_en2de_09", + "group": "trans_en2de", + "concept": "rain", + "prompt": "The German word for \"sky\" is \"Himmel\". The German word for \"rain\" is \"", + "form_a": " Regen", + "form_b": " rain", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Regen", + " rain", + " Wand", + " wall" + ], + "null_targets": [ + " Wand", + " wall" + ], + "answer": "Regen" + }, + { + "id": "trans_en2de_12", + "group": "trans_en2de", + "concept": "door", + "prompt": "The German word for \"sky\" is \"Himmel\". The German word for \"door\" is \"", + "form_a": " Tür", + "form_b": " door", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Tür", + " door", + " Schuh", + " shoe" + ], + "null_targets": [ + " Schuh", + " shoe" + ], + "answer": "Tür" + }, + { + "id": "trans_de2en_15", + "group": "trans_de2en", + "concept": "train", + "prompt": "Das englische Wort für \"Himmel\" ist \"sky\". Das englische Wort für \"Zug\" ist \"", + "form_a": "train", + "form_b": " Zug", + "lang_a": "en", + "lang_b": "de", + "targets": [ + "train", + " Zug", + " Karte", + " cup" + ], + "null_targets": [ + " Karte", + " cup" + ], + "answer": "train" + }, + { + "id": "trans_de2en_16", + "group": "trans_de2en", + "concept": "city", + "prompt": "Das englische Wort für \"Himmel\" ist \"sky\". Das englische Wort für \"Stadt\" ist \"", + "form_a": "city", + "form_b": " Stadt", + "lang_a": "en", + "lang_b": "de", + "targets": [ + "city", + " Stadt", + " Karte", + " needle" + ], + "null_targets": [ + " Karte", + " needle" + ], + "answer": "city" + }, + { + "id": "trans_de2en_17", + "group": "trans_de2en", + "concept": "forest", + "prompt": "Das englische Wort für \"Himmel\" ist \"sky\". Das englische Wort für \"Wald\" ist \"", + "form_a": "forest", + "form_b": " Wald", + "lang_a": "en", + "lang_b": "de", + "targets": [ + "forest", + " Wald", + " Karte", + " ticket" + ], + "null_targets": [ + " Karte", + " ticket" + ], + "answer": "forest" + }, + { + "id": "trans_de2en_18", + "group": "trans_de2en", + "concept": "sea", + "prompt": "Das englische Wort für \"Himmel\" ist \"sky\". Das englische Wort für \"Meer\" ist \"", + "form_a": "sea", + "form_b": " Meer", + "lang_a": "en", + "lang_b": "de", + "targets": [ + "sea", + " Meer", + " Salz", + " soap" + ], + "null_targets": [ + " Salz", + " soap" + ], + "answer": "sea" + }, + { + "id": "trans_de2en_19", + "group": "trans_de2en", + "concept": "fire", + "prompt": "Das englische Wort für \"Himmel\" ist \"sky\". Das englische Wort für \"Feuer\" ist \"", + "form_a": "fire", + "form_b": " Feuer", + "lang_a": "en", + "lang_b": "de", + "targets": [ + "fire", + " Feuer", + " Salz", + " fork" + ], + "null_targets": [ + " Salz", + " fork" + ], + "answer": "fire" + }, + { + "id": "trans_de2en_20", + "group": "trans_de2en", + "concept": "newspaper", + "prompt": "Das englische Wort für \"Himmel\" ist \"sky\". Das englische Wort für \"Zeitung\" ist \"", + "form_a": " newspaper", + "form_b": " Zeitung", + "lang_a": "en", + "lang_b": "de", + "targets": [ + " newspaper", + " Zeitung", + " Salz", + " salt" + ], + "null_targets": [ + " Salz", + " salt" + ], + "answer": "newspaper" + }, + { + "id": "trans_de2en_21", + "group": "trans_de2en", + "concept": "kitchen", + "prompt": "Das englische Wort für \"Himmel\" ist \"sky\". Das englische Wort für \"Küche\" ist \"", + "form_a": " kitchen", + "form_b": " Küche", + "lang_a": "en", + "lang_b": "de", + "targets": [ + " kitchen", + " Küche", + " Wand", + " wall" + ], + "null_targets": [ + " Wand", + " wall" + ], + "answer": "kitchen" + }, + { + "id": "trans_de2en_22", + "group": "trans_de2en", + "concept": "sister", + "prompt": "Das englische Wort für \"Himmel\" ist \"sky\". Das englische Wort für \"Schwester\" ist \"", + "form_a": " sister", + "form_b": " Schwester", + "lang_a": "en", + "lang_b": "de", + "targets": [ + " sister", + " Schwester", + " Schuh", + " shoe" + ], + "null_targets": [ + " Schuh", + " shoe" + ], + "answer": "sister" + }, + { + "id": "trans_en2de_15", + "group": "trans_en2de", + "concept": "train", + "prompt": "The German word for \"sky\" is \"Himmel\". The German word for \"train\" is \"", + "form_a": " Zug", + "form_b": " train", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Zug", + " train", + " Karte", + " cup" + ], + "null_targets": [ + " Karte", + " cup" + ], + "answer": "Zug" + }, + { + "id": "trans_en2de_16", + "group": "trans_en2de", + "concept": "city", + "prompt": "The German word for \"sky\" is \"Himmel\". The German word for \"city\" is \"", + "form_a": "Stadt", + "form_b": " city", + "lang_a": "de", + "lang_b": "en", + "targets": [ + "Stadt", + " city", + " Karte", + " needle" + ], + "null_targets": [ + " Karte", + " needle" + ], + "answer": "Stadt" + }, + { + "id": "trans_en2de_17", + "group": "trans_en2de", + "concept": "forest", + "prompt": "The German word for \"sky\" is \"Himmel\". The German word for \"forest\" is \"", + "form_a": " Wald", + "form_b": " forest", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Wald", + " forest", + " Karte", + " ticket" + ], + "null_targets": [ + " Karte", + " ticket" + ], + "answer": "Wald" + }, + { + "id": "trans_en2de_18", + "group": "trans_en2de", + "concept": "sea", + "prompt": "The German word for \"sky\" is \"Himmel\". The German word for \"sea\" is \"", + "form_a": " Meer", + "form_b": " sea", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Meer", + " sea", + " Salz", + " soap" + ], + "null_targets": [ + " Salz", + " soap" + ], + "answer": "Meer" + }, + { + "id": "trans_en2de_19", + "group": "trans_en2de", + "concept": "fire", + "prompt": "The German word for \"sky\" is \"Himmel\". The German word for \"fire\" is \"", + "form_a": " Feuer", + "form_b": " fire", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Feuer", + " fire", + " Salz", + " fork" + ], + "null_targets": [ + " Salz", + " fork" + ], + "answer": "Feuer" + }, + { + "id": "trans_en2de_20", + "group": "trans_en2de", + "concept": "newspaper", + "prompt": "The German word for \"sky\" is \"Himmel\". The German word for \"newspaper\" is \"", + "form_a": " Zeitung", + "form_b": " newspaper", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Zeitung", + " newspaper", + " Salz", + " salt" + ], + "null_targets": [ + " Salz", + " salt" + ], + "answer": "Zeitung" + }, + { + "id": "trans_en2de_21", + "group": "trans_en2de", + "concept": "kitchen", + "prompt": "The German word for \"sky\" is \"Himmel\". The German word for \"kitchen\" is \"", + "form_a": " Küche", + "form_b": " kitchen", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Küche", + " kitchen", + " Wand", + " wall" + ], + "null_targets": [ + " Wand", + " wall" + ], + "answer": "Küche" + }, + { + "id": "trans_en2de_22", + "group": "trans_en2de", + "concept": "sister", + "prompt": "The German word for \"sky\" is \"Himmel\". The German word for \"sister\" is \"", + "form_a": " Schwester", + "form_b": " sister", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Schwester", + " sister", + " Schuh", + " shoe" + ], + "null_targets": [ + " Schuh", + " shoe" + ], + "answer": "Schwester" + }, + { + "id": "cloze_de_01", + "group": "cloze_de", + "concept": "pen", + "prompt": "Zum Schreiben benutze ich einen", + "form_a": " Stift", + "form_b": " pen", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Stift", + " pen", + " Karte", + " cup" + ], + "null_targets": [ + " Karte", + " cup" + ], + "answer": "Stift" + }, + { + "id": "cloze_de_03", + "group": "cloze_de", + "concept": "station", + "prompt": "Der Zug hält pünktlich am", + "form_a": " Bahnhof", + "form_b": " station", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Bahnhof", + " station", + " Karte", + " needle" + ], + "null_targets": [ + " Karte", + " needle" + ], + "answer": "Bahnhof" + }, + { + "id": "cloze_de_04", + "group": "cloze_de", + "concept": "door", + "prompt": "Ich klopfe dreimal und öffne dann die", + "form_a": " Tür", + "form_b": " door", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Tür", + " door", + " Karte", + " ticket" + ], + "null_targets": [ + " Karte", + " ticket" + ], + "answer": "Tür" + }, + { + "id": "cloze_de_05", + "group": "cloze_de", + "concept": "book", + "prompt": "Vor dem Schlafen liest das Kind noch ein", + "form_a": " Buch", + "form_b": " book", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Buch", + " book", + " Salz", + " soap" + ], + "null_targets": [ + " Salz", + " soap" + ], + "answer": "Buch" + }, + { + "id": "cloze_de_06", + "group": "cloze_de", + "concept": "tree", + "prompt": "Der Vogel sitzt hoch oben auf dem", + "form_a": " Baum", + "form_b": " tree", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Baum", + " tree", + " Salz", + " fork" + ], + "null_targets": [ + " Salz", + " fork" + ], + "answer": "Baum" + }, + { + "id": "cloze_de_07", + "group": "cloze_de", + "concept": "snow", + "prompt": "Im Winter fällt auf den Bergen viel", + "form_a": " Schnee", + "form_b": " snow", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Schnee", + " snow", + " Salz", + " salt" + ], + "null_targets": [ + " Salz", + " salt" + ], + "answer": "Schnee" + }, + { + "id": "cloze_de_09", + "group": "cloze_de", + "concept": "key", + "prompt": "Ich schließe die Wohnung ab mit dem", + "form_a": " Schlüssel", + "form_b": " key", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Schlüssel", + " key", + " Wand", + " wall" + ], + "null_targets": [ + " Wand", + " wall" + ], + "answer": "Schlüssel" + }, + { + "id": "cloze_de_10", + "group": "cloze_de", + "concept": "horse", + "prompt": "Auf dem Bauernhof reitet das Mädchen ein", + "form_a": " Pferd", + "form_b": " horse", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Pferd", + " horse", + " Schuh", + " shoe" + ], + "null_targets": [ + " Schuh", + " shoe" + ], + "answer": "Pferd" + }, + { + "id": "cloze_de_13", + "group": "cloze_de", + "concept": "school", + "prompt": "Nach den Ferien gehen die Kinder wieder in die", + "form_a": " Schule", + "form_b": " school", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Schule", + " school", + " Karte", + " cup" + ], + "null_targets": [ + " Karte", + " cup" + ], + "answer": "Schule" + }, + { + "id": "cloze_de_14", + "group": "cloze_de", + "concept": "church", + "prompt": "Sonntags läuten die Glocken der alten", + "form_a": " Kirche", + "form_b": " church", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Kirche", + " church", + " Karte", + " needle" + ], + "null_targets": [ + " Karte", + " needle" + ], + "answer": "Kirche" + }, + { + "id": "cloze_de_15", + "group": "cloze_de", + "concept": "train", + "prompt": "Wir fahren in den Urlaub bequem mit dem", + "form_a": " Zug", + "form_b": " train", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Zug", + " train", + " Karte", + " ticket" + ], + "null_targets": [ + " Karte", + " ticket" + ], + "answer": "Zug" + }, + { + "id": "cloze_de_16", + "group": "cloze_de", + "concept": "forest", + "prompt": "Am Wochenende wandern wir stundenlang durch den", + "form_a": " Wald", + "form_b": " forest", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Wald", + " forest", + " Salz", + " soap" + ], + "null_targets": [ + " Salz", + " soap" + ], + "answer": "Wald" + }, + { + "id": "cloze_de_17", + "group": "cloze_de", + "concept": "ship", + "prompt": "Der Kapitän steuert das große", + "form_a": " Schiff", + "form_b": " ship", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Schiff", + " ship", + " Salz", + " fork" + ], + "null_targets": [ + " Salz", + " fork" + ], + "answer": "Schiff" + }, + { + "id": "cloze_de_18", + "group": "cloze_de", + "concept": "money", + "prompt": "Für die Fahrkarte brauche ich noch etwas", + "form_a": " Geld", + "form_b": " money", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Geld", + " money", + " Salz", + " salt" + ], + "null_targets": [ + " Salz", + " salt" + ], + "answer": "Geld" + }, + { + "id": "cloze_en_01", + "group": "cloze_en", + "concept": "key", + "prompt": "I unlock the front door with my", + "form_a": " key", + "form_b": " Schlüssel", + "lang_a": "en", + "lang_b": "de", + "targets": [ + " key", + " Schlüssel", + " Wand", + " wall" + ], + "null_targets": [ + " Wand", + " wall" + ], + "answer": "key" + }, + { + "id": "cloze_en_02", + "group": "cloze_en", + "concept": "horse", + "prompt": "At the farm the girl rides a brown", + "form_a": " horse", + "form_b": " Pferd", + "lang_a": "en", + "lang_b": "de", + "targets": [ + " horse", + " Pferd", + " Schuh", + " shoe" + ], + "null_targets": [ + " Schuh", + " shoe" + ], + "answer": "horse" + }, + { + "id": "cloze_en_03", + "group": "cloze_en", + "concept": "window", + "prompt": "She waved at us through the kitchen", + "form_a": " window", + "form_b": " Fenster", + "lang_a": "en", + "lang_b": "de", + "targets": [ + " window", + " Fenster", + " Karte", + " cup" + ], + "null_targets": [ + " Karte", + " cup" + ], + "answer": "window" + }, + { + "id": "cloze_en_04", + "group": "cloze_en", + "concept": "tree", + "prompt": "The bird built its nest high up in the", + "form_a": " tree", + "form_b": " Baum", + "lang_a": "en", + "lang_b": "de", + "targets": [ + " tree", + " Baum", + " Karte", + " needle" + ], + "null_targets": [ + " Karte", + " needle" + ], + "answer": "tree" + }, + { + "id": "cloze_en_05", + "group": "cloze_en", + "concept": "table", + "prompt": "He put the plates down on the wooden", + "form_a": " table", + "form_b": " Tisch", + "lang_a": "en", + "lang_b": "de", + "targets": [ + " table", + " Tisch", + " Karte", + " ticket" + ], + "null_targets": [ + " Karte", + " ticket" + ], + "answer": "table" + }, + { + "id": "cloze_en_06", + "group": "cloze_en", + "concept": "dog", + "prompt": "The postman was chased down the street by a barking", + "form_a": " dog", + "form_b": " Hund", + "lang_a": "en", + "lang_b": "de", + "targets": [ + " dog", + " Hund", + " Salz", + " soap" + ], + "null_targets": [ + " Salz", + " soap" + ], + "answer": "dog" + }, + { + "id": "cloze_en_07", + "group": "cloze_en", + "concept": "rain", + "prompt": "Take an umbrella, the forecast says heavy", + "form_a": " rain", + "form_b": " Regen", + "lang_a": "en", + "lang_b": "de", + "targets": [ + " rain", + " Regen", + " Salz", + " fork" + ], + "null_targets": [ + " Salz", + " fork" + ], + "answer": "rain" + }, + { + "id": "cloze_en_10", + "group": "cloze_en", + "concept": "mountain", + "prompt": "They spent all summer climbing the highest", + "form_a": " mountain", + "form_b": " Berg", + "lang_a": "en", + "lang_b": "de", + "targets": [ + " mountain", + " Berg", + " Salz", + " salt" + ], + "null_targets": [ + " Salz", + " salt" + ], + "answer": "mountain" + }, + { + "id": "cloze_en_12", + "group": "cloze_en", + "concept": "bird", + "prompt": "High above the field we watched a small", + "form_a": " bird", + "form_b": " Vogel", + "lang_a": "en", + "lang_b": "de", + "targets": [ + " bird", + " Vogel", + " Wand", + " wall" + ], + "null_targets": [ + " Wand", + " wall" + ], + "answer": "bird" + }, + { + "id": "cloze_en_13", + "group": "cloze_en", + "concept": "train", + "prompt": "We traveled to the capital on a high-speed", + "form_a": " train", + "form_b": " Zug", + "lang_a": "en", + "lang_b": "de", + "targets": [ + " train", + " Zug", + " Schuh", + " shoe" + ], + "null_targets": [ + " Schuh", + " shoe" + ], + "answer": "train" + }, + { + "id": "cloze_en_14", + "group": "cloze_en", + "concept": "money", + "prompt": "He paid for the tickets with borrowed", + "form_a": " money", + "form_b": " Geld", + "lang_a": "en", + "lang_b": "de", + "targets": [ + " money", + " Geld", + " Karte", + " cup" + ], + "null_targets": [ + " Karte", + " cup" + ], + "answer": "money" + }, + { + "id": "cloze_en_15", + "group": "cloze_en", + "concept": "voice", + "prompt": "The choir admired her beautiful singing", + "form_a": " voice", + "form_b": " Stimme", + "lang_a": "en", + "lang_b": "de", + "targets": [ + " voice", + " Stimme", + " Karte", + " needle" + ], + "null_targets": [ + " Karte", + " needle" + ], + "answer": "voice" + }, + { + "id": "cloze_en_16", + "group": "cloze_en", + "concept": "newspaper", + "prompt": "Every morning my grandfather reads the", + "form_a": " newspaper", + "form_b": " Zeitung", + "lang_a": "en", + "lang_b": "de", + "targets": [ + " newspaper", + " Zeitung", + " Karte", + " ticket" + ], + "null_targets": [ + " Karte", + " ticket" + ], + "answer": "newspaper" + }, + { + "id": "cloze_en_17", + "group": "cloze_en", + "concept": "war", + "prompt": "After years of peace the country slid into", + "form_a": " war", + "form_b": " Krieg", + "lang_a": "en", + "lang_b": "de", + "targets": [ + " war", + " Krieg", + " Salz", + " soap" + ], + "null_targets": [ + " Salz", + " soap" + ], + "answer": "war" + }, + { + "id": "cloze_en_18", + "group": "cloze_en", + "concept": "city", + "prompt": "From the hilltop you can see the whole", + "form_a": " city", + "form_b": " Stadt", + "lang_a": "en", + "lang_b": "de", + "targets": [ + " city", + " Stadt", + " Salz", + " fork" + ], + "null_targets": [ + " Salz", + " fork" + ], + "answer": "city" + }, + { + "id": "fact_de_02", + "group": "fact_de", + "concept": "dog", + "prompt": "Der beste Freund des Menschen ist der", + "form_a": " Hund", + "form_b": " dog", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Hund", + " dog", + " Salz", + " salt" + ], + "null_targets": [ + " Salz", + " salt" + ], + "answer": "Hund" + }, + { + "id": "fact_de_05", + "group": "fact_de", + "concept": "snow", + "prompt": "Weiß und kalt fällt im Winter der", + "form_a": " Schnee", + "form_b": " snow", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Schnee", + " snow", + " Wand", + " wall" + ], + "null_targets": [ + " Wand", + " wall" + ], + "answer": "Schnee" + }, + { + "id": "fact_de_07", + "group": "fact_de", + "concept": "horse", + "prompt": "Vor die Kutsche spannt man ein", + "form_a": " Pferd", + "form_b": " horse", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Pferd", + " horse", + " Schuh", + " shoe" + ], + "null_targets": [ + " Schuh", + " shoe" + ], + "answer": "Pferd" + }, + { + "id": "fact_de_09", + "group": "fact_de", + "concept": "week", + "prompt": "Sieben Tage sind eine", + "form_a": " Woche", + "form_b": " week", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Woche", + " week", + " Karte", + " cup" + ], + "null_targets": [ + " Karte", + " cup" + ], + "answer": "Woche" + }, + { + "id": "fact_de_10", + "group": "fact_de", + "concept": "year", + "prompt": "Zwölf Monate sind ein", + "form_a": " Jahr", + "form_b": " year", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Jahr", + " year", + " Karte", + " needle" + ], + "null_targets": [ + " Karte", + " needle" + ], + "answer": "Jahr" + }, + { + "id": "fact_de_11", + "group": "fact_de", + "concept": "tooth", + "prompt": "Der Zahnarzt zieht den kranken", + "form_a": " Zahn", + "form_b": " tooth", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Zahn", + " tooth", + " Karte", + " ticket" + ], + "null_targets": [ + " Karte", + " ticket" + ], + "answer": "Zahn" + }, + { + "id": "fact_de_12", + "group": "fact_de", + "concept": "island", + "prompt": "Ein Stück Land mitten im Meer nennt man eine", + "form_a": " Insel", + "form_b": " island", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Insel", + " island", + " Salz", + " soap" + ], + "null_targets": [ + " Salz", + " soap" + ], + "answer": "Insel" + }, + { + "id": "fact_de_13", + "group": "fact_de", + "concept": "church", + "prompt": "Am Sonntag gehen viele Menschen in die", + "form_a": " Kirche", + "form_b": " church", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Kirche", + " church", + " Salz", + " fork" + ], + "null_targets": [ + " Salz", + " fork" + ], + "answer": "Kirche" + }, + { + "id": "exemplar_de_01", + "group": "exemplar_de", + "concept": "dog", + "prompt": "Ein Beispiel für ein Haustier ist der", + "form_a": " Hund", + "form_b": " dog", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Hund", + " dog", + " Salz", + " salt" + ], + "null_targets": [ + " Salz", + " salt" + ], + "answer": "Hund" + }, + { + "id": "exemplar_de_02", + "group": "exemplar_de", + "concept": "table", + "prompt": "Ein Beispiel für ein Möbelstück ist der", + "form_a": " Tisch", + "form_b": " table", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Tisch", + " table", + " Wand", + " wall" + ], + "null_targets": [ + " Wand", + " wall" + ], + "answer": "Tisch" + }, + { + "id": "exemplar_de_04", + "group": "exemplar_de", + "concept": "doctor", + "prompt": "Ein Beispiel für einen Beruf ist der", + "form_a": " Arzt", + "form_b": " doctor", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Arzt", + " doctor", + " Schuh", + " shoe" + ], + "null_targets": [ + " Schuh", + " shoe" + ], + "answer": "Arzt" + }, + { + "id": "exemplar_de_05", + "group": "exemplar_de", + "concept": "teacher", + "prompt": "Ein Beispiel für einen Beruf ist der", + "form_a": " Lehrer", + "form_b": " teacher", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Lehrer", + " teacher", + " Karte", + " cup" + ], + "null_targets": [ + " Karte", + " cup" + ], + "answer": "Lehrer" + }, + { + "id": "exemplar_de_08", + "group": "exemplar_de", + "concept": "bird", + "prompt": "Ein Beispiel für ein Tier mit Flügeln ist der", + "form_a": " Vogel", + "form_b": " bird", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Vogel", + " bird", + " Karte", + " needle" + ], + "null_targets": [ + " Karte", + " needle" + ], + "answer": "Vogel" + }, + { + "id": "exemplar_de_09", + "group": "exemplar_de", + "concept": "head", + "prompt": "Ein Beispiel für ein Körperteil ist der", + "form_a": " Kopf", + "form_b": " head", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Kopf", + " head", + " Karte", + " ticket" + ], + "null_targets": [ + " Karte", + " ticket" + ], + "answer": "Kopf" + }, + { + "id": "exemplar_de_10", + "group": "exemplar_de", + "concept": "sea", + "prompt": "Ein Beispiel für ein Gewässer ist das", + "form_a": " Meer", + "form_b": " sea", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Meer", + " sea", + " Salz", + " soap" + ], + "null_targets": [ + " Salz", + " soap" + ], + "answer": "Meer" + }, + { + "id": "exemplar_de_11", + "group": "exemplar_de", + "concept": "train", + "prompt": "Ein Beispiel für ein Verkehrsmittel ist der", + "form_a": " Zug", + "form_b": " train", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Zug", + " train", + " Salz", + " fork" + ], + "null_targets": [ + " Salz", + " fork" + ], + "answer": "Zug" + }, + { + "id": "exemplar_de_12", + "group": "exemplar_de", + "concept": "church", + "prompt": "Ein Beispiel für ein Gebäude ist die", + "form_a": " Kirche", + "form_b": " church", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Kirche", + " church", + " Salz", + " salt" + ], + "null_targets": [ + " Salz", + " salt" + ], + "answer": "Kirche" + }, + { + "id": "exemplar_de_13", + "group": "exemplar_de", + "concept": "sister", + "prompt": "Ein Beispiel für ein Familienmitglied ist die", + "form_a": " Schwester", + "form_b": " sister", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Schwester", + " sister", + " Wand", + " wall" + ], + "null_targets": [ + " Wand", + " wall" + ], + "answer": "Schwester" + }, + { + "id": "ctrl_digit_01", + "group": "ctrl_digit", + "concept": "seven", + "prompt": "3 + 4 =", + "form_a": "7", + "form_b": "7", + "lang_a": "de", + "lang_b": "en", + "targets": [ + "7", + "7", + " Schuh", + " shoe" + ], + "null_targets": [ + " Schuh", + " shoe" + ], + "answer": "7" + }, + { + "id": "ctrl_digit_02", + "group": "ctrl_digit", + "concept": "nine", + "prompt": "5 + 4 =", + "form_a": "9", + "form_b": "9", + "lang_a": "de", + "lang_b": "en", + "targets": [ + "9", + "9", + " Karte", + " cup" + ], + "null_targets": [ + " Karte", + " cup" + ], + "answer": "9" + }, + { + "id": "ctrl_digit_03", + "group": "ctrl_digit", + "concept": "eight", + "prompt": "6 + 2 =", + "form_a": "8", + "form_b": "8", + "lang_a": "de", + "lang_b": "en", + "targets": [ + "8", + "8", + " Karte", + " needle" + ], + "null_targets": [ + " Karte", + " needle" + ], + "answer": "8" + }, + { + "id": "ctrl_digit_04", + "group": "ctrl_digit", + "concept": "six", + "prompt": "2 + 4 =", + "form_a": "6", + "form_b": "6", + "lang_a": "de", + "lang_b": "en", + "targets": [ + "6", + "6", + " Karte", + " ticket" + ], + "null_targets": [ + " Karte", + " ticket" + ], + "answer": "6" + }, + { + "id": "ctrl_digit_05", + "group": "ctrl_digit", + "concept": "five", + "prompt": "1 + 4 =", + "form_a": "5", + "form_b": "5", + "lang_a": "de", + "lang_b": "en", + "targets": [ + "5", + "5", + " Salz", + " soap" + ], + "null_targets": [ + " Salz", + " soap" + ], + "answer": "5" + }, + { + "id": "ctrl_digit_06", + "group": "ctrl_digit", + "concept": "four", + "prompt": "1 + 3 =", + "form_a": "4", + "form_b": "4", + "lang_a": "de", + "lang_b": "en", + "targets": [ + "4", + "4", + " Salz", + " fork" + ], + "null_targets": [ + " Salz", + " fork" + ], + "answer": "4" + }, + { + "id": "ctrl_propn_01", + "group": "ctrl_propn", + "concept": "paris", + "prompt": "Die Hauptstadt von Frankreich heißt", + "form_a": " Paris", + "form_b": " Paris", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Paris", + " Paris", + " Salz", + " salt" + ], + "null_targets": [ + " Salz", + " salt" + ], + "answer": "Paris" + }, + { + "id": "ctrl_propn_02", + "group": "ctrl_propn", + "concept": "berlin", + "prompt": "Die Hauptstadt von Deutschland heißt", + "form_a": " Berlin", + "form_b": " Berlin", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Berlin", + " Berlin", + " Wand", + " wall" + ], + "null_targets": [ + " Wand", + " wall" + ], + "answer": "Berlin" + }, + { + "id": "ctrl_propn_03", + "group": "ctrl_propn", + "concept": "london", + "prompt": "Die Hauptstadt von England heißt", + "form_a": " London", + "form_b": " London", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " London", + " London", + " Schuh", + " shoe" + ], + "null_targets": [ + " Schuh", + " shoe" + ], + "answer": "London" + }, + { + "id": "ctrl_propn_04", + "group": "ctrl_propn", + "concept": "madrid", + "prompt": "Die Hauptstadt von Spanien heißt", + "form_a": " Madrid", + "form_b": " Madrid", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Madrid", + " Madrid", + " Karte", + " cup" + ], + "null_targets": [ + " Karte", + " cup" + ], + "answer": "Madrid" + }, + { + "id": "ctrl_propn_05", + "group": "ctrl_propn", + "concept": "oslo", + "prompt": "Die Hauptstadt von Norwegen heißt", + "form_a": " Oslo", + "form_b": " Oslo", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Oslo", + " Oslo", + " Karte", + " needle" + ], + "null_targets": [ + " Karte", + " needle" + ], + "answer": "Oslo" + }, + { + "id": "ctrl_propn_06", + "group": "ctrl_propn", + "concept": "lima", + "prompt": "Die Hauptstadt von Peru heißt", + "form_a": " Lima", + "form_b": " Lima", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " Lima", + " Lima", + " Karte", + " ticket" + ], + "null_targets": [ + " Karte", + " ticket" + ], + "answer": "Lima" + } + ] +} \ No newline at end of file diff --git a/data/cross_lens/cross_lens_prompts_en_de_exclusion_report.json b/data/cross_lens/cross_lens_prompts_en_de_exclusion_report.json new file mode 100644 index 0000000..7ae6f7d --- /dev/null +++ b/data/cross_lens/cross_lens_prompts_en_de_exclusion_report.json @@ -0,0 +1,230 @@ +{ + "tokenizer": "Qwen/Qwen3.5-9B", + "n_candidates": 132, + "n_kept": 102, + "n_cross": 90, + "n_ctrl": 12, + "per_group": { + "antonym_de": 9, + "cloze_de": 14, + "cloze_en": 15, + "ctrl_digit": 6, + "ctrl_propn": 6, + "exemplar_de": 10, + "fact_de": 8, + "trans_de2en": 17, + "trans_en2de": 17 + }, + "excluded": [ + { + "id": "antonym_de_05", + "concept": "light", + "de": "leicht", + "en": "light", + "reason": "cognate (lev=2)" + }, + { + "id": "antonym_de_09", + "concept": "weak", + "de": "schwach", + "en": "weak", + "reason": "not single-token: de " + }, + { + "id": "antonym_de_10", + "concept": "dark", + "de": "dunkel", + "en": "dark", + "reason": "not single-token: de " + }, + { + "id": "antonym_de_13", + "concept": "new", + "de": "neu", + "en": "new", + "reason": "cognate (lev=1)" + }, + { + "id": "antonym_de_14", + "concept": "quiet", + "de": "leise", + "en": "quiet", + "reason": "not single-token: de " + }, + { + "id": "trans_de2en_04", + "concept": "chair", + "de": "Stuhl", + "en": "chair", + "reason": "not single-token: de " + }, + { + "id": "trans_de2en_10", + "concept": "cheese", + "de": "Käse", + "en": "cheese", + "reason": "not single-token: de " + }, + { + "id": "trans_de2en_11", + "concept": "egg", + "de": "Ei", + "en": "egg", + "reason": "cognate (lev=2)" + }, + { + "id": "trans_de2en_13", + "concept": "flower", + "de": "Blume", + "en": "flower", + "reason": "not single-token: de " + }, + { + "id": "trans_de2en_14", + "concept": "knife", + "de": "Messer", + "en": "knife", + "reason": "not single-token: de " + }, + { + "id": "trans_en2de_04", + "concept": "chair", + "de": "Stuhl", + "en": "chair", + "reason": "not single-token: de " + }, + { + "id": "trans_en2de_10", + "concept": "cheese", + "de": "Käse", + "en": "cheese", + "reason": "not single-token: de " + }, + { + "id": "trans_en2de_11", + "concept": "egg", + "de": "Ei", + "en": "egg", + "reason": "cognate (lev=2)" + }, + { + "id": "trans_en2de_13", + "concept": "flower", + "de": "Blume", + "en": "flower", + "reason": "not single-token: de " + }, + { + "id": "trans_en2de_14", + "concept": "knife", + "de": "Messer", + "en": "knife", + "reason": "not single-token: de " + }, + { + "id": "cloze_de_02", + "concept": "knife", + "de": "Messer", + "en": "knife", + "reason": "not single-token: de " + }, + { + "id": "cloze_de_08", + "concept": "coffee", + "de": "Kaffee", + "en": "coffee", + "reason": "cognate (lev=2)" + }, + { + "id": "cloze_de_11", + "concept": "cheese", + "de": "Käse", + "en": "cheese", + "reason": "not single-token: de " + }, + { + "id": "cloze_de_12", + "concept": "flower", + "de": "Blume", + "en": "flower", + "reason": "not single-token: de " + }, + { + "id": "cloze_en_08", + "concept": "cheese", + "de": "Käse", + "en": "cheese", + "reason": "not single-token: de " + }, + { + "id": "cloze_en_09", + "concept": "egg", + "de": "Ei", + "en": "egg", + "reason": "cognate (lev=2)" + }, + { + "id": "cloze_en_11", + "concept": "knife", + "de": "Messer", + "en": "knife", + "reason": "not single-token: de " + }, + { + "id": "fact_de_01", + "concept": "lion", + "de": "Löwe", + "en": "lion", + "reason": "not single-token: de " + }, + { + "id": "fact_de_03", + "concept": "bee", + "de": "Biene", + "en": "bee", + "reason": "cognate (lev=2)" + }, + { + "id": "fact_de_04", + "concept": "moon", + "de": "Mond", + "en": "moon", + "reason": "cognate (lev=2)" + }, + { + "id": "fact_de_06", + "concept": "egg", + "de": "Ei", + "en": "egg", + "reason": "cognate (lev=2)" + }, + { + "id": "fact_de_08", + "concept": "cheese", + "de": "Käse", + "en": "cheese", + "reason": "not single-token: de " + }, + { + "id": "exemplar_de_03", + "concept": "juice", + "de": "Saft", + "en": "juice", + "reason": "not single-token: de " + }, + { + "id": "exemplar_de_06", + "concept": "coat", + "de": "Mantel", + "en": "coat", + "reason": "not single-token: de " + }, + { + "id": "exemplar_de_07", + "concept": "shirt", + "de": "Hemd", + "en": "shirt", + "reason": "not single-token: de " + } + ] +} \ No newline at end of file diff --git a/data/cross_lens/cross_lens_prompts_en_zh.json b/data/cross_lens/cross_lens_prompts_en_zh.json new file mode 100644 index 0000000..b6882da --- /dev/null +++ b/data/cross_lens/cross_lens_prompts_en_zh.json @@ -0,0 +1,105 @@ +{ + "note": "Cross-lens prompt bank v2 (scaled + controlled). Groups: antonym_zh, trans_zh2en, trans_en2zh, cloze_zh, cloze_en, fact_zh, exemplar_zh (cross-script positives); ctrl_digit, ctrl_propn (script-invariant controls, expect NO lens-dependent flip). form_a/form_b = the two canonical surface forms of the concept (answer-script first); null_targets = script-matched unrelated tokens for the within-lens shared-feature floor. All targets (incl. nulls) are decomposed by the runner; aggregation is offline (agg_cross_lens_v2.py).", + "prompts": [ + {"id": "ant_01", "group": "antonym_zh", "concept": "big", "prompt": "“高”的反义词是“低”。“小”的反义词是“", "form_a": "大", "form_b": " big", "targets": ["大", " big", " large", "桌", " chair"], "null_targets": ["桌", " chair"], "answer": "大"}, + {"id": "ant_02", "group": "antonym_zh", "concept": "slow", "prompt": "“高”的反义词是“低”。“快”的反义词是“", "form_a": "慢", "form_b": " slow", "targets": ["慢", " slow", "门", " door"], "null_targets": ["门", " door"], "answer": "慢"}, + {"id": "ant_03", "group": "antonym_zh", "concept": "low", "prompt": "“大”的反义词是“小”。“高”的反义词是“", "form_a": "低", "form_b": " low", "targets": ["低", " low", "纸", " paper"], "null_targets": ["纸", " paper"], "answer": "低"}, + {"id": "ant_04", "group": "antonym_zh", "concept": "short", "prompt": "“大”的反义词是“小”。“长”的反义词是“", "form_a": "短", "form_b": " short", "targets": ["短", " short", "鞋", " shoe"], "null_targets": ["鞋", " shoe"], "answer": "短"}, + {"id": "ant_05", "group": "antonym_zh", "concept": "old", "prompt": "“大”的反义词是“小”。“新”的反义词是“", "form_a": "旧", "form_b": " old", "targets": ["旧", " old", "盐", " salt"], "null_targets": ["盐", " salt"], "answer": "旧"}, + {"id": "ant_06", "group": "antonym_zh", "concept": "hot", "prompt": "“大”的反义词是“小”。“冷”的反义词是“", "form_a": "热", "form_b": " hot", "targets": ["热", " hot", "墙", " wall"], "null_targets": ["墙", " wall"], "answer": "热"}, + {"id": "ant_07", "group": "antonym_zh", "concept": "bad", "prompt": "“大”的反义词是“小”。“好”的反义词是“", "form_a": "坏", "form_b": " bad", "targets": ["坏", " bad", "针", " needle"], "null_targets": ["针", " needle"], "answer": "坏"}, + {"id": "ant_08", "group": "antonym_zh", "concept": "few", "prompt": "“大”的反义词是“小”。“多”的反义词是“", "form_a": "少", "form_b": " few", "targets": ["少", " few", " less", "伞", " umbrella"], "null_targets": ["伞", " umbrella"], "answer": "少"}, + {"id": "ant_09", "group": "antonym_zh", "concept": "weak", "prompt": "“大”的反义词是“小”。“强”的反义词是“", "form_a": "弱", "form_b": " weak", "targets": ["弱", " weak", "票", " ticket"], "null_targets": ["票", " ticket"], "answer": "弱"}, + {"id": "ant_10", "group": "antonym_zh", "concept": "white", "prompt": "“大”的反义词是“小”。“黑”的反义词是“", "form_a": "白", "form_b": " white", "targets": ["白", " white", "杯", " cup"], "null_targets": ["杯", " cup"], "answer": "白"}, + {"id": "ant_11", "group": "antonym_zh", "concept": "late", "prompt": "“大”的反义词是“小”。“早”的反义词是“", "form_a": "晚", "form_b": " late", "targets": ["晚", " late", "桌", " table"], "null_targets": ["桌", " table"], "answer": "晚"}, + {"id": "ant_12", "group": "antonym_zh", "concept": "near", "prompt": "“大”的反义词是“小”。“远”的反义词是“", "form_a": "近", "form_b": " near", "targets": ["近", " near", " close", "盐", " wall"], "null_targets": ["盐", " wall"], "answer": "近"}, + {"id": "ant_13", "group": "antonym_zh", "concept": "heavy", "prompt": "“大”的反义词是“小”。“轻”的反义词是“", "form_a": "重", "form_b": " heavy", "targets": ["重", " heavy", "门", " paper"], "null_targets": ["门", " paper"], "answer": "重"}, + {"id": "ant_14", "group": "antonym_zh", "concept": "difficult", "prompt": "“大”的反义词是“小”。“难”的反义词是“", "form_a": "易", "form_b": " easy", "targets": ["易", " easy", "鞋", " door"], "null_targets": ["鞋", " door"], "answer": "易"}, + + {"id": "tz2e_01", "group": "trans_zh2en", "concept": "water", "prompt": "“天”的英文是“sky”。“水”的英文是“", "form_a": "water", "form_b": "水", "targets": ["water", "水", " water", "桌", "chair"], "null_targets": ["桌", "chair"], "answer": "water"}, + {"id": "tz2e_02", "group": "trans_zh2en", "concept": "fire", "prompt": "“天”的英文是“sky”。“火”的英文是“", "form_a": "fire", "form_b": "火", "targets": ["fire", "火", " fire", "门", "door"], "null_targets": ["门", "door"], "answer": "fire"}, + {"id": "tz2e_03", "group": "trans_zh2en", "concept": "cat", "prompt": "“天”的英文是“sky”。“猫”的英文是“", "form_a": "cat", "form_b": "猫", "targets": ["cat", "猫", " cat", "纸", "paper"], "null_targets": ["纸", "paper"], "answer": "cat"}, + {"id": "tz2e_04", "group": "trans_zh2en", "concept": "book", "prompt": "“天”的英文是“sky”。“书”的英文是“", "form_a": "book", "form_b": "书", "targets": ["book", "书", " book", "盐", "salt"], "null_targets": ["盐", "salt"], "answer": "book"}, + {"id": "tz2e_05", "group": "trans_zh2en", "concept": "mountain", "prompt": "“天”的英文是“sky”。“山”的英文是“", "form_a": "mountain", "form_b": "山", "targets": ["mountain", "山", " mountain", "鞋", "shoe"], "null_targets": ["鞋", "shoe"], "answer": "mountain"}, + {"id": "tz2e_06", "group": "trans_zh2en", "concept": "moon", "prompt": "“天”的英文是“sky”。“月”的英文是“", "form_a": "moon", "form_b": "月", "targets": ["moon", "月", " moon", "墙", "wall"], "null_targets": ["墙", "wall"], "answer": "moon"}, + {"id": "tz2e_07", "group": "trans_zh2en", "concept": "rain", "prompt": "“天”的英文是“sky”。“雨”的英文是“", "form_a": "rain", "form_b": "雨", "targets": ["rain", "雨", " rain", "针", "needle"], "null_targets": ["针", "needle"], "answer": "rain"}, + {"id": "tz2e_08", "group": "trans_zh2en", "concept": "heart", "prompt": "“天”的英文是“sky”。“心”的英文是“", "form_a": "heart", "form_b": "心", "targets": ["heart", "心", " heart", "伞", "umbrella"], "null_targets": ["伞", "umbrella"], "answer": "heart"}, + {"id": "tz2e_09", "group": "trans_zh2en", "concept": "fish", "prompt": "“天”的英文是“sky”。“鱼”的英文是“", "form_a": "fish", "form_b": "鱼", "targets": ["fish", "鱼", " fish", "票", "ticket"], "null_targets": ["票", "ticket"], "answer": "fish"}, + {"id": "tz2e_10", "group": "trans_zh2en", "concept": "dog", "prompt": "“天”的英文是“sky”。“狗”的英文是“", "form_a": "dog", "form_b": "狗", "targets": ["dog", "狗", " dog", "杯", "cup"], "null_targets": ["杯", "cup"], "answer": "dog"}, + + {"id": "te2z_01", "group": "trans_en2zh", "concept": "water", "prompt": "The Chinese word for \"sky\" is \"天\". The Chinese word for \"water\" is \"", "form_a": "水", "form_b": " water", "targets": ["水", " water", "water", "桌", " chair"], "null_targets": ["桌", " chair"], "answer": "水"}, + {"id": "te2z_02", "group": "trans_en2zh", "concept": "fire", "prompt": "The Chinese word for \"sky\" is \"天\". The Chinese word for \"fire\" is \"", "form_a": "火", "form_b": " fire", "targets": ["火", " fire", "fire", "门", " door"], "null_targets": ["门", " door"], "answer": "火"}, + {"id": "te2z_03", "group": "trans_en2zh", "concept": "cat", "prompt": "The Chinese word for \"sky\" is \"天\". The Chinese word for \"cat\" is \"", "form_a": "猫", "form_b": " cat", "targets": ["猫", " cat", "cat", "纸", " paper"], "null_targets": ["纸", " paper"], "answer": "猫"}, + {"id": "te2z_04", "group": "trans_en2zh", "concept": "book", "prompt": "The Chinese word for \"sky\" is \"天\". The Chinese word for \"book\" is \"", "form_a": "书", "form_b": " book", "targets": ["书", " book", "book", "盐", " salt"], "null_targets": ["盐", " salt"], "answer": "书"}, + {"id": "te2z_05", "group": "trans_en2zh", "concept": "mountain", "prompt": "The Chinese word for \"sky\" is \"天\". The Chinese word for \"mountain\" is \"", "form_a": "山", "form_b": " mountain", "targets": ["山", " mountain", "mountain", "鞋", " shoe"], "null_targets": ["鞋", " shoe"], "answer": "山"}, + {"id": "te2z_06", "group": "trans_en2zh", "concept": "moon", "prompt": "The Chinese word for \"sky\" is \"天\". The Chinese word for \"moon\" is \"", "form_a": "月", "form_b": " moon", "targets": ["月", " moon", "moon", "墙", " wall"], "null_targets": ["墙", " wall"], "answer": "月"}, + {"id": "te2z_07", "group": "trans_en2zh", "concept": "rain", "prompt": "The Chinese word for \"sky\" is \"天\". The Chinese word for \"rain\" is \"", "form_a": "雨", "form_b": " rain", "targets": ["雨", " rain", "rain", "针", " needle"], "null_targets": ["针", " needle"], "answer": "雨"}, + {"id": "te2z_08", "group": "trans_en2zh", "concept": "heart", "prompt": "The Chinese word for \"sky\" is \"天\". The Chinese word for \"heart\" is \"", "form_a": "心", "form_b": " heart", "targets": ["心", " heart", "heart", "伞", " umbrella"], "null_targets": ["伞", " umbrella"], "answer": "心"}, + {"id": "te2z_09", "group": "trans_en2zh", "concept": "fish", "prompt": "The Chinese word for \"sky\" is \"天\". The Chinese word for \"fish\" is \"", "form_a": "鱼", "form_b": " fish", "targets": ["鱼", " fish", "fish", "票", " ticket"], "null_targets": ["票", " ticket"], "answer": "鱼"}, + {"id": "te2z_10", "group": "trans_en2zh", "concept": "dog", "prompt": "The Chinese word for \"sky\" is \"天\". The Chinese word for \"dog\" is \"", "form_a": "狗", "form_b": " dog", "targets": ["狗", " dog", "dog", "杯", " cup"], "null_targets": ["杯", " cup"], "answer": "狗"}, + + {"id": "clz_01", "group": "cloze_zh", "concept": "water", "prompt": "他很渴，想喝一杯", "form_a": "水", "form_b": " water", "targets": ["水", " water", "门", " door"], "null_targets": ["门", " door"], "answer": "水"}, + {"id": "clz_02", "group": "cloze_zh", "concept": "cold", "prompt": "冬天的天气非常", "form_a": "冷", "form_b": " cold", "targets": ["冷", " cold", "纸", " paper"], "null_targets": ["纸", " paper"], "answer": "冷"}, + {"id": "clz_03", "group": "cloze_zh", "concept": "hot", "prompt": "夏天的天气非常", "form_a": "热", "form_b": " hot", "targets": ["热", " hot", "盐", " salt"], "null_targets": ["盐", " salt"], "answer": "热"}, + {"id": "clz_04", "group": "cloze_zh", "concept": "rain", "prompt": "他打开窗户，看见外面正在下", "form_a": "雨", "form_b": " rain", "targets": ["雨", " rain", "鞋", " shoe"], "null_targets": ["鞋", " shoe"], "answer": "雨"}, + {"id": "clz_05", "group": "cloze_zh", "concept": "big", "prompt": "大象的身体非常", "form_a": "大", "form_b": " big", "targets": ["大", " big", " large", "墙", " wall"], "null_targets": ["墙", " wall"], "answer": "大"}, + {"id": "clz_06", "group": "cloze_zh", "concept": "small", "prompt": "老鼠的身体非常", "form_a": "小", "form_b": " small", "targets": ["小", " small", "针", " needle"], "null_targets": ["针", " needle"], "answer": "小"}, + {"id": "clz_07", "group": "cloze_zh", "concept": "fly", "prompt": "鸟在天空中", "form_a": "飞", "form_b": " fly", "targets": ["飞", " fly", " flying", "票", " ticket"], "null_targets": ["票", " ticket"], "answer": "飞"}, + {"id": "clz_08", "group": "cloze_zh", "concept": "swim", "prompt": "鱼在水里", "form_a": "游", "form_b": " swim", "targets": ["游", " swim", "伞", " umbrella"], "null_targets": ["伞", " umbrella"], "answer": "游"}, + {"id": "clz_09", "group": "cloze_zh", "concept": "sleep", "prompt": "晚上十一点，他很困，想去床上", "form_a": "睡", "form_b": " sleep", "targets": ["睡", " sleep", "杯", " cup"], "null_targets": ["杯", " cup"], "answer": "睡觉"}, + {"id": "clz_10", "group": "cloze_zh", "concept": "green", "prompt": "春天来了，草和树叶都变成了", "form_a": "绿", "form_b": " green", "targets": ["绿", " green", "桌", " table"], "null_targets": ["桌", " table"], "answer": "绿"}, + {"id": "clz_11", "group": "cloze_zh", "concept": "fast", "prompt": "猎豹跑得非常", "form_a": "快", "form_b": " fast", "targets": ["快", " fast", "门", " wall"], "null_targets": ["门", " wall"], "answer": "快"}, + {"id": "clz_12", "group": "cloze_zh", "concept": "slow", "prompt": "乌龟走得非常", "form_a": "慢", "form_b": " slow", "targets": ["慢", " slow", "纸", " needle"], "null_targets": ["纸", " needle"], "answer": "慢"}, + {"id": "clz_13", "group": "cloze_zh", "concept": "sweet", "prompt": "糖的味道很", "form_a": "甜", "form_b": " sweet", "targets": ["甜", " sweet", "鞋", " door"], "null_targets": ["鞋", " door"], "answer": "甜"}, + {"id": "clz_14", "group": "cloze_zh", "concept": "sour", "prompt": "柠檬的味道很", "form_a": "酸", "form_b": " sour", "targets": ["酸", " sour", "墙", " paper"], "null_targets": ["墙", " paper"], "answer": "酸"}, + {"id": "clz_15", "group": "cloze_zh", "concept": "water", "prompt": "冰融化以后就变成了", "form_a": "水", "form_b": " water", "targets": ["水", " water", "针", " ticket"], "null_targets": ["针", " ticket"], "answer": "水"}, + {"id": "clz_16", "group": "cloze_zh", "concept": "moon", "prompt": "白天天上有太阳，晚上天上有", "form_a": "月", "form_b": " moon", "targets": ["月", " moon", "盐", " cup"], "null_targets": ["盐", " cup"], "answer": "月亮"}, + {"id": "clz_17", "group": "cloze_zh", "concept": "eat", "prompt": "他饿了，想吃", "form_a": "饭", "form_b": " food", "targets": ["饭", " food", "票", " umbrella"], "null_targets": ["票", " umbrella"], "answer": "饭"}, + {"id": "clz_18", "group": "cloze_zh", "concept": "red", "prompt": "血的颜色是", "form_a": "红", "form_b": " red", "targets": ["红", " red", "伞", " chair"], "null_targets": ["伞", " chair"], "answer": "红"}, + + {"id": "cle_01", "group": "cloze_en", "concept": "cold", "prompt": "In winter the weather is very", "form_a": " cold", "form_b": "冷", "targets": [" cold", "冷", " door", "门"], "null_targets": [" door", "门"], "answer": "cold"}, + {"id": "cle_02", "group": "cloze_en", "concept": "water", "prompt": "He was thirsty, so he drank a glass of", "form_a": " water", "form_b": "水", "targets": [" water", "水", " paper", "纸"], "null_targets": [" paper", "纸"], "answer": "water"}, + {"id": "cle_03", "group": "cloze_en", "concept": "big", "prompt": "An elephant's body is extremely", "form_a": " large", "form_b": "大", "targets": [" large", " big", "大", " salt", "盐"], "null_targets": [" salt", "盐"], "answer": "large"}, + {"id": "cle_04", "group": "cloze_en", "concept": "small", "prompt": "Compared to an elephant, a mouse is very", "form_a": " small", "form_b": "小", "targets": [" small", "小", " shoe", "鞋"], "null_targets": [" shoe", "鞋"], "answer": "small"}, + {"id": "cle_05", "group": "cloze_en", "concept": "slow", "prompt": "The opposite of fast is", "form_a": " slow", "form_b": "慢", "targets": [" slow", "慢", " wall", "墙"], "null_targets": [" wall", "墙"], "answer": "slow"}, + {"id": "cle_06", "group": "cloze_en", "concept": "sweet", "prompt": "Sugar tastes very", "form_a": " sweet", "form_b": "甜", "targets": [" sweet", "甜", " needle", "针"], "null_targets": [" needle", "针"], "answer": "sweet"}, + {"id": "cle_07", "group": "cloze_en", "concept": "moon", "prompt": "At night, the sky is lit by the", "form_a": " moon", "form_b": "月", "targets": [" moon", "月", " ticket", "票"], "null_targets": [" ticket", "票"], "answer": "moon"}, + {"id": "cle_08", "group": "cloze_en", "concept": "water", "prompt": "Fish can only live in the", "form_a": " water", "form_b": "水", "targets": [" water", "水", " umbrella", "伞"], "null_targets": [" umbrella", "伞"], "answer": "water"}, + {"id": "cle_09", "group": "cloze_en", "concept": "rain", "prompt": "He opened his umbrella because it had started to", "form_a": " rain", "form_b": "雨", "targets": [" rain", "雨", " cup", "杯"], "null_targets": [" cup", "杯"], "answer": "rain"}, + {"id": "cle_10", "group": "cloze_en", "concept": "fast", "prompt": "A cheetah can run extremely", "form_a": " fast", "form_b": "快", "targets": [" fast", "快", " table", "桌"], "null_targets": [" table", "桌"], "answer": "fast"}, + + {"id": "fac_01", "group": "fact_zh", "concept": "Paris", "prompt": "中国的首都是北京。法国的首都是", "form_a": "巴黎", "form_b": " Paris", "targets": ["巴黎", " Paris", "桌", " chair"], "null_targets": ["桌", " chair"], "answer": "巴黎"}, + {"id": "fac_02", "group": "fact_zh", "concept": "Tokyo", "prompt": "中国的首都是北京。日本的首都是", "form_a": "东京", "form_b": " Tokyo", "targets": ["东京", " Tokyo", "门", " door"], "null_targets": ["门", " door"], "answer": "东京"}, + {"id": "fac_03", "group": "fact_zh", "concept": "London", "prompt": "中国的首都是北京。英国的首都是", "form_a": "伦敦", "form_b": " London", "targets": ["伦敦", " London", "纸", " paper"], "null_targets": ["纸", " paper"], "answer": "伦敦"}, + {"id": "fac_04", "group": "fact_zh", "concept": "Washington", "prompt": "中国的首都是北京。美国的首都是", "form_a": "华盛顿", "form_b": " Washington", "targets": ["华盛顿", " Washington", "盐", " salt"], "null_targets": ["盐", " salt"], "answer": "华盛顿"}, + {"id": "fac_05", "group": "fact_zh", "concept": "Moscow", "prompt": "中国的首都是北京。俄罗斯的首都是", "form_a": "莫斯科", "form_b": " Moscow", "targets": ["莫斯科", " Moscow", "鞋", " shoe"], "null_targets": ["鞋", " shoe"], "answer": "莫斯科"}, + {"id": "fac_06", "group": "fact_zh", "concept": "Berlin", "prompt": "中国的首都是北京。德国的首都是", "form_a": "柏林", "form_b": " Berlin", "targets": ["柏林", " Berlin", "墙", " wall"], "null_targets": ["墙", " wall"], "answer": "柏林"}, + {"id": "fac_07", "group": "fact_zh", "concept": "Rome", "prompt": "中国的首都是北京。意大利的首都是", "form_a": "罗马", "form_b": " Rome", "targets": ["罗马", " Rome", "针", " needle"], "null_targets": ["针", " needle"], "answer": "罗马"}, + {"id": "fac_08", "group": "fact_zh", "concept": "Madrid", "prompt": "中国的首都是北京。西班牙的首都是", "form_a": "马德里", "form_b": " Madrid", "targets": ["马德里", " Madrid", "伞", " umbrella"], "null_targets": ["伞", " umbrella"], "answer": "马德里"}, + {"id": "fac_09", "group": "fact_zh", "concept": "Cairo", "prompt": "中国的首都是北京。埃及的首都是", "form_a": "开罗", "form_b": " Cairo", "targets": ["开罗", " Cairo", "票", " ticket"], "null_targets": ["票", " ticket"], "answer": "开罗"}, + {"id": "fac_10", "group": "fact_zh", "concept": "Seoul", "prompt": "中国的首都是北京。韩国的首都是", "form_a": "首尔", "form_b": " Seoul", "targets": ["首尔", " Seoul", "杯", " cup"], "null_targets": ["杯", " cup"], "answer": "首尔"}, + + {"id": "exm_01", "group": "exemplar_zh", "concept": "apple", "prompt": "一种常见的水果是", "form_a": "苹果", "form_b": " apple", "targets": ["苹果", " apple", "门", " door"], "null_targets": ["门", " door"], "answer": "苹果"}, + {"id": "exm_02", "group": "exemplar_zh", "concept": "cat", "prompt": "一种常见的宠物是", "form_a": "猫", "form_b": " cat", "targets": ["猫", " cat", "狗", " dog", "纸"], "null_targets": ["纸"], "answer": "猫"}, + {"id": "exm_03", "group": "exemplar_zh", "concept": "red", "prompt": "一种常见的颜色是", "form_a": "红", "form_b": " red", "targets": ["红", " red", "盐", " salt"], "null_targets": ["盐", " salt"], "answer": "红色"}, + {"id": "exm_04", "group": "exemplar_zh", "concept": "tea", "prompt": "一种常见的饮料是", "form_a": "茶", "form_b": " tea", "targets": ["茶", " tea", "鞋", " shoe"], "null_targets": ["鞋", " shoe"], "answer": "茶"}, + {"id": "exm_05", "group": "exemplar_zh", "concept": "car", "prompt": "一种常见的交通工具是", "form_a": "汽车", "form_b": " car", "targets": ["汽车", " car", "墙", " wall"], "null_targets": ["墙", " wall"], "answer": "汽车"}, + {"id": "exm_06", "group": "exemplar_zh", "concept": "football", "prompt": "一种常见的体育运动是", "form_a": "足球", "form_b": " football", "targets": ["足球", " football", " soccer", "针", " needle"], "null_targets": ["针", " needle"], "answer": "足球"}, + {"id": "exm_07", "group": "exemplar_zh", "concept": "winter", "prompt": "一年有四个季节：春、夏、秋、", "form_a": "冬", "form_b": " winter", "targets": ["冬", " winter", "票", " ticket"], "null_targets": ["票", " ticket"], "answer": "冬"}, + {"id": "exm_08", "group": "exemplar_zh", "concept": "rose", "prompt": "一种常见的花是", "form_a": "玫瑰", "form_b": " rose", "targets": ["玫瑰", " rose", "伞", " umbrella"], "null_targets": ["伞", " umbrella"], "answer": "玫瑰"}, + + {"id": "dig_01", "group": "ctrl_digit", "concept": "7", "prompt": "三加四等于", "form_a": "七", "form_b": " seven", "targets": ["七", " seven", " 7", "7"], "null_targets": [], "answer": "七"}, + {"id": "dig_02", "group": "ctrl_digit", "concept": "10", "prompt": "二乘五等于", "form_a": "十", "form_b": " ten", "targets": ["十", " ten", " 10", "10"], "null_targets": [], "answer": "十"}, + {"id": "dig_03", "group": "ctrl_digit", "concept": "9", "prompt": "Q: What is 2 + 2? A: 4\nQ: What is 6 + 3? A:", "form_a": " 9", "form_b": "九", "targets": [" 9", "九", " nine", "9"], "null_targets": [], "answer": "9"}, + {"id": "dig_04", "group": "ctrl_digit", "concept": "7", "prompt": "九减二等于", "form_a": "七", "form_b": " seven", "targets": ["七", " seven", " 7", "7"], "null_targets": [], "answer": "七"}, + {"id": "dig_05", "group": "ctrl_digit", "concept": "8", "prompt": "Q: What is 2 + 2? A: 4\nQ: What is 4 times 2? A:", "form_a": " 8", "form_b": "八", "targets": [" 8", "八", " eight", "8"], "null_targets": [], "answer": "8"}, + {"id": "dig_06", "group": "ctrl_digit", "concept": "10", "prompt": "五加五等于", "form_a": "十", "form_b": " ten", "targets": ["十", " ten", " 10", "10"], "null_targets": [], "answer": "十"}, + + {"id": "prn_01", "group": "ctrl_propn", "concept": "NASA", "prompt": "美国国家航空航天局的英文缩写是", "form_a": "NASA", "form_b": " NASA", "targets": ["NASA", " NASA"], "null_targets": [], "answer": "NASA"}, + {"id": "prn_02", "group": "ctrl_propn", "concept": "DNA", "prompt": "脱氧核糖核酸的英文缩写是", "form_a": "DNA", "form_b": " DNA", "targets": ["DNA", " DNA"], "null_targets": [], "answer": "DNA"}, + {"id": "prn_03", "group": "ctrl_propn", "concept": "NBA", "prompt": "美国职业篮球联赛的英文缩写是", "form_a": "NBA", "form_b": " NBA", "targets": ["NBA", " NBA"], "null_targets": [], "answer": "NBA"}, + {"id": "prn_04", "group": "ctrl_propn", "concept": "iPhone", "prompt": "苹果公司生产的智能手机叫", "form_a": "iPhone", "form_b": " iPhone", "targets": ["iPhone", " iPhone"], "null_targets": [], "answer": "iPhone"}, + {"id": "prn_05", "group": "ctrl_propn", "concept": "AI", "prompt": "人工智能的英文缩写是", "form_a": "AI", "form_b": " AI", "targets": ["AI", " AI"], "null_targets": [], "answer": "AI"}, + {"id": "prn_06", "group": "ctrl_propn", "concept": "WHO", "prompt": "世界卫生组织的英文缩写是", "form_a": "WHO", "form_b": " WHO", "targets": ["WHO", " WHO"], "null_targets": [], "answer": "WHO"} + ] +} diff --git a/data/cross_lens/cross_lens_three_lens_prompts.json b/data/cross_lens/cross_lens_three_lens_prompts.json new file mode 100644 index 0000000..9322f33 --- /dev/null +++ b/data/cross_lens/cross_lens_three_lens_prompts.json @@ -0,0 +1,55 @@ +{ + "note": "One-prompt-per-family bank for the worked example read by three fitted lenses at once (tab:app-cross-lens-de-antonym-layers). ant_01 and antonym_de_01 are copied from cross_lens_prompts_en_zh.json / cross_lens_prompts_en_de.json with extended target lists covering the English, Chinese and German answer forms plus the two nulls. Run scripts/run/run_cross_lens_readouts.py once per lens with --n-positions 1 --decompose-top1 --top-feats 10, then scripts/analyze/cross_lens_three_lens_prompt.py.", + "prompts": [ + { + "id": "ant_01", + "group": "antonym_zh", + "concept": "big", + "prompt": "“高”的反义词是“低”。“小”的反义词是“", + "form_a": "大", + "form_b": " big", + "targets": [ + "大", + "大的", + " big", + " large", + " groß", + " große", + " Groß", + "桌", + " chair" + ], + "null_targets": [ + "桌", + " chair" + ], + "answer": "大" + }, + { + "id": "antonym_de_01", + "group": "antonym_de", + "concept": "big", + "prompt": "Das Gegenteil von \"hoch\" ist \"tief\". Das Gegenteil von \"klein\" ist \"", + "form_a": " groß", + "form_b": " big", + "lang_a": "de", + "lang_b": "en", + "targets": [ + " groß", + " große", + " Groß", + " big", + " large", + "大", + "大的", + " Salz", + " salt" + ], + "null_targets": [ + " Salz", + " salt" + ], + "answer": "groß" + } + ] +} diff --git a/data/wsd/wsd_sense_anchors.json b/data/wsd/wsd_sense_anchors.json new file mode 100644 index 0000000..43bf4a5 --- /dev/null +++ b/data/wsd/wsd_sense_anchors.json @@ -0,0 +1,101 @@ +{ + "_meta": { + "purpose": "Frozen sense-anchor token sets for weights-only feature typing (CoarseWSD-20). Typing rule: beta_f(sense) = mean anchor code (one-vs-rest contrast); a feature is sense-typed when its max one-vs-rest beta exceeds tau from the random-anchor null.", + "rule": "Per sense: the class-map sense name plus curated common-word associates; the ambiguous target itself is never an anchor; no anchor is reused across senses of the same word; anchors are filtered per tokenizer to leading-space single tokens at runtime and the surviving counts are reported.", + "frozen": "2026-08-01, before any occurrence-mode run was executed", + "sense_ids": "match CoarseWSD-20 class_map.txt per word" + }, + "apple": { + "0": ["iPhone", "Mac", "software", "computer", "technology", "company"], + "1": ["fruit", "orchard", "pear", "cherry", "juice", "ripe"] + }, + "arm": { + "0": ["processor", "chip", "architecture", "Intel", "silicon", "computing"], + "1": ["shoulder", "elbow", "hand", "leg", "wrist", "limb"] + }, + "bank": { + "0": ["money", "loan", "deposit", "financial", "credit", "account"], + "1": ["river", "shore", "stream", "flood", "water", "erosion"] + }, + "bass": { + "0": ["guitar", "riff", "band", "rock", "amplifier"], + "1": ["singer", "tenor", "choir", "opera", "voice"], + "2": ["cello", "orchestra", "violin", "symphony", "concerto"] + }, + "bow": { + "0": ["ship", "stern", "hull", "vessel", "sailors"], + "1": ["arrow", "archer", "quiver", "shooting", "hunting"], + "2": ["violin", "strings", "cello", "rosin", "instrument"] + }, + "chair": { + "0": ["committee", "president", "board", "meeting", "elected"], + "1": ["table", "furniture", "seat", "cushion", "wooden"] + }, + "club": { + "0": ["society", "members", "organization", "association", "founded"], + "1": ["nightclub", "dancing", "bar", "DJ", "party"], + "2": ["weapon", "stick", "bat", "blunt", "wielded"] + }, + "crane": { + "0": ["machinery", "lifting", "construction", "tower", "load", "hoist"], + "1": ["bird", "heron", "stork", "wetland", "feathers", "migratory"] + }, + "deck": { + "0": ["ship", "sailors", "hull", "aboard", "vessel"], + "1": ["porch", "patio", "balcony", "wooden", "backyard"] + }, + "digit": { + "0": ["number", "numeral", "decimal", "integer", "arithmetic"], + "1": ["finger", "toe", "thumb", "hand", "limb"] + }, + "hood": { + "0": ["villain", "comics", "Marvel", "superhero", "character"], + "1": ["car", "engine", "bonnet", "vehicle", "windshield"], + "2": ["cloak", "jacket", "head", "garment", "sweatshirt"] + }, + "java": { + "0": ["Indonesia", "island", "coffee", "Jakarta", "volcano"], + "1": ["programming", "software", "code", "compiler", "language"] + }, + "mole": { + "0": ["burrow", "rodent", "digging", "fur", "underground"], + "1": ["spy", "agent", "intelligence", "undercover", "espionage"], + "2": ["molar", "chemistry", "gram", "concentration", "atoms"], + "3": ["sauce", "Mexican", "chili", "chocolate", "cuisine"], + "4": ["pier", "harbor", "breakwater", "jetty", "seawall"] + }, + "pitcher": { + "0": ["baseball", "innings", "strikeout", "mound", "batter"], + "1": ["jug", "water", "pour", "ceramic", "vessel"] + }, + "pound": { + "0": ["kilogram", "weight", "ounce", "mass", "grams"], + "1": ["sterling", "currency", "pence", "British", "dollar"] + }, + "seal": { + "0": ["marine", "whiskers", "blubber", "colony", "swimming"], + "1": ["singer", "album", "song", "Grammy", "vocalist"], + "2": ["emblem", "stamp", "official", "insignia", "wax"], + "3": ["gasket", "valve", "leak", "pressure", "rubber"] + }, + "spring": { + "0": ["water", "geyser", "mineral", "well", "flowing"], + "1": ["summer", "winter", "autumn", "season", "blossom"], + "2": ["coil", "mechanical", "tension", "metal", "compression"] + }, + "square": { + "0": ["rectangle", "triangle", "geometry", "shape", "sides"], + "1": ["company", "corporation", "payments", "games"], + "2": ["plaza", "town", "market", "city", "fountain"], + "3": ["squared", "multiplication", "integer", "root", "arithmetic"] + }, + "trunk": { + "0": ["tree", "bark", "branches", "roots", "wood"], + "1": ["car", "luggage", "rear", "vehicle", "storage"], + "2": ["torso", "body", "abdomen", "muscles"] + }, + "yard": { + "0": ["meter", "feet", "inches", "garden", "lawn"], + "1": ["mast", "sail", "rigging", "spar", "ship"] + } +} diff --git a/docs/DATA.md b/docs/DATA.md index 9c59cc0..0dc78b7 100644 --- a/docs/DATA.md +++ b/docs/DATA.md @@ -15,6 +15,15 @@ A record is one **case** — a prompt, a position to score, two targets (`A`, `B`), and an expected side. Every script that consumes a bank reads the same schema; the banks are model-agnostic. +Two further paper-input sets live beside them. [`data/cross_lens/`](../data/cross_lens/) +holds the EN–ZH and EN–DE prompt banks of the cross-lens study (one record per +prompt: two surface forms, an unrelated null target, family and control tags) +together with the EN–DE cognate-exclusion report and the one-prompt bank of the +three-lens worked example; schema and provenance are in +[`data/cross_lens/README.md`](../data/cross_lens/README.md). [`data/wsd/`](../data/wsd/) +holds the optional sense-gloss anchors that `scripts/run/run_wsd_feature_alignment.py` +accepts through `--anchors`; the paper's CoarseWSD-20 runs did not use them. + ## 2. Extracted readouts (per-model artifacts) Run `scripts/data/extract_model_readout.py --model-id --out .pt` @@ -36,7 +45,7 @@ itself comes straight from the model's `lm_head`. ## 3. SAE checkpoints (pretrained dictionaries) The selected dictionaries cited in the paper are published on the **Hugging Face -Hub** at [`matteohe/sparse-readout-prism`](https://huggingface.co/matteohe/sparse-readout-prism), +Hub** at [`hematteo/sparse-readout-prism`](https://huggingface.co/hematteo/sparse-readout-prism), laid out as `//checkpoint.pt` (download with `huggingface_hub.hf_hub_download`). The manifest in [`configs/registries/exp2_selected_sae_checkpoints.yaml`](../configs/registries/exp2_selected_sae_checkpoints.yaml) diff --git a/docs/REPRODUCE.md b/docs/REPRODUCE.md index 5644aea..873304c 100644 --- a/docs/REPRODUCE.md +++ b/docs/REPRODUCE.md @@ -14,6 +14,7 @@ rendered separately in the paper's LaTeX source from these metrics. ```bash uv sync # install Python deps +uv sync --extra lens # + jlens (Jacobian lens), only for the cross-lens study ``` Hardware: the headline results were produced on a single A40 (an HPC @@ -32,10 +33,25 @@ External data: `scripts/data/build_query_banks.py --native-source c4` (its provenance / revision is pinned in `data/query_banks/qwen_gemma_result1_model_native_prompts_c4_manifest.json`). +- The direct-geometry baseline runs (Appendix "Direct-Geometry Alternatives") + read that same C4 slice as their `model_native` bank; regenerate it with + `scripts/data/build_query_banks.py --native-source c4 --seed 0 --max-native 10000` + and pass `--model-native-file qwen_gemma_result1_model_native_prompts_c4.jsonl`. +- The cross-lens study fits its lenses on seeded C4 prompts (`allenai/c4`, + configs `en` / `zh` / `de`, streaming, `shuffle(seed=0)`) dumped by + `scripts/run/fit_jlens.py prompts`; the dumps are raw C4 text and are not + shipped. The lenses need the `lens` extra (`uv sync --extra lens`). +- CoarseWSD-20 (Loureiro, Rezaee, Pilehvar and Camacho-Collados, 2021, + *Computational Linguistics* 47(2)) for the sense-labelled evaluation: clone + and pass the checkout (or its + `data/CoarseWSD-20`) as `--data-root` to `scripts/run/run_wsd_feature_alignment.py`. +- The three-seed dictionaries behind the stability appendix are not on the Hub; + train them from `configs/sweeps/qwen35_2b_seedvar_base.yaml` (see the + `configs/sweeps/README.md` table for the per-cell overrides). 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) (`//checkpoint.pt`). `configs/registries/exp2_selected_sae_checkpoints.yaml` records the canonical manifest (each checkpoint's model, width, `k`, metrics, @@ -125,16 +141,16 @@ check. | Paper artifact | Script | Config / inputs | |---|---|---| -| Method schematics (`fig:readout-score-schematic` / `readout_query_three_panel`, score-accounting TikZ) | — (paper-side LaTeX/TikZ; no producing script in this repo) | — | +| Method schematics (`fig:method-schematic` / `readout_query_one_column`, score-accounting TikZ) | — (paper-side LaTeX/TikZ; no producing script in this repo) | — | | §4 main fidelity table (Qwen 0.8B/2B/9B) | `scripts/eval/task_fidelity_evaluation.py` | `configs/registries/exp2_selected_sae_checkpoints.yaml` | | §4 main case-study four-panel figure (`fig:selected-readout-prism-examples`: bug/error, bug/insect, bark/dog, bass/trout) | `scripts/figures/compute_sae_paper_examples.py` | hardcoded example sets (`--example-set section43_polysemy_margin`, the default; `EXAMPLE_SETS` in `research/qwen_example_prompts.py`), plus a required `--checkpoint` | | §1 intro figure: logit-lens vs Sparse Readout Prism comparison (`fig:main-lens-prism-comparison`) | `scripts/figures/compute_lens_prism_comparison.py` | defaults reproduce the paper case (software-context prompt, `--target-a " bug"`, `--target-b " insect"`, `--k 256`); `--checkpoint` points at the Qwen3.5-2B 32x/k256 selected SAE in `configs/registries/exp2_selected_sae_checkpoints.yaml`. Computes readout states live (reuses `collect_readout_states_batched` from `sparse_readout_prism.research.qwen_readout`; no precomputed cache) | | Table / metrics behind: benchmark-derived task targets | `scripts/run/run_benchmark_derived_query_suite.py`, `scripts/figures/compute_benchmark_task_group_examples.py` | `--out-dir results/benchmark_derived_query_suite_task_group_v2_qwen_20260523` (the suite uses task-group contrasts — the paper's `task_group_v2` design — and does not derive the dir from it; pass `--out-dir` explicitly) | | Metrics behind: selected-score / family / selection queries (`fig:qwen-general-readout-scores-basic`, `fig:app-qwen-general-readout-scores-extra`) | `scripts/figures/compute_general_readout_queries.py` | query specs are built in (the paper's jury/guilty panels plus the abstention-family contrast); `--checkpoint` points at the Qwen3.5-2B 32x/k256 selected SAE in `configs/registries/result1_query_fidelity_cluster.yaml` | | Metrics behind: literature-prompt all-layer trace | `scripts/figures/compute_all_layer_literature_prompt.py` | requires `--checkpoint` (the Qwen3.5-2B 32x/k256 selected SAE); prompt/model defaults reproduce the paper trace | -| Baseline comparison metrics (Sparse RP vs nulls) | `scripts/run/run_readout_baseline_comparisons.py` | outputs persisted as CSV/JSON for the paper-side baseline figure | -| Result-1 five-model query-fidelity tables / §4.5 distributional readout metrics | `scripts/run/run_query_fidelity_bank.py` | `configs/registries/result1_query_fidelity_cluster.yaml` + the curated/model-native banks under `data/query_banks/` | -| Appendix J: lexical-edit / readout-side control stress test (`tab:lexical-control-primary-methods`, `tab:lexical-control-cross-model-results`) | `scripts/run/run_qwen_profanity_suppression_eval.py` | `--checkpoint `; the primary-methods table reads `baseline_comparison.csv` / `candidate_constrained_summary.csv` (the eval emits these tables only — no paper figure). For the cross-model table, rerun once per model with the checkpoints in `configs/registries/result1_query_fidelity_cluster.yaml` (Qwen3.5-2B fidelity 32x/k256; R1-Distill-Qwen-7B fidelity 32x/k256; Ministral-3-8B-Base strict_budget 16x/k128) | +| `tab:app-direct-geometry-grid`; SRP / best-alternative / lead columns of `tab:readout-score-fidelity-summary` | `scripts/run/run_readout_baseline_comparisons.py` | one run per softcap-free readout, `--model` in Qwen3.5-0.8B, Qwen3.5-2B, Qwen3.5-9B, Ministral-3-8B-Base, R1-Distill-Qwen-7B, R1-Distill-Llama-8B: `--registry configs/registries/result1_query_fidelity_cluster.yaml --bank-dir data/query_banks --banks curated_ab,case_candidates,model_native --model-native-file qwen_gemma_result1_model_native_prompts_c4.jsonl --operating-point fidelity --methods sparse_rp,nearest_row_ridge_top128,knn_basis_top128,row_cluster_d16384_k256,row_cluster_d65536_k256,row_cluster_hard_d65536,pca_256 --max-native 500 --max-curated 320 --max-cases 40 --max-len 64 --seed 0 --out-root results/direct_geometry_runs/`; coverage = `accepted_rate` with `accepted_ci_lo/hi` in `baseline_by_method.csv` (unfloored rho_0, 400 base-case resamples). The C4 model-native slice is regenerated with `scripts/data/build_query_banks.py --native-source c4` (see §0) | +| Result-1 five-model query-fidelity tables / §4.5 distributional readout metrics; per-model reconstruction with cluster-bootstrap intervals (`tab:app-fidelity-cis`) | `scripts/run/run_query_fidelity_bank.py` | `configs/registries/result1_query_fidelity_cluster.yaml` + the curated/model-native banks under `data/query_banks/`. `tab:app-fidelity-cis` reads, per model, `sign_agreement` (+`_ci_lo/_ci_hi`), `median_rho5` (+CI) and `pass_rho5_lt_0.50` (+CI) from the group metrics: `rho5` is the floored relative error \|m_exact − m_sparse\| / (\|m_exact\| + 0.5) and every interval comes from the same 400 base-case resamples | +| Appendix J lexical-edit / readout-side control stress test (`tab:lexical-control-stress-test`, `tab:lexical-control-primary-methods`, `tab:lexical-control-cross-model-results`) | `scripts/run/run_qwen_profanity_suppression_eval.py` | `--checkpoint --model-id --out-dir --device cuda --dtype bfloat16 --batch-size 8 --prompt-limit 16 --pair-limit 9 --max-open-prompts 6 --num-samples-per-prompt 2 --max-new-tokens 24` (the original 16-prompt × 9-pair grid; `baseline_comparison.csv` / `candidate_constrained_summary.csv` at scale 16 give the primary-methods table). Cross-model rows: Qwen3.5-2B and DeepSeek-R1-Distill-Qwen-7B at 32x/k256, Ministral-3-8B-Base at 16x/k128 (`configs/registries/result1_query_fidelity_cluster.yaml`). Note: with the extended scale grid, `choose_operating_point` can select a scale above 16; read the scale-16 rows from `candidate_constrained_summary.csv` | | Appendix K sweep / selection tables (`tab:app-k-model-finalists`, `tab:app-model-suite-current`, and the other `tab:app-k-*` sweep ledgers) | — (no figure ships; numbers hand-transcribed from the audited literals in [`data/appendix/appendix_k_sweep_tables.json`](../data/appendix/appendix_k_sweep_tables.json)) | `configs/sweeps/arch_frontier_160m.yaml`, `configs/sweeps/paper_phase2_frontier.yaml` | | Appendix L: global readout-feature profiles (shell/java, port/cell) | `scripts/analyze/analyze_readout_feature_directions.py` | requires `--checkpoint` (the Qwen3.5-2B 32x/k256 selected SAE); `--sae-id qwen2b_k256 --context-set interesting_domains --case-ids shell_beach,shell_terminal,java_coffee,java_programming --top-directions 20 --out-dir results/qwen_pre_token_readout_directions_shell_java` (rerun with the `port_harbor,port_network,cell_prison,cell_biology` case-ids and `--out-dir results/qwen_pre_token_readout_directions_port_cell` — the output basenames are fixed, so the two runs MUST use distinct `--out-dir`s or the second overwrites the first); persists the per-direction metrics behind the global-readout bars | | Appendix M: metrics behind feature-resolved DLA, vocab-mean verify contrast | `scripts/figures/compute_prism_dla.py` | requires `--checkpoint` (the Qwen3.5-2B 32x/k256 selected SAE); `--contrast-mode vocab_mean --out-dir results/qwen2b_32x_prism_dla_verify_token_layer_readout` (output basenames are the fixed `qwen2b_32x_prism_dla_*` prefix; the dir distinguishes this vocab-mean run from the default `_verify_assume`) | @@ -144,6 +160,31 @@ check. | Appendix K: W_U row-norm tail stats (`tab:app-k-row-norm-tail`) | `scripts/analyze/compute_row_norm_tail_stats.py` | `--registry configs/registries/result1_query_fidelity_cluster.yaml --models Qwen3.5-9B,Ministral-3-8B-Base,R1-Distill-Qwen-7B,R1-Distill-Llama-8B` (or `--artifact "label=.pt"`); centred row norms, no GPU/checkpoint | | Appendix K: dense / negative-control diagnostic (`omp_diagnostic` in [`data/appendix/appendix_k_sweep_tables.json`](../data/appendix/appendix_k_sweep_tables.json)) | `scripts/eval/dense_control_diagnostic.py` | `--w-u .pt --checkpoint --setting "Qwen-9B 32x, k=256" --n-rows 20000 --bootstrap 1000` (encoder vs LS/NNLS-on-support vs signed/nonneg OMP vs dense rank-k, scored by rowEV) | | Appendix L: feature-label audit (`tab:app-main-case-study-feature-audit` data; `tab:app-qualitative-feature-label-audit` counts) | `scripts/analyze/audit_feature_labels.py` | `substrate --feature-ids 36,4095,… --model-id Qwen/Qwen3.5-2B` emits per-feature top rows; `aggregate --annotations data/audit/feature_label_audit_annotations.csv --validate-against data/audit/feature_label_audit.json` tallies the counts (classification is human; see note below) | +| Matched-KL frontier metrics (`fig:lexical-matched-kl-frontier`) | `scripts/run/run_qwen_profanity_suppression_eval.py` | once per model (Qwen/Qwen3.5-2B, Qwen/Qwen3.5-0.8B, Qwen/Qwen3.5-9B, deepseek-ai/DeepSeek-R1-Distill-Qwen-7B; 32x/k256 checkpoints): `--checkpoint --model-id --out-dir --device cuda --dtype bfloat16 --batch-size 8 --max-open-prompts 0` (full 20-prompt × 17-pair × 14-scale × 8-method grid); the frontier plots the `split == heldout` rows of `candidate_constrained_summary.csv` (`median_kl_bits` vs `mean_bad_prob_reduction` / `candidate_flip_rate` per method and scale) | +| Paired differences at matched KL quoted with `fig:lexical-matched-kl-frontier` (SRP − mean-row / PCA rank-1 / PCA rank-4; term- and prompt-clustered 95% CIs) | `scripts/eval/paired_matched_kl_bootstrap.py` | `--input qwen2b=/candidate_constrained_rows.csv --input qwen0p8b=... --input qwen9b=... --input r1qwen7b=... --out paired_matched_kl_results.json --n-boot 10000 --seed 0` (default `--target-kl 0.02 0.05 0.1 0.2`); CPU, ~17 s for four models. Reproduces the paper's paired results exactly from the paper CSVs | +| `tab:app-cross-model-nulls` | `scripts/run/run_readout_baseline_comparisons.py` | one run per model Qwen3.5-0.8B/2B/9B, Gemma-4-E2B, Gemma-4-E4B with `--methods sparse_rp,shuffled_row_code,random_support_same_magnitudes` (all other flags as in the row above); coverage = `accepted_rate`, sign = `sign_agreement` in `baseline_by_method.csv` | +| `fig:app-robustness-baseline-comparisons` (Qwen3.5-2B reference panel; the 848-cluster CI paragraph of the fidelity-results subsection) | `scripts/run/run_readout_baseline_comparisons.py` | `--model Qwen3.5-2B --methods sparse_rp,shuffled_row_code,random_support_same_magnitudes,pca_64,pca_256,pca_1024,nearest_row_ridge_top128` (other flags as above); `baseline_by_method.csv`, `baseline_by_method_margin_bin.csv`, `baseline_feature_compactness.csv` | +| `tab:app-error-tails`, `tab:app-error-tail-margins` (and the per-model mean/p95/max ranges quoted in that subsection) | `scripts/eval/analyze_error_tails.py` | `--input-root results/direct_geometry_runs --input-prefix --model-tags qwen0p8b,qwen2b,qwen9b,ministral8b,r1qwen7b,r1llama8b --primary-model-tag qwen2b --out-dir results/error_tails --n-boot 10000 --seed 20260711` over the six direct-geometry run dirs; `summary_pooled_by_method.csv` (8,009 contrasts per method), the `sparse_rp` block of `summary_pooled_by_method_margin_bin.csv`, `summary_by_model_method.csv`. Verified to reproduce both tables from the paper's raw run CSVs | +| `tab:app-nearest-rows` | `scripts/analyze/nearest_rows_baseline.py` | `--model-id Qwen/Qwen3.5-2B --revision 15852e8c16360a2fea060d615a32b45270f8a8fc --contrast bug,insect --contrast bug,error --top-n 12 --out-dir results/nearest_rows_qwen2b` (full-vocabulary centering mean, contrast tokens excluded); the `centered_*` columns of `nearest_rows_table.csv` are the table, `nearest_rows.json` holds the raw-cosine ranking too | +| Table 2 predicted-vs-realized readout-side changes, one row per model (`tab:causal-validation-summary`; per-model rows and protocol in `app:causal-validation`) | `scripts/eval/run_causal_contribution_validation.py` | once per model, fidelity operating point (32x, k=256): `--model-id --checkpoint /k256_32x/checkpoint.pt")> --bank-dir data/query_banks --out-dir results/causal_ --device cuda --dtype bfloat16 --seed 0`, with defaults `--banks curated_ab,case_candidates,model_native --max-native 300 --top-features 10 --random-per-case 10 --rho-gate 0.5 --n-boot 2000 --max-len 64`. Model → Hub dir: `Qwen/Qwen3.5-0.8B`→`qwen3.5-0.8b`, `Qwen/Qwen3.5-2B`→`qwen3.5-2b`, `Qwen/Qwen3.5-9B`→`qwen3.5-9b`, `mistralai/Ministral-3-8B-Base-2512`→`ministral-3-8b`, `deepseek-ai/DeepSeek-R1-Distill-Qwen-7B`→`r1-distill-qwen-7b`, `deepseek-ai/DeepSeek-R1-Distill-Llama-8B`→`r1-distill-llama-8b`. Table columns come from `summary.json`: r²/CI/slope/n ← `gated_predicted`, Random ← `gated_random_control.r2`; `ungated_predicted` is the all-pairs comparison. The `model_native` bank carries no A/B targets and contributes zero contrasts (a missing file is logged and skipped) | +| Appendix harness self-test line (`app:causal-validation`: r²=1.000, slope 1.04, zero-scoring controls) | `scripts/eval/run_causal_contribution_validation.py` | `--self-test` (CPU, no model/checkpoint; prints `r2=1.000 slope=1.044`, three PASS lines) | +| Appendix I: cross-seed stability of contrast explanations (`tab:app-cross-seed-stability`) | `scripts/eval/cross_seed_stability.py` | three same-recipe checkpoints trained from `configs/sweeps/qwen35_2b_seedvar_base.yaml`; `--dictionaries results/qwen35_2b_seedvar/qwen2b_d65536_k256_s0/checkpoint.pt …_s1/checkpoint.pt …_s2/checkpoint.pt --w-u "$SRP_ARCHIVE_ROOT/data/qwen35-2b/qwen35_2b.pt" --bank data/query_banks/qwen_gemma_result1_curated_ab.jsonl --width-tag 32x --out results/seed_stability/cross_seed_stability_32x.json` (defaults `--max-contrasts 150 --seed 0 --top-m 8 --top-r 12 --n-sample 4096 --n-hidden 512`; 16x column: the `qwen2b_d32768_k128_s{0,1,2}` checkpoints, `--width-tag 16x`) | +| Appendix I: cross-seed feature-group matching incl. the below-null direction check (`tab:app-feature-group-matching`) | `scripts/eval/feature_group_matching.py` | same dictionaries / `--w-u` / `--bank` as above; `--width-tag 32x --direction-check --out results/seed_stability/feature_group_matching_32x.json` (defaults `--n-query 100 --n-null 500 --greedy-pool 200 --greedy-k 3 --min-group-size 4 --direction-null 500 --direction-seed 0`; 16x: the k128 checkpoints, `--width-tag 16x`) | +| Appendix I: held-out recovery of the stable core (`tab:app-loo-core-recovery`, plus the 0.444 cross-recipe figure) | `scripts/eval/loo_core_recovery.py` | same dictionaries / `--w-u` / `--bank`; `--width-tag 32x --reference-dict --out results/seed_stability/loo_core_recovery_32x.json` (defaults `--side-n 96 --n-clusters 16384 --kmeans-iters 12 --kmeans-seed 0 --n-boot 2000 --bootstrap-seed 1`; 16x: k128 checkpoints, `--width-tag 16x`, no `--reference-dict`; k-means wants a GPU/MPS) | +| Appendix K: sparsity budget vs dictionary usage (`tab:app-low-k-dead-features`) | `scripts/train/train_readout_sae_from_config.py` | multi-seed rows: `--config configs/sweeps/qwen35_2b_seedvar_base.yaml` with `--set factorizer.d_features= --set factorizer.k= --set evaluation.k= --set run.seed= --set run.init_seed= --set data.data_seed= --set run.output_dir=results/qwen35_2b_seedvar/qwen2b_d_k_s --set run.name=qwen2b_d_k_s` over cells (32768,256,{0,1,2}) (65536,128,{0,1,2}) (32768,64,{1,2}) (plus the Appendix I family (65536,256,{0,1,2}) (32768,128,{0,1,2})); single-seed rows: `--config configs/sweeps/lowk_qwen35_2b_16x_base.yaml` and `--config configs/sweeps/lowk_qwen35_0p8b_16x_base.yaml` with `--set factorizer.k=<32 or 64> --set evaluation.k=<32 or 64>` and `run.name`/`run.output_dir`; dead / rare / top-1 / KL read from each cell's `metrics.json` | +| Appendix K: Qwen3.5-0.8B 32x/k256 seed window (`tab:app-k-qwen08b-seed-window`) | `scripts/train/train_readout_sae_from_config.py` | `--config configs/sweeps/qwen35_0p8b_paper_topk_32x_k256_s1.yaml` and `--config configs/sweeps/qwen35_0p8b_paper_topk_32x_k256_s2.yaml`, no `--set`; seed 0 row is the archived `qwen0p8b_k256_32x` dictionary in `configs/registries/exp2_selected_sae_checkpoints.yaml` | +| Appendix "Sense-Labelled Evaluation": frozen CoarseWSD-20 readout bundles and coverage statistics (`app:sense-labelled-evaluation` data paragraph: contexts scored, single-token coverage, row-gate coverage, median target rank) | `scripts/run/run_wsd_feature_alignment.py` | `--dataset coarsewsd20 --data-root --model-id --checkpoint --out-dir results/wsd_core/coarsewsd20_ --splits train,test --batch-size 8 (2B) / 4 (9B, R1-Llama) --n-boot 5000 --seed 0`; writes `representations.pt` (input to the two rows below), `audit.json`, `scoring_summary.json`, `metrics.json`; `--analyze-only --out-dir --n-boot 5000 --seed 0` recomputes `metrics.json` without inference | +| Sense alignment on CoarseWSD-20 (`tab:app-sense-alignment`; group-size 1/2/4/8 sweep in the appendix prose) | `scripts/analyze/analyze_wsd_sense_groups.py` | `--bundle results/wsd_core/coarsewsd20_/representations.pt --out results/wsd_sense_groups/coarsewsd20___srp.json --group-sizes 1,2,4,8 --selector mean_diff --n-null 200 --n-boot 2000 --seed 0` (defaults: `--basis srp`, train-standardized); the table reads `_g8.json` fields `word_mean_balanced`, `word_mean_null_balanced`, `word_mean_majority_balanced`, `words_beating_null_p95_balanced`, `words_beating_majority_balanced`, `word_mean_full_account_balanced`, `word_mean_hidden_balanced`; CPU only. Verified against the paper run outputs cell for cell | +| Appendix "Sense-Labelled Evaluation", classifier-framing paragraph (projections vs signed contributions, row-coefficient shuffles, random features) | `scripts/analyze/analyze_wsd_classifier_framing.py` | `--bundle results/wsd_core/coarsewsd20_/representations.pt --out results/wsd_classifier_framing/coarsewsd20_.json --ks 5,10,20,50 --primary-k 10 --primary-encoding weighted --n-boot 5000 --n-null-seeds 20 --seed 0` (all defaults); the paragraph quotes `primary.methods..accuracy`; CPU only | +| Fitted lenses for the cross-lens study (inputs to all `fig:cross-lens-*` / `tab:app-cross-lens-*`) | `scripts/run/fit_jlens.py` | `prompts --prompts-json /c4_prompts_en_seed0.json --n-prompts 1000 --seed 0` (default `--min-chars 800 --c4-config en`); `prompts … c4_prompts_zh_seed0.json --n-prompts 1000 --seed 0 --min-chars 300 --c4-config zh`; same with `--c4-config de`; `smoke --model-id Qwen/Qwen3.5-0.8B --prompts-json --out /smoke_0p8b.lens.pt --dim-batch 16`; then per corpus `fit --model-id Qwen/Qwen3.5-9B --prompts-json --n-prompts 100 --shard $i --num-shards 4 --n-layers 12 --dim-batch 8 --ckpt-dir --out-dir ` for i in 0..3 and `merge --shard-dir --num-shards 4 --out /qwen35_9b_jlens_{en,zh,de}_seed0_n100.pt`; `--n-prompts 300` for the `*_n300.pt` refits. Needs `uv sync --extra lens` | +| Ridge translators (`tab:app-cross-lens-extension` rows 3–4, `fig:cross-lens-extension`) | `scripts/run/fit_ridge_lens.py` | `--model-id Qwen/Qwen3.5-9B --prompts-json /c4_prompts_en_seed0.json --n-prompts 100 --holdout 10 --layers-from /qwen35_9b_jlens_en_seed0_n100.pt --ckpt-dir --tag ridge_en --out /qwen35_9b_ridgelens_en_seed0_n100.pt`; repeat with `c4_prompts_zh_seed0.json --tag ridge_zh --out …ridgelens_zh_seed0_n100.pt` (layers still from the EN Jacobian lens). Holdout R² and deep top-1 agreement land in `.report.json` | +| Readout dumps behind every cross-lens artifact | `scripts/run/run_cross_lens_readouts.py` | once per (lens, bank): `--model-id Qwen/Qwen3.5-9B --lens --sae --k 128 --prompts data/cross_lens/cross_lens_prompts_en_zh.json --n-positions 1 --decompose-top1 --out /en_zh__.json` for the EN/ZH Jacobian n=100 and n=300 lenses and the EN/ZH ridge translators; `--prompts data/cross_lens/cross_lens_prompts_en_de.json` for the EN/DE Jacobian n=100 and n=300 lenses; three-lens example: `--prompts data/cross_lens/cross_lens_three_lens_prompts.json --n-positions 1 --decompose-top1 --top-feats 10` under the EN, ZH, DE n=100 lenses (optionally `--lens identity --layers-from `) | +| `tab:app-cross-lens-families`; `tab:app-cross-lens-extension` rows 1, 3, 4, 5; `fig:cross-lens-extension` (EN–ZH cells) | `scripts/eval/aggregate_cross_lens_en_zh.py` | `--prompts data/cross_lens/cross_lens_prompts_en_zh.json --seed 0 --out /_summary.json [--rows-csv …]` with `--lens-a/--lens-b` = (jlens_en_n100, jlens_zh_n100) main study; (ridge_en, ridge_zh) ridge row; (jlens_en_n100, ridge_en) Jacobian-vs-ridge row; (jlens_en_n300, jlens_zh_n300) n=300 row. Per-family agreement/CI and the "Lens-only split" column are `per_group` and `lens_only_split`; floors are `null_within_lens` (/159) and `null_shuffle_cross_lens` (/240). Verified to reproduce the paper run's summary key for key | +| `tab:app-cross-lens-extension` rows 2, 6; `fig:cross-lens-extension` (EN–DE cells); second-language-pair per-family counts, 73/90 and 76/90 divergence | `scripts/eval/aggregate_cross_lens_en_de.py` | `--lens-a /en_de__jlens_en_n100.json --lens-b /en_de__jlens_de_n100.json --prompts data/cross_lens/cross_lens_prompts_en_de.json --seed 0 --out /en_de_jlens_n100_summary.json`; same with the `_n300` dumps. Divergence is `divergence_rate`; floors /204 and /270 | +| EN–DE bank + cognate-exclusion report (inputs to the EN–DE cells) | `scripts/data/build_cross_lens_de_bank.py` | `--tokenizer Qwen/Qwen3.5-9B --out data/cross_lens/cross_lens_prompts_en_de.json --report data/cross_lens/cross_lens_prompts_en_de_exclusion_report.json` (shipped outputs are in `data/cross_lens/`) | +| `tab:app-cross-lens-antonym-layers` (and the f112 contributions quoted beside it) | `scripts/analyze/cross_lens_antonym_layers.py` | `run --model-id Qwen/Qwen3.5-9B --lens /qwen35_9b_jlens_en_seed0_n100.pt --sae --k 128 --out /antonym_layers_en.json` (default `--prompt '"小"的反义词是"'`); same with the ZH lens → `antonym_layers_zh.json`; then `table --dump EN=/antonym_layers_en.json --dump ZH=/antonym_layers_zh.json --layers 24,26,29,final --out-csv /antonym_layers_table.csv --out-features-csv /antonym_layers_features.csv` | +| `tab:app-cross-lens-de-antonym-layers` (and the groß/large/big feature ids and 15.16/15.60/22.19 logits in the prose) | `scripts/analyze/cross_lens_three_lens_prompt.py` | `--dump EN=/three_lens__jlens_en.json --dump ZH=/three_lens__jlens_zh.json --dump DE=/three_lens__jlens_de.json --prompt-id antonym_de_01 --layers 24,26,29 --out-csv /three_lens_top1.csv --out-targets-csv /three_lens_targets.csv --out-features-csv /three_lens_features.csv` (table = `top1_token`, `top1_lens_logit`; quoted groß logits = `original_logit` column) | +| `fig:cross-lens-butterfly` (metrics only) | `scripts/figures/compute_cross_lens_shared_feature.py` | `--dump EN=/en_zh__jlens_en_n100.json --dump ZH=/en_zh__jlens_zh_n100.json --prompt-id fac_03 --layer 26 --out-csv /cross_lens_shared_feature.csv --out-json /cross_lens_shared_feature.json` (top-1 logits 38.5 / 39.2, f23180 shares 0.752 / 0.629, largest other +1.69 / +1.55) | ### Notes and paper-only artifacts @@ -152,7 +193,7 @@ lexical-edit tables are transcribed by hand from the CSV outputs of the scripts above. A few artifacts have no producing script and are authored in the paper's LaTeX source: -- **Schematics** — `fig:readout-score-schematic` (`readout_query_three_panel.pdf`) and the +- **Schematics** — `fig:method-schematic` (`readout_query_one_column.pdf`) and the score-accounting TikZ schematic are hand-drawn LaTeX/TikZ; there is no repo deliverable. - **Feature-label audit counts** (`tab:app-qualitative-feature-label-audit`) — the @@ -180,6 +221,13 @@ LaTeX source: uv run pytest -q ``` +The ported experiment code carries its own CPU tests: additive identity of the +weighted row-kNN and k-means centroid baselines, the causal-validation harness +self-test, the paired matched-KL bootstrap on synthetic rows, the cross-lens +aggregators and table builders on stub dumps, sense-group selection and +balanced accuracy on a synthetic bundle, and the seed-stability contrast +pipeline on synthetic dictionaries. + The test suite is fast and CI-safe (no GPU, no model loads): identity correctness of the decomposition, schema invariants of the task-fidelity evaluator, checkpoint-registry resolution, the data-loading branches diff --git a/pyproject.toml b/pyproject.toml index 488e53f..5bf6fc3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "sparse-readout-prism" -version = "0.1.0" +version = "0.2.0" description = "Sparse Readout Prism: a sparse LM-head basis that factorizes language-model unembedding rows into reusable readout features and decomposes selected logits into signed feature contributions plus an explicit residual." readme = "README.md" license = { file = "LICENSE" } @@ -39,6 +39,8 @@ dependencies = [ "scipy>=1.15.3,<2", # scripts/eval/dense_control_diagnostic.py (scipy.optimize.nnls) "datasets>=4.8.4,<5", # scripts/data/extract_model_readout.py + build_query_banks.py "pyyaml>=6,<7", # config loading in utils.py + research/registry.py (was only transitive) + "pandas>=2.2,<3", # scripts/eval/analyze_error_tails.py + analyze_task_grounded_fidelity.py (CSV aggregation) + "scikit-learn>=1.5,<2", # scripts/run/run_wsd_feature_alignment.py + scripts/analyze/analyze_wsd_*.py ] [project.urls] @@ -51,6 +53,10 @@ Documentation = "https://github.com/hematteo/sparse-readout-prism#documentation" requires = ["hatchling"] build-backend = "hatchling.build" +[tool.hatch.metadata] +# Required for the git-sourced `lens` extra above. +allow-direct-references = true + [tool.hatch.build.targets.wheel] packages = ["src/sparse_readout_prism"] @@ -63,6 +69,12 @@ dev = [ "pytest>=7", "ruff>=0.7", ] +# Jacobian-lens reference implementation (Apache-2.0, not on PyPI). Needed only +# by the cross-lens study scripts (scripts/run/fit_jlens.py and friends), which +# import it lazily so the rest of the repo installs and tests without it. +lens = [ + "jlens @ git+https://github.com/anthropics/jacobian-lens", +] [tool.ruff] line-length = 120 diff --git a/scripts/README.md b/scripts/README.md index 721ee98..64eccc3 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -19,11 +19,11 @@ figure / table → script + config map; every script below appears in it. | Bucket | Purpose | |--------------|------------------------------------------------------------------------------------------------------------------| | `train/` | Train readout-row factorizers (TopK, Matryoshka, JumpReLU, ...). | -| `run/` | End-to-end experiment suites (query-fidelity bank, baseline/null comparisons, lexical-edit / readout-side control stress test). | -| `eval/` | Evaluate trained factorizers (task-fidelity evaluation; dense / negative-control row-reconstruction diagnostic). | -| `analyze/` | Inspect readout-feature projection directions at the final hidden state; W_U row-norm tail stats; feature-label audit. | -| `figures/` | Compute and persist the metrics behind the paper figures from `results/` artefacts (metric / eval tables, example panels, lens grids, decompositions). This repo is metrics-only: it persists those metrics as CSV/JSON/.pt; the figures themselves are rendered separately in the paper's LaTeX source from these artefacts. | -| `data/` | Extract readout + hidden states from a model, build query banks, mine same-token polysemy context pairs. | +| `run/` | End-to-end experiment suites: query-fidelity bank, direct-geometry baseline / null comparisons, lexical-edit / readout-side control stress test (with the matched-KL frontier controls), Jacobian-lens and ridge-translator fitting, cross-lens readout dumps, CoarseWSD-20 representation bundles. | +| `eval/` | Evaluate trained factorizers and aggregate run outputs: task-fidelity evaluation; dense / negative-control row-reconstruction diagnostic; causal validation of feature contributions; pooled error tails; paired matched-KL bootstrap; EN–ZH / EN–DE cross-lens aggregators; cross-seed stability, feature-group matching and leave-one-out core recovery. | +| `analyze/` | Inspect readout-feature projection directions at the final hidden state; W_U row-norm tail stats; feature-label audit; nearest-rows reading of selected directions; cross-lens antonym-layer and three-lens tables; sense-group alignment and the classifier framing on CoarseWSD-20. | +| `figures/` | Compute and persist the metrics behind the paper figures from `results/` artefacts (metric / eval tables, example panels, lens grids, decompositions, the cross-lens shared-feature panel). This repo is metrics-only: it persists those metrics as CSV/JSON/.pt; the figures themselves are rendered separately in the paper's LaTeX source from these artefacts. | +| `data/` | Extract readout + hidden states from a model, build query banks, mine same-token polysemy context pairs, build the EN–DE cross-lens bank with its cognate-exclusion report. | ## Output conventions @@ -34,6 +34,11 @@ gitignored). Pass `--paper-dir ''` to skip the mirror. Run manifests record provenance (command line, args, git commit, package/torch versions) via `sparse_readout_prism.research.run_io.run_provenance`. +The cross-lens scripts (`run/fit_jlens.py`, `run/fit_ridge_lens.py`, +`run/run_cross_lens_readouts.py`, `analyze/cross_lens_antonym_layers.py run`) +import the Jacobian-lens reference implementation lazily; install it with +`uv sync --extra lens`. Every other script runs from the base environment. + The earlier exploratory families (causal / attention-head circuits, CoT-faithfulness and agentic-injection probes, the 160M transcoder/SAE lineage, the classical / tuned proto-token-lens renderers, and the per-snapshot diff --git a/scripts/analyze/analyze_wsd_classifier_framing.py b/scripts/analyze/analyze_wsd_classifier_framing.py new file mode 100644 index 0000000..ac9392d --- /dev/null +++ b/scripts/analyze/analyze_wsd_classifier_framing.py @@ -0,0 +1,669 @@ +#!/usr/bin/env python3 +"""Truncated-display sense classifier on CoarseWSD-20 readout bundles. + +Backs the "classifier framing" paragraph of the sense-labelled evaluation +appendix (``app:sense-labelled-evaluation``). Per word, a nearest-centroid +classifier is fit on the train split over the ``--primary-k`` features with +the largest absolute value, retaining each displayed term's signed value +(``weighted``) or only its sign (``signed``), and evaluated on the held-out +split. Four sources are compared under the same encoding: + + srp signed contributions z_{v,i} * p_i(h) + projection_only unweighted projections p_i(h) + shuffled_srp row coefficients permuted within the word (per-word seed) + random_srp_features k feature positions drawn at random per word + +The paragraph quotes ``primary.methods..accuracy`` (projections +against contributions, with the shuffled and random sources alongside) at +``--primary-encoding weighted --primary-k 10``; ``balanced_accuracy`` and +``macro_f1`` are written next to it, with per-word bootstrap intervals, +paired per-word differences against the ``srp`` source and a null-seed +sensitivity for the two randomized sources. Secondary curves cover every +``--ks`` value under both encodings. Feature-level readings are restricted to +the local score gate: sign preservation and +|residual| / (|centered exact score| + 0.5) < 0.5. + +AmbiStory bundles go through the same encoding path (story display against +the two gloss-anchor displays); the paper uses CoarseWSD-20 only. + +Reads the ``representations.pt`` bundles written by +``scripts/run/run_wsd_feature_alignment.py``; CPU-only, deterministic under +``--seed``. Paper run, once per model bundle (Qwen3.5-2B, Qwen3.5-9B, +DeepSeek-R1-Distill-Llama-8B), all values at their defaults: + + python scripts/analyze/analyze_wsd_classifier_framing.py \ + --bundle results/wsd_core/coarsewsd20_qwen2b/representations.pt \ + --out results/wsd_classifier_framing/coarsewsd20_qwen2b.json \ + --ks 5,10,20,50 --primary-k 10 --primary-encoding weighted \ + --n-boot 5000 --n-null-seeds 20 --seed 0 +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from collections import defaultdict +from pathlib import Path +from typing import Any + +import numpy as np +import torch + + +METHODS = ("srp", "projection_only", "shuffled_srp", "random_srp_features") +ENCODINGS = ("signed", "weighted") + + +def stable_seed(text: str, seed: int) -> int: + digest = hashlib.sha1(text.encode("utf-8")).hexdigest()[:8] + return (int(digest, 16) + seed) % (2**32) + + +def write_json(path: Path, value: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(value, indent=2, ensure_ascii=False) + "\n") + + +def score_gate(bundle: dict[str, Any]) -> tuple[np.ndarray, dict[str, float]]: + contribution = bundle["contribution"].float().numpy() + exact = bundle["exact_logit"].float().numpy() + reconstructed = bundle["reconstructed_logit"].float().numpy() + feature_sum = contribution.sum(axis=1) + base = reconstructed - feature_sum + centered_exact = exact - base + residual = exact - reconstructed + rho = np.abs(residual) / (np.abs(centered_exact) + 0.5) + sign_match = np.sign(centered_exact) == np.sign(feature_sum) + passed = sign_match & (rho < 0.5) + return passed, { + "n_items": int(len(passed)), + "n_pass": int(passed.sum()), + "pass_fraction": float(passed.mean()), + "sign_match_fraction": float(sign_match.mean()), + "median_rho_plus_0p5": float(np.median(rho)), + } + + +def validate_feature_alignment(bundle: dict[str, Any]) -> None: + metadata = bundle["metadata"] + feature_ids = bundle["feature_ids"].numpy() + first_by_target: dict[str, np.ndarray] = {} + for row, ids in zip(metadata, feature_ids): + target = str(row["target"]) + if target in first_by_target: + if not np.array_equal(first_by_target[target], ids): + raise ValueError(f"feature coordinates change within target {target!r}") + else: + first_by_target[target] = ids.copy() + + +def source_matrix(bundle: dict[str, Any], method: str, seed: int) -> np.ndarray: + projection = bundle["projection"].float().numpy() + contribution = bundle["contribution"].float().numpy() + beta = bundle["beta"].float().numpy() + metadata = bundle["metadata"] + if method == "srp" or method == "random_srp_features": + return contribution + if method == "projection_only": + return projection + if method != "shuffled_srp": + raise KeyError(method) + shuffled = np.empty_like(contribution) + by_target: dict[str, list[int]] = defaultdict(list) + for i, row in enumerate(metadata): + by_target[str(row["target"])].append(i) + for target, indices in by_target.items(): + rng = np.random.default_rng(stable_seed(target, seed)) + permutation = rng.permutation(beta.shape[1]) + idx = np.asarray(indices) + shuffled[idx] = projection[idx] * beta[idx][:, permutation] + return shuffled + + +def encode_top_features( + bundle: dict[str, Any], + method: str, + k: int, + encoding: str, + seed: int, +) -> np.ndarray: + matrix = source_matrix(bundle, method, seed) + n_rows, width = matrix.shape + if not 1 <= k <= width: + raise ValueError(f"invalid k={k} for width={width}") + selected = np.empty((n_rows, k), dtype=np.int64) + if method == "random_srp_features": + by_target: dict[str, list[int]] = defaultdict(list) + for i, row in enumerate(bundle["metadata"]): + by_target[str(row["target"])].append(i) + for target, rows in by_target.items(): + rng = np.random.default_rng(stable_seed(f"random:{target}", seed)) + choice = rng.permutation(width)[:k] + selected[np.asarray(rows)] = choice[None, :] + else: + selected = np.argpartition(np.abs(matrix), width - k, axis=1)[:, -k:] + + encoded = np.zeros_like(matrix, dtype=np.float32) + row_index = np.arange(n_rows)[:, None] + values = matrix[row_index, selected] + if encoding == "signed": + values = np.sign(values) + elif encoding != "weighted": + raise KeyError(encoding) + encoded[row_index, selected] = values + norms = np.linalg.norm(encoded, axis=1, keepdims=True) + return encoded / np.maximum(norms, 1e-12) + + +def safe_spearman(x: list[float], y: list[float]) -> float: + from scipy.stats import spearmanr + + if len(x) < 3 or np.ptp(x) < 1e-8 or np.ptp(y) < 1e-8: + return float("nan") + return float(spearmanr(x, y).statistic) + + +def cluster_bootstrap_delta( + rows: list[dict[str, Any]], + cluster_key: str, + statistic, + n_boot: int, + seed: int, +) -> list[float]: + grouped: dict[str, list[dict[str, Any]]] = defaultdict(list) + for row in rows: + grouped[str(row[cluster_key])].append(row) + keys = sorted(grouped) + rng = np.random.default_rng(seed) + values: list[float] = [] + for _ in range(n_boot): + sample: list[dict[str, Any]] = [] + for selected in rng.choice(keys, size=len(keys), replace=True): + sample.extend(grouped[str(selected)]) + value = float(statistic(sample)) + if np.isfinite(value): + values.append(value) + return values + + +def confidence_interval(values: list[float]) -> list[float]: + if not values: + return [float("nan"), float("nan")] + return [float(np.percentile(values, 2.5)), float(np.percentile(values, 97.5))] + + +def ambistory_rows(bundle: dict[str, Any], encoded: np.ndarray, gate: np.ndarray) -> list[dict[str, Any]]: + metadata = bundle["metadata"] + anchor_index = { + (str(row["target"]), str(row["anchor_key"])): i for i, row in enumerate(metadata) if row["kind"] == "anchor" + } + rows: list[dict[str, Any]] = [] + for i, row in enumerate(metadata): + if row["kind"] != "story" or not gate[i]: + continue + anchor_ids = [anchor_index.get((str(row["target"]), str(key))) for key in row["anchor_keys"]] + if any(index is None or not gate[int(index)] for index in anchor_ids): + continue + similarities = [float(encoded[i] @ encoded[int(index)]) for index in anchor_ids] + human_margin = float(row["human_margin"]) + predicted_margin = similarities[0] - similarities[1] + rows.append( + { + "item_id": str(row["item_id"]), + "setup_id": str(row["setup_id"]), + "target": str(row["target"]), + "human_margin": human_margin, + "predicted_margin": predicted_margin, + "eligible_gap1": int(abs(human_margin) >= 1.0), + "predicted_tie": int(abs(predicted_margin) < 1e-12), + "correct": int(human_margin != 0 and np.sign(human_margin) == np.sign(predicted_margin)), + "tie_half_score": ( + 0.5 + if abs(predicted_margin) < 1e-12 + else float(human_margin != 0 and np.sign(human_margin) == np.sign(predicted_margin)) + ), + } + ) + return rows + + +def summarize_ambistory_rows(rows: list[dict[str, Any]]) -> dict[str, Any]: + eligible = [row for row in rows if row["eligible_gap1"]] + eligible_non_ties = [row for row in eligible if not row["predicted_tie"]] + non_ties = [row for row in rows if row["human_margin"] != 0] + return { + "n_contexts": len(rows), + "n_setups": len({row["setup_id"] for row in rows}), + "n_targets": len({row["target"] for row in rows}), + "preference_spearman": safe_spearman( + [row["human_margin"] for row in rows], + [row["predicted_margin"] for row in rows], + ), + "preferred_sense_accuracy_gap1": ( + float(np.mean([row["tie_half_score"] for row in eligible])) if eligible else float("nan") + ), + "preferred_sense_accuracy_gap1_non_abstain": ( + float(np.mean([row["correct"] for row in eligible_non_ties])) if eligible_non_ties else float("nan") + ), + "preferred_sense_coverage_gap1": (len(eligible_non_ties) / len(eligible) if eligible else float("nan")), + "n_gap1": len(eligible), + "preferred_sense_accuracy_non_tie": ( + float(np.mean([row["correct"] for row in non_ties])) if non_ties else float("nan") + ), + } + + +def paired_ambistory_comparison( + srp_rows: list[dict[str, Any]], + baseline_rows: list[dict[str, Any]], + n_boot: int, + seed: int, +) -> dict[str, Any]: + baseline_by_id = {row["item_id"]: row for row in baseline_rows} + paired = [] + for srp in srp_rows: + baseline = baseline_by_id.get(srp["item_id"]) + if baseline is None: + continue + paired.append( + { + "setup_id": srp["setup_id"], + "human_margin": srp["human_margin"], + "eligible_gap1": srp["eligible_gap1"], + "srp_margin": srp["predicted_margin"], + "baseline_margin": baseline["predicted_margin"], + "srp_correct": srp["correct"], + "baseline_correct": baseline["correct"], + "srp_tie_half_score": srp["tie_half_score"], + "baseline_tie_half_score": baseline["tie_half_score"], + } + ) + + def rho_delta(sample): + human = [row["human_margin"] for row in sample] + return safe_spearman(human, [row["srp_margin"] for row in sample]) - safe_spearman( + human, [row["baseline_margin"] for row in sample] + ) + + eligible = [row for row in paired if row["eligible_gap1"]] + + def accuracy_delta(sample): + return float(np.mean([row["srp_tie_half_score"] for row in sample])) - float( + np.mean([row["baseline_tie_half_score"] for row in sample]) + ) + + return { + "n_contexts": len(paired), + "preference_spearman_difference": float(rho_delta(paired)), + "preference_spearman_difference_ci": confidence_interval( + cluster_bootstrap_delta(paired, "setup_id", rho_delta, n_boot, seed) + ), + "preferred_sense_accuracy_gap1_difference": float(accuracy_delta(eligible)), + "preferred_sense_accuracy_gap1_difference_ci": confidence_interval( + cluster_bootstrap_delta(eligible, "setup_id", accuracy_delta, n_boot, seed + 1) + ), + } + + +def bootstrap_ambistory_method(rows: list[dict[str, Any]], n_boot: int, seed: int) -> dict[str, Any]: + eligible = [row for row in rows if row["eligible_gap1"]] + + def rho(sample): + return safe_spearman( + [row["human_margin"] for row in sample], + [row["predicted_margin"] for row in sample], + ) + + def accuracy(sample): + return float(np.mean([row["tie_half_score"] for row in sample])) + + return { + **summarize_ambistory_rows(rows), + "preference_spearman_ci": confidence_interval(cluster_bootstrap_delta(rows, "setup_id", rho, n_boot, seed)), + "preferred_sense_accuracy_gap1_ci": confidence_interval( + cluster_bootstrap_delta(eligible, "setup_id", accuracy, n_boot, seed + 1) + ), + } + + +def summarize_null_values(values: list[float], observed: float) -> dict[str, Any]: + array = np.asarray(values, dtype=float) + return { + "n_seeds": len(values), + "mean": float(array.mean()), + "std": float(array.std()), + "min": float(array.min()), + "max": float(array.max()), + "q05_q95": [float(v) for v in np.quantile(array, [0.05, 0.95])], + "observed_srp": float(observed), + "fraction_null_at_least_observed": float(np.mean(array >= observed)), + } + + +def ambistory_null_seed_sensitivity( + bundle: dict[str, Any], + gate: np.ndarray, + primary_k: int, + primary_encoding: str, + n_null_seeds: int, + seed: int, + observed: dict[str, Any], +) -> dict[str, Any]: + result: dict[str, Any] = {} + for method in ("shuffled_srp", "random_srp_features"): + rho_values: list[float] = [] + accuracy_values: list[float] = [] + for offset in range(n_null_seeds): + encoded = encode_top_features(bundle, method, primary_k, primary_encoding, seed + offset) + summary = summarize_ambistory_rows(ambistory_rows(bundle, encoded, gate)) + rho_values.append(float(summary["preference_spearman"])) + accuracy_values.append(float(summary["preferred_sense_accuracy_gap1"])) + result[method] = { + "preference_spearman": summarize_null_values(rho_values, float(observed["preference_spearman"])), + "preferred_sense_accuracy_gap1": summarize_null_values( + accuracy_values, float(observed["preferred_sense_accuracy_gap1"]) + ), + } + return result + + +def analyze_ambistory( + bundle: dict[str, Any], + ks: list[int], + primary_k: int, + primary_encoding: str, + n_boot: int, + n_null_seeds: int, + seed: int, +) -> dict[str, Any]: + gate, gate_summary = score_gate(bundle) + metrics: dict[str, Any] = { + "dataset": "ambistory", + "question": "Do top contributing features identify the human-preferred sense?", + "fidelity_gate": gate_summary, + "results": {}, + "primary": { + "encoding": primary_encoding, + "k": primary_k, + "methods": {}, + "comparisons": {}, + }, + } + primary_rows: dict[str, list[dict[str, Any]]] = {} + for encoding in ENCODINGS: + metrics["results"][encoding] = {} + for k in ks: + metrics["results"][encoding][str(k)] = {} + for method in METHODS: + encoded = encode_top_features(bundle, method, k, encoding, seed) + rows = ambistory_rows(bundle, encoded, gate) + metrics["results"][encoding][str(k)][method] = summarize_ambistory_rows(rows) + if encoding == primary_encoding and k == primary_k: + primary_rows[method] = rows + for method_index, method in enumerate(METHODS): + metrics["primary"]["methods"][method] = bootstrap_ambistory_method( + primary_rows[method], n_boot, seed + 50 + method_index * 2 + ) + for comparison_index, baseline in enumerate(METHODS[1:]): + metrics["primary"]["comparisons"][f"srp_minus_{baseline}"] = paired_ambistory_comparison( + primary_rows["srp"], + primary_rows[baseline], + n_boot, + seed + 100 + comparison_index * 10, + ) + metrics["primary"]["null_seed_sensitivity"] = ambistory_null_seed_sensitivity( + bundle, + gate, + primary_k, + primary_encoding, + n_null_seeds, + seed, + metrics["primary"]["methods"]["srp"], + ) + return metrics + + +def coarse_per_word(bundle: dict[str, Any], encoded: np.ndarray, gate: np.ndarray) -> dict[str, dict[str, float]]: + from sklearn.metrics import accuracy_score, balanced_accuracy_score, f1_score + + metadata = bundle["metadata"] + words = sorted({str(row["target"]) for row in metadata}) + per_word: dict[str, dict[str, float]] = {} + for word in words: + train_idx = [ + i for i, row in enumerate(metadata) if str(row["target"]) == word and row["split"] == "train" and gate[i] + ] + test_idx = [ + i for i, row in enumerate(metadata) if str(row["target"]) == word and row["split"] == "test" and gate[i] + ] + if not train_idx or not test_idx: + continue + train_y = np.asarray([str(metadata[i]["sense"]) for i in train_idx]) + test_y = np.asarray([str(metadata[i]["sense"]) for i in test_idx]) + senses = sorted(set(train_y)) + if len(senses) < 2 or not set(test_y).issubset(set(senses)): + continue + centroids = [] + for sense in senses: + centroid = encoded[np.asarray(train_idx)[train_y == sense]].mean(axis=0) + centroid /= max(np.linalg.norm(centroid), 1e-12) + centroids.append(centroid) + scores = encoded[test_idx] @ np.stack(centroids).T + predictions = np.asarray([senses[index] for index in scores.argmax(axis=1)]) + per_word[word] = { + "n_train": len(train_idx), + "n_test": len(test_idx), + "n_senses": len(senses), + "accuracy": float(accuracy_score(test_y, predictions)), + "balanced_accuracy": float(balanced_accuracy_score(test_y, predictions)), + "macro_f1": float(f1_score(test_y, predictions, average="macro")), + } + return per_word + + +def summarize_coarse(per_word: dict[str, dict[str, float]]) -> dict[str, Any]: + return { + "n_words": len(per_word), + "n_train": int(sum(row["n_train"] for row in per_word.values())), + "n_test": int(sum(row["n_test"] for row in per_word.values())), + **{ + metric: float(np.mean([row[metric] for row in per_word.values()])) + for metric in ("accuracy", "balanced_accuracy", "macro_f1") + }, + } + + +def paired_word_comparison( + srp: dict[str, dict[str, float]], + baseline: dict[str, dict[str, float]], + n_boot: int, + seed: int, +) -> dict[str, Any]: + words = sorted(set(srp) & set(baseline)) + rng = np.random.default_rng(seed) + result: dict[str, Any] = {"n_words": len(words)} + for metric in ("accuracy", "balanced_accuracy", "macro_f1"): + differences = np.asarray([srp[word][metric] - baseline[word][metric] for word in words], dtype=float) + boot = [float(rng.choice(differences, size=len(differences), replace=True).mean()) for _ in range(n_boot)] + result[f"{metric}_difference"] = float(differences.mean()) + result[f"{metric}_difference_ci"] = confidence_interval(boot) + result[f"{metric}_positive_words"] = int((differences > 0).sum()) + return result + + +def bootstrap_coarse_method(per_word: dict[str, dict[str, float]], n_boot: int, seed: int) -> dict[str, Any]: + result = summarize_coarse(per_word) + values = np.asarray([row["macro_f1"] for row in per_word.values()], dtype=float) + rng = np.random.default_rng(seed) + boot = [float(rng.choice(values, size=len(values), replace=True).mean()) for _ in range(n_boot)] + result["macro_f1_ci"] = confidence_interval(boot) + return result + + +def coarse_null_seed_sensitivity( + bundle: dict[str, Any], + gate: np.ndarray, + primary_k: int, + primary_encoding: str, + n_null_seeds: int, + seed: int, + observed: dict[str, Any], +) -> dict[str, Any]: + result: dict[str, Any] = {} + for method in ("shuffled_srp", "random_srp_features"): + values: list[float] = [] + for offset in range(n_null_seeds): + encoded = encode_top_features(bundle, method, primary_k, primary_encoding, seed + offset) + summary = summarize_coarse(coarse_per_word(bundle, encoded, gate)) + values.append(float(summary["macro_f1"])) + result[method] = {"macro_f1": summarize_null_values(values, float(observed["macro_f1"]))} + return result + + +def analyze_coarse( + bundle: dict[str, Any], + ks: list[int], + primary_k: int, + primary_encoding: str, + n_boot: int, + n_null_seeds: int, + seed: int, +) -> dict[str, Any]: + gate, gate_summary = score_gate(bundle) + metrics: dict[str, Any] = { + "dataset": "coarsewsd20", + "question": "Do top contributing features predict held-out word senses?", + "fidelity_gate": gate_summary, + "results": {}, + "primary": { + "encoding": primary_encoding, + "k": primary_k, + "methods": {}, + "comparisons": {}, + }, + } + primary_per_word: dict[str, dict[str, dict[str, float]]] = {} + for encoding in ENCODINGS: + metrics["results"][encoding] = {} + for k in ks: + metrics["results"][encoding][str(k)] = {} + for method in METHODS: + encoded = encode_top_features(bundle, method, k, encoding, seed) + per_word = coarse_per_word(bundle, encoded, gate) + metrics["results"][encoding][str(k)][method] = { + "aggregate": summarize_coarse(per_word), + "per_word": per_word, + } + if encoding == primary_encoding and k == primary_k: + primary_per_word[method] = per_word + for method_index, method in enumerate(METHODS): + metrics["primary"]["methods"][method] = bootstrap_coarse_method( + primary_per_word[method], n_boot, seed + 150 + method_index + ) + for comparison_index, baseline in enumerate(METHODS[1:]): + metrics["primary"]["comparisons"][f"srp_minus_{baseline}"] = paired_word_comparison( + primary_per_word["srp"], + primary_per_word[baseline], + n_boot, + seed + 200 + comparison_index, + ) + metrics["primary"]["null_seed_sensitivity"] = coarse_null_seed_sensitivity( + bundle, + gate, + primary_k, + primary_encoding, + n_null_seeds, + seed, + metrics["primary"]["methods"]["srp"], + ) + return metrics + + +def self_test() -> int: + matrix = np.asarray([[1.0, -4.0, 3.0, 0.1]], dtype=np.float32) + bundle = { + "projection": torch.from_numpy(matrix), + "contribution": torch.from_numpy(matrix), + "beta": torch.ones_like(torch.from_numpy(matrix)), + "feature_ids": torch.arange(4)[None], + "metadata": [{"target": "bank"}], + } + signed = encode_top_features(bundle, "srp", 2, "signed", 0) + assert signed[0, 1] < 0 and signed[0, 2] > 0 + assert signed[0, 0] == 0 and signed[0, 3] == 0 + assert np.isclose(np.linalg.norm(signed[0]), 1.0) + weighted = encode_top_features(bundle, "srp", 2, "weighted", 0) + assert abs(weighted[0, 1]) > abs(weighted[0, 2]) > 0 + validate_feature_alignment(bundle) + print("self-test passed") + return 0 + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--bundle", type=Path, help="representations.pt from run_wsd_feature_alignment.py") + parser.add_argument("--out", type=Path, help="metrics JSON") + parser.add_argument("--ks", default="5,10,20,50") + parser.add_argument("--primary-k", type=int, default=10) + parser.add_argument("--primary-encoding", choices=ENCODINGS, default="weighted") + parser.add_argument("--n-boot", type=int, default=5000) + parser.add_argument("--n-null-seeds", type=int, default=20) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--self-test", action="store_true") + return parser.parse_args() + + +def main() -> int: + args = parse_args() + if args.self_test: + return self_test() + if args.bundle is None or args.out is None: + raise SystemExit("--bundle and --out are required") + ks = [int(value.strip()) for value in args.ks.split(",") if value.strip()] + if args.primary_k not in ks: + raise SystemExit("--primary-k must be included in --ks") + bundle = torch.load(args.bundle, map_location="cpu", weights_only=False) + validate_feature_alignment(bundle) + dataset = str(bundle["run"]["dataset"]) + if dataset == "ambistory": + metrics = analyze_ambistory( + bundle, + ks, + args.primary_k, + args.primary_encoding, + args.n_boot, + args.n_null_seeds, + args.seed, + ) + elif dataset == "coarsewsd20": + metrics = analyze_coarse( + bundle, + ks, + args.primary_k, + args.primary_encoding, + args.n_boot, + args.n_null_seeds, + args.seed, + ) + else: + raise ValueError(f"unsupported dataset {dataset!r}") + metrics["run"] = bundle["run"] + metrics["analysis"] = { + "bundle": str(args.bundle), + "ks": ks, + "primary_k": args.primary_k, + "primary_encoding": args.primary_encoding, + "n_boot": args.n_boot, + "n_null_seeds": args.n_null_seeds, + "seed": args.seed, + "methods": list(METHODS), + "encodings": list(ENCODINGS), + } + write_json(args.out, metrics) + print(f"wrote {args.out}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/analyze/analyze_wsd_sense_groups.py b/scripts/analyze/analyze_wsd_sense_groups.py new file mode 100644 index 0000000..da23919 --- /dev/null +++ b/scripts/analyze/analyze_wsd_sense_groups.py @@ -0,0 +1,521 @@ +#!/usr/bin/env python3 +"""Per-sense feature groups on CoarseWSD-20 readout bundles (``tab:app-sense-alignment``). + +For each ambiguous word, select ``--group-size`` dictionary features per sense +on the train split, then assign each held-out context to the sense whose group +carries the larger summed contribution. Feature scales are standardized by +train-split statistics alone (``--raw-scale`` disables this); no test label +enters the selection. Summary fields, all balanced accuracy averaged over +words, and the table row each one backs: + + word_mean_balanced Selected groups + word_mean_null_balanced Permutation null (train sense labels permuted + ``--n-null`` times, the selection repeated each time) + word_mean_majority_balanced Majority + words_beating_null_p95_balanced Words above null p95 + words_beating_majority_balanced Words above majority + word_mean_full_account_balanced Full account, no selection (nearest centroid on the + full train-standardized contribution vector) + word_mean_hidden_balanced Hidden state (nearest centroid on the full + train-standardized decoded state) + +Per-word entries carry the same quantities plus stratified bootstrap intervals, +unbalanced accuracy, the selected feature positions and the local score gate +(|exact - reconstructed| / (|exact| + 0.5) < 0.5) coverage. + +``--selector mean_diff`` (default) ranks features by the train-mean contribution +gap between a sense and the other senses; ``top1_freq`` ranks by how often a +feature carries the largest contribution on that sense's train instances. +Senses are processed in descending train count and a position taken by an +earlier sense is skipped, so groups are disjoint. + +``--basis`` swaps the dictionary contributions for a direct-W_U-geometry +account (``pca256`` / ``ridge128`` / ``knn128``, sized by ``--pca-rank`` / +``--neighbor-k`` / ``--ridge-lam``) under the identical selection and +prediction protocol, rebuilt from the bundle's stored decoded states with no +model forward pass. These need the unembedding matrix, read from ``--w-u`` (an +extraction payload with ``W_U_orig``, as written by +``scripts/data/extract_model_readout.py``) or, failing that, from the local +Hugging Face cache for ``--model-id``. The paper's table reports ``srp`` only. + +Reads the ``representations.pt`` bundles written by +``scripts/run/run_wsd_feature_alignment.py``; CPU-only, deterministic under +``--seed``. Paper run, once per model bundle (Qwen3.5-2B, Qwen3.5-9B, +DeepSeek-R1-Distill-Llama-8B), writing ``_g{1,2,4,8}.json``: + + python scripts/analyze/analyze_wsd_sense_groups.py \ + --bundle results/wsd_core/coarsewsd20_qwen2b/representations.pt \ + --out results/wsd_sense_groups/coarsewsd20_qwen2b__srp.json \ + --group-sizes 1,2,4,8 --selector mean_diff --n-null 200 --n-boot 2000 --seed 0 + +The table reads the ``_g8`` file; the group-size sweep in the appendix prose +reads all four. +""" + +from __future__ import annotations + +import argparse +import json +import os +from collections import Counter +from pathlib import Path + +import numpy as np +import torch + + +def load_wu_from_artifact(path: Path) -> torch.Tensor: + """Load the unembedding matrix (V, d_model) fp32 from an extraction payload (``W_U_orig`` / ``W_U``).""" + payload = torch.load(path, map_location="cpu", weights_only=True) + w = payload.get("W_U_orig", payload.get("W_U")) + if w is None: + raise KeyError(f"{path}: no 'W_U_orig'/'W_U' key (keys: {list(payload)})") + return w.float() + + +def load_wu(model_id: str) -> torch.Tensor: + """Load the LM-head weight (V, d_model) fp32 from the local Hugging Face cache. + + Reads only the tensor from safetensors; falls back to the tied embedding + when no lm_head.weight exists (e.g. Qwen3.5-2B, tie_word_embeddings). The + cache root follows ``HF_HUB_CACHE``, then ``HF_HOME/hub``, then + ``~/.cache/huggingface/hub``. + """ + from glob import glob + + from safetensors import safe_open + + hub = os.environ.get("HF_HUB_CACHE") + if not hub: + hub = str(Path(os.environ.get("HF_HOME") or (Path.home() / ".cache" / "huggingface")) / "hub") + cache = Path(hub) / f"models--{model_id.replace('/', '--')}" + snaps = sorted(cache.glob("snapshots/*")) + if not snaps: + raise FileNotFoundError(f"no local snapshot for {model_id}") + snap = snaps[-1] + shards = sorted(glob(str(snap / "*.safetensors"))) + for suffix in ("lm_head.weight", "embed_tokens.weight"): + for shard in shards: + with safe_open(shard, framework="pt") as f: + for key in f.keys(): + if key.endswith(suffix): + w = f.get_tensor(key).float() + print(f"[wu] {model_id}: {key} {tuple(w.shape)} from {Path(shard).name}") + return w + raise KeyError(f"no lm_head/embed_tokens weight in {snap}") + + +def geometry_contributions( + basis: str, + W: torch.Tensor, # (V, d_model) fp32 + token_ids: list[int], + hidden_by_word: dict[int, torch.Tensor], # token_id -> (n_word, d_model) fp32 + pca_rank: int, + neighbor_k: int, + ridge_lam: float, + torch_seed: int, +) -> dict[int, dict[str, np.ndarray]]: + """Per-word contribution matrices for a direct-geometry basis. + + Math is line-for-line the vectorized form of readout_baselines.py: + rows are mean-centered and unit-normalized, the target row is rebuilt in + the basis, and contribution_j(h) = row_norm * coeff_j * (h . basis_j). + Returns {token_id: {contribution (n, F), exact (n,), recon (n,)}}. + """ + row_mean = W.mean(dim=0) # (d_model,) + centered = W - row_mean + row_norm_all = centered.norm(dim=1).clamp_min(1e-8) + row_normalized_all = centered / row_norm_all[:, None] + del centered + + v_pca = None + if basis == "pca256": + torch.manual_seed(torch_seed) + q = min(pca_rank + 16, min(row_normalized_all.shape) - 1) + _, _, vt = torch.svd_lowrank(row_normalized_all, q=q, niter=4) + v_pca = vt[:, :pca_rank].T.contiguous() # (pca_rank, d_model) + + out: dict[int, dict[str, np.ndarray]] = {} + for tid in token_ids: + x = row_normalized_all[tid] # (d_model,) + row_norm = row_norm_all[tid] + w_row = W[tid] + h = hidden_by_word[tid] # (n, d_model) + if basis == "pca256": + coeffs = v_pca @ x # (pca_rank,) + contribution = row_norm * (h @ v_pca.T) * coeffs # (n, pca_rank) + else: + sims = row_normalized_all @ x # (V,) + sims[tid] = -float("inf") # exclude self + top = torch.topk(sims, k=neighbor_k) + basis_rows = row_normalized_all[top.indices] # (neighbor_k, d_model) + if basis == "ridge128": + a = basis_rows @ basis_rows.T + ridge_lam * torch.eye(neighbor_k) + coeffs = torch.linalg.solve(a, basis_rows @ x) # (neighbor_k,) + elif basis == "knn128": + w = top.values.clamp_min(0.0) + w = w / w.sum().clamp_min(1e-8) + mean_nbr = w @ basis_rows # (d_model,) + gamma = (x @ mean_nbr) / (mean_nbr @ mean_nbr).clamp_min(1e-8) + coeffs = gamma * w + else: + raise ValueError(f"unknown basis {basis}") + contribution = row_norm * (h @ basis_rows.T) * coeffs # (n, neighbor_k) + exact = h @ w_row # (n,) + recon = h @ row_mean + contribution.sum(dim=1) # (n,) + out[tid] = { + "contribution": contribution.numpy(), + "exact": exact.numpy(), + "recon": recon.numpy(), + } + return out + + +def select_anchors( + train_contribution: np.ndarray, + senses: list[str], + train_y: np.ndarray, + selector: str = "mean_diff", + group_size: int = 1, +) -> dict[str, list[int]]: + """Greedy per-sense feature-group assignment (group_size=1: one feature). + + selector="top1_freq": the feature most often carrying the largest + contribution on that sense's train instances. + selector="mean_diff": the feature with the largest train-mean contribution + gap between this sense and the other senses. + Senses are processed in descending train count; a feature position taken by + an earlier sense is skipped. Ties break by position index for determinism. + """ + order = sorted(senses, key=lambda s: (-int((train_y == s).sum()), s)) + taken: set[int] = set() + anchors: dict[str, list[int]] = {} + top1 = train_contribution.argmax(axis=1) + for sense in order: + mask = train_y == sense + if selector == "top1_freq": + counts = Counter(top1[mask].tolist()) + ranked = [pos for pos, _ in sorted(counts.items(), key=lambda kv: (-kv[1], kv[0]))] + else: + gap = train_contribution[mask].mean(axis=0) - train_contribution[~mask].mean(axis=0) + ranked = np.argsort(-gap).tolist() + group: list[int] = [] + for pos in ranked: + if pos not in taken: + group.append(int(pos)) + taken.add(int(pos)) + if len(group) >= group_size: + break + if not group: # every candidate taken: fall back to first free position + free = [p for p in range(train_contribution.shape[1]) if p not in taken] + group = [free[0] if free else int(top1[mask][0])] + taken.add(group[0]) + anchors[sense] = group + return anchors + + +def word_accuracy( + contribution: np.ndarray, anchors: dict[str, int], senses: list[str], gold: np.ndarray +) -> tuple[float, np.ndarray]: + anchor_matrix = np.stack([contribution[:, anchors[s]].sum(axis=1) for s in senses], axis=1) # (n, n_senses) + predictions = np.array([senses[i] for i in anchor_matrix.argmax(axis=1)]) + return float((predictions == gold).mean()), predictions + + +def balanced_accuracy(predictions: np.ndarray, gold: np.ndarray) -> float: + """Mean per-sense recall over the senses present in gold.""" + recalls = [float((predictions[gold == s] == s).mean()) for s in np.unique(gold)] + return float(np.mean(recalls)) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + parser.add_argument( + "--bundle", type=Path, required=True, help="representations.pt from run_wsd_feature_alignment.py" + ) + parser.add_argument( + "--out", type=Path, required=True, help="summary JSON (with --group-sizes: _g.json each)" + ) + parser.add_argument("--n-boot", type=int, default=2000) + parser.add_argument("--n-null", type=int, default=200) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--selector", choices=["top1_freq", "mean_diff"], default="mean_diff") + parser.add_argument("--group-size", type=int, default=1) + parser.add_argument( + "--group-sizes", + default=None, + help="comma list (e.g. 1,2,4,8): sweep group sizes in one run, writing _g.json each; " + "shares the basis/W_U preparation across sizes", + ) + parser.add_argument( + "--basis", + choices=["srp", "pca256", "ridge128", "knn128"], + default="srp", + help="srp: bundle contributions; others: direct-W_U geometry controls", + ) + parser.add_argument("--model-id", default=None, help="override bundle run.model_id for --basis controls") + parser.add_argument( + "--w-u", + type=Path, + default=None, + help="extraction payload (.pt with W_U_orig) for --basis controls; default: local HF cache of --model-id", + ) + parser.add_argument("--pca-rank", type=int, default=256) + parser.add_argument("--neighbor-k", type=int, default=128) + parser.add_argument("--ridge-lam", type=float, default=1e-3) + parser.add_argument( + "--raw-scale", + action="store_true", + help="skip train-statistic standardization of contributions (default: standardize)", + ) + args = parser.parse_args() + + bundle = torch.load(args.bundle, map_location="cpu", weights_only=False) + metadata = bundle["metadata"] + contribution = bundle["contribution"].float().numpy() # (n_items, k_support) + exact = bundle["exact_logit"].numpy() + recon = bundle["reconstructed_logit"].numpy() + + if args.basis != "srp": + if args.w_u is not None: + wu = load_wu_from_artifact(args.w_u) + else: + wu = load_wu(args.model_id or bundle["run"]["model_id"]) + hidden = bundle["hidden"].float() # (n_items, d_model) + word_tid: dict[str, int] = {} + word_idx: dict[str, list[int]] = {} + for i, row in enumerate(metadata): + word_tid[row["target"]] = int(row["token_id"]) + word_idx.setdefault(row["target"], []).append(i) + hidden_by_word = {word_tid[w]: hidden[np.array(idx)] for w, idx in word_idx.items()} + geo = geometry_contributions( + args.basis, + wu, + [word_tid[w] for w in sorted(word_idx)], + hidden_by_word, + args.pca_rank, + args.neighbor_k, + args.ridge_lam, + args.seed, + ) + n_feat = args.pca_rank if args.basis == "pca256" else args.neighbor_k + contribution = np.zeros((len(metadata), n_feat), dtype=np.float32) + geo_exact = np.zeros(len(metadata), dtype=np.float32) + recon = np.zeros(len(metadata), dtype=np.float32) + for w, idx in word_idx.items(): + g = geo[word_tid[w]] + contribution[np.array(idx)] = g["contribution"] + geo_exact[np.array(idx)] = g["exact"] + recon[np.array(idx)] = g["recon"] + # W_U extraction sanity: fp32 rebuild of h.w_t vs the runtime logit + exact_gap = float(np.max(np.abs(geo_exact - exact))) + print(f"[wu] max |exact(rebuilt) - exact(bundle)| = {exact_gap:.4f}") + exact = geo_exact + + rho = np.abs(exact - recon) / (np.abs(exact) + 0.5) # floored rho, paper convention + gate = rho < 0.5 + + group_sizes = [int(x) for x in args.group_sizes.split(",")] if args.group_sizes else [args.group_size] + + def ncm_balanced(train_x: np.ndarray, test_x: np.ndarray, train_y: np.ndarray, test_y: np.ndarray) -> float: + """Nearest class-centroid on train-standardized features: the + no-selection reference (how much sense information the representation + carries when the whole account is used).""" + mu = train_x.mean(axis=0) + sd = train_x.std(axis=0) + 1e-6 + ztr = (train_x - mu) / sd + zte = (test_x - mu) / sd + senses = sorted(set(train_y)) + cents = np.stack([ztr[train_y == s].mean(axis=0) for s in senses]) + dist = ((zte[:, None, :] - cents[None]) ** 2).sum(axis=-1) # (n_test, n_senses) + pred = np.array([senses[i] for i in dist.argmin(axis=1)]) + return balanced_accuracy(pred, test_y) + + hidden_all = bundle["hidden"].float().numpy() # (n_items, d_model) + words = sorted({row["target"] for row in metadata}) + word_cache: dict[str, dict] = {} + for word in words: + train_idx = np.array([i for i, r in enumerate(metadata) if r["target"] == word and r["split"] == "train"]) + test_idx = np.array([i for i, r in enumerate(metadata) if r["target"] == word and r["split"] == "test"]) + if len(train_idx) == 0 or len(test_idx) == 0: + continue + train_y = np.array([metadata[i]["sense"] for i in train_idx]) + test_y = np.array([metadata[i]["sense"] for i in test_idx]) + senses = sorted(set(train_y)) + if len(senses) < 2 or not set(test_y).issubset(set(senses)): + continue + + train_c = contribution[train_idx] + test_c = contribution[test_idx] + if not args.raw_scale: + # Features have incommensurate contribution scales; standardize by + # train-split statistics (identically for every basis, no leakage) + # so the cross-feature argmax compares deviations, not magnitudes. + mu = train_c.mean(axis=0) + sd = train_c.std(axis=0) + 1e-6 + train_c = (train_c - mu) / sd + test_c = (test_c - mu) / sd + word_cache[word] = { + "train_idx": train_idx, + "test_idx": test_idx, + "train_y": train_y, + "test_y": test_y, + "senses": senses, + "train_c": train_c, + "test_c": test_c, + # basis-independent references, computed once per word + "full_account_balanced": ncm_balanced(contribution[train_idx], contribution[test_idx], train_y, test_y), + "hidden_balanced": ncm_balanced(hidden_all[train_idx], hidden_all[test_idx], train_y, test_y), + } + + for group_size in group_sizes: + run_group(args, group_size, word_cache, gate) + return 0 + + +def run_group(args, group_size: int, word_cache: dict, gate: np.ndarray) -> None: + rng = np.random.default_rng(args.seed) + per_word: dict[str, dict] = {} + pooled_correct: list[int] = [] + pooled_correct_gated: list[int] = [] + + for word, wc in word_cache.items(): + train_idx, test_idx = wc["train_idx"], wc["test_idx"] + train_y, test_y = wc["train_y"], wc["test_y"] + senses = wc["senses"] + train_c, test_c = wc["train_c"], wc["test_c"] + + anchors = select_anchors(train_c, senses, train_y, args.selector, group_size) + acc, predictions = word_accuracy(test_c, anchors, senses, test_y) + bal_acc = balanced_accuracy(predictions, test_y) + majority = max(senses, key=lambda s: int((train_y == s).sum())) + majority_acc = float((test_y == majority).mean()) + majority_bal = balanced_accuracy(np.full_like(test_y, majority), test_y) + + # bootstrap CI over test instances (stratified by sense for balanced acc) + boot = np.empty(args.n_boot) + boot_bal = np.empty(args.n_boot) + correct = (predictions == test_y).astype(float) + sense_pools = {s: np.flatnonzero(test_y == s) for s in np.unique(test_y)} + for b in range(args.n_boot): + sample = rng.integers(0, len(correct), len(correct)) + boot[b] = correct[sample].mean() + strat = np.concatenate([pool[rng.integers(0, len(pool), len(pool))] for pool in sense_pools.values()]) + boot_bal[b] = balanced_accuracy(predictions[strat], test_y[strat]) + ci = [float(np.percentile(boot, 2.5)), float(np.percentile(boot, 97.5))] + bal_ci = [float(np.percentile(boot_bal, 2.5)), float(np.percentile(boot_bal, 97.5))] + + # label-shuffle null: re-select anchors under permuted train labels + null_accs = np.empty(args.n_null) + null_bals = np.empty(args.n_null) + for n in range(args.n_null): + shuffled_y = rng.permutation(train_y) + null_anchors = select_anchors(train_c, senses, shuffled_y, args.selector, group_size) + null_accs[n], null_pred = word_accuracy(test_c, null_anchors, senses, test_y) + null_bals[n] = balanced_accuracy(null_pred, test_y) + + test_gate = gate[test_idx] + gated_acc = float(correct[test_gate].mean()) if test_gate.any() else float("nan") + gated_bal = balanced_accuracy(predictions[test_gate], test_y[test_gate]) if test_gate.any() else float("nan") + pooled_correct.extend(correct.astype(int).tolist()) + pooled_correct_gated.extend(correct[test_gate].astype(int).tolist()) + + per_word[word] = { + "n_train": int(len(train_idx)), + "n_test": int(len(test_idx)), + "n_senses": len(senses), + "anchors": {s: [int(p) for p in anchors[s]] for s in senses}, + "accuracy": acc, + "accuracy_ci": ci, + "balanced_accuracy": bal_acc, + "balanced_accuracy_ci": bal_ci, + "majority_accuracy": majority_acc, + "majority_balanced_accuracy": majority_bal, + "chance_balanced_accuracy": 1.0 / len(senses), + "null_accuracy_mean": float(null_accs.mean()), + "null_accuracy_p95": float(np.percentile(null_accs, 95)), + "null_balanced_mean": float(null_bals.mean()), + "null_balanced_p95": float(np.percentile(null_bals, 95)), + "beats_null_p95": bool(acc > np.percentile(null_accs, 95)), + "beats_null_p95_balanced": bool(bal_acc > np.percentile(null_bals, 95)), + "gate_fraction": float(test_gate.mean()), + "gated_accuracy": gated_acc, + "gated_balanced_accuracy": gated_bal, + "full_account_balanced": wc["full_account_balanced"], + "hidden_balanced": wc["hidden_balanced"], + } + + names = sorted(per_word) + word_acc = np.array([per_word[w]["accuracy"] for w in names]) + word_bal = np.array([per_word[w]["balanced_accuracy"] for w in names]) + word_majority = np.array([per_word[w]["majority_accuracy"] for w in names]) + word_majority_bal = np.array([per_word[w]["majority_balanced_accuracy"] for w in names]) + word_chance_bal = np.array([per_word[w]["chance_balanced_accuracy"] for w in names]) + word_null = np.array([per_word[w]["null_accuracy_mean"] for w in names]) + word_null_bal = np.array([per_word[w]["null_balanced_mean"] for w in names]) + boot_mean = np.empty(args.n_boot) + boot_mean_bal = np.empty(args.n_boot) + for b in range(args.n_boot): + sample = rng.integers(0, len(word_acc), len(word_acc)) + boot_mean[b] = word_acc[sample].mean() + boot_mean_bal[b] = word_bal[sample].mean() + + summary = { + "bundle": str(args.bundle), + "basis": args.basis, + "selector": args.selector, + "group_size": group_size, + "seed": args.seed, + "n_words": len(names), + "word_mean_accuracy": float(word_acc.mean()), + "word_mean_accuracy_ci": [ + float(np.percentile(boot_mean, 2.5)), + float(np.percentile(boot_mean, 97.5)), + ], + "standardized": not args.raw_scale, + "word_mean_balanced": float(word_bal.mean()), + "word_mean_balanced_ci": [ + float(np.percentile(boot_mean_bal, 2.5)), + float(np.percentile(boot_mean_bal, 97.5)), + ], + "word_mean_majority": float(word_majority.mean()), + "word_mean_majority_balanced": float(word_majority_bal.mean()), + "word_mean_chance_balanced": float(word_chance_bal.mean()), + "word_mean_null": float(word_null.mean()), + "word_mean_null_balanced": float(word_null_bal.mean()), + "words_beating_null_p95": int(sum(per_word[w]["beats_null_p95"] for w in names)), + "words_beating_null_p95_balanced": int(sum(per_word[w]["beats_null_p95_balanced"] for w in names)), + "words_beating_majority_balanced": int( + sum(per_word[w]["balanced_accuracy"] > per_word[w]["majority_balanced_accuracy"] for w in names) + ), + "word_mean_full_account_balanced": float(np.mean([per_word[w]["full_account_balanced"] for w in names])), + "word_mean_hidden_balanced": float(np.mean([per_word[w]["hidden_balanced"] for w in names])), + "pooled_accuracy": float(np.mean(pooled_correct)), + "pooled_gated_accuracy": float(np.mean(pooled_correct_gated)), + "gate_fraction_overall": float(len(pooled_correct_gated) / max(len(pooled_correct), 1)), + "per_word": per_word, + } + + out = args.out if args.group_sizes is None else args.out.with_name(f"{args.out.stem}_g{group_size}.json") + out.parent.mkdir(parents=True, exist_ok=True) + tmp = out.with_suffix(".tmp") + tmp.write_text(json.dumps(summary, indent=2)) + os.replace(tmp, out) + print( + f"{args.bundle.parent.name} [{args.basis} g={group_size}]: words={summary['n_words']} " + f"bal acc={summary['word_mean_balanced']:.3f} " + f"CI[{summary['word_mean_balanced_ci'][0]:.3f},{summary['word_mean_balanced_ci'][1]:.3f}] " + f"(chance={summary['word_mean_chance_balanced']:.3f} " + f"null={summary['word_mean_null_balanced']:.3f} " + f"beat_null={summary['words_beating_null_p95_balanced']}/{summary['n_words']}) | " + f"raw acc={summary['word_mean_accuracy']:.3f} " + f"majority={summary['word_mean_majority']:.3f} " + f"gated acc={summary['pooled_gated_accuracy']:.3f} " + f"(gate {summary['gate_fraction_overall']:.2f}) | " + f"full acct={summary['word_mean_full_account_balanced']:.3f} " + f"hidden={summary['word_mean_hidden_balanced']:.3f}" + ) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/analyze/cross_lens_antonym_layers.py b/scripts/analyze/cross_lens_antonym_layers.py new file mode 100644 index 0000000..76eb705 --- /dev/null +++ b/scripts/analyze/cross_lens_antonym_layers.py @@ -0,0 +1,310 @@ +#!/usr/bin/env python3 +"""The published antonym prompt, layer by layer, under one fitted lens (``tab:app-cross-lens-antonym-layers``). + +``run`` reads the prompt ``"小"的反义词是"`` (model answer 大) through one fitted +lens on Qwen3.5-9B and records, per fitted layer and at the final position: the +top-10 tokens with softmax shares under (a) the lens transport and (b) the +identity transport (logit lens, same final norm), plus the Sparse Readout Prism +decomposition of the English and Chinese answer forms (' big', 'big', '大', with +the ' small' / '小' contrasts) on the transported state, with the top signed +feature contributions itemised. It is run once per lens. ``table`` then assembles +the layer x lens table of top-1 token and softmax share (and the dominant feature +of each probed form under each lens) from the per-lens JSON files. + +Decomposition basis: the k=128 seed-0 readout dictionary for Qwen3.5-9B (8x, +D=32768), Hugging Face ``hematteo/sparse-readout-prism`` file +``qwen3.5-9b/k128_8x/checkpoint.pt``. ``jlens`` is imported lazily +(``uv sync --extra lens``). + +Paper runs:: + + cross_lens_antonym_layers.py run --model-id Qwen/Qwen3.5-9B \ + --lens /qwen35_9b_jlens_en_seed0_n100.pt --sae --k 128 \ + --out /antonym_layers_en.json + cross_lens_antonym_layers.py run --model-id Qwen/Qwen3.5-9B \ + --lens /qwen35_9b_jlens_zh_seed0_n100.pt --sae --k 128 \ + --out /antonym_layers_zh.json + cross_lens_antonym_layers.py table --dump EN=/antonym_layers_en.json \ + --dump ZH=/antonym_layers_zh.json --layers 24,26,29,final \ + --out-csv /antonym_layers_table.csv --out-features-csv /antonym_layers_features.csv +""" + +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path + +import torch +import torch.nn.functional as F + +from sparse_readout_prism.utils import write_csv + + +def log(msg): + print(f"[antonym] {msg}", flush=True) + + +def _import_jlens(): + try: + import jlens + except ImportError as e: + raise ImportError( + "jlens (the Jacobian-lens reference implementation, Apache-2.0, " + "github.com/anthropics/jacobian-lens) is not installed; install it with `uv sync --extra lens`" + ) from e + return jlens + + +def load_sae(checkpoint): + """Raw TopK dictionary tensors from a training checkpoint (decoder rows unit-normalised).""" + ckpt = torch.load(checkpoint, map_location="cpu", weights_only=True) + state = ckpt["model_state_dict"] + decoder = state["decoder"].float().contiguous() # (d_features, d_model) + encoder_w = state["encoder.weight"].float().contiguous() # (d_features, d_model) + encoder_b = state["encoder.bias"].float().contiguous() # (d_features,) + decoder = decoder / decoder.norm(dim=1, keepdim=True).clamp_min(1e-8) + return decoder, encoder_w, encoder_b + + +def encode_topk(x, encoder_w, encoder_b, k): + acts = F.relu(x @ encoder_w.T + encoder_b) + values, indices = torch.topk(acts, k=min(k, acts.shape[-1]), dim=-1) + code = torch.zeros_like(acts) + code.scatter_(dim=-1, index=indices, src=values) + return code + + +@torch.no_grad() +def cmd_run(args) -> None: + import transformers + + jlens = _import_jlens() + + hf = transformers.AutoModelForCausalLM.from_pretrained(args.model_id, dtype=torch.bfloat16).cuda().eval() + tok = transformers.AutoTokenizer.from_pretrained(args.model_id) + model = jlens.from_hf(hf, tok) + lens = jlens.JacobianLens.load(args.lens) + layers = sorted(lens.jacobians) + J = {l: lens.jacobians[l].float().cuda() for l in layers} + W = hf.get_output_embeddings().weight.detach().float().cuda() # (vocab, d_model) + row_mean = W.mean(dim=0) + final_norm = hf.model.norm + decoder, encoder_w, encoder_b = load_sae(Path(args.sae)) + decoder, encoder_w, encoder_b = decoder.cuda(), encoder_w.cuda(), encoder_b.cuda() + + captured = {} + + def make_hook(idx): + def hook(module, inputs, output): + captured[("out", idx)] = (output[0] if isinstance(output, tuple) else output).detach() + + return hook + + handles = [model.layers[l].register_forward_hook(make_hook(l)) for l in layers] + + def norm_hook(module, inputs, output): + if output.ndim == 3: + captured["final_norm_out"] = output.detach() + + handles.append(final_norm.register_forward_hook(norm_hook)) + + per_layer, final_logits, _ = lens.apply(model, args.prompt, positions=[-1]) + h_ref = captured["final_norm_out"][0, -1].float().clone() + final_logits = final_logits.float().squeeze().cuda() + greedy = tok.decode([int(final_logits.argmax())]) + log(f"prompt={args.prompt!r} greedy={greedy!r}") + + def top10(logits): + pr = torch.softmax(logits, -1) + v, i = torch.topk(pr, 10) + return [[tok.decode([t]), round(float(x) * 100, 1)] for t, x in zip(i.tolist(), v.tolist())] + + def probe_ids(): + out = {} + for text in (" big", "big", "大", " small", "小"): + ids = tok(text, add_special_tokens=False)["input_ids"] + if len(ids) == 1: + out[text] = ids[0] + return out + + probes = probe_ids() + log(f"single-token probes: {list(probes)}") + + def decompose(h_state, token_id): + W_row = W[token_id] + centered = W_row - row_mean + norm = centered.norm().clamp_min(1e-8) + code = encode_topk((centered / norm)[None, :], encoder_w, encoder_b, args.k)[0] + contributions = norm * code * (h_state @ decoder.T) + active = torch.nonzero(contributions != 0).flatten() + top = active[contributions[active].abs().argsort(descending=True)[:16]] + return { + "original_logit": float(h_state @ W_row), + "base": float(h_state @ row_mean), + "feature_sum": float(contributions.sum()), + "residual": float(h_state @ W_row) - float(h_state @ row_mean) - float(contributions.sum()), + "top_features": [{"id": int(f), "contribution": float(contributions[f])} for f in top], + } + + rows, seen = {}, set() + for l in layers: + h = captured[("out", l)][0, -1].float() + h_j = final_norm((J[l] @ h).to(torch.bfloat16)).float() + h_id = final_norm(h.to(torch.bfloat16)).float() # logit lens: identity transport + entry = { + "jlens_top10": top10(h_j @ W.T), + "logitlens_top10": top10(h_id @ W.T), + "targets": {t: decompose(h_j, i) for t, i in probes.items()}, + } + rows[str(l)] = entry + for t in entry["targets"].values(): + seen.update(f["id"] for f in t["top_features"]) + rows["final"] = { + "jlens_top10": top10(final_logits), + "logitlens_top10": top10(final_logits), + "targets": {t: decompose(h_ref, i) for t, i in probes.items()}, + } + for t in rows["final"]["targets"].values(): + seen.update(f["id"] for f in t["top_features"]) + for h in handles: + h.remove() + + fids = torch.tensor(sorted(seen), dtype=torch.long) + enc, bias = encoder_w[fids], encoder_b[fids] + best_s = torch.full((len(fids), 10), -float("inf"), device=W.device) + best_i = torch.zeros((len(fids), 10), dtype=torch.long, device=W.device) + for s in range(0, W.shape[0], 8192): + rws = W[s : s + 8192] + c = rws - row_mean + x = c / c.norm(dim=1, keepdim=True).clamp_min(1e-8) + sc = F.relu(x @ enc.T + bias).T + m = torch.cat([best_s, sc], 1) + mi = torch.cat([best_i, torch.arange(s, s + rws.shape[0], device=W.device).expand(len(fids), -1)], 1) + best_s, keep = torch.topk(m, k=10, dim=1) + best_i = torch.gather(mi, 1, keep) + labels = { + int(f): [tok.decode([t]) for t, sc in zip(best_i[j].tolist(), best_s[j].tolist()) if sc > 0] + for j, f in enumerate(fids.tolist()) + } + + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + payload = { + "model": args.model_id, + "lens": args.lens, + "sae": args.sae, + "prompt": args.prompt, + "greedy": greedy, + "layers": layers, + "rows": rows, + "feature_top_tokens": labels, + } + tmp = out.with_suffix(".json.tmp") + tmp.write_text(json.dumps(payload, ensure_ascii=False)) + os.replace(tmp, out) + + for L in [str(l) for l in layers] + ["final"]: + jl = " · ".join(f"{t}({p}%)" for t, p in rows[L]["jlens_top10"][:3]) + ll = " · ".join(f"{t}({p}%)" for t, p in rows[L]["logitlens_top10"][:3]) + log(f"L{L:>5} lens: {jl} | logit-lens: {ll}") + log(f"wrote -> {out}") + + +def parse_dump_args(specs: list[str]) -> list[tuple[str, Path]]: + """Parse repeated ``LABEL=path`` arguments, preserving order.""" + out = [] + for spec in specs: + label, sep, path = spec.partition("=") + if not sep or not label or not path: + raise ValueError(f"--dump expects LABEL=path, got {spec!r}") + out.append((label, Path(path))) + return out + + +def antonym_layer_table(dumps: dict[str, dict], layers: list[str]) -> tuple[list[dict], list[dict]]: + """Layer x lens rows of top-1 token / softmax share, and per-probe dominant features. + + ``dumps`` maps a lens label to the JSON payload written by ``run``. Layers are + the string keys of ``rows`` (fitted layer indices and ``"final"``). + """ + table, features = [], [] + for L in layers: + for label, d in dumps.items(): + r = d["rows"][L] + lens_top, logit_top = r["jlens_top10"][0], r["logitlens_top10"][0] + table.append( + { + "layer": L, + "lens": label, + "top1_token": lens_top[0], + "top1_softmax_pct": lens_top[1], + "logit_lens_top1_token": logit_top[0], + "logit_lens_top1_softmax_pct": logit_top[1], + "top3_tokens": " | ".join(t for t, _ in r["jlens_top10"][:3]), + } + ) + for target, dec in r["targets"].items(): + top = dec["top_features"][0] if dec["top_features"] else {"id": None, "contribution": None} + features.append( + { + "layer": L, + "lens": label, + "target": target, + "original_logit": dec["original_logit"], + "feature_sum": dec["feature_sum"], + "residual": dec["residual"], + "dominant_feature": top["id"], + "dominant_contribution": top["contribution"], + } + ) + return table, features + + +def cmd_table(args) -> None: + dumps = {label: json.loads(path.read_text()) for label, path in parse_dump_args(args.dump)} + first = next(iter(dumps.values())) + layers = [s.strip() for s in args.layers.split(",") if s.strip()] if args.layers else list(first["rows"]) + table, features = antonym_layer_table(dumps, layers) + labels = list(dumps) + print(f"prompt={first['prompt']!r}") + print("layer " + "".join(f"{lb:>28s}" for lb in labels)) + for L in layers: + cells = [t for t in table if t["layer"] == L] + print(f"{L:>5s} " + "".join(f"{c['top1_token']!r} {c['top1_softmax_pct']}%".rjust(28) for c in cells)) + if args.out_csv: + write_csv(args.out_csv, table) + print(f"wrote {args.out_csv}") + if args.out_features_csv: + write_csv(args.out_features_csv, features) + print(f"wrote {args.out_features_csv}") + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + sub = p.add_subparsers(dest="cmd", required=True) + + sp = sub.add_parser("run", help="read the prompt through one fitted lens (GPU)") + sp.add_argument("--model-id", default="Qwen/Qwen3.5-9B") + sp.add_argument("--lens", required=True, help="fitted lens .pt (JacobianLens container)") + sp.add_argument("--sae", required=True, help="readout SAE checkpoint.pt (paper: Qwen3.5-9B k=128 seed 0, 8x)") + sp.add_argument("--k", type=int, default=128) + sp.add_argument("--prompt", default='"小"的反义词是"') + sp.add_argument("--out", required=True, help="output JSON") + sp.set_defaults(fn=cmd_run) + + sp = sub.add_parser("table", help="assemble the layer x lens table from per-lens JSON files") + sp.add_argument("--dump", action="append", required=True, help="LABEL=path, repeatable (e.g. EN=..., ZH=...)") + sp.add_argument("--layers", default="24,26,29,final", help="comma-separated layer keys (default: paper rows)") + sp.add_argument("--out-csv", default=None, help="top-1 token and softmax share per layer x lens") + sp.add_argument("--out-features-csv", default=None, help="dominant feature per layer x lens x probed form") + sp.set_defaults(fn=cmd_table) + + args = p.parse_args(argv) + args.fn(args) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/analyze/cross_lens_three_lens_prompt.py b/scripts/analyze/cross_lens_three_lens_prompt.py new file mode 100644 index 0000000..d38040d --- /dev/null +++ b/scripts/analyze/cross_lens_three_lens_prompt.py @@ -0,0 +1,184 @@ +#!/usr/bin/env python3 +"""One prompt's frozen states read by several fitted lenses (``tab:app-cross-lens-de-antonym-layers``). + +Offline over readout dumps from ``scripts/run/run_cross_lens_readouts.py``, one per +lens, all produced from the same one-prompt bank +(``data/cross_lens/cross_lens_three_lens_prompts.json``). For the chosen prompt id +and layers the script tabulates, per lens: the top-1 token and its transported +logit (the table), the full top-5 reading, the lens logit, rank and decomposed +(fp32) logit of every probed target (e.g. ' groß', ' large', ' big', '大'), and +the dominant readout feature of each target's decomposition and of the lens's +own top-1 token. + +Paper run (the German antonym prompt ``antonym_de_01``, answer ``groß``, read by the +English-, Chinese- and German-fitted 100-prompt Jacobian lenses; dumps produced +with ``--n-positions 1 --decompose-top1 --top-feats 10``):: + + cross_lens_three_lens_prompt.py --dump EN=/three_lens__jlens_en.json \ + --dump ZH=/three_lens__jlens_zh.json --dump DE=/three_lens__jlens_de.json \ + --prompt-id antonym_de_01 --layers 24,26,29 \ + --out-csv /three_lens_top1.csv --out-targets-csv /three_lens_targets.csv \ + --out-features-csv /three_lens_features.csv + +Adding ``--dump LOGIT=/three_lens__identity.json`` (the ``--lens identity`` +run) includes the lens-free reading as a further column. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +from sparse_readout_prism.utils import write_csv + +POS = "-1" + + +def parse_dump_args(specs: list[str]) -> list[tuple[str, Path]]: + """Parse repeated ``LABEL=path`` arguments, preserving order.""" + out = [] + for spec in specs: + label, sep, path = spec.partition("=") + if not sep or not label or not path: + raise ValueError(f"--dump expects LABEL=path, got {spec!r}") + out.append((label, Path(path))) + return out + + +def select_record(dump: dict, prompt_id: str | None) -> dict: + records = dump["records"] + if prompt_id is None: + if len(records) != 1: + raise ValueError(f"dump has {len(records)} records; pass --prompt-id") + return records[0] + for r in records: + if r["id"] == prompt_id: + return r + raise KeyError(f"prompt id {prompt_id!r} not in dump") + + +def three_lens_tables( + records: dict[str, dict], layers: list[str], targets: list[str] | None +) -> tuple[list[dict], list[dict], list[dict]]: + """Per layer x lens: top-1 rows, per-target rows, and dominant-feature rows. + + ``records`` maps a lens label to the prompt's record in that lens's dump. + """ + top1_rows, target_rows, feature_rows = [], [], [] + for L in layers: + for label, rec in records.items(): + e = rec["layers"][L][POS] + tok1, logit1 = e["top5"][0] + top1_rows.append( + { + "layer": L, + "lens": label, + "top1_token": tok1, + "top1_lens_logit": logit1, + "top5_tokens": " | ".join(t for t, _ in e["top5"]), + } + ) + if e.get("top1_decomp"): + d = e["top1_decomp"] + f0 = d["top_features"][0] if d["top_features"] else {"id": None, "contribution": None} + feature_rows.append( + { + "layer": L, + "lens": label, + "token": tok1, + "role": "top1", + "original_logit": d["original_logit"], + "feature_sum": d["feature_sum"], + "residual": d["residual"], + "dominant_feature": f0["id"], + "dominant_contribution": f0["contribution"], + } + ) + probe = targets if targets is not None else list(e["targets"]) + for t in probe: + d = e["targets"].get(t) + if d is None: + continue + f0 = d["top_features"][0] if d["top_features"] else {"id": None, "contribution": None} + target_rows.append( + { + "layer": L, + "lens": label, + "target": t, + "lens_logit": d["lens_logit"], + "lens_rank": d["lens_rank"], + "original_logit": d["original_logit"], + } + ) + feature_rows.append( + { + "layer": L, + "lens": label, + "token": t, + "role": "target", + "original_logit": d["original_logit"], + "feature_sum": d["feature_sum"], + "residual": d["residual"], + "dominant_feature": f0["id"], + "dominant_contribution": f0["contribution"], + } + ) + return top1_rows, target_rows, feature_rows + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument( + "--dump", action="append", required=True, help="LABEL=path, repeatable (e.g. EN=..., ZH=..., DE=...)" + ) + p.add_argument("--prompt-id", default=None, help="record id to tabulate (default: the only record)") + p.add_argument("--layers", default="24,26,29", help="comma-separated layer keys (default: paper rows)") + p.add_argument("--targets", default=None, help="comma-separated probed targets (default: all in the dump)") + p.add_argument("--out-csv", default=None, help="top-1 token and transported logit per layer x lens") + p.add_argument("--out-targets-csv", default=None, help="lens logit and rank per layer x lens x target") + p.add_argument("--out-features-csv", default=None, help="dominant feature per layer x lens x token") + args = p.parse_args(argv) + + dumps = parse_dump_args(args.dump) + records = {label: select_record(json.loads(path.read_text()), args.prompt_id) for label, path in dumps} + layers = [s.strip() for s in args.layers.split(",") if s.strip()] + targets = [s for s in args.targets.split(",")] if args.targets else None + top1_rows, target_rows, feature_rows = three_lens_tables(records, layers, targets) + + first = next(iter(records.values())) + labels = list(records) + print( + f"prompt {first['id']} ({first['group']}), answer {first['answer']!r}, continuation {first['continuation']!r}" + ) + print("layer " + "".join(f"{lb:>24s}" for lb in labels)) + for L in layers: + cells = [r for r in top1_rows if r["layer"] == L] + print(f"{L:>5s} " + "".join(f"{c['top1_token']!r} {c['top1_lens_logit']:.1f}".rjust(24) for c in cells)) + probe = sorted({r["target"] for r in target_rows}, key=lambda t: [r["target"] for r in target_rows].index(t)) + for L in layers: + print(f"-- layer {L}: lens logit (rank) per target") + for t in probe: + cells = [r for r in target_rows if r["layer"] == L and r["target"] == t] + print(f" {t!r:14s}" + "".join(f"{c['lens_logit']:.2f} (r{c['lens_rank']})".rjust(24) for c in cells)) + for r in feature_rows: + if r["layer"] == L and r["role"] == "top1": + print( + f" {r['lens']:>6s} top1={r['token']!r:12s} f{r['dominant_feature']} " + f"{r['dominant_contribution']:+.2f} (sum {r['feature_sum']:.1f}, resid {r['residual']:.2f})" + ) + + if args.out_csv: + write_csv(args.out_csv, top1_rows) + print(f"wrote {args.out_csv}") + if args.out_targets_csv: + write_csv(args.out_targets_csv, target_rows) + print(f"wrote {args.out_targets_csv}") + if args.out_features_csv: + write_csv(args.out_features_csv, feature_rows) + print(f"wrote {args.out_features_csv}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/analyze/nearest_rows_baseline.py b/scripts/analyze/nearest_rows_baseline.py new file mode 100644 index 0000000..898f834 --- /dev/null +++ b/scripts/analyze/nearest_rows_baseline.py @@ -0,0 +1,269 @@ +#!/usr/bin/env python3 +"""Nearest unembedding rows of selected contrast directions (``tab:app-nearest-rows``). + +The zero-fitting reading of a selected readout direction q = w_a - w_b ranks +the vocabulary rows of W_U by |cosine(w_t, q)| and lists the top tokens, with +the two contrast tokens themselves excluded. Rows are ranked both row-centered +(w_t - mu, the SRP preprocessing; q is unchanged by centering) and raw, and the +signed cosine is reported so the supported side of the contrast is visible +(positive supports the first token). + +Paper run: Qwen3.5-2B at HF revision 15852e8c16360a2fea060d615a32b45270f8a8fc +(tied embeddings, lm_head.weight == embed_tokens.weight, (248320, 2048)), the +contrasts " bug" - " insect" and " bug" - " error" as leading-space single +tokens, top 12 rows by |cosine|: + + uv run python scripts/analyze/nearest_rows_baseline.py \\ + --model-id Qwen/Qwen3.5-2B --revision 15852e8c16360a2fea060d615a32b45270f8a8fc \\ + --contrast bug,insect --contrast bug,error --top-n 12 \\ + --out-dir results/nearest_rows_qwen2b + +The centering mean is the full-vocabulary mean of W_U +(``data.resolve_row_mean`` with no checkpoint), which is what the table used; +``--checkpoint`` centers with a dictionary's stored training mean instead. +``--w-u`` reads W_U from an extraction artifact (``W_U_orig`` written by +``scripts/data/extract_model_readout.py``) instead of loading the model; the +tokenizer is still loaded from ``--model-id``. No random numbers are used. + +Outputs: ``nearest_rows.json`` (token report, per-contrast centered and raw +rankings, provenance), ``nearest_rows.csv`` (long form: contrast, variant, +rank, token_id, token, cosine) and ``nearest_rows_table.csv`` (the table +layout: per rank the centered and raw token with cosines rounded to two +decimals and glosses for the non-Latin tokens that occur in the paper's lists). +""" + +from __future__ import annotations + +import argparse +import csv +import json +from pathlib import Path + +import torch + +from sparse_readout_prism.data import resolve_row_mean +from sparse_readout_prism.research.qwen_readout import load_qwen_model +from sparse_readout_prism.research.run_io import run_provenance +from sparse_readout_prism.utils import find_lm_head_with_path + +# Glosses for the non-Latin tokens that enter the paper's top-12 lists. +GLOSS = { + "昆虫": "insect", + "故障": "malfunction", + "误差": "(measurement) error", + "错误": "error", + "漏洞": "vulnerability", + " насеком": "insect- (ru)", + " насекомых": "insects (ru)", + " ошибки": "errors (ru)", +} + + +def single_token_id(tokenizer, text: str) -> int | None: + ids = tokenizer.encode(text, add_special_tokens=False) + return int(ids[0]) if len(ids) == 1 else None + + +def token_report(tokenizer, words: list[str]) -> dict[str, dict]: + """Single-token check for each contrast word with and without a leading space.""" + report: dict[str, dict] = {} + for w in words: + with_sp = single_token_id(tokenizer, " " + w) + without_sp = single_token_id(tokenizer, w) + report[w] = { + "with_leading_space": {"id": with_sp, "single_token": with_sp is not None}, + "without_leading_space": {"id": without_sp, "single_token": without_sp is not None}, + } + return report + + +def resolve_contrast_ids(tokenizer, words: list[str], *, leading_space: bool) -> dict[str, int]: + ids: dict[str, int] = {} + for w in words: + text = (" " + w) if leading_space else w + tid = single_token_id(tokenizer, text) + if tid is None: + raise ValueError(f"{text!r} is not a single token") + if tokenizer.decode([tid]) != text: + raise ValueError(f"{text!r} does not round-trip: id {tid} decodes to {tokenizer.decode([tid])!r}") + ids[w] = tid + return ids + + +@torch.no_grad() +def top_rows(tokenizer, q: torch.Tensor, unit_rows: torch.Tensor, exclude: set[int], n: int) -> list[dict]: + """Top-``n`` rows of ``unit_rows`` by |cosine| with ``q``, excluded ids ranked last.""" + q_unit = q / q.norm() + cos = unit_rows @ q_unit # (V,) + score = cos.abs().clone() + for i in exclude: + score[i] = -1.0 + idx = torch.topk(score, n).indices + return [ + { + "rank": r + 1, + "token_id": int(i), + "token": tokenizer.decode([int(i)]), + "cosine": round(float(cos[i]), 4), + } + for r, i in enumerate(idx.tolist()) + ] + + +@torch.no_grad() +def nearest_rows( + W_U: torch.Tensor, + row_mean: torch.Tensor, + tokenizer, + contrasts: list[tuple[str, str]], + *, + top_n: int, + leading_space: bool = True, +) -> tuple[dict[str, dict], dict[str, dict], dict[str, int]]: + """Centered and raw |cosine| rankings for each ``(a, b)`` contrast direction w_a - w_b.""" + words = list(dict.fromkeys(w for pair in contrasts for w in pair)) + report = token_report(tokenizer, words) + ids = resolve_contrast_ids(tokenizer, words, leading_space=leading_space) + W_c = W_U - row_mean # (V, d) centered + W_c_unit = W_c / W_c.norm(dim=1, keepdim=True).clamp_min(1e-8) + W_raw_unit = W_U / W_U.norm(dim=1, keepdim=True).clamp_min(1e-8) + prefix = " " if leading_space else "" + out: dict[str, dict] = {} + for a, b in contrasts: + name = f"{a}_minus_{b}" + q = W_U[ids[a]] - W_U[ids[b]] # row_mean cancels: centered q == raw q + exclude = {ids[a], ids[b]} + out[name] = { + "definition": f"w_[{prefix}{a}] - w_[{prefix}{b}] (ids {ids[a]}, {ids[b]})", + "q_norm": round(float(q.norm()), 4), + "cos_q_centered_rows": { + f"{prefix}{a}": round(float(W_c_unit[ids[a]] @ (q / q.norm())), 4), + f"{prefix}{b}": round(float(W_c_unit[ids[b]] @ (q / q.norm())), 4), + }, + "top_centered": top_rows(tokenizer, q, W_c_unit, exclude, top_n), + "top_raw": top_rows(tokenizer, q, W_raw_unit, exclude, top_n), + } + return out, report, ids + + +def table_rows(contrast_results: dict[str, dict]) -> list[dict]: + """Side-by-side centered / raw listing per rank, cosines to two decimals, glosses attached.""" + rows: list[dict] = [] + for name, c in contrast_results.items(): + for cen, raw in zip(c["top_centered"], c["top_raw"]): + rows.append( + { + "contrast": name, + "rank": cen["rank"], + "centered_cosine": f"{cen['cosine']:+.2f}", + "centered_token": cen["token"], + "centered_gloss": GLOSS.get(cen["token"], ""), + "raw_cosine": f"{raw['cosine']:+.2f}", + "raw_token": raw["token"], + "raw_gloss": GLOSS.get(raw["token"], ""), + } + ) + return rows + + +def parse_contrast(spec: str) -> tuple[str, str]: + parts = [p.strip() for p in spec.split(",")] + if len(parts) != 2 or not all(parts): + raise argparse.ArgumentTypeError(f"contrast must be 'a,b', got {spec!r}") + return parts[0], parts[1] + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--model-id", default="Qwen/Qwen3.5-2B") + ap.add_argument( + "--revision", default=None, help="HF weight revision to pin (paper: 15852e8c16360a2fea060d615a32b45270f8a8fc)" + ) + ap.add_argument("--w-u", type=Path, default=None, help="extraction artifact with W_U_orig; skips the model load") + ap.add_argument( + "--checkpoint", type=Path, default=None, help="dictionary checkpoint whose stored row_mean centers the rows" + ) + ap.add_argument( + "--contrast", + type=parse_contrast, + action="append", + default=None, + help="'a,b' for the direction w_a - w_b (repeatable; default bug,insect and bug,error)", + ) + ap.add_argument( + "--no-leading-space", action="store_true", help="resolve the contrast words without a leading space" + ) + ap.add_argument("--top-n", type=int, default=12) + ap.add_argument("--out-dir", type=Path, required=True) + args = ap.parse_args() + contrasts = args.contrast or [("bug", "insect"), ("bug", "error")] + args.out_dir.mkdir(parents=True, exist_ok=True) + + if args.w_u is not None: + from transformers import AutoTokenizer + + art = torch.load(args.w_u, map_location="cpu", weights_only=True) + W_U = art.get("W_U_orig", art.get("W_U")).float() + w_u_source = f"artifact W_U_orig: {args.w_u}" + tokenizer = AutoTokenizer.from_pretrained(args.model_id, revision=args.revision) + else: + model, tokenizer = load_qwen_model(args.model_id, torch.float32, revision=args.revision) + lm_head, lm_head_path = find_lm_head_with_path(model) + W_U = lm_head.weight.detach().float().cpu() + w_u_source = f"{lm_head_path}.weight (tie_word_embeddings={getattr(model.config, 'tie_word_embeddings', None)})" + del model + V, d = W_U.shape + print(f"W_U: {V} x {d} from {w_u_source}", flush=True) + + ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=True) if args.checkpoint else None + row_mean = resolve_row_mean(W_U, ckpt=ckpt) + + contrast_results, report, ids = nearest_rows( + W_U, row_mean, tokenizer, contrasts, top_n=args.top_n, leading_space=not args.no_leading_space + ) + prefix = "" if args.no_leading_space else " " + results = { + "model_id": args.model_id, + "revision": args.revision, + "w_u_source": w_u_source, + "vocab": V, + "d_model": d, + "tokenization": "single tokens " + ", ".join(repr(prefix + w) for w in ids), + "token_report": report, + "centering": ( + "row_mean from checkpoint " + str(args.checkpoint) + if ckpt is not None and ckpt.get("row_mean") is not None + else "mu = W_U.mean(dim=0), full-vocabulary mean (preprocess_rows centering)" + ), + "ranking": f"top-{args.top_n} by |cosine|; signed cosine reported; contrast tokens excluded", + "contrasts": contrast_results, + "provenance": run_provenance(args), + } + (args.out_dir / "nearest_rows.json").write_text(json.dumps(results, indent=2, ensure_ascii=False, default=str)) + + with (args.out_dir / "nearest_rows.csv").open("w", newline="", encoding="utf-8") as f: + w = csv.writer(f) + w.writerow(["contrast", "variant", "rank", "token_id", "token", "cosine"]) + for name, c in contrast_results.items(): + for variant in ("top_centered", "top_raw"): + for row in c[variant]: + w.writerow([name, variant, row["rank"], row["token_id"], repr(row["token"]), row["cosine"]]) + + table = table_rows(contrast_results) + with (args.out_dir / "nearest_rows_table.csv").open("w", newline="", encoding="utf-8") as f: + w = csv.DictWriter(f, fieldnames=list(table[0].keys())) + w.writeheader() + w.writerows(table) + + for name, c in contrast_results.items(): + print(f"\n=== {name} ({c['definition']}) ===") + for variant in ("top_centered", "top_raw"): + print(f" -- {variant} --") + for row in c[variant]: + print(f" {row['rank']:>2} {row['cosine']:+.4f} {row['token']!r}") + print(f"\nWrote {args.out_dir / 'nearest_rows.json'}, nearest_rows.csv and nearest_rows_table.csv", flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/data/build_cross_lens_de_bank.py b/scripts/data/build_cross_lens_de_bank.py new file mode 100644 index 0000000..2990c2d --- /dev/null +++ b/scripts/data/build_cross_lens_de_bank.py @@ -0,0 +1,727 @@ +#!/usr/bin/env python3 +"""Build the English-German cross-lens prompt bank and its cognate-exclusion report. + +Writes ``data/cross_lens/cross_lens_prompts_en_de.json`` (90 cross-lens prompts in +seven families plus 12 controls) and the companion exclusion report, the inputs +of the second-language-pair cells of the cross-lens study +(``tab:app-cross-lens-extension`` rows "EN--DE, Jacobian" and "EN--DE, Jacobian, +n=300"; ``fig:cross-lens-extension``). The bank mirrors the EN-ZH bank with the +script confound removed: both languages are Latin-script, so the surface-language +call is lexical and the cognate filter below is load-bearing. + +Construction rules, fixed before any lens was fitted: + - Every decomposed surface (form_a, form_b, nulls) must be a single token + under the Qwen tokenizer; the builder tries the context-appropriate spacing + first (no leading space inside a quoted frame, a leading space otherwise), + then the alternative, and drops the item if neither is a single token. + - Cognate exclusion: after NFKD diacritic stripping, ss/eszett folding and + casefolding, items with Levenshtein(form_a, form_b) <= 2 or identical forms + are excluded. Exclusions are reported, not hidden (30 of 132 candidates in + the shipped bank). + - null_targets are language-matched unrelated nouns [de_null, en_null], + assigned round-robin, skipping any null equal to either form. + +Paper run:: + + build_cross_lens_de_bank.py --tokenizer Qwen/Qwen3.5-9B \ + --out data/cross_lens/cross_lens_prompts_en_de.json \ + --report data/cross_lens/cross_lens_prompts_en_de_exclusion_report.json +""" + +from __future__ import annotations + +import argparse +import json +import unicodedata +from pathlib import Path + +# (group, concept, prompt, de_surface, en_surface, answer_lang) +# answer_lang: which side is the in-context gold continuation (form_a). +CANDIDATES = [ + # --- antonym_de: German few-shot antonym completion, German answer --- + ("antonym_de", "big", 'Das Gegenteil von "hoch" ist "tief". Das Gegenteil von "klein" ist "', "groß", "big", "de"), + ( + "antonym_de", + "small", + 'Das Gegenteil von "hoch" ist "tief". Das Gegenteil von "groß" ist "', + "klein", + "small", + "de", + ), + ( + "antonym_de", + "slow", + 'Das Gegenteil von "hoch" ist "tief". Das Gegenteil von "schnell" ist "', + "langsam", + "slow", + "de", + ), + ( + "antonym_de", + "bright", + 'Das Gegenteil von "hoch" ist "tief". Das Gegenteil von "dunkel" ist "', + "hell", + "bright", + "de", + ), + ( + "antonym_de", + "light", + 'Das Gegenteil von "hoch" ist "tief". Das Gegenteil von "schwer" ist "', + "leicht", + "light", + "de", + ), + ( + "antonym_de", + "bad", + 'Das Gegenteil von "hoch" ist "tief". Das Gegenteil von "gut" ist "', + "schlecht", + "bad", + "de", + ), + ( + "antonym_de", + "empty", + 'Das Gegenteil von "hoch" ist "tief". Das Gegenteil von "voll" ist "', + "leer", + "empty", + "de", + ), + ( + "antonym_de", + "dry", + 'Das Gegenteil von "hoch" ist "tief". Das Gegenteil von "nass" ist "', + "trocken", + "dry", + "de", + ), + ( + "antonym_de", + "weak", + 'Das Gegenteil von "hoch" ist "tief". Das Gegenteil von "stark" ist "', + "schwach", + "weak", + "de", + ), + ( + "antonym_de", + "dark", + 'Das Gegenteil von "hoch" ist "tief". Das Gegenteil von "hell" ist "', + "dunkel", + "dark", + "de", + ), + ("antonym_de", "late", 'Das Gegenteil von "hoch" ist "tief". Das Gegenteil von "früh" ist "', "spät", "late", "de"), + ( + "antonym_de", + "cheap", + 'Das Gegenteil von "hoch" ist "tief". Das Gegenteil von "teuer" ist "', + "billig", + "cheap", + "de", + ), + ("antonym_de", "new", 'Das Gegenteil von "hoch" ist "tief". Das Gegenteil von "alt" ist "', "neu", "new", "de"), + ( + "antonym_de", + "quiet", + 'Das Gegenteil von "hoch" ist "tief". Das Gegenteil von "laut" ist "', + "leise", + "quiet", + "de", + ), + # --- trans_de2en: German metalinguistic frame, English answer --- + ( + "trans_de2en", + "dog", + 'Das englische Wort für "Himmel" ist "sky". Das englische Wort für "Hund" ist "', + "Hund", + "dog", + "en", + ), + ( + "trans_de2en", + "tree", + 'Das englische Wort für "Himmel" ist "sky". Das englische Wort für "Baum" ist "', + "Baum", + "tree", + "en", + ), + ( + "trans_de2en", + "table", + 'Das englische Wort für "Himmel" ist "sky". Das englische Wort für "Tisch" ist "', + "Tisch", + "table", + "en", + ), + ( + "trans_de2en", + "chair", + 'Das englische Wort für "Himmel" ist "sky". Das englische Wort für "Stuhl" ist "', + "Stuhl", + "chair", + "en", + ), + ( + "trans_de2en", + "window", + 'Das englische Wort für "Himmel" ist "sky". Das englische Wort für "Fenster" ist "', + "Fenster", + "window", + "en", + ), + ( + "trans_de2en", + "horse", + 'Das englische Wort für "Himmel" ist "sky". Das englische Wort für "Pferd" ist "', + "Pferd", + "horse", + "en", + ), + ( + "trans_de2en", + "bird", + 'Das englische Wort für "Himmel" ist "sky". Das englische Wort für "Vogel" ist "', + "Vogel", + "bird", + "en", + ), + ( + "trans_de2en", + "mountain", + 'Das englische Wort für "Himmel" ist "sky". Das englische Wort für "Berg" ist "', + "Berg", + "mountain", + "en", + ), + ( + "trans_de2en", + "rain", + 'Das englische Wort für "Himmel" ist "sky". Das englische Wort für "Regen" ist "', + "Regen", + "rain", + "en", + ), + ( + "trans_de2en", + "cheese", + 'Das englische Wort für "Himmel" ist "sky". Das englische Wort für "Käse" ist "', + "Käse", + "cheese", + "en", + ), + ( + "trans_de2en", + "egg", + 'Das englische Wort für "Himmel" ist "sky". Das englische Wort für "Ei" ist "', + "Ei", + "egg", + "en", + ), + ( + "trans_de2en", + "door", + 'Das englische Wort für "Himmel" ist "sky". Das englische Wort für "Tür" ist "', + "Tür", + "door", + "en", + ), + ( + "trans_de2en", + "flower", + 'Das englische Wort für "Himmel" ist "sky". Das englische Wort für "Blume" ist "', + "Blume", + "flower", + "en", + ), + ( + "trans_de2en", + "knife", + 'Das englische Wort für "Himmel" ist "sky". Das englische Wort für "Messer" ist "', + "Messer", + "knife", + "en", + ), + # --- trans_en2de: English metalinguistic frame, German answer --- + ( + "trans_en2de", + "dog", + 'The German word for "sky" is "Himmel". The German word for "dog" is "', + "Hund", + "dog", + "de", + ), + ( + "trans_en2de", + "tree", + 'The German word for "sky" is "Himmel". The German word for "tree" is "', + "Baum", + "tree", + "de", + ), + ( + "trans_en2de", + "table", + 'The German word for "sky" is "Himmel". The German word for "table" is "', + "Tisch", + "table", + "de", + ), + ( + "trans_en2de", + "chair", + 'The German word for "sky" is "Himmel". The German word for "chair" is "', + "Stuhl", + "chair", + "de", + ), + ( + "trans_en2de", + "window", + 'The German word for "sky" is "Himmel". The German word for "window" is "', + "Fenster", + "window", + "de", + ), + ( + "trans_en2de", + "horse", + 'The German word for "sky" is "Himmel". The German word for "horse" is "', + "Pferd", + "horse", + "de", + ), + ( + "trans_en2de", + "bird", + 'The German word for "sky" is "Himmel". The German word for "bird" is "', + "Vogel", + "bird", + "de", + ), + ( + "trans_en2de", + "mountain", + 'The German word for "sky" is "Himmel". The German word for "mountain" is "', + "Berg", + "mountain", + "de", + ), + ( + "trans_en2de", + "rain", + 'The German word for "sky" is "Himmel". The German word for "rain" is "', + "Regen", + "rain", + "de", + ), + ( + "trans_en2de", + "cheese", + 'The German word for "sky" is "Himmel". The German word for "cheese" is "', + "Käse", + "cheese", + "de", + ), + ("trans_en2de", "egg", 'The German word for "sky" is "Himmel". The German word for "egg" is "', "Ei", "egg", "de"), + ( + "trans_en2de", + "door", + 'The German word for "sky" is "Himmel". The German word for "door" is "', + "Tür", + "door", + "de", + ), + ( + "trans_en2de", + "flower", + 'The German word for "sky" is "Himmel". The German word for "flower" is "', + "Blume", + "flower", + "de", + ), + ( + "trans_en2de", + "knife", + 'The German word for "sky" is "Himmel". The German word for "knife" is "', + "Messer", + "knife", + "de", + ), + ( + "trans_de2en", + "train", + 'Das englische Wort für "Himmel" ist "sky". Das englische Wort für "Zug" ist "', + "Zug", + "train", + "en", + ), + ( + "trans_de2en", + "city", + 'Das englische Wort für "Himmel" ist "sky". Das englische Wort für "Stadt" ist "', + "Stadt", + "city", + "en", + ), + ( + "trans_de2en", + "forest", + 'Das englische Wort für "Himmel" ist "sky". Das englische Wort für "Wald" ist "', + "Wald", + "forest", + "en", + ), + ( + "trans_de2en", + "sea", + 'Das englische Wort für "Himmel" ist "sky". Das englische Wort für "Meer" ist "', + "Meer", + "sea", + "en", + ), + ( + "trans_de2en", + "fire", + 'Das englische Wort für "Himmel" ist "sky". Das englische Wort für "Feuer" ist "', + "Feuer", + "fire", + "en", + ), + ( + "trans_de2en", + "newspaper", + 'Das englische Wort für "Himmel" ist "sky". Das englische Wort für "Zeitung" ist "', + "Zeitung", + "newspaper", + "en", + ), + ( + "trans_de2en", + "kitchen", + 'Das englische Wort für "Himmel" ist "sky". Das englische Wort für "Küche" ist "', + "Küche", + "kitchen", + "en", + ), + ( + "trans_de2en", + "sister", + 'Das englische Wort für "Himmel" ist "sky". Das englische Wort für "Schwester" ist "', + "Schwester", + "sister", + "en", + ), + ( + "trans_en2de", + "train", + 'The German word for "sky" is "Himmel". The German word for "train" is "', + "Zug", + "train", + "de", + ), + ( + "trans_en2de", + "city", + 'The German word for "sky" is "Himmel". The German word for "city" is "', + "Stadt", + "city", + "de", + ), + ( + "trans_en2de", + "forest", + 'The German word for "sky" is "Himmel". The German word for "forest" is "', + "Wald", + "forest", + "de", + ), + ( + "trans_en2de", + "sea", + 'The German word for "sky" is "Himmel". The German word for "sea" is "', + "Meer", + "sea", + "de", + ), + ( + "trans_en2de", + "fire", + 'The German word for "sky" is "Himmel". The German word for "fire" is "', + "Feuer", + "fire", + "de", + ), + ( + "trans_en2de", + "newspaper", + 'The German word for "sky" is "Himmel". The German word for "newspaper" is "', + "Zeitung", + "newspaper", + "de", + ), + ( + "trans_en2de", + "kitchen", + 'The German word for "sky" is "Himmel". The German word for "kitchen" is "', + "Küche", + "kitchen", + "de", + ), + ( + "trans_en2de", + "sister", + 'The German word for "sky" is "Himmel". The German word for "sister" is "', + "Schwester", + "sister", + "de", + ), + # --- cloze_de: German sentence completion, German noun answer --- + ("cloze_de", "pen", "Zum Schreiben benutze ich einen", "Stift", "pen", "de"), + ("cloze_de", "knife", "Ich schneide das Brot mit einem scharfen", "Messer", "knife", "de"), + ("cloze_de", "station", "Der Zug hält pünktlich am", "Bahnhof", "station", "de"), + ("cloze_de", "door", "Ich klopfe dreimal und öffne dann die", "Tür", "door", "de"), + ("cloze_de", "book", "Vor dem Schlafen liest das Kind noch ein", "Buch", "book", "de"), + ("cloze_de", "tree", "Der Vogel sitzt hoch oben auf dem", "Baum", "tree", "de"), + ("cloze_de", "snow", "Im Winter fällt auf den Bergen viel", "Schnee", "snow", "de"), + ("cloze_de", "coffee", "Zum Frühstück trinke ich immer eine Tasse", "Kaffee", "coffee", "de"), + ("cloze_de", "key", "Ich schließe die Wohnung ab mit dem", "Schlüssel", "key", "de"), + ("cloze_de", "horse", "Auf dem Bauernhof reitet das Mädchen ein", "Pferd", "horse", "de"), + ("cloze_de", "cheese", "Zur Brotzeit essen wir Brot mit", "Käse", "cheese", "de"), + ("cloze_de", "flower", "Zum Geburtstag schenke ich ihr eine", "Blume", "flower", "de"), + ("cloze_de", "school", "Nach den Ferien gehen die Kinder wieder in die", "Schule", "school", "de"), + ("cloze_de", "church", "Sonntags läuten die Glocken der alten", "Kirche", "church", "de"), + ("cloze_de", "train", "Wir fahren in den Urlaub bequem mit dem", "Zug", "train", "de"), + ("cloze_de", "forest", "Am Wochenende wandern wir stundenlang durch den", "Wald", "forest", "de"), + ("cloze_de", "ship", "Der Kapitän steuert das große", "Schiff", "ship", "de"), + ("cloze_de", "money", "Für die Fahrkarte brauche ich noch etwas", "Geld", "money", "de"), + # --- cloze_en: English sentence completion, English answer --- + ("cloze_en", "key", "I unlock the front door with my", "Schlüssel", "key", "en"), + ("cloze_en", "horse", "At the farm the girl rides a brown", "Pferd", "horse", "en"), + ("cloze_en", "window", "She waved at us through the kitchen", "Fenster", "window", "en"), + ("cloze_en", "tree", "The bird built its nest high up in the", "Baum", "tree", "en"), + ("cloze_en", "table", "He put the plates down on the wooden", "Tisch", "table", "en"), + ("cloze_en", "dog", "The postman was chased down the street by a barking", "Hund", "dog", "en"), + ("cloze_en", "rain", "Take an umbrella, the forecast says heavy", "Regen", "rain", "en"), + ("cloze_en", "cheese", "The mouse in the cartoon always steals the", "Käse", "cheese", "en"), + ("cloze_en", "egg", "For breakfast she boiled herself an", "Ei", "egg", "en"), + ("cloze_en", "mountain", "They spent all summer climbing the highest", "Berg", "mountain", "en"), + ("cloze_en", "knife", "He spread the butter with a blunt", "Messer", "knife", "en"), + ("cloze_en", "bird", "High above the field we watched a small", "Vogel", "bird", "en"), + ("cloze_en", "train", "We traveled to the capital on a high-speed", "Zug", "train", "en"), + ("cloze_en", "money", "He paid for the tickets with borrowed", "Geld", "money", "en"), + ("cloze_en", "voice", "The choir admired her beautiful singing", "Stimme", "voice", "en"), + ("cloze_en", "newspaper", "Every morning my grandfather reads the", "Zeitung", "newspaper", "en"), + ("cloze_en", "war", "After years of peace the country slid into", "Krieg", "war", "en"), + ("cloze_en", "city", "From the hilltop you can see the whole", "Stadt", "city", "en"), + # --- fact_de: German general knowledge, German common-noun answer --- + ("fact_de", "lion", "Der König der Tiere ist der", "Löwe", "lion", "de"), + ("fact_de", "dog", "Der beste Freund des Menschen ist der", "Hund", "dog", "de"), + ("fact_de", "bee", "Honig wird gemacht von der fleißigen", "Biene", "bee", "de"), + ("fact_de", "moon", "Nachts sieht man am Himmel die Sterne und den", "Mond", "moon", "de"), + ("fact_de", "snow", "Weiß und kalt fällt im Winter der", "Schnee", "snow", "de"), + ("fact_de", "egg", "Die Henne legt jeden Morgen ein", "Ei", "egg", "de"), + ("fact_de", "horse", "Vor die Kutsche spannt man ein", "Pferd", "horse", "de"), + ("fact_de", "cheese", "Aus Milch macht man Butter und", "Käse", "cheese", "de"), + ("fact_de", "week", "Sieben Tage sind eine", "Woche", "week", "de"), + ("fact_de", "year", "Zwölf Monate sind ein", "Jahr", "year", "de"), + ("fact_de", "tooth", "Der Zahnarzt zieht den kranken", "Zahn", "tooth", "de"), + ("fact_de", "island", "Ein Stück Land mitten im Meer nennt man eine", "Insel", "island", "de"), + ("fact_de", "church", "Am Sonntag gehen viele Menschen in die", "Kirche", "church", "de"), + # --- exemplar_de: German category exemplars, German answer --- + ("exemplar_de", "dog", "Ein Beispiel für ein Haustier ist der", "Hund", "dog", "de"), + ("exemplar_de", "table", "Ein Beispiel für ein Möbelstück ist der", "Tisch", "table", "de"), + ("exemplar_de", "juice", "Ein Beispiel für ein Getränk ist der", "Saft", "juice", "de"), + ("exemplar_de", "doctor", "Ein Beispiel für einen Beruf ist der", "Arzt", "doctor", "de"), + ("exemplar_de", "teacher", "Ein Beispiel für einen Beruf ist der", "Lehrer", "teacher", "de"), + ("exemplar_de", "coat", "Ein Beispiel für ein Kleidungsstück ist der", "Mantel", "coat", "de"), + ("exemplar_de", "shirt", "Ein Beispiel für ein Kleidungsstück ist das", "Hemd", "shirt", "de"), + ("exemplar_de", "bird", "Ein Beispiel für ein Tier mit Flügeln ist der", "Vogel", "bird", "de"), + ("exemplar_de", "head", "Ein Beispiel für ein Körperteil ist der", "Kopf", "head", "de"), + ("exemplar_de", "sea", "Ein Beispiel für ein Gewässer ist das", "Meer", "sea", "de"), + ("exemplar_de", "train", "Ein Beispiel für ein Verkehrsmittel ist der", "Zug", "train", "de"), + ("exemplar_de", "church", "Ein Beispiel für ein Gebäude ist die", "Kirche", "church", "de"), + ("exemplar_de", "sister", "Ein Beispiel für ein Familienmitglied ist die", "Schwester", "sister", "de"), + # --- controls: no language flip expected --- + ("ctrl_digit", "seven", "3 + 4 =", "7", "7", "de"), + ("ctrl_digit", "nine", "5 + 4 =", "9", "9", "de"), + ("ctrl_digit", "eight", "6 + 2 =", "8", "8", "de"), + ("ctrl_digit", "six", "2 + 4 =", "6", "6", "de"), + ("ctrl_digit", "five", "1 + 4 =", "5", "5", "de"), + ("ctrl_digit", "four", "1 + 3 =", "4", "4", "de"), + ("ctrl_propn", "paris", "Die Hauptstadt von Frankreich heißt", "Paris", "Paris", "de"), + ("ctrl_propn", "berlin", "Die Hauptstadt von Deutschland heißt", "Berlin", "Berlin", "de"), + ("ctrl_propn", "london", "Die Hauptstadt von England heißt", "London", "London", "de"), + ("ctrl_propn", "madrid", "Die Hauptstadt von Spanien heißt", "Madrid", "Madrid", "de"), + ("ctrl_propn", "oslo", "Die Hauptstadt von Norwegen heißt", "Oslo", "Oslo", "de"), + ("ctrl_propn", "lima", "Die Hauptstadt von Peru heißt", "Lima", "Lima", "de"), +] + +DE_NULLS = ["Salz", "Wand", "Schuh", "Tasse", "Nadel", "Karte", "Seife", "Gabel"] +EN_NULLS = ["salt", "wall", "shoe", "cup", "needle", "ticket", "soap", "fork"] + +CROSS_GROUPS = { + "antonym_de", + "trans_de2en", + "trans_en2de", + "cloze_de", + "cloze_en", + "fact_de", + "exemplar_de", +} + +NOTE = ( + "Cross-lens prompt bank EN-DE. Same schema and metrics as " + "cross_lens_prompts_en_zh.json; both languages Latin-script, so the " + "surface-language call is lexical and cognates (normalized Levenshtein " + "<= 2) are excluded by construction. form_a = in-context answer form, " + "form_b = other-language form; null_targets = [de, en] language-matched " + "unrelated tokens. Built by scripts/data/build_cross_lens_de_bank.py; " + "exclusions in the companion report cross_lens_prompts_en_de_exclusion_report.json." +) + + +def norm(s: str) -> str: + s = s.strip().casefold().replace("ß", "ss") + s = unicodedata.normalize("NFKD", s) + return "".join(c for c in s if not unicodedata.combining(c)) + + +def levenshtein(a: str, b: str) -> int: + if len(a) < len(b): + a, b = b, a + prev = list(range(len(b) + 1)) + for i, ca in enumerate(a, 1): + cur = [i] + for j, cb in enumerate(b, 1): + cur.append(min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (ca != cb))) + prev = cur + return prev[-1] + + +def single_token(tok, surface: str, prefer_space: bool) -> str | None: + """Return the single-token spacing variant of `surface`, or None.""" + variants = [" " + surface, surface] if prefer_space else [surface, " " + surface] + for v in variants: + if len(tok.encode(v, add_special_tokens=False)) == 1: + return v + return None + + +def build_bank(tok) -> tuple[list[dict], list[dict]]: + """Apply the construction rules to CANDIDATES; returns (prompts, excluded).""" + prompts, excluded = [], [] + counters: dict[str, int] = {} + null_i = 0 + for group, concept, prompt, de, en, answer_lang in CANDIDATES: + counters[group] = counters.get(group, 0) + 1 + pid = f"{group}_{counters[group]:02d}" + is_ctrl = group.startswith("ctrl_") + + if not is_ctrl and (norm(de) == norm(en) or levenshtein(norm(de), norm(en)) <= 2): + excluded.append( + { + "id": pid, + "concept": concept, + "de": de, + "en": en, + "reason": f"cognate (lev={levenshtein(norm(de), norm(en))})", + } + ) + continue + + # In-context continuation: quoted frames take no leading space, plain + # sentence frames do. The answer-language surface is form_a. + quoted = prompt.rstrip().endswith('"') + de_tok = single_token(tok, de, prefer_space=not quoted if answer_lang == "de" else True) + en_tok = single_token(tok, en, prefer_space=not quoted if answer_lang == "en" else True) + if de_tok is None or en_tok is None: + excluded.append( + { + "id": pid, + "concept": concept, + "de": de, + "en": en, + "reason": "not single-token: " + + ("de " if de_tok is None else "") + + ("en" if en_tok is None else ""), + } + ) + continue + + form_a, form_b = (de_tok, en_tok) if answer_lang == "de" else (en_tok, de_tok) + + # Language-matched nulls, skipping collisions with either form. + nulls = [] + for pool, prefer in ((DE_NULLS, True), (EN_NULLS, True)): + for step in range(len(pool)): + cand = pool[(null_i + step) % len(pool)] + if norm(cand) in (norm(de), norm(en)): + continue + cand_tok = single_token(tok, cand, prefer_space=prefer) + if cand_tok is not None: + nulls.append(cand_tok) + break + null_i += 1 + if len(nulls) < 2: + excluded.append({"id": pid, "concept": concept, "de": de, "en": en, "reason": "no single-token null pair"}) + continue + + prompts.append( + { + "id": pid, + "group": group, + "concept": concept, + "prompt": prompt, + "form_a": form_a, + "form_b": form_b, + "lang_a": answer_lang, + "lang_b": "en" if answer_lang == "de" else "de", + "targets": [form_a, form_b, *nulls], + "null_targets": nulls, + "answer": form_a.strip(), + } + ) + return prompts, excluded + + +def main(argv: list[str] | None = None) -> int: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--tokenizer", default="Qwen/Qwen3.5-9B") + ap.add_argument("--out", required=True, help="output bank JSON") + ap.add_argument("--report", required=True, help="output exclusion report JSON") + args = ap.parse_args(argv) + + from transformers import AutoTokenizer + + tok = AutoTokenizer.from_pretrained(args.tokenizer) + prompts, excluded = build_bank(tok) + + n_cross = sum(1 for p in prompts if p["group"] in CROSS_GROUPS) + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + out.write_text(json.dumps({"note": NOTE, "prompts": prompts}, ensure_ascii=False, indent=1)) + report = { + "tokenizer": args.tokenizer, + "n_candidates": len(CANDIDATES), + "n_kept": len(prompts), + "n_cross": n_cross, + "n_ctrl": len(prompts) - n_cross, + "per_group": {g: sum(1 for p in prompts if p["group"] == g) for g in sorted({p["group"] for p in prompts})}, + "excluded": excluded, + } + report_path = Path(args.report) + report_path.parent.mkdir(parents=True, exist_ok=True) + report_path.write_text(json.dumps(report, ensure_ascii=False, indent=1)) + print( + f"[bank] kept {len(prompts)} ({n_cross} cross + {len(prompts) - n_cross} ctrl), " + f"excluded {len(excluded)} of {len(CANDIDATES)} candidates" + ) + for e in excluded: + print(f" - {e['id']} {e['de']}/{e['en']}: {e['reason']}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/eval/aggregate_cross_lens_en_de.py b/scripts/eval/aggregate_cross_lens_en_de.py new file mode 100644 index 0000000..e731387 --- /dev/null +++ b/scripts/eval/aggregate_cross_lens_en_de.py @@ -0,0 +1,265 @@ +#!/usr/bin/env python3 +"""Aggregate an English-German cross-lens sweep: same metrics as the EN-ZH aggregator, lexical language call. + +Reads two readout dumps from ``scripts/run/run_cross_lens_readouts.py`` (the same +frozen Qwen3.5-9B states read through the English- and German-fitted lenses) and +the EN-DE prompt bank, and computes the EN-DE rows of +``tab:app-cross-lens-extension`` / ``fig:cross-lens-extension`` and the per-family +counts of the second-language-pair paragraph: cross-lens dominant-feature +agreement (majority of the mid-band layers {21, 24, 26, 29}, Wilson 95% CIs), +within-lens cross-form agreement, the unrelated-token and shuffled-pairing null +floors, and the token divergence rate (the two lenses' top-1 strings differ on a +majority of mid-band layers). + +Both languages share the Latin script, so the surface-language call for a top-1 +token is lexical rather than by script. In order: (1) membership in the bank +item's own gold forms (the form in the item's language); (2) an umlaut / eszett +cue means German; (3) membership in small embedded EN / DE common-word lists +(casefolded, diacritic-normalised); else OTHER. The per-prompt call is the +mid-band plurality vote (a 2-2 tie resolves to the alphabetically first label, +DE < EN < OTHER). The headline agreement metric never uses the language call; +it is a descriptive diagnostic. + +Paper runs (``--prompts data/cross_lens/cross_lens_prompts_en_de.json --seed 0``):: + + # 100-prompt lenses: extension-matrix row 2 + aggregate_cross_lens_en_de.py --lens-a /en_de__jlens_en_n100.json \ + --lens-b /en_de__jlens_de_n100.json --out /en_de_jlens_n100_summary.json + # 300-prompt refits: row 6 + aggregate_cross_lens_en_de.py --lens-a /en_de__jlens_en_n300.json \ + --lens-b /en_de__jlens_de_n300.json --out /en_de_jlens_n300_summary.json +""" + +from __future__ import annotations + +import argparse +import json +import math +import random +import unicodedata +from pathlib import Path + +from sparse_readout_prism.utils import write_csv + +MIDBAND = ["21", "24", "26", "29"] +POS = "-1" +CROSS_GROUPS = ( + "antonym_de", + "trans_de2en", + "trans_en2de", + "cloze_de", + "cloze_en", + "fact_de", + "exemplar_de", +) + +EN_WORDS = set( + """the of and to in is was for on as with by at from it an be this that or + are not his her their they them he she you we all one two more no yes big + small slow bright light bad empty dry weak dark late cheap new quiet dog + tree table chair window horse bird mountain rain cheese egg door flower + knife train city forest sea fire newspaper kitchen sister pen station book + snow coffee key head money voice war week year tooth island church school + ship path car work child woman brother song game time clock street village + meat fruit sugar eye mouth lake sky water day night house man word for is + """.split() +) +DE_WORDS = set( + """der die das und ist war für auf als mit von bei aus es ein eine sein + nicht sie er wir ihr alle einen einem einer dem den groß klein langsam + hell leicht schlecht leer trocken schwach dunkel spät billig neu leise + hund baum tisch stuhl fenster pferd vogel berg regen käse ei tür blume + messer zug stadt wald meer feuer zeitung küche schwester stift bahnhof + buch schnee kaffee schlüssel kopf geld stimme krieg woche jahr zahn insel + kirche schule schiff weg wagen arbeit kind frau bruder lied spiel zeit + uhr straße dorf fleisch obst zucker auge mund see himmel wasser tag nacht + haus mann wort löwe biene mond arzt lehrer mantel hemd saft + """.split() +) + + +def wilson(k, n, z=1.96): + if n == 0: + return (0.0, 0.0, 1.0) + p = k / n + d = 1 + z * z / n + c = (p + z * z / (2 * n)) / d + h = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / d + return (p, max(0.0, c - h), min(1.0, c + h)) + + +def norm(s): + s = s.strip().casefold().replace("ß", "ss") + s = unicodedata.normalize("NFKD", s) + return "".join(c for c in s if not unicodedata.combining(c)) + + +def lang_of(tok, item): + t = tok.strip() + if not t: + return "OTHER" + n = norm(t) + forms = { + "de": item["form_a"] if item["lang_a"] == "de" else item["form_b"], + "en": item["form_a"] if item["lang_a"] == "en" else item["form_b"], + } + if n == norm(forms["de"]): + return "DE" + if n == norm(forms["en"]): + return "EN" + if any(c in t for c in "äöüÄÖÜß"): + return "DE" + in_de, in_en = n in DE_WORDS, n in EN_WORDS + if in_de and not in_en: + return "DE" + if in_en and not in_de: + return "EN" + return "OTHER" + + +def dom_feat(rec, layer, target): + t = rec["layers"].get(layer, {}).get(POS, {}).get("targets", {}).get(target) + if not t or not t.get("top_features"): + return None + return t["top_features"][0]["id"] + + +def majority_same(rec_a, rec_b, target_a, target_b): + same = total = 0 + for L in MIDBAND: + fa, fb = dom_feat(rec_a, L, target_a), dom_feat(rec_b, L, target_b) + if fa is None or fb is None: + continue + total += 1 + same += int(fa == fb) + if total == 0: + return None + return same >= (total + 1) // 2 + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--lens-a", required=True, help="dump under the English-fitted lens (output keys *_en)") + p.add_argument("--lens-b", required=True, help="dump under the German-fitted lens (output keys *_de)") + p.add_argument("--prompts", required=True, help="EN-DE prompt bank JSON") + p.add_argument("--seed", type=int, default=0, help="seed for the shuffled-pairing null") + p.add_argument("--out", required=True, help="summary JSON (per-group rates, floors, per-prompt rows)") + p.add_argument("--rows-csv", default=None, help="optional per-prompt rows as CSV") + args = p.parse_args(argv) + + ar = {r["id"]: r for r in json.loads(Path(args.lens_a).read_text())["records"]} + br = {r["id"]: r for r in json.loads(Path(args.lens_b).read_text())["records"]} + bank_list = json.loads(Path(args.prompts).read_text())["prompts"] + assert len({v["id"] for v in bank_list}) == len(bank_list), "duplicate prompt ids in bank" + bank = {v["id"]: v for v in bank_list} + ids = sorted(set(ar) & set(br) & set(bank)) + print(f"[agg] {len(ids)} prompts present in both dumps") + + rows = [] + for rid in ids: + v, a, b = bank[rid], ar[rid], br[rid] + fa, fb = v["form_a"], v["form_b"] + row = {"id": rid, "group": v["group"], "concept": v.get("concept", "")} + row["cross_lens_pass"] = majority_same(a, b, fa, fa) + row["cross_form_en"] = majority_same(a, a, fa, fb) + row["cross_form_de"] = majority_same(b, b, fa, fb) + nulls = [majority_same(a, a, fa, nt) for nt in v.get("null_targets", [])] + nulls = [x for x in nulls if x is not None] + row["null_hits"] = sum(nulls) + row["null_total"] = len(nulls) + # top-1 language under each lens (mid-band vote), lexical call + for tag, rec in (("en", a), ("de", b)): + langs = [] + for L in MIDBAND: + top5 = rec["layers"].get(L, {}).get(POS, {}).get("top5") + if top5: + langs.append(lang_of(top5[0][0], v)) + # Plurality vote; a 2-2 tie resolves to the alphabetically first label + # (DE < EN < OTHER), fixed so the vote is deterministic. + row[f"top1_lang_{tag}"] = max(sorted(set(langs)), key=langs.count) if langs else "NA" + # divergence: do the two lenses' top-1 token strings differ (mid-band vote)? + diff = tot = 0 + for L in MIDBAND: + ta = a["layers"].get(L, {}).get(POS, {}).get("top5") + tb = b["layers"].get(L, {}).get(POS, {}).get("top5") + if ta and tb: + tot += 1 + diff += int(ta[0][0] != tb[0][0]) + row["token_diverges"] = (diff >= (tot + 1) // 2) if tot else None + rows.append(row) + + rng = random.Random(args.seed) + cross_ids = [r["id"] for r in rows if r["group"] in CROSS_GROUPS] + shuffle_hits = shuffle_total = 0 + for rid in cross_ids: + others = [x for x in cross_ids if bank[x]["concept"] != bank[rid]["concept"]] + for oid in rng.sample(others, min(3, len(others))): + res = majority_same(ar[rid], br[oid], bank[rid]["form_a"], bank[oid]["form_a"]) + if res is not None: + shuffle_total += 1 + shuffle_hits += int(res) + + def rate(sel): + vals = [r for r in rows if sel(r) and r["cross_lens_pass"] is not None] + k = sum(r["cross_lens_pass"] for r in vals) + return k, len(vals), wilson(k, len(vals)) + + summary: dict = {"per_group": {}, "rows": rows} + print("\n=== CROSS-LENS DOMINANT-FEATURE AGREEMENT (majority of mid-band) ===") + for grp in sorted({r["group"] for r in rows}): + k, n, (pt, lo, hi) = rate(lambda r, g=grp: r["group"] == g) + summary["per_group"][grp] = {"pass": k, "n": n, "rate": pt, "ci": [lo, hi]} + print(f" {grp:14s} {k:3d}/{n:<3d} {pt:.2f} [{lo:.2f}, {hi:.2f}]") + k, n, (pt, lo, hi) = rate(lambda r: r["group"] in CROSS_GROUPS) + summary["headline"] = {"pass": k, "n": n, "rate": pt, "ci": [lo, hi]} + print(f" {'ALL CROSS':14s} {k:3d}/{n:<3d} {pt:.2f} [{lo:.2f}, {hi:.2f}]") + + cf_en = [r["cross_form_en"] for r in rows if r["group"] in CROSS_GROUPS and r["cross_form_en"] is not None] + cf_de = [r["cross_form_de"] for r in rows if r["group"] in CROSS_GROUPS and r["cross_form_de"] is not None] + nk = sum(r["null_hits"] for r in rows) + nn = sum(r["null_total"] for r in rows) + summary["cross_form_en"] = {"pass": sum(cf_en), "n": len(cf_en)} + summary["cross_form_de"] = {"pass": sum(cf_de), "n": len(cf_de)} + summary["null_within_lens"] = {"pass": nk, "n": nn, "rate": wilson(nk, nn)[0]} + summary["null_shuffle_cross_lens"] = { + "pass": shuffle_hits, + "n": shuffle_total, + "rate": wilson(shuffle_hits, shuffle_total)[0], + } + div = [r["token_diverges"] for r in rows if r["group"] in CROSS_GROUPS and r["token_diverges"] is not None] + summary["divergence_rate"] = {"diverging": sum(div), "n": len(div)} + print("\n=== ONE FEATURE CARRIES BOTH FORMS (within-lens) ===") + print(f" EN lens: {sum(cf_en)}/{len(cf_en)} DE lens: {sum(cf_de)}/{len(cf_de)}") + print("\n=== NULL FLOORS ===") + print(f" (a) form vs unrelated token, within-lens: {nk}/{nn} ({wilson(nk, nn)[0]:.2f})") + print( + f" (b) cross-lens shuffled prompts: {shuffle_hits}/{shuffle_total} " + f"({wilson(shuffle_hits, shuffle_total)[0]:.2f})" + ) + print("\n=== TOKEN DIVERGENCE (descriptive) ===") + print(f" top-1 differs between lenses on {sum(div)}/{len(div)} cross prompts") + print("\n=== LANGUAGE-FOLLOWS-LENS (top-1 lexical call, mid-band vote) ===") + summary["lens_only_split"] = {} + for grp in sorted({r["group"] for r in rows}): + sel = [r for r in rows if r["group"] == grp] + flip = sum(1 for r in sel if r["top1_lang_en"] == "EN" and r["top1_lang_de"] == "DE") + summary["lens_only_split"][grp] = {"pass": flip, "n": len(sel)} + print(f" {grp:14s} EN-lens=EN & DE-lens=DE on {flip}/{len(sel)}") + cross_sel = [r for r in rows if r["group"] in CROSS_GROUPS] + flip = sum(1 for r in cross_sel if r["top1_lang_en"] == "EN" and r["top1_lang_de"] == "DE") + summary["lens_only_split"]["all_cross"] = {"pass": flip, "n": len(cross_sel)} + print(f" {'ALL CROSS':14s} EN-lens=EN & DE-lens=DE on {flip}/{len(cross_sel)}") + + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + with out.open("w", encoding="utf-8") as f: + json.dump(summary, f, ensure_ascii=False, indent=1) + print(f"\n[agg] wrote {out}") + if args.rows_csv: + write_csv(args.rows_csv, rows) + print(f"[agg] wrote {args.rows_csv}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/eval/aggregate_cross_lens_en_zh.py b/scripts/eval/aggregate_cross_lens_en_zh.py new file mode 100644 index 0000000..21ca8f7 --- /dev/null +++ b/scripts/eval/aggregate_cross_lens_en_zh.py @@ -0,0 +1,229 @@ +#!/usr/bin/env python3 +"""Aggregate an English-Chinese cross-lens sweep: dominant-feature agreement, null floors, surface split. + +Reads two readout dumps from ``scripts/run/run_cross_lens_readouts.py`` (the same +frozen Qwen3.5-9B states read through two lenses) and the EN-ZH prompt bank, and +computes the numbers behind ``tab:app-cross-lens-families`` and the EN-ZH rows of +``tab:app-cross-lens-extension`` / ``fig:cross-lens-extension``: + + - per-family and pooled cross-lens agreement: the dominant feature (largest + |contribution| in the decomposition of the prompt's answer form ``form_a``) + is the same feature id under both lenses on a majority of the mid-band + layers {21, 24, 26, 29} at the final position, with Wilson 95% CIs; + - within-lens cross-form agreement: ``form_a`` and ``form_b`` share a dominant + feature under the same lens ("one feature carries both surface forms"); + - two null floors: (a) unrelated token, ``form_a`` vs each ``null_targets`` + entry within the slot-A lens; (b) shuffled pairing, prompt i's ``form_a`` + under lens A vs prompt j's ``form_a`` under lens B for three seeded draws of + j with a different concept; + - the lens-only surface split: the script (Latin / CJK) of each lens's top-1 + token by mid-band plurality vote (a 2-2 tie resolves to the alphabetically + first label, CJK < LATIN < OTHER), and the count of prompts read Latin by + lens A and CJK by lens B (the "Lens-only split" column); + - control families (``ctrl_*``) reported separately, never pooled into the + headline. + +Slot semantics: ``--lens-a`` is the English-fitted lens and ``--lens-b`` the +Chinese-fitted lens (output keys ``*_en`` / ``*_zh``). For the construction +comparison the two slots hold the English-fitted Jacobian lens and the +English-fitted ridge translator; the key names are unchanged. + +Paper runs (all with ``--prompts data/cross_lens/cross_lens_prompts_en_zh.json --seed 0``):: + + # main study: extension-matrix row 1 and tab:app-cross-lens-families + aggregate_cross_lens_en_zh.py --lens-a /en_zh__jlens_en_n100.json \ + --lens-b /en_zh__jlens_zh_n100.json --out /en_zh_jlens_n100_summary.json + # ridge translators: row 3 + aggregate_cross_lens_en_zh.py --lens-a /en_zh__ridge_en_n100.json \ + --lens-b /en_zh__ridge_zh_n100.json --out /en_zh_ridge_n100_summary.json + # Jacobian vs ridge on the English corpus: row 4 + aggregate_cross_lens_en_zh.py --lens-a /en_zh__jlens_en_n100.json \ + --lens-b /en_zh__ridge_en_n100.json --out /en_jlens_vs_ridge_summary.json + # 300-prompt refits: row 5 + aggregate_cross_lens_en_zh.py --lens-a /en_zh__jlens_en_n300.json \ + --lens-b /en_zh__jlens_zh_n300.json --out /en_zh_jlens_n300_summary.json +""" + +from __future__ import annotations + +import argparse +import json +import math +import random +import unicodedata +from pathlib import Path + +from sparse_readout_prism.utils import write_csv + +MIDBAND = ["21", "24", "26", "29"] +POS = "-1" +CROSS_GROUPS = ( + "antonym_zh", + "trans_zh2en", + "trans_en2zh", + "cloze_zh", + "cloze_en", + "fact_zh", + "exemplar_zh", +) + + +def wilson(k, n, z=1.96): + if n == 0: + return (0.0, 0.0, 1.0) + p = k / n + d = 1 + z * z / n + c = (p + z * z / (2 * n)) / d + h = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / d + return (p, max(0.0, c - h), min(1.0, c + h)) + + +def script_of(tok): + for ch in tok: + if ch.strip() == "" or not ch.isalnum(): + continue + name = unicodedata.name(ch, "") + if "CJK" in name: + return "CJK" + if "LATIN" in name: + return "LATIN" + return "OTHER" + + +def dom_feat(rec, layer, target): + t = rec["layers"].get(layer, {}).get(POS, {}).get("targets", {}).get(target) + if not t or not t.get("top_features"): + return None + return t["top_features"][0]["id"] + + +def majority_same(rec_a, rec_b, target_a, target_b): + """Majority-of-midband agreement between dom(target_a in rec_a) and dom(target_b in rec_b).""" + same = total = 0 + for L in MIDBAND: + fa, fb = dom_feat(rec_a, L, target_a), dom_feat(rec_b, L, target_b) + if fa is None or fb is None: + continue + total += 1 + same += int(fa == fb) + if total == 0: + return None + return same >= (total + 1) // 2 + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--lens-a", required=True, help="dump under the English-fitted lens (output keys *_en)") + p.add_argument("--lens-b", required=True, help="dump under the Chinese-fitted lens (output keys *_zh)") + p.add_argument("--prompts", required=True, help="EN-ZH prompt bank JSON") + p.add_argument("--seed", type=int, default=0, help="seed for the shuffled-pairing null") + p.add_argument("--out", required=True, help="summary JSON (per-group rates, floors, per-prompt rows)") + p.add_argument("--rows-csv", default=None, help="optional per-prompt rows as CSV") + args = p.parse_args(argv) + + enr = {r["id"]: r for r in json.loads(Path(args.lens_a).read_text())["records"]} + zhr = {r["id"]: r for r in json.loads(Path(args.lens_b).read_text())["records"]} + bank = {v["id"]: v for v in json.loads(Path(args.prompts).read_text())["prompts"]} + ids = sorted(set(enr) & set(zhr) & set(bank)) + print(f"[agg] {len(ids)} prompts present in both dumps") + + rows = [] + for rid in ids: + v, e, z = bank[rid], enr[rid], zhr[rid] + fa, fb = v["form_a"], v["form_b"] + row = {"id": rid, "group": v["group"], "concept": v.get("concept", "")} + # headline: same dominant feature for the SAME concept token, lens A vs lens B + row["cross_lens_pass"] = majority_same(e, z, fa, fa) + # one feature carries both surface forms, within each lens + row["cross_form_en"] = majority_same(e, e, fa, fb) + row["cross_form_zh"] = majority_same(z, z, fa, fb) + # null floor (a): form_a vs unrelated token, within lens A + nulls = [majority_same(e, e, fa, nt) for nt in v.get("null_targets", [])] + nulls = [x for x in nulls if x is not None] + row["null_hits"] = sum(nulls) + row["null_total"] = len(nulls) + # surface script of top-1 under each lens (mid-band vote) + for tag, rec in (("en", e), ("zh", z)): + scripts = [] + for L in MIDBAND: + top5 = rec["layers"].get(L, {}).get(POS, {}).get("top5") + if top5: + scripts.append(script_of(top5[0][0])) + # Plurality vote; a 2-2 tie resolves to the alphabetically first label + # (CJK < LATIN < OTHER), fixed so the vote is deterministic. + row[f"top1_script_{tag}"] = max(sorted(set(scripts)), key=scripts.count) if scripts else "NA" + rows.append(row) + + # null floor (b): shuffled pairing, form_a of prompt i under lens A vs form_a of j under lens B + rng = random.Random(args.seed) + cross_ids = [r["id"] for r in rows if r["group"] in CROSS_GROUPS] + shuffle_hits = shuffle_total = 0 + for rid in cross_ids: + others = [x for x in cross_ids if bank[x]["concept"] != bank[rid]["concept"]] + for oid in rng.sample(others, min(3, len(others))): + res = majority_same(enr[rid], zhr[oid], bank[rid]["form_a"], bank[oid]["form_a"]) + if res is not None: + shuffle_total += 1 + shuffle_hits += int(res) + + def rate(sel): + vals = [r for r in rows if sel(r) and r["cross_lens_pass"] is not None] + k = sum(r["cross_lens_pass"] for r in vals) + return k, len(vals), wilson(k, len(vals)) + + summary: dict = {"per_group": {}, "rows": rows} + print("\n=== CROSS-LENS DOMINANT-FEATURE AGREEMENT (majority of mid-band) ===") + for grp in sorted({r["group"] for r in rows}): + k, n, (pt, lo, hi) = rate(lambda r, g=grp: r["group"] == g) + summary["per_group"][grp] = {"pass": k, "n": n, "rate": pt, "ci": [lo, hi]} + print(f" {grp:14s} {k:3d}/{n:<3d} {pt:.2f} [{lo:.2f}, {hi:.2f}]") + k, n, (pt, lo, hi) = rate(lambda r: r["group"] in CROSS_GROUPS) + summary["headline"] = {"pass": k, "n": n, "rate": pt, "ci": [lo, hi]} + print(f" {'ALL CROSS':14s} {k:3d}/{n:<3d} {pt:.2f} [{lo:.2f}, {hi:.2f}]") + + cf_en = [r["cross_form_en"] for r in rows if r["group"] in CROSS_GROUPS and r["cross_form_en"] is not None] + cf_zh = [r["cross_form_zh"] for r in rows if r["group"] in CROSS_GROUPS and r["cross_form_zh"] is not None] + nk = sum(r["null_hits"] for r in rows) + nn = sum(r["null_total"] for r in rows) + summary["cross_form_en"] = {"pass": sum(cf_en), "n": len(cf_en)} + summary["cross_form_zh"] = {"pass": sum(cf_zh), "n": len(cf_zh)} + summary["null_within_lens"] = {"pass": nk, "n": nn, "rate": wilson(nk, nn)[0]} + summary["null_shuffle_cross_lens"] = { + "pass": shuffle_hits, + "n": shuffle_total, + "rate": wilson(shuffle_hits, shuffle_total)[0], + } + print("\n=== ONE FEATURE CARRIES BOTH FORMS (within-lens) ===") + print(f" EN lens: {sum(cf_en)}/{len(cf_en)} ZH lens: {sum(cf_zh)}/{len(cf_zh)}") + print("\n=== NULL FLOORS ===") + print(f" (a) form vs unrelated token, within-lens: {nk}/{nn} ({wilson(nk, nn)[0]:.2f})") + print( + f" (b) cross-lens shuffled prompts: {shuffle_hits}/{shuffle_total} " + f"({wilson(shuffle_hits, shuffle_total)[0]:.2f})" + ) + + print("\n=== LANGUAGE-FOLLOWS-LENS (top-1 script, mid-band vote) ===") + summary["lens_only_split"] = {} + for grp in sorted({r["group"] for r in rows}): + sel = [r for r in rows if r["group"] == grp] + flip = sum(1 for r in sel if r["top1_script_en"] == "LATIN" and r["top1_script_zh"] == "CJK") + summary["lens_only_split"][grp] = {"pass": flip, "n": len(sel)} + print(f" {grp:14s} EN=Latin & ZH=CJK on {flip}/{len(sel)}") + cross_sel = [r for r in rows if r["group"] in CROSS_GROUPS] + flip = sum(1 for r in cross_sel if r["top1_script_en"] == "LATIN" and r["top1_script_zh"] == "CJK") + summary["lens_only_split"]["all_cross"] = {"pass": flip, "n": len(cross_sel)} + print(f" {'ALL CROSS':14s} EN=Latin & ZH=CJK on {flip}/{len(cross_sel)}") + + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + with out.open("w", encoding="utf-8") as f: + json.dump(summary, f, ensure_ascii=False, indent=1) + print(f"\n[agg] wrote {out}") + if args.rows_csv: + write_csv(args.rows_csv, rows) + print(f"[agg] wrote {args.rows_csv}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/eval/analyze_error_tails.py b/scripts/eval/analyze_error_tails.py new file mode 100644 index 0000000..d314d70 --- /dev/null +++ b/scripts/eval/analyze_error_tails.py @@ -0,0 +1,391 @@ +#!/usr/bin/env python3 +"""Mean and tail statistics of selected-score reconstruction error, per method and model. + +Consumes completed baseline-comparison runs (the ``baseline_query_rows.csv`` +written by ``scripts/run/run_readout_baseline_comparisons.py``, one directory +per model, laid out as ``//``) and +recomputes every per-row statistic from the raw margins rather than trusting +the runner's summary columns: + + absolute_error = |m_exact - m_recon| + rho = absolute_error / (|m_exact| + 0.5) floored relative error rho_0.5 + rho_0 = absolute_error / max(|m_exact|, 1e-6) unfloored rho_0 (the runner's residual_direct) + sign_match = (m_exact > 0) == (m_recon > 0) the runner's convention + sign_match_strict= sign_match with an exact-zero m_exact counted as disagreement + accepted = sign_match_strict and rho < 0.5 + margin_bin = |m_exact| in {<0.5, 0.5-1, 1-2, >=2} the runner's four bins + margin_bin_fine = |m_exact| in {<0.5, 0.5-1, 1-2, 2-5, >=5} + +``margin_bin`` and ``sign_match`` are recomputed with definitions identical to +the runner's input columns of the same name; ``rho_0`` equals ``residual_direct``. + +Outputs (CSV/JSON only): + + summary_by_method.csv primary-model rows, per method: mean/median/p90/p95/p99/max of + rho and absolute_error, sign agreement, rho<0.5, accepted rate, + cluster-bootstrap 95% CIs (clusters = base_case_id) + summary_by_model.csv SRP rows per model, same columns and CIs + summary_by_margin_bin.csv SRP rows per (model, margin_bin_fine), no CIs + summary_by_family.csv SRP rows per (model, family), no CIs + paired_method_differences.csv primary model, SRP minus each baseline on the same cases + (mean rho, p95 rho, sign agreement, accepted), paired cluster bootstrap + tail_rows.csv primary-model SRP rows in the top 1% by rho, joined to the tokenizer audit + summary_pooled_by_method.csv all models pooled, per method: mean/p95/max absolute error, median + rho_0, sign agreement, coverage (sign_match and rho_0<0.5) + summary_by_model_method.csv the same six statistics per (model, method) + summary_pooled_by_method_margin_bin.csv + the same statistics plus bin share, per (method, margin_bin), pooled + summary.json provenance, input hashes, row counts, tail threshold + +Paper artifacts. ``tab:app-error-tails`` is ``summary_pooled_by_method.csv`` +(8,009 contrasts per method over the six softcap-free readouts) and +``tab:app-error-tail-margins`` is the ``sparse_rp`` block of +``summary_pooled_by_method_margin_bin.csv``; the per-model ranges quoted in the +same appendix come from ``summary_by_model_method.csv``. The floored-rho +summaries with bootstrap intervals are the distribution-wide companion to +``tab:app-direct-geometry-grid`` on the primary model. + +Paper run (inputs are the six direct-geometry runs of the baseline script, one +directory per model tag qwen0p8b, qwen2b, qwen9b, ministral8b, r1qwen7b, +r1llama8b; ``--input-prefix`` names any common directory prefix, e.g. +``run_`` for ``run_qwen2b/``): + + uv run python scripts/eval/analyze_error_tails.py \\ + --input-root results/direct_geometry_runs \\ + --out-dir results/error_tails --n-boot 10000 --seed 20260711 +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import platform +import sys +from pathlib import Path + +import numpy as np +import pandas as pd + +from sparse_readout_prism.research.run_io import run_provenance + +MODEL_TAGS = ( + "qwen0p8b", + "qwen2b", + "qwen9b", + "ministral8b", + "r1qwen7b", + "r1llama8b", +) +MARGIN_BINS_FINE = [-np.inf, 0.5, 1.0, 2.0, 5.0, np.inf] +MARGIN_LABELS_FINE = ["<0.5", "0.5-1", "1-2", "2-5", ">=5"] +MARGIN_BINS = [-np.inf, 0.5, 1.0, 2.0, np.inf] +MARGIN_LABELS = ["<0.5", "0.5-1", "1-2", ">=2"] +ANCHOR = "sparse_rp" + + +def sha256(path: Path) -> str: + h = hashlib.sha256() + with path.open("rb") as f: + for chunk in iter(lambda: f.read(1024 * 1024), b""): + h.update(chunk) + return h.hexdigest() + + +def add_row_metrics(rows: pd.DataFrame) -> pd.DataFrame: + """Recompute the per-row error statistics from ``exact_margin`` / ``sparse_margin``.""" + required = { + "model_tag", + "method", + "base_case_id", + "case_id", + "family", + "exact_margin", + "sparse_margin", + } + missing = required - set(rows.columns) + if missing: + raise ValueError(f"raw rows missing columns: {sorted(missing)}") + exact = pd.to_numeric(rows["exact_margin"], errors="raise").to_numpy(float) + recon = pd.to_numeric(rows["sparse_margin"], errors="raise").to_numpy(float) + err = np.abs(exact - recon) + rows["absolute_error"] = err + rows["rho"] = err / (np.abs(exact) + 0.5) + rows["rho_0"] = err / np.maximum(np.abs(exact), 1e-6) + rows["exact_zero"] = np.isclose(exact, 0.0, atol=1e-12) + rows["near_zero"] = np.abs(exact) < 0.5 + rows["sign_match"] = (exact > 0) == (recon > 0) + rows["sign_match_strict"] = (~rows["exact_zero"]) & (np.signbit(exact) == np.signbit(recon)) + rows["accepted"] = rows["sign_match_strict"] & (rows["rho"] < 0.5) + rows["covered"] = rows["sign_match"] & (rows["rho_0"] < 0.5) + rows["margin_bin"] = pd.cut(np.abs(exact), bins=MARGIN_BINS, labels=MARGIN_LABELS, right=False).astype(str) + rows["margin_bin_fine"] = pd.cut( + np.abs(exact), bins=MARGIN_BINS_FINE, labels=MARGIN_LABELS_FINE, right=False + ).astype(str) + return rows + + +def load_rows(input_root: Path, prefix: str, model_tags: list[str]) -> tuple[pd.DataFrame, list[dict]]: + frames: list[pd.DataFrame] = [] + inputs: list[dict] = [] + for tag in model_tags: + p = input_root / f"{prefix}{tag}" / "baseline_query_rows.csv" + if not p.exists(): + continue + df = pd.read_csv(p) + df["model_tag"] = tag + frames.append(df) + inputs.append({"path": str(p.resolve()), "sha256": sha256(p), "rows": len(df)}) + if not frames: + raise FileNotFoundError(f"no {prefix}*/baseline_query_rows.csv under {input_root}") + rows = pd.concat(frames, ignore_index=True) + return add_row_metrics(rows), inputs + + +def point_metrics(df: pd.DataFrame) -> dict[str, float | int]: + rho = df["rho"].to_numpy(float) + ae = df["absolute_error"].to_numpy(float) + out: dict[str, float | int] = {"n_rows": int(len(df))} + for prefix, arr in (("rho", rho), ("absolute_error", ae)): + out[f"{prefix}_mean"] = float(np.mean(arr)) + out[f"{prefix}_median"] = float(np.median(arr)) + for q in (90, 95, 99): + out[f"{prefix}_p{q}"] = float(np.percentile(arr, q)) + out[f"{prefix}_max"] = float(np.max(arr)) + out["sign_agreement"] = float(df["sign_match_strict"].mean()) + out["rho_lt_0p5"] = float((df["rho"] < 0.5).mean()) + out["accepted_rate"] = float(df["accepted"].mean()) + out["exact_zero_n"] = int(df["exact_zero"].sum()) + out["near_zero_n"] = int(df["near_zero"].sum()) + return out + + +def bootstrap_ci(df: pd.DataFrame, n_boot: int, seed: int) -> dict[str, float]: + clusters = {str(k): idx.to_numpy() for k, idx in df.groupby("base_case_id", sort=True).groups.items()} + keys = sorted(clusters) + if len(keys) < 2 or n_boot <= 0: + return {} + rng = np.random.default_rng(seed) + rho = df["rho"].to_numpy(float) + sign = df["sign_match_strict"].to_numpy(float) + accepted = df["accepted"].to_numpy(float) + stats = np.empty((n_boot, 5), dtype=float) + for b in range(n_boot): + sampled = rng.integers(0, len(keys), size=len(keys)) + idx = np.concatenate([clusters[keys[i]] for i in sampled]) + rb = rho[idx] + stats[b] = ( + rb.mean(), + np.median(rb), + np.percentile(rb, 95), + sign[idx].mean(), + accepted[idx].mean(), + ) + names = ("rho_mean", "rho_median", "rho_p95", "sign_agreement", "accepted_rate") + out: dict[str, float] = {} + for j, name in enumerate(names): + out[f"{name}_ci_lo"] = float(np.percentile(stats[:, j], 2.5)) + out[f"{name}_ci_hi"] = float(np.percentile(stats[:, j], 97.5)) + return out + + +def summarize( + df: pd.DataFrame, + group_cols: list[str], + *, + bootstrap: bool, + n_boot: int, + seed: int, +) -> pd.DataFrame: + records: list[dict] = [] + grouper = group_cols[0] if len(group_cols) == 1 else group_cols + for key, part in df.groupby(grouper, sort=True, observed=True): + key = (key,) if not isinstance(key, tuple) else key + rec = dict(zip(group_cols, key)) + rec.update(point_metrics(part)) + if bootstrap: + group_seed = seed + int(hashlib.sha1(str(key).encode()).hexdigest()[:7], 16) + rec.update(bootstrap_ci(part.reset_index(drop=True), n_boot, group_seed)) + records.append(rec) + return pd.DataFrame(records) + + +def paired_differences(qdf: pd.DataFrame, n_boot: int, seed: int) -> pd.DataFrame: + keys = ["bank", "base_case_id", "case_id", "query"] + cols = keys + ["rho", "sign_match_strict", "accepted"] + a = qdf[qdf.method == ANCHOR][cols].copy() + results: list[dict] = [] + for method in sorted(set(qdf.method) - {ANCHOR}): + b = qdf[qdf.method == method][cols].copy() + m = a.merge(b, on=keys, suffixes=("_srp", "_base"), validate="one_to_one") + if m.empty: + continue + clusters = {str(k): idx.to_numpy() for k, idx in m.groupby("base_case_id", sort=True).groups.items()} + ckeys = sorted(clusters) + rng = np.random.default_rng(seed + int(hashlib.sha1(method.encode()).hexdigest()[:7], 16)) + + def diffs(frame: pd.DataFrame) -> np.ndarray: + return np.array( + [ + frame.rho_srp.mean() - frame.rho_base.mean(), + np.percentile(frame.rho_srp, 95) - np.percentile(frame.rho_base, 95), + frame.sign_match_strict_srp.mean() - frame.sign_match_strict_base.mean(), + frame.accepted_srp.mean() - frame.accepted_base.mean(), + ] + ) + + point = diffs(m) + boots = np.empty((n_boot, 4), float) + for i in range(n_boot): + sample = rng.integers(0, len(ckeys), size=len(ckeys)) + idx = np.concatenate([clusters[ckeys[j]] for j in sample]) + boots[i] = diffs(m.iloc[idx]) + rec: dict[str, float | str | int] = {"baseline": method, "n_pairs": len(m)} + for j, name in enumerate(("mean_rho", "p95_rho", "sign_agreement", "accepted_rate")): + rec[f"delta_srp_minus_baseline_{name}"] = float(point[j]) + rec[f"{name}_ci_lo"] = float(np.percentile(boots[:, j], 2.5)) + rec[f"{name}_ci_hi"] = float(np.percentile(boots[:, j], 97.5)) + results.append(rec) + return pd.DataFrame(results) + + +def attach_token_audit(tail: pd.DataFrame, audit_path: Path) -> pd.DataFrame: + if not audit_path.exists() or tail.empty: + return tail + audit = pd.read_csv(audit_path) + if "case_id" not in audit.columns: + return tail + keep = [c for c in ("token_id", "reason", "side", "target") if c in audit.columns] + grouped = audit.groupby("case_id", sort=False)[keep].agg(lambda s: ";".join(dict.fromkeys(map(str, s.dropna())))) + grouped = grouped.add_prefix("token_audit_").reset_index() + return tail.merge(grouped, on="case_id", how="left") + + +# --------------------------------------------------------------------------- # +# Paper-table summaries: pooled over models, unfloored rho_0, runner conventions +# --------------------------------------------------------------------------- # + + +def coverage_metrics(df: pd.DataFrame) -> dict[str, float | int]: + """Mean/p95/max absolute error, median rho_0, sign agreement and coverage, + the six statistics of the error-tail tables (unfloored rho_0, the runner's + sign convention).""" + ae = df["absolute_error"].to_numpy(float) + return { + "n_rows": int(len(df)), + "n_models": int(df["model_tag"].nunique()), + "absolute_error_mean": float(np.mean(ae)), + "absolute_error_p95": float(np.percentile(ae, 95)), + "absolute_error_max": float(np.max(ae)), + "rho_0_median": float(df["rho_0"].median()), + "sign_agreement": float(df["sign_match"].mean()), + "coverage": float(df["covered"].mean()), + } + + +def summarize_coverage(df: pd.DataFrame, group_cols: list[str], *, share_within: str | None = None) -> pd.DataFrame: + """``coverage_metrics`` per group. ``share_within`` names a column whose + groups are the denominator of a ``share`` column (bin share within a method).""" + records: list[dict] = [] + grouper = group_cols[0] if len(group_cols) == 1 else group_cols + totals = df.groupby(share_within, sort=True).size() if share_within else None + for key, part in df.groupby(grouper, sort=True, observed=True): + key = (key,) if not isinstance(key, tuple) else key + rec = dict(zip(group_cols, key)) + if totals is not None: + rec["share"] = float(len(part) / totals[rec[share_within]]) + rec.update(coverage_metrics(part)) + records.append(rec) + out = pd.DataFrame(records) + if "margin_bin" in group_cols and not out.empty: + order = {lbl: i for i, lbl in enumerate(MARGIN_LABELS)} + out = out.sort_values( + [c for c in group_cols if c != "margin_bin"] + ["margin_bin"], + key=lambda s: s.map(order) if s.name == "margin_bin" else s, + ) + out = out.reset_index(drop=True) + return out + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument( + "--input-root", type=Path, required=True, help="directory holding /baseline_query_rows.csv" + ) + ap.add_argument("--input-prefix", type=str, default="", help="common prefix of the per-model run directories") + ap.add_argument( + "--model-tags", + type=str, + default=",".join(MODEL_TAGS), + help="comma-separated model tags; missing directories are skipped", + ) + ap.add_argument( + "--primary-model-tag", + type=str, + default="qwen2b", + help="model for the per-method / paired / tail analyses (the paper's Qwen3.5-2B)", + ) + ap.add_argument("--out-dir", type=Path, required=True) + ap.add_argument("--n-boot", type=int, default=10_000) + ap.add_argument("--seed", type=int, default=20260711) + args = ap.parse_args() + args.out_dir.mkdir(parents=True, exist_ok=True) + + model_tags = [t.strip() for t in args.model_tags.split(",") if t.strip()] + rows, inputs = load_rows(args.input_root, args.input_prefix, model_tags) + qdf = rows[rows.model_tag == args.primary_model_tag].copy() + if qdf.empty: + raise ValueError(f"rows for the primary model tag {args.primary_model_tag!r} are required") + + by_method = summarize(qdf, ["method"], bootstrap=True, n_boot=args.n_boot, seed=args.seed) + srp_rows = rows[rows.method == ANCHOR] + by_model = summarize(srp_rows, ["model_tag"], bootstrap=True, n_boot=args.n_boot, seed=args.seed) + by_margin = summarize(srp_rows, ["model_tag", "margin_bin_fine"], bootstrap=False, n_boot=0, seed=args.seed) + by_family = summarize(srp_rows, ["model_tag", "family"], bootstrap=False, n_boot=0, seed=args.seed) + paired = paired_differences(qdf, args.n_boot, args.seed) + + srp_q = qdf[qdf.method == ANCHOR].copy() + threshold = float(srp_q.rho.quantile(0.99)) + tail = srp_q[srp_q.rho >= threshold].copy().sort_values("rho", ascending=False) + tail["tail_definition"] = f"{args.primary_model_tag}_srp_top1pct_rho_ge_{threshold:.9g}" + tail = attach_token_audit( + tail, args.input_root / f"{args.input_prefix}{args.primary_model_tag}" / "tokenizer_audit.csv" + ) + + pooled_by_method = summarize_coverage(rows, ["method"]) + by_model_method = summarize_coverage(rows, ["model_tag", "method"]) + pooled_by_method_bin = summarize_coverage(rows, ["method", "margin_bin"], share_within="method") + + by_method.to_csv(args.out_dir / "summary_by_method.csv", index=False) + by_model.to_csv(args.out_dir / "summary_by_model.csv", index=False) + by_margin.to_csv(args.out_dir / "summary_by_margin_bin.csv", index=False) + by_family.to_csv(args.out_dir / "summary_by_family.csv", index=False) + paired.to_csv(args.out_dir / "paired_method_differences.csv", index=False) + tail.to_csv(args.out_dir / "tail_rows.csv", index=False) + pooled_by_method.to_csv(args.out_dir / "summary_pooled_by_method.csv", index=False) + by_model_method.to_csv(args.out_dir / "summary_by_model_method.csv", index=False) + pooled_by_method_bin.to_csv(args.out_dir / "summary_pooled_by_method_margin_bin.csv", index=False) + + manifest = { + **run_provenance(args), + "python": sys.version, + "platform": platform.platform(), + "seed": args.seed, + "n_boot": args.n_boot, + "rho_definition": "abs(exact-reconstructed)/(abs(exact)+0.5)", + "rho_0_definition": "abs(exact-reconstructed)/max(abs(exact),1e-6)", + "coverage_definition": "sign_match and rho_0 < 0.5", + "inputs": inputs, + "model_tags": sorted(rows.model_tag.unique()), + "rows_total": len(rows), + "primary_model_tag": args.primary_model_tag, + "primary_model_rows": len(qdf), + "tail_threshold": threshold, + } + (args.out_dir / "summary.json").write_text(json.dumps(manifest, indent=2, default=str)) + print(json.dumps(manifest, indent=2, default=str), flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/eval/cross_seed_stability.py b/scripts/eval/cross_seed_stability.py new file mode 100644 index 0000000..2612916 --- /dev/null +++ b/scripts/eval/cross_seed_stability.py @@ -0,0 +1,247 @@ +#!/usr/bin/env python3 +"""Cross-seed stability of contrast explanations (Appendix I, ``tab:app-cross-seed-stability``). + +Given several dictionaries trained from one recipe with different seeds (same +width, sparsity and preprocessing), decompose every curated single-token +contrast ``q = w_A - w_B`` with each dictionary and ask whether the same token +families land on the same side of the contrast across seeds. For each seed a +contrast yields row-side coefficients ``beta_i = rn_A z_A,i - rn_B z_B,i``; the +top-M positive and top-M negative features summarise each side by the token +set of their top-R centred unembedding rows. + +Computes exactly five quantities: + + same-side Jaccard J(T_pos^s1, T_pos^s2), J(T_neg^s1, T_neg^s2) per seed pair and contrast + cross-side leakage J(T_pos^s1, T_neg^s2) should be ~0 + cross-contrast null J(T_pos^s1(c), T_pos^s2(c')), c' != c random null anchor + matched decoder cosine nearest-neighbour decoder cosine between seeds, for a + random feature sample and for the features the contrasts use + matched-projection corr correlation between projections of held-out hidden states + (``h_LN`` in the extraction payload) onto each used feature + and onto its nearest decoder match in the other seed + +plus the fraction of contrasts whose mean same-side Jaccard exceeds the +cross-contrast null's 90th percentile. Everything is local to the extraction +payload and the checkpoints; no model forward pass. + +Dictionary-family discipline: pass dictionaries from ONE recipe (the +seed-variation family trained from ``configs/sweeps/qwen35_2b_seedvar_base.yaml``). +The script refuses dictionaries that differ in width or ``k``. Cross-recipe +agreement is weaker and is reported separately (Appendix I). + +Paper runs (Qwen3.5-2B, three seeds per width, curated A/B bank, all defaults): + + uv run python scripts/eval/cross_seed_stability.py \\ + --dictionaries results/qwen35_2b_seedvar/qwen2b_d65536_k256_s0/checkpoint.pt \\ + results/qwen35_2b_seedvar/qwen2b_d65536_k256_s1/checkpoint.pt \\ + results/qwen35_2b_seedvar/qwen2b_d65536_k256_s2/checkpoint.pt \\ + --w-u "$SRP_ARCHIVE_ROOT/data/qwen35-2b/qwen35_2b.pt" \\ + --bank data/query_banks/qwen_gemma_result1_curated_ab.jsonl \\ + --width-tag 32x --out results/seed_stability/cross_seed_stability_32x.json + + (16x: the ``qwen2b_d32768_k128_s{0,1,2}`` checkpoints, ``--width-tag 16x``.) + +Under the Qwen3.5 tokenizer the 150-contrast cap resolves to 63 unique +single-token pairs of the curated bank (126 signed side-sets). +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np +import torch + +from sparse_readout_prism.research.seed_stability import ( + center_rows, + contrast_features, + feature_token_set, + jaccard, + load_contrast_pairs, + load_dictionary, + load_readout, +) +from sparse_readout_prism.utils import to_jsonable + + +def _stats(v) -> dict: + v = np.asarray(v, dtype=float) + return dict(mean=float(v.mean()), median=float(np.median(v)), p10=float(np.percentile(v, 10)), n=int(len(v))) + + +def _unit(d: torch.Tensor) -> torch.Tensor: + return d / d.norm(dim=1, keepdim=True).clamp_min(1e-8) + + +def run( + dicts: list, + labels: list[int], + W: torch.Tensor, + h_LN: torch.Tensor | None, + tok, + pairs: list[tuple], + rng: np.random.Generator, + *, + top_m: int, + top_r: int, + n_sample: int, + n_hidden: int, + width_tag: str | None, +) -> dict: + W_c, rn, W_n = center_rows(W) + n_seeds = len(dicts) + cache: dict[int, str] = {} + + # per-seed per-contrast side token-sets + used features + side_sets: list[list[tuple[set, set]]] = [[] for _ in dicts] # [seed][ci] = (T_pos, T_neg) + used_feats: list[set[int]] = [set() for _ in dicts] + for si, d in enumerate(dicts): + dec = d[0] + for _a, _b, ia, ib in pairs: + top_pos, top_neg = contrast_features(W_n, rn, d, ia, ib, top_m) + used_feats[si].update(top_pos + top_neg) + side_sets[si].append( + ( + feature_token_set(top_pos, dec, W_c, tok, top_r, cache), + feature_token_set(top_neg, dec, W_c, tok, top_r, cache), + ) + ) + print(f"seed {labels[si]}: decomposed {len(pairs)} contrasts", flush=True) + + # explanation-level reproduction + same_side, cross_side, cross_contrast = [], [], [] + for s1 in range(n_seeds): + for s2 in range(s1 + 1, n_seeds): + for ci in range(len(pairs)): + p1, n1 = side_sets[s1][ci] + p2, n2 = side_sets[s2][ci] + same_side.append(jaccard(p1, p2)) + same_side.append(jaccard(n1, n2)) + cross_side.append(jaccard(p1, n2)) + cj = int(rng.integers(len(pairs) - 1)) + cj = cj + 1 if cj >= ci else cj + cross_contrast.append(jaccard(p1, side_sets[s2][cj][0])) + + # fraction of contrasts whose same-side overlap beats the cross-contrast null p90 + null90 = float(np.percentile(cross_contrast, 90)) + per_contrast = np.asarray(same_side).reshape(-1, 2).mean(axis=1) # mean of pos/neg per (pair, ci) + frac_above_null = float((per_contrast > null90).mean()) + + # basis-level nearest-neighbour cosine + basis = {} + for s1 in range(n_seeds): + for s2 in range(s1 + 1, n_seeds): + d1, d2 = dicts[s1][0], dicts[s2][0] + d1n, d2n = _unit(d1), _unit(d2) + samp = torch.tensor(rng.choice(d1.shape[0], min(n_sample, d1.shape[0]), replace=False)) + nn_rand = (d1n[samp] @ d2n.T).max(dim=1).values + used = torch.tensor(sorted(used_feats[s1])) + nn_used = (d1n[used] @ d2n.T).max(dim=1).values + basis[f"s{labels[s1]}-s{labels[s2]}"] = dict( + nn_cos_random=_stats(nn_rand.tolist()), nn_cos_used=_stats(nn_used.tolist()) + ) + + # projection correlation on held-out hidden states + proj = {} + if h_LN is not None: + H = h_LN.reshape(-1, h_LN.shape[-1]).float() # (n_states, d_model) + H = H[torch.tensor(rng.choice(H.shape[0], min(n_hidden, H.shape[0]), replace=False))] + for s1 in range(n_seeds): + for s2 in range(s1 + 1, n_seeds): + d1, d2 = dicts[s1][0], dicts[s2][0] + used = torch.tensor(sorted(used_feats[s1])) + match = (_unit(d1[used]) @ _unit(d2).T).argmax(dim=1) + p1 = H @ d1[used].T + p2 = H @ d2[match].T + r = torch.corrcoef(torch.stack([p1.flatten(), p2.flatten()]))[0, 1] + proj[f"s{labels[s1]}-s{labels[s2]}"] = float(r) + + return dict( + width=width_tag, + k=int(dicts[0][3]), + n_contrasts=len(pairs), + top_m=top_m, + top_r=top_r, + same_side=_stats(same_side), + cross_side=_stats(cross_side), + cross_contrast_null=_stats(cross_contrast), + null_p90=null90, + frac_contrasts_above_null_p90=frac_above_null, + basis_nn_cosine=basis, + matched_projection_corr=proj, + ) + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument( + "--dictionaries", + nargs="+", + type=Path, + required=True, + help="checkpoint.pt of each seed, one recipe (paper: seeds 0, 1, 2 of one width)", + ) + ap.add_argument("--seed-labels", nargs="+", type=int, default=None, help="labels for output keys (default 0..n-1)") + ap.add_argument("--w-u", type=Path, required=True, help="extraction payload {W_U_orig, h_LN, ...}") + ap.add_argument("--bank", type=Path, required=True, help="curated A/B JSONL bank (target_a / target_b)") + ap.add_argument("--tokenizer", default="Qwen/Qwen3.5-2B", help="HF tokenizer id or local path") + ap.add_argument("--width-tag", default=None, help='label stored in the output, e.g. "32x"') + ap.add_argument("--max-contrasts", type=int, default=150) + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--top-m", type=int, default=8, help="features per side") + ap.add_argument("--top-r", type=int, default=12, help="rows per feature for the token summary") + ap.add_argument("--n-sample", type=int, default=4096, help="random features for the basis-level NN cosine") + ap.add_argument("--n-hidden", type=int, default=512, help="held-out hidden states for the projection check") + ap.add_argument("--out", type=Path, required=True, help="output JSON") + args = ap.parse_args() + + if len(args.dictionaries) < 2: + ap.error("--dictionaries needs at least two checkpoints") + labels = args.seed_labels if args.seed_labels is not None else list(range(len(args.dictionaries))) + if len(labels) != len(args.dictionaries): + ap.error("--seed-labels must have one entry per dictionary") + + rng = np.random.default_rng(args.seed) + W, h_LN = load_readout(args.w_u) + from transformers import AutoTokenizer + + tok = AutoTokenizer.from_pretrained(args.tokenizer) + + dicts = [load_dictionary(p) for p in args.dictionaries] + for lab, d in zip(labels, dicts): + print(f"seed {lab}: D={d[0].shape[0]} k={d[3]}") + if len({(d[0].shape[0], d[3]) for d in dicts}) != 1: + raise SystemExit("dictionaries differ in width or k; pass checkpoints from one dictionary family") + + pairs = load_contrast_pairs(args.bank, tok, rng, args.max_contrasts) + print(f"contrasts: {len(pairs)} unique single-token pairs") + + out = run( + dicts, + labels, + W, + h_LN, + tok, + pairs, + rng, + top_m=args.top_m, + top_r=args.top_r, + n_sample=args.n_sample, + n_hidden=args.n_hidden, + width_tag=args.width_tag, + ) + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(to_jsonable(out), indent=2) + "\n") + print(json.dumps({k: v for k, v in out.items() if k != "basis_nn_cosine"}, indent=2)) + print( + "basis NN-cosine (used features):", + {k: round(v["nn_cos_used"]["median"], 3) for k, v in out["basis_nn_cosine"].items()}, + ) + print(f"wrote {args.out}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/eval/feature_group_matching.py b/scripts/eval/feature_group_matching.py new file mode 100644 index 0000000..fb7363e --- /dev/null +++ b/scripts/eval/feature_group_matching.py @@ -0,0 +1,367 @@ +#!/usr/bin/env python3 +"""Cross-seed feature-group matching (Appendix I, ``tab:app-feature-group-matching``). + +Contrast-level explanations reproduce across seeds while feature indices +reshuffle. This script asks whether the feature-level token GROUPS themselves +have an equivalent in another seed's dictionary, i.e. whether the grouping +structure is re-indexed rather than seed noise (cf. Paulo & Belrose 2025, +arXiv:2501.16615). + +Protocol + * Query population: features USED by the curated contrast bank in the query + seed (top-M positive + top-M negative per contrast, the same construction + as ``cross_seed_stability.py``); ``--n-query`` are sampled per seed pair + among those with at least ``--min-group-size`` tokens. + * Group of a feature = token strings of its top-R centred unembedding rows + by decoder similarity (identical to the token summary used elsewhere). + * Match: best-of-D search over ALL candidate-seed features by Jaccard of + token groups. Secondary: greedy union of up to ``--greedy-k`` candidate + groups drawn from the ``--greedy-pool`` best single matches (feature + splitting), scored by recall of the query group. + * Null: ``--n-null`` pseudo-groups of top-R tokens drawn from the empirical + token-frequency distribution over the candidate seed's groups, run through + the IDENTICAL best-of-D search, which prices in the inflation of searching + a 65,536-entry dictionary. + * Criteria, fixed before the analysis: a query group reproduces if its best + single Jaccard exceeds the null's 99th percentile; strong equivalence is a + best single Jaccard of at least 0.5. Full distributions are stored either way. + +``--direction-check`` adds the direction-level check of the below-null tail: +for every query group at or below the null p99, the maximum decoder cosine +over the full candidate dictionary, against a null of ``--direction-null`` +random unit directions passed through the same maximum. It writes +``max_decoder_cosine`` into those per-query records and a ``direction_check`` +summary per seed pair (counts above the direction null p99 and at cosine >= 0.5). + +Seed pairs are all ordered pairs (query, candidate) with query index below +candidate index, so three dictionaries give s0->s1, s0->s2, s1->s2. + +Dictionary-family discipline: pass dictionaries from ONE recipe (the +seed-variation family trained from ``configs/sweeps/qwen35_2b_seedvar_base.yaml``); +the script refuses dictionaries that differ in width or ``k``. + +Paper runs (Qwen3.5-2B, three seeds per width, curated A/B bank, all defaults): + + uv run python scripts/eval/feature_group_matching.py \\ + --dictionaries results/qwen35_2b_seedvar/qwen2b_d65536_k256_s0/checkpoint.pt \\ + results/qwen35_2b_seedvar/qwen2b_d65536_k256_s1/checkpoint.pt \\ + results/qwen35_2b_seedvar/qwen2b_d65536_k256_s2/checkpoint.pt \\ + --w-u "$SRP_ARCHIVE_ROOT/data/qwen35-2b/qwen35_2b.pt" \\ + --bank data/query_banks/qwen_gemma_result1_curated_ab.jsonl \\ + --width-tag 32x --direction-check \\ + --out results/seed_stability/feature_group_matching_32x.json + + (16x: the ``qwen2b_d32768_k128_s{0,1,2}`` checkpoints, ``--width-tag 16x``.) +""" + +from __future__ import annotations + +import argparse +import itertools +import json +from collections import Counter +from pathlib import Path + +import numpy as np +import torch + +from sparse_readout_prism.research.seed_stability import ( + center_rows, + contrast_features, + load_contrast_pairs, + load_dictionary, + load_readout, + token_strings, +) +from sparse_readout_prism.utils import resolve_device, to_jsonable + + +def all_feature_groups(dec: torch.Tensor, W_c: torch.Tensor, tok, device, cache: dict, top_r: int) -> list[frozenset]: + """Token group (frozenset of strings) of every feature: its top-R centred rows.""" + D = dec.shape[0] + Wt = W_c.T.to(device) # (d_model, vocab) + top_rows = torch.empty((D, top_r), dtype=torch.long) + for lo in range(0, D, 2048): + chunk = dec[lo : lo + 2048].to(device) + sims = chunk @ Wt # (chunk, vocab) + top_rows[lo : lo + 2048] = sims.topk(top_r, dim=1).indices.cpu() + del sims, chunk + return [frozenset(token_strings(tok, top_rows[f].tolist(), cache)) for f in range(D)] + + +def best_matches(query: frozenset, inv_index: dict, groups: list[frozenset], greedy_pool: int, greedy_k: int): + """``(best_jaccard, best_feature, greedy recall@1..greedy_k)`` over all candidate groups.""" + counts: Counter = Counter() + for t in query: + for f in inv_index.get(t, ()): + counts[f] += 1 + if not counts: + return 0.0, -1, [0.0] * greedy_k + scored = sorted( + ((inter / (len(query) + len(groups[f]) - inter), f) for f, inter in counts.items()), + reverse=True, + ) + best_j, best_f = scored[0] + pool = [f for _, f in scored[:greedy_pool]] + covered: set = set() + recalls = [] + for _ in range(greedy_k): + gain_best, f_best = 0, None + for f in pool: + gain = len((query & groups[f]) - covered) + if gain > gain_best: + gain_best, f_best = gain, f + if f_best is not None: + covered |= query & groups[f_best] + pool.remove(f_best) + recalls.append(len(covered) / max(1, len(query))) + return best_j, best_f, recalls + + +def _stats(v) -> dict: + v = np.asarray(v, dtype=float) + return dict( + mean=float(v.mean()), + median=float(np.median(v)), + p90=float(np.percentile(v, 90)), + p99=float(np.percentile(v, 99)), + n=int(len(v)), + ) + + +def _unit(d: torch.Tensor) -> torch.Tensor: + return d / d.norm(dim=1, keepdim=True).clamp_min(1e-8) + + +def run( + dicts: list, + labels: list[int], + W: torch.Tensor, + tok, + pairs: list[tuple], + rng: np.random.Generator, + device, + *, + top_m: int, + top_r: int, + n_query: int, + n_null: int, + greedy_pool: int, + greedy_k: int, + min_group_size: int, + width_tag: str | None, + direction_check: bool, + direction_null: int, + direction_seed: int, +) -> dict: + W_c, rn, W_n = center_rows(W) + cache: dict[int, str] = {} + + # token groups of every feature, every seed + groups = [all_feature_groups(d[0], W_c, tok, device, cache, top_r) for d in dicts] + print("feature groups computed", flush=True) + + # used features per seed (same construction as cross_seed_stability.py) + used: list[set[int]] = [set() for _ in dicts] + for si, d in enumerate(dicts): + for _a, _b, ia, ib in pairs: + top_pos, top_neg = contrast_features(W_n, rn, d, ia, ib, top_m) + used[si].update(top_pos) + used[si].update(top_neg) + print("used features:", [len(u) for u in used]) + + seed_pairs = list(itertools.combinations(range(len(dicts)), 2)) # (query seed, candidate seed) + unit_decs = [_unit(d[0]) for d in dicts] + results: dict[str, dict] = {} + for qs, cs in seed_pairs: + cand_groups = groups[cs] + inv: dict[str, list[int]] = {} + for f, g in enumerate(cand_groups): + for t in g: + inv.setdefault(t, []).append(f) + + # frequency-matched null token pool + tok_freq = Counter(t for g in cand_groups for t in g) + pool_toks = list(tok_freq) + pool_p = np.array([tok_freq[t] for t in pool_toks], dtype=float) + pool_p /= pool_p.sum() + + # queries: sampled used features of the query seed with large-enough groups + cands = [f for f in sorted(used[qs]) if len(groups[qs][f]) >= min_group_size] + qs_feats = rng.choice(cands, size=min(n_query, len(cands)), replace=False) + + q_best, q_rec, q_cos, q_records = [], [], [], [] + dq, dc = unit_decs[qs], unit_decs[cs] + for f in qs_feats: + f = int(f) + bj, bf, recalls = best_matches(groups[qs][f], inv, cand_groups, greedy_pool, greedy_k) + q_best.append(bj) + q_rec.append(recalls) + q_cos.append(float(dq[f] @ dc[bf]) if bf >= 0 else 0.0) + q_records.append( + dict( + feature=f, + query_tokens=sorted(groups[qs][f]), + best_jaccard=round(bj, 4), + best_feature=int(bf), + best_match_tokens=sorted(cand_groups[bf]) if bf >= 0 else [], + recall3=round(recalls[-1], 4), + cosine=round(q_cos[-1], 4), + ) + ) + + n_best, n_rec = [], [] + for _ in range(n_null): + g = frozenset(rng.choice(pool_toks, size=top_r, replace=False, p=pool_p)) + bj, _, recalls = best_matches(g, inv, cand_groups, greedy_pool, greedy_k) + n_best.append(bj) + n_rec.append(recalls) + + null99 = float(np.percentile(n_best, 99)) + rec3 = [r[-1] for r in q_rec] + n_rec3 = [r[-1] for r in n_rec] + res = dict( + n_query=int(len(qs_feats)), + best_jaccard=_stats(q_best), + null_best_jaccard=_stats(n_best), + frac_above_null_p99=float((np.array(q_best) > null99).mean()), + frac_jaccard_ge_050=float((np.array(q_best) >= 0.5).mean()), + frac_jaccard_ge_025=float((np.array(q_best) >= 0.25).mean()), + recall_at_1=_stats([r[0] for r in q_rec]), + recall_at_3=_stats(rec3), + null_recall_at_3=_stats(n_rec3), + frac_recall3_above_null_p99=float((np.array(rec3) > np.percentile(n_rec3, 99)).mean()), + matched_decoder_cosine=_stats(q_cos), + null99=null99, + per_query=q_records, + ) + key = f"s{labels[qs]}->s{labels[cs]}" + results[key] = res + print( + f"{key}: best-J med {res['best_jaccard']['median']:.3f} " + f"(null p99 {null99:.3f}), >null {res['frac_above_null_p99']:.2f}, " + f">=0.5 {res['frac_jaccard_ge_050']:.2f}, recall@3 med {res['recall_at_3']['median']:.3f} " + f"(null {res['null_recall_at_3']['median']:.3f})", + flush=True, + ) + + if direction_check: + drng = np.random.default_rng(direction_seed) + d_model = dicts[0][0].shape[1] + for qs, cs in seed_pairs: + key = f"s{labels[qs]}->s{labels[cs]}" + r = results[key] + fails = [q for q in r["per_query"] if q["best_jaccard"] <= r["null99"]] + dq, dc = unit_decs[qs], unit_decs[cs] + g = torch.tensor(drng.standard_normal((direction_null, d_model)), dtype=torch.float32) + g = g / g.norm(dim=1, keepdim=True) + null_max = (g @ dc.T).max(dim=1).values + cos99 = float(np.percentile(null_max.numpy(), 99)) + n_dir = n_half = 0 + for q in fails: + maxcos = float((dq[q["feature"]] @ dc.T).max()) + q["max_decoder_cosine"] = round(maxcos, 4) + n_dir += int(maxcos > cos99) + n_half += int(q["max_decoder_cosine"] >= 0.5) + r["direction_check"] = dict( + n_below_null_p99=len(fails), + n_null=direction_null, + direction_null_p99=cos99, + n_above_direction_null_p99=n_dir, + n_decoder_cosine_ge_0p5=n_half, + ) + print( + f"{key}: {len(fails)} token-level failures | direction-null p99 = {cos99:.3f} | " + f"above direction null: {n_dir}/{len(fails)} | cosine >= 0.5: {n_half}/{len(fails)}" + ) + for q in sorted(fails, key=lambda x: -x["max_decoder_cosine"]): + mark = "DIR-MATCH" if q["max_decoder_cosine"] > cos99 else "no match " + print( + f" [{mark}] feat {q['feature']} maxcos={q['max_decoder_cosine']:.2f} " + f"bestJ={q['best_jaccard']:.2f} :: {', '.join(q['query_tokens'][:8])}" + ) + + return dict(width=width_tag, top_r=top_r, top_m=top_m, n_null=n_null, pairs=results) + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument( + "--dictionaries", + nargs="+", + type=Path, + required=True, + help="checkpoint.pt of each seed, one recipe (paper: seeds 0, 1, 2 of one width)", + ) + ap.add_argument("--seed-labels", nargs="+", type=int, default=None, help="labels for output keys (default 0..n-1)") + ap.add_argument("--w-u", type=Path, required=True, help="extraction payload {W_U_orig, ...}") + ap.add_argument("--bank", type=Path, required=True, help="curated A/B JSONL bank (target_a / target_b)") + ap.add_argument("--tokenizer", default="Qwen/Qwen3.5-2B", help="HF tokenizer id or local path") + ap.add_argument("--width-tag", default=None, help='label stored in the output, e.g. "32x"') + ap.add_argument("--max-contrasts", type=int, default=150) + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--top-m", type=int, default=8, help="features per side when collecting used features") + ap.add_argument("--top-r", type=int, default=12, help="rows per feature -> token group") + ap.add_argument("--n-query", type=int, default=100, help="query features sampled per seed pair") + ap.add_argument("--n-null", type=int, default=500, help="frequency-matched pseudo-groups per seed pair") + ap.add_argument("--greedy-pool", type=int, default=200, help="best single matches eligible for the greedy union") + ap.add_argument("--greedy-k", type=int, default=3, help="maximum features in the greedy union (recall@k)") + ap.add_argument("--min-group-size", type=int, default=4, help="minimum tokens in a query group") + ap.add_argument("--direction-check", action="store_true", help="direction-level check of the below-null tail") + ap.add_argument("--direction-null", type=int, default=500, help="random unit directions for the direction null") + ap.add_argument("--direction-seed", type=int, default=0) + ap.add_argument("--device", type=str, default=None, help="device for the all-feature row scoring") + ap.add_argument("--out", type=Path, required=True, help="output JSON") + args = ap.parse_args() + + if len(args.dictionaries) < 2: + ap.error("--dictionaries needs at least two checkpoints") + labels = args.seed_labels if args.seed_labels is not None else list(range(len(args.dictionaries))) + if len(labels) != len(args.dictionaries): + ap.error("--seed-labels must have one entry per dictionary") + + rng = np.random.default_rng(args.seed) + device = resolve_device(args.device) + print(f"device={device}") + W, _ = load_readout(args.w_u) + from transformers import AutoTokenizer + + tok = AutoTokenizer.from_pretrained(args.tokenizer) + + dicts = [load_dictionary(p) for p in args.dictionaries] + for lab, d in zip(labels, dicts): + print(f"seed {lab}: D={d[0].shape[0]} k={d[3]}") + if len({(d[0].shape[0], d[3]) for d in dicts}) != 1: + raise SystemExit("dictionaries differ in width or k; pass checkpoints from one dictionary family") + + pairs = load_contrast_pairs(args.bank, tok, rng, args.max_contrasts) + print(f"contrasts: {len(pairs)} unique single-token pairs") + + out = run( + dicts, + labels, + W, + tok, + pairs, + rng, + device, + top_m=args.top_m, + top_r=args.top_r, + n_query=args.n_query, + n_null=args.n_null, + greedy_pool=args.greedy_pool, + greedy_k=args.greedy_k, + min_group_size=args.min_group_size, + width_tag=args.width_tag, + direction_check=args.direction_check, + direction_null=args.direction_null, + direction_seed=args.direction_seed, + ) + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(to_jsonable(out), indent=2) + "\n") + print(f"wrote {args.out}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/eval/loo_core_recovery.py b/scripts/eval/loo_core_recovery.py new file mode 100644 index 0000000..4ce294f --- /dev/null +++ b/scripts/eval/loo_core_recovery.py @@ -0,0 +1,298 @@ +#!/usr/bin/env python3 +"""Held-out recovery of the stable grouping core (Appendix I, ``tab:app-loo-core-recovery``). + +Which methods recover reproducible grouping structure, on identical targets? +For each curated contrast and each side, every seed's dictionary yields a side +token-set (top-M features per sign, top-R centred rows per feature). For each +held-out seed ``i`` the core is the set of tokens present in the side sets of +ALL remaining seeds (leave-one-out construction; with three seeds, 2 of 2), and + + srp_loo_recall = |core_i & set_i| / |core_i| held-out dictionary, same recipe + knn_loo_recall = |core_i & K| / |core_i| weighted-kNN side set on the rows + cluster_loo_recall = |core_i & CL| / |core_i| hard k-means clustering of the rows + reference recall = |core_i & R| / |core_i| optional dictionary from another recipe + +all scored on the same cores. The kNN side set of a target row is its top +``--side-n`` cosine neighbours among the centred, normalised rows; the cluster +side set is the target's k-means cluster (spherical k-means with +``--n-clusters`` centroids, ``--kmeans-iters`` iterations, seeded), truncated to +the ``--side-n`` members closest to the centroid. 95% CIs are a +contrast-clustered bootstrap (resample contrasts, keep all their cells, +``--n-boot`` replicates) for each method's mean recall and for the SRP/kNN ratio. + +Dictionary-family discipline: ``--dictionaries`` must come from ONE recipe (the +seed-variation family trained from ``configs/sweeps/qwen35_2b_seedvar_base.yaml``); +the script refuses dictionaries that differ in width or ``k``. A dictionary +from a different recipe enters only through ``--reference-dict`` and is scored +on the same cores without contributing to them (Appendix I, cross-recipe +disclosure). + +Paper runs (Qwen3.5-2B, three seeds per width, curated A/B bank, all defaults): + + uv run python scripts/eval/loo_core_recovery.py \\ + --dictionaries results/qwen35_2b_seedvar/qwen2b_d65536_k256_s0/checkpoint.pt \\ + results/qwen35_2b_seedvar/qwen2b_d65536_k256_s1/checkpoint.pt \\ + results/qwen35_2b_seedvar/qwen2b_d65536_k256_s2/checkpoint.pt \\ + --w-u "$SRP_ARCHIVE_ROOT/data/qwen35-2b/qwen35_2b.pt" \\ + --bank data/query_banks/qwen_gemma_result1_curated_ab.jsonl \\ + --width-tag 32x \\ + --reference-dict "$SRP_ARCHIVE_ROOT/results/qwen35_2b_sae/converge/topk_d65536_32x_k256_rowseeded_hybrid_lamramp_s0/checkpoint.pt" \\ + --out results/seed_stability/loo_core_recovery_32x.json + + (16x: the ``qwen2b_d32768_k128_s{0,1,2}`` checkpoints, ``--width-tag 16x``, + no ``--reference-dict``.) + +The k-means fit over the full vocabulary wants a GPU (or Apple MPS); the rest +runs on CPU. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np +import torch + +from sparse_readout_prism.research.seed_stability import ( + center_rows, + contrast_features, + feature_token_set, + load_contrast_pairs, + load_dictionary, + load_readout, + token_strings, +) +from sparse_readout_prism.utils import resolve_device, to_jsonable + + +def kmeans_unit(X: torch.Tensor, n: int, iters: int, seed: int, device) -> tuple[torch.Tensor, torch.Tensor]: + """Spherical k-means on unit rows: ``(centroids (n, d), assignment (N,))`` on CPU.""" + g = torch.Generator().manual_seed(seed) + Xd = X.to(device) + C = Xd[torch.randperm(X.shape[0], generator=g)[:n].to(device)].clone() + ones = torch.ones(X.shape[0], device=device) + assign = torch.empty(X.shape[0], dtype=torch.long, device=device) + for it in range(iters): + for s in range(0, X.shape[0], 4096): + assign[s : s + 4096] = (Xd[s : s + 4096] @ C.T).argmax(dim=1) + C_new = torch.zeros_like(C) + cnt = torch.zeros(n, device=device) + C_new.index_add_(0, assign, Xd) + cnt.index_add_(0, assign, ones) + dead = cnt == 0 + C = C_new / cnt.clamp_min(1.0)[:, None] + nd = int(dead.sum()) + if nd: + C[dead] = Xd[torch.randperm(X.shape[0], generator=g)[:nd].to(device)] + C = C / C.norm(dim=1, keepdim=True).clamp_min(1e-8) + print(f" kmeans iter {it + 1}/{iters} (dead={nd})", flush=True) + return C.cpu(), assign.cpu() + + +def run( + dicts: list, + W: torch.Tensor, + tok, + pairs: list[tuple], + device, + *, + top_m: int, + top_r: int, + side_n: int, + n_clusters: int, + kmeans_iters: int, + kmeans_seed: int, + n_boot: int, + bootstrap_seed: int, + reference_dict=None, + width_tag: str | None, +) -> dict: + W_c, rn, W_n = center_rows(W) + cache: dict[int, str] = {} + n_seeds = len(dicts) + + def dict_side_sets(d) -> list[tuple[set, set]]: + out = [] + for _a, _b, ia, ib in pairs: + top_pos, top_neg = contrast_features(W_n, rn, d, ia, ib, top_m) + out.append( + ( + feature_token_set(top_pos, d[0], W_c, tok, top_r, cache), + feature_token_set(top_neg, d[0], W_c, tok, top_r, cache), + ) + ) + return out + + side_sets = [] + for si, d in enumerate(dicts): + side_sets.append(dict_side_sets(d)) + print(f"seed {si} side sets done", flush=True) + + knn_sets = [] + for _a, _b, ia, ib in pairs: + sides = [] + for i in (ia, ib): + sims = W_n @ W_n[i] + sims[i] = -1 + sides.append(token_strings(tok, torch.topk(sims, side_n).indices.tolist(), cache)) + knn_sets.append(tuple(sides)) + + ref_sets = None + if reference_dict is not None: + ref_sets = dict_side_sets(reference_dict) + print("reference dictionary side sets done", flush=True) + + print(f"fitting k-means n_clusters={n_clusters} on {device} ...", flush=True) + C, assign = kmeans_unit(W_n, n_clusters, kmeans_iters, kmeans_seed, device) + + def cluster_side(i: int) -> set[str]: + members = (assign == assign[i]).nonzero().flatten() + if len(members) > side_n: + sims = W_n[members] @ C[assign[i]] + members = members[torch.topk(sims, side_n).indices] + return token_strings(tok, members.tolist(), cache) + + cluster_sets = [(cluster_side(ia), cluster_side(ib)) for _a, _b, ia, ib in pairs] + print("cluster side sets done", flush=True) + + srp_loo, knn_loo, cl_loo, ref_loo, core_sizes, cell_ci = [], [], [], [], [], [] + for ci in range(len(pairs)): + for side in (0, 1): + sets = [side_sets[si][ci][side] for si in range(n_seeds)] + K = knn_sets[ci][side] + CL = cluster_sets[ci][side] + for held in range(n_seeds): + others = [sets[j] for j in range(n_seeds) if j != held] + core = set.intersection(*others) + if not core: + continue + cell_ci.append(ci) + core_sizes.append(len(core)) + srp_loo.append(len(core & sets[held]) / len(core)) + knn_loo.append(len(core & K) / len(core)) + cl_loo.append(len(core & CL) / len(core)) + if ref_sets is not None: + ref_loo.append(len(core & ref_sets[ci][side]) / len(core)) + + def m(v): + return round(float(np.mean(v)), 4) if v else None + + # contrast-clustered bootstrap CIs (resample contrasts, keep all their cells) + cell_idx = np.array(cell_ci) + arrs = dict(srp=np.array(srp_loo), knn=np.array(knn_loo), cluster=np.array(cl_loo)) + brng = np.random.default_rng(bootstrap_seed) + n_c = len(pairs) + boots: dict[str, list[float]] = {k: [] for k in arrs} + boots["ratio_srp_knn"] = [] + for _ in range(n_boot): + samp = brng.integers(0, n_c, n_c) + mask_counts = np.bincount(samp, minlength=n_c) + w = mask_counts[cell_idx].astype(float) + if w.sum() == 0: + continue + means = {k: float((v * w).sum() / w.sum()) for k, v in arrs.items()} + for k in arrs: + boots[k].append(means[k]) + boots["ratio_srp_knn"].append(means["srp"] / max(means["knn"], 1e-9)) + + def ci95(k: str) -> list[float]: + v = np.array(boots[k]) + return [round(float(np.percentile(v, 2.5)), 4), round(float(np.percentile(v, 97.5)), 4)] + + return dict( + width=width_tag, + n_loo_cells=len(srp_loo), + core_median_size=float(np.median(core_sizes)), + srp_heldout_seed_recall=m(srp_loo), + srp_ci95=ci95("srp"), + knn_recall_same_cores=m(knn_loo), + knn_ci95=ci95("knn"), + cluster_recall_same_cores=m(cl_loo), + cluster_ci95=ci95("cluster"), + paper_dict_recall_same_cores=m(ref_loo), + ratio_srp_over_knn=round(float(np.mean(srp_loo) / np.mean(knn_loo)), 3), + ratio_ci95=ci95("ratio_srp_knn"), + ) + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument( + "--dictionaries", + nargs="+", + type=Path, + required=True, + help="checkpoint.pt of each seed, one recipe, at least three (paper: seeds 0, 1, 2 of one width)", + ) + ap.add_argument("--w-u", type=Path, required=True, help="extraction payload {W_U_orig, ...}") + ap.add_argument("--bank", type=Path, required=True, help="curated A/B JSONL bank (target_a / target_b)") + ap.add_argument("--tokenizer", default="Qwen/Qwen3.5-2B", help="HF tokenizer id or local path") + ap.add_argument("--width-tag", default=None, help='label stored in the output, e.g. "32x"') + ap.add_argument( + "--reference-dict", + type=Path, + default=None, + help="optional dictionary from another recipe, scored on the same cores (paper: 32x run only)", + ) + ap.add_argument("--max-contrasts", type=int, default=150) + ap.add_argument("--seed", type=int, default=0, help="contrast-order seed") + ap.add_argument("--top-m", type=int, default=8, help="features per side") + ap.add_argument("--top-r", type=int, default=12, help="rows per feature for the token summary") + ap.add_argument("--side-n", type=int, default=96, help="tokens per kNN / cluster side set") + ap.add_argument("--n-clusters", type=int, default=16384) + ap.add_argument("--kmeans-iters", type=int, default=12) + ap.add_argument("--kmeans-seed", type=int, default=0) + ap.add_argument("--n-boot", type=int, default=2000) + ap.add_argument("--bootstrap-seed", type=int, default=1) + ap.add_argument("--device", type=str, default=None, help="device for the k-means fit") + ap.add_argument("--out", type=Path, required=True, help="output JSON") + args = ap.parse_args() + + if len(args.dictionaries) < 3: + ap.error("--dictionaries needs at least three checkpoints (leave-one-out cores need two remaining seeds)") + + rng = np.random.default_rng(args.seed) + device = resolve_device(args.device) + W, _ = load_readout(args.w_u) + from transformers import AutoTokenizer + + tok = AutoTokenizer.from_pretrained(args.tokenizer) + + dicts = [load_dictionary(p) for p in args.dictionaries] + for si, d in enumerate(dicts): + print(f"seed {si}: D={d[0].shape[0]} k={d[3]}") + if len({(d[0].shape[0], d[3]) for d in dicts}) != 1: + raise SystemExit("dictionaries differ in width or k; pass checkpoints from one dictionary family") + reference = load_dictionary(args.reference_dict) if args.reference_dict is not None else None + + pairs = load_contrast_pairs(args.bank, tok, rng, args.max_contrasts) + print(f"contrasts: {len(pairs)} unique single-token pairs") + + out = run( + dicts, + W, + tok, + pairs, + device, + top_m=args.top_m, + top_r=args.top_r, + side_n=args.side_n, + n_clusters=args.n_clusters, + kmeans_iters=args.kmeans_iters, + kmeans_seed=args.kmeans_seed, + n_boot=args.n_boot, + bootstrap_seed=args.bootstrap_seed, + reference_dict=reference, + width_tag=args.width_tag, + ) + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(to_jsonable(out), indent=2) + "\n") + print(json.dumps(to_jsonable(out), indent=2)) + print(f"wrote {args.out}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/eval/paired_matched_kl_bootstrap.py b/scripts/eval/paired_matched_kl_bootstrap.py new file mode 100644 index 0000000..cfe3cc3 --- /dev/null +++ b/scripts/eval/paired_matched_kl_bootstrap.py @@ -0,0 +1,197 @@ +#!/usr/bin/env python3 +"""Paired cluster bootstrap for the matched-KL edit frontier. + +Backs the paired differences quoted with ``fig:lexical-matched-kl-frontier`` +(Appendix J, "Matched-KL Frontier and Cross-Model Outcome"): the readout SAE +direction edit minus each label-free category-direction baseline in held-out +suppression, compared at matched distributional cost. + +The marginal CIs in ``candidate_constrained_summary.csv`` treat each method's +mean independently, but every method is evaluated on the SAME held-out +candidates, so the test for "SRP beats a baseline at matched KL" is a paired +one: per-candidate differences, resampled by cluster. + +Input: one ``candidate_constrained_rows.csv`` per model, as written by +``scripts/run/run_qwen_profanity_suppression_eval.py``; only ``split == heldout`` +rows are used. For each target KL: pick, per method, the swept scale whose +median full-vocabulary KL (bits) is log-nearest to the target (within a factor +of 2); align rows by case id; form ``d_i = y_SRP,i - y_base,i``; cluster-bootstrap +the mean of ``d``. Clusters: held-out term (``bad``; 12 clusters, +primary/conservative) and prompt (20; secondary). Outcomes: +``bad_prob_reduction`` (primary) and ``flip`` (secondary). Comparisons with fewer +than 20 aligned candidates are skipped. + +Output: a JSON list with one record per (model, target KL, baseline, outcome): +the matched scales and their median KLs, ``mean_diff``, ``ci_term`` and +``ci_prompt`` (95% percentile intervals). + +Paper run (10,000 resamples, seed 0, target KL 0.02 / 0.05 / 0.10 / 0.20):: + + python scripts/eval/paired_matched_kl_bootstrap.py \\ + --input qwen2b=/candidate_constrained_rows.csv \\ + --input qwen0p8b=/candidate_constrained_rows.csv \\ + --input qwen9b=/candidate_constrained_rows.csv \\ + --input r1qwen7b=/candidate_constrained_rows.csv \\ + --out paired_matched_kl_results.json --n-boot 10000 --seed 0 +""" + +from __future__ import annotations + +import argparse +import csv +import json +from pathlib import Path + +import numpy as np + +BASELINES = ["mean_row_direction", "pca_group_rank4", "pca_group_direction"] +SRP = "feature_suppression" +DEFAULT_TARGET_KLS = [0.02, 0.05, 0.1, 0.2] +MIN_PAIRED_N = 20 + + +def load_rows(path: Path) -> list[dict]: + with path.open() as f: + return [r for r in csv.DictReader(f) if r["split"] == "heldout"] + + +def scale_for_kl(rows: list[dict], method: str, target: float): + """Swept scale whose median KL is log-nearest to target (within 2x).""" + kls: dict[float, list[float]] = {} + for r in rows: + if r["method"] == method: + kls.setdefault(float(r["scale"]), []).append(float(r["full_vocab_kl_bits"])) + best, best_d, best_med = None, None, None + for s, v in kls.items(): + med = float(np.median(v)) + if med <= 0: + continue + d = abs(np.log(med) - np.log(target)) + if best_d is None or d < best_d: + best, best_d, best_med = s, d, med + if best is None or best_d > np.log(2.0): + return None, None + return best, best_med + + +def paired_diff(rows, method_a, scale_a, method_b, scale_b, ykey): + a = {r["id"]: float(r[ykey]) for r in rows if r["method"] == method_a and float(r["scale"]) == scale_a} + b = {r["id"]: float(r[ykey]) for r in rows if r["method"] == method_b and float(r["scale"]) == scale_b} + meta = {r["id"]: (r["bad"], r["prompt"]) for r in rows if r["method"] == method_a and float(r["scale"]) == scale_a} + ids = sorted(set(a) & set(b)) + d = np.array([a[i] - b[i] for i in ids]) + term = np.array([meta[i][0] for i in ids]) + prompt = np.array([meta[i][1] for i in ids]) + return d, term, prompt + + +def cluster_ci(d: np.ndarray, clusters: np.ndarray, n_boot: int, seed: int): + keys = np.unique(clusters) + idx_by = {k: np.where(clusters == k)[0] for k in keys} + rng = np.random.default_rng(seed) + means = [] + for _ in range(n_boot): + pick = rng.choice(len(keys), size=len(keys), replace=True) + sel = np.concatenate([idx_by[keys[p]] for p in pick]) + means.append(d[sel].mean()) + means = np.array(means) + return float(np.percentile(means, 2.5)), float(np.percentile(means, 97.5)) + + +def paired_comparisons( + model: str, + rows: list[dict], + *, + target_kls: list[float], + n_boot: int, + seed: int, + verbose: bool = True, +) -> list[dict]: + out_rows: list[dict] = [] + for target in target_kls: + s_srp, kl_srp = scale_for_kl(rows, SRP, target) + if s_srp is None: + continue + for base in BASELINES: + s_b, kl_b = scale_for_kl(rows, base, target) + if s_b is None: + continue + for ykey, tag in [("bad_prob_reduction", "dp"), ("flip", "flip")]: + d, term, prompt = paired_diff(rows, SRP, s_srp, base, s_b, ykey) + if len(d) < MIN_PAIRED_N: + continue + lo_t, hi_t = cluster_ci(d, term, n_boot, seed) + lo_p, hi_p = cluster_ci(d, prompt, n_boot, seed) + sig = "*" if (lo_t > 0 or hi_t < 0) else " " + out_rows.append( + dict( + model=model, + target_kl=target, + baseline=base, + outcome=tag, + n=len(d), + srp_scale=s_srp, + srp_kl=round(kl_srp, 4), + base_scale=s_b, + base_kl=round(kl_b, 4), + mean_diff=round(float(d.mean()), 4), + ci_term=[round(lo_t, 4), round(hi_t, 4)], + ci_prompt=[round(lo_p, 4), round(hi_p, 4)], + ) + ) + if verbose and tag == "dp": + print( + f" KL~{target:<5} vs {base:<22} d={d.mean():+.4f} " + f"CI_term[{lo_t:+.4f},{hi_t:+.4f}]{sig} CI_prompt[{lo_p:+.4f},{hi_p:+.4f}] " + f"(SRP@{s_srp} kl={kl_srp:.3f} vs @{s_b} kl={kl_b:.3f}, n={len(d)})" + ) + return out_rows + + +def parse_input(spec: str) -> tuple[str, Path]: + if "=" not in spec: + raise argparse.ArgumentTypeError(f"--input expects MODEL=path/to/candidate_constrained_rows.csv, got {spec!r}") + model, path = spec.split("=", 1) + model = model.strip() + if not model: + raise argparse.ArgumentTypeError(f"--input has an empty model label: {spec!r}") + return model, Path(path) + + +def main(argv: list[str] | None = None) -> int: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument( + "--input", + type=parse_input, + action="append", + required=True, + metavar="MODEL=CSV", + help="Model label and its candidate_constrained_rows.csv; repeat once per model.", + ) + ap.add_argument("--out", type=Path, required=True, help="Output JSON (list of paired comparisons).") + ap.add_argument("--n-boot", type=int, default=10_000, help="Cluster-bootstrap resamples (paper: 10000).") + ap.add_argument("--seed", type=int, default=0, help="numpy default_rng seed, reused for every CI (paper: 0).") + ap.add_argument( + "--target-kl", + type=float, + nargs="+", + default=DEFAULT_TARGET_KLS, + help="Matched-KL targets in bits (paper: 0.02 0.05 0.1 0.2).", + ) + args = ap.parse_args(argv) + + out_rows: list[dict] = [] + for model, path in args.input: + rows = load_rows(path) + print(f"\n=== {model}: {len(rows)} held-out rows ===") + out_rows.extend( + paired_comparisons(model, rows, target_kls=list(args.target_kl), n_boot=args.n_boot, seed=args.seed) + ) + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(out_rows, indent=2)) + print(f"\nwrote {args.out} ({len(out_rows)} comparisons)") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/eval/run_causal_contribution_validation.py b/scripts/eval/run_causal_contribution_validation.py new file mode 100644 index 0000000..fbecd58 --- /dev/null +++ b/scripts/eval/run_causal_contribution_validation.py @@ -0,0 +1,441 @@ +#!/usr/bin/env python3 +"""Predicted versus realized readout-side changes (causal validation of feature contributions). + +Backs ``tab:causal-validation-summary`` (one row per model) and the protocol +appendix ``app:causal-validation``, whose quoted self-test numbers come from +``--self-test``. + +For a selected contrast q = w_A - w_B (two rows of the unembedding matrix W_U) +at decoded state h, the decomposition assigns feature i the contribution + + c_i = beta_i(q) * (h . d_i), beta_i(q) = rn_A z_A,i - rn_B z_B,i + +where z are the TopK codes of the centered, per-row normalized rows, rn their +centered norms, and d_i the unit-norm decoder direction. The test ablates the +decoded state along d_i, + + h' = h - (h . d_i) d_i, + +and measures the realized margin change with the dense LM head, + + m(h') - m(h) = -(h . d_i) (q . d_i) + +(exact, because the readout is linear in h; no new forward pass, so the claim +is scoped to the readout side). The prediction uses the sparse code beta_i(q) +while the measurement uses the true projection q . d_i, so sparse-code +misattribution, decoder non-orthogonality and the row residual all break the +agreement: slope ~ 1 with high r^2 is falsifiable, not an identity. + +Per contrast the top ``--top-features`` features by |c_i| are tested plus +``--random-per-case`` random features outside that set as controls. The stored +``delta_real`` is the ablated mass +(h . d_i)(q . d_i), i.e. minus the realized +margin change, so it is directly comparable to c_i. + +Slope orientation: ``fit_r2_slope(x=c_pred, y=delta_real)`` fits the realized +change on the predicted contribution through the origin (y = s x), so the slope +is realized-on-predicted. Slope > 1 means the realized change exceeds the +prediction; slope < 1 means it falls short. + +``summary.json`` subsets: ``gated_predicted`` (covered contrasts, rho below +``--rho-gate`` with sign agreement; the table's r^2, CI, slope and n), +``ungated_predicted`` (all contrasts) and ``gated_random_control`` (the table's +"Random" column). CIs are 95% cluster bootstraps over base cases. + +Inputs: W_U is read from the live ``lm_head`` of ``--model-id``; +``--checkpoint`` is a TopK dictionary in the trainer schema (the released +dictionaries at hematteo/sparse-readout-prism, laid out as +``//checkpoint.pt``); ``--bank-dir`` holds the JSONL +query banks (``data/query_banks``). The ``model_native`` bank is the C4 slice +(``qwen_gemma_result1_model_native_prompts_c4.jsonl``, regenerated locally by +``scripts/data/build_query_banks.py --native-source c4``); its rows carry no A/B +targets, so it contributes no contrasts and a missing file is logged and +skipped. The ~260 contrasts per model come from ``curated_ab`` (240 targeted +rows) and ``case_candidates`` (36), minus targets that do not tokenize to a +single token. + +Outputs: ``causal_rows.csv`` (one row per prediction-realization pair), +``summary.json`` and ``manifest.json`` (run provenance). + +Paper run: six models, each at the fidelity operating point (32x width, +k=256; Hub path ``/k256_32x/checkpoint.pt``), every other flag at its +default:: + + uv run python scripts/eval/run_causal_contribution_validation.py \\ + --model-id --checkpoint \\ + --bank-dir data/query_banks --out-dir results/causal_ \\ + --device cuda --dtype bfloat16 --seed 0 + + MODEL_ID Hub checkpoint TAG + Qwen/Qwen3.5-0.8B qwen3.5-0.8b/k256_32x/checkpoint.pt qwen0p8b + Qwen/Qwen3.5-2B qwen3.5-2b/k256_32x/checkpoint.pt qwen2b + Qwen/Qwen3.5-9B qwen3.5-9b/k256_32x/checkpoint.pt qwen9b + mistralai/Ministral-3-8B-Base-2512 ministral-3-8b/k256_32x/checkpoint.pt ministral8b + deepseek-ai/DeepSeek-R1-Distill-Qwen-7B r1-distill-qwen-7b/k256_32x/checkpoint.pt r1qwen7b + deepseek-ai/DeepSeek-R1-Distill-Llama-8B r1-distill-llama-8b/k256_32x/checkpoint.pt r1llama8b + +Defaults the paper run relied on: ``--banks curated_ab,case_candidates,model_native +--max-native 300 --top-features 10 --random-per-case 10 --rho-gate 0.5 +--n-boot 2000 --max-len 64``. + +``--self-test`` runs the synthetic residual-free check of the math core on CPU +(no model, no checkpoint) and exits. +""" + +from __future__ import annotations + +import argparse +import hashlib +import time +from pathlib import Path + +import numpy as np +import torch + +from sparse_readout_prism.data import center_normalize_rows +from sparse_readout_prism.factorizers import TopKSAE, load_factorizer +from sparse_readout_prism.research.qwen_readout import encode_topk +from sparse_readout_prism.research.run_io import load_bank, run_provenance +from sparse_readout_prism.utils import find_lm_head, load_causal_lm, resolve_device, set_seed, write_csv, write_json + +# bank name -> (file under --bank-dir, default cap on base cases) +BANK_FILES = { + "curated_ab": ("qwen_gemma_result1_curated_ab.jsonl", 320), + "case_candidates": ("qwen_gemma_result1_case_candidates.jsonl", 40), + "model_native": ("qwen_gemma_result1_model_native_prompts_c4.jsonl", 300), +} + + +def log(msg: str) -> None: + print(f"[causal {time.strftime('%H:%M:%S')}] {msg}", flush=True) + + +# --------------------------------------------------------------------------- # +# Banks and tokens +# --------------------------------------------------------------------------- # + + +def load_banks(bank_dir: Path, banks: list[str], caps: dict[str, int]) -> list[dict]: + """Load the requested banks, capping each by base case; a missing file is skipped.""" + rows: list[dict] = [] + for b in banks: + fname, default_cap = BANK_FILES[b] + path = bank_dir / fname + if not path.exists(): + log(f"bank missing, skipping: {path}") + continue + kept = load_bank(path, caps.get(b, default_cap)) + for r in kept: + r["bank"] = b + rows.extend(kept) + log(f"bank {b}: {len(kept)} records ({len({r['base_case_id'] for r in kept})} base cases)") + return rows + + +def single_token_id(tok, term: str): + """Token id of ``term`` if it (bare first, then space-prefixed) is a single token, else None.""" + for variant in (term, " " + term): + ids = tok.encode(variant, add_special_tokens=False) + if len(ids) == 1: + return ids[0] + return None + + +# --------------------------------------------------------------------------- # +# Core math (unit-testable without a model) +# --------------------------------------------------------------------------- # + + +def causal_pairs_for_contrast( + h: torch.Tensor, # (d,) + q: torch.Tensor, # (d,) = w_A - w_B (raw) + beta: torch.Tensor, # (D,) raw-space code difference (rn_A*z_A - rn_B*z_B) + W_dec: torch.Tensor, # (D, d) unit rows + top_k: int, + rng: np.random.Generator, + n_random: int, +): + """Return (feature_id, c_pred, delta_real, is_random) tuples.""" + proj = W_dec @ h # (D,) h . d_i + c = beta * proj # predicted signed contribution per feature + qd = W_dec @ q # (D,) q . d_i + realized = proj * qd # = -(delta m) under ablation; sign matches c's claim + top = torch.topk(c.abs(), k=min(top_k, c.numel())).indices + active = set(top.tolist()) + out = [(int(i), float(c[i]), float(realized[i]), 0) for i in top] + # random control: features NOT in the predicted-top set (mostly beta=0) + pool = rng.choice(c.numel(), size=min(4 * n_random, c.numel()), replace=False) + ctrl = [int(i) for i in pool if int(i) not in active][:n_random] + out += [(int(i), float(c[i]), float(realized[i]), 1) for i in ctrl] + return out + + +def fit_r2_slope(x: np.ndarray, y: np.ndarray) -> tuple[float, float]: + """OLS y = s*x through the origin and the r2 of that fit. + + Called as ``fit_r2_slope(x=c_pred, y=delta_real)``: the slope is realized + change on predicted contribution, so slope > 1 means the realized change + exceeds the prediction. + """ + if len(x) < 3 or float(np.dot(x, x)) == 0.0: + return float("nan"), float("nan") + s = float(np.dot(x, y) / np.dot(x, x)) + resid = y - s * x + denom = float(np.sum((y - y.mean()) ** 2)) + r2 = 1.0 - float(np.sum(resid**2)) / denom if denom > 0 else float("nan") + return r2, s + + +def cluster_bootstrap_r2(rows: list[dict], n_boot: int, seed: int) -> tuple[float, float]: + """95% percentile interval of r2 under a cluster bootstrap over ``row['cluster']``.""" + by_cluster: dict[str, list[tuple[float, float]]] = {} + for r in rows: + by_cluster.setdefault(r["cluster"], []).append((r["c_pred"], r["delta_real"])) + keys = sorted(by_cluster) + rng = np.random.default_rng(seed) + stats = [] + for _ in range(n_boot): + sample = rng.choice(len(keys), size=len(keys), replace=True) + xs, ys = [], [] + for si in sample: + for cx, cy in by_cluster[keys[si]]: + xs.append(cx) + ys.append(cy) + r2, _ = fit_r2_slope(np.array(xs), np.array(ys)) + if np.isfinite(r2): + stats.append(r2) + if not stats: + return float("nan"), float("nan") + return float(np.percentile(stats, 2.5)), float(np.percentile(stats, 97.5)) + + +def self_test_metrics(seed: int = 0) -> dict: + """Synthetic residual-free check of the math core (no model). + + A random unit-row dictionary, a random decoded state, and a contrast built + exactly in the dictionary span from a known 3-sparse beta. Returns the + fitted statistics, the pass flags and the synthetic inputs so callers can + run further identity checks. + """ + torch.manual_seed(seed) + d, D = 64, 512 + W_dec = torch.randn(D, d) + W_dec = W_dec / W_dec.norm(dim=1, keepdim=True) + h = torch.randn(d) + rng = np.random.default_rng(seed) + # Construct q exactly in the dictionary span with known sparse beta: + beta = torch.zeros(D) + beta[[3, 40, 100]] = torch.tensor([2.0, -1.5, 0.7]) + q = beta @ W_dec # residual-free contrast + pairs = causal_pairs_for_contrast(h, q, beta, W_dec, top_k=3, rng=rng, n_random=8) + pred = np.array([p[1] for p in pairs if p[3] == 0]) + real = np.array([p[2] for p in pairs if p[3] == 0]) + r2, slope = fit_r2_slope(pred, real) + # With orthonormal-ish random dirs and residual-free q, agreement is high + # but NOT exact (off-diagonal d_i . d_j != 0) -- assert strong, not perfect. + ok1 = r2 > 0.9 and 0.7 < slope < 1.3 + ctrl_pred = np.array([abs(p[1]) for p in pairs if p[3] == 1]) + ctrl_max_abs_pred = float(ctrl_pred.max(initial=0.0)) + ok2 = ctrl_max_abs_pred < 1e-6 # random features carry ~0 prediction + # Identity check: realized values equal the algebraic form + i0 = pairs[0][0] + lhs = pairs[0][2] + rhs = float((W_dec[i0] @ h) * (W_dec[i0] @ q)) + identity_abs_err = abs(lhs - rhs) + ok3 = identity_abs_err < 1e-4 + return { + "r2": r2, + "slope": slope, + "ctrl_max_abs_pred": ctrl_max_abs_pred, + "identity_abs_err": identity_abs_err, + "checks": [ + ("r2/slope on span-constructed q", ok1), + ("random controls predict ~0", ok2), + ("realized identity", ok3), + ], + "pairs": pairs, + "h": h, + "q": q, + "beta": beta, + "W_dec": W_dec, + } + + +def self_test() -> int: + m = self_test_metrics() + for name, ok in m["checks"]: + print(f" [{'PASS' if ok else 'FAIL'}] {name}") + print(f" (r2={m['r2']:.3f} slope={m['slope']:.3f})") + return 0 if all(ok for _, ok in m["checks"]) else 1 + + +# --------------------------------------------------------------------------- # +# Main +# --------------------------------------------------------------------------- # + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--model-id", type=str, help="HF model id; W_U is read from its live lm_head") + ap.add_argument("--checkpoint", type=Path, help="TopK dictionary checkpoint.pt (trainer schema)") + ap.add_argument("--bank-dir", type=Path, help="directory holding the JSONL query banks") + ap.add_argument("--banks", type=str, default="curated_ab,case_candidates,model_native") + ap.add_argument("--max-native", type=int, default=300, help="base-case cap for the model_native bank") + ap.add_argument("--out-dir", type=Path) + ap.add_argument("--device", default="cuda") + ap.add_argument("--dtype", choices=["bfloat16", "float32"], default="bfloat16") + ap.add_argument("--top-features", type=int, default=10) + ap.add_argument("--random-per-case", type=int, default=10) + ap.add_argument("--max-len", type=int, default=64) + ap.add_argument("--rho-gate", type=float, default=0.5) + ap.add_argument("--n-boot", type=int, default=2000) + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--limit", type=int, default=None, help="debug: cap contrasts") + ap.add_argument("--self-test", action="store_true") + args = ap.parse_args() + + if args.self_test: + return self_test() + if not (args.model_id and args.checkpoint and args.bank_dir and args.out_dir): + ap.error("--model-id, --checkpoint, --bank-dir, --out-dir required (or --self-test)") + bad = [b for b in args.banks.split(",") if b.strip() and b.strip() not in BANK_FILES] + if bad: + ap.error(f"unknown bank(s) {bad}; have {list(BANK_FILES)}") + args.out_dir.mkdir(parents=True, exist_ok=True) + set_seed(args.seed) + rng = np.random.default_rng(args.seed) + device = resolve_device(args.device) + + log(f"loading model {args.model_id}") + dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float32 + model, tok = load_causal_lm(args.model_id, dtype=dtype, device_map=None) + model.to(device) + lm_head = find_lm_head(model) + W_U = lm_head.weight.detach().float().to(device) # (V, d) + V, d = W_U.shape + row_mean = W_U.mean(dim=0) # (d,) centering mean of the live readout + log(f"W_U ({V}, {d}) from live lm_head") + + log(f"loading dictionary {args.checkpoint}") + ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=True) + sae = load_factorizer(ckpt, device=device, freeze=True) + if not isinstance(sae, TopKSAE): + raise ValueError(f"expected a TopK-family dictionary, got architecture {sae.architecture!r}") + k = int(sae.k) + W_dec = sae.decoder.detach().float() # (D, d) + enc_w = sae.encoder.weight.detach().float() # (D, d) + enc_b = sae.encoder.bias.detach().float() # (D,) + dec_norm = W_dec.norm(dim=1) + W_dec_unit = W_dec / dec_norm[:, None].clamp_min(1e-8) # (D, d) unit rows + ckpt_row_mean = ckpt.get("row_mean") + if ckpt_row_mean is not None: + gap = float((ckpt_row_mean.float().to(device) - row_mean).abs().max()) + log(f"checkpoint row_mean vs live W_U mean: max abs gap {gap:.3e} (the live mean is used)") + log(f"dictionary D={W_dec.shape[0]} k={k} (decoder norms {dec_norm.min():.3f}-{dec_norm.max():.3f})") + + banks = [b.strip() for b in args.banks.split(",") if b.strip()] + caps = {"model_native": args.max_native} + records = load_banks(args.bank_dir, banks, caps) + if args.limit: + records = records[: args.limit] + + # hidden-state capture: lm_head input at final position + cache: dict[str, torch.Tensor] = {} + grabbed: list[torch.Tensor] = [] + + def pre_hook(_m, inputs): + grabbed.append((inputs[0] if isinstance(inputs, tuple) else inputs).detach()) + + hook = lm_head.register_forward_pre_hook(pre_hook) + + @torch.no_grad() + def hidden_for(prompt: str) -> torch.Tensor: + if prompt not in cache: + grabbed.clear() + enc = tok(prompt, return_tensors="pt", truncation=True, max_length=args.max_len).to(device) + model(**enc) + cache[prompt] = grabbed[-1][0, -1, :].float() # (d,) + if len(cache) > 4096: + cache.pop(next(iter(cache))) + return cache[prompt] + + @torch.no_grad() + def raw_beta(a_id: int, b_id: int) -> torch.Tensor: + rn, x = center_normalize_rows(W_U[[a_id, b_id]], row_mean) # (2,), (2, d) + codes = encode_topk(x, enc_w, enc_b, k) # (2, D) + return rn[0] * codes[0] - rn[1] * codes[1] # raw-space beta_i(q) + + out_rows: list[dict] = [] + n_done = n_skip = 0 + for rec in records: + a = rec.get("target_a") + b = rec.get("target_b") + prompt = rec.get("prompt") + if not (a and b and prompt): + n_skip += 1 + continue + a_id, b_id = single_token_id(tok, a), single_token_id(tok, b) + if a_id is None or b_id is None or a_id == b_id: + n_skip += 1 + continue + h = hidden_for(prompt) + q = (W_U[a_id] - W_U[b_id]).float() + beta = raw_beta(a_id, b_id) + # coverage quantities (margin reconstruction; offset cancels for a difference) + m_exact = float(h @ q) + m_recon = float((beta * (W_dec_unit @ h)).sum()) + rho = abs(m_exact - m_recon) / max(abs(m_exact), 1e-6) + gated = int(rho < args.rho_gate and (m_exact > 0) == (m_recon > 0)) + cluster = rec.get("base_case_id") or rec.get("case_id") or hashlib.sha1(prompt.encode()).hexdigest()[:10] + pairs = causal_pairs_for_contrast(h, q, beta, W_dec_unit, args.top_features, rng, args.random_per_case) + for fid, c_pred, delta_real, is_rand in pairs: + out_rows.append( + dict( + case_id=rec.get("case_id", ""), + cluster=cluster, + bank=rec.get("bank", ""), + target_a=a, + target_b=b, + feature_id=fid, + c_pred=c_pred, + delta_real=delta_real, + is_random=is_rand, + m_exact=m_exact, + rho=rho, + gated=gated, + ) + ) + n_done += 1 + if n_done % 100 == 0: + log(f"contrasts {n_done}/{len(records)} (skipped {n_skip})") + hook.remove() + log(f"done: {n_done} contrasts, {n_skip} skipped, {len(out_rows)} feature pairs") + + write_csv(args.out_dir / "causal_rows.csv", out_rows, write_empty=True) + + summary = { + "model_id": args.model_id, + "n_contrasts": n_done, + "k": k, + "top_features": args.top_features, + "seed": args.seed, + } + for subset, sel in [ + ("gated_predicted", lambda r: r["gated"] and not r["is_random"]), + ("ungated_predicted", lambda r: not r["is_random"]), + ("gated_random_control", lambda r: r["gated"] and r["is_random"]), + ]: + rows = [r for r in out_rows if sel(r)] + x = np.array([r["c_pred"] for r in rows]) + y = np.array([r["delta_real"] for r in rows]) + r2, slope = fit_r2_slope(x, y) + lo, hi = cluster_bootstrap_r2(rows, args.n_boot, args.seed) + summary[subset] = dict(n=len(rows), r2=r2, slope=slope, r2_ci=[lo, hi]) + log(f"{subset}: n={len(rows)} r2={r2:.4f} slope={slope:.3f} CI[{lo:.3f},{hi:.3f}]") + write_json(summary, args.out_dir / "summary.json") + write_json(run_provenance(args), args.out_dir / "manifest.json") + log(f"wrote {args.out_dir}/causal_rows.csv, summary.json and manifest.json") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/figures/compute_cross_lens_shared_feature.py b/scripts/figures/compute_cross_lens_shared_feature.py new file mode 100644 index 0000000..172270a --- /dev/null +++ b/scripts/figures/compute_cross_lens_shared_feature.py @@ -0,0 +1,143 @@ +#!/usr/bin/env python3 +"""Metrics behind ``fig:cross-lens-butterfly``: two lenses, two surface tokens, one dominant readout feature. + +Offline over two (or more) readout dumps from +``scripts/run/run_cross_lens_readouts.py`` produced with ``--decompose-top1`` on the +same bank. For one prompt id and one layer the script extracts, per lens: the +top-1 token the lens reports and its logit, the feature sum of that token's +decomposition, the dominant feature (largest |contribution|) with its share of +the feature sum, and the largest contribution from any other feature. It then +reports whether the dominant feature is shared across the lenses, each lens's +share carried by the shared feature, and the feature's top unembedding rows from +the dump's ``feature_top_tokens`` labels. No rendering; the figure is drawn in +the paper source from these numbers. + +Paper run (factual-recall prompt ``fac_03``, "The capital of China is Beijing. The +capital of the UK is", layer 26, English- and Chinese-fitted 100-prompt Jacobian +lenses):: + + compute_cross_lens_shared_feature.py --dump EN=/en_zh__jlens_en_n100.json \ + --dump ZH=/en_zh__jlens_zh_n100.json --prompt-id fac_03 --layer 26 \ + --out-csv /cross_lens_shared_feature.csv --out-json /cross_lens_shared_feature.json +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +from sparse_readout_prism.utils import write_csv, write_json + +POS = "-1" + + +def parse_dump_args(specs: list[str]) -> list[tuple[str, Path]]: + """Parse repeated ``LABEL=path`` arguments, preserving order.""" + out = [] + for spec in specs: + label, sep, path = spec.partition("=") + if not sep or not label or not path: + raise ValueError(f"--dump expects LABEL=path, got {spec!r}") + out.append((label, Path(path))) + return out + + +def lens_top1_summary(record: dict, layer: str) -> dict: + """Top-1 token, its logit, and the dominant-feature share of its decomposition at one layer.""" + e = record["layers"][layer][POS] + if "top1_decomp" not in e: + raise KeyError("dump lacks top1_decomp; rerun run_cross_lens_readouts.py with --decompose-top1") + tok1, logit1 = e["top5"][0] + d = e["top1_decomp"] + feats = d["top_features"] + dom = feats[0] if feats else {"id": None, "contribution": 0.0} + others = [f["contribution"] for f in feats[1:]] + fsum = d["feature_sum"] + return { + "top1_token": tok1, + "top1_lens_logit": logit1, + "top1_original_logit": d["original_logit"], + "feature_sum": fsum, + "residual": d["residual"], + "dominant_feature": dom["id"], + "dominant_contribution": dom["contribution"], + "dominant_share_of_feature_sum": dom["contribution"] / fsum if fsum else None, + "largest_other_contribution": max(others) if others else None, + "contributions": {int(f["id"]): float(f["contribution"]) for f in feats}, + } + + +def shared_feature_metrics(records: dict[str, dict], layer: str, labels_by_lens: dict[str, dict]) -> dict: + """Per-lens top-1 summaries plus the shared dominant feature and its share under each lens.""" + per_lens = {label: lens_top1_summary(rec, layer) for label, rec in records.items()} + dominant = {s["dominant_feature"] for s in per_lens.values()} + shared = next(iter(dominant)) if len(dominant) == 1 else None + shared_share = {} + if shared is not None: + for label, s in per_lens.items(): + c = s["contributions"].get(shared, 0.0) + shared_share[label] = c / s["feature_sum"] if s["feature_sum"] else None + top_tokens = None + if shared is not None: + for labels in labels_by_lens.values(): + if str(shared) in labels: + top_tokens = labels[str(shared)] + break + return { + "layer": layer, + "per_lens": per_lens, + "shared_dominant_feature": shared, + "shared_feature_share_of_feature_sum": shared_share, + "shared_feature_top_tokens": top_tokens, + } + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--dump", action="append", required=True, help="LABEL=path, repeatable (e.g. EN=..., ZH=...)") + p.add_argument("--prompt-id", default="fac_03", help="record id (paper: fac_03)") + p.add_argument("--layer", default="26", help="layer key (paper: 26)") + p.add_argument("--out-csv", default=None, help="one row per lens") + p.add_argument("--out-json", default=None, help="full metrics including the shared-feature block") + args = p.parse_args(argv) + + records, labels_by_lens = {}, {} + for label, path in parse_dump_args(args.dump): + dump = json.loads(path.read_text()) + matches = [r for r in dump["records"] if r["id"] == args.prompt_id] + if not matches: + raise KeyError(f"prompt id {args.prompt_id!r} not in {path}") + records[label] = matches[0] + labels_by_lens[label] = dump.get("feature_top_tokens", {}) + + metrics = shared_feature_metrics(records, args.layer, labels_by_lens) + metrics["prompt_id"] = args.prompt_id + rows = [] + for label, s in metrics["per_lens"].items(): + row = {"lens": label, "prompt_id": args.prompt_id, "layer": args.layer} + row.update({k: v for k, v in s.items() if k != "contributions"}) + row["shared_dominant_feature"] = metrics["shared_dominant_feature"] + row["shared_feature_share_of_feature_sum"] = metrics["shared_feature_share_of_feature_sum"].get(label) + rows.append(row) + share = s["dominant_share_of_feature_sum"] + print( + f"{label:>6s}: top1={s['top1_token']!r} logit={s['top1_lens_logit']:.1f} " + f"dominant=f{s['dominant_feature']} share={share:.3f} " + f"largest_other={s['largest_other_contribution']:+.2f} feature_sum={s['feature_sum']:.2f}" + ) + print( + f"shared dominant feature: {metrics['shared_dominant_feature']} top rows: {metrics['shared_feature_top_tokens']}" + ) + + if args.out_csv: + write_csv(args.out_csv, rows) + print(f"wrote {args.out_csv}") + if args.out_json: + write_json(metrics, args.out_json) + print(f"wrote {args.out_json}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/run/fit_jlens.py b/scripts/run/fit_jlens.py new file mode 100644 index 0000000..692bdd7 --- /dev/null +++ b/scripts/run/fit_jlens.py @@ -0,0 +1,303 @@ +#!/usr/bin/env python3 +"""Fit the Jacobian lens (jlens) on a seeded C4 prompt sample: prompts, smoke, fit, merge. + +Produces the fitted lenses behind the cross-lens study (``fig:cross-lens-butterfly``, +``fig:cross-lens-extension``, ``tab:app-cross-lens-families``, +``tab:app-cross-lens-extension``, ``tab:app-cross-lens-antonym-layers``, +``tab:app-cross-lens-de-antonym-layers``): one Jacobian lens per fitting corpus +(English, Chinese and German C4) with identical seeds, source layers and procedure, +at two fitting budgets (100 and 300 prompts, both read from the head of the same +seeded dump). + +Subcommands, run in this order: + + prompts seeded streaming C4 prompt dump (idempotent, atomic write) + smoke tiny fit on a small model with finite / save-load / apply checks + fit one prompt shard on one GPU, resumable through the jlens checkpoint + merge n_prompts-weighted merge of the shard lenses plus a half-vs-full + convergence report per layer + +The estimator is the reference implementation of the Jacobian lens (``jlens``, +Apache-2.0, https://github.com/anthropics/jacobian-lens); this script only +orchestrates sharding and storage. Jacobians accumulate in fp32 and are saved in +fp16. Install the dependency with ``uv sync --extra lens``; the import is lazy so +the rest of the repository installs and tests without it. + +Source layers are ``pick_source_layers(n_layers, --n-layers)``: evenly spaced over +5-95% depth, excluding layer 0 and the final (target) layer. For Qwen3.5-9B and +``--n-layers 12`` this gives {2, 5, 7, 10, 12, 14, 17, 19, 21, 24, 26, 29}. + +Paper runs (Qwen3.5-9B, four shards on four GPUs, ``--dim-batch 8``):: + + # seeded prompt dumps, 1000 prompts each (the English dump keeps the default --min-chars 800) + fit_jlens.py prompts --prompts-json /c4_prompts_en_seed0.json --n-prompts 1000 --seed 0 + fit_jlens.py prompts --prompts-json /c4_prompts_zh_seed0.json --n-prompts 1000 --seed 0 \ + --min-chars 300 --c4-config zh + fit_jlens.py prompts --prompts-json /c4_prompts_de_seed0.json --n-prompts 1000 --seed 0 \ + --min-chars 300 --c4-config de + # smoke on a small model before the 9B fit + fit_jlens.py smoke --model-id Qwen/Qwen3.5-0.8B --prompts-json /c4_prompts_en_seed0.json \ + --out /smoke_0p8b.lens.pt --dim-batch 16 + # one shard per GPU (i in 0..3), then merge; --n-prompts 100 for the main lenses and + # --n-prompts 300 for the larger-corpus refits + fit_jlens.py fit --model-id Qwen/Qwen3.5-9B --prompts-json /c4_prompts_en_seed0.json \ + --n-prompts 100 --shard $i --num-shards 4 --n-layers 12 --dim-batch 8 \ + --ckpt-dir --out-dir + fit_jlens.py merge --shard-dir --num-shards 4 --out /qwen35_9b_jlens_en_seed0_n100.pt + +The same fit + merge with the zh / de dumps produce the Chinese- and German-fitted +lenses (``qwen35_9b_jlens_zh_seed0_n100.pt``, ``qwen35_9b_jlens_de_seed0_n100.pt``), +and with ``--n-prompts 300`` the three ``*_n300.pt`` refits. +""" + +from __future__ import annotations + +import argparse +import gc +import json +import os +import time +from pathlib import Path + +import torch + +from sparse_readout_prism.utils import git_commit + + +def log(msg: str) -> None: + print(f"[fit_jlens] {msg}", flush=True) + + +def _import_jlens(): + try: + import jlens + except ImportError as e: + raise ImportError( + "jlens (the Jacobian-lens reference implementation, Apache-2.0, " + "github.com/anthropics/jacobian-lens) is not installed; install it with `uv sync --extra lens`" + ) from e + return jlens + + +def load_model(model_id: str): + import transformers + + jlens = _import_jlens() + hf = ( + transformers.AutoModelForCausalLM.from_pretrained( + model_id, torch_dtype=torch.bfloat16, attn_implementation="sdpa" + ) + .cuda() + .eval() + ) + tok = transformers.AutoTokenizer.from_pretrained(model_id) + return jlens.from_hf(hf, tok) + + +def pick_source_layers(n_layers: int, n_pick: int) -> list[int]: + # Evenly spaced over 5-95% depth; excludes layer 0 and the final layer + # (target). Deterministic given (n_layers, n_pick). + fracs = [(i + 1) / (n_pick + 1) for i in range(n_pick)] + layers = sorted({max(1, min(n_layers - 2, round(f * (n_layers - 1)))) for f in fracs}) + return layers + + +def cmd_prompts(args) -> None: + out = Path(args.prompts_json) + if out.exists(): + log(f"[resume] prompts exist: {out}") + return + from datasets import load_dataset + + config = args.dataset_config if args.dataset_config is not None else args.c4_config + ds = load_dataset(args.dataset, config, split=args.dataset_split, streaming=True, revision=args.revision) + ds = ds.shuffle(seed=args.seed, buffer_size=10_000) + + prompts = [] + scanned = 0 + for ex in ds: + scanned += 1 + text = (ex[args.text_field] or "").strip() + if len(text) < args.min_chars: + continue + prompts.append(text[:4000]) # jlens.fit truncates to its max_seq_len tokens + if len(prompts) >= args.n_prompts: + break + if len(prompts) < args.n_prompts: + raise RuntimeError( + f"only {len(prompts)}/{args.n_prompts} prompts after scanning {scanned} records " + f"({args.dataset} {config}); widen the stream or relax --min-chars" + ) + out.parent.mkdir(parents=True, exist_ok=True) + payload = { + "source": f"{args.dataset} {config} streaming", + "dataset": args.dataset, + "dataset_config": config, + "dataset_split": args.dataset_split, + "revision": args.revision, + "text_field": args.text_field, + "records_scanned": scanned, + "seed": args.seed, + "min_chars": args.min_chars, + "n": len(prompts), + "prompts": prompts, + } + tmp = out.with_suffix(".json.tmp") + tmp.write_text(json.dumps(payload)) + os.replace(tmp, out) + log(f"wrote {len(prompts)} prompts -> {out}") + + +def read_prompts(path: str, n: int | None = None, start: int = 0) -> list[str]: + prompts = json.loads(Path(path).read_text())["prompts"] + if start < 0: + raise ValueError(f"start must be non-negative, got {start}") + prompts = prompts[start:] + return prompts[:n] if n else prompts + + +def cmd_smoke(args) -> None: + jlens = _import_jlens() + + model = load_model(args.model_id) + prompts = read_prompts(args.prompts_json, 2) + layers = pick_source_layers(model.n_layers, 4) + t0 = time.time() + lens = jlens.fit(model, prompts, source_layers=layers, dim_batch=args.dim_batch) + dt = time.time() - t0 + for layer, J in lens.jacobians.items(): + assert torch.isfinite(J).all(), f"non-finite Jacobian at layer {layer}" + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + lens.save(str(out)) + jlens.JacobianLens.load(str(out)) + per_layer, final_logits, _ = lens.apply(model, "The capital of France is", positions=[-1]) + assert torch.isfinite(final_logits).all() + for layer, logits in per_layer.items(): + assert torch.isfinite(logits).all(), f"non-finite lens logits at layer {layer}" + log( + f"smoke OK: d_model={model.d_model}, layers={layers}, " + f"{dt:.1f}s for 2 prompts ({dt / (2 * model.d_model) * 1e3:.2f} ms per prompt-dim)" + ) + + +def cmd_fit(args) -> None: + jlens = _import_jlens() + + prompts = read_prompts(args.prompts_json, args.n_prompts, start=args.start) + shard_prompts = prompts[args.shard :: args.num_shards] + model = load_model(args.model_id) + layers = pick_source_layers(model.n_layers, args.n_layers) + ckpt = Path(args.ckpt_dir) / f"shard{args.shard}.ckpt.pt" + ckpt.parent.mkdir(parents=True, exist_ok=True) + out = Path(args.out_dir) / f"shard{args.shard}.lens.pt" + if out.exists(): + log(f"[resume] shard {args.shard} lens exists: {out}") + return + log( + f"git={git_commit()} model={args.model_id} shard={args.shard}/{args.num_shards} " + f"prompt_slice=[{args.start},{args.start + args.n_prompts}) prompts={len(shard_prompts)} " + f"layers={layers} dim_batch={args.dim_batch} ckpt={ckpt}" + ) + t0 = time.time() + # Retry OUTSIDE the except block: the exception traceback pins the failed + # attempt's autograd graph, so an in-except retry OOMs against ghost memory. + lens = None + for dim_batch in (args.dim_batch, args.dim_batch // 2, args.dim_batch // 4): + if dim_batch < 2: + break + oom = False + try: + lens = jlens.fit( + model, shard_prompts, source_layers=layers, dim_batch=dim_batch, checkpoint_path=str(ckpt), resume=True + ) + except torch.cuda.OutOfMemoryError: + oom = True + if not oom: + break + log(f"OOM at dim_batch={dim_batch}; freeing and halving") + gc.collect() + torch.cuda.empty_cache() + if lens is None: + raise RuntimeError("all dim_batch fallbacks OOMed") + for layer, J in lens.jacobians.items(): + assert torch.isfinite(J).all(), f"non-finite Jacobian at layer {layer}" + tmp = str(out) + ".tmp" + lens.save(tmp) + os.replace(tmp, out) + log(f"shard {args.shard} done: {lens.n_prompts} prompts in {(time.time() - t0) / 3600:.2f} h -> {out}") + + +def cmd_merge(args) -> None: + jlens = _import_jlens() + + final = Path(args.out) + if final.exists(): + log(f"[resume] merged lens exists: {final}") + return + lenses = [jlens.JacobianLens.load(str(Path(args.shard_dir) / f"shard{i}.lens.pt")) for i in range(args.num_shards)] + merged = jlens.JacobianLens.merge(lenses) + # Convergence report: first half of the shards vs all of them. Small + # relative deltas mean the prompt budget has saturated. + half = jlens.JacobianLens.merge(lenses[: max(1, args.num_shards // 2)]) + for layer in sorted(merged.jacobians): + rel = (half.jacobians[layer].float() - merged.jacobians[layer].float()).norm() + rel = rel / merged.jacobians[layer].float().norm() + log(f"convergence layer {layer}: ||J_half - J_full|| / ||J_full|| = {rel:.4f}") + final.parent.mkdir(parents=True, exist_ok=True) + tmp = str(final) + ".tmp" + merged.save(tmp) + os.replace(tmp, final) + log(f"merged {sum(lens.n_prompts for lens in lenses)} prompts, {len(merged.jacobians)} layers -> {final}") + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + sub = p.add_subparsers(dest="cmd", required=True) + + sp = sub.add_parser("prompts", help="seeded streaming C4 prompt dump") + sp.add_argument("--prompts-json", required=True, help="output manifest ({'prompts': [...], ...})") + sp.add_argument("--n-prompts", type=int, default=1000) + sp.add_argument("--seed", type=int, default=0) + sp.add_argument("--min-chars", type=int, default=800) + sp.add_argument("--c4-config", default="en", help="C4 language config, e.g. en, zh, de") + sp.add_argument("--dataset", default="allenai/c4", help="HF dataset id") + sp.add_argument("--dataset-config", default=None, help="dataset config; falls back to --c4-config") + sp.add_argument("--dataset-split", default="train") + sp.add_argument("--revision", default=None, help="optional dataset revision pin (the paper runs used none)") + sp.add_argument("--text-field", default="text", help="record field holding the document text") + sp.set_defaults(fn=cmd_prompts) + + sp = sub.add_parser("smoke", help="tiny fit with finite / save-load / apply checks") + sp.add_argument("--model-id", required=True) + sp.add_argument("--prompts-json", required=True) + sp.add_argument("--out", required=True) + sp.add_argument("--dim-batch", type=int, default=16) + sp.set_defaults(fn=cmd_smoke) + + sp = sub.add_parser("fit", help="fit one prompt shard on one GPU") + sp.add_argument("--model-id", required=True) + sp.add_argument("--prompts-json", required=True) + sp.add_argument("--n-prompts", type=int, required=True) + sp.add_argument("--start", type=int, default=0, help="offset into the frozen prompt manifest") + sp.add_argument("--shard", type=int, required=True) + sp.add_argument("--num-shards", type=int, required=True) + sp.add_argument("--n-layers", type=int, default=12) + sp.add_argument("--dim-batch", type=int, default=16) + sp.add_argument("--ckpt-dir", required=True) + sp.add_argument("--out-dir", required=True) + sp.set_defaults(fn=cmd_fit) + + sp = sub.add_parser("merge", help="merge shard lenses into one lens file") + sp.add_argument("--shard-dir", required=True) + sp.add_argument("--num-shards", type=int, required=True) + sp.add_argument("--out", required=True) + sp.set_defaults(fn=cmd_merge) + + args = p.parse_args(argv) + args.fn(args) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/run/fit_ridge_lens.py b/scripts/run/fit_ridge_lens.py new file mode 100644 index 0000000..b3e513b --- /dev/null +++ b/scripts/run/fit_ridge_lens.py @@ -0,0 +1,281 @@ +#!/usr/bin/env python3 +"""Fit corpus-conditioned ridge-translator lenses, one per source layer (tuned-lens style). + +The second lens family of the cross-lens study (``tab:app-cross-lens-extension`` +rows "EN--ZH, ridge translator" and "Jacobian vs ridge, EN"; +``fig:cross-lens-extension``). Per source layer ``l`` the script solves the +bias-free ridge regression + + W_l = argmin_W || X_l W^T - Y ||_F^2 + lambda ||W||_F^2 + +where ``X_l`` are the block-output residuals at layer ``l`` and ``Y`` the +final-block residuals of the same tokens, over the same seeded C4 prompts used to +fit the corresponding Jacobian lens (``scripts/run/fit_jlens.py``). The result is +stored in the JacobianLens container (transport = residual @ W.T, fp16), so +``scripts/run/run_cross_lens_readouts.py`` and the aggregators consume it +unchanged. + +Protocol: + - source layers and d_model are copied from a reference Jacobian lens file + (``--layers-from``); + - fit on prompts [0, n_prompts - holdout), lambda selected per layer on the + holdout tail by translation R^2, then refit on all n_prompts prompts at the + selected lambda; + - lambda grid: g * tr(XtX) / (N * d) for g in {1e-3, 1e-2, 1e-1, 1, 10}; + - diagnostics reported, not enforced: per-layer holdout R^2, and top-1 + agreement between decode(transport(h_l)) and the model's final logits on the + last 5 holdout prompts at the two deepest fitted layers (written to + ``.report.json``). + +Resumable: the second-moment accumulators checkpoint every 10 prompts (atomic), +and the script exits early if ``--out`` already exists. ``jlens`` is imported +lazily (``uv sync --extra lens``). + +Paper runs (Qwen3.5-9B, layer set copied from the English-fitted Jacobian lens):: + + fit_ridge_lens.py --model-id Qwen/Qwen3.5-9B --prompts-json /c4_prompts_en_seed0.json \ + --n-prompts 100 --holdout 10 --layers-from /qwen35_9b_jlens_en_seed0_n100.pt \ + --ckpt-dir --tag ridge_en --out /qwen35_9b_ridgelens_en_seed0_n100.pt + fit_ridge_lens.py --model-id Qwen/Qwen3.5-9B --prompts-json /c4_prompts_zh_seed0.json \ + --n-prompts 100 --holdout 10 --layers-from /qwen35_9b_jlens_en_seed0_n100.pt \ + --ckpt-dir --tag ridge_zh --out /qwen35_9b_ridgelens_zh_seed0_n100.pt +""" + +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path + +import torch + +LAMBDA_GRID = [1e-3, 1e-2, 1e-1, 1.0, 10.0] + + +def log(msg: str) -> None: + print(f"[ridge_lens] {msg}", flush=True) + + +def _import_jlens(): + try: + import jlens + except ImportError as e: + raise ImportError( + "jlens (the Jacobian-lens reference implementation, Apache-2.0, " + "github.com/anthropics/jacobian-lens) is not installed; install it with `uv sync --extra lens`" + ) from e + return jlens + + +def _activation_recorder(): + _import_jlens() + from jlens.hooks import ActivationRecorder + + return ActivationRecorder + + +def load_model(model_id: str): + import transformers + + jlens = _import_jlens() + hf = ( + transformers.AutoModelForCausalLM.from_pretrained( + model_id, torch_dtype=torch.bfloat16, attn_implementation="sdpa" + ) + .cuda() + .eval() + ) + tok = transformers.AutoTokenizer.from_pretrained(model_id) + return jlens.from_hf(hf, tok) + + +@torch.no_grad() +def accumulate(model, prompts, layers, final_layer, max_seq_len, ckpt_path, every=10): + """One pass over prompts, accumulating XtX / XtY / YtY traces per layer.""" + ActivationRecorder = _activation_recorder() + + record_at = sorted(set(layers) | {final_layer}) + + state = None + if ckpt_path.exists(): + state = torch.load(ckpt_path, map_location="cpu", weights_only=True) + log(f"[resume] accumulator checkpoint at prompt {state['next_prompt']}") + + acc = None + start = 0 + if state is not None: + acc = {l: {k: v.cuda() for k, v in per.items()} for l, per in state["acc"].items()} + start = state["next_prompt"] + + for pi in range(start, len(prompts)): + with ActivationRecorder(model.layers, at=record_at) as rec: + input_ids = model.encode(prompts[pi], max_length=max_seq_len) + model.forward(input_ids) + acts = {i: rec.activations[i].detach() for i in record_at} + y = acts[final_layer][0].float() # (seq, d_model) + if acc is None: + d = y.shape[1] + acc = { + l: { + "xtx": torch.zeros(d, d, device="cuda"), + "xty": torch.zeros(d, d, device="cuda"), + "ytr": torch.zeros((), device="cuda"), + "n": torch.zeros((), device="cuda"), + } + for l in layers + } + for l in layers: + x = acts[l][0].float() # (seq, d_model) + acc[l]["xtx"] += x.T @ x + acc[l]["xty"] += x.T @ y + acc[l]["ytr"] += (y * y).sum() + acc[l]["n"] += y.shape[0] + if (pi + 1) % every == 0 or pi == len(prompts) - 1: + tmp = ckpt_path.with_suffix(".tmp") + torch.save( + { + "next_prompt": pi + 1, + "acc": {l: {k: v.cpu() for k, v in per.items()} for l, per in acc.items()}, + }, + tmp, + ) + os.replace(tmp, ckpt_path) + log(f"prompt {pi + 1}/{len(prompts)} accumulated") + return acc + + +def solve(xtx: torch.Tensor, xty: torch.Tensor, lam: float) -> torch.Tensor: + d = xtx.shape[0] + reg = xtx + lam * torch.eye(d, device=xtx.device) + return torch.linalg.solve(reg, xty).T.contiguous() # transport = x @ W.T + + +def r2(acc_h, W) -> float: + """Holdout R^2 from accumulated second moments (vs zero baseline).""" + xtx, xty, ytr = acc_h["xtx"], acc_h["xty"], acc_h["ytr"] + sse = torch.einsum("ij,jk,ik->", W, xtx, W) - 2.0 * torch.einsum("ij,ji->", W, xty) + ytr + return float(1.0 - sse / ytr) + + +def main(argv: list[str] | None = None) -> int: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--model-id", default="Qwen/Qwen3.5-9B") + ap.add_argument("--prompts-json", required=True, help="seeded prompt dump from fit_jlens.py prompts") + ap.add_argument("--n-prompts", type=int, default=100) + ap.add_argument("--holdout", type=int, default=10) + ap.add_argument("--layers-from", required=True, help="reference Jacobian lens .pt (source layers, d_model)") + ap.add_argument("--max-seq-len", type=int, default=512) + ap.add_argument("--ckpt-dir", required=True, help="directory for the resumable accumulator checkpoints") + ap.add_argument("--out", required=True, help="output lens .pt (JacobianLens container)") + ap.add_argument("--tag", default="ridge", help="checkpoint file prefix inside --ckpt-dir") + args = ap.parse_args(argv) + + out = Path(args.out) + if out.exists(): + log(f"[resume] {out} exists; done") + return 0 + + ref = torch.load(args.layers_from, map_location="cpu", weights_only=True) + layers = sorted(int(l) for l in ref["source_layers"]) + d_model = int(ref["d_model"]) + log(f"layers={layers} d_model={d_model}") + + payload = json.loads(Path(args.prompts_json).read_text()) + prompts = payload["prompts"][: args.n_prompts] + n_fit = len(prompts) - args.holdout + log(f"{len(prompts)} prompts ({n_fit} fit + {args.holdout} holdout)") + + model = load_model(args.model_id) + final_layer = model.n_layers - 1 + ckpt_dir = Path(args.ckpt_dir) + ckpt_dir.mkdir(parents=True, exist_ok=True) + + acc_fit = accumulate( + model, + prompts[:n_fit], + layers, + final_layer, + args.max_seq_len, + ckpt_dir / f"{args.tag}_fit.ckpt", + ) + acc_hold = accumulate( + model, + prompts[n_fit:], + layers, + final_layer, + args.max_seq_len, + ckpt_dir / f"{args.tag}_hold.ckpt", + ) + + report, W_final = {}, {} + for l in layers: + scale = float(acc_fit[l]["xtx"].diagonal().sum() / (acc_fit[l]["n"] * d_model)) + best = None + for g in LAMBDA_GRID: + lam = g * scale + W = solve(acc_fit[l]["xtx"], acc_fit[l]["xty"], lam) + score = r2(acc_hold[l], W) + if best is None or score > best[2]: + best = (g, lam, score) + g, lam, score = best + # Refit on fit+holdout at the selected lambda (same lambda scale). + xtx = acc_fit[l]["xtx"] + acc_hold[l]["xtx"] + xty = acc_fit[l]["xty"] + acc_hold[l]["xty"] + W_final[l] = solve(xtx, xty, lam) + report[l] = {"lambda_mult": g, "lambda": lam, "holdout_r2": score} + log(f"layer {l}: g={g} holdout_R2={score:.4f}") + + # Diagnostic: top-1 agreement of decoded transported states vs model + # logits on the last 5 holdout prompts, two deepest fitted layers. + ActivationRecorder = _activation_recorder() + + deep = layers[-2:] + agree = {l: [0, 0] for l in deep} + with torch.no_grad(): + for prompt in prompts[-5:]: + with ActivationRecorder(model.layers, at=sorted(set(deep) | {final_layer})) as rec: + input_ids = model.encode(prompt, max_length=args.max_seq_len) + model.forward(input_ids) + acts = {i: rec.activations[i].detach() for i in rec.activations} + ml = model.unembed(acts[final_layer][0].float()).argmax(-1) + for l in deep: + tl = model.unembed((acts[l][0].float() @ W_final[l].T)).argmax(-1) + agree[l][0] += int((tl == ml.to(tl.device)).sum()) + agree[l][1] += int(ml.numel()) + for l in deep: + k, n = agree[l] + report[l]["deep_top1_agreement"] = k / max(n, 1) + log(f"layer {l}: transported-vs-model top1 agreement {k}/{n} = {k / max(n, 1):.3f}") + + out.parent.mkdir(parents=True, exist_ok=True) + tmp = out.with_suffix(".pt.tmp") + torch.save( + { + "J": {l: W_final[l].cpu().to(torch.float16) for l in layers}, + "n_prompts": len(prompts), + "source_layers": layers, + "d_model": d_model, + }, + tmp, + ) + os.replace(tmp, out) + Path(str(out) + ".report.json").write_text( + json.dumps( + { + "model_id": args.model_id, + "prompts_json": args.prompts_json, + "n_prompts": len(prompts), + "holdout": args.holdout, + "lambda_grid": LAMBDA_GRID, + "layers": {str(l): report[l] for l in layers}, + }, + indent=1, + ) + ) + log(f"saved {out}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/run/run_cross_lens_readouts.py b/scripts/run/run_cross_lens_readouts.py new file mode 100644 index 0000000..18619f5 --- /dev/null +++ b/scripts/run/run_cross_lens_readouts.py @@ -0,0 +1,322 @@ +#!/usr/bin/env python3 +"""Per-layer lens readings and readout decompositions over a prompt bank (raw dump). + +The producer for every cross-lens artifact: ``fig:cross-lens-butterfly``, +``fig:cross-lens-extension``, ``tab:app-cross-lens-families``, +``tab:app-cross-lens-extension`` and ``tab:app-cross-lens-de-antonym-layers``. +For each prompt in the bank the script records the greedy continuation, then per +fitted layer and per scored position: the lens's top-5 tokens with their logits, +and the Sparse Readout Prism decomposition of every probed target token's +transported logit (base term, feature sum, residual and the top signed feature +contributions), optionally also of the lens's own top-1 token +(``--decompose-top1``). Aggregation is offline: +``scripts/eval/aggregate_cross_lens_en_zh.py``, +``scripts/eval/aggregate_cross_lens_en_de.py``, +``scripts/analyze/cross_lens_three_lens_prompt.py`` and +``scripts/figures/compute_cross_lens_shared_feature.py`` all read this dump. + +Transport. ``--lens`` is a fitted lens ``.pt`` in the JacobianLens container (a +Jacobian lens from ``scripts/run/fit_jlens.py`` or a ridge translator from +``scripts/run/fit_ridge_lens.py``): the hooked block output ``h_l`` at each scored +position is transported with ``J_l @ h_l``, passed through the model's final norm, +and read through the LM head. ``--lens identity`` skips the transport (a plain +logit-lens reading on the same layer set, copied from ``--layers-from``). A +position-alignment guard checks on the first prompt that the hooked state +reproduces the lens's own logits at every layer. + +Decomposition basis. ``--sae`` is the k=128 seed-0 readout dictionary for +Qwen3.5-9B (8x, D=32768), Hugging Face ``hematteo/sparse-readout-prism`` file +``qwen3.5-9b/k128_8x/checkpoint.pt``, a different operating point from the main +tables' 32x/k=256 dictionaries. Rows of the LM head are centred on the vocabulary +mean and unit-normalised before encoding, and the contribution of feature ``i`` to +the score of token ``t`` is ``||W_t - mu|| * code_i(t) * (h . d_i)``. + +Bank schema: ``{"prompts": [{"id", "group", "prompt", "targets": [...], ...}]}`` +(see ``data/cross_lens/README.md``). Every target string is probed by the id of +its first token. ``jlens`` is imported lazily (``uv sync --extra lens``). + +Paper runs (Qwen3.5-9B, one run per lens and bank, final prompt position only):: + + run_cross_lens_readouts.py --model-id Qwen/Qwen3.5-9B --lens --sae \ + --k 128 --prompts data/cross_lens/cross_lens_prompts_en_zh.json --n-positions 1 \ + --decompose-top1 --out /en_zh__.json + + Lenses: the English-, Chinese- and German-fitted Jacobian lenses (n=100 and + n=300) and the English- and Chinese-fitted ridge translators. Banks: + ``data/cross_lens/cross_lens_prompts_en_zh.json`` (EN-ZH cells) and + ``data/cross_lens/cross_lens_prompts_en_de.json`` (EN-DE cells). The three-lens + worked example reads ``data/cross_lens/cross_lens_three_lens_prompts.json`` with + ``--top-feats 10`` under the EN, ZH and DE n=100 lenses, plus + ``--lens identity --layers-from `` for the lens-free reading. +""" + +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path + +import torch +import torch.nn.functional as F + +from sparse_readout_prism.research.run_io import run_provenance +from sparse_readout_prism.utils import write_json + + +def log(msg: str) -> None: + print(f"[readouts] {msg}", flush=True) + + +def _import_jlens(): + try: + import jlens + except ImportError as e: + raise ImportError( + "jlens (the Jacobian-lens reference implementation, Apache-2.0, " + "github.com/anthropics/jacobian-lens) is not installed; install it with `uv sync --extra lens`" + ) from e + return jlens + + +def load_sae(checkpoint: Path): + """Raw TopK dictionary tensors from a training checkpoint (decoder rows unit-normalised).""" + ckpt = torch.load(checkpoint, map_location="cpu", weights_only=True) + state = ckpt["model_state_dict"] + decoder = state["decoder"].float().contiguous() # (d_features, d_model) + encoder_w = state["encoder.weight"].float().contiguous() # (d_features, d_model) + encoder_b = state["encoder.bias"].float().contiguous() # (d_features,) + decoder = decoder / decoder.norm(dim=1, keepdim=True).clamp_min(1e-8) + return decoder, encoder_w, encoder_b + + +def encode_topk(x, encoder_w, encoder_b, k): + acts = F.relu(x @ encoder_w.T + encoder_b) + values, indices = torch.topk(acts, k=min(k, acts.shape[-1]), dim=-1) + code = torch.zeros_like(acts) + code.scatter_(dim=-1, index=indices, src=values) + return code + + +@torch.no_grad() +def feature_top_tokens(W, row_mean, feature_ids, encoder_w, encoder_b, tokenizer, top_tokens=12, chunk=8192): + """Top unembedding rows (by encoder activation) for each feature id, as decoded strings.""" + if not feature_ids: + return {} + device = W.device + fids = torch.tensor(sorted(feature_ids), dtype=torch.long) + enc = encoder_w[fids].to(device) + bias = encoder_b[fids].to(device) + best_scores = torch.full((len(fids), top_tokens), -float("inf"), device=device) + best_ids = torch.zeros((len(fids), top_tokens), dtype=torch.long, device=device) + for start in range(0, W.shape[0], chunk): + rows = W[start : start + chunk].float() + centered = rows - row_mean + x = centered / centered.norm(dim=1, keepdim=True).clamp_min(1e-8) + scores = F.relu(x @ enc.T + bias).T + merged = torch.cat([best_scores, scores], dim=1) + ids = torch.arange(start, start + rows.shape[0], device=device).expand(len(fids), -1) + merged_ids = torch.cat([best_ids, ids], dim=1) + best_scores, keep = torch.topk(merged, k=top_tokens, dim=1) + best_ids = torch.gather(merged_ids, 1, keep) + return { + int(fid): [tokenizer.decode([t]) for t, s in zip(best_ids[i].tolist(), best_scores[i].tolist()) if s > 0] + for i, fid in enumerate(fids.tolist()) + } + + +@torch.no_grad() +def run(args) -> dict: + import transformers + + jlens = _import_jlens() + + hf = transformers.AutoModelForCausalLM.from_pretrained(args.model_id, dtype=torch.bfloat16).cuda().eval() + tok = transformers.AutoTokenizer.from_pretrained(args.model_id) + model = jlens.from_hf(hf, tok) + # Lens-free control: J = I at every layer, i.e. no transport at all. This + # is the floor the cross-lens agreement has to beat: how often two surface + # forms share a dominant readout feature when no lens is involved. + identity = args.lens == "identity" + if identity: + ref = jlens.JacobianLens.load(args.layers_from) + layers = sorted(ref.jacobians) + del ref + lens, J = None, None + else: + lens = jlens.JacobianLens.load(args.lens) + layers = sorted(lens.jacobians) + J = {l: lens.jacobians[l].float().cuda() for l in layers} + W = hf.get_output_embeddings().weight.detach().float().cuda() # (vocab, d_model) + row_mean = W.mean(dim=0) + final_norm = hf.model.norm + decoder, encoder_w, encoder_b = load_sae(Path(args.sae)) + decoder, encoder_w, encoder_b = decoder.cuda(), encoder_w.cuda(), encoder_b.cuda() + + spec = json.loads(Path(args.prompts).read_text())["prompts"] + positions = list(range(-args.n_positions, 0)) + log(f"{len(spec)} prompts, layers={layers}, positions={positions}") + + captured = {} + + def make_hook(idx): + def hook(module, inputs, output): + captured[idx] = (output[0] if isinstance(output, tuple) else output).detach() + + return hook + + handles = [model.layers[l].register_forward_hook(make_hook(l)) for l in layers] + + def decompose(h_state, token_id): + W_row = W[token_id] + centered = W_row - row_mean + norm = centered.norm().clamp_min(1e-8) + code = encode_topk((centered / norm)[None, :], encoder_w, encoder_b, args.k)[0] + contributions = norm * code * (h_state @ decoder.T) + base = float(h_state @ row_mean) + feat_sum = float(contributions.sum()) + original = float(h_state @ W_row) + active = torch.nonzero(contributions != 0).flatten() + top = active[contributions[active].abs().argsort(descending=True)[: args.top_feats]] + return { + "original_logit": original, + "base": base, + "feature_sum": feat_sum, + "residual": original - base - feat_sum, + "top_features": [{"id": int(f), "contribution": float(contributions[f])} for f in top], + } + + records = [] + seen = set() + for vi, v in enumerate(spec): + ids = tok(v["prompt"], return_tensors="pt").input_ids.cuda() + greedy = hf.generate(ids, max_new_tokens=8, do_sample=False)[0, ids.shape[1] :] + continuation = tok.decode(greedy) + # First token id of each target string (multi-token targets probed by head). + target_ids = {} + for t in v["targets"]: + t_ids = tok(t, add_special_tokens=False).input_ids + if t_ids: + target_ids[t] = t_ids[0] + + if identity: + with torch.no_grad(): + hf(ids) + per_layer = {} + for l in layers: + hs = torch.stack([final_norm(captured[l][0, pos].to(torch.bfloat16)).float() for pos in positions]) + per_layer[l] = hs @ W.T + else: + per_layer, final_logits, _ = lens.apply(model, v["prompt"], positions=positions) + rec = { + "id": v["id"], + "group": v["group"], + "continuation": continuation, + "intermediate": v.get("intermediate"), + "answer": v.get("answer"), + "targets": {t: int(i) for t, i in target_ids.items()}, + "layers": {}, + } + for l in layers: + ll = per_layer[l].float().cuda() # (n_positions, vocab) + lrec = {} + for pi, pos in enumerate(positions): + src = captured[l][0, pos].float() + transported = src if identity else J[l] @ src + h_t = final_norm(transported.to(torch.bfloat16)).float() + logits_p = ll[pi] + if vi == 0: + # Position-alignment guard: the hooked state at `pos` must + # reproduce lens.apply's logits at slot `pi`, else the + # decomposition and the lens fields are silently misaligned. + rel = float((h_t @ W.T - logits_p).abs().max() / logits_p.abs().max().clamp_min(1e-6)) + if rel > 5e-2: + raise RuntimeError(f"position misalignment at layer {l} pos {pos}: rel={rel:.3f}") + top5 = logits_p.topk(5) + entry = { + "top5": [ + (tok.decode([t]), round(float(s), 2)) + for t, s in zip(top5.indices.tolist(), top5.values.tolist()) + ], + "targets": {}, + } + if args.decompose_top1: + d1 = decompose(h_t, int(top5.indices[0])) + d1["top_features"] = d1["top_features"][:8] + entry["top1_decomp"] = d1 + seen.update(f["id"] for f in d1["top_features"]) + for t, tid in target_ids.items(): + d = decompose(h_t, tid) + d["lens_logit"] = float(logits_p[tid]) + d["lens_rank"] = int((logits_p > logits_p[tid]).sum()) + entry["targets"][t] = d + seen.update(f["id"] for f in d["top_features"]) + lrec[str(pos)] = entry + rec["layers"][str(l)] = lrec + records.append(rec) + log(f"{vi + 1}/{len(spec)}: {v['id']} -> {continuation!r}") + + for h in handles: + h.remove() + log(f"labeling {len(seen)} unique features") + labels = feature_top_tokens(W, row_mean, seen, encoder_w, encoder_b, tok) + + return { + "model": args.model_id, + "lens": args.lens, + "sae": args.sae, + "positions": positions, + "layers": layers, + "records": records, + "feature_top_tokens": labels, + } + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--model-id", default="Qwen/Qwen3.5-9B") + p.add_argument( + "--lens", + required=True, + help="path to a fitted lens .pt (JacobianLens container), or the literal 'identity' for the " + "lens-free control (no transport; plain logit-lens readout)", + ) + p.add_argument( + "--layers-from", + default=None, + help="required with --lens identity: a fitted lens .pt whose layer set the control copies, " + "so the comparison is layer-matched", + ) + p.add_argument("--sae", required=True, help="readout SAE checkpoint.pt (paper: Qwen3.5-9B k=128 seed 0, 8x)") + p.add_argument("--k", type=int, default=128, help="active features per row code") + p.add_argument("--prompts", required=True, help="prompt bank JSON ({'prompts': [...]})") + p.add_argument("--n-positions", type=int, default=4, help="score the last n prompt positions") + p.add_argument("--decompose-top1", action="store_true", help="also decompose the lens's own top-1 token") + p.add_argument("--top-feats", type=int, default=10, help="signed feature contributions kept per decomposition") + p.add_argument("--out", required=True, help="output dump JSON") + args = p.parse_args(argv) + if args.lens == "identity" and not args.layers_from: + p.error("--lens identity requires --layers-from ") + + payload = run(args) + + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + tmp = out.with_suffix(".json.tmp") + tmp.write_text(json.dumps(payload, ensure_ascii=False)) + os.replace(tmp, out) + write_json( + { + "n_records": len(payload["records"]), + "layers": payload["layers"], + "positions": payload["positions"], + "provenance": run_provenance(args), + }, + out.with_suffix(".manifest.json"), + ) + log(f"wrote {len(payload['records'])} prompts -> {out}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/run/run_query_fidelity_bank.py b/scripts/run/run_query_fidelity_bank.py index f208b9b..9417dd8 100644 --- a/scripts/run/run_query_fidelity_bank.py +++ b/scripts/run/run_query_fidelity_bank.py @@ -522,18 +522,27 @@ def _group_metrics(rows: list[dict]) -> dict: rd = np.array([r["residual_direct"] for r in rows], float) ar = np.array([r["abs_residual"] for r in rows], float) sign = np.array([r["sign_match"] for r in rows], bool) + # Floored relative error rho_0.5 = |m_exact - m_sparse| / (|m_exact| + 0.5): the + # per-model reconstruction table with cluster-bootstrap intervals + # (tab:app-fidelity-cis) reports its median and its < 0.5 rate alongside sign. + ex_arr = np.array(ex, float) + sp_arr = np.array(sp, float) + rho5 = np.abs(ex_arr - sp_arr) / (np.abs(ex_arr) + 0.5) bases = sorted({r["base_case_id"] for r in rows}) - # bootstrap by base_case_id + # bootstrap by base_case_id (one draw sequence; every statistic below is + # evaluated on the same resamples) rng = np.random.default_rng(0) by_base: dict[str, list[int]] = {} for i, r in enumerate(rows): by_base.setdefault(r["base_case_id"], []).append(i) - boot_sign, boot_med = [], [] + boot_sign, boot_med, boot_med5, boot_pass5 = [], [], [], [] for _ in range(400): pick = rng.choice(bases, size=len(bases), replace=True) idx = [i for b in pick for i in by_base[b]] boot_sign.append(sign[idx].mean()) boot_med.append(np.median(rd[idx])) + boot_med5.append(np.median(rho5[idx])) + boot_pass5.append((rho5[idx] < 0.5).mean()) return { "n_base_cases": len(bases), "n_rows": len(rows), @@ -551,6 +560,12 @@ def _group_metrics(rows: list[dict]) -> dict: "pass_rd_lt_0.25": float((rd < 0.25).mean()), "pass_rd_lt_0.50": float((rd < 0.50).mean()), "pass_rd_lt_1.00": float((rd < 1.00).mean()), + "median_rho5": float(np.median(rho5)), + "median_rho5_ci_lo": float(np.percentile(boot_med5, 2.5)), + "median_rho5_ci_hi": float(np.percentile(boot_med5, 97.5)), + "pass_rho5_lt_0.50": float((rho5 < 0.5).mean()), + "pass_rho5_lt_0.50_ci_lo": float(np.percentile(boot_pass5, 2.5)), + "pass_rho5_lt_0.50_ci_hi": float(np.percentile(boot_pass5, 97.5)), "sign_flips": int((~sign).sum()), "tiny_margin_rows": int(sum(r["tiny_margin"] for r in rows)), } diff --git a/scripts/run/run_qwen_profanity_suppression_eval.py b/scripts/run/run_qwen_profanity_suppression_eval.py index d24812b..a4c069f 100644 --- a/scripts/run/run_qwen_profanity_suppression_eval.py +++ b/scripts/run/run_qwen_profanity_suppression_eval.py @@ -14,13 +14,43 @@ not evaluate toxicity or alignment. Produces the data behind the Appendix-J lexical-edit / readout-side control -stress test (paper Section 5 edit-stress, ``sec:srp-edit-stress``): +stress test (paper Section 5 edit-stress, ``sec:srp-edit-stress``; +``tab:lexical-control-stress-test`` lists the checks): ``baseline_comparison.csv`` / ``candidate_constrained_summary.csv`` -> ``tab:lexical-control-primary-methods`` (single model, primary methods at the selected operating point). Rerun once per model (Qwen3.5-2B, DeepSeek-R1-Distill-Qwen-7B at 32x/k256; Ministral-3-8B at 16x/k128 -- the checkpoints enumerated in ``configs/registries/result1_query_fidelity_cluster.yaml``) for ``tab:lexical-control-cross-model-results``. See ``docs/REPRODUCE.md``. + +Matched-KL frontier (``fig:lexical-matched-kl-frontier``, Appendix J +"Matched-KL Frontier and Cross-Model Outcome"). The candidate grid is 20 +prompts x 17 profane-to-reference pairs (5 discovery + up to 12 held-out pairs; +pairs whose terms fail the single-token filter are dropped per tokenizer) x 14 +scales (0.5 .. 64) x 8 methods: none, the readout SAE direction edit +(``feature_suppression``), discovery / oracle token bias, norm-matched random +features, and three label-free category-direction controls built from the five +discovery rows of W_U (``mean_row_direction``, ``pca_group_direction`` = PCA +rank-1, ``pca_group_rank4`` = the full rank of the discovery rows). The +frontier's metrics are the ``split == heldout`` rows of +``candidate_constrained_summary.csv`` (``median_kl_bits`` against +``mean_bad_prob_reduction`` / ``candidate_flip_rate`` per method and scale); +the paired differences at matched KL are computed from +``candidate_constrained_rows.csv`` by ``scripts/eval/paired_matched_kl_bootstrap.py``. + +Paper command line, once per model (Qwen/Qwen3.5-2B, Qwen/Qwen3.5-0.8B, +Qwen/Qwen3.5-9B, deepseek-ai/DeepSeek-R1-Distill-Qwen-7B; 32x/k256 checkpoints):: + + python scripts/run/run_qwen_profanity_suppression_eval.py \ + --checkpoint --model-id --out-dir \ + --device cuda --dtype bfloat16 --batch-size 8 --max-open-prompts 0 + +The primary-methods / cross-model tables were produced on the original grid +(16 prompts x 9 pairs, 7200 rows over the 10 scales up to 16 and 5 methods); +``--prompt-limit 16 --pair-limit 9`` restricts this script to that grid (the +added prompts and pairs are appended after the originals), and that run also +passed ``--max-open-prompts 6 --num-samples-per-prompt 2 --max-new-tokens 24`` +for the open-generation pilot. """ from __future__ import annotations @@ -140,6 +170,14 @@ "fucked": "fuck", "sucks": "suck", "bastard": "insult", + "fucking": "fuck", + "shitty": "shit_crap", + "damned": "hell_damn", + "asshole": "insult", + "bitch": "insult", + "piss": "insult", + "pissed": "insult", + "bloody": "hell_damn", } PAIRS = [ ("hell", "heck"), @@ -151,6 +189,18 @@ ("fucked", "messed"), ("sucks", "bad"), ("bastard", "fool"), + # Expanded held-out pairs (statistical power for the matched-KL comparison): + # all bad terms are outside DISCOVERY_TERMS, so every added pair lands in + # the held-out split. Pairs whose terms fail the single-token filter are + # dropped per-tokenizer by build_eval_cases_with_tokenizer. + ("fucking", "really"), + ("shitty", "lousy"), + ("damned", "doomed"), + ("asshole", "idiot"), + ("bitch", "witch"), + ("piss", "whine"), + ("pissed", "angry"), + ("bloody", "total"), ] PROMPTS = [ "What the", @@ -169,8 +219,15 @@ "You really", "They totally", "The angry gamer yelled", + "Honestly this whole thing is", + "Stop being such a", + "The whole project went to", + "Frankly the update was", ] -SCALES = [0.5, 1, 2, 3, 4, 6, 8, 10, 12, 16] +# Scales 24-64 extend the rank-1 category-direction baselines to the KL range +# the 10-direction SRP edit reaches, so the matched-KL frontier is compared +# over a shared KL range rather than truncated at the baselines' low end. +SCALES = [0.5, 1, 2, 3, 4, 6, 8, 10, 12, 16, 24, 32, 48, 64] BENIGN_PROMPTS = [ 'Define the word "hell" in a dictionary style:', 'In the sentence "what the hell", the word hell is', @@ -284,8 +341,9 @@ def load_all(args: argparse.Namespace) -> Loaded: decoder = decoder.to(tensor_device) encoder_w = encoder_w.to(tensor_device) encoder_b = encoder_b.to(tensor_device) - label_ids = sorted({fid for ids in FEATURE_SETS.values() for fid in ids}) - log(f"tensor_device={tensor_device}; labeling {len(label_ids)} FEATURE_SETS ids") + n_features = decoder.shape[0] + label_ids = sorted({fid for ids in FEATURE_SETS.values() for fid in ids if fid < n_features}) + log(f"tensor_device={tensor_device}; labeling {len(label_ids)} FEATURE_SETS ids (D={n_features})") labels = display_label_features( W=W_U, row_mean=row_mean, @@ -341,6 +399,39 @@ def suppress_vec(decoder: torch.Tensor, feature_ids: list[int], scale: float) -> return -float(scale) * directions.sum(dim=0) +def category_direction_units( + W_U: torch.Tensor, discovery_ids: list[int], row_mean: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Mean-row and top-PCA unit directions from the discovery tokens' unembedding + rows, centered by the shared row mean (parity with the readout SAE, which is + trained on row-mean-centered rows). Both point toward the discovery side, so + subtracting a positive multiple suppresses those tokens (matching the sign + convention of ``suppress_vec``). These are the label-free "mean row / PCA + directions of the same token group" controls.""" + idx = torch.tensor(discovery_ids, dtype=torch.long, device=W_U.device) + rows = W_U[idx].float() # (n_disc, d_model) + centered = rows - row_mean.to(rows.dtype) # deviations from the global row mean + m = centered.mean(dim=0) # (d_model,) + mean_dir = m / m.norm().clamp_min(1e-12) + # Top principal component about the global-mean origin (uncentered SVD; rank + # <= n_disc). Deterministic: torch.linalg.svd, not the randomized pca_lowrank. + _u, _s, vh = torch.linalg.svd(centered, full_matrices=False) + pc = vh[0] # (d_model,) + if torch.dot(pc, mean_dir) < 0: # orient toward the discovery ("profane") side + pc = -pc + pca_dir = pc / pc.norm().clamp_min(1e-12) + # PC set of rank <= 4 (capped at the number of discovery rows): the + # capacity-fair "PCA directions" of the token group, each unit-norm and + # oriented toward the discovery side. + n_pc = min(4, vh.shape[0]) + pca_dirs = vh[:n_pc].clone() + for j in range(n_pc): + if torch.dot(pca_dirs[j], mean_dir) < 0: + pca_dirs[j] = -pca_dirs[j] + pca_dirs[j] = pca_dirs[j] / pca_dirs[j].norm().clamp_min(1e-12) + return mean_dir, pca_dir, pca_dirs + + def random_feature_ids(n_features: int, n_total: int, seed: int = 17) -> list[int]: rng = random.Random(seed) return rng.sample(range(n_total), k=n_features) @@ -371,6 +462,9 @@ def apply_method( scale: float, token_ids: list[int], audited_ids: list[int], + mean_dir: torch.Tensor | None = None, + pca_dir: torch.Tensor | None = None, + pca_dirs: torch.Tensor | None = None, ) -> tuple[torch.Tensor, dict[str, float]]: if method == "none": return logits.clone(), {"token_bias_strength": 0.0, "intervention_norm": 0.0} @@ -388,6 +482,30 @@ def apply_method( "token_bias_strength": 0.0, "intervention_norm": float(v.norm().item()), } + if method == "mean_row_direction": + if mean_dir is None: + raise ValueError("mean_row_direction requires mean_dir") + v = -float(scale) * mean_dir + return logits + v @ W_U.T, { + "token_bias_strength": 0.0, + "intervention_norm": float(v.norm().item()), + } + if method == "pca_group_direction": + if pca_dir is None: + raise ValueError("pca_group_direction requires pca_dir") + v = -float(scale) * pca_dir + return logits + v @ W_U.T, { + "token_bias_strength": 0.0, + "intervention_norm": float(v.norm().item()), + } + if method == "pca_group_rank4": + if pca_dirs is None: + raise ValueError("pca_group_rank4 requires pca_dirs") + v = -float(scale) * pca_dirs.sum(dim=0) # summed unit PCs, as suppress_vec + return logits + v @ W_U.T, { + "token_bias_strength": 0.0, + "intervention_norm": float(v.norm().item()), + } if method.startswith("token_logit_bias"): v_ref = suppress_vec(decoder, feature_ids, scale) feature_delta = v_ref @ W_U.T @@ -567,6 +685,9 @@ def candidate_rows( discovery_ids = [token_ids[t] for t in DISCOVERY_TERMS if t in token_ids] all_profane_ids = [token_ids[t] for t in PROFANITY_TERMS if t in token_ids] random_ids = random_feature_ids(len(feature_ids), loaded.decoder.shape[0], seed=args.random_seed) + mean_dir, pca_dir, pca_dirs = (None, None, None) + if discovery_ids: + mean_dir, pca_dir, pca_dirs = category_direction_units(loaded.W_U, discovery_ids, loaded.row_mean) unique_prompts = sorted({case["prompt"] for case in cases}) log(f"computing next-token logits for {len(unique_prompts)} unique prompts in batches of {args.batch_size}") prompt_logits = next_logits_batch( @@ -615,6 +736,9 @@ def candidate_rows( scale=scale, token_ids=ids, audited_ids=audited_ids, + mean_dir=mean_dir, + pca_dir=pca_dir, + pca_dirs=pca_dirs, ) edit_margin = float((edited[bad_id] - edited[good_id]).item()) edit_prob = pair_prob(edited, bad_id, good_id) @@ -1048,6 +1172,7 @@ def main() -> int: write_csv(args.out_dir / "feature_token_rows.csv", feature_token_rows) write_csv(args.out_dir / "feature_summary.csv", feature_summary) + n_features = loaded.decoder.shape[0] selected_features = select_discovery_features(feature_token_rows, args.selected_feature_count) if not selected_features: if not args.allow_feature_set_fallback: @@ -1057,9 +1182,10 @@ def main() -> int: "for this model — pass --allow-feature-set-fallback only if that is the " "checkpoint you are evaluating" ) - selected_features = FEATURE_SETS["all_narrow"] - selected_features = selected_features[: args.selected_feature_count] - log(f"selected discovery features: {selected_features}") + # Qwen-2B ids; only valid on dictionaries wide enough to contain them. + selected_features = [f for f in FEATURE_SETS["all_narrow"] if f < n_features] + selected_features = [f for f in selected_features if f < n_features][: args.selected_feature_count] + log(f"selected discovery features (D={n_features}): {selected_features}") cases = build_eval_cases_with_tokenizer(loaded.tokenizer, token_ids) if args.prompt_limit is not None: @@ -1074,6 +1200,9 @@ def main() -> int: "token_logit_bias_discovery", "token_logit_bias_oracle", "random_feature_suppression", + "mean_row_direction", + "pca_group_direction", + "pca_group_rank4", ] log(f"evaluating {len(cases)} candidate-constrained rows across {len(methods)} methods") cand = candidate_rows(loaded, cases, token_ids, selected_features, methods, SCALES, args) diff --git a/scripts/run/run_readout_baseline_comparisons.py b/scripts/run/run_readout_baseline_comparisons.py index ce58e33..ae6b6dd 100644 --- a/scripts/run/run_readout_baseline_comparisons.py +++ b/scripts/run/run_readout_baseline_comparisons.py @@ -14,6 +14,37 @@ random_support_same_magnitudes - support-randomization null pca_64, pca_256, pca_1024 - dense PCA references nearest_row_ridge_top128 - lexical / token-prototype reference + knn_basis_top128 - zero-fit weighted kNN over the nearest rows + row_cluster_d_k - k-means centroid dictionary, top-k restricted LS code + row_cluster_hard_d - single-centroid (hard cluster assignment) variant + +The last three are the direct-geometry alternatives of the paper's baseline +table (`tab:readout-score-fidelity-summary`, full grid +`tab:app-direct-geometry-grid`, error tails `tab:app-error-tails` via +`scripts/eval/analyze_error_tails.py`); the first seven form the Qwen3.5-2B +reference panel (`fig:app-robustness-baseline-comparisons`) and the two nulls +the cross-model null table (`tab:app-cross-model-nulls`). + +Paper runs (all with --bank-dir data/query_banks --banks curated_ab,case_candidates,model_native +--model-native-file qwen_gemma_result1_model_native_prompts_c4.jsonl --operating-point fidelity +--max-native 500 --max-curated 320 --max-cases 40 --max-len 64 --seed 0). The model-native bank of +these runs is the C4-sampled slice (regenerate it as in data/query_banks/README.md; the first 500 +base cases in file order are kept), so each model's full logit-difference bank is ~1,350 rows over +~850 base-case clusters (500 C4 prompts x 2 native queries + the curated and case-candidate rows): + + # Qwen3.5-2B reference panel (fig:app-robustness-baseline-comparisons) + --model Qwen3.5-2B --methods sparse_rp,shuffled_row_code,random_support_same_magnitudes,\ + pca_64,pca_256,pca_1024,nearest_row_ridge_top128 + # Null controls, one run per model in Qwen3.5-0.8B/2B/9B, Gemma-4-E2B/E4B (tab:app-cross-model-nulls) + --model --methods sparse_rp,shuffled_row_code,random_support_same_magnitudes + # Direct-geometry grid, one run per softcap-free readout: Qwen3.5-0.8B/2B/9B, + # Ministral-3-8B-Base, R1-Distill-Qwen-7B, R1-Distill-Llama-8B (tab:app-direct-geometry-grid) + --model --methods sparse_rp,nearest_row_ridge_top128,knn_basis_top128,\ + row_cluster_d16384_k256,row_cluster_d65536_k256,row_cluster_hard_d65536,pca_256 + +An uncurated audit population (random base cases, top-20 signed SRP contributions +per row) uses `--methods sparse_rp --max-native 50 --max-curated 80 --max-cases 20 +--seed 20260712 --sample-bases-random --audit-export-top 20`. Output schema: cells/__/{rows.jsonl, summary_cell.json, done.json} @@ -49,8 +80,9 @@ from sparse_readout_prism.factorizers import TopKSAE, load_factorizer from sparse_readout_prism.research.cell_metrics import coverage_stats from sparse_readout_prism.research.run_io import ( + EPS, group_rows as _by, - load_bank, + load_bank as _load_bank_prefix, margin_row_stats, resolve_ab_case, run_provenance, @@ -518,6 +550,181 @@ def decompose_row( ) +# --------------------------------------------------------------------------- # +# F. Zero-fit k-nearest-row basis +# --------------------------------------------------------------------------- # + + +class KNNBasisMethod(BaselineMethod): + """Top-K nearest centered-normalized vocabulary rows as a *zero-fit* basis: + reconstruct the target row from its neighbours weighted by their cosine + similarity, with no fitted regression. Probes the bare neighbourhood + structure (the top-k nearest rows of each token as a feature basis) rather + than a regression onto it -- the zero-fit counterpart of the fitted + NearestRowRidgeMethod. + + Design note: a plain sum of cosine-weighted neighbours over 128 rows + over-reconstructs the target by an order of magnitude, so the zero-fit + estimate is the similarity-weighted neighbourhood mean rescaled by one + scalar, the target's projection onto that mean (no per-neighbour fit).""" + + name = "knn_basis" + + def __init__(self, *, top_k: int = 128, **kw) -> None: + super().__init__(**kw) + self.top_k = int(top_k) + self.name = f"knn_basis_top{self.top_k}" + + @torch.no_grad() + def decompose_row( + self, + h: torch.Tensor, + row_idx: int, + exclude_ids: Optional[Iterable[int]] = None, + ) -> BaselineDecomposition: + x = self.row_normalized_all[row_idx] # (d_model,) + sims = self.row_normalized_all @ x # (V,) + excl = {int(row_idx)} + if exclude_ids is not None: + excl.update(int(i) for i in exclude_ids) + mask = torch.zeros_like(sims, dtype=torch.bool) + mask[list(excl)] = True + sims = sims.masked_fill(mask, -float("inf")) + top = torch.topk(sims, k=self.top_k) + nbr_ids = top.indices # (top_k,) + w = top.values.clamp_min(0.0) + w = w / w.sum().clamp_min(1e-8) # (top_k,) convex weights + X = self.row_normalized_all[nbr_ids] # (top_k, d_model) + mean_nbr = w @ X # (d_model,) weighted neighbourhood mean + gamma = (x @ mean_nbr) / (mean_nbr @ mean_nbr).clamp_min(1e-8) + coeffs = gamma * w # per-neighbour coefficients; decoded = coeffs @ X + decoded = coeffs @ X # (d_model,) rescaled neighbourhood estimate + W_row = self.W[row_idx].to(self.device) + row_norm = self.row_norm_all[row_idx] + reconstructed_row = self.row_mean + row_norm * decoded + residual_row = W_row - reconstructed_row + h_nbr = X @ h # (top_k,) + feature_contributions = row_norm * coeffs * h_nbr # (top_k,) + base_term = h @ self.row_mean + feature_sum = feature_contributions.sum() + residual_term = h @ residual_row + original_logit = h @ W_row + reconstructed_logit = base_term + feature_sum + identity_error = original_logit - (base_term + feature_sum + residual_term) + return BaselineDecomposition( + original_logit=original_logit, + base_term=base_term, + feature_contributions=feature_contributions, + feature_sum=feature_sum, + residual_term=residual_term, + reconstructed_logit=reconstructed_logit, + identity_error=identity_error, + active_feature_indices=nbr_ids.detach().cpu(), + method=self.name, + ) + + +# --------------------------------------------------------------------------- # +# G. K-means row-cluster basis +# --------------------------------------------------------------------------- # + + +def _torch_kmeans_unit(X: torch.Tensor, n_clusters: int, seed: int, iters: int = 12, chunk: int = 4096) -> torch.Tensor: + """Lloyd's k-means on unit-norm rows (cosine assignment), on-device. + Returns unit-normalized centroids (n_clusters, d). Dead centroids are + reseeded from random rows each iteration. Deterministic given `seed`.""" + V = X.shape[0] + g = torch.Generator().manual_seed(seed) + C = X[torch.randperm(V, generator=g)[:n_clusters].to(X.device)].clone() + ones = torch.ones(V, device=X.device) + for _ in range(iters): + assign = torch.empty(V, dtype=torch.long, device=X.device) + for s in range(0, V, chunk): + assign[s : s + chunk] = (X[s : s + chunk] @ C.T).argmax(dim=1) + C_new = torch.zeros_like(C) + count = torch.zeros(n_clusters, device=X.device) + C_new.index_add_(0, assign, X) + count.index_add_(0, assign, ones) + dead = count == 0 + C = C_new / count.clamp_min(1.0)[:, None] + n_dead = int(dead.sum()) + if n_dead: + ridx = torch.randperm(V, generator=g)[:n_dead].to(X.device) + C[dead] = X[ridx] + C = C / C.norm(dim=1, keepdim=True).clamp_min(1e-8) + return C + + +class RowClusterMethod(BaselineMethod): + """Sparsity-matched k-means centroid dictionary: fit `n_clusters` unit + centroids on the centered/unit rows, then code each row over its + top-`code_k` centroids by |cosine| via a restricted least-squares solve -- + the same dictionary size D and active budget k as the SRP factorizer, with + centroids in place of learned decoder directions. `hard=True` is the + literal clustering reading: the row's own single cluster centroid + reconstructs it. + + Design note: an earlier version of this baseline solved a full + least-squares projection over n_clusters = d_model centroids, a full-rank + reprojection of the row space that reconstructs any row near-perfectly by + construction; it was redesigned so the active budget stays far below + d_model. Do not reintroduce the full-span variant.""" + + name = "row_cluster" + + def __init__(self, *, n_clusters: int, code_k: int = 256, hard: bool = False, seed: int = 0, **kw) -> None: + super().__init__(**kw) + self.n_clusters = int(n_clusters) + self.hard = bool(hard) + self.code_k = 1 if hard else min(int(code_k), self.n_clusters) + self.name = f"row_cluster_hard_d{self.n_clusters}" if hard else f"row_cluster_d{self.n_clusters}_k{self.code_k}" + self.C = _torch_kmeans_unit(self.row_normalized_all, self.n_clusters, seed) + self._eye = torch.eye(self.code_k, device=self.device) + + @torch.no_grad() + def decompose_row( + self, + h: torch.Tensor, + row_idx: int, + exclude_ids: Optional[Iterable[int]] = None, + ) -> BaselineDecomposition: + x = self.row_normalized_all[row_idx] # (d_model,) + sims = self.C @ x # (n_clusters,) + if self.hard: + sel = sims.argmax().unsqueeze(0) # (1,) the row's own cluster + Ck = self.C[sel] # (1, d_model) + coeffs = sims[sel] # projection onto that centroid + else: + sel = torch.topk(sims.abs(), k=self.code_k).indices # (code_k,) + Ck = self.C[sel] # (code_k, d_model) + A = Ck @ Ck.T + 1e-4 * self._eye + coeffs = torch.linalg.solve(A, Ck @ x) # (code_k,) + decoded = coeffs @ Ck # (d_model,) + W_row = self.W[row_idx].to(self.device) + row_norm = self.row_norm_all[row_idx] + reconstructed_row = self.row_mean + row_norm * decoded + residual_row = W_row - reconstructed_row + h_c = Ck @ h # (code_k,) + feature_contributions = row_norm * coeffs * h_c # (code_k,) + base_term = h @ self.row_mean + feature_sum = feature_contributions.sum() + residual_term = h @ residual_row + original_logit = h @ W_row + reconstructed_logit = base_term + feature_sum + identity_error = original_logit - (base_term + feature_sum + residual_term) + return BaselineDecomposition( + original_logit=original_logit, + base_term=base_term, + feature_contributions=feature_contributions, + feature_sum=feature_sum, + residual_term=residual_term, + reconstructed_logit=reconstructed_logit, + identity_error=identity_error, + active_feature_indices=sel.detach().cpu(), + method=self.name, + ) + + # --------------------------------------------------------------------------- # # Factory # --------------------------------------------------------------------------- # @@ -541,6 +748,9 @@ def build_method( 'random_support_same_magnitudes' 'pca_64', 'pca_256', 'pca_1024' 'nearest_row_ridge_top128' + 'knn_basis_top128' + 'row_cluster_d65536_k256', 'row_cluster_d16384_k256' + 'row_cluster_hard_d65536' """ common = dict( W=W, @@ -567,6 +777,16 @@ def build_method( if spec.startswith("nearest_row_ridge_top"): top = int(spec.split("nearest_row_ridge_top", 1)[1]) return NearestRowRidgeMethod(top_k=top, **common) + if spec.startswith("knn_basis_top"): + top = int(spec.split("knn_basis_top", 1)[1]) + return KNNBasisMethod(top_k=top, **common) + if spec.startswith("row_cluster_hard_d"): + n = int(spec.split("row_cluster_hard_d", 1)[1]) + return RowClusterMethod(n_clusters=n, hard=True, seed=seed, **common) + if spec.startswith("row_cluster_d"): + body = spec.split("row_cluster_d", 1)[1] # "_k" + d_str, k_str = body.split("_k", 1) + return RowClusterMethod(n_clusters=int(d_str), code_k=int(k_str), seed=seed, **common) raise ValueError(f"unknown method spec: {spec}") @@ -662,6 +882,23 @@ def log(msg: str) -> None: print(f"[rbase {time.strftime('%H:%M:%S')}] {msg}", flush=True) +def load_bank(path: Path, cap: int | None, *, random_seed: int | None = None) -> list[dict]: + """Load a JSONL bank capped by *base case*. With ``random_seed`` None the cap + keeps the file-prefix base cases (the shared ``run_io.load_bank``); otherwise + it keeps a seeded random sample of base cases (``--sample-bases-random``), + rows staying in file order either way.""" + if random_seed is None: + return _load_bank_prefix(path, cap) + rows = [json.loads(line) for line in path.read_text().splitlines() if line.strip()] + if cap is not None and cap < len(rows): + bases = list(dict.fromkeys(r["base_case_id"] for r in rows)) + rng = np.random.default_rng(random_seed) + indices = rng.choice(len(bases), size=min(cap, len(bases)), replace=False) + keep_bases = [bases[i] for i in indices] + rows = [r for r in rows if r["base_case_id"] in set(keep_bases)] + return rows + + def _group_metrics(rows: list[dict]) -> dict: if not rows: return {} @@ -815,6 +1052,27 @@ def emit(base_id, case_id, family, query, mr, expected_side): row["expected_side"] = expected_side row["softcap"] = softcap row.update(coverage_stats(mr["feat_margin"], mr["sparse"])) + if args.audit_export_top > 0 and method_spec == "sparse_rp": + fm = mr["feat_margin"].detach().float().cpu() + order = torch.argsort(fm.abs(), descending=True) + total_abs = float(fm.abs().sum().item()) + keep = order[: min(args.audit_export_top, len(order))] + cumulative = 0.0 + features = [] + for fid in keep.tolist(): + contribution = float(fm[fid].item()) + mass = abs(contribution) / total_abs if total_abs > EPS else 0.0 + cumulative += mass + features.append( + { + "feature_id": int(fid), + "contribution": contribution, + "abs_mass_fraction": mass, + "cumulative_abs_mass_fraction": cumulative, + } + ) + row["audit_features"] = features + row["audit_top_abs_mass_fraction"] = cumulative rows.append(row) # Process bank rows. @@ -905,6 +1163,13 @@ def main() -> int: help="comma-separated method specs", ) ap.add_argument("--banks", type=str, default="curated_ab,case_candidates,model_native") + ap.add_argument( + "--model-native-file", + type=str, + default="qwen_gemma_result1_model_native_prompts.jsonl", + help="model_native bank file name under --bank-dir (paper baseline runs: the C4 slice " + "qwen_gemma_result1_model_native_prompts_c4.jsonl)", + ) ap.add_argument("--max-native", type=int, default=500) ap.add_argument("--max-curated", type=int, default=320) ap.add_argument("--max-cases", type=int, default=40) @@ -912,6 +1177,17 @@ def main() -> int: ap.add_argument("--max-len", type=int, default=64) ap.add_argument("--wu-tol", type=float, default=0.05) ap.add_argument("--seed", type=int, default=0) + ap.add_argument( + "--sample-bases-random", + action="store_true", + help="apply each bank cap to a seeded random base-case sample instead of the file prefix", + ) + ap.add_argument( + "--audit-export-top", + type=int, + default=0, + help="for sparse_rp, serialize this many largest-|contribution| features per decomposition", + ) ap.add_argument("--dry-run", action="store_true") args = ap.parse_args() @@ -931,7 +1207,7 @@ def main() -> int: bank_files = { "model_native": ( - "qwen_gemma_result1_model_native_prompts.jsonl", + args.model_native_file, args.max_native, ), "curated_ab": ( @@ -956,7 +1232,9 @@ def main() -> int: if not path.exists(): log(f"bank file missing, skipping {bname}: {path}") continue - banks[bname] = load_bank(path, cap) + bank_offsets = {"curated_ab": 11, "case_candidates": 23, "model_native": 37} + sample_seed = args.seed + bank_offsets[bname] if args.sample_bases_random else None + banks[bname] = load_bank(path, cap, random_seed=sample_seed) if not banks: log("no banks loaded; nothing to do") return 1 diff --git a/scripts/run/run_wsd_feature_alignment.py b/scripts/run/run_wsd_feature_alignment.py new file mode 100644 index 0000000..dc445cf --- /dev/null +++ b/scripts/run/run_wsd_feature_alignment.py @@ -0,0 +1,1236 @@ +#!/usr/bin/env python3 +"""Frozen sense-labelled readout representations for Sparse Readout Prism. + +Scores an ambiguous target word's selected logit at the final position of a +cloze prompt, decomposes it through a trained readout dictionary (top-k codes +of the centred, unit-normalised unembedding row; contribution_i = beta_i * +p_i(h)), and freezes the per-context representation bundle +(``representations.pt``: decoded state, projections, contributions, row +coefficients, feature ids, exact and reconstructed logits, target rank) that +the two sense analyses consume: + +* ``scripts/analyze/analyze_wsd_sense_groups.py`` -- per-sense feature groups, + permutation null, majority and no-selection references + (``tab:app-sense-alignment``, appendix ``app:sense-labelled-evaluation``). +* ``scripts/analyze/analyze_wsd_classifier_framing.py`` -- truncated-display + centroid classifier over signed contributions, unweighted projections, + row-coefficient shuffles and random features (the "classifier framing" + paragraph of the same appendix). + +Its own outputs back that appendix's data and coverage statements: +``audit.json`` (contexts kept and single-token coverage per vocabulary: 20/20 +words on the Qwen3.5 vocabularies, 13/20 on DeepSeek-R1-Distill-Llama-8B), +``scoring_summary.json`` (contexts scored, median target rank, row-gate +coverage) and ``metrics.json`` (full-vector centroid references computed on +the same bundle, with paired per-word bootstrap differences). + +Datasets +-------- +* CoarseWSD-20 (Loureiro, Rezaee, Pilehvar and Camacho-Collados, 2021, + Computational Linguistics 47(2), doi:10.1162/coli_a_00405) -- the dataset + used in the paper. ``--data-root`` is a checkout of + https://github.com/danlou/bert-disambiguation; the loader looks for + ``/data/CoarseWSD-20//{train,test}.{data,gold}.txt`` plus + ``class_map.txt`` (```` or ``/CoarseWSD-20`` also work). The + shipped train/test split is kept. The target occurrence is replaced by + ``[BLANK]`` and the sentence is wrapped in ``COARSEWSD_TEMPLATE``; the + target's leading-space token is scored at the prompt's final position, so + it is a genuine next-token readout even when the disambiguating evidence + originally followed the target. +* AmbiStory (SemEval-2026 Task 5): graded sense plausibility against + gloss-prompt anchors. Served by the same scoring path; not used in the paper. + +Paper runs (one per model, ``--seed 0``; the checkpoints are the 32x / k=256 +selected dictionaries for Qwen3.5-2B, Qwen3.5-9B and +DeepSeek-R1-Distill-Llama-8B; ``--batch-size 8`` for the 2B model and ``4`` +for the other two): + + python scripts/run/run_wsd_feature_alignment.py \ + --dataset coarsewsd20 --data-root \ + --model-id Qwen/Qwen3.5-2B --checkpoint \ + --out-dir results/wsd_core/coarsewsd20_qwen2b \ + --splits train,test --batch-size 8 --n-boot 5000 --seed 0 + +``--analyze-only --out-dir --n-boot 5000 --seed 0`` recomputes +``metrics.json`` from a frozen bundle without model inference. + +Gold labels are never used to fit the dictionary or to select features. +``--position-mode occurrence`` (score the sentence prefix so the target +occurrence itself is the next token) and ``--anchors`` (a frozen sense-anchor +token set, shipped as ``data/wsd/wsd_sense_anchors.json``, whose codes are +embedded in the bundle for an offline feature-typing analysis not included in +this repository) are variants the paper does not report; the paper's runs use +the cloze position and pass no anchors. +""" + +from __future__ import annotations + +import argparse +import ast +import hashlib +import itertools +import json +import platform +import re +import sys +import time +from collections import defaultdict +from pathlib import Path +from typing import Any, Iterable + +import numpy as np +import torch + +from sparse_readout_prism.research.run_io import run_provenance +from sparse_readout_prism.utils import find_lm_head, load_causal_lm, set_seed + + +AMBISTORY_STORY_TEMPLATE = """Read the story and recover the missing word. + +{story} + +The missing word is:""" + +AMBISTORY_GLOSS_TEMPLATE = """Read the meaning and name the matching word. + +Meaning: {gloss} + +The matching word is:""" + +COARSEWSD_TEMPLATE = """Read the sentence and recover the missing word. + +{sentence} + +The missing word is:""" + + +def log(message: str) -> None: + print(f"[wsd {time.strftime('%H:%M:%S')}] {message}", flush=True) + + +def stable_id(*parts: str, n: int = 16) -> str: + text = "\x1f".join(parts) + return hashlib.sha1(text.encode("utf-8")).hexdigest()[:n] + + +def write_json(path: Path, value: Any) -> None: + path.write_text(json.dumps(value, indent=2, ensure_ascii=False) + "\n") + + +def write_jsonl(path: Path, rows: Iterable[dict[str, Any]]) -> None: + with path.open("w") as handle: + for row in rows: + handle.write(json.dumps(row, ensure_ascii=False) + "\n") + + +def mask_exact_target(text: str, target: str) -> str | None: + """Replace the first case-insensitive whole target occurrence.""" + pattern = re.compile(rf"(? str | None: + masked = mask_exact_target(sentence.strip(), target) + if masked is None: + return None + pieces = [precontext.strip(), masked] + if ending.strip(): + pieces.append(ending.strip()) + return "\n".join(p for p in pieces if p) + + +def load_ambistory(root: Path, splits: list[str], limit: int | None) -> tuple[list[dict], dict]: + records: list[dict] = [] + anchors: dict[tuple[str, str], dict] = {} + skipped = defaultdict(int) + split_counts: dict[str, int] = {} + + for split in splits: + path = root / f"{split}.json" + if not path.exists(): + raise FileNotFoundError(path) + raw = json.loads(path.read_text()) + rows = list(raw.values()) if isinstance(raw, dict) else list(raw) + grouped: dict[tuple[str, str, str, str], list[dict]] = defaultdict(list) + for row in rows: + key = ( + str(row.get("homonym", "")).strip(), + str(row.get("precontext", "")).strip(), + str(row.get("sentence", "")).strip(), + str(row.get("ending", "")).strip(), + ) + grouped[key].append(row) + + made = 0 + for (target, precontext, sentence, ending), sense_rows in sorted(grouped.items()): + if len(sense_rows) != 2: + skipped["not_two_senses"] += 1 + continue + story = make_story(precontext, sentence, ending, target) + if story is None: + skipped["target_not_found"] += 1 + continue + ordered = sorted(sense_rows, key=lambda r: str(r.get("judged_meaning", ""))) + glosses = [str(r.get("judged_meaning", "")).strip() for r in ordered] + if not all(glosses) or glosses[0] == glosses[1]: + skipped["bad_gloss_pair"] += 1 + continue + averages = [float(r["average"]) for r in ordered] + stdevs = [float(r.get("stdev", 0.0)) for r in ordered] + setup_id = stable_id(target, precontext, sentence) + context_id = stable_id(target, precontext, sentence, ending) + anchor_keys = [] + for gloss in glosses: + akey = stable_id(target, gloss) + anchor_keys.append(akey) + anchors[(target, gloss)] = { + "item_id": f"anchor:{akey}", + "kind": "anchor", + "dataset": "ambistory", + "split": split, + "target": target, + "gloss": gloss, + "anchor_key": akey, + "prompt": AMBISTORY_GLOSS_TEMPLATE.format(gloss=gloss), + } + records.append( + { + "item_id": f"story:{split}:{context_id}", + "kind": "story", + "dataset": "ambistory", + "split": split, + "target": target, + "setup_id": setup_id, + "context_id": context_id, + "has_ending": bool(ending.strip()), + "glosses": glosses, + "anchor_keys": anchor_keys, + "ratings": averages, + "rating_stdevs": stdevs, + "human_margin": averages[0] - averages[1], + "prompt": AMBISTORY_STORY_TEMPLATE.format(story=story), + } + ) + made += 1 + if limit is not None and len(records) >= limit: + break + split_counts[split] = made + if limit is not None and len(records) >= limit: + break + + # Only score anchors actually referenced by retained story records. + used = {key for row in records for key in row["anchor_keys"]} + anchor_rows = [row for row in anchors.values() if row["anchor_key"] in used] + items = anchor_rows + records + audit = { + "dataset": "ambistory", + "splits": splits, + "n_story_contexts": len(records), + "n_anchors": len(anchor_rows), + "n_setups": len({r["setup_id"] for r in records}), + "n_targets": len({r["target"] for r in records}), + "split_story_counts": split_counts, + "skipped": dict(skipped), + } + return items, audit + + +def parse_class_map(path: Path) -> dict[str, Any]: + text = path.read_text().strip() + try: + value = ast.literal_eval(text) + if isinstance(value, dict): + return {str(k): v for k, v in value.items()} + except (ValueError, SyntaxError): + pass + mapping: dict[str, Any] = {} + for line in text.splitlines(): + fields = line.strip().split("\t") + if len(fields) >= 2: + mapping[fields[0]] = fields[1] + return mapping + + +def locate_coarse_root(root: Path) -> Path: + candidates = [root, root / "data" / "CoarseWSD-20", root / "CoarseWSD-20"] + for candidate in candidates: + if candidate.exists() and any(candidate.glob("*/train.data.txt")): + return candidate + raise FileNotFoundError(f"could not find CoarseWSD-20 below {root}") + + +def resolve_target_position(tokens: list[str], raw_pos: int, word: str) -> int | None: + candidates = [raw_pos, raw_pos - 1] + for pos in candidates: + if 0 <= pos < len(tokens) and tokens[pos].lower() == word.lower(): + return pos + exact = [i for i, token in enumerate(tokens) if token.lower() == word.lower()] + return exact[0] if len(exact) == 1 else None + + +def load_coarsewsd( + root: Path, + splits: list[str], + limit: int | None, + limit_per_word: int | None, + position_mode: str = "cloze", +) -> tuple[list[dict], dict]: + data_root = locate_coarse_root(root) + records: list[dict] = [] + skipped = defaultdict(int) + per_word_split: dict[str, dict[str, int]] = defaultdict(dict) + + for word_dir in sorted(p for p in data_root.iterdir() if p.is_dir()): + word = word_dir.name + cmap_path = word_dir / "class_map.txt" + class_map = parse_class_map(cmap_path) if cmap_path.exists() else {} + for split in splits: + data_path = word_dir / f"{split}.data.txt" + gold_path = word_dir / f"{split}.gold.txt" + if not data_path.exists() or not gold_path.exists(): + continue + lines = data_path.read_text().splitlines() + labels = gold_path.read_text().splitlines() + if len(lines) != len(labels): + raise ValueError(f"length mismatch: {data_path} and {gold_path}") + candidates: list[dict] = [] + for local_idx, (line, label) in enumerate(zip(lines, labels)): + fields = line.strip().split() + if len(fields) < 3: + skipped["short_line"] += 1 + continue + try: + raw_pos = int(fields[0]) + except ValueError: + skipped["bad_position"] += 1 + continue + tokens = fields[1:] + pos = resolve_target_position(tokens, raw_pos, word) + if pos is None: + skipped["target_position_unresolved"] += 1 + continue + if position_mode == "occurrence": + # Score the real occurrence: the prompt is the sentence + # prefix, so the target's leading-space token is the + # genuine next token at the scored position. + if pos < 1: + skipped["no_left_context"] += 1 + continue + prompt = " ".join(tokens[:pos]) + else: + tokens[pos] = "[BLANK]" + sentence = " ".join(tokens) + prompt = COARSEWSD_TEMPLATE.format(sentence=sentence) + sense = label.strip() + candidates.append( + { + "item_id": f"coarse:{word}:{split}:{local_idx}", + "kind": "context", + "dataset": "coarsewsd20", + "split": split, + "target": word, + "sense": sense, + "sense_name": class_map.get(sense, sense), + "prefix_words": pos, + "prompt": prompt, + } + ) + # A prefix cap can contain only the majority sense because the + # released files are sense-grouped. For pilots, round-robin across + # senses while keeping the original order within each sense. + if limit_per_word is not None and len(candidates) > limit_per_word: + by_sense: dict[str, list[dict]] = defaultdict(list) + for candidate in candidates: + by_sense[str(candidate["sense"])].append(candidate) + balanced: list[dict] = [] + for offset in range(max(len(v) for v in by_sense.values())): + for sense in sorted(by_sense): + if offset < len(by_sense[sense]): + balanced.append(by_sense[sense][offset]) + if len(balanced) >= limit_per_word: + break + if len(balanced) >= limit_per_word: + break + candidates = balanced + if limit is not None: + candidates = candidates[: max(0, limit - len(records))] + records.extend(candidates) + made = len(candidates) + per_word_split[word][split] = made + if limit is not None and len(records) >= limit: + break + if limit is not None and len(records) >= limit: + break + + audit = { + "dataset": "coarsewsd20", + "data_root": str(data_root), + "splits": splits, + "n_contexts": len(records), + "n_targets": len({r["target"] for r in records}), + "per_word_split": dict(per_word_split), + "skipped": dict(skipped), + } + return records, audit + + +def load_sae(path: Path) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: + ckpt = torch.load(path, map_location="cpu", weights_only=False) + if "W_dec" in ckpt: + w_dec = ckpt["W_dec"].float() + w_enc = ckpt["W_enc"].float() + b_enc = ckpt.get("b_enc") + k = int(ckpt.get("k") or ckpt.get("config", {}).get("k") or 256) + else: + state = ckpt.get("model_state_dict") or ckpt.get("state_dict") + cfg = ckpt.get("factorizer") or ckpt.get("config", {}).get("factorizer", {}) + if state is None: + raise KeyError(f"unrecognized SAE checkpoint: {path}") + w_dec = state["decoder"].float() + w_enc = state["encoder.weight"].float().T.contiguous() + b_enc = state.get("encoder.bias") + k = int(cfg.get("k", 256)) + if w_enc.shape == w_dec.shape: + w_enc = w_enc.T.contiguous() + if w_enc.shape != (w_dec.shape[1], w_dec.shape[0]): + raise ValueError(f"encoder {tuple(w_enc.shape)} incompatible with decoder {tuple(w_dec.shape)}") + if b_enc is None: + b_enc = torch.zeros(w_dec.shape[0]) + return w_dec, w_enc, b_enc.float(), k + + +def load_model(model_id: str, device: str, dtype: str): + """Frozen causal LM + tokenizer; left padding keeps the scored position last in the batch.""" + torch_dtype = torch.bfloat16 if dtype == "bfloat16" else torch.float32 + model, tokenizer = load_causal_lm(model_id, dtype=torch_dtype, device_map=None) + model.to(device) + if tokenizer.pad_token_id is None: + tokenizer.pad_token = tokenizer.eos_token + tokenizer.padding_side = "left" + return model, tokenizer + + +def resolve_single_token(tokenizer, target: str) -> tuple[int, str] | None: + # The prompt ends in ':', so the natural continuation includes a leading + # space. Requiring this exact continuation avoids silently scoring a token + # form that generation would not use at the prompt boundary. + candidate = " " + target + ids = tokenizer.encode(candidate, add_special_tokens=False) + return (int(ids[0]), candidate) if len(ids) == 1 else None + + +def prepare_target_codes( + targets: list[str], + tokenizer, + w_u: torch.Tensor, + row_mean: torch.Tensor, + w_enc: torch.Tensor, + b_enc: torch.Tensor, + w_dec_unit: torch.Tensor, + k: int, +) -> tuple[dict[str, dict[str, Any]], dict[str, Any]]: + info: dict[str, dict[str, Any]] = {} + skipped: dict[str, str] = {} + for target in sorted(set(targets)): + resolved = resolve_single_token(tokenizer, target) + if resolved is None: + skipped[target] = "multi_token" + continue + token_id, continuation = resolved + row = w_u[token_id] - row_mean + row_norm = row.norm().clamp_min(1e-8) + acts = torch.relu((row / row_norm) @ w_enc + b_enc) + values, indices = torch.topk(acts, k=min(k, acts.numel())) + beta = row_norm * values + recon = beta @ w_dec_unit[indices] + row_error = float((row - recon).norm() / row_norm) + row_cos = float(torch.nn.functional.cosine_similarity(row[None], recon[None]).item()) + info[target] = { + "token_id": token_id, + "continuation": continuation, + "feature_ids": indices.detach(), + "beta": beta.detach(), + "row_relative_error": row_error, + "row_cosine": row_cos, + "n_positive_codes": int((values > 0).sum().item()), + } + audit = { + "n_targets_requested": len(set(targets)), + "n_targets_single_token": len(info), + "single_token_fraction": len(info) / max(len(set(targets)), 1), + "skipped_targets": skipped, + "targets": { + target: {key: value for key, value in target_info.items() if key not in {"feature_ids", "beta"}} + for target, target_info in info.items() + }, + } + return info, audit + + +def score_items( + items: list[dict], + model, + tokenizer, + lm_head, + w_u: torch.Tensor, + row_mean: torch.Tensor, + w_dec_unit: torch.Tensor, + target_info: dict[str, dict[str, Any]], + batch_size: int, + max_length: int, + device: str, +) -> dict[str, Any]: + kept = [item for item in items if item["target"] in target_info] + metadata: list[dict] = [] + hidden_rows: list[torch.Tensor] = [] + projection_rows: list[torch.Tensor] = [] + contribution_rows: list[torch.Tensor] = [] + beta_rows: list[torch.Tensor] = [] + feature_rows: list[torch.Tensor] = [] + exact_logits: list[float] = [] + reconstructed_logits: list[float] = [] + target_logprobs: list[float] = [] + target_ranks: list[int] = [] + grabbed: list[torch.Tensor] = [] + + def pre_hook(_module, inputs): + grabbed.append(inputs[0].detach()) + + hook = lm_head.register_forward_pre_hook(pre_hook) + try: + with torch.inference_mode(): + for start in range(0, len(kept), batch_size): + batch = kept[start : start + batch_size] + enc = tokenizer( + [row["prompt"] for row in batch], + return_tensors="pt", + padding=True, + truncation=True, + max_length=max_length, + ).to(device) + grabbed.clear() + output = model(**enc, use_cache=False) + if not grabbed: + raise RuntimeError("LM-head hook did not capture a hidden state") + hidden = grabbed[-1][:, -1, :].float() + logits = output.logits[:, -1, :].float() + for offset, row in enumerate(batch): + tinfo = target_info[row["target"]] + feature_ids = tinfo["feature_ids"] + beta = tinfo["beta"] + projection = w_dec_unit[feature_ids] @ hidden[offset] + contribution = beta * projection + token_id = int(tinfo["token_id"]) + exact = float(logits[offset, token_id]) + reconstructed = float(hidden[offset] @ row_mean + contribution.sum()) + logprob = float(logits[offset, token_id] - torch.logsumexp(logits[offset], dim=0)) + rank = int((logits[offset] > logits[offset, token_id]).sum().item() + 1) + meta = {key: value for key, value in row.items() if key != "prompt"} + meta.update( + { + "prompt_sha1": stable_id(row["prompt"], n=20), + "token_id": token_id, + "continuation": tinfo["continuation"], + "row_relative_error": tinfo["row_relative_error"], + "row_cosine": tinfo["row_cosine"], + } + ) + metadata.append(meta) + hidden_rows.append(hidden[offset].half().cpu()) + projection_rows.append(projection.half().cpu()) + contribution_rows.append(contribution.half().cpu()) + beta_rows.append(beta.half().cpu()) + feature_rows.append(feature_ids.int().cpu()) + exact_logits.append(exact) + reconstructed_logits.append(reconstructed) + target_logprobs.append(logprob) + target_ranks.append(rank) + del output, logits, hidden + if torch.cuda.is_available(): + torch.cuda.empty_cache() + if (start // batch_size + 1) % 20 == 0 or start + batch_size >= len(kept): + log(f"scored {min(start + batch_size, len(kept))}/{len(kept)}") + finally: + hook.remove() + + if not metadata: + raise RuntimeError("no tokenizable items were scored") + return { + "metadata": metadata, + "hidden": torch.stack(hidden_rows), + "projection": torch.stack(projection_rows), + "contribution": torch.stack(contribution_rows), + "beta": torch.stack(beta_rows), + "feature_ids": torch.stack(feature_rows), + "exact_logit": torch.tensor(exact_logits, dtype=torch.float32), + "reconstructed_logit": torch.tensor(reconstructed_logits, dtype=torch.float32), + "target_logprob": torch.tensor(target_logprobs, dtype=torch.float32), + "target_rank": torch.tensor(target_ranks, dtype=torch.int32), + } + + +def l2_normalize(matrix: np.ndarray) -> np.ndarray: + norms = np.linalg.norm(matrix, axis=1, keepdims=True) + return matrix / np.maximum(norms, 1e-12) + + +def method_matrices(bundle: dict[str, Any], seed: int) -> dict[str, np.ndarray]: + hidden = bundle["hidden"].float().numpy() + projection = bundle["projection"].float().numpy() + contribution = bundle["contribution"].float().numpy() + beta = bundle["beta"].float().numpy() + metadata = bundle["metadata"] + shuffled = np.empty_like(contribution) + for target in sorted({row["target"] for row in metadata}): + idx = [i for i, row in enumerate(metadata) if row["target"] == target] + local_seed = seed + int(stable_id(target, n=8), 16) + rng = np.random.default_rng(local_seed) + permutation = rng.permutation(beta.shape[1]) + shuffled[idx] = projection[idx] * beta[idx][:, permutation] + return { + "srp": l2_normalize(contribution), + "support_projection": l2_normalize(projection), + "hidden": l2_normalize(hidden), + "shuffled_srp": l2_normalize(shuffled), + "token_only": l2_normalize(beta), + } + + +def safe_spearman(x: list[float] | np.ndarray, y: list[float] | np.ndarray) -> float: + from scipy.stats import spearmanr + + if len(x) < 3 or np.ptp(x) < 1e-8 or np.ptp(y) < 1e-8: + return float("nan") + return float(spearmanr(x, y).statistic) + + +def cluster_bootstrap( + rows: list[dict[str, Any]], + cluster_key: str, + stat_fn, + n_boot: int, + seed: int, +) -> list[float]: + by_cluster: dict[str, list[dict[str, Any]]] = defaultdict(list) + for row in rows: + by_cluster[str(row[cluster_key])].append(row) + keys = sorted(by_cluster) + rng = np.random.default_rng(seed) + values: list[float] = [] + for _ in range(n_boot): + sample: list[dict[str, Any]] = [] + for selected in rng.choice(keys, size=len(keys), replace=True): + sample.extend(by_cluster[str(selected)]) + value = float(stat_fn(sample)) + if np.isfinite(value): + values.append(value) + return values + + +def ci(values: list[float]) -> list[float]: + if not values: + return [float("nan"), float("nan")] + return [float(np.percentile(values, 2.5)), float(np.percentile(values, 97.5))] + + +def analyze_ambistory(bundle: dict[str, Any], out_dir: Path, n_boot: int, seed: int) -> dict: + metadata = bundle["metadata"] + matrices = method_matrices(bundle, seed) + anchor_index = {(row["target"], row["anchor_key"]): i for i, row in enumerate(metadata) if row["kind"] == "anchor"} + story_indices = [i for i, row in enumerate(metadata) if row["kind"] == "story"] + predictions: list[dict[str, Any]] = [] + metrics: dict[str, Any] = {"dataset": "ambistory", "methods": {}, "comparisons": {}} + rows_by_method: dict[str, list[dict[str, Any]]] = {} + shifts_by_method: dict[str, list[dict[str, Any]]] = {} + + for method, matrix in matrices.items(): + rows: list[dict[str, Any]] = [] + for i in story_indices: + meta = metadata[i] + anchor_ids = [anchor_index.get((meta["target"], key)) for key in meta["anchor_keys"]] + if any(index is None for index in anchor_ids): + continue + similarities = [float(matrix[i] @ matrix[int(index)]) for index in anchor_ids] + pred_margin = similarities[0] - similarities[1] + human_margin = float(meta["human_margin"]) + row = { + "method": method, + "item_id": meta["item_id"], + "setup_id": meta["setup_id"], + "context_id": meta["context_id"], + "split": meta["split"], + "target": meta["target"], + "has_ending": bool(meta["has_ending"]), + "human_margin": human_margin, + "pred_margin": pred_margin, + "similarity_0": similarities[0], + "similarity_1": similarities[1], + "correct": int(human_margin != 0 and np.sign(pred_margin) == np.sign(human_margin)), + "eligible_gap1": int(abs(human_margin) >= 1.0), + "row_relative_error": float(meta["row_relative_error"]), + } + rows.append(row) + if method == "srp": + predictions.append(row) + + def rho_fn(sample): + return safe_spearman([r["human_margin"] for r in sample], [r["pred_margin"] for r in sample]) + + eligible = [r for r in rows if r["eligible_gap1"]] + non_ties = [r for r in rows if r["human_margin"] != 0] + rho = rho_fn(rows) + rho_boot = cluster_bootstrap(rows, "setup_id", rho_fn, n_boot, seed) + accuracy_gap1 = float(np.mean([r["correct"] for r in eligible])) if eligible else float("nan") + acc_boot = ( + cluster_bootstrap( + eligible, + "setup_id", + lambda sample: np.mean([r["correct"] for r in sample]), + n_boot, + seed + 1, + ) + if eligible + else [] + ) + + # Controlled context change: representation distance should grow with + # the change in human sense-preference margin within the same setup. + shift_rows: list[dict[str, Any]] = [] + by_setup: dict[str, list[dict[str, Any]]] = defaultdict(list) + by_item = {r["item_id"]: r for r in rows} + for i in story_indices: + meta = metadata[i] + if meta["item_id"] in by_item: + by_setup[meta["setup_id"]].append({"i": i, **by_item[meta["item_id"]]}) + reversal_total = reversal_correct = 0 + for setup_id, setup_rows in by_setup.items(): + for left, right in itertools.combinations(setup_rows, 2): + distance = float(1.0 - matrix[left["i"]] @ matrix[right["i"]]) + human_shift = abs(left["human_margin"] - right["human_margin"]) + shift_rows.append({"setup_id": setup_id, "distance": distance, "human_shift": human_shift}) + if left["human_margin"] * right["human_margin"] < 0: + reversal_total += 1 + if left["pred_margin"] * right["pred_margin"] < 0: + reversal_correct += 1 + shift_rho = safe_spearman([r["human_shift"] for r in shift_rows], [r["distance"] for r in shift_rows]) + shift_boot = ( + cluster_bootstrap( + shift_rows, + "setup_id", + lambda sample: safe_spearman([r["human_shift"] for r in sample], [r["distance"] for r in sample]), + n_boot, + seed + 2, + ) + if shift_rows + else [] + ) + metrics["methods"][method] = { + "n_contexts": len(rows), + "n_setups": len({r["setup_id"] for r in rows}), + "preference_spearman": rho, + "preference_spearman_ci": ci(rho_boot), + "preferred_sense_accuracy_gap1": accuracy_gap1, + "preferred_sense_accuracy_gap1_ci": ci(acc_boot), + "n_gap1": len(eligible), + "preferred_sense_accuracy_non_tie": ( + float(np.mean([r["correct"] for r in non_ties])) if non_ties else float("nan") + ), + "profile_shift_spearman": shift_rho, + "profile_shift_spearman_ci": ci(shift_boot), + "n_shift_pairs": len(shift_rows), + "sense_reversal_consistency": (reversal_correct / reversal_total if reversal_total else float("nan")), + "n_reversal_pairs": reversal_total, + } + rows_by_method[method] = rows + shifts_by_method[method] = shift_rows + + # Confidence intervals for the scientific comparisons must be paired: all + # methods score the same stories, so resample setups once and take the + # within-resample difference. Independent per-method intervals do not test + # whether SRP itself differs from a control. + for comparison_index, baseline in enumerate(("shuffled_srp", "support_projection", "hidden")): + srp_rows = rows_by_method["srp"] + baseline_rows = rows_by_method[baseline] + if [row["item_id"] for row in srp_rows] != [row["item_id"] for row in baseline_rows]: + raise ValueError(f"unaligned AmbiStory rows for srp and {baseline}") + paired_rows = [ + { + "setup_id": left["setup_id"], + "human_margin": left["human_margin"], + "eligible_gap1": left["eligible_gap1"], + "srp_pred_margin": left["pred_margin"], + "baseline_pred_margin": right["pred_margin"], + "srp_correct": left["correct"], + "baseline_correct": right["correct"], + } + for left, right in zip(srp_rows, baseline_rows) + ] + + def preference_delta(sample): + human = [row["human_margin"] for row in sample] + return safe_spearman(human, [row["srp_pred_margin"] for row in sample]) - safe_spearman( + human, [row["baseline_pred_margin"] for row in sample] + ) + + eligible_paired = [row for row in paired_rows if row["eligible_gap1"]] + + def accuracy_delta(sample): + return np.mean([row["srp_correct"] for row in sample]) - np.mean( + [row["baseline_correct"] for row in sample] + ) + + srp_shifts = shifts_by_method["srp"] + baseline_shifts = shifts_by_method[baseline] + if [row["setup_id"] for row in srp_shifts] != [row["setup_id"] for row in baseline_shifts] or not np.allclose( + [row["human_shift"] for row in srp_shifts], + [row["human_shift"] for row in baseline_shifts], + ): + raise ValueError(f"unaligned AmbiStory shift rows for srp and {baseline}") + paired_shifts = [ + { + "setup_id": left["setup_id"], + "human_shift": left["human_shift"], + "srp_distance": left["distance"], + "baseline_distance": right["distance"], + } + for left, right in zip(srp_shifts, baseline_shifts) + ] + + def shift_delta(sample): + human = [row["human_shift"] for row in sample] + return safe_spearman(human, [row["srp_distance"] for row in sample]) - safe_spearman( + human, [row["baseline_distance"] for row in sample] + ) + + comparison_seed = seed + 100 + comparison_index * 10 + metrics["comparisons"][f"srp_minus_{baseline}"] = { + "preference_spearman_difference": float(preference_delta(paired_rows)), + "preference_spearman_difference_ci": ci( + cluster_bootstrap(paired_rows, "setup_id", preference_delta, n_boot, comparison_seed) + ), + "preferred_sense_accuracy_gap1_difference": float(accuracy_delta(eligible_paired)), + "preferred_sense_accuracy_gap1_difference_ci": ci( + cluster_bootstrap( + eligible_paired, + "setup_id", + accuracy_delta, + n_boot, + comparison_seed + 1, + ) + ), + "profile_shift_spearman_difference": float(shift_delta(paired_shifts)), + "profile_shift_spearman_difference_ci": ci( + cluster_bootstrap( + paired_shifts, + "setup_id", + shift_delta, + n_boot, + comparison_seed + 2, + ) + ), + } + write_jsonl(out_dir / "predictions.jsonl", predictions) + return metrics + + +def sampled_pair_auc(matrix: np.ndarray, labels: np.ndarray, seed: int, max_pairs: int = 20000) -> float: + from sklearn.metrics import roc_auc_score + + n = len(labels) + if n < 2: + return float("nan") + total = n * (n - 1) // 2 + rng = np.random.default_rng(seed) + if total <= max_pairs: + pairs = list(itertools.combinations(range(n), 2)) + else: + chosen: set[tuple[int, int]] = set() + while len(chosen) < max_pairs: + left = int(rng.integers(0, n)) + right = int(rng.integers(0, n - 1)) + if right >= left: + right += 1 + chosen.add((min(left, right), max(left, right))) + pairs = list(chosen) + gold = np.array([int(labels[i] == labels[j]) for i, j in pairs]) + if len(np.unique(gold)) < 2: + return float("nan") + scores = np.array([float(matrix[i] @ matrix[j]) for i, j in pairs]) + return float(roc_auc_score(gold, scores)) + + +def fit_train_clusters( + train_matrix: np.ndarray, + test_matrix: np.ndarray, + n_clusters: int, + seed: int, +) -> np.ndarray: + """Fit unsupervised clusters on train only and assign held-out examples.""" + from sklearn.cluster import KMeans + + if n_clusters < 1: + raise ValueError("n_clusters must be positive") + if len(train_matrix) < n_clusters: + raise ValueError(f"need at least {n_clusters} train rows, received {len(train_matrix)}") + if np.max(np.ptp(train_matrix, axis=0)) < 1e-8: + return np.zeros(len(test_matrix), dtype=int) + model = KMeans(n_clusters=n_clusters, n_init=20, random_state=seed) + model.fit(train_matrix) + return model.predict(test_matrix) + + +def analyze_coarsewsd(bundle: dict[str, Any], out_dir: Path, n_boot: int, seed: int) -> dict: + from sklearn.metrics import ( + accuracy_score, + adjusted_rand_score, + balanced_accuracy_score, + f1_score, + normalized_mutual_info_score, + ) + + metadata = bundle["metadata"] + matrices = method_matrices(bundle, seed) + words = sorted({row["target"] for row in metadata}) + metrics: dict[str, Any] = {"dataset": "coarsewsd20", "methods": {}, "comparisons": {}} + srp_predictions: list[dict[str, Any]] = [] + per_word_by_method: dict[str, dict[str, dict[str, float]]] = {} + + for method, matrix in matrices.items(): + per_word: dict[str, dict[str, float]] = {} + for word in words: + train_idx = [i for i, row in enumerate(metadata) if row["target"] == word and row["split"] == "train"] + test_idx = [i for i, row in enumerate(metadata) if row["target"] == word and row["split"] == "test"] + if not train_idx or not test_idx: + continue + train_y = np.array([metadata[i]["sense"] for i in train_idx]) + test_y = np.array([metadata[i]["sense"] for i in test_idx]) + senses = sorted(set(train_y)) + if len(senses) < 2 or not set(test_y).issubset(set(senses)): + continue + centroids = [] + for sense in senses: + centroid = matrix[np.array(train_idx)[train_y == sense]].mean(axis=0) + centroid /= max(np.linalg.norm(centroid), 1e-12) + centroids.append(centroid) + centroid_matrix = np.stack(centroids) + scores = matrix[test_idx] @ centroid_matrix.T + predictions = np.array([senses[i] for i in scores.argmax(axis=1)]) + n_clusters = len(senses) + train_matrix = matrix[train_idx] + test_matrix = matrix[test_idx] + cluster_pred = fit_train_clusters( + train_matrix, + test_matrix, + n_clusters=n_clusters, + seed=seed, + ) + word_metrics = { + "n_train": len(train_idx), + "n_test": len(test_idx), + "n_senses": len(senses), + "accuracy": float(accuracy_score(test_y, predictions)), + "balanced_accuracy": float(balanced_accuracy_score(test_y, predictions)), + "macro_f1": float(f1_score(test_y, predictions, average="macro")), + "pairwise_auc": sampled_pair_auc(matrix[test_idx], test_y, seed), + "ari": float(adjusted_rand_score(test_y, cluster_pred)), + "nmi": float(normalized_mutual_info_score(test_y, cluster_pred)), + "row_relative_error": float(metadata[test_idx[0]]["row_relative_error"]), + } + per_word[word] = word_metrics + if method == "srp": + for local, idx in enumerate(test_idx): + srp_predictions.append( + { + "item_id": metadata[idx]["item_id"], + "target": word, + "gold_sense": str(test_y[local]), + "predicted_sense": str(predictions[local]), + "correct": int(test_y[local] == predictions[local]), + "target_rank": int(bundle["target_rank"][idx]), + "target_logprob": float(bundle["target_logprob"][idx]), + } + ) + metric_names = ["accuracy", "balanced_accuracy", "macro_f1", "pairwise_auc", "ari", "nmi"] + aggregate: dict[str, Any] = {"n_words": len(per_word)} + rng = np.random.default_rng(seed) + word_names = sorted(per_word) + for metric_name in metric_names: + values = np.array([per_word[word][metric_name] for word in word_names], dtype=float) + values = values[np.isfinite(values)] + aggregate[metric_name] = float(values.mean()) if len(values) else float("nan") + boot = [] + if len(values): + for _ in range(n_boot): + boot.append(float(rng.choice(values, size=len(values), replace=True).mean())) + aggregate[f"{metric_name}_ci"] = ci(boot) + gated_words = [w for w in word_names if per_word[w]["row_relative_error"] < 0.5] + aggregate["row_gate_words"] = len(gated_words) + aggregate["row_gate_fraction"] = len(gated_words) / max(len(word_names), 1) + aggregate["row_gate"] = { + metric_name: ( + float(np.mean([per_word[w][metric_name] for w in gated_words])) if gated_words else float("nan") + ) + for metric_name in metric_names + } + metrics["methods"][method] = {"aggregate": aggregate, "per_word": per_word} + per_word_by_method[method] = per_word + + # CoarseWSD words are the independent units. Bootstrap paired per-word + # differences so the intervals answer whether SRP beats each control. + metric_names = ["accuracy", "balanced_accuracy", "macro_f1", "pairwise_auc", "ari", "nmi"] + for comparison_index, baseline in enumerate(("shuffled_srp", "support_projection", "hidden")): + common_words = sorted(set(per_word_by_method["srp"]) & set(per_word_by_method[baseline])) + comparison: dict[str, Any] = {"n_words": len(common_words)} + rng = np.random.default_rng(seed + 200 + comparison_index) + for metric_name in metric_names: + differences = np.array( + [ + per_word_by_method["srp"][word][metric_name] - per_word_by_method[baseline][word][metric_name] + for word in common_words + ], + dtype=float, + ) + differences = differences[np.isfinite(differences)] + comparison[f"{metric_name}_difference"] = float(differences.mean()) if len(differences) else float("nan") + boot = ( + [float(rng.choice(differences, size=len(differences), replace=True).mean()) for _ in range(n_boot)] + if len(differences) + else [] + ) + comparison[f"{metric_name}_difference_ci"] = ci(boot) + metrics["comparisons"][f"srp_minus_{baseline}"] = comparison + write_jsonl(out_dir / "predictions.jsonl", srp_predictions) + return metrics + + +def summarize_scoring(bundle: dict[str, Any]) -> dict[str, Any]: + metadata = bundle["metadata"] + exact = bundle["exact_logit"].numpy() + recon = bundle["reconstructed_logit"].numpy() + ranks = bundle["target_rank"].numpy() + return { + "n_scored": len(metadata), + "n_targets": len({row["target"] for row in metadata}), + "median_target_rank": float(np.median(ranks)), + "target_top1_fraction": float(np.mean(ranks == 1)), + "target_top10_fraction": float(np.mean(ranks <= 10)), + "median_absolute_logit_residual": float(np.median(np.abs(exact - recon))), + "mean_absolute_logit_residual": float(np.mean(np.abs(exact - recon))), + "row_gate_fraction_items": float(np.mean([row["row_relative_error"] < 0.5 for row in metadata])), + } + + +def self_test() -> int: + assert mask_exact_target("They followed the track.", "track") == "They followed the [BLANK]." + assert mask_exact_target("a tracker", "track") is None + assert resolve_target_position(["the", "bank", "closed"], 1, "bank") == 1 + assert resolve_target_position(["the", "bank", "closed"], 2, "bank") == 1 + matrix = l2_normalize(np.array([[1.0, 0.0], [0.9, 0.1], [0.0, 1.0], [0.1, 0.9]])) + auc = sampled_pair_auc(matrix, np.array([0, 0, 1, 1]), seed=0) + assert auc > 0.99 + train = np.array([[0.0, 0.0], [0.1, 0.0], [10.0, 10.0], [9.9, 10.0]]) + test = np.array([[0.05, 0.0], [9.95, 10.0]]) + cluster_pred = fit_train_clusters(train, test, n_clusters=2, seed=0) + assert cluster_pred[0] != cluster_pred[1] + assert safe_spearman([1, 2, 3], [2, 4, 6]) > 0.99 + print("self-test passed") + return 0 + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--dataset", choices=("ambistory", "coarsewsd20"), help="the paper uses coarsewsd20") + parser.add_argument( + "--data-root", + type=Path, + help="coarsewsd20: a bert-disambiguation checkout (or its data/CoarseWSD-20 dir); ambistory: dir of .json", + ) + parser.add_argument("--model-id", help="Hugging Face model id") + parser.add_argument("--checkpoint", type=Path, help="trained readout dictionary checkpoint (.pt)") + parser.add_argument("--out-dir", type=Path, help="output dir: representations.pt, audit/scoring/metrics/run_config") + parser.add_argument( + "--splits", default=None, help="comma-separated; defaults dev for AmbiStory, train,test for CoarseWSD" + ) + parser.add_argument("--batch-size", type=int, default=8) + parser.add_argument("--max-length", type=int, default=384) + parser.add_argument("--limit", type=int, default=None) + parser.add_argument("--limit-per-word", type=int, default=None) + parser.add_argument( + "--position-mode", + choices=("cloze", "occurrence"), + default="cloze", + help="cloze: mask target + cue; occurrence: score the sentence " + "prefix so the target occurrence is the genuine next token " + "(coarsewsd20 only)", + ) + parser.add_argument( + "--anchors", + type=Path, + default=None, + help="frozen sense-anchor JSON; anchor rows and sparse codes are " + "embedded in the bundle for offline weights-only feature typing", + ) + parser.add_argument("--device", default="cuda") + parser.add_argument("--dtype", choices=("bfloat16", "float32"), default="bfloat16") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--n-boot", type=int, default=2000) + parser.add_argument("--analyze-only", action="store_true") + parser.add_argument("--self-test", action="store_true") + return parser.parse_args() + + +def main() -> int: + args = parse_args() + if args.self_test: + return self_test() + required = (args.dataset, args.data_root, args.model_id, args.checkpoint, args.out_dir) + if not args.analyze_only and not all(required): + raise SystemExit("dataset, data-root, model-id, checkpoint, and out-dir are required") + if args.analyze_only and args.out_dir is None: + raise SystemExit("--out-dir is required with --analyze-only") + + args.out_dir.mkdir(parents=True, exist_ok=True) + representation_path = args.out_dir / "representations.pt" + if args.analyze_only: + bundle = torch.load(representation_path, map_location="cpu", weights_only=False) + else: + set_seed(args.seed) + if args.splits: + splits = [part.strip() for part in args.splits.split(",") if part.strip()] + else: + splits = ["dev"] if args.dataset == "ambistory" else ["train", "test"] + if args.dataset == "ambistory": + if args.position_mode != "cloze": + raise SystemExit("--position-mode occurrence supports coarsewsd20 only") + items, data_audit = load_ambistory(args.data_root, splits, args.limit) + else: + items, data_audit = load_coarsewsd( + args.data_root, splits, args.limit, args.limit_per_word, args.position_mode + ) + log(f"loaded {len(items)} items ({data_audit})") + log(f"loading {args.model_id}") + model, tokenizer = load_model(args.model_id, args.device, args.dtype) + if args.position_mode == "occurrence": + # Long prefixes must keep their tail (the context nearest the + # scored position), so truncate from the left. + tokenizer.truncation_side = "left" + lm_head = find_lm_head(model) + w_u = lm_head.weight.detach().float().to(args.device) + row_mean = w_u.mean(dim=0) + w_dec, w_enc, b_enc, k = load_sae(args.checkpoint) + w_dec = w_dec.to(args.device) + w_dec_unit = w_dec / w_dec.norm(dim=1, keepdim=True).clamp_min(1e-8) + w_enc = w_enc.to(args.device) + b_enc = b_enc.to(args.device) + if w_dec.shape[1] != w_u.shape[1]: + raise ValueError(f"dictionary/model width mismatch: {w_dec.shape[1]} vs {w_u.shape[1]}") + target_info, token_audit = prepare_target_codes( + [row["target"] for row in items], + tokenizer, + w_u, + row_mean, + w_enc, + b_enc, + w_dec_unit, + k, + ) + anchor_bundle = None + if args.anchors is not None: + anchors_cfg = json.loads(args.anchors.read_text()) + anchor_words = sorted( + { + anchor + for word, senses in anchors_cfg.items() + if not word.startswith("_") + for anchor_list in senses.values() + for anchor in anchor_list + } + ) + anchor_info, anchor_audit = prepare_target_codes( + anchor_words, + tokenizer, + w_u, + row_mean, + w_enc, + b_enc, + w_dec_unit, + k, + ) + anchor_bundle = { + "config": anchors_cfg, + "codes": { + word: { + "token_id": info["token_id"], + "feature_ids": info["feature_ids"].cpu(), + "beta": info["beta"].cpu(), + } + for word, info in anchor_info.items() + }, + "rows": {word: w_u[info["token_id"]].cpu().half() for word, info in anchor_info.items()}, + "audit": anchor_audit, + } + log( + "anchor coverage " + f"{anchor_audit['n_targets_single_token']}/" + f"{anchor_audit['n_targets_requested']} single-token" + ) + del w_enc, b_enc, w_dec + audit = {**data_audit, "tokenization_and_rows": token_audit, "k": k} + write_json(args.out_dir / "audit.json", audit) + log(f"single-token coverage {token_audit['n_targets_single_token']}/{token_audit['n_targets_requested']}") + bundle = score_items( + items, + model, + tokenizer, + lm_head, + w_u, + row_mean, + w_dec_unit, + target_info, + args.batch_size, + args.max_length, + args.device, + ) + bundle["run"] = { + "dataset": args.dataset, + "model_id": args.model_id, + "checkpoint": str(args.checkpoint), + "splits": splits, + "seed": args.seed, + "k": k, + "batch_size": args.batch_size, + "max_length": args.max_length, + "position_mode": args.position_mode, + } + if anchor_bundle is not None: + bundle["anchors"] = anchor_bundle + torch.save(bundle, representation_path) + write_json(args.out_dir / "scoring_summary.json", summarize_scoring(bundle)) + + metrics = ( + analyze_ambistory(bundle, args.out_dir, args.n_boot, args.seed) + if bundle["run"]["dataset"] == "ambistory" + else analyze_coarsewsd(bundle, args.out_dir, args.n_boot, args.seed) + ) + metrics["run"] = bundle["run"] + metrics["scoring"] = summarize_scoring(bundle) + write_json(args.out_dir / "metrics.json", metrics) + run_config = { + "argv": sys.argv, + "python": sys.version, + "platform": platform.platform(), + "torch": torch.__version__, + "numpy": np.__version__, + "story_template": AMBISTORY_STORY_TEMPLATE, + "gloss_template": AMBISTORY_GLOSS_TEMPLATE, + "coarsewsd_template": COARSEWSD_TEMPLATE, + "provenance": run_provenance(args), + } + write_json(args.out_dir / "run_config.json", run_config) + log(f"wrote results to {args.out_dir}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/sparse_readout_prism/research/seed_stability.py b/src/sparse_readout_prism/research/seed_stability.py new file mode 100644 index 0000000..735972f --- /dev/null +++ b/src/sparse_readout_prism/research/seed_stability.py @@ -0,0 +1,150 @@ +"""Shared pipeline for the cross-seed stability analyses (Appendix I). + +Backs three entry points: ``scripts/eval/cross_seed_stability.py`` +(``tab:app-cross-seed-stability``), ``scripts/eval/feature_group_matching.py`` +(``tab:app-feature-group-matching``) and ``scripts/eval/loo_core_recovery.py`` +(``tab:app-loo-core-recovery``). All three compare token summaries produced by +independently seeded dictionaries of one training recipe, so the steps that +produce those summaries must be identical across the scripts and live here: + +* single-token A/B contrasts parsed from a curated JSONL bank (the bare form + is tried before the space-prefixed form, the first occurrence of a pair + wins, the list is shuffled once with the caller's generator and capped), +* rows centred against the full-vocabulary mean of ``W_U`` and per-row + normalised, +* TopK codes from a dictionary's encoder and the contrast coefficient + ``beta = rn_A * z_A - rn_B * z_B`` for the contrast ``w_A - w_B``, +* the top-M features per sign, and the token set of a feature's top-R + centred rows (decoded, stripped, lower-cased, printable strings only). + +Checkpoints are read as raw tensors (``decoder`` / ``encoder.weight`` / +``encoder.bias`` from ``model_state_dict``) so the scripts also accept the +legacy ``W_dec`` / ``W_enc`` / ``b_enc`` layout of an archived reference +dictionary. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import torch + +Dictionary = tuple[torch.Tensor, torch.Tensor, torch.Tensor, int] + + +def load_dictionary(path: str | Path) -> Dictionary: + """``(decoder, encoder_w, encoder_b, k)`` as float32 CPU tensors. + + ``decoder`` and ``encoder_w`` are ``(d_features, d_model)``; ``k`` is read + from the checkpoint's factorizer config (default 256 when absent). + """ + ckpt = torch.load(path, map_location="cpu", weights_only=True) + if "W_dec" in ckpt: # legacy raw-tensor layout + dec = ckpt["W_dec"].float() + enc_w = ckpt["W_enc"].float().T.contiguous() + bias = ckpt.get("b_enc") + enc_b = bias.float() if bias is not None else torch.zeros(enc_w.shape[0]) + k = int(ckpt.get("k") or ckpt.get("config", {}).get("k") or 256) + return dec, enc_w, enc_b, k + state = ckpt.get("model_state_dict") or ckpt.get("state_dict") + if state is None: + raise KeyError(f"{path}: no model_state_dict / state_dict (or legacy W_dec) in checkpoint") + dec = state["decoder"].float() + enc_w = state["encoder.weight"].float() + enc_b = state.get("encoder.bias", torch.zeros(enc_w.shape[0])).float() + cfg = ckpt.get("factorizer") or ckpt.get("config", {}).get("factorizer", {}) + return dec, enc_w, enc_b, int(cfg.get("k", 256)) + + +def load_readout(path: str | Path) -> tuple[torch.Tensor, torch.Tensor | None]: + """``(W_U, h_LN)`` from an extraction payload ``{W_U_orig, h_LN, ...}``; ``h_LN`` may be None.""" + payload = torch.load(path, map_location="cpu", weights_only=True) + W = payload.get("W_U_orig", payload.get("W_U")) + if W is None: + raise KeyError(f"{path}: no W_U_orig / W_U in payload") + return W.float(), payload.get("h_LN") + + +def center_rows(W: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Full-vocabulary centring and per-row normalisation: ``(W_c, row_norms, W_n)``.""" + W_c = W - W.mean(0) + rn = W_c.norm(dim=1).clamp_min(1e-8) + return W_c, rn, W_c / rn[:, None] + + +def topk_codes(x: torch.Tensor, enc_w: torch.Tensor, enc_b: torch.Tensor, k: int) -> torch.Tensor: + """ReLU encoder activations, zeroed below the k-th largest value of each row.""" + acts = torch.relu(x @ enc_w.T + enc_b) + if k < acts.shape[-1]: + thresh = acts.topk(k, dim=-1).values[..., -1:] + acts = torch.where(acts >= thresh, acts, torch.zeros_like(acts)) + return acts + + +def single_token_id(tok, term: str) -> int | None: + """Token id of ``term`` if it (or ``" " + term``) is a single token, else None.""" + for variant in (term, " " + term): + ids = tok.encode(variant, add_special_tokens=False) + if len(ids) == 1: + return int(ids[0]) + return None + + +def load_contrast_pairs(bank_path: str | Path, tok, rng: np.random.Generator, max_contrasts: int) -> list[tuple]: + """Unique single-token ``(a, b, id_a, id_b)`` contrasts of a bank, shuffled with ``rng`` and capped.""" + recs = [json.loads(line) for line in Path(bank_path).read_text().splitlines() if line.strip()] + pairs: list[tuple] = [] + seen: set[tuple[str, str]] = set() + for r in recs: + a, b = r.get("target_a"), r.get("target_b") + if not a or not b or (a, b) in seen: + continue + ia, ib = single_token_id(tok, a), single_token_id(tok, b) + if ia is None or ib is None or ia == ib: + continue + seen.add((a, b)) + pairs.append((a, b, ia, ib)) + rng.shuffle(pairs) + return pairs[:max_contrasts] + + +def contrast_features( + W_n: torch.Tensor, rn: torch.Tensor, dictionary: Dictionary, ia: int, ib: int, top_m: int +) -> tuple[list[int], list[int]]: + """Top-M positive and top-M negative feature ids of the contrast row ``ia`` minus row ``ib``.""" + _, enc_w, enc_b, k = dictionary + z = topk_codes(torch.stack([W_n[ia], W_n[ib]]), enc_w, enc_b, k) # (2, d_features) + beta = rn[ia] * z[0] - rn[ib] * z[1] + return torch.topk(beta, top_m).indices.tolist(), torch.topk(-beta, top_m).indices.tolist() + + +def token_strings(tok, ids, cache: dict[int, str] | None = None) -> set[str]: + """Decoded, stripped, lower-cased, printable token strings of ``ids``.""" + out: set[str] = set() + for r in ids: + r = int(r) + s = None if cache is None else cache.get(r) + if s is None: + s = tok.decode([r]).strip().lower() + if cache is not None: + cache[r] = s + if s and s.isprintable(): + out.add(s) + return out + + +def feature_token_set( + fids, dec: torch.Tensor, W_c: torch.Tensor, tok, top_r: int, cache: dict[int, str] | None = None +) -> set[str]: + """Union over ``fids`` of the token strings of each feature's top-R centred rows.""" + out: set[str] = set() + for f in fids: + rows = torch.topk(W_c @ dec[f], top_r).indices.tolist() + out |= token_strings(tok, rows, cache) + return out + + +def jaccard(a, b) -> float: + return len(a & b) / max(1, len(a | b)) diff --git a/tests/test_baseline_and_dla.py b/tests/test_baseline_and_dla.py index 25c5e99..4ffcdbd 100644 --- a/tests/test_baseline_and_dla.py +++ b/tests/test_baseline_and_dla.py @@ -49,14 +49,23 @@ def test_baseline_methods_preserve_additive_identity() -> None: W, row_mean, sae = _make_sae_inputs() device = torch.device("cpu") h = torch.randn(8) - for spec in ("sparse_rp", "shuffled_row_code", "random_support_same_magnitudes"): + # spec -> expected number of active atoms (None = not checked). + expected_active = { + "sparse_rp": None, + "shuffled_row_code": 4, # nulls keep the SAE's k-sparsity, only the support/assignment changes + "random_support_same_magnitudes": 4, + "knn_basis_top4": 4, # the four nearest rows + "row_cluster_d16_k4": 4, # top-4 of 16 centroids + "row_cluster_hard_d16": 1, # the row's own centroid + } + for spec, n_active in expected_active.items(): method = baseline.build_method(spec, W=W, row_mean=row_mean, device=device, sae=sae, k=4, seed=0) dec = method.decompose_row(h, row_idx=3) # original_logit == base + feature_sum + residual to float precision. assert dec.identity_error.abs().item() < 1e-4, spec - if spec != "sparse_rp": - # nulls keep the SAE's k-sparsity, only the support/assignment changes. - assert dec.active_feature_indices.numel() == 4, spec + assert torch.isfinite(dec.reconstructed_logit), spec + if n_active is not None: + assert dec.active_feature_indices.numel() == n_active, spec def test_build_method_rejects_bad_specs() -> None: diff --git a/tests/test_causal_contribution_validation.py b/tests/test_causal_contribution_validation.py new file mode 100644 index 0000000..85c8cda --- /dev/null +++ b/tests/test_causal_contribution_validation.py @@ -0,0 +1,98 @@ +"""Correctness tests for scripts/eval/run_causal_contribution_validation.py (loaded by path). + +Pins the math core behind ``tab:causal-validation-summary`` on the script's own +synthetic residual-free case (a contrast built exactly in the span of a random +unit-row dictionary from a known 3-sparse beta): + +* predicted contributions track the realized changes (r2 ~ 1, slope ~ 1), +* random-feature controls carry zero predicted contribution, so they have no + explanatory power (the through-origin fit is undefined for them), +* the additive identity holds: the contributions sum to the exact margin, each + stored realized change equals (h . d_i)(q . d_i), and ablating h along d_i + moves the dense margin by exactly minus that amount. + +CPU only, no model, well under a second. +""" + +from __future__ import annotations + +import importlib.util +import math +import sys +from pathlib import Path + +import numpy as np + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT / "src")) + + +def _load(relpath: str, name: str): + spec = importlib.util.spec_from_file_location(name, ROOT / relpath) + mod = importlib.util.module_from_spec(spec) + sys.modules[name] = mod + spec.loader.exec_module(mod) + return mod + + +causal = _load("scripts/eval/run_causal_contribution_validation.py", "causal_contribution_validation") + + +def test_self_test_passes_and_recovers_unit_agreement() -> None: + m = causal.self_test_metrics() + assert all(ok for _, ok in m["checks"]), m["checks"] + # Residual-free contrast: the fit through the origin recovers the prediction. + assert abs(m["r2"] - 1.0) < 1e-2 + assert abs(m["slope"] - 1.0) < 0.1 + assert causal.self_test() == 0 + + +def test_random_controls_have_no_explanatory_power() -> None: + m = causal.self_test_metrics() + ctrl = [p for p in m["pairs"] if p[3] == 1] + assert len(ctrl) == 8 + # Controls lie outside the code support, so their predicted contribution is 0 ... + assert m["ctrl_max_abs_pred"] < 1e-6 + x = np.array([p[1] for p in ctrl]) + y = np.array([p[2] for p in ctrl]) + # ... and a through-origin fit of realized on predicted is undefined (no r2 to claim). + r2, slope = causal.fit_r2_slope(x, y) + assert math.isnan(r2) and math.isnan(slope) + + +def test_additive_identity_and_ablation_sign() -> None: + m = causal.self_test_metrics() + h, q, beta, W_dec = m["h"], m["q"], m["beta"], m["W_dec"] + # Sum of all per-feature contributions equals the exact margin (residual-free q). + m_exact = float(h @ q) + m_recon = float((beta * (W_dec @ h)).sum()) + assert abs(m_exact - m_recon) < 1e-4 * max(1.0, abs(m_exact)) + for fid, c_pred, delta_real, is_random in m["pairs"]: + d_i = W_dec[fid] + # Stored realized change is the algebraic form (h . d_i)(q . d_i). + assert abs(delta_real - float((d_i @ h) * (d_i @ q))) < 1e-4 + # Ablating h along the unit direction moves the dense margin by exactly -delta_real. + h_abl = h - (h @ d_i) * d_i + assert abs(float(h_abl @ q - h @ q) + delta_real) < 1e-4 + if not is_random: + # Predicted contribution is beta_i times the projection, by construction. + assert abs(c_pred - float(beta[fid] * (d_i @ h))) < 1e-5 + + +def test_fit_r2_slope_orientation() -> None: + # y = 2x exactly: slope is realized-on-predicted, so it must read 2 (not 0.5). + x = np.array([1.0, 2.0, 3.0, 4.0]) + r2, slope = causal.fit_r2_slope(x, 2.0 * x) + assert abs(slope - 2.0) < 1e-12 and abs(r2 - 1.0) < 1e-12 + assert all(math.isnan(v) for v in causal.fit_r2_slope(np.zeros(4), x)) + + +def test_cluster_bootstrap_ci_brackets_point_estimate() -> None: + m = causal.self_test_metrics() + rows = [{"cluster": f"c{i % 3}", "c_pred": p[1], "delta_real": p[2]} for i, p in enumerate(m["pairs"]) if p[3] == 0] + rows = rows * 3 # three copies so every resample keeps >= 3 points per fit + lo, hi = causal.cluster_bootstrap_r2(rows, n_boot=50, seed=0) + assert 0.0 <= lo <= hi <= 1.0 + assert lo > 0.9 + # Same seed, same interval. + assert (lo, hi) == causal.cluster_bootstrap_r2(rows, n_boot=50, seed=0) diff --git a/tests/test_cross_lens_scripts.py b/tests/test_cross_lens_scripts.py new file mode 100644 index 0000000..6e4814b --- /dev/null +++ b/tests/test_cross_lens_scripts.py @@ -0,0 +1,303 @@ +"""CPU-only checks for the cross-lens study scripts (no GPU, no model, no jlens). + +Covers the offline aggregation and table-assembly paths on synthetic dumps: + * scripts/eval/aggregate_cross_lens_en_zh.py (wilson, script_of, mid-band vote, end-to-end) + * scripts/eval/aggregate_cross_lens_en_de.py (lexical language call, divergence) + * scripts/data/build_cross_lens_de_bank.py (normalisation, Levenshtein, bank rules) + * scripts/analyze/cross_lens_three_lens_prompt.py + * scripts/figures/compute_cross_lens_shared_feature.py + * scripts/analyze/cross_lens_antonym_layers.py (table subcommand) + * scripts/run/fit_jlens.py (source-layer picker) +""" + +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] + + +def _load(mod_name: str, rel_path: str): + spec = importlib.util.spec_from_file_location(mod_name, ROOT / rel_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +agg_zh = _load("aggregate_cross_lens_en_zh", "scripts/eval/aggregate_cross_lens_en_zh.py") +agg_de = _load("aggregate_cross_lens_en_de", "scripts/eval/aggregate_cross_lens_en_de.py") +bank_de = _load("build_cross_lens_de_bank", "scripts/data/build_cross_lens_de_bank.py") +three = _load("cross_lens_three_lens_prompt", "scripts/analyze/cross_lens_three_lens_prompt.py") +shared = _load("compute_cross_lens_shared_feature", "scripts/figures/compute_cross_lens_shared_feature.py") +antonym = _load("cross_lens_antonym_layers", "scripts/analyze/cross_lens_antonym_layers.py") +fit_jlens = _load("fit_jlens", "scripts/run/fit_jlens.py") + +MIDBAND = ["21", "24", "26", "29"] + + +# --------------------------------------------------------------------------- # +# synthetic dump builders +# --------------------------------------------------------------------------- # + + +def _decomp(feats, lens_logit=None, rank=None): + d = { + "original_logit": sum(c for _, c in feats) + 0.5, + "base": -0.1, + "feature_sum": sum(c for _, c in feats), + "residual": 0.6, + "top_features": [{"id": i, "contribution": c} for i, c in feats], + } + if lens_logit is not None: + d["lens_logit"] = lens_logit + d["lens_rank"] = rank + return d + + +def _record(pid, group, top1, dom_by_target, top1_feats=None): + """One prompt record whose mid-band layers all carry the same readings.""" + layers = {} + for L in ["2", *MIDBAND]: + entry = { + "top5": [[top1, 20.0], ["x", 1.0]], + "targets": {t: _decomp([(fid, 5.0), (fid + 1000, 0.5)], 10.0, 0) for t, fid in dom_by_target.items()}, + } + if top1_feats is not None: + entry["top1_decomp"] = _decomp(top1_feats) + layers[L] = {"-1": entry} + return { + "id": pid, + "group": group, + "continuation": top1, + "intermediate": None, + "answer": top1.strip(), + "targets": {t: i for i, t in enumerate(dom_by_target)}, + "layers": layers, + } + + +def _dump(records, labels=None): + return { + "model": "m", + "lens": "l", + "sae": "s", + "positions": [-1], + "layers": [2, 21, 24, 26, 29], + "records": records, + "feature_top_tokens": labels or {}, + } + + +# --------------------------------------------------------------------------- # +# EN-ZH aggregator +# --------------------------------------------------------------------------- # + + +def test_wilson_matches_reference_values(): + p, lo, hi = agg_zh.wilson(77, 80) + assert abs(p - 0.9625) < 1e-9 + assert round(lo, 2) == 0.90 and round(hi, 2) == 0.99 + assert agg_zh.wilson(0, 0) == (0.0, 0.0, 1.0) + + +def test_script_of_distinguishes_scripts(): + assert agg_zh.script_of(" London") == "LATIN" + assert agg_zh.script_of("伦敦") == "CJK" + assert agg_zh.script_of(" 42") == "OTHER" + + +def test_majority_same_uses_midband_vote(): + a = _record("p", "g", "t", {"f": 7}) + b = _record("p", "g", "t", {"f": 7}) + assert agg_zh.majority_same(a, b, "f", "f") is True + # flip two of four mid-band layers: 2/4 still counts as a majority (>= ceil(n/2)) + for L in ["21", "24"]: + b["layers"][L]["-1"]["targets"]["f"]["top_features"][0]["id"] = 99 + assert agg_zh.majority_same(a, b, "f", "f") is True + b["layers"]["26"]["-1"]["targets"]["f"]["top_features"][0]["id"] = 99 + assert agg_zh.majority_same(a, b, "f", "f") is False + assert agg_zh.majority_same(a, b, "missing", "f") is None + + +def test_en_zh_aggregator_end_to_end(tmp_path): + bank = { + "prompts": [ + { + "id": "a1", + "group": "antonym_zh", + "concept": "big", + "form_a": "大", + "form_b": " big", + "null_targets": ["桌"], + }, + { + "id": "a2", + "group": "antonym_zh", + "concept": "slow", + "form_a": "慢", + "form_b": " slow", + "null_targets": ["门"], + }, + {"id": "d1", "group": "ctrl_digit", "concept": "7", "form_a": "七", "form_b": " seven", "null_targets": []}, + ] + } + en = _dump( + [ + _record("a1", "antonym_zh", " big", {"大": 1, " big": 1, "桌": 50}), + _record("a2", "antonym_zh", " slow", {"慢": 2, " slow": 3, "门": 60}), + _record("d1", "ctrl_digit", " 7", {"七": 4, " seven": 4}), + ] + ) + zh = _dump( + [ + _record("a1", "antonym_zh", "大", {"大": 1, " big": 1, "桌": 50}), + _record("a2", "antonym_zh", "慢", {"慢": 9, " slow": 3, "门": 60}), + _record("d1", "ctrl_digit", "七", {"七": 4, " seven": 4}), + ] + ) + (tmp_path / "en.json").write_text(json.dumps(en, ensure_ascii=False)) + (tmp_path / "zh.json").write_text(json.dumps(zh, ensure_ascii=False)) + (tmp_path / "bank.json").write_text(json.dumps(bank, ensure_ascii=False)) + out = tmp_path / "summary.json" + rc = agg_zh.main( + [ + "--lens-a", + str(tmp_path / "en.json"), + "--lens-b", + str(tmp_path / "zh.json"), + "--prompts", + str(tmp_path / "bank.json"), + "--out", + str(out), + "--rows-csv", + str(tmp_path / "rows.csv"), + ] + ) + assert rc == 0 + s = json.loads(out.read_text()) + assert s["headline"] == {"pass": 1, "n": 2, "rate": 0.5, "ci": s["headline"]["ci"]} + assert s["per_group"]["ctrl_digit"]["pass"] == 1 + assert s["cross_form_en"] == {"pass": 1, "n": 2} + assert s["null_within_lens"]["pass"] == 0 and s["null_within_lens"]["n"] == 2 + assert s["lens_only_split"]["antonym_zh"] == {"pass": 2, "n": 2} + assert s["lens_only_split"]["all_cross"] == {"pass": 2, "n": 2} + assert (tmp_path / "rows.csv").read_text().startswith("id,group,concept,cross_lens_pass") + + +# --------------------------------------------------------------------------- # +# EN-DE aggregator and bank builder +# --------------------------------------------------------------------------- # + + +def test_lang_of_lexical_call(): + item = {"form_a": " groß", "form_b": " big", "lang_a": "de", "lang_b": "en"} + assert agg_de.lang_of(" groß", item) == "DE" + assert agg_de.lang_of("big", item) == "EN" + assert agg_de.lang_of(" Küche", item) == "DE" # umlaut cue + assert agg_de.lang_of(" hund", item) == "DE" # word list + assert agg_de.lang_of(" window", item) == "EN" + assert agg_de.lang_of(" xyzzy", item) == "OTHER" + assert agg_de.lang_of(" ", item) == "OTHER" + + +def test_norm_and_levenshtein_cognate_rule(): + assert bank_de.norm("Käse") == "kase" + assert bank_de.norm("groß") == "gross" + assert bank_de.levenshtein("neu", "new") == 1 + assert bank_de.levenshtein("leicht", "light") == 2 + assert bank_de.levenshtein("hund", "dog") > 2 + + +def test_build_bank_applies_rules_with_stub_tokenizer(): + class Tok: + def encode(self, s, add_special_tokens=False): + # Everything is one token except any surface containing 'schw'. + return [0, 1] if "schw" in s.lower() else [0] + + prompts, excluded = bank_de.build_bank(Tok()) + ids = {p["id"] for p in prompts} + reasons = {e["id"]: e["reason"] for e in excluded} + assert reasons["antonym_de_13"] == "cognate (lev=1)" # neu / new + assert reasons["antonym_de_09"].startswith("not single-token: de") # schwach + assert "antonym_de_01" in ids + first = next(p for p in prompts if p["id"] == "antonym_de_01") + assert first["form_a"] == "groß" and first["form_b"] == " big" # quoted frame: no leading space + assert first["targets"] == [first["form_a"], first["form_b"], *first["null_targets"]] + assert all(p["lang_b"] != p["lang_a"] for p in prompts) + + +# --------------------------------------------------------------------------- # +# three-lens table, shared-feature metrics, antonym table +# --------------------------------------------------------------------------- # + + +def test_three_lens_tables(tmp_path): + recs = { + "EN": _record("q", "antonym_de", "large", {" groß": 6764, " large": 12474}, top1_feats=[(12474, 9.0)]), + "DE": _record("q", "antonym_de", " groß", {" groß": 6764, " large": 12474}, top1_feats=[(6764, 8.0)]), + } + top1, targets, feats = three.three_lens_tables(recs, ["24", "29"], None) + assert [r["top1_token"] for r in top1] == ["large", " groß", "large", " groß"] + assert {r["target"] for r in targets} == {" groß", " large"} + dom = {(r["lens"], r["token"]): r["dominant_feature"] for r in feats if r["layer"] == "24"} + assert dom[("EN", "large")] == 12474 and dom[("DE", " groß")] == 6764 + for label, rec in recs.items(): + (tmp_path / f"{label}.json").write_text(json.dumps(_dump([rec]), ensure_ascii=False)) + rc = three.main( + [ + "--dump", + f"EN={tmp_path / 'EN.json'}", + "--dump", + f"DE={tmp_path / 'DE.json'}", + "--layers", + "24,29", + "--out-csv", + str(tmp_path / "t.csv"), + ] + ) + assert rc == 0 and (tmp_path / "t.csv").exists() + + +def test_shared_feature_metrics(): + recs = { + "EN": _record("f", "fact_zh", " London", {"伦敦": 1}, top1_feats=[(23180, 30.0), (26030, 1.7), (5, 0.3)]), + "ZH": _record("f", "fact_zh", "伦敦", {"伦敦": 1}, top1_feats=[(23180, 24.0), (5150, 1.5)]), + } + m = shared.shared_feature_metrics(recs, "26", {"EN": {"23180": [" London", "伦敦"]}}) + assert m["shared_dominant_feature"] == 23180 + assert abs(m["per_lens"]["EN"]["dominant_share_of_feature_sum"] - 30.0 / 32.0) < 1e-9 + assert m["per_lens"]["EN"]["largest_other_contribution"] == 1.7 + assert m["shared_feature_share_of_feature_sum"]["ZH"] == 24.0 / 25.5 + assert m["shared_feature_top_tokens"] == [" London", "伦敦"] + recs["ZH"]["layers"]["26"]["-1"]["top1_decomp"]["top_features"][0]["id"] = 7 + assert shared.shared_feature_metrics(recs, "26", {})["shared_dominant_feature"] is None + + +def test_antonym_layer_table(): + def payload(tok, pct, fid): + rows = { + L: { + "jlens_top10": [[tok, pct], ["x", 1.0]], + "logitlens_top10": [["y", 50.0]], + "targets": {"大": _decomp([(fid, 12.5)]), "big": _decomp([(fid, 5.0)])}, + } + for L in ["24", "29", "final"] + } + return {"prompt": "p", "rows": rows} + + table, feats = antonym.antonym_layer_table( + {"EN": payload("large", 39.7, 112), "ZH": payload("大的", 36.4, 112)}, ["24", "final"] + ) + assert [(r["layer"], r["lens"], r["top1_token"], r["top1_softmax_pct"]) for r in table][:2] == [ + ("24", "EN", "large", 39.7), + ("24", "ZH", "大的", 36.4), + ] + assert all(r["dominant_feature"] == 112 for r in feats) + assert antonym.parse_dump_args(["EN=a.json", "ZH=b.json"]) == [("EN", Path("a.json")), ("ZH", Path("b.json"))] + + +def test_pick_source_layers_matches_paper_layer_set(): + assert fit_jlens.pick_source_layers(32, 12) == [2, 5, 7, 10, 12, 14, 17, 19, 21, 24, 26, 29] + assert fit_jlens.pick_source_layers(12, 4) == [2, 4, 7, 9] diff --git a/tests/test_error_tails_and_nearest_rows.py b/tests/test_error_tails_and_nearest_rows.py new file mode 100644 index 0000000..a2f562e --- /dev/null +++ b/tests/test_error_tails_and_nearest_rows.py @@ -0,0 +1,175 @@ +"""CPU-only tests for the error-tail and nearest-row analyses: + +* scripts/eval/analyze_error_tails.py (tab:app-error-tails, tab:app-error-tail-margins) +* scripts/analyze/nearest_rows_baseline.py (tab:app-nearest-rows) + +Synthetic inputs only: hand-built baseline_query_rows.csv files +and a random W_U with a fake tokenizer; no model, no network. +""" + +from __future__ import annotations + +import importlib.util +import json +import sys +from pathlib import Path + +import numpy as np +import pandas as pd +import torch + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT / "src")) + + +def _load(mod_name: str, rel_path: str): + spec = importlib.util.spec_from_file_location(mod_name, ROOT / rel_path) + mod = importlib.util.module_from_spec(spec) + sys.modules[mod_name] = mod + spec.loader.exec_module(mod) + return mod + + +tails = _load("analyze_error_tails", "scripts/eval/analyze_error_tails.py") +nrb = _load("nearest_rows_baseline", "scripts/analyze/nearest_rows_baseline.py") + + +# --------------------------------------------------------------------------- # +# error tails +# --------------------------------------------------------------------------- # + + +def _fake_run_dir(root: Path, tag: str, seed: int) -> None: + rng = np.random.default_rng(seed) + rows = [] + for method, scale in (("sparse_rp", 0.2), ("nearest_row_ridge_top128", 1.0)): + for i in range(60): + exact = float(rng.normal(0, 3)) + rows.append( + { + "model": tag, + "method": method, + "bank": "curated_ab", + "family": "fam_a" if i % 2 else "fam_b", + "base_case_id": f"base_{i // 3}", + "case_id": f"case_{i}", + "query": "single_token", + "exact_margin": exact, + "sparse_margin": exact + float(rng.normal(0, scale)), + } + ) + d = root / f"run_{tag}" + d.mkdir(parents=True) + pd.DataFrame(rows).to_csv(d / "baseline_query_rows.csv", index=False) + pd.DataFrame([{"case_id": "case_0", "token_id": 1, "reason": "ok", "side": "A", "target": " x"}]).to_csv( + d / "tokenizer_audit.csv", index=False + ) + + +def test_error_tails_row_metrics_and_pooled_tables(tmp_path: Path, monkeypatch) -> None: + for tag, seed in (("m1", 0), ("m2", 1)): + _fake_run_dir(tmp_path / "in", tag, seed) + out = tmp_path / "out" + monkeypatch.setattr( + sys, + "argv", + [ + "analyze_error_tails.py", + "--input-root", + str(tmp_path / "in"), + "--input-prefix", + "run_", + "--model-tags", + "m1,m2", + "--primary-model-tag", + "m1", + "--out-dir", + str(out), + "--n-boot", + "25", + "--seed", + "3", + ], + ) + assert tails.main() == 0 + + rows, _ = tails.load_rows(tmp_path / "in", "run_", ["m1", "m2"]) + exact = rows.exact_margin.to_numpy() + recon = rows.sparse_margin.to_numpy() + # Definitions: floored rho, unfloored rho_0 == the runner's residual_direct convention. + assert np.allclose(rows.rho, np.abs(exact - recon) / (np.abs(exact) + 0.5)) + assert np.allclose(rows.rho_0, np.abs(exact - recon) / np.maximum(np.abs(exact), 1e-6)) + assert set(rows.margin_bin) <= set(tails.MARGIN_LABELS) + assert set(rows.margin_bin_fine) <= set(tails.MARGIN_LABELS_FINE) + assert {"absolute_error", "rho", "rho_0", "sign_match", "sign_match_strict", "accepted", "covered"} <= set( + rows.columns + ) + + pooled = pd.read_csv(out / "summary_pooled_by_method.csv").set_index("method") + assert pooled.loc["sparse_rp", "n_rows"] == 120 and pooled.loc["sparse_rp", "n_models"] == 2 + srp = rows[rows.method == "sparse_rp"] + assert np.isclose(pooled.loc["sparse_rp", "absolute_error_mean"], srp.absolute_error.mean()) + assert np.isclose(pooled.loc["sparse_rp", "coverage"], (srp.sign_match & (srp.rho_0 < 0.5)).mean()) + # The tighter method has the better tail on every pooled statistic. + assert pooled.loc["sparse_rp", "absolute_error_p95"] < pooled.loc["nearest_row_ridge_top128", "absolute_error_p95"] + + bins = pd.read_csv(out / "summary_pooled_by_method_margin_bin.csv") + srp_bins = bins[bins.method == "sparse_rp"] + assert np.isclose(srp_bins.share.sum(), 1.0) + assert list(srp_bins.margin_bin) == [b for b in tails.MARGIN_LABELS if b in set(srp_bins.margin_bin)] + + by_method = pd.read_csv(out / "summary_by_method.csv") + assert {"rho_p95", "absolute_error_max", "accepted_rate_ci_lo"} <= set(by_method.columns) + paired = pd.read_csv(out / "paired_method_differences.csv") + assert paired.n_pairs.iloc[0] == 60 + manifest = json.loads((out / "summary.json").read_text()) + assert manifest["rows_total"] == 240 and "command" in manifest + + +# --------------------------------------------------------------------------- # +# nearest rows +# --------------------------------------------------------------------------- # + + +class _FakeTokenizer: + """Vocabulary of single-character tokens; ' x' is one token, 'x' is one token.""" + + def __init__(self, vocab: list[str]): + self.vocab = vocab + self.index = {t: i for i, t in enumerate(vocab)} + + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: + return [self.index[text]] if text in self.index else [0, 0] + + def decode(self, ids: list[int]) -> str: + return "".join(self.vocab[i] for i in ids) + + +def test_nearest_rows_excludes_contrast_tokens_and_ranks_by_abs_cosine() -> None: + vocab = [" bug", " insect", " error", "bug", "Bug", " bugs", "昆虫", "z1", "z2", "z3", "z4", "z5"] + tok = _FakeTokenizer(vocab) + gen = torch.Generator().manual_seed(0) + d = 64 # wide enough that random rows have small |cosine| with q, planted rows large + W = torch.randn(len(vocab), d, generator=gen) + # Plant structure: 'bug'/'Bug' near ' bug', '昆虫' near ' insect' (the other side of q). + W[3] = W[0] + 0.1 * torch.randn(d, generator=gen) + W[4] = W[0] + 0.1 * torch.randn(d, generator=gen) + W[6] = W[1] + 0.1 * torch.randn(d, generator=gen) + row_mean = W.mean(dim=0) + res, report, ids = nrb.nearest_rows(W, row_mean, tok, [("bug", "insect")], top_n=4) + assert ids == {"bug": 0, "insect": 1} + assert report["insect"]["without_leading_space"]["single_token"] is False + top = res["bug_minus_insect"]["top_centered"] + listed = {r["token_id"] for r in top} + assert not listed & {0, 1}, "contrast tokens must be excluded" + assert [r["rank"] for r in top] == [1, 2, 3, 4] + abs_cos = [abs(r["cosine"]) for r in top] + assert abs_cos == sorted(abs_cos, reverse=True) + assert {3, 4, 6} <= listed # the planted rows lead the |cosine| ranking + ins = next(r for r in top if r["token_id"] == 6) + assert ins["cosine"] < 0 # insect-side rows carry a negative signed cosine + assert all(r["cosine"] > 0 for r in top if r["token_id"] in (3, 4)) + table = nrb.table_rows(res) + assert len(table) == 4 and table[0]["contrast"] == "bug_minus_insect" + assert all(t["centered_cosine"].startswith(("+", "-")) for t in table) + assert any(t["centered_gloss"] == "insect" for t in table) == any(t["centered_token"] == "昆虫" for t in table) diff --git a/tests/test_lexical_control_directions.py b/tests/test_lexical_control_directions.py new file mode 100644 index 0000000..f42426c --- /dev/null +++ b/tests/test_lexical_control_directions.py @@ -0,0 +1,136 @@ +"""Unit checks for the label-free category-direction controls of the lexical-edit +stress test (scripts/run/run_qwen_profanity_suppression_eval.py, loaded by path): + +- ``category_direction_units`` returns unit-norm directions oriented toward the + discovery side (non-negative dot with the mean-row direction) and is + deterministic across calls (torch.linalg.svd, not a randomized PCA); +- the ``mean_row_direction`` / ``pca_group_direction`` / ``pca_group_rank4`` + methods of ``apply_method`` produce finite logits with a nonzero + ``intervention_norm`` on a tiny synthetic W_U, and lower the discovery tokens' + logits relative to the vocabulary average (the suppression sign convention). +""" + +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path + +import pytest +import torch + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT / "src")) + + +def _load(relpath: str, name: str): + spec = importlib.util.spec_from_file_location(name, ROOT / relpath) + mod = importlib.util.module_from_spec(spec) + sys.modules[name] = mod # dataclass resolution needs the module registered + spec.loader.exec_module(mod) + return mod + + +prof = _load("scripts/run/run_qwen_profanity_suppression_eval.py", "qwen_profanity_suppression_eval") + +NEW_METHODS = ("mean_row_direction", "pca_group_direction", "pca_group_rank4") + + +def _synthetic(vocab: int = 40, d_model: int = 12, d_features: int = 16): + gen = torch.Generator().manual_seed(0) + W_U = torch.randn(vocab, d_model, generator=gen) + row_mean = W_U.mean(dim=0) + decoder = torch.randn(d_features, d_model, generator=gen) + decoder = decoder / decoder.norm(dim=1, keepdim=True) + logits = torch.randn(vocab, generator=gen) + discovery_ids = [1, 4, 7, 9, 11] + return W_U, row_mean, decoder, logits, discovery_ids + + +def test_category_direction_units_are_unit_norm_and_oriented() -> None: + W_U, row_mean, _decoder, _logits, discovery_ids = _synthetic() + mean_dir, pca_dir, pca_dirs = prof.category_direction_units(W_U, discovery_ids, row_mean) + + assert mean_dir.shape == (W_U.shape[1],) + assert pca_dir.shape == (W_U.shape[1],) + assert pca_dirs.shape == (min(4, len(discovery_ids)), W_U.shape[1]) + assert abs(mean_dir.norm().item() - 1.0) < 1e-5 + assert abs(pca_dir.norm().item() - 1.0) < 1e-5 + assert torch.allclose(pca_dirs.norm(dim=1), torch.ones(pca_dirs.shape[0]), atol=1e-5) + + # Oriented toward the discovery side: non-negative alignment with the mean row. + assert torch.dot(mean_dir, pca_dir).item() >= 0.0 + for j in range(pca_dirs.shape[0]): + assert torch.dot(mean_dir, pca_dirs[j]).item() >= 0.0 + # pca_dir is the first row of the rank-<=4 set. + assert torch.allclose(pca_dir, pca_dirs[0], atol=1e-6) + # Mean direction is the normalized centered mean of the discovery rows. + m = (W_U[discovery_ids] - row_mean).mean(dim=0) + assert torch.allclose(mean_dir, m / m.norm(), atol=1e-6) + + +def test_category_direction_units_deterministic() -> None: + W_U, row_mean, _decoder, _logits, discovery_ids = _synthetic() + first = prof.category_direction_units(W_U, discovery_ids, row_mean) + second = prof.category_direction_units(W_U, discovery_ids, row_mean) + for a, b in zip(first, second): + assert torch.equal(a, b) + + +@pytest.mark.parametrize("method", NEW_METHODS) +def test_new_methods_produce_finite_logits_and_nonzero_norm(method: str) -> None: + W_U, row_mean, decoder, logits, discovery_ids = _synthetic() + mean_dir, pca_dir, pca_dirs = prof.category_direction_units(W_U, discovery_ids, row_mean) + scale = 2.0 + edited, meta = prof.apply_method( + logits=logits, + W_U=W_U, + decoder=decoder, + method=method, + feature_ids=[0, 1, 2], + random_ids=[3, 4, 5], + scale=scale, + token_ids=[], + audited_ids=[], + mean_dir=mean_dir, + pca_dir=pca_dir, + pca_dirs=pca_dirs, + ) + assert edited.shape == logits.shape + assert torch.isfinite(edited).all() + assert meta["intervention_norm"] > 0.0 + assert meta["token_bias_strength"] == 0.0 + assert not torch.allclose(edited, logits) + if method != "pca_group_rank4": + # A single unit direction scaled by ``scale``. + assert abs(meta["intervention_norm"] - scale) < 1e-5 + # Suppression sign convention: the discovery tokens' logits drop relative to + # the vocabulary-average change. + delta = edited - logits + assert delta[discovery_ids].mean().item() < delta.mean().item() + + +@pytest.mark.parametrize("method", NEW_METHODS) +def test_new_methods_require_their_direction(method: str) -> None: + W_U, _row_mean, decoder, logits, _discovery_ids = _synthetic() + with pytest.raises(ValueError): + prof.apply_method( + logits=logits, + W_U=W_U, + decoder=decoder, + method=method, + feature_ids=[0, 1, 2], + random_ids=[3, 4, 5], + scale=1.0, + token_ids=[], + audited_ids=[], + ) + + +def test_grid_constants_match_paper_protocol() -> None: + assert len(prof.PROMPTS) == 20 + assert len(prof.PAIRS) == 17 + heldout = [bad for bad, _good in prof.PAIRS if bad not in prof.DISCOVERY_TERMS] + assert len(heldout) == 12 + assert prof.SCALES == [0.5, 1, 2, 3, 4, 6, 8, 10, 12, 16, 24, 32, 48, 64] + assert all(bad in prof.PAIR_TO_FEATURE_SET for bad, _good in prof.PAIRS) diff --git a/tests/test_paired_matched_kl_bootstrap.py b/tests/test_paired_matched_kl_bootstrap.py new file mode 100644 index 0000000..302073a --- /dev/null +++ b/tests/test_paired_matched_kl_bootstrap.py @@ -0,0 +1,130 @@ +"""Smoke test for scripts/eval/paired_matched_kl_bootstrap.py on a synthetic +candidate_constrained_rows.csv: the matched scales are the log-nearest median +KLs, the paired mean difference equals the direct per-candidate mean, both +clusterings (held-out term and prompt) are reported, and the run is seeded. +""" + +from __future__ import annotations + +import csv +import importlib.util +import json +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] + + +def _load(relpath: str, name: str): + spec = importlib.util.spec_from_file_location(name, ROOT / relpath) + mod = importlib.util.module_from_spec(spec) + sys.modules[name] = mod + spec.loader.exec_module(mod) + return mod + + +boot = _load("scripts/eval/paired_matched_kl_bootstrap.py", "paired_matched_kl_bootstrap") + +TERMS = ["damn", "crap", "bullshit"] +PROMPTS = ["Oh", "Holy", "This is", "You are such a", "That really", "I am so", "You really"] +# Per-method KL per scale: SRP reaches KL 0.05 at scale 2, the baseline at scale 4. +KL = { + "feature_suppression": {1.0: 0.02, 2.0: 0.05, 4.0: 0.2}, + "mean_row_direction": {1.0: 0.004, 2.0: 0.01, 4.0: 0.05, 8.0: 0.2}, +} + + +def _write_rows(path: Path) -> None: + fields = [ + "id", + "prompt", + "bad", + "good", + "split", + "method", + "scale", + "bad_prob_reduction", + "flip", + "full_vocab_kl_bits", + ] + rows = [] + for ti, term in enumerate(TERMS): + for pi, prompt in enumerate(PROMPTS): + cid = f"c{ti}_{pi}" + for method, kls in KL.items(): + for scale, kl in kls.items(): + # Deterministic outcome: SRP suppresses 0.1 more than the baseline. + base = 0.05 * (ti + 1) + 0.01 * pi + dp = base + (0.1 if method == "feature_suppression" else 0.0) + rows.append( + { + "id": cid, + "prompt": prompt, + "bad": term, + "good": "ref", + "split": "heldout", + "method": method, + "scale": scale, + "bad_prob_reduction": dp, + "flip": int(method == "feature_suppression"), + "full_vocab_kl_bits": kl, + } + ) + # A discovery-split row that must be ignored. + rows.append({**rows[0], "id": "disc", "split": "discovery", "bad_prob_reduction": 99.0}) + with path.open("w", newline="") as f: + w = csv.DictWriter(f, fieldnames=fields) + w.writeheader() + w.writerows(rows) + + +def test_paired_bootstrap_on_synthetic_rows(tmp_path: Path) -> None: + csv_path = tmp_path / "candidate_constrained_rows.csv" + _write_rows(csv_path) + out = tmp_path / "paired.json" + rc = boot.main( + ["--input", f"toy={csv_path}", "--out", str(out), "--n-boot", "50", "--seed", "0", "--target-kl", "0.05", "0.2"] + ) + assert rc == 0 + results = json.loads(out.read_text()) + # Only mean_row_direction is present -> 2 targets x 1 baseline x 2 outcomes. + assert len(results) == 4 + assert {r["model"] for r in results} == {"toy"} + by_key = {(r["target_kl"], r["outcome"]): r for r in results} + + r = by_key[(0.05, "dp")] + assert r["baseline"] == "mean_row_direction" + assert r["srp_scale"] == 2.0 and r["base_scale"] == 4.0 + assert r["srp_kl"] == 0.05 and r["base_kl"] == 0.05 + assert r["n"] == len(TERMS) * len(PROMPTS) + assert abs(r["mean_diff"] - 0.1) < 1e-9 + # Constant per-candidate difference -> both cluster CIs collapse onto it. + assert all(abs(v - 0.1) < 1e-9 for v in r["ci_term"]) + assert all(abs(v - 0.1) < 1e-9 for v in r["ci_prompt"]) + assert by_key[(0.05, "flip")]["mean_diff"] == 1.0 + assert by_key[(0.2, "dp")]["srp_scale"] == 4.0 and by_key[(0.2, "dp")]["base_scale"] == 8.0 + + # Seeded: a second run reproduces the file byte for byte. + out2 = tmp_path / "paired2.json" + boot.main(["--input", f"toy={csv_path}", "--out", str(out2), "--n-boot", "50", "--target-kl", "0.05", "0.2"]) + assert out2.read_text() == out.read_text() + + +def test_scale_for_kl_respects_two_x_window() -> None: + rows = [ + {"method": "m", "scale": "1", "full_vocab_kl_bits": "0.5"}, + {"method": "m", "scale": "2", "full_vocab_kl_bits": "1.0"}, + ] + assert boot.scale_for_kl(rows, "m", 0.02) == (None, None) + scale, med = boot.scale_for_kl(rows, "m", 0.6) + assert scale == 1.0 and med == 0.5 + + +def test_cluster_ci_seeded_and_bracketing() -> None: + import numpy as np + + d = np.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]) + clusters = np.array(["a", "a", "b", "b", "c", "c"]) + lo, hi = boot.cluster_ci(d, clusters, n_boot=200, seed=3) + assert boot.cluster_ci(d, clusters, n_boot=200, seed=3) == (lo, hi) + assert 0.0 <= lo <= d.mean() <= hi <= 5.0 diff --git a/tests/test_seed_stability_scripts.py b/tests/test_seed_stability_scripts.py new file mode 100644 index 0000000..5bd2dc9 --- /dev/null +++ b/tests/test_seed_stability_scripts.py @@ -0,0 +1,203 @@ +"""CPU smoke tests for the cross-seed stability scripts on synthetic inputs. + +Three tiny dictionaries, a 40-row readout and a stub tokenizer; no network, no +GPU. Identical dictionaries must reproduce each other exactly (same-side +Jaccard 1, best-single Jaccard 1, held-out core recall 1), which pins the +shared pipeline in ``sparse_readout_prism.research.seed_stability``. +""" + +from __future__ import annotations + +import importlib.util +import json +import sys +from pathlib import Path + +import numpy as np +import torch + +from sparse_readout_prism.research.seed_stability import ( + load_contrast_pairs, + load_dictionary, + single_token_id, + topk_codes, +) + +ROOT = Path(__file__).resolve().parents[1] +V, D_MODEL, D_FEAT, K = 40, 8, 64, 4 + + +class _StubTokenizer: + def __init__(self, vocab: list[str]) -> None: + self.vocab = vocab + self.index = {w: i for i, w in enumerate(vocab)} + + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: + w = text.strip() + return [self.index[w]] if w in self.index else [0, 1] + + def decode(self, ids: list[int]) -> str: + return " " + self.vocab[int(ids[0])] + + +def _load_script(name: str): + path = ROOT / "scripts" / "eval" / f"{name}.py" + spec = importlib.util.spec_from_file_location(f"_seed_stability_{name}", path) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def _inputs(tmp_path: Path, identical: bool): + g = torch.Generator().manual_seed(0) + vocab = [f"w{i:02d}" for i in range(V)] + W = torch.randn(V, D_MODEL, generator=g) + h_LN = torch.randn(3, 5, D_MODEL, generator=g) + dicts = [] + for s in range(3): + gs = torch.Generator().manual_seed(100 if identical else 100 + s) + dec = torch.randn(D_FEAT, D_MODEL, generator=gs) + dec = dec / dec.norm(dim=1, keepdim=True) + ckpt = { + "model_state_dict": {"decoder": dec, "encoder.weight": dec.clone(), "encoder.bias": torch.zeros(D_FEAT)}, + "factorizer": {"architecture": "topk", "d_features": D_FEAT, "k": K}, + } + path = tmp_path / f"ckpt_s{s}.pt" + torch.save(ckpt, path) + dicts.append(load_dictionary(path)) + recs = [{"target_a": vocab[2 * i], "target_b": vocab[2 * i + 1]} for i in range(12)] + recs.append({"target_a": "notaword", "target_b": vocab[0]}) # multi-token target: skipped + recs.append({"target_a": vocab[0], "target_b": vocab[1]}) # duplicate pair: skipped + bank = tmp_path / "bank.jsonl" + bank.write_text("\n".join(json.dumps(r) for r in recs) + "\n") + return _StubTokenizer(vocab), W, h_LN, dicts, bank + + +def test_dictionary_and_bank_helpers(tmp_path: Path) -> None: + tok, W, _, dicts, bank = _inputs(tmp_path, identical=False) + dec, enc_w, enc_b, k = dicts[0] + assert dec.shape == (D_FEAT, D_MODEL) and enc_w.shape == (D_FEAT, D_MODEL) and enc_b.shape == (D_FEAT,) + assert k == K + assert single_token_id(tok, "w03") == 3 and single_token_id(tok, "notaword") is None + pairs = load_contrast_pairs(bank, tok, np.random.default_rng(0), max_contrasts=10) + assert len(pairs) == 10 and len({(a, b) for a, b, _, _ in pairs}) == 10 + assert pairs == load_contrast_pairs(bank, tok, np.random.default_rng(0), max_contrasts=10) + codes = topk_codes(W[:5], enc_w, enc_b, K) + assert codes.shape == (5, D_FEAT) and bool(((codes > 0).sum(dim=1) <= K).all()) + + +def test_cross_seed_identical_dictionaries_reproduce(tmp_path: Path) -> None: + tok, W, h_LN, dicts, bank = _inputs(tmp_path, identical=True) + mod = _load_script("cross_seed_stability") + rng = np.random.default_rng(0) + pairs = load_contrast_pairs(bank, tok, rng, 150) + out = mod.run( + dicts, [0, 1, 2], W, h_LN, tok, pairs, rng, top_m=3, top_r=4, n_sample=16, n_hidden=8, width_tag="test" + ) + assert out["n_contrasts"] == 12 and out["k"] == K + assert out["same_side"]["mean"] == 1.0 + assert out["frac_contrasts_above_null_p90"] == 1.0 + for v in out["basis_nn_cosine"].values(): + assert abs(v["nn_cos_used"]["median"] - 1.0) < 1e-5 + for r in out["matched_projection_corr"].values(): + assert abs(r - 1.0) < 1e-5 + + +def test_cross_seed_different_seeds_in_range(tmp_path: Path) -> None: + tok, W, h_LN, dicts, bank = _inputs(tmp_path, identical=False) + mod = _load_script("cross_seed_stability") + rng = np.random.default_rng(0) + pairs = load_contrast_pairs(bank, tok, rng, 150) + out = mod.run( + dicts, [0, 1, 2], W, h_LN, tok, pairs, rng, top_m=3, top_r=4, n_sample=16, n_hidden=8, width_tag="test" + ) + for key in ("same_side", "cross_side", "cross_contrast_null"): + assert 0.0 <= out[key]["mean"] <= 1.0 + assert out["same_side"]["n"] == 3 * 12 * 2 # seed pairs x contrasts x sides + assert set(out["basis_nn_cosine"]) == {"s0-s1", "s0-s2", "s1-s2"} + assert set(out["matched_projection_corr"]) == {"s0-s1", "s0-s2", "s1-s2"} + + +def _run_matching(mod, tok, W, dicts, bank): + rng = np.random.default_rng(0) + pairs = load_contrast_pairs(bank, tok, rng, 150) + return mod.run( + dicts, + [0, 1, 2], + W, + tok, + pairs, + rng, + torch.device("cpu"), + top_m=3, + top_r=4, + n_query=10, + n_null=20, + greedy_pool=20, + greedy_k=3, + min_group_size=4, + width_tag="test", + direction_check=True, + direction_null=20, + direction_seed=0, + ) + + +def test_feature_group_matching_identical_dictionaries(tmp_path: Path) -> None: + tok, W, _, dicts, bank = _inputs(tmp_path, identical=True) + out = _run_matching(_load_script("feature_group_matching"), tok, W, dicts, bank) + assert set(out["pairs"]) == {"s0->s1", "s0->s2", "s1->s2"} + for res in out["pairs"].values(): + assert res["n_query"] == 10 and len(res["per_query"]) == 10 + assert res["best_jaccard"]["median"] == 1.0 + assert res["frac_jaccard_ge_050"] == 1.0 + assert res["recall_at_3"]["median"] == 1.0 + # best match may be a tied duplicate group (source tie-break), so test the group, not the index + assert all(q["best_match_tokens"] == q["query_tokens"] for q in res["per_query"]) + assert all(-1.0 <= q["cosine"] <= 1.0 for q in res["per_query"]) + assert res["direction_check"]["n_below_null_p99"] == 0 + + +def test_feature_group_matching_direction_check_annotates_tail(tmp_path: Path) -> None: + tok, W, _, dicts, bank = _inputs(tmp_path, identical=False) + out = _run_matching(_load_script("feature_group_matching"), tok, W, dicts, bank) + for res in out["pairs"].values(): + check = res["direction_check"] + fails = [q for q in res["per_query"] if q["best_jaccard"] <= res["null99"]] + assert check["n_below_null_p99"] == len(fails) + assert all("max_decoder_cosine" in q for q in fails) + assert all("max_decoder_cosine" not in q for q in res["per_query"] if q["best_jaccard"] > res["null99"]) + assert 0 <= check["n_decoder_cosine_ge_0p5"] <= len(fails) + assert 0.0 <= res["null_best_jaccard"]["p99"] <= 1.0 + + +def test_loo_core_recovery_identical_dictionaries(tmp_path: Path) -> None: + tok, W, _, dicts, bank = _inputs(tmp_path, identical=True) + mod = _load_script("loo_core_recovery") + rng = np.random.default_rng(0) + pairs = load_contrast_pairs(bank, tok, rng, 150) + out = mod.run( + dicts, + W, + tok, + pairs, + torch.device("cpu"), + top_m=3, + top_r=4, + side_n=6, + n_clusters=8, + kmeans_iters=2, + kmeans_seed=0, + n_boot=50, + bootstrap_seed=1, + reference_dict=dicts[0], + width_tag="test", + ) + assert out["n_loo_cells"] == 12 * 2 * 3 # contrasts x sides x held-out seeds + assert out["srp_heldout_seed_recall"] == 1.0 + assert out["srp_ci95"] == [1.0, 1.0] + assert out["paper_dict_recall_same_cores"] == 1.0 + assert 0.0 <= out["knn_recall_same_cores"] <= 1.0 + assert 0.0 <= out["cluster_recall_same_cores"] <= 1.0 + assert out["ratio_srp_over_knn"] >= 1.0 diff --git a/tests/test_wsd_sense_groups.py b/tests/test_wsd_sense_groups.py new file mode 100644 index 0000000..2e11292 --- /dev/null +++ b/tests/test_wsd_sense_groups.py @@ -0,0 +1,201 @@ +"""CPU smoke tests for the CoarseWSD-20 sense analyses on a synthetic bundle. + +* scripts/analyze/analyze_wsd_sense_groups.py (tab:app-sense-alignment) +* scripts/analyze/analyze_wsd_classifier_framing.py (classifier-framing paragraph) + +The bundle mimics ``representations.pt`` from ``run_wsd_feature_alignment.py``: +two words, two or three senses, 16 dictionary positions, with a strong and a +weaker planted discriminative position per sense. No model, no network, well under a second. +""" + +from __future__ import annotations + +import importlib.util +import json +import sys +from pathlib import Path + +import numpy as np +import torch + +ROOT = Path(__file__).resolve().parents[1] + + +def _load(mod_name: str, rel_path: str): + spec = importlib.util.spec_from_file_location(mod_name, ROOT / rel_path) + mod = importlib.util.module_from_spec(spec) + sys.modules[mod_name] = mod + spec.loader.exec_module(mod) + return mod + + +sg = _load("analyze_wsd_sense_groups", "scripts/analyze/analyze_wsd_sense_groups.py") +cf = _load("analyze_wsd_classifier_framing", "scripts/analyze/analyze_wsd_classifier_framing.py") + +WIDTH = 16 +D_HIDDEN = 8 +# word -> {sense: (planted contribution position, n_train, n_test)} +DESIGN = { + "bank": {"0": (3, 30, 15), "1": (7, 10, 5)}, + "seal": {"0": (1, 20, 10), "1": (9, 12, 6), "2": (13, 8, 4)}, +} + + +def make_bundle(seed: int = 0) -> dict: + rng = np.random.default_rng(seed) + metadata, contribution, projection, hidden = [], [], [], [] + for word, senses in DESIGN.items(): + beta = rng.uniform(0.5, 2.0, size=WIDTH).astype(np.float32) + for sense, (pos, n_train, n_test) in senses.items(): + for split, n in (("train", n_train), ("test", n_test)): + for i in range(n): + p = rng.normal(0.0, 0.3, size=WIDTH).astype(np.float32) + p[pos] += 3.0 + p[pos + 1] += 1.5 # weaker second cue so group sizes 1 and 2 are both recoverable + h = rng.normal(0.0, 0.3, size=D_HIDDEN).astype(np.float32) + h[int(sense) % D_HIDDEN] += 3.0 + metadata.append( + { + "item_id": f"coarse:{word}:{split}:{sense}:{i}", + "kind": "context", + "dataset": "coarsewsd20", + "split": split, + "target": word, + "sense": sense, + "token_id": 100 + len(word), + "row_relative_error": 0.2, + } + ) + projection.append(p) + contribution.append(beta * p) + hidden.append(h) + contribution = torch.tensor(np.stack(contribution)) + n = contribution.shape[0] + return { + "metadata": metadata, + "hidden": torch.tensor(np.stack(hidden)), + "projection": torch.tensor(np.stack(projection)), + "contribution": contribution, + "beta": torch.stack([torch.ones(WIDTH)] * n), + "feature_ids": torch.stack([torch.arange(WIDTH, dtype=torch.int32)] * n), + "exact_logit": contribution.sum(dim=1) + 5.0, + "reconstructed_logit": contribution.sum(dim=1) + 5.0, + "target_logprob": torch.zeros(n), + "target_rank": torch.ones(n, dtype=torch.int32), + "run": {"dataset": "coarsewsd20", "model_id": "synthetic", "k": WIDTH, "seed": seed}, + } + + +def _word_arrays(bundle: dict, word: str, split: str) -> tuple[np.ndarray, np.ndarray]: + idx = [i for i, r in enumerate(bundle["metadata"]) if r["target"] == word and r["split"] == split] + x = bundle["contribution"].numpy()[idx] + y = np.array([bundle["metadata"][i]["sense"] for i in idx]) + return x, y + + +def _standardize(train_c: np.ndarray) -> np.ndarray: + """Train-statistic standardization, as the script applies before selection.""" + return (train_c - train_c.mean(axis=0)) / (train_c.std(axis=0) + 1e-6) + + +def test_select_anchors_recovers_planted_positions_and_is_disjoint(): + bundle = make_bundle() + projection = bundle["projection"].numpy() + for word, senses in DESIGN.items(): + idx = [i for i, r in enumerate(bundle["metadata"]) if r["target"] == word and r["split"] == "train"] + train_y = np.array([bundle["metadata"][i]["sense"] for i in idx]) + names = sorted(senses) + # On unscaled projections the strong cue is the single best feature under both selectors. + for selector in ("mean_diff", "top1_freq"): + anchors = sg.select_anchors(projection[idx], names, train_y, selector, 1) + assert {s: anchors[s][0] for s in names} == {s: senses[s][0] for s in names}, (word, selector) + # On train-standardized contributions (the script's default input) a group of two is the planted pair; + # standardization equalizes the two cues' gaps, so only the set is determined. + train_c, _ = _word_arrays(bundle, word, "train") + groups = sg.select_anchors(_standardize(train_c), names, train_y, "mean_diff", 2) + assert {s: set(groups[s]) for s in names} == {s: {senses[s][0], senses[s][0] + 1} for s in names}, word + groups = sg.select_anchors(_standardize(train_c), names, train_y, "mean_diff", 4) + flat = [p for s in names for p in groups[s]] + assert all(len(groups[s]) == 4 for s in names) + assert len(flat) == len(set(flat)), "groups must not share positions" + + +def test_group_prediction_and_balanced_accuracy(): + bundle = make_bundle() + train_c, train_y = _word_arrays(bundle, "bank", "train") + test_c, test_y = _word_arrays(bundle, "bank", "test") + anchors = sg.select_anchors(train_c, ["0", "1"], train_y, "mean_diff", 2) + acc, pred = sg.word_accuracy(test_c, anchors, ["0", "1"], test_y) + assert acc == 1.0 and sg.balanced_accuracy(pred, test_y) == 1.0 + # Balanced accuracy is the mean per-sense recall: all-majority on a 15/5 split is 0.5. + assert sg.balanced_accuracy(np.full_like(test_y, "0"), test_y) == 0.5 + + +def test_sense_groups_end_to_end_sweep(tmp_path, monkeypatch): + bundle_path = tmp_path / "representations.pt" + torch.save(make_bundle(), bundle_path) + out = tmp_path / "out" / "synthetic__srp.json" + argv = [ + "analyze_wsd_sense_groups.py", + "--bundle", + str(bundle_path), + "--out", + str(out), + "--group-sizes", + "1,2", + "--n-boot", + "20", + "--n-null", + "10", + "--seed", + "0", + ] + monkeypatch.setattr(sys, "argv", argv) + assert sg.main() == 0 + for g in (1, 2): + path = out.with_name(f"synthetic__srp_g{g}.json") + assert path.exists(), path + summary = json.loads(path.read_text()) + assert summary["group_size"] == g and summary["basis"] == "srp" and summary["standardized"] is True + assert summary["n_words"] == 2 and set(summary["per_word"]) == set(DESIGN) + assert summary["word_mean_balanced"] > 0.95 + assert summary["word_mean_null_balanced"] < 0.8 + assert summary["word_mean_majority_balanced"] == (1 / 2 + 1 / 3) / 2 + assert summary["words_beating_majority_balanced"] == 2 + assert summary["words_beating_null_p95_balanced"] == 2 + assert summary["word_mean_full_account_balanced"] > 0.95 + assert summary["word_mean_hidden_balanced"] > 0.95 + assert summary["gate_fraction_overall"] == 1.0 + for word, senses in DESIGN.items(): + pw = summary["per_word"][word] + assert pw["n_senses"] == len(senses) + assert all(len(pw["anchors"][s]) == g for s in senses) + assert pw["chance_balanced_accuracy"] == 1 / len(senses) + # Single-size runs write --out verbatim and are reproducible under the seed. + single = tmp_path / "single.json" + monkeypatch.setattr( + sys, "argv", argv[:5] + ["--out", str(single), "--group-size", "1", "--n-boot", "20", "--n-null", "10"] + ) + assert sg.main() == 0 + a = json.loads(single.read_text()) + b = json.loads(out.with_name("synthetic__srp_g1.json").read_text()) + assert a["per_word"] == b["per_word"] + + +def test_classifier_framing_self_test_and_coarse_path(tmp_path): + assert cf.self_test() == 0 + bundle = make_bundle() + cf.validate_feature_alignment(bundle) + gate, gate_summary = cf.score_gate(bundle) + assert gate.all() and gate_summary["pass_fraction"] == 1.0 + metrics = cf.analyze_coarse( + bundle, ks=[2, 4], primary_k=2, primary_encoding="weighted", n_boot=5, n_null_seeds=2, seed=0 + ) + assert set(metrics["primary"]["methods"]) == set(cf.METHODS) + srp = metrics["primary"]["methods"]["srp"] + assert srp["n_words"] == 2 and srp["accuracy"] > 0.95 and srp["balanced_accuracy"] > 0.95 + # Projections and contributions differ only by a per-word rescaling of the coordinates, + # which top-|value| selection is not invariant to, so both are reported separately. + assert "projection_only" in metrics["results"]["weighted"]["4"] + assert set(metrics["primary"]["comparisons"]) == {f"srp_minus_{m}" for m in cf.METHODS[1:]} + assert set(metrics["primary"]["null_seed_sensitivity"]) == {"shuffled_srp", "random_srp_features"} diff --git a/uv.lock b/uv.lock index ffdf60f..aac8d1e 100644 --- a/uv.lock +++ b/uv.lock @@ -192,6 +192,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c7/0d/67e5b4109ea4a837e80daa87c2c696711955e40449a97e8926672534def2/click-8.4.1-py3-none-any.whl", hash = "sha256:482be17c6991b8c19c5429a1e995d9b0efdbb63172824c41f99965dc0ade8ec2", size = 116639, upload-time = "2026-05-22T04:08:35.26Z" }, ] +[[package]] +name = "cloudpickle" +version = "3.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/27/fb/576f067976d320f5f0114a8d9fa1215425441bb35627b1993e5afd8111e5/cloudpickle-3.1.2.tar.gz", hash = "sha256:7fda9eb655c9c230dab534f1983763de5835249750e85fbcef43aaa30a9a2414", size = 22330, upload-time = "2025-11-03T09:25:26.604Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl", hash = "sha256:9acb47f6afd73f60dc1df93bb801b472f05ff42fa6c84167d25cb206be1fbf4a", size = 22228, upload-time = "2025-11-03T09:25:25.534Z" }, +] + [[package]] name = "colorama" version = "0.4.6" @@ -464,6 +473,29 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, ] +[[package]] +name = "jlens" +version = "0.1.0" +source = { git = "https://github.com/anthropics/jacobian-lens#581d398613e5602a5af361e1c34d3a92ea82ba8e" } +dependencies = [ + { name = "huggingface-hub" }, + { name = "numpy" }, + { name = "torch" }, + { name = "transformers" }, +] + +[[package]] +name = "joblib" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cloudpickle" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d5/1d/537ab090f302b838943a1b56497dd53059b9a9b46a074936470173a2e207/joblib-1.6.0.tar.gz", hash = "sha256:2ccc96785b12046c08fd6d55839c12857831b54a3c1673ffadd2f04bfc4eda03", size = 327903, upload-time = "2026-08-31T09:39:04.122Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/18/53/84099323c2ec4be98d935f63c033ac4151ee83836ca1050ede3b3aadf155/joblib-1.6.0-py3-none-any.whl", hash = "sha256:3dbbf9f6e4b592a2357b854608e980fe6390d131d7a82f011a377ef2ebef7aba", size = 306115, upload-time = "2026-08-31T09:39:02.298Z" }, +] + [[package]] name = "markdown-it-py" version = "4.2.0" @@ -587,6 +619,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7e/82/69e539c4c2027f1e1697e09aaa2449243085a0edf81ae2c6341e84d769b6/multiprocess-0.70.19-py39-none-any.whl", hash = "sha256:0d4b4397ed669d371c81dcd1ef33fd384a44d6c3de1bd0ca7ac06d837720d3c5", size = 133477, upload-time = "2026-01-19T06:47:38.619Z" }, ] +[[package]] +name = "narwhals" +version = "2.25.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/6f/7b/6248dada39781db1ab3ebf08943080df0796098515a87f6f8696d14ec744/narwhals-2.25.0.tar.gz", hash = "sha256:62c036c810662bf7820b7737077176313bc59350eeeefb808510f388c743e4b2", size = 677076, upload-time = "2026-08-20T18:10:15.454Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/eb/dc/55481808fd70ef1567cf13540ffd4702af3f74b112e35427564b03f79c2d/narwhals-2.25.0-py3-none-any.whl", hash = "sha256:1f0f403e8c7e4463cde9bfe78b12fdd809e3ae3dda6d9b2f802934fb9c7a6a8f", size = 467373, upload-time = "2026-08-20T18:10:13.834Z" }, +] + [[package]] name = "networkx" version = "3.6.1" @@ -1100,6 +1141,33 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5d/e6/ec8471c8072382cb91233ba7267fd931219753bb43814cbc71757bfd4dab/safetensors-0.7.0-cp38-abi3-win_amd64.whl", hash = "sha256:d1239932053f56f3456f32eb9625590cc7582e905021f94636202a864d470755", size = 341380, upload-time = "2025-11-19T15:18:44.427Z" }, ] +[[package]] +name = "scikit-learn" +version = "1.9.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "joblib" }, + { name = "narwhals" }, + { name = "numpy" }, + { name = "scipy" }, + { name = "threadpoolctl" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/fa/6f/37092bdb25f712817231799fc5674d8e704066a8a70c1d2d40517e18b4ab/scikit_learn-1.9.0.tar.gz", hash = "sha256:8833266989d3a5110178a9fae30783675460724d0e1efb13b14901d2c660c557", size = 7750767, upload-time = "2026-06-02T11:54:32.706Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f5/be/e844fd9586e66540a15b71924d17a6cbc1bb749e81ddd0a796bcdba4c055/scikit_learn-1.9.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:9db6f4d34e68c8899e4cab27fdf8eafe6ed21f2ba52ceb25ea250cd237f8e47b", size = 8789686, upload-time = "2026-06-02T11:53:05.439Z" }, + { url = "https://files.pythonhosted.org/packages/42/e2/ff880f62677a17d035817d543cb0fc8727d01eccbee81c5f7fc733a9d856/scikit_learn-1.9.0-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:f401448645a3e7bc115aa3c094097865155b34bff1cba8101857d9104e99074c", size = 8256782, upload-time = "2026-06-02T11:53:08.904Z" }, + { url = "https://files.pythonhosted.org/packages/25/64/eb40435e1a508ab1b4e284ce43ae80f6a162e5be5e38ed5a6fab467a9ea4/scikit_learn-1.9.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:fd3a8ef0c758555a3b23c03adaa858af32f7736785ded50ad5991f59c4ed03fa", size = 8992419, upload-time = "2026-06-02T11:53:11.551Z" }, + { url = "https://files.pythonhosted.org/packages/8d/da/4810a28e473185429e45a57eebcc91fc991b33d889cc0676063e671db03d/scikit_learn-1.9.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f7e254636164090da847715a27f8e5478feb98c40a9e0ee90cbd277de9e5ceb8", size = 9281411, upload-time = "2026-06-02T11:53:15.063Z" }, + { url = "https://files.pythonhosted.org/packages/3b/67/be3d369f40d8178ba3bd86635d132e08cb5329b023e4669d9426d84bc007/scikit_learn-1.9.0-cp311-cp311-win_amd64.whl", hash = "sha256:5dc1818c77575d149e25fce9ef82dd7b7263ae372f03494158668ad632a69759", size = 8272736, upload-time = "2026-06-02T11:53:18.108Z" }, + { url = "https://files.pythonhosted.org/packages/37/79/a733f02dc2118da7e77a134b34f39f40201a353311b011d20859d2db3556/scikit_learn-1.9.0-cp311-cp311-win_arm64.whl", hash = "sha256:366652351f092b219c248f1e72821e841960a63d8f358f1dcfd54dc1cbdbbc28", size = 7919564, upload-time = "2026-06-02T11:53:21.2Z" }, + { url = "https://files.pythonhosted.org/packages/ac/20/75f915ff375d6249e6550ac740fdbbd66159a068fd3af1400ff62036b07a/scikit_learn-1.9.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:2bd41b0d201bc81575531b96b713d3eb5e5f50fb0b82101ff0f92294fdc236ac", size = 8741122, upload-time = "2026-06-02T11:53:24.08Z" }, + { url = "https://files.pythonhosted.org/packages/cc/d5/2b5148f2279196775e1db2aeb85d14b70ac80e7e32b3b28e7ebeafb0901d/scikit_learn-1.9.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:5be45aa4a42a68a533913a6ed736cf309de2226411c79ef8d609a5456f1939b1", size = 8261512, upload-time = "2026-06-02T11:53:27.183Z" }, + { url = "https://files.pythonhosted.org/packages/a0/ee/5adbc77656b71f9456a2f5a7a9fdb4bcf9207a6b962889f1c2f9323afa4e/scikit_learn-1.9.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5e50ed4da51974e86e940690e9a3d82e729b62b5a49f7c9bac534d515d39d86f", size = 8837603, upload-time = "2026-06-02T11:53:30.328Z" }, + { url = "https://files.pythonhosted.org/packages/6c/c2/63fdda36c56437eeb44aaf9493c8bcd62ce230ab1598924fc626ffbfa943/scikit_learn-1.9.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:056c92bb67ad4c28463c2f2653d9701449201e7e7a9e94e321be0f71c4fef2b8", size = 9132097, upload-time = "2026-06-02T11:53:33.456Z" }, + { url = "https://files.pythonhosted.org/packages/83/a4/c8e67227c680e2259c8864ae72ff48b06e16a6f51253a22167aa02a8aa4e/scikit_learn-1.9.0-cp312-cp312-win_amd64.whl", hash = "sha256:4306775fad04cc4b472a1b15af1ae9cede1540fbfcc17fbce3767cd8dc7ae283", size = 8211173, upload-time = "2026-06-02T11:53:36.602Z" }, + { url = "https://files.pythonhosted.org/packages/cf/fd/3c0863792e98e67e9184aa4029288a175935eb65443afcd30d4f143450cf/scikit_learn-1.9.0-cp312-cp312-win_arm64.whl", hash = "sha256:26e22435f63bcdcf396b574273f29f13dd531f5ea035801f5be10ba1540a4e60", size = 7867451, upload-time = "2026-06-02T11:53:39.075Z" }, +] + [[package]] name = "scipy" version = "1.17.1" @@ -1160,13 +1228,15 @@ wheels = [ [[package]] name = "sparse-readout-prism" -version = "0.1.0" +version = "0.2.0" source = { editable = "." } dependencies = [ { name = "accelerate" }, { name = "datasets" }, { name = "numpy" }, + { name = "pandas" }, { name = "pyyaml" }, + { name = "scikit-learn" }, { name = "scipy" }, { name = "torch" }, { name = "transformers" }, @@ -1177,20 +1247,26 @@ dev = [ { name = "pytest" }, { name = "ruff" }, ] +lens = [ + { name = "jlens" }, +] [package.metadata] requires-dist = [ { name = "accelerate", specifier = ">=1.0,<2" }, { name = "datasets", specifier = ">=4.8.4,<5" }, + { name = "jlens", marker = "extra == 'lens'", git = "https://github.com/anthropics/jacobian-lens" }, { name = "numpy", specifier = ">=1.26,<3" }, + { name = "pandas", specifier = ">=2.2,<3" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=7" }, { name = "pyyaml", specifier = ">=6,<7" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.7" }, + { name = "scikit-learn", specifier = ">=1.5,<2" }, { name = "scipy", specifier = ">=1.15.3,<2" }, { name = "torch", specifier = ">=2.5,<3" }, { name = "transformers", specifier = ">=5.0,<6" }, ] -provides-extras = ["dev"] +provides-extras = ["dev", "lens"] [[package]] name = "sympy" @@ -1204,6 +1280,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, ] +[[package]] +name = "threadpoolctl" +version = "3.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/b7/4d/08c89e34946fce2aec4fbb45c9016efd5f4d7f24af8e5d93296e935631d8/threadpoolctl-3.6.0.tar.gz", hash = "sha256:8ab8b4aa3491d812b623328249fab5302a68d2d71745c8a4c719a2fcaba9f44e", size = 21274, upload-time = "2025-03-13T13:49:23.031Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/32/d5/f9a850d79b0851d1d4ef6456097579a9005b31fea68726a4ae5f2d82ddd9/threadpoolctl-3.6.0-py3-none-any.whl", hash = "sha256:43a0b8fd5a2928500110039e43a5eed8480b918967083ea48dc3ab9f13c4a7fb", size = 18638, upload-time = "2025-03-13T13:49:21.846Z" }, +] + [[package]] name = "tokenizers" version = "0.22.2"