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.
• 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, 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.

Raises:

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

__init__

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