Reference J-lens loader, WorkspaceBench cosine readout, docs cleanup

CLAUDE.md

@@ -1,6 +1,6 @@
 # CLAUDE.md
 
-Guidance for Claude Code (claude.ai/code) when working in this repository.
+This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
 
 ## Project overview
 
@@ -23,8 +23,8 @@ Token Search"). Pivotal tokens are one of three scales it searches.
 
 ## The invariant that matters most
 
-**A latent event's `score` is a J-lens readout probability. It is NOT a
-probability delta.**
+**A latent event's `score` is a J-lens readout score (a probability, or a cosine
+for `jlens_cosine`). It is NOT a probability delta.**
 
 - Latent events have `prob_delta`, `prob_before`, `prob_after`, and
   `is_positive` set to `None`, deliberately. Do not fill them in.
@@ -41,6 +41,13 @@ probability delta.**
   A single 0.5 floor across both scales keeps the filler meta-token `" the"`
   (readout 0.92) and throws away a pivotal token worth +0.45. This has already
   been shipped once by accident; do not reintroduce it.
+  `tests/test_regressions.py` pins this and the other bugs listed under "Things
+  that will bite you".
+- The same rule holds one level down: `jlens_cosine` scores are cosines, while
+  `jlens`/`logit_lens` scores are probabilities. `min_readout_score`,
+  `most_surfaced()` and exporter `min_score` raise on a mix, and `summary()`
+  reports `readout_score_by_scale`. Scales are defined by `READOUT_SCORE_SCALE`
+  in `pts/events.py`; register any new readout method there.
 - Never count latent events as positive or negative — they have no valence.
 - `logit_lens` readouts are weaker evidence than `jlens` ones; keep
   `readout_method` visible so they can be filtered apart.
@@ -57,7 +64,8 @@ which is the main way this project can mislead people. Tests in
 pip install -e .           # core + model deps
 pip install -e '.[all]'    # + sentence-transformers, math-verify, pytest
 
-pytest tests/ -q           # 66 tests, ~15s, uses a tiny random model
+pytest tests/ -q           # 94 tests, ~20s, uses a tiny random model
+pytest tests/test_jlens.py::test_causal_mask_holds -q   # a single test
 
 pts run --granularity token|sentence|latent|all --model M --output-path events.jsonl
 pts fit-jlens --model M --output-path ./jlens/m          # calibrate the Jacobian lens
@@ -84,13 +92,17 @@ pts push --input-path X --hf-repo-id user/repo
 | `pts/searchers/base.py` | Model loading, prompt formatting, the probability cache. |
 | `pts/searchers/{token,sentence,latent,reasoning}.py` | The four searchers. |
 | `pts/oracle.py`, `pts/dataset.py` | Success evaluation and dataset loading (largely unchanged). |
+| `pts/verification.py` | Arithmetic CoT verification used by Sentence PTS. Imports torch at module level. |
+| `pts/cli.py` | Every `pts` subcommand and its flags. |
 | `pts/exporters.py` | All output formats + dataset cards. |
 | `pts/core.py`, `pts/storage.py`, `pts/thought_anchors.py` | legacy compatibility shims. |
 
 `import pts` must **not** require torch or transformers. The schema, storage,
 classification, and linking layers are pure Python; model-touching code is
 imported lazily via `__getattr__`. Keep it that way.
 
+`research/` and `visualizer/` are standalone and excluded from the package.
+
 ## The J-lens
 
 `J_l = E[∂h_final,t' / ∂h_l,t]`, averaged over source positions `t`, all later
@@ -108,9 +120,18 @@ row — the paper's `O(n × d_model)` cost.
 `torch.autograd.functional.jacobian`, and `test_causal_mask_holds` checks the
 assumption directly. Do not weaken those tests.
 
