From 8a4d8fb3d62f8d6298f5a25d0f704b35995d8302 Mon Sep 17 00:00:00 2001 From: aidxhxr Date: Mon, 9 Nov 2026 14:56:11 -0500 Subject: [PATCH 1/2] feat(jacobian_lens): import J++ artifact provenance and the parameters.jacobians schema JacobianLens.load() dropped the top-level `provenance` dict that Neuronpedia's converted J++ lenses carry, so the released Qwen3.5-0.8B artifact loaded with empty metadata, and the J++ Lens package's own `parameters.jacobians` + `config` schema was rejected. - Keep recorded provenance verbatim under metadata["provenance"] and normalise model name/revision, method, lrp_mode, target layer, relative end-transport layer, capture point and stream reduction onto the metadata keys validate_model() reads. Explicit `metadata` keys win; no revision is invented when none was recorded; unreloadable leaves are dropped and named in dropped_fields. - Accept the J++ package schema (parameters.jacobians, source_layers, config) and the {"J", "config"} variant; read d_model and num_prompts_trained_on from the config; mark the lens converted_from="jpp_lens_artifact". Only lens_type="jacobian" loads. - validate_model() resolves relative_end_transport_layer against n_layers and checks recorded capture_point="block_output" and stream_reduce="none". - merge() refuses imported J++ maps (method="jpp" or the sentinel): they are expert-combined, not prompt means. - Tests with synthetic payloads for both schemas, round trips, merge refusal and the new validate_model checks; docs section on published J++ lenses. Part 1 of #1991. --- docs/source/content/jacobian_lens_fitting.md | 54 +++- tests/unit/tools/test_jacobian_lens.py | 23 ++ tests/unit/tools/test_jacobian_lens_import.py | 265 ++++++++++++++++ .../tools/analysis/jacobian_lens.py | 287 +++++++++++++++++- 4 files changed, 610 insertions(+), 19 deletions(-) diff --git a/docs/source/content/jacobian_lens_fitting.md b/docs/source/content/jacobian_lens_fitting.md index 3f617ca0ff..d44e4bc733 100644 --- a/docs/source/content/jacobian_lens_fitting.md +++ b/docs/source/content/jacobian_lens_fitting.md @@ -521,10 +521,58 @@ The full-vocabulary dictionary is cached on the model's device and is vocabulary ## Importing an existing lens -`JacobianLens.load()` accepts two file schemas: the standard artifact format (four +`JacobianLens.load()` accepts three file schemas: the standard artifact format (four official keys — `J`, `n_prompts`, `source_layers`, `d_model`) written by -`JacobianLens.save()` and compatible with the Anthropic reference package, and the -fit-checkpoint format described below. +`JacobianLens.save()` and compatible with the Anthropic reference package, the J++ +Lens package's own format, and the fit-checkpoint format described below. + +### Published J++ lenses and artifact provenance + +The J++ lenses on `neuronpedia/jacobian-lens` (`/jpp/...`) use the standard +artifact keys plus a top-level `provenance` dict that records the source, the method, +the fitting configuration and the activation convention the maps read. `load()` keeps +that dict verbatim under `metadata["provenance"]` and normalises the recorded values +onto the keys the rest of the class already reads: + +| Normalised key | Read from | Checked by `validate_model()` | +|---|---|---| +| `model_name` | `provenance.model_name`, `provenance.hf_model_name`, or `source_config.hf_model_name` | yes, against `model.cfg.model_name` | +| `model_revision` | `provenance.model_revision` / `hf_model_revision`, only when recorded | yes, against the booted HF commit | +| `method`, `lrp_mode` | `provenance` or `source_config` | no (informational; `method == "jpp"` blocks `merge()`) | +| `target_layer` | `provenance.target_layer` (absolute block index, as in the released R-Lens files) | yes, must be the final block | +| `relative_end_transport_layer` | `source_config.relative_end_transport_layer` (J++ package convention) | yes, resolved against `n_layers` when no `target_layer` is recorded | +| `capture_point` | `provenance.capture_point` | yes, must be `"block_output"` | +| `stream_reduce` | `provenance.stream_reduce` | yes, must be `"none"` (single-stream contract) | + +Nothing is invented: a file without a recorded revision loads without `model_revision`, +and an explicit `metadata` dict in the same file wins over the provenance. Leaves that +`weights_only=True` could not reload (tensors, non-string keys) are dropped from the +copied provenance and named in `dropped_fields`. + +```python +lens = JacobianLens.from_pretrained( + "neuronpedia/jacobian-lens", + filename="qwen3.5-0.8b/jpp/Salesforce-wikitext/Qwen3.5-0.8B_jpp_lens.pt", + revision="e166deb5ad540932100ef803c8c9bcc2a5453694", + model=model, +) +lens.metadata["method"] # "jpp" +lens.metadata["model_name"] # "Qwen/Qwen3.5-0.8B" +lens.metadata["provenance"]["source_config"]["lrp_mode"] # "rlens" +``` + +The authors' original files (`koayon/jpp-lenses`, written by the +[`jpp_lens`](https://github.com/safety-research/jpp_lens) package's `BaseLens.save`) +store the maps under `parameters.jacobians` with a `config` dict instead of +`n_prompts`/`d_model`. `load()` reads `d_model` and `num_prompts_trained_on` from the +config, keeps it under `metadata["provenance"]["source_config"]`, normalises the same +keys as above, and sets `converted_from: "jpp_lens_artifact"`. Only +`lens_type == "jacobian"` files load; a logit-lens file raises. + +J++ maps combine clustered expert Jacobians with learned signed weights, so they are +not prompt means and carry no additive statistics. `merge()` therefore refuses any lens +whose metadata records `method == "jpp"` or `converted_from == "jpp_lens_artifact"`. +Readout and interventions work unchanged. ### Fit checkpoint format diff --git a/tests/unit/tools/test_jacobian_lens.py b/tests/unit/tools/test_jacobian_lens.py index dd382d24ce..42eef85a26 100644 --- a/tests/unit/tools/test_jacobian_lens.py +++ b/tests/unit/tools/test_jacobian_lens.py @@ -775,6 +775,29 @@ def test_validate_model_rejects_nonfinal_target_metadata(toy_model: _ToyBridge) lens.validate_model(toy_model) +def test_validate_model_resolves_relative_end_transport_layer(toy_model: _ToyBridge) -> None: + with pytest.raises(ValueError, match=r"final.?layer"): + _lens(metadata={"relative_end_transport_layer": -2}).validate_model(toy_model) + final = _lens(metadata={"relative_end_transport_layer": -1}) + assert final.validate_model(toy_model) is final + + +def test_validate_model_checks_recorded_capture_convention(toy_model: _ToyBridge) -> None: + with pytest.raises(ValueError, match="capture_point"): + _lens(metadata={"capture_point": "attention_output"}).validate_model(toy_model) + with pytest.raises(ValueError, match="stream_reduce"): + _lens(metadata={"stream_reduce": "mean"}).validate_model(toy_model) + matching = _lens( + metadata={ + "model_name": "toy-bridge", + "capture_point": "block_output", + "stream_reduce": "none", + "relative_end_transport_layer": -1, + } + ) + assert matching.validate_model(toy_model) is matching + + def test_readout_bounds_checks(toy_model: _ToyBridge, fitted_lens: JacobianLens) -> None: with pytest.raises(ValueError, match="out of range"): fitted_lens.readout(toy_model, "a toy prompt", layers=[99], use_jacobian=False) diff --git a/tests/unit/tools/test_jacobian_lens_import.py b/tests/unit/tools/test_jacobian_lens_import.py index 78825a4ef8..b07f7f81ff 100644 --- a/tests/unit/tools/test_jacobian_lens_import.py +++ b/tests/unit/tools/test_jacobian_lens_import.py @@ -323,3 +323,268 @@ def test_merge_weighted_average_correct(self, tmp_path): torch.full((d_model, d_model), expected_mean), atol=1e-5, ) + + +# --------------------------------------------------------------------------- +# J++ artifacts: converted (J + provenance) and native (parameters.jacobians) +# --------------------------------------------------------------------------- + + +def _author_config( + *, + d_model=8, + n_layers=3, + n_prompts=10, + lrp_mode="rlens", + lens_type="jacobian", + hf_model_name="Qwen/Qwen3.5-0.8B", + relative_end_transport_layer=-1, +): + """A ``LensConfig.to_dict()`` payload as written by the J++ Lens package.""" + return { + "hf_model_name": hf_model_name, + "checkpoint_name": "recipe_lens", + "lens_type": lens_type, + "checkpoint_every_n_prompts": 4, + "creation_time": "2026-10-07_00-51-51", + "artifacts_base_dir": "artifacts", + "checkpoint_path": "artifacts/2026-10-07/recipe_lens", + "source_layers": list(range(n_layers)), + "relative_end_transport_layer": relative_end_transport_layer, + "d_model": d_model, + "num_prompts_trained_on": n_prompts, + "next_prompt_idx": 0, + "jacobian_rows_per_pass": 16, + "max_seq_len": 128, + "skip_first_n_positions": 16, + "num_clusters": 0, + "cluster_projection_dim": 0, + "lrp_mode": lrp_mode, + } + + +def _converted_provenance(config, **overrides): + """The ``provenance`` dict Neuronpedia's convert-jpp-lens.py writes.""" + provenance = { + "source": "neuronpedia fit, koayon/jpp_lens recipe, 96 WikiText prompts", + "method": "jpp", + "lrp_mode": config["lrp_mode"], + "source_config": config, + "capture_point": "block_output", + "stream_reduce": "none", + "stream_index": None, + } + provenance.update(overrides) + return provenance + + +def _jacobians(n_layers=3, d_model=8): + return {i: torch.randn(d_model, d_model) for i in range(n_layers)} + + +def _save_converted_artifact(path, jacobians, provenance, *, n_prompts=10, metadata=None): + d_model = next(iter(jacobians.values())).shape[0] + payload = { + "J": {layer: matrix.half() for layer, matrix in jacobians.items()}, + "n_prompts": n_prompts, + "source_layers": sorted(jacobians), + "d_model": d_model, + "provenance": provenance, + } + if metadata is not None: + payload["metadata"] = metadata + torch.save(payload, path) + + +def _save_author_artifact(path, jacobians, config): + torch.save( + { + "parameters": { + "jacobians": {layer: matrix.half() for layer, matrix in jacobians.items()} + }, + "source_layers": sorted(jacobians), + "config": config, + }, + path, + ) + + +class TestLoadConvertedProvenance: + def test_provenance_preserved_in_metadata(self, tmp_path): + p = str(tmp_path / "jpp.pt") + provenance = _converted_provenance(_author_config()) + _save_converted_artifact(p, _jacobians(), provenance) + lens = JacobianLens.load(p) + assert lens.metadata["provenance"] == provenance + + def test_recorded_values_normalized_without_inventing_revision(self, tmp_path): + p = str(tmp_path / "jpp.pt") + _save_converted_artifact(p, _jacobians(), _converted_provenance(_author_config())) + lens = JacobianLens.load(p) + assert lens.metadata["model_name"] == "Qwen/Qwen3.5-0.8B" + assert lens.metadata["method"] == "jpp" + assert lens.metadata["lrp_mode"] == "rlens" + assert lens.metadata["capture_point"] == "block_output" + assert lens.metadata["stream_reduce"] == "none" + assert lens.metadata["relative_end_transport_layer"] == -1 + assert "model_revision" not in lens.metadata + assert "target_layer" not in lens.metadata + assert lens.n_prompts == 10 + assert lens.d_model == 8 + + def test_recorded_revision_and_target_layer_harvested(self, tmp_path): + p = str(tmp_path / "rlens.pt") + provenance = _converted_provenance( + _author_config(), target_layer=1, model_revision="2fc0636" + ) + _save_converted_artifact(p, _jacobians(), provenance) + lens = JacobianLens.load(p) + assert lens.metadata["target_layer"] == 1 + assert lens.metadata["model_revision"] == "2fc0636" + + def test_explicit_metadata_wins_over_provenance(self, tmp_path): + p = str(tmp_path / "jpp.pt") + provenance = _converted_provenance(_author_config()) + _save_converted_artifact( + p, _jacobians(), provenance, metadata={"model_name": "explicit", "corpus": "wiki"} + ) + lens = JacobianLens.load(p) + assert lens.metadata["model_name"] == "explicit" + assert lens.metadata["corpus"] == "wiki" + assert lens.metadata["method"] == "jpp" + assert lens.metadata["provenance"] == provenance + + def test_unsafe_provenance_leaves_dropped_and_recorded(self, tmp_path): + p = str(tmp_path / "jpp.pt") + provenance = _converted_provenance(_author_config(), fit_stats=torch.randn(3)) + _save_converted_artifact(p, _jacobians(), provenance) + lens = JacobianLens.load(p) + assert "fit_stats" not in lens.metadata["provenance"] + assert any("fit_stats" in field for field in lens.metadata["dropped_fields"]) + assert lens.metadata["provenance"]["method"] == "jpp" + assert lens.metadata["model_name"] == "Qwen/Qwen3.5-0.8B" + + def test_save_load_roundtrip_preserves_provenance(self, tmp_path): + p = str(tmp_path / "jpp.pt") + q = str(tmp_path / "resaved.pt") + _save_converted_artifact(p, _jacobians(), _converted_provenance(_author_config())) + lens = JacobianLens.load(p) + lens.save(q) + reloaded = JacobianLens.load(q) + assert reloaded.metadata == lens.metadata + assert reloaded.n_prompts == lens.n_prompts + for layer in lens.source_layers: + torch.testing.assert_close(reloaded.jacobians[layer], lens.jacobians[layer]) + + def test_merge_refuses_jpp_maps(self, tmp_path): + p1, p2 = str(tmp_path / "a.pt"), str(tmp_path / "b.pt") + provenance = _converted_provenance(_author_config()) + _save_converted_artifact(p1, _jacobians(), provenance) + _save_converted_artifact(p2, _jacobians(), provenance) + with pytest.raises(ValueError, match="additive"): + JacobianLens.merge([JacobianLens.load(p1), JacobianLens.load(p2)]) + + def test_plain_artifact_without_provenance_unchanged(self, tmp_path): + p = str(tmp_path / "plain.pt") + _save_artifact(p, metadata={"model_name": "gpt2"}) + lens = JacobianLens.load(p) + assert lens.metadata == {"model_name": "gpt2"} + + +class TestLoadAuthorSchema: + def test_loads_same_jacobians_as_native_schema(self, tmp_path): + jacobians = _jacobians() + native, author = str(tmp_path / "native.pt"), str(tmp_path / "author.pt") + torch.save( + { + "J": {layer: matrix.half() for layer, matrix in jacobians.items()}, + "n_prompts": 10, + "source_layers": sorted(jacobians), + "d_model": 8, + }, + native, + ) + _save_author_artifact(author, jacobians, _author_config()) + native_lens = JacobianLens.load(native) + author_lens = JacobianLens.load(author) + assert author_lens.source_layers == native_lens.source_layers + assert author_lens.d_model == native_lens.d_model + assert author_lens.n_prompts == native_lens.n_prompts + for layer in native_lens.source_layers: + torch.testing.assert_close(author_lens.jacobians[layer], native_lens.jacobians[layer]) + assert author_lens.jacobians[layer].dtype == torch.float32 + + def test_config_normalized_into_metadata(self, tmp_path): + p = str(tmp_path / "author.pt") + config = _author_config() + _save_author_artifact(p, _jacobians(), config) + lens = JacobianLens.load(p) + assert lens.metadata["converted_from"] == "jpp_lens_artifact" + assert lens.metadata["provenance"]["source_config"] == config + assert lens.metadata["model_name"] == "Qwen/Qwen3.5-0.8B" + assert lens.metadata["lrp_mode"] == "rlens" + assert lens.metadata["relative_end_transport_layer"] == -1 + assert "model_revision" not in lens.metadata + assert "method" not in lens.metadata + + def test_j_with_config_variant(self, tmp_path): + p = str(tmp_path / "j_config.pt") + jacobians = _jacobians() + torch.save( + { + "J": {layer: matrix.half() for layer, matrix in jacobians.items()}, + "config": _author_config(n_prompts=7), + }, + p, + ) + lens = JacobianLens.load(p) + assert lens.d_model == 8 + assert lens.n_prompts == 7 + assert lens.metadata["model_name"] == "Qwen/Qwen3.5-0.8B" + assert lens.metadata["converted_from"] == "jpp_lens_artifact" + + def test_d_model_inferred_when_config_omits_it(self, tmp_path): + p = str(tmp_path / "author.pt") + config = _author_config() + del config["d_model"] + _save_author_artifact(p, _jacobians(d_model=4), config) + assert JacobianLens.load(p).d_model == 4 + + def test_rejects_non_jacobian_lens_type(self, tmp_path): + p = str(tmp_path / "logit.pt") + torch.save( + { + "parameters": {}, + "source_layers": [0, 1], + "config": _author_config(lens_type="logit"), + }, + p, + ) + with pytest.raises(ValueError, match="lens_type='logit'"): + JacobianLens.load(p) + + def test_rejects_parameters_without_jacobians(self, tmp_path): + p = str(tmp_path / "no_jac.pt") + torch.save( + {"parameters": {"other": {0: 1.0}}, "source_layers": [0], "config": _author_config()}, + p, + ) + with pytest.raises(ValueError, match="parameters.jacobians"): + JacobianLens.load(p) + + def test_save_load_roundtrip(self, tmp_path): + p, q = str(tmp_path / "author.pt"), str(tmp_path / "resaved.pt") + _save_author_artifact(p, _jacobians(), _author_config()) + lens = JacobianLens.load(p) + lens.save(q) + reloaded = JacobianLens.load(q) + assert reloaded.metadata == lens.metadata + for layer in lens.source_layers: + torch.testing.assert_close(reloaded.jacobians[layer], lens.jacobians[layer]) + + def test_merge_refuses_imported_author_artifacts(self, tmp_path): + p1, p2 = str(tmp_path / "a.pt"), str(tmp_path / "b.pt") + _save_author_artifact(p1, _jacobians(), _author_config()) + _save_author_artifact(p2, _jacobians(), _author_config()) + with pytest.raises(ValueError, match="additive"): + JacobianLens.merge([JacobianLens.load(p1), JacobianLens.load(p2)]) diff --git a/transformer_lens/tools/analysis/jacobian_lens.py b/transformer_lens/tools/analysis/jacobian_lens.py index 5d1d21912a..bc4d1bcc49 100644 --- a/transformer_lens/tools/analysis/jacobian_lens.py +++ b/transformer_lens/tools/analysis/jacobian_lens.py @@ -195,6 +195,18 @@ def _registry_revision() -> str: # (rather than nested under a "metadata" key). _CHECKPOINT_FLAT_PROVENANCE: frozenset = frozenset({"model_name", "model_revision", "corpus"}) +# Sentinel recorded on lenses imported from the J++ Lens package's own file +# schema (safety-research/jpp_lens). Published J++ maps combine clustered expert +# Jacobians with learned weights, so they carry no prompt-mean statistics that +# merge() could average; the package's ``lrp_mode``/``num_clusters`` fields cannot +# tell such a map apart from a plain fit, hence the conservative marker. +_JPP_ARTIFACT_SENTINEL = "jpp_lens_artifact" +# Recorded ``method`` values whose maps are not prompt means. +_NON_ADDITIVE_METHODS: frozenset = frozenset({"jpp"}) +# Activation-capture convention this lens reads: single-stream block outputs. +_CAPTURE_POINT_BLOCK_OUTPUT = "block_output" +_STREAM_REDUCE_NONE = "none" + @dataclass class JacobianLensReadout: @@ -269,8 +281,10 @@ class JacobianLens: n_prompts: Number of prompts averaged into the fit. d_model: Residual stream width the lens was fitted for. metadata: Optional provenance (model name, TransformerLens version, fit - hyperparameters). Preserved by :meth:`save`/:meth:`load`; artifacts - from the reference implementation load with empty metadata. + hyperparameters). Preserved by :meth:`save`/:meth:`load`; plain + artifacts from the reference implementation load with empty + metadata, while converted J++/R-Lens files keep their recorded + ``provenance`` plus the normalised keys :meth:`validate_model` reads. """ def __init__( @@ -348,12 +362,30 @@ def save(self, path: str, *, dtype: torch.dtype = torch.float16) -> None: def load(cls: Type[_LensT], path: str) -> _LensT: """Load a lens artifact or fit checkpoint saved in a supported schema. - Two file schemas are accepted: + Three file schemas are accepted: **Artifact** (the reference format, written by :meth:`save` or the Anthropic reference package): must contain a ``J`` key mapping layer indices to transport matrices, plus ``n_prompts``, ``d_model``, and an - optional ``metadata`` dict. + optional ``metadata`` dict. Converted published files (Neuronpedia's + J++ lenses, the released R-Lens files) add a top-level ``provenance`` + dict; it is kept verbatim under ``metadata["provenance"]`` and its + recorded ``hf_model_name``/``model_name``, ``model_revision``, + ``method``, ``lrp_mode``, ``target_layer``, + ``relative_end_transport_layer``, ``capture_point`` and + ``stream_reduce`` are normalised onto the metadata keys + :meth:`validate_model` checks. Explicit ``metadata`` keys take + precedence. No fitting revision is invented when none was recorded. + + **J++ Lens package artifact** (``BaseLens.save`` in + ``safety-research/jpp_lens``): ``parameters.jacobians`` holds the + per-layer matrices, ``config`` the package's ``LensConfig`` + (``hf_model_name``, ``d_model``, ``num_prompts_trained_on``, + ``relative_end_transport_layer``, ``lrp_mode``, ...). The config is kept + under ``metadata["provenance"]["source_config"]`` and normalised the + same way; ``converted_from: "jpp_lens_artifact"`` is added so + :meth:`merge` refuses these maps, which are expert-combined rather than + prompt means. Only ``lens_type == "jacobian"`` files load. **Fit checkpoint** (running-sum format, written by the reference implementation's ``write_checkpoint()`` during fitting): must contain @@ -400,20 +432,100 @@ def load(cls: Type[_LensT], path: str) -> _LensT: """ payload = torch.load(path, map_location="cpu", weights_only=True) if "J" in payload: - return cls( - {int(layer): matrix for layer, matrix in payload["J"].items()}, - n_prompts=int(payload.get("n_prompts", 0)), - d_model=int(payload["d_model"]), - metadata=payload.get("metadata"), - ) + return cls._from_artifact_payload(path, payload) + if "parameters" in payload: + return cls._from_jpp_payload(path, payload) if "jacobian_sum" in payload: return cls._from_checkpoint_payload(path, payload) raise ValueError( f"{path} does not look like a Jacobian lens artifact or fit checkpoint. " - "Expected a 'J' key (artifact) or 'jacobian_sum' key (fit checkpoint). " + "Expected a 'J' key (artifact), a 'parameters' key (J++ Lens package " + "artifact), or a 'jacobian_sum' key (fit checkpoint). " "See JacobianLens.load() for the supported file schemas." ) + @classmethod + def _from_artifact_payload(cls: Type[_LensT], path: str, payload: Dict[str, Any]) -> _LensT: + """Build a lens from a ``J``-keyed artifact, importing any recorded provenance. + + Converted published files (Neuronpedia's ``convert-jpp-lens.py``, the + released R-Lens files) add a top-level ``provenance`` dict next to the + four official keys; the J++ Lens package also writes ``{"J", "config"}`` + for a J-lens with its own ``LensConfig``. Both are kept verbatim under + ``metadata["provenance"]`` and normalised onto the keys + :meth:`validate_model` reads. A plain artifact loads exactly as before. + """ + jacobians = {int(layer): matrix for layer, matrix in payload["J"].items()} + config = payload.get("config") + provenance = payload.get("provenance") + converted_from: Optional[str] = None + if not isinstance(provenance, dict): + provenance = None + if provenance is None and isinstance(config, dict): + provenance = {"source_config": config} + converted_from = _JPP_ARTIFACT_SENTINEL + d_model = payload.get("d_model") + n_prompts = payload.get("n_prompts") + if isinstance(config, dict): + if d_model is None: + d_model = config.get("d_model") + if n_prompts is None: + n_prompts = config.get("num_prompts_trained_on") + if d_model is None: + d_model = _matrix_width(path, jacobians) + explicit = payload.get("metadata") + return cls( + jacobians, + n_prompts=int(n_prompts or 0), + d_model=int(d_model), + metadata=_import_artifact_metadata( + explicit if isinstance(explicit, dict) else None, + provenance, + converted_from=converted_from, + ), + ) + + @classmethod + def _from_jpp_payload(cls: Type[_LensT], path: str, payload: Dict[str, Any]) -> _LensT: + """Build a lens from the J++ Lens package's ``BaseLens.save`` schema. + + The file holds ``parameters.jacobians`` (``{layer: [d, d]}``), the + ``source_layers`` list, and ``config`` (the package's ``LensConfig``, + which carries ``hf_model_name``, ``d_model``, ``num_prompts_trained_on``, + ``relative_end_transport_layer`` and ``lrp_mode``). The config is kept + under ``metadata["provenance"]["source_config"]`` and the lens is marked + ``converted_from: "jpp_lens_artifact"`` so :meth:`merge` refuses it. + """ + config = payload.get("config") + if not isinstance(config, dict): + config = {} + lens_type = config.get("lens_type", "jacobian") + if lens_type != "jacobian": + raise ValueError( + f"{path} holds a lens_type={lens_type!r} lens from the J++ Lens package; " + "JacobianLens loads only lens_type='jacobian' artifacts, which store " + "their maps under parameters.jacobians." + ) + parameters = payload["parameters"] + raw_jacobians = parameters.get("jacobians") if isinstance(parameters, dict) else None + if not isinstance(raw_jacobians, dict): + raise ValueError( + f"{path} has no parameters.jacobians mapping; the J++ Lens package stores " + "one [d_model, d_model] transport matrix per source layer there." + ) + jacobians = {int(layer): matrix for layer, matrix in raw_jacobians.items()} + d_model = config.get("d_model") + if d_model is None: + d_model = _matrix_width(path, jacobians) + return cls( + jacobians, + n_prompts=int(config.get("num_prompts_trained_on") or 0), + d_model=int(d_model), + metadata=_import_artifact_metadata( + None, {"source_config": config}, converted_from=_JPP_ARTIFACT_SENTINEL + ), + ) + @classmethod def _from_checkpoint_payload(cls: Type[_LensT], path: str, payload: Dict[str, Any]) -> _LensT: """Reconstruct a JacobianLens from a fit-checkpoint payload. @@ -575,15 +687,17 @@ def merge(cls, lenses: Sequence["JacobianLens"]) -> "JacobianLens": Provenance must match across shards (apart from ``n_prompts``), so a merge cannot silently relabel matrices fitted with different models, corpora, dtypes, or recorded estimator settings. The merged count - replaces the per-shard count. + replaces the per-shard count. Lenses imported from J++ artifacts + (``method == "jpp"`` or ``converted_from == "jpp_lens_artifact"``) are + refused: their maps are expert-combined, not prompt means. Args: lenses: Lenses that agree exactly on ``source_layers`` and ``d_model``. Raises: - ValueError: On an empty sequence, mismatched lenses, or shards from - different estimators. + ValueError: On an empty sequence, mismatched lenses, shards from + different estimators, or imported J++ maps. """ if not lenses: raise ValueError("cannot merge an empty sequence of lenses") @@ -596,8 +710,9 @@ def merge(cls, lenses: Sequence["JacobianLens"]) -> "JacobianLens": f"invalid shards: {invalid_counts}" ) first = lenses[0] - for lens in lenses: + for index, lens in enumerate(lenses): _validate_metadata(lens.metadata) + _require_additive_statistics(lens, index) first_provenance = { key: value for key, value in first.metadata.items() if key != "n_prompts" } @@ -680,13 +795,32 @@ def validate_model(self, model: Any) -> "JacobianLens": f"lens has source layers {out_of_range} outside the model's " f"0..{final_layer - 1} source range — this lens belongs to a different model." ) - target_layer = int(self.metadata.get("target_layer", final_layer)) + recorded_target = self.metadata.get("target_layer") + if recorded_target is None: + # The J++ Lens package records the target relative to the end of the model. + relative_end = self.metadata.get("relative_end_transport_layer") + recorded_target = final_layer if relative_end is None else n_layers + int(relative_end) + target_layer = int(recorded_target) if target_layer != final_layer: raise ValueError( f"lens targets layer {target_layer}, but readout supports only the " f"published final-layer convention ({final_layer}); refit without " "a custom target layer." ) + capture_point = self.metadata.get("capture_point") + if capture_point is not None and capture_point != _CAPTURE_POINT_BLOCK_OUTPUT: + raise ValueError( + f"lens records capture_point={capture_point!r}, but readout reads block " + f"outputs (blocks.{{layer}}.hook_out, capture_point=" + f"{_CAPTURE_POINT_BLOCK_OUTPUT!r}); this lens was fitted on a different " + "activation." + ) + stream_reduce = self.metadata.get("stream_reduce") + if stream_reduce is not None and stream_reduce != _STREAM_REDUCE_NONE: + raise ValueError( + f"lens records stream_reduce={stream_reduce!r}, but JacobianLens supports " + f"single-stream block outputs only (stream_reduce={_STREAM_REDUCE_NONE!r})." + ) return self # ------------------------------------------------------------------ # @@ -1921,6 +2055,127 @@ def _require_raw_bridge(model: Any, *, estimator: str) -> None: ) +def _matrix_width(path: str, jacobians: Dict[int, torch.Tensor]) -> int: + """Infer ``d_model`` from the first transport matrix of an artifact without one.""" + if not jacobians: + raise ValueError(f"{path} holds no transport matrices; at least one layer is required.") + return int(next(iter(jacobians.values())).shape[0]) + + +def _sanitize_metadata_value(value: Any, path: str, dropped: List[str]) -> Tuple[bool, Any]: + """Return ``(keep, cleaned)``, dropping leaves ``weights_only=True`` cannot reload. + + Dropped leaves are named in ``dropped`` with their dotted path so an imported + provenance dict survives with only the offending entries removed. + """ + if value is None or type(value) in (bool, int, float, str): + return True, value + if isinstance(value, torch.Tensor): + dropped.append(f"{path}: shape={tuple(value.shape)} dtype={value.dtype}") + return False, None + if type(value) in (list, tuple): + items = [] + for index, item in enumerate(value): + keep, cleaned = _sanitize_metadata_value(item, f"{path}[{index}]", dropped) + if keep: + items.append(cleaned) + return True, type(value)(items) + if type(value) is dict: + cleaned_dict: Dict[str, Any] = {} + for key, item in value.items(): + if type(key) is not str: + dropped.append(f"{path}.{key!r}: non-string key") + continue + keep, cleaned = _sanitize_metadata_value(item, f"{path}.{key}", dropped) + if keep: + cleaned_dict[key] = cleaned + return True, cleaned_dict + dropped.append(f"{path}: {type(value).__name__}") + return False, None + + +def _normalize_artifact_provenance(provenance: Dict[str, Any]) -> Dict[str, Any]: + """Map recorded artifact provenance onto the metadata keys this class reads. + + Values are looked up on the provenance dict first and then on its nested + ``source_config`` (the J++ Lens package's ``LensConfig``). Only recorded + values are copied: a missing fitting revision stays missing rather than + being guessed from the model. + """ + source_config = provenance.get("source_config") + sources: List[Dict[str, Any]] = [provenance] + if isinstance(source_config, dict): + sources.append(source_config) + + def recorded(*names: str) -> Any: + for mapping in sources: + for name in names: + value = mapping.get(name) + if value is not None: + return value + return None + + normalized: Dict[str, Any] = {} + model_name = recorded("model_name", "hf_model_name") + if isinstance(model_name, str) and model_name: + normalized["model_name"] = model_name + revision = recorded("model_revision", "hf_model_revision") + if isinstance(revision, str) and revision: + normalized["model_revision"] = revision + for key in ("method", "lrp_mode", "capture_point", "stream_reduce"): + value = recorded(key) + if isinstance(value, str): + normalized[key] = value + for key in ("target_layer", "relative_end_transport_layer"): + value = recorded(key) + if isinstance(value, int) and not isinstance(value, bool): + normalized[key] = value + return normalized + + +def _import_artifact_metadata( + explicit: Optional[Dict[str, Any]], + provenance: Optional[Dict[str, Any]], + *, + converted_from: Optional[str], +) -> Optional[Dict[str, Any]]: + """Combine an artifact's ``metadata`` dict with its recorded provenance. + + Explicit ``metadata`` keys win; the raw provenance is kept under + ``metadata["provenance"]`` and its normalised keys fill the gaps. With no + provenance and no import marker the explicit dict is returned untouched so + plain artifacts load exactly as before. + """ + if provenance is None and converted_from is None: + return explicit + metadata: Dict[str, Any] = dict(explicit or {}) + dropped: List[str] = [] + if provenance is not None: + _, clean_provenance = _sanitize_metadata_value(provenance, "provenance", dropped) + metadata.setdefault("provenance", clean_provenance) + for key, value in _normalize_artifact_provenance(clean_provenance).items(): + metadata.setdefault(key, value) + if converted_from is not None: + metadata["converted_from"] = converted_from + if dropped: + metadata["dropped_fields"] = list(metadata.get("dropped_fields", [])) + dropped + return metadata + + +def _require_additive_statistics(lens: "JacobianLens", index: int) -> None: + """Refuse to merge maps that are not prompt means of Jacobians.""" + method = lens.metadata.get("method") + converted_from = lens.metadata.get("converted_from") + if method in _NON_ADDITIVE_METHODS or converted_from == _JPP_ARTIFACT_SENTINEL: + raise ValueError( + f"lens {index} was imported from a J++ Lens artifact (method={method!r}, " + f"converted_from={converted_from!r}). Its maps combine clustered expert " + "Jacobians with learned weights and carry no additive prompt statistics, so " + "merge() cannot average them with other shards; combine such fits with the " + "J++ Lens package instead." + ) + + def _validate_metadata(metadata: Dict[str, Any]) -> None: """Reject values that ``torch.load(weights_only=True)`` cannot reload.""" From 820159d45956a5e3ff8af406ba5d15a9e6999250 Mon Sep 17 00:00:00 2001 From: aidxhxr Date: Tue, 10 Nov 2026 10:18:12 -0500 Subject: [PATCH 2/2] fix(jacobian_lens): treat d_model 0 as unset and check source_layers against the jacobians --- tests/unit/tools/test_jacobian_lens_import.py | 42 +++++++++++++++-- .../tools/analysis/jacobian_lens.py | 45 +++++++++++++++---- 2 files changed, 75 insertions(+), 12 deletions(-) diff --git a/tests/unit/tools/test_jacobian_lens_import.py b/tests/unit/tools/test_jacobian_lens_import.py index b07f7f81ff..df6340d519 100644 --- a/tests/unit/tools/test_jacobian_lens_import.py +++ b/tests/unit/tools/test_jacobian_lens_import.py @@ -190,7 +190,12 @@ def test_flat_provenance_harvested(self, tmp_path): """Checkpoints that store model_name / corpus at top level are handled.""" p = str(tmp_path / "ckpt.pt") _save_checkpoint( - p, flat_provenance={"model_name": "llama3", "corpus": "pile", "model_revision": "abc"} + p, + flat_provenance={ + "model_name": "llama3", + "corpus": "pile", + "model_revision": "abc", + }, ) lens = JacobianLens.load(p) assert lens.metadata.get("model_name") == "llama3" @@ -256,7 +261,11 @@ def test_merge_two_checkpoints_succeeds(self, tmp_path): assert merged.n_prompts == l1.n_prompts + l2.n_prompts def test_merge_two_artifacts_succeeds(self, tmp_path): - shared_meta = {"model_name": "gpt2", "corpus": "wiki", "transformer_lens_fit": True} + shared_meta = { + "model_name": "gpt2", + "corpus": "wiki", + "transformer_lens_fit": True, + } p1, p2 = str(tmp_path / "a1.pt"), str(tmp_path / "a2.pt") _save_artifact(p1, metadata=dict(shared_meta)) _save_artifact(p2, metadata=dict(shared_meta)) @@ -446,7 +455,10 @@ def test_explicit_metadata_wins_over_provenance(self, tmp_path): p = str(tmp_path / "jpp.pt") provenance = _converted_provenance(_author_config()) _save_converted_artifact( - p, _jacobians(), provenance, metadata={"model_name": "explicit", "corpus": "wiki"} + p, + _jacobians(), + provenance, + metadata={"model_name": "explicit", "corpus": "wiki"}, ) lens = JacobianLens.load(p) assert lens.metadata["model_name"] == "explicit" @@ -550,6 +562,24 @@ def test_d_model_inferred_when_config_omits_it(self, tmp_path): _save_author_artifact(p, _jacobians(d_model=4), config) assert JacobianLens.load(p).d_model == 4 + def test_d_model_zero_in_config_falls_back_to_matrix_width(self, tmp_path): + path = tmp_path / "lens.pt" + _save_author_artifact(path, _jacobians(d_model=8), _author_config(d_model=0)) + + lens = JacobianLens.load(str(path)) + + assert lens.d_model == 8 + + def test_rejects_source_layers_that_disagree_with_jacobians(self, tmp_path): + path = tmp_path / "lens.pt" + _save_author_artifact(path, _jacobians(n_layers=3), _author_config()) + payload = torch.load(path) + payload["source_layers"] = [0, 1] + torch.save(payload, path) + + with pytest.raises(ValueError, match="source_layers"): + JacobianLens.load(str(path)) + def test_rejects_non_jacobian_lens_type(self, tmp_path): p = str(tmp_path / "logit.pt") torch.save( @@ -566,7 +596,11 @@ def test_rejects_non_jacobian_lens_type(self, tmp_path): def test_rejects_parameters_without_jacobians(self, tmp_path): p = str(tmp_path / "no_jac.pt") torch.save( - {"parameters": {"other": {0: 1.0}}, "source_layers": [0], "config": _author_config()}, + { + "parameters": {"other": {0: 1.0}}, + "source_layers": [0], + "config": _author_config(), + }, p, ) with pytest.raises(ValueError, match="parameters.jacobians"): diff --git a/transformer_lens/tools/analysis/jacobian_lens.py b/transformer_lens/tools/analysis/jacobian_lens.py index bc4d1bcc49..34928a4322 100644 --- a/transformer_lens/tools/analysis/jacobian_lens.py +++ b/transformer_lens/tools/analysis/jacobian_lens.py @@ -471,7 +471,7 @@ def _from_artifact_payload(cls: Type[_LensT], path: str, payload: Dict[str, Any] d_model = config.get("d_model") if n_prompts is None: n_prompts = config.get("num_prompts_trained_on") - if d_model is None: + if not d_model: d_model = _matrix_width(path, jacobians) explicit = payload.get("metadata") return cls( @@ -514,8 +514,16 @@ def _from_jpp_payload(cls: Type[_LensT], path: str, payload: Dict[str, Any]) -> "one [d_model, d_model] transport matrix per source layer there." ) jacobians = {int(layer): matrix for layer, matrix in raw_jacobians.items()} + source_layers = payload.get("source_layers") + if source_layers is not None: + listed = sorted(int(layer) for layer in source_layers) + if listed != sorted(jacobians): + raise ValueError( + f"{path} lists source_layers {listed} but parameters.jacobians holds " + f"layers {sorted(jacobians)}" + ) d_model = config.get("d_model") - if d_model is None: + if not d_model: d_model = _matrix_width(path, jacobians) return cls( jacobians, @@ -1393,7 +1401,10 @@ def fraction_of_variance( } per_position = {layer: torch.tensor(fractions[layer]) for layer in resolved_layers} return JSpaceVarianceProfile( - layers=resolved_layers, median=median, pooled=pooled, per_position=per_position + layers=resolved_layers, + median=median, + pooled=pooled, + per_position=per_position, ) # ------------------------------------------------------------------ # @@ -1577,7 +1588,9 @@ def swap_hooks( raise ValueError(f"clamp_to is missing activation {hook_name!r}") clean = clamp_to[hook_name] _validate_residual_activation( - clean, d_model=model.cfg.d_model, hook_name=f"clamp_to[{hook_name!r}]" + clean, + d_model=model.cfg.d_model, + hook_name=f"clamp_to[{hook_name!r}]", ) normalized = _normalize_positions(positions, clean.shape[1]) clean_selected = clean if positions is None else clean[:, normalized, :] @@ -1626,7 +1639,10 @@ def transform( return selected.float() + delta hooks.append( - (hook_name, _make_intervention_hook(transform, positions, model.cfg.d_model)) + ( + hook_name, + _make_intervention_hook(transform, positions, model.cfg.d_model), + ) ) return hooks @@ -1648,7 +1664,12 @@ def swap_clamp_hooks( stacklevel=2, ) return self.swap_hooks( - model, source_token, target_token, layers, positions=positions, clamp_to=clean_cache + model, + source_token, + target_token, + layers, + positions=positions, + clamp_to=clean_cache, ) def coordinate_patch_hooks( @@ -1853,7 +1874,11 @@ def fit( n_prompts=n_done, d_model=options.d_model, metadata=_build_fit_metadata( - model, options, estimator=ESTIMATOR_JACOBIAN, n_prompts=n_done, metadata=metadata + model, + options, + estimator=ESTIMATOR_JACOBIAN, + n_prompts=n_done, + metadata=metadata, ), ) @@ -2272,7 +2297,11 @@ def hook_fn( def _distinct_ids( - model: Any, source_token: ConceptToken, target_token: ConceptToken, *, operation: str + model: Any, + source_token: ConceptToken, + target_token: ConceptToken, + *, + operation: str, ) -> Tuple[int, int]: """Resolve a source/target pair, refusing one that names the same token twice.""" source_id, target_id = to_token_ids(model, [source_token, target_token])