{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Module 20: Multimodal"
      ],
      "id": "da9c6b31-2884-4405-b6a6-16b613c83453"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [],
      "id": "eb6ebccd-2297-4cd4-b789-cb88440c89d4"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "dd379617-abe0-4924-bb2b-f233c4213509"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Introduction\n",
        "\n",
        "Every module so far ends at a model that only ever reads **token IDs**:\n",
        "tokenize (m03), embed (m04), attend (m05), predict the next token\n",
        "(m06–m08). The whole book is text. This module opens the second modality\n",
        "— **pixels** — and answers how one model can relate a picture to the\n",
        "words that describe it.\n",
        "\n",
        "A **multimodal model** processes more than one kind of input (here:\n",
        "images *and* text) in a shared representation. There are only two\n",
        "genuinely new ideas, and everything else is the transformer you already\n",
        "built:\n",
        "\n",
        "- **Pixels become tokens.** Cut an image into a grid of fixed-size\n",
        "  **patches**, flatten each, and project it into the same vector space\n",
        "  token embeddings live in. A patch is a “visual word.” This is a\n",
        "  **Vision Transformer** (ViT).\n",
        "- **A shared space.** Encode an image and a caption to two vectors and\n",
        "  train so **matched** pairs sit close and **mismatched** pairs far —\n",
        "  the **CLIP** contrastive objective. This is what makes “a photo of a\n",
        "  cat” land next to 🐱.\n",
        "\n",
        "Why it matters for LLMs:\n",
        "\n",
        "- Modern frontier models (GPT-4o, Gemini, Claude, Llama 4) are natively\n",
        "  multimodal — they see images by turning them into token-like patch\n",
        "  embeddings.\n",
        "- CLIP-style alignment is the backbone of text-to-image search,\n",
        "  zero-shot image classification, and the vision encoders bolted onto\n",
        "  today’s chat models.\n",
        "\n",
        "### What You’ll Learn\n",
        "\n",
        "After this module, you can:\n",
        "\n",
        "- **Patchify** an image into a sequence a transformer can read, and\n",
        "  build a ViT **patch-embedding** layer (patch projection + positional\n",
        "  embedding + CLS token).\n",
        "- Explain why images and text must be **L2-normalized** into a shared\n",
        "  space, and build the **CLIP contrastive loss** — a symmetric\n",
        "  cross-entropy over the batch similarity matrix — from scratch.\n",
        "- Measure alignment with **image→text retrieval accuracy**, and *watch*\n",
        "  a tiny CLIP pull two modalities together until the similarity matrix’s\n",
        "  diagonal lights up.\n",
        "- Reason about the **temperature** (logit scale) and the\n",
        "  batch-as-negatives trick, and name the pitfalls (skipping\n",
        "  normalization, one-sided loss, unclamped τ).\n",
        "\n",
        "### Prerequisites\n",
        "\n",
        "This module requires familiarity with:\n",
        "\n",
        "- [Module 04: Embeddings](../m04_embeddings/lesson.qmd) — patch\n",
        "  embeddings *are* token embeddings for images; the positional-embedding\n",
        "  idea carries over directly.\n",
        "- [Module 05: Attention](../m05_attention/lesson.qmd) — a ViT is a\n",
        "  transformer encoder over the patch sequence.\n",
        "- [Module 07: Training](../m07_training/lesson.qmd) — the contrastive\n",
        "  loss is a cross-entropy, optimized exactly as you trained the GPT.\n",
        "- [Module 17: Evaluation](../m17_evaluation/lesson.qmd) — retrieval\n",
        "  accuracy is the metric that says whether the two modalities actually\n",
        "  aligned.\n",
        "\n",
        "## Intuition: Patches as Visual Words\n",
        "\n",
        "A language model never sees letters — it sees a sequence of token\n",
        "vectors. To feed an image to the *same* machinery, we need to turn it\n",
        "into a sequence of vectors too. The Vision Transformer’s answer is\n",
        "disarmingly simple: **chop the image into a grid of small square\n",
        "patches, and treat each patch as one token.**\n",
        "\n",
        "A 32×32 image with 16×16 patches becomes a 2×2 grid — four “visual\n",
        "words.” Each patch (a little block of pixels) is flattened into a vector\n",
        "and linearly projected to the model’s embedding dimension `d`, exactly\n",
        "like a token ID is looked up into a `d`-dim embedding in m04. Add a\n",
        "**positional embedding** so the model knows where each patch sat\n",
        "(patches have no inherent order once flattened), prepend a learnable\n",
        "**\\[CLS\\]** token whose final vector summarizes the whole image, and you\n",
        "have a sequence the transformer reads with no other changes.\n",
        "\n",
        "Walk the pipeline one step at a time:"
      ],
      "id": "07d6caec-c98d-4d8d-a524-1c0cd18edccc"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "5d4d65e9-0f7f-4914-a4f5-8772340fbd59"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "5b36597a-f93a-47d6-b4e8-75fb49899c8d"
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [],
      "source": [],
      "id": "e608204c"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "925cdae4-72dc-486c-9701-d1f063062eb5"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## The Math: Patchify & Patch Embeddings\n",
        "\n",
        "Let an image be a tensor of shape $(C, H, W)$ — channels, height, width.\n",
        "With a patch size $P$ that divides $H$ and $W$, patchify reshapes it\n",
        "into\n",
        "\n",
        "$$(C, H, W) \\;\\longrightarrow\\; (N,\\; P^2 C), \\qquad N = \\frac{H}{P}\\cdot\\frac{W}{P},$$\n",
        "\n",
        "a sequence of $N$ patch vectors, each of length $P^2 C$. A linear layer\n",
        "with weight $W_e \\in \\mathbb{R}^{(P^2 C)\\times d}$ projects every patch\n",
        "to the embedding size, we add a learned positional embedding\n",
        "$E_{\\text{pos}}$, and prepend a learned CLS token $x_{\\text{cls}}$:\n",
        "\n",
        "$$Z = \\big[\\, x_{\\text{cls}};\\; X W_e \\,\\big] + E_{\\text{pos}}\n",
        "\\;\\in\\; \\mathbb{R}^{(N+1)\\times d}.$$\n",
        "\n",
        "That is the entire ViT input layer — and it is line-for-line the same\n",
        "shape story as m04’s token embeddings, only the “lookup” is a linear\n",
        "projection of pixels. From $Z$ onward it is a standard transformer\n",
        "encoder (m05/m06).\n",
        "\n",
        "## Code: Patch Embeddings from Scratch\n",
        "\n",
        "The implementation lives in `multimodal.py`. `patchify` does the reshape\n",
        "(it is a pure tensor permutation — no learned parameters, and lossless):"
      ],
      "id": "d8cec8da-6d68-4005-ae6d-8771902d8581"
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "image (3, 32, 32)  ->  patches (4, 768)\n",
            "N = 4 patches, each 768 numbers (16·16·3)"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "from multimodal import patchify, num_patches\n",
        "\n",
        "image = torch.randn(3, 32, 32)          # one RGB image\n",
        "patches = patchify(image, patch_size=16)\n",
        "print(f\"image {tuple(image.shape)}  ->  patches {tuple(patches.shape)}\")\n",
        "print(f\"N = {num_patches(32, 32, 16)} patches, each {patches.shape[1]} numbers (16·16·3)\")"
      ],
      "id": "23fef0d6"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "`PatchEmbedding` wraps the projection, the positional embedding, and the\n",
        "CLS token into the layer a ViT starts with:"
      ],
      "id": "728e2e9d-805d-471d-9828-eb9edadc2487"
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "patch tokens: (4, 5, 64)   (4 images, 4 patches + 1 CLS, d=64)\n",
            "pooled image vector: (4, 64)"
          ]
        }
      ],
      "source": [
        "from multimodal import PatchEmbedding\n",
        "\n",
        "patch_embed = PatchEmbedding(image_size=32, patch_size=16, in_channels=3, embed_dim=64)\n",
        "batch = torch.randn(4, 3, 32, 32)       # 4 images\n",
        "tokens = patch_embed(batch)\n",
        "print(f\"patch tokens: {tuple(tokens.shape)}   (4 images, 4 patches + 1 CLS, d=64)\")\n",
        "\n",
        "image_vector = patch_embed.pool(tokens) # the CLS row summarizes each image\n",
        "print(f\"pooled image vector: {tuple(image_vector.shape)}\")"
      ],
      "id": "6dd197d6"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "> **Key Insight**\n",
        ">\n",
        "> A patch embedding is a token embedding. Once an image is a sequence of\n",
        "> `d`-dim vectors, **nothing downstream knows or cares that it came from\n",
        "> pixels** — the same attention, the same blocks, the same training.\n",
        "> Multimodality is mostly a question of getting every modality into one\n",
        "> shared sequence of vectors.\n",
        "\n",
        "## Intuition: One Space for Two Modalities\n",
        "\n",
        "Turning an image into tokens lets a transformer *process* it, but it\n",
        "does not yet connect the image to language. For that we need image\n",
        "vectors and text vectors to live in **one shared space**, where\n",
        "closeness means “these describe the same thing.”\n",
        "\n",
        "CLIP builds that space with a beautifully simple training signal. Take a\n",
        "big batch of $(image, caption)$ pairs. Encode all images to vectors and\n",
        "all captions to vectors. Now form every possible pairing — an\n",
        "$N \\times N$ grid of similarities. The **matched** pairs are exactly the\n",
        "diagonal; every off-diagonal cell is a **mismatch**. Train to make the\n",
        "diagonal win: pull each image toward its own caption and push it away\n",
        "from the other $N-1$ captions in the batch, which serve as free\n",
        "**negatives**."
      ],
      "id": "5140050b-3567-426e-b04f-6450d526c760"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "7aef6f73-dd28-4e4d-96cb-f4380ae0dca3"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "18199dfe-a88a-41b2-bead-19fd42676f90"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "e89e23e4-6074-4028-a8b6-aed27de4e680"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## The Math: Contrastive Alignment (CLIP)\n",
        "\n",
        "Let $I \\in \\mathbb{R}^{N\\times d}$ be the batch of image embeddings and\n",
        "$T \\in \\mathbb{R}^{N\\times d}$ the paired text embeddings. First\n",
        "**L2-normalize** each row, $\\hat I_i = I_i / \\lVert I_i\\rVert$, so a dot\n",
        "product is a cosine similarity. The scaled similarity matrix is\n",
        "\n",
        "$$\\text{logits} = s\\,\\hat I \\hat T^{\\top} \\;\\in\\; \\mathbb{R}^{N\\times N},\n",
        "\\qquad s = \\exp(t),$$\n",
        "\n",
        "where $t$ is a **learned temperature** (CLIP stores it in log space and\n",
        "clamps it). Row $i$ scores image $i$ against all $N$ captions; the\n",
        "correct label is $i$ (the diagonal). So it is just cross-entropy — done\n",
        "in **both** directions and averaged:\n",
        "\n",
        "$$\\mathcal{L} = \\tfrac12\\Big(\n",
        "\\underbrace{\\text{CE}(\\text{logits},\\, y)}_{\\text{image}\\to\\text{text}} +\n",
        "\\underbrace{\\text{CE}(\\text{logits}^{\\top},\\, y)}_{\\text{text}\\to\\text{image}}\n",
        "\\Big),\n",
        "\\qquad y = (0, 1, \\dots, N-1).$$\n",
        "\n",
        "The symmetry matters: the image→text term makes each image find its\n",
        "caption; the text→image term makes each caption find its image. Drop one\n",
        "and retrieval in that direction suffers.\n",
        "\n",
        "## Code: CLIP from Scratch\n",
        "\n",
        "All of it is a handful of lines in `multimodal.py`. First the\n",
        "normalize-then-score step, and the fact that a matched batch wins its\n",
        "diagonal:"
      ],
      "id": "737a427c-fbd0-4eee-a1ff-e77fa5f934f2"
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "norms before: [1.0, 0.9999999403953552, 0.9999999403953552] (all 1.0)\n",
            "argmax of each row: [0, 1, 2, 3, 4, 5] (diagonal 0..5 wins)"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "from multimodal import l2_normalize, clip_logits, clip_loss, contrastive_accuracy\n",
        "\n",
        "torch.manual_seed(0)\n",
        "img = torch.randn(6, 32)                 # 6 image vectors\n",
        "txt = img + 0.05 * torch.randn(6, 32)    # captions that (roughly) match\n",
        "\n",
        "print(\"norms before:\", l2_normalize(img).norm(dim=-1)[:3].tolist(), \"(all 1.0)\")\n",
        "logits = clip_logits(img, txt, logit_scale=10.0)\n",
        "print(\"argmax of each row:\", logits.argmax(dim=1).tolist(), \"(diagonal 0..5 wins)\")"
      ],
      "id": "da46cae2"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The symmetric loss and the retrieval metric:"
      ],
      "id": "9fd684d7-cd78-4152-b000-2cf1f89449fe"
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "loss  matched=0.000   scrambled=16.693\n",
            "retrieval accuracy (matched):   1.00\n",
            "retrieval accuracy (scrambled): 0.00"
          ]
        }
      ],
      "source": [
        "matched   = clip_loss(img, txt)                 # paired correctly\n",
        "scrambled = clip_loss(img, txt.flip(0))         # every pair wrong\n",
        "print(f\"loss  matched={matched:.3f}   scrambled={scrambled:.3f}\")\n",
        "print(f\"retrieval accuracy (matched):   {contrastive_accuracy(img, txt):.2f}\")\n",
        "print(f\"retrieval accuracy (scrambled): {contrastive_accuracy(img, txt.flip(0)):.2f}\")"
      ],
      "id": "bd193b46"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "`CLIPModel` packages two projection encoders and the learned\n",
        "temperature. (Here the encoders are small linear heads over feature\n",
        "vectors so the demo is fast; in a real CLIP the image encoder is the ViT\n",
        "you built above and the text encoder is a Transformer from m06 — the\n",
        "contrastive machinery is identical.)"
      ],
      "id": "ff102903-9ee4-4ed6-9fb7-5682874e04c0"
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "learned temperature (logit scale): 14.29  (CLIP init = 1/0.07)\n",
            "similarity matrix shape: (6, 6)"
          ]
        }
      ],
      "source": [
        "from multimodal import CLIPModel\n",
        "\n",
        "model = CLIPModel(image_feat_dim=32, text_feat_dim=32, embed_dim=16)\n",
        "print(f\"learned temperature (logit scale): {model.logit_scale():.2f}  (CLIP init = 1/0.07)\")\n",
        "print(f\"similarity matrix shape: {tuple(model(img, txt).shape)}\")"
      ],
      "id": "614d8024"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Watch Two Modalities Align\n",
        "\n",
        "The payoff: train the tiny CLIP and *see* the shared space form.\n",
        "`demonstrate_clip` makes synthetic pairs — each concept gets an\n",
        "image-feature view and a text-feature view — that start unrelated, then\n",
        "minimizes the contrastive loss. The similarity matrix begins as noise\n",
        "and its **diagonal lights up** as matched pairs are pulled together;\n",
        "retrieval accuracy climbs from chance toward 1.0."
      ],
      "id": "119a98d7-e0a1-42f1-8d6b-ba20567c5f39"
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "============================================================\n",
            "CLIP: aligning two modalities (synthetic)\n",
            "============================================================\n",
            "  N=12 pairs, embed_dim=12, 300 steps\n",
            "\n",
            "    step      loss   top-1 acc\n",
            "       0    7.9794       0.000\n",
            "      75    0.0031       1.000\n",
            "     150    0.0011       1.000\n",
            "     225    0.0005       1.000\n",
            "     299    0.0003       1.000\n",
            "\n",
            "  Retrieval accuracy: 0.000 -> 1.000  (chance = 0.083)"
          ]
        }
      ],
      "source": [
        "from multimodal import demonstrate_clip\n",
        "\n",
        "result = demonstrate_clip(n_concepts=12, steps=300, seed=0)"
      ],
      "id": "31cc4322"
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {},
      "outputs": [],
      "source": [],
      "id": "e19a4f10"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Drag the slider to step through training. Watch the muddy grid resolve\n",
        "into a bright diagonal — that diagonal *is* the alignment:"
      ],
      "id": "237dd3ef-8ad8-48c9-8250-a8c50050f191"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "4fd9d8a9-a895-4a92-89bd-a7f513c0488a"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "4d10508d-bc14-4d32-8b3f-789561854735"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "89fe08c3-dbf3-406b-bf07-62b75bec897e"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "> **Try This**\n",
        ">\n",
        "> 1.  **Fewer steps**: rerun `demonstrate_clip(steps=20)` — the diagonal\n",
        ">     only half-forms, and accuracy stalls below 1.0. Alignment needs\n",
        ">     enough gradient steps.\n",
        "> 2.  **Harder batch**: bump `n_concepts` to 32. More in-batch negatives\n",
        ">     make each classification harder — the same reason CLIP trained\n",
        ">     with enormous batch sizes.\n",
        "> 3.  **More noise**: raise `noise=1.5`. The two views share less, so\n",
        ">     the ceiling drops.\n",
        "\n",
        "## Interactive Exploration: Temperature\n",
        "\n",
        "The temperature $s = \\exp(t)$ scales the cosine similarities *before*\n",
        "the softmax, which controls how sharply the model must separate the\n",
        "matched pair from the rest. Small scale (high temperature) → a soft,\n",
        "forgiving distribution; large scale (low temperature) → a peaky one that\n",
        "demands the diagonal dominate. CLIP **learns** this value, and\n",
        "**clamps** it so it cannot run away and make the logits explode.\n",
        "\n",
        "Drag the logit scale and watch the softmax over one image’s similarities\n",
        "to five captions (the matched caption is the first bar):"
      ],
      "id": "9669e6ab-fed1-4899-8658-6cf94a52583d"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "472ef48c-9853-4a11-b46b-5ff3c5b3227f"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "0b9ca938-421d-45aa-b6cb-c8b3ac564306"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "> **Key Insight**\n",
        ">\n",
        "> Temperature does not change *which* caption is most similar — it\n",
        "> changes how confidently the loss insists on it. Too low a scale and\n",
        "> every pair looks equally good (no learning signal); too high and a\n",
        "> single hard example dominates the gradient. The learned-and-clamped\n",
        "> value is CLIP threading that needle.\n",
        "\n",
        "## Common Pitfalls\n",
        "\n",
        "> **Forgetting to normalize**\n",
        ">\n",
        "> Skip the L2 step and your “cosine” similarity is a raw dot product —\n",
        "> it grows with vector *length*, so a long, off-topic vector can\n",
        "> outscore a short, on-topic one. Always normalize both towers before\n",
        "> the similarity.\n",
        "\n",
        "> **Using a one-sided loss**\n",
        ">\n",
        "> Only `CE(logits, y)` (image→text) trains images to find captions, not\n",
        "> captions to find images — text→image retrieval will lag. The loss must\n",
        "> be **symmetric**.\n",
        "\n",
        "> **An unclamped learned temperature**\n",
        ">\n",
        "> `logit_scale` is learned, and gradient pressure pushes it up (sharper\n",
        "> = lower loss). Left unclamped it can blow up and destabilize training\n",
        "> — CLIP clamps `exp(t)` at 100. `CLIPModel.logit_scale()` does the\n",
        "> clamp for you.\n",
        "\n",
        "> **Position-embedding count must match**\n",
        ">\n",
        "> `PatchEmbedding` allocates `n_patches + 1` positions (the `+1` is\n",
        "> CLS). Change the image or patch size and the count changes; a mismatch\n",
        "> is a silent shape bug.\n",
        "\n",
        "> **Small batches weaken the signal**\n",
        ">\n",
        "> The batch *is* the negative set — with `N` items each image sees only\n",
        "> `N−1` negatives. Tiny batches make the task too easy to learn a\n",
        "> discriminative space, which is why real CLIP used batches in the tens\n",
        "> of thousands.\n",
        "\n",
        "## Exercises\n",
        "\n",
        "### Exercise 1: Non-square patches / different sizes\n",
        "\n",
        "`patchify` requires the patch size to divide the image. Write a check\n",
        "that, given an image and a patch size, either returns the patch count or\n",
        "explains why it doesn’t tile evenly."
      ],
      "id": "980e5132-cb6d-41ad-9141-efd34e75d3f2"
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "(True, 196)\n",
            "(False, '16 does not divide 30x32')"
          ]
        }
      ],
      "source": [
        "from multimodal import num_patches\n",
        "\n",
        "def can_patchify(height: int, width: int, patch_size: int):\n",
        "    \"\"\"Return (True, N) if it tiles evenly, else (False, reason).\"\"\"\n",
        "    # Your implementation here\n",
        "    if height % patch_size or width % patch_size:\n",
        "        return (False, f\"{patch_size} does not divide {height}x{width}\")\n",
        "    return (True, num_patches(height, width, patch_size))\n",
        "\n",
        "print(can_patchify(224, 224, 16))   # (True, 196)\n",
        "print(can_patchify(30, 32, 16))     # (False, ...)"
      ],
      "id": "56337308"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Exercise 2: CLS vs. mean pooling\n",
        "\n",
        "`PatchEmbedding.pool` takes the CLS token. Write a mean-pool alternative\n",
        "(average the patch rows, excluding CLS) and compare the shapes."
      ],
      "id": "0ab6c9e3-de2f-4a52-8dca-debeff7d8317"
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CLS pool:  (2, 32)\n",
            "mean pool: (2, 32)"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "from multimodal import PatchEmbedding\n",
        "\n",
        "pe = PatchEmbedding(image_size=32, patch_size=16, embed_dim=32)\n",
        "tokens = pe(torch.randn(2, 3, 32, 32))\n",
        "\n",
        "def mean_pool(tokens):\n",
        "    \"\"\"Average the patch tokens (skip index 0, the CLS token).\"\"\"\n",
        "    # Your implementation here\n",
        "    return tokens[:, 1:].mean(dim=1)\n",
        "\n",
        "print(\"CLS pool: \", tuple(pe.pool(tokens).shape))\n",
        "print(\"mean pool:\", tuple(mean_pool(tokens).shape))"
      ],
      "id": "550b1d50"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Exercise 3: Break normalization\n",
        "\n",
        "Reimplement `clip_logits` *without* the L2 normalization and confirm\n",
        "that scaling one image’s features by 10× wrongly changes its ranking."
      ],
      "id": "64c30293-d710-4c80-ab1d-23b37a92f7ff"
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "normalized argmax row 0: 0 (still 0)"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "from multimodal import clip_logits\n",
        "\n",
        "torch.manual_seed(0)\n",
        "img = torch.randn(4, 16)\n",
        "txt = img.clone()\n",
        "# Your implementation here: compare normalized vs un-normalized when img[0] *= 10\n",
        "scaled = img.clone(); scaled[0] *= 10.0\n",
        "print(\"normalized argmax row 0:\", clip_logits(scaled, txt)[0].argmax().item(), \"(still 0)\")"
      ],
      "id": "a0c10715"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Summary\n",
        "\n",
        "Key takeaways:\n",
        "\n",
        "1.  **A patch is a token.** `patchify` turns a $(C, H, W)$ image into a\n",
        "    sequence of $N = (H/P)(W/P)$ patch vectors; a linear projection +\n",
        "    positional embedding + CLS token (`PatchEmbedding`) makes it exactly\n",
        "    the input a transformer already reads. Multimodality is mostly “get\n",
        "    every modality into one sequence of vectors.”\n",
        "2.  **A shared space is learned by contrast.** CLIP encodes images and\n",
        "    text to vectors, L2-normalizes them, and trains the batch similarity\n",
        "    matrix so the **diagonal** (matched pairs) wins — a **symmetric\n",
        "    cross-entropy** with the in-batch mismatches as free negatives.\n",
        "3.  **Normalize, then score.** Cosine similarity needs unit vectors; the\n",
        "    raw dot product would rank by length. `l2_normalize` before\n",
        "    `clip_logits`, always.\n",
        "4.  **Temperature is learned and clamped.** The logit scale $\\exp(t)$\n",
        "    sets how sharply the loss separates the matched pair; CLIP learns it\n",
        "    and clamps it at 100.\n",
        "5.  **The metric is retrieval accuracy.** `contrastive_accuracy` — does\n",
        "    each image’s nearest caption end up being its own? — climbs from\n",
        "    chance to 1.0 as the space forms, and *is* the alignment made\n",
        "    measurable.\n",
        "\n",
        "## What’s Next\n",
        "\n",
        "You now have both halves of a multimodal model: a vision front-end that\n",
        "turns pixels into tokens, and a contrastive objective that binds vision\n",
        "and language into one space. The natural next step is **fusion** —\n",
        "letting a language model *attend* to image features (cross-attention, as\n",
        "in Flamingo) so it can describe, answer questions about, and reason over\n",
        "what it sees, rather than only retrieve. That turns CLIP-style alignment\n",
        "into a full vision-language *generator*.\n",
        "\n",
        "### Going Deeper\n",
        "\n",
        "**Core Papers:**\n",
        "\n",
        "- [Learning Transferable Visual Models From Natural Language Supervision\n",
        "  (CLIP)](https://arxiv.org/abs/2103.00020) — Radford et al. (2021):\n",
        "  400M image-text pairs, the symmetric contrastive objective and learned\n",
        "  temperature built here, zero-shot transfer.\n",
        "- [An Image is Worth 16×16 Words\n",
        "  (ViT)](https://arxiv.org/abs/2010.11929) — Dosovitskiy et al. (2020):\n",
        "  patchify + linear projection + position embeddings\n",
        "  - CLS token — the vision front-end of this module.\n",
        "- [Flamingo: a Visual Language Model for Few-Shot\n",
        "  Learning](https://arxiv.org/abs/2204.14198) — Alayrac et al. (2022):\n",
        "  gated cross-attention fusing vision features into a frozen LM — the\n",
        "  “what’s next” fusion path.\n",
        "\n",
        "**Practical Resources:**\n",
        "\n",
        "- [ALIGN](https://arxiv.org/abs/2102.05918) — Jia et al. (2021): the\n",
        "  same contrastive recipe at billion-pair, noisy-caption scale.\n",
        "- [BLIP](https://arxiv.org/abs/2201.12086) — Li et al. (2022):\n",
        "  bootstrapped captioning that unifies contrastive alignment with\n",
        "  generation."
      ],
      "id": "9416e570-276b-412e-a063-900caf99b9ad"
    }
  ],
  "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"
    }
  }
}