diff --git a/docs/source/content/analysis_tools.md b/docs/source/content/analysis_tools.md index e5346ee5c..b9683ecf5 100644 --- a/docs/source/content/analysis_tools.md +++ b/docs/source/content/analysis_tools.md @@ -20,7 +20,7 @@ the [4.0 migration guide](migrating_to_v4.md) for existing `HookedTransformer` c | What does each layer "predict", and where does a token's rank change? | **Logit Lens** | A prompt or activation cache, optional target tokens, top-k or a vocabulary subset → per-layer logits, log-probabilities, ranks and entropy. | What the residual stream reads as through the real final norm and unembedding; the final entry reproduces the model's own logits. | | What vocabulary directions appear in an MLP's gradient factors? | **Backward Lens** | A prompt, one target token, and selected layers → gradient factors and vocabulary rankings. | A diagnostic of the forward inputs and backward signals composing a weight gradient. | | How can I read or edit residuals through a fitted transport map? | **Jacobian Lens** | A matching lens artifact and model, plus prompts or activations → vocabulary readouts, decompositions, or interventions. | Readouts under the fitted map; causal effects require running and measuring an intervention. | -| Does a subspace at one hook point carry a variable the output depends on? | **Interchange Interventions** | A hook point, base and source inputs, a basis, and a task readout → task logits with the base's subspace coordinates set to the source's. | The effect of that intervention on the chosen pairs; compare with a norm-matched placebo and random bases before reading it as use. | +| Does a subspace at one hook point carry a variable the output depends on? | **Interchange Interventions** | A hook point, base and source inputs, a basis (fixed, or learned by DAS from pairs labelled by a high-level model), and a task readout → task logits with the base's subspace coordinates set to the source's. | The effect of that intervention on the chosen pairs; compare with a norm-matched placebo and random bases before reading it as use. | | How much do two subspaces overlap? | **Projection Kernel** | Two subspace bases, or Bridge attention-head weight spaces → overlap scores and principal-angle information. | Shared geometric support, independent of the choice of basis within each subspace. | | What are a head's QK/OV singular directions? | **SVD Circuits: head decomposition** | Model weights and a layer/head index → singular values, vectors, and degeneracy reports. | The linear structure of the head's weight maps. | @@ -43,7 +43,7 @@ available hooks, and the meaning of the selected tensor axes. | Logit Lens | Targets TransformerBridge, raw or compatibility mode. Needs a decoder with an `ln_final` component (or no final norm) and `W_U`; encoder-decoder and encoder-only bridges are rejected, as are quantized unembeddings. Hybrid SSM stacks work with `incl_mid=False`. Run the model in evaluation mode. | One cached forward pass over the residual hooks, or reuse of a cache. Logits are produced in chunks of `chunk_size` rows; full-vocabulary output is materialized only when no `targets`, `top_k` or `vocab` selector is given. | | Backward Lens | Requires a raw decoder-only TransformerBridge with dense, non-gated MLPs and without compatibility mode, such as GPT-2 or Pythia/GPT-NeoX. Use one target token and check the restrictions in the [tool guide](backward_lens.md). | Gradient computation plus vocabulary projections for selected layers and positions; retaining full logits increases memory. | | Jacobian Lens | Requires a fresh, causal decoder-only Bridge with raw HF weights, without compatibility mode or weight processing. Validate the lens against the model. Fitting additionally requires all modules in evaluation mode. | Loading an existing artifact avoids fitting. The ordinary fitting estimator uses one forward and `ceil(d_model / dim_batch)` backwards per prompt; larger batches increase memory. | -| Interchange Interventions | Any model with the hook API (TransformerBridge or `HookedRootModule`). Activations at the hook must be `[batch, d_model]`, or `[batch, pos, d_model]` with one position index per example; base and source rows are paired. | One forward pass per `batch_size` sources to collect activations, then one per `batch_size` bases per intervention with `intervened_logits`; both run without gradients, so memory is bounded by `batch_size`. `run_with_intervention` is a single differentiable pass over all the bases it is given. | +| Interchange Interventions | Any model with the hook API (TransformerBridge or `HookedRootModule`). Activations at the hook must be `[batch, d_model]`, or `[batch, pos, d_model]` with one position index per example; base and source rows are paired. | One forward pass per `batch_size` sources to collect activations, then one per `batch_size` bases per intervention with `intervened_logits`; both run without gradients, so memory is bounded by `batch_size`. `run_with_intervention` is a single differentiable pass over all the bases it is given. DAS adds a forward and backward pass per batch per epoch; `micro_batch_size` bounds its memory. | | Projection Kernel | The numerical API accepts finite, real floating-point matrices via orthonormal bases. The attention-head wrapper requires Bridge weights with compatible dimensions and ranks. | Basis extraction uses SVD. All-head comparisons allocate basis stacks and a pairwise score grid; they can be large despite requiring no forward pass. | | SVD head decomposition | Uses Bridge per-block weight accessors; requires accessible, compatible `W_Q`, `W_K`, `W_V`, and `W_O` for the selected head. No compatibility mode or activation cache is needed. | Factored QK/OV SVD, with rank bounded by `d_head`; the dense `d_model × d_model` product is not materialized. | diff --git a/docs/source/content/interchange_interventions.md b/docs/source/content/interchange_interventions.md index 635aa86ec..0564378f8 100644 --- a/docs/source/content/interchange_interventions.md +++ b/docs/source/content/interchange_interventions.md @@ -33,25 +33,39 @@ from transformer_lens.tools.analysis import ( bridge = TransformerBridge.boot_transformers("gpt2") answers = [bridge.to_single_token(" is"), bridge.to_single_token(" are")] task_logits = lambda tokens: bridge(tokens)[:, -1, answers] - -base = bridge.to_tokens(["The key near the cabinet", "The keys near the cabinet"]) -source = base.flip(0) # each base takes the other prompt's subject number -subject = torch.full((2,), 2) # the subject token, after the BOS token hook = "blocks.1.hook_out" -source_acts = collect_activations(bridge, hook, source, task_logits, positions=subject) -basis = mean_difference_basis(source_acts, labels=[1, 0]) + +def pairs(nouns): + """Bases and sources differ only in the subject's number; the target is the source's verb.""" + singular = bridge.to_tokens([f"The {n} near the cabinet" for n in nouns]) + plural = bridge.to_tokens([f"The {n}s near the cabinet" for n in nouns]) + base, source = torch.cat([singular, plural]), torch.cat([plural, singular]) + targets = (torch.arange(len(base)) < len(nouns)).long() # " are" when the source is plural + subject = torch.full((len(base),), 2) # the subject token, after the BOS token and "The" + return base, source, targets, subject + + +fit_base, fit_source, fit_targets, fit_at = pairs(["key", "dog", "cat", "book", "car", "boy"]) +test_base, test_source, test_targets, test_at = pairs(["girl", "bird", "tree", "house"]) + +# Fit the basis on the fit pairs' sources, labelled by subject number ... +fit_acts = collect_activations(bridge, hook, fit_source, task_logits, positions=fit_at) +basis = mean_difference_basis(fit_acts, labels=fit_targets.tolist()) + +# ... and measure the interchange accuracy (IIA) on held-out pairs +test_acts = collect_activations(bridge, hook, test_source, task_logits, positions=test_at) logits = intervened_logits( - bridge, hook, SubspaceInterchange(basis), base, source_acts, task_logits, positions=subject + bridge, hook, SubspaceInterchange(basis), test_base, test_acts, task_logits, positions=test_at ) -print(task_logits(base).argmax(-1), logits.argmax(-1)) # tensor([0, 1]) tensor([1, 0]) +print((logits.argmax(-1) == test_targets).float().mean()) # tensor(1.) ``` `task_logits` runs the model in one forward pass and reads the task out of its output, so the tools do not assume next-token prediction; the hook must fire exactly once in that pass. Activations at the hook must be `[batch, d_model]`, or `[batch, pos, d_model]` with one position -index per example. A real experiment fits the basis on held-out pairs, not on the pairs it -evaluates. +index per example. Fit a basis on one set of pairs and evaluate it on pairs it was not fit on, as +here. `collect_activations` and `intervened_logits` run in batches of `batch_size` without gradients and return CPU tensors, so memory does not grow with the number of pairs. `run_with_intervention` is @@ -69,6 +83,36 @@ parameters. Run every basis on the same pairs. A basis whose IIA does not clearly exceed a random basis of the same rank is not evidence that the subspace carries the variable. +## Learn the subspace with DAS + +Distributed alignment search (DAS; Geiger et al., 2024) learns the basis instead of fixing it. +`fit_das` trains an orthonormal basis so that the interchange makes the model produce, for each +base/source pair, the class the high-level model predicts: + +```python +from transformer_lens.tools.analysis import fit_das + +result = fit_das( + bridge, hook, fit_base, fit_source, fit_targets, task_logits, rank=1, + base_positions=fit_at, source_positions=fit_at, epochs=10, +) +logits = intervened_logits( + bridge, hook, SubspaceInterchange(result.subspace), test_base, test_acts, task_logits, + positions=test_at, +) +print((logits.argmax(-1) == test_targets).float().mean()) # tensor(1.) +``` + +`targets[i]` is the high-level model's output for base `i` with the variable taken from source +`i`. Fit on one set of pairs and evaluate the IIA on pairs the basis was not trained on. Include +training pairs where several variables differ between base and source, labelled by the high-level +model's interchange: otherwise the basis can learn whatever the training pairs happen to move +together. The model's parameters are frozen during training and restored afterwards; +`micro_batch_size` accumulates gradients over smaller parts of each batch to save memory. + +DAS is expressive enough to fit arbitrary label mappings and to find effects in randomly +initialised models (Arora et al., 2024), so its IIA needs controls trained the same way. + ## Placebo `NormMatchedPlacebo(reference, basis)` moves each base exactly as far as `reference` does, but @@ -86,5 +130,9 @@ IIA beside the placebo and the baselines. ## References - Geiger, Lu, Icard and Potts (2021). Causal Abstractions of Neural Networks. NeurIPS. +- Arora, Jurafsky and Potts (2024). CausalGym: Benchmarking causal interpretability methods on + linguistic tasks. ACL. +- Geiger, Wu, Potts, Icard and Goodman (2024). Finding Alignments Between Interpretable Causal + Variables and Distributed Neural Representations. CLeaR. - Makelov, Lange and Nanda (2024). Is This the Subspace You Are Looking for? An Interpretability Illusion for Subspace Activation Patching. ICLR. diff --git a/tests/integration/model_bridge/test_interchange.py b/tests/integration/model_bridge/test_interchange.py index 5c4b79d81..74ad81abc 100644 --- a/tests/integration/model_bridge/test_interchange.py +++ b/tests/integration/model_bridge/test_interchange.py @@ -6,6 +6,7 @@ NormMatchedPlacebo, SubspaceInterchange, collect_activations, + fit_das, intervened_logits, random_basis, ) @@ -62,3 +63,59 @@ def test_the_placebo_matches_the_reference_displacement_on_real_activations(gpt2 base, source = acts, acts.flip(0) moved = (placebo(base, source) - base).norm(dim=-1) assert torch.allclose(moved, (reference(base, source) - base).norm(dim=-1), rtol=1e-4) + + +NOUNS = "key dog cat book car boy girl bird tree house road table door letter song plan".split() + + +def _number_pairs(bridge, nouns): + """Base and source differ only in the subject's number; the target is the source's verb.""" + singular = bridge.to_tokens([f"The {n} near the cabinet" for n in nouns]) + plural = bridge.to_tokens([f"The {n}s near the cabinet" for n in nouns]) + # BOS, "The", noun, "near", "the", "cabinet": every noun and its plural is one token + assert singular.shape[1] == plural.shape[1] == 6 + base = torch.cat([singular, plural]) + source = torch.cat([plural, singular]) + targets = torch.cat([torch.ones(len(nouns)), torch.zeros(len(nouns))]).long() + return base, source, targets + + +def test_das_finds_subject_number_on_held_out_nouns(gpt2_bridge): + answers = [gpt2_bridge.to_single_token(" is"), gpt2_bridge.to_single_token(" are")] + task_logits = lambda toks: gpt2_bridge(toks)[:, -1, answers] + hook = "blocks.1.hook_out" + base, source, targets = _number_pairs(gpt2_bridge, NOUNS[:10]) + subject = torch.full((len(base),), 2) # after the BOS token and "The" + result = fit_das( + gpt2_bridge, + hook, + base, + source, + targets, + task_logits, + rank=1, + base_positions=subject, + source_positions=subject, + epochs=10, + batch_size=20, + lr=1e-2, + ) + + base, source, targets = _number_pairs(gpt2_bridge, NOUNS[10:]) + subject = torch.full((len(base),), 2) + source_acts = collect_activations(gpt2_bridge, hook, source, task_logits, positions=subject) + + def accuracy(basis): + logits = intervened_logits( + gpt2_bridge, + hook, + SubspaceInterchange(basis), + base, + source_acts, + task_logits, + positions=subject, + ) + return (logits.argmax(-1) == targets).float().mean().item() + + assert accuracy(result.subspace) >= 0.9 + assert accuracy(random_basis(gpt2_bridge.cfg.d_model, 1, seed=0)) <= 0.2 diff --git a/tests/unit/tools/test_interchange.py b/tests/unit/tools/test_interchange.py index 6b016c122..f037d556f 100644 --- a/tests/unit/tools/test_interchange.py +++ b/tests/unit/tools/test_interchange.py @@ -8,9 +8,12 @@ from transformer_lens.HookedRootModule import HookedRootModule from transformer_lens.tools.analysis import SubspaceInterchange as exported_interchange from transformer_lens.tools.analysis.interchange import ( + LearnedSubspace, NormMatchedPlacebo, SubspaceInterchange, + _model_device, collect_activations, + fit_das, intervened_logits, mean_difference_basis, probe_basis, @@ -32,6 +35,7 @@ def __init__(self, seq: bool = False): self.hook_mid = HookPoint() self.seq = seq self.setup() + self.eval() def forward(self, x: torch.Tensor) -> torch.Tensor: h = self.hook_mid(self.inp(x)) @@ -155,6 +159,7 @@ def __init__(self): super().__init__() self.hook_pattern = HookPoint() self.setup() + self.eval() def forward(self, x: torch.Tensor) -> torch.Tensor: return self.hook_pattern(x).flatten(1)[:, :2] @@ -203,6 +208,7 @@ def __init__(self): super().__init__() self.hook_unused = HookPoint() self.setup() + self.eval() x, source, full = ( torch.randn(2, D_MODEL), @@ -227,3 +233,136 @@ def test_intervened_logits_matches_one_pass_in_bounded_batches(): assert not batched.requires_grad and batched.device.type == "cpu" with pytest.raises(ValueError, match="one row per base"): intervened_logits(model, "hook_mid", intervention, base, acts[:-1], model) + + +class _Planted(HookedRootModule): + """The output depends on the hidden state only through one planted direction.""" + + def __init__(self, direction: torch.Tensor): + super().__init__() + self.register_buffer("direction", direction / direction.norm()) + self.scale = nn.Parameter(torch.tensor(4.0)) + self.hook_mid = HookPoint() + self.setup() + self.eval() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + score = self.scale * (self.hook_mid(x) @ self.direction) + return torch.stack([-score, score], dim=-1) + + +def _planted_pairs(n: int = 256, seed: int = 0): + generator = torch.Generator().manual_seed(seed) + direction = torch.randn(D_MODEL, generator=generator) + model = _Planted(direction) + base = torch.randn(n, D_MODEL, generator=generator) + source = torch.randn(n, D_MODEL, generator=generator) + targets = (source @ model.direction > 0).long() # the variable comes from the source + return model, base, source, targets + + +def test_learned_subspace_stays_orthonormal_and_checks_its_rank(): + subspace = LearnedSubspace(D_MODEL, 3) + assert subspace.basis.shape == (D_MODEL, 3) and _orthonormal(subspace.basis.detach()) + with pytest.raises(ValueError, match="rank must be in"): + LearnedSubspace(D_MODEL, 0) + with pytest.raises(ValueError, match="rank must be in"): + LearnedSubspace(D_MODEL, D_MODEL + 1) + + +def test_das_recovers_a_planted_direction(): + model, base, source, targets = _planted_pairs() + result = fit_das(model, "hook_mid", base, source, targets, model, rank=1, epochs=30, lr=5e-2) + cosine = (result.basis[:, 0] @ model.direction).abs() + assert cosine > 0.95 + assert result.losses[-1] < result.losses[0] + + _, test_base, test_source, _ = _planted_pairs(seed=1) + source_acts = collect_activations(model, "hook_mid", test_source, model) + logits = intervened_logits( + model, "hook_mid", SubspaceInterchange(result.subspace), test_base, source_acts, model + ) + # The planted direction differs per seed, so evaluate on the training model's direction + held_out_targets = (test_source @ model.direction > 0).long() + assert (logits.argmax(-1) == held_out_targets).float().mean() > 0.9 + + +def test_splitting_batches_trains_the_same_way(): + model, base, source, targets = _planted_pairs() + kwargs = dict(rank=1, epochs=30, lr=5e-2, batch_size=64) + sizes: dict[str, list[int]] = {"whole": [], "split": []} + + def recording(name): + def task_logits(x): + sizes[name].append(len(x)) + return model(x) + + return task_logits + + whole = fit_das(model, "hook_mid", base, source, targets, recording("whole"), **kwargs) + split = fit_das( + model, "hook_mid", base, source, targets, recording("split"), micro_batch_size=16, **kwargs + ) + # Only a split run keeps every forward pass to the micro-batch size + assert max(sizes["whole"]) == 64 and max(sizes["split"]) == 16 + # Gradients agree up to rounding; Adam's sign-like first steps can amplify that slightly + assert abs(whole.losses[0] - split.losses[0]) < 1e-2 + for result in (whole, split): + assert (result.basis[:, 0] @ model.direction).abs() > 0.95 + + +def test_das_leaves_the_callers_random_stream_alone(): + model, base, source, targets = _planted_pairs(n=32) + torch.manual_seed(123) + expected = torch.rand(3) + torch.manual_seed(123) + first = fit_das(model, "hook_mid", base, source, targets, model, rank=1, epochs=1, seed=5) + assert torch.equal(torch.rand(3), expected) + again = fit_das(model, "hook_mid", base, source, targets, model, rank=1, epochs=1, seed=5) + assert torch.equal(first.basis, again.basis) # the seed still fixes the result + + +def test_das_never_reseeds_a_device_generator(monkeypatch): + def reseeded(*args, **kwargs): + raise AssertionError("fit_das reseeded a global or device generator") + + for module, name in ( + (torch, "manual_seed"), + (torch.random, "manual_seed"), + (torch.cuda, "manual_seed"), + (torch.cuda, "manual_seed_all"), + (torch.mps, "manual_seed"), + ): + monkeypatch.setattr(module, name, reseeded) + model, base, source, targets = _planted_pairs(n=32) + fit_das(model, "hook_mid", base, source, targets, model, rank=1, epochs=1, seed=5) + + +def test_the_subspace_defaults_to_the_models_device(): + assert _model_device(nn.Linear(2, 2, device="meta")) == torch.device("meta") + assert _model_device(_Pattern()) == torch.device("cpu") # no parameters + model, base, source, targets = _planted_pairs(n=32) + result = fit_das(model, "hook_mid", base, source, targets, model, rank=1, epochs=1) + assert result.basis.device == model.scale.device + + +def test_models_in_training_mode_are_rejected(): + model, base, source, targets = _planted_pairs(n=8) + model.train() + acts = torch.randn(8, D_MODEL) + full = SubspaceInterchange(torch.eye(D_MODEL)) + with pytest.raises(ValueError, match="collect_activations requires .* evaluation mode"): + collect_activations(model, "hook_mid", source, model) + with pytest.raises(ValueError, match="intervened_logits requires .* evaluation mode"): + intervened_logits(model, "hook_mid", full, base, acts, model) + with pytest.raises(ValueError, match="fit_das requires .* evaluation mode"): + fit_das(model, "hook_mid", base, source, targets, model, rank=1) + + +def test_das_leaves_the_model_untouched(): + model, base, source, targets = _planted_pairs(n=32) + before = model.scale.detach().clone() + fit_das(model, "hook_mid", base, source, targets, model, rank=1, epochs=2) + assert model.scale.requires_grad and torch.equal(model.scale.detach(), before) + with pytest.raises(ValueError, match="same length"): + fit_das(model, "hook_mid", base, source[:-1], targets, model, rank=1) diff --git a/transformer_lens/tools/analysis/__init__.py b/transformer_lens/tools/analysis/__init__.py index c624c13b5..db9b450bf 100644 --- a/transformer_lens/tools/analysis/__init__.py +++ b/transformer_lens/tools/analysis/__init__.py @@ -19,8 +19,9 @@ analysis. - interchange: Interchange interventions on linear subspaces: set a base's activation inside a subspace to its value on a source input, with - difference-of-means, probe and random bases, a norm-matched placebo, and - batched, gradient-free evaluation. + difference-of-means, probe and random bases, a norm-matched placebo, + batched gradient-free evaluation, and distributed alignment search (DAS) + to learn the subspace. - jacobian_lens: The Jacobian lens (J-lens) -- per-layer causal transport to the output vocabulary basis, with loading of published lens artifacts, native fitting, readouts, interventions, J-space sparse decomposition, and @@ -78,9 +79,12 @@ get_act_patch_direct_path_all_sources, ) from transformer_lens.tools.analysis.interchange import ( + DASResult, + LearnedSubspace, NormMatchedPlacebo, SubspaceInterchange, collect_activations, + fit_das, intervened_logits, mean_difference_basis, probe_basis, @@ -166,6 +170,7 @@ "BackwardLensResult", "BenchmarkCorpus", "CoordinatePatch", + "DASResult", "DegenerateDirectionError", "DirectLogitAttribution", "EdgeAttributionConfig", @@ -181,6 +186,7 @@ "JSpaceVarianceProfile", "JacobianLens", "JacobianLensReadout", + "LearnedSubspace", "LinearGradientFactors", "LogitLensResult", "LogitReadout", @@ -210,6 +216,7 @@ "direct_logit_attribution", "estimate_occupancy", "faithfulness", + "fit_das", "fit_sparse_probe", "get_act_patch_direct_path", "get_act_patch_direct_path_all_sources", diff --git a/transformer_lens/tools/analysis/interchange.py b/transformer_lens/tools/analysis/interchange.py index bd7cc2bbb..2182d3093 100644 --- a/transformer_lens/tools/analysis/interchange.py +++ b/transformer_lens/tools/analysis/interchange.py @@ -13,7 +13,10 @@ The module provides the intervention, fixed-direction bases to intervene on (the difference of class means, a linear probe's weights, a random basis), and a norm-matched random-subspace placebo. -An interchange accuracy is not evidence on its own: compare it with the placebo and the baselines. +:func:`fit_das` learns the basis instead, by distributed alignment search (DAS; Geiger et al., +2024): it trains an orthonormal basis so that the interchange produces the high-level model's +predicted outputs. An interchange accuracy is not evidence on its own: compare it with the placebo +and the baselines. Works with any model exposing the hook API (``TransformerBridge`` or a ``HookedRootModule``). Activations at the hook must be ``[batch, d_model]``, or ``[batch, pos, d_model]`` with one @@ -23,12 +26,16 @@ class means, a linear probe's weights, a random basis), and a norm-matched rando from __future__ import annotations from dataclasses import dataclass -from typing import Any, Callable, Hashable, Iterator, Optional, Sequence, Union +from typing import Any, Callable, Hashable, Iterator, Optional, Sequence, Union, cast import numpy as np import torch +import torch.nn.functional as F from jaxtyping import Float, Int +from torch import nn +from torch.nn.utils.parametrizations import orthogonal +from transformer_lens.tools.analysis._model_state import require_eval_mode from transformer_lens.tools.analysis.projection_kernel import orthonormal_subspace TaskLogits = Callable[[torch.Tensor], Float[torch.Tensor, "batch n_classes"]] @@ -99,20 +106,47 @@ def probe_basis( return orthonormal_subspace((w - w.mean(0)).T, rank=w.shape[0] - 1).basis +class LearnedSubspace(nn.Module): + """A trainable rank-``rank`` subspace of ``d_model`` whose basis stays orthonormal. + + Args: + d_model: Dimension of the activation. + rank: Number of basis directions, in ``[1, d_model]``. + """ + + def __init__(self, d_model: int, rank: int): + super().__init__() + if not 0 < rank <= d_model: + raise ValueError(f"rank must be in [1, {d_model}], got {rank}") + self.proj = orthogonal(nn.Linear(rank, d_model, bias=False)) + + @property + def basis(self) -> Float[torch.Tensor, "d_model rank"]: + return cast(torch.Tensor, self.proj.weight) + + +Basis = Union[torch.Tensor, LearnedSubspace] + + +def _basis(basis: Basis, like: torch.Tensor) -> torch.Tensor: + b = basis.basis if isinstance(basis, LearnedSubspace) else basis + return b.to(like.device, like.dtype) + + @dataclass(frozen=True) class SubspaceInterchange: """Set the base's coordinates inside ``basis`` to the source's. Attributes: - basis: Orthonormal basis, ``[d_model, rank]``. + basis: Orthonormal basis, ``[d_model, rank]``, or a :class:`LearnedSubspace`. strength: 1 interchanges the coordinates; other values scale the move, 0 leaves the base. """ - basis: torch.Tensor + basis: Basis strength: float = 1.0 def __call__(self, base: torch.Tensor, source: torch.Tensor) -> torch.Tensor: - B = self.basis.to(base.device, base.dtype) + B = _basis(self.basis, base) return base + self.strength * ((source - base) @ B) @ B.T @@ -202,9 +236,11 @@ def collect_activations( Activations, ``[batch, d_model]``. Raises: - ValueError: If the activation does not match the positions given (or not given). + ValueError: If any part of the model is in training mode, or the activation does not match + the positions given (or not given). RuntimeError: If the hook does not fire exactly once per forward pass. """ + require_eval_mode(model, operation="collect_activations") out = [] for idx in _chunks(len(inputs), batch_size): store: list[torch.Tensor] = [] @@ -302,7 +338,11 @@ def intervened_logits( Returns: Task logits under the intervention, ``[batch, n_classes]``. + + Raises: + ValueError: If any part of the model is in training mode, or the inputs do not pair up. """ + require_eval_mode(model, operation="intervened_logits") if len(source_activations) != len(base_inputs): raise ValueError("source_activations must have one row per base input") out = [] @@ -318,3 +358,139 @@ def intervened_logits( ) out.append(logits.detach().cpu()) return torch.cat(out) + + +def _model_device(model: Any) -> torch.device: + parameter = next(iter(model.parameters()), None) + return parameter.device if parameter is not None else torch.device("cpu") + + +@dataclass(frozen=True) +class DASResult: + """A subspace learned by :func:`fit_das`. + + Attributes: + subspace: The learned subspace; pass it to :class:`SubspaceInterchange`. + losses: Mean training loss per epoch. + hook_name: The hook point it was trained at. + rank: Its rank. + """ + + subspace: LearnedSubspace + losses: tuple[float, ...] + hook_name: str + rank: int + + @property + def basis(self) -> Float[torch.Tensor, "d_model rank"]: + return self.subspace.basis.detach() + + +def fit_das( + model: Any, + hook_name: str, + base_inputs: torch.Tensor, + source_inputs: torch.Tensor, + targets: Int[torch.Tensor, "batch"], + task_logits: TaskLogits, + *, + rank: int, + base_positions: Optional[Int[torch.Tensor, "batch"]] = None, + source_positions: Optional[Int[torch.Tensor, "batch"]] = None, + epochs: int = 20, + batch_size: int = 128, + lr: float = 1e-2, + strength: float = 1.0, + seed: int = 0, + device: Optional[Union[str, torch.device]] = None, + micro_batch_size: Optional[int] = None, +) -> DASResult: + """Learn the subspace whose interchange makes the model output ``targets`` (DAS). + + ``targets[i]`` is the class the high-level causal model predicts for base ``i`` with the + aligned variable set to its value in source ``i``. The basis is trained with cross-entropy + on the task logits under the interchange. Model parameters are frozen during the call and + their ``requires_grad`` flags restored afterwards. + + Train on pairs where several variables differ between base and source, labelled by the + high-level model's interchange, as well as on pairs that differ only in the aligned variable: + otherwise the subspace can learn whatever the pairs happen to move together. + + Args: + model: A model exposing the hook API. + hook_name: The hook point to align. + base_inputs: Base inputs, one row per pair. + source_inputs: Source inputs, one row per pair. + targets: The high-level model's output class for each pair under the interchange. + task_logits: Runs the model on a batch of inputs and returns task logits. + rank: Dimension of the subspace. + base_positions: For ``[batch, pos, d_model]`` activations, the base position per pair. + source_positions: Likewise for the sources. + epochs: Passes over the pairs. + batch_size: Pairs per optimisation step. + lr: Adam learning rate. + strength: Interchange strength during training. + seed: Seed for the basis initialisation and the batch order. + device: Device for the subspace's parameters; by default the device of the model's + parameters, or the CPU for a model without any. + micro_batch_size: Split each batch into parts this size and accumulate gradients: the same + gradients, up to floating-point rounding, in less memory. + + Returns: + The learned subspace with its training losses. + + Raises: + ValueError: If any part of the model is in training mode, or the inputs and targets differ + in length. + """ + require_eval_mode(model, operation="fit_das") + if not (len(base_inputs) == len(source_inputs) == len(targets)): + raise ValueError("base_inputs, source_inputs and targets must have the same length") + rng = np.random.default_rng(seed) + micro = micro_batch_size or batch_size + source_acts = collect_activations( + model, hook_name, source_inputs, task_logits, positions=source_positions, batch_size=micro + ) + # The subspace is built on the CPU, so seed only the CPU generator, inside a fork that restores + # it; torch.manual_seed would also reseed every CUDA, MPS and XPU generator + with torch.random.fork_rng(devices=[]): + torch.default_generator.manual_seed(seed) + subspace = LearnedSubspace(source_acts.shape[1], rank) + subspace = subspace.to(device if device is not None else _model_device(model)) + step = SubspaceInterchange(subspace, strength) + optimizer = torch.optim.Adam(subspace.parameters(), lr=lr) + labels = torch.as_tensor(targets, dtype=torch.long) + + flags = [p.requires_grad for p in model.parameters()] + model.requires_grad_(False) + losses = [] + try: + for _ in range(epochs): + order = torch.as_tensor(rng.permutation(len(labels))) + total = 0.0 + for chunk in _chunks(len(order), batch_size): + batch = order[chunk] + optimizer.zero_grad() + for part in _chunks(len(batch), micro): + idx = batch[part] + logits = run_with_intervention( + model, + hook_name, + step, + base_inputs[idx], + source_acts[idx], + task_logits, + positions=_at(base_positions, idx), + ) + # Summed, then divided by the whole batch, so split batches give equal gradients + loss = F.cross_entropy( + logits, labels[idx].to(logits.device), reduction="sum" + ) / len(batch) + loss.backward() + total += loss.item() * len(batch) + optimizer.step() + losses.append(total / len(order)) + finally: + for p, flag in zip(model.parameters(), flags): + p.requires_grad_(flag) + return DASResult(subspace.eval(), tuple(losses), hook_name, rank)