-No reference code was released with the workspace paper — this is written from
-the equations and has **not** been validated against the authors' results. Say so
-when writing docs. "Meta-token" is our term, not the paper's.
+`fit` was written from the equations before Anthropic released reference code
+(`anthropics/jacobian-lens`). Its estimator differs: the reference skips the first
+16 positions and the last one, and sums over `t' >= t` where we average. Our
+fitted lenses have **not** been validated against theirs. Say so when writing
+docs. "Meta-token" is our term, not the paper's.
+
+`JLens.load` also reads reference `.pt` lenses (`{"J": {layer: [d,d]}, ...}`),
+locally or as `hf://<org>/<repo>/<file.pt>`. Neuronpedia hosts ~40 at
+`neuronpedia/jacobian-lens`. They use our convention (`J @ h` on decoder-block
+outputs, target = final block), and `lm_head(final_norm(J h))` is the reference
+readout. Loading checks width and depth against the model. `jlens_cosine` is
+WorkspaceBench's readout, not the paper's.
 
 ## Things that will bite you
 

README.md

@@ -89,7 +89,9 @@ pts run --granularity all --model Qwen/Qwen3-0.6B \
         --readout-method jlens --jlens-path ./jlens --output-path events.jsonl
 ```
 
-No J-lens yet? `--readout-method logit_lens` needs no calibration. It is the same
+No J-lens yet? `--jlens-path` also takes a reference lens straight from the Hub,
+e.g. `hf://neuronpedia/jacobian-lens/qwen3-1.7b/jlens/Salesforce-wikitext/Qwen3-1.7B_jacobian_lens.pt`.
+Or use `--readout-method logit_lens`, which needs no calibration. It is the same
 readout with `J = I`, and it is a weaker signal. See
 [docs/latent_pts.md](docs/latent_pts.md).
 
