The Logit Lens in PyTorch

import torch
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
from transformer_lens import HookedTransformer

torch.set_grad_enabled(False)
plt.rcParams["figure.dpi"] = 120      # sharper text in the notebook
plt.rcParams["font.family"] = ["DejaVu Sans", "Kohinoor Devanagari", "Arial Unicode MS"]   # later fonts = fallbacks for missing glyphs

device = "mps" if torch.backends.mps.is_available() else "cpu"
print(f"Using device: {device}")

model_name = "gpt2"
# later:  "gpt2-medium", "gpt2-large", "gpt2-xl"

model = HookedTransformer.from_pretrained_no_processing(model_name, device=device)
model.eval()

cfg = model.cfg
print(f"layers:     {cfg.n_layers}")
print(f"d_model:    {cfg.d_model}")
print(f"vocab size: {cfg.d_vocab}")
print(f"max tokens: {cfg.n_ctx}")

W_E = model.W_E
W_U = model.W_U
print("W_E:", W_E.shape, W_E.dtype, W_E.device)
print("W_U:", W_U.shape, W_U.dtype, W_U.device)

# GPT-2 ties the embedding and unembedding
print("W_U == W_E.T ?", torch.equal(W_U, W_E.T))

n_params = sum(p.numel() for p in model.parameters())
print(f"parameters: {n_params:,}   (W_E counted twice → minus {W_E.numel():,} = {n_params - W_E.numel():,})")

text = (
    "Sometimes, when people say plasma, they mean a state of matter. "
    "Other times, when people say plasma, they mean"
)

# text = (
#     "Nepal ist ein Himalaya-Staat mit sehr vielen Bergen."
# )
max_tokens = 200

tokens = model.to_tokens(text, prepend_bos=False)[:, :max_tokens]   # [1, seq]
seq = tokens.shape[1]

print("shape: ", tokens.shape)
print("dtype: ", tokens.dtype)
print("device:", tokens.device)
print(tokens[0, :8])
print(model.to_str_tokens(tokens[0]))

import functools
import unicodedata
from transformers.models.gpt2.tokenization_gpt2 import bytes_to_unicode

_byte_decoder = {ch: b for b, ch in bytes_to_unicode().items()}


def token_bytes(token_id) -> bytes:
    # The raw UTF-8 bytes a token stands for.
    return bytes(_byte_decoder[ch] for ch in model.tokenizer.convert_ids_to_tokens(int(token_id)))


def utf8_len(lead: int) -> int:
    # How many bytes a character has, read from its first byte.
    return 1 if lead < 0x80 else 2 if lead < 0xE0 else 3 if lead < 0xF0 else 4


def label_inputs(ids) -> list[str]:
    # One label per input token. '·' marks where a character is cut: 'न·' then '·न'.
    pieces = [token_bytes(i) for i in ids]
    text = b"".join(pieces).decode("utf-8", errors="replace")
    spans, pos = [], 0                      # (start_byte, end_byte, char) for each character
    for ch in text:
        n = len(ch.encode("utf-8"))
        spans.append((pos, pos + n, ch))
        pos += n
    labels, start = [], 0
    for piece in pieces:
        end = start + len(piece)
        label = ""
        for s, e, ch in spans:
            if s < end and e > start:       # this character overlaps this token
                label += ("·" if s < start else "") + ch + ("·" if e > end else "")
        labels.append(label)
        start = end
    return labels


@functools.lru_cache(maxsize=None)
def start_range(frag: bytes) -> str:
    # First…last readable character whose UTF-8 bytes start with frag, e.g. b'\xe0\xa4' → 'ऄ…ह'.
    n = utf8_len(frag[0])
    lo, hi = {2: (0x80, 0x800), 3: (0x800, 0x10000), 4: (0x10000, 0x110000)}.get(n, (0, 0))
    chars = [chr(cp) for cp in range(lo, hi)
             if unicodedata.category(chr(cp))[0] in "LNPS"            # letters, numbers, punct, symbols
             and chr(cp).encode("utf-8").startswith(frag)]           # (skips marks, unassigned, surrogates)
    return f"{chars[0]}…{chars[-1]}" if chars else frag.hex(" ")


