entoscope.lens

Logit, tuned and parameter lens.

The operator, its input, configs, enums, collectors, method base and factory are importable from entoscope.lens; the operator, its input, configs and enums from entoscope too.

class entoscope.lens.config.LensMethodConfig(*, type: str)

Bases: MethodConfig

The method config of a lens that takes no fields, such as logit_lens.

A lens with fields of its own names a dataclass extending this one.

classmethod factory() → type[LensMethodFactory]

Name the factory lens methods register with.

Returns:

LensMethodFactory.

class entoscope.lens.config.LensConfig(*, config_id: str | None = None, method: LensMethodConfig, top_k: int = 1, source: ModelLocatorConfig)

Bases: AbstractConfig

What a lens run reads and how.

In a YAML file method is the mapping read strictly into the lens’s method config class, or just the lens’s type for its defaults, so a file naming a lens nothing registered or a field the lens does not have fails as it is read.

Variables:

• method (entoscope.lens.config.LensMethodConfig) – The lens’s method config, whose type names the lens.

• top_k (int) – How many tokens to keep per row.

• source (entoscope.model.config.ModelLocatorConfig) – Where to read. A lens that reads activations reads one position per row, and one that reads weights reads no position.

Raises:

ValueError – If top_k is below 1, or source does not name the positions the lens reads.

class entoscope.lens.input.LensInput(config: LensConfig, prompt: str | None = None)

Bases: OperationInput

A lens config and the prompt it reads.

Variables:

• config (entoscope.lens.config.LensConfig) – The lens config.

• prompt (str | None) – The prompt whose activations the lens reads, None exactly when the lens reads the model’s weights instead, as parameter_lens does.

Raises:

ValueError – If the prompt is missing for a lens that reads activations, or given to one that reads weights.

prepare(model: TransformerBridge | CachedModel, *, index_only: bool = False) → LensInput

Resolve the index a run iterates, then gather the activations it reads on top of it.

The index needs no forward pass, so index_only lets an operation be planned against a model without spending the compute. An input that gathers no activations, its method running the model itself, is ready once its index is.

Parameters:

• model – The model the run reads.

• index_only – Whether to stop after the index, before any forward pass.

Returns:

The input itself, prepared in place.

activation_names() → list[str]

Name the transformer activations the lens reads, which its run caches and a plan reports.

Returns:

The transformer activation names, sorted, none for a lens that reads the model’s weights.

property tokens: Int[Tensor, 'batch pos']

The prompt’s tokens, once the index is prepared.

property cache: ActivationCache

The activations the index reads, once the input is prepared.

property index: Index[ModelLocatorConfig]

The rows the lens reads, once the index is prepared.

to(device: str | device = 'cpu') → LensInput

Move the tensors the input holds to device.

Parameters:

device – The device to move to.

Returns:

The input itself, moved in place.

class entoscope.lens.lens.Lens(model: TransformerBridge | CachedModel, collector_class: type[LensCollector] = LensVisualizationCollector)

Bases: Operator

Reads every index element out onto the vocabulary through the lens config.method names.

One forward pass caches the activations the index reads, none at all for a lens that reads only the weights, and every row becomes its element’s top-k tokens.

Parameters:

• model – The model to read.

• collector_class – The collector to write the results with. The default also draws result.html, and LensCollector writes the CSV files alone.

apply(lens_input: LensInput, save_dir: str | None = None, return_collector: bool = False, verbose: bool = True) → LensCollector | None

Read every row of the input’s index out onto the vocabulary.

The method is built before any forward pass, so a config that names something the run cannot do fails before the model is touched. Where the results were saved is logged through the entoscope logger.

Parameters:

• lens_input – The config and prompt to run, prepared in place.

• save_dir – The directory to write result.csv, index.csv, config.json and the page into, created if missing. None writes nothing.

• return_collector – Whether to return the collector.

• verbose – Whether to show a progress bar.

