{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "dt50gf9tfxg",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/andyrdt/puzzles/blob/main/04_2026/starter_notebook.ipynb)\n",
    "\n",
    "# Monthly Algorithmic Challenge — April 2026: Max of List\n",
    "\n",
    "*Inspired by Callum McDougall's [ARENA Monthly Algorithmic Challenges](https://learn.arena.education/chapter1_transformer_interp/monthly_algorithmic/).*\n",
    "\n",
    "## Overview\n",
    "\n",
    "We've trained two small attention-only transformers to solve the same task: **given a list of 5 numbers, predict the maximum.**\n",
    "\n",
    "Both models achieve 100% test accuracy. Your challenge: **reverse-engineer the algorithm each model has learned.**\n",
    "\n",
    "### The two models\n",
    "\n",
    "| | Model 1 (easier) | Model 2 (harder) |\n",
    "|---|---|---|\n",
    "| **Numbers** | 0–9 | 0–99 |\n",
    "| **Tokenization** | One token per number | Two digit tokens per number (e.g. 42 → `4`, `2`) |\n",
    "| **Layers** | 1 | 2 |\n",
    "| **Heads** | 4 | 4 |\n",
    "| **`d_model`** | 64 | 64 |\n",
    "| **Parameters** | 18,944 | 35,712 |\n",
    "| **Vocab** | 14 tokens (0-9 + BOS, SEP, ANS, EOS) | 14 tokens (digits 0-9 + BOS, SEP, ANS, EOS) |\n",
    "\n",
    "Both models are trained and evaluated on lists of length 5.\n",
    "\n",
    "### Architecture\n",
    "\n",
    "Both models are **attention-only** causal transformers — no MLPs, no LayerNorm. The architecture is:\n",
    "\n",
    "```\n",
    "token_embedding + positional_embedding\n",
    "    → attention layer(s) with residual connections\n",
    "    → linear unembed → logits\n",
    "```\n",
    "\n",
    "Positional embeddings are learned (`nn.Embedding`). Each attention layer has `n_heads` independent heads, each computing Q, K, V projections, with a shared W_O output projection. The model returns both logits and per-layer attention patterns.\n",
    "\n",
    "### What constitutes a good solution?\n",
    "\n",
    "- **Describe the mechanism** the model uses to find the max.\n",
    "- **Provide evidence** via attention pattern visualizations, ablation experiments, activation patching, direct logit attribution, or other relevant techniques.\n",
    "- **Present your solution clearly** by utilizing the markdown cells, and presenting clean figures.\n",
    "\n",
    "We recommend starting with Model 1 — it's simple enough that you should be able to fully explain every weight matrix."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7r1ybe7gomm",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "yb8jccpx7ni",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q torch einops nnsight==0.6.3 huggingface_hub matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "byipxtlkn",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json, importlib\n",
    "from pathlib import Path\n",
    "\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "from nnsight import NNsight\n",
    "from huggingface_hub import hf_hub_download\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(f\"Using device: {device}\")\n",
    "\n",
    "# Download model definition and both sets of weights from HuggingFace\n",
    "model_py_path = hf_hub_download(\"andyrdt/04_2026_puzzle_1a\", \"model.py\")\n",
    "spec = importlib.util.spec_from_file_location(\"model\", model_py_path)\n",
    "model_module = importlib.util.module_from_spec(spec)\n",
    "spec.loader.exec_module(model_module)\n",
    "AttentionOnlyTransformer = model_module.AttentionOnlyTransformer\n",
    "\n",
    "# Pre-download both models so later cells don't need network access\n",
    "config_1_path = hf_hub_download(\"andyrdt/04_2026_puzzle_1a\", \"config.json\")\n",
    "weights_1_path = hf_hub_download(\"andyrdt/04_2026_puzzle_1a\", \"model.pt\")\n",
    "config_2_path = hf_hub_download(\"andyrdt/04_2026_puzzle_1b\", \"config.json\")\n",
    "weights_2_path = hf_hub_download(\"andyrdt/04_2026_puzzle_1b\", \"model.pt\")\n",
    "\n",
    "print(\"Downloaded model definition + weights for both models.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "emc3umcjsco",
   "metadata": {},
   "source": [
    "## Helper functions\n",
    "\n",
    "Tokenization helpers and a utility to visualize attention patterns. You'll use these throughout."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "qj51ou1xp9m",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── Model 1 vocab (numbers 0-9, each is its own token) ──\n",
    "NUM_RANGE_1 = 10\n",
    "BOS_1, SEP_1, ANS_1, EOS_1 = 10, 11, 12, 13\n",
    "VOCAB_SIZE_1 = 14\n",
    "TOKEN_NAMES_1 = {10: \"BOS\", 11: \"SEP\", 12: \"ANS\", 13: \"EOS\"}\n",
    "\n",
    "def tokenize_1(nums: list[int]) -> list[int]:\n",
    "    \"\"\"Tokenize a list of numbers for Model 1.\n",
    "    Example: [3, 7, 2] -> [BOS, 3, SEP, 7, SEP, 2, ANS]\"\"\"\n",
    "    tokens = [BOS_1]\n",
    "    for i, n in enumerate(nums):\n",
    "        tokens.append(n)\n",
    "        if i < len(nums) - 1:\n",
    "            tokens.append(SEP_1)\n",
    "    tokens.append(ANS_1)\n",
    "    return tokens\n",
    "\n",
    "def token_labels_1(tokens: list[int]) -> list[str]:\n",
    "    return [TOKEN_NAMES_1.get(t, str(t)) for t in tokens]\n",
    "\n",
    "\n",
    "# ── Model 2 vocab (digits 0-9, two per number) ──\n",
    "BOS_2, SEP_2, ANS_2, EOS_2 = 10, 11, 12, 13\n",
    "VOCAB_SIZE_2 = 14\n",
    "TOKEN_NAMES_2 = {10: \"BOS\", 11: \"SEP\", 12: \"ANS\", 13: \"EOS\"}\n",
    "\n",
    "def tokenize_2(nums: list[int]) -> list[int]:\n",
    "    \"\"\"Tokenize a list of numbers for Model 2.\n",
    "    Example: [42, 7, 85] -> [BOS, 4, 2, SEP, 0, 7, SEP, 8, 5, ANS]\"\"\"\n",
    "    tokens = [BOS_2]\n",
    "    for i, n in enumerate(nums):\n",
    "        tokens.append(n // 10)\n",
    "        tokens.append(n % 10)\n",
    "        if i < len(nums) - 1:\n",
    "            tokens.append(SEP_2)\n",
    "    tokens.append(ANS_2)\n",
    "    return tokens\n",
    "\n",
    "def token_labels_2(tokens: list[int]) -> list[str]:\n",
    "    return [TOKEN_NAMES_2.get(t, str(t)) for t in tokens]\n",
    "\n",
    "\n",
    "# ── Attention visualization ──\n",
    "def _format_attn_ax(ax, attn_matrix, token_labels, title):\n",
    "    \"\"\"Format a single attention heatmap axis.\"\"\"\n",
    "    ax.imshow(attn_matrix, cmap=\"Blues\", vmin=0, vmax=1)\n",
    "    ax.set_xticks(range(len(token_labels)))\n",
    "    ax.set_xticklabels(token_labels, rotation=45, ha=\"right\",\n",
    "                       rotation_mode=\"anchor\", fontsize=7)\n",
    "    ax.set_yticks(range(len(token_labels)))\n",
    "    ax.set_yticklabels(token_labels, fontsize=7)\n",
    "    ax.set_title(title, fontsize=10)\n",
    "\n",
    "\n",
    "def plot_attention(attn_patterns, token_labels, title=\"Attention\"):\n",
    "    \"\"\"Plot attention heatmaps.\n",
    "\n",
    "    Args:\n",
    "        attn_patterns: tensor of shape (n_heads, seq, seq) or list of such tensors\n",
    "                       (one per layer).\n",
    "        token_labels: list of strings for tick labels.\n",
    "        title: plot title.\n",
    "    \"\"\"\n",
    "    if isinstance(attn_patterns, list):\n",
    "        n_layers = len(attn_patterns)\n",
    "        n_heads = attn_patterns[0].shape[0]\n",
    "        fig, axes = plt.subplots(n_layers, n_heads, figsize=(4 * n_heads, 3.5 * n_layers))\n",
    "        if n_layers == 1:\n",
    "            axes = [axes]\n",
    "        for layer_idx in range(n_layers):\n",
    "            attn = attn_patterns[layer_idx].detach().cpu().numpy()\n",
    "            for h in range(n_heads):\n",
    "                _format_attn_ax(axes[layer_idx][h], attn[h], token_labels,\n",
    "                                f\"L{layer_idx}H{h}\")\n",
    "    else:\n",
    "        attn = attn_patterns.detach().cpu().numpy()\n",
    "        n_heads = attn.shape[0]\n",
    "        fig, axes = plt.subplots(1, n_heads, figsize=(4 * n_heads, 3.5))\n",
    "        if n_heads == 1:\n",
    "            axes = [axes]\n",
    "        for h in range(n_heads):\n",
    "            _format_attn_ax(axes[h], attn[h], token_labels, f\"H{h}\")\n",
    "\n",
    "    plt.suptitle(title)\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "print(\"Helpers loaded.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "sc97d3i8oid",
   "metadata": {},
   "source": [
    "---\n",
    "# Model 1: Max of list (0–9), 1-layer attention-only\n",
    "\n",
    "**Task**: Given 5 numbers from 0–9, predict the maximum.\n",
    "\n",
    "**Input format**: `[BOS] n1 [SEP] n2 [SEP] n3 [SEP] n4 [SEP] n5 [ANS]`\n",
    "\n",
    "**Output**: At the `[ANS]` position, the model should predict the max value. Then it should predict `[EOS]`.\n",
    "\n",
    "**Example**: `[BOS] 3 [SEP] 7 [SEP] 2 [SEP] 5 [SEP] 1 [ANS]` → model predicts `7`\n",
    "\n",
    "This model has a single attention layer with 4 heads and no MLPs. The entire computation is: embed → one multi-head attention layer (with residual) → unembed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "yf3l88qh97h",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load Model 1\n",
    "config_1 = json.loads(Path(config_1_path).read_text())\n",
    "raw_model_1 = AttentionOnlyTransformer.from_config(config_1[\"model\"])\n",
    "raw_model_1.load_state_dict(torch.load(weights_1_path, map_location=device, weights_only=True))\n",
    "raw_model_1.eval().to(device)\n",
    "# `NNsight` wraps the model for tracing/interventions, but still exposes the\n",
    "# underlying module tree and parameters for normal inspection.\n",
    "model_1 = NNsight(raw_model_1)\n",
    "\n",
    "print(f\"Model 1 config: {config_1['model']}\")\n",
    "print(f\"Parameters: {sum(p.numel() for p in model_1.parameters()):,}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "xoiyoz1at0i",
   "metadata": {},
   "source": [
    "### Verifying Model 1 works\n",
    "\n",
    "Let's run a few examples and check the model gets them right. We show the full output distribution at the ANS position."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0o4qtnz13lkd",
   "metadata": {},
   "outputs": [],
   "source": [
    "examples_1 = [\n",
    "    [3, 7, 2, 5, 1],\n",
    "    [0, 0, 0, 0, 0],\n",
    "    [9, 1, 8, 2, 7],\n",
    "    [1, 2, 3, 4, 5],\n",
    "]\n",
    "\n",
    "fig, axes = plt.subplots(1, len(examples_1), figsize=(4 * len(examples_1), 3))\n",
    "\n",
    "for idx, nums in enumerate(examples_1):\n",
    "    tokens = tokenize_1(nums)\n",
    "    x = torch.tensor([tokens], device=device)\n",
    "\n",
    "    logits, _ = model_1(x)\n",
    "\n",
    "    # Softmax over number tokens at the ANS position (last token)\n",
    "    probs = torch.softmax(logits[0, -1, :NUM_RANGE_1], dim=-1).detach().cpu()\n",
    "    pred = probs.argmax().item()\n",
    "    true_max = max(nums)\n",
    "\n",
    "    ax = axes[idx]\n",
    "    colors = [\"green\" if i == true_max else \"lightgray\" for i in range(NUM_RANGE_1)]\n",
    "    ax.bar(range(NUM_RANGE_1), probs.numpy(), color=colors)\n",
    "    ax.set_xticks(range(NUM_RANGE_1))\n",
    "    ax.set_ylim(0, 1.1)\n",
    "    ax.set_xlabel(\"Token\")\n",
    "    ax.set_ylabel(\"P(token)\")\n",
    "    status = \"correct\" if pred == true_max else \"WRONG\"\n",
    "    ax.set_title(f\"{nums}\\npred={pred}, true={true_max} ({status})\")\n",
    "\n",
    "plt.suptitle(\"Model 1: output distribution at ANS position\", fontsize=13)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "kydgj8pzhqm",
   "metadata": {},
   "source": [
    "### Example: attention patterns for Model 1\n",
    "\n",
    "Here's what the 4 attention heads look like on a single input. The ANS row (bottom) is where the model reads from the input to make its prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19nh7y44td1",
   "metadata": {},
   "outputs": [],
   "source": [
    "nums = [3, 7, 2, 5, 1]\n",
    "tokens = tokenize_1(nums)\n",
    "x = torch.tensor([tokens], device=device)\n",
    "\n",
    "_, attn_patterns = model_1(x)\n",
    "\n",
    "# `attn_patterns` is a list with one tensor per layer.\n",
    "# Each tensor has shape: (batch, n_heads, seq, seq)\n",
    "attn = attn_patterns[0][0]  # layer 0, batch 0\n",
    "plot_attention(attn, token_labels_1(tokens), title=f\"Model 1: {nums} (max={max(nums)})\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "pvp1ae9wbr",
   "metadata": {},
   "source": [
    "### Your turn!\n",
    "\n",
    "You now have `model_1` (an `NNsight`-wrapped model) and `raw_model_1` (the underlying plain `nn.Module`).\n",
    "\n",
    "The wrapped model still exposes parameters and submodules for direct inspection, while also supporting `NNsight` tracing and interventions.\n",
    "\n",
    "Your goal is to try and understand how the model works.\n",
    "\n",
    "Good luck!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2v1h3znl2op",
   "metadata": {},
   "source": [
    "---\n",
    "# Model 2: Max of list (0–99), 2-layer attention-only, digit tokenization\n",
    "\n",
    "**Task**: Given 5 numbers from 0–99, predict the maximum.\n",
    "\n",
    "**Tokenization**: Each number is split into two digit tokens (tens, ones), always zero-padded. So `42` → tokens `4`, `2` and `7` → tokens `0`, `7`.\n",
    "\n",
    "**Input format**: `[BOS] d1t d1o [SEP] d2t d2o [SEP] ... d5t d5o [ANS]`\n",
    "\n",
    "**Output**: At `[ANS]`, model predicts the tens digit of the max. Then the ones digit. Then `[EOS]`.\n",
    "\n",
    "**Example**: `[BOS] 4 2 [SEP] 1 7 [SEP] 8 5 [SEP] 0 3 [SEP] 6 1 [ANS]` → model predicts `8` then `5`\n",
    "\n",
    "This model has **2 attention layers** with 4 heads each. The interesting question is how the layers divide the work — a 1-layer model can learn the tens digit (100% accuracy) but plateaus at ~40% for the ones digit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "145kvifo56y8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load Model 2\n",
    "config_2 = json.loads(Path(config_2_path).read_text())\n",
    "raw_model_2 = AttentionOnlyTransformer.from_config(config_2[\"model\"])\n",
    "raw_model_2.load_state_dict(torch.load(weights_2_path, map_location=device, weights_only=True))\n",
    "raw_model_2.eval().to(device)\n",
    "model_2 = NNsight(raw_model_2)\n",
    "\n",
    "print(f\"Model 2 config: {config_2['model']}\")\n",
    "print(f\"Parameters: {sum(p.numel() for p in model_2.parameters()):,}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ls905thepu",
   "metadata": {},
   "source": [
    "### Verifying Model 2 works\n",
    "\n",
    "For Model 2, the output is two tokens: tens digit then ones digit. We feed the input up to `[ANS]`, get the tens prediction, then feed that back to get the ones prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "pxuozl2wke",
   "metadata": {},
   "outputs": [],
   "source": [
    "examples_2 = [\n",
    "    [42, 17, 85, 3, 61],\n",
    "    [99, 0, 50, 25, 75],\n",
    "    [87, 86, 85, 84, 83],\n",
    "    [9, 19, 29, 39, 49],\n",
    "]\n",
    "\n",
    "for nums in examples_2:\n",
    "    tokens = tokenize_2(nums)\n",
    "    x = torch.tensor([tokens], device=device)\n",
    "\n",
    "    # Get tens digit prediction (at ANS position)\n",
    "    logits_tens, _ = model_2(x)\n",
    "    pred_tens = logits_tens[0, -1, :10].argmax().item()\n",
    "\n",
    "    # Feed tens digit back, get ones digit prediction\n",
    "    tokens_ext = tokens + [pred_tens]\n",
    "    x_ext = torch.tensor([tokens_ext], device=device)\n",
    "    logits_ones, _ = model_2(x_ext)\n",
    "    pred_ones = logits_ones[0, -1, :10].argmax().item()\n",
    "\n",
    "    pred_num = pred_tens * 10 + pred_ones\n",
    "    true_max = max(nums)\n",
    "    status = \"correct\" if pred_num == true_max else \"WRONG\"\n",
    "    print(f\"  {nums} → predicted {pred_num:2d}, true max {true_max:2d}  [{status}]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6o58cfq6yl",
   "metadata": {},
   "source": [
    "### Example: attention patterns for Model 2\n",
    "\n",
    "With 2 layers, we can see how the model builds up its computation across layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "hou1733h968",
   "metadata": {},
   "outputs": [],
   "source": [
    "nums = [42, 17, 85, 3, 61]\n",
    "tokens = tokenize_2(nums)\n",
    "x = torch.tensor([tokens], device=device)\n",
    "\n",
    "_, attn_patterns = model_2(x)\n",
    "\n",
    "# `attn_patterns` is a list with one tensor per layer.\n",
    "# Each tensor has shape: (batch, n_heads, seq, seq)\n",
    "attn_l0 = attn_patterns[0][0]  # layer 0, batch 0\n",
    "attn_l1 = attn_patterns[1][0]  # layer 1, batch 0\n",
    "\n",
    "plot_attention([attn_l0, attn_l1], token_labels_2(tokens),\n",
    "               title=f\"Model 2: {nums} (max={max(nums)})\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "hrgmzhjxrz9",
   "metadata": {},
   "source": [
    "### Your turn!\n",
    "\n",
    "You now have `model_2` (an `NNsight`-wrapped model) and `raw_model_2` (the underlying plain `nn.Module`).\n",
    "\n",
    "As above, the wrapped model still exposes parameters and submodules for direct inspection, while also supporting `NNsight` tracing and interventions.\n",
    "\n",
    "Your goal is to try and understand how the model works.\n",
    "\n",
    "Good luck!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}