{ "cells": [ { "cell_type": "markdown", "id": "8eab4211", "metadata": {}, "source": [ "# Looking inside a protein language model\n", "## Attention, contacts, and causal probing of ESM-2\n", "\n", "**Format:** 75-minute hands-on session  |  **Level:** upper-level undergraduate  |  **Runs on:** a laptop CPU (no GPU needed) or Google Colab\n", "\n", "ESM-2 is a transformer trained on ~65 million protein sequences with a single objective: *fill in masked amino acids*. It never sees a 3D structure. Yet its internal attention patterns turn out to encode which residues touch each other in the folded protein. Today we will open the model up and ask **where** that knowledge lives and **how we can tell**.\n", "\n", "### Learning objectives\n", "By the end of this session you should be able to:\n", "1. Explain the masked-language-modelling objective and what an attention head computes.\n", "2. Extract attention maps from ESM-2 and show that a small number of heads encode residue–residue contacts.\n", "3. Distinguish **correlational** evidence (attention patterns, probes) from **causal** evidence (ablations), and explain why both are needed.\n", "4. Train a linear probe for secondary structure and evaluate it with proper controls (baselines, held-out proteins, shuffled labels).\n", "\n", "### Schedule\n", "| Time | Part | Exercises |\n", "|---|---|---|\n", "| 0–5 min | 0. Setup | – |\n", "| 5–12 | 1. ESM-2 as a masked language model | Ex 1 |\n", "| 12–25 | 2. Opening the hood: attention maps | Ex 2, 3 |\n", "| 25–42 | 3. Which heads \"see\" 3D contacts? | Ex 4, 5, 6 |\n", "| 42–60 | 4. Causal interventions: ablating heads and MLPs | Ex 7, 8, 9 |\n", "| 60–72 | 5. Probing for secondary structure | Ex 10 |\n", "| 72–75 | Wrap-up and discussion | – |\n", "| (homework) | Optional extensions A–C | – |\n", "\n", "### How to work through this notebook\n", "* Cells marked **Exercise** contain blanks written as `___`. Replace each `___` with code, then run the cell. Most exercises are followed by a **self-check** that prints `OK` when your answer is right.\n", "* Stuck? Expand the **Hint** below an exercise. Don't spend more than ~3 minutes stuck — ask a neighbour.\n", "* Run the cells **in order**; later parts reuse variables from earlier ones.\n", "* Questions marked **Discuss** are for your group; jot short answers in the markdown cells provided." ] }, { "cell_type": "markdown", "id": "630fe5df", "metadata": {}, "source": [ "---\n", "## Part 0 · Setup (5 min)\n", "\n", "We use Meta's `fair-esm` package and the 35-million-parameter ESM-2 model (`esm2_t12_35M_UR50D`: **12 layers × 20 attention heads**, hidden size 480). It is small enough to run on a CPU in well under a second per protein, but large enough to show the phenomena we care about." ] }, { "cell_type": "code", "execution_count": null, "id": "73db898d", "metadata": {}, "outputs": [], "source": [ "import os\n", "os.environ.setdefault(\"KMP_DUPLICATE_LIB_OK\", \"TRUE\") # avoids an OpenMP clash on some Windows/conda installs\n", "\n", "import importlib.util, subprocess, sys\n", "if importlib.util.find_spec(\"esm\") is None: # install fair-esm if missing (e.g. on Colab)\n", " subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", \"fair-esm\"])\n", "\n", "import time, contextlib, urllib.request\n", "import numpy as np\n", "import torch\n", "import matplotlib.pyplot as plt\n", "import esm\n", "import sklearn\n", "from sklearn.linear_model import LogisticRegression\n", "from sklearn.preprocessing import StandardScaler\n", "from sklearn.pipeline import make_pipeline\n", "from scipy.stats import rankdata\n", "\n", "torch.manual_seed(0); np.random.seed(0)\n", "torch.set_grad_enabled(False) # we never train the transformer today\n", "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", "plt.rcParams.update({\"figure.dpi\": 110, \"axes.spines.top\": False, \"axes.spines.right\": False})\n", "print(f\"torch {torch.__version__} | fair-esm {esm.__version__} | scikit-learn {sklearn.__version__} | device: {device}\")" ] }, { "cell_type": "code", "execution_count": null, "id": "280329af", "metadata": {}, "outputs": [], "source": [ "t0 = time.time()\n", "model, alphabet = esm.pretrained.esm2_t12_35M_UR50D() # downloads ~140 MB the first time\n", "model = model.eval().to(device)\n", "batch_converter = alphabet.get_batch_converter()\n", "\n", "N_LAYERS, N_HEADS, D_MODEL = model.num_layers, model.attention_heads, model.embed_dim\n", "HEAD_DIM = D_MODEL // N_HEADS\n", "AAS = \"ACDEFGHIKLMNPQRSTVWY\"\n", "n_params = sum(p.numel() for p in model.parameters())\n", "print(f\"Loaded in {time.time()-t0:.1f} s: {N_LAYERS} layers x {N_HEADS} heads, d_model={D_MODEL}, \"\n", " f\"head_dim={HEAD_DIM}, {n_params/1e6:.1f} M parameters\")\n", "print(model.layers[0])" ] }, { "cell_type": "markdown", "id": "f33cf76e", "metadata": {}, "source": [ "### A one-minute refresher on the architecture\n", "\n", "A protein of length $L$ becomes $L+2$ tokens: a start token `` (index 0), one token per residue, and an end token ``. Each token carries a vector in the **residue stream** $x \\in \\mathbb{R}^{480}$. Every layer adds two updates to that stream:\n", "\n", "$$x \\leftarrow x + \\text{Attn}(\\text{LN}(x)), \\qquad x \\leftarrow x + \\text{MLP}(\\text{LN}(x))$$\n", "\n", "Inside `Attn`, each of the 20 heads computes an $L\\times L$ **attention matrix**\n", "\n", "$$A^{(h)} = \\text{softmax}\\!\\left(\\frac{Q^{(h)} K^{(h)\\top}}{\\sqrt{d_{head}}}\\right), \\qquad \\text{head output} = A^{(h)} V^{(h)},$$\n", "\n", "where row $i$ says *how much residue $i$ reads from every other residue*. The 20 head outputs are concatenated and mixed by a linear layer (`out_proj`). The MLP (`fc1 → GELU → fc2`) then acts on each position separately.\n", "\n", "Two things to keep in mind for later:\n", "* the **attention pattern** $A$ (where a head looks) and the **head output** $AV$ (what it writes) are different objects;\n", "* everything a layer writes is *added* to the residue stream, so later layers can read it — or ignore it.\n", "\n", "The helper below runs the model once and returns everything we need." ] }, { "cell_type": "code", "execution_count": null, "id": "c2409c01", "metadata": {}, "outputs": [], "source": [ "def run_esm(seq):\n", " \"\"\"Run ESM-2 on one sequence. Returns logits, attention maps, per-layer representations and ESM's own contact map.\"\"\"\n", " _, _, toks = batch_converter([(\"protein\", seq)])\n", " out = model(toks.to(device), repr_layers=list(range(N_LAYERS + 1)),\n", " need_head_weights=True, return_contacts=True)\n", " return {\n", " \"tokens\": toks[0],\n", " \"logits\": out[\"logits\"][0].cpu(), # (L+2, vocab)\n", " \"attentions\": out[\"attentions\"][0].cpu().numpy(), # (layers, heads, L+2, L+2) -- includes /\n", " \"contacts\": out[\"contacts\"][0].cpu().numpy(), # (L, L) ESM's built-in contact predictor\n", " \"reps\": {l: r[0, 1:-1].cpu().numpy() # layer -> (L, d_model); layer 0 = embeddings\n", " for l, r in out[\"representations\"].items()},\n", " }\n", "\n", "# Human ubiquitin (76 aa) -- our running example\n", "ub = \"MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG\"\n", "out = run_esm(ub)\n", "print(\"attentions:\", out[\"attentions\"].shape, \"| contacts:\", out[\"contacts\"].shape, \"| layer-12 reps:\", out[\"reps\"][12].shape)" ] }, { "cell_type": "markdown", "id": "4c9c2b33", "metadata": {}, "source": [ "---\n", "## Part 1 · ESM-2 as a masked language model (7 min)\n", "\n", "During training, ~15% of residues are hidden behind a `` token and the model must predict them from context. Let's play that game ourselves.\n", "\n", "### Exercise 1 — predict a masked residue\n", "Complete `predict_masked` so that it replaces residue `pos` (0-indexed position in the sequence) with the mask token. Remember that token 0 is ``, so residue `i` sits at **token index `i + 1`**. The mask token's index is `alphabet.mask_idx`." ] }, { "cell_type": "code", "execution_count": null, "id": "65cd5c68", "metadata": {}, "outputs": [], "source": [ "def predict_masked(seq, pos, top_k=5):\n", " \"\"\"Mask residue `pos` (0-indexed) and return the model's top-k amino acids with probabilities.\"\"\"\n", " _, _, toks = batch_converter([(\"protein\", seq)])\n", " masked = toks.clone()\n", " # TODO: put the mask token at the right token index\n", " masked[0, ___] = ___\n", " logits = model(masked.to(device))[\"logits\"][0, pos + 1].cpu()\n", " aa_idx = [alphabet.get_idx(a) for a in AAS]\n", " p = torch.softmax(logits[aa_idx], dim=-1) # renormalise over the 20 amino acids\n", " order = torch.argsort(p, descending=True)[:top_k]\n", " return [(AAS[i], round(float(p[i]), 3)) for i in order]\n", "\n", "# --- self-check ---\n", "assert ub[43] == \"I\", \"residue 44 of ubiquitin is Ile44 (0-indexed position 43)\"\n", "_, _, _toks = batch_converter([(\"protein\", ub)])\n", "_ref = torch.softmax(model(_toks.to(device))[\"logits\"][0, 44].cpu(), -1) # unmasked prediction, for comparison\n", "res = predict_masked(ub, 43)\n", "assert len(res) == 5 and all(0 <= p <= 1 for _, p in res)\n", "assert not np.isclose(res[0][1], float(_ref[alphabet.get_idx(res[0][0])]) / float(_ref[[alphabet.get_idx(a) for a in AAS]].sum())), \\\n", " \"your prediction equals the unmasked one -- is the mask in the right place?\"\n", "print(\"OK\")\n", "for name, pos in [(\"Ile44 (hydrophobic patch)\", 43), (\"Lys48 (polyubiquitin linkage)\", 47),\n", " (\"Gly76 (C-terminal tail)\", 75), (\"Asp52 (surface loop)\", 51)]:\n", " print(f\"{name:32s} true={ub[pos]} predicted={predict_masked(ub, pos)}\")" ] }, { "cell_type": "markdown", "id": "89f4bd4d", "metadata": {}, "source": [ "
Hint\n", "\n", "`masked[0, pos + 1] = alphabet.mask_idx` — row 0 because the batch has one sequence.\n", "
\n", "\n", "**Discuss (1 min).** For which residues is the model confident, and for which is it unsure? Why might a buried hydrophobic position or a glycine in a tight turn be easier to predict than a surface residue? *(This same calculation is the basis of zero-shot mutation-effect prediction; see Optional Extension B.)*\n", "\n", "*Your notes:*" ] }, { "cell_type": "markdown", "id": "2767eae9", "metadata": {}, "source": [ "---\n", "## Part 2 · Opening the hood: attention maps (13 min)\n", "\n", "`out[\"attentions\"]` has shape `(layers, heads, L+2, L+2)`: one attention matrix per head, **including** the `` and `` tokens.\n", "\n", "### Exercise 2 — strip the special tokens\n", "Slice the array so that only residue-to-residue attention remains, giving shape `(12, 20, 76, 76)`." ] }, { "cell_type": "code", "execution_count": null, "id": "19900bc9", "metadata": {}, "outputs": [], "source": [ "A_full = out[\"attentions\"] # (12, 20, 78, 78)\n", "# TODO: drop the first () and last () token along BOTH of the last two axes\n", "A = ___\n", "\n", "# --- self-check ---\n", "assert A.shape == (N_LAYERS, N_HEADS, len(ub), len(ub)), A.shape\n", "print(\"OK shape:\", A.shape)\n", "print(\"Row sums of A_full (should be 1):\", A_full[0, 0].sum(-1)[:4].round(3))\n", "print(\"Row sums of A (residue part only):\", A[0, 0].sum(-1)[:4].round(3))" ] }, { "cell_type": "markdown", "id": "7e1d2700", "metadata": {}, "source": [ "
Hint\n", "\n", "Use `1:-1` on axes 2 and 3: `A_full[:, :, 1:-1, 1:-1]`.\n", "
\n", "\n", "**Discuss.** Why do the rows of `A` no longer sum to 1? Where did the missing attention go? *(Look at the next figure.)*\n", "\n", "Let's look at a sample of heads from early, middle and late layers." ] }, { "cell_type": "code", "execution_count": null, "id": "7589e22a", "metadata": {}, "outputs": [], "source": [ "show_layers, show_heads = [0, 3, 6, 9, 11], [0, 5, 8, 13, 14]\n", "fig, axes = plt.subplots(len(show_layers), len(show_heads), figsize=(11, 11))\n", "for r, l in enumerate(show_layers):\n", " for c, h in enumerate(show_heads):\n", " ax = axes[r, c]\n", " ax.imshow(A_full[l, h], cmap=\"viridis\", vmin=0, vmax=np.percentile(A_full[l, h], 99))\n", " ax.set_xticks([]); ax.set_yticks([])\n", " ax.set_title(f\"L{l} H{h}\", fontsize=9)\n", "fig.suptitle(\"Ubiquitin attention maps (including row/col 0 and last)\", y=0.995)\n", "fig.tight_layout()" ] }, { "cell_type": "markdown", "id": "ff8c15d8", "metadata": {}, "source": [ "You should see a \"zoo\" of patterns: heads that look at **neighbouring residues** (bright off-diagonals), heads that dump attention on the **``/`` tokens** (bright first/last column, often called an *attention sink*), and a few late-layer heads with **scattered off-diagonal spots**. Let's quantify the first two types for every head.\n", "\n", "### Exercise 3 — a head \"zoo\" census\n", "For each head compute:\n", "* `prev_mass`: the average attention residue $i$ pays to residue $i-1$, i.e. the mean of entries $A[i, i-1]$;\n", "* `next_mass`: the average attention to residue $i+1$;\n", "* `sink_mass`: the average attention paid to `` or `` (computed for you).\n", "\n", "`np.diagonal(X, offset=k, axis1=-2, axis2=-1)` returns the entries $X[i, i+k]$ along the last two axes." ] }, { "cell_type": "code", "execution_count": null, "id": "20106975", "metadata": {}, "outputs": [], "source": [ "# TODO: choose the right diagonal offsets\n", "prev_mass = np.diagonal(A, offset=___, axis1=-2, axis2=-1).mean(-1) # (12, 20): mean A[i, i-1]\n", "next_mass = np.diagonal(A, offset=___, axis1=-2, axis2=-1).mean(-1) # (12, 20): mean A[i, i+1]\n", "sink_mass = (A_full[:, :, 1:-1, 0] + A_full[:, :, 1:-1, -1]).mean(-1) # (12, 20)\n", "\n", "# --- self-check ---\n", "assert np.isclose(prev_mass[2, 3], np.mean([A[2, 3, i, i - 1] for i in range(1, len(ub))]))\n", "assert np.isclose(next_mass[2, 3], np.mean([A[2, 3, i, i + 1] for i in range(len(ub) - 1)]))\n", "print(\"OK\")\n", "\n", "fig, axes = plt.subplots(1, 3, figsize=(13, 3.6))\n", "for ax, M, title in zip(axes, [prev_mass, next_mass, sink_mass],\n", " [\"attention to residue i-1\", \"attention to residue i+1\", \"attention to /\"]):\n", " im = ax.imshow(M, cmap=\"magma\", aspect=\"auto\", vmin=0, vmax=1)\n", " ax.set_xlabel(\"head\"); ax.set_ylabel(\"layer\"); ax.set_title(title)\n", "fig.colorbar(im, ax=axes, shrink=0.8, label=\"mean attention\")\n", "for name, M in [(\"previous-token\", prev_mass), (\"next-token\", next_mass), (\"sink\", sink_mass)]:\n", " l, h = np.unravel_index(M.argmax(), M.shape)\n", " print(f\"strongest {name:15s} head: L{l} H{h} ({M[l, h]:.2f})\")" ] }, { "cell_type": "markdown", "id": "47756c99", "metadata": {}, "source": [ "
Hint\n", "\n", "$A[i, i-1]$ lies one step *below* the main diagonal, so `offset=-1`; $A[i, i+1]$ is `offset=1`.\n", "
\n", "\n", "**Discuss.** Positional heads are easy to explain: proteins have strong local sequence statistics. Sink heads are a known transformer phenomenon (a head with \"nothing useful to do\" parks its attention on a default token). Neither type tells us about 3D structure — for that we need ground truth." ] }, { "cell_type": "markdown", "id": "5165731e", "metadata": {}, "source": [ "---\n", "## Part 3 · Which heads \"see\" 3D contacts? (17 min)\n", "\n", "We now compare attention maps against real structures from the Protein Data Bank. Two residues are defined to be **in contact** if their C$\\beta$ atoms (C$\\alpha$ for glycine) are closer than **8 Å**.\n", "\n", "The next cell contains a small PDB reader (no need to study it) and downloads 19 small, well-studied proteins. We split them by protein into a **training set** (used to pick heads or fit models) and a **test set** (used only for evaluation). Note that we removed obvious homologues across the split (e.g. two protein G domains) — keep that in mind for Part 5." ] }, { "cell_type": "code", "execution_count": null, "id": "60959eba", "metadata": {}, "outputs": [], "source": [ "#@title Helper: download PDB files and extract sequence, C-beta coordinates and secondary structure\n", "THREE2ONE = dict(ALA=\"A\", ARG=\"R\", ASN=\"N\", ASP=\"D\", CYS=\"C\", GLN=\"Q\", GLU=\"E\", GLY=\"G\", HIS=\"H\", ILE=\"I\",\n", " LEU=\"L\", LYS=\"K\", MET=\"M\", PHE=\"F\", PRO=\"P\", SER=\"S\", THR=\"T\", TRP=\"W\", TYR=\"Y\", VAL=\"V\")\n", "os.makedirs(\"pdb_cache\", exist_ok=True)\n", "\n", "def fetch_pdb(pdb_id):\n", " path = f\"pdb_cache/{pdb_id}.pdb\"\n", " if not os.path.exists(path):\n", " urllib.request.urlretrieve(f\"https://files.rcsb.org/download/{pdb_id}.pdb\", path)\n", " with open(path) as fh:\n", " return fh.read().splitlines()\n", "\n", "def load_chain(pdb_id, chain=\"A\"):\n", " \"\"\"Return sequence, C-beta coords (C-alpha for Gly), and 3-state secondary structure (H/E/C) from HELIX/SHEET records.\"\"\"\n", " residues, order, helices, strands = {}, [], [], []\n", " for ln in fetch_pdb(pdb_id):\n", " rec = ln[:6].strip()\n", " if rec == \"ENDMDL\": # NMR entries: keep the first model only\n", " break\n", " if rec == \"HELIX\" and ln[19] == chain:\n", " helices.append((int(ln[21:25]), int(ln[33:37])))\n", " elif rec == \"SHEET\" and ln[21] == chain:\n", " strands.append((int(ln[22:26]), int(ln[33:37])))\n", " elif rec == \"ATOM\" and ln[21] == chain and ln[16] in \" A\" and ln[17:20] in THREE2ONE:\n", " key = (int(ln[22:26]), ln[26])\n", " if key not in residues:\n", " residues[key] = {\"aa\": THREE2ONE[ln[17:20]]}\n", " order.append(key)\n", " residues[key][ln[12:16].strip()] = np.array([float(ln[30:38]), float(ln[38:46]), float(ln[46:54])])\n", " order = [k for k in order if \"CA\" in residues[k]]\n", " nums = [k[0] for k in order]\n", " ss = \"\".join(\"H\" if any(a <= n <= b for a, b in helices) else\n", " \"E\" if any(a <= n <= b for a, b in strands) else \"C\" for n in nums)\n", " return {\"seq\": \"\".join(residues[k][\"aa\"] for k in order),\n", " \"cb\": np.stack([residues[k].get(\"CB\", residues[k][\"CA\"]) for k in order]),\n", " \"ss\": ss}\n", "\n", "TRAIN = {\"1UBQ\": \"A\", \"1PGA\": \"A\", \"1SHG\": \"A\", \"1ENH\": \"A\", \"2CI2\": \"I\", \"1MBN\": \"A\",\n", " \"1TEN\": \"A\", \"5PTI\": \"A\", \"3CHY\": \"A\", \"2LZM\": \"A\", \"2TRX\": \"A\", \"1FKB\": \"A\"}\n", "TEST = {\"7RSA\": \"A\", \"1BTA\": \"A\", \"1APS\": \"A\", \"1HRC\": \"A\", \"1LMB\": \"3\", \"1CSP\": \"A\", \"1RX2\": \"A\"}\n", "prots = {pid: load_chain(pid, ch) for pid, ch in {**TRAIN, **TEST}.items()}\n", "for pid, d in prots.items():\n", " print(f\"{pid} ({'train' if pid in TRAIN else 'test '}) L={len(d['seq']):4d} \"\n", " f\"helix={d['ss'].count('H')/len(d['ss']):.0%} strand={d['ss'].count('E')/len(d['ss']):.0%}\")\n", "assert prots[\"1UBQ\"][\"seq\"] == ub" ] }, { "cell_type": "markdown", "id": "b5dbeb7c", "metadata": {}, "source": [ "### Exercise 4 — contact map from coordinates\n", "Given an `(L, 3)` array of C$\\beta$ coordinates, return an `(L, L)` integer array that is 1 where the distance is below `threshold` and 0 otherwise. Use NumPy broadcasting: `cb[:, None, :] - cb[None, :, :]` has shape `(L, L, 3)`." ] }, { "cell_type": "code", "execution_count": null, "id": "dc118329", "metadata": {}, "outputs": [], "source": [ "def contact_map(cb, threshold=8.0):\n", " # TODO: pairwise Euclidean distances, then threshold\n", " dist = ___\n", " return (dist < threshold).astype(int)\n", "\n", "# --- self-check on a toy example: points at x = 0, 5, 20 Angstrom ---\n", "toy = np.array([[0., 0, 0], [5, 0, 0], [20, 0, 0]])\n", "assert (contact_map(toy) == np.array([[1, 1, 0], [1, 1, 0], [0, 0, 1]])).all(), contact_map(toy)\n", "print(\"OK\")\n", "for d in prots.values():\n", " d[\"true\"] = contact_map(d[\"cb\"])" ] }, { "cell_type": "markdown", "id": "865243ba", "metadata": {}, "source": [ "
Hint\n", "\n", "`np.linalg.norm(cb[:, None, :] - cb[None, :, :], axis=-1)`\n", "
\n", "\n", "### From attention maps to contact scores\n", "Contacts are symmetric but attention is not, so we **symmetrise** each head ($A + A^\\top$). We then apply the **average product correction (APC)**, borrowed from coevolution analysis: it subtracts the part of each entry explained by row and column averages, removing \"this residue is attended to by everyone\" background:\n", "\n", "$$\\text{APC}(F)_{ij} = F_{ij} - \\frac{F_{i\\cdot}\\,F_{\\cdot j}}{F_{\\cdot\\cdot}}$$\n", "\n", "This is exactly the preprocessing ESM's own contact predictor uses. We run ESM on every protein once and cache the results (~30 s)." ] }, { "cell_type": "code", "execution_count": null, "id": "912aa652", "metadata": {}, "outputs": [], "source": [ "def apc(F):\n", " row, col, tot = F.sum(-1, keepdims=True), F.sum(-2, keepdims=True), F.sum((-1, -2), keepdims=True)\n", " return F - row * col / tot\n", "\n", "def head_features(attn_full):\n", " \"\"\"(layers, heads, L+2, L+2) attention -> (layers, heads, L, L) symmetrised, APC-corrected contact scores.\"\"\"\n", " a = attn_full[..., 1:-1, 1:-1]\n", " return apc(a + np.swapaxes(a, -1, -2))\n", "\n", "t0 = time.time()\n", "for pid, d in prots.items():\n", " o = run_esm(d[\"seq\"])\n", " d[\"hf\"] = head_features(o[\"attentions\"]).astype(np.float32)\n", " d[\"esm_contacts\"] = o[\"contacts\"]\n", " d[\"reps\"] = o[\"reps\"]\n", "print(f\"cached {len(prots)} proteins in {time.time()-t0:.1f} s\")" ] }, { "cell_type": "markdown", "id": "a64a9427", "metadata": {}, "source": [ "### Exercise 5 — precision at L\n", "The standard metric is **precision@L**: rank all residue pairs with sequence separation $|i-j| \\ge 6$ by predicted score, keep the top $L$ pairs, and report the fraction that are true contacts. (Pairs closer than 6 in sequence are trivially in contact and are excluded.)" ] }, { "cell_type": "code", "execution_count": null, "id": "d1785a5a", "metadata": {}, "outputs": [], "source": [ "def precision_at_L(pred, true, min_sep=6):\n", " L = true.shape[0]\n", " i, j = np.triu_indices(L, k=min_sep) # all pairs with j - i >= min_sep\n", " scores, labels = pred[i, j], true[i, j]\n", " # TODO: indices of the L highest-scoring pairs, then the fraction of those that are true contacts\n", " top = ___\n", " return ___\n", "\n", "# --- self-check ---\n", "d = prots[\"1UBQ\"]\n", "assert np.isclose(precision_at_L(d[\"true\"].astype(float), d[\"true\"]), 1.0), \"a perfect predictor should score 1\"\n", "i, j = np.triu_indices(len(ub), k=6)\n", "print(f\"OK random-guess baseline for ubiquitin: {d['true'][i, j].mean():.3f}\")\n", "print(f\"ESM's built-in contact predictor on ubiquitin: P@L = {precision_at_L(d['esm_contacts'], d['true']):.2f}\")" ] }, { "cell_type": "markdown", "id": "637e5fff", "metadata": {}, "source": [ "
Hint\n", "\n", "`top = np.argsort(-scores)[:L]` and `return labels[top].mean()`.\n", "
\n", "\n", "Now score **every one of the 240 heads** on every training protein." ] }, { "cell_type": "code", "execution_count": null, "id": "d25c648b", "metadata": {}, "outputs": [], "source": [ "P_head = np.array([[[precision_at_L(prots[p][\"hf\"][l, h], prots[p][\"true\"]) for h in range(N_HEADS)]\n", " for l in range(N_LAYERS)] for p in TRAIN]) # (n_train, 12, 20)\n", "P_mean = P_head.mean(0)\n", "\n", "fig, ax = plt.subplots(figsize=(8, 4.2))\n", "im = ax.imshow(P_mean, cmap=\"viridis\", aspect=\"auto\")\n", "ax.set_xlabel(\"head\"); ax.set_ylabel(\"layer\"); ax.set_xticks(range(N_HEADS)); ax.set_yticks(range(N_LAYERS))\n", "ax.set_title(\"Contact precision@L of each attention head (mean over training proteins)\")\n", "fig.colorbar(im, ax=ax, label=\"precision@L\")\n", "\n", "ranked = np.dstack(np.unravel_index(np.argsort(-P_mean.ravel()), P_mean.shape))[0]\n", "print(\"Top 6 heads:\", [f\"L{l}H{h}={P_mean[l, h]:.2f}\" for l, h in ranked[:6]])\n", "print(f\"Median head: {np.median(P_mean):.3f}\")\n", "best_l, best_h = ranked[0]" ] }, { "cell_type": "code", "execution_count": null, "id": "0cd4d874", "metadata": {}, "outputs": [], "source": [ "def plot_pair(ax, pred, true, title, L_top=None):\n", " \"\"\"Upper triangle: top-L predicted pairs (|i-j|>=6). Lower triangle: true contact map.\"\"\"\n", " L = true.shape[0]\n", " i, j = np.triu_indices(L, k=6)\n", " top = np.argsort(-pred[i, j])[: (L_top or L)]\n", " ax.imshow(np.tril(true), cmap=\"Greys\", vmin=0, vmax=1.5)\n", " hit = true[i[top], j[top]] == 1\n", " ax.scatter(j[top][hit], i[top][hit], s=6, c=\"tab:blue\", label=\"correct\")\n", " ax.scatter(j[top][~hit], i[top][~hit], s=6, c=\"tab:red\", label=\"wrong\")\n", " ax.set_title(f\"{title}\\nP@L = {precision_at_L(pred, true):.2f}\", fontsize=10)\n", " ax.set_xticks([]); ax.set_yticks([])\n", "\n", "d = prots[\"1UBQ\"]\n", "worst_l, worst_h = ranked[len(ranked) // 2]\n", "fig, axes = plt.subplots(1, 3, figsize=(12, 4.3))\n", "plot_pair(axes[0], d[\"hf\"][worst_l, worst_h], d[\"true\"], f\"median head L{worst_l}H{worst_h}\")\n", "plot_pair(axes[1], d[\"hf\"][best_l, best_h], d[\"true\"], f\"best single head L{best_l}H{best_h}\")\n", "plot_pair(axes[2], d[\"esm_contacts\"], d[\"true\"], \"ESM contact head (all 240 heads)\")\n", "axes[2].legend(loc=\"lower left\", fontsize=8, markerscale=2)\n", "fig.suptitle(\"Ubiquitin: lower triangle = true contacts (grey); upper = top-L predictions\", y=1.02)\n", "fig.tight_layout()" ] }, { "cell_type": "markdown", "id": "b8b24f96", "metadata": {}, "source": [ "**Discuss.** A *single* attention head, never trained on structures, picks out a large fraction of real contacts. Where in the network do the best heads sit? Why might contact information emerge late rather than early?\n", "\n", "### Exercise 6 — how many heads do you actually need?\n", "ESM's built-in contact predictor is a logistic regression over all 240 APC-corrected head maps. Let's fit our own, but with an **L1 penalty**, which drives most weights to exactly zero. Smaller `C` means a stronger penalty and fewer heads. Each training example is one residue pair; its 240 features are that pair's score in every head.\n", "\n", "Fill in the fitting line, then run the sweep over `C`." ] }, { "cell_type": "code", "execution_count": null, "id": "102baa6f", "metadata": {}, "outputs": [], "source": [ "def pair_dataset(pids, min_sep=6):\n", " X, y = [], []\n", " for p in pids:\n", " d = prots[p]; i, j = np.triu_indices(len(d[\"seq\"]), k=min_sep)\n", " X.append(d[\"hf\"][:, :, i, j].reshape(N_LAYERS * N_HEADS, -1).T) # (n_pairs, 240)\n", " y.append(d[\"true\"][i, j])\n", " return np.concatenate(X), np.concatenate(y)\n", "\n", "def make_l1_logreg(C):\n", " \"\"\"L1-penalised logistic regression (handles the scikit-learn >= 1.8 API change).\"\"\"\n", " major, minor = (int(x) for x in sklearn.__version__.split(\".\")[:2])\n", " if (major, minor) >= (1, 8):\n", " return LogisticRegression(l1_ratio=1.0, C=C, solver=\"liblinear\", max_iter=2000)\n", " return LogisticRegression(penalty=\"l1\", C=C, solver=\"liblinear\", max_iter=2000)\n", "\n", "def test_precision(weights_240):\n", " \"\"\"Mean P@L on the TEST proteins when contact score = weighted sum of head maps.\"\"\"\n", " w = weights_240.reshape(N_LAYERS, N_HEADS)\n", " return np.mean([precision_at_L(np.tensordot(w, prots[p][\"hf\"], axes=([0, 1], [0, 1])), prots[p][\"true\"])\n", " for p in TEST])\n", "\n", "X_tr, y_tr = pair_dataset(TRAIN)\n", "print(f\"training pairs: {X_tr.shape[0]:,} features: {X_tr.shape[1]} contact fraction: {y_tr.mean():.3f}\")\n", "\n", "sweep = []\n", "for C in [0.005, 0.01, 0.02, 0.05, 0.15, 0.5]:\n", " clf = make_l1_logreg(C)\n", " # TODO: fit the classifier on the training pairs\n", " ___\n", " w = clf.coef_[0]\n", " heads_used = [f\"L{k // N_HEADS}H{k % N_HEADS}\" for k in np.flatnonzero(w)]\n", " sweep.append((C, len(heads_used), test_precision(w)))\n", " print(f\"C={C:<6} heads used={len(heads_used):3d} test P@L={sweep[-1][2]:.3f} \"\n", " f\"{heads_used if len(heads_used) <= 8 else ''}\")\n", "\n", "esm_test = np.mean([precision_at_L(prots[p][\"esm_contacts\"], prots[p][\"true\"]) for p in TEST])\n", "print(f\"\\nESM's own contact head (240 heads, trained on more data): test P@L = {esm_test:.3f}\")" ] }, { "cell_type": "markdown", "id": "2627dffd", "metadata": {}, "source": [ "
Hint\n", "\n", "`clf.fit(X_tr, y_tr)`\n", "
" ] }, { "cell_type": "code", "execution_count": null, "id": "8e5effbb", "metadata": {}, "outputs": [], "source": [ "W_esm = model.contact_head.regression.weight.detach().cpu().numpy().reshape(N_LAYERS, N_HEADS)\n", "fig, axes = plt.subplots(1, 2, figsize=(12, 3.8))\n", "pts = sorted({(s[1], round(s[2], 4)) for s in sweep if s[1] > 0})\n", "axes[0].plot([p[0] for p in pts], [p[1] for p in pts], \"o-\", color=\"tab:blue\", label=\"L1 probe (ours)\")\n", "axes[0].axhline(esm_test, ls=\"--\", color=\"grey\", label=\"ESM contact head (240 heads)\")\n", "axes[0].set_xscale(\"log\"); axes[0].set_ylim(0, 0.6); axes[0].set_xlabel(\"number of heads with non-zero weight\")\n", "axes[0].set_ylabel(\"test precision@L\"); axes[0].legend(); axes[0].set_title(\"Sparsity vs accuracy\")\n", "vmax = np.abs(W_esm).max()\n", "im = axes[1].imshow(W_esm, cmap=\"RdBu_r\", vmin=-vmax, vmax=vmax, aspect=\"auto\")\n", "axes[1].set_xticks(range(N_HEADS)); axes[1].set_yticks(range(N_LAYERS))\n", "axes[1].set_xlabel(\"head\"); axes[1].set_ylabel(\"layer\"); axes[1].set_title(\"Weights of ESM's contact head\")\n", "fig.colorbar(im, ax=axes[1])\n", "fig.tight_layout()" ] }, { "cell_type": "markdown", "id": "6993ccaa", "metadata": {}, "source": [ "**Discuss.**\n", "1. How few heads do you need to get most of the performance? What does this say about how \"localised\" the contact information is?\n", "2. Do the heads chosen by the L1 probe match the high-weight heads in ESM's contact head? Why might two equally good models choose *different* heads? *(Hint: correlated features.)*\n", "3. Everything so far is **correlational**: these heads' patterns *line up* with contacts. Does that prove the model *uses* contact information to predict amino acids?\n", "\n", "*Your notes:*" ] }, { "cell_type": "markdown", "id": "6d7361e4", "metadata": {}, "source": [ "---\n", "## Part 4 · Causal interventions: ablating heads and MLPs (18 min)\n", "\n", "To make causal claims we must **intervene**: switch a component off and measure what changes. We will use two very different notions of \"this head matters\".\n", "\n", "**(A) Readout importance.** If we delete a head's map from the *input of the contact predictor*, how much worse do the predicted contacts get? This asks: *does our decoder rely on this head?*\n", "\n", "**(B) Causal importance inside the network.** If we silence a head's *output* during the forward pass, every downstream layer sees a different residue stream and computes different attention. This asks: *does the network rely on this component to build its internal contact map?*\n", "\n", "### Exercise 7 — readout ablation\n", "First, check that we can reproduce ESM's contact head by hand: a weighted sum of our APC'd head maps plus a bias, passed through a sigmoid. Then zero out the weights of the `k` heads with the largest |weight| and compare against zeroing `k` **random** heads (a control!)." ] }, { "cell_type": "code", "execution_count": null, "id": "abdc7cf4", "metadata": {}, "outputs": [], "source": [ "b_esm = model.contact_head.regression.bias.item()\n", "\n", "def contacts_from_weights(hf, W, b=b_esm):\n", " logits = np.tensordot(W, hf, axes=([0, 1], [0, 1])) + b\n", " return 1 / (1 + np.exp(-np.clip(logits, -50, 50)))\n", "\n", "d = prots[\"1UBQ\"]\n", "print(\"max |hand-built - ESM| =\", np.abs(contacts_from_weights(d[\"hf\"], W_esm) - d[\"esm_contacts\"]).max())\n", "\n", "def readout_precision(W, pids=list(prots)):\n", " return np.mean([precision_at_L(contacts_from_weights(prots[p][\"hf\"], W), prots[p][\"true\"]) for p in pids])\n", "\n", "rng = np.random.default_rng(0)\n", "by_weight = np.argsort(-np.abs(W_esm).ravel()) # flat head indices, most important first\n", "ks, top_curve, rand_curve = [0, 1, 2, 3, 5, 8, 12, 20, 40], [], []\n", "for k in ks:\n", " W_top = W_esm.copy().ravel()\n", " # TODO: set the weights of the k most important heads to zero\n", " W_top[___] = 0\n", " top_curve.append(readout_precision(W_top.reshape(N_LAYERS, N_HEADS)))\n", " rand_vals = []\n", " for rep in range(3): # average over 3 random draws\n", " W_rand = W_esm.copy().ravel()\n", " W_rand[rng.choice(N_LAYERS * N_HEADS, size=k, replace=False)] = 0\n", " rand_vals.append(readout_precision(W_rand.reshape(N_LAYERS, N_HEADS)))\n", " rand_curve.append(np.mean(rand_vals))\n", "\n", "# --- self-check ---\n", "assert np.isclose(top_curve[0], rand_curve[0]) and top_curve[-1] < top_curve[0]\n", "print(\"OK\")\n", "fig, ax = plt.subplots(figsize=(6, 3.8))\n", "ax.plot(ks, top_curve, \"o-\", label=\"remove highest-|weight| heads\")\n", "ax.plot(ks, rand_curve, \"s--\", label=\"remove random heads (control)\")\n", "ax.set_xlabel(\"number of heads removed from the readout\"); ax.set_ylabel(\"precision@L (all 19 proteins)\")\n", "ax.legend(); ax.set_title(\"Readout ablation\")\n", "fig.tight_layout()" ] }, { "cell_type": "markdown", "id": "d72ef15e", "metadata": {}, "source": [ "
Hint\n", "\n", "`by_weight[:k]` holds the flat indices of the top-k heads, so `W_top[by_weight[:k]] = 0`.\n", "
\n", "\n", "**Discuss.** Removing the single most important head barely changes anything, yet removing ~5 collapses performance. What does that tell you about **redundancy** between heads? Why is the random-head curve an essential control?" ] }, { "cell_type": "markdown", "id": "bd92eed5", "metadata": {}, "source": [ "### Hooks: editing the forward pass\n", "PyTorch lets us attach a function to any module that runs during the forward pass and can modify its input or output. Inside each attention block, the concatenated head outputs enter `out_proj`, so a **forward pre-hook** on `out_proj` sees a tensor of shape `(T, B, D_MODEL)` — note that `fair-esm` puts the *token* axis first. The 480 features are 20 consecutive blocks of 24, one per head:\n", "\n", "```\n", "[ head 0 (24 dims) | head 1 (24 dims) | ... | head 19 (24 dims) ]\n", "```\n", "\n", "### Exercise 8 — write a head-ablation hook\n", "Complete the hook so that the selected heads' outputs are replaced by zeros (`mode=\"zero\"`) or by their average over residue positions (`mode=\"mean\"`). Reshape to `(T, B, N_HEADS, HEAD_DIM)`, edit, and reshape back." ] }, { "cell_type": "code", "execution_count": null, "id": "8bae98d1", "metadata": {}, "outputs": [], "source": [ "def make_head_ablation_hook(heads, mode=\"zero\"):\n", " def hook(module, args):\n", " x = args[0] # (T, B, D_MODEL) -- concatenated head outputs\n", " T, B, E = x.shape\n", " # TODO: reshape so that heads get their own axis\n", " x = x.reshape(___).clone()\n", " if mode == \"zero\":\n", " # TODO: zero the chosen heads (all tokens, all batch items, all head dims)\n", " x[___] = 0.0\n", " elif mode == \"mean\": # replace with the average over residue tokens\n", " x[1:-1, :, heads, :] = x[1:-1, :, heads, :].mean(0, keepdim=True)\n", " return (x.reshape(T, B, E),)\n", " return hook\n", "\n", "def make_mlp_ablation_hook():\n", " def hook(module, args, output): # forward hook on fc2: zero the whole MLP update\n", " return torch.zeros_like(output)\n", " return hook\n", "\n", "@contextlib.contextmanager\n", "def ablate(heads_by_layer=None, mlp_layers=(), mode=\"zero\"):\n", " \"\"\"Usage: with ablate({11: [13]}): out = run_esm(seq)\"\"\"\n", " handles = []\n", " for l, heads in (heads_by_layer or {}).items():\n", " handles.append(model.layers[l].self_attn.out_proj.register_forward_pre_hook(make_head_ablation_hook(heads, mode)))\n", " for l in mlp_layers:\n", " handles.append(model.layers[l].fc2.register_forward_hook(make_mlp_ablation_hook()))\n", " try:\n", " yield\n", " finally:\n", " for h in handles:\n", " h.remove() # always clean up, even if something fails\n", "\n", "# --- self-check 1: hook with no heads changes nothing ---\n", "base = run_esm(ub)\n", "with ablate({5: []}):\n", " same = run_esm(ub)\n", "assert np.allclose(base[\"contacts\"], same[\"contacts\"], atol=1e-5)\n", "# --- self-check 2: ablating a head changes the output ---\n", "with ablate({5: [0, 1, 2]}):\n", " diff = run_esm(ub)\n", "assert not torch.allclose(base[\"logits\"], diff[\"logits\"]), \"ablation had no effect -- check your indexing\"\n", "print(\"OK\")" ] }, { "cell_type": "markdown", "id": "4614ed91", "metadata": {}, "source": [ "
Hint\n", "\n", "`x.reshape(T, B, N_HEADS, HEAD_DIM)` and `x[:, :, heads, :] = 0.0`.\n", "
\n", "\n", "### A puzzle\n", "Ablate the **output** of the best contact head (found in Part 3) and compare ESM's predicted contacts and its amino-acid logits before and after." ] }, { "cell_type": "code", "execution_count": null, "id": "685974b8", "metadata": {}, "outputs": [], "source": [ "with ablate({int(best_l): [int(best_h)]}):\n", " abl = run_esm(ub)\n", "print(f\"ablated L{best_l}H{best_h}\")\n", "print(\"max change in predicted contacts:\", float(np.abs(abl[\"contacts\"] - base[\"contacts\"]).max()))\n", "print(\"max change in logits: \", float((abl[\"logits\"] - base[\"logits\"]).abs().max()))" ] }, { "cell_type": "markdown", "id": "7623d3eb", "metadata": {}, "source": [ "**Discuss (important!).** The contact map does not change at all, although the logits do. Explain why. *(Two reasons: what does the contact head read — the head's attention pattern or its output? And which layers could be affected by a change in layer 11's output?)*\n", "\n", "*Your notes:*\n", "\n", "### Component scan: attention blocks vs MLPs\n", "To find where the contact signal is *built*, we silence a whole attention block (all 20 heads) or a whole MLP in one layer at a time and measure two things on six proteins:\n", "* contact **precision@L** of ESM's contact head (structure signal), and\n", "* **masked-LM loss** (cross-entropy on 15% randomly masked residues; lower is better) — the task the model was actually trained on.\n", "\n", "This takes about a minute on a laptop CPU." ] }, { "cell_type": "code", "execution_count": null, "id": "24de670a", "metadata": {}, "outputs": [], "source": [ "def masked_lm_loss(seq, n_masks=4, frac=0.15, seed=0):\n", " \"\"\"Average cross-entropy over n_masks random masking patterns (batched into one forward pass).\"\"\"\n", " rng = np.random.default_rng(seed)\n", " _, _, toks = batch_converter([(\"protein\", seq)])\n", " batch, targets = toks.repeat(n_masks, 1), []\n", " for k in range(n_masks):\n", " pos = rng.choice(len(seq), size=max(1, int(frac * len(seq))), replace=False) + 1\n", " batch[k, pos] = alphabet.mask_idx\n", " targets.append(pos)\n", " logp = torch.log_softmax(model(batch.to(device))[\"logits\"].cpu(), dim=-1)\n", " return -np.mean([logp[k, targets[k], toks[0, targets[k]]].mean().item() for k in range(n_masks)])\n", "\n", "EVAL = [\"1UBQ\", \"1TEN\", \"3CHY\", \"2TRX\", \"1CSP\", \"1APS\"]\n", "def evaluate(pids=EVAL):\n", " P = np.mean([precision_at_L(run_esm(prots[p][\"seq\"])[\"contacts\"], prots[p][\"true\"]) for p in pids])\n", " loss = np.mean([masked_lm_loss(prots[p][\"seq\"]) for p in pids])\n", " return P, loss\n", "\n", "t0 = time.time()\n", "base_P, base_loss = evaluate()\n", "scan = {\"attn\": [], \"mlp\": []}\n", "for l in range(N_LAYERS):\n", " with ablate(heads_by_layer={l: list(range(N_HEADS))}):\n", " scan[\"attn\"].append(evaluate())\n", " with ablate(mlp_layers=[l]):\n", " scan[\"mlp\"].append(evaluate())\n", "print(f\"scan finished in {time.time()-t0:.0f} s | baseline P@L={base_P:.3f}, masked-LM loss={base_loss:.3f}\")\n", "\n", "fig, axes = plt.subplots(1, 2, figsize=(12, 3.8))\n", "x = np.arange(N_LAYERS)\n", "for ax, idx, base_val, ylabel in [(axes[0], 0, base_P, \"contact precision@L\"), (axes[1], 1, base_loss, \"masked-LM loss\")]:\n", " ax.bar(x - 0.2, [s[idx] for s in scan[\"attn\"]], 0.4, label=\"ablate attention block\", color=\"tab:blue\")\n", " ax.bar(x + 0.2, [s[idx] for s in scan[\"mlp\"]], 0.4, label=\"ablate MLP\", color=\"tab:orange\")\n", " ax.axhline(base_val, color=\"k\", ls=\"--\", lw=1, label=\"no ablation\")\n", " ax.set_xticks(x); ax.set_xlabel(\"layer\"); ax.set_ylabel(ylabel)\n", "axes[0].set_title(\"Structure signal (higher = better)\"); axes[1].set_title(\"Training objective (lower = better)\")\n", "handles, labels = axes[0].get_legend_handles_labels()\n", "fig.legend(handles, labels, loc=\"upper center\", ncol=3, fontsize=9, bbox_to_anchor=(0.5, 1.07))\n", "fig.tight_layout()" ] }, { "cell_type": "markdown", "id": "0752cef5", "metadata": {}, "source": [ "**Discuss.**\n", "1. Which single component is most essential overall? Which layer's attention block matters most for contacts?\n", "2. Does a component that matters for contacts always matter for the masked-LM loss, and vice versa? What would it mean if they diverged?\n", "3. Why does ablating the layer-11 attention block leave contacts unchanged but hurt the loss?\n", "\n", "### Exercise 9 — is any *single* head in layer 1 responsible?\n", "Silencing all of layer 1's attention wrecks contact prediction. Loop over its 20 heads and ablate them **one at a time** (contacts only, ~15 s)." ] }, { "cell_type": "code", "execution_count": null, "id": "8b7c5d63", "metadata": {}, "outputs": [], "source": [ "def contact_precision(pids=EVAL):\n", " return np.mean([precision_at_L(run_esm(prots[p][\"seq\"])[\"contacts\"], prots[p][\"true\"]) for p in pids])\n", "\n", "single = []\n", "for h in range(N_HEADS):\n", " # TODO: ablate only head h of layer 1, then measure contact precision\n", " with ablate(___):\n", " single.append(contact_precision())\n", "\n", "with ablate({1: list(range(N_HEADS))}, mode=\"mean\"):\n", " layer1_mean = contact_precision()\n", "\n", "fig, ax = plt.subplots(figsize=(8, 3.4))\n", "ax.bar(range(N_HEADS), single, color=\"tab:blue\")\n", "ax.axhline(base_P, color=\"k\", ls=\"--\", lw=1, label=\"no ablation\")\n", "ax.axhline(scan[\"attn\"][1][0], color=\"tab:red\", ls=\"-\", lw=3, alpha=0.6, label=f\"all 20 heads, zero ({scan['attn'][1][0]:.2f})\")\n", "ax.axhline(layer1_mean, color=\"tab:purple\", ls=\"--\", lw=1.5, label=f\"all 20 heads, mean ({layer1_mean:.2f})\")\n", "ax.set_xticks(range(N_HEADS)); ax.set_xlabel(\"head ablated in layer 1\"); ax.set_ylabel(\"contact precision@L\"); ax.set_ylim(0, 0.8)\n", "ax.legend(fontsize=8, loc=\"lower right\"); ax.set_title(\"Single-head vs whole-layer ablation\")\n", "fig.tight_layout()\n", "print(\"largest single-head drop:\", round(base_P - min(single), 3))" ] }, { "cell_type": "markdown", "id": "bf252f96", "metadata": {}, "source": [ "
Hint\n", "\n", "`with ablate({1: [h]}):`\n", "
\n", "\n", "**Discuss.**\n", "1. No single head matters, but all of them together are essential. Name two mechanisms that could produce this (think *redundancy* and *distributed computation*). What does this imply for \"one head at a time\" importance scores, which are common in interpretability papers?\n", "2. **Zero vs mean ablation.** Setting an activation to zero can push the network into states it never saw during training, so a large effect might reflect \"broken input\" rather than \"lost information\". Mean ablation is a gentler control. Do the two agree here?\n", "\n", "*Your notes:*" ] }, { "cell_type": "markdown", "id": "692d7ae0", "metadata": {}, "source": [ "---\n", "## Part 5 · Probing for secondary structure (12 min)\n", "\n", "A **linear probe** is a simple classifier trained on frozen internal representations. If a linear model can read a property out of layer $\\ell$, we say the property is *linearly decodable* there. We will probe each layer's residue representations for 3-state secondary structure (**H** helix, **E** strand, **C** other) taken from the PDB annotations.\n", "\n", "Good probing practice requires controls:\n", "* **majority-class baseline** — accuracy from always guessing the most common class;\n", "* **input baseline** — a probe on the amino-acid identities alone (a window of ±3 neighbours), i.e. what you can get *without* the transformer;\n", "* **held-out proteins** — train and test on *different* proteins;\n", "* **shuffled labels** — the probe should fail if labels are meaningless.\n", "\n", "### Exercise 10 — layer-wise probes\n", "Complete the loop: build the train/test matrices for layer `l`, fit the probe, and record test accuracy." ] }, { "cell_type": "code", "execution_count": null, "id": "3e55b88c", "metadata": {}, "outputs": [], "source": [ "SS2I = {\"H\": 0, \"E\": 1, \"C\": 2}\n", "\n", "def ss_dataset(pids, layer):\n", " X = np.concatenate([prots[p][\"reps\"][layer] for p in pids])\n", " y = np.concatenate([[SS2I[c] for c in prots[p][\"ss\"]] for p in pids])\n", " return X, y\n", "\n", "def make_probe():\n", " return make_pipeline(StandardScaler(), LogisticRegression(C=0.003, max_iter=3000))\n", "\n", "probe_acc = []\n", "for l in range(N_LAYERS + 1): # layer 0 = token embeddings\n", " # TODO: training data from TRAIN proteins, test data from TEST proteins\n", " X_tr_l, y_tr_l = ___\n", " X_te_l, y_te_l = ___\n", " probe = make_probe()\n", " # TODO: fit on train, score (accuracy) on test\n", " ___\n", " probe_acc.append(___)\n", "\n", "# --- self-check ---\n", "assert len(probe_acc) == N_LAYERS + 1 and all(0 <= a <= 1 for a in probe_acc)\n", "print(\"OK test accuracy by layer:\", np.round(probe_acc, 2))" ] }, { "cell_type": "markdown", "id": "caa87372", "metadata": {}, "source": [ "
Hint\n", "\n", "`ss_dataset(TRAIN, l)`, `ss_dataset(TEST, l)`, `probe.fit(X_tr_l, y_tr_l)`, `probe.score(X_te_l, y_te_l)`.\n", "
\n", "\n", "Now the controls (given)." ] }, { "cell_type": "code", "execution_count": null, "id": "4e109678", "metadata": {}, "outputs": [], "source": [ "def onehot_window(pids, w=3):\n", " X, y = [], []\n", " for p in pids:\n", " s = prots[p][\"seq\"]; oh = np.zeros((len(s) + 2 * w, 20))\n", " for i, c in enumerate(s):\n", " oh[i + w, AAS.index(c)] = 1\n", " X.append(np.concatenate([oh[k:k + len(s)] for k in range(2 * w + 1)], axis=1))\n", " y.append([SS2I[c] for c in prots[p][\"ss\"]])\n", " return np.concatenate(X), np.concatenate(y)\n", "\n", "best_layer = int(np.argmax(probe_acc))\n", "_, y_te = ss_dataset(TEST, 0)\n", "majority = np.bincount(y_te).max() / len(y_te)\n", "Xo_tr, yo_tr = onehot_window(TRAIN); Xo_te, yo_te = onehot_window(TEST)\n", "onehot_acc = LogisticRegression(max_iter=3000).fit(Xo_tr, yo_tr).score(Xo_te, yo_te)\n", "\n", "X_tr_b, y_tr_b = ss_dataset(TRAIN, best_layer); X_te_b, y_te_b = ss_dataset(TEST, best_layer)\n", "shuffled_acc = make_probe().fit(X_tr_b, np.random.default_rng(0).permutation(y_tr_b)).score(X_te_b, y_te_b)\n", "\n", "# Leakage demo: split RESIDUES at random instead of splitting PROTEINS\n", "from sklearn.model_selection import train_test_split\n", "X_all, y_all = ss_dataset(list(prots), best_layer)\n", "Xa, Xb, ya, yb = train_test_split(X_all, y_all, test_size=0.35, random_state=0)\n", "leaky_acc = make_probe().fit(Xa, ya).score(Xb, yb)\n", "\n", "fig, ax = plt.subplots(figsize=(7.5, 4))\n", "ax.plot(range(N_LAYERS + 1), probe_acc, \"o-\", color=\"tab:blue\", label=\"ESM layer probe (held-out proteins)\")\n", "ax.axhline(majority, color=\"grey\", ls=\":\", label=f\"majority class ({majority:.2f})\")\n", "ax.axhline(onehot_acc, color=\"tab:green\", ls=\"--\", label=f\"amino-acid window ±3 ({onehot_acc:.2f})\")\n", "ax.axhline(shuffled_acc, color=\"tab:red\", ls=\"-.\", label=f\"shuffled labels, layer {best_layer} ({shuffled_acc:.2f})\")\n", "ax.scatter([best_layer], [leaky_acc], marker=\"*\", s=180, color=\"tab:orange\", zorder=5,\n", " label=f\"random residue split, layer {best_layer} ({leaky_acc:.2f})\")\n", "ax.set_xlabel(\"layer (0 = embeddings)\"); ax.set_ylabel(\"3-state accuracy\"); ax.set_ylim(0.2, 0.9)\n", "ax.legend(fontsize=8, loc=\"lower right\"); ax.set_title(\"Where is secondary structure linearly decodable?\")\n", "fig.tight_layout()" ] }, { "cell_type": "markdown", "id": "7fecd4d2", "metadata": {}, "source": [ "**Discuss.**\n", "1. At which depth does secondary structure become decodable? Compare with where the contact heads lived (Part 3).\n", "2. Why is the random-residue split so much more optimistic? Which number would you report in a paper?\n", "3. A probe succeeding shows that the information is *present* in a representation. Does it show the model *uses* it? Design an experiment, using tools from Part 4, that would test whether ESM uses its secondary-structure representation when predicting masked residues.\n", "\n", "*Your notes:*" ] }, { "cell_type": "markdown", "id": "4dc62b39", "metadata": {}, "source": [ "---\n", "## Wrap-up (3 min)\n", "\n", "**What we found**\n", "* Attention heads in a sequence-only model form a zoo: positional heads, sink heads, and a few late-layer heads whose patterns match 3D contacts.\n", "* The **readable** contact signal is concentrated: one or a few heads give most of the precision, matching a 240-head regression.\n", "* **Causal** experiments tell a complementary story: the late heads *display* the contact signal, but silencing them inside the network changes nothing downstream, whereas silencing earlier components (MLP 0, the layer-1 attention block, mid-layer MLPs) are what build the signal, degrades it, in a distributed, redundant way that single-head ablation misses.\n", "* Secondary structure is linearly decodable in later layers, well above input-only baselines — but only honest splits and controls make that claim trustworthy.\n", "\n", "**Interpretability lessons that transfer beyond proteins**\n", "1. Attention patterns, probes and weights are *correlational*; use interventions for causal claims.\n", "2. Always compare against controls (random components, shuffled labels, input-only baselines).\n", "3. Single-component ablations can underestimate importance when computation is redundant.\n", "4. How you split data (by protein, by family) can matter more than which model you use.\n", "\n", "**Further reading**\n", "* Rao, Meier, Sercu, Ovchinnikov & Rives (2021). *Transformer protein language models are unsupervised structure learners.* ICLR.\n", "* Vig et al. (2021). *BERTology meets biology: interpreting attention in protein language models.* ICLR.\n", "* Lin et al. (2023). *Evolutionary-scale prediction of atomic-level protein structure with a language model.* Science (ESM-2 / ESMFold).\n", "* Zhang et al. (2024). *Protein language models learn evolutionary statistics of interacting sequence motifs.* PNAS.\n", "* Hewitt & Liang (2019). *Designing and interpreting probes with control tasks.* EMNLP.\n", "* Simon & Zou (2024). *InterPLM: discovering interpretable features in protein language models via sparse autoencoders.* bioRxiv.\n", "* Meier et al. (2021). *Language models enable zero-shot prediction of the effects of mutations on protein function.* NeurIPS." ] }, { "cell_type": "markdown", "id": "cf9a26de", "metadata": {}, "source": [ "---\n", "## Optional extensions (homework or for fast groups)\n", "\n", "### A. Are there \"helix neurons\"?\n", "Each MLP has 1,920 hidden neurons (the output of `fc1 → GELU`). Find the neuron in a layer whose activation best separates helix from non-helix residues on the **training** proteins (by ROC-AUC), then check how well that same neuron does on the **test** proteins." ] }, { "cell_type": "code", "execution_count": null, "id": "d6dfb697", "metadata": {}, "outputs": [], "source": [ "def mlp_neuron_acts(seq, layer):\n", " \"\"\"Hidden activations (L, 1920) of the MLP in `layer`, captured with a pre-hook on fc2.\"\"\"\n", " store = {}\n", " handle = model.layers[layer].fc2.register_forward_pre_hook(lambda m, args: store.update(a=args[0][1:-1, 0].cpu().numpy()))\n", " try:\n", " run_esm(seq)\n", " finally:\n", " handle.remove()\n", " return store[\"a\"]\n", "\n", "def auc_per_column(X, y):\n", " \"\"\"Vectorised ROC-AUC of every column of X for binary labels y (Mann-Whitney U).\"\"\"\n", " R = rankdata(X, axis=0); n1 = y.sum(); n0 = len(y) - n1\n", " return (R[y == 1].sum(0) - n1 * (n1 + 1) / 2) / (n1 * n0)\n", "\n", "layer = 10\n", "acts_tr = np.concatenate([mlp_neuron_acts(prots[p][\"seq\"], layer) for p in TRAIN])\n", "acts_te = np.concatenate([mlp_neuron_acts(prots[p][\"seq\"], layer) for p in TEST])\n", "helix_tr = np.concatenate([[c == \"H\" for c in prots[p][\"ss\"]] for p in TRAIN]).astype(int)\n", "helix_te = np.concatenate([[c == \"H\" for c in prots[p][\"ss\"]] for p in TEST]).astype(int)\n", "\n", "auc_tr = auc_per_column(acts_tr, helix_tr)\n", "n_best = int(np.argmax(np.abs(auc_tr - 0.5)))\n", "auc_te = auc_per_column(acts_te[:, [n_best]], helix_te)[0]\n", "print(f\"layer {layer}: best helix neuron #{n_best} train AUC={auc_tr[n_best]:.2f} test AUC={auc_te:.2f}\")\n", "\n", "d = prots[\"1MBN\"]; a = mlp_neuron_acts(d[\"seq\"], layer)[:, n_best]\n", "fig, ax = plt.subplots(figsize=(11, 2.8))\n", "ax.plot(a, color=\"k\", lw=1)\n", "for i, c in enumerate(d[\"ss\"]):\n", " if c == \"H\": ax.axvspan(i - 0.5, i + 0.5, color=\"tab:red\", alpha=0.15, lw=0)\n", " if c == \"E\": ax.axvspan(i - 0.5, i + 0.5, color=\"tab:blue\", alpha=0.15, lw=0)\n", "ax.set_xlabel(\"residue (myoglobin)\"); ax.set_ylabel(f\"neuron {n_best}\"); ax.set_title(\"Neuron activation (red shading = helix)\")\n", "fig.tight_layout()" ] }, { "cell_type": "markdown", "id": "936b34a8", "metadata": {}, "source": [ "**Questions.** The neuron was *selected* on the training proteins out of 1,920 candidates. How much does its AUC drop on the test proteins, and why should you expect some drop (think multiple comparisons)? Repeat for layers 5 and 9. If individual neurons are weak but the whole layer is strongly decodable, what does that suggest about how features are represented (*superposition*, *polysemantic neurons*)? This motivates **sparse autoencoders**, which try to recover cleaner features from a layer (see InterPLM in the reading list).\n", "\n", "### B. Zero-shot mutation effects, and which components they depend on\n", "The **masked-marginal score** of a mutation $\\text{wt}\\to\\text{mt}$ at position $i$ is $\\log p(\\text{mt}) - \\log p(\\text{wt})$ with residue $i$ masked. Negative scores mean the model finds the mutant less plausible; these scores correlate with experimental fitness and stability measurements. Below we compute the full single-mutant landscape for a stretch of ubiquitin, then ask whether it survives ablating an MLP." ] }, { "cell_type": "code", "execution_count": null, "id": "64290c6a", "metadata": {}, "outputs": [], "source": [ "def masked_marginal_landscape(seq, positions):\n", " \"\"\"(len(positions), 20) matrix of log p(mt) - log p(wt) with each position masked in turn (batched).\"\"\"\n", " _, _, toks = batch_converter([(\"protein\", seq)])\n", " batch = toks.repeat(len(positions), 1)\n", " for k, pos in enumerate(positions):\n", " batch[k, pos + 1] = alphabet.mask_idx\n", " logp = torch.log_softmax(model(batch.to(device))[\"logits\"].cpu(), dim=-1)\n", " aa_idx = [alphabet.get_idx(a) for a in AAS]\n", " S = np.zeros((len(positions), 20))\n", " for k, pos in enumerate(positions):\n", " lp = logp[k, pos + 1]\n", " S[k] = (lp[aa_idx] - lp[alphabet.get_idx(seq[pos])]).numpy()\n", " return S\n", "\n", "positions = list(range(40, 76)) # residues 41-76 (includes the Ile44 patch and the C-terminal tail)\n", "S_base = masked_marginal_landscape(ub, positions)\n", "with ablate(mlp_layers=[4]):\n", " S_abl = masked_marginal_landscape(ub, positions)\n", "\n", "fig, axes = plt.subplots(1, 2, figsize=(13, 4.2), sharey=True)\n", "for ax, S, title in [(axes[0], S_base, \"intact model\"), (axes[1], S_abl, \"MLP 4 ablated\")]:\n", " im = ax.imshow(S.T, cmap=\"RdBu\", vmin=-8, vmax=8, aspect=\"auto\")\n", " ax.set_xticks(range(0, len(positions), 5)); ax.set_xticklabels([f\"{ub[p]}{p+1}\" for p in positions[::5]])\n", " ax.set_yticks(range(20)); ax.set_yticklabels(list(AAS)); ax.set_title(f\"Masked-marginal scores, {title}\")\n", "fig.colorbar(im, ax=axes, label=\"log p(mt) - log p(wt)\")\n", "r = np.corrcoef(S_base.ravel(), S_abl.ravel())[0, 1]\n", "print(f\"correlation between intact and ablated landscapes: r = {r:.2f}\")" ] }, { "cell_type": "markdown", "id": "79b4c0fb", "metadata": {}, "source": [ "**Questions.** Which positions are most mutation-intolerant? Do they make sense structurally (core residues, the Ile44 hydrophobic patch, the Gly75–Gly76 tail)? Which layer's MLP, when ablated, changes the landscape the most? In protein engineering, one might look for mutations elsewhere in the protein that *raise* the score of a destabilising variant (a candidate compensatory mutation). How would you test such a prediction experimentally, and what are the limits of trusting a language-model score for stability?\n", "\n", "### C. Does scale change the picture?\n", "Swap in a larger model and rerun Part 3. Everything is written in terms of `N_LAYERS` and `N_HEADS`, so only the loader changes:\n", "\n", "```python\n", "model, alphabet = esm.pretrained.esm2_t30_150M_UR50D() # 30 layers x 20 heads, ~600 MB\n", "model = model.eval().to(device); batch_converter = alphabet.get_batch_converter()\n", "N_LAYERS, N_HEADS, D_MODEL = model.num_layers, model.attention_heads, model.embed_dim; HEAD_DIM = D_MODEL // N_HEADS\n", "```\n", "Do the best contact heads still sit near the top of the network? Does the best single head get better? (The 8M-parameter `esm2_t6_8M_UR50D` is a useful comparison in the other direction.)" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 5 }