{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=attention-scoring-md-1 fw=all -->\n",
    "\n",
    "# Attention Scoring and Masking\n",
    "\n",
    "In that section we computed attention weights from\n",
    "distance-based kernels such as the Gaussian. Distances are slightly more\n",
    "expensive to compute than dot products, and once the softmax\n",
    "the equation guarantees nonnegative normalized weights, we\n",
    "are free to pick any *scoring function* $a(\\mathbf{q}, \\mathbf{k})$ we like.\n",
    "This section settles the choice used by essentially every modern\n",
    "architecture—the scaled dot product—and builds the two pieces of machinery\n",
    "that make it practical: a masked softmax that lets one batched computation\n",
    "serve sequences of different lengths (and autoregressive models that must\n",
    "not look ahead), and batched matrix multiplication. We close with the story\n",
    "of where these ideas came from: the machine translation problem that turned\n",
    "\"learning to align\" into the attention mechanism."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "attention-scoring-attention-scoring-and-masking"
   },
   "source": [
    "%matplotlib inline\n",
    "from d2l import jax as d2l\n",
    "from flax import nnx\n",
    "import jax\n",
    "from jax import numpy as jnp\n",
    "import math"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=attention-scoring-md-2 fw=all -->\n",
    "\n",
    "## Dot-Product Attention\n",
    "\n",
    "Let's review the attention function (without exponentiation) arising from\n",
    "the Gaussian kernel:\n",
    "\n",
    "$$\n",
    "a(\\mathbf{q}, \\mathbf{k}_i) = -\\frac{1}{2} \\|\\mathbf{q} - \\mathbf{k}_i\\|^2  = \\mathbf{q}^\\top \\mathbf{k}_i -\\frac{1}{2} \\|\\mathbf{k}_i\\|^2  -\\frac{1}{2} \\|\\mathbf{q}\\|^2.\n",
    "$$\n",
    "\n",
    "First, the final term depends on $\\mathbf{q}$ only, so it is identical for\n",
    "all keys, and the softmax normalization the equation\n",
    "removes it entirely. Second, if the key norms $\\|\\mathbf{k}_i\\|$ are all\n",
    "equal, the middle term drops out the same way, and the Gaussian kernel and\n",
    "the dot product induce identical attention weights. In general the norms are\n",
    "not equal, and dropping the term is a modeling decision rather than an\n",
    "approximation: we adopt the dot product $\\mathbf{q}^\\top \\mathbf{k}_i$ as a\n",
    "compatibility function in its own right — one that learned query and key\n",
    "representations can shape freely — with the Gaussian expansion surviving as\n",
    "the exact special case of equal key norms.\n",
    "\n",
    "One adjustment is still needed, to keep the magnitude of the scores under\n",
    "control. Assume that all elements of the query $\\mathbf{q} \\in \\mathbb{R}^d$\n",
    "and the key $\\mathbf{k}_i \\in \\mathbb{R}^d$ are independent random variables\n",
    "with zero mean and unit variance. The dot product of the two vectors then\n",
    "has zero mean and variance $d$: the typical score grows like $\\sqrt{d}$ for\n",
    "no reason other than vector length. Rescaling by $1/\\sqrt{d}$ keeps the\n",
    "score variance at $1$ regardless of dimension, and yields the *scaled\n",
    "dot-product attention* scoring function of the Transformer\n",
    "[@Vaswani.Shazeer.Parmar.ea.2017]:\n",
    "\n",
    "$$ a(\\mathbf{q}, \\mathbf{k}_i) = \\mathbf{q}^\\top \\mathbf{k}_i / \\sqrt{d}.$$\n",
    "\n",
    "The attention weights are obtained, as always, with the softmax:\n",
    "\n",
    "$$\\alpha(\\mathbf{q}, \\mathbf{k}_i) = \\mathrm{softmax}(a(\\mathbf{q}, \\mathbf{k}_i)) = \\frac{\\exp(\\mathbf{q}^\\top \\mathbf{k}_i / \\sqrt{d})}{\\sum_{j=1}^m \\exp(\\mathbf{q}^\\top \\mathbf{k}_j / \\sqrt{d})}.$$\n",
    "\n",
    "### Softmax Saturation and the $1/\\sqrt{d}$ Factor\n",
    "\n",
    "Why insist on unit score variance? The softmax saturates: once one score\n",
    "exceeds the others by a large margin, the winning weight approaches $1$, the\n",
    "rest approach $0$, and, since the Jacobian of the softmax is\n",
    "\n",
    "$$\n",
    "\\frac{\\partial \\boldsymbol{\\alpha}}{\\partial \\mathbf{a}} = \\mathrm{diag}(\\boldsymbol{\\alpha}) - \\boldsymbol{\\alpha} \\boldsymbol{\\alpha}^\\top,\n",
    "$$\n",
    "\n",
    "which tends to the zero matrix as $\\boldsymbol{\\alpha}$ approaches a one-hot\n",
    "vector, gradients stop flowing to everything upstream of the scores. The\n",
    "softmax's gradient never vanishes exactly, as we noted in\n",
    "that section, but it can come arbitrarily close. Let's\n",
    "measure both effects rather than take them on faith. We draw random queries\n",
    "and keys with unit-variance entries, compute attention over $64$ candidate\n",
    "keys, and record the entropy of the resulting weight distribution as the\n",
    "dimension $d$ grows, with and without the $1/\\sqrt{d}$ factor."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "attention-scoring-softmax-saturation-and-the-1-sqrt-d-factor-1"
   },
   "source": [
    "def attention_stats(d, num_keys=64, num_queries=256, scaled=True, seed=0):\n",
    "    key_q, key_k = jax.random.split(jax.random.key(seed))\n",
    "    q = jax.random.normal(key_q, (num_queries, d))\n",
    "    keys = jax.random.normal(key_k, (num_queries, num_keys, d))\n",
    "    scores = (keys @ q[..., None]).squeeze(-1)\n",
    "    if scaled:\n",
    "        scores = scores / math.sqrt(d)\n",
    "    alpha = jax.nn.softmax(scores, axis=-1)\n",
    "    entropy = -(alpha * jnp.log(alpha + 1e-12)).sum(-1).mean()\n",
    "    # Frobenius norm of the softmax Jacobian diag(alpha) - alpha alpha^T\n",
    "    jac = (alpha[..., None] * jnp.eye(num_keys)\n",
    "           - alpha[..., :, None] * alpha[..., None, :])\n",
    "    return float(entropy), float(jnp.sqrt((jac**2).sum((-2, -1))).mean())\n",
    "\n",
    "ds = [2**k for k in range(2, 10)]\n",
    "entropies = [[attention_stats(d, scaled=s)[0] for d in ds]\n",
    "             for s in (False, True)]\n",
    "d2l.plot(ds, entropies, 'dimension d', 'entropy (nats)',\n",
    "         legend=['unscaled', 'scaled by 1/sqrt(d)'], xscale='log')"
   ],
   "execution_count": 2,
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "image/svg+xml": "<?xml version=\"1.0\" encoding=\"utf-8\" standalone=\"no\"?>\n<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n  \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n<svg xmlns:xlink=\"http://www.w3.org/1999/xlink\" width=\"236.740625pt\" height=\"183.35625pt\" viewBox=\"0 0 236.740625 183.35625\" xmlns=\"http://www.w3.org/2000/svg\" version=\"1.1\">\n <metadata>\n  <rdf:RDF xmlns:dc=\"http://purl.org/dc/elements/1.1/\" xmlns:cc=\"http://creativecommons.org/ns#\" xmlns:rdf=\"http://www.w3.org/1999/02/22-rdf-syntax-ns#\">\n   <cc:Work>\n    <dc:type rdf:resource=\"http://purl.org/dc/dcmitype/StillImage\"/>\n    <dc:date>2026-07-19T13:51:06.822285</dc:date>\n    <dc:format>image/svg+xml</dc:format>\n    <dc:creator>\n     <cc:Agent>\n      <dc:title>Matplotlib v3.10.8, https://matplotlib.org/</dc:title>\n     </cc:Agent>\n    </dc:creator>\n   </cc:Work>\n  </rdf:RDF>\n </metadata>\n <defs>\n  <style type=\"text/css\">*{stroke-linejoin: round; stroke-linecap: butt}</style>\n </defs>\n <g id=\"figure_1\">\n  <g id=\"patch_1\">\n   <path d=\"M 0 183.35625 \nL 236.740625 183.35625 \nL 236.740625 0 \nL 0 0 \nz\n\" style=\"fill: #ffffff\"/>\n  </g>\n  <g id=\"axes_1\">\n   <g id=\"patch_2\">\n    <path d=\"M 34.240625 145.8 \nL 229.540625 145.8 \nL 229.540625 7.2 \nL 34.240625 7.2 \nz\n\" style=\"fill: #ffffff\"/>\n   </g>\n   <g id=\"matplotlib.axis_1\">\n    <g id=\"xtick_1\">\n     <g id=\"line2d_1\">\n      <path d=\"M 76.646801 145.8 \nL 76.646801 7.2 \n\" clip-path=\"url(#p192ec952d7)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_2\">\n      <defs>\n       <path id=\"m5bc33b46a9\" d=\"M 0 0 \nL 0 3.5 \n\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </defs>\n      <g>\n       <use xlink:href=\"#m5bc33b46a9\" x=\"76.646801\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_1\">\n      <!-- $\\mathdefault{10^{1}}$ -->\n      <g transform=\"translate(67.846801 160.398438) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-31\" d=\"M 794 531 \nL 1825 531 \nL 1825 4091 \nL 703 3866 \nL 703 4441 \nL 1819 4666 \nL 2450 4666 \nL 2450 531 \nL 3481 531 \nL 3481 0 \nL 794 0 \nL 794 531 \nz\n\" transform=\"scale(0.015625)\"/>\n        <path id=\"DejaVuSans-30\" d=\"M 2034 4250 \nQ 1547 4250 1301 3770 \nQ 1056 3291 1056 2328 \nQ 1056 1369 1301 889 \nQ 1547 409 2034 409 \nQ 2525 409 2770 889 \nQ 3016 1369 3016 2328 \nQ 3016 3291 2770 3770 \nQ 2525 4250 2034 4250 \nz\nM 2034 4750 \nQ 2819 4750 3233 4129 \nQ 3647 3509 3647 2328 \nQ 3647 1150 3233 529 \nQ 2819 -91 2034 -91 \nQ 1250 -91 836 529 \nQ 422 1150 422 2328 \nQ 422 3509 836 4129 \nQ 1250 4750 2034 4750 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-31\" transform=\"translate(0 0.684375)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0.684375)\"/>\n       <use xlink:href=\"#DejaVuSans-31\" transform=\"translate(128.203125 38.965625) scale(0.7)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_2\">\n     <g id=\"line2d_3\">\n      <path d=\"M 160.902977 145.8 \nL 160.902977 7.2 \n\" clip-path=\"url(#p192ec952d7)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_4\">\n      <g>\n       <use xlink:href=\"#m5bc33b46a9\" x=\"160.902977\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_2\">\n      <!-- $\\mathdefault{10^{2}}$ -->\n      <g transform=\"translate(152.102977 160.398438) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-32\" d=\"M 1228 531 \nL 3431 531 \nL 3431 0 \nL 469 0 \nL 469 531 \nQ 828 903 1448 1529 \nQ 2069 2156 2228 2338 \nQ 2531 2678 2651 2914 \nQ 2772 3150 2772 3378 \nQ 2772 3750 2511 3984 \nQ 2250 4219 1831 4219 \nQ 1534 4219 1204 4116 \nQ 875 4013 500 3803 \nL 500 4441 \nQ 881 4594 1212 4672 \nQ 1544 4750 1819 4750 \nQ 2544 4750 2975 4387 \nQ 3406 4025 3406 3419 \nQ 3406 3131 3298 2873 \nQ 3191 2616 2906 2266 \nQ 2828 2175 2409 1742 \nQ 1991 1309 1228 531 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-31\" transform=\"translate(0 0.765625)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0.765625)\"/>\n       <use xlink:href=\"#DejaVuSans-32\" transform=\"translate(128.203125 39.046875) scale(0.7)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_3\">\n     <g id=\"line2d_5\">\n      <defs>\n       <path id=\"m4e70513e8a\" d=\"M 0 0 \nL 0 2 \n\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </defs>\n      <g>\n       <use xlink:href=\"#m4e70513e8a\" x=\"43.117898\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_4\">\n     <g id=\"line2d_6\">\n      <g>\n       <use xlink:href=\"#m4e70513e8a\" x=\"51.283165\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_5\">\n     <g id=\"line2d_7\">\n      <g>\n       <use xlink:href=\"#m4e70513e8a\" x=\"57.954674\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_6\">\n     <g id=\"line2d_8\">\n      <g>\n       <use xlink:href=\"#m4e70513e8a\" x=\"63.595354\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_7\">\n     <g id=\"line2d_9\">\n      <g>\n       <use xlink:href=\"#m4e70513e8a\" x=\"68.481534\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_8\">\n     <g id=\"line2d_10\">\n      <g>\n       <use xlink:href=\"#m4e70513e8a\" x=\"72.79145\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_9\">\n     <g id=\"line2d_11\">\n      <g>\n       <use xlink:href=\"#m4e70513e8a\" x=\"102.010438\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_10\">\n     <g id=\"line2d_12\">\n      <g>\n       <use xlink:href=\"#m4e70513e8a\" x=\"116.847214\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_11\">\n     <g id=\"line2d_13\">\n      <g>\n       <use xlink:href=\"#m4e70513e8a\" x=\"127.374074\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_12\">\n     <g id=\"line2d_14\">\n      <g>\n       <use xlink:href=\"#m4e70513e8a\" x=\"135.539341\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_13\">\n     <g id=\"line2d_15\">\n      <g>\n       <use xlink:href=\"#m4e70513e8a\" x=\"142.21085\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_14\">\n     <g id=\"line2d_16\">\n      <g>\n       <use xlink:href=\"#m4e70513e8a\" x=\"147.851531\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_15\">\n     <g id=\"line2d_17\">\n      <g>\n       <use xlink:href=\"#m4e70513e8a\" x=\"152.73771\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_16\">\n     <g id=\"line2d_18\">\n      <g>\n       <use xlink:href=\"#m4e70513e8a\" x=\"157.047626\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_17\">\n     <g id=\"line2d_19\">\n      <g>\n       <use xlink:href=\"#m4e70513e8a\" x=\"186.266614\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_18\">\n     <g id=\"line2d_20\">\n      <g>\n       <use xlink:href=\"#m4e70513e8a\" x=\"201.10339\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_19\">\n     <g id=\"line2d_21\">\n      <g>\n       <use xlink:href=\"#m4e70513e8a\" x=\"211.63025\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_20\">\n     <g id=\"line2d_22\">\n      <g>\n       <use xlink:href=\"#m4e70513e8a\" x=\"219.795517\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_21\">\n     <g id=\"line2d_23\">\n      <g>\n       <use xlink:href=\"#m4e70513e8a\" x=\"226.467026\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_3\">\n     <!-- dimension d -->\n     <g transform=\"translate(101.226562 174.076563) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-64\" d=\"M 2906 2969 \nL 2906 4863 \nL 3481 4863 \nL 3481 0 \nL 2906 0 \nL 2906 525 \nQ 2725 213 2448 61 \nQ 2172 -91 1784 -91 \nQ 1150 -91 751 415 \nQ 353 922 353 1747 \nQ 353 2572 751 3078 \nQ 1150 3584 1784 3584 \nQ 2172 3584 2448 3432 \nQ 2725 3281 2906 2969 \nz\nM 947 1747 \nQ 947 1113 1208 752 \nQ 1469 391 1925 391 \nQ 2381 391 2643 752 \nQ 2906 1113 2906 1747 \nQ 2906 2381 2643 2742 \nQ 2381 3103 1925 3103 \nQ 1469 3103 1208 2742 \nQ 947 2381 947 1747 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-69\" d=\"M 603 3500 \nL 1178 3500 \nL 1178 0 \nL 603 0 \nL 603 3500 \nz\nM 603 4863 \nL 1178 4863 \nL 1178 4134 \nL 603 4134 \nL 603 4863 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-6d\" d=\"M 3328 2828 \nQ 3544 3216 3844 3400 \nQ 4144 3584 4550 3584 \nQ 5097 3584 5394 3201 \nQ 5691 2819 5691 2113 \nL 5691 0 \nL 5113 0 \nL 5113 2094 \nQ 5113 2597 4934 2840 \nQ 4756 3084 4391 3084 \nQ 3944 3084 3684 2787 \nQ 3425 2491 3425 1978 \nL 3425 0 \nL 2847 0 \nL 2847 2094 \nQ 2847 2600 2669 2842 \nQ 2491 3084 2119 3084 \nQ 1678 3084 1418 2786 \nQ 1159 2488 1159 1978 \nL 1159 0 \nL 581 0 \nL 581 3500 \nL 1159 3500 \nL 1159 2956 \nQ 1356 3278 1631 3431 \nQ 1906 3584 2284 3584 \nQ 2666 3584 2933 3390 \nQ 3200 3197 3328 2828 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-65\" d=\"M 3597 1894 \nL 3597 1613 \nL 953 1613 \nQ 991 1019 1311 708 \nQ 1631 397 2203 397 \nQ 2534 397 2845 478 \nQ 3156 559 3463 722 \nL 3463 178 \nQ 3153 47 2828 -22 \nQ 2503 -91 2169 -91 \nQ 1331 -91 842 396 \nQ 353 884 353 1716 \nQ 353 2575 817 3079 \nQ 1281 3584 2069 3584 \nQ 2775 3584 3186 3129 \nQ 3597 2675 3597 1894 \nz\nM 3022 2063 \nQ 3016 2534 2758 2815 \nQ 2500 3097 2075 3097 \nQ 1594 3097 1305 2825 \nQ 1016 2553 972 2059 \nL 3022 2063 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-6e\" d=\"M 3513 2113 \nL 3513 0 \nL 2938 0 \nL 2938 2094 \nQ 2938 2591 2744 2837 \nQ 2550 3084 2163 3084 \nQ 1697 3084 1428 2787 \nQ 1159 2491 1159 1978 \nL 1159 0 \nL 581 0 \nL 581 3500 \nL 1159 3500 \nL 1159 2956 \nQ 1366 3272 1645 3428 \nQ 1925 3584 2291 3584 \nQ 2894 3584 3203 3211 \nQ 3513 2838 3513 2113 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-73\" d=\"M 2834 3397 \nL 2834 2853 \nQ 2591 2978 2328 3040 \nQ 2066 3103 1784 3103 \nQ 1356 3103 1142 2972 \nQ 928 2841 928 2578 \nQ 928 2378 1081 2264 \nQ 1234 2150 1697 2047 \nL 1894 2003 \nQ 2506 1872 2764 1633 \nQ 3022 1394 3022 966 \nQ 3022 478 2636 193 \nQ 2250 -91 1575 -91 \nQ 1294 -91 989 -36 \nQ 684 19 347 128 \nL 347 722 \nQ 666 556 975 473 \nQ 1284 391 1588 391 \nQ 1994 391 2212 530 \nQ 2431 669 2431 922 \nQ 2431 1156 2273 1281 \nQ 2116 1406 1581 1522 \nL 1381 1569 \nQ 847 1681 609 1914 \nQ 372 2147 372 2553 \nQ 372 3047 722 3315 \nQ 1072 3584 1716 3584 \nQ 2034 3584 2315 3537 \nQ 2597 3491 2834 3397 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-6f\" d=\"M 1959 3097 \nQ 1497 3097 1228 2736 \nQ 959 2375 959 1747 \nQ 959 1119 1226 758 \nQ 1494 397 1959 397 \nQ 2419 397 2687 759 \nQ 2956 1122 2956 1747 \nQ 2956 2369 2687 2733 \nQ 2419 3097 1959 3097 \nz\nM 1959 3584 \nQ 2709 3584 3137 3096 \nQ 3566 2609 3566 1747 \nQ 3566 888 3137 398 \nQ 2709 -91 1959 -91 \nQ 1206 -91 779 398 \nQ 353 888 353 1747 \nQ 353 2609 779 3096 \nQ 1206 3584 1959 3584 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-20\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-64\"/>\n      <use xlink:href=\"#DejaVuSans-69\" transform=\"translate(63.476562 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6d\" transform=\"translate(91.259766 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(188.671875 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6e\" transform=\"translate(250.195312 0)\"/>\n      <use xlink:href=\"#DejaVuSans-73\" transform=\"translate(313.574219 0)\"/>\n      <use xlink:href=\"#DejaVuSans-69\" transform=\"translate(365.673828 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6f\" transform=\"translate(393.457031 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6e\" transform=\"translate(454.638672 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(518.017578 0)\"/>\n      <use xlink:href=\"#DejaVuSans-64\" transform=\"translate(549.804688 0)\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"matplotlib.axis_2\">\n    <g id=\"ytick_1\">\n     <g id=\"line2d_24\">\n      <path d=\"M 34.240625 110.709583 \nL 229.540625 110.709583 \n\" clip-path=\"url(#p192ec952d7)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_25\">\n      <defs>\n       <path id=\"m714b48ce25\" d=\"M 0 0 \nL -3.5 0 \n\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </defs>\n      <g>\n       <use xlink:href=\"#m714b48ce25\" x=\"34.240625\" y=\"110.709583\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_4\">\n      <!-- 1 -->\n      <g transform=\"translate(20.878125 114.508802) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-31\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_2\">\n     <g id=\"line2d_26\">\n      <path d=\"M 34.240625 74.791556 \nL 229.540625 74.791556 \n\" clip-path=\"url(#p192ec952d7)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_27\">\n      <g>\n       <use xlink:href=\"#m714b48ce25\" x=\"34.240625\" y=\"74.791556\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_5\">\n      <!-- 2 -->\n      <g transform=\"translate(20.878125 78.590774) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-32\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_3\">\n     <g id=\"line2d_28\">\n      <path d=\"M 34.240625 38.873528 \nL 229.540625 38.873528 \n\" clip-path=\"url(#p192ec952d7)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_29\">\n      <g>\n       <use xlink:href=\"#m714b48ce25\" x=\"34.240625\" y=\"38.873528\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_6\">\n      <!-- 3 -->\n      <g transform=\"translate(20.878125 42.672747) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-33\" d=\"M 2597 2516 \nQ 3050 2419 3304 2112 \nQ 3559 1806 3559 1356 \nQ 3559 666 3084 287 \nQ 2609 -91 1734 -91 \nQ 1441 -91 1130 -33 \nQ 819 25 488 141 \nL 488 750 \nQ 750 597 1062 519 \nQ 1375 441 1716 441 \nQ 2309 441 2620 675 \nQ 2931 909 2931 1356 \nQ 2931 1769 2642 2001 \nQ 2353 2234 1838 2234 \nL 1294 2234 \nL 1294 2753 \nL 1863 2753 \nQ 2328 2753 2575 2939 \nQ 2822 3125 2822 3475 \nQ 2822 3834 2567 4026 \nQ 2313 4219 1838 4219 \nQ 1578 4219 1281 4162 \nQ 984 4106 628 3988 \nL 628 4550 \nQ 988 4650 1302 4700 \nQ 1616 4750 1894 4750 \nQ 2613 4750 3031 4423 \nQ 3450 4097 3450 3541 \nQ 3450 3153 3228 2886 \nQ 3006 2619 2597 2516 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-33\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_7\">\n     <!-- entropy (nats) -->\n     <g transform=\"translate(14.798438 112.129688) rotate(-90) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-74\" d=\"M 1172 4494 \nL 1172 3500 \nL 2356 3500 \nL 2356 3053 \nL 1172 3053 \nL 1172 1153 \nQ 1172 725 1289 603 \nQ 1406 481 1766 481 \nL 2356 481 \nL 2356 0 \nL 1766 0 \nQ 1100 0 847 248 \nQ 594 497 594 1153 \nL 594 3053 \nL 172 3053 \nL 172 3500 \nL 594 3500 \nL 594 4494 \nL 1172 4494 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-72\" d=\"M 2631 2963 \nQ 2534 3019 2420 3045 \nQ 2306 3072 2169 3072 \nQ 1681 3072 1420 2755 \nQ 1159 2438 1159 1844 \nL 1159 0 \nL 581 0 \nL 581 3500 \nL 1159 3500 \nL 1159 2956 \nQ 1341 3275 1631 3429 \nQ 1922 3584 2338 3584 \nQ 2397 3584 2469 3576 \nQ 2541 3569 2628 3553 \nL 2631 2963 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-70\" d=\"M 1159 525 \nL 1159 -1331 \nL 581 -1331 \nL 581 3500 \nL 1159 3500 \nL 1159 2969 \nQ 1341 3281 1617 3432 \nQ 1894 3584 2278 3584 \nQ 2916 3584 3314 3078 \nQ 3713 2572 3713 1747 \nQ 3713 922 3314 415 \nQ 2916 -91 2278 -91 \nQ 1894 -91 1617 61 \nQ 1341 213 1159 525 \nz\nM 3116 1747 \nQ 3116 2381 2855 2742 \nQ 2594 3103 2138 3103 \nQ 1681 3103 1420 2742 \nQ 1159 2381 1159 1747 \nQ 1159 1113 1420 752 \nQ 1681 391 2138 391 \nQ 2594 391 2855 752 \nQ 3116 1113 3116 1747 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-79\" d=\"M 2059 -325 \nQ 1816 -950 1584 -1140 \nQ 1353 -1331 966 -1331 \nL 506 -1331 \nL 506 -850 \nL 844 -850 \nQ 1081 -850 1212 -737 \nQ 1344 -625 1503 -206 \nL 1606 56 \nL 191 3500 \nL 800 3500 \nL 1894 763 \nL 2988 3500 \nL 3597 3500 \nL 2059 -325 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-28\" d=\"M 1984 4856 \nQ 1566 4138 1362 3434 \nQ 1159 2731 1159 2009 \nQ 1159 1288 1364 580 \nQ 1569 -128 1984 -844 \nL 1484 -844 \nQ 1016 -109 783 600 \nQ 550 1309 550 2009 \nQ 550 2706 781 3412 \nQ 1013 4119 1484 4856 \nL 1984 4856 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-61\" d=\"M 2194 1759 \nQ 1497 1759 1228 1600 \nQ 959 1441 959 1056 \nQ 959 750 1161 570 \nQ 1363 391 1709 391 \nQ 2188 391 2477 730 \nQ 2766 1069 2766 1631 \nL 2766 1759 \nL 2194 1759 \nz\nM 3341 1997 \nL 3341 0 \nL 2766 0 \nL 2766 531 \nQ 2569 213 2275 61 \nQ 1981 -91 1556 -91 \nQ 1019 -91 701 211 \nQ 384 513 384 1019 \nQ 384 1609 779 1909 \nQ 1175 2209 1959 2209 \nL 2766 2209 \nL 2766 2266 \nQ 2766 2663 2505 2880 \nQ 2244 3097 1772 3097 \nQ 1472 3097 1187 3025 \nQ 903 2953 641 2809 \nL 641 3341 \nQ 956 3463 1253 3523 \nQ 1550 3584 1831 3584 \nQ 2591 3584 2966 3190 \nQ 3341 2797 3341 1997 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-29\" d=\"M 513 4856 \nL 1013 4856 \nQ 1481 4119 1714 3412 \nQ 1947 2706 1947 2009 \nQ 1947 1309 1714 600 \nQ 1481 -109 1013 -844 \nL 513 -844 \nQ 928 -128 1133 580 \nQ 1338 1288 1338 2009 \nQ 1338 2731 1133 3434 \nQ 928 4138 513 4856 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-65\"/>\n      <use xlink:href=\"#DejaVuSans-6e\" transform=\"translate(61.523438 0)\"/>\n      <use xlink:href=\"#DejaVuSans-74\" transform=\"translate(124.902344 0)\"/>\n      <use xlink:href=\"#DejaVuSans-72\" transform=\"translate(164.111328 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6f\" transform=\"translate(202.974609 0)\"/>\n      <use xlink:href=\"#DejaVuSans-70\" transform=\"translate(264.15625 0)\"/>\n      <use xlink:href=\"#DejaVuSans-79\" transform=\"translate(327.632812 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(386.8125 0)\"/>\n      <use xlink:href=\"#DejaVuSans-28\" transform=\"translate(418.599609 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6e\" transform=\"translate(457.613281 0)\"/>\n      <use xlink:href=\"#DejaVuSans-61\" transform=\"translate(520.992188 0)\"/>\n      <use xlink:href=\"#DejaVuSans-74\" transform=\"translate(582.271484 0)\"/>\n      <use xlink:href=\"#DejaVuSans-73\" transform=\"translate(621.480469 0)\"/>\n      <use xlink:href=\"#DejaVuSans-29\" transform=\"translate(673.580078 0)\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"line2d_30\">\n    <path d=\"M 43.117898 44.608444 \nL 68.481534 73.419019 \nL 93.84517 96.657417 \nL 119.208807 114.574563 \nL 144.572443 124.782397 \nL 169.93608 133.178864 \nL 195.299716 137.863579 \nL 220.663352 139.5 \n\" clip-path=\"url(#p192ec952d7)\" style=\"fill: none; stroke: #1f77b4; stroke-width: 1.5; stroke-linecap: square\"/>\n   </g>\n   <g id=\"line2d_31\">\n    <path d=\"M 43.117898 13.5 \nL 68.481534 14.490432 \nL 93.84517 14.806588 \nL 119.208807 14.689893 \nL 144.572443 14.523572 \nL 169.93608 14.361987 \nL 195.299716 14.841219 \nL 220.663352 14.720285 \n\" clip-path=\"url(#p192ec952d7)\" style=\"fill: none; stroke-dasharray: 5.55,2.4; stroke-dashoffset: 0; stroke: #bf00bf; stroke-width: 1.5\"/>\n   </g>\n   <g id=\"patch_3\">\n    <path d=\"M 34.240625 145.8 \nL 34.240625 7.2 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_4\">\n    <path d=\"M 229.540625 145.8 \nL 229.540625 7.2 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_5\">\n    <path d=\"M 34.240625 145.8 \nL 229.540625 145.8 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_6\">\n    <path d=\"M 34.240625 7.2 \nL 229.540625 7.2 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"legend_1\">\n    <g id=\"patch_7\">\n     <path d=\"M 96.329688 92.678125 \nL 222.540625 92.678125 \nQ 224.540625 92.678125 224.540625 90.678125 \nL 224.540625 62.321875 \nQ 224.540625 60.321875 222.540625 60.321875 \nL 96.329688 60.321875 \nQ 94.329688 60.321875 94.329688 62.321875 \nL 94.329688 90.678125 \nQ 94.329688 92.678125 96.329688 92.678125 \nz\n\" style=\"fill: #ffffff; opacity: 0.8; stroke: #cccccc; stroke-linejoin: miter\"/>\n    </g>\n    <g id=\"line2d_32\">\n     <path d=\"M 98.329688 68.420313 \nL 108.329688 68.420313 \nL 118.329688 68.420313 \n\" style=\"fill: none; stroke: #1f77b4; stroke-width: 1.5; stroke-linecap: square\"/>\n    </g>\n    <g id=\"text_8\">\n     <!-- unscaled -->\n     <g transform=\"translate(126.329688 71.920313) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-75\" d=\"M 544 1381 \nL 544 3500 \nL 1119 3500 \nL 1119 1403 \nQ 1119 906 1312 657 \nQ 1506 409 1894 409 \nQ 2359 409 2629 706 \nQ 2900 1003 2900 1516 \nL 2900 3500 \nL 3475 3500 \nL 3475 0 \nL 2900 0 \nL 2900 538 \nQ 2691 219 2414 64 \nQ 2138 -91 1772 -91 \nQ 1169 -91 856 284 \nQ 544 659 544 1381 \nz\nM 1991 3584 \nL 1991 3584 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-63\" d=\"M 3122 3366 \nL 3122 2828 \nQ 2878 2963 2633 3030 \nQ 2388 3097 2138 3097 \nQ 1578 3097 1268 2742 \nQ 959 2388 959 1747 \nQ 959 1106 1268 751 \nQ 1578 397 2138 397 \nQ 2388 397 2633 464 \nQ 2878 531 3122 666 \nL 3122 134 \nQ 2881 22 2623 -34 \nQ 2366 -91 2075 -91 \nQ 1284 -91 818 406 \nQ 353 903 353 1747 \nQ 353 2603 823 3093 \nQ 1294 3584 2113 3584 \nQ 2378 3584 2631 3529 \nQ 2884 3475 3122 3366 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-6c\" d=\"M 603 4863 \nL 1178 4863 \nL 1178 0 \nL 603 0 \nL 603 4863 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-75\"/>\n      <use xlink:href=\"#DejaVuSans-6e\" transform=\"translate(63.378906 0)\"/>\n      <use xlink:href=\"#DejaVuSans-73\" transform=\"translate(126.757812 0)\"/>\n      <use xlink:href=\"#DejaVuSans-63\" transform=\"translate(178.857422 0)\"/>\n      <use xlink:href=\"#DejaVuSans-61\" transform=\"translate(233.837891 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6c\" transform=\"translate(295.117188 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(322.900391 0)\"/>\n      <use xlink:href=\"#DejaVuSans-64\" transform=\"translate(384.423828 0)\"/>\n     </g>\n    </g>\n    <g id=\"line2d_33\">\n     <path d=\"M 98.329688 83.098438 \nL 108.329688 83.098438 \nL 118.329688 83.098438 \n\" style=\"fill: none; stroke-dasharray: 5.55,2.4; stroke-dashoffset: 0; stroke: #bf00bf; stroke-width: 1.5\"/>\n    </g>\n    <g id=\"text_9\">\n     <!-- scaled by 1/sqrt(d) -->\n     <g transform=\"translate(126.329688 86.598438) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-62\" d=\"M 3116 1747 \nQ 3116 2381 2855 2742 \nQ 2594 3103 2138 3103 \nQ 1681 3103 1420 2742 \nQ 1159 2381 1159 1747 \nQ 1159 1113 1420 752 \nQ 1681 391 2138 391 \nQ 2594 391 2855 752 \nQ 3116 1113 3116 1747 \nz\nM 1159 2969 \nQ 1341 3281 1617 3432 \nQ 1894 3584 2278 3584 \nQ 2916 3584 3314 3078 \nQ 3713 2572 3713 1747 \nQ 3713 922 3314 415 \nQ 2916 -91 2278 -91 \nQ 1894 -91 1617 61 \nQ 1341 213 1159 525 \nL 1159 0 \nL 581 0 \nL 581 4863 \nL 1159 4863 \nL 1159 2969 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-2f\" d=\"M 1625 4666 \nL 2156 4666 \nL 531 -594 \nL 0 -594 \nL 1625 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-71\" d=\"M 947 1747 \nQ 947 1113 1208 752 \nQ 1469 391 1925 391 \nQ 2381 391 2643 752 \nQ 2906 1113 2906 1747 \nQ 2906 2381 2643 2742 \nQ 2381 3103 1925 3103 \nQ 1469 3103 1208 2742 \nQ 947 2381 947 1747 \nz\nM 2906 525 \nQ 2725 213 2448 61 \nQ 2172 -91 1784 -91 \nQ 1150 -91 751 415 \nQ 353 922 353 1747 \nQ 353 2572 751 3078 \nQ 1150 3584 1784 3584 \nQ 2172 3584 2448 3432 \nQ 2725 3281 2906 2969 \nL 2906 3500 \nL 3481 3500 \nL 3481 -1331 \nL 2906 -1331 \nL 2906 525 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-73\"/>\n      <use xlink:href=\"#DejaVuSans-63\" transform=\"translate(52.099609 0)\"/>\n      <use xlink:href=\"#DejaVuSans-61\" transform=\"translate(107.080078 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6c\" transform=\"translate(168.359375 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(196.142578 0)\"/>\n      <use xlink:href=\"#DejaVuSans-64\" transform=\"translate(257.666016 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(321.142578 0)\"/>\n      <use xlink:href=\"#DejaVuSans-62\" transform=\"translate(352.929688 0)\"/>\n      <use xlink:href=\"#DejaVuSans-79\" transform=\"translate(416.40625 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(475.585938 0)\"/>\n      <use xlink:href=\"#DejaVuSans-31\" transform=\"translate(507.373047 0)\"/>\n      <use xlink:href=\"#DejaVuSans-2f\" transform=\"translate(570.996094 0)\"/>\n      <use xlink:href=\"#DejaVuSans-73\" transform=\"translate(604.6875 0)\"/>\n      <use xlink:href=\"#DejaVuSans-71\" transform=\"translate(656.787109 0)\"/>\n      <use xlink:href=\"#DejaVuSans-72\" transform=\"translate(720.263672 0)\"/>\n      <use xlink:href=\"#DejaVuSans-74\" transform=\"translate(761.376953 0)\"/>\n      <use xlink:href=\"#DejaVuSans-28\" transform=\"translate(800.585938 0)\"/>\n      <use xlink:href=\"#DejaVuSans-64\" transform=\"translate(839.599609 0)\"/>\n      <use xlink:href=\"#DejaVuSans-29\" transform=\"translate(903.076172 0)\"/>\n     </g>\n    </g>\n   </g>\n  </g>\n </g>\n <defs>\n  <clipPath id=\"p192ec952d7\">\n   <rect x=\"34.240625\" y=\"7.2\" width=\"195.3\" height=\"138.6\"/>\n  </clipPath>\n </defs>\n</svg>\n"
     },
     "metadata": {}
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=attention-scoring-md-3 fw=all -->\n",
    "\n",
    "A uniform distribution over $64$ keys has entropy $\\ln 64 \\approx 4.2$ nats.\n",
    "The scaled scores hold their entropy essentially constant, at about $3.7$\n",
    "nats, across two orders of magnitude in $d$: the weight distribution keeps\n",
    "the same moderate sharpness no matter how wide the vectors are. Without\n",
    "scaling, the entropy collapses as $d$ grows—below $0.2$ nats by $d = 512$,\n",
    "meaning the average query puts a weight of more than $0.9$ on a single key\n",
    "purely because the vectors are long. The gradient tells the same story:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "attention-scoring-softmax-saturation-and-the-1-sqrt-d-factor-2"
   },
   "source": [
    "for d in (4, 64, 512):\n",
    "    _, jac_unscaled = attention_stats(d, scaled=False)\n",
    "    _, jac_scaled = attention_stats(d, scaled=True)\n",
    "    print(f'd = {d:3d}: Jacobian norm {jac_unscaled:.3f} (unscaled), '\n",
    "          f'{jac_scaled:.3f} (scaled)')"
   ],
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "d =   4: Jacobian norm 0.258 (unscaled), 0.178 (scaled)\nd =  64: Jacobian norm 0.218 (unscaled), 0.184 (scaled)\nd = 512: Jacobian norm 0.103 (unscaled), 0.186 (scaled)\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=attention-scoring-md-4 fw=all -->\n",
    "\n",
    "With scaling, the Jacobian norm is the same at every dimension. Without it,\n",
    "the norm decays as the softmax saturates—by $d = 512$ it has fallen to\n",
    "about half the scaled value and is still shrinking. A model built on\n",
    "unscaled scores at realistic dimensions would start its life with\n",
    "near-one-hot attention and barely any gradient with which to fix it. One\n",
    "division by $\\sqrt{d}$ removes the problem, which is why it is part of the\n",
    "definition the equation rather than a tuning trick.\n",
    "\n",
    "## Masking\n",
    "\n",
    "Attention as defined so far attends to *every* key. In practice we routinely\n",
    "need it not to, for one of two reasons. The first is *padding*: sequences of\n",
    "different lengths end up in the same minibatch, padded with dummy tokens to\n",
    "a common length (here shown as `<blank>`):\n",
    "\n",
    "```\n",
    "Dive  into  Deep    Learning\n",
    "Learn to    code    <blank>\n",
    "Hello world <blank> <blank>\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=attention-scoring-md-5 fw=all -->\n",
    "\n",
    "Padding tokens carry no meaning, and no query should waste weight on them.\n",
    "The second reason is *causality*: a language model trained to predict the\n",
    "next token computes outputs for all positions of a sequence in parallel, and\n",
    "the query at position $t$ must not see keys at positions beyond $t$—\n",
    "otherwise the model can copy the very future it is being trained to predict.\n",
    "Both cases call for the same operation: restrict the attention sum\n",
    "$\\sum_{i=1}^n \\alpha(\\mathbf{q}, \\mathbf{k}_i) \\mathbf{v}_i$ to a valid\n",
    "prefix $\\sum_{i=1}^l \\alpha(\\mathbf{q}, \\mathbf{k}_i) \\mathbf{v}_i$ with\n",
    "$l \\leq n$, where $l$ depends on the sequence (padding) or on the query\n",
    "position (causality).\n",
    "\n",
    "### The Masked Softmax Operation\n",
    "\n",
    "The operation is common enough to have a name, the *masked softmax*, and a\n",
    "standard implementation trick. Rather than branching on which keys are\n",
    "valid—conditionals are poison for the heavily optimized batched kernels\n",
    "that attention runs on—we overwrite the scores of the invalid positions\n",
    "with a very negative number before the softmax, so that their weights come\n",
    "out as zero.\n",
    "\n",
    "The choice of \"very negative\" deserves care. Older codebases used a literal\n",
    "constant such as $-10^{6}$. In single precision that works; in the half\n",
    "precisions that modern training runs in, it does not. The float16 format\n",
    "tops out near $6.5 \\times 10^4$, so $-10^6$ silently overflows, and in\n",
    "bfloat16 a merely-large constant may fail to fully suppress a weight once\n",
    "genuine scores are large themselves. Writing literal $-\\infty$ masks\n",
    "exactly, but if every key of some query is masked the softmax returns NaN\n",
    "and poisons the training run. The dtype-safe idiom, which we adopt, masks\n",
    "with the most negative *finite* value of the score's dtype\n",
    "(`torch.finfo(X.dtype).min` and `jnp.finfo(X.dtype).min`, respectively): the\n",
    "masked weights are exactly zero at any precision, and a fully masked query\n",
    "degrades to a uniform distribution instead of NaN — though a uniform average\n",
    "over *invalid* values is still garbage, so callers must guarantee that every\n",
    "query keeps at least one valid key."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "attention-scoring-the-masked-softmax-operation-1"
   },
   "source": [
    "def masked_softmax(X, valid_lens):\n",
    "    \"\"\"Perform softmax operation by masking elements on the last axis.\"\"\"\n",
    "    # X: 3D tensor, valid_lens: 1D or 2D tensor\n",
    "    if valid_lens is None:\n",
    "        return jax.nn.softmax(X, axis=-1)\n",
    "    shape = X.shape\n",
    "    if valid_lens.ndim == 1:\n",
    "        valid_lens = jnp.repeat(valid_lens, shape[1])\n",
    "    else:\n",
    "        valid_lens = valid_lens.reshape(-1)\n",
    "    mask = jnp.arange(shape[-1])[None, :] < valid_lens[:, None]\n",
    "    # Most negative finite score: exactly zero weight after the softmax,\n",
    "    # at any precision, without the NaN risk of literal -inf\n",
    "    X = jnp.where(mask, X.reshape(-1, shape[-1]), jnp.finfo(X.dtype).min)\n",
    "    return jax.nn.softmax(X.reshape(shape), axis=-1)"
   ],
   "execution_count": 4,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=attention-scoring-md-6 fw=all -->\n",
    "\n",
    "To illustrate how this function works, consider a minibatch of two examples\n",
    "with two queries and four keys each, where the valid lengths are $2$ and\n",
    "$3$, respectively. All weights beyond the valid length come out as zero,\n",
    "and each row still sums to $1$:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "attention-scoring-the-masked-softmax-operation-2"
   },
   "source": [
    "masked_softmax(jax.random.uniform(jax.random.key(0), (2, 2, 4)),\n",
    "               jnp.array([2, 3]))"
   ],
   "execution_count": 5,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array([[[0.4922724 , 0.5077276 , 0.        , 0.        ],\n        [0.5997409 , 0.40025908, 0.        , 0.        ]],\n\n       [[0.3531402 , 0.19856434, 0.4482954 , 0.        ],\n        [0.31042197, 0.2873157 , 0.40226236, 0.        ]]], dtype=float32)"
     },
     "metadata": {},
     "execution_count": 5
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=attention-scoring-md-7 fw=all -->\n",
    "\n",
    "For finer control we can pass a two-dimensional tensor of valid lengths,\n",
    "one per query:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "attention-scoring-the-masked-softmax-operation-3"
   },
   "source": [
    "masked_softmax(jax.random.uniform(jax.random.key(1), (2, 2, 4)),\n",
    "               jnp.array([[1, 3], [2, 4]]))"
   ],
   "execution_count": 6,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array([[[1.        , 0.        , 0.        , 0.        ],\n        [0.35672572, 0.2543091 , 0.38896522, 0.        ]],\n\n       [[0.43679553, 0.5632045 , 0.        , 0.        ],\n        [0.2580636 , 0.2509437 , 0.258296  , 0.23269679]]], dtype=float32)"
     },
     "metadata": {},
     "execution_count": 6
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=attention-scoring-md-8 fw=all -->\n",
    "\n",
    "### Causal Masking\n",
    "\n",
    "Per-query valid lengths are exactly what causality needs: for a sequence of\n",
    "length $n$, the query at position $t$ may attend to keys $1, \\ldots, t$, so\n",
    "the valid lengths are simply $(1, 2, \\ldots, n)$, shared by every sequence\n",
    "in the batch. The resulting attention pattern is lower triangular:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "attention-scoring-causal-masking"
   },
   "source": [
    "scores = jax.random.normal(jax.random.key(2), (1, 6, 6))\n",
    "causal_lens = jnp.arange(1, 7)[None, :]  # query t sees keys 1..t\n",
    "d2l.show_heatmaps(masked_softmax(scores, causal_lens)[None],\n",
    "                  xlabel='Keys', ylabel='Queries')"
   ],
   "execution_count": 7,
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "image/svg+xml": "<?xml version=\"1.0\" encoding=\"utf-8\" standalone=\"no\"?>\n<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n  \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n<svg xmlns:xlink=\"http://www.w3.org/1999/xlink\" width=\"193.43925pt\" height=\"156.35625pt\" viewBox=\"0 0 193.43925 156.35625\" xmlns=\"http://www.w3.org/2000/svg\" version=\"1.1\">\n <metadata>\n  <rdf:RDF xmlns:dc=\"http://purl.org/dc/elements/1.1/\" xmlns:cc=\"http://creativecommons.org/ns#\" xmlns:rdf=\"http://www.w3.org/1999/02/22-rdf-syntax-ns#\">\n   <cc:Work>\n    <dc:type rdf:resource=\"http://purl.org/dc/dcmitype/StillImage\"/>\n    <dc:date>2026-07-19T13:51:09.141918</dc:date>\n    <dc:format>image/svg+xml</dc:format>\n    <dc:creator>\n     <cc:Agent>\n      <dc:title>Matplotlib v3.10.8, https://matplotlib.org/</dc:title>\n     </cc:Agent>\n    </dc:creator>\n   </cc:Work>\n  </rdf:RDF>\n </metadata>\n <defs>\n  <style type=\"text/css\">*{stroke-linejoin: round; stroke-linecap: butt}</style>\n </defs>\n <g id=\"figure_1\">\n  <g id=\"patch_1\">\n   <path d=\"M -0 156.35625 \nL 193.43925 156.35625 \nL 193.43925 0 \nL -0 0 \nz\n\" style=\"fill: #ffffff\"/>\n  </g>\n  <g id=\"axes_1\">\n   <g id=\"patch_2\">\n    <path d=\"M 34.240625 118.8 \nL 145.840625 118.8 \nL 145.840625 7.2 \nL 34.240625 7.2 \nz\n\" style=\"fill: #ffffff\"/>\n   </g>\n   <g clip-path=\"url(#p2d9181945e)\">\n    <image xlink:href=\"data:image/png;base64,\niVBORw0KGgoAAAANSUhEUgAAAJsAAACbCAYAAAB1YemMAAACUklEQVR4nO3dMWpVURSG0X1fUkTUShC0MKiFaawEwcJOsBEnYeF4hAjaZgSxE1unIGInERGtBEFs5D1HcT+CrDWBv/nY3eEsvx7f203g4ouXxcwsV66tP3Jwef2NmZndtpk5+5DsbJIVGLEREhsZsZERGxmxkREbGbGRERsZsZERGxmxkREbGbGRERsZsZERGxmxkREbGbGRERsZsZERGxmxkREbmeXv6XHyIn4OLiQzmwdPV9/Yvj1ZfWNmZrn/KNnZXD1sdpIVGLEREhsZsZERGxmxkREbGbGRERsZsZERGxmxkREbGbGRERsZsZERGxmxkREbGbGRERsZsZERGxmxkREbmf25fisZ2tw4Sna2b16vvrH7/m31jZmZvSfPkp2Ky0ZGbGTERkZsZMRGRmxkxEZGbGTERkZsZMRGRmxkxEZGbGTERkZsZMRGRmxkxEZGbGTERkZsZMRGRmxkxEZm+Xj7ZvJH/NH7d8XMbL98Wn1j7+7D1Tf+Ry4bGbGRERsZsZERGxmxkREbGbGRERsZsZERGxmxkREbGbGRERsZsZERGxmxkREbGbGRERsZsZERGxmxkdn//PNPMnTnx1my4wHx+eWykREbGbGRERsZsZERGxmxkREbGbGRERsZsZERGxmxkREbGbGRERsZsZERGxmxkREbGbGRERsZsZERGxmxkVmez6Xk2+5Xv78WM5xjLhsZsZERGxmxkREbGbGRERsZsZERGxmxkREbGbGRERsZsZERGxmxkREbGbGRERsZsZERGxmxkREbGbGRERuZf52OJt1Fd/xHAAAAAElFTkSuQmCC\" id=\"imagefd5fc5fab4\" transform=\"scale(1 -1) translate(0 -111.6)\" x=\"34.240625\" y=\"-7.2\" width=\"111.6\" height=\"111.6\"/>\n   </g>\n   <g id=\"matplotlib.axis_1\">\n    <g id=\"xtick_1\">\n     <g id=\"line2d_1\">\n      <defs>\n       <path id=\"md296e2af64\" d=\"M 0 0 \nL 0 3.5 \n\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </defs>\n      <g>\n       <use xlink:href=\"#md296e2af64\" x=\"43.540625\" y=\"118.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_1\">\n      <!-- 0 -->\n      <g transform=\"translate(40.359375 133.398438) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-30\" d=\"M 2034 4250 \nQ 1547 4250 1301 3770 \nQ 1056 3291 1056 2328 \nQ 1056 1369 1301 889 \nQ 1547 409 2034 409 \nQ 2525 409 2770 889 \nQ 3016 1369 3016 2328 \nQ 3016 3291 2770 3770 \nQ 2525 4250 2034 4250 \nz\nM 2034 4750 \nQ 2819 4750 3233 4129 \nQ 3647 3509 3647 2328 \nQ 3647 1150 3233 529 \nQ 2819 -91 2034 -91 \nQ 1250 -91 836 529 \nQ 422 1150 422 2328 \nQ 422 3509 836 4129 \nQ 1250 4750 2034 4750 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-30\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_2\">\n     <g id=\"line2d_2\">\n      <g>\n       <use xlink:href=\"#md296e2af64\" x=\"80.740625\" y=\"118.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_2\">\n      <!-- 2 -->\n      <g transform=\"translate(77.559375 133.398438) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-32\" d=\"M 1228 531 \nL 3431 531 \nL 3431 0 \nL 469 0 \nL 469 531 \nQ 828 903 1448 1529 \nQ 2069 2156 2228 2338 \nQ 2531 2678 2651 2914 \nQ 2772 3150 2772 3378 \nQ 2772 3750 2511 3984 \nQ 2250 4219 1831 4219 \nQ 1534 4219 1204 4116 \nQ 875 4013 500 3803 \nL 500 4441 \nQ 881 4594 1212 4672 \nQ 1544 4750 1819 4750 \nQ 2544 4750 2975 4387 \nQ 3406 4025 3406 3419 \nQ 3406 3131 3298 2873 \nQ 3191 2616 2906 2266 \nQ 2828 2175 2409 1742 \nQ 1991 1309 1228 531 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-32\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_3\">\n     <g id=\"line2d_3\">\n      <g>\n       <use xlink:href=\"#md296e2af64\" x=\"117.940625\" y=\"118.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_3\">\n      <!-- 4 -->\n      <g transform=\"translate(114.759375 133.398438) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-34\" d=\"M 2419 4116 \nL 825 1625 \nL 2419 1625 \nL 2419 4116 \nz\nM 2253 4666 \nL 3047 4666 \nL 3047 1625 \nL 3713 1625 \nL 3713 1100 \nL 3047 1100 \nL 3047 0 \nL 2419 0 \nL 2419 1100 \nL 313 1100 \nL 313 1709 \nL 2253 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-34\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_4\">\n     <!-- Keys -->\n     <g transform=\"translate(78.371094 147.076562) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-4b\" d=\"M 628 4666 \nL 1259 4666 \nL 1259 2694 \nL 3353 4666 \nL 4166 4666 \nL 1850 2491 \nL 4331 0 \nL 3500 0 \nL 1259 2247 \nL 1259 0 \nL 628 0 \nL 628 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-65\" d=\"M 3597 1894 \nL 3597 1613 \nL 953 1613 \nQ 991 1019 1311 708 \nQ 1631 397 2203 397 \nQ 2534 397 2845 478 \nQ 3156 559 3463 722 \nL 3463 178 \nQ 3153 47 2828 -22 \nQ 2503 -91 2169 -91 \nQ 1331 -91 842 396 \nQ 353 884 353 1716 \nQ 353 2575 817 3079 \nQ 1281 3584 2069 3584 \nQ 2775 3584 3186 3129 \nQ 3597 2675 3597 1894 \nz\nM 3022 2063 \nQ 3016 2534 2758 2815 \nQ 2500 3097 2075 3097 \nQ 1594 3097 1305 2825 \nQ 1016 2553 972 2059 \nL 3022 2063 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-79\" d=\"M 2059 -325 \nQ 1816 -950 1584 -1140 \nQ 1353 -1331 966 -1331 \nL 506 -1331 \nL 506 -850 \nL 844 -850 \nQ 1081 -850 1212 -737 \nQ 1344 -625 1503 -206 \nL 1606 56 \nL 191 3500 \nL 800 3500 \nL 1894 763 \nL 2988 3500 \nL 3597 3500 \nL 2059 -325 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-73\" d=\"M 2834 3397 \nL 2834 2853 \nQ 2591 2978 2328 3040 \nQ 2066 3103 1784 3103 \nQ 1356 3103 1142 2972 \nQ 928 2841 928 2578 \nQ 928 2378 1081 2264 \nQ 1234 2150 1697 2047 \nL 1894 2003 \nQ 2506 1872 2764 1633 \nQ 3022 1394 3022 966 \nQ 3022 478 2636 193 \nQ 2250 -91 1575 -91 \nQ 1294 -91 989 -36 \nQ 684 19 347 128 \nL 347 722 \nQ 666 556 975 473 \nQ 1284 391 1588 391 \nQ 1994 391 2212 530 \nQ 2431 669 2431 922 \nQ 2431 1156 2273 1281 \nQ 2116 1406 1581 1522 \nL 1381 1569 \nQ 847 1681 609 1914 \nQ 372 2147 372 2553 \nQ 372 3047 722 3315 \nQ 1072 3584 1716 3584 \nQ 2034 3584 2315 3537 \nQ 2597 3491 2834 3397 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-4b\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(60.576172 0)\"/>\n      <use xlink:href=\"#DejaVuSans-79\" transform=\"translate(122.099609 0)\"/>\n      <use xlink:href=\"#DejaVuSans-73\" transform=\"translate(181.279297 0)\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"matplotlib.axis_2\">\n    <g id=\"ytick_1\">\n     <g id=\"line2d_4\">\n      <defs>\n       <path id=\"m3932ae2917\" d=\"M 0 0 \nL -3.5 0 \n\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </defs>\n      <g>\n       <use xlink:href=\"#m3932ae2917\" x=\"34.240625\" y=\"16.5\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_5\">\n      <!-- 0 -->\n      <g transform=\"translate(20.878125 20.299219) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-30\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_2\">\n     <g id=\"line2d_5\">\n      <g>\n       <use xlink:href=\"#m3932ae2917\" x=\"34.240625\" y=\"53.7\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_6\">\n      <!-- 2 -->\n      <g transform=\"translate(20.878125 57.499219) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-32\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_3\">\n     <g id=\"line2d_6\">\n      <g>\n       <use xlink:href=\"#m3932ae2917\" x=\"34.240625\" y=\"90.9\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_7\">\n      <!-- 4 -->\n      <g transform=\"translate(20.878125 94.699219) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-34\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_8\">\n     <!-- Queries -->\n     <g transform=\"translate(14.798437 82.307031) rotate(-90) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-51\" d=\"M 2522 4238 \nQ 1834 4238 1429 3725 \nQ 1025 3213 1025 2328 \nQ 1025 1447 1429 934 \nQ 1834 422 2522 422 \nQ 3209 422 3611 934 \nQ 4013 1447 4013 2328 \nQ 4013 3213 3611 3725 \nQ 3209 4238 2522 4238 \nz\nM 3406 84 \nL 4238 -825 \nL 3475 -825 \nL 2784 -78 \nQ 2681 -84 2626 -87 \nQ 2572 -91 2522 -91 \nQ 1538 -91 948 567 \nQ 359 1225 359 2328 \nQ 359 3434 948 4092 \nQ 1538 4750 2522 4750 \nQ 3503 4750 4090 4092 \nQ 4678 3434 4678 2328 \nQ 4678 1516 4351 937 \nQ 4025 359 3406 84 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-75\" d=\"M 544 1381 \nL 544 3500 \nL 1119 3500 \nL 1119 1403 \nQ 1119 906 1312 657 \nQ 1506 409 1894 409 \nQ 2359 409 2629 706 \nQ 2900 1003 2900 1516 \nL 2900 3500 \nL 3475 3500 \nL 3475 0 \nL 2900 0 \nL 2900 538 \nQ 2691 219 2414 64 \nQ 2138 -91 1772 -91 \nQ 1169 -91 856 284 \nQ 544 659 544 1381 \nz\nM 1991 3584 \nL 1991 3584 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-72\" d=\"M 2631 2963 \nQ 2534 3019 2420 3045 \nQ 2306 3072 2169 3072 \nQ 1681 3072 1420 2755 \nQ 1159 2438 1159 1844 \nL 1159 0 \nL 581 0 \nL 581 3500 \nL 1159 3500 \nL 1159 2956 \nQ 1341 3275 1631 3429 \nQ 1922 3584 2338 3584 \nQ 2397 3584 2469 3576 \nQ 2541 3569 2628 3553 \nL 2631 2963 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-69\" d=\"M 603 3500 \nL 1178 3500 \nL 1178 0 \nL 603 0 \nL 603 3500 \nz\nM 603 4863 \nL 1178 4863 \nL 1178 4134 \nL 603 4134 \nL 603 4863 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-51\"/>\n      <use xlink:href=\"#DejaVuSans-75\" transform=\"translate(78.710938 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(142.089844 0)\"/>\n      <use xlink:href=\"#DejaVuSans-72\" transform=\"translate(203.613281 0)\"/>\n      <use xlink:href=\"#DejaVuSans-69\" transform=\"translate(244.726562 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(272.509766 0)\"/>\n      <use xlink:href=\"#DejaVuSans-73\" transform=\"translate(334.033203 0)\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"patch_3\">\n    <path d=\"M 34.240625 118.8 \nL 34.240625 7.2 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_4\">\n    <path d=\"M 145.840625 118.8 \nL 145.840625 7.2 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_5\">\n    <path d=\"M 34.240625 118.8 \nL 145.840625 118.8 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_6\">\n    <path d=\"M 34.240625 7.2 \nL 145.840625 7.2 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n  </g>\n  <g id=\"axes_2\">\n   <g id=\"patch_7\">\n    <path d=\"M 152.815625 104.58 \nL 156.973625 104.58 \nL 156.973625 21.42 \nL 152.815625 21.42 \nz\n\" style=\"fill: #ffffff\"/>\n   </g>\n   <image xlink:href=\"data:image/png;base64,\niVBORw0KGgoAAAANSUhEUgAAAAYAAABzCAYAAAC7I3M6AAAAv0lEQVR4nNWVQQrDQAwDXfD/v9pLT1l71Q/sBBQMpTmukMbWhuSlz1txeDK0T+d3wrYFmyFm9BTjZsFjVY/g6GgSyKFaLsOHt82gKPlT+QvaJUbVWNS/waldjqKXGh24oB9lO/RL+AOBGGoSFjLsqDm4mj6W/lSFDDcK4ZvGbXTQ36BJQAa0Hlk2nAWCF081xsBKWiQcjyNykwM6jCx0nM8jO2wH74GbQ5S/IMIXMbASnAodFwllR/mOy76Puam+t3oOUronbv0AAAAASUVORK5CYII=\" id=\"imagec2f70b4f38\" transform=\"scale(1 -1) translate(0 -82.8)\" x=\"152.64\" y=\"-21.6\" width=\"4.32\" height=\"82.8\"/>\n   <g id=\"matplotlib.axis_3\"/>\n   <g id=\"matplotlib.axis_4\">\n    <g id=\"ytick_4\">\n     <g id=\"line2d_7\">\n      <defs>\n       <path id=\"m691de02ab5\" d=\"M 0 0 \nL 3.5 0 \n\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </defs>\n      <g>\n       <use xlink:href=\"#m691de02ab5\" x=\"156.973625\" y=\"104.58\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_9\">\n      <!-- 0.00 -->\n      <g transform=\"translate(163.973625 108.379219) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-2e\" d=\"M 684 794 \nL 1344 794 \nL 1344 0 \nL 684 0 \nL 684 794 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-30\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(95.410156 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(159.033203 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_5\">\n     <g id=\"line2d_8\">\n      <g>\n       <use xlink:href=\"#m691de02ab5\" x=\"156.973625\" y=\"83.79\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_10\">\n      <!-- 0.25 -->\n      <g transform=\"translate(163.973625 87.589219) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-35\" d=\"M 691 4666 \nL 3169 4666 \nL 3169 4134 \nL 1269 4134 \nL 1269 2991 \nQ 1406 3038 1543 3061 \nQ 1681 3084 1819 3084 \nQ 2600 3084 3056 2656 \nQ 3513 2228 3513 1497 \nQ 3513 744 3044 326 \nQ 2575 -91 1722 -91 \nQ 1428 -91 1123 -41 \nQ 819 9 494 109 \nL 494 744 \nQ 775 591 1075 516 \nQ 1375 441 1709 441 \nQ 2250 441 2565 725 \nQ 2881 1009 2881 1497 \nQ 2881 1984 2565 2268 \nQ 2250 2553 1709 2553 \nQ 1456 2553 1204 2497 \nQ 953 2441 691 2322 \nL 691 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-30\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-32\" transform=\"translate(95.410156 0)\"/>\n       <use xlink:href=\"#DejaVuSans-35\" transform=\"translate(159.033203 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_6\">\n     <g id=\"line2d_9\">\n      <g>\n       <use xlink:href=\"#m691de02ab5\" x=\"156.973625\" y=\"63\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_11\">\n      <!-- 0.50 -->\n      <g transform=\"translate(163.973625 66.799219) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-30\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-35\" transform=\"translate(95.410156 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(159.033203 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_7\">\n     <g id=\"line2d_10\">\n      <g>\n       <use xlink:href=\"#m691de02ab5\" x=\"156.973625\" y=\"42.21\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_12\">\n      <!-- 0.75 -->\n      <g transform=\"translate(163.973625 46.009219) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-37\" d=\"M 525 4666 \nL 3525 4666 \nL 3525 4397 \nL 1831 0 \nL 1172 0 \nL 2766 4134 \nL 525 4134 \nL 525 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-30\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-37\" transform=\"translate(95.410156 0)\"/>\n       <use xlink:href=\"#DejaVuSans-35\" transform=\"translate(159.033203 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_8\">\n     <g id=\"line2d_11\">\n      <g>\n       <use xlink:href=\"#m691de02ab5\" x=\"156.973625\" y=\"21.42\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_13\">\n      <!-- 1.00 -->\n      <g transform=\"translate(163.973625 25.219219) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-31\" d=\"M 794 531 \nL 1825 531 \nL 1825 4091 \nL 703 3866 \nL 703 4441 \nL 1819 4666 \nL 2450 4666 \nL 2450 531 \nL 3481 531 \nL 3481 0 \nL 794 0 \nL 794 531 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-31\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(95.410156 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(159.033203 0)\"/>\n      </g>\n     </g>\n    </g>\n   </g>\n   <g id=\"LineCollection_1\"/>\n   <g id=\"patch_8\">\n    <path d=\"M 152.815625 104.58 \nL 154.894625 104.58 \nL 156.973625 104.58 \nL 156.973625 21.42 \nL 154.894625 21.42 \nL 152.815625 21.42 \nL 152.815625 104.58 \nz\n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n  </g>\n </g>\n <defs>\n  <clipPath id=\"p2d9181945e\">\n   <rect x=\"34.240625\" y=\"7.2\" width=\"111.6\" height=\"111.6\"/>\n  </clipPath>\n </defs>\n</svg>\n"
     },
     "metadata": {}
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=attention-scoring-md-9 fw=all -->\n",
    "\n",
    "On the attention side, this triangular mask is the key difference between a\n",
    "model that merely reads a sequence and one that can be trained, in parallel\n",
    "over all positions, to generate it — generation also needs the shifted\n",
    "next-token objective and a decoding loop. The mask will accompany us through\n",
    "every decoder in the chapters ahead.\n",
    "\n",
    "### Composing Masks\n",
    "\n",
    "`valid_lens` describes prefixes, which cover the two cases above, but the\n",
    "general interface is a boolean tensor: entry $(i, j)$ says whether query $i$\n",
    "may attend to key $j$. Every requirement takes this form — padding excludes\n",
    "keys beyond the sequence length, causality excludes keys beyond the query,\n",
    "and structural patterns such as the attention windows of\n",
    "that section exclude by distance — and a key must\n",
    "survive *all* requirements at once, so masks compose by logical AND.\n",
    "Broadcasting keeps the bookkeeping cheap: a padding mask has shape\n",
    "$(\\textrm{batch}, 1, \\textrm{keys})$, a causal mask\n",
    "$(1, \\textrm{queries}, \\textrm{keys})$, and their AND broadcasts to the full\n",
    "$(\\textrm{batch}, \\textrm{queries}, \\textrm{keys})$ without materializing\n",
    "either input per example. Applying the composite is the same idiom as\n",
    "before: overwrite the excluded scores with the dtype's most negative finite\n",
    "value, then softmax."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "attention-scoring-composing-masks"
   },
   "source": [
    "valid_lens, n = jnp.array([6, 3]), 6\n",
    "j = jnp.arange(n)\n",
    "padding = (j[None, :] < valid_lens[:, None])[:, None, :]  # (batch, 1, key)\n",
    "causal = (j[None, :] <= j[:, None])[None, :, :]         # (1, query, key)\n",
    "mask = padding & causal                                 # (batch, query, key)\n",
    "scores = jax.random.normal(jax.random.key(0), (2, n, n))\n",
    "weights = jax.nn.softmax(\n",
    "    jnp.where(mask, scores, jnp.finfo(scores.dtype).min), axis=-1)\n",
    "d2l.show_heatmaps(weights[:, None], xlabel='Keys', ylabel='Queries')"
   ],
   "execution_count": 8,
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "image/svg+xml": "<?xml version=\"1.0\" encoding=\"utf-8\" standalone=\"no\"?>\n<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n  \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n<svg xmlns:xlink=\"http://www.w3.org/1999/xlink\" width=\"144.83925pt\" height=\"183.35625pt\" viewBox=\"0 0 144.83925 183.35625\" xmlns=\"http://www.w3.org/2000/svg\" version=\"1.1\">\n <metadata>\n  <rdf:RDF xmlns:dc=\"http://purl.org/dc/elements/1.1/\" xmlns:cc=\"http://creativecommons.org/ns#\" xmlns:rdf=\"http://www.w3.org/1999/02/22-rdf-syntax-ns#\">\n   <cc:Work>\n    <dc:type rdf:resource=\"http://purl.org/dc/dcmitype/StillImage\"/>\n    <dc:date>2026-07-19T13:51:10.855785</dc:date>\n    <dc:format>image/svg+xml</dc:format>\n    <dc:creator>\n     <cc:Agent>\n      <dc:title>Matplotlib v3.10.8, https://matplotlib.org/</dc:title>\n     </cc:Agent>\n    </dc:creator>\n   </cc:Work>\n  </rdf:RDF>\n </metadata>\n <defs>\n  <style type=\"text/css\">*{stroke-linejoin: round; stroke-linecap: butt}</style>\n </defs>\n <g id=\"figure_1\">\n  <g id=\"patch_1\">\n   <path d=\"M 0 183.35625 \nL 144.83925 183.35625 \nL 144.83925 0 \nL 0 0 \nz\n\" style=\"fill: #ffffff\"/>\n  </g>\n  <g id=\"axes_1\">\n   <g id=\"patch_2\">\n    <path d=\"M 34.240625 70.2 \nL 97.240625 70.2 \nL 97.240625 7.2 \nL 34.240625 7.2 \nz\n\" style=\"fill: #ffffff\"/>\n   </g>\n   <g clip-path=\"url(#pffacbb297e)\">\n    <image xlink:href=\"data:image/png;base64,\niVBORw0KGgoAAAANSUhEUgAAAFgAAABYCAYAAABxlTA0AAABhklEQVR4nO3bPUoEQRRF4aphEkEDERTMNFdjA8HMVRi7BCNdh6k7MBZXMMGAICiCBiIayeAfg8F0u4Xm2gdRzpc/XnGoqKBqM3luS6gZX0RzdWk1XVnar2k8WxeXs51vL/HOQTypTgwMMzDMwDADwwwMMzDMwDADwwwMMzDMwDADwwwMq7O7cfxcWcKnw+bsNF5Zt3fj2cHGTrZzbiHfGU+qEwPDDAwzMMzAMAPDDAwzMMzAMAPDDAwzMMzAsDq7HcWvaXVlrc+zdHK+vhnP7j1c93iSbrzBMAPDDAwzMMzAMAPDDAwzMMzAMAPDDAwzMMzAMAPDhu3TfTzcvk+iuenxUbzzN54cf8IbDDMwzMAwA8MMDDMwzMAwA8MMDDMwzMAwA8MMDDMwbFg+XuPhq/3DaG7r5jLe+dd4g2EGhhkYZmCYgWEGhhkYZmCYgWEGhhkYZmCYgWH1oMzHf5VPPh/7PMu/5A2GGRhmYJiBYQaGGRhmYJiBYQaGGRhmYJiBYQaGGRj2DUPNL9Om11UYAAAAAElFTkSuQmCC\" id=\"image35fa0b54e1\" transform=\"scale(1 -1) translate(0 -63.36)\" x=\"34.240625\" y=\"-6.84\" width=\"63.36\" height=\"63.36\"/>\n   </g>\n   <g id=\"matplotlib.axis_1\">\n    <g id=\"xtick_1\">\n     <g id=\"line2d_1\">\n      <defs>\n       <path id=\"m2d8f6b5fd7\" d=\"M 0 0 \nL 0 3.5 \n\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </defs>\n      <g>\n       <use xlink:href=\"#m2d8f6b5fd7\" x=\"39.490625\" y=\"70.2\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_2\">\n     <g id=\"line2d_2\">\n      <g>\n       <use xlink:href=\"#m2d8f6b5fd7\" x=\"91.990625\" y=\"70.2\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n    </g>\n   </g>\n   <g id=\"matplotlib.axis_2\">\n    <g id=\"ytick_1\">\n     <g id=\"line2d_3\">\n      <defs>\n       <path id=\"mcaaa74b8c2\" d=\"M 0 0 \nL -3.5 0 \n\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </defs>\n      <g>\n       <use xlink:href=\"#mcaaa74b8c2\" x=\"34.240625\" y=\"12.45\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_1\">\n      <!-- 0 -->\n      <g transform=\"translate(20.878125 16.249219) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-30\" d=\"M 2034 4250 \nQ 1547 4250 1301 3770 \nQ 1056 3291 1056 2328 \nQ 1056 1369 1301 889 \nQ 1547 409 2034 409 \nQ 2525 409 2770 889 \nQ 3016 1369 3016 2328 \nQ 3016 3291 2770 3770 \nQ 2525 4250 2034 4250 \nz\nM 2034 4750 \nQ 2819 4750 3233 4129 \nQ 3647 3509 3647 2328 \nQ 3647 1150 3233 529 \nQ 2819 -91 2034 -91 \nQ 1250 -91 836 529 \nQ 422 1150 422 2328 \nQ 422 3509 836 4129 \nQ 1250 4750 2034 4750 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-30\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_2\">\n     <g id=\"line2d_4\">\n      <g>\n       <use xlink:href=\"#mcaaa74b8c2\" x=\"34.240625\" y=\"33.45\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_2\">\n      <!-- 2 -->\n      <g transform=\"translate(20.878125 37.249219) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-32\" d=\"M 1228 531 \nL 3431 531 \nL 3431 0 \nL 469 0 \nL 469 531 \nQ 828 903 1448 1529 \nQ 2069 2156 2228 2338 \nQ 2531 2678 2651 2914 \nQ 2772 3150 2772 3378 \nQ 2772 3750 2511 3984 \nQ 2250 4219 1831 4219 \nQ 1534 4219 1204 4116 \nQ 875 4013 500 3803 \nL 500 4441 \nQ 881 4594 1212 4672 \nQ 1544 4750 1819 4750 \nQ 2544 4750 2975 4387 \nQ 3406 4025 3406 3419 \nQ 3406 3131 3298 2873 \nQ 3191 2616 2906 2266 \nQ 2828 2175 2409 1742 \nQ 1991 1309 1228 531 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-32\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_3\">\n     <g id=\"line2d_5\">\n      <g>\n       <use xlink:href=\"#mcaaa74b8c2\" x=\"34.240625\" y=\"54.45\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_3\">\n      <!-- 4 -->\n      <g transform=\"translate(20.878125 58.249219) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-34\" d=\"M 2419 4116 \nL 825 1625 \nL 2419 1625 \nL 2419 4116 \nz\nM 2253 4666 \nL 3047 4666 \nL 3047 1625 \nL 3713 1625 \nL 3713 1100 \nL 3047 1100 \nL 3047 0 \nL 2419 0 \nL 2419 1100 \nL 313 1100 \nL 313 1709 \nL 2253 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-34\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_4\">\n     <!-- Queries -->\n     <g transform=\"translate(14.798438 58.007031) rotate(-90) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-51\" d=\"M 2522 4238 \nQ 1834 4238 1429 3725 \nQ 1025 3213 1025 2328 \nQ 1025 1447 1429 934 \nQ 1834 422 2522 422 \nQ 3209 422 3611 934 \nQ 4013 1447 4013 2328 \nQ 4013 3213 3611 3725 \nQ 3209 4238 2522 4238 \nz\nM 3406 84 \nL 4238 -825 \nL 3475 -825 \nL 2784 -78 \nQ 2681 -84 2626 -87 \nQ 2572 -91 2522 -91 \nQ 1538 -91 948 567 \nQ 359 1225 359 2328 \nQ 359 3434 948 4092 \nQ 1538 4750 2522 4750 \nQ 3503 4750 4090 4092 \nQ 4678 3434 4678 2328 \nQ 4678 1516 4351 937 \nQ 4025 359 3406 84 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-75\" d=\"M 544 1381 \nL 544 3500 \nL 1119 3500 \nL 1119 1403 \nQ 1119 906 1312 657 \nQ 1506 409 1894 409 \nQ 2359 409 2629 706 \nQ 2900 1003 2900 1516 \nL 2900 3500 \nL 3475 3500 \nL 3475 0 \nL 2900 0 \nL 2900 538 \nQ 2691 219 2414 64 \nQ 2138 -91 1772 -91 \nQ 1169 -91 856 284 \nQ 544 659 544 1381 \nz\nM 1991 3584 \nL 1991 3584 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-65\" d=\"M 3597 1894 \nL 3597 1613 \nL 953 1613 \nQ 991 1019 1311 708 \nQ 1631 397 2203 397 \nQ 2534 397 2845 478 \nQ 3156 559 3463 722 \nL 3463 178 \nQ 3153 47 2828 -22 \nQ 2503 -91 2169 -91 \nQ 1331 -91 842 396 \nQ 353 884 353 1716 \nQ 353 2575 817 3079 \nQ 1281 3584 2069 3584 \nQ 2775 3584 3186 3129 \nQ 3597 2675 3597 1894 \nz\nM 3022 2063 \nQ 3016 2534 2758 2815 \nQ 2500 3097 2075 3097 \nQ 1594 3097 1305 2825 \nQ 1016 2553 972 2059 \nL 3022 2063 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-72\" d=\"M 2631 2963 \nQ 2534 3019 2420 3045 \nQ 2306 3072 2169 3072 \nQ 1681 3072 1420 2755 \nQ 1159 2438 1159 1844 \nL 1159 0 \nL 581 0 \nL 581 3500 \nL 1159 3500 \nL 1159 2956 \nQ 1341 3275 1631 3429 \nQ 1922 3584 2338 3584 \nQ 2397 3584 2469 3576 \nQ 2541 3569 2628 3553 \nL 2631 2963 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-69\" d=\"M 603 3500 \nL 1178 3500 \nL 1178 0 \nL 603 0 \nL 603 3500 \nz\nM 603 4863 \nL 1178 4863 \nL 1178 4134 \nL 603 4134 \nL 603 4863 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-73\" d=\"M 2834 3397 \nL 2834 2853 \nQ 2591 2978 2328 3040 \nQ 2066 3103 1784 3103 \nQ 1356 3103 1142 2972 \nQ 928 2841 928 2578 \nQ 928 2378 1081 2264 \nQ 1234 2150 1697 2047 \nL 1894 2003 \nQ 2506 1872 2764 1633 \nQ 3022 1394 3022 966 \nQ 3022 478 2636 193 \nQ 2250 -91 1575 -91 \nQ 1294 -91 989 -36 \nQ 684 19 347 128 \nL 347 722 \nQ 666 556 975 473 \nQ 1284 391 1588 391 \nQ 1994 391 2212 530 \nQ 2431 669 2431 922 \nQ 2431 1156 2273 1281 \nQ 2116 1406 1581 1522 \nL 1381 1569 \nQ 847 1681 609 1914 \nQ 372 2147 372 2553 \nQ 372 3047 722 3315 \nQ 1072 3584 1716 3584 \nQ 2034 3584 2315 3537 \nQ 2597 3491 2834 3397 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-51\"/>\n      <use xlink:href=\"#DejaVuSans-75\" transform=\"translate(78.710938 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(142.089844 0)\"/>\n      <use xlink:href=\"#DejaVuSans-72\" transform=\"translate(203.613281 0)\"/>\n      <use xlink:href=\"#DejaVuSans-69\" transform=\"translate(244.726562 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(272.509766 0)\"/>\n      <use xlink:href=\"#DejaVuSans-73\" transform=\"translate(334.033203 0)\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"patch_3\">\n    <path d=\"M 34.240625 70.2 \nL 34.240625 7.2 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_4\">\n    <path d=\"M 97.240625 70.2 \nL 97.240625 7.2 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_5\">\n    <path d=\"M 34.240625 70.2 \nL 97.240625 70.2 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_6\">\n    <path d=\"M 34.240625 7.2 \nL 97.240625 7.2 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n  </g>\n  <g id=\"axes_2\">\n   <g id=\"patch_7\">\n    <path d=\"M 34.240625 145.8 \nL 97.240625 145.8 \nL 97.240625 82.8 \nL 34.240625 82.8 \nz\n\" style=\"fill: #ffffff\"/>\n   </g>\n   <g clip-path=\"url(#pd97db22228)\">\n    <image xlink:href=\"data:image/png;base64,\niVBORw0KGgoAAAANSUhEUgAAAFgAAABYCAYAAABxlTA0AAABaUlEQVR4nO3cMUoDUQBF0fnJlBYSbFK4gIghZAU2FsI0swYtsiLBNqtIIy4gRQTtrLW2CDYhJKltH7kMCff0jz9cfjUMU3Y/X/sq9Dq5i3b372/pkVUZDONtF3pdP8C5MzDMwDADwwwMMzDMwDADwwwMMzDMwDADwwwMMzCsbGYP8evKcjOKdv32MT2yKlfX8bYL3mCYgWEGhhkYZmCYgWEGhhkYZmCYgWEGhhkYZmBY3RvfxuPdx2c2bOMjT443GGZgmIFhBoYZGGZgmIFhBoYZGGZgmIFhBoYZGGZgWL1drvJx00S7/fo3PtOP//SPgWEGhhkYZmCYgWEGhhkYZmCYgWEGhhkYZmCYgWF1fzqOx+vnebS7XDzFZ54abzDMwDADwwwMMzDMwDADwwwMMzDMwDADwwwMMzCszKqL+M9/L3/fx3yWs+QNhhkYZmCYgWEGhhkYZmCYgWEGhhkYZmCYgWEGhhkYdgBaQRyge+ToDwAAAABJRU5ErkJggg==\" id=\"image43334ab3ee\" transform=\"scale(1 -1) translate(0 -63.36)\" x=\"34.240625\" y=\"-82.44\" width=\"63.36\" height=\"63.36\"/>\n   </g>\n   <g id=\"matplotlib.axis_3\">\n    <g id=\"xtick_3\">\n     <g id=\"line2d_6\">\n      <g>\n       <use xlink:href=\"#m2d8f6b5fd7\" x=\"39.490625\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_5\">\n      <!-- 0 -->\n      <g transform=\"translate(36.309375 160.398438) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-30\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_4\">\n     <g id=\"line2d_7\">\n      <g>\n       <use xlink:href=\"#m2d8f6b5fd7\" x=\"91.990625\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_6\">\n      <!-- 5 -->\n      <g transform=\"translate(88.809375 160.398438) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-35\" d=\"M 691 4666 \nL 3169 4666 \nL 3169 4134 \nL 1269 4134 \nL 1269 2991 \nQ 1406 3038 1543 3061 \nQ 1681 3084 1819 3084 \nQ 2600 3084 3056 2656 \nQ 3513 2228 3513 1497 \nQ 3513 744 3044 326 \nQ 2575 -91 1722 -91 \nQ 1428 -91 1123 -41 \nQ 819 9 494 109 \nL 494 744 \nQ 775 591 1075 516 \nQ 1375 441 1709 441 \nQ 2250 441 2565 725 \nQ 2881 1009 2881 1497 \nQ 2881 1984 2565 2268 \nQ 2250 2553 1709 2553 \nQ 1456 2553 1204 2497 \nQ 953 2441 691 2322 \nL 691 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-35\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_7\">\n     <!-- Keys -->\n     <g transform=\"translate(54.071094 174.076563) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-4b\" d=\"M 628 4666 \nL 1259 4666 \nL 1259 2694 \nL 3353 4666 \nL 4166 4666 \nL 1850 2491 \nL 4331 0 \nL 3500 0 \nL 1259 2247 \nL 1259 0 \nL 628 0 \nL 628 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-79\" d=\"M 2059 -325 \nQ 1816 -950 1584 -1140 \nQ 1353 -1331 966 -1331 \nL 506 -1331 \nL 506 -850 \nL 844 -850 \nQ 1081 -850 1212 -737 \nQ 1344 -625 1503 -206 \nL 1606 56 \nL 191 3500 \nL 800 3500 \nL 1894 763 \nL 2988 3500 \nL 3597 3500 \nL 2059 -325 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-4b\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(60.576172 0)\"/>\n      <use xlink:href=\"#DejaVuSans-79\" transform=\"translate(122.099609 0)\"/>\n      <use xlink:href=\"#DejaVuSans-73\" transform=\"translate(181.279297 0)\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"matplotlib.axis_4\">\n    <g id=\"ytick_4\">\n     <g id=\"line2d_8\">\n      <g>\n       <use xlink:href=\"#mcaaa74b8c2\" x=\"34.240625\" y=\"88.05\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_8\">\n      <!-- 0 -->\n      <g transform=\"translate(20.878125 91.849219) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-30\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_5\">\n     <g id=\"line2d_9\">\n      <g>\n       <use xlink:href=\"#mcaaa74b8c2\" x=\"34.240625\" y=\"109.05\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_9\">\n      <!-- 2 -->\n      <g transform=\"translate(20.878125 112.849219) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-32\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_6\">\n     <g id=\"line2d_10\">\n      <g>\n       <use xlink:href=\"#mcaaa74b8c2\" x=\"34.240625\" y=\"130.05\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_10\">\n      <!-- 4 -->\n      <g transform=\"translate(20.878125 133.849219) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-34\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_11\">\n     <!-- Queries -->\n     <g transform=\"translate(14.798437 133.607031) rotate(-90) scale(0.1 -0.1)\">\n      <use xlink:href=\"#DejaVuSans-51\"/>\n      <use xlink:href=\"#DejaVuSans-75\" transform=\"translate(78.710938 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(142.089844 0)\"/>\n      <use xlink:href=\"#DejaVuSans-72\" transform=\"translate(203.613281 0)\"/>\n      <use xlink:href=\"#DejaVuSans-69\" transform=\"translate(244.726562 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(272.509766 0)\"/>\n      <use xlink:href=\"#DejaVuSans-73\" transform=\"translate(334.033203 0)\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"patch_8\">\n    <path d=\"M 34.240625 145.8 \nL 34.240625 82.8 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_9\">\n    <path d=\"M 97.240625 145.8 \nL 97.240625 82.8 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_10\">\n    <path d=\"M 34.240625 145.8 \nL 97.240625 145.8 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_11\">\n    <path d=\"M 34.240625 82.8 \nL 97.240625 82.8 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n  </g>\n  <g id=\"axes_3\">\n   <g id=\"patch_12\">\n    <path d=\"M 104.215625 118.08 \nL 108.373625 118.08 \nL 108.373625 34.92 \nL 104.215625 34.92 \nz\n\" style=\"fill: #ffffff\"/>\n   </g>\n   <image xlink:href=\"data:image/png;base64,\niVBORw0KGgoAAAANSUhEUgAAAAYAAABzCAYAAAC7I3M6AAAAx0lEQVR4nL2VSQ7EMAgEicT/vzqXOSUs+QBlqSPLObrVFGCIr/7/2obPrXI6N7cuVagRsXQAvNmhZ4UF6unK8HjUrLAlqTosCa6HQoEKbD2UnlXENgYLchNTnkQMpY/PhynRu0sFMhz3gxhyqMV+yPCNBXJ3cc+xQGSAUMRAhw7vhL9on4DLTSxKt+jxQiEwlMxIGZ5xgEF1oFD0pKIw35+ZJzrmc/MyhBODHBvh6EDGPLkrB8yCeZDj2RcKHfcBBrbkxq1V4S8NNhBOPsFR5gAAAABJRU5ErkJggg==\" id=\"imagebf26c682fe\" transform=\"scale(1 -1) translate(0 -82.8)\" x=\"104.4\" y=\"-34.56\" width=\"4.32\" height=\"82.8\"/>\n   <g id=\"matplotlib.axis_5\"/>\n   <g id=\"matplotlib.axis_6\">\n    <g id=\"ytick_7\">\n     <g id=\"line2d_11\">\n      <defs>\n       <path id=\"mc82d5e44f7\" d=\"M 0 0 \nL 3.5 0 \n\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </defs>\n      <g>\n       <use xlink:href=\"#mc82d5e44f7\" x=\"108.373625\" y=\"118.08\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_12\">\n      <!-- 0.00 -->\n      <g transform=\"translate(115.373625 121.879219) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-2e\" d=\"M 684 794 \nL 1344 794 \nL 1344 0 \nL 684 0 \nL 684 794 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-30\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(95.410156 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(159.033203 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_8\">\n     <g id=\"line2d_12\">\n      <g>\n       <use xlink:href=\"#mc82d5e44f7\" x=\"108.373625\" y=\"97.29\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_13\">\n      <!-- 0.25 -->\n      <g transform=\"translate(115.373625 101.089219) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-30\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-32\" transform=\"translate(95.410156 0)\"/>\n       <use xlink:href=\"#DejaVuSans-35\" transform=\"translate(159.033203 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_9\">\n     <g id=\"line2d_13\">\n      <g>\n       <use xlink:href=\"#mc82d5e44f7\" x=\"108.373625\" y=\"76.5\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_14\">\n      <!-- 0.50 -->\n      <g transform=\"translate(115.373625 80.299219) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-30\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-35\" transform=\"translate(95.410156 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(159.033203 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_10\">\n     <g id=\"line2d_14\">\n      <g>\n       <use xlink:href=\"#mc82d5e44f7\" x=\"108.373625\" y=\"55.71\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_15\">\n      <!-- 0.75 -->\n      <g transform=\"translate(115.373625 59.509219) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-37\" d=\"M 525 4666 \nL 3525 4666 \nL 3525 4397 \nL 1831 0 \nL 1172 0 \nL 2766 4134 \nL 525 4134 \nL 525 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-30\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-37\" transform=\"translate(95.410156 0)\"/>\n       <use xlink:href=\"#DejaVuSans-35\" transform=\"translate(159.033203 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_11\">\n     <g id=\"line2d_15\">\n      <g>\n       <use xlink:href=\"#mc82d5e44f7\" x=\"108.373625\" y=\"34.92\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_16\">\n      <!-- 1.00 -->\n      <g transform=\"translate(115.373625 38.719219) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-31\" d=\"M 794 531 \nL 1825 531 \nL 1825 4091 \nL 703 3866 \nL 703 4441 \nL 1819 4666 \nL 2450 4666 \nL 2450 531 \nL 3481 531 \nL 3481 0 \nL 794 0 \nL 794 531 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-31\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(95.410156 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(159.033203 0)\"/>\n      </g>\n     </g>\n    </g>\n   </g>\n   <g id=\"LineCollection_1\"/>\n   <g id=\"patch_13\">\n    <path d=\"M 104.215625 118.08 \nL 106.294625 118.08 \nL 108.373625 118.08 \nL 108.373625 34.92 \nL 106.294625 34.92 \nL 104.215625 34.92 \nL 104.215625 118.08 \nz\n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n  </g>\n </g>\n <defs>\n  <clipPath id=\"pffacbb297e\">\n   <rect x=\"34.240625\" y=\"7.2\" width=\"63\" height=\"63\"/>\n  </clipPath>\n  <clipPath id=\"pd97db22228\">\n   <rect x=\"34.240625\" y=\"82.8\" width=\"63\" height=\"63\"/>\n  </clipPath>\n </defs>\n</svg>\n"
     },
     "metadata": {}
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=attention-scoring-md-10 fw=all -->\n",
    "\n",
    "The first sequence shows the plain causal triangle; the second is cut off\n",
    "at its valid length of $3$, the intersection of both constraints. Composition\n",
    "sharpens the fully-masked hazard flagged above: masks that are harmless\n",
    "alone can leave some query with an empty intersection, so the guarantee of\n",
    "at least one valid key per query must hold for the *composite*. The same\n",
    "machinery handles *packed sequences* — several documents concatenated into\n",
    "one training row — by ANDing the causal mask with a block-diagonal mask\n",
    "that keeps each document from attending into its neighbors.\n",
    "\n",
    "## Batched Attention\n",
    "\n",
    "### Batch Matrix Multiplication\n",
    "\n",
    "Attention is computed on minibatches of queries, keys, and values, so we\n",
    "need to multiply batches of matrices by one another. Assume that\n",
    "\n",
    "$$\n",
    "\\mathbf{Q} = [\\mathbf{Q}_1, \\mathbf{Q}_2, \\ldots, \\mathbf{Q}_n]  \\in \\mathbb{R}^{n \\times a \\times b}, \\qquad\n",
    "\\mathbf{K} = [\\mathbf{K}_1, \\mathbf{K}_2, \\ldots, \\mathbf{K}_n]  \\in \\mathbb{R}^{n \\times b \\times c}.\n",
    "$$\n",
    "\n",
    "Then the batch matrix multiplication (BMM) computes one matrix product per\n",
    "batch element,\n",
    "\n",
    "$$\\textrm{BMM}(\\mathbf{Q}, \\mathbf{K}) = [\\mathbf{Q}_1 \\mathbf{K}_1, \\mathbf{Q}_2 \\mathbf{K}_2, \\ldots, \\mathbf{Q}_n \\mathbf{K}_n] \\in \\mathbb{R}^{n \\times a \\times c}.$$\n",
    "\n",
    "Let's see this in action in a deep learning framework:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "attention-scoring-batch-matrix-multiplication"
   },
   "source": [
    "Q = jnp.ones((2, 3, 4))\n",
    "K = jnp.ones((2, 4, 6))\n",
    "d2l.check_shape(jax.lax.batch_matmul(Q, K), (2, 3, 6))"
   ],
   "execution_count": 9,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=attention-scoring-md-11 fw=all -->\n",
    "\n",
    "### The DotProductAttention Class\n",
    "\n",
    "Now we can state scaled dot-product attention in the form in which it is\n",
    "actually computed. For $n$ queries and $m$ key--value pairs, with queries\n",
    "and keys of length $d$ and values of length $v$, stack the queries into\n",
    "$\\mathbf{Q} \\in \\mathbb{R}^{n \\times d}$, the keys into $\\mathbf{K} \\in\n",
    "\\mathbb{R}^{m \\times d}$, and the values into $\\mathbf{V} \\in\n",
    "\\mathbb{R}^{m \\times v}$. Then the entire attention computation is two\n",
    "matrix products and a softmax:\n",
    "\n",
    "$$ \\mathrm{softmax}\\left(\\frac{\\mathbf{Q} \\mathbf{K}^\\top }{\\sqrt{d}}\\right) \\mathbf{V} \\in \\mathbb{R}^{n \\times v}.$$\n",
    "\n",
    "Requiring queries and keys to share the length $d$ is no real restriction:\n",
    "a learned matrix $\\mathbf{M}$ turns $\\mathbf{q}^\\top \\mathbf{k}$ into\n",
    "$\\mathbf{q}^\\top \\mathbf{M} \\mathbf{k}$ and translates between spaces of\n",
    "different dimension (an exercise below). Applied to a minibatch,\n",
    "the equation uses the batch matrix multiplication of\n",
    "the equation twice. The implementation applies dropout to\n",
    "the attention weights as regularization, and stores the weights for\n",
    "visualization:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "attention-scoring-the-dotproductattention-class-1"
   },
   "source": [
    "class DotProductAttention(nnx.Module):\n",
    "    \"\"\"Scaled dot product attention.\"\"\"\n",
    "    def __init__(self, dropout, rngs=None):\n",
    "        rngs = nnx.Rngs(dropout=0) if rngs is None else rngs\n",
    "        self.dropout = nnx.Dropout(dropout, rngs=rngs)\n",
    "\n",
    "    # Shape of queries: (batch_size, no. of queries, d)\n",
    "    # Shape of keys: (batch_size, no. of key-value pairs, d)\n",
    "    # Shape of values: (batch_size, no. of key-value pairs, value dimension)\n",
    "    # Shape of valid_lens: (batch_size,) or (batch_size, no. of queries)\n",
    "    def __call__(self, queries, keys, values, valid_lens=None):\n",
    "        d = queries.shape[-1]\n",
    "        # Swap the last two dimensions of keys with keys.swapaxes(1, 2)\n",
    "        scores = queries@(keys.swapaxes(1, 2)) / math.sqrt(d)\n",
    "        attention_weights = masked_softmax(scores, valid_lens)\n",
    "        # NNX idiom: return (output, weights); PyTorch stores weights on self\n",
    "        return self.dropout(attention_weights) @ values, attention_weights"
   ],
   "execution_count": 10,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=attention-scoring-md-12 fw=all -->\n",
    "\n",
    "To see the class at work we use a minibatch of size $2$, with $10$ keys and\n",
    "values of dimension $4$, a single $2$-dimensional query per example, and\n",
    "valid lengths of $2$ and $6$. The output is one $4$-dimensional row per\n",
    "query:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "attention-scoring-the-dotproductattention-class-2"
   },
   "source": [
    "queries = jax.random.normal(jax.random.key(0), (2, 1, 2))\n",
    "keys = jax.random.normal(jax.random.key(1), (2, 10, 2))\n",
    "values = jax.random.normal(jax.random.key(2), (2, 10, 4))\n",
    "valid_lens = jnp.array([2, 6])\n",
    "\n",
    "attention = DotProductAttention(dropout=0.5)\n",
    "output, attention_weights = nnx.view(\n",
    "    attention, deterministic=True)(queries, keys, values, valid_lens)\n",
    "d2l.check_shape(output, (2, 1, 4))"
   ],
   "execution_count": 11,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=attention-scoring-md-13 fw=all -->\n",
    "\n",
    "The stored attention weights confirm that the mask did its job: weights\n",
    "vanish beyond the second and sixth key, respectively."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "attention-scoring-the-dotproductattention-class-3"
   },
   "source": [
    "d2l.show_heatmaps(attention_weights.reshape((1, 1, 2, 10)),\n",
    "                  xlabel='Keys', ylabel='Queries')"
   ],
   "execution_count": 12,
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "image/svg+xml": "<?xml version=\"1.0\" encoding=\"utf-8\" standalone=\"no\"?>\n<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n  \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n<svg xmlns:xlink=\"http://www.w3.org/1999/xlink\" width=\"187.07675pt\" height=\"103.438906pt\" viewBox=\"0 0 187.07675 103.438906\" xmlns=\"http://www.w3.org/2000/svg\" version=\"1.1\">\n <metadata>\n  <rdf:RDF xmlns:dc=\"http://purl.org/dc/elements/1.1/\" xmlns:cc=\"http://creativecommons.org/ns#\" xmlns:rdf=\"http://www.w3.org/1999/02/22-rdf-syntax-ns#\">\n   <cc:Work>\n    <dc:type rdf:resource=\"http://purl.org/dc/dcmitype/StillImage\"/>\n    <dc:date>2026-07-19T13:51:13.607498</dc:date>\n    <dc:format>image/svg+xml</dc:format>\n    <dc:creator>\n     <cc:Agent>\n      <dc:title>Matplotlib v3.10.8, https://matplotlib.org/</dc:title>\n     </cc:Agent>\n    </dc:creator>\n   </cc:Work>\n  </rdf:RDF>\n </metadata>\n <defs>\n  <style type=\"text/css\">*{stroke-linejoin: round; stroke-linecap: butt}</style>\n </defs>\n <g id=\"figure_1\">\n  <g id=\"patch_1\">\n   <path d=\"M -0 103.438906 \nL 187.07675 103.438906 \nL 187.07675 0 \nL -0 0 \nz\n\" style=\"fill: #ffffff\"/>\n  </g>\n  <g id=\"axes_1\">\n   <g id=\"patch_2\">\n    <path d=\"M 34.240625 59.94 \nL 145.840625 59.94 \nL 145.840625 37.62 \nL 34.240625 37.62 \nz\n\" style=\"fill: #ffffff\"/>\n   </g>\n   <g clip-path=\"url(#pb009b26535)\">\n    <image xlink:href=\"data:image/png;base64,\niVBORw0KGgoAAAANSUhEUgAAAJsAAAAfCAYAAADwQL9CAAAAuklEQVR4nO3aoQ0CQRRF0ZnVODwSFIKEArBUgqAf+qAUFAYFCY7QAMnSw0zyyIZz/MsfcbNq6+d8Gkur96t5Wkop4/PRvK27fdftcr+1315tuk4Py23XfqqGXz+A/yE2YsRGjNiIERsxYiNGbMSIjRixESM2YsRGjNiIERsxYiNGbMTUQ5k1/892XMy7jq+vl6490+LLRozYiBEbMWIjRmzEiI0YsREjNmLERozYiBEbMWIjRmzEiI2YL5vSDwMOfRR4AAAAAElFTkSuQmCC\" id=\"image8b61c4e87a\" transform=\"scale(1 -1) translate(0 -22.32)\" x=\"34.240625\" y=\"-37.62\" width=\"111.6\" height=\"22.32\"/>\n   </g>\n   <g id=\"matplotlib.axis_1\">\n    <g id=\"xtick_1\">\n     <g id=\"line2d_1\">\n      <defs>\n       <path id=\"m8a0c1b934f\" d=\"M 0 0 \nL 0 3.5 \n\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </defs>\n      <g>\n       <use xlink:href=\"#m8a0c1b934f\" x=\"39.820625\" y=\"59.94\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_1\">\n      <!-- 0 -->\n      <g transform=\"translate(36.639375 74.538438) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-30\" d=\"M 2034 4250 \nQ 1547 4250 1301 3770 \nQ 1056 3291 1056 2328 \nQ 1056 1369 1301 889 \nQ 1547 409 2034 409 \nQ 2525 409 2770 889 \nQ 3016 1369 3016 2328 \nQ 3016 3291 2770 3770 \nQ 2525 4250 2034 4250 \nz\nM 2034 4750 \nQ 2819 4750 3233 4129 \nQ 3647 3509 3647 2328 \nQ 3647 1150 3233 529 \nQ 2819 -91 2034 -91 \nQ 1250 -91 836 529 \nQ 422 1150 422 2328 \nQ 422 3509 836 4129 \nQ 1250 4750 2034 4750 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-30\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_2\">\n     <g id=\"line2d_2\">\n      <g>\n       <use xlink:href=\"#m8a0c1b934f\" x=\"95.620625\" y=\"59.94\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_2\">\n      <!-- 5 -->\n      <g transform=\"translate(92.439375 74.538438) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-35\" d=\"M 691 4666 \nL 3169 4666 \nL 3169 4134 \nL 1269 4134 \nL 1269 2991 \nQ 1406 3038 1543 3061 \nQ 1681 3084 1819 3084 \nQ 2600 3084 3056 2656 \nQ 3513 2228 3513 1497 \nQ 3513 744 3044 326 \nQ 2575 -91 1722 -91 \nQ 1428 -91 1123 -41 \nQ 819 9 494 109 \nL 494 744 \nQ 775 591 1075 516 \nQ 1375 441 1709 441 \nQ 2250 441 2565 725 \nQ 2881 1009 2881 1497 \nQ 2881 1984 2565 2268 \nQ 2250 2553 1709 2553 \nQ 1456 2553 1204 2497 \nQ 953 2441 691 2322 \nL 691 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-35\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_3\">\n     <!-- Keys -->\n     <g transform=\"translate(78.371094 88.216563) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-4b\" d=\"M 628 4666 \nL 1259 4666 \nL 1259 2694 \nL 3353 4666 \nL 4166 4666 \nL 1850 2491 \nL 4331 0 \nL 3500 0 \nL 1259 2247 \nL 1259 0 \nL 628 0 \nL 628 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-65\" d=\"M 3597 1894 \nL 3597 1613 \nL 953 1613 \nQ 991 1019 1311 708 \nQ 1631 397 2203 397 \nQ 2534 397 2845 478 \nQ 3156 559 3463 722 \nL 3463 178 \nQ 3153 47 2828 -22 \nQ 2503 -91 2169 -91 \nQ 1331 -91 842 396 \nQ 353 884 353 1716 \nQ 353 2575 817 3079 \nQ 1281 3584 2069 3584 \nQ 2775 3584 3186 3129 \nQ 3597 2675 3597 1894 \nz\nM 3022 2063 \nQ 3016 2534 2758 2815 \nQ 2500 3097 2075 3097 \nQ 1594 3097 1305 2825 \nQ 1016 2553 972 2059 \nL 3022 2063 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-79\" d=\"M 2059 -325 \nQ 1816 -950 1584 -1140 \nQ 1353 -1331 966 -1331 \nL 506 -1331 \nL 506 -850 \nL 844 -850 \nQ 1081 -850 1212 -737 \nQ 1344 -625 1503 -206 \nL 1606 56 \nL 191 3500 \nL 800 3500 \nL 1894 763 \nL 2988 3500 \nL 3597 3500 \nL 2059 -325 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-73\" d=\"M 2834 3397 \nL 2834 2853 \nQ 2591 2978 2328 3040 \nQ 2066 3103 1784 3103 \nQ 1356 3103 1142 2972 \nQ 928 2841 928 2578 \nQ 928 2378 1081 2264 \nQ 1234 2150 1697 2047 \nL 1894 2003 \nQ 2506 1872 2764 1633 \nQ 3022 1394 3022 966 \nQ 3022 478 2636 193 \nQ 2250 -91 1575 -91 \nQ 1294 -91 989 -36 \nQ 684 19 347 128 \nL 347 722 \nQ 666 556 975 473 \nQ 1284 391 1588 391 \nQ 1994 391 2212 530 \nQ 2431 669 2431 922 \nQ 2431 1156 2273 1281 \nQ 2116 1406 1581 1522 \nL 1381 1569 \nQ 847 1681 609 1914 \nQ 372 2147 372 2553 \nQ 372 3047 722 3315 \nQ 1072 3584 1716 3584 \nQ 2034 3584 2315 3537 \nQ 2597 3491 2834 3397 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-4b\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(60.576172 0)\"/>\n      <use xlink:href=\"#DejaVuSans-79\" transform=\"translate(122.099609 0)\"/>\n      <use xlink:href=\"#DejaVuSans-73\" transform=\"translate(181.279297 0)\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"matplotlib.axis_2\">\n    <g id=\"ytick_1\">\n     <g id=\"line2d_3\">\n      <defs>\n       <path id=\"m38bdfd003c\" d=\"M 0 0 \nL -3.5 0 \n\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </defs>\n      <g>\n       <use xlink:href=\"#m38bdfd003c\" x=\"34.240625\" y=\"43.2\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_4\">\n      <!-- 0 -->\n      <g transform=\"translate(20.878125 46.999219) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-30\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_2\">\n     <g id=\"line2d_4\">\n      <g>\n       <use xlink:href=\"#m38bdfd003c\" x=\"34.240625\" y=\"54.36\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_5\">\n      <!-- 1 -->\n      <g transform=\"translate(20.878125 58.159219) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-31\" d=\"M 794 531 \nL 1825 531 \nL 1825 4091 \nL 703 3866 \nL 703 4441 \nL 1819 4666 \nL 2450 4666 \nL 2450 531 \nL 3481 531 \nL 3481 0 \nL 794 0 \nL 794 531 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-31\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_6\">\n     <!-- Queries -->\n     <g transform=\"translate(14.798437 68.087031) rotate(-90) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-51\" d=\"M 2522 4238 \nQ 1834 4238 1429 3725 \nQ 1025 3213 1025 2328 \nQ 1025 1447 1429 934 \nQ 1834 422 2522 422 \nQ 3209 422 3611 934 \nQ 4013 1447 4013 2328 \nQ 4013 3213 3611 3725 \nQ 3209 4238 2522 4238 \nz\nM 3406 84 \nL 4238 -825 \nL 3475 -825 \nL 2784 -78 \nQ 2681 -84 2626 -87 \nQ 2572 -91 2522 -91 \nQ 1538 -91 948 567 \nQ 359 1225 359 2328 \nQ 359 3434 948 4092 \nQ 1538 4750 2522 4750 \nQ 3503 4750 4090 4092 \nQ 4678 3434 4678 2328 \nQ 4678 1516 4351 937 \nQ 4025 359 3406 84 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-75\" d=\"M 544 1381 \nL 544 3500 \nL 1119 3500 \nL 1119 1403 \nQ 1119 906 1312 657 \nQ 1506 409 1894 409 \nQ 2359 409 2629 706 \nQ 2900 1003 2900 1516 \nL 2900 3500 \nL 3475 3500 \nL 3475 0 \nL 2900 0 \nL 2900 538 \nQ 2691 219 2414 64 \nQ 2138 -91 1772 -91 \nQ 1169 -91 856 284 \nQ 544 659 544 1381 \nz\nM 1991 3584 \nL 1991 3584 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-72\" d=\"M 2631 2963 \nQ 2534 3019 2420 3045 \nQ 2306 3072 2169 3072 \nQ 1681 3072 1420 2755 \nQ 1159 2438 1159 1844 \nL 1159 0 \nL 581 0 \nL 581 3500 \nL 1159 3500 \nL 1159 2956 \nQ 1341 3275 1631 3429 \nQ 1922 3584 2338 3584 \nQ 2397 3584 2469 3576 \nQ 2541 3569 2628 3553 \nL 2631 2963 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-69\" d=\"M 603 3500 \nL 1178 3500 \nL 1178 0 \nL 603 0 \nL 603 3500 \nz\nM 603 4863 \nL 1178 4863 \nL 1178 4134 \nL 603 4134 \nL 603 4863 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-51\"/>\n      <use xlink:href=\"#DejaVuSans-75\" transform=\"translate(78.710938 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(142.089844 0)\"/>\n      <use xlink:href=\"#DejaVuSans-72\" transform=\"translate(203.613281 0)\"/>\n      <use xlink:href=\"#DejaVuSans-69\" transform=\"translate(244.726562 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(272.509766 0)\"/>\n      <use xlink:href=\"#DejaVuSans-73\" transform=\"translate(334.033203 0)\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"patch_3\">\n    <path d=\"M 34.240625 59.94 \nL 34.240625 37.62 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_4\">\n    <path d=\"M 145.840625 59.94 \nL 145.840625 37.62 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_5\">\n    <path d=\"M 34.240625 59.94 \nL 145.840625 59.94 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_6\">\n    <path d=\"M 34.240625 37.62 \nL 145.840625 37.62 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n  </g>\n  <g id=\"axes_2\">\n   <g id=\"patch_7\">\n    <path d=\"M 152.815625 90.36 \nL 156.973625 90.36 \nL 156.973625 7.2 \nL 152.815625 7.2 \nz\n\" style=\"fill: #ffffff\"/>\n   </g>\n   <image xlink:href=\"data:image/png;base64,\niVBORw0KGgoAAAANSUhEUgAAAAYAAAB0CAYAAACmJkOCAAAAyElEQVR4nM2WSwrEMAxDXcj9zzqbblp/egG/gEII02WEJMt2Qq+6f2XNN6yyO7dhGQQQo1qLFY9ij30AB5Rz6B7uIqNQKkiKAUwuAyRV8e6SmjTxgAcy5An+qTkD+vrQNVgw15sI5qWXi909kLwCpTCgDPgB83dfuXpVFDAD3vakR18H3KEqWJKJVCBAyTEgSukBw2WP9thsJP0zOAEoFcjoz22EUQ5kyFVhwH5BZx4MkBQ2Ee7mhAEjX2E88sxRCluyYg4MquoD0XoQUGMmzEcAAAAASUVORK5CYII=\" id=\"image57b22d4f91\" transform=\"scale(1 -1) translate(0 -83.52)\" x=\"152.64\" y=\"-6.48\" width=\"4.32\" height=\"83.52\"/>\n   <g id=\"matplotlib.axis_3\"/>\n   <g id=\"matplotlib.axis_4\">\n    <g id=\"ytick_3\">\n     <g id=\"line2d_5\">\n      <defs>\n       <path id=\"m6e68e98c18\" d=\"M 0 0 \nL 3.5 0 \n\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </defs>\n      <g>\n       <use xlink:href=\"#m6e68e98c18\" x=\"156.973625\" y=\"90.36\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_7\">\n      <!-- 0.0 -->\n      <g transform=\"translate(163.973625 94.159219) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-2e\" d=\"M 684 794 \nL 1344 794 \nL 1344 0 \nL 684 0 \nL 684 794 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-30\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(95.410156 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_4\">\n     <g id=\"line2d_6\">\n      <g>\n       <use xlink:href=\"#m6e68e98c18\" x=\"156.973625\" y=\"61.676682\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_8\">\n      <!-- 0.2 -->\n      <g transform=\"translate(163.973625 65.475901) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-32\" d=\"M 1228 531 \nL 3431 531 \nL 3431 0 \nL 469 0 \nL 469 531 \nQ 828 903 1448 1529 \nQ 2069 2156 2228 2338 \nQ 2531 2678 2651 2914 \nQ 2772 3150 2772 3378 \nQ 2772 3750 2511 3984 \nQ 2250 4219 1831 4219 \nQ 1534 4219 1204 4116 \nQ 875 4013 500 3803 \nL 500 4441 \nQ 881 4594 1212 4672 \nQ 1544 4750 1819 4750 \nQ 2544 4750 2975 4387 \nQ 3406 4025 3406 3419 \nQ 3406 3131 3298 2873 \nQ 3191 2616 2906 2266 \nQ 2828 2175 2409 1742 \nQ 1991 1309 1228 531 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-30\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-32\" transform=\"translate(95.410156 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_5\">\n     <g id=\"line2d_7\">\n      <g>\n       <use xlink:href=\"#m6e68e98c18\" x=\"156.973625\" y=\"32.993365\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_9\">\n      <!-- 0.4 -->\n      <g transform=\"translate(163.973625 36.792584) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-34\" d=\"M 2419 4116 \nL 825 1625 \nL 2419 1625 \nL 2419 4116 \nz\nM 2253 4666 \nL 3047 4666 \nL 3047 1625 \nL 3713 1625 \nL 3713 1100 \nL 3047 1100 \nL 3047 0 \nL 2419 0 \nL 2419 1100 \nL 313 1100 \nL 313 1709 \nL 2253 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-30\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-34\" transform=\"translate(95.410156 0)\"/>\n      </g>\n     </g>\n    </g>\n   </g>\n   <g id=\"LineCollection_1\"/>\n   <g id=\"patch_8\">\n    <path d=\"M 152.815625 90.36 \nL 154.894625 90.36 \nL 156.973625 90.36 \nL 156.973625 7.2 \nL 154.894625 7.2 \nL 152.815625 7.2 \nL 152.815625 90.36 \nz\n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n  </g>\n </g>\n <defs>\n  <clipPath id=\"pb009b26535\">\n   <rect x=\"34.240625\" y=\"37.62\" width=\"111.6\" height=\"22.32\"/>\n  </clipPath>\n </defs>\n</svg>\n"
     },
     "metadata": {}
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=attention-scoring-md-14 fw=all -->\n",
    "\n",
    "## From Alignment to Attention\n",
    "\n",
    "Where did all of this come from? Not from databases, but from machine\n",
    "translation. Around 2014, the leading neural approach encoded a source\n",
    "sentence with an RNN into a single fixed-size state vector and decoded the\n",
    "translation from that vector with a second RNN (we build such\n",
    "encoder--decoder models in full in that section). The design has\n",
    "the flaw this chapter opened with: one fixed-size vector must carry an\n",
    "entire sentence, and for long sentences it cannot. Translation quality\n",
    "degraded visibly with sentence length. @Graves.2013 had faced a\n",
    "version of this problem when generating handwriting from text, and solved\n",
    "it with a differentiable model that *aligned* each output pen stroke with a\n",
    "position in the source text—though with the constraint that the alignment\n",
    "could only move forward, an assumption borrowed from decoding in speech\n",
    "recognition [@rabiner1993fundamentals].\n",
    "\n",
    "@Bahdanau.Cho.Bengio.2014 removed the constraint. Their translation\n",
    "model kept the two RNNs but gave the decoder a new capability at every\n",
    "step: use the current decoder state as a *query* against all encoder\n",
    "states, which serve as keys and values, and feed the resulting weighted\n",
    "summary—a fresh one per output token—into the next prediction. The paper's\n",
    "title called the idea \"jointly learning to align and translate\", and the\n",
    "learned weights behaved exactly like the soft alignments of classical\n",
    "statistical translation: mostly monotone along the diagonal, with clean\n",
    "departures where the two languages order words differently, as sketched in\n",
    "the figure. Nothing forced the model to align; the behavior\n",
    "emerged from training. This is the attention mechanism of\n",
    "the equation in its original habitat, and its impact\n",
    "reached far beyond translation.\n",
    "\n",
    "![Soft alignment between an English sentence and its French translation, in the style of a learned attention map (darker cells indicate larger weight; schematic). The alignment is mostly monotone, but \"black cat\" maps to \"chat noir\" with the order reversed, and both \"était\" and \"assis\" draw on \"sat\".](../img/mdl-attention-alignment.svg)\n",
    "\n",
    "One detail differed from the scoring function we settled on above. The\n",
    "decoder state and the encoder states were vectors of different sizes, so\n",
    "instead of a dot product, @Bahdanau.Cho.Bengio.2014 scored with a\n",
    "small one-hidden-layer MLP, now known as *additive attention*:\n",
    "\n",
    "$$a(\\mathbf{q}, \\mathbf{k}) = \\mathbf{w}_v^\\top \\tanh(\\mathbf{W}_q \\mathbf{q} + \\mathbf{W}_k \\mathbf{k}) \\in \\mathbb{R},$$\n",
    "\n",
    "where $\\mathbf{W}_q \\in \\mathbb{R}^{h \\times q}$, $\\mathbf{W}_k \\in\n",
    "\\mathbb{R}^{h \\times k}$, and $\\mathbf{w}_v \\in \\mathbb{R}^{h}$ are learned.\n",
    "The two projections embed queries and keys into a shared $h$-dimensional\n",
    "space, and $\\mathbf{w}_v$ reads a score off the sum. It takes only a few\n",
    "lines to compute—here scoring three $20$-dimensional queries against six\n",
    "$2$-dimensional keys, dimensions no dot product could pair up:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "attention-scoring-from-alignment-to-attention"
   },
   "source": [
    "key_q, key_k, k1, k2, k3 = jax.random.split(jax.random.key(0), 5)\n",
    "queries = jax.random.normal(key_q, (3, 20))\n",
    "keys = jax.random.normal(key_k, (6, 2))\n",
    "num_hiddens = 8\n",
    "W_q = jax.random.normal(k1, (num_hiddens, 20)) / math.sqrt(20)\n",
    "W_k = jax.random.normal(k2, (num_hiddens, 2)) / math.sqrt(2)\n",
    "w_v = jax.random.normal(k3, (num_hiddens,)) / math.sqrt(num_hiddens)\n",
    "features = jnp.tanh((queries @ W_q.T)[:, None, :]\n",
    "                    + (keys @ W_k.T)[None, :, :])\n",
    "scores = features @ w_v\n",
    "d2l.check_shape(scores, (3, 6))\n",
    "jax.nn.softmax(scores, axis=-1)"
   ],
   "execution_count": 13,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array([[0.13390937, 0.3772766 , 0.17127123, 0.1027277 , 0.06276076,\n        0.15205443],\n       [0.19218282, 0.31524566, 0.14511336, 0.08953383, 0.07662669,\n        0.18129766],\n       [0.15669951, 0.3378386 , 0.15281527, 0.11713327, 0.09428051,\n        0.1412329 ]], dtype=float32)"
     },
     "metadata": {},
     "execution_count": 13
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=attention-scoring-md-15 fw=all -->\n",
    "\n",
    "Additive attention held its own for a few years, but the outcome was\n",
    "decided by hardware: a dot product between all queries and all keys is a\n",
    "single matrix multiplication, the one operation accelerators are built\n",
    "around, while the additive score requires materializing an\n",
    "$n \\times m \\times h$ tensor of hidden activations. When a learned metric\n",
    "is wanted, projecting queries and keys with learned matrices *before* a dot\n",
    "product achieves it at matmul speed—which is precisely the form attention\n",
    "takes inside the Transformer, whose authors then discarded the RNN\n",
    "scaffolding altogether and kept attention as the only mechanism relating\n",
    "sequence positions [@Vaswani.Shazeer.Parmar.ea.2017]. The next\n",
    "sections follow that road: first attention with multiple heads, then what\n",
    "replaces the RNN's notion of position.\n",
    "\n",
    "## Summary\n",
    "\n",
    "Softmax normalization turns any scoring function into valid attention\n",
    "weights, and the scoring function of choice is the scaled dot product\n",
    "the equation: it is a single batched matrix\n",
    "multiplication, and the $1/\\sqrt{d}$ factor holds the score variance at $1$\n",
    "so that the softmax neither saturates nor starves its own gradient as the\n",
    "dimension grows—an effect we measured directly. Masking makes the same\n",
    "batched computation respect variable sequence lengths and causal\n",
    "structure: overwrite invalid scores with the most negative finite value of\n",
    "the dtype (not a hard-coded constant, which breaks in half precision)\n",
    "before the softmax; arbitrary boolean requirements — padding, causality,\n",
    "structure — compose into one mask by logical AND. `DotProductAttention` packages scoring, masking,\n",
    "dropout on the weights, and value pooling in a dozen lines that the rest of\n",
    "this book reuses. Additive attention, the original scoring function of the\n",
    "translation models that started the field, survives as history and as a\n",
    "reminder that attention weights are learned soft alignments.\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Implement distance-based attention by modifying the\n",
    "   `DotProductAttention` code. You only need the squared norms of the keys\n",
    "   $\\|\\mathbf{k}_i\\|^2$ for an efficient implementation.\n",
    "1. Modify dot-product attention to allow for queries and keys of different\n",
    "   dimensionalities by employing a matrix $\\mathbf{M}$ to adjust dimensions,\n",
    "   scoring with $\\mathbf{q}^\\top \\mathbf{M} \\mathbf{k}$.\n",
    "1. How does the computational cost of the equation scale with\n",
    "   the dimensionality of keys, queries, and values, and with their number?\n",
    "   What about the memory bandwidth requirements?\n",
    "1. Derive the softmax Jacobian $\\mathrm{diag}(\\boldsymbol{\\alpha}) -\n",
    "   \\boldsymbol{\\alpha}\\boldsymbol{\\alpha}^\\top$ and verify it numerically\n",
    "   against automatic differentiation on a random score vector. Compute its\n",
    "   Frobenius norm when $\\boldsymbol{\\alpha}$ is one-hot and when it is\n",
    "   uniform over $m$ keys. Which regime does the saturation experiment\n",
    "   approach as $d$ grows without scaling?\n",
    "1. What does `masked_softmax` return for a query whose valid length is $0$?\n",
    "   Compare the behavior of masking with the most negative finite value\n",
    "   against masking with literal $-\\infty$. In what situations can a fully\n",
    "   masked query arise in practice, and what would you do about it?\n",
    "1. Count the parameters and floating-point operations needed to score $n$\n",
    "   queries against $m$ keys with additive attention (hidden size $h$) and\n",
    "   with scaled dot-product attention (shared dimension $d$). Implement a\n",
    "   batched version of the additive score and time both variants for\n",
    "   $d = h = 64$ and $d = h = 256$. Which is faster on your hardware, and\n",
    "   why?"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