def dangling_bytes(ids) -> bytes:
    # Bytes at the end of the input that don't yet form a complete character.
    b = b"".join(token_bytes(i) for i in ids)
    for k in range(4):
        try:
            b[:len(b) - k].decode("utf-8")
            return b[len(b) - k:]
        except UnicodeDecodeError:
            pass
    return b""


def label_prediction(prev_ids, pred_id) -> str:
    # Label for a token predicted right after prev_ids.
    b = token_bytes(pred_id)
    try:
        return b.decode("utf-8")                                     # 1. complete by itself
    except UnicodeDecodeError:
        pass
    try:
        return "·" + (dangling_bytes(prev_ids) + b).decode("utf-8")  # 2. finishes the input's half-character
    except UnicodeDecodeError:
        pass
    starts = [k for k, x in enumerate(b) if x >= 0xC0]                 # 3. ends by starting a new character
    if starts:
        i = starts[-1]                                               # the last character is the unfinished one
        try:
            return b[:i].decode("utf-8") + start_range(b[i:]) + "·"
        except UnicodeDecodeError:
            pass
    return b.hex(" ")                                                # 4. orphan continuation byte


# Demo: decode() vs. honest labels
demo = model.to_tokens("नेपाल एक", prepend_bos=False)[0].tolist()
print("decode():", [model.tokenizer.decode([t]) for t in demo])
print("labels:  ", label_inputs(demo))
print("predict e0 a4 after 'न'    :", label_prediction(demo[:2], demo[0]))
print("predict a8    after 'e0 a4':", label_prediction(demo[:1], demo[1]))
print()
print("this text:", label_inputs(tokens[0].tolist()))

logits, cache = model.run_with_cache(tokens)

print("logits:", logits.shape)
print("number of cached activations:", len(cache))
print(list(cache.keys())[:14])

x0 = cache["resid_pre", 0]                         # [1, seq, 768]
manual_x0 = model.W_E[tokens] + model.W_pos[:seq]  # [1, seq, 768] + [seq, 768] → broadcast
print("resid_pre 0 == W_E[tokens] + W_pos ?", torch.allclose(x0, manual_x0))
print("resid_post 0 == resid_pre 1 ?       ", torch.equal(cache["resid_post", 0], cache["resid_pre", 1]))

ids = tokens[0].tolist()
in_labels = label_inputs(ids)
preds = logits[0].argmax(dim=-1).tolist()          # [seq] — top token id per position

for i in range(seq):
    pred = label_prediction(ids[:i + 1], preds[i])  # prediction i comes right after inputs 0..i
    if i + 1 < seq:
        actual, mark = in_labels[i + 1], ("✓" if preds[i] == ids[i + 1] else "✗")
    else:
        actual, mark = "???", ""
    print(f"{i:2d}  {in_labels[i]!r:14} → {pred!r:12} actual: {actual!r:12} {mark}")

gen = tokens.clone()
for _ in range(5):
    next_logits = model(gen)[0, -1]                # last position only → [vocab]
    next_id = next_logits.argmax()                 # 0-dim tensor
    gen = torch.cat([gen, next_id.view(1, 1)], dim=1)

print(model.to_string(gen[0]))

n_layers = model.cfg.n_layers

resid_list = [cache["resid_pre", 0]] + [cache["resid_post", L] for L in range(n_layers)]
resid = torch.stack(resid_list, dim=0)[:, 0]       # [n_layers+1, seq, d_model]

layer_names = ["h_in"] + [f"h{L}_out" for L in range(n_layers)]

print("resid:", resid.shape)
print(layer_names)

norms = resid.norm(dim=-1)                         # [n_layers+1, seq]
print("avg length per layer:", norms.mean(dim=-1)) # [n_layers+1]

pos = 5                                            # the first ' plasma'
for name, n in zip(layer_names, norms[:, pos]):
    print(f"{name:8} {n.item():8.1f}")

def plain_norm(x: torch.Tensor, eps: float = 1e-5) -> torch.Tensor:
    # Normalize the last dim to mean 0, std 1. No learned parameters.
    mean = x.mean(dim=-1, keepdim=True)                    # [..., 1]
    var = ((x - mean) ** 2).mean(dim=-1, keepdim=True)     # [..., 1]
    return (x - mean) / torch.sqrt(var + eps)

