{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=transformer-block-md-1 fw=all -->\n",
    "\n",
    "# The Transformer Block\n",
    "\n",
    "The previous chapter built attention: a layer that lets every position read\n",
    "from every other, weighted by learned relevance. Attention alone, however, is\n",
    "not what anyone deploys. Deployed models stack a composite unit — the\n",
    "*transformer block* — in which attention supplies communication between\n",
    "positions and a position-wise feed-forward network supplies computation\n",
    "within each one, both writing their results into a shared residual stream.\n",
    "This section builds that unit. We settle the two design decisions that\n",
    "separate a 2017 block from a 2026 one: where the normalization layers go\n",
    "(with an experiment at initialization that shows why nearly every modern\n",
    "model moved them), and what the feed-forward network computes (with a\n",
    "matched-budget race between the classic MLP and its gated successor). The\n",
    "product is a single configurable `TransformerBlock` class whose flags span\n",
    "a decade of architectures; the next section stacks it into a language model,\n",
    "and the rest of the chapter never builds another block."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "transformer-block-the-transformer-block"
   },
   "source": [
    "%matplotlib inline\n",
    "from d2l import torch as d2l\n",
    "import time\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=transformer-block-md-2 fw=all -->\n",
    "\n",
    "## The Anatomy of a Block\n",
    "\n",
    "A transformer block is two sublayers wired to one highway.\n",
    "that section introduced the *residual stream* view:\n",
    "each token carries a running vector $\\mathbf{x} \\in \\mathbb{R}^d$ from layer\n",
    "to layer, and sublayers do not replace it — they read from it, compute, and\n",
    "*add* their result back. The block's two sublayers divide the labor:\n",
    "\n",
    "- **Attention** is the only place where positions interact. Each token\n",
    "  queries the others as in that section and adds the\n",
    "  retrieved mixture to its stream. Everything the model knows about\n",
    "  *context* enters here.\n",
    "- The **feed-forward network** (FFN) acts on each position separately: the\n",
    "  same two- or three-matrix MLP, applied token by token, with no\n",
    "  interaction between positions. Everything the model computes *from* a\n",
    "  token's accumulated state — and, as we will see, most of its parameters —\n",
    "  lives here.\n",
    "\n",
    "Writing $\\mathrm{Attn}$ for multi-head self-attention and $\\mathrm{Norm}$\n",
    "for a normalization layer, the modern (*pre-norm*) block computes\n",
    "\n",
    "$$\n",
    "\\mathbf{h} = \\mathbf{x} + \\mathrm{Attn}(\\mathrm{Norm}_1(\\mathbf{x})), \\qquad\n",
    "\\mathbf{y} = \\mathbf{h} + \\mathrm{FFN}(\\mathrm{Norm}_2(\\mathbf{h})).\n",
    "$$\n",
    "\n",
    "Both sublayers preserve the shape $(n, d)$: the block maps a sequence of\n",
    "$d$-dimensional vectors to a sequence of the same length and width, which\n",
    "is what makes stacking dozens of copies possible. Note also what the block\n",
    "does *not* contain: positions. Attention is permutation equivariant\n",
    "(that section), and the block does nothing to\n",
    "change that — positional information is the model's job, not the block's,\n",
    "a division of responsibilities the next section exploits.\n",
    "\n",
    "the figure shows the wiring, in the two normalization\n",
    "arrangements whose difference is our first order of business.\n",
    "\n",
    "![The transformer block: attention and a position-wise FFN read from the residual stream and add their results back. Left, the 2017 arrangement — normalization sits on the stream itself, after each addition. Right, the modern arrangement — normalization moves onto each branch, and the stream runs uninterrupted from input to output.](../img/mdl-transformers-block-anatomy.svg)\n",
    "\n",
    "## Where the Normalization Goes\n",
    "\n",
    "Residual connections and normalization entered the block together in 2017\n",
    "[@Vaswani.Shazeer.Parmar.ea.2017]: without some control on scale, the\n",
    "sum of dozens of sublayer outputs drifts, and training deep stacks becomes\n",
    "a tuning nightmare. The normalizer of choice is *layer normalization*\n",
    "[@Ba.Kiros.Hinton.2016], which standardizes each token's vector across\n",
    "its $d$ features — unlike the batch normalization of\n",
    "that section, it involves no batch statistics, so it behaves\n",
    "identically in training and inference and is indifferent to sequence\n",
    "length. What the 2017 paper fixed, and what a few years of painful\n",
    "experience revised, is *where* it goes.\n",
    "\n",
    "### Two Arrangements\n",
    "\n",
    "The original block normalizes the stream *after* each addition\n",
    "(\"post-LN\"):\n",
    "\n",
    "$$\n",
    "\\mathbf{h} = \\mathrm{Norm}_1(\\mathbf{x} + \\mathrm{Attn}(\\mathbf{x})), \\qquad\n",
    "\\mathbf{y} = \\mathrm{Norm}_2(\\mathbf{h} + \\mathrm{FFN}(\\mathbf{h})).\n",
    "$$\n",
    "\n",
    "The modern block the equation normalizes each branch's *input*\n",
    "(\"pre-LN\") and leaves the stream alone. In code the two differ by two\n",
    "lines. The minimal block below takes the arrangement as a flag; its\n",
    "attention is the `d2l.MultiHeadAttention` of\n",
    "that section and its FFN is the classic\n",
    "two-matrix MLP (we upgrade both ingredients later in this section)."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "transformer-block-two-arrangements"
   },
   "source": [
    "class MiniBlock(nn.Module):\n",
    "    \"\"\"Attention + MLP on a residual stream; norm placement as a flag.\"\"\"\n",
    "    def __init__(self, num_hiddens, pre_norm):\n",
    "        super().__init__()\n",
    "        self.pre_norm = pre_norm\n",
    "        self.attention = d2l.MultiHeadAttention(num_hiddens, 4, dropout=0)\n",
    "        self.W_1 = nn.Linear(num_hiddens, 4 * num_hiddens)\n",
    "        self.W_2 = nn.Linear(4 * num_hiddens, num_hiddens)\n",
    "        self.norm1 = nn.LayerNorm(num_hiddens)\n",
    "        self.norm2 = nn.LayerNorm(num_hiddens)\n",
    "\n",
    "    def ffn(self, X):\n",
    "        return self.W_2(F.gelu(self.W_1(X)))\n",
    "\n",
    "    def forward(self, X):\n",
    "        if self.pre_norm:\n",
    "            Y = self.norm1(X)\n",
    "            X = X + self.attention(Y, Y, Y, None)\n",
    "            return X + self.ffn(self.norm2(X))\n",
    "        X = self.norm1(X + self.attention(X, X, X, None))\n",
    "        return self.norm2(X + self.ffn(X))"
   ],
   "execution_count": 2,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=transformer-block-md-3 fw=all -->\n",
    "\n",
    "### Signal Propagation at Initialization\n",
    "\n",
    "Which arrangement is better? We can get a long way without training\n",
    "anything. A freshly initialized network is a fixed random function, so we\n",
    "can *measure* how it treats signals: stack $N$ blocks, push a random\n",
    "sequence through, and watch what survives — deterministically, in seconds,\n",
    "with no seed lottery. We initialize every weight matrix as the original\n",
    "transformer did (Xavier initialization) and track three quantities: the\n",
    "scale of the residual stream, how *distinct* the tokens remain from one\n",
    "another, and how much gradient each block's attention receives from a\n",
    "generic loss at the top.\n",
    "\n",
    "For the distinctness measure we use the cross-token spread\n",
    "$\\|\\mathbf{H} - \\bar{\\mathbf{H}}\\| / \\|\\mathbf{H}\\|$, where\n",
    "$\\bar{\\mathbf{H}}$ copies the average token to every position: it is $1$\n",
    "when tokens are unrelated and $0$ when they have collapsed onto a single\n",
    "point. One caution learned the hard way: after a final normalization layer\n",
    "the naive probe loss $\\|\\mathbf{H}\\|^2$ is constant by construction, so we\n",
    "probe with a generic linear readout $\\langle \\mathbf{R}, \\mathbf{H}\n",
    "\\rangle$ for fixed random $\\mathbf{R}$ instead."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "transformer-block-signal-propagation-at-initialization"
   },
   "source": [
    "def signal_stats(pre_norm, num_blks=32, num_hiddens=256):\n",
    "    torch.manual_seed(0)\n",
    "    blks = nn.ModuleList([MiniBlock(num_hiddens, pre_norm)\n",
    "                          for _ in range(num_blks)])\n",
    "    torch.manual_seed(1)\n",
    "    X = torch.randn(2, 64, num_hiddens)\n",
    "    H = X\n",
    "    for blk in blks:\n",
    "        H = blk(H)                      # materialize the lazy layers\n",
    "    for p in blks.parameters():\n",
    "        if p.ndim == 2:                 # the 2017 initialization\n",
    "            nn.init.xavier_uniform_(p)\n",
    "    H, spreads, rms = X, [], []\n",
    "    for blk in blks:\n",
    "        H = blk(H)\n",
    "        spreads.append(\n",
    "            ((H - H.mean(1, keepdim=True)).norm() / H.norm()).item())\n",
    "        rms.append(H.pow(2).mean().sqrt().item())\n",
    "    (H * torch.randn_like(H)).mean().backward()  # generic readout loss\n",
    "    grads = [blk.attention.W_q.weight.grad.norm().item() for blk in blks]\n",
    "    proj = {n: getattr(blks[-1].attention, n).weight.grad.norm().item()\n",
    "            for n in ('W_q', 'W_k', 'W_v', 'W_o')}\n",
    "    return spreads, rms, grads, proj\n",
    "\n",
    "stats = {pre: signal_stats(pre) for pre in (False, True)}\n",
    "for pre, name in ((False, 'post-LN'), (True, 'pre-LN')):\n",
    "    spreads, rms, grads, proj = stats[pre]\n",
    "    print(f'{name:>8}: stream RMS at k=32: {rms[-1]:5.2f},  '\n",
    "          'token spread at k=8,16,32: '\n",
    "          + ' '.join(f'{spreads[k-1]:8.1e}' for k in (8, 16, 32)))\n",
    "    print('          grad norms at k=32: '\n",
    "          + '  '.join(f'{n} {g:7.1e}' for n, g in proj.items()))\n",
    "d2l.plot(torch.arange(1, 33), [stats[False][2], stats[True][2]],\n",
    "         'block index k', 'query projection gradient norm',\n",
    "         legend=['post-LN', 'pre-LN'], yscale='log')"
   ],
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      " post-LN: stream RMS at k=32:  1.00,  token spread at k=8,16,32:  3.7e-01  3.6e-02  1.2e-04\n          grad norms at k=32: W_q 2.9e-09  W_k 2.6e-09  W_v 5.8e-02  W_o 5.9e-02\n  pre-LN: stream RMS at k=32:  5.25,  token spread at k=8,16,32:  7.4e-01  5.2e-01  3.7e-01\n          grad norms at k=32: W_q 1.9e-03  W_k 1.9e-03  W_v 8.1e-02  W_o 8.4e-02\n"
     ]
    },
    {
     "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=\"253.878125pt\" height=\"183.35625pt\" viewBox=\"0 0 253.878125 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-19T14:04:50.366023</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 253.878125 183.35625 \nL 253.878125 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 51.378125 145.8 \nL 246.678125 145.8 \nL 246.678125 7.2 \nL 51.378125 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 54.528125 145.8 \nL 54.528125 7.2 \n\" clip-path=\"url(#pad48f39ba2)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_2\">\n      <defs>\n       <path id=\"mbbb72a4c8b\" d=\"M 0 0 \nL 0 3.5 \n\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </defs>\n      <g>\n       <use xlink:href=\"#mbbb72a4c8b\" x=\"54.528125\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_1\">\n      <!-- 0 -->\n      <g transform=\"translate(51.346875 160.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_3\">\n      <path d=\"M 111.800852 145.8 \nL 111.800852 7.2 \n\" clip-path=\"url(#pad48f39ba2)\" 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=\"#mbbb72a4c8b\" x=\"111.800852\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_2\">\n      <!-- 10 -->\n      <g transform=\"translate(105.438352 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       </defs>\n       <use xlink:href=\"#DejaVuSans-31\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_3\">\n     <g id=\"line2d_5\">\n      <path d=\"M 169.07358 145.8 \nL 169.07358 7.2 \n\" clip-path=\"url(#pad48f39ba2)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_6\">\n      <g>\n       <use xlink:href=\"#mbbb72a4c8b\" x=\"169.07358\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_3\">\n      <!-- 20 -->\n      <g transform=\"translate(162.71108 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-32\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_4\">\n     <g id=\"line2d_7\">\n      <path d=\"M 226.346307 145.8 \nL 226.346307 7.2 \n\" clip-path=\"url(#pad48f39ba2)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_8\">\n      <g>\n       <use xlink:href=\"#mbbb72a4c8b\" x=\"226.346307\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_4\">\n      <!-- 30 -->\n      <g transform=\"translate(219.983807 160.398438) 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       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_5\">\n     <!-- block index k -->\n     <g transform=\"translate(116.007031 174.076563) 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-6c\" d=\"M 603 4863 \nL 1178 4863 \nL 1178 0 \nL 603 0 \nL 603 4863 \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-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-6b\" d=\"M 581 4863 \nL 1159 4863 \nL 1159 1991 \nL 2875 3500 \nL 3609 3500 \nL 1753 1863 \nL 3688 0 \nL 2938 0 \nL 1159 1709 \nL 1159 0 \nL 581 0 \nL 581 4863 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-20\" 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-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-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-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-78\" d=\"M 3513 3500 \nL 2247 1797 \nL 3578 0 \nL 2900 0 \nL 1881 1375 \nL 863 0 \nL 184 0 \nL 1544 1831 \nL 300 3500 \nL 978 3500 \nL 1906 2253 \nL 2834 3500 \nL 3513 3500 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-62\"/>\n      <use xlink:href=\"#DejaVuSans-6c\" transform=\"translate(63.476562 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6f\" transform=\"translate(91.259766 0)\"/>\n      <use xlink:href=\"#DejaVuSans-63\" transform=\"translate(152.441406 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6b\" transform=\"translate(207.421875 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(265.332031 0)\"/>\n      <use xlink:href=\"#DejaVuSans-69\" transform=\"translate(297.119141 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6e\" transform=\"translate(324.902344 0)\"/>\n      <use xlink:href=\"#DejaVuSans-64\" transform=\"translate(388.28125 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(451.757812 0)\"/>\n      <use xlink:href=\"#DejaVuSans-78\" transform=\"translate(511.53125 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(570.710938 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6b\" transform=\"translate(602.498047 0)\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"matplotlib.axis_2\">\n    <g id=\"ytick_1\">\n     <g id=\"line2d_9\">\n      <path d=\"M 51.378125 128.268868 \nL 246.678125 128.268868 \n\" clip-path=\"url(#pad48f39ba2)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_10\">\n      <defs>\n       <path id=\"mfcd78c0be9\" d=\"M 0 0 \nL -3.5 0 \n\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </defs>\n      <g>\n       <use xlink:href=\"#mfcd78c0be9\" x=\"51.378125\" y=\"128.268868\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_6\">\n      <!-- $\\mathdefault{10^{-8}}$ -->\n      <g transform=\"translate(20.878125 132.068087) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-2212\" d=\"M 678 2272 \nL 4684 2272 \nL 4684 1741 \nL 678 1741 \nL 678 2272 \nz\n\" transform=\"scale(0.015625)\"/>\n        <path id=\"DejaVuSans-38\" d=\"M 2034 2216 \nQ 1584 2216 1326 1975 \nQ 1069 1734 1069 1313 \nQ 1069 891 1326 650 \nQ 1584 409 2034 409 \nQ 2484 409 2743 651 \nQ 3003 894 3003 1313 \nQ 3003 1734 2745 1975 \nQ 2488 2216 2034 2216 \nz\nM 1403 2484 \nQ 997 2584 770 2862 \nQ 544 3141 544 3541 \nQ 544 4100 942 4425 \nQ 1341 4750 2034 4750 \nQ 2731 4750 3128 4425 \nQ 3525 4100 3525 3541 \nQ 3525 3141 3298 2862 \nQ 3072 2584 2669 2484 \nQ 3125 2378 3379 2068 \nQ 3634 1759 3634 1313 \nQ 3634 634 3220 271 \nQ 2806 -91 2034 -91 \nQ 1263 -91 848 271 \nQ 434 634 434 1313 \nQ 434 1759 690 2068 \nQ 947 2378 1403 2484 \nz\nM 1172 3481 \nQ 1172 3119 1398 2916 \nQ 1625 2713 2034 2713 \nQ 2441 2713 2670 2916 \nQ 2900 3119 2900 3481 \nQ 2900 3844 2670 4047 \nQ 2441 4250 2034 4250 \nQ 1625 4250 1398 4047 \nQ 1172 3844 1172 3481 \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-2212\" transform=\"translate(128.203125 39.046875) scale(0.7)\"/>\n       <use xlink:href=\"#DejaVuSans-38\" transform=\"translate(186.855469 39.046875) scale(0.7)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_2\">\n     <g id=\"line2d_11\">\n      <path d=\"M 51.378125 94.360028 \nL 246.678125 94.360028 \n\" clip-path=\"url(#pad48f39ba2)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_12\">\n      <g>\n       <use xlink:href=\"#mfcd78c0be9\" x=\"51.378125\" y=\"94.360028\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_7\">\n      <!-- $\\mathdefault{10^{-6}}$ -->\n      <g transform=\"translate(20.878125 98.159247) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-36\" d=\"M 2113 2584 \nQ 1688 2584 1439 2293 \nQ 1191 2003 1191 1497 \nQ 1191 994 1439 701 \nQ 1688 409 2113 409 \nQ 2538 409 2786 701 \nQ 3034 994 3034 1497 \nQ 3034 2003 2786 2293 \nQ 2538 2584 2113 2584 \nz\nM 3366 4563 \nL 3366 3988 \nQ 3128 4100 2886 4159 \nQ 2644 4219 2406 4219 \nQ 1781 4219 1451 3797 \nQ 1122 3375 1075 2522 \nQ 1259 2794 1537 2939 \nQ 1816 3084 2150 3084 \nQ 2853 3084 3261 2657 \nQ 3669 2231 3669 1497 \nQ 3669 778 3244 343 \nQ 2819 -91 2113 -91 \nQ 1303 -91 875 529 \nQ 447 1150 447 2328 \nQ 447 3434 972 4092 \nQ 1497 4750 2381 4750 \nQ 2619 4750 2861 4703 \nQ 3103 4656 3366 4563 \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-2212\" transform=\"translate(128.203125 39.046875) scale(0.7)\"/>\n       <use xlink:href=\"#DejaVuSans-36\" transform=\"translate(186.855469 39.046875) scale(0.7)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_3\">\n     <g id=\"line2d_13\">\n      <path d=\"M 51.378125 60.451188 \nL 246.678125 60.451188 \n\" clip-path=\"url(#pad48f39ba2)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_14\">\n      <g>\n       <use xlink:href=\"#mfcd78c0be9\" x=\"51.378125\" y=\"60.451188\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_8\">\n      <!-- $\\mathdefault{10^{-4}}$ -->\n      <g transform=\"translate(20.878125 64.250407) 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-31\" transform=\"translate(0 0.684375)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0.684375)\"/>\n       <use xlink:href=\"#DejaVuSans-2212\" transform=\"translate(128.203125 38.965625) scale(0.7)\"/>\n       <use xlink:href=\"#DejaVuSans-34\" transform=\"translate(186.855469 38.965625) scale(0.7)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_4\">\n     <g id=\"line2d_15\">\n      <path d=\"M 51.378125 26.542348 \nL 246.678125 26.542348 \n\" clip-path=\"url(#pad48f39ba2)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_16\">\n      <g>\n       <use xlink:href=\"#mfcd78c0be9\" x=\"51.378125\" y=\"26.542348\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_9\">\n      <!-- $\\mathdefault{10^{-2}}$ -->\n      <g transform=\"translate(20.878125 30.341567) scale(0.1 -0.1)\">\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-2212\" transform=\"translate(128.203125 39.046875) scale(0.7)\"/>\n       <use xlink:href=\"#DejaVuSans-32\" transform=\"translate(186.855469 39.046875) scale(0.7)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_10\">\n     <!-- query projection gradient norm -->\n     <g transform=\"translate(14.798438 154.799219) rotate(-90) scale(0.1 -0.1)\">\n      <defs>\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       <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-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-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-6a\" d=\"M 603 3500 \nL 1178 3500 \nL 1178 -63 \nQ 1178 -731 923 -1031 \nQ 669 -1331 103 -1331 \nL -116 -1331 \nL -116 -844 \nL 38 -844 \nQ 366 -844 484 -692 \nQ 603 -541 603 -63 \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-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-67\" d=\"M 2906 1791 \nQ 2906 2416 2648 2759 \nQ 2391 3103 1925 3103 \nQ 1463 3103 1205 2759 \nQ 947 2416 947 1791 \nQ 947 1169 1205 825 \nQ 1463 481 1925 481 \nQ 2391 481 2648 825 \nQ 2906 1169 2906 1791 \nz\nM 3481 434 \nQ 3481 -459 3084 -895 \nQ 2688 -1331 1869 -1331 \nQ 1566 -1331 1297 -1286 \nQ 1028 -1241 775 -1147 \nL 775 -588 \nQ 1028 -725 1275 -790 \nQ 1522 -856 1778 -856 \nQ 2344 -856 2625 -561 \nQ 2906 -266 2906 331 \nL 2906 616 \nQ 2728 306 2450 153 \nQ 2172 0 1784 0 \nQ 1141 0 747 490 \nQ 353 981 353 1791 \nQ 353 2603 747 3093 \nQ 1141 3584 1784 3584 \nQ 2172 3584 2450 3431 \nQ 2728 3278 2906 2969 \nL 2906 3500 \nL 3481 3500 \nL 3481 434 \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-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      </defs>\n      <use xlink:href=\"#DejaVuSans-71\"/>\n      <use xlink:href=\"#DejaVuSans-75\" transform=\"translate(63.476562 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(126.855469 0)\"/>\n      <use xlink:href=\"#DejaVuSans-72\" transform=\"translate(188.378906 0)\"/>\n      <use xlink:href=\"#DejaVuSans-79\" transform=\"translate(229.492188 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(288.671875 0)\"/>\n      <use xlink:href=\"#DejaVuSans-70\" transform=\"translate(320.458984 0)\"/>\n      <use xlink:href=\"#DejaVuSans-72\" transform=\"translate(383.935547 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6f\" transform=\"translate(422.798828 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6a\" transform=\"translate(483.980469 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(511.763672 0)\"/>\n      <use xlink:href=\"#DejaVuSans-63\" transform=\"translate(573.287109 0)\"/>\n      <use xlink:href=\"#DejaVuSans-74\" transform=\"translate(628.267578 0)\"/>\n      <use xlink:href=\"#DejaVuSans-69\" transform=\"translate(667.476562 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6f\" transform=\"translate(695.259766 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6e\" transform=\"translate(756.441406 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(819.820312 0)\"/>\n      <use xlink:href=\"#DejaVuSans-67\" transform=\"translate(851.607422 0)\"/>\n      <use xlink:href=\"#DejaVuSans-72\" transform=\"translate(915.083984 0)\"/>\n      <use xlink:href=\"#DejaVuSans-61\" transform=\"translate(956.197266 0)\"/>\n      <use xlink:href=\"#DejaVuSans-64\" transform=\"translate(1017.476562 0)\"/>\n      <use xlink:href=\"#DejaVuSans-69\" transform=\"translate(1080.953125 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(1108.736328 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6e\" transform=\"translate(1170.259766 0)\"/>\n      <use xlink:href=\"#DejaVuSans-74\" transform=\"translate(1233.638672 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(1272.847656 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6e\" transform=\"translate(1304.634766 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6f\" transform=\"translate(1368.013672 0)\"/>\n      <use xlink:href=\"#DejaVuSans-72\" transform=\"translate(1429.195312 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6d\" transform=\"translate(1468.558594 0)\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"line2d_17\">\n    <path d=\"M 60.255398 27.856478 \nL 65.98267 29.098191 \nL 71.709943 29.823714 \nL 77.437216 30.541783 \nL 83.164489 29.764184 \nL 88.891761 32.07777 \nL 94.619034 35.305306 \nL 100.346307 39.573361 \nL 106.07358 44.119789 \nL 111.800852 45.705096 \nL 117.528125 51.387526 \nL 123.255398 54.295109 \nL 128.98267 59.424388 \nL 134.709943 65.069332 \nL 140.437216 67.769687 \nL 146.164489 72.921295 \nL 151.891761 76.830822 \nL 157.619034 81.153849 \nL 163.346307 86.121157 \nL 169.07358 89.597523 \nL 174.800852 94.205434 \nL 180.528125 98.859342 \nL 186.255398 103.328213 \nL 191.98267 109.394334 \nL 197.709943 115.486061 \nL 203.437216 120.725045 \nL 209.164489 124.323948 \nL 214.891761 129.065474 \nL 220.619034 134.643397 \nL 226.346307 137.352213 \nL 232.07358 139.5 \nL 237.800852 137.489551 \n\" clip-path=\"url(#pad48f39ba2)\" style=\"fill: none; stroke: #1f77b4; stroke-width: 1.5; stroke-linecap: square\"/>\n   </g>\n   <g id=\"line2d_18\">\n    <path d=\"M 60.255398 13.5 \nL 65.98267 14.068661 \nL 71.709943 15.240044 \nL 77.437216 16.185742 \nL 83.164489 17.558333 \nL 88.891761 18.640406 \nL 94.619034 19.981379 \nL 100.346307 20.537085 \nL 106.07358 21.470486 \nL 111.800852 23.240852 \nL 117.528125 24.737179 \nL 123.255398 25.190943 \nL 128.98267 27.826078 \nL 134.709943 27.089669 \nL 140.437216 27.832704 \nL 146.164489 29.26848 \nL 151.891761 30.291524 \nL 157.619034 31.034173 \nL 163.346307 31.216005 \nL 169.07358 31.850804 \nL 174.800852 32.020397 \nL 180.528125 33.62008 \nL 186.255398 33.630514 \nL 191.98267 33.181505 \nL 197.709943 34.162495 \nL 203.437216 34.870345 \nL 209.164489 34.964076 \nL 214.891761 36.387512 \nL 220.619034 36.977268 \nL 226.346307 37.604728 \nL 232.07358 37.088316 \nL 237.800852 38.601321 \n\" clip-path=\"url(#pad48f39ba2)\" 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 51.378125 145.8 \nL 51.378125 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 246.678125 145.8 \nL 246.678125 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 51.378125 145.8 \nL 246.678125 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 51.378125 7.2 \nL 246.678125 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 58.378125 140.8 \nL 128.635938 140.8 \nQ 130.635938 140.8 130.635938 138.8 \nL 130.635938 110.44375 \nQ 130.635938 108.44375 128.635938 108.44375 \nL 58.378125 108.44375 \nQ 56.378125 108.44375 56.378125 110.44375 \nL 56.378125 138.8 \nQ 56.378125 140.8 58.378125 140.8 \nz\n\" style=\"fill: #ffffff; opacity: 0.8; stroke: #cccccc; stroke-linejoin: miter\"/>\n    </g>\n    <g id=\"line2d_19\">\n     <path d=\"M 60.378125 116.542188 \nL 70.378125 116.542188 \nL 80.378125 116.542188 \n\" style=\"fill: none; stroke: #1f77b4; stroke-width: 1.5; stroke-linecap: square\"/>\n    </g>\n    <g id=\"text_11\">\n     <!-- post-LN -->\n     <g transform=\"translate(88.378125 120.042188) scale(0.1 -0.1)\">\n      <defs>\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-2d\" d=\"M 313 2009 \nL 1997 2009 \nL 1997 1497 \nL 313 1497 \nL 313 2009 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-4c\" d=\"M 628 4666 \nL 1259 4666 \nL 1259 531 \nL 3531 531 \nL 3531 0 \nL 628 0 \nL 628 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-4e\" d=\"M 628 4666 \nL 1478 4666 \nL 3547 763 \nL 3547 4666 \nL 4159 4666 \nL 4159 0 \nL 3309 0 \nL 1241 3903 \nL 1241 0 \nL 628 0 \nL 628 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-70\"/>\n      <use xlink:href=\"#DejaVuSans-6f\" transform=\"translate(63.476562 0)\"/>\n      <use xlink:href=\"#DejaVuSans-73\" transform=\"translate(124.658203 0)\"/>\n      <use xlink:href=\"#DejaVuSans-74\" transform=\"translate(176.757812 0)\"/>\n      <use xlink:href=\"#DejaVuSans-2d\" transform=\"translate(215.966797 0)\"/>\n      <use xlink:href=\"#DejaVuSans-4c\" transform=\"translate(252.050781 0)\"/>\n      <use xlink:href=\"#DejaVuSans-4e\" transform=\"translate(307.763672 0)\"/>\n     </g>\n    </g>\n    <g id=\"line2d_20\">\n     <path d=\"M 60.378125 131.220313 \nL 70.378125 131.220313 \nL 80.378125 131.220313 \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_12\">\n     <!-- pre-LN -->\n     <g transform=\"translate(88.378125 134.720313) scale(0.1 -0.1)\">\n      <use xlink:href=\"#DejaVuSans-70\"/>\n      <use xlink:href=\"#DejaVuSans-72\" transform=\"translate(63.476562 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(102.339844 0)\"/>\n      <use xlink:href=\"#DejaVuSans-2d\" transform=\"translate(163.863281 0)\"/>\n      <use xlink:href=\"#DejaVuSans-4c\" transform=\"translate(199.947266 0)\"/>\n      <use xlink:href=\"#DejaVuSans-4e\" transform=\"translate(255.660156 0)\"/>\n     </g>\n    </g>\n   </g>\n  </g>\n </g>\n <defs>\n  <clipPath id=\"pad48f39ba2\">\n   <rect x=\"51.378125\" 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=transformer-block-md-4 fw=all -->\n",
    "\n",
    "The two arrangements produce different networks at the same parameter\n",
    "values. Post-LN renormalizes the stream after every addition, so its RMS is\n",
    "pinned at $1$ at every depth; pre-LN only ever adds, so its stream grows\n",
    "like the square root of the depth (the block outputs are nearly independent\n",
    "at initialization, and variances add). That growth is not a defect — it is\n",
    "a protection. Each pre-LN branch reads the stream through its own\n",
    "normalization, so as the stream grows, each successive block's *relative*\n",
    "contribution shrinks, and the function stays tame.\n",
    "\n",
    "The post-LN stack has no such damping, and the token-spread column shows\n",
    "what happens instead. At initialization attention is nearly uniform, so\n",
    "every attention sublayer pulls each token toward the average token; post-LN\n",
    "then renormalizes, locking in the contraction, and the spread decays\n",
    "*geometrically* — by block 32 the tokens are identical to about one part in\n",
    "ten thousand. Pre-LN dilutes the same contraction by the growing stream,\n",
    "and the spread falls only polynomially, still at $0.4$ after 32 blocks.\n",
    "This is the *rank collapse* of @Dong.Cordonnier.Loukas.2021 caught\n",
    "in the act, and the gradients show why it matters for learning: a head\n",
    "attending over identical tokens returns the same mixture *whatever its\n",
    "query and key weights*, so those weights stop receiving gradient. The\n",
    "block-32 norms printed above locate the starvation precisely: post-LN's\n",
    "query and key projections receive gradients of order $10^{-9}$, six\n",
    "orders of magnitude below their pre-LN counterparts, while its value and\n",
    "output projections — which transform the one surviving token rather than\n",
    "compare tokens — come in around $0.06$, as healthy as pre-LN's (the FFN\n",
    "parameters, fed by the same residual stream, are likewise unaffected). It\n",
    "is specifically the query/key pathway, the part that decides *where* to\n",
    "attend, that dies; the plot shows the same starvation deepening block by\n",
    "block down the post-LN stack, while pre-LN tapers gently. (The effect\n",
    "softens under today's smaller initializations — one of several crutches,\n",
    "along with learning-rate warmup, that kept deep post-LN models trainable\n",
    "[@xiong2020layer].)\n",
    "\n",
    "The practical verdict matched this picture: post-LN transformers diverge\n",
    "without carefully tuned warmup, pre-LN models train without incident at the\n",
    "same depth [@xiong2020layer], and essentially every model since GPT-2\n",
    "has been pre-norm [@Radford.Wu.Child.ea.2019]. We make pre-norm our\n",
    "default; the next section, which has a trainable model, will show the\n",
    "post-LN failure not at initialization but live, during training.\n",
    "\n",
    "### RMSNorm\n",
    "\n",
    "Modern models also simplified the normalizer itself. LayerNorm standardizes\n",
    "each token vector by subtracting its mean and dividing by its standard\n",
    "deviation, then rescales (and re-shifts) feature-wise. *RMSNorm*\n",
    "[@Zhang.Sennrich.2019] drops the centering and the shift, keeping only\n",
    "the part that controls scale:\n",
    "\n",
    "$$\n",
    "\\mathrm{RMSNorm}(\\mathbf{x}) = \\frac{\\mathbf{x}}{\\sqrt{\\tfrac{1}{d} \\sum_{i=1}^{d} x_i^2 + \\epsilon}} \\odot \\boldsymbol{\\gamma}.\n",
    "$$\n",
    "\n",
    "It is a one-line function, and both frameworks ship it; let's verify our\n",
    "implementation against the built-in and time the two normalizers."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "transformer-block-rmsnorm"
   },
   "source": [
    "def rms_norm(X, weight, eps=1e-6):\n",
    "    return X / torch.sqrt((X * X).mean(-1, keepdim=True) + eps) * weight\n",
    "\n",
    "device = d2l.try_gpu()\n",
    "X = torch.randn(64, 512, 1024, device=device)\n",
    "layernorm = nn.LayerNorm(1024, device=device)\n",
    "rmsnorm = nn.RMSNorm(1024, eps=1e-6, device=device)\n",
    "print(f'ours vs. built-in: max deviation '\n",
    "      f'{(rms_norm(X, rmsnorm.weight) - rmsnorm(X)).abs().max():.1e}')\n",
    "for name, f in (('LayerNorm', layernorm), ('RMSNorm', rmsnorm)):\n",
    "    for _ in range(10):\n",
    "        f(X)                       # warmup\n",
    "    if device.type == 'cuda':\n",
    "        torch.cuda.synchronize()\n",
    "    t0 = time.time()\n",
    "    for _ in range(200):\n",
    "        f(X)\n",
    "    if device.type == 'cuda':\n",
    "        torch.cuda.synchronize()\n",
    "    print(f'{name}: {(time.time() - t0) / 200 * 1e3:.3f} ms')"
   ],
   "execution_count": 4,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "ours vs. built-in: max deviation 9.5e-07\nLayerNorm: 0.681 ms\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "RMSNorm: 0.652 ms\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=transformer-block-md-5 fw=all -->\n",
    "\n",
    "The timing is close to an anticlimax: on a modern GPU both normalizers are\n",
    "memory-bound, and what dropping the mean subtraction saves depends on the\n",
    "framework's kernel — nothing measurable in our PyTorch run, roughly a third\n",
    "in our JAX run, and either way microseconds in blocks that spend their time\n",
    "elsewhere. What RMSNorm actually buys is less machinery — no centering, no\n",
    "bias parameter, one statistic instead of two — at comparable quality, which\n",
    "is why Llama and most open models since use it [@touvron2023llama]. The\n",
    "lesson is about defaults rather than speed: when a component's job is \"keep\n",
    "the scale near one\", the simplest layer that does the job wins.\n",
    "\n",
    "### Normalizing Queries and Keys\n",
    "\n",
    "One more normalization migrated *inside* attention. Nothing in the block\n",
    "above controls the size of the attention logits\n",
    "$\\mathbf{q}^\\top \\mathbf{k} / \\sqrt{d}$ themselves: if training inflates\n",
    "the queries and keys, the softmax saturates, its Jacobian vanishes\n",
    "(that section), and training can spike or\n",
    "stall — an instability first met at vision-transformer scale, where\n",
    "attention entropy was observed collapsing in billion-parameter runs\n",
    "[@Dehghani.Djolonga.Mustafa.ea.2023]. *QK-norm*\n",
    "[@Henry.Dachapally.Pawar.ea.2020] applies RMSNorm to the queries and\n",
    "keys per head, right before the dot product, so the logit scale is pinned\n",
    "by construction rather than by hope. It costs two extra norms per layer and\n",
    "has become a standard ingredient of 2025-era models (Gemma 3, Qwen3, and\n",
    "OLMo 2 all ship it). OLMo 2 is also the one prominent re-litigation of this\n",
    "section's verdict: it pairs QK-norm with norms placed *after* each sublayer\n",
    "— though still off the stream's identity path, conceding the main lesson of\n",
    "the experiment above [@OLMo.2025]. Gemma 3 hedges in the other\n",
    "direction and normalizes each branch on both sides, before *and* after the\n",
    "sublayer [@Gemma.Team.2025]. the figure lays the four\n",
    "placements now in circulation side by side; note that only the original\n",
    "post-LN interrupts the stream itself, which is why it alone collapses in\n",
    "the experiment above. We leave QK-norm as an exercise rather than a flag;\n",
    "it drops into the block through the `attn_factory` hook introduced below.\n",
    "\n",
    "![Four placements of normalization around one sublayer; the FFN sublayer is treated identically, and the residual stream runs bottom to top. Post-LN normalizes the stream itself after each addition; pre-LN normalizes what the branch reads; OLMo 2 normalizes what the branch writes; Gemma 3 normalizes the branch on both sides. Only post-LN touches the stream's identity path.](../img/mdl-transformers-norm-taxonomy.svg)\n",
    "\n",
    "## The Feed-Forward Network\n",
    "\n",
    "Between the attention sublayers sits the block's other half: a small MLP\n",
    "applied to each position independently,\n",
    "\n",
    "$$\n",
    "\\mathrm{FFN}(\\mathbf{x}) = \\mathbf{W}_2\\, \\phi(\\mathbf{W}_1 \\mathbf{x}),\n",
    "$$\n",
    "\n",
    "with $\\mathbf{W}_1 \\in \\mathbb{R}^{4d \\times d}$ and $\\mathbf{W}_2 \\in\n",
    "\\mathbb{R}^{d \\times 4d}$ in the classic configuration. The factor-4\n",
    "expansion is a convention the field has never found strong reason to\n",
    "revisit; what it implies is worth noticing: the FFN holds $8d^2$ parameters\n",
    "against attention's $4d^2$, so about two thirds of a block's parameters sit\n",
    "in these two matrices. When people say a transformer's knowledge lives\n",
    "mostly in its FFNs, this ratio is the accounting behind the claim.\n",
    "\n",
    "Two upgrades separate the 2017 FFN from today's. The activation $\\phi$\n",
    "moved from ReLU to *GELU* [@Hendrycks.Gimpel.2016], a smoothed relative\n",
    "that ends the debate about the kink at zero (GPT onward\n",
    "[@Radford.Narasimhan.Salimans.ea.2018]). Then @Shazeer.2020\n",
    "revived an older idea — gating — and found that a *gated linear unit* with\n",
    "the SiLU activation beat both:\n",
    "\n",
    "$$\n",
    "\\mathrm{SwiGLU}(\\mathbf{x}) = \\mathbf{W}_2 \\big(\\mathrm{SiLU}(\\mathbf{W}_g \\mathbf{x}) \\odot \\mathbf{W}_1 \\mathbf{x}\\big), \\qquad \\mathrm{SiLU}(z) = z \\cdot \\mathrm{sigmoid}(z).\n",
    "$$\n",
    "\n",
    "Instead of gating a linear transform through a fixed nonlinearity, SwiGLU\n",
    "computes *two* linear views of the input and multiplies them, one squashed\n",
    "into a soft gate, the other passed through untouched — the value is\n",
    "transmitted at full strength wherever the gate is open, rather than bent\n",
    "through the activation. The multiplicative interaction is the same trick\n",
    "gates played in LSTMs, here compressed into a single layer. A third matrix\n",
    "means more parameters at the same width, so fair comparisons shrink the\n",
    "hidden width to $\\tfrac{8}{3} d$, making three matrices of the gated FFN\n",
    "cost the same $8d^2$ as the classic one, up to rounding. Both variants, in\n",
    "one class:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "transformer-block-the-feed-forward-network-1"
   },
   "source": [
    "class FeedForward(nn.Module):\n",
    "    \"\"\"Position-wise FFN: GELU MLP or SwiGLU at matched parameter count.\"\"\"\n",
    "    def __init__(self, num_hiddens, act='swiglu', bias=False):\n",
    "        super().__init__()\n",
    "        assert act in ('swiglu', 'gelu'), f'unknown act: {act!r}'\n",
    "        self.act = act\n",
    "        if act == 'gelu':\n",
    "            width = 4 * num_hiddens\n",
    "        else:  # 'swiglu': three matrices; width 8d/3 matches the MLP budget\n",
    "            width = round(8 * num_hiddens / 3)\n",
    "            self.W_g = nn.Linear(num_hiddens, width, bias=bias)\n",
    "        self.W_1 = nn.Linear(num_hiddens, width, bias=bias)\n",
    "        self.W_2 = nn.Linear(width, num_hiddens, bias=bias)\n",
    "\n",
    "    def forward(self, X):\n",
    "        if self.act == 'gelu':\n",
    "            return self.W_2(F.gelu(self.W_1(X), approximate='tanh'))\n",
    "        return self.W_2(F.silu(self.W_g(X)) * self.W_1(X))"
   ],
   "execution_count": 5,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=transformer-block-md-6 fw=all -->\n",
    "\n",
    "(The GELU here is the tanh approximation, the exact variant GPT-2 shipped\n",
    "with — a detail that will matter when we load its weights in\n",
    "that section.) The parameter match is worth checking rather than\n",
    "trusting:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "transformer-block-the-feed-forward-network-2"
   },
   "source": [
    "for act in ('gelu', 'swiglu'):\n",
    "    ffn = FeedForward(256, act)\n",
    "    print(f'{act:>7}: {sum(p.numel() for p in ffn.parameters())} parameters')"
   ],
   "execution_count": 6,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "   gelu: 524288 parameters\n swiglu: 524544 parameters\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=transformer-block-md-7 fw=all -->\n",
    "\n",
    "Whether the gate is *worth* anything at equal cost is an empirical\n",
    "question, and a matched-budget race settles it — but only once we have a\n",
    "model to train. That is the closing experiment of this section.\n",
    "\n",
    "## A Configurable Block\n",
    "\n",
    "We now assemble the block for real, and we build it the way this chapter\n",
    "will use it: every design decision above becomes a constructor argument.\n",
    "`norm` selects LayerNorm or RMSNorm, `act` selects the FFN, `pre_norm`\n",
    "selects the arrangement, and two factory hooks let later sections swap\n",
    "whole sublayers — a mixture-of-experts FFN, or a cache-friendly\n",
    "grouped-query attention — without touching this class again. The\n",
    "constructor validates its flag strings: a typo must fail loudly, not\n",
    "silently select an architecture. The default configuration (pre-norm,\n",
    "RMSNorm, SwiGLU) is the one you would find inside an open-weights model\n",
    "released this year; `TransformerBlock(num_hiddens, num_heads, norm='layer',\n",
    "act='gelu', pre_norm=False, bias=True)` is a compact post-norm block in the\n",
    "2017 encoder arrangement (with one anachronism kept for simplicity: GELU,\n",
    "where the 2017 original used ReLU)."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "transformer-block-a-configurable-block"
   },
   "source": [
    "class TransformerBlock(nn.Module):\n",
    "    \"\"\"Configurable transformer block: attention + FFN on a residual\n",
    "    stream.\"\"\"\n",
    "    def __init__(self, num_hiddens, num_heads, dropout=0, norm='rms',\n",
    "                 act='swiglu', pre_norm=True, bias=False, attn_factory=None,\n",
    "                 ffn_factory=None):\n",
    "        super().__init__()\n",
    "        assert norm in ('rms', 'layer'), f'unknown norm: {norm!r}'\n",
    "        assert num_hiddens % num_heads == 0\n",
    "        self.pre_norm = pre_norm\n",
    "        make_norm = nn.RMSNorm if norm == 'rms' else nn.LayerNorm\n",
    "        self.norm1, self.norm2 = make_norm(num_hiddens), make_norm(num_hiddens)\n",
    "        self.attention = (d2l.MultiHeadAttention(num_hiddens, num_heads,\n",
    "                                                 dropout, bias=bias)\n",
    "                          if attn_factory is None else attn_factory())\n",
    "        self.ffn = (FeedForward(num_hiddens, act, bias=bias)\n",
    "                    if ffn_factory is None else ffn_factory())\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, X, valid_lens=None):\n",
    "        if self.pre_norm:\n",
    "            Y = self.norm1(X)\n",
    "            X = X + self.dropout(self.attention(Y, Y, Y, valid_lens))\n",
    "            return X + self.dropout(self.ffn(self.norm2(X)))\n",
    "        X = self.norm1(X + self.dropout(self.attention(X, X, X, valid_lens)))\n",
    "        return self.norm2(X + self.dropout(self.ffn(X)))"
   ],
   "execution_count": 7,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=transformer-block-md-8 fw=all -->\n",
    "\n",
    "### Shapes and Parameters\n",
    "\n",
    "The block preserves its input shape whatever the flags, and its parameter\n",
    "census confirms the accounting from above: $4d^2$ in attention, $8d^2$ in\n",
    "the FFN (however the activation slices it), and a rounding error of\n",
    "normalization weights — about $12d^2$ per block, two thirds of it FFN."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "transformer-block-shapes-and-parameters-1"
   },
   "source": [
    "X = torch.ones(2, 10, 256)\n",
    "blk = TransformerBlock(256, num_heads=8)\n",
    "blk2017 = TransformerBlock(256, num_heads=8, norm='layer', act='gelu',\n",
    "                           pre_norm=False, bias=True)\n",
    "d2l.check_shape(blk(X), X.shape)\n",
    "d2l.check_shape(blk2017(X), X.shape)\n",
    "count = lambda m: sum(p.numel() for p in m.parameters())\n",
    "print(f'attention {count(blk.attention)}, ffn {count(blk.ffn)}, '\n",
    "      f'total {count(blk)} = {count(blk) / 256 ** 2:.2f} d^2')"
   ],
   "execution_count": 8,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "attention 262144, ffn 524544, total 787200 = 12.01 d^2\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=transformer-block-md-9 fw=all -->\n",
    "\n",
    "One block is not yet a language model, but it is close enough to train.\n",
    "The same `valid_lens` mechanism that masked padding in\n",
    "that section makes the block *causal*: pass\n",
    "each query a valid length equal to its own position, and every token\n",
    "attends only backward. A dozen lines wrap an embedding, a stack of blocks,\n",
    "and a tied output head (as in that section) around\n",
    "that trick:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "transformer-block-shapes-and-parameters-2"
   },
   "source": [
    "class CharLM(nn.Module):\n",
    "    \"\"\"Minimal char LM: causal masking via per-query valid lengths.\"\"\"\n",
    "    def __init__(self, vocab_size, num_hiddens=128, num_heads=4, num_blks=2,\n",
    "                 act='swiglu', max_len=128):\n",
    "        super().__init__()\n",
    "        self.token_emb = nn.Embedding(vocab_size, num_hiddens)\n",
    "        self.pos_emb = nn.Embedding(max_len, num_hiddens)\n",
    "        for emb in (self.token_emb, self.pos_emb):\n",
    "            nn.init.normal_(emb.weight, std=0.02)\n",
    "        self.blks = nn.ModuleList([\n",
    "            TransformerBlock(num_hiddens, num_heads, act=act)\n",
    "            for _ in range(num_blks)])\n",
    "        self.norm = nn.RMSNorm(num_hiddens)\n",
    "\n",
    "    def forward(self, X):\n",
    "        B, T = X.shape\n",
    "        H = self.token_emb(X) + self.pos_emb(\n",
    "            torch.arange(T, device=X.device))\n",
    "        causal = torch.arange(1, T + 1, device=X.device).repeat(B, 1)\n",
    "        for blk in self.blks:\n",
    "            H = blk(H, causal)\n",
    "        return F.linear(self.norm(H), self.token_emb.weight)"
   ],
   "execution_count": 9,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=transformer-block-md-10 fw=all -->\n",
    "\n",
    "### The Flags at Work: GELU versus SwiGLU\n",
    "\n",
    "Now the promised race. Same data (the character-level Time Machine corpus\n",
    "of that section), same model, same seed, same optimizer and\n",
    "learning rate, same 600 steps — the only difference is the `act` flag, and\n",
    "the parameter counts match to a tenth of a percent."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "transformer-block-the-flags-at-work-gelu-versus-swiglu"
   },
   "source": [
    "data = d2l.TimeMachine(batch_size=64, num_steps=128, tokenization='char',\n",
    "                       num_train=100000, num_val=3000)\n",
    "for act in ('gelu', 'swiglu'):\n",
    "    torch.manual_seed(0)\n",
    "    model = CharLM(len(data.vocab), act=act)\n",
    "    losses = d2l.train_lm(model, data,\n",
    "                          torch.optim.AdamW(model.parameters(), lr=1e-3,\n",
    "                                            weight_decay=0.0), 600)\n",
    "    print(f'{act:>7}: loss at step 200/400/600: ' + '/'.join(\n",
    "        f'{sum(losses[k-100:k]) / 100:.2f}' for k in (200, 400, 600)))"
   ],
   "execution_count": 10,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "   gelu: loss at step 200/400/600: 2.26/1.65/1.31\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      " swiglu: loss at step 200/400/600: 2.24/1.51/1.17\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=transformer-block-md-11 fw=all -->\n",
    "\n",
    "The gated FFN ends more than a tenth of a nat ahead — a margin that holds up\n",
    "across seeds (rerunning with seeds 1 and 2 moves each number by a couple of\n",
    "hundredths, not the gap) and, more importantly, agrees in direction with\n",
    "@Shazeer.2020's systematic sweep and with the consistent choice of\n",
    "the major model families since Llama. It is a modest, real improvement of\n",
    "the kind that architecture progress is actually made of: no single dramatic\n",
    "win, but a percent here and a percent there, at equal cost, compounding.\n",
    "\n",
    "## Summary\n",
    "\n",
    "A transformer block is two sublayers writing into a residual stream:\n",
    "attention communicates between positions, the position-wise FFN computes\n",
    "within each one and holds about two thirds of the parameters. Where the\n",
    "normalization goes decides how deep stacks behave at initialization: the\n",
    "post-LN arrangement of 2017 renormalizes the stream after each addition,\n",
    "which pins its scale but lets near-uniform attention collapse the tokens\n",
    "geometrically — the query and key projections at the top of a 32-block\n",
    "stack receive gradients six orders of magnitude smaller than at the\n",
    "bottom, while the value, output, and FFN weights keep ordinary gradients;\n",
    "it is the where-to-attend pathway that starves.\n",
    "Pre-norm moves the normalizer onto each branch, lets the stream grow like\n",
    "the square root of the depth, and keeps every block trainable; it has been\n",
    "the default since GPT-2. RMSNorm keeps only the scale statistic (same\n",
    "measured cost, less machinery), and QK-norm extends the same discipline to\n",
    "the attention logits. In the FFN, SwiGLU replaces the fixed nonlinearity\n",
    "with a learned soft gate; at matched parameter count it wins by a small,\n",
    "seed-stable margin on our character model, consistent with its\n",
    "near-universal adoption. The section's product is `TransformerBlock`:\n",
    "normalization,\n",
    "activation, and arrangement as flags, attention and FFN as swappable\n",
    "factories — the single unit from which this chapter builds a GPT, a\n",
    "KV-cached decoder, an encoder, a vision transformer, and a\n",
    "mixture-of-experts model.\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. At initialization the block outputs are nearly independent, so the\n",
    "   pre-norm stream RMS should grow like $\\sqrt{1 + 2N}$ after $N$ blocks\n",
    "   with two sublayers each — up to the sublayers' output scale. Extend the\n",
    "   signal-propagation experiment to fit the growth exponent from a log-log\n",
    "   plot. How close is it to $1/2$, and why does the attention sublayer\n",
    "   contribute less than the FFN sublayer early in the stack?\n",
    "2. OLMo 2 places the normalization *after* each sublayer but *off* the\n",
    "   stream: $\\mathbf{h} = \\mathbf{x} +\n",
    "   \\mathrm{Norm}(\\mathrm{Attn}(\\mathbf{x}))$. Add this third arrangement to\n",
    "   `MiniBlock` and rerun the experiment. Which of post-LN's pathologies\n",
    "   does it avoid, which of pre-LN's properties does it give up, and what\n",
    "   happens to the stream RMS?\n",
    "3. The signal-propagation experiment uses the 2017 Xavier initialization.\n",
    "   Rerun it with each framework's default initialization and with all\n",
    "   weights scaled by an extra factor of $0.5$. How does the collapse depth\n",
    "   change? Relate your observations to why post-LN models of GPT-1's era\n",
    "   (12 blocks) were trainable at all.\n",
    "4. Implement two more of @Shazeer.2020's gated variants by changing\n",
    "   one line of `FeedForward`: ReGLU (ReLU gate) and GEGLU (GELU gate). Race\n",
    "   all four at matched parameters on `CharLM`. Do your rankings agree with\n",
    "   the paper's? What does the spread between the gated variants, compared\n",
    "   with the gap to the ungated MLP, tell you about *which* design choice\n",
    "   carries the improvement?\n",
    "5. Work out the exact parameter count of `TransformerBlock(d, h)` for both\n",
    "   `act` settings, including normalization weights, and check it against\n",
    "   the census cell. Real models round the SwiGLU width up to a multiple of\n",
    "   256 for hardware efficiency; how far from $8d^2$ does the Llama-7B\n",
    "   configuration ($d = 4096$, FFN width $11008$) land?\n",
    "6. Implement QK-norm as an `attn_factory`: wrap `d2l.MultiHeadAttention`\n",
    "   so that queries and keys pass through an RMSNorm (one per head\n",
    "   dimension) before the dot product. At initialization, measure the\n",
    "   standard deviation of the attention logits with and without it as you\n",
    "   scale `num_hiddens` from 64 to 1024. Which curve does\n",
    "   the equation's $1/\\sqrt{d}$ scaling predict, and\n",
    "   which do you observe?"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-pytorch",
   "display_name": "d2l (PyTorch)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
