{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Installing packages:\n",
      "\t.package(path: \"/home/jupyter/notebooks/swift/FastaiNotebook_02a_why_sqrt5\")\n",
      "\t\tFastaiNotebook_02a_why_sqrt5\n",
      "With SwiftPM flags: []\n",
      "Working in: /tmp/tmp1on2fw8s/swift-install\n",
      "[1/6] Compiling FastaiNotebook_02a_why_sqrt5 01_matmul.swift\n",
      "[2/6] Compiling FastaiNotebook_02a_why_sqrt5 02_fully_connected.swift\n",
      "[3/6] Compiling FastaiNotebook_02a_why_sqrt5 02a_why_sqrt5.swift\n",
      "[4/6] Compiling FastaiNotebook_02a_why_sqrt5 00_load_data.swift\n",
      "[5/6] Compiling FastaiNotebook_02a_why_sqrt5 01a_fastai_layers.swift\n",
      "[6/7] Merging module FastaiNotebook_02a_why_sqrt5\n",
      "[7/8] Compiling jupyterInstalledPackages jupyterInstalledPackages.swift\n",
      "[8/9] Merging module jupyterInstalledPackages\n",
      "[9/9] Linking libjupyterInstalledPackages.so\n",
      "Initializing Swift...\n",
      "Installation complete!\n"
     ]
    }
   ],
   "source": [
    "%install-location $cwd/swift-install\n",
    "%install '.package(path: \"$cwd/FastaiNotebook_02a_why_sqrt5\")' FastaiNotebook_02a_why_sqrt5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "//export\n",
    "import Path\n",
    "import TensorFlow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import FastaiNotebook_02a_why_sqrt5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Why you need a good init"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "var x = TF(randomNormal: [512, 1])\n",
    "let a = TF(randomNormal: [512,512])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in 0..<100 { x = a • x }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "▿ 2 elements\n",
       "  - .0 : nan(0x1fffff)\n",
       "  - .1 : nan(0x1fffff)\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(x.mean(),x.std())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "var x = TF(randomNormal: [512, 1])\n",
    "let a = TF(randomNormal: [512,512])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "27\r\n"
     ]
    }
   ],
   "source": [
    "for i in 0..<100 {\n",
    "    x = a • x\n",
    "    if x.std().scalarized().isNaN {\n",
    "        print(i)\n",
    "        break\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "var x = TF(randomNormal: [512, 1])\n",
    "let a = TF(randomNormal: [512,512]) * 0.01"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in 0..<100 { x = a • x }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "▿ 2 elements\n",
       "  - .0 : 0.0\n",
       "  - .1 : 0.0\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(x.mean(),x.std())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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 `matmul(A,x).norm() <= M*x.norm()` so dividing A by this M insures you don't overflow. You can still vanish with this)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The magic number for scaling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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/sqrt(n_in)` where `n_in` is the number of inputs of our matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "var x = TF(randomNormal: [512, 1])\n",
    "let a = TF(randomNormal: [512,512]) / sqrt(512)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in 0..<100 { x = a • x }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "▿ 2 elements\n",
       "  - mean : 0.061417937\n",
       "  - std : 1.4370023\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(mean: x.mean(), std: x.std())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And indeed it works. Note that this magic number isn't very far from the 0.01 we had earlier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.044194173824159216\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "1 / sqrt(512)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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 = matmul(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",
    "for i in 0..<a.shape[0] {\n",
    "    for k in 0..<b.shape[1] {\n",
    "        y[i][0] += a[i][k] * x[k][0]\n",
    "    }\n",
    "}\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": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "▿ 2 elements\n",
       "  - mean : -0.013813974\n",
       "  - std : 0.9675647\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "var x = TF(randomNormal: [512, 1])\n",
    "(mean: x.mean(), std: x.std())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "▿ 2 elements\n",
       "  - .0 : 0.037324883\n",
       "  - .1 : 512.3489\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "var mean = Float()\n",
    "var sqr = Float()\n",
    "for i in 0..<100 {\n",
    "    let x = TF(randomNormal: [512, 1])\n",
    "    let a = TF(randomNormal: [512, 512])\n",
    "    let y = a • x\n",
    "    mean += y.mean().scalarized()\n",
    "    sqr  += pow(y, 2).mean().scalarized()\n",
    "}\n",
    "(mean/100, sqr/100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "▿ 2 elements\n",
       "  - .0 : -0.00084701786\n",
       "  - .1 : 0.9972176\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "var mean = Float()\n",
    "var sqr = Float()\n",
    "for i in 0..<10000 {\n",
    "    let x = TF(randomNormal: [])\n",
    "    let a = TF(randomNormal: [])\n",
    "    let y = a * x\n",
    "    mean += y.scalarized()\n",
    "    sqr  += pow(y, 2).scalarized()\n",
    "}\n",
    "(mean/10000,sqrt(sqr/10000))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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 `sqrt(512)` being our magic number.\n",
    "\n",
    "If we scale the weights of the matrix `a` and divide them by this `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": {},
   "source": [
    "### Adding ReLU in the mix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can reproduce the previous experiment with a ReLU, to see that this time, the mean shifts and the variance becomes 0.5. This time the magic number will be `math.sqrt(2/512)` to properly scale the weights of the matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "▿ 2 elements\n",
       "  - mean : 0.33342138\n",
       "  - sqrt : 0.537505\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "var mean = Float()\n",
    "var sqr = Float()\n",
    "for i in 0..<10000 {\n",
    "    let x = TF(randomNormal: [])\n",
    "    let a = TF(randomNormal: [])\n",
    "    var y = (a*x).scalarized()\n",
    "    y = y < 0 ? 0 : y\n",
    "    mean += y\n",
    "    sqr  += pow(y, 2)\n",
    "}\n",
    "(mean: mean/10000, sqrt: sqr/10000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can double check by running the experiment on the whole matrix product. The variance becomes 512/2 this time:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "▿ 2 elements\n",
       "  - mean : 8.947798\n",
       "  - sqrt : 252.94208\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "var mean = Float()\n",
    "var sqr = Float()\n",
    "for i in 0..<100 {\n",
    "    let x = TF(randomNormal: [512, 1])\n",
    "    let a = TF(randomNormal: [512, 512])\n",
    "    var y = a • x\n",
    "    y = max(y, TF(zeros: y.shape))\n",
    "    mean += y.mean().scalarized()\n",
    "    sqr  += pow(y, 2).mean().scalarized()\n",
    "}\n",
    "(mean: mean/100, sqrt: sqr/100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or that scaling the coefficient with the magic number gives us a scale of 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "▿ 2 elements\n",
       "  - mean : 0.5630086\n",
       "  - sqrt : 0.99951345\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "var mean = Float()\n",
    "var sqr = Float()\n",
    "for i in 0..<100 {\n",
    "    let x = TF(randomNormal: [512, 1])\n",
    "    let a = TF(randomNormal: [512, 512]) * sqrt(2/512)\n",
    "    var y = a • x\n",
    "    y = max(y, TF(zeros: y.shape))\n",
    "    mean += y.mean().scalarized()\n",
    "    sqr  += pow(y, 2).mean().scalarized()\n",
    "}\n",
    "(mean: mean/100, sqrt: sqr/100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Swift",
   "language": "swift",
   "name": "swift"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}