normed = plain_norm(resid)                                 # [n_layers+1, seq, d_model]
v = normed[3, 5]
print("shape:", normed.shape)
print(f"mean: {v.mean().item():.6f}   std: {v.std(unbiased=False).item():.6f}")
print("length at each layer (token 5):", normed.norm(dim=-1)[:, 5])

ln = model.ln_final
print("ln_final.w:", ln.w.shape, "  ln_final.b:", ln.b.shape)
print(f"w ranges from {ln.w.min().item():.3f} to {ln.w.max().item():.3f}")

manual = plain_norm(resid[-1]) * ln.w + ln.b
print("ln_final == plain_norm * w + b ?", torch.allclose(ln(resid[-1]), manual, atol=1e-5))

h_out_logits = ln(resid[-1]) @ model.W_U + model.b_U       # [seq, vocab]
print("h_out == model logits ?", torch.allclose(h_out_logits, logits[0], atol=1e-3))

lens_logits = plain_norm(resid) @ model.W_U                    # [n_layers+1, seq, vocab]  (b_U is all zeros for GPT-2)
h_out_logits = model.ln_final(resid[-1]) @ model.W_U + model.b_U

all_logits = torch.cat([lens_logits, h_out_logits.unsqueeze(0)], dim=0)   # [n_layers+2, seq, vocab]
all_probs = all_logits.softmax(dim=-1)
all_tops = all_logits.argmax(dim=-1)                          # [n_layers+2, seq]
all_layer_names = layer_names + ["h_out"]

print("logits:", all_logits.shape, " probs:", all_probs.shape, " top:", all_tops.shape)
print("probs sum to 1?", all_probs.sum(dim=-1)[0, :3])

agree = (all_tops[-2] == all_tops[-1]).float().mean().item()
print(f"h11_out and h_out agree on top-1 at {agree:.0%} of positions")

model.to_str_tokens((model.ln_final.b @ model.W_U).topk(8).indices)

def show_position(pos: int):
    ids = tokens[0].tolist()
    labels = label_inputs(ids)
    nxt = labels[pos + 1] if pos + 1 < len(ids) else "???"
    print(f"position {pos}: input {labels[pos]!r}, actual next {nxt!r}")
    for k, name in enumerate(all_layer_names):
        t = all_tops[k, pos].item()
        p = all_probs[k, pos, t].item()
        mark = "✓" if pos + 1 < len(ids) and t == ids[pos + 1] else ""
        print(f"  {name:8} {label_prediction(ids[:pos + 1], t)!r:14} p={p:.3f} {mark}")

show_position(4)
print()
show_position(19)

def logit_lens(resid: torch.Tensor, norm: str = "plain") -> torch.Tensor:
    # resid: [n_layers+1, seq, d_model] → logits: [n_layers+2, seq, d_vocab]
    if norm == "plain":
        normed = plain_norm(resid)
    elif norm == "ln_final":
        normed = model.ln_final(resid)
    else:
        raise ValueError(f"unknown norm: {norm!r}")

    lens_logits = normed @ model.W_U + model.b_U
    final_logits = model.ln_final(resid[-1]) @ model.W_U + model.b_U
    return torch.cat([lens_logits, final_logits.unsqueeze(0)], dim=0)

def ranks_of(lens_logits: torch.Tensor, target_ids: torch.Tensor) -> torch.Tensor:
    # lens_logits: [rows, cols, vocab], target_ids: [cols] → ranks: [rows, cols] (1 = top)
    idx = target_ids.view(1, -1, 1).expand(lens_logits.shape[0], -1, 1)
    target_logit = lens_logits.gather(dim=-1, index=idx)
    return (lens_logits > target_logit).sum(dim=-1) + 1

targets = tokens[0, 1:]                                        # [seq-1]

