transformer_lens.tools.analysis.interchange module¶

Interchange interventions on linear subspaces.

An interchange intervention runs the model on a base input while setting the base’s activation at one hook point, inside a subspace, to its value on a source input:

h' = h_base + strength * (h_source - h_base) @ B @ B.T

B is an orthonormal [d_model, rank] basis, the convention of orthonormal_subspace(). The base keeps its own coordinates outside the subspace. If the model then produces the output a high-level causal model predicts for the base with one variable taken from the source, the subspace carries that variable at that hook point (Geiger et al., 2021).

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. neuron_basis() turns a set of MLP neurons, such as the features a sparse probe selects, into a basis, and random_neuron_basis() draws random neurons as its placebo. 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. DAS fits arbitrary label mappings and finds effects in randomly initialised models (Arora et al., 2024), so the module also provides the controls to train the same way: arbitrary_targets() for an arbitrary-label control, and random_init_bridge() or reinitialized() for an untrained copy of the model. Compare a learned subspace with both, and 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 position index per example.

class transformer_lens.tools.analysis.interchange.DASResult(subspace: LearnedSubspace, losses: tuple[float, ...], hook_name: str, rank: int)¶

Bases: object

A subspace learned by fit_das().

subspace¶

The learned subspace; pass it to SubspaceInterchange.

Type:: transformer_lens.tools.analysis.interchange.LearnedSubspace

losses¶

Mean training loss per epoch.

Type:: tuple[float, …]

hook_name¶

The hook point it was trained at.

Type:: str

property basis: Float[Tensor, 'd_model rank']¶

hook_name: str¶

losses: tuple[float, ...]¶

rank: int¶

subspace: LearnedSubspace¶

transformer_lens.tools.analysis.interchange.Intervention¶

Maps (base activations, source activations), both [batch, d_model], to new base activations.

alias of Callable[[Tensor, Tensor], Tensor]

class transformer_lens.tools.analysis.interchange.LearnedSubspace(d_model: int, rank: int)¶

Bases: Module

A trainable rank-rank subspace of d_model whose basis stays orthonormal.

Parameters:: rank – Number of basis directions, in [1, d_model].

property basis: Float[Tensor, 'd_model rank']¶

transformer_lens.tools.analysis.interchange.Neurons¶

a list, a NumPy array, or a tensor such as a sparse probe’s selected_features.

Type:: Integer neuron indices

alias of Sequence[int] | ndarray | Tensor

class transformer_lens.tools.analysis.interchange.NormMatchedPlacebo(reference: SubspaceInterchange, basis: Tensor)¶

Bases: object

Move each base exactly as far as reference does, but inside a random subspace.

An effect that the placebo also produces comes from displacing the activation by that much, not from the aligned subspace.

reference¶

The intervention whose displacement is matched.

Type:: transformer_lens.tools.analysis.interchange.SubspaceInterchange

basis¶

Orthonormal basis of the placebo subspace, usually from random_basis() with the reference’s rank, or from random_neuron_basis() for a reference built by neuron_basis().

Type:: torch.Tensor

basis: Tensor¶

reference: SubspaceInterchange¶

class transformer_lens.tools.analysis.interchange.SubspaceInterchange(basis: Tensor | LearnedSubspace, strength: float = 1.0)¶

Bases: object

Set the base’s coordinates inside basis to the source’s.

basis¶

Orthonormal basis, [d_model, rank], or a LearnedSubspace.

Type:: torch.Tensor | transformer_lens.tools.analysis.interchange.LearnedSubspace

strength¶

1 interchanges the coordinates; other values scale the move, 0 leaves the base.

Type:: float

basis: Tensor | LearnedSubspace¶

strength: float = 1.0¶

transformer_lens.tools.analysis.interchange.TaskLogits¶

Runs the model on a batch of inputs in one forward pass and returns task logits, e.g. lambda tokens: model(tokens)[:, -1, answer_token_ids].

alias of Callable[[Tensor], Float[Tensor, 'batch n_classes']]

transformer_lens.tools.analysis.interchange.arbitrary_targets(keys: Sequence[Hashable], n_classes: int, seed: int) → Int[Tensor, 'batch']¶

