Bayesian Deep Learning part1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MLSS2019: Bayesian Deep Learning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this tutorial we will learn what basic building blocks are needed\n",
    "to endow (deep) neural networks with uncertainty estimates."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The plan of the tutorial\n",
    "1. [Setup and imports](#Setup-and-imports)\n",
    "2. [Easy uncertainty in networks](#Easy-uncertainty-in-networks)\n",
    "   1. [Bayesification via dropout and weight decay](#Bayesification-via-dropout-and-weight-decay)\n",
    "   2. [Implementing function sampling with the DropoutLinear Layer](#Implementing-function-sampling-with-the-DropoutLinear-Layer)\n",
    "   3. [Implementing-DropoutLinear](#Implementing-DropoutLinear)\n",
    "   4. [Comparing sample functions to point-estimates](#Comparing-sample-functions-to-point-estimates)\n",
    "3. [(optional) Dropout $2$-d Convolutional layer](#(optional)-Dropout-$2$-d-Convolutional-layer)\n",
    "4. [(optional) A brief reminder on Bayesian and Variational Inference](#(optional)-A-brief-reminder-on-Bayesian-and-Variational-Inference)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**(note)**\n",
    "* to view documentation on something  type in `something?` (with one question mark)\n",
    "* to view code of something type in `something??` (with two question marks)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup and imports"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section we import necessary modules and functions and\n",
    "define the computational device."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we install some boilerplate service code for this tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q --upgrade git+https://github.com/ivannz/mlss2019-bayesian-deep-learning.git"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, numpy for computing, matplotlib for plotting and tqdm for progress bars."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tqdm\n",
    "import numpy as np\n",
    "\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For deep learning stuff will be using [pytorch](https://pytorch.org/).\n",
    "\n",
    "If you are unfamiliar with it, it is basically like `numpy` with autograd,\n",
    "stricter data type enforcement, native GPU support, and tools for building\n",
    "training and serializing models.\n",
    "<!-- (and with `axis` argument replaced with `dim` :) -->\n",
    "\n",
    "There are good introductory tutorials on `pytorch`, like this\n",
    "[one](https://pytorch.org/tutorials/beginner/deep_learning_60min_blitz.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will need some functionality from scikit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import confusion_matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we import the boilerplate code.\n",
    "\n",
    "* a procedure that implements a minibatch SGD **fit** loop\n",
    "* a function, that **evaluates** the model on the provided dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mlss2019bdl import fit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "# pseudocode\n",
    "def fit(model, dataset, criterion, ...):\n",
    "    for epoch in epochs:\n",
    "        for batch in dataset:\n",
    "            loss = criterion(model, batch)  # forward pass\n",
    "\n",
    "            grad = loss.backward()          # gradient via back propagation\n",
    "\n",
    "            adam_step(grad)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mlss2019bdl import predict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "# pseudocode\n",
    "def predict(model, dataset, ...):\n",
    "    for input_batch in dataset:\n",
    "        output.append(model(input_batch))  # forward pass\n",
    "    \n",
    "    return concatenate(output)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Easy uncertainty in networks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate the initial small dataset $S_0 = (x_i, y_i)_{i=1}^{m_0}$\n",
    "with $y_i = g(x_i)$, $x_i$ on a regular-spaced grid, and $\n",
    "g\n",
    "    \\colon \\mathbb{R} \\to \\mathbb{R}\n",
    "    \\colon x \\mapsto \\tfrac{x^2}4 + \\sin \\frac\\pi2 x\n",
    "$.\n",
    "<!--\n",
    "`dataset_from_numpy` **converts** numpy arrays into torch tensors,\n",
    "**places** them on the specified compute device, **and packages**\n",
    "into a dataset\n",
    "-->"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mlss2019bdl import dataset_from_numpy\n",
    "\n",
    "X_train = np.linspace(-6.0, +6.0, num=20)[:, np.newaxis]\n",
    "y_train = np.sin(X_train * np.pi / 2) + 0.25 * X_train**2\n",
    "\n",
    "train = dataset_from_numpy(X_train, y_train, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_domain = np.linspace(-10., +10., num=251)[:, np.newaxis]\n",
    "\n",
    "domain = dataset_from_numpy(X_domain, device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Suppose we have the following model: a 3-layer fully connected\n",
    "network with LeakyReLU activations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.nn import Linear, Sequential\n",
    "from torch.nn import LeakyReLU\n",
    "\n",
    "\n",
    "model = Sequential(\n",
    "    Linear(1, 512, bias=True),\n",
    "    LeakyReLU(),\n",
    "\n",
    "    Linear(512, 512, bias=True),\n",
    "    LeakyReLU(),\n",
    "\n",
    "    Linear(512, 1, bias=True),\n",
    ")\n",
    "\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We fit our model on `train` using MSE loss and $\\ell_2$ penalty on\n",
    "weights (`weight_decay`):\n",
    "$$\n",
    "    \\tfrac1{2 m} \\|f_\\omega(x) - y\\|_2^2 + \\lambda \\|\\omega\\|_2^2\n",
    "    \\,, $$\n",
    "where $\\omega$ are all the learnable parameters of the network $f_\\omega$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fit, ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "fit(model, train, criterion=\"mse\", n_epochs=2000, verbose=True, weight_decay=1e-3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "..., compute the predictions, ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred = predict(model, domain)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "..., and plot them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(12, 5))\n",
    "\n",
    "ax.scatter(X_train, y_train, c=\"black\", s=40, label=\"train\")\n",
    "\n",
    "ax.plot(X_domain, y_pred.numpy(), c=\"C0\", lw=2, label=\"prediction\")\n",
    "\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This model seems to fit the train set adequately well. However, there is no\n",
    "way to assess how confident this model is with respect to its predictions.\n",
    "Indeed, the prediction $\\hat{y}_x = f_\\omega(x)$ is is a deterministic function\n",
    "of the input $x$ and the learnt parameters $\\omega$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `Bayesification` via dropout and weight decay"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One inexpensive way to make any network into a stochastic function of its\n",
    "input is to add dropout before any parameterized layer like `linear`\n",
    "or `convolutional`, [Hinton et al. 2012](https://arxiv.org/abs/1207.0580).\n",
    "Essentially, dropout applies a Bernoulli mask to the features of the input.\n",
    "\n",
    "In [Gal, Y. (2016)](http://www.cs.ox.ac.uk/people/yarin.gal/website/thesis/thesis.pdf)\n",
    "it has been shown that a simple, somewhat ad-hoc approach of\n",
    "adding uncertainty quantification to networks through dropout,\n",
    "coupled with $\\ell_2$ weight penalty, is a special case of Variational Inference."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For input\n",
    "$\n",
    "    x\\in \\mathbb{R}^{[\\mathrm{in}]}\n",
    "$ the dropout layer acts like this:\n",
    "\n",
    "$$\n",
    "    y_j = x_j \\, m_j\n",
    "    \\,, $$\n",
    "\n",
    "where $m\\in \\mathbb{R}^{[\\mathrm{in}]}$ with $\n",
    "m_j \\sim \\pi_p(m_j)\n",
    "    = \\mathcal{Ber}\\bigl(\\bigl\\{0, \\tfrac1{1-p}\\bigr\\}, 1-p\\bigr)\n",
    "$,\n",
    "i.e. equals $\\tfrac1{1-p}$ with probability $1-p$ and $0$ otherwise."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### (task) Always Active Dropout"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Useful methods:\n",
    "* `torch.rand(d1, ..., dn)` -- draw $d_1\\times \\ldots \\times d_n$ tensor of uniform rv-s\n",
    "* `torch.rand_like(other)` -- draw a tensor of uniform rv-s with the shape, data type and device as `other`\n",
    "\n",
    "\n",
    "* `torch.bernoulli(pi)` -- draw tensor $t$ with independent $\n",
    "t_\\alpha \\sim \\mathcal{Ber}\\bigl(\\{0, 1\\}, \\pi_\\alpha\\bigr)\n",
    "$ for each index $\\alpha$\n",
    "* `torch.full((d1, ..., dn), v)` -- a $d_1\\times \\ldots \\times d_n$ tensor with the same value $v$\n",
    "\n",
    "\n",
    "* `Tensor.to(other)` -- assume move `Tensor` to the device of the `other` and cast to its data type."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.nn import Module\n",
    "\n",
    "class ActiveDropout(Module):\n",
    "    # all building blocks of networks are inherited from Module!\n",
    "\n",
    "    def __init__(self, p=0.5):\n",
    "        super().__init__()  # init the base class\n",
    "\n",
    "        self.p = p\n",
    "\n",
    "    def forward(self, input):\n",
    "        ## Exercise: implement feature dropout on input\n",
    "        #  self.p - contains the specified dropout rate\n",
    "        \n",
    "        mask = torch.rand_like(input) > self.p\n",
    "        return input * mask.to(input) / (1 - self.p)\n",
    "\n",
    "        # prob = torch.full_like(input, 1 - self.p)\n",
    "        # return input * torch.bernoulli(prob) / prob\n",
    "\n",
    "        # return F.dropout(input, self.p, True)\n",
    "\n",
    "        pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### (task) Rebuilding the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's recreate the model above with this freshly minted dropout layer.\n",
    "Then fit and plot it's prediction uncertainty due to forward pass stochasticity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_model(p=0.5):\n",
    "    \"\"\"Build a model with dropout layers' rate set to `p`.\"\"\"\n",
    "\n",
    "    return Sequential(\n",
    "        ## Exercise: Use `ActiveDropout` before linear layers of our\n",
    "        #  first network. Note that dropping out inputs is not a good idea\n",
    "\n",
    "        Linear(1, 512, bias=True),\n",
    "        LeakyReLU(),\n",
    "\n",
    "        ActiveDropout(p),\n",
    "        Linear(512, 512, bias=True),\n",
    "        LeakyReLU(),\n",
    "\n",
    "        ActiveDropout(p),\n",
    "        Linear(512, 1, bias=True),\n",
    "\n",
    "        # pass\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = build_model(p=0.5)\n",
    "\n",
    "model.to(device)\n",
    "\n",
    "fit(model, train, criterion=\"mse\", n_epochs=2000, verbose=True,\n",
    "    weight_decay=1e-3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Sampling the random output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's take the test sample $\\tilde{S} = (\\tilde{x}_i)_{i=1}^m \\in \\mathcal{X}$\n",
    "and repeat the stochastic forward pass $B$ times at each $x\\in \\tilde{S}$:\n",
    "\n",
    "* for $b = 1 .. B$ do:\n",
    "\n",
    "  1. draw $y_{bi} \\sim f_\\omega(\\tilde{x}_i)$ for $i = 1 .. m$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def point_estimate(model, dataset, n_samples=1, verbose=False):\n",
    "    \"\"\"Draw pointwise samples with stochastic forward pass.\"\"\"\n",
    "\n",
    "    outputs = []\n",
    "    for sample in tqdm.tqdm(range(n_samples), disable=not verbose):\n",
    "\n",
    "        outputs.append(predict(model, dataset))\n",
    "\n",
    "    return torch.stack(outputs, dim=0)\n",
    "\n",
    "\n",
    "samples = point_estimate(model, domain, n_samples=101, verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The approximate $95\\%$ confidence band of predictions is..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(12, 5))\n",
    "ax.scatter(X_train, y_train, c=\"black\", s=40, label=\"train\")\n",
    "\n",
    "mean, std = samples.mean(dim=0).numpy(), samples.std(dim=0).numpy()\n",
    "ax.plot(X_domain, mean + 1.96 * std, c=\"k\")\n",
    "ax.plot(X_domain, mean - 1.96 * std, c=\"k\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Implementing function sampling with the DropoutLinear Layer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's inspect the draws $y_{bi}$ as $B$ functional samples:\n",
    "$(x_i, y_{bi})_{i=1}^m$ - the $b$-th sample path. Below we\n",
    "plot $5$ random paths."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = point_estimate(model, domain, n_samples=101, verbose=True)\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(12, 5))\n",
    "\n",
    "ax.scatter(X_train, y_train, c=\"black\", s=40, label=\"train\")\n",
    "ax.plot(X_domain[:, 0], samples[:5, :, 0].numpy().T, c=\"C0\", lw=1, alpha=0.25);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is clear that they are very erratic!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Computing stochastic forward passes with a new mask each time is equivalent\n",
    "to drawing new **independent** prediction from for each point $x\\in \\tilde{S}$,\n",
    "without considering that, in fact, at adjacent points the predictions should\n",
    "be correlated. If we were interested in uncertainty at some particular point,\n",
    "this would be okay: **fast and simple**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, if we are interested in the uncertainty of an integral **path-dependent**\n",
    "measure of the whole estimated function, or are doing **optimization** of\n",
    "the unknown true function taking estimation uncertainty into account, then\n",
    "this clearly erratic behaviour of paths is undesirable. Ex. see\n",
    "[blog: Gal, Y. 2016](http://www.cs.ox.ac.uk/people/yarin.gal/website/blog_2248.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to implement some extra functionality on top of `pytorch`,\n",
    "in order to draw realizations from the induced distribution over\n",
    "functions, defined by a network, i.e. $\n",
    "\\bigl\\{\n",
    "    f_\\omega\\colon \\mathcal{X}\\to\\mathcal{Y}\n",
    "\\bigr\\}_{\\omega \\sim q(\\omega)}\n",
    "$\n",
    "where $q(\\omega)$ is a distribution over the parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One of the design approaches is to allow layers\n",
    "to cache random draws of their parameters for reuse\n",
    "in all subsequent forward passes, until this is no\n",
    "longer needed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Freeze/unfreeze interface"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a base **trait-class** `FreezableWeight` that adds interface\n",
    "for freezing and unfreezing layer's random **weight** parameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FreezableWeight(Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.unfreeze()\n",
    "\n",
    "    def unfreeze(self):\n",
    "        self.register_buffer(\"frozen_weight\", None)\n",
    "\n",
    "    def is_frozen(self):\n",
    "        \"\"\"Check if a frozen weight is available.\"\"\"\n",
    "        return isinstance(self.frozen_weight, torch.Tensor)\n",
    "\n",
    "    def freeze(self):\n",
    "        \"\"\"Sample from the parameter distribution and freeze.\"\"\"\n",
    "        raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we declare a pair of functions:\n",
    "* `freeze()` instructs each compatible layer of the model to **sample and freeze** its randomness\n",
    "* `unfreeze()` requests the layers to **undo** this"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def freeze(model):\n",
    "    for layer in model.modules():\n",
    "        if isinstance(layer, FreezableWeight):\n",
    "            layer.freeze()\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def unfreeze(model):\n",
    "    for layer in model.modules():\n",
    "        if isinstance(layer, FreezableWeight):\n",
    "            layer.unfreeze()\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### (task) Sampling realizations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The algorithm to sample a random function is:\n",
    "* for $b = 1... B$ do:\n",
    "\n",
    "  1. draw an independent realization $f_b\\colon \\mathcal{X} \\to \\mathcal{Y}$\n",
    "  with from the process $\\{f_\\omega\\}_{\\omega \\sim q(\\omega)}$\n",
    "  2. get $\\hat{y}_{bi} = f_b(\\tilde{x}_i)$ for $i=1 .. m$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_function(model, dataset, n_samples=1, verbose=False):\n",
    "    \"\"\"Draw a realization of a random function.\"\"\"\n",
    "\n",
    "    ## Exercise: code a function similar to `point_estimate()`,\n",
    "    ##  that collects the predictions from `frozen` models. Don't\n",
    "    ##  forget to unfreeze before returning.\n",
    "\n",
    "    outputs = []\n",
    "    for _ in tqdm.tqdm(range(n_samples), disable=not verbose):\n",
    "        freeze(model)\n",
    "\n",
    "        outputs.append(predict(model, dataset))\n",
    "\n",
    "    unfreeze(model)\n",
    "\n",
    "    return torch.stack(outputs, dim=0)\n",
    "\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**(note)** although the internal loop in both functions looks\n",
    "similar they, conceptually the functions differ:\n",
    "<strong>\n",
    "```python\n",
    "def point_estimate(f, S):\n",
    "    for x in S:\n",
    "        for w from f.q:  # different w for different x\n",
    "            yield f(x, w)\n",
    "\n",
    "\n",
    "def sample_function(f, S):\n",
    "    for w from f.q:\n",
    "        for x in S:      # same w for different x (thanks to freeze)\n",
    "            yield f(x, w)\n",
    "```\n",
    "</strong>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Implementing `DropoutLinear`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we will merge `ActiveDropout` and `Linear` layers into one, which\n",
    "\n",
    "1. (on forward pass) **drops out** the inputs, if necessary, and **applies** the linear (affine) transform\n",
    "2. (on freeze) **randomly zeros** columns in a copy of the the weight matrix $W$\n",
    "\n",
    "Preferably, we will try to preserve interface, so that the resulting\n",
    "object is backwards compatible with `Linear`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This way we would be able to draw realizations from the induced\n",
    "distribution over functions defined by the network $\n",
    "\\bigl\\{\n",
    "    f_\\omega\\colon \\mathcal{X}\\to\\mathcal{Y}\n",
    "\\bigr\\}_{\\omega \\sim q(\\omega)}\n",
    "$\n",
    "where $q(\\omega)$ a distribution over the network parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### (task) Fused dropout-linear operation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "On the inputs into a linear layer dropout acts like this: for input\n",
    "$\n",
    "    x\\in \\mathbb{R}^{[\\mathrm{in}]}\n",
    "$ and layer weights $\n",
    "    W\\in \\mathbb{R}^{[\\mathrm{out}] \\times [\\mathrm{in}]}\n",
    "$\n",
    "and bias $\n",
    "    b\\in \\mathbb{R}^{[\\mathrm{out}]}\n",
    "$ the resulting effect is\n",
    "\n",
    "$$\n",
    "    \\tilde{x} = x \\odot m\n",
    "    \\,, \\\\\n",
    "    y = \\tilde{x} W^\\top + b\n",
    "%     = b + \\sum_i x_i m_i W_i\n",
    "    \\,, $$\n",
    "\n",
    "where $\\odot$ is the elementwise product and $m\\in \\mathbb{R}^{[\\mathrm{in}]}$\n",
    "with $m_j \\sim \\pi_p(m_j) = \\mathcal{Ber}\\bigl(\\bigl\\{0, \\tfrac1{1-p}\\bigr\\}, 1-p\\bigr)$,\n",
    "i.e. equals $\\tfrac1{1-p}$ with probability $1-p$ and $0$ otherwise."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let\n",
    "$\n",
    "    x\\in \\mathbb{R}^{[\\mathrm{in}]}\n",
    "$, $\n",
    "    W\\in \\mathbb{R}^{[\\mathrm{out}] \\times [\\mathrm{in}]}\n",
    "$\n",
    "and $\n",
    "    b\\in \\mathbb{R}^{[\\mathrm{out}]}\n",
    "$. Let's use the following `torch`'s functions:\n",
    "\n",
    "* `F.dropout(x, p, on/off)` -- independent Bernoulli dropout $x\\mapsto x\\odot m$\n",
    "  for $m\\sim \\mathcal{Ber}\\bigl(\\bigl\\{0, \\tfrac1{1-p}\\bigr\\}, 1-p\\bigr)$\n",
    "\n",
    "* `F.linear(x, W, b)` -- affine transformation $x \\mapsto x W^\\top + b$\n",
    "\n",
    "**(note)** the `.weight` of a linear layer in `pytorch` is an $\n",
    "{\n",
    "    [\\mathrm{out}]\n",
    "    \\times [\\mathrm{in}]\n",
    "}\n",
    "$ matrix.\n",
    "\n",
    "<!-- `pytorch` has a function for this `F.dropout(input, p, training)`. It multiplies\n",
    "each element of the `input` tensor by an independent Bernoulli rv. The argument\n",
    "`p` has the same meaning as above. The boolean argument `training` toggles the\n",
    "effect: if `False` then the input is returned as-is, otherwise the mask is applied. -->"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def DropoutLinear_forward(self, input):\n",
    "    ## Exercise: If not frozen, then apply always active dropout,\n",
    "    #  then linear transformation. If frozen, apply the transform\n",
    "    #  using the frozen weight\n",
    "\n",
    "    # linear with frozen weight\n",
    "    if self.is_frozen():\n",
    "        return F.linear(input, self.frozen_weight, self.bias)\n",
    "\n",
    "    # stochastic pass as in `ActiveDropout` + Linear\n",
    "    input = F.dropout(input, self.p, True)\n",
    "\n",
    "    return F.linear(input, self.weight, self.bias)\n",
    "    # return super().forward(F.dropout(input, self.p, True))\n",
    "\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Parameter freezer for our custom layer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For input\n",
    "$\n",
    "    x\\in \\mathbb{R}^{[\\mathrm{in}]}\n",
    "$ and a layer parameters $\n",
    "    W\\in \\mathbb{R}^{[\\mathrm{out}] \\times [\\mathrm{in}]}\n",
    "$\n",
    "and $\n",
    "    b\\in \\mathbb{R}^{[\\mathrm{out}]}\n",
    "$ the effect in `DropoutLinear` is\n",
    "\n",
    "$$\n",
    "    y_j\n",
    "        = \\bigl[(x \\odot m) W^\\top + b\\bigr]_j\n",
    "        = b_j + \\sum_i x_i m_i W_{ji}\n",
    "        = b_j + \\sum_i x_i \\breve{W}_{ji}\n",
    "    \\,, $$\n",
    "\n",
    "where the each column of $\\breve{W}_i$ is, independently, either\n",
    "$\\mathbf{0} \\in \\mathbb{R}^{[\\mathrm{out}]}$ with probability $p$ or\n",
    "some (learnable) vector in $\\mathbb{R}^{[\\mathrm{out}]}$\n",
    "\n",
    "$$\n",
    "    \\breve{W}_i \\sim\n",
    "\\begin{cases}\n",
    "    \\mathbf{0}\n",
    "        & \\text{ w. prob } p \\,, \\\\\n",
    "    \\tfrac1{1-p} M_i\n",
    "        & \\text{ w. prob } 1-p \\,.\n",
    "\\end{cases}\n",
    "$$\n",
    "\n",
    "Thus the multiplicative effect of the random mask $m$ on $x$ can be\n",
    "equivalently seen as a random **on/off** switch effect on the\n",
    "**columns** of the matrix $W$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def DropoutLinear_freeze(self):\n",
    "    \"\"\"Apply dropout with rate `p` to columns of `weight` and freeze it.\"\"\"\n",
    "    # we leverage torch's broadcasting semantics and draw a one-row\n",
    "    #  mask binary mask, that we later multiply the weight by.\n",
    "\n",
    "    # let's draw the new weight\n",
    "    with torch.no_grad():\n",
    "        prob = torch.full_like(self.weight[:1, :], 1 - self.p)\n",
    "        feature_mask = torch.bernoulli(prob) / prob\n",
    "\n",
    "        frozen_weight = self.weight * feature_mask\n",
    "\n",
    "    # and store it\n",
    "    self.register_buffer(\"frozen_weight\", frozen_weight)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assemble the blocks into a layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DropoutLinear(Linear, FreezableWeight):\n",
    "    \"\"\"Linear layer with dropout on inputs.\"\"\"\n",
    "    def __init__(self, in_features, out_features, bias=True, p=0.5):\n",
    "        super().__init__(in_features, out_features, bias=bias)\n",
    "\n",
    "        self.p = p\n",
    "\n",
    "    forward = DropoutLinear_forward\n",
    "\n",
    "    freeze = DropoutLinear_freeze"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparing sample functions to point-estimates "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's rewrite the model builder function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_model(p=0.5):\n",
    "    \"\"\"Build a model with the custom layer and dropout rate set to `p`.\"\"\"\n",
    "\n",
    "    return Sequential(\n",
    "        ## Exercise: Plug-in `DropoutLinear` layer into our second network.\n",
    "\n",
    "        Linear(1, 512, bias=True),\n",
    "        LeakyReLU(),\n",
    "\n",
    "        DropoutLinear(512, 512, bias=True , p=p),\n",
    "        LeakyReLU(),\n",
    "\n",
    "        DropoutLinear(512, 1, bias=True, p=p),\n",
    "\n",
    "        # pass\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create a new instance and retrain the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = build_model(p=0.5)\n",
    "model.to(device)\n",
    "\n",
    "fit(model, train, criterion=\"mse\", n_epochs=2000, verbose=True, weight_decay=1e-3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "... and obtain two estimates: pointwise and functional."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples_pe = point_estimate(model, domain, n_samples=51, verbose=True)\n",
    "samples_sf = sample_function(model, domain, n_samples=51, verbose=True)\n",
    "\n",
    "samples_pe.shape, samples_sf.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "(torch.Size([51, 251, 1]), torch.Size([51, 251, 1]))\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's compare <span style=\"color:#1f77b4\">**point estimates**</span>\n",
    "with <span style=\"color:#ff7f0e\">**function sampling**</span>."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(12, 5))\n",
    "\n",
    "ax.plot(X_domain[:, 0], samples_pe[:10, :, 0].numpy().T,\n",
    "        c=\"C1\", lw=1, alpha=0.5)\n",
    "\n",
    "ax.plot(X_domain[:, 0], samples_sf[:10, :, 0].numpy().T,\n",
    "        c=\"C0\", lw=2, alpha=0.5)\n",
    "\n",
    "ax.scatter(X_train, y_train, c=\"black\", s=40,\n",
    "           label=\"train\", zorder=+10);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(12, 5))\n",
    "\n",
    "ax.scatter(X_train, y_train, c=\"black\", s=40, label=\"train\")\n",
    "\n",
    "mean, std = samples_sf.mean(dim=0).numpy(), samples_sf.std(dim=0).numpy()\n",
    "ax.plot(X_domain, mean + 1.96 * std, c=\"C0\")\n",
    "ax.plot(X_domain, mean - 1.96 * std, c=\"C0\");\n",
    "\n",
    "mean, std = samples_pe.mean(dim=0).numpy(), samples_pe.std(dim=0).numpy()\n",
    "ax.plot(X_domain, mean + 1.96 * std, c=\"C1\")\n",
    "ax.plot(X_domain, mean - 1.96 * std, c=\"C1\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Pros of `point-estimate`:\n",
    "* uses stochastic forward passes -- no need to for extra code and classes\n",
    "\n",
    "Cons of `point-estimate`:\n",
    "* samples from the predictive distribution at adjacent inputs are independent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**(note)**\n",
    "The parameter distribution of the layer we've built is\n",
    "\n",
    "$$\n",
    "    q(\\omega\\mid \\theta)\n",
    "        = \\prod_i q(\\omega_i\\mid \\theta_i)\n",
    "        = \\prod_i \\bigl\\{\n",
    "            p \\delta_{\\mathbf{0}} (\\omega_i)\n",
    "            + (1 - p) \\delta_{\\tfrac1{1-p} \\theta_i}(\\omega_i)\n",
    "        \\bigr\\}\n",
    "    \\,, $$\n",
    "\n",
    "where $\\omega_i$ is the $i$-th column of $\\omega$, $\\delta_a$ is a\n",
    "**point-mass** distribution at $a$, and $\\theta$ is the learnt\n",
    "approximate posterior mean of $\\omega$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Under benign assumptions and certain relaxations\n",
    "[Gal, Y. 2016 (eq. (6.3) p.109, Prop. 4 p.149)](http://www.cs.ox.ac.uk/people/yarin.gal/website/thesis/thesis.pdf)\n",
    "has shown that a deep network with dropout rate $p$\n",
    "and $\\ell_2$ weight penalty (`weight_decay`) performs (doubly)\n",
    "**stochastic variational inference** with the following stochastic\n",
    "approximate **evidence lower bound**: for the dataset $D = (x_i, y_i)_i$\n",
    "of size $N = \\lvert D \\rvert$ and random batches $B$ of size\n",
    "$\\lvert B \\rvert = m$\n",
    "\n",
    "$$\n",
    "    \\frac1{N} \\Bigl( \\underbrace{\n",
    "        \\mathbb{E}_{\\omega\\sim q(\\omega\\mid \\theta)} \\log p(D \\mid \\omega)\n",
    "        - KL\\bigl(q(\\omega\\mid \\theta) \\big\\| \\pi(\\omega) \\bigr)\n",
    "    }_{ELBO(\\theta)} \\Bigr)\n",
    "    \\approx \\frac1{\\lvert B \\rvert}\n",
    "        \\sum_{i\\in B} \\log p(y_i \\mid x_i, \\omega^{(1)}_i, \\ldots, \\omega^{(L)}_i)\n",
    "        - \\sum_{l=1}^L\n",
    "            \\frac{1-p^{(l)}}{2 s^2 N} \\|\\theta^{(l)}\\|_2^2\n",
    "%             - [\\mathrm{in}_{(l)}] \\, \\mathbb{H}(\\mathcal{Ber}(p^{(l)}))\n",
    "%         + \\mathrm{const}\n",
    "\\,, $$\n",
    "where $\\omega_i^{(l)}$ are independently drawn from $q(\\omega \\mid \\theta)$\n",
    "(one random draw per element in $B$) and $s^2$ is the prior variance."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Thus `weight_decay` should be decreasing with $p$ and $N$:\n",
    "$$ \\lambda = \\frac{1-p}{2 s^2 N} \\,. $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Question(s) (to ponder in your spare time)\n",
    "\n",
    "* what happens to the confidence bands, when you increase the number\n",
    "  of path-wise and pointwise samples?\n",
    "\n",
    "* what will happen if you change the dropout rate $p$ and keep `n_epochs` at 2000?\n",
    "\n",
    "* what happens if for $p=\\tfrac12$ we use much less `n_epochs`?\n",
    "\n",
    "* how does different settings of `weight_decay` affect the bands?\n",
    "\n",
    "Try to rebuild the model with different $p \\in (0, 1)$ using `build_model(p)`, use\n",
    "`fit(..., n_epochs=...)`, and then plot the predictive bands."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mlss2019bdl.plotting import plot1d_bands\n",
    "\n",
    "# model = fit(build_model(p=...), train, n_epochs=..., weight_decay=..., criterion=\"mse\")\n",
    "# plot1d_bands(sample_function(model, domain, n_samples=101), c=\"C0\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_a = fit(build_model(p=0.15), train, criterion=\"mse\", n_epochs=2000, weight_decay=1e-3)\n",
    "\n",
    "model_z = fit(build_model(p=0.75), train, criterion=\"mse\", n_epochs=2000, weight_decay=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(12, 5))\n",
    "\n",
    "samples_a = sample_function(model_a, domain, n_samples=101)\n",
    "samples_z = sample_function(model_z, domain, n_samples=101)\n",
    "\n",
    "plot1d_bands(X_domain, samples_a.transpose(0, 2), c=\"r\")\n",
    "plot1d_bands(X_domain, samples_z.transpose(0, 2), c=\"b\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_a = fit(build_model(p=0.50), train, criterion=\"mse\", n_epochs=20, weight_decay=1e-3)\n",
    "\n",
    "model_z = fit(build_model(p=0.50), train, criterion=\"mse\", n_epochs=200, weight_decay=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(12, 5))\n",
    "\n",
    "samples_a = sample_function(model_a, domain, n_samples=101)\n",
    "samples_z = sample_function(model_z, domain, n_samples=101)\n",
    "\n",
    "plot1d_bands(X_domain, samples_a.transpose(0, 2), c=\"r\")\n",
    "plot1d_bands(X_domain, samples_z.transpose(0, 2), c=\"b\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_a = fit(build_model(p=0.50), train, criterion=\"mse\", n_epochs=2000, weight_decay=1e-5)\n",
    "\n",
    "model_z = fit(build_model(p=0.50), train, criterion=\"mse\", n_epochs=2000, weight_decay=1e-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(12, 5))\n",
    "\n",
    "samples_a = sample_function(model_a, domain, n_samples=101)\n",
    "\n",
    "samples_z = sample_function(model_z, domain, n_samples=101)\n",
    "\n",
    "plot1d_bands(X_domain, samples_a.transpose(0, 2), c=\"r\")\n",
    "plot1d_bands(X_domain, samples_z.transpose(0, 2), c=\"b\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_a = fit(build_model(p=0.10), train, criterion=\"mse\", n_epochs=2000, weight_decay=1e-3)\n",
    "\n",
    "model_z = fit(build_model(p=0.90), train, criterion=\"mse\", n_epochs=2000, weight_decay=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(12, 5))\n",
    "\n",
    "samples_a = sample_function(model_a, domain, n_samples=101)\n",
    "\n",
    "samples_z = sample_function(model_z, domain, n_samples=101)\n",
    "\n",
    "plot1d_bands(X_domain, samples_a.transpose(0, 2), c=\"r\")\n",
    "plot1d_bands(X_domain, samples_z.transpose(0, 2), c=\"b\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (optional) Dropout $2$-d Convolutional layer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Typically, in convolutional neural networks the dropout acts upon the feature\n",
    "(channel) information and not on the spatial dimensions. Thus entire channels\n",
    "are dropped out and for $\n",
    "    x \\in \\mathbb{R}^{\n",
    "        [\\mathrm{in}]\n",
    "        \\times h\n",
    "        \\times w}\n",
    "$ and $\n",
    "    y \\in \\mathbb{R}^{\n",
    "        [\\mathrm{out}]\n",
    "        \\times h'\n",
    "        \\times w'}\n",
    "$ the full effect of the `Dropout+Conv2d` layer is\n",
    "\n",
    "$$\n",
    "    y_{lij} = ((x \\odot m) \\ast W_l)_{ij} + b_l\n",
    "        = b_l + \\sum_k \\sum_{pq} x_{k i_p j_q} m_k W_{lkpq}\n",
    "    \\,, \\tag{conv-2d} $$\n",
    "    \n",
    "where i.i.d $m_k \\sim \\mathcal{Ber}\\bigl(\\bigl\\{0, \\tfrac1{1-p}\\bigr\\}, 1-p\\bigr)$,\n",
    "and indices $i_p$ and $j_q$ represent the spatial location in $x$ that correspond\n",
    "to the $p$ and $q$ elements in the kernel $\n",
    "    W\\in \\mathbb{R}^{\n",
    "        [\\mathrm{out}]\n",
    "        \\times [\\mathrm{in}]\n",
    "        \\times h\n",
    "        \\times w}\n",
    "$ relative to $(i, j)$ coordinates in $y$.\n",
    "The exact values of $i_p$ and $j_q$ depend on the configuration of the\n",
    "convolutional layer, e.g. stride, kernel size and dilation.\n",
    "\n",
    "**(note)** Informative illustrations on the effects of convolution\n",
    "parameters can be found in [Convolution arithmetic](https://github.com/vdumoulin/conv_arithmetic) \n",
    "repo."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## (optional) A brief reminder on Bayesian and Variational Inference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Bayesian Inference is a principled framework of reasoning about uncertainty.\n",
    "\n",
    "In Bayesian Inference (**BI**) we *assume* that the observation\n",
    "data $D$ follows a *model* $m$ with data generating distribution\n",
    "$p(D\\mid m, \\omega)$ *governed by unknown parameters* $\\omega$.\n",
    "The goal of **BI** is to reason about the model and/or its parameters,\n",
    "and new data given the observed data $D$ and our assumptions, i.e\n",
    "to seek the **posterior** parameter and predictive distributions:\n",
    "\n",
    "$$\\begin{align}\n",
    "    p(d \\mid D, m)\n",
    "        % &= \\mathbb{E}_{\n",
    "        %     \\omega \\sim p(\\omega \\mid D, m)\n",
    "        % } p(d \\mid D, \\omega, m)\n",
    "        &= \\int p(d \\mid D, \\omega, m) p(\\omega \\mid D, m) d\\omega\n",
    "    \\,, \\\\\n",
    "    p(\\omega \\mid D, m)\n",
    "        &= \\frac{p(D\\mid \\omega, m) \\, \\pi(\\omega \\mid m)}{p(D\\mid m)}\n",
    "    \\,.\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "* the **prior** distribution $\\pi(\\omega \\mid m)$ reflects our belief\n",
    "  before having made the observations\n",
    "\n",
    "* the data distribution $p(D \\mid \\omega, m)$ reflects our assumptions\n",
    "  about the data generating process, and determines the parameter\n",
    "  **likelihood** (Gaussian, Categorical, Poisson)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unless the distributions and likelihoods are conjugate, posterior in\n",
    "Bayesian inference is typically intractable and it is common to resort\n",
    "to **Variational Inference** or **Monte Carlo** approximations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This key idea of this approach is to seek an approximation $q(\\omega)$\n",
    "to the intractable posterior $p(\\omega \\mid D, m)$, via a variational\n",
    "optimization problem over some tractable family of distributions $\\mathcal{Q}$:\n",
    "\n",
    "$$\n",
    "    q^*(\\omega)\n",
    "        \\in \\arg \\min_{q\\in \\mathcal{Q}} \\mathrm{KL}(q(\\omega) \\| p(\\omega \\mid D, m))\n",
    "    \\,, $$\n",
    "\n",
    "where the Kullback-Leibler divergence between $P$ and $Q$ ($P\\ll Q$)\n",
    "with densities $p$ and $q$, respectively, is given by\n",
    "\n",
    "$$\n",
    "    \\mathrm{KL}(q(\\omega) \\| p(\\omega))\n",
    "%         = \\mathbb{E}_{\\omega \\sim Q} \\log \\tfrac{dQ}{dP}(\\omega)\n",
    "        = \\mathbb{E}_{\\omega \\sim q(\\omega)}\n",
    "            \\log \\tfrac{q(\\omega)}{p(\\omega)}\n",
    "    \\,. \\tag{kl-div} $$\n",
    "\n",
    "\n",
    "Note that the family of variational approximations $\\mathcal{Q}$ can be\n",
    "structured **arbitrarily**: point-mass, products, mixture, dependent on\n",
    "input, having mixed hierarchical structure, -- any valid distribution."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Although computing the divergence w.r.t. the unknown posterior\n",
    "is still hard and intractable, it is possible to do away with it\n",
    "through the following identity, which is based on the Bayes rule.\n",
    "\n",
    "For **any** $q(\\omega) \\ll p(\\omega \\mid D; \\phi)$ and any model $m$\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "    \\overbrace{\n",
    "        \\log p(D \\mid m)\n",
    "    }^{\\text{evidence}}\n",
    "        &= \\overbrace{\n",
    "            \\mathbb{E}_{\\omega \\sim q} \\log p(D\\mid \\omega, m)\n",
    "        }^{\\text{expected conditional likelihood}}\n",
    "        - \\overbrace{\n",
    "            \\mathrm{KL}(q(\\omega)\\| \\pi(\\omega \\mid m))\n",
    "        }^{\\text{proximity to prior belief}}\n",
    "        \\\\\n",
    "        &+ \\underbrace{\n",
    "            \\mathrm{KL}(q(\\omega)\\| p(\\omega \\mid D, m))\n",
    "        }_{\\text{posterior approximation}}\n",
    "\\end{align}\n",
    "    \\,. \\tag{master-identity}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Instead of minimizing the divergence of the approximation from the posterior,\n",
    "we maximize the **Evidence Lower Bound** with respect to $q(\\omega)$:\n",
    "\n",
    "$$\n",
    "    q^* \\in\n",
    "    \\arg\\max_{q\\in Q}\n",
    "        \\mathcal{L}(q) = \n",
    "            \\mathbb{E}_{\\omega \\sim q} \\log p(D\\mid \\omega, m)\n",
    "            - \\mathrm{KL}(q(\\omega)\\| \\pi(\\omega \\mid m))\n",
    "    \\,. \\tag{max-ELBO} $$\n",
    "\n",
    "* the expected $\\log$-likelihood favours $q$ that place their mass on\n",
    "parameters $\\omega$ that explain $D$ under the specified model $m$.\n",
    "\n",
    "* the negative KL-divergence discourages the approximation $q$\n",
    "from straying too far away from to the prior belief $\\pi$ under $m$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We usually consider the following setup (conditioning on model $m$ is omitted):\n",
    "* the likelihood factorizes $\n",
    "p(D \\mid \\omega)\n",
    "    = \\prod_i p(y_i, x_i \\mid \\omega)\n",
    "    \\propto \\prod_i p(y_i \\mid x_i, \\omega)\n",
    "$\n",
    "for $D = (x_i, y_i)_{i=1}^N$\n",
    "\n",
    "* the approximation is parameterized by $\\theta$: $q(\\omega\\mid \\theta)$\n",
    "\n",
    "* the prior on $\\omega$ itself depends on hyper-parameters $\\lambda$, that\n",
    "  can be fixed, or variable ($\\pi(\\omega \\mid \\lambda)$)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this case the variational objective (evidence lower bound)\n",
    "\n",
    "$$\n",
    "    \\log p(D\\mid \\lambda )\n",
    "        \\geq \\mathcal{L}(\\theta, \\lambda)\n",
    "            = \\mathbb{E}_{\\omega \\sim q(\\omega \\mid \\theta)}\n",
    "                \\sum_i \\log p_\\phi(y_i \\mid x_i, \\omega)\n",
    "            - KL(q(\\omega \\mid \\theta) \\| \\pi(\\omega \\mid \\lambda))\n",
    "    $$\n",
    "\n",
    "is maximized with respect to $\\theta$ (to approximate the posterior)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Priors can be\n",
    "* *subjective*, i.e. reflecting prior beliefs (but not arbitrary),\n",
    "* *objective*, i.e. reflecting our lack of knowledge,\n",
    "* *empirical*, i.e. learnt from data (we also optimize over hyper-parameters $\\lambda$)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The stochastic variant of ELBO is formed by randomly batching\n",
    "the dataset $D$:\n",
    "\n",
    "$$\n",
    "    \\mathcal{L}(\\theta, \\lambda)\n",
    "        \\approx \\mathcal{L}_\\mathrm{SGVB}(\\theta, \\lambda)\n",
    "        = \\lvert D \\rvert \\biggl(\n",
    "            \\tfrac1{\\lvert B \\rvert}\n",
    "                \\sum_{b \\in B} \\mathbb{E}_{\\omega \\sim q(\\omega \\mid \\theta)}\n",
    "                    \\log p(y_b \\mid x_b, \\omega)\n",
    "        \\biggr)\n",
    "        - KL(q(\\omega \\mid \\theta) \\| \\pi(\\omega \\mid \\lambda))\n",
    "    \\,. $$\n",
    "\n",
    "* Stochastic optimization follows noisy unbiased gradient estimates, which are\n",
    "usually cheap, allow escaping from local optima, and optimize the objective in\n",
    "expectation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to get a gradient of $\n",
    "    F_\\theta = \\mathbb{E}_{\\omega \\sim q(\\omega \\mid \\theta)} f(\\omega)\n",
    "$ w.r.t $\\theta$ we use either:\n",
    "\n",
    "###### (REINFORCE)\n",
    "$\n",
    "\\nabla_\\theta F_\\theta\n",
    "    = \\mathbb{E}_{\\omega \\sim q(\\omega \\mid \\theta)}\n",
    "        (f(\\omega) - b_\\theta) \\nabla_\\theta \\log q(\\omega \\mid \\theta)\n",
    "$\n",
    "* for some $b_\\theta$ that is used to control variance\n",
    "\n",
    "###### (reparameterization)\n",
    "$\n",
    "\\nabla_\\theta F_\\theta\n",
    "    = \\nabla_\\theta \\mathbb{E}_{\\varepsilon \\sim q(\\varepsilon)}\n",
    "        f(g(\\theta; \\varepsilon))\n",
    "    = \\mathbb{E}_{\\varepsilon \\sim q(\\varepsilon)}\n",
    "        \\nabla_\\theta g(\\theta; \\varepsilon)\n",
    "            \\nabla_\\omega f(\\omega) \\big\\vert_{\\omega = g(\\theta; \\varepsilon)}\n",
    "$\n",
    "* when there are $q$ and differentiable $g$ such that sampling from\n",
    "$q(\\omega \\mid \\theta)$ is equivalent to $\\omega = g(\\theta; \\varepsilon)$\n",
    "with $\\varepsilon \\sim q(\\varepsilon)$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The variational approximation might yield high dimensional integrals,\n",
    "which are slow/prohibitive to compute. To make the computations faster\n",
    "without foregoing much of the precision, we may use Monte Carlo methods:\n",
    "\n",
    "$$\n",
    "    \\mathbb{E}_{\\omega \\sim q(\\omega\\mid \\theta)} \\, f(\\omega)\n",
    "        \\overset{\\text{MC}}{\\approx}\n",
    "            \\frac1{\\lvert \\mathcal{W}\\rvert}\n",
    "                \\sum_{\\omega \\in \\mathcal{W}} f(\\omega)\n",
    "    \\,,\n",
    "$$\n",
    "\n",
    "where $\\mathcal{W} = (\\omega_b)_{b=1}^B$ is a sample of independent draws\n",
    "from $q(\\omega\\mid \\theta)$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we also approximate the expectation in the gradient of ELBO\n",
    "via Monte Carlo we get **doubly stochastic variational objective**:\n",
    "\n",
    "$$\n",
    "    \\nabla_\\theta \\mathcal{L}_\\mathrm{DSVB}(\\theta, \\lambda)\n",
    "        \\approx\n",
    "            \\lvert D \\rvert \\biggl(\n",
    "                \\tfrac1{\\lvert B \\rvert}\n",
    "                    \\sum_{b \\in B}\n",
    "                        \\mathop{gradient}(x_b, y_b)\n",
    "            \\biggr)\n",
    "            - \\nabla_\\theta KL(q(\\omega \\mid \\theta) \\| \\pi(\\omega \\mid \\lambda))\n",
    "    \\,, $$\n",
    "\n",
    "where `gradient` is $\n",
    "    \\nabla_\\theta\n",
    "        \\mathbb{E}_{\\omega \\sim q(\\omega \\mid \\theta)}\n",
    "            \\log p(y \\mid x, \\omega)\n",
    "$ using one of the approaches above, typically approximated using\n",
    "one independent draw of $\\omega$ per $b\\in B$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use a similar sampling approach to compute the gradient of the divergence term."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A good overview of Bayesian Inference can be found at [bdl101.ml](http://bdl101.ml/),\n",
    "in [this lecture](http://mlg.eng.cam.ac.uk/zoubin/talks/lect1bayes.pdf),\n",
    "[this paper](https://arxiv.org/abs/1206.7051.pdf), or\n",
    "[this review](https://arxiv.org/abs/1601.00670.pdf),\n",
    "among other great resources. It is also possible to consult\n",
    "the references at [wiki](https://en.wikipedia.org/wiki/Bayesian_inference)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can estimate the divergence term in the ELBO\n",
    "with Monte Carlo, or, for example, for the predictive distribution\n",
    "we have\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "    \\mathbb{E}_{y\\sim p(y\\mid x, D, m)} \\, g(y)\n",
    "        &\\overset{\\text{BI}}{=}\n",
    "            \\mathbb{E}_{\\omega\\sim p(\\omega \\mid D, m)}\n",
    "                \\mathbb{E}_{y\\sim p(y\\mid x, \\omega, D, m)} \\, g(y) \n",
    "        \\\\\n",
    "        &\\overset{\\text{VI}}{\\approx}\n",
    "            \\mathbb{E}_{\\omega\\sim q(\\omega)}\n",
    "                \\mathbb{E}_{y\\sim p(y\\mid x, \\omega, D, m)} \\, g(y)\n",
    "        \\\\\n",
    "        &\\overset{\\text{MC}}{\\approx}\n",
    "%             \\hat{\\mathbb{E}}_{\\omega \\sim \\mathcal{W}}\n",
    "%                 \\mathbb{E}_{y\\sim p(y\\mid x, \\omega, D, m)} \\, g(y)\n",
    "            \\frac1{\\lvert \\mathcal{W}\\rvert} \\sum_{\\omega \\in \\mathcal{W}}\n",
    "                \\mathbb{E}_{y\\sim p(y\\mid x, \\omega, D, m)} \\, g(y)\n",
    "    \\,,\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "where $\\mathcal{W} = (\\omega_b)_{b=1}^B \\sim q(\\omega)$\n",
    "-- iid samples from the variational approximation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  }
 ],
 "metadata": {
  "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}