transformer_lens.tools.analysis.logit_lens module¶

Logit lens and vocabulary readout.

The logit lens reads every intermediate residual-stream state through the model’s own final norm and unembedding, so each layer can be asked “what would the model predict if it stopped here?”. logit_lens() does that in one call over a cache’s accumulated residual stream; logit_readout() is the kernel underneath it and reads any d_model vector (a head output, a steering direction, a sparse-coder decoder column) through the same path.

Both go through the real ln_final module, W_U, b_U and the adapter’s post-unembedding transform (Gemma-style softcap, Cohere logit scale), so the final stack entry reproduces the model’s own logits on a raw bridge and the model’s log-probabilities in compatibility mode (center_unembed shifts raw logits by a per-position constant). Logits are produced chunk by chunk, so [layers, batch, pos, d_vocab] is never materialized unless the full vocabulary is explicitly requested.

Example:

from transformer_lens.model_bridge import TransformerBridge
from transformer_lens.tools.analysis import logit_lens

model = TransformerBridge.boot_transformers("gpt2", device="cpu")
result = logit_lens(model, "The Eiffel Tower is in the city of", targets=" Paris")
for label, rank in zip(result.labels, result.rank_trajectory(" Paris").tolist()):
    print(f"{label:>12}: rank {rank}")

# Rank of the actual next token at every position and layer: a per-row target.
tokens = model.to_tokens("The Eiffel Tower is in the city of Paris")
nxt = logit_lens(model, tokens[:, :-1], positions=None, targets=tokens[:, 1:])
nxt.rank_trajectory()  # [layers, batch, pos]

class transformer_lens.tools.analysis.logit_lens.LogitLensResult(readout: LogitReadout, labels: List[str], has_batch_dim: bool, target_names: List[str | None] | None = None)¶

Bases: object

Result of logit_lens().

readout¶

The LogitReadout over the accumulated stack. Its leading axis is the stack entry, aligned with labels; the remaining leading dims are [batch, pos] after the requested positions / batch_slice (an integer selection drops its axis).

Type:: transformer_lens.tools.analysis.logit_lens.LogitReadout

labels¶

Stack entry labels from accumulated_resid ("0_pre", "3_mid", "final_post").

Type:: List[str]

has_batch_dim¶

Whether a batch axis follows the stack axis.

Type:: bool

target_names¶

For shared targets, the string passed for each target (None where an id was passed), so rank_trajectory() accepts the strings again.

Type:: List[str | None] | None

has_batch_dim: bool¶

labels: List[str]¶

rank_trajectory(target: str | int | None = None) → Tensor¶

Full-vocabulary rank of one target at every stack entry, [layer, ...].

target is the string passed to targets or a token id; None selects the only target when exactly one was requested, and is the only form for per-row targets.

readout: LogitReadout¶

target_names: List[str | None] | None = None¶

to_dataframe(tokenizer: Any | None = None, top_k: int | None = None) → Any¶

One row per (stack entry, batch, position, token) as a pandas DataFrame.

Rows within a (layer, batch, pos) group are sorted by descending value and top_k keeps the leading ones. Full-vocabulary "logits" readouts go through transformer_lens.utilities.logits_utils.logits_to_df() (columns logit, log_prob, probability); every other readout carries token_index, optional token_string and one value column named like logits_to_df names it (logit / log_prob / probability).

top_tokens(tokenizer: Any) → Dict[str, Any]¶: Decoded top-k tokens per stack entry, {label: nested lists}.

property values: Tensor¶: Shortcut for readout.values.

class transformer_lens.tools.analysis.logit_lens.LogitReadout(values: Float[Tensor, '... selected'], return_type: str, entropy: Float[Tensor, '...'], top_ids: Int[Tensor, '... k'] | None = None, vocab_ids: Int[Tensor, 'v'] | None = None, target_ids: Tensor | None = None, target_ranks: Int[Tensor, '... t'] | None = None, per_row_targets: bool = False, applied_ln: bool = True, applied_output_transform: bool = True)¶

Bases: object

Result of logit_readout().

Exactly one of the three selectors expands the trailing axis of values: targets gives [..., n_targets] ([..., 1] for per-row targets), top_k gives [..., k] with top_ids, vocab gives [..., len(vocab)]; with none of them the trailing axis is the full vocabulary. entropy and target_ranks are always computed over the full vocabulary, whatever was selected.

values¶

The requested quantity (see return_type) with the leading dims of the input vectors and the selected trailing axis.

Type:: jaxtyping.Float[Tensor, ‘… selected’]

return_type¶

"logits", "log_probs" or "probs".

Type:: str

entropy¶

Full-vocabulary entropy of the softmax, in nats, [...].

Type:: jaxtyping.Float[Tensor, ‘…’]

top_ids¶

Token ids aligned with values when top_k was used.

Type:: jaxtyping.Int[Tensor, ‘… k’] | None

vocab_ids¶

The requested subset ids when vocab was used.

Type:: jaxtyping.Int[Tensor, ‘v’] | None

