Steps — Logit Lens

LogitLens step: project per-layer residuals onto the vocabulary.

Tracks how next-token predictions evolve across the depth of the network by applying the model’s final norm and unembedding to each recorded layer’s output. Cross-architecture access is provided by nnterp’s standardized ln_final and lm_head modules.

LogitLensResult

Per-layer next-token probability distributions.

Attributes:

• all_probs: Tensor [n_layers, n_inputs, seq, vocab] of softmax outputs, or
• None. Kept only when the step ran with store_full_probs=True;
• it is n_layers * n_inputs * seq * vocab floats, which is tens to
• hundreds of gigabytes on a real model, so it is off by default. The
• reduced fields below (and the standard heatmap) do not need it.
• max_probs: Tensor [n_layers, n_inputs, seq] of argmax probabilities.
• predicted_tokens: Tensor [n_layers, n_inputs, seq] of argmax token IDs.
• predicted_words: Decoded tokens, shape [n_layers][n_inputs][seq].
• input_words: Decoded input tokens, shape [n_inputs][seq].
• attention_mask: Tensor [n_inputs, seq] marking real (1) vs padding (0)
• positions; useful for masking padded columns at visualization time.
• addresses: Component addresses (one :class:Node per layer) the rows
• of the tensors correspond to. The lens reads each layer’s block
• output, so these are resid_post nodes.

__init__

def __init__(self, all_probs: Tensor | None, max_probs: Tensor, predicted_tokens: Tensor, predicted_words: list[list[list[str]]], input_words: list[list[str]], attention_mask: Tensor, addresses: list[Node]) -> None:

LogitLens

Compute logit-lens probabilities across layers.

Runs a single nnsight trace on the batched prompts, captures full-sequence residual outputs at each requested layer, and projects them through the model’s standardized final norm and unembedding (via nnterp).

Reads from results:

• results['prompts']: PromptBatch

Writes to results:

• results['logit_lens']: LogitLensResult

Args:

• model: MuranoModel to record from.
• layers: Layer indices to record, or "all" for every layer.
• store_full_probs: Keep the full [n_layers, n_inputs, seq, vocab]
• softmax on the result. Off by default: that tensor is tens to
• hundreds of gigabytes on a real model (every layer, position, and
• vocabulary entry), and the reduced max_probs / predicted_tokens
• / predicted_words fields (and the standard heatmap) do not need
• it. Turn it on only for a targeted, small run that inspects the full
• distribution.

Raises:

• ValueError: If layers is a string other than "all".

__init__

def __init__(self, model: ModelBackend, layers: list[int] | str = 'all', store_full_probs: bool = False):