Returns:

The collector, whose result holds every row’s top-k tokens and probabilities, if return_collector is set, and None otherwise.

plan(lens_input: LensInput, save_dir: str | None = None, verbose: bool = True) → OperationPlan

Describe the run apply would make, without a forward pass and without touching save_dir.

The input’s index is resolved exactly as the run resolves it, and kept for apply.

Parameters:

• lens_input – The config and prompt the run would use.

• save_dir – The directory the run would write into.

• verbose – Whether to also log the plan.

Returns:

The plan – one step per index row and one forward pass, or none for a lens that reads weights.

class entoscope.lens.collector.LensCollector(operation_input: InputT, model: TransformerBridge | CachedModel, save_dir: str | None = None)

Bases: FrameCollector[LensInput]

One row per index element: its top-k tokens, best first.

Rank i fills token_{i} (the token id), str_token_{i} (its string, escaped, see str_token on reading it back) and probability_{i} (the softmax over the element’s logits). Writes result.csv, index.csv and config.json, and LensVisualizationCollector adds result.html.

collect(logits: Tensor, index: int)

Rank one element’s tokens.

Parameters:

• logits – The element’s logits over the vocabulary, [d_vocab] or [1, d_vocab].

• index – The element’s row in the index.

class entoscope.lens.collector.LensVisualizationCollector(operation_input: InputT, model: TransformerBridge | CachedModel, save_dir: str | None = None)

Bases: VisualizationCollector[LensInput, LensFigure], LensCollector

A LensCollector that also draws its result as an interactive Plotly result.html.

Its figures are keyed by top-k rank, from 0 (the best).

figure_class

alias of LensFigure

Interactive Plotly rendering of a lens run.

The grid is the run’s own index, layer up and position (or head, or neuron) across, coloured by a rank’s probability, with that rank’s token printed in the cell and the whole top-k in hover. One grid is drawn per rank, so a run’s second and third guesses are as visible as its first.

class entoscope.lens.visualization.LensFigure(result: DataFrame, operation_input: InputT, model: TransformerBridge | CachedModel, index: Index)

Bases: GridFigure[LensInput, int]

The page of a lens run: one grid per top-k rank, each panel named by its rank from 0.

property title: str

The page’s title.

panels() → list[int]

Name the page’s panels.

Returns:

The panels, in page order, each drawn as its own plot.

heading(rank: int) → str

Name one panel’s plot.

Parameters:

panel – One of panels().

Returns:

The plot’s heading, shown above the subtitle line.

subtitle(rank: int) → str

Describe what one panel’s plot holds, under its heading.

Parameters:

panel – One of panels().

Returns:

The subtitle, as HTML.

value_column(rank: int) → str

Name the result column one panel colours its cells by.

Parameters:

panel – One of panels().

Returns:

A column of the result frame.

value_title(rank: int) → str

Title one panel’s colour bar.

Parameters:

panel – One of panels().

Returns:

The title, as text, by default the value column’s name.

bar_yaxis(rank: int) → dict

Lay out the value axis of one panel drawn as a bar chart.

Parameters:

panel – One of panels().

Returns:

plotly yaxis settings, by default only the axis title.

cell_text(rank: int) → list[str]

Choose the label printed inside each cell of one panel.

The size cap is applied by figure, so a subclass answers for its own result alone.

Parameters:

panel – One of panels().

Returns:

One label per result row, or None to leave the cells bare, the default.

str_tokens() → dict[str, list[str]]

Read the prompt’s token strings, to label the position axes with.

Read once at construction. A figure whose index carries no position, or a run with no prompt to read one from, labels its ticks with the coordinates alone.

Returns:

The token strings for each position column of the result, keyed by that column. Empty by default.

hover(rank: int, position: int, coordinate: str) → str

Write what one cell shows on hover.

Parameters:

• panel – One of panels().

• position – The result row the cell is built from.