target_ids¶

Shared targets as [n_targets]; per-row targets as a tensor with the leading dims of the vectors (one id per row).

Type:: torch.Tensor | None

target_ranks¶

Zero-based full-vocabulary rank of each target (0 is the argmax; ties count only strictly larger logits), aligned with values.

Type:: jaxtyping.Int[Tensor, ‘… t’] | None

per_row_targets¶

Whether targets was given per row (a tensor).

Type:: bool

applied_ln¶

Whether the model’s final norm was applied to the vectors.

Type:: bool

applied_output_transform¶

Whether the adapter’s post-unembedding transform was invoked (it is the identity for most architectures).

Type:: bool

applied_ln: bool = True¶

applied_output_transform: bool = True¶

decode(tokenizer: Any) → Any¶

Decode the trailing-axis token ids to strings, keeping the leading structure.

top_k and per-row readouts decode per row; shared targets / vocab readouts decode the id list once. Full-vocabulary readouts have no id list to decode.

entropy: Float[Tensor, '...']¶

per_row_targets: bool = False¶

return_type: str¶

row_ids(index: Tuple[int, ...]) → Tensor | None¶: Trailing-axis token ids for one row of values (None = full vocabulary).

property selected_ids: Tensor | None¶: Token ids along the trailing axis of values when it is not the full vocabulary.

target_ids: Tensor | None = None¶

target_ranks: Int[Tensor, '... t'] | None = None¶

top_ids: Int[Tensor, '... k'] | None = None¶

values: Float[Tensor, '... selected']¶

vocab_ids: Int[Tensor, 'v'] | None = None¶

Read the accumulated residual stream through the final norm and unembedding.

Builds the stack with accumulated_resid(), optionally applies a per-entry lens translator, applies the final norm through apply_ln_to_stack() with recomputed statistics, and hands the result to logit_readout(). The final stack entry reproduces the model’s logits (raw bridge) or log-probabilities (compatibility mode).

Parameters:

model – A TransformerBridge, or an ActivationCache (then input must be omitted).
input – Prompt(s) or token tensor to run when no cache is given. Only the residual-stream hooks are cached. A batch of strings is padded by the tokenizer, so with the default positions=-1 the shorter prompts are read at a pad position; pass a token tensor (or per-prompt positions) when prompts differ in length.
cache – An existing cache to read instead of running the model.
layers – Indices into the accumulated stack (negative allowed); None keeps every entry.
incl_mid – Include resid_mid entries. Raises on hybrid stacks whose SSM blocks have no resid_mid.
positions – Position selection (Slice semantics); default last token.
batch_slice – Batch selection (Slice semantics).
lens – Optional f(entry, index) -> entry applied to each stack entry before the norm, e.g. a tuned-lens translator. Combine with enable_grad=True to fit or attribute through it.
targets – As in logit_readout(); a per-row tensor is aligned with the stack’s [batch, pos] dims after positions and batch_slice are applied ([batch, pos] next-token ids with positions=None, [batch] with the default last position).
apply_ln – As in logit_readout().
vocab – As in logit_readout().
return_type – As in logit_readout().
top_k – As in logit_readout().
chunk_size – As in logit_readout().
enable_grad – As in logit_readout().

Returns:

A LogitLensResult.

Read d_model vectors through the model’s final norm and unembedding.

Computes transform(ln_final(x) @ W_U + b_U) for every vector, where transform is the adapter’s post-unembedding transform, in chunks of chunk_size rows so at most chunk_size × d_vocab logits exist at once. Everything derived from the full vocabulary (the softmax normalizer, entropy, target_ranks) is computed per chunk and discarded; only the selected columns are kept.

Parameters:

model – A TransformerBridge in eval mode.
vectors – Residual-stream-shaped vectors, any leading dims.
apply_ln – Apply the model’s ln_final module first (recomputed statistics per vector). False reads the raw vectors and does not reproduce the model.
vocab – Token ids to keep along the trailing axis.
targets – Tokens whose values and full-vocabulary ranks to report. Strings (single tokens) or ints, or a sequence of them, are shared across every row. A tensor is per row: it is right-aligned and broadcast to the vectors’ leading dims (for [batch, pos, d_model] vectors pass [batch, pos] next-token ids, or [batch, 1] for one id per example) and the trailing axis of values has length 1. Note the convention differs from ActivationCache.logit_attrs, which reads a 1-D tensor as per-example first; here alignment is purely right-to-left.
return_type – "logits", "log_probs" or "probs".
top_k – Keep the k largest-logit tokens per vector, with their ids.
chunk_size – Rows of vectors processed per unembedding matmul.
enable_grad – Keep autograd on (for trainable lenses and attribution); default runs under torch.no_grad().

Returns:

A LogitReadout.

Raises:

ValueError – On a bad return_type, more than one selector, a chunk_size below 1, a wrong trailing dimension, empty vectors, per-row targets that do not broadcast, or a model whose final norm cannot be applied exactly.
NotImplementedError – For encoder-decoder and encoder-only bridges.