{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-1 fw=all -->\n",
    "\n",
    "# Initialization\n",
    "\n",
    "Every parameter in a network has a value before training starts: a fully\n",
    "connected layer mapping 512 inputs to 256 outputs brings a $256 \\times 512$\n",
    "weight matrix filled with random numbers. that section explains why the scale of those\n",
    "numbers decides whether a deep network trains at all, and derives the Xavier\n",
    "[@Glorot.Bengio.2010] and He [@He.Zhang.Ren.ea.2015] schemes that\n",
    "keep signal variance stable across layers. This section is the practical\n",
    "companion: what the library does when you say nothing, how to impose a scheme\n",
    "of your own on a whole model or a single block, what transformer-era code adds\n",
    "on top of Xavier and He, and how to write an initializer that no menu provides."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "init-initialization"
   },
   "source": [
    "import math\n",
    "from mxnet import init, np, npx\n",
    "from mxnet.gluon import nn\n",
    "npx.set_np()"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-2 fw=all -->\n",
    "\n",
    "## Defaults and When to Override Them\n",
    "\n",
    "You can usually ignore initialization because the default is sensible."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-3 fw=mxnet -->\n",
    "\n",
    "Gluon separates construction from initialization: `nn.Dense(256)` carries no\n",
    "numbers until you call `initialize()`, and when the input width is left\n",
    "unspecified even that call only records your choice, with the draw deferred to\n",
    "the first forward pass (that section). What it draws is the legacy\n",
    "exception in the table below: every weight comes from $U(-0.07, 0.07)$\n",
    "regardless of the layer's shape, a fixed constant rather than a fan-aware\n",
    "bound, and the bias starts at zero. A uniform distribution on $[-b, b]$ has\n",
    "standard deviation $b/\\sqrt{3}$, so all three claims are cheap to check\n",
    "against a fresh layer (passing `in_units` pins the input width, so nothing is\n",
    "deferred here):"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "init-defaults-and-when-to-override-them"
   },
   "source": [
    "np.random.seed(0)\n",
    "layer = nn.Dense(256, in_units=512)\n",
    "layer.initialize()\n",
    "w = layer.weight.data()\n",
    "print(f'range [{float(w.min()):.4f}, {float(w.max()):.4f}], '\n",
    "      f'predicted bound 0.0700')\n",
    "print(f'std {float(w.std()):.4f}, predicted {0.07 / math.sqrt(3):.4f}')\n",
    "print(f'bias all zero: {bool((layer.bias.data() == 0).all())}')"
   ],
   "execution_count": 2,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "range [-0.0700, 0.0700], predicted bound 0.0700\nstd 0.0404, predicted 0.0404\nbias all zero: True\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-4 fw=all -->\n",
    "\n",
    "Other libraries make different but equally fan-aware choices, with one legacy\n",
    "exception:\n",
    "\n",
    "| Library | Default for a dense layer's weight |\n",
    "|:--|:--|\n",
    "| PyTorch | $U(-1/\\sqrt{n_\\textrm{in}},\\, 1/\\sqrt{n_\\textrm{in}})$, bias likewise |\n",
    "| Flax (JAX) | LeCun normal: truncated normal, std $1/\\sqrt{n_\\textrm{in}}$; zero bias |\n",
    "| Keras (TensorFlow) | Glorot uniform: $U(\\pm\\sqrt{6/(n_\\textrm{in} + n_\\textrm{out})})$; zero bias |\n",
    "| MXNet Gluon | $U(-0.07, 0.07)$, not fan-aware, so override it; zero bias |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-5 fw=mxnet -->\n",
    "\n",
    "When should you override? Here the first candidate is the default itself:\n",
    "$\\pm 0.07$ was tuned for the hidden widths of an earlier era and ignores\n",
    "fan-in, so a layer much wider than those starts with too much variance, a\n",
    "much narrower one with too little, and either error compounds across depth by\n",
    "the analysis of that section. Beyond that, the usual\n",
    "situations recur: a deep network without normalization layers, where variance\n",
    "compounds even when each layer is right (we demonstrate this below);\n",
    "reproducing a paper whose results depend on its initialization recipe; and\n",
    "architecture-specific corrections such as the residual scaling later in this\n",
    "section. The lazy fine print runs the other way from PyTorch's: with the\n",
    "input width unspecified, `initialize()` merely records which initializer to\n",
    "use, and the draw happens at the first forward pass, once the layer knows its\n",
    "fan-in."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-6 fw=all -->\n",
    "\n",
    "## Applying Initializers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-7 fw=mxnet -->\n",
    "\n",
    "To impose a scheme we hand an `Initializer` object to the model:\n",
    "`net.initialize(init=..., force_reinit=True)` walks every parameter of the\n",
    "block tree of that section and runs the initializer on\n",
    "it (`force_reinit` lifts the guard against silently clobbering a model that\n",
    "already has values). The object is a default, not a decree: a parameter\n",
    "constructed with its own initializer keeps it, and every `Dense` bias is\n",
    "constructed with `zeros`, which is why the normal draw below leaves the\n",
    "biases at zero. Dispatching by layer type, the pattern PyTorch users write,\n",
    "is possible (iterate over `net` and test `isinstance`), but Gluon's native\n",
    "unit of selection is the subtree: every block, and every individual\n",
    "parameter, has an `initialize` method of its own. We show both, starting with\n",
    "the whole net."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "init-applying-initializers-1"
   },
   "source": [
    "net = nn.Sequential()\n",
    "net.add(nn.Dense(8, activation='relu'))\n",
    "net.add(nn.Dense(1))\n",
    "net.initialize(init=init.Normal(sigma=0.01))\n",
    "\n",
    "X = np.ones((2, 4))\n",
    "net(X)  # the input width was deferred, so the draw happens here\n",
    "net[0].weight.data()[0], net[0].bias.data()[0]"
   ],
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(array([ 0.00490368,  0.00765541, -0.00184343,  0.00680209]), array(0.))"
     },
     "metadata": {},
     "execution_count": 3
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-8 fw=mxnet -->\n",
    "\n",
    "Mixing schemes per block is `initialize` on the subtree: `net[0]` is a block\n",
    "and `net[0].weight` is a parameter, and both accept an initializer directly,\n",
    "with `force_reinit=True` because we are overwriting live values. Below we\n",
    "give the first layer Xavier initialization (that section) and set\n",
    "the second to a constant, a poor idea for training by the symmetry argument\n",
    "of that section, but it makes the mechanics visible:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "init-applying-initializers-2"
   },
   "source": [
    "net[0].weight.initialize(init=init.Xavier(), force_reinit=True)\n",
    "net[1].initialize(init=init.Constant(42), force_reinit=True)\n",
    "net[0].weight.data()[0], net[1].weight.data()"
   ],
   "execution_count": 4,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(array([ 0.37153655, -0.47844273, -0.4485028 , -0.5750185 ]),\n array([[42., 42., 42., 42., 42., 42., 42., 42.]]))"
     },
     "metadata": {},
     "execution_count": 4
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-9 fw=mxnet -->\n",
    "\n",
    "There is no second route to learn, because `force_reinit=True` already is the\n",
    "route for a model with live values: the two cells above re-initialized the\n",
    "same network three times. Whole net, block, or single parameter, the call is\n",
    "the same, and nothing else is needed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-10 fw=all -->\n",
    "\n",
    "## Modern Schemes: Truncation, Depth, and Zeros\n",
    "\n",
    "Xavier and He set the variance of a single layer. The schemes below, standard\n",
    "in transformer codebases, adjust what happens in the distribution's tails,\n",
    "across depth, and at a block's start.\n",
    "\n",
    "### Truncated Normals\n",
    "\n",
    "A Gaussian gets the variance right, but its tails are unbounded. That is\n",
    "harmless for one draw and a near-certainty at scale: among the $10^8$ weights\n",
    "of a BERT-sized model, dozens land beyond five standard deviations. Large\n",
    "draws also consume disproportionate dynamic range once a model is cast to low\n",
    "precision (that section). BERT and implementations in the ViT\n",
    "lineage use truncated-normal initialization\n",
    "[@Devlin.Chang.Lee.ea.2018; @Dosovitskiy.Beyer.Kolesnikov.ea.2021]: the\n",
    "tails are cut off at a fixed multiple of the nominal standard deviation. Raw\n",
    "truncated-normal initializers do not necessarily restore the variance removed\n",
    "with the tails; fan-aware variance-scaling initializers often do."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-11 fw=mxnet -->\n",
    "\n",
    "`mxnet.init` ships no truncated normal; the menu ends just before the scheme\n",
    "we want. The escape hatch is the one this section closes with: an initializer\n",
    "is a subclass of `init.Initializer` that fills an array inside its\n",
    "`_init_weight` method, so the missing entry costs a dozen lines. We draw\n",
    "normal values and redraw every entry that lands outside two standard\n",
    "deviations until none does, the same discard-and-redraw semantics as Keras,\n",
    "and keep the drawing function separate so it can double as the demonstration:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "init-truncated-normals"
   },
   "source": [
    "def trunc_normal(sigma, shape):\n",
    "    w = np.random.normal(0, sigma, shape)\n",
    "    tails = np.abs(w) > 2 * sigma\n",
    "    while bool(tails.any()):\n",
    "        w = np.where(tails, np.random.normal(0, sigma, shape), w)\n",
    "        tails = np.abs(w) > 2 * sigma\n",
    "    return w\n",
    "\n",
    "class TruncatedNormal(init.Initializer):\n",
    "    def __init__(self, sigma=0.02):\n",
    "        super().__init__(sigma=sigma)\n",
    "        self.sigma = sigma\n",
    "\n",
    "    def _init_weight(self, name, arr):\n",
    "        arr[:] = trunc_normal(self.sigma, arr.shape)\n",
    "\n",
    "np.random.seed(0)\n",
    "w = np.random.normal(0, 0.02, (1000, 1000))\n",
    "print(f'normal:    std {float(w.std()):.4f}, '\n",
    "      f'max weight {float(np.abs(w).max()):.4f}')\n",
    "w = trunc_normal(0.02, (1000, 1000))\n",
    "print(f'truncated: std {float(w.std()):.4f}, '\n",
    "      f'max weight {float(np.abs(w).max()):.4f}')"
   ],
   "execution_count": 5,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "normal:    std 0.0200, max weight 0.0919\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "truncated: std 0.0176, max weight 0.0400\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-12 fw=all -->\n",
    "\n",
    "A million plain-normal draws produce entries near $5\\sigma = 0.1$; the\n",
    "truncated version guarantees $|w| \\leq 0.04$. Its printed standard deviation\n",
    "dips slightly below the nominal 0.02 because truncation removes tail mass;\n",
    "practice ignores the difference.\n",
    "\n",
    "### Scaling Down Residual Branches\n",
    "\n",
    "A residual network computes $\\mathbf{x}_{k+1} = \\mathbf{x}_k +\n",
    "f_k(\\mathbf{x}_k)$, so its output is the input plus the sum of $N$ block\n",
    "contributions. If every block is initialized identically, each contribution\n",
    "has variance proportional to that of the already-inflated stream it reads, and\n",
    "the stream's variance compounds geometrically with depth. This is the\n",
    "depth-wise cousin of the layer-wise explosion in\n",
    "that section. GPT-2's fix is to shrink only the *last*\n",
    "linear layer of each residual branch, the output projection that writes into\n",
    "the stream, by a factor of $1/\\sqrt{N}$: with $N$ roughly independent\n",
    "contributions each scaled down that way, the sum's variance stays $O(1)$\n",
    "regardless of depth [@Radford.Wu.Child.ea.2019]. For an $L$-layer\n",
    "transformer the published recipe is a normal with std $0.02/\\sqrt{2L}$ on the\n",
    "residual projections — the number of contributions is $N = 2L$, because each\n",
    "layer writes to the stream twice, once from attention and once from the MLP.\n",
    "Unlike Xavier or He, the base 0.02 is not derived from fan-in; it is an\n",
    "empirical constant that worked at GPT-2's widths and stuck.\n",
    "\n",
    "the figure draws the two regimes side by side.\n",
    "\n",
    "![A residual stream with additive block contributions, unscaled (left) versus scaled by 1/sqrt(N) (right). Under the independent-contribution approximation, unscaled variance grows with N and scaling keeps the sum O(1). In an unnormalized stack, each branch reads an already inflated stream, so growth can be faster, as the experiment below shows.](../img/bg-residual-stream.svg)\n",
    "\n",
    "### Starting a Block at Zero\n",
    "\n",
    "Zero-initializing *every* weight is fatal: all units in a layer compute the\n",
    "same output, receive the same gradient, and stay identical forever\n",
    "(that section). Zero-initializing just the *last* layer\n",
    "of a residual block is a different and useful move. The branch then\n",
    "contributes exactly nothing, each block is the identity map, and the network\n",
    "starts as a shallow function whose depth switches on gradually during\n",
    "training. Symmetry is not a problem because the branch's earlier layers keep\n",
    "their random weights. On the first backward pass, the zeroed projection\n",
    "receives a nonzero gradient because its input is nonzero, while the earlier\n",
    "layers receive zero gradient through that projection. After the projection\n",
    "moves off zero, gradient reaches the whole branch. Keep this scheme distinct\n",
    "from the previous one: GPT-2 makes every residual projection *small but\n",
    "nonzero*, whereas zero-init makes one layer *exactly zero* so the block starts\n",
    "as an exact identity.\n",
    "\n",
    "### Watching the Variance Compound\n",
    "\n",
    "Claims about variance at depth are cheap to test. We reuse a compact residual\n",
    "block (it repeats the one from that section) and stack\n",
    "it $N$ deep:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "init-watching-the-variance-compound-1"
   },
   "source": [
    "he = init.MSRAPrelu(factor_type='in', slope=0.0)\n",
    "\n",
    "class ResidualBlock(nn.Block):\n",
    "    def __init__(self, d, out_init=he):\n",
    "        super().__init__()\n",
    "        self.body = nn.Sequential()\n",
    "        self.body.add(nn.Dense(4 * d, activation='relu',\n",
    "                               weight_initializer=he))\n",
    "        self.body.add(nn.Dense(d, weight_initializer=out_init))\n",
    "\n",
    "    def forward(self, X):\n",
    "        return X + self.body(X)"
   ],
   "execution_count": 6,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-13 fw=mxnet -->\n",
    "\n",
    "Every dense layer gets He initialization through its `weight_initializer`\n",
    "constructor argument, which sets the per-parameter initializer that\n",
    "`net.initialize()` respects. The menu entry is `init.MSRAPrelu`, named after\n",
    "the He et al. paper's lab: with `factor_type='in'` and `slope=0.0` (the PReLU\n",
    "slope, zero for a plain ReLU) it draws a Gaussian with variance\n",
    "$2/n_\\textrm{in}$, exactly He normal. The treatment is a constructor argument\n",
    "as well: the block takes its output projection's initializer, so each of the\n",
    "three treatments is just a different initializer. Leaving it alone means He\n",
    "again and zeroing is `init.Zero()`, but scaling by $1/\\sqrt{N}$ has no menu\n",
    "entry, so once more we subclass `Initializer`: the instance stores the depth,\n",
    "and its `_init_weight` shrinks a fresh He draw in place."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "init-watching-the-variance-compound-2"
   },
   "source": [
    "class ScaledHe(init.Initializer):\n",
    "    def __init__(self, num_blocks):\n",
    "        super().__init__(num_blocks=num_blocks)\n",
    "        self.num_blocks = num_blocks\n",
    "\n",
    "    def _init_weight(self, name, arr):\n",
    "        he._init_weight(name, arr)\n",
    "        arr *= self.num_blocks ** -0.5\n",
    "\n",
    "def output_std(num_blocks, out_init):\n",
    "    np.random.seed(0)\n",
    "    net = nn.Sequential()\n",
    "    net.add(*[ResidualBlock(64, out_init=out_init(num_blocks))\n",
    "              for _ in range(num_blocks)])\n",
    "    net.initialize()  # deferred: the draws happen at the forward pass\n",
    "    X = np.random.normal(0, 1, (256, 64))\n",
    "    return float(net(X).std())"
   ],
   "execution_count": 7,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-14 fw=all -->\n",
    "\n",
    "One forward pass on unit-variance inputs per depth and treatment:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "init-watching-the-variance-compound-3"
   },
   "source": [
    "tweaks = {'default': lambda n: he,\n",
    "          'scaled': ScaledHe,\n",
    "          'zero': lambda n: init.Zero()}\n",
    "print(f'{\"N\":>3}' + ''.join(f'{name:>10}' for name in tweaks))\n",
    "for n in (2, 8, 32):\n",
    "    stds = (output_std(n, tweak) for tweak in tweaks.values())\n",
    "    print(f'{n:>3}' + ''.join(f'{s:>10.3g}' for s in stds))"
   ],
   "execution_count": 8,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "  N   default    scaled      zero\n  2      2.94      1.99     0.997\n  8        85      2.49     0.997\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      " 32  4.75e+07      2.65     0.997\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-15 fw=all -->\n",
    "\n",
    "The default column multiplies by a roughly constant factor per block in this\n",
    "unnormalized stack, reaching tens of millions by $N=32$. Such initial\n",
    "activations make optimization impractical. The scaled column stays near a\n",
    "small constant over the depths tested: scaling cuts each branch's variance by\n",
    "a factor of $N$, following the independent-contribution approximation above.\n",
    "The zero column reproduces the input's standard deviation, since every block\n",
    "is the identity. This forward-pass test does not prove how an optimizer will\n",
    "behave, but it catches an unusable initialization before training begins.\n",
    "\n",
    "## Custom Initializers\n",
    "\n",
    "Occasionally the menu has nothing you need."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-16 fw=mxnet -->\n",
    "\n",
    "By now this is familiar: an initializer is a subclass of `init.Initializer`\n",
    "whose `_init_weight` method fills the array it is handed, in place, and this\n",
    "section has already written two (the truncated normal and the depth-scaled\n",
    "He). Writing one is no harder than using one."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-17 fw=all -->\n",
    "\n",
    "Suppose, to make the point vividly, we want weights distributed as\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "    w \\sim \\begin{cases}\n",
    "        U(5, 10) & \\textrm{ with probability } \\frac{1}{4} \\\\\n",
    "            0    & \\textrm{ with probability } \\frac{1}{2} \\\\\n",
    "        U(-10, -5) & \\textrm{ with probability } \\frac{1}{4}\n",
    "    \\end{cases}\n",
    "\\end{aligned}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-18 fw=mxnet -->\n",
    "\n",
    "Draw uniformly from $U(-10, 10)$, then zero every entry of magnitude below 5.\n",
    "Handing an instance to `net.initialize` makes it official:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "init-custom-initializers-1"
   },
   "source": [
    "class MyInit(init.Initializer):\n",
    "    def _init_weight(self, name, data):\n",
    "        data[:] = np.random.uniform(-10, 10, data.shape)\n",
    "        data *= np.abs(data) >= 5\n",
    "\n",
    "net = nn.Sequential()\n",
    "net.add(nn.Dense(8, activation='relu'))\n",
    "net.add(nn.Dense(1))\n",
    "net.initialize(MyInit())\n",
    "net(X)\n",
    "net[0].weight.data()[:2]"
   ],
   "execution_count": 9,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "array([[ 9.935757 , -0.       ,  0.       , -7.146276 ],\n       [-0.       , -0.       , -5.6079783, -6.314646 ]])"
     },
     "metadata": {},
     "execution_count": 9
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-19 fw=mxnet -->\n",
    "\n",
    "No guard is needed on the way in, because gradients are only recorded inside\n",
    "an `autograd.record()` block: outside one, indexed assignment is a plain\n",
    "array write. `weight.data()` hands back the array the parameter holds on the\n",
    "device, so the same door handles one-off surgery, such as offsetting a whole\n",
    "matrix or pinning a single entry:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "init-custom-initializers-2"
   },
   "source": [
    "net[0].weight.data()[:] += 1\n",
    "net[0].weight.data()[0, 0] = 42\n",
    "net[0].weight.data()[0]"
   ],
   "execution_count": 10,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "array([42.      ,  1.      ,  1.      , -6.146276])"
     },
     "metadata": {},
     "execution_count": 10
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-20 fw=mxnet -->\n",
    "\n",
    "The ordering caveat of the lazy world (that section) splits by\n",
    "route here too: `initialize` may run before shapes are known, since it only\n",
    "records which initializer to use and the draw waits for the first forward\n",
    "pass, but `weight.data()` raises a `DeferredInitializationError` until that\n",
    "pass has happened. The cells above respected the order by running `net(X)`\n",
    "before touching `net[0].weight.data()`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-21 fw=all -->\n",
    "\n",
    "## Summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-22 fw=mxnet -->\n",
    "\n",
    "Gluon initializes parameters when you ask it to: `initialize()` records the\n",
    "plan, the first forward pass supplies the shapes, and the default draw,\n",
    "$U(-0.07, 0.07)$, is a legacy constant rather than a fan-aware rule, so\n",
    "override it more readily than elsewhere. The mechanism is one pattern: an\n",
    "`Initializer` object handed to `initialize` on the net, a block, or a single\n",
    "parameter, with `force_reinit=True` to overwrite live values."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-23 fw=all -->\n",
    "\n",
    "On top of Xavier and He,\n",
    "transformer-era code truncates normal tails to bound outliers, shrinks each\n",
    "residual output projection by $1/\\sqrt{N}$ so the stream's variance stays\n",
    "$O(1)$ at any depth, and zero-initializes a block's last layer to start it as\n",
    "the identity."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-24 fw=mxnet -->\n",
    "\n",
    "Anything the menu lacks, truncation included, is an `init.Initializer`\n",
    "subclass with a few lines of array code in its `_init_weight`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-25 fw=all -->\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Instrument the residual stack: record the standard deviation of the\n",
    "   activation after every block (run the stack one block at a time, or\n",
    "   capture per-block activations with the tools of that section)\n",
    "   for the default and scaled treatments at $N=32$, and plot it against\n",
    "   depth. Which curve matches the geometric-growth prediction?\n",
    "1. Zero-initialize *all* layers of every block instead of just the output\n",
    "   projection. The forward pass still returns the input, but what can the\n",
    "   network learn? Work out which parameters receive a nonzero gradient, and\n",
    "   relate the answer to the symmetry-breaking argument of\n",
    "   that section.\n",
    "1. Write an initializer that fills each parameter from a dictionary keyed by\n",
    "   parameter name (walk `net.named_parameters()` in PyTorch, `net.weights`\n",
    "   in TensorFlow, `nnx.iter_modules(net)` in JAX, `net.collect_params()`\n",
    "   in MXNet). You have re-invented part of checkpoint loading,\n",
    "   which that section covers.\n",
    "1. For a normal distribution truncated at $\\pm 2\\sigma$: what fraction of\n",
    "   draws does the clip discard, and by how much does it shrink the standard\n",
    "   deviation? Verify both numbers against the printed output of the truncation\n",
    "   demo above."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-mxnet",
   "display_name": "d2l (MXNet)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