• coordinate – The cell’s laid-out coordinates, from cell_labels.

Returns:

The hover text, as HTML.

class entoscope.lens.method.lens.AbstractLensMethod(model: TransformerBridge | CachedModel, config: LensConfig)

Bases: ABC, Generic[LensMethodConfigT]

The base of every lens: reads each index element out onto the vocabulary.

A lens is built once per run, from the run’s config alone, so a lens that handles several kinds of activation tells them apart through index.activation() itself. Its type argument names the dataclass of its method config, extending LensMethodConfig, which a config file’s method is read into. A lens naming none takes LensMethodConfig itself, which no field fits.

Parameters:

• model – The model the lens reads.

• config – The run’s config.

Variables:

• reads_activations (ClassVar[bool]) – False for a lens that reads only the model’s weights, which takes no prompt and has no position axis.

• method_config – The run’s method config, typed by the lens’s type argument.

abstractmethod extract(lens_input: LensInput, index: IndexElement[ModelLocatorConfig]) → Float[Tensor, '...']

Read the tensor one element names.

Parameters:

• lens_input – The prepared input, whose cache holds the activations the index reads.

• index – The element.

Returns:

The element’s activation, or the weight it names.

Raises:

ValueError – If the lens cannot read the element.

abstractmethod project(source: Float[Tensor, '...'], index: IndexElement[ModelLocatorConfig]) → Float[Tensor, '... d_vocab']

Map one element’s tensor onto the vocabulary.

The collector ranks the tokens of a single slice, so an element projects to a single d_vocab vector.

Parameters:

• source – What extract returned for the element.

• index – The element.

Returns:

The element’s logits over the vocabulary.

Raises:

ValueError – If the lens cannot read the element.

entoscope.lens.method.logit_lens.projection_layer(model: TransformerBridge | CachedModel, index: IndexElement[ModelLocatorConfig]) → int

Find the single layer whose weights project an element.

Parameters:

• model – The model the element is read from.

• index – The element.

Returns:

The element’s layer.

Raises:

ValueError – If the element names no single layer, or is a key or a value of grouped-query attention, which has fewer heads than the projection weights.

class entoscope.lens.method.logit_lens.LogitLens(model: TransformerBridge | CachedModel, config: LensConfig)

Bases: AbstractLensMethod[LensMethodConfigT]

Reads an element out the way the model reads out its final residual stream.

A residual-stream state, or a sublayer’s own input or output, is unembedded as it is. An MLP-hidden or attention head activation is first written into the residual stream through the weight that writes it there, and a query or key is read against the embedding through the other side of its head.

extract(lens_input: LensInput, index: IndexElement[ModelLocatorConfig]) → Float[Tensor, '...']

Read the element’s activation out of the prompt’s cache.

Parameters:

• lens_input – The prepared input.

• index – The element.

Returns:

The activation, sliced at the element.

project(activation: Float[Tensor, '...'], index: IndexElement[ModelLocatorConfig]) → Float[Tensor, '... d_vocab']

Read an activation out through the model’s final norm and unembedding.

Parameters:

• activation – The element’s activation.

• index – The element.

Returns:

The element’s logits over the vocabulary.

Raises:

ValueError – If the activation is neither in the residual stream’s basis nor an MLP-hidden or attention head activation, or names no single layer where it needs one.

entoscope.lens.method.tuned_lens.resource_id(model: TransformerBridge | CachedModel) → str

Name the pretrained tuned lens that belongs to a model.

Parameters:

model – The model.

Returns:

The model’s tokenizer name, or its model name when it has no tokenizer name.

Raises:

ValueError – If the model has neither.

entoscope.lens.method.tuned_lens.load_translators(model: TransformerBridge | CachedModel, resource: str) → ModuleList

Load the pretrained translators of a tuned lens.

They are checked against the model they translate for and placed on its device and dtype, which is why a lens loads its own rather than sharing them across models.

