Logit Lens › Implementation

# uv pip install transformers torch numpy matplotlib
import warnings

from transformers import GPT2LMHeadModel, GPT2Tokenizer
from transformers import logging as hf_logging

# Suppress HuggingFace verbosity and warnings
hf_logging.set_verbosity_error()
warnings.filterwarnings("ignore")

# Load GPT-2 (small, 12 layers, d=768, V=50257)
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
model = GPT2LMHeadModel.from_pretrained("gpt2")
model.eval()

# Extract and copy the unembedding matrix and layer norm weights
# .clone() ensures we get contiguous copies (avoids mmap alignment issues on ARM)
W_U = model.lm_head.weight.detach().float().clone()  # (V, d)
ln_weight = model.transformer.ln_f.weight.detach().float().clone()
ln_bias = model.transformer.ln_f.bias.detach().float().clone()

def get_hidden_states(prompt):
    """Run forward pass and return all hidden states."""
    inputs = tokenizer(prompt, return_tensors="pt")
    input_ids = inputs["input_ids"]
    tokens = tokenizer.convert_ids_to_tokens(input_ids[0])

    with torch.no_grad():
        outputs = model.transformer(
            input_ids,
            output_hidden_states=True,
        )

    # hidden_states is a tuple of (n_layers + 1) tensors, each shape (1, seq_len, d)
    # .clone() ensures contiguous memory (avoids ARM alignment faults with mmap weights)
    hidden_states = tuple(h.float().clone() for h in outputs.hidden_states)
    return hidden_states, tokens, input_ids

def apply_layer_norm(x, weight, bias, eps=1e-5):
    """Apply layer normalization manually."""
    mean = x.mean(dim=-1, keepdim=True)
    var = x.var(dim=-1, keepdim=True, unbiased=False)
    x_norm = (x - mean) / (var + eps).sqrt()
    return x_norm * weight + bias


def logit_lens_decode(
    hidden_state, W_U, ln_weight, ln_bias, apply_final_norm=True
):
    """
    Apply the logit lens to a single hidden state tensor.

    Args:
        hidden_state: (seq_len, d) tensor of hidden states
        W_U: (V, d) unembedding matrix
        ln_weight, ln_bias: LayerNorm parameters

    Returns:
        logits: (seq_len, V) logit tensor
        probs: (seq_len, V) probability tensor
    """
    # Hugging Face returns the last hidden state after the model's final
    # LayerNorm. Earlier states still need that normalization for the lens.
    h_normed = (
        apply_layer_norm(hidden_state, ln_weight, ln_bias)
        if apply_final_norm
        else hidden_state
    )
    # Multiply by unembedding: (seq_len, d) x (d, V) -> (seq_len, V)
    logits = h_normed @ W_U.T
    probs = torch.softmax(logits, dim=-1)
    return logits, probs

def get_top_tokens_per_layer(prompt, position=-1, top_k=5):
    """
    For a given prompt and token position, get the top-k predictions
    at every layer using the logit lens.
    """
    hidden_states, tokens, input_ids = get_hidden_states(prompt)

    results = []
    for layer_idx, hs in enumerate(hidden_states):
        # hs shape: (1, seq_len, d) -> take the target position
        h = hs[0, position, :]  # (d,)
        _, probs = logit_lens_decode(
            h.unsqueeze(0),
            W_U,
            ln_weight,
            ln_bias,
            apply_final_norm=layer_idx < len(hidden_states) - 1,
        )
        probs = probs[0]  # (V,)

        top_probs, top_indices = torch.topk(probs, top_k)
        top_tokens = tokenizer.convert_ids_to_tokens(top_indices.tolist())

        results.append(
            {
                "layer": layer_idx,
                "top_tokens": top_tokens,
                "top_probs": top_probs.tolist(),
            }
        )

    return results, tokens

prompt = "The Eiffel Tower is located in the city of"
results, tokens = get_top_tokens_per_layer(prompt, position=-1, top_k=5)

def logit_lens_heatmap_data(prompt):
    """
    For each (layer, position) pair, get the top-1 predicted token and its probability.
    """
    hidden_states, tokens, input_ids = get_hidden_states(prompt)
    n_layers = len(hidden_states)
    seq_len = len(tokens)

    top_token_matrix = []
    prob_matrix = np.zeros((n_layers, seq_len))

    for layer_idx, hs in enumerate(hidden_states):
        h = hs[0]  # (seq_len, d)
        _, probs = logit_lens_decode(
            h,
            W_U,
            ln_weight,
            ln_bias,
            apply_final_norm=layer_idx < len(hidden_states) - 1,
        )

        top_probs, top_indices = torch.topk(probs, 1, dim=-1)
        top_tokens_row = [
            tokenizer.convert_ids_to_tokens([idx.item()])[0]
            for idx in top_indices[:, 0]
        ]

        top_token_matrix.append(top_tokens_row)
        prob_matrix[layer_idx] = top_probs[:, 0].numpy()

    return top_token_matrix, prob_matrix, tokens


top_token_matrix, prob_matrix, tokens_hm = logit_lens_heatmap_data(prompt)

def track_token_probs_across_layers(prompt, candidate_tokens, position=-1):
    """
    Track the probability of specific candidate tokens across all layers.
    Returns a dict mapping each token to its probability at each layer.
    """
    hidden_states, tokens, input_ids = get_hidden_states(prompt)

    # Get token IDs for candidates
    candidate_ids = {
        tok: tokenizer.encode(tok, add_special_tokens=False)[0]
        for tok in candidate_tokens
    }

    layer_probs = {tok: [] for tok in candidate_tokens}

    for layer_idx, hs in enumerate(hidden_states):
        h = hs[0, position, :]  # (d,)
        _, probs = logit_lens_decode(
            h.unsqueeze(0),
            W_U,
            ln_weight,
            ln_bias,
            apply_final_norm=layer_idx < len(hidden_states) - 1,
        )
        probs = probs[0]
        for tok, tok_id in candidate_ids.items():
            if tok_id != tokenizer.unk_token_id:
                layer_probs[tok].append(probs[tok_id].item())
            else:
                layer_probs[tok].append(0.0)

    return layer_probs, tokens


# Track candidates across layers for the Eiffel Tower prompt
candidates = [" Paris", " London", " France", " Berlin"]
layer_probs, _ = track_token_probs_across_layers(
    prompt, candidates, position=-1
)