@@ -162,8 +164,6 @@ and compatible with that work, not the same as it.
 ## Datasets
 
 - [codelion/Qwen3-0.6B-pts](https://huggingface.co/datasets/codelion/Qwen3-0.6B-pts)
-- [codelion/Qwen3-0.6B-pts-thought-anchors](https://huggingface.co/datasets/codelion/Qwen3-0.6B-pts-thought-anchors)
-- [codelion/Qwen3-0.6B-pts-steering-vectors](https://huggingface.co/datasets/codelion/Qwen3-0.6B-pts-steering-vectors)
 - [codelion/DeepSeek-R1-Distill-Qwen-1.5B-pts](https://huggingface.co/datasets/codelion/DeepSeek-R1-Distill-Qwen-1.5B-pts)
 
 ## References

docs/dataset_schema.md

@@ -96,7 +96,7 @@ a hypothetical. It is what the first implementation did.
 |---|---|---|
 | `search_method` | string | `token_pts` \| `sentence_pts` \| `latent_pts` \| `latent_pts_enrichment` |
 | `intervention_type` | string? | `append_token` \| `replace_sentence` \| `remove_sentence` |
-| `readout_method` | string? | Latent only: `jlens` \| `logit_lens`. **`logit_lens` is weaker evidence**, so filter on this. |
+| `readout_method` | string? | Latent only: `jlens` \| `jlens_cosine` \| `logit_lens`. **`logit_lens` is weaker evidence**, so filter on this. `jlens_cosine` scores a cosine, not a probability, so never compare its `score` with the others. |
 
 ### Classification and links
 

docs/latent_pts.md

@@ -57,6 +57,14 @@ Jacobian for all source positions at the same time. That is the paper's stated
 cost of `O(n x d_model)` backward passes. Dividing row `t` by the count of
 `t' >= t` turns the sum into the mean the definition asks for.
 
+Anthropic's reference code, [`anthropics/jacobian-lens`](https://github.com/anthropics/jacobian-lens),
+came out after `fit` was written, and its estimator differs in two ways: it
+skips the first 16 positions (attention sinks) and the last one, and it keeps
+the sum over `t' >= t` instead of dividing it into a mean. So a lens from
+`pts fit-jlens` is not the same matrix as a reference lens, and ours has not
+been validated against theirs. If one exists for your model, prefer the
+reference lens (below).
+
 `tests/test_jlens.py` checks this against a brute-force
 `torch.autograd.functional.jacobian`. The two agree to about 1e-8, and a separate
 test checks the causal-mask assumption directly.
@@ -70,6 +78,17 @@ It is also a weaker signal. A workspace claim is a claim about future influence,
 and the logit lens does not measure future influence. Events read this way are
 tagged `readout_method: "logit_lens"` so you can filter them out.
 
+### The cosine readout
+
+[WorkspaceBench](https://github.com/camilablank/workspace-bench) reads its J-lens
+arm differently: it ranks tokens by the cosine between `h_l` and each token's
+J-lens vector `J_l^T w_t`, with no final norm and no softmax. That stops tokens
+with large J-lens vectors from dominating every readout.
+`--readout-method jlens_cosine` does the same. Its `score` is a cosine in
+[-1, 1], not a probability, so `EventStorage` refuses to rank or threshold it
+together with `jlens` or `logit_lens` events. Filter to one `readout_method`
+first.
+
 ## Usage
 
 ### 1. Fit a J-lens (once per model)
@@ -88,6 +107,24 @@ the lens is a property of the model rather than the task, so the source dataset
 works fine with no labels. Pass `--calibration-file` for your own text, one
 sequence per line. Lower `--basis-chunk` if you run out of memory.
 
+### 1b. Or use a reference lens
+
+Neuronpedia hosts reference-code lenses for about 40 models (GPT-2, Gemma 2/3/4,
+Llama 3.1/3.3, Qwen3/3.5/3.6, OLMo 3 and others) at
+[`neuronpedia/jacobian-lens`](https://huggingface.co/neuronpedia/jacobian-lens).
+`--jlens-path` takes one directly as `hf://<org>/<repo>/<file.pt>`, or a local
+`.pt` file in the same format:
+
+```bash
+pts enrich --input-path events.jsonl --output-path events_latent.jsonl \
+  --model Qwen/Qwen3-1.7B --with-latent --readout-method jlens \
+  --jlens-path hf://neuronpedia/jacobian-lens/qwen3-1.7b/jlens/Salesforce-wikitext/Qwen3-1.7B_jacobian_lens.pt
+```
+
+These lenses cover every layer below the last, so any workspace layer works. A
+lens whose width or depth does not fit the model is refused. When the
+`config.yaml` next to it names a different model, you get a warning.
+
 ### 2. Enrich a dataset you already have
 
 This is the main path. It reuses curated token and sentence datasets instead of
@@ -113,52 +150,16 @@ pts run --granularity all --model Qwen/Qwen3-0.6B \
         --output-path events.jsonl
 ```
 
-## What to keep in mind
-
-Each of these is a way the output can mislead you if you forget it.
-
-**A latent score is a readout probability, not a Δ-probability.** It says how
-strongly the lens surfaces a token, not how much that token changed the answer.
-It is not comparable to the `prob_delta` on token and sentence events.
-`prob_delta` and `is_positive` stay `null` on latent events on purpose.
-
-**Latent events are observational.** Enrichment reports what the lens sees. It
-does not show that a meta-token caused an emitted event. That would take an
-intervention: steer or ablate the direction and re-measure success.
-
-**Readouts are noisy.** Neither lens is guaranteed to be faithful to what the
-model actually represents. A high-scoring `verify` might be an artifact of the
-unembedding geometry rather than a real concept.
-
-**This is an independent reimplementation.** No code was released with the
-workspace paper. The math here comes from the published equations and is checked
-for internal correctness, but it has not been validated against the authors'
-results, so do not report PTS numbers as reproducing theirs.
-
-**Links are heuristics.** `linked_event_ids` come from a weighted score (query
-match, context overlap, category agreement, timing), not a verified causal path.
-Run `--shuffle-control`: if the observed link scores do not clearly beat the
-shuffled baseline, the structure is not above chance.
-
-**Records are model-specific.** A pivotal token in one model tells you nothing
-about another model's workspace. Enriching with a different model is refused
-unless you pass `--allow-model-mismatch`, and even then both model ids are stored.
-
-## What would make it convincing
-
-The claim under test is that emitted pivotal tokens and thought-anchor sentences
-are often preceded by latent meta-tokens in the workspace. In rough order of
-strength, the evidence that would back it up:
-
-1. **Intervention.** Steer or ablate a meta-token direction and show success
-   probability moves.
-2. **Lead time with a control.** Show a category-matching meta-token appears `k`
-   tokens before the emitted event more often than chance, using
-   `--shuffle-control` as the baseline.
-3. **Consistent chains.** Show `latent: verification -> token: " Wait" ->
-   sentence: "Let me check..."` holds across many queries, not just anecdotes.
-4. **J-lens beats logit-lens.** If the effect is just as strong with
-   `--readout-method logit_lens`, it is not about future influence, and so not
-   about a workspace.
-
-Until at least (2) holds with a control, treat the output as exploratory.
+## Notes
+
+- A latent event's `score` is a readout score, not a Δ-probability: a
+  probability for `jlens` and `logit_lens`, a cosine for `jlens_cosine`. It is
+  not comparable to the `prob_delta` on token and sentence events, which is why
+  `prob_delta` and `is_positive` are `null` on latent events.
+- Latent events are observational. The lens shows what a meta-token leans toward,
+  not that it caused anything downstream. A causal claim would need an
+  intervention (steer or ablate, then re-measure).
+- Run `--shuffle-control`. If the link scores do not beat the shuffled baseline,
+  the structure is not above chance.
+- Records are model-specific. Enriching with a different model needs
+  `--allow-model-mismatch`.

pts/__init__.py

@@ -75,6 +75,7 @@
     "TokenExporter": ("pts.exporters", "TokenExporter"),
     "EventExporter": ("pts.exporters", "EventExporter"),
     "JLens": ("pts.latent.jlens", "JLens"),
+    "JLensCosine": ("pts.latent.jlens", "JLensCosine"),
     "Readout": ("pts.latent.jlens", "Readout"),
 }
 

pts/cli.py

@@ -538,13 +538,21 @@ def _add_latent_args(p) -> None:
     p.add_argument(
         "--readout-method",
         default="logit_lens",
-        choices=["jlens", "logit_lens"],
+        choices=["jlens", "jlens_cosine", "logit_lens"],
         help="How to read meta-tokens out of the workspace. 'jlens' needs "
-             "--jlens-path (fit one with `pts fit-jlens`). 'logit_lens' needs no "
+             "--jlens-path. 'jlens_cosine' ranks by cosine to each token's J-lens "
+             "vector, as WorkspaceBench does; its score is a cosine, not a "
+             "probability. 'logit_lens' needs no "
              "calibration but is weaker evidence: it reads what an activation "
              "would say now, not what it pushes the model to say later.",
     )
-    p.add_argument("--jlens-path", default=None, help="Directory holding fitted J-lens matrices")
+    p.add_argument(
+        "--jlens-path",
+        default=None,
+        help="A `pts fit-jlens` directory, a reference jacobian-lens .pt file, or "
+             "hf://<org>/<repo>/<file.pt> (e.g. a Neuronpedia lens from "
+             "hf://neuronpedia/jacobian-lens/...)",
+    )
     p.add_argument(
         "--workspace-layers",
         nargs="+",
@@ -565,7 +573,9 @@ def _add_latent_args(p) -> None:
     # Not --top-k: `pts run` already uses that for sampling.
     p.add_argument("--readout-top-k", type=int, default=25,
                    help="How many vocabulary tokens to read out per position")
-    p.add_argument("--min-score", type=float, default=0.01, help="Minimum readout score to keep")
+    p.add_argument("--min-score", type=float, default=0.01,
+                   help="Minimum readout score to keep (a probability, or a cosine "
+                        "for jlens_cosine)")
     p.add_argument("--keep-per-position", type=int, default=3,
                    help="Max meta-token events kept per position/layer")
     p.add_argument("--link-threshold", type=float, default=0.5, help="Minimum link score")

pts/event_storage.py

@@ -17,6 +17,7 @@
     EVENT_SENTENCE,
     EVENT_TOKEN,
     from_any_record,
+    readout_score_scale,
 )
 
 logger = logging.getLogger(__name__)
@@ -175,7 +176,13 @@ def filter(
         So the two scales get their own thresholds: ``min_prob_delta`` for
         emitted events, ``min_readout_score`` for latent ones. Each applies only
         to the scale it belongs to and leaves the other untouched.
+
+        Latent scores are not all on one scale either: ``jlens_cosine`` scores a
+        cosine, the others a probability. ``min_readout_score`` raises if the
+        latent events here mix the two; filter by ``readout_method`` first.
         """
+        if min_readout_score is not None:
+            self._single_readout_scale("min_readout_score")
         result = EventStorage()
 
         for evt in self.events:
@@ -235,8 +242,25 @@ def most_surfaced(self, n: int = 10) -> List[CausalReasoningEvent]:
         the token matters causally.
         """
         latent = self.by_event_type(EVENT_LATENT)
+        self._single_readout_scale("most_surfaced()")
         return sorted(latent, key=lambda e: e.score, reverse=True)[:n]
 
+    def _latent_scales(self) -> Dict[str, List[CausalReasoningEvent]]:
+        scales: Dict[str, List[CausalReasoningEvent]] = {}
+        for e in self.by_event_type(EVENT_LATENT):
+            scales.setdefault(readout_score_scale(e.readout_method), []).append(e)
+        return scales
+
+    def _single_readout_scale(self, operation: str) -> None:
+        scales = self._latent_scales()
+        if len(scales) > 1:
+            raise ValueError(
+                f"{operation} would compare latent scores on different scales "
+                f"({', '.join(sorted(scales))}): a jlens_cosine score is a cosine, "
+                "not a probability. Filter to one readout_method first, e.g. "
+                "filter(criteria={'readout_method': 'jlens'})."
+            )
+
     def queries(self) -> List[str]:
         seen = []
         for e in self.events:
@@ -262,6 +286,7 @@ def summary(self) -> Dict[str, Any]:
 
         emitted = [e for e in self.events if e.prob_delta is not None]
         latent = self.by_event_type(EVENT_LATENT)
+        scales = self._latent_scales()
 
         summary = {
             "total_events": len(self.events),
@@ -281,11 +306,23 @@ def summary(self) -> Dict[str, Any]:
             "max_abs_prob_delta": (
                 max(abs(e.prob_delta) for e in emitted) if emitted else None
             ),
-            # Latent only: these are readout probabilities, on a different scale.
+            # Latent only, and only when every readout is on one scale: these
+            # are readout scores, not probability deltas.
             "average_readout_score": (
-                sum(e.score for e in latent) / len(latent) if latent else None
+                sum(e.score for e in latent) / len(latent)
+                if latent and len(scales) == 1 else None
+            ),
+            "max_readout_score": (
+                max(e.score for e in latent) if latent and len(scales) == 1 else None
             ),
-            "max_readout_score": max((e.score for e in latent), default=None),
+            "readout_score_by_scale": {
+                scale: {
+                    "count": len(evts),
+                    "average": sum(e.score for e in evts) / len(evts),
+                    "max": max(e.score for e in evts),
+                }
+                for scale, evts in scales.items()
+            },
         }
         # the legacy format's get_anchor_summary reads these names; they mean emitted-only.
         summary["average_score"] = summary["average_abs_prob_delta"]

pts/events.py

@@ -54,6 +54,29 @@
     EVENT_SENTENCE: VISIBILITY_EMITTED,
 }
 
+# Readout methods for latent events, and the scale each one's ``score`` is on.
+# ``jlens`` and ``logit_lens`` score a softmax probability; ``jlens_cosine``
+# scores a cosine in [-1, 1]. Scores on different scales must never be ranked
+# or thresholded together.
+READOUT_JLENS = "jlens"
+READOUT_LOGIT_LENS = "logit_lens"
+READOUT_JLENS_COSINE = "jlens_cosine"
+
+SCORE_SCALE_PROBABILITY = "probability"
+SCORE_SCALE_COSINE = "cosine"
+
+READOUT_SCORE_SCALE = {
+    READOUT_JLENS: SCORE_SCALE_PROBABILITY,
+    READOUT_LOGIT_LENS: SCORE_SCALE_PROBABILITY,
+    READOUT_JLENS_COSINE: SCORE_SCALE_COSINE,
+}
+
+
+def readout_score_scale(readout_method: Optional[str]) -> str:
+    """The scale a latent event's ``score`` is on. Records that predate the
+    field were all softmax readouts, so a missing method means probability."""
+    return READOUT_SCORE_SCALE.get(readout_method or READOUT_JLENS, SCORE_SCALE_PROBABILITY)
+
 
 def _now() -> str:
     return time.strftime("%Y-%m-%dT%H:%M:%S")