From 6ecd50aa3c11562f264c11e584aaf5501c7b5230 Mon Sep 17 00:00:00 2001 From: Allan Scott Date: Sat, 13 Jun 2026 19:17:55 -0400 Subject: [PATCH 1/3] propose(readout-aligned-geometry): build the dictionary Gram on the basis the model reads through MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Polygram's geometry-tracks-behaviour claim (Spearman 0.637) is built on RAW decoder vectors (clustered.py:165 k-means input + centroid/variance/gamma-PCA), the "deceptively compressible" basis. The behaviourally-relevant object is the readout-aligned decision geometry — now a measured, model-general fact: powered R2 (tau_star_powered.py, 20k tokens, 3 seeds) shows a trained rank-r projection onto the readout-aligned subspace beats frozen SVD on open-class next-token recovery at every rung (GPT-2 +52pp, Pythia-70m +31pp, 160m +40pp). Proposes a `readout-aligned` geometric profile (new polygram/geometry/readout_aligned.py mirroring clustered.py) that projects decoder vectors onto the top-r SVD subspace of gain*U BEFORE clustering/knob assignment; from_sae_lens(u_matrix=, gain=) threads the readout geometry (required for the profile; caller supplies U/gain, no host load in core). Q-Orca emission transparent. Gate: head-to-head Spearman on GPT-2-small, WIN >= 0.70 vs 0.637, or NO-IMPROVEMENT (basis isn't the limit -> Reckoning #3 prune more likely). Proposal/ tasks/spec only. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../proposal.md | 96 +++++++++++++++++++ .../specs/readout-aligned-profile/spec.md | 51 ++++++++++ .../tasks.md | 47 +++++++++ 3 files changed, 194 insertions(+) create mode 100644 openspec/changes/add-readout-aligned-geometry-profile/proposal.md create mode 100644 openspec/changes/add-readout-aligned-geometry-profile/specs/readout-aligned-profile/spec.md create mode 100644 openspec/changes/add-readout-aligned-geometry-profile/tasks.md diff --git a/openspec/changes/add-readout-aligned-geometry-profile/proposal.md b/openspec/changes/add-readout-aligned-geometry-profile/proposal.md new file mode 100644 index 0000000..4acaff4 --- /dev/null +++ b/openspec/changes/add-readout-aligned-geometry-profile/proposal.md @@ -0,0 +1,96 @@ +# Readout-aligned dictionary geometry — build the Gram on the basis the model reads through + +Add a `readout-aligned` geometric profile that builds Polygram's dictionary geometry (clustering + knob +assignment) on feature directions **projected onto the model's readout subspace**, instead of on the **raw +decoder vectors** it uses today. This is the highest-leverage test of Polygram's load-bearing claim — and it is +now backed by hard cross-model evidence, not a hunch. + +## Why + +Polygram's load-bearing claim is **behavioural**: the predicted dictionary geometry must track real co-firing / +ablation-KL — measured as `Spearman(Polygram-overlap, co-firing Jaccard) = 0.637` +(`docs/research/behavioural-scaleup-probe.md`). But that geometry is built on **raw decoder vectors**: +`polygram/geometry/clustered.py` feeds the raw per-feature projections straight into k-means +(`clustered.py:165`), the cluster centroids (`:176`), the residual-variance fidelity (`:177`), and the +per-cluster-PCA γ assignment (`:182`). The raw decoder/activation spectrum is exactly the **deceptively +compressible** basis the research program warns against — the *behaviourally* relevant object is the +**readout-aligned decision geometry** (the directions the model's argmax actually reads through), not raw +decoder norm. + +That used to be a hypothesis; it is now a **measured, model-general** fact. A powered cross-model run +(fieldrun `tau_star_powered.py`, wikitext 20k tokens, 3 seeds) shows a trained rank-`r` projection onto the +**readout-aligned** subspace beats the frozen SVD lens on open-class next-token recovery at every rung: +**GPT-2 +52pp, Pythia-70m +31pp, Pythia-160m +40pp** (open-class R@32, all 6 model×rank cells win). The +readout-aligned directions, not the energy-weighted decoder spectrum, govern behaviour — across architectures. + +**So if Polygram fits its geometry in the raw-decoder basis, its 0.637 Spearman is partly measuring the +basis, not the model.** Fixing the basis is the single highest-leverage change to the geometry-tracks-behaviour +claim — and it is a fair, falsifiable test of **Reckoning #3** (is the quantum Gram load-bearing, or a +removable scaffold?): if the *right* basis raises the Spearman, the Gram earns its keep; if it doesn't move, +the predictive power isn't basis-limited and the prune becomes more likely. + +## What + +### 1. A `readout-aligned` geometric profile — `polygram/geometry/readout_aligned.py` + +Mirror `clustered.py`'s `ClusteredKnobAssignment` / `clustered()` factory with a `ReadoutAlignedKnobAssignment` +/ `readout_aligned()` that, **before** clustering + knob assignment, projects the decoder vectors onto the +**readout subspace**: + +``` +readout subspace R = top-r right singular directions of (gain ⊙ U) # U: (vocab, d_model) unembed +proj_readout = projections @ Rᵀ # (n_features, r) +# then k-means / centroids / residual-variance / γ-PCA run on proj_readout, not the raw projections +``` + +This is the same readout-aligned construction R2 validated (the unembed the model reads through, optionally +weighted by the final-norm `gain`). The strategy's `.assign(projections, ...)` does the projection internally, +so the `from_sae_lens` call site (`sae_import.py:813`) is otherwise unchanged. Register `readout_aligned()` in +`polygram/geometry/__init__.py` alongside `clustered()` / `uniform_sphere()`. + +### 2. Thread the readout geometry in — `from_sae_lens(u_matrix=, gain=)` + +`from_sae_lens` already takes `profile=`. Add net-new optional `u_matrix: np.ndarray | None` (the host unembed, +`(vocab, d_model)`) and `gain: np.ndarray | float | None` (final-norm gain, default ones). When +`profile="readout-aligned"`, `u_matrix` is **required** — raise a clear `ValueError` if missing. Polygram does +**not** load a host model itself (keeps the core numpy-only / low-dep); the caller supplies `U`/`gain` (the +behavioural extra's `runtime._load_host_model` already exposes a host for the probe). + +### 3. Before/after behavioural comparison + +Extend `examples/behavioural_gram_scaleup.py` to run **both** profiles (`clustered` vs +`readout-aligned`, same features/SAE/host) and emit a `polygram_overlap_readout_aligned` column into +`docs/research/data/scaleup_pairs.csv`; record the head-to-head Spearman table in +`docs/research/behavioural-scaleup-probe.md` (no new page). + +## Scope / what this is NOT + +- **Geometry / knob-assignment only.** The profile changes how β/γ are assigned; the encoding's structural form + and the **Q-Orca emission are transparent** to it (`_qorca_emit.py` / `Dictionary.gram()` consume the final + knob values, not the projection vectors). No change to the quantum machine structure. +- **Caller supplies `U`/`gain`.** No host-model load inside the Polygram core (no new hard dep on transformers). +- **v1 host: GPT-2-small** (the existing scaleup-probe substrate). Cross-arch (Pythia, via the now-merged + forge adapter) is a follow-up once the GPT-2 result lands. +- **`clustered` stays the default.** `readout-aligned` is opt-in; existing dictionaries/tests are byte-identical + unless the profile is selected. + +## Falsifiable acceptance gate (descriptive, both outcomes first-class) + +Re-run the behavioural scaleup probe on GPT-2-small head-to-head (raw `clustered` vs `readout-aligned`), same +panels: + +- **WIN:** readout-aligned raises `Spearman(Polygram-overlap, co-firing Jaccard)` to **≥ 0.70** (from 0.637) → + readout-alignment is the right basis; the Gram's predictive power was basis-limited (and the quantum lens + earns its keep on the right basis). +- **NO IMPROVEMENT:** Spearman does not move → the Gram's predictive power is **not** basis-limited; an + equally-valuable result that makes **Reckoning #3**'s prune (quantum lens = removable scaffold) more likely. + +Either way the verdict is the head-to-head table, descriptive; no necessity claims. + +## Related + +- `docs/research/behavioural-scaleup-probe.md` — the 0.637 baseline this targets. +- Manifesto Reckoning #3 (quantum lens load-bearing?) — this is a fair test of it. +- fieldrun R2 / `tau_star_powered.py` — the powered, model-general evidence that readout-aligned directions + govern behaviour; sae-forge `add-capability-ceiling-diagnostic` — the consumer (readout-aligned atom + *selection* is the action when its `selection_gap` is large). diff --git a/openspec/changes/add-readout-aligned-geometry-profile/specs/readout-aligned-profile/spec.md b/openspec/changes/add-readout-aligned-geometry-profile/specs/readout-aligned-profile/spec.md new file mode 100644 index 0000000..2059cd5 --- /dev/null +++ b/openspec/changes/add-readout-aligned-geometry-profile/specs/readout-aligned-profile/spec.md @@ -0,0 +1,51 @@ +# readout-aligned-profile Specification (delta) + +## ADDED Requirements + +### Requirement: Readout-aligned geometric profile + +Polygram SHALL provide a `readout-aligned` geometric profile (registered alongside `clustered` / +`uniform_sphere`) that builds the dictionary geometry — clustering, cluster centroids, residual-variance +fidelity, and γ assignment — on feature directions **projected onto the model's readout subspace**, not on the +raw decoder vectors. + +The readout subspace SHALL be the top-`r` right singular directions of `gain ⊙ U` (`U`: the host unembed, +`(vocab, d_model)`; `gain`: the final-norm gain, `(d_model,)`, default ones), and the geometry SHALL be built +on `projections @ Rᵀ`. The projection SHALL occur **inside** the profile's knob-assignment strategy, leaving +the `from_sae_lens` call site and the `GeometricProfile`/registry plumbing structurally unchanged. + +`get_profile("readout-aligned")` SHALL resolve to this profile. The `clustered` profile SHALL remain the +default; a run that does not select `readout-aligned` SHALL be byte-identical to current behaviour. + +#### Scenario: readout-aligned geometry differs from raw-decoder geometry + +- **GIVEN** a feature set and a host unembed `U` whose readout subspace re-orders the raw decoder directions +- **WHEN** a dictionary is built with `profile="readout-aligned"` and again with `profile="clustered"` +- **THEN** the two SHALL produce different cluster assignments / β/γ knobs (the geometry is built on a + different basis), while the `clustered` result is unchanged from the pre-change behaviour + +#### Scenario: deterministic given the readout inputs + +- **GIVEN** identical `(projections, U, gain, seed)` +- **WHEN** the readout-aligned profile assigns knobs +- **THEN** the assignment SHALL be identical across runs + +### Requirement: `from_sae_lens` accepts the readout geometry + +`from_sae_lens` SHALL accept optional `u_matrix: np.ndarray | None` (`(vocab, d_model)`) and +`gain: np.ndarray | float | None` (default ones). When the resolved profile is `readout-aligned`, `u_matrix` +SHALL be **required** and 2-D with its second axis equal to the feature dimension; absent or malformed, +construction SHALL raise a clear `ValueError`. Polygram SHALL NOT load a host model itself to obtain `U`/`gain` +— the caller supplies them (keeping the core dependency-light). + +#### Scenario: readout-aligned profile without a host unembed is rejected + +- **GIVEN** `profile="readout-aligned"` and `u_matrix=None` +- **WHEN** `from_sae_lens(...)` is called +- **THEN** it SHALL raise `ValueError` naming the missing `u_matrix` + +#### Scenario: default profile needs no readout inputs + +- **GIVEN** no `profile` (or `profile="clustered"`) and no `u_matrix` +- **WHEN** `from_sae_lens(...)` is called +- **THEN** it SHALL succeed and produce a result byte-identical to the pre-change behaviour diff --git a/openspec/changes/add-readout-aligned-geometry-profile/tasks.md b/openspec/changes/add-readout-aligned-geometry-profile/tasks.md new file mode 100644 index 0000000..c6e7c07 --- /dev/null +++ b/openspec/changes/add-readout-aligned-geometry-profile/tasks.md @@ -0,0 +1,47 @@ +# Implementation tasks + +## 0. Design pre-locks (blocking) + +- [ ] 0.1 Lock the readout subspace construction: `R = top-r right singular directions of (gain ⊙ U)` + (`U`: `(vocab, d_model)` unembed; `gain`: `(d_model,)` final-norm, default ones). `r` defaults to a + capped rank (e.g. `min(64, d_model)`); the geometry is built on `projections @ Rᵀ`. Confirm this matches + the R2 / sae-forge `_readout_aligned_order` construction. +- [ ] 0.2 Lock that the projection happens **inside** `ReadoutAlignedKnobAssignment.assign(...)` so the + `from_sae_lens` call site (`sae_import.py:813`) and the `GeometricProfile`/registry plumbing are unchanged. +- [ ] 0.3 Lock scope: knob assignment (β/γ) only; Q-Orca emission + `Dictionary.gram()` consume final knobs and + are transparent. No host-model load in the core. + +## 1. `polygram/geometry/readout_aligned.py` — the profile + +- [ ] 1.1 `ReadoutAlignedKnobAssignment` mirroring `ClusteredKnobAssignment` (`clustered.py:143`), carrying + `u_matrix` + `gain` (+ `readout_rank`); `.assign(projections, ...)` projects onto the readout subspace, then + runs the existing k-means / centroid / residual-variance / γ-PCA path on the projected vectors. +- [ ] 1.2 `ReadoutAlignedFidelity` (implements the `GeometricFidelity` protocol) + `readout_aligned()` factory + (mirroring `clustered()` at `clustered.py:258`). +- [ ] 1.3 Register `readout_aligned()` in `polygram/geometry/__init__.py` (+ `__all__`); ensure + `get_profile("readout-aligned")` resolves. + +## 2. `polygram/sae_import.py` — thread `u_matrix` / `gain` + +- [ ] 2.1 Add `u_matrix: np.ndarray | None = None`, `gain: np.ndarray | float | None = None` to `from_sae_lens` + (`:615`). When the resolved profile is `readout-aligned`, **require** `u_matrix` (2-D, second axis == d_model) + — raise a clear `ValueError` otherwise. Inject `u_matrix`/`gain` into the strategy before `.assign` (`:813`). +- [ ] 2.2 Default profile unchanged (`clustered`); a run without `profile="readout-aligned"` is byte-identical. + +## 3. Tests + +- [ ] 3.1 Synthetic: a feature set whose raw-decoder clustering differs from its readout-aligned clustering + (construct `U` so the readout subspace re-orders the directions) → the two profiles assign different β/γ / + cluster labels; `readout-aligned` without `u_matrix` raises `ValueError`. +- [ ] 3.2 `clustered` default path byte-identical (existing geometry tests green); `get_profile("readout-aligned")` + registered. +- [ ] 3.3 Determinism: same `(projections, U, gain, seed)` → identical assignment. + +## 4. Behavioural gate (blocking — the decomposition IS the result) + +- [ ] 4.1 Extend `examples/behavioural_gram_scaleup.py` to run `clustered` vs `readout-aligned` head-to-head on + GPT-2-small (supply `U`/`gain` from the host the probe already loads); emit + `polygram_overlap_readout_aligned` into `docs/research/data/scaleup_pairs.csv`. +- [ ] 4.2 Record the head-to-head Spearman table in `docs/research/behavioural-scaleup-probe.md` (in place, no + new page). **Descriptive verdict, pre-committed both ways:** WIN (Spearman ≥ 0.70) or NO-IMPROVEMENT + (basis isn't the limit → Reckoning #3 prune more likely). No necessity claims. From 5cc840a87bea2b574d53f0097280158889fb2306 Mon Sep 17 00:00:00 2001 From: Allan Scott Date: Sat, 13 Jun 2026 19:27:30 -0400 Subject: [PATCH 2/3] feat(readout-aligned): implement the readout-aligned geometry profile (tasks 1-3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements the profile from the proposal + addresses the review's edge cases: - polygram/geometry/readout_aligned.py: ReadoutAlignedKnobAssignment + ReadoutAlignedFidelity + readout_aligned() factory. readout_subspace(U, gain, d_model, rank) = top-r right singular dirs of gain*U via np.linalg.svd(..., full_matrices=False)[2][:r]; r capped at min(rank, d_model, vocab) (so readout_rank > rank(gain*U) is safe); gain scalar->broadcast / vector / None->ones. The knob assignment projects (n_features, d_model) -> (n_features, r) then DELEGATES to the calibrated ClusteredKnobAssignment logic in readout space (k-means / beta / per-cluster-PCA gamma), and the fidelity's cosine-overlap target is computed in the same projected space (consistent). readout_rank is tunable (default 64). - readout_variance_captured(): the review's suggested diagnostic — fraction of decoder Frobenius energy in the top-r readout subspace (quantifies how readout-aligned the SAE features already are). - from_sae_lens(u_matrix=, gain=, readout_rank=): inject the caller-supplied host unembed into the registry placeholder strategy+fidelity via dataclasses.replace; u_matrix REQUIRED for the profile (ValueError else), 2-D validated. Polygram loads no host itself. Default (clustered) path byte-identical. - Registered in geometry/__init__.py. tests/test_readout_aligned_profile.py (10): registration, validation, readout!=raw, determinism, edge cases, variance diagnostic. 109 clustered/profile/uniform regression tests unaffected. Q-Orca emission transparent (consumes final knobs). Behavioural head-to-head gate (task 4) next. Co-Authored-By: Claude Opus 4.8 (1M context) --- polygram/geometry/__init__.py | 11 ++ polygram/geometry/readout_aligned.py | 178 ++++++++++++++++++++++++++ polygram/sae_import.py | 34 +++++ tests/test_readout_aligned_profile.py | 117 +++++++++++++++++ 4 files changed, 340 insertions(+) create mode 100644 polygram/geometry/readout_aligned.py create mode 100644 tests/test_readout_aligned_profile.py diff --git a/polygram/geometry/__init__.py b/polygram/geometry/__init__.py index d87d81e..ddc1828 100644 --- a/polygram/geometry/__init__.py +++ b/polygram/geometry/__init__.py @@ -19,6 +19,12 @@ ) from polygram.geometry.learned_axis_assignment import LearnedKnobAssignment from polygram.geometry.profile import GeometricProfile +from polygram.geometry.readout_aligned import ( + ReadoutAlignedFidelity, + ReadoutAlignedKnobAssignment, + readout_aligned, + readout_variance_captured, +) from polygram.geometry.objectives import ( behavioural_objective, pearson_objective, @@ -44,6 +50,7 @@ # Register built-ins at import time. register_profile(clustered()) register_profile(uniform_sphere()) +register_profile(readout_aligned()) __all__ = [ "ClusteredKnobAssignment", @@ -54,6 +61,8 @@ "LearnedAxisObjective", "LearnedKnobAssignment", "RankRecallAtKFidelity", + "ReadoutAlignedFidelity", + "ReadoutAlignedKnobAssignment", "TierPreservationFidelity", "UniformSphereKnobAssignment", "available_profiles", @@ -61,6 +70,8 @@ "clustered", "get_profile", "pearson_objective", + "readout_aligned", + "readout_variance_captured", "register_profile", "spearman_objective", "uniform_sphere", diff --git a/polygram/geometry/readout_aligned.py b/polygram/geometry/readout_aligned.py new file mode 100644 index 0000000..037bf83 --- /dev/null +++ b/polygram/geometry/readout_aligned.py @@ -0,0 +1,178 @@ +"""`readout-aligned` profile — build the dictionary geometry on the basis the +model *reads through*, not on raw decoder vectors. + +Before the existing `clustered` k-means + β/γ assignment, project the decoder +vectors onto the **readout subspace**: the top-`r` right singular directions of +`gain ⊙ U` (the host unembed, optionally weighted by the final-norm gain). The +powered cross-model R2 result (fieldrun `tau_star_powered.py`) shows this +subspace — not the energy-weighted decoder spectrum — is what governs the +model's argmax (GPT-2 +52pp / Pythia-70m +31pp / Pythia-160m +40pp open-class +R@32 vs the frozen SVD lens). So fitting the geometry here, rather than on raw +decoder norm, tests whether Polygram's geometry-tracks-behaviour Spearman was +partly measuring the (wrong) basis. + +Scope: the projection is a pre-processing of the per-feature projections; the +clustering, β spread, per-cluster-PCA γ, and the fidelity's cosine-overlap target +all run in the projected space, so the resulting knobs (and therefore the Gram / +Q-Orca emission, which consume the knobs) reflect readout geometry. Nothing about +the encoding structure or the quantum machine changes. `u_matrix`/`gain` are +supplied by the caller (Polygram stays numpy-light — no host load in the core). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import numpy as np + +from polygram.geometry.clustered import ClusteredKnobAssignment, TierPreservationFidelity +from polygram.geometry.profile import GeometricProfile +from polygram.geometry.protocols import KnobAssignmentResult + +if TYPE_CHECKING: + from polygram.dictionary import Dictionary + +_DEFAULT_READOUT_RANK = 64 + + +def _resolve_gain(gain: "np.ndarray | float | None", d_model: int) -> np.ndarray: + """Final-norm gain → a `(d_model,)` vector. None → ones; scalar → broadcast.""" + if gain is None: + return np.ones(d_model, dtype=float) + g = np.asarray(gain, dtype=float) + if g.ndim == 0: + return np.full(d_model, float(g)) + if g.shape != (d_model,): + raise ValueError( + f"readout-aligned gain must be a scalar or a ({d_model},) vector " + f"matching d_model; got shape {g.shape}" + ) + return g + + +def readout_subspace( + u_matrix: np.ndarray, gain: "np.ndarray | float | None", d_model: int, rank: int +) -> np.ndarray: + """Return `Vt` `(r, d_model)` — the top-`r` right singular directions of + `gain ⊙ U`. `r = min(rank, d_model, vocab)` (so `rank > rank(gain⊙U)` is safe). + `R = Vt.T`; projecting a row vector is `v @ Vt.T`.""" + U = np.asarray(u_matrix, dtype=float) + if U.ndim != 2 or U.shape[1] != d_model: + raise ValueError( + f"readout-aligned u_matrix must be (vocab, d_model) with " + f"d_model={d_model}; got {U.shape}" + ) + weighted = U * _resolve_gain(gain, d_model)[None, :] # (vocab, d_model) + r = int(min(max(1, rank), d_model, U.shape[0])) + vt = np.linalg.svd(weighted, full_matrices=False)[2] # (min(vocab,d_model), d_model) + return vt[:r] + + +def _project(projections: np.ndarray, vt: np.ndarray) -> np.ndarray: + """`(n_features, d_model) @ (r, d_model)ᵀ -> (n_features, r)`.""" + return np.asarray(projections, dtype=float) @ vt.T + + +def readout_variance_captured( + projections: np.ndarray, + u_matrix: np.ndarray, + gain: "np.ndarray | float | None" = None, + rank: int = _DEFAULT_READOUT_RANK, +) -> float: + """Diagnostic: fraction of decoder-vector Frobenius energy captured by the + top-`r` readout subspace = ‖P·proj‖²_F / ‖proj‖²_F ∈ [0, 1]. Near 1 ⇒ the SAE + features already live in the readout subspace (little to gain); low ⇒ raw and + readout geometries diverge sharply.""" + proj = np.asarray(projections, dtype=float) + total = float(np.sum(proj ** 2)) + if total < 1e-12: + return 1.0 + vt = readout_subspace(u_matrix, gain, proj.shape[1], rank) + captured = float(np.sum(_project(proj, vt) ** 2)) + return float(np.clip(captured / total, 0.0, 1.0)) + + +@dataclass(frozen=True) +class ReadoutAlignedKnobAssignment: + """k-means + β/γ assignment (the `clustered` logic) run in the readout + subspace. `u_matrix`/`gain` are injected by `from_sae_lens` at import time; + a `None` `u_matrix` at assign time is a usage error.""" + + beta_range: tuple[float, float] = (-0.5, 0.5) + u_matrix: np.ndarray | None = None + gain: "np.ndarray | float | None" = None + readout_rank: int = _DEFAULT_READOUT_RANK + + def assign( + self, + projections: np.ndarray, + feature_names: list[str], + *, + n_clusters: int | None, + gamma_range: tuple[float, float], + assign_gamma: bool, + seed: int, + assign_amp_knobs: bool = False, + assign_phase_knobs: bool = False, + encoding: object = None, + ) -> KnobAssignmentResult: + if self.u_matrix is None: + raise ValueError( + "readout-aligned profile: u_matrix is unset. Pass " + "from_sae_lens(..., profile='readout-aligned', u_matrix=U) so " + "the readout subspace can be built." + ) + vt = readout_subspace( + self.u_matrix, self.gain, np.asarray(projections).shape[1], self.readout_rank + ) + proj_readout = _project(projections, vt) + # Delegate to the calibrated clustered logic — in the projected space. + return ClusteredKnobAssignment(self.beta_range).assign( + proj_readout, + feature_names, + n_clusters=n_clusters, + gamma_range=gamma_range, + assign_gamma=assign_gamma, + seed=seed, + assign_amp_knobs=assign_amp_knobs, + assign_phase_knobs=assign_phase_knobs, + encoding=encoding, + ) + + +@dataclass(frozen=True) +class ReadoutAlignedFidelity: + """`tier_preservation` Pearson fidelity computed on the readout-projected + cosine-overlap target (consistent with the readout-space knobs).""" + + u_matrix: np.ndarray | None = None + gain: "np.ndarray | float | None" = None + readout_rank: int = _DEFAULT_READOUT_RANK + + def compute( + self, projections: np.ndarray, dictionary: "Dictionary" + ) -> float | None: + if self.u_matrix is None: + return None + vt = readout_subspace( + self.u_matrix, self.gain, np.asarray(projections).shape[1], self.readout_rank + ) + proj_readout = _project(projections, vt) + return TierPreservationFidelity().compute(proj_readout, dictionary) + + +def readout_aligned() -> GeometricProfile: + """Built-in profile: cluster + assign knobs in the model's readout subspace. + + Registered with placeholder (`u_matrix=None`) strategies; `from_sae_lens` + injects the caller-supplied `u_matrix`/`gain`/`readout_rank` via + `dataclasses.replace` and raises if `u_matrix` is missing. + """ + return GeometricProfile( + name="readout-aligned", + knob_assignment=ReadoutAlignedKnobAssignment(), + geometric_fidelity=ReadoutAlignedFidelity(), + default_n_clusters=2, + default_gamma_range=(-0.25, 0.25), + ) diff --git a/polygram/sae_import.py b/polygram/sae_import.py index 512dd41..59e44eb 100644 --- a/polygram/sae_import.py +++ b/polygram/sae_import.py @@ -625,6 +625,9 @@ def from_sae_lens( gamma_range: tuple[float, float] | None = None, config: "SAEImportConfig | None" = None, profile: "str | GeometricProfile | None" = None, + u_matrix: "np.ndarray | None" = None, + gain: "np.ndarray | float | None" = None, + readout_rank: int = 64, clustered: bool | None = None, block_formation: "BlockFormation | None" = None, assign_amp_knobs: bool | None = None, @@ -682,6 +685,37 @@ def from_sae_lens( cfg = config if config is not None else SAEImportConfig() resolved_profile = _resolve_profile(profile, cfg) + # readout-aligned profile (add-readout-aligned-geometry-profile): inject the + # caller-supplied host unembed into the (registry placeholder) strategy + + # fidelity. Polygram does not load a host itself — `u_matrix` is required. + if resolved_profile.name == "readout-aligned": + import dataclasses + + if u_matrix is None: + raise ValueError( + "from_sae_lens(profile='readout-aligned'): u_matrix is required " + "(the host unembed, shape (vocab, d_model)). Supply u_matrix=... " + "and optionally gain=... (final-norm gain); Polygram does not load " + "a host model itself." + ) + u_arr = np.asarray(u_matrix, dtype=float) + if u_arr.ndim != 2: + raise ValueError( + f"from_sae_lens: u_matrix must be 2-D (vocab, d_model); got " + f"{u_arr.ndim}-D shape {u_arr.shape}" + ) + resolved_profile = dataclasses.replace( + resolved_profile, + knob_assignment=dataclasses.replace( + resolved_profile.knob_assignment, + u_matrix=u_arr, gain=gain, readout_rank=int(readout_rank), + ), + geometric_fidelity=dataclasses.replace( + resolved_profile.geometric_fidelity, + u_matrix=u_arr, gain=gain, readout_rank=int(readout_rank), + ), + ) + # `encoding_partition` (add-encoding-partition Phase 1) — accepted # to lock the API surface that sae-forge's `add-block-structured-sae` # consumes. Phase 2 (the actual per-block dispatch) lives in diff --git a/tests/test_readout_aligned_profile.py b/tests/test_readout_aligned_profile.py new file mode 100644 index 0000000..05409a6 --- /dev/null +++ b/tests/test_readout_aligned_profile.py @@ -0,0 +1,117 @@ +"""Tests for the `readout-aligned` geometric profile (add-readout-aligned-geometry-profile). + +Builds the dictionary geometry in the model's readout subspace (top-`r` SVD of `gain ⊙ U`) instead of on raw +decoder vectors. Synthetic fixtures only (no host download): a random unembed `U` whose readout subspace +re-orders the decoder directions, so the readout-aligned clustering provably differs from the raw one. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from polygram import from_sae_lens +from polygram.geometry import available_profiles, readout_variance_captured +from polygram.geometry.readout_aligned import ReadoutAlignedKnobAssignment, readout_subspace +from polygram.sae_import import SAEFeatureRecord + + +def _records(seed: int = 0, n: int = 24, d_model: int = 32): + rng = np.random.default_rng(seed) + proj = rng.standard_normal((n, d_model)) + proj /= np.linalg.norm(proj, axis=1, keepdims=True) + recs = { + i: SAEFeatureRecord(feature_id=i, name=f"feat_{i}", projection=proj[i].astype(float)) + for i in range(n) + } + # Select a small subset so from_sae_lens returns a flat Dictionary (with `.features`), + # not an auto-promoted ClusteredDictionary (mirrors test_uniform_sphere_profile). + return recs, list(range(8)), proj + + +def _unembed(seed: int = 1, vocab: int = 200, d_model: int = 32): + return np.random.default_rng(seed).standard_normal((vocab, d_model)) + + +def test_readout_aligned_is_registered(): + assert "readout-aligned" in available_profiles() + + +def test_from_sae_lens_requires_u_matrix(): + records, ids, _ = _records() + with pytest.raises(ValueError, match="u_matrix is required"): + from_sae_lens(records, ids, profile="readout-aligned") + + +def test_from_sae_lens_rejects_non_2d_u_matrix(): + records, ids, _ = _records() + with pytest.raises(ValueError, match="2-D"): + from_sae_lens(records, ids, profile="readout-aligned", u_matrix=np.ones(32)) + + +def test_readout_aligned_builds_and_reports_profile(): + records, ids, _ = _records() + U = _unembed(d_model=32) + d, rep = from_sae_lens(records, ids, profile="readout-aligned", u_matrix=U) + assert rep.profile == "readout-aligned" + assert len(d.features) == len(ids) + # fidelity (if computed) is a valid correlation or NaN, never out of range + assert rep.geometric_fidelity is None or np.isnan(rep.geometric_fidelity) \ + or -1.0 <= rep.geometric_fidelity <= 1.0 + + +def test_readout_aligned_differs_from_clustered(): + """The whole point: building geometry in the readout subspace changes the knobs vs the raw basis.""" + records, ids, _ = _records() + U = _unembed(d_model=32) + d_ro, _ = from_sae_lens(records, ids, profile="readout-aligned", u_matrix=U, readout_rank=4) + d_cl, _ = from_sae_lens(records, ids, profile="clustered") + betas_ro = [f.beta for f in d_ro.features] + betas_cl = [f.beta for f in d_cl.features] + assert betas_ro != betas_cl + + +def test_default_path_needs_no_u_matrix_and_stays_clustered(): + records, ids, _ = _records() + _, rep = from_sae_lens(records, ids) + assert rep.profile == "clustered" + + +def test_readout_subspace_edge_cases(): + U = _unembed(vocab=50, d_model=16) + # rank capped at min(rank, d_model, vocab) + assert readout_subspace(U, None, 16, 999).shape == (16, 16) + # gain scalar and vector both accepted, shape (r, d_model) + assert readout_subspace(U, 2.0, 16, 4).shape == (4, 16) + assert readout_subspace(U, np.ones(16), 16, 4).shape == (4, 16) + # wrong d_model raises + with pytest.raises(ValueError, match="vocab, d_model"): + readout_subspace(U, None, 8, 4) + + +def test_variance_captured_in_unit_range_and_monotone_in_rank(): + _, _, proj = _records(d_model=32) + U = _unembed(d_model=32) + vc_small = readout_variance_captured(proj, U, rank=2) + vc_big = readout_variance_captured(proj, U, rank=16) + assert 0.0 <= vc_small <= vc_big <= 1.0 # more rank captures at least as much + + +def test_assign_is_deterministic(): + _, _, proj = _records(d_model=32) + U = _unembed(d_model=32) + ka = ReadoutAlignedKnobAssignment(u_matrix=U, readout_rank=4) + kw = dict(n_clusters=2, gamma_range=(-0.25, 0.25), assign_gamma=True, seed=0) + a = ka.assign(proj, [f"f{i}" for i in range(len(proj))], **kw) + b = ka.assign(proj, [f"f{i}" for i in range(len(proj))], **kw) + assert a.cluster_per_feature == b.cluster_per_feature + assert a.betas == b.betas and a.gammas == b.gammas + + +def test_assign_without_u_matrix_raises(): + _, _, proj = _records(d_model=32) + with pytest.raises(ValueError, match="u_matrix is unset"): + ReadoutAlignedKnobAssignment().assign( + proj, [f"f{i}" for i in range(len(proj))], + n_clusters=2, gamma_range=(-0.25, 0.25), assign_gamma=False, seed=0, + ) From 9625a641ab143492e4f85de4ff9292ff7bfdcd94 Mon Sep 17 00:00:00 2001 From: Allan Scott Date: Sat, 13 Jun 2026 19:35:46 -0400 Subject: [PATCH 3/3] =?UTF-8?q?gate(readout-aligned):=20behavioural=20head?= =?UTF-8?q?-to-head=20=E2=80=94=20NO=20IMPROVEMENT;=20readout=20basis=20is?= =?UTF-8?q?=20wrong=20for=20co-firing?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ran the head-to-head behavioural gate (task 4). Extended behavioural_gram_scaleup.py with --profile {clustered,readout-aligned} --readout-rank R (+ HF-cache SAE fallback, --sae-path), supplying GPT-2's tied unembed + ln_f gain as u_matrix/gain. RESULT (blocks.10, Spearman(Polygram, co-firing Jaccard)): clustered 0.640; readout r=64 0.267; r=256 0.067; r=768(full) 0.640. Readout-alignment STRICTLY HURTS — full rank just recovers the baseline (a full-rank readout projection is an orthonormal rotation and k-means is rotation-invariant), every truncation degrades. Mechanism: co-firing is ENCODER-side (which features fire together, full-residual geometry); the readout subspace is DECODE-side (argmax). Readout-alignment is right for the decode tax (R2 +52/+31/+40pp) but WRONG for co-firing -> this VINDICATES the raw-decoder geometry here (0.640 is not a basis artifact; Reckoning #3's basis-limited concern doesn't apply). The profile ships as a correct tool for decode-relevant geometry; a decode-side behavioural metric is where it would help. Writeup in proposal "Gate RESULT" + behavioural-scaleup-probe.md. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/research/behavioural-scaleup-probe.md | 21 +++++++++ examples/behavioural_gram_scaleup.py | 47 ++++++++++++++++++- .../proposal.md | 30 ++++++++++++ 3 files changed, 96 insertions(+), 2 deletions(-) diff --git a/docs/research/behavioural-scaleup-probe.md b/docs/research/behavioural-scaleup-probe.md index db4325c..f5855ee 100644 --- a/docs/research/behavioural-scaleup-probe.md +++ b/docs/research/behavioural-scaleup-probe.md @@ -5,6 +5,27 @@ > [`tech-debt-backlog` §4.4](../../openspec/changes/tech-debt-backlog/tasks.md). > Reproducible via `python examples/behavioural_gram_scaleup.py`. +## Readout-aligned head-to-head (2026-06-13, `add-readout-aligned-geometry-profile`) + +Does building the geometry in GPT-2's **readout subspace** (top-`r` SVD of `gain⊙U`) instead of on raw decoder +vectors raise the Spearman? **No — it strictly hurts.** Head-to-head on `blocks.10` +(`--profile {clustered,readout-aligned} --readout-rank R`): + +| profile / rank | Spearman(Polygram, Jaccard) | +|---|---| +| `clustered` (raw decoder) | **0.640** | +| `readout-aligned` r=64 | 0.267 | +| `readout-aligned` r=256 | 0.067 | +| `readout-aligned` r=768 (full) | **0.640** (= baseline) | + +Full-rank readout-projection is an orthonormal rotation (k-means is rotation-invariant) so it reproduces +`clustered` exactly; every truncation degrades. **Mechanism:** co-firing Jaccard is an **encoder-side** +phenomenon (which features fire together, in the full-residual geometry); the readout subspace is **decode-side** +(what the argmax reads through). Readout-alignment is the right basis for the *decode* tax (R2: +52/+31/+40pp +open-class R@32) but the **wrong** basis for *co-firing*. So this **vindicates** the raw-decoder geometry for +this claim — 0.640 is not a basis artifact. The `readout-aligned` profile ships as a tool for decode-relevant +geometry; a decode-side behavioural metric (e.g. logit-attribution overlap) is where it would be expected to help. + ## Context PR #20 ([§4.2 / `behavioural-gram-probe.md`](behavioural-gram-probe.md)) diff --git a/examples/behavioural_gram_scaleup.py b/examples/behavioural_gram_scaleup.py index a3cc73a..09f4177 100644 --- a/examples/behavioural_gram_scaleup.py +++ b/examples/behavioural_gram_scaleup.py @@ -354,6 +354,8 @@ def _run_probe( seed_candidates: list[int], min_firing_rate: float, progress: bool, + profile: str = "clustered", + readout_rank: int = 64, ) -> dict: deps = _import_torch_and_transformers() if deps is None: @@ -430,8 +432,17 @@ def _capture_hook(module, args): if progress: print("Building Polygram Dictionary via from_sae_lens...") records = load_sae_safetensors(sae_path, feature_ids=selected_ids) + fsl_kwargs: dict = {} + if profile == "readout-aligned": + # Readout subspace from GPT-2's (tied) unembed + final-norm gain. + fsl_kwargs = { + "profile": "readout-aligned", + "u_matrix": model.lm_head.weight.detach().cpu().numpy().astype(np.float64), + "gain": model.transformer.ln_f.weight.detach().cpu().numpy().astype(np.float64), + "readout_rank": readout_rank, + } d_mps, report = from_sae_lens( - records, selected_ids, assign_gamma=True, name="ScaleupBlocks10", + records, selected_ids, assign_gamma=True, name="ScaleupBlocks10", **fsl_kwargs, ) g_polygram = np.abs(d_mps.gram()) ** 2 @@ -743,6 +754,26 @@ def main(argv: list[str] | None = None) -> None: parser.add_argument( "--quiet", action="store_true", help="suppress progress prints" ) + parser.add_argument( + "--profile", + default="clustered", + help=( + "geometric profile for the Polygram Dictionary. 'readout-aligned' " + "builds the geometry in GPT-2's readout subspace (unembed + final-" + "norm gain) — the head-to-head test of whether the basis is what " + "limits the Polygram→behaviour Spearman." + ), + ) + parser.add_argument( + "--readout-rank", type=int, default=64, + help="rank of the readout subspace for profile=readout-aligned.", + ) + parser.add_argument( + "--sae-path", + type=Path, + default=None, + help="override the SAE checkpoint path (defaults to ./scratch then the HF cache).", + ) args = parser.parse_args(argv) if not (2 <= args.n_features <= MAX_FEATURES): @@ -751,9 +782,19 @@ def main(argv: list[str] | None = None) -> None: f"got {args.n_features}" ) - sae_path = Path( + sae_path = args.sae_path or Path( "./scratch/real-sae/blocks.10.hook_resid_pre/sae_weights.safetensors" ) + if not sae_path.exists(): + import glob + import os + + hits = glob.glob(os.path.expanduser( + "~/.cache/huggingface/hub/models--jbloom--GPT2-Small-SAEs-Reformatted/" + "snapshots/*/blocks.10.hook_resid_pre/sae_weights.safetensors" + )) + if hits: + sae_path = Path(hits[0]) if not sae_path.exists(): print( f"behavioural_gram_scaleup: SAE checkpoint not found at " @@ -773,6 +814,8 @@ def main(argv: list[str] | None = None) -> None: seed_candidates=list(args.seed_candidates), min_firing_rate=args.min_firing_rate, progress=not args.quiet, + profile=args.profile, + readout_rank=args.readout_rank, ) if "skipped" in report: _print_report(report, {}) diff --git a/openspec/changes/add-readout-aligned-geometry-profile/proposal.md b/openspec/changes/add-readout-aligned-geometry-profile/proposal.md index 4acaff4..b03b6ac 100644 --- a/openspec/changes/add-readout-aligned-geometry-profile/proposal.md +++ b/openspec/changes/add-readout-aligned-geometry-profile/proposal.md @@ -87,6 +87,36 @@ panels: Either way the verdict is the head-to-head table, descriptive; no necessity claims. +### Gate RESULT (2026-06-13) — NO IMPROVEMENT; the readout basis is *wrong for co-firing* + +Implemented + run head-to-head on GPT-2-small `blocks.10` (`examples/behavioural_gram_scaleup.py +--profile {clustered,readout-aligned} --readout-rank R`): + +| profile / rank | `Spearman(Polygram-overlap, co-firing Jaccard)` | +|---|---| +| `clustered` (raw decoder, baseline) | **0.640** | +| `readout-aligned`, rank 64 (default) | 0.267 | +| `readout-aligned`, rank 256 | 0.067 | +| `readout-aligned`, rank 768 (full `d_model`) | **0.640** (exactly the baseline) | + +**Readout-alignment does NOT raise the Spearman — it strictly *hurts*, recovering the baseline only at full +rank.** The full-rank result is the key tell: a full-rank readout projection is just an orthonormal rotation +(plus gain weighting), and k-means is rotation-invariant — so it reproduces `clustered` exactly (0.640). Every +*truncated* rank discards residual dimensions and degrades the co-firing signal. + +**The mechanism — and the real lesson — is a basis/metric mismatch:** Polygram's behavioural metric is +**co-firing** (which features *activate together* — an **encoder-side** phenomenon living in the full-residual +feature geometry), whereas R2's readout subspace is **decode-side** (the directions the model's *argmax* reads +through). Readout-alignment is the right basis for the *decode* tax (R2's +52/+31/+40pp) but the **wrong** basis +for *co-firing*. So this is the **NO-IMPROVEMENT** branch — and it **vindicates** Polygram's raw-decoder +geometry for the co-firing claim: the 0.640 is *not* a basis artifact, and Reckoning #3's basis-limited concern +does not apply here. + +**What ships:** the `readout-aligned` profile is implemented, tested, and correct — it is a legitimate tool for +*decode-relevant* geometry questions; it is simply **not** the right basis for the co-firing Gram, which this +gate establishes cleanly. (A future probe with a *decode-side* behavioural metric — e.g. logit-attribution +overlap — is where readout-alignment would be expected to help.) + ## Related - `docs/research/behavioural-scaleup-probe.md` — the 0.637 baseline this targets.