Parameters:

• model – The model the translators translate for.

• resource – A local directory holding config.json and params.pt, or a resource id on the tuned-lens hub space.

Returns:

One frozen translator per residual state but the last.

Raises:

ValueError – If the translators were trained for a model of another depth or width.

class entoscope.lens.method.tuned_lens.TunedLensConfig(*, type: str, translators: str | None = None)

Bases: LensMethodConfig

The method config of tuned_lens.

Variables:

translators (str | None) – The pretrained translators, a resource id on the tuned-lens hub space or a local directory holding config.json and params.pt. Unset, the model’s own name is used.

class entoscope.lens.method.tuned_lens.TunedLens(model: TransformerBridge | CachedModel, config: LensConfig)

Bases: LogitLens[TunedLensConfig]

The logit lens with each residual state’s trained translator applied before the model’s read-out.

The translators are loaded the first time an element is projected.

Parameters:

• model – The model the lens reads.

• config – The run’s config.

Raises:

ValueError – If no translators are named and the model has no name to find them by.

property translators: ModuleList

The pretrained translators, one per residual state but the last.

class entoscope.lens.method.parameter_lens.ProjectionSpace(*values)

Bases: StrEnum

The vocabulary space a weight is read in: W_U for unembedding, W_E transposed for embedding.

The two coincide on a model that ties its embedding and unembedding, as GPT-2 does with HuggingFace’s weights, until enable_compatibility_mode folds LayerNorm in and centers the unembedding.

class entoscope.lens.method.parameter_lens.ParameterLensConfig(*, type: str, projection: ProjectionSpace = ProjectionSpace.UNEMBEDDING, scale: float = 1.0)

Bases: LensMethodConfig

The method config of parameter_lens.

Variables:

• projection (entoscope.lens.method.parameter_lens.ProjectionSpace) – The vocabulary space a weight is read in.

• scale (float) – What the weight is multiplied by first. A negative scale reads it against its own direction, and the magnitude acts as a temperature.

class entoscope.lens.method.parameter_lens.ParameterLens(model: TransformerBridge | CachedModel, config: LensConfig)

Bases: AbstractLensMethod[ParameterLensConfig]

Reads a weight vector onto the vocabulary: what a neuron or a head dimension reads or writes.

A unit is read through the weight ACTIVATION_WEIGHT names for its activation: the one computing it off the residual stream, or, for one no weight computes (mlp.post, attn.z), the one writing it back. It reads the model’s weights rather than a prompt’s activations, so a run takes no prompt.

extract(lens_input: LensInput, index: IndexElement[ModelLocatorConfig]) → Float[Tensor, 'd_model']

Read the weight vectors of the units an element names, summed into one direction.

Parameters:

• lens_input – The input, whose prompt is unused.

• index – The element.

Returns:

The summed weight vector.

Raises:

ValueError – If the activation names no weight vector, as a residual-stream one does not, or the element names no single layer.

project(weight: Float[Tensor, 'd_model'], index: IndexElement[ModelLocatorConfig]) → Float[Tensor, 'd_vocab']

Read a weight vector in the configured vocabulary space, scaled first.

Parameters:

• weight – The element’s summed weight vector.

• index – The element.

Returns:

The element’s logits over the vocabulary.

class entoscope.lens.method.factory.LensMethodFactory

Bases: RegistryFactory[AbstractLensMethod]

Resolves the lens a config’s method.type names and builds it for a run.

It holds logit_lens, tuned_lens, parameter_lens and every lens registered with register.

classmethod create(model: TransformerBridge | CachedModel, config: LensConfig) → AbstractLensMethod

Build the lens a config names, once per run.

Parameters:

• model – The model the lens reads.

• config – The run’s config, whose method.type names the lens.

Returns:

The lens.

Raises:

ValueError – If no lens is registered under the name, or the lens rejects the config.