02b_initializing_jcat.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "4307uKhRx8ob"
   },
   "source": [
    "### Modified version of notebook 02b_initializing.ipynb; see the cell below titled \"What just happened?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "f0fQuvk2x8oj"
   },
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "4WR9qspex8ox"
   },
   "source": [
    "### Why you need a good init"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "L850vNGAx8o0"
   },
   "source": [
    "To understand why initialization is important in a neural net, we'll focus on the basic operation you have there: matrix multiplications. So let's just take a vector `x`, and a matrix `a` initialized randomly, then multiply them 100 times (as if we had 100 layers). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "B5XQUuEnx8o5"
   },
   "outputs": [],
   "source": [
    "x = torch.randn(512)\n",
    "a = torch.randn(512,512)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "VyNySkt_x8pD"
   },
   "outputs": [],
   "source": [
    "for i in range(100): x = a @ x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "c_5uElqJx8pO",
    "outputId": "ad880b3f-5953-4796-abab-bb1a6549f890"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(nan), tensor(nan))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.mean(),x.std()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "kDZT0N-9x8pY"
   },
   "source": [
    "The problem you'll get with that is activation explosion: very soon, your activations will go to nan. We can even ask the loop to break when that first happens:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "yAtw4lsex8pa"
   },
   "outputs": [],
   "source": [
    "x = torch.randn(512)\n",
    "a = torch.randn(512,512)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "XXqC-sfex8pe"
   },
   "outputs": [],
   "source": [
    "for i in range(100): \n",
    "    x = a @ x\n",
    "    if x.std() != x.std(): break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "IL4-47SVx8pk",
    "outputId": "1c0839ab-fe82-4fd2-b79a-687f459d30ce"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "28"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "i"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "IRaDjr6zx8pq"
   },
   "source": [
    "It only takes around 30 multiplications! On the other hand, if you initialize your activations with a scale that is too low, then you'll get another problem:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "zuen8l8gx8ps"
   },
   "outputs": [],
   "source": [
    "x = torch.randn(512)\n",
    "a = torch.randn(512,512) * 0.01"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "8AnLR0Stx8px"
   },
   "outputs": [],
   "source": [
    "for i in range(100): x = a @ x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "4SgcNqvNx8p1",
    "outputId": "0ee4d0b9-ac50-458d-aabb-12c0e72bb851"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(0.), tensor(0.))"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.mean(),x.std()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "MDLRsH-Lx8p5"
   },
   "source": [
    "Here, every activation vanished to 0. So to avoid that problem, people have come with several strategies to initialize their weight matrices, such as:\n",
    "- use a standard deviation that will make sure x and Ax have exactly the same scale\n",
    "- use an orthogonal matrix to initialize the weight (orthogonal matrices have the special property that they preserve the L2 norm, so x and Ax would have the same sum of squares in that case)\n",
    "- use [spectral normalization](https://arxiv.org/pdf/1802.05957.pdf) on the matrix A  (the spectral norm of A is the least possible number M such that `torch.norm(A@x) <= M*torch.norm(x)` so dividing A by this M insures you don't overflow. You can still vanish with this)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "D_BCMmGvx8p6"
   },
   "source": [
    "### The magic number for scaling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "E8R2YMszx8p7"
   },
   "source": [
    "Here we will focus on the first one, which is the Xavier initialization. It tells us that we should use a scale equal to `1/math.sqrt(n_in)` where `n_in` is the number of inputs of our matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "1P0YIIA_x8p8"
   },
   "outputs": [],
   "source": [
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "NhsR07eSx8qB"
   },
   "outputs": [],
   "source": [
    "# multiply weights by the scale factor for Xavier initialization\n",
    "x = torch.randn(512,dtype=torch.float64)\n",
    "a = torch.randn(512,512,dtype=torch.float64) / math.sqrt(512)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Si-B2YhOx8qH"
   },
   "outputs": [],
   "source": [
    "# Now we will simulate a deep network by iteratively applying a linear layer to the input\n",
    "# note that for n_iter much larger than 100, mean and std start to blow up!\n",
    "n_layers = 50\n",
    "for i in range(n_layers): x = a @ x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "A19J8dEbx8qL",
    "outputId": "157560ff-ade3-4475-9798-123076a8a736"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(0.0704, dtype=torch.float64), tensor(1.6490, dtype=torch.float64))"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.mean(),x.std()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "JsQ7ZP5dx8qR"
   },
   "source": [
    "And indeed it works. That is, it works as long as n_layers doesn't get larger than about 50. Why the mean and std blow up for n_layers larger than 50 is not clear to me at this point. Note that this magic number 1/ math.sqrt(512)  isn't very far from the 0.01 we had earlier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "J2u2sld0x8qR",
    "outputId": "e630c413-ed1e-4f8f-da27-d92a60b1654f"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.044194173824159216"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "1/ math.sqrt(512)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "J4642Gh5x8qY"
   },
   "source": [
    "But where does it come from? It's not that mysterious if you remember the definition of the matrix multiplication. When we do `y = a @ x`, the coefficients of `y` are defined by\n",
    "\n",
    "$$y_{i} = a_{i,0} x_{0} + a_{i,1} x_{1} + \\cdots + a_{i,n-1} x_{n-1} = \\sum_{k=0}^{n-1} a_{i,k} x_{k}$$\n",
    "\n",
    "or in code:\n",
    "```\n",
    "y[i] = sum([c*d for c,d in zip(a[i], x)])\n",
    "```\n",
    "\n",
    "Now at the very beginning, our `x` vector has a mean of roughly 0. and a standard deviation of roughly 1. (since we picked it that way)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ZkQumnMix8qZ",
    "outputId": "583db46a-c94a-4670-c963-1642e8ea5051"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(-0.0674), tensor(1.0235))"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.randn(512)\n",
    "x.mean(), x.std()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Bbz39OPxx8qc"
   },
   "source": [
    "NB: This is why it's extremely important to normalize your inputs in Deep Learning, the initialization rules have been designed with inputs that have a mean 0. and a standard deviation of 1.\n",
    "\n",
    "If you need a refresher from your statistics course, the mean is the sum of all the elements divided by the number of elements (a basic average). The standard deviation shows whether the data points stay close to the mean or are far away from it. It's computed by the following formula:\n",
    "\n",
    "$$\\sigma = \\sqrt{\\frac{1}{n}\\left[(x_{0}-m)^{2} + (x_{1}-m)^{2} + \\cdots + (x_{n-1}-m)^{2}\\right]}$$\n",
    "\n",
    "where m is the mean and $\\sigma$ (the greek letter sigma) is the standard deviation. To avoid that square root, we also often consider a quantity called the variance, which is $\\sigma$ squared. \n",
    "\n",
    "Here we have a mean of 0, so the variance is just the mean of x squared, and the standard deviation is its square root.\n",
    "\n",
    "If we go back to `y = a @ x` and assume that we chose weights for `a` that also have a mean of 0, we can compute the variance of `y` quite easily. Since it's random, and we may fall on bad numbers, we repeat the operation 100 times."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "ekZgRXXrx8qe",
    "outputId": "8961dc97-1e91-4df5-e7da-6eec59e12b1e"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(-0.0029423062165346912, 512.098355930394)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mean,sqr = 0.,0.\n",
    "n_iter = 1000\n",
    "for i in range(n_iter):\n",
    "    x = torch.randn(512,dtype=torch.float64)\n",
    "    a = torch.randn(512, 512,dtype=torch.float64)\n",
    "    y = a @ x\n",
    "    # note .item() extracts a float from a tensor\n",
    "    mean += y.mean().item()\n",
    "    sqr  += y.pow(2).mean().item()\n",
    "mean/n_iter,sqr/n_iter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "SZAF7gX_x8qi"
   },
   "source": [
    "Now that looks very close to the dimension of our matrix 512. And that's no coincidence! When you compute y, you sum 512 product of one element of a by one element of x. So what's the mean and the standard deviation of such a product of one element of `a` by one element of `x`? We can show mathematically that as long as the elements in `a` and the elements in `x` are independent, the mean is 0 and the std is 1.\n",
    "\n",
    "This can also be seen experimentally:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "SXXxedB2x8qj",
    "outputId": "340c7110-e843-499a-a108-08f6b064af38"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.004454867896552917, 1.0046282618339664)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mean,sqr = 0.,0.\n",
    "for i in range(10000):\n",
    "    x = torch.randn(1)\n",
    "    a = torch.randn(1)\n",
    "    y = a*x\n",
    "    mean += y.item()\n",
    "    sqr  += y.pow(2).item()\n",
    "mean/10000,math.sqrt(sqr/10000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "7LxGBOOPx8qo"
   },
   "source": [
    "Then we sum 512 of those things that have a mean of zero, and a variance of 1, so we get something that has a mean of 0, and variance of 512. To go to the standard deviation, we have to add a square root, hence `math.sqrt(512)` being our magic number.\n",
    "\n",
    "If we scale the weights of the matrix `a` and divide them by this `math.sqrt(512)`, it will give us a `y` of scale 1, and repeating the product as many times as we want and it won't overflow or vanish."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "4xZ3s36Mx8qp"
   },
   "source": [
    "### Linear transformation of a scalar input, followed by a ReLU"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "OekGJMFcx8qr"
   },
   "source": [
    "We can reproduce the previous experiment with a linear transformation followed by a ReLU,. We see that this time, the mean shifts and the variance becomes 0.5. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "UqZ0oWFex8qs",
    "outputId": "f77a5cbd-17b5-44c8-9a6c-aeadb01d0a1f"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.3210805317212011, 0.5095136192303175)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mean,sqr = 0.,0.\n",
    "n_iter = 100000\n",
    "for i in range(n_iter):\n",
    "    # intitialize input and weights\n",
    "    x = torch.randn(1)\n",
    "    a = torch.randn(1)\n",
    "    # linear transformation\n",
    "    y = a*x\n",
    "    # ReLU\n",
    "    y = 0 if y < 0 else y.item()\n",
    "    mean += y\n",
    "    sqr  += y** 2\n",
    "mean/n_iter,sqr/n_iter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "JlalJ9xax8qv"
   },
   "source": [
    "### What just happened?\n",
    "\n",
    "If you're like me, you want to understand the mathemagic behind the empirically determined results for $mean$ and $sqr$ in the previous cell. Why did the mean shift, and why did the expectation of y**2 change to 0.5? With a bit of effort, we can find out, and it will be good practice for those who want to develop the ability to read papers.\n",
    "\n",
    "On the other hand if you'd rather not slog through a a few screen of equations, you can accept the empirical result and move on to the bottom of this cell and read the few lines below the TL;DR...\n",
    "\n",
    "We start with $a$ and $x$, which are $\\textbf{independent Gaussian random variables}$ with zero mean and unit variance, and $y = a \\cdot x$\n",
    "\n",
    "The effect of the Relu layer is that $y$ is drawn from a $\\textbf{mixture distribution}$: half the time $x \\lt 0$ so $y$ is 'clamped' to (i.e. replaced by) zero; the other half the time $x \\ge 0$, so $y$ is drawn from a positive half-Gaussian (ref. https://en.wikipedia.org/wiki/Half-normal_distribution).\n",
    "\n",
    "How can we calculate $mean = E[y]$, the $\\textbf{expectation value of y}$? \n",
    "\n",
    "Accounting for the $\\textbf{mixture distribution}$, we have\n",
    "\n",
    "$mean = E[y] = 0.5 \\cdot E[0] + 0.5 E[a \\cdot x] = 0.5 \\cdot 0 + 0.5 E[a] \\cdot E[x] = 0.5 E[x]^2$, the last step follows because $a$ and $x$ are independent random variables which have the same distribution. So to calculate $mean$, we need to evaluate $E[x]$, the $\\textbf{expectation value}$ of $x$ under the half-Gaussian distribution. \n",
    "\n",
    "The $\\textbf{normalized probability density function}$ for the half-Gaussian is\n",
    "\n",
    "$ p_{half}(x) =  \\frac{1}{\\sigma}\\sqrt{\\frac{2}{\\pi}} \\cdot exp{\\frac{-x^2}{2\\sigma^2}} $, where support is on the positive half-line $0 \\leq x < \\infty$\n",
    "\n",
    "By $\\textbf{normalized}$, we mean that $\\int_0^{\\infty} p_{half}(x)dx = \\int_0^{\\infty}\\frac{1}{\\sigma}\\sqrt{\\frac{2}{\\pi}} \\cdot exp{\\frac{-x^2}{2\\sigma^2}dx} = 1$\n",
    "\n",
    "$ E[x] = \\frac{1}{\\sigma}\\sqrt{\\frac{2}{\\pi}} \\cdot \\int_0^{\\infty} x \\cdot exp{\\frac{-x^2}{2\\sigma^2}dx}   = \\sigma \\sqrt{\\frac{2}{\\pi}} = \\sqrt{\\frac{2}{\\pi}}$, since $\\sigma = 1$\n",
    "\n",
    "Therefore $mean = E[y] = 0.5E[x]^2 = \\frac{1}{\\pi} \\approx 0.318$\n",
    "\n",
    "Next we turn to the calculation of $sqr$, the $\\textbf{expectation value}$ of $y^2$. \n",
    "\n",
    "$E[y^2] = E[a \\cdot x \\cdot a \\cdot x] = E[a^2] \\cdot E[x^2] = E[x^2]^2$; the last step again follows because $a$ and $x$ are independent random variables.\n",
    "\n",
    "Let's evaluate $E[x^2]$, the expectation value of $x^2$ under the half-Gaussian distribution: \n",
    "\n",
    "$ E[x^2] = \\frac{1}{\\sigma}\\sqrt{\\frac{2}{\\pi}} \\cdot \\int_0^{\\infty} x^2 \\cdot exp{\\frac{-x^2}{2\\sigma^2}dx} $\n",
    "\n",
    "Reflection of the half-Gaussian distribution about the vertical axis (i.e. extending to the negative half-line) produces a full Gaussian, apart from a scale factor:\n",
    "$\\int_{-\\infty}^{\\infty} p_{half}(x)dx = 2$, so that $0.5 \\cdot \\int_{-\\infty}^{\\infty} p_{half}(x)dx = \\int_{-\\infty}^{\\infty} p_{full}(x)dx = 1$, so \n",
    "\n",
    "$ p_{full}(x) =  0.5 \\cdot p_{half}(x)$ \n",
    "\n",
    "The variance of a zero mean unit normal (full) Gaussian is 1 by definition:\n",
    "\n",
    "$\\int_{-\\infty}^{\\infty} x^2 \\cdot p_{full}(x)dx = 1$\n",
    "\n",
    "By symmetry we can write this as $2 \\cdot \\int_{0}^{\\infty} x^2 \\cdot p_{full}(x)dx = 2\\cdot \\int_{0}^{\\infty} x^2 \\cdot 0.5 \\cdot p_{half}(x)dx = \\int_{0}^{\\infty} x^2 \\cdot p_{half}(x)dx = 1$\n",
    "\n",
    "But the last expression is just the expectation value of $x^2$ under the half-Gaussian distribution; so we've shown that $E[x^2] = 1$\n",
    "\n",
    "Now we can compute $E[y^2]$, from our above work: $E[y^2]= E[x^2]^2 = 1$\n",
    "\n",
    "Finally, taking account that $E[y^2]$ is to be computed under the $\\textbf{mixture distribution}$:\n",
    "\n",
    "$sqr = 0.5 \\cdot 0 + 0.5 E[y^2] = 0. + 0.5 \\cdot 1 = 0.5$\n",
    "\n",
    "TL;DR\n",
    "\n",
    "We've shown analytically that the output of a linear layer followed by a Relu has $E[mean] = \\frac{1}{\\pi}$ and $E[sqr] = 0.5$, if the inputs and weights are initialized by sampling from a unit normal Gaussian.\n",
    "\n",
    "This agrees with the empirical results from the previous cell.\n",
    "\n",
    "One could easily extend this derivation for the (mathematically simpler) case of initializing with a $\\textbf{uniform distribution}$. Try it!\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "-SGtA9RVx8qw"
   },
   "source": [
    "### Linear matrix transformation of a vector input, followed by a ReLU\n",
    "What happens when we run the experiment on the whole matrix product, instead of on the product of scalars?\n",
    "\n",
    "$sqr$ is the sum of $512$ independent random variables, each of which has expectation value of 0.5. Since the expectation value of a sum of random variables is the sum of the expectation values , we should get $E[sqr]  = 512 \\cdot 0.5 = 256$. Indeed, this turns out to be the case.\n",
    "\n",
    "However, $mean$ is also the sum of $512$ independent random variables, each of which has expectation of $\\frac{1}{\\pi}$, as we showed above. Therefore,  $mean$ should scale by $512$ to become $E[mean] =\\frac{512}{\\pi} \\approx 163$. Instead we find that $mean \\approx 9$. I cannot explain this; I'm still turning it over in my mind to understand why."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "C1a_ZHH2x8q1",
    "outputId": "821d0dd2-c17c-4e4d-f6cd-ff2581c64176"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "512 9.035336665630341 256.6674785644531\n"
     ]
    }
   ],
   "source": [
    "# Since sqr is the sum of 512 random variables, each of value 0.5, I expect it to scale by a factor of 512\n",
    "#    to become 512*0.5 = 256. This is the case\n",
    "# I also expected the mean to scale by 512, so that mean = 512/pi ~ 163; I cannot explain why empirically, mean ~ 9\n",
    "n_iter = 10000\n",
    "n_dim = 512\n",
    "mean,sqr = 0.,0.\n",
    "for i in range(n_iter):\n",
    "    x = torch.randn(n_dim)\n",
    "    a = torch.randn(n_dim,n_dim)\n",
    "    y = a@x\n",
    "    y = y.clamp(min==0)\n",
    "    mean += y.mean().item()\n",
    "    # Note the corrected \n",
    "    sqr  += y.pow(2).mean().item()\n",
    "print(n_dim,mean/n_iter,sqr/n_iter)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Gd0FaEmUx8q-"
   },
   "source": [
    "### We can scale the weights by a 'magic number' $\\sqrt\\frac{2}{512}$ chosen to make $E[sqr] = 1$ \n",
    "\n",
    "\n",
    "The expectation  value of the product of a scalar $a$ with a random variable $X$ is \n",
    "\n",
    "$E(aX) = a \\cdot E(X)$\n",
    "\n",
    "\n",
    "With this rule, we can evaluate the expectation value of $(aX)^2$:\n",
    "\n",
    "$E[(aX)^2] = E[a^2 \\cdot X^2] = a^2 \\cdot E[X^2]$. So scaling the weights by a factor of $a$ scales this expectation value by $a^2$\n",
    "\n",
    "The variance of the product of a scalar $a$ and a random variable $X$ is:\n",
    "\n",
    "$var(aX)  = a^2 \\cdot var(X) $ \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "m7c3Bf4Jx8q-",
    "outputId": "be14cb50-8f0c-4bd3-8aed-e980e9b30c2a"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.5986050982549787, 1.0889559798707833)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# check: sqr should scale by 2/512 --> 0.5*2/512 = 1\n",
    "# check: mean should scale by sqrt(2/512) --> 9.016*sqrt(2/512) ~ 0.56\n",
    "mean,sqr = 0.,0.\n",
    "n_iter = 1000\n",
    "n_dim = 512\n",
    "for i in range(n_iter):\n",
    "    x = torch.randn(n_dim)\n",
    "    a = torch.randn(n_dim) * math.sqrt(2/n_dim)\n",
    "    y = a @ x\n",
    "    y = y.clamp(min=0)\n",
    "    mean += y.mean().item()\n",
    "    sqr  += y.pow(2).mean().item()\n",
    "mean/n_iter,sqr/n_iter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "GypIEQhiJX6j"
   },
   "source": [
    "### Let's see how well the new initialization we designed for a linear layer followed by a ReLU activation works: \n",
    "We simulate a deep network by repeatedly applying a linear layer followed by a ReLU activation. After each ReLU, randomly initialize the weights with a standard normal distribution multiplied by our \"magic number\". \n",
    "\n",
    "We find that the mean and variance remain fairly stable for up to about $10$ to  $20$  layers. \n",
    "\n",
    "However I'd expected $mean$ and $sqr$ to be stable over **hundreds if not thousands of layers**, and this is clearly not the case. I'm at a loss to explain why.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "Kpb7UM82GE_x",
    "outputId": "dac75965-3af7-440c-82c3-51fd5f7020ba"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.5753585098577015, 1.146278557410716)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# what happens when we repeatedly apply a linear layer followed by a ReLU\n",
    "# with our new initialization?\n",
    "# \n",
    "\n",
    "# control parameters\n",
    "n_layers = 10\n",
    "n_dim = 512\n",
    "n_iter = 100\n",
    "\n",
    "# initialize accumulators\n",
    "yy = 0\n",
    "zz = 0\n",
    "\n",
    "\n",
    "# outer loop over n_iter: repeat the experiment n_iter times and take\n",
    "#     avearages of mean and sqr\n",
    "for j in range(n_iter):\n",
    "\n",
    "    # initialize accumulators\n",
    "    mean, sqr = 0.,0.\n",
    "    \n",
    "    # intitialize weights with the \"magic number\" to impose unit variance\n",
    "    x = torch.randn(n_dim, dtype=torch.float64)\n",
    "    a = torch.randn(n_dim,n_dim,dtype=torch.float64) * math.sqrt(2/n_dim)\n",
    "    \n",
    "    # simulate the effect of a deep network by repeatedly applying\n",
    "    #     a linear layer followed by a ReLU\n",
    "    for i in range(n_layers): \n",
    "        x = a @ x\n",
    "        x = x.clamp(min=0)\n",
    "        mean = x.mean().item()\n",
    "        sqr  = x.pow(2).mean().item()\n",
    "\n",
    "    # accumulate sums of mean and sqr\n",
    "    yy += mean\n",
    "    zz += sqr\n",
    "\n",
    "# averages of mean and sqr over n_iter experiments    \n",
    "yy/n_iter,zz/n_iter\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "QnkZHjNex8rC"
   },
   "source": [
    "The math behind is a tiny bit more complex, and you can find everything in the [Kaiming](https://arxiv.org/abs/1502.01852) and the [Xavier](http://proceedings.mlr.press/v9/glorot10a.html) paper but this gives the intuition behind those results."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "02b_initializing_jcat.ipynb",
   "provenance": [],
   "version": "0.3.2"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}