ProbAI_2022_BNN_Tutorial_regression.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/bryagd/DiskANN/blob/main/ProbAI_2022_BNN_Tutorial_regression.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# A Hands-on Tutorial for Bayesian Neural Networks\n",
        "\n",
        "(Part 1: regression)\n",
        "\n",
        "[Yingzhen Li](http://yingzhenli.net/home/en/)\n",
        "\n",
        "(As part of the BNN lecture at [ProbAI 2022](https://probabilistic.ai/))\n",
        "\n",
        "In this tutorial, you will implement various Bayesian neural network methods based on variational inference.\n",
        "\n",
        "We will go through regression tasks to see the applications of uncertainty estimation in practice, including a case study on Bayesian optimisation.\n",
        "\n",
        "**How to use this tutorial notebook?**\n",
        "\n",
        "*   Read the descriptions in the text;\n",
        "*   Fill in the missing code whenever you see a block that looks like below:\n",
        "\n",
        "```\n",
        "### beginning of your code ###\n",
        "[insert your own code here]\n",
        "### end of your code ###\n",
        "```\n",
        "There will be hints provided in the code blocks as well to guide you through.\n",
        "\n",
        "Let us set up the required packages, and then, enjoy 😊"
      ],
      "metadata": {
        "id": "Z97rDy1VjK4I"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%matplotlib inline\n",
        "\n",
        "import numpy as np\n",
        "import math\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.distributions as dist\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "\n",
        "EPS = 1e-5  # define a small constant for numerical stability control"
      ],
      "metadata": {
        "id": "1Zt00SKnNgyR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# What is a Bayesian neural network?\n",
        "\n",
        "In short, a Bayesian neural network (BNN) is a neural network that use (approximate) Bayesian inference for uncertainty estimation. For example, we can treat the NN parameters as random variables and infer them using (approximate) Bayesian posterior inference.\n",
        "\n",
        "For the mathematical foundation, see [lecture notes here](http://yingzhenli.net/home/pdf/ProbAI2022_lecture_note.pdf)."
      ],
      "metadata": {
        "id": "uJqCF8BXD7dO"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Implementing the BNN with mean-field variational inference (MFVI)\n",
        "\n",
        "In this part, you will implement the Bayes-by-backprop algorithm by [Blundell et al. (2015)](https://arxiv.org/abs/1505.05424), which is doing mean-field variational inference (MFVI) in weight space for Bayesian neural networks.\n",
        "\n",
        "For the mathematical foundation, see [lecture notes here](http://yingzhenli.net/home/pdf/ProbAI2022_lecture_note.pdf).\n"
      ],
      "metadata": {
        "id": "-eYv-w6TNmW4"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Implementing the ```MFVILinear``` class and construct a BNN\n",
        "\n",
        "We will implement the MFVI approach in a modular way, so we need a \"BNN layer\" module ```MFVILinear``` to start with. This will allow us to stack-up many network layers efficiently. Then we will implement the training & testing functions, and test the model's performance with a regression example.\n",
        "\n",
        "In below block you will implment a BNN layer with mean-field variational inference (MFVI), and more specifically:\n",
        "\n",
        "* the initialisation of the variational parameters\n",
        "* the forward pass with Monte Carlo\n",
        "* the $KL[q||p]$ regularisation term for one layer, i.e., $KL[q(W) || p(W)] + KL[q(b) || p(b)]$"
      ],
      "metadata": {
        "id": "kOQ75uGIkpW1"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dpydUZNVjCRV"
      },
      "outputs": [],
      "source": [
        "class MFVILinear(nn.Module):\n",
        "    \"\"\"Applies a linear transformation to the incoming data: y = xW^T + b, where\n",
        "    the weight W and bias b are sampled from the q distribution.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, dim_in, dim_out, prior_weight_std=1.0, prior_bias_std=1.0, init_std=0.05,\n",
        "                 sqrt_width_scaling=False, device=None, dtype=None):\n",
        "        factory_kwargs = {'device': device, 'dtype': dtype}\n",
        "        super(MFVILinear, self).__init__()\n",
        "        self.dim_in = dim_in  # dimension of network layer input\n",
        "        self.dim_out = dim_out  # dimension of network layer output\n",
        "\n",
        "        # define the trainable variational parameters for q distribtuion\n",
        "        # first define and initialise the mean parameters\n",
        "        self.weight_mean = nn.Parameter(torch.empty((dim_out, dim_in), **factory_kwargs))\n",
        "        self.bias_mean = nn.Parameter(torch.empty(dim_out, **factory_kwargs))\n",
        "        self._weight_std_param = nn.Parameter(torch.empty((dim_out, dim_in), **factory_kwargs))\n",
        "        self._bias_std_param = nn.Parameter(torch.empty(dim_out, **factory_kwargs))\n",
        "        self.reset_parameters(init_std)\n",
        "\n",
        "        # define the prior parameters (for prior p, assume the mean is 0)\n",
        "        prior_mean = 0.0\n",
        "        if sqrt_width_scaling:  # prior variance scales as 1/dim_in\n",
        "            prior_weight_std /= self.dim_in ** 0.5\n",
        "        # prior parameters are registered as constants\n",
        "        self.register_buffer('prior_weight_mean', torch.full_like(self.weight_mean, prior_mean))\n",
        "        self.register_buffer('prior_weight_std', torch.full_like(self._weight_std_param, prior_weight_std))\n",
        "        self.register_buffer('prior_bias_mean', torch.full_like(self.bias_mean, prior_mean))\n",
        "        self.register_buffer('prior_bias_std', torch.full_like(self._bias_std_param, prior_bias_std))\n",
        "\n",
        "    def extra_repr(self):\n",
        "        s = \"dim_in={}, dim_in={}, bias=True\".format(self.dim_in, self.dim_out)\n",
        "        weight_std = self.prior_weight_std.data.flatten()[0]\n",
        "        if torch.allclose(weight_std, self.prior_weight_std):\n",
        "            s += f\", weight prior std={weight_std.item():.2f}\"\n",
        "        bias_std = self.prior_bias_std.flatten()[0]\n",
        "        if torch.allclose(bias_std, self.prior_bias_std):\n",
        "            s += f\", bias prior std={bias_std.item():.2f}\"\n",
        "        return s\n",
        "\n",
        "    def reset_parameters(self, init_std=0.05):\n",
        "        nn.init.kaiming_uniform_(self.weight_mean, a=math.sqrt(5))\n",
        "        bound = self.dim_in ** -0.5\n",
        "        nn.init.uniform_(self.bias_mean, -bound, bound)\n",
        "        ### begin of your code ###\n",
        "        # hints: Notice that standard deviation > 0 always so we cannot directly treat\n",
        "        # the standard deviation as a free parameter. Instead we parameterise a free-form\n",
        "        # parameter self._weight_std_param and then transform it to the standard deviation\n",
        "        # in the function \"weight_std(self)\" below. Note that we want to initialise these\n",
        "        # parameters so that standard deviation of q(W) (and q(b)) is initialised as init_std.\n",
        "\n",
        "        self._weight_std_param.data =\n",
        "        self._bias_std_param.data =\n",
        "\n",
        "        ### end of your code ###\n",
        "\n",
        "    # define the q distribution standard deviations with property decorator\n",
        "    @property\n",
        "    def weight_std(self):\n",
        "        ### begin of your code ###\n",
        "        # hints: depending on how you define self._weight_std_param above, the way\n",
        "        # you convert self._weight_std_param to weight_std (standard deviation of q(W))\n",
        "        # will be different.\n",
        "\n",
        "        weight_std =\n",
        "\n",
        "        ### end of your code ###\n",
        "        return weight_std\n",
        "\n",
        "    @property\n",
        "    def bias_std(self):\n",
        "        ### begin of your code ###\n",
        "\n",
        "        bias_std =\n",
        "\n",
        "        ### end of your code ###\n",
        "        return bias_std\n",
        "\n",
        "    # KL divergence KL[q||p] between two Gaussians\n",
        "    def kl_divergence(self):\n",
        "        q_weight = dist.Normal(self.weight_mean, self.weight_std)\n",
        "        p_weight = dist.Normal(self.prior_weight_mean, self.prior_weight_std)\n",
        "        kl = dist.kl_divergence(q_weight, p_weight).sum()\n",
        "        q_bias = dist.Normal(self.bias_mean, self.bias_std)\n",
        "        p_bias = dist.Normal(self.prior_bias_mean, self.prior_bias_std)\n",
        "        kl += dist.kl_divergence(q_bias, p_bias).sum()\n",
        "        return kl\n",
        "\n",
        "    # forward pass with Monte Carlo (MC) sampling\n",
        "    def forward(self, input):\n",
        "        weight = self._normal_sample(self.weight_mean, self.weight_std)\n",
        "        bias = self._normal_sample(self.bias_mean, self.bias_std)\n",
        "        return F.linear(input, weight, bias)\n",
        "\n",
        "    def _normal_sample(self, mean, std):\n",
        "        ### begin of your code ###\n",
        "        # please implement the sampling process for a factorised Gaussian\n",
        "\n",
        "        ### end of your code ###\n",
        "\n",
        "# construct a BNN\n",
        "def make_mfvi_bnn(layer_sizes, activation='LeakyReLU', **layer_kwargs):\n",
        "    nonlinearity = getattr(nn, activation)() if isinstance(activation, str) else activation\n",
        "    net = nn.Sequential()\n",
        "    for i, (dim_in, dim_out) in enumerate(zip(layer_sizes[:-1], layer_sizes[1:])):\n",
        "        net.add_module(f'MFVILinear{i}', MFVILinear(dim_in, dim_out, **layer_kwargs))\n",
        "        if i < len(layer_sizes) - 2:\n",
        "            net.add_module(f'Nonlinarity{i}', nonlinearity)\n",
        "    return net"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "We have just defined the ```MFVILinear``` class and ```make_mfvi_bnn``` for network construction.\n",
        "\n",
        "In below we define the training procedure of a BNN, mainly by defining the training objective. This also involves helper functions that collects the KL regularizers accross layers as well as a prediction function.\n",
        "\n",
        "In particular, as we work with regression data, we will use Gaussian likelihood, i.e., $p(y | x, \\theta) = N(y; f_{\\theta}(x), \\sigma^2)$ where $f_{\\theta}$ is the neural network and $\\sigma^2$ is the output variance parameter that we will also optimise."
      ],
      "metadata": {
        "id": "2Pg74n4fjKEV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# collect the kl divergence for all MFVILinear layers\n",
        "def kl_divergence(bnn):\n",
        "    kl = 0.0\n",
        "    for module in bnn:\n",
        "        if hasattr(module, 'kl_divergence'):\n",
        "            kl = kl + module.kl_divergence()\n",
        "    return kl\n",
        "\n",
        "# define the training function which minimises the negative ELBO\n",
        "# note that the (tempered) negative ELBO = negative log-likliehood of dataset + beta * KL\n",
        "# where the neg. log-likelihood *per datapoint* is computed with a function\n",
        "# nll = data_loss_func(y, y_pred)\n",
        "# this data_loss_func is to be defined later, depending on the learning task\n",
        "def train_step(net, opt, data_loss_func, dataloader, N_data, beta=1.0):\n",
        "    for _, (x, y) in enumerate(dataloader):\n",
        "        x = x.to(device); y = y.to(device)\n",
        "        opt.zero_grad() # opt is the optimiser\n",
        "        y_pred = net(x)\n",
        "        ### begin of your code ###\n",
        "        # notice we might use mini-batches, so be careful for the right data-count!\n",
        "\n",
        "        nll = data_loss_func(y, y_pred).mean()\n",
        "        kl = kl_divergence(net)\n",
        "        loss =\n",
        "\n",
        "        ### end of your code ###\n",
        "\n",
        "        loss.backward()\n",
        "        opt.step()\n",
        "    return nll, kl\n",
        "\n",
        "# now define the data_loss_func that will be used later\n",
        "# in detail, we will define a Gaussian log-likelihood function below, and then define\n",
        "# data_loss_func = lambda y, y_pred: -gauss_loglik(y, y_pred, log_noise_var)\n",
        "# where log_noise_var is the output noise (log) variance which is a seperate model parameter\n",
        "# define gaussian log-likelihood\n",
        "def gauss_loglik(y, y_pred, log_noise_var):\n",
        "    # y should have shape as (batch_size, y_dim)\n",
        "    # y_pred should have shape as (batch_size, y_dim) or (K, batch_size, y_dim)\n",
        "    # where K is the number of MC samples\n",
        "    # this function should return per-data loss of shape (batch_size,) or (K, batch_size)\n",
        "    ### begin of your code ###\n",
        "    # hint: consult with your textbook or wikipedia for the Gaussian distribution form\n",
        "\n",
        "    ll =\n",
        "\n",
        "    ### end of your code ###\n",
        "    return ll\n",
        "\n",
        "# define the prediction function with Monte Carlo sampling using K samples\n",
        "def predict(bnn, x_test, K=1):\n",
        "    y_pred = []\n",
        "    for _ in range(K):\n",
        "        y_pred.append(bnn(x_test))\n",
        "    # shape (K, batch_size, y_dim) or (batch_size, y_dim) if K = 1\n",
        "    return torch.stack(y_pred, dim=0).squeeze(0)\n",
        "\n",
        "# define the error metrics: RMSE and test negative log-likelihood\n",
        "# in below functions y_pred should have shape (K, batch_size, y_dim) or (batch_size, y_dim)\n",
        "def rmse(y, y_pred):\n",
        "    if len(y_pred.shape) > 2: # using K > 1 MC samples\n",
        "        y_pred = y_pred.mean(0)\n",
        "    return (y - y_pred).pow(2).sum(-1).mean().sqrt()\n",
        "\n",
        "def test_nll(y, y_pred, data_loss_func):\n",
        "    nll = data_loss_func(y, y_pred)  # with shape (batch_size) or (K, batch_size)\n",
        "    if len(nll) == 2:  # using K > 1 MC samples\n",
        "        nll = -torch.logsumexp(-nll, dim=0) + math.log(nll.shape[0]) # Bayesian predictive average\n",
        "    return nll.mean()"
      ],
      "metadata": {
        "id": "CgTaXt3NnKSi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Testing: a simple regression task\n",
        "We will run a simple 1-D regression task to test the implemented BNN, in terms of its predictive performance and uncertainty estimation.\n",
        "\n",
        "Let's get some synthetic data:\n",
        "\n"
      ],
      "metadata": {
        "id": "StOG-wgEnv9t"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def ground_truth_func(x):\n",
        "    return np.sin(x * math.pi / 2 + 0.8) * np.exp(-0.1 * np.abs(x)) + 0.1 * x\n",
        "\n",
        "def gen_data(N_data, ground_truth_func, noise_std=None):\n",
        "    # generate the training dataset, note here we will make data into 2 clusters\n",
        "    x1 = np.random.randn(int(N_data/2), 1) * 0.5 + 2.0\n",
        "    x2 = np.random.randn(int(N_data/2), 1) * 0.5 - 2.0\n",
        "    x = np.concatenate([x1, x2], axis=0)\n",
        "    y = ground_truth_func(x)\n",
        "    if noise_std is not None and noise_std > EPS:\n",
        "        # assume homogeneous noise setting, i.e., \"homoscedasticity\"\n",
        "        y += np.random.randn(y.shape[0], y.shape[1]) * noise_std\n",
        "    return x, y\n",
        "\n",
        "def normalise_data(x, mean, std):\n",
        "    return (x - mean) / std\n",
        "\n",
        "def unnormalise_data(x, mean, std):\n",
        "    return x * std + mean\n",
        "\n",
        "class regression_data(Dataset):\n",
        "     def __init__(self, x, y, normalise=True):\n",
        "         super(regression_data, self).__init__()\n",
        "         self.update_data(x, y, normalise)\n",
        "\n",
        "     def __len__(self):\n",
        "         return self.x.shape[0]\n",
        "\n",
        "     def __getitem__(self, index):\n",
        "         x = torch.tensor(self.x[index]).float()\n",
        "         y = torch.tensor(self.y[index]).float()\n",
        "         return x, y\n",
        "\n",
        "     def update_data(self, x, y, normalise=True, update_stats=True):\n",
        "         assert x.shape[0] == y.shape[0]\n",
        "         self.x = x\n",
        "         self.y = y\n",
        "         # normalise data\n",
        "         self.normalise = normalise\n",
        "         if update_stats:\n",
        "             self.x_mean = self.x.mean(0) if normalise else 0.0\n",
        "             self.x_std = self.x.std(0) if normalise else 1.0\n",
        "             self.y_mean = self.y.mean(0) if normalise else 0.0\n",
        "             self.y_std = self.y.std(0) if normalise else 1.0\n",
        "         if self.normalise:\n",
        "             self.x = normalise_data(self.x, self.x_mean, self.x_std)\n",
        "             self.y = normalise_data(self.y, self.y_mean, self.y_std)\n",
        "\n",
        "N_data = 100\n",
        "noise_std = 0.1\n",
        "x_train, y_train = gen_data(N_data, ground_truth_func, noise_std)\n",
        "dataset = regression_data(x_train, y_train)\n",
        "dataloader = DataLoader(dataset, batch_size=100, shuffle=True)\n",
        "\n",
        "# plot the training data and ground truth\n",
        "x_test = np.arange(np.min(x_train) - 1.0, np.max(x_train)+1.0, 0.01)[:, np.newaxis]\n",
        "y_test = ground_truth_func(x_test)\n",
        "plt.plot(x_train, y_train, 'ro', label='data')\n",
        "plt.plot(x_test, y_test, 'k-', label='ground-truth')\n",
        "plt.legend()\n",
        "plt.title('ground-truth function')\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "QDjtmJQ8Lglu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We set-up the hyper-parameters and build a BNN to be trained.\n",
        "\n"
      ],
      "metadata": {
        "id": "2PqirrxvLXOs"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "device = 'cuda:0' if torch.cuda.is_available() else 'cpu'\n",
        "x_dim, y_dim = x_train.shape[1], y_train.shape[1]\n",
        "# build a BNN, with hidden layer width = h_dim\n",
        "h_dim = 50\n",
        "layer_sizes = [x_dim, h_dim, h_dim, y_dim]\n",
        "# you can change the activation function here or define your own customized activation\n",
        "activation=nn.GELU()\n",
        "# you can change the prior parameters as you wish\n",
        "layer_kwargs = {'prior_weight_std': 1.0,\n",
        "                'prior_bias_std': 1.0,\n",
        "                'sqrt_width_scaling': False,\n",
        "                'init_std': 0.05,\n",
        "                'device': device}\n",
        "mfvi_regression_net = make_mfvi_bnn(layer_sizes, activation=activation, **layer_kwargs)\n",
        "# we assume a Gaussian likelihood with homogeneuous noise\n",
        "log_noise_var = nn.Parameter(torch.ones(size=(), device=device)*-3.0)\n",
        "# print out the BNN settings\n",
        "print(\"BNN architecture: \\n\", mfvi_regression_net)\n",
        "\n",
        "# plot the BNN prior in function space\n",
        "K = 50  # number of Monte Carlos samples used in test time\n",
        "x_test_norm = normalise_data(x_test, dataset.x_mean, dataset.x_std)\n",
        "x_test_norm = torch.tensor(x_test_norm, ).float().to(device)\n",
        "\n",
        "def to_numpy(x):\n",
        "    return x.detach().cpu().numpy() # convert a torch tensor to a numpy array\n",
        "\n",
        "def get_regression_results(net, x, K, log_noise_var=None):\n",
        "    y_pred = predict(net, x, K=K)  # shape (K, N_test, y_dim)\n",
        "    y_pred_mean = y_pred.mean(0)\n",
        "    if log_noise_var is not None:\n",
        "        # note here the preditive std needs to count for output noise variance\n",
        "        y_pred_std = (y_pred.var(0) + torch.exp(log_noise_var)).sqrt()\n",
        "    else:\n",
        "        y_pred_std = y_pred.std(0)\n",
        "    # unnormalise\n",
        "    y_pred_mean = unnormalise_data(to_numpy(y_pred_mean), dataset.y_mean, dataset.y_std)\n",
        "    y_pred_std = unnormalise_data(to_numpy(y_pred_std), 0.0, dataset.y_std)\n",
        "    return y_pred_mean, y_pred_std\n",
        "\n",
        "# plot the BNN prior and ground truth\n",
        "def plot_regression(x_train, y_train, x_test, y_pred_mean, y_pred_std_noiseless, y_pred_std, title=''):\n",
        "    plt.plot(x_train, y_train, 'ro', label='data')\n",
        "    plt.plot(x_test, y_test, 'k-', label='ground-truth')\n",
        "    plt.plot(x_test, y_pred_mean, 'b-', label='prediction mean')\n",
        "    # plot the uncertainty as +- 2 * std\n",
        "    # first for the total uncertainty (model/epistemic + data/aleatoric)\n",
        "    plt.fill_between(x_test[:,0], y_pred_mean[:,0]-2*y_pred_std[:,0],\n",
        "                     y_pred_mean[:,0]+2*y_pred_std[:,0],\n",
        "                     color='c', alpha=0.3, label='total uncertainty')\n",
        "    # then for the model/epistemic uncertainty only\n",
        "    plt.fill_between(x_test[:,0], y_pred_mean[:,0]-2*y_pred_std_noiseless[:,0],\n",
        "                     y_pred_mean[:,0]+2*y_pred_std_noiseless[:,0],\n",
        "                     color='b', alpha=0.3, label='model uncertainty')\n",
        "    plt.legend()\n",
        "    plt.title(title)\n",
        "    plt.show()\n",
        "\n",
        "y_pred_mean, y_pred_std_noiseless = get_regression_results(mfvi_regression_net, x_test_norm, K)\n",
        "model_noise_std = unnormalise_data(to_numpy(torch.exp(0.5*log_noise_var)), 0.0, dataset.y_std)\n",
        "y_pred_std = np.sqrt(y_pred_std_noiseless ** 2 + model_noise_std**2)\n",
        "plot_regression(x_train, y_train, x_test, y_pred_mean, y_pred_std_noiseless, y_pred_std,\n",
        "                title='BNN init (before training, MFVI)')\n",
        "print(model_noise_std, noise_std, y_pred_std_noiseless.mean())"
      ],
      "metadata": {
        "id": "SIIq7s3KVBbX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now define the data loss (i.e., negative log-likelihood) and train the BNN."
      ],
      "metadata": {
        "id": "_iKDaUqiXuRi"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# define the training function\n",
        "def train_network(net, opt, dataloader, data_loss_func, learning_rate=1e-3,\n",
        "                  N_epochs=2000, beta=1.0, verbose=True):\n",
        "    net.train()\n",
        "    logs = []\n",
        "    for i in range(N_epochs):\n",
        "        nll, kl = train_step(net, opt, data_loss_func, dataloader,\n",
        "                          N_data=len(dataloader.dataset), beta=beta)\n",
        "        logs.append([to_numpy(nll), to_numpy(kl)])\n",
        "        if (i+1) % 100 == 0 and verbose:\n",
        "            print(\"Epoch {}, nll={}, kl={}\".format(i+1, logs[-1][0], logs[-1][1]))\n",
        "    return np.array(logs)\n",
        "\n",
        "# start training\n",
        "learning_rate = 1e-3\n",
        "params = list(mfvi_regression_net.parameters()) + [log_noise_var]\n",
        "opt = torch.optim.Adam(params, lr=learning_rate)\n",
        "# define the regression loss: negative gaussian log-likelihood\n",
        "data_loss_func = lambda y, y_pred: -gauss_loglik(y, y_pred, log_noise_var)\n",
        "# hyper-parameters of training\n",
        "beta = 1.0\n",
        "N_epochs = 2000\n",
        "# the training loop starts\n",
        "logs = train_network(mfvi_regression_net, opt, dataloader, data_loss_func,\n",
        "                     beta=beta, verbose=True, N_epochs=N_epochs)\n",
        "\n",
        "# plot the training curve\n",
        "def plot_training_loss(logs, beta):\n",
        "    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(8, 4))\n",
        "    ax1.plot(np.arange(logs.shape[0]), logs[:, 0], 'r-', label='nll')\n",
        "    ax2.plot(np.arange(logs.shape[0]), logs[:, 1], 'r-', label='KL')\n",
        "    ax1.legend()\n",
        "    ax2.legend()\n",
        "    ax1.set_xlabel('epoch')\n",
        "    ax2.set_xlabel('epoch')\n",
        "    ax1.set_title('ELBO (beta={})'.format(beta))\n",
        "    ax2.set_title('ELBO (beta={})'.format(beta))\n",
        "    plt.show()\n",
        "\n",
        "plot_training_loss(logs, beta)"
      ],
      "metadata": {
        "id": "z-rD6aj-TvLZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Training is now finished, let's plot the predictions again, but now it shows the mean & std of Bayesian predictive inference."
      ],
      "metadata": {
        "id": "T2TDVt0-hPTi"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "y_pred_mean, y_pred_std_noiseless = get_regression_results(mfvi_regression_net, x_test_norm, K)\n",
        "model_noise_std = unnormalise_data(to_numpy(torch.exp(0.5*log_noise_var)), 0.0, dataset.y_std)\n",
        "y_pred_std = np.sqrt(y_pred_std_noiseless ** 2 + model_noise_std**2)\n",
        "plot_regression(x_train, y_train, x_test, y_pred_mean, y_pred_std_noiseless, y_pred_std,\n",
        "                title='BNN approx. posterior (MFVI)')\n",
        "print(model_noise_std, noise_std, y_pred_std_noiseless.mean())"
      ],
      "metadata": {
        "id": "dDWVAvPYibn3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**playground**\n",
        "\n",
        "You can test various hyper-parameter settings and see how they would affect the regression results.\n",
        "\n",
        "In particular, you can go back to previous code blocks and change the settings e.g,:\n",
        "\n",
        "1. Change the depth of the network, by changing the list of ```layer_sizes```;\n",
        "2. Activation functions: try e.g., ```nn.Tanh()``` or ```nn.ReLU()``` or the below ```Sine``` activation (copy-paste the code below to the place when you define the activation function);\n",
        "3.   Training epochs: try setting different values for ```N_epochs``` when calling ```train_network()```, e.g., ```N_epochs = 1000``` vs ```N_epochs = 10000```, and see what happens to the training curves;\n",
        "4.   Tempered ELBO: try using different values of ```beta```, e.g., ```beta = 0.1``` vs ```beta = 10```;\n",
        "5.   Prior settings: try ```sqrt_width_scaling=True```;\n",
        "6.   Initialisation of the q standard deviation ```init_std```.\n",
        "\n",
        "Come up with you own conclusions: which settings work well?"
      ],
      "metadata": {
        "id": "7D07ALhSrpSu"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# copy-paste below code to the appropriate place, and call it by\n",
        "# activation = Sine()\n",
        "class Sine(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(Sine, self).__init__()\n",
        "    def forward(self, x):\n",
        "        return torch.sin(x)"
      ],
      "metadata": {
        "id": "x_ZNoqFvkl2-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Using other q distributions?"
      ],
      "metadata": {
        "id": "r4k3n8ihJ1uO"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Hopefully from above tests, you can see that MFVI works but it needs careful settings of the hyper-parameters (i.e., tricks that are mostly hidden in paper appendices).\n",
        "\n",
        "A natural question is to ask whether other design of the q distribution would return better results. Within the Gaussian family, Gaussians with full-rank covariance matrices are more expressive than factorised Gaussians. But at the same time, it requires much more variational parameters if we were to use full covariance Gaussians for every layer. For example, a hidden layer with ```dim_in = dim_out = 50``` would need $50 \\times 50 = 2500$ parameters for parameterising the mean, but $\\sum_{i=1}^{50\\times 50} i = 3126250$ parameters for parameterising the (symmetric) full covariance matrix! Therefore, when selecting the q distribution family, one needs to also consider the computational & memory costs for such approximation.\n",
        "\n",
        "In below we will implement 2 \"economic\" solutions:\n",
        "\n",
        "1.   The so-called \"last-layer BNN\" approach: only apply full-covariance Gaussian posterior approximation to the last layer of the network, and use MLE/MAP solutions for the previous layers.\n",
        "2.   Monte Carlo dropout (MC-dropout).\n",
        "\n",
        "For mathematical foundations, see [lecture notes here](http://yingzhenli.net/home/pdf/ProbAI2022_lecture_note.pdf).\n"
      ],
      "metadata": {
        "id": "f9GetssjlunI"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Last-layer BNN with full-covariace Gaussian approximation"
      ],
      "metadata": {
        "id": "ZztkJAkyNZUV"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "For the \"last-layer BNN\" method we need to use a full-covariance Gaussian posterior approximation for the output layer of the network.\n",
        "\n",
        "In below block you will implment a BNN layer with variational inference and Gaussian q with **full covariance matrix**, and more specifically:\n",
        "\n",
        "* the initialisation of the variational parameters\n",
        "* the forward pass with Monte Carlo\n",
        "* the $KL[q||p]$ regularisation term for one layer"
      ],
      "metadata": {
        "id": "zEmpAu4gWUGI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class FullCovGaussianLinear(nn.Module):\n",
        "    \"\"\"Applies a linear transformation to the incoming data: y = xW^T + b, where\n",
        "    the weight W and bias b are sampled from the q distribution.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, dim_in, dim_out, prior_weight_std=1.0, prior_bias_std=1.0, init_std=0.5,\n",
        "                 sqrt_width_scaling=False, device=None, dtype=None):\n",
        "        factory_kwargs = {'device': device, 'dtype': dtype}\n",
        "        super(FullCovGaussianLinear, self).__init__()\n",
        "        self.dim_in = dim_in  # dimension of network layer input\n",
        "        self.dim_out = dim_out  # dimension of network layer output\n",
        "\n",
        "        # define the trainable variational parameters for q distribtuion\n",
        "        # first define and initialise the mean parameters\n",
        "        self.weight_mean = nn.Parameter(torch.empty((dim_out, dim_in), **factory_kwargs))\n",
        "        self.bias_mean = nn.Parameter(torch.empty(dim_out, **factory_kwargs))\n",
        "        # then define and initialise the parameters for covariance matrix\n",
        "        # note here that we concatenate all the parameters in W and b and work with their covariances\n",
        "        # in particular, we would like to parameterise the Cholesky factor L of the covariance\n",
        "        # so that LL^T = covariance matrix of q(W, b)\n",
        "        num_params = dim_out * dim_in + dim_out # total number of parameters\n",
        "        self._cov_diag = nn.Parameter(torch.empty((num_params,), **factory_kwargs))\n",
        "        self._cov_tril = nn.Parameter(torch.empty((num_params, num_params), **factory_kwargs))\n",
        "        self.reset_parameters()\n",
        "\n",
        "        # define the prior parameters (for prior p, assume the mean is 0)\n",
        "        prior_mean = 0.0\n",
        "        if sqrt_width_scaling:  # prior variance scales as 1/dim_in\n",
        "            prior_weight_std /= self.dim_in ** 0.5\n",
        "        self.prior_weight_std = prior_weight_std\n",
        "        self.prior_bias_std = prior_bias_std\n",
        "        # prior parameters are registered as constants\n",
        "        self.register_buffer(\"prior_mean\", torch.full((num_params,), prior_mean, **factory_kwargs))\n",
        "        prior_weight_std = torch.full_like(self.weight_mean, prior_weight_std)\n",
        "        prior_bias_std = torch.full_like(self.bias_mean, prior_bias_std)\n",
        "        prior_std_diag = torch.concat((prior_weight_std.flatten(), prior_bias_std.flatten()))\n",
        "        self.register_buffer('prior_scale_tril', prior_std_diag.diag_embed())\n",
        "\n",
        "    def extra_repr(self):\n",
        "        s = \"dim_in={}, dim_in={}, bias=True\".format(self.dim_in, self.dim_out)\n",
        "        s += f\", weight prior std={self.prior_weight_std:.2f}\"\n",
        "        s += f\", bias prior std={self.prior_bias_std:.2f}\"\n",
        "        return s\n",
        "\n",
        "    def reset_parameters(self, init_std=0.5):\n",
        "        nn.init.kaiming_uniform_(self.weight_mean, a=math.sqrt(5))\n",
        "        bound = self.dim_in ** -0.5\n",
        "        nn.init.uniform_(self.bias_mean, -bound, bound)\n",
        "        _init_std_param = math.log(math.expm1(init_std))\n",
        "        self._cov_diag.data = torch.full_like(self._cov_diag.data, _init_std_param)\n",
        "        self._cov_tril.data = torch.full_like(self._cov_tril.data, 0.0)\n",
        "\n",
        "    # define the q distribution standard deviations with property decorator\n",
        "    @property\n",
        "    def mean(self):\n",
        "        # flatten the weight matrix into a vector\n",
        "        return torch.concat((self.weight_mean.flatten(), self.bias_mean.flatten()))\n",
        "\n",
        "    @property\n",
        "    def scale_tril(self):\n",
        "        # this returns the cholesky decomposition L of the covariance: Cov = LL^T\n",
        "        return F.softplus(self._cov_diag).diagflat() + torch.tril(self._cov_tril, diagonal=-1)\n",
        "\n",
        "    # KL divergence KL[q||p] between two Gaussians\n",
        "    def kl_divergence(self):\n",
        "        q = dist.MultivariateNormal(self.mean, scale_tril=self.scale_tril)\n",
        "        p = dist.MultivariateNormal(self.prior_mean, scale_tril=self.prior_scale_tril)\n",
        "        kl = dist.kl_divergence(q, p).sum()\n",
        "        return kl\n",
        "\n",
        "    # forward pass with Monte Carlo (MC) sampling\n",
        "    def forward(self, input):\n",
        "        weight, bias = self._normal_sample(self.mean, self.scale_tril)\n",
        "        return F.linear(input, weight, bias)\n",
        "\n",
        "    def _normal_sample(self, mean, scale_tril):\n",
        "        sample = mean + scale_tril @ torch.randn_like(mean)\n",
        "        # chunk out the weight and bias\n",
        "        weight_vec, bias = torch.split(sample, [self.dim_in * self.dim_out, self.dim_out])\n",
        "        return weight_vec.reshape(self.dim_out, self.dim_in), bias\n",
        "\n",
        "# now also define a function to make a network with only last layer using variational inference\n",
        "def make_last_layer_bnn(layer_sizes, activation='LeakyReLU', **layer_kwargs):\n",
        "    nonlinearity = getattr(nn, activation)() if isinstance(activation, str) else activation\n",
        "    net = nn.Sequential()\n",
        "    linear_layer_kwargs = {}\n",
        "    if 'device' in layer_kwargs:\n",
        "        linear_layer_kwargs['device'] = layer_kwargs['device']\n",
        "    for i, (dim_in, dim_out) in enumerate(zip(layer_sizes[:-2], layer_sizes[1:-1])):\n",
        "        net.add_module(f'Linear{i}', nn.Linear(dim_in, dim_out, **linear_layer_kwargs))\n",
        "        if i < len(layer_sizes) - 2:\n",
        "            net.add_module(f'Nonlinarity{i}', nonlinearity)\n",
        "    # last layer using variational inference with full covariance matrix Gaussian\n",
        "    net.add_module(f'FullCovGaussianLinear{i+1}',\n",
        "                   FullCovGaussianLinear(layer_sizes[-2], layer_sizes[-1], **layer_kwargs))\n",
        "    return net"
      ],
      "metadata": {
        "id": "Hhi3VIdSNQ8r"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now we can run the same regression test again and see how the above full-covariance approximation method compare with the MFVI approach above."
      ],
      "metadata": {
        "id": "shfw6zNCXhj4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# build a last-layer-BNN with 2-hidden-layers, with hidden layer width = h_dim\n",
        "h_dim = 50\n",
        "layer_sizes = [x_dim, h_dim, h_dim, y_dim]\n",
        "# you can change the activation function here\n",
        "activation=nn.GELU()\n",
        "# you can change the prior parameters as you wish\n",
        "layer_kwargs = {'prior_weight_std': 1.0,\n",
        "                'prior_bias_std': 1.0,\n",
        "                'sqrt_width_scaling': False,\n",
        "                'init_std': 0.5,\n",
        "                'device': device}\n",
        "last_layer_bnn_regression_net = make_last_layer_bnn(layer_sizes, activation=activation, **layer_kwargs)\n",
        "# we assume a Gaussian likelihood with homogeneuous noise\n",
        "last_layer_bnn_log_noise_var = nn.Parameter(torch.ones(size=(), device=device)*-3.0)\n",
        "# print out the BNN settings\n",
        "print(\"BNN architecture: \\n\", last_layer_bnn_regression_net)\n",
        "\n",
        "y_pred_mean, y_pred_std_noiseless = get_regression_results(last_layer_bnn_regression_net, x_test_norm, K)\n",
        "model_noise_std = unnormalise_data(to_numpy(torch.exp(0.5*last_layer_bnn_log_noise_var)), 0.0, dataset.y_std)\n",
        "y_pred_std = np.sqrt(y_pred_std_noiseless ** 2 + model_noise_std**2)\n",
        "plot_regression(x_train, y_train, x_test, y_pred_mean, y_pred_std_noiseless, y_pred_std,\n",
        "                title='BNN init (before training, last-layer BNN)')\n",
        "print(model_noise_std, noise_std, y_pred_std_noiseless.mean())"
      ],
      "metadata": {
        "id": "8AHiqu5eXpC_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now start training for this \"last-layer-BNN\" network.\n"
      ],
      "metadata": {
        "id": "w2FjUc8XZKFh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# start training\n",
        "learning_rate = 1e-3\n",
        "params = list(last_layer_bnn_regression_net.parameters()) + [last_layer_bnn_log_noise_var]\n",
        "last_layer_bnn_opt = torch.optim.Adam(params, lr=learning_rate)\n",
        "data_loss_func = lambda y, y_pred: -gauss_loglik(y, y_pred, last_layer_bnn_log_noise_var)\n",
        "\n",
        "beta = 1.0\n",
        "N_epochs = 2000\n",
        "logs = train_network(last_layer_bnn_regression_net, last_layer_bnn_opt,\n",
        "                     dataloader, data_loss_func, beta=beta, verbose=True, N_epochs=N_epochs)\n",
        "# plot the training curve\n",
        "plot_training_loss(logs, beta)"
      ],
      "metadata": {
        "id": "WPERGYPbZKl-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Training is now finished.\n",
        "\n",
        "Let's plot the predictions again, and compare it with the MFVI results above."
      ],
      "metadata": {
        "id": "UWLGn6sxZmnC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# plot the prediction results\n",
        "y_pred_mean, y_pred_std_noiseless = get_regression_results(last_layer_bnn_regression_net, x_test_norm, K)\n",
        "model_noise_std = unnormalise_data(to_numpy(torch.exp(0.5*last_layer_bnn_log_noise_var)), 0.0, dataset.y_std)\n",
        "y_pred_std = np.sqrt(y_pred_std_noiseless ** 2 + model_noise_std**2)\n",
        "plot_regression(x_train, y_train, x_test, y_pred_mean, y_pred_std_noiseless, y_pred_std,\n",
        "                title='BNN approx. posterior (last-layer BNN)')\n",
        "print(model_noise_std, noise_std, y_pred_std_noiseless.mean())"
      ],
      "metadata": {
        "id": "9C7CmwUfcpui"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Similar to the MFVI tests, you can also test various hyper-parameter settings and see how they would affect the regression results using \"last-layer BNN\"."
      ],
      "metadata": {
        "id": "smpFt88Gi0FS"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Monte Carlo dropout (MC-dropout)"
      ],
      "metadata": {
        "id": "9YwGM-0WQhTx"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "For the mathematical foundation, see [lecture notes here](http://yingzhenli.net/home/pdf/ProbAI2022_lecture_note.pdf).\n",
        "\n",
        "For the MC-dropout method with dropout probability $p$, the q distribution is a mixture of delta measures. Specifically, for a layer with parameters $\\{W, b \\}, W \\in \\mathbb{R}^{d_{out} \\times d_{in}}, b \\in \\mathbb{R}^{d_{out}}$, the q distribution is\n",
        "\n",
        "$q(W, b) = \\prod_{i=1}^{d_{out}}q(W_i, b_i),$\n",
        "\n",
        "$q(W_i, b_i) = p \\delta(W_i = 0, b_i = 0) + (1 - p) \\delta(W_i = M_i, b_i = m_i),$\n",
        "\n",
        "with variational parameter for the layer as $\\{M, m \\}, M \\in \\mathbb{R}^{d_{out} \\times d_{in}}, m \\in \\mathbb{R}^{d_{out}}$. This means there are two equivalent way to compute $\\mathbb{E}_q[\\log p(y|x, \\theta)]$ with Monte Carlo: for the forward pass of a layer, below computations are equivalent:\n",
        "\n",
        "1.   **drop weights**: sample $W, b \\sim q(W, b)$, then compute $y = xW^T + b$;\n",
        "2.   **drop units**: compute $\\hat{y} = x M^T + m$, then apply dropout $y = \\text{dropout}(\\hat{y}; p)$.\n",
        "\n",
        "The $KL[q||p]$ regularizer for MC-dropout is ill-defined for Gaussian prior $p(\\theta)$. In practice this is replaced by an $\\ell_2$ regularizer, i.e., $\\frac{1-p}{2\\sigma^2_{\\text{prior}}}|| M ||_2^2$ for the weight variational parameter $M$ (and similarly for $m$).\n",
        "\n",
        "In below block you will implment a BNN layer with MC-dropout, and more specifically:\n",
        "\n",
        "* the forward pass with Monte Carlo: you are asked to implement both \"drop weight\" and the \"drop unit\" approach.\n",
        "* the $KL[q||p]$ regularisation term for one layer: use the $\\ell_2$ regularizer"
      ],
      "metadata": {
        "id": "Vfg88KutVyPb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class MCDropoutLinear(nn.Module):\n",
        "    \"\"\"Applies a linear transformation to the incoming data: y = xW^T + b, where\n",
        "    the weight W and bias b are sampled from the q distribution.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, dim_in, dim_out, prior_weight_std=1.0, prior_bias_std=1.0, dropout_prob=0.1,\n",
        "                 sqrt_width_scaling=False, device=None, dtype=None):\n",
        "        factory_kwargs = {'device': device, 'dtype': dtype}\n",
        "        super(MCDropoutLinear, self).__init__()\n",
        "        self.dim_in = dim_in  # dimension of network layer input\n",
        "        self.dim_out = dim_out  # dimension of network layer output\n",
        "        self.dropout_prob = dropout_prob\n",
        "        # the boolean indicator for the dropout implementation\n",
        "        # if true then run the \"drop weight\" forward pass, else run \"drop unit\"\n",
        "        self.drop_weight = False\n",
        "\n",
        "        # define the trainable variational parameters for q distribtuion\n",
        "        # first define and initialise the mean parameters\n",
        "        self.weight = nn.Parameter(torch.empty((dim_out, dim_in), **factory_kwargs))\n",
        "        self.bias = nn.Parameter(torch.empty(dim_out, **factory_kwargs))\n",
        "        self.reset_parameters()\n",
        "\n",
        "        # define the prior parameters (for prior p, assume the mean is 0)\n",
        "        prior_mean = 0.0\n",
        "        if sqrt_width_scaling:  # prior variance scales as 1/dim_in\n",
        "            prior_weight_std /= self.dim_in ** 0.5\n",
        "        # prior parameters are registered as constants\n",
        "        self.register_buffer('prior_weight_mean', torch.full_like(self.weight, prior_mean))\n",
        "        self.register_buffer('prior_weight_std', torch.full_like(self.weight, prior_weight_std))\n",
        "        self.register_buffer('prior_bias_mean', torch.full_like(self.bias, prior_mean))\n",
        "        self.register_buffer('prior_bias_std', torch.full_like(self.bias, prior_bias_std))\n",
        "\n",
        "    def extra_repr(self):\n",
        "        s = \"dim_in={}, dim_in={}, bias=True\".format(self.dim_in, self.dim_out)\n",
        "        weight_std = self.prior_weight_std.data.flatten()[0]\n",
        "        if torch.allclose(weight_std, self.prior_weight_std):\n",
        "            s += f\", weight prior std={weight_std.item():.2f}\"\n",
        "        bias_std = self.prior_bias_std.flatten()[0]\n",
        "        if torch.allclose(bias_std, self.prior_bias_std):\n",
        "            s += f\", bias prior std={bias_std.item():.2f}\"\n",
        "        return s\n",
        "\n",
        "    def reset_parameters(self, init_std=0.5):\n",
        "        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))\n",
        "        bound = self.dim_in ** -0.5\n",
        "        nn.init.uniform_(self.bias, -bound, bound)\n",
        "\n",
        "    def set_dropout_mode(self, drop_weight=False):\n",
        "        self.drop_weight = drop_weight\n",
        "\n",
        "    # KL divergence KL[q||p]: for MCDropout this returns an l2 regulariser\n",
        "    def kl_divergence(self):\n",
        "        ### begin of your code ###\n",
        "        # hints: using an appropriate l2 regulariser as the kl\n",
        "\n",
        "        kl_weight =\n",
        "        kl_bias =\n",
        "\n",
        "        ### end of your code ###\n",
        "        return kl_weight + kl_bias\n",
        "\n",
        "    # forward pass with Monte Carlo (MC) sampling\n",
        "    def forward(self, input):\n",
        "        ### begin of your code ###\n",
        "        # hints: your can implement the dropout method using an appropriate masking matrix\n",
        "        if self.drop_weight:  # drop weight method\n",
        "            out =\n",
        "\n",
        "        else: # drop unit method\n",
        "            out =\n",
        "\n",
        "        ### end of your code ###\n",
        "        return out\n",
        "\n",
        "# construct a BNN\n",
        "def make_mcdropout_bnn(layer_sizes, activation='LeakyReLU', **layer_kwargs):\n",
        "    nonlinearity = getattr(nn, activation)() if isinstance(activation, str) else activation\n",
        "    net = nn.Sequential()\n",
        "    for i, (dim_in, dim_out) in enumerate(zip(layer_sizes[:-1], layer_sizes[1:])):\n",
        "        net.add_module(f'MCDropoutLinear{i}', MCDropoutLinear(dim_in, dim_out, **layer_kwargs))\n",
        "        if i < len(layer_sizes) - 2:\n",
        "            net.add_module(f'Nonlinarity{i}', nonlinearity)\n",
        "    return net"
      ],
      "metadata": {
        "id": "SeuiRYwNn2Or"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now we can run the same regression test again and see how the MC-dropout method compare with the other approaches above."
      ],
      "metadata": {
        "id": "jgKHmufGU4-g"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# build a BNN with 2-hidden-layers, with hidden layer width = h_dim\n",
        "h_dim = 50\n",
        "layer_sizes = [x_dim, h_dim, h_dim, y_dim]\n",
        "# you can change the activation function here\n",
        "activation=nn.GELU()\n",
        "# you can change the prior parameters as you wish\n",
        "layer_kwargs = {'prior_weight_std': 1.0,\n",
        "                'prior_bias_std': 1.0,\n",
        "                'sqrt_width_scaling': False,\n",
        "                'dropout_prob': 0.1,\n",
        "                'device': device}\n",
        "mcdropout_bnn_regression_net = make_mcdropout_bnn(layer_sizes, activation=activation, **layer_kwargs)\n",
        "# we assume a Gaussian likelihood with homogeneuous noise\n",
        "mcdropout_bnn_log_noise_var = nn.Parameter(torch.ones(size=(), device=device)*-3.0)\n",
        "# print out the BNN settings\n",
        "print(\"BNN architecture: \\n\", mcdropout_bnn_regression_net)\n",
        "\n",
        "def set_dropout_mode(bnn, drop_weight=False):\n",
        "    print(\"set the dropout mode: drop_weight={}\".format(drop_weight))\n",
        "    for module in bnn:\n",
        "        if hasattr(module, 'set_dropout_mode'):\n",
        "            module.set_dropout_mode(drop_weight)\n",
        "\n",
        "set_dropout_mode(mcdropout_bnn_regression_net, drop_weight=True)\n",
        "y_pred_mean, y_pred_std_noiseless = get_regression_results(mcdropout_bnn_regression_net, x_test_norm, K)\n",
        "model_noise_std = unnormalise_data(to_numpy(torch.exp(0.5*mcdropout_bnn_log_noise_var)), 0.0, dataset.y_std)\n",
        "y_pred_std = np.sqrt(y_pred_std_noiseless ** 2 + model_noise_std**2)\n",
        "plot_regression(x_train, y_train, x_test, y_pred_mean, y_pred_std_noiseless, y_pred_std,\n",
        "                title='BNN init (before training, MC-dropout)')\n",
        "print(model_noise_std, noise_std, y_pred_std_noiseless.mean())"
      ],
      "metadata": {
        "id": "ulX_qTfjm-JU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now start training for this \"MC-dropout BNN\" network."
      ],
      "metadata": {
        "id": "s13GEP1XVIkq"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# start training\n",
        "learning_rate = 1e-3\n",
        "params = list(mcdropout_bnn_regression_net.parameters()) + [mcdropout_bnn_log_noise_var]\n",
        "mcdropout_bnn_opt = torch.optim.Adam(params, lr=learning_rate)\n",
        "data_loss_func = lambda y, y_pred: -gauss_loglik(y, y_pred, mcdropout_bnn_log_noise_var)\n",
        "set_dropout_mode(mcdropout_bnn_regression_net, drop_weight=False)\n",
        "\n",
        "beta = 1.0\n",
        "N_epochs = 2000\n",
        "logs = train_network(mcdropout_bnn_regression_net, mcdropout_bnn_opt,\n",
        "                     dataloader, data_loss_func, beta=beta, verbose=True, N_epochs=N_epochs)\n",
        "# plot the training curve\n",
        "plot_training_loss(logs, beta)"
      ],
      "metadata": {
        "id": "W-j_7rDhnGdL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Training is now finished.\n",
        "\n",
        "Let's plot the predictions again, and compare it with the results above."
      ],
      "metadata": {
        "id": "oCFsxSHsVTeA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# plot the prediction results\n",
        "set_dropout_mode(mcdropout_bnn_regression_net, drop_weight=True)\n",
        "y_pred_mean, y_pred_std_noiseless = get_regression_results(mcdropout_bnn_regression_net, x_test_norm, K)\n",
        "model_noise_std = unnormalise_data(to_numpy(torch.exp(0.5*mcdropout_bnn_log_noise_var)), 0.0, dataset.y_std)\n",
        "y_pred_std = np.sqrt(y_pred_std_noiseless ** 2 + model_noise_std**2)\n",
        "plot_regression(x_train, y_train, x_test, y_pred_mean, y_pred_std_noiseless, y_pred_std,\n",
        "                title='BNN approx. posterior (MC-dropout)')\n",
        "print(model_noise_std, noise_std, y_pred_std_noiseless.mean())"
      ],
      "metadata": {
        "id": "sFThZlR6nNKG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Similar to previous tests, you can also test various hyper-parameter settings and see how they would affect the regression results using MC-dropout.\n",
        "\n",
        "In particular, try setting different value for ```dropout_prob``` in ```layer_kwargs```.\n",
        "\n",
        "Also you can try using \"drop unit\" forward pass when visualising the results, by using\n",
        "\n",
        "```set_dropout_mode(mcdropout_bnn_regression_net, drop_weight=False)```\n",
        "\n",
        "You can also try using \"drop weight\" forward pass when training:\n",
        "\n",
        "```set_dropout_mode(mcdropout_bnn_regression_net, drop_weight=True)```\n",
        "\n",
        "By doing so you'll see some visual differences in terms of learning curves and posterior predictive mean/std visualisations. Can you explain that?"
      ],
      "metadata": {
        "id": "MgqfF_nqVZcw"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Case study 1: Bayesian optimisation"
      ],
      "metadata": {
        "id": "4gyMwQLznjrr"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### What is Bayesian optimisation (BO)?\n",
        "\n",
        "(Feel free to skip this part if you already understand the idea.)\n",
        "\n",
        "Assume you are interested in optimising a function $x^* = \\arg\\max f_0(x)$, but you don't actually know the analytic form of $f(x)$. Rather, you only have access of it as a \"black-box\", where you provide an input $x$ to this \"black-box\", and it will return you a (noisy version of) output $y = f_0(x) + \\text{noise}$.\n",
        "\n",
        "Bayesian optimisation (BO) is a class of methods that tackle this challenge.\n",
        "\n",
        "To motivate BO, let's imagine you already have a set of datapoints $D = \\{ (x_n, y_n) \\}_{n=1}^N$ collected by sending the queries $x_1, ..., x_N$ to the above mentioned \"black-box\". Then we can fit a *surrogate function* $f_{\\theta}(x)$ to data. In large data limit $N \\rightarrow +\\infty$, with flexible enough model, we expect $f_{\\theta}(x)$ to be very close, if not identical, to the ground-truth function $f_0(x)$. In such case we can solve the optimisation problem by finding $x^* = \\arg\\max f_{\\theta}(x)$ instead, which is tractable.\n",
        "\n",
        "However, in practice we don't have a big dataset to train the surrogate model. This is expecially the case if the \"black-box\" corresponds to an expensive experiment (e.g., training a Transformer network where $x$ represents the hyper-parameter settings). In such scenario $f_{\\theta}(x)$ will be quite different from $f(x)$ in most of the unseen input locations.\n",
        "\n",
        "The key idea of BO is to optimise $f_0$ using \"helps\" from this surrogate by taking the uncertainty of model fitting into account, and it aims to find the optimum of $f_0$ with least amount of queries to the expensive \"black-box\". Specifically, there are 3 ingredients of an BO method:\n",
        "\n",
        "1.   Acquisition function: using the current estimated surrogate function $f_{\\theta}(x)$ and its uncertainty estimate to compute an acquisition function;\n",
        "2.   Query the \"black-box\": find the next input to query by maximising the acquisition function;\n",
        "3.   Surrogate model update: given new queried result and historical data, update the surrogate model and its uncertainty estimate.\n",
        "\n",
        "At the beginning since the model is uncertain, a good acquisition function will take uncertainty into account and encourage \"exploration\", i.e., querying inputs at different locations.\n",
        "\n",
        "As we collect more data, with proper Bayesian posterior updates, the surrogate model $f_{\\theta}$ will become closer and closer to the ground-truth function $f_0$ and the uncertainty will be reduced.\n",
        "\n",
        "So at some point, the model will become certain about its output, and a good acquisition function will also enable \"exploitation\" at this stage, to seek for the optimum of the surrogate function $f_{\\theta}$ as to solve the original optimization problem."
      ],
      "metadata": {
        "id": "3kDCCbrqiMEJ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### The Upper Confidence Bound (UCB) method\n",
        "\n",
        "The UCB is an acquisition function that uses both mean prediction and uncertainty. Specifically, assume the surrogate model provides both mean $m(x)$ and standard deviation $\\sigma(x)$ for a given input $x$, then the UCB acquisition function is the following:\n",
        "\n",
        "$$UCB(x) = m(x) + \\beta \\sigma(x).$$\n",
        "\n",
        "And the query procedure will pick the next query input as $x = \\arg\\max UCB(x)$.\n",
        "\n",
        "Initially, $\\sigma(x)$ can be quite large for many regions, meaning that UCB will mainly explore. As we collect more data, $\\sigma(x)$ will decrease around the regions that have been searched, and this allows the algorithm to\n",
        "\n",
        "1.   ignore some searched regions where the model confidently thinks their function value is small;\n",
        "2.   exploit some other searched regions where the model confidently believes the optimum might be there;\n",
        "3.   explore some other promising regions that has not been searched before.\n",
        "\n",
        "Here $\\beta$ a hyper-parameter specified by the user at a time, to achieve a desired balance between exploration ($\\sigma(x)$) and exploitation ($m(x)$). When $\\beta = 0$, it means we trust the surrogate model and exploit on that. When $\\beta$ is large, we allow the query process to focus on regions that have large $\\sigma(x)$ value (where the model is most uncertainty about). For optimal BO, this $\\beta$ coefficient will decrease during time; for simplicity, in this demo we only consider a fixed value of $\\beta$."
      ],
      "metadata": {
        "id": "bVVZAjrRiQAw"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "First we implement the ```ucb``` acquisition function as mentioned above."
      ],
      "metadata": {
        "id": "jIVBSDOViWXa"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# the UCB acquisition function\n",
        "def ucb(net, x, beta, K):\n",
        "    ### begin of your code ###\n",
        "    # hint: do K times of forward passes and then compute the mean and std\n",
        "    # final results should have shape (N, 1) with x having shape (N, dim_x)\n",
        "\n",
        "    y_pred_mean =\n",
        "    y_pred_std =\n",
        "\n",
        "    ### end of your code ###\n",
        "    return y_pred_mean + math.sqrt(beta) * y_pred_std, y_pred_mean, y_pred_std\n",
        "\n",
        "# some helper function\n",
        "def get_next_query(net, x, acquisition_func, **kwargs):\n",
        "    # search for the next query point in x using the acquisition function\n",
        "    # and here we also assume that x_dim = 1\n",
        "    f, y_pred_mean, y_pred_std = acquisition_func(net, x, **kwargs)  # shape (N, 1)\n",
        "    next_query = x[torch.argmax(f[:, 0])].unsqueeze(0)\n",
        "    return next_query, f, y_pred_mean, y_pred_std\n",
        "\n",
        "def query_oracle(x, oracle_func, noise_std=None):\n",
        "    # given input, query the oracle function (potentially with noise)\n",
        "    y = oracle_func(x)\n",
        "    if noise_std is not None and noise_std > EPS:\n",
        "        y += np.random.randn(y.shape[0], y.shape[1]) * noise_std\n",
        "    return y\n",
        "\n",
        "def reset_parameters(bnn, init_std):\n",
        "    print(\"re-initialise the parameters with init_std={}\".format(init_std))\n",
        "    for module in bnn:\n",
        "        if hasattr(module, 'reset_parameters'):\n",
        "            module.reset_parameters(init_std)"
      ],
      "metadata": {
        "id": "uEZGG31Gnoak"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now let's implement the main loop of BO."
      ],
      "metadata": {
        "id": "DI1pQKESiahp"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# now define the Bayesian optimisation search\n",
        "def bayesian_optimisation(net, log_noise_var, oracle_func, acquisition_func, x, x_init, T,\n",
        "                          noise_std = None, init_std = 0.1, N_epochs = 1000, **kwargs):\n",
        "    # we start from a number of observations to initialise the dataset\n",
        "    y_init = query_oracle(x_init, oracle_func, noise_std)\n",
        "    observations = regression_data(x_init, y_init)\n",
        "    # preparing the query candidates\n",
        "    x_norm = normalise_data(x, observations.x_mean, observations.x_std)\n",
        "    x_tensor = torch.tensor(x_norm).float().to(device)\n",
        "    # now start the loop of BO for T steps\n",
        "    acquisition_func_list = []\n",
        "    mean_func_list = []; std_func_list = []\n",
        "    for t in range(T):\n",
        "        dataloader = DataLoader(observations, batch_size=min(100, len(observations)), shuffle=True)\n",
        "        # train the network\n",
        "        print(\"update BNN posterior...\")\n",
        "        params = list(net.parameters()) + [log_noise_var]\n",
        "        opt = torch.optim.Adam(params, lr=1e-3)\n",
        "        data_loss_func = lambda y, y_pred: -gauss_loglik(y, y_pred, log_noise_var)\n",
        "        # if using MC-dropout, use \"drop unit\" mode for training\n",
        "        set_dropout_mode(net, drop_weight=False)  # won't do anything if not usig MC-dropout\n",
        "        train_network(net, opt, dataloader, data_loss_func, N_epochs=N_epochs, verbose=True)\n",
        "        # now get next acquisition point and query\n",
        "        print(\"query next point...\")\n",
        "        # if using MC-dropout, use \"drop weight\" mode for computing the UCB acquisition function\n",
        "        set_dropout_mode(net, drop_weight=True)   # won't do anything if not usig MC-dropout\n",
        "        x_next, f, y_pred_mean, y_pred_std = get_next_query(net, x_tensor, acquisition_func, **kwargs)\n",
        "        x_next = unnormalise_data(to_numpy(x_next), observations.x_mean, observations.x_std)\n",
        "        y_next = query_oracle(x_next, oracle_func, noise_std)\n",
        "        # record the selections to be visualised\n",
        "        acquisition_func_list.append(unnormalise_data(to_numpy(f), observations.y_mean, observations.y_std))\n",
        "        mean_func_list.append(unnormalise_data(to_numpy(y_pred_mean), observations.y_mean, observations.y_std))\n",
        "        std_func_list.append(unnormalise_data(to_numpy(y_pred_std), 0.0, observations.y_std))\n",
        "        # add new observation to the dataset\n",
        "        x_init = np.concatenate((x_init, x_next), axis=0)\n",
        "        y_init = np.concatenate((y_init, y_next), axis=0)\n",
        "        observations.update_data(x_init, y_init, update_stats=False)\n",
        "        # reset the network for next round's training\n",
        "        init_std_ = init_std * 0.9 ** (t+1) if init_std is not None else None\n",
        "        reset_parameters(net, init_std=init_std_)\n",
        "        log_noise_var.data = torch.full_like(log_noise_var.data, -3.0)\n",
        "\n",
        "    return acquisition_func_list, mean_func_list, std_func_list, observations\n",
        "\n",
        "# define plot function\n",
        "def plot_bo_one_step(x_init, y_init, x, y_true, x_query, y_query, acquisition_value,\n",
        "                    mean_value, std_value, beta, title):\n",
        "    plt.figure()\n",
        "    plt.plot(x, y_true, 'k-', label='ground truth')\n",
        "    plt.plot(x, mean_value, 'b-', label='prediction mean')\n",
        "    plt.fill_between(x[:,0], mean_value[:,0]-math.sqrt(beta)*std_value[:,0],\n",
        "                     mean_value[:,0]+math.sqrt(beta)*std_value[:,0],\n",
        "                     color='b', alpha=0.3, label='epistemic uncertainty')\n",
        "    plt.plot(x, acquisition_value, 'g-', label='acquisition func.')\n",
        "    plt.plot(x_init, y_init, 'ro', label='observed data')\n",
        "    plt.plot(x_query, y_query, 'go', label='query')\n",
        "    plt.axvline(x_query, ymax=y_query, color='g', linestyle='--')\n",
        "    plt.legend()\n",
        "    plt.title(title)\n",
        "    plt.show()"
      ],
      "metadata": {
        "id": "ne3pSg4TTrSm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "With all the helper functions set, let's test BO with MFVI-BNN on optimising a ground-truth function defined below."
      ],
      "metadata": {
        "id": "HRU1vDmnidXN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# now start the test\n",
        "# first define the ground-truth/oracle function that we want to optimize\n",
        "def ground_truth_func_func(x):\n",
        "    return np.sin(x * math.pi / 2 + 0.8) * np.exp(-0.1 * np.abs(x)) + 0.1 * x\n",
        "oracle_func = ground_truth_func\n",
        "\n",
        "# set-up BO hyper-parameters\n",
        "acquisition_func = ucb\n",
        "ucb_params = {'beta': 1.0, 'K': 50}\n",
        "noise_std = 0.1\n",
        "x_min = -3; x_max = 3; dx = 0.01\n",
        "x = np.arange(x_min, x_max+EPS, dx)[:, np.newaxis]\n",
        "y_true = query_oracle(x, oracle_func, noise_std=None)\n",
        "# choose some random query locations to start with\n",
        "N_init = 5\n",
        "x_init = np.copy(x); np.random.shuffle(x_init); x_init = x_init[:N_init]\n",
        "T = 10  # run T steps of BO search\n",
        "\n",
        "# then set-up the BNN\n",
        "x_dim, y_dim = x.shape[1], y_true.shape[1]\n",
        "h_dim = 50\n",
        "layer_sizes = [x_dim, h_dim, h_dim, y_dim]\n",
        "activation=nn.GELU()\n",
        "# MFVI-BNN hyper-parameters\n",
        "layer_kwargs = {'prior_weight_std': 1.0,\n",
        "                'prior_bias_std': 1.0,\n",
        "                'sqrt_width_scaling': False,\n",
        "                'init_std': 0.1,\n",
        "                'device': device}\n",
        "bo_net = make_mfvi_bnn(layer_sizes, activation=activation, **layer_kwargs)\n",
        "# we assume a Gaussian likelihood with homogeneuous noise\n",
        "bo_log_noise_var = nn.Parameter(torch.ones(size=(), device=device)*-3.0)\n",
        "# print out the BNN settings\n",
        "print(\"BNN architecture: \\n\", bo_net)\n",
        "\n",
        "# start the BO process!\n",
        "acquisition_func_list, mean_func_list, std_func_list, observations = \\\n",
        "    bayesian_optimisation(bo_net, bo_log_noise_var, oracle_func, acquisition_func, x, x_init, T,\n",
        "                          noise_std=noise_std, init_std=layer_kwargs['init_std'], N_epochs=1000,\n",
        "                          **ucb_params)\n",
        "\n",
        "# plot results\n",
        "def plot_bo_results(acquisition_func_list, mean_func_list, std_func_list, observations, T, x, y_true):\n",
        "    for t in range(T):\n",
        "        x_init = unnormalise_data(observations.x[:N_init+t], observations.x_mean, observations.x_std)\n",
        "        y_init = unnormalise_data(observations.y[:N_init+t], observations.y_mean, observations.y_std)\n",
        "        x_query = unnormalise_data(observations.x[N_init+t], observations.x_mean, observations.x_std)\n",
        "        y_query = unnormalise_data(observations.y[N_init+t], observations.y_mean, observations.y_std)\n",
        "        acquisition_value = acquisition_func_list[t]\n",
        "        mean_value = mean_func_list[t]; std_value = std_func_list[t]\n",
        "        title = 'BO with UCB, t={}, beta={}'.format(t+1, ucb_params['beta'])\n",
        "        plot_bo_one_step(x_init, y_init, x, y_true, x_query, y_query, acquisition_value,\n",
        "                         mean_value, std_value, beta = ucb_params['beta'], title=title)\n",
        "plot_bo_results(acquisition_func_list, mean_func_list, std_func_list, observations, T, x, y_true)"
      ],
      "metadata": {
        "id": "ZTy4QJCAkze5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's test again but with some other settings of your choice.\n",
        "\n",
        "First, you can try using a BNN with MC-dropout, by e.g., copy-paste the code in below block to appropriate place.\n",
        "\n",
        "You can also try other hyper-parameters for the BNN, e.g., changing the activation function, number of layers, ```N_epochs```, initialisation of the q standard deviation ```init_std``` if using Gaussian approximations, ```dropout_prob``` if using MC-dropout, etc."
      ],
      "metadata": {
        "id": "vCsJROl2pM7o"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# MC-dropout BNN hyper-parameters\n",
        "layer_kwargs = {'prior_weight_std': 1.0,\n",
        "                'prior_bias_std': 1.0,\n",
        "                'sqrt_width_scaling': False,\n",
        "                'dropout_prob': 0.1,\n",
        "                'device': device}\n",
        "bo_net = make_mcdropout_bnn(layer_sizes, activation=activation, **layer_kwargs)\n",
        "layer_kwargs['init_std'] = None"
      ],
      "metadata": {
        "id": "nReyeRH8aFZt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Second, you can try BO for a different ground-truth function. You can define a new ground-truth function ```new_ground_truth(x)``` and set ```oracle_func = new_ground_truth``` to define the oracle to query. An example is provided below."
      ],
      "metadata": {
        "id": "KqnMnyfGZ3Rw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# copy-paste some of the below code to the appropriate place\n",
        "# example: quadratic function\n",
        "def quadratic_func(x):\n",
        "    return -2 * x ** 2 + 3 * x + 0.3\n",
        "oracle_func = quadratic_func"
      ],
      "metadata": {
        "id": "gb5NpdDUpSWn"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}