{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Module 21: Interpretability"
      ],
      "id": "05d02efa-3724-4c92-a63f-9627c1fe3e04"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [],
      "id": "91270a5c-4ef1-43c0-99b2-ad596ff83117"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "40f4aae8-9d88-402e-9028-5c82ad92a1de"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Introduction\n",
        "\n",
        "Every module so far taught you to *build* a language model. This one\n",
        "teaches you to *look inside* one while it runs. **Interpretability** is\n",
        "the study of what a trained network is actually computing — turning a\n",
        "pile of matrices into mechanisms you can name.\n",
        "\n",
        "We start with the single most buildable interpretability tool: the\n",
        "**logit lens** (nostalgebraist, 2020). Recall that a decoder-only\n",
        "transformer keeps one vector per position — the **residual stream** —\n",
        "and refines it block by block, only *reading it out* to vocabulary\n",
        "logits at the very end:\n",
        "\n",
        "    logits = lm_head(ln_final(x))        # the model's own read-out head\n",
        "\n",
        "The logit lens applies that *same* read-out head to the residual stream\n",
        "at **every** layer, not just the last. Each intermediate state becomes a\n",
        "distribution over the vocabulary — *what would the model predict if it\n",
        "had to answer at this depth?* Watching that guess sharpen from layer to\n",
        "layer is the closest thing to watching the model think.\n",
        "\n",
        "Why it matters for LLMs:\n",
        "\n",
        "- **The model already predicts, early.** On many models the answer is\n",
        "  legible in the residual stream several layers before the output — the\n",
        "  later layers refine, not decide. The lens makes that visible.\n",
        "- **It reuses machinery you already built.** No new training, no model\n",
        "  surgery: the m06 `GPTModel` already exposes its per-layer hidden\n",
        "  states and ties its read-out to the token embedding. The lens is ten\n",
        "  lines on top.\n",
        "- **It is the honest start of a big field.** The tuned lens, direct\n",
        "  logit attribution, and circuit analysis all begin from “decode the\n",
        "  residual stream.” We build the first *two* from scratch here — the\n",
        "  logit lens, then direct logit attribution — and the rest have a\n",
        "  foothold.\n",
        "\n",
        "### What You’ll Learn\n",
        "\n",
        "After this module, you can:\n",
        "\n",
        "- Explain the **residual stream** as a running prediction the model\n",
        "  refines and reads out once.\n",
        "- Apply the model’s **own** `ln_final` + `lm_head` to any layer to get a\n",
        "  distribution — and see why the last layer’s lens equals the model’s\n",
        "  output *exactly*.\n",
        "- Measure a prediction’s **rank**, **entropy**, and **commit depth**\n",
        "  across layers, and track the **residual-stream norm**.\n",
        "- Train a tiny model to memorize a sentence and **watch its guess\n",
        "  crystallize** with depth.\n",
        "- State honestly what the logit lens can and cannot tell you (the\n",
        "  **bias** the tuned lens fixes).\n",
        "- Decompose a logit into an **exact sum of per-component contributions**\n",
        "  (direct logit attribution) and see *which* attention head or FFN wrote\n",
        "  the answer.\n",
        "- Use the **logit difference** to isolate what separates two candidate\n",
        "  tokens, and know why frozen-LN attribution is exact bookkeeping rather\n",
        "  than a causal claim.\n",
        "\n",
        "### Prerequisites\n",
        "\n",
        "This module requires familiarity with:\n",
        "\n",
        "- [Module 6: Transformer](../m06_transformer/lesson.qmd) — the residual\n",
        "  stream, `ln_final`, `lm_head`, and weight tying we read out.\n",
        "- [Module 7: Training](../m07_training/lesson.qmd) — we train a tiny\n",
        "  model to memorize one sentence for the demonstration.\n",
        "- [Module 8: Generation](../m08_generation/lesson.qmd) — logits →\n",
        "  softmax → a next-token prediction, applied here layer by layer.\n",
        "\n",
        "## Intuition: The Residual Stream as a Running Guess\n",
        "\n",
        "Picture the model as an assembly line. A token enters as an embedding\n",
        "vector. Each transformer block **adds** something to that vector —\n",
        "attention pulls in context, the FFN transforms it — but crucially, every\n",
        "block *writes back into the same running vector*. That vector, carried\n",
        "from the embedding all the way to the output, is the **residual\n",
        "stream**. It is the model’s working memory for that position.\n",
        "\n",
        "At the end of the line, one operation converts the vector into a\n",
        "prediction: normalize it (`ln_final`), then compare it against every\n",
        "token’s embedding (`lm_head`, whose weights *are* the token embeddings —\n",
        "that’s weight tying from m06). Tokens whose embedding points the same\n",
        "way as the residual stream get high logits.\n",
        "\n",
        "Here is the move. That converter — normalize, then compare to token\n",
        "embeddings — doesn’t care *which* residual-stream vector you hand it. So\n",
        "hand it the **half-finished** vector after block 1, after block 2, and\n",
        "so on. Each gives a distribution: the model’s best guess *so far*. Early\n",
        "layers are vague; later layers sharpen. The logit lens is exactly this —\n",
        "reading out the residual stream before it’s done.\n",
        "\n",
        "> **Key Insight**\n",
        ">\n",
        "> The logit lens adds **nothing** to the model. It reuses the model’s\n",
        "> own final normalization and unembedding, applied to states those\n",
        "> layers never actually see. That reuse is the whole trick — and, as\n",
        "> we’ll be careful to say later, also its main limitation.\n",
        "\n",
        "## The Math: Reading Out Any Layer\n",
        "\n",
        "Let $h_k$ be the residual stream after block $k$ (with $h_0$ the raw\n",
        "token+position embedding, and $L$ blocks total, so $h_L$ is the final\n",
        "state). The model’s output is\n",
        "\n",
        "$$\\text{logits} = \\text{lm\\_head}\\big(\\text{ln\\_final}(h_L)\\big).$$\n",
        "\n",
        "The **logit lens** at layer $k$ applies the identical read-out to $h_k$:\n",
        "\n",
        "$$\\text{lens}_k = \\text{lm\\_head}\\big(\\text{ln\\_final}(h_k)\\big),\n",
        "\\qquad\n",
        "p_k = \\text{softmax}(\\text{lens}_k).$$\n",
        "\n",
        "Two consequences fall straight out:\n",
        "\n",
        "1.  **The last layer is exact.** By definition $\\text{lens}_L$ is the\n",
        "    model’s own output — not an approximation, the *same tensor*. This\n",
        "    is the anchor: any correct lens implementation reproduces\n",
        "    `model(tokens)` at the final layer bit for bit.\n",
        "2.  **Weight tying makes the lens a similarity.** Since `lm_head.weight`\n",
        "    is the token-embedding matrix $E$,\n",
        "    $\\text{lens}_k = \\text{ln\\_final}(h_k)\\,E^\\top$ — the logit for\n",
        "    token $t$ is the (normalized) residual stream’s **dot product with\n",
        "    token $t$’s embedding**. The lens asks: *which token does this\n",
        "    vector look like?*\n",
        "\n",
        "From the per-layer distribution $p_k$ we read three numbers:\n",
        "\n",
        "- **Rank** of a target token $t$: how many tokens outrank it, plus one.\n",
        "  The shallowest layer where the rank hits 1 is the model’s **commit\n",
        "  depth** for that answer.\n",
        "- **Entropy** $H(p_k) = -\\sum_t p_k(t)\\log_2 p_k(t)$ (bits): high =\n",
        "  unsure, and it typically falls with depth as the guess sharpens.\n",
        "- **Residual-stream norm** $\\lVert h_k\\rVert$: it grows with depth,\n",
        "  because every block *adds* to the stream.\n",
        "\n",
        "### Step by Step\n",
        "\n",
        "Step through the read-out the lens performs at each layer:"
      ],
      "id": "22da87c8-f4a3-419f-a14f-d851ab9548ea"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "5d403386-2a69-416e-801c-f11baea549dc"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "c8b4b05e-642f-4042-8578-5923932e937c"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "b232d894-418e-4f33-9f77-c9b611bdfd8b"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Code: The Logit Lens from Scratch\n",
        "\n",
        "The whole lens lives in `interpretability.py`. Its heart is one function\n",
        "— `unembed` — the model’s read-out head applied to any hidden state:\n",
        "\n",
        "``` python\n",
        "def unembed(model, hidden):\n",
        "    return model.lm_head(model.ln_final(hidden))\n",
        "```\n",
        "\n",
        "Load a model and confirm the anchor: the lens at the final layer *is*\n",
        "the model’s output."
      ],
      "id": "177cc3c3-8be0-408a-8572-1e3dba5e0546"
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "lens tensor shape : (5, 1, 6, 48)\n",
            "final lens == model output : True"
          ]
        }
      ],
      "source": [
        "import importlib.util\n",
        "import sys\n",
        "from pathlib import Path\n",
        "\n",
        "import torch\n",
        "\n",
        "spec = importlib.util.spec_from_file_location(\n",
        "    \"interpretability\", Path(\"interpretability.py\").resolve()\n",
        ")\n",
        "I = importlib.util.module_from_spec(spec)\n",
        "sys.modules[\"interpretability\"] = I\n",
        "spec.loader.exec_module(I)\n",
        "\n",
        "torch.manual_seed(0)\n",
        "model = I.GPTModel(vocab_size=48, embed_dim=64, num_heads=4, num_layers=4, max_seq_len=32)\n",
        "tokens = torch.randint(0, 48, (1, 6))\n",
        "\n",
        "lens = I.layer_logits(model, tokens)          # (num_layers+1, batch, seq, vocab)\n",
        "model.eval()\n",
        "with torch.no_grad():\n",
        "    output = model(tokens)\n",
        "\n",
        "print(f\"lens tensor shape : {tuple(lens.shape)}\")\n",
        "print(f\"final lens == model output : {torch.equal(lens[-1], output)}\")"
      ],
      "id": "229f666a"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "`layer_logits` runs one forward pass with `return_hidden_states=True`,\n",
        "then unembeds each of the `num_layers + 1` residual-stream checkpoints\n",
        "(the embeddings plus one per block). The last slab equals\n",
        "`model(tokens)` exactly — that is the correctness anchor from the math\n",
        "above, made a `torch.equal`.\n",
        "\n",
        "Now the reader-facing view: the top predictions at the last position,\n",
        "layer by layer. (This model is untrained, so the tokens are arbitrary\n",
        "ids — we’re checking the *mechanism*; the meaning comes next.)"
      ],
      "id": "d43eaec4-f533-403c-93f0-6a786f83c9cf"
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            " emb: 38(0.04), 17(0.03), 18(0.03)\n",
            "  L1: 32(0.03), 12(0.03), 15(0.03)\n",
            "  L2: 32(0.03), 2(0.03), 39(0.03)\n",
            "  L3: 32(0.03), 6(0.03), 4(0.03)\n",
            "  L4: 32(0.03), 33(0.03), 39(0.02)"
          ]
        }
      ],
      "source": [
        "rows = I.logit_lens(model, tokens, position=-1, top_k=3)\n",
        "for layer, row in zip([\"emb\", \"L1\", \"L2\", \"L3\", \"L4\"], rows):\n",
        "    shown = \", \".join(f\"{tok}({p:.2f})\" for tok, p in row)\n",
        "    print(f\"{layer:>4}: {shown}\")"
      ],
      "id": "4754293a"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "And the three diagnostics — rank of a chosen token, entropy per layer,\n",
        "and the residual-stream norm:"
      ],
      "id": "28257300-7001-4b7b-aa10-86360dc04253"
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "rank of final top token : [7, 1, 1, 1, 1]\n",
            "entropy (bits) per layer : [5.56, 5.57, 5.57, 5.57, 5.57]\n",
            "residual norm per layer  : [0.2, 6.3, 8.6, 11.0, 14.2]"
          ]
        }
      ],
      "source": [
        "target = int(lens[-1, 0, -1].argmax())        # the model's final top token\n",
        "print(\"rank of final top token :\", I.prediction_rank(lens, target, position=-1))\n",
        "print(\"entropy (bits) per layer :\", [round(e, 2) for e in I.prediction_entropy(lens, position=-1)])\n",
        "\n",
        "model.eval()\n",
        "with torch.no_grad():\n",
        "    _, hidden = model(tokens, return_hidden_states=True)\n",
        "print(\"residual norm per layer  :\", [round(n, 1) for n in I.residual_norms(hidden)])"
      ],
      "id": "03ceb17a"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The final rank is always 1 (that token *is* the argmax by construction),\n",
        "and the residual norm climbs with depth because each block adds to the\n",
        "stream. On an untrained model the earlier layers are noise — to see the\n",
        "lens do something meaningful, we need a model that has actually learned.\n",
        "\n",
        "## Watch a Prediction Crystallize\n",
        "\n",
        "An untrained model’s intermediate lens is noise, so “watch the guess\n",
        "sharpen” would be a name-drop. Instead we make it real: train a tiny GPT\n",
        "for a couple of seconds to **memorize one sentence**, then lens its\n",
        "prediction for the missing final word. `demonstrate_logit_lens` does\n",
        "exactly this — build a tiny model, overfit it on\n",
        "\n",
        "> `the cat sat on mat and the dog ran to the park`\n",
        "\n",
        "feed it every word but the last, and read out the final position, where\n",
        "the answer is **park**."
      ],
      "id": "e6537679-d93b-4652-a32a-cedecce13390"
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [],
      "source": [],
      "id": "6f708e4d"
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "prompt : <bos> the cat sat on mat and the dog ran to the\n",
            "answer : park (loss after training: 0.0011)\n",
            "\n",
            "layer  top token  p(top)  rank(park)  H bits\n",
            "  emb        the    1.00           5    0.02\n",
            "   L1       park    0.94           1    0.50\n",
            "   L2       park    1.00           1    0.06\n",
            "   L3       park    1.00           1    0.02\n",
            "   L4       park    1.00           1    0.01"
          ]
        }
      ],
      "source": [
        "trace = I.demonstrate_logit_lens(seed=0)\n",
        "print(\"prompt :\", \" \".join(trace[\"prompt\"]))\n",
        "print(\"answer :\", trace[\"target\"][\"token\"], f\"(loss after training: {trace['final_loss']:.4f})\")\n",
        "print()\n",
        "print(f\"{'layer':>5} {'top token':>10} {'p(top)':>7} {'rank(park)':>11} {'H bits':>7}\")\n",
        "for lab, row, rank, ent, p in zip(\n",
        "    trace[\"layers\"], trace[\"top_k\"], trace[\"target_rank\"],\n",
        "    trace[\"entropy_bits\"], trace[\"target_prob\"]\n",
        "):\n",
        "    print(f\"{lab:>5} {row[0]['token']:>10} {row[0]['prob']:>7.2f} {rank:>11} {ent:>7.2f}\")"
      ],
      "id": "6254c341"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Read that table top to bottom. At the **embeddings** (`emb`), the lens\n",
        "just echoes the *current* input token, “the” — before any block has run,\n",
        "the residual stream is essentially the token embedding, and weight tying\n",
        "reads it back as itself, so “park” sits far down the ranking. Then a\n",
        "**single block** is enough: by `L1` the model has already surfaced\n",
        "**park** as its top guess with probability past 0.9. The remaining\n",
        "blocks don’t change the answer — they just make it certain, driving the\n",
        "probability to essentially 1.0 and the entropy toward 0. The model\n",
        "*decided* early and *committed* late.\n",
        "\n",
        "> **Key Insight**\n",
        ">\n",
        "> The interesting computation happens in the **first block or two**; the\n",
        "> rest is refinement. This “predict early, sharpen late” shape is\n",
        "> exactly what the logit lens was invented to reveal — and why looking\n",
        "> only at the output hides where the model actually did the work.\n",
        "\n",
        "The heatmap makes the whole trajectory legible at once: each row is a\n",
        "layer, each column a token the model considered, and the cell brightness\n",
        "is the probability. The final answer lights up as a bright column that\n",
        "switches on after the first block; the rank and entropy lines confirm\n",
        "the “decide early, commit late” story."
      ],
      "id": "b8490208-c4c8-4ab1-a325-d4f0db65f8e1"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "92d51dc1-c3ff-48d4-aac7-ca4cf36987ef"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Interactive Exploration\n",
        "\n",
        "Drive the layer index yourself. Pick a depth and see that layer’s full\n",
        "top-k distribution as a bar chart — the answer token highlighted —\n",
        "alongside its rank, probability, entropy, and residual-stream norm. Step\n",
        "from the embeddings to the output and watch **park** climb from nowhere\n",
        "to certainty."
      ],
      "id": "e9e195d2-1983-48cb-a1b4-0251f522d2f8"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "573b0499-d385-4a47-8377-743a94ad4b2a"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "e9908bac-6dbe-4c4d-9054-d0bb3990fff8"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "> **Try This**\n",
        ">\n",
        "> 1.  **Find the commit depth.** Step until `rank(\"park\")` first hits 1\n",
        ">     — that’s the layer where the model committed. How many blocks did\n",
        ">     it take?\n",
        "> 2.  **Watch entropy collapse.** Note the entropy at `emb` versus the\n",
        ">     output. Where does the biggest drop happen?\n",
        "> 3.  **Watch the norm grow.** The residual-stream norm climbs every\n",
        ">     layer even after the answer is locked in — the later blocks keep\n",
        ">     writing, they just stop *changing the ranking*.\n",
        "\n",
        "# Direct Logit Attribution: Who Wrote the Answer?\n",
        "\n",
        "The logit lens told us *when* the model’s guess for **park** appeared:\n",
        "it jumped to rank 1 after the first block. But it never told us *which\n",
        "part of the model* put it there. Was it an attention head pulling “park”\n",
        "in from earlier in the sentence? An FFN? The lens shows the running\n",
        "total; it cannot break that total into the individual writes that\n",
        "produced it.\n",
        "\n",
        "**Direct logit attribution (DLA)** does exactly that. It splits a\n",
        "token’s logit into one number per component — the embeddings, and each\n",
        "block’s attention and FFN write — so you can point at the head or FFN\n",
        "that pushed the answer up. Where the lens is a *cumulative* read-out,\n",
        "DLA is the *marginal* contribution of each write, attributed to\n",
        "attention versus FFN. Two views of the same event: the lens showed the\n",
        "answer crystallize; DLA shows which write made it win.\n",
        "\n",
        "## The Math: A Logit is a Sum of Contributions\n",
        "\n",
        "The whole method rests on one fact you built in m06: **the residual\n",
        "stream is a sum.** A pre-norm block adds to the stream twice —\n",
        "$x \\leftarrow x + \\text{attn}(\\text{ln}_1(x))$, then\n",
        "$x \\leftarrow x + \\text{ffn}(\\text{ln}_2(x))$ — so the final state\n",
        "unrolls into\n",
        "\n",
        "$$x_\\text{final} = \\underbrace{x_0}_{\\text{embed}} + \\sum_{k=1}^{L}\\Big(\\underbrace{a_k}_{\\text{attn}_k} + \\underbrace{f_k}_{\\text{ffn}_k}\\Big).$$\n",
        "\n",
        "The model reads that out with\n",
        "$\\text{logits} = \\text{lm\\_head}(\\text{ln\\_final}(x_\\text{final}))$. The\n",
        "read-out is *almost* linear — the only nonlinearity is `ln_final`, which\n",
        "divides by the per-position standard deviation $\\sigma$ of its input. So\n",
        "we **freeze** $\\sigma$ (and, for LayerNorm, the mean) at the value the\n",
        "true $x_\\text{final}$ produces. Frozen, `ln_final` is an *affine* map,\n",
        "and an affine map distributes over a sum. Writing $W_U$ for the\n",
        "unembedding ($=E^\\top$, weight tying), $\\gamma,\\beta$ for the LayerNorm\n",
        "gain and bias, and $\\bar\\mu_c$ for a component’s own feature-mean:\n",
        "\n",
        "$$\\text{logit}[t] \\;=\\; \\sum_{c}\\; \\underbrace{W_U[t]\\cdot\\!\\Big(\\gamma \\odot \\tfrac{c-\\bar\\mu_c}{\\sigma}\\Big)}_{\\text{contribution of component }c} \\;+\\; \\underbrace{W_U[t]\\cdot\\beta}_{\\text{bias}}.$$\n",
        "\n",
        "Each bracketed term is the **direct contribution** of one write to token\n",
        "$t$’s logit, and — this is the anchor — they sum *exactly* back to the\n",
        "model’s real logit. It is the DLA analogue of the lens’s\n",
        "`layer_logits[-1] == model(x)`: an identity a correct implementation\n",
        "reproduces to floating-point rounding.\n",
        "\n",
        "> **Key Insight**\n",
        ">\n",
        "> The residual stream being a **sum**, and the read-out being (once the\n",
        "> norm is frozen) **linear**, is the entire reason attribution works.\n",
        "> Every interpretability method that “reads a direction out of the\n",
        "> residual stream” — DLA, the tuned lens, activation patching — leans on\n",
        "> this same linearity.\n",
        "\n",
        "### Step by Step\n",
        "\n",
        "Step through how a single logit is decomposed into component\n",
        "contributions:"
      ],
      "id": "82db2bb7-7c5d-4e60-bab1-2eff2e815689"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "a3f44acf-8875-4bd9-ae0d-0841feadc5d9"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "d7d63f5c-d0d8-4de3-9641-ff8649cc6485"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "64d1d32a-d8ae-44b1-93df-b090329854e9"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Code: Attribution from Scratch\n",
        "\n",
        "Everything lives in `dla.py`. The first function decomposes the residual\n",
        "stream into its additive writes; the second reads any one of them out\n",
        "with the norm frozen. Load the module and confirm the two structural\n",
        "facts the math promised — the writes sum to the final stream, and the\n",
        "per-component contributions sum to the model’s true logit."
      ],
      "id": "06d8abba-74a2-4ba6-845f-845dabcade24"
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "components      : ['embed', 'L1.attn', 'L1.ffn', 'L2.attn', 'L2.ffn', 'L3.attn', 'L3.ffn']\n",
            "writes sum to x : True"
          ]
        }
      ],
      "source": [
        "import importlib.util\n",
        "import sys\n",
        "from pathlib import Path\n",
        "\n",
        "import torch\n",
        "\n",
        "spec = importlib.util.spec_from_file_location(\"dla\", Path(\"dla.py\").resolve())\n",
        "DLA = importlib.util.module_from_spec(spec)\n",
        "sys.modules[\"dla\"] = DLA\n",
        "spec.loader.exec_module(DLA)\n",
        "\n",
        "torch.manual_seed(0)\n",
        "model = DLA.GPTModel(vocab_size=48, embed_dim=64, num_heads=4, num_layers=3, max_seq_len=32)\n",
        "tokens = torch.randint(0, 48, (1, 6))\n",
        "\n",
        "labels, comps, x_final = DLA.residual_components(model, tokens)\n",
        "print(\"components      :\", labels)\n",
        "print(\"writes sum to x :\", torch.allclose(comps.sum(0), x_final, atol=1e-5))"
      ],
      "id": "6c024687"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "`residual_components` re-runs the model’s *own* sublayers, capturing\n",
        "each write — `embed`, then `L1.attn`, `L1.ffn`, `L2.attn`, … — and they\n",
        "sum back to the exact residual stream `ln_final` reads out. Now\n",
        "attribute a token’s logit:"
      ],
      "id": "fd302cf2-e52b-459f-9872-3420181b543b"
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "target token id : 11\n",
            "     embed: -0.004\n",
            "   L1.attn: +0.072\n",
            "    L1.ffn: +0.005\n",
            "   L2.attn: +0.084\n",
            "    L2.ffn: -0.001\n",
            "   L3.attn: +0.240\n",
            "    L3.ffn: +0.000\n",
            "      bias: +0.000\n",
            "\n",
            "Σ contributions + bias = 0.3957\n",
            "model's true logit     = 0.3957\n",
            "exact reconstruction   : True"
          ]
        }
      ],
      "source": [
        "out = DLA.direct_logit_attribution(model, tokens, position=-1)\n",
        "print(f\"target token id : {out['target']['id']}\")\n",
        "for lab, c in zip(out[\"labels\"], out[\"contributions\"]):\n",
        "    print(f\"  {lab:>8}: {c:+.3f}\")\n",
        "print(f\"  {'bias':>8}: {out['bias']:+.3f}\")\n",
        "print(f\"\\nΣ contributions + bias = {out['reconstructed']:.4f}\")\n",
        "print(f\"model's true logit     = {out['model_logit']:.4f}\")\n",
        "print(f\"exact reconstruction   : {abs(out['reconstructed'] - out['model_logit']) < 1e-3}\")"
      ],
      "id": "03585756"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "That final line is the correctness anchor: the decomposition is\n",
        "**lossless**. Every logit the model produces is, exactly, the sum of\n",
        "what each component wrote plus the constant bias. (This model is\n",
        "untrained, so the contributions are arbitrary — the *mechanism* is what\n",
        "we’re checking. The meaning comes when we attribute a trained model’s\n",
        "real prediction below.)\n",
        "\n",
        "### The logit difference: what separates two candidates\n",
        "\n",
        "A raw logit mixes two things: how much the model likes this token, and a\n",
        "shared “confidence” push that raises *every* token. To isolate the\n",
        "first, compare the answer against a **distractor** — attribute the\n",
        "*logit difference*\n",
        "$\\text{logit}[\\text{target}] - \\text{logit}[\\text{distractor}]$. This\n",
        "projects each write onto the direction\n",
        "$W_U[\\text{target}] - W_U[\\text{distractor}]$, so any part of a write\n",
        "that lifts both tokens equally drops out."
      ],
      "id": "58f5edc1-6290-4396-ab7f-81e6a41d2ab3"
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "contribution to logit[3] - logit[7]:\n",
            "     embed: +0.011\n",
            "   L1.attn: -0.000\n",
            "    L1.ffn: -0.009\n",
            "   L2.attn: +0.024\n",
            "    L2.ffn: -0.002\n",
            "   L3.attn: +0.040\n",
            "    L3.ffn: +0.002\n",
            "      bias: +0.000\n",
            "reconstructed = 0.0670  vs true diff = 0.0670"
          ]
        }
      ],
      "source": [
        "diff = DLA.logit_diff_attribution(model, tokens, target_id=3, distractor_id=7)\n",
        "print(\"contribution to logit[3] - logit[7]:\")\n",
        "for lab, c in zip(diff[\"labels\"], diff[\"contributions\"]):\n",
        "    print(f\"  {lab:>8}: {c:+.3f}\")\n",
        "print(f\"  {'bias':>8}: {diff['bias']:+.3f}\")\n",
        "print(f\"reconstructed = {diff['reconstructed']:.4f}  vs true diff = {diff['model_logit_diff']:.4f}\")"
      ],
      "id": "a2d8a7d8"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The `bias` term is the LayerNorm $\\beta$’s per-token offset,\n",
        "$W_U[\\text{target}]\\cdot\\beta - W_U[\\text{distractor}]\\cdot\\beta$. On\n",
        "this *fresh* model it reads `0.000` — because `nn.LayerNorm` initializes\n",
        "$\\beta$ to zero — which is a handy sanity check. But $\\beta$ is a\n",
        "learned parameter: once the model trains it drifts, and the term becomes\n",
        "genuinely nonzero (you’ll see it below). It is *per-token*, not a global\n",
        "constant, so it does **not** cancel in the difference. We report it so\n",
        "the reconstruction stays exact — honesty over a tidy-but-wrong “it\n",
        "cancels.”\n",
        "\n",
        "## Which Head Wrote It?\n",
        "\n",
        "Attention is not one write — it is a sum over heads. A block’s attention\n",
        "output is `out_proj(concat(head₀, …, head_{H-1}))`, and because\n",
        "concatenation lays the heads side by side and `out_proj` is linear, head\n",
        "$h$’s contribution is `out_proj` applied to *only* head $h$’s slice of\n",
        "the concatenation. So we can split a block’s attention contribution\n",
        "across its heads and ask the sharpest interpretability question there\n",
        "is: **which head wrote the answer?**"
      ],
      "id": "aae84956-4b03-4c33-b5bc-68d87fefaa79"
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "layer 0 attention, per-head contribution to the target logit:\n",
            "  head 0: -0.039\n",
            "  head 1: +0.042\n",
            "  head 2: +0.021\n",
            "  head 3: +0.048\n",
            "  bias   : +0.000\n",
            "\n",
            "Σ heads + bias      = 0.0719\n",
            "whole L1.attn write = 0.0719\n",
            "match               : True"
          ]
        }
      ],
      "source": [
        "heads = DLA.attention_head_attribution(model, tokens, layer=0, target_id=int(out[\"target\"][\"id\"]))\n",
        "print(f\"layer {heads['layer']} attention, per-head contribution to the target logit:\")\n",
        "for h, c in enumerate(heads[\"head_contributions\"]):\n",
        "    print(f\"  head {h}: {c:+.3f}\")\n",
        "print(f\"  bias   : {heads['bias']:+.3f}\")\n",
        "print(f\"\\nΣ heads + bias      = {heads['total']:.4f}\")\n",
        "print(f\"whole L1.attn write = {heads['attn_contribution']:.4f}\")\n",
        "print(f\"match               : {abs(heads['total'] - heads['attn_contribution']) < 1e-3}\")"
      ],
      "id": "ceeb2c34"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The head contributions plus the shared `out_proj` bias sum to the\n",
        "block’s whole attention contribution — the same per-head decomposition\n",
        "Elhage et al. use to name the heads inside a circuit.\n",
        "\n",
        "## Watch the Components Vote\n",
        "\n",
        "Now the payoff. `demonstrate_dla` trains the tiny GPT to memorize the\n",
        "*same* sentence the logit lens used, then attributes its prediction for\n",
        "the missing final word **park** — against the distractor **sun**. Where\n",
        "the lens showed **park** climb to rank 1 after block 1, DLA shows the\n",
        "individual writes that lifted it."
      ],
      "id": "c5a3da2c-33e2-4b8e-ba26-6a2abd430cc2"
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {},
      "outputs": [],
      "source": [],
      "id": "a1435aba"
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "prompt : <bos> the cat sat on mat and the dog ran to the\n",
            "answer : park (loss after training: 0.0011)\n",
            "\n",
            "component    → park\n",
            "    embed    -0.015\n",
            "  L1.attn    +1.688\n",
            "   L1.ffn    +0.247\n",
            "  L2.attn    +1.559\n",
            "   L2.ffn    +0.991\n",
            "  L3.attn    +1.315\n",
            "   L3.ffn    +1.206\n",
            "  L4.attn    +0.601\n",
            "   L4.ffn    +1.687\n",
            "     bias    -0.035\n",
            "\n",
            "reconstructed = 9.243  ==  model logit = 9.243"
          ]
        }
      ],
      "source": [
        "trace = DLA.demonstrate_dla(seed=0)\n",
        "print(\"prompt :\", \" \".join(trace[\"prompt\"]))\n",
        "print(\"answer :\", trace[\"target\"][\"token\"],\n",
        "      f\"(loss after training: {trace['final_loss']:.4f})\")\n",
        "print()\n",
        "print(f\"{'component':>9} {'→ park':>9}\")\n",
        "for lab, c in zip(trace[\"labels\"], trace[\"target_logit\"][\"contributions\"]):\n",
        "    print(f\"{lab:>9} {c:>+9.3f}\")\n",
        "tl = trace[\"target_logit\"]\n",
        "print(f\"{'bias':>9} {tl['bias']:>+9.3f}\")\n",
        "print(f\"\\nreconstructed = {tl['reconstructed']:.3f}  ==  model logit = {tl['model_logit']:.3f}\")"
      ],
      "id": "640240c2"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Two things jump out of that column. First, the **embeddings contribute\n",
        "almost nothing** to “park”: the raw input at the final position is the\n",
        "word “the”, whose embedding points at “the”, not “park” — the model has\n",
        "to *compute* the answer. Second, **every block writes a positive push**,\n",
        "but the biggest ones come from **attention in the early layers** (which\n",
        "pull “park” in from context) and the **FFN in the last layer** (which\n",
        "sharpens it) — not from any single place. The bar chart below makes the\n",
        "votes visible; green pushes “park” up, red would push it down, and the\n",
        "biggest contributor is ringed."
      ],
      "id": "ed621bd0-b614-4c63-9229-a38cf1c91161"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "e6e5494b-933f-468d-98a1-d60626dc848f"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "And the per-head split of each layer’s attention — the finest grain DLA\n",
        "reaches here. Each row is a layer; each cell is one head’s push on\n",
        "“park”, so a bright cell is a head that did real work."
      ],
      "id": "648124bc-e1e0-47e5-b025-bd787d303000"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "76bc341a-5dd2-42b6-b1db-3fba11a38782"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Interactive Exploration: Attribution\n",
        "\n",
        "Switch between the two views. **Target logit** shows what pushed “park”\n",
        "up in absolute terms; **logit difference** shows what pushed “park”\n",
        "*over* the distractor “sun” — the same bars, but with the shared “raise\n",
        "everything” component removed. Watch which writes stay large (they\n",
        "genuinely separate the candidates) and which shrink (they were just\n",
        "raising confidence)."
      ],
      "id": "53746e1f-1313-4d1a-8b1d-7ce59b31fa9d"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "8475dea1-a6d4-48eb-a6b7-03ebc5cc6164"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "0c7e72c4-1a47-4116-967d-79d2020b8dcf"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "> **Try This**\n",
        ">\n",
        "> 1.  **Find the author.** In “target logit” view, which single\n",
        ">     component has the tallest green bar? Is it attention or an FFN,\n",
        ">     and in which layer?\n",
        "> 2.  **Switch to the difference.** Flip to “logit difference (park −\n",
        ">     sun)”. Which bars barely move (they separate the candidates) and\n",
        ">     which shrink (they were just raising every token)?\n",
        "> 3.  **Check the books.** Read the header line: Σ contributions ± bias\n",
        ">     always equals the model’s true value. Attribution here is exact,\n",
        ">     not approximate.\n",
        "\n",
        "## Common Pitfalls\n",
        "\n",
        "> **The logit lens is a biased probe, not ground truth**\n",
        ">\n",
        "> The lens applies the *final* `ln_final` + `lm_head` to intermediate\n",
        "> states those layers were never trained to be read by. On GPT-2-style\n",
        "> models it works well, but on others it can be misleading — an early\n",
        "> layer’s “prediction” may be an artifact of forcing the final head onto\n",
        "> an unfinished vector. Treat the lens as *suggestive*. The **tuned\n",
        "> lens** (Belrose et al., 2023) fixes this by *learning* a small affine\n",
        "> probe per layer; it is the rigorous successor.\n",
        "\n",
        "> **Norm growth is not meaning**\n",
        ">\n",
        "> The residual-stream norm rising with depth is mostly bookkeeping —\n",
        "> blocks add to the stream, so it grows. A larger norm does **not** mean\n",
        "> a layer is “more important.” Use rank and entropy for what the model\n",
        "> *predicts*; use norm only as a sanity signal.\n",
        "\n",
        "> **Softmax before you compare probabilities**\n",
        ">\n",
        "> Raw lens logits are not comparable across layers — their scale drifts.\n",
        "> Always `softmax` first (as `logit_lens` and the diagnostics do) before\n",
        "> reading a probability or an entropy.\n",
        "\n",
        "> **Read the right position**\n",
        ">\n",
        "> Next-token predictions live at the **last** position of the prompt.\n",
        "> Lensing an interior position tells you what the model would predict\n",
        "> *there*, which is a different question. `position=-1` is the usual\n",
        "> choice.\n",
        "\n",
        "> **Frozen-LN attribution is exact, but it is a convention, not a\n",
        "> counterfactual**\n",
        ">\n",
        "> DLA is *exactly* additive only because we freeze `ln_final`’s scale at\n",
        "> the true `x_final`. A component’s attributed push is therefore “its\n",
        "> effect on the logit *given the normalization the whole model\n",
        "> produced*.” It is **not** what you’d get by deleting that component\n",
        "> and re-running — the norm would then rescale everything. Treat DLA as\n",
        "> an honest bookkeeping of the final logit, and reach for **activation\n",
        "> patching** when you need a true causal (delete-and-rerun) answer.\n",
        "\n",
        "> **Attribution is not causation**\n",
        ">\n",
        "> A component with a big positive contribution *wrote toward* the\n",
        "> answer, but it may have done so *because* an earlier component set up\n",
        "> the residual stream that way. DLA tells you the final tally, not the\n",
        "> chain of cause. Naming a **circuit** means tracing those dependencies\n",
        "> (which head reads which earlier write), not just ranking bars.\n",
        "\n",
        "> **Prefer the logit difference for a decision**\n",
        ">\n",
        "> A raw-logit attribution mixes “likes this token” with a shared “raise\n",
        "> every token” push. When you care *why the model picked A over B*,\n",
        "> attribute the **logit difference** — it removes the common component\n",
        "> and keeps only what separated them.\n",
        "\n",
        "## Exercises\n",
        "\n",
        "### Exercise 1: Commit depth\n",
        "\n",
        "The **commit depth** is the shallowest layer where a target token\n",
        "reaches rank 1. Compute it from a trace’s `target_rank` list."
      ],
      "id": "8a233eaa-0088-463f-acf2-c4e5036b2e42"
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "target ranks per layer: [5, 1, 1, 1, 1]\n",
            "'park' commits at layer index 1 (L1)"
          ]
        }
      ],
      "source": [
        "trace = I.demonstrate_logit_lens(seed=0)\n",
        "\n",
        "def commit_depth(ranks):\n",
        "    # Your implementation here: return the first index where rank == 1 (else None).\n",
        "    for i, r in enumerate(ranks):\n",
        "        if r == 1:\n",
        "            return i\n",
        "    return None\n",
        "\n",
        "depth = commit_depth(trace[\"target_rank\"])\n",
        "print(\"target ranks per layer:\", trace[\"target_rank\"])\n",
        "print(f\"'{trace['target']['token']}' commits at layer index {depth} \"\n",
        "      f\"({trace['layers'][depth]})\")"
      ],
      "id": "bcc2e7e3"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Exercise 2: Prediction sharpening\n",
        "\n",
        "Compute the **entropy drop** between consecutive layers and find the\n",
        "block that sharpens the prediction most. Does it line up with the commit\n",
        "depth?"
      ],
      "id": "29577d90-2591-4d54-a63e-41e7f4437ac2"
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "entropy drop after each block: [-0.477, 0.435, 0.039, 0.007]"
          ]
        }
      ],
      "source": [
        "ent = trace[\"entropy_bits\"]\n",
        "drops = [round(ent[i] - ent[i + 1], 3) for i in range(len(ent) - 1)]\n",
        "print(\"entropy drop after each block:\", drops)\n",
        "# Your turn: which block index has the largest drop? Compare it to commit_depth above."
      ],
      "id": "988ab06f"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Exercise 3: Reconstruct the logit from its parts\n",
        "\n",
        "The DLA anchor is that the contributions plus the bias equal the model’s\n",
        "true logit. Verify it yourself, and find the single component that\n",
        "pushed the answer hardest."
      ],
      "id": "3032dcd3-c485-4db2-a307-84610aa14884"
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "reconstructed = 9.243  vs model = 9.243\n",
            "biggest contributor to 'park': L1.attn (+1.69)"
          ]
        }
      ],
      "source": [
        "trace = DLA.demonstrate_dla(seed=0)\n",
        "contribs = trace[\"target_logit\"][\"contributions\"]\n",
        "bias = trace[\"target_logit\"][\"bias\"]\n",
        "\n",
        "# Your implementation: sum the contributions and the bias, and compare to the model.\n",
        "total = sum(contribs) + bias\n",
        "print(f\"reconstructed = {total:.3f}  vs model = {trace['target_logit']['model_logit']:.3f}\")\n",
        "\n",
        "top = trace[\"labels\"][contribs.index(max(contribs))]\n",
        "print(f\"biggest contributor to '{trace['target']['token']}': {top} (+{max(contribs):.2f})\")"
      ],
      "id": "0962807f"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Exercise 4: The most important head\n",
        "\n",
        "Attribute the answer to a chosen layer’s attention heads and find the\n",
        "head that wrote the most. Does the same head dominate in every layer?"
      ],
      "id": "edb31300-ba67-4a3d-8c49-b8977c477c00"
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "L1: heads = [0.29, 0.59, 0.59, 0.21]  →  top head = 2\n",
            "L2: heads = [1.04, 0.24, 0.18, 0.1]  →  top head = 0\n",
            "L3: heads = [0.24, 0.32, 0.68, 0.07]  →  top head = 2\n",
            "L4: heads = [0.3, 0.18, -0.02, 0.15]  →  top head = 0"
          ]
        }
      ],
      "source": [
        "model = DLA._train_toy_model(seed=0, train_steps=400, lr=3e-3,\n",
        "                             embed_dim=64, num_layers=4, num_heads=4)\n",
        "prompt = torch.tensor([DLA.TOY_SENTENCE[:-1]])\n",
        "answer = DLA.TOY_SENTENCE[-1]\n",
        "\n",
        "for layer in range(4):\n",
        "    h = DLA.attention_head_attribution(model, prompt, layer=layer, target_id=answer)\n",
        "    hc = h[\"head_contributions\"]\n",
        "    print(f\"L{layer+1}: heads = {[round(c, 2) for c in hc]}  →  top head = {hc.index(max(hc))}\")\n",
        "# Your turn: is it always the same head index, or does the \"author\" move by layer?"
      ],
      "id": "8554f0d7"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Exercise 5: Target logit vs logit difference\n",
        "\n",
        "Compare a component’s raw contribution to its contribution *after*\n",
        "removing a distractor. Which components shrink most — the ones that were\n",
        "only raising overall confidence?"
      ],
      "id": "0f62a942-37c0-4f23-b39f-b6d84b85b8c3"
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "component      raw     diff  shrunk by\n",
            "    embed    -0.01    +0.04      -0.06\n",
            "  L1.attn    +1.69    +1.60      +0.08\n",
            "   L1.ffn    +0.25    +0.28      -0.03\n",
            "  L2.attn    +1.56    +1.52      +0.04\n",
            "   L2.ffn    +0.99    +1.05      -0.06\n",
            "  L3.attn    +1.31    +1.25      +0.06\n",
            "   L3.ffn    +1.21    +1.31      -0.10\n",
            "  L4.attn    +0.60    +0.75      -0.15\n",
            "   L4.ffn    +1.69    +1.73      -0.04"
          ]
        }
      ],
      "source": [
        "raw = DLA.direct_logit_attribution(model, prompt, target_id=answer)\n",
        "diff = DLA.logit_diff_attribution(model, prompt, target_id=answer,\n",
        "                                  distractor_id=DLA.TOY_VOCAB.index(\"sun\"))\n",
        "print(f\"{'component':>9} {'raw':>8} {'diff':>8} {'shrunk by':>10}\")\n",
        "for lab, r, d in zip(raw[\"labels\"], raw[\"contributions\"], diff[\"contributions\"]):\n",
        "    print(f\"{lab:>9} {r:>+8.2f} {d:>+8.2f} {r - d:>+10.2f}\")\n",
        "# Your turn: a large \"shrunk by\" means that write mostly raised confidence, not the choice."
      ],
      "id": "2c16489a"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Summary\n",
        "\n",
        "Key takeaways:\n",
        "\n",
        "1.  **The residual stream is a running prediction.** A transformer\n",
        "    refines one vector per position through its blocks and reads it out\n",
        "    to logits only at the end.\n",
        "2.  **The logit lens reads it out early.** Apply the model’s own\n",
        "    `ln_final` + `lm_head` to any layer’s hidden state to get that\n",
        "    layer’s guess — no new training, no model surgery.\n",
        "3.  **The last layer’s lens is the model’s output, exactly.** That\n",
        "    identity (`layer_logits(model, x)[-1] == model(x)`) is the\n",
        "    correctness anchor.\n",
        "4.  **Weight tying makes the lens a similarity.** A token’s lens logit\n",
        "    is the normalized residual stream’s dot product with that token’s\n",
        "    embedding — *which token does this vector look like?*\n",
        "5.  **Models decide early and commit late.** On our memorizing model the\n",
        "    answer reaches rank 1 after a single block, then the remaining\n",
        "    blocks only sharpen probability toward 1 and drive entropy toward 0.\n",
        "6.  **Rank, entropy, and norm are the read-outs.** Rank gives commit\n",
        "    depth, entropy gives confidence, and the residual norm grows with\n",
        "    depth (bookkeeping, not meaning).\n",
        "7.  **The lens is biased, and honest about it.** It forces the final\n",
        "    head onto states it never saw; the tuned lens learns a per-layer\n",
        "    probe to do it right.\n",
        "8.  **A logit is a sum of component contributions.** Because the\n",
        "    residual stream is a sum and the read-out is affine (once\n",
        "    `ln_final`’s scale is frozen), a token’s logit decomposes *exactly*\n",
        "    into one contribution per write — embeddings, and each block’s\n",
        "    attention and FFN. `Σ contributions + bias == model logit`.\n",
        "9.  **Direct logit attribution says who wrote the answer.** DLA turns\n",
        "    that decomposition into a signed bar per component, and — since\n",
        "    attention is a sum over heads — down to *which head* did the work.\n",
        "    On our memorizing model the embeddings contribute ~0, and\n",
        "    early-layer attention plus the last FFN write “park”.\n",
        "10. **Use the logit difference to isolate a decision.** Attributing\n",
        "    `logit[target] − logit[distractor]` projects onto\n",
        "    `W_U[target] − W_U[distractor]`, removing the shared “raise\n",
        "    everything” push and keeping only what *separated* the candidates.\n",
        "\n",
        "## What’s Next\n",
        "\n",
        "You have gone from *what* the residual stream predicts (the logit lens)\n",
        "to *which component wrote it* (direct logit attribution) — both by\n",
        "decoding the same running vector. The frontier from here trades exact\n",
        "bookkeeping for causal claims:\n",
        "\n",
        "- The **tuned lens** (Belrose et al., 2023) — a *learned* per-layer\n",
        "  probe that fixes the logit lens’s bias.\n",
        "- **Activation patching** — replace one component’s activation with\n",
        "  another run’s and measure the effect, the causal counterpart to DLA’s\n",
        "  bookkeeping.\n",
        "- **Induction heads and circuits** (Olsson et al., 2022; Elhage et\n",
        "  al., 2021) — trace *which* head reads *which* earlier write, naming\n",
        "  the mechanism DLA only ranks.\n",
        "- **Sparse autoencoders / features** — decompose a single write into\n",
        "  interpretable directions, going below the per-component grain this\n",
        "  module reaches.\n",
        "\n",
        "Each begins exactly where this module ends — reading structure out of\n",
        "the residual stream.\n",
        "\n",
        "### Going Deeper\n",
        "\n",
        "**Core Ideas:**\n",
        "\n",
        "- [interpreting GPT: the logit\n",
        "  lens](https://www.lesswrong.com/posts/AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens)\n",
        "  — nostalgebraist (2020), the original logit lens.\n",
        "- [A Mathematical Framework for Transformer\n",
        "  Circuits](https://transformer-circuits.pub/2021/framework/index.html)\n",
        "  — Elhage et al. (2021), the residual-stream view, the unembedding as a\n",
        "  read-out, and **direct logit attribution**.\n",
        "- [Eliciting Latent Predictions from Transformers with the Tuned\n",
        "  Lens](https://arxiv.org/abs/2303.08112) — Belrose et al. (2023), the\n",
        "  learned per-layer probe that fixes the logit lens’s bias.\n",
        "- [In-context Learning and Induction\n",
        "  Heads](https://arxiv.org/abs/2209.11895) — Olsson et al. (2022), the\n",
        "  mechanism behind in-context learning.\n",
        "\n",
        "**Practical Resources:**\n",
        "\n",
        "- [A Comprehensive Mechanistic Interpretability Explainer &\n",
        "  Glossary](https://www.neelnanda.io/mechanistic-interpretability/glossary)\n",
        "  — Neel Nanda, on the logit-difference view and direct logit\n",
        "  attribution in practice."
      ],
      "id": "3611cdce-db84-450d-a034-4448a859e8a7"
    }
  ],
  "nbformat": 4,
  "nbformat_minor": 5,
  "metadata": {
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "path": "/opt/buildhome/.asdf/installs/python/3.11.14/share/jupyter/kernels/python3"
    },
    "language_info": {
      "name": "python",
      "codemirror_mode": {
        "name": "ipython",
        "version": "3"
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.14"
    }
  }
}