for norm in ["plain", "ln_final"]:
    L = logit_lens(resid, norm)                                # [n_layers+2, seq, vocab]
    top = L.argmax(dim=-1)
    agree = (top == top[-1]).float().mean(dim=-1)              # [n_layers+2]
    ranks = ranks_of(L[:, :-1], targets)                       # [n_layers+2, seq-1]

    print(f"--- {norm} ---")
    for k, name in enumerate(all_layer_names):
        print(f"  {name:8} agree={agree[k].item():.2f}   rank of ' plasma' @ pos 19: {ranks[k, 19].item():>6}")

def run_example(text: str, norm: str = "ln_final", max_tokens: int = 200, prepend_bos: bool = False):
    # Returns tokens [1, seq], lens_logits [n_layers+2, seq, vocab], layer_names (list of str)
    tokens = model.to_tokens(text, prepend_bos=prepend_bos)[:, :max_tokens]

    keep = lambda name: name == "blocks.0.hook_resid_pre" or name.endswith("hook_resid_post")
    _, cache = model.run_with_cache(tokens, names_filter=keep)

    n_layers = model.cfg.n_layers
    resid = torch.stack([cache["resid_pre", 0]] + [cache["resid_post", L] for L in range(n_layers)])[:, 0]
    assert resid.shape == (n_layers + 1, tokens.shape[1], model.cfg.d_model), resid.shape

    layer_names = ["h_in"] + [f"h{L}_out" for L in range(n_layers)] + ["h_out"]
    return tokens, logit_lens(resid, norm), layer_names

clean = lambda s: repr(s)[1:-1]
MONO = ["DejaVu Sans Mono", "Kohinoor Devanagari", "Arial Unicode MS"]   # monospace + fallbacks


def cell_figsize(n_rows: int, n_cols: int, lines_per_cell: int = 1, chars_per_line: int = 8,
                 fontsize: int = 8, pad: tuple = (2.5, 2.5)) -> tuple:
    # Figure size in inches so every cell fits its text.
    char_w = 0.6 * fontsize / 72
    line_h = 1.3 * fontsize / 72
    col_w = chars_per_line * char_w + 0.2
    row_h = lines_per_cell * line_h + 0.15
    return (n_cols * col_w + pad[0], n_rows * row_h + pad[1])    # pad = room for labels & colorbar


def label_axes(ax, fig, im, layer_names, bottom_labels, top_labels, top_title, cbar_label, title):
    ax.set_yticks(range(len(layer_names)), layer_names)
    ax.set_xticks(range(len(bottom_labels)), [clean(t) for t in bottom_labels], rotation=60, ha="right")
    ax.set_xlabel("input token")
    top_axis = ax.secondary_xaxis("top")
    top_axis.set_xticks(range(len(top_labels)), [clean(t) for t in top_labels], rotation=60, ha="left")
    top_axis.set_xlabel(top_title)
    fig.colorbar(im, ax=ax, label=cbar_label, fraction=0.03, pad=0.02)
    ax.set_title(title)
    fig.tight_layout()

def plot_lens_topk(tokens: torch.Tensor, lens_logits: torch.Tensor, layer_names: list[str],
                   k: int = 3, start: int = 0, end: int | None = None,
                   title: str = "", fontsize: int = 8, max_chars: int = 10):
    assert len(layer_names) == lens_logits.shape[0], f"{len(layer_names)} names for {lens_logits.shape[0]} rows"
    ids = tokens[0].tolist()
    in_labels = label_inputs(ids)
    end = end or len(ids)

    probs = lens_logits[:, start:end].softmax(dim=-1)       # [rows, cols, vocab]
    top_p, top_id = probs.topk(k, dim=-1)                   # [rows, cols, k] each, sorted high→low
    top_p, top_id = top_p.cpu().numpy(), top_id.cpu().numpy()
    n_rows, n_cols, _ = top_p.shape
    nxt = [in_labels[i + 1] if i + 1 < len(ids) else "???" for i in range(start, end)]
    nxt_ids = [ids[i + 1] if i + 1 < len(ids) else -1 for i in range(start, end)]

    fig, ax = plt.subplots(figsize=cell_figsize(n_rows, n_cols, lines_per_cell=k,
                                                chars_per_line=max_chars + 6, fontsize=fontsize))
    im = ax.imshow(top_p[:, :, 0], cmap="Blues", vmin=0, vmax=1, aspect="auto", origin="lower")

    for r in range(n_rows):
        for c in range(n_cols):
            lines, hit = [], False
            for j in range(k):
                t = int(top_id[r, c, j])
                word = label_prediction(ids[:start + c + 1], t)     # context: inputs up to this position
                is_next = t == nxt_ids[c]
                hit |= is_next
                lines.append(f"{clean(word)[:max_chars]:<{max_chars}} {top_p[r, c, j]:.2f}{'✓' if is_next else ' '}")
            ax.text(c, r, "\n".join(lines), ha="center", va="center",
                    fontsize=fontsize, family=MONO,
                    color="white" if top_p[r, c, 0] > 0.6 else "black",
                    fontweight="bold" if (hit and k == 1) else "normal")

    label_axes(ax, fig, im, layer_names, in_labels[start:end], nxt, "actual next token",
               "probability of #1 token", title or f"top-{k} logit lens  (✓ = actual next token)")
    plt.show()

