Examples¶
Every operation has an experiment script and a YAML file of configs. A script runs one entry of the
YAML file beside it, picked by --config-id, so save the two side by side:
python lens.py # the file's default entry
python lens.py --config-id tuned_lens # another entry
python patching.py --config-id head_sweep --dry-run # plan it, without running the model
python lens.py --help # every option
A run writes its results into runs/<config-id>/; open the .html page there in a browser. Every
script also takes --model-name (default gpt2), --output-dir (default runs) and --config-path,
and all but custom_lens.py take --cache-dir to keep the model’s activations on disk between runs.
The prompts are options too, and the YAML files explain every entry in their comments.
Experiment scripts shows how a script turns its config into options.
Lens¶
lens.py¶
"""Run one entry of lens.yaml on a prompt.
uv run python examples/lens.py --config-id tuned_lens
uv run python examples/lens.py --config-id parameter_lens --dry-run
Results land in runs/<config-id>/: open result.html in a browser.
"""
from dataclasses import dataclass, field
from pathlib import Path
from transformer_lens.model_bridge import TransformerBridge
from entoscope import CachedModel, ConfigFactory, Lens, LensConfig, LensInput, experiment
from entoscope.lens import LensMethodFactory
CONFIGS = Path(__file__).with_name("lens.yaml")
@dataclass
class Config:
config_id: str = field(default="logit_lens", metadata={"help": "the entry of the config file to run"})
config_path: str = field(default=str(CONFIGS), metadata={"help": "a YAML file"})
prompt: str = "The capital of France is"
model_name: str = "gpt2"
output_dir: str = field(default="runs", metadata={"help": "results go into <output-dir>/<config-id>"})
cache_dir: str | None = field(default=None, metadata={"help": "keep the model's activations on disk here"})
dry_run: bool = field(default=False, metadata={"help": "log what the run would do, without running it"})
@experiment
def main(config: Config):
"""Run one lens config on a prompt."""
model = CachedModel(TransformerBridge.boot_transformers(config.model_name), config.cache_dir)
model.eval()
lens_config = ConfigFactory[LensConfig](config.config_path).create(config.config_id)
# A lens that reads the weights rather than the activations takes no prompt.
reads_activations = LensMethodFactory.get(lens_config.method.type).reads_activations
lens_input = LensInput(lens_config, config.prompt if reads_activations else None)
save_dir = f"{config.output_dir}/{config.config_id}"
if config.dry_run:
Lens(model).plan(lens_input, save_dir)
else:
Lens(model).apply(lens_input, save_dir)
lens.yaml¶
# Lens configs for lens.py: `uv run python examples/lens.py --config-id <entry>`.
#
# Every entry is a LensConfig. `method` names the lens - a mapping with `type` and the lens's own
# fields, or just the type for its defaults - `top_k` how many tokens a row keeps, and `source`
# where in the model the lens reads.
# The residual stream after every layer at every position: the classic logit-lens grid, drawn with
# position across and layer up.
logit_lens:
method: logit_lens
top_k: 5
source:
activation: resid.post
location: [layer, position]
# MLP-hidden activations are projected through the layer's W_out, so a row reads what the indexed
# neurons write into the residual stream. Keep the neuron axis filtered: unfiltered it is d_mlp rows
# per layer and position.
neuron_logit_lens:
method: logit_lens
top_k: 5
source:
activation: mlp.post
location: [layer, neuron, position]
filter:
- neuron: 373
position: -1
# attn.z is projected through the head's W_O, so every row reads what one head writes into the
# residual stream - here at the last position.
head_logit_lens:
method: logit_lens
top_k: 5
source:
activation: attn.z
location: [layer, head, position]
filter:
- position: -1
# The tuned lens reads the same activations through each layer's trained translator, which maps the
# residual stream onto the last layer's basis before unembedding it. The translators are downloaded
# from the tuned-lens hub space for the model's own name unless `translators` names them: a resource
# id there, or a local directory holding config.json and params.pt.
tuned_lens:
method: tuned_lens
top_k: 5
source:
activation: resid.post
location: [layer, position]
# The parameter lens reads the weights instead of the activations, so it runs without a prompt.
# `mlp.post` names the value vector a neuron writes into the residual stream (a row of W_out) - here
# neuron 373's, once per layer. `position` is not a location of this lens: weights have no position.
parameter_lens:
method: parameter_lens
top_k: 5
source:
activation: mlp.post
location: [layer, neuron]
filter:
- neuron: 373
# Attention heads are read one subhead at a time: every dimension of the last layer's head 7, through
# the head's W_O. `projection: embedding` reads the vectors against the embedding matrix, as Dar et
# al. do, and a negative `scale` would read them against their own direction.
subhead_parameter_lens:
method:
type: parameter_lens
projection: embedding
top_k: 5
source:
activation: attn.z
location: [layer, head, neuron]
filter:
- layer: -1
head: 7
# Explicit coordinates name exact points instead of a grid; each may carry its own activation. A
# one-dimensional index like this one is drawn as a bar chart.
explicit_coordinates:
method: logit_lens
top_k: 5
source:
coordinates:
- {activation: resid.mid, layer: -1, position: -1}
- {activation: attn.out, layer: -1, position: -1}
- {activation: mlp.out, layer: -1, position: -1}
- {activation: resid.post, layer: -1, position: -1}
Activation patching¶
patching.py¶
"""Run one entry of patching.yaml on a pair of prompts.
uv run python examples/patching.py --config-id head_sweep
uv run python examples/patching.py --config-id zero_ablation --source-prompt ""
uv run python examples/patching.py --config-id attribution --dry-run
Results land in runs/<config-id>/: open result.html in a browser. Prompts of your own come with
answers of their own, named for the metrics as --tokens correct:" Rome" distractor:" Madrid" other:" Paris".
"""
from dataclasses import dataclass, field
from pathlib import Path
from transformer_lens.model_bridge import TransformerBridge
from entoscope import CachedModel, ConfigFactory, Patcher, PatchingConfig, PatchingInput, experiment
CONFIGS = Path(__file__).with_name("patching.yaml")
@dataclass
class Config:
config_id: str = field(default="residual_sweep", metadata={"help": "the entry of the config file to run"})
config_path: str = field(default=str(CONFIGS), metadata={"help": "a YAML file"})
source_prompt: str = field(
default="The capital of France is", metadata={"help": "the run patched from; empty for none"}
)
target_prompt: str = field(default="The capital of Spain is", metadata={"help": "the run patched into"})
tokens: dict[str, str] = field(
default_factory=lambda: {"correct": " Paris", "distractor": " Madrid", "other": " Rome"},
metadata={
"help": "the tokens the metrics read, as name:token, by the names the configs give them; "
"they replace the defaults, so name every token the config reads"
},
)
batch_size: int = field(default=1, metadata={"help": "patched runs per forward pass ('exact' only)"})
model_name: str = "gpt2"
output_dir: str = field(default="runs", metadata={"help": "results go into <output-dir>/<config-id>"})
cache_dir: str | None = field(default=None, metadata={"help": "keep the model's activations on disk here"})
dry_run: bool = field(default=False, metadata={"help": "log what the run would do, without running it"})
@experiment
def main(config: Config):
"""Run one patching config on a source and a target prompt."""
model = CachedModel(TransformerBridge.boot_transformers(config.model_name), config.cache_dir)
model.eval()
patching_input = PatchingInput(
ConfigFactory[PatchingConfig](config.config_path).create(config.config_id),
# A source that writes a value or noise reads no source run, so it may go without a prompt.
config.source_prompt or None,
config.target_prompt,
tokens=config.tokens,
)
patcher = Patcher(model, batch_size=config.batch_size)
save_dir = f"{config.output_dir}/{config.config_id}"
if config.dry_run:
patcher.plan(patching_input, save_dir)
else:
patcher.apply(patching_input, save_dir)
patching.yaml¶
# Patching configs for patching.py: `uv run python examples/patching.py --config-id <entry>`.
#
# A patching run writes activations of the source run into the target run, one patched run per row
# of the index, and reads metrics off every patched run. The script's defaults patch "The capital of
# France is" (source) into "The capital of Spain is" (target), and give the tokens the metrics read
# their names: 'correct' is " Paris", 'distractor' " Madrid" and 'other' " Rome".
#
# `metric` names each scalar read off a patched run, written to the result column '<name>_metric'. A
# metric gives its `type`, the `tokens` it reads, by those names, and the comparison it is reported under.
# 'logit_diff' and 'logprob_diff' compare their first token with their second, 'probability' and
# 'log_prob' read a single token, and 'kl_divergence' scores the whole distribution against the source
# run's and reads none. The names are roles rather than the tokens themselves, so an entry runs on
# other prompts once `--tokens` names their answers. Every metric comes off the same forward pass, so
# asking for several costs nothing, and each column is drawn as its own plot.
#
# The comparison says what a patched run is reported against. 'absolute' is the metric of the patched
# run itself; 'indirect_effect' subtracts the unpatched target run, which is activation patching
# stated as a causal quantity; 'normalized' divides that by the source run's own distance from the
# target, putting 0 at 'the patch changed nothing' and 1 at 'the patch reproduced the source run'.
# The residual stream before every layer at every position: where the information that tells the
# two prompts apart lives, layer by layer.
residual_sweep:
metric:
logit_diff:
type: logit_diff
tokens: [correct, distractor]
comparison: normalized
kl_divergence:
type: kl_divergence
comparison: indirect_effect
method: exact
patch:
- source: &residual
activation: resid.pre
location: [layer, position]
type: replace
target: *residual
# Several metrics of one type, each on tokens of its own: the source run's answer against the target
# run's, and against a third city neither prompt is about, whose own chance is read as well. Each
# metric's columns carry its name.
competing_answers:
metric:
against_target:
type: logit_diff
tokens: [correct, distractor]
comparison: normalized
against_other:
type: logit_diff
tokens: [correct, other]
comparison: indirect_effect
other:
type: probability
tokens: [other]
comparison: indirect_effect
method: exact
patch:
- source: *residual
target: *residual
# Every head's output, at every position at once. A comparison sits on its metric because the right
# one differs per metric: both columns come off the one run.
head_sweep:
metric:
logit_diff:
type: logit_diff
tokens: [correct, distractor]
comparison: normalized
probability:
type: probability
tokens: [correct]
comparison: indirect_effect
method: exact
patch:
- source: &heads
activation: attn.z
location: [layer, head]
type: replace
target: *heads
# One head's attention pattern, entry by entry: `dest_pos` is the token attending and `src_pos` the
# token it attends to. A patch writes the source run's attention weight alone, the head still mixing
# the target run's values. This sweeps the last token's weight on every key; without the `dest_pos`
# filter it sweeps the whole pattern, drawn as the pattern itself (src_pos across, dest_pos up), but the
# entries with src_pos past dest_pos are masked to 0 in both runs, so about half its runs score nothing.
pattern_entries:
metric:
logit_diff:
type: logit_diff
tokens: [correct, distractor]
comparison: normalized
method: exact
patch:
- source: &pattern
activation: attn.pattern
location: [layer, head, dest_pos, src_pos]
filter:
- layer: 9
head: 8
dest_pos: -1
type: replace
target: *pattern
# Zero ablation: the source names no location, only a value (0.0 by default), so nothing is read from
# a source run - pass `--source-prompt ""` to run without one. Reported as an indirect effect this is
# 'm(x | do(a = 0)) - m(x)'. 'normalized' and 'kl_divergence' score against the source run, so neither
# is available without one.
zero_ablation:
metric:
logit_diff:
type: logit_diff
tokens: [correct, distractor]
comparison: indirect_effect
method: exact
patch:
- source:
type: value_replace
target:
activation: mlp.out
location: [layer]
type: replace
# Causal-tracing corruption: every token's embedding gets Gaussian noise at three times the
# embedding's own spread, one position at a time. `embed` sits outside the blocks, so it has no layer.
noise_tracing:
metric:
logit_diff:
type: logit_diff
tokens: [correct, distractor]
comparison: indirect_effect
method: exact
patch:
- source:
activation: embed
location: [position]
type: gaussian_noise
target:
activation: embed
location: [position]
type: replace
# Attribution patching: the first-order expansion of the effect the residual sweep measures. One
# forward and one backward pass score every row of the index, so the cost is flat in the size of the
# sweep rather than one run per row - at the price of a linearity assumption a large intervention
# leaves behind.
attribution:
metric:
logit_diff:
type: logit_diff
tokens: [correct, distractor]
comparison: indirect_effect
method: attribution
patch:
- source: *residual
target: *residual
# Integrated gradients: the same estimate with the gradient averaged over `interpolation_steps` points
# on the path from the source activation to the target one, which recovers effects a single gradient
# at the clean point hides. The path runs between the two runs, so this needs a source prompt, and two
# prompts of the same token length.
integrated_gradients:
metric:
logit_diff:
type: logit_diff
tokens: [correct, distractor]
comparison: indirect_effect
method:
type: integrated_gradients
interpolation_steps: 10
patch:
- source: *residual
target: *residual
# Explicit coordinates: source coordinate i is read from the source run and written at target
# coordinate i of the target run, one patched run per pair. Each coordinate may carry its own
# activation instead of the element-level one.
explicit_coordinates:
metric:
logit_diff:
type: logit_diff
tokens: [correct, distractor]
comparison: absolute
method: exact
patch:
- source:
coordinates:
- {activation: resid.post, layer: -2, position: -1}
- {activation: mlp.out, layer: -1, position: -1}
type: replace
target:
coordinates:
- {activation: resid.post, layer: -1, position: -1}
- {activation: mlp.out, layer: -1, position: -1}
type: replace
# Two elements patched together: each patched run applies one row of every element. The element with
# the most rows lays the result out; an element with a single row is applied on every run, and the result
# page names it in its subtitle.
heads_with_fixed_mlp:
metric:
logit_diff:
type: logit_diff
tokens: [correct, distractor]
comparison: normalized
method: exact
patch:
- source: &late_heads
activation: attn.z
location: [layer, head]
filter:
- layer: {range: [8, 11]}
type: replace
target: *late_heads
- source: &first_mlp
activation: mlp.out
location: [layer]
filter:
- layer: 0
type: replace
target: *first_mlp
Probing¶
probe.py¶
"""Run one entry of probe.yaml on positive and negative example prompts.
uv run python examples/probe.py --config-id last_token_by_layer
uv run python examples/probe.py --positive "I loved it" "Great fun" --negative "I hated it" "Dull"
Results land in runs/<config-id>/: the scores in result.csv, the fitted probes in probes.pkl.
"""
from dataclasses import dataclass, field
from pathlib import Path
from transformer_lens.model_bridge import TransformerBridge
from entoscope import CachedModel, ConfigFactory, ProbeConfig, ProbeInput, Prober, experiment
CONFIGS = Path(__file__).with_name("probe.yaml")
POSITIVE = [
"The movie was absolutely wonderful",
"I loved every minute of it",
"What a delightful surprise this was",
"The food was fresh and delicious",
"She gave a brilliant performance",
"This is the best day of my life",
"The staff were kind and helpful",
"I would happily do it again",
]
NEGATIVE = [
"The movie was dreadful and boring",
"I hated every minute of it",
"What a miserable waste of time",
"The food was cold and tasteless",
"She gave a painful performance",
"This is the worst day of my life",
"The staff were rude and careless",
"I would never do it again",
]
@dataclass
class Config:
config_id: str = field(default="residual_by_layer", metadata={"help": "the entry of the config file to run"})
config_path: str = field(default=str(CONFIGS), metadata={"help": "a YAML file"})
positive: list[str] = field(default_factory=lambda: list(POSITIVE), metadata={"help": "prompts labelled 1"})
negative: list[str] = field(default_factory=lambda: list(NEGATIVE), metadata={"help": "prompts labelled 0"})
model_name: str = "gpt2"
output_dir: str = field(default="runs", metadata={"help": "results go into <output-dir>/<config-id>"})
cache_dir: str | None = field(default=None, metadata={"help": "keep the model's activations on disk here"})
dry_run: bool = field(default=False, metadata={"help": "log what the run would do, without running it"})
@experiment
def main(config: Config):
"""Train and score one probe per location of a probe config."""
model = CachedModel(TransformerBridge.boot_transformers(config.model_name), config.cache_dir)
model.eval()
probe_input = ProbeInput(
ConfigFactory[ProbeConfig](config.config_path).create(config.config_id), config.positive, config.negative
)
save_dir = f"{config.output_dir}/{config.config_id}"
if config.dry_run:
Prober(model).plan(probe_input, save_dir)
else:
Prober(model).apply(probe_input, save_dir)
probe.yaml¶
# Probe configs for probe.py: `uv run python examples/probe.py --config-id <entry>`.
#
# A probing run trains one classifier per row of the index on the activations of the positive and the
# negative example prompts, and scores it on those same examples - there is no held-out split, so read
# the scores as how separable the classes are at each location, not as generalization.
#
# `pool` reduces the axes a row leaves whole: here `position`, since the prompts differ in length.
# Every axis left unpooled is flattened into the feature vector. `layer` is never a tensor axis and
# `neuron` is the feature axis, so neither can be pooled.
# The mean residual stream after every layer. The mean includes the BOS position, which is the same
# for every prompt.
residual_by_layer:
method: linear
pool: {position: mean}
metric: [accuracy, f1, auroc]
source:
activation: resid.post
location: [layer]
# The last token's residual stream, which has read the whole prompt, scored by every metric.
last_token_by_layer:
method:
type: linear
max_iter: 2000
seed: 0
pool: {position: last}
metric: [accuracy, precision, recall, f1, auroc]
source:
activation: resid.post
location: [layer]
# One probe per head of the later layers, on the head's output averaged over positions: d_head
# features per probe.
head_outputs:
method: linear
pool: {position: mean}
metric: [accuracy, auroc]
source:
activation: attn.z
location: [layer, head]
filter:
- layer: {range: [8, 11]}
Circuit discovery¶
circuit.py¶
"""Run one entry of circuit.yaml on some prompts.
uv run python examples/circuit.py
uv run python examples/circuit.py --config-id strong_routes --prompts "The Eiffel Tower is in"
Results land in runs/<config-id>/: open circuit_<i>.html in a browser, one page per prompt.
"""
from dataclasses import dataclass, field
from pathlib import Path
from transformer_lens.model_bridge import TransformerBridge
from entoscope import CachedModel, CircuitConfig, CircuitDiscoverer, CircuitInput, ConfigFactory, experiment
CONFIGS = Path(__file__).with_name("circuit.yaml")
@dataclass
class Config:
config_id: str = field(default="information_flow_routes", metadata={"help": "the entry of the config file to run"})
config_path: str = field(default=str(CONFIGS), metadata={"help": "a YAML file"})
prompts: list[str] = field(
default_factory=lambda: ["The capital of France is", "The Eiffel Tower is in"],
metadata={"help": "one circuit is discovered per prompt"},
)
model_name: str = "gpt2"
output_dir: str = field(default="runs", metadata={"help": "results go into <output-dir>/<config-id>"})
cache_dir: str | None = field(default=None, metadata={"help": "keep the model's activations on disk here"})
dry_run: bool = field(default=False, metadata={"help": "log what the run would do, without running it"})
@experiment
def main(config: Config):
"""Discover the circuit behind each prompt."""
model = CachedModel(TransformerBridge.boot_transformers(config.model_name), config.cache_dir)
model.eval()
circuit_input = CircuitInput(
ConfigFactory[CircuitConfig](config.config_path).create(config.config_id), config.prompts
)
save_dir = f"{config.output_dir}/{config.config_id}"
if config.dry_run:
CircuitDiscoverer(model).plan(circuit_input, save_dir)
else:
CircuitDiscoverer(model).apply(circuit_input, save_dir)
circuit.yaml¶
# Circuit-discovery configs for circuit.py: `uv run python examples/circuit.py --config-id <entry>`.
#
# A circuit spans the whole model, so there is no `source` to locate: an entry is only the method.
# Information Flow Routes is licensed for non-commercial use only (see the README).
# The method with its defaults: contributions under 0.04 dropped, routes traced back from the
# next-token prediction.
information_flow_routes:
method: information_flow_routes
# A sparser picture: the full graph keeps the same edges, while only edges above 0.1 survive into the
# routes to the prediction.
strong_routes:
method:
type: information_flow_routes
threshold: 0.04
path_threshold: 0.1
# The routes to the predictions at the last two positions, merged into one graph.
last_two_predictions:
method:
type: information_flow_routes
prediction_positions: [-2, -1]
A custom lens¶
A lens the script registers itself, as Custom methods describes, and the configs that name it.
custom_lens.py¶
"""Register a lens of your own and run it from custom_lens.yaml.
uv run python examples/custom_lens.py --config-id warm
The lens is registered when this module runs, above ``main``: an experiment script runs ``main`` as
it is decorated, and a config naming the lens is checked against the registry as it is read.
"""
from dataclasses import dataclass, field
from pathlib import Path
from transformer_lens.model_bridge import TransformerBridge
from entoscope import ConfigFactory, Lens, LensConfig, LensInput, LensMethodConfig, experiment
from entoscope.lens import LensMethodFactory
from entoscope.lens.method.logit_lens import LogitLens
CONFIGS = Path(__file__).with_name("custom_lens.yaml")
@dataclass(kw_only=True)
class TemperatureLensConfig(LensMethodConfig):
temperature: float = 1.0
@LensMethodFactory.register("temperature_lens")
class TemperatureLens(LogitLens[TemperatureLensConfig]):
"""The logit lens, with its logits divided by a temperature before the collector's softmax."""
def project(self, activation, index):
return super().project(activation, index) / self.method_config.temperature
@dataclass
class Config:
config_id: str = field(default="warm", metadata={"help": "the entry of the config file to run"})
config_path: str = field(default=str(CONFIGS), metadata={"help": "a YAML file"})
prompt: str = "The capital of France is"
model_name: str = "gpt2"
output_dir: str = field(default="runs", metadata={"help": "results go into <output-dir>/<config-id>"})
dry_run: bool = field(default=False, metadata={"help": "log what the run would do, without running it"})
@experiment
def main(config: Config):
"""Run the temperature lens on a prompt."""
model = TransformerBridge.boot_transformers(config.model_name)
model.eval()
lens_input = LensInput(ConfigFactory[LensConfig](config.config_path).create(config.config_id), config.prompt)
save_dir = f"{config.output_dir}/{config.config_id}"
if config.dry_run:
Lens(model).plan(lens_input, save_dir)
else:
Lens(model).apply(lens_input, save_dir)
custom_lens.yaml¶
# Configs naming the lens custom_lens.py registers: `uv run python examples/custom_lens.py --config-id <entry>`.
#
# A registered method is named like a built-in one. Its fields are read strictly into its method
# config, so a field it does not have, or a lens nothing registered, fails as the file is read.
# Logits divided by 2 before the softmax: the same ranking, flatter probabilities.
warm:
method:
type: temperature_lens
temperature: 2.0
top_k: 5
source:
activation: resid.post
location: [layer, position]
# Just the type: the lens with its defaults, which is the logit lens itself.
plain:
method: temperature_lens
top_k: 5
source:
activation: resid.post
location: [layer, position]