diff --git a/demos/Jacobian_Lens_Demo.ipynb b/demos/Jacobian_Lens_Demo.ipynb index f608bf97d..3ea7dcc3b 100644 --- a/demos/Jacobian_Lens_Demo.ipynb +++ b/demos/Jacobian_Lens_Demo.ipynb @@ -1726,12 +1726,137 @@ "print(\"model output:\", top_j[model.cfg.n_layers - 1][-1][0])\n" ] }, + { + "cell_type": "markdown", + "id": "jpp-intro", + "metadata": {}, + "source": [ + "## 7. Pretrained J++ lenses and filtered readouts\n", + "\n", + "[J++ Lens](https://github.com/safety-research/jpp_lens) (Ayonrinde & Lindsey, 2026) combines filtered Jacobian estimates into one fixed matrix per layer. We can read its published maps through the same `JacobianLens`; the LRP rules used while fitting are not needed at inference. This smaller example uses the converted Qwen3.5-0.8B artifact from [Neuronpedia](https://huggingface.co/neuronpedia/jacobian-lens), with the authors' provenance retained in `lens.metadata[\"provenance\"]`.\n", + "\n", + "The tokenizer helper excludes individually decoded tokens without a Unicode letter or digit, scanning the whole LM-head vocabulary including padding rows. `keep_token_ids` can retain scored symbols such as `+`. Special tokens with letters in their decoded spelling are retained by this policy.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "jpp-load", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "method: jpp\n", + "excluded vocabulary ids: 6133\n" + ] + } + ], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "import gc\n", + "import torch\n", + "\n", + "from transformer_lens.model_bridge import TransformerBridge\n", + "from transformer_lens.tools.analysis import JacobianLens\n", + "from transformer_lens.utilities.tokenize_utils import readout_token_exclusions\n", + "\n", + "# Release the previous model before loading the smaller J++ example.\n", + "if \"lens\" in globals():\n", + " lens.clear_device_cache()\n", + "for name in (\"result_jlens\", \"result_logitlens\", \"lens\", \"model\"):\n", + " globals().pop(name, None)\n", + "gc.collect()\n", + "if torch.cuda.is_available():\n", + " torch.cuda.empty_cache()\n", + "\n", + "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", + "model = TransformerBridge.boot_transformers(\n", + " \"Qwen/Qwen3.5-0.8B\",\n", + " revision=\"2fc06364715b967f1860aea9cf38778875588b17\",\n", + " device=device,\n", + " dtype=torch.float32,\n", + ").eval()\n", + "lens = JacobianLens.from_pretrained(\n", + " \"neuronpedia/jacobian-lens\",\n", + " filename=\"qwen3.5-0.8b/jpp/Salesforce-wikitext/Qwen3.5-0.8B_jpp_lens.pt\",\n", + " revision=\"e166deb5ad540932100ef803c8c9bcc2a5453694\",\n", + " model=model,\n", + ")\n", + "assert lens.metadata[\"provenance\"][\"method\"] == \"jpp\"\n", + "excluded = readout_token_exclusions(model.tokenizer, vocab_size=model.W_U.shape[-1])\n", + "print(\"method:\", lens.metadata[\"provenance\"][\"method\"])\n", + "print(\"excluded vocabulary ids:\", len(excluded))\n" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "jpp-filter", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "layer 4\n", + " raw J++: [' ,', ' \\n\\n', ' \",', ' ##', ' ,\"']\n", + " filtered J++: ['<|im_end|>', ' a', ' the', ' either', ' The']\n", + " filtered logit lens: ['менно', ' Nations', ' Panas', 'ң', ' geweest']\n", + "layer 12\n", + " raw J++: [' ____', ' ___', ' ______', ' ?', ' __']\n", + " filtered J++: [' nation', ' named', ' national', ' Eleven', ' Americans']\n", + " filtered logit lens: ['ありますか', ' gì', 'voie', '怎样的', '其中之一']\n", + "layer 20\n", + " raw J++: ['____', ':', '巴黎', '?', '___']\n", + " filtered J++: ['巴黎', '法国', 'Paris', 'France', '哪个']\n", + " filtered logit lens: ['哪一个', ' capitals', '巴黎', '哪个', 'elia']\n" + ] + } + ], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "prompt = \"法国首都是巴黎。埃菲尔铁塔所在的城市是\"\n", + "readout_options = dict(\n", + " layers=[4, 12, 20], positions=[-1], top_k=5, return_full_logits=True\n", + ")\n", + "raw = lens.readout(model, prompt, **readout_options)\n", + "filtered = lens.readout(model, prompt, exclude_token_ids=excluded, **readout_options)\n", + "logit_baseline = lens.readout(\n", + " model, prompt, use_jacobian=False, exclude_token_ids=excluded, **readout_options\n", + ")\n", + "assert filtered.exclude_token_ids == excluded\n", + "assert torch.equal(raw.model_logits, filtered.model_logits)\n", + "assert torch.equal(raw.model_topk_indices, filtered.model_topk_indices)\n", + "for layer in readout_options[\"layers\"]:\n", + " assert torch.equal(raw.lens_logits[layer], filtered.lens_logits[layer])\n", + " assert not set(filtered.lens_topk_indices[layer].flatten().tolist()).intersection(\n", + " excluded\n", + " )\n", + " print(f\"layer {layer}\")\n", + " print(\" raw J++:\", raw.top_tokens(model.tokenizer, k=5)[layer][0])\n", + " print(\" filtered J++:\", filtered.top_tokens(model.tokenizer, k=5)[layer][0])\n", + " print(\n", + " \" filtered logit lens:\",\n", + " logit_baseline.top_tokens(model.tokenizer, k=5)[layer][0],\n", + " )\n" + ] + }, + { + "cell_type": "markdown", + "id": "jpp-policy", + "metadata": {}, + "source": [ + "Filtering changes the ranking universe, not the raw lens logits or the model's predictions. The same exclusions are applied to the logit-lens baseline. These readouts illustrate how to use the published maps; comparing lens quality requires a fixed task evaluation rather than judging a few attractive token lists.\n" + ] + }, { "cell_type": "markdown", "id": "00df913e", "metadata": {}, "source": [ - "## 7. Fitting your own lens\n", + "## 8. Fitting your own lens\n", "\n", "For models without a published artifact, `JacobianLens.fit` reproduces the reference estimator on\n", "a raw `TransformerBridge` — one forward and `ceil(d_model / dim_batch)` backward passes per\n", diff --git a/tests/unit/tools/test_jacobian_lens_filtered_readout.py b/tests/unit/tools/test_jacobian_lens_filtered_readout.py new file mode 100644 index 000000000..fe9e60e54 --- /dev/null +++ b/tests/unit/tools/test_jacobian_lens_filtered_readout.py @@ -0,0 +1,141 @@ +"""Converted lens provenance and optional vocabulary exclusions.""" + +import pytest +import torch + +from tests.unit.tools.conftest import D_MODEL, D_VOCAB, N_LAYERS, _lens +from transformer_lens.tools.analysis import JacobianLens + + +def test_converted_provenance_round_trips_without_becoming_model_checks(tmp_path, toy_model): + provenance = { + "method": "jpp", + "source_config": {"hf_model_name": "a/mirror"}, + "capture_point": "block_output", + "stream_reduce": "none", + } + path = tmp_path / "converted.pt" + torch.save( + { + "J": {0: torch.eye(D_MODEL)}, + "d_model": D_MODEL, + "n_prompts": 96, + "provenance": provenance, + "metadata": {"note": "retained"}, + }, + path, + ) + lens = JacobianLens.load(str(path)) + assert lens.metadata == {"note": "retained", "provenance": provenance} + lens.validate_model(toy_model) + lens.save(str(path)) + assert JacobianLens.load(str(path)).metadata == lens.metadata + + +def test_different_imported_provenance_uses_existing_merge_guard(tmp_path): + lenses = [] + for method in ("jpp", "jacobian"): + path = tmp_path / f"{method}.pt" + torch.save( + { + "J": {0: torch.eye(D_MODEL)}, + "d_model": D_MODEL, + "n_prompts": 2, + "provenance": {"method": method}, + }, + path, + ) + lenses.append(JacobianLens.load(str(path))) + with pytest.raises(ValueError, match="provenance"): + JacobianLens.merge(lenses) + merged = JacobianLens.merge([lenses[0], lenses[0]]) + assert merged.metadata["provenance"] == {"method": "jpp"} + + +@pytest.mark.parametrize("use_jacobian", [True, False]) +def test_exclusions_rank_before_topk_and_leave_raw_outputs_unchanged(toy_model, use_jacobian): + lens = _lens() + kwargs = dict( + layers=[0, N_LAYERS - 1], + positions=[-1], + top_k=3, + return_full_logits=True, + use_jacobian=use_jacobian, + ) + baseline = lens.readout(toy_model, "a prompt", **kwargs) + excluded = sorted({int(ids[0, 0]) for ids in baseline.lens_topk_indices.values()}) + filtered = lens.readout(toy_model, "a prompt", exclude_token_ids=excluded + excluded, **kwargs) + assert filtered.exclude_token_ids == excluded + assert baseline.exclude_token_ids == [] + torch.testing.assert_close(filtered.model_logits, baseline.model_logits, rtol=0, atol=0) + assert torch.equal(filtered.model_topk_indices, baseline.model_topk_indices) + assert torch.equal(filtered.model_topk_values, baseline.model_topk_values) + for layer in baseline.lens_logits: + original = baseline.lens_logits[layer] + assert torch.equal(filtered.lens_logits[layer], original) + ranked = original.clone() + ranked[:, excluded] = -torch.inf + expected = ranked.topk(3, dim=-1) + assert torch.equal(filtered.lens_topk_indices[layer], expected.indices) + assert torch.equal(filtered.lens_topk_values[layer], original.gather(-1, expected.indices)) + + +def test_empty_exclusions_preserve_default_results(toy_model): + lens = _lens() + baseline = lens.readout(toy_model, "a prompt", top_k=3) + actual = lens.readout(toy_model, "a prompt", top_k=3, exclude_token_ids=[]) + for layer in baseline.lens_topk_indices: + assert torch.equal(actual.lens_topk_indices[layer], baseline.lens_topk_indices[layer]) + assert torch.equal(actual.lens_topk_values[layer], baseline.lens_topk_values[layer]) + + +def test_excluded_ids_cannot_win_negative_infinity_ties(toy_model, monkeypatch): + from types import SimpleNamespace + + import transformer_lens.tools.analysis.jacobian_lens as module + + original = toy_model.run_with_cache + + def infinite_logits(*args, **kwargs): + logits, cache = original(*args, **kwargs) + return torch.full_like(logits, -torch.inf), cache + + monkeypatch.setattr(toy_model, "run_with_cache", infinite_logits) + monkeypatch.setattr( + module, + "logit_readout", + lambda *args, **kwargs: SimpleNamespace(values=torch.full((1, D_VOCAB), -torch.inf)), + ) + result = _lens().readout( + toy_model, + "prompt", + layers=[0, N_LAYERS - 1], + positions=[-1], + top_k=2, + exclude_token_ids=list(range(D_VOCAB - 2)), + ) + for ids in result.lens_topk_indices.values(): + assert set(ids.flatten().tolist()) == {D_VOCAB - 2, D_VOCAB - 1} + + +def test_conflicting_provenance_is_not_silently_overwritten(tmp_path): + path = tmp_path / "conflict.pt" + torch.save( + { + "J": {0: torch.eye(D_MODEL)}, + "d_model": D_MODEL, + "metadata": {"provenance": {"method": "j"}}, + "provenance": {"method": "jpp"}, + }, + path, + ) + with pytest.raises(ValueError, match="conflicting"): + JacobianLens.load(str(path)) + + +@pytest.mark.parametrize( + "excluded", [[-1], [D_VOCAB], [True], list(range(D_VOCAB)), list(range(D_VOCAB - 1))] +) +def test_invalid_or_insufficient_exclusions_raise(toy_model, excluded): + with pytest.raises(ValueError, match="exclude|eligible"): + _lens().readout(toy_model, "a prompt", top_k=2, exclude_token_ids=excluded) diff --git a/tests/unit/utilities/test_readout_token_exclusions.py b/tests/unit/utilities/test_readout_token_exclusions.py new file mode 100644 index 000000000..1d0d0b201 --- /dev/null +++ b/tests/unit/utilities/test_readout_token_exclusions.py @@ -0,0 +1,41 @@ +"""Unicode-aware readout filtering over the full output vocabulary.""" + +import pytest + +from transformer_lens.utilities import tokenize_utils + + +class Tokenizer: + words = ["", "\n", "!!!", "心脏", "é", "٣", "<|im_end|>", "+", " word"] + + def decode(self, ids): + return self.words[ids[0]] if ids[0] < len(self.words) else "" + + +class BatchTokenizer(Tokenizer): + def batch_decode(self, rows): + return [self.decode(row) for row in rows] + + +@pytest.mark.parametrize("tokenizer", [Tokenizer(), BatchTokenizer()]) +def test_unicode_padding_and_special_tokens(tokenizer): + assert tokenize_utils.readout_token_exclusions(tokenizer, 11) == [0, 1, 2, 7, 9, 10] + assert tokenize_utils.readout_token_exclusions(tokenizer, 11, keep_token_ids=[7, 7, 3]) == [ + 0, + 1, + 2, + 9, + 10, + ] + + +@pytest.mark.parametrize("width", [0, -1, True]) +def test_invalid_vocabulary_size(width): + with pytest.raises(ValueError, match="vocab_size"): + tokenize_utils.readout_token_exclusions(Tokenizer(), width) + + +@pytest.mark.parametrize("ids", [[-1], [11], [True]]) +def test_invalid_retained_ids(ids): + with pytest.raises(ValueError, match="keep_token_ids"): + tokenize_utils.readout_token_exclusions(Tokenizer(), 11, keep_token_ids=ids) diff --git a/transformer_lens/tools/analysis/jacobian_lens.py b/transformer_lens/tools/analysis/jacobian_lens.py index 5d1d21912..635afed08 100644 --- a/transformer_lens/tools/analysis/jacobian_lens.py +++ b/transformer_lens/tools/analysis/jacobian_lens.py @@ -61,7 +61,7 @@ """ import warnings -from dataclasses import dataclass +from dataclasses import dataclass, field from importlib.metadata import version from typing import ( Any, @@ -83,7 +83,11 @@ from tqdm.auto import tqdm from transformer_lens.ActivationCache import ActivationCache -from transformer_lens.tools.analysis._model_state import require_eval_mode, to_token_ids +from transformer_lens.tools.analysis._model_state import ( + require_eval_mode, + require_int, + to_token_ids, +) from transformer_lens.tools.analysis.jacobian_lens_coordinate_patch import ( CoordinatePatch, solve_coordinate_patch, @@ -222,6 +226,8 @@ class JacobianLensReadout: with the ``pos`` axis of retained top-k and optional full logits. use_jacobian: Whether the Jacobian transport was applied (``False`` = logit lens). + exclude_token_ids: + Sorted, deduplicated vocabulary ids excluded from lens rankings only. """ lens_topk_values: Dict[int, Float[torch.Tensor, "pos k"]] @@ -233,6 +239,7 @@ class JacobianLensReadout: use_jacobian: bool = True lens_logits: Optional[Dict[int, Float[torch.Tensor, "pos d_vocab"]]] = None model_logits: Optional[Float[torch.Tensor, "pos d_vocab"]] = None + exclude_token_ids: List[int] = field(default_factory=list) def top_tokens(self, tokenizer: Any, k: int = 5) -> Dict[int, List[List[str]]]: """Decode the top-``k`` tokens per layer and position. @@ -355,6 +362,10 @@ def load(cls: Type[_LensT], path: str) -> _LensT: indices to transport matrices, plus ``n_prompts``, ``d_model``, and an optional ``metadata`` dict. + Converted workspace lenses may also carry top-level ``provenance``. + It is preserved under ``metadata["provenance"]`` for save/load round + trips, without turning source model names into validation constraints. + **Fit checkpoint** (running-sum format, written by the reference implementation's ``write_checkpoint()`` during fitting): must contain a ``jacobian_sum`` key mapping layer indices to *running-sum* matrices @@ -400,11 +411,18 @@ def load(cls: Type[_LensT], path: str) -> _LensT: """ payload = torch.load(path, map_location="cpu", weights_only=True) if "J" in payload: + metadata = dict(payload.get("metadata") or {}) + if "provenance" in payload: + _validate_metadata(metadata) + _validate_metadata({"provenance": payload["provenance"]}) + if "provenance" in metadata and metadata["provenance"] != payload["provenance"]: + raise ValueError("Artifact has conflicting metadata and top-level provenance") + metadata["provenance"] = payload["provenance"] 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"), + metadata=metadata, ) if "jacobian_sum" in payload: return cls._from_checkpoint_payload(path, payload) @@ -741,6 +759,7 @@ def readout( use_jacobian: bool = True, top_k: int = DEFAULT_TOP_K, return_full_logits: bool = False, + exclude_token_ids: Optional[Sequence[int]] = None, ) -> JacobianLensReadout: """Read per-layer vocabulary logits for a prompt. @@ -765,6 +784,9 @@ def readout( return_full_logits: Also retain full vocabulary tensors on CPU. This is opt-in because a 64-token Gemma readout across all layers is roughly 1.7 GB. + exclude_token_ids: Vocabulary ids to exclude before lens top-k, + including the final-layer and logit-lens readouts. Model top-k + and full logits stay unfiltered. Duplicates are ignored. Returns: A :class:`JacobianLensReadout`. @@ -808,6 +830,31 @@ def readout( f"{selected_model_logits.shape[-1]}" ) model_topk = selected_model_logits.topk(top_k, dim=-1) + excluded = sorted( + { + require_int( + token, + "exclude_token_ids entry", + minimum=0, + maximum=selected_model_logits.shape[-1] - 1, + ) + for token in (() if exclude_token_ids is None else exclude_token_ids) + } + ) + eligible = selected_model_logits.shape[-1] - len(excluded) + if top_k > eligible: + raise ValueError( + f"top_k={top_k} exceeds the {eligible} eligible tokens after exclusions" + ) + eligible_ids = None + if excluded: + keep = torch.ones( + selected_model_logits.shape[-1], + dtype=torch.bool, + device=selected_model_logits.device, + ) + keep[excluded] = False + eligible_ids = torch.arange(keep.numel(), device=keep.device)[keep] full_model_logits = selected_model_logits.cpu() if return_full_logits else None lens_topk_values: Dict[int, torch.Tensor] = {} lens_topk_indices: Dict[int, torch.Tensor] = {} @@ -815,7 +862,6 @@ def readout( for layer in layers: if layer == final_layer: layer_logits = selected_model_logits - layer_topk = model_topk else: activation = cache[hook_names[layer]] _validate_residual_activation( @@ -827,9 +873,18 @@ def readout( transported = self.transport(residual, layer) if use_jacobian else residual.float() # fp32 only after the output transform: the model soft-caps in W_U's dtype. layer_logits = logit_readout(model, transported, apply_ln=True).values.float() - layer_topk = layer_logits.topk(top_k, dim=-1) + if eligible_ids is not None: + ids = eligible_ids.to(layer_logits.device) + ranked = layer_logits.index_select(-1, ids) + layer_topk = ranked.topk(top_k, dim=-1) + topk_indices = ids[layer_topk.indices] + else: + layer_topk = ( + model_topk if layer == final_layer else layer_logits.topk(top_k, dim=-1) + ) + topk_indices = layer_topk.indices lens_topk_values[layer] = layer_topk.values.cpu() - lens_topk_indices[layer] = layer_topk.indices.cpu() + lens_topk_indices[layer] = topk_indices.cpu() if full_lens_logits is not None: if layer == final_layer: assert full_model_logits is not None @@ -846,6 +901,7 @@ def readout( use_jacobian=use_jacobian, lens_logits=full_lens_logits, model_logits=full_model_logits, + exclude_token_ids=excluded, ) @torch.no_grad() diff --git a/transformer_lens/utilities/__init__.py b/transformer_lens/utilities/__init__.py index 99cea18ba..7be8faef1 100644 --- a/transformer_lens/utilities/__init__.py +++ b/transformer_lens/utilities/__init__.py @@ -87,6 +87,7 @@ get_input_with_manually_prepended_bos, get_tokenizer_with_bos, get_tokens_with_bos_removed, + readout_token_exclusions, tokenize_and_concatenate, ) from .typed_module_list import TypedModuleList diff --git a/transformer_lens/utilities/tokenize_utils.py b/transformer_lens/utilities/tokenize_utils.py index da18ba97a..b34c10767 100644 --- a/transformer_lens/utilities/tokenize_utils.py +++ b/transformer_lens/utilities/tokenize_utils.py @@ -7,7 +7,7 @@ import os from copy import deepcopy -from typing import Any, Optional, cast +from typing import Any, Optional, Sequence, cast import einops import numpy as np @@ -20,6 +20,42 @@ from transformer_lens.utilities.tensors import get_cumsum_along_dim +def readout_token_exclusions( + tokenizer: Any, vocab_size: int, *, keep_token_ids: Sequence[int] = () +) -> list[int]: + """Ids whose individual decode contains no Unicode letter or digit. + + Scan the LM-head width, not ``len(tokenizer)``, to include padded output + rows that decode to an empty string. Special tokens are treated like any + other decoded text. ``keep_token_ids`` spares scored symbols (e.g. ``+``). + This builds a ranking policy; it does not modify the tokenizer or logits. + """ + if ( + isinstance(vocab_size, bool) + or not isinstance(vocab_size, (int, np.integer)) + or vocab_size < 1 + ): + raise ValueError("vocab_size must be a positive integer") + kept = set() + for token in keep_token_ids: + if ( + isinstance(token, bool) + or not isinstance(token, (int, np.integer)) + or not 0 <= token < vocab_size + ): + raise ValueError("keep_token_ids must contain integer ids within vocab_size") + kept.add(int(token)) + if callable(getattr(tokenizer, "batch_decode", None)): + decoded = tokenizer.batch_decode([[token] for token in range(vocab_size)]) + else: + decoded = [tokenizer.decode([token]) for token in range(vocab_size)] + return [ + token + for token, text in enumerate(decoded) + if token not in kept and not any(char.isalnum() for char in text) + ] + + def tokenize_and_concatenate( dataset: Dataset | IterableDataset, tokenizer: PreTrainedTokenizerBase,