def fmt_rank(r: int) -> str:
    # 1..999 as-is, then 1.6k, 24k
    if r < 1000:
        return str(r)
    if r < 10_000:
        return f"{r / 1000:.1f}k"
    return f"{r // 1000}k"


def plot_ranks(tokens: torch.Tensor, lens_logits: torch.Tensor, layer_names: list[str],
               start: int = 0, end: int | None = None, against: str = "final",
               title: str = "", fontsize: int = 8):
    ids = tokens[0].tolist()
    in_labels = label_inputs(ids)
    seq = len(ids)

    if against == "final":                                   # model's own final top-1
        end = end or seq
        target_ids = lens_logits[-1, start:end].argmax(dim=-1)
    elif against == "truth":                                 # actual next token (last pos has none)
        end = min(end or seq, seq - 1)
        target_ids = tokens[0, start + 1 : end + 1]
    else:
        raise ValueError(f"against must be 'final' or 'truth', got {against!r}")

    ranks = ranks_of(lens_logits[:, start:end], target_ids).cpu().numpy()
    if against == "final":
        target_strs = [label_prediction(ids[:start + c + 1], t) for c, t in enumerate(target_ids.tolist())]
    else:
        target_strs = in_labels[start + 1 : end + 1]
    n_rows, n_cols = ranks.shape

    norm = LogNorm(vmin=1, vmax=model.cfg.d_vocab)
    fig, ax = plt.subplots(figsize=cell_figsize(n_rows, n_cols, chars_per_line=5, fontsize=fontsize))
    im = ax.imshow(ranks, cmap="Blues_r", norm=norm, aspect="auto", origin="lower")

    for r in range(n_rows):
        for c in range(n_cols):
            ax.text(c, r, fmt_rank(int(ranks[r, c])), ha="center", va="center", fontsize=fontsize,
                    color="white" if norm(ranks[r, c]) < 0.35 else "black")

    label_axes(ax, fig, im, layer_names, in_labels[start:end], target_strs,
               "model's final prediction" if against == "final" else "actual next token",
               "rank (1 = top, log scale)", title)
    plt.show()

tokens, L, names = run_example(text, norm="ln_final")
print("lens logits:", L.shape)

plot_lens_topk(tokens, L, names, k=1, start=12, end=24, title="top-1 logit lens (ln_final) — gpt2")

plot_lens_topk(tokens, L, names, k=5, start=16, end=24)

plot_ranks(tokens, L, names, against="final", title="rank of the model's final prediction (ln_final)")

plot_ranks(tokens, L, names, against="truth", title="rank of the actual next token (ln_final)")

# Same text, plain-norm lens: compare the ' say' → ' plasma' column
tokens_p, L_plain, names_p = run_example(text, norm="plain")
plot_ranks(tokens_p, L_plain, names_p, against="final", title="rank of the model's final prediction (plain norm)")

tokens_ne, L_ne, names_ne = run_example("नेपाल एक सुन्दर देश हो।", norm="ln_final")
print("tokens:", tokens_ne.shape[1], "   labels:", label_inputs(tokens_ne[0].tolist()))
plot_lens_topk(tokens_ne, L_ne, names_ne, k=2, start=0, end=12, title="top-2 logit lens on Nepali (ln_final)")