A fixed pseudo-random class per key: targets for the arbitrary-label control.

Train fit_das() on these instead of the high-level model’s targets, with the same pairs, rank and budget. The interchange accuracy it reaches shows how much accuracy the method’s expressivity alone produces; a learned subspace should clearly exceed it.

Parameters:: keys – One key per pair identifying its source, e.g. the source text.
Returns:: A class in [0, n_classes) per key; equal keys always get equal classes.

transformer_lens.tools.analysis.interchange.collect_activations(model: Any, hook_name: str, inputs: Tensor, task_logits: Callable[[Tensor], Float[Tensor, 'batch n_classes']], *, positions: Int[Tensor, 'batch'] | None = None, batch_size: int = 256) → Float[Tensor, 'batch d_model']¶

Activations at hook_name, at one position per example if given, returned on the CPU.

Parameters:

model – A model exposing the hook API.
inputs – Model inputs, one row per example.
task_logits – Runs the model on a batch of inputs.
positions – For [batch, pos, d_model] activations, the position to read per example.
batch_size – Rows per forward pass.

Returns:

Activations, [batch, d_model].

Raises:

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.

transformer_lens.tools.analysis.interchange.fit_das(model: Any, hook_name: str, base_inputs: Tensor, source_inputs: Tensor, targets: Int[Tensor, 'batch'], task_logits: Callable[[Tensor], Float[Tensor, 'batch n_classes']], *, rank: int, base_positions: Int[Tensor, 'batch'] | None = None, source_positions: Int[Tensor, 'batch'] | None = None, epochs: int = 20, batch_size: int = 128, lr: float = 0.01, strength: float = 1.0, seed: int = 0, device: str | device | None = None, micro_batch_size: int | None = 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.

Parameters:

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.

transformer_lens.tools.analysis.interchange.intervened_logits(model: Any, hook_name: str, intervention: Callable[[Tensor, Tensor], Tensor], base_inputs: Tensor, source_activations: Float[Tensor, 'batch d_model'], task_logits: Callable[[Tensor], Float[Tensor, 'batch n_classes']], *, positions: Int[Tensor, 'batch'] | None = None, batch_size: int = 256) → Float[Tensor, 'batch n_classes']¶

Task logits under intervention, in batches without gradients, returned on the CPU.

The evaluation counterpart of run_with_intervention(), with the same arguments and checks: memory is bounded by batch_size rather than by the number of pairs.

Parameters:

model – A model exposing the hook API.
intervention – Maps base and source activations to new base activations.
source_activations – Activations from the source inputs at the same hook point, one row per base.
task_logits – Runs the model on a batch of inputs.
positions – For [batch, pos, d_model] activations, the position to intervene on per base.
batch_size – Rows per forward pass.

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.

transformer_lens.tools.analysis.interchange.mean_difference_basis(activations: Float[Tensor, 'batch d_model'], labels: Sequence[Hashable]) → Float[Tensor, 'd_model rank']¶

The span of the centred class means: for two classes, the difference-in-means direction.

Parameters:

activations – Activations, one row per example.
labels – The variable’s value for each example; at least two distinct values.

Returns:

Orthonormal basis of rank n_classes - 1, [d_model, n_classes - 1].

Raises:

ValueError – If fewer than two classes are present.

transformer_lens.tools.analysis.interchange.neuron_basis(d_mlp: int, neurons: Sequence[int] | ndarray | Tensor) → Float[Tensor, 'd_mlp k']¶

The coordinate basis of a set of MLP neurons, e.g. the features a sparse probe selects.

Interchanging it sets exactly those neurons to the source’s values, at a neuron hook such as blocks.N.mlp.hook_post. Compare it with random_neuron_basis() of the same size: a random direction is the wrong placebo for a claim about neurons, because it spreads over every neuron.

Parameters:

d_mlp – Number of neurons at the hook.
neurons – Integer indices of the neurons, each in [0, d_mlp) and none repeated.

Returns:

Orthonormal basis of unit vectors, [d_mlp, k], in the order given.

Raises:

ValueError – If no neuron is given, or an index is not an integer, out of range or repeated.

transformer_lens.tools.analysis.interchange.probe_basis(weights: Float[Tensor, 'n_rows d_model']) → Float[Tensor, 'd_model rank']¶

The directions of a linear probe’s weights, in activation space.

Undo any input scaling first: for a probe fit on standardised inputs, divide each weight by its feature’s standard deviation.

Parameters:: weights – Probe weights; one row for a binary probe, one row per class otherwise.
Returns:: the weight direction for one row, the span of the centred rows (rank n_rows - 1) for several.
Return type:: Orthonormal basis

transformer_lens.tools.analysis.interchange.random_basis(d_model: int, rank: int, seed: int) → Float[Tensor, 'd_model rank']¶

A seeded random orthonormal basis: the baseline any aligned basis should beat.

Parameters:: seed – Seed for the Gaussian matrix the basis is drawn from.
Returns:: Orthonormal basis, [d_model, rank].

transformer_lens.tools.analysis.interchange.random_init_bridge(model_name: str, *, seed: int, dtype: dtype = torch.float32, device: str | device | None = None, **boot_kwargs: Any) → Any¶

A TransformerBridge with the architecture’s HuggingFace initialisation and no training.

The random-init control: train fit_das() on it with the same pairs, rank and budget as on the trained model. Effects DAS finds here come from the method and the architecture, not from anything the model learned. Built with load_weights=False, so the model class and prepared configuration match the trained bridge.

Parameters:

model_name – Model name or alias, as for TransformerBridge.boot_transformers.
device – Device for the bridge; by default the device boot_transformers would choose.
**boot_kwargs – Passed on to TransformerBridge.boot_transformers, e.g. revision.

Returns:

The untrained TransformerBridge, in evaluation mode.

transformer_lens.tools.analysis.interchange.random_neuron_basis(d_mlp: int, k: int, seed: int, exclude: Sequence[int] | ndarray | Tensor | None = None) → Float[Tensor, 'd_mlp k']¶

The coordinate basis of k neurons drawn at random: the placebo for neuron_basis().

Parameters:

d_mlp – Number of neurons at the hook.
seed – Seed for the draw; the caller’s random stream is untouched.
exclude – Neurons the draw must avoid, usually the ones the placebo is compared with.

Returns:

Orthonormal basis of unit vectors, [d_mlp, k].

Raises:

ValueError – If fewer than k neurons remain to draw from, or exclude holds an invalid index.

transformer_lens.tools.analysis.interchange.reinitialized(model: Module, seed: int) → Module¶

A copy of model with freshly initialised parameters: the random-init control.

For models built from standard nn layers, such as a HookedRootModule: every module with reset_parameters is reset, so every parameter must belong to one. Buffers, such as fixed architectural matrices, are kept. For HuggingFace models use random_init_bridge(), which applies the architecture’s own initialisation scheme.

Parameters:: model – The trained model; it is not modified.
Returns:: The untrained copy, on the device of the model’s parameters, in evaluation mode.
Raises:: ValueError – If a parameter belongs to no module with reset_parameters, so it would keep its trained value.

transformer_lens.tools.analysis.interchange.run_with_intervention(model: Any, hook_name: str, intervention: Callable[[Tensor, Tensor], Tensor], base_inputs: Tensor, source_activations: Float[Tensor, 'batch d_model'], task_logits: Callable[[Tensor], Float[Tensor, 'batch n_classes']], *, positions: Int[Tensor, 'batch'] | None = None) → Float[Tensor, 'batch n_classes']¶

Task logits for base_inputs with intervention applied at hook_name, in one pass.

Differentiable with respect to any parameters of intervention, for training an intervention. To evaluate one over many pairs in bounded memory, use intervened_logits().

Parameters:

model – A model exposing the hook API.
intervention – Maps base and source activations to new base activations, e.g. a SubspaceInterchange.
source_activations – Activations from the source inputs at the same hook point, one row per base, e.g. from collect_activations().
task_logits – Runs the model on a batch of inputs.
positions – For [batch, pos, d_model] activations, the position to intervene on per base.

Returns:

Task logits under the intervention.

Raises:

ValueError – If the activation does not match the positions given (or not given).
RuntimeError – If the hook does not fire exactly once in the forward pass.