Tutorial 3.1 KL_divergence.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "**ML for LFT in a Nutshell**\n",
        "----\n",
        "\n",
        "Recall the following map from Prof. Cheng's lecture this morning.\n",
        "![picture](https://drive.google.com/uc?export=view&id=14eZSrjzaRHuz4Rtf3hAqyJ0uu6QW2dTS)\n",
        "Our endgoal is to generate samples from a target distribution $p(\\phi) = \\frac{e^{-S[\\phi]}}{Z}$.\n",
        "* Action $S[\\phi]$ in exponent is determined by the physics problem at hand. E.g. the action for interacting scalar field theory on lattice is  $$S[\\phi] = \\sum\\limits_{x,y \\in V_L} \\phi_x \\, \\Delta_{x,y} \\, \\phi_y + \\sum\\limits_{x \\in V_L} (m^2 \\phi_x^2 + \\lambda \\phi^4_x) ,$$ at leading order.\n",
        "* $Z$ is an appropriate normalization constant.\n",
        "\n",
        "Let us consider the simplest ML approch to generate target distribution $p(\\phi)$.\n",
        "* Initialize a Neural Network with parameters $\\theta$ sampled from some initial distribution $Q_{\\theta}$, and inputs are $x, y \\in V_L$. Then outputs are drawn from initial distribution $Q(x)$.\n",
        "* In general, the target distribution $p(\\phi)$ requires entirely different choice for network parameter distributions. Let us label this intended distribution as $P_{\\theta}$.\n",
        "* Network outputs for parameters sampled from $P_{\\theta}$, and inputs $x, y \\in V_L$, are labelled as $P(x)$.\n",
        "* Training updates distribution $Q_{\\theta}$ until the difference between distributions $P(x)$ and $Q(x)$ is minimized.\n",
        "\n",
        "**We need a good loss function**\n",
        "----\n",
        "\n",
        "* An ideal loss function will minimize differences between two probability distributions, of the forms $P(x) \\sim e^{- S_P(x)}$ and $Q(x) \\sim e^{- S_Q(x)}$.\n",
        "* If the loss function involves logarithms of $P(x)$ and $Q(x)$, then evaluating derivatives of the loss function becomes easier.\n",
        "\n",
        "KL or Kullback-Leibler divergence makes a good loss function, for problems of this nature.\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "jzRB8VqdKB4M"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Kullback-Leibler diveregnce:**\n",
        "---\n",
        "\n",
        "It's a statistical measure of how much the probability distributions of two events differ from each other, or how much entropy is gained or lost in changing one probability distributon to the other.\n",
        "\n",
        "\n",
        "> Suppose, $Q(x)$ and $P(x)$ denote the probability distributions of *approximate* and *target* distributions $Q$ and $P$ on some space $x$. Then the KL divergence loss from $Q$ to $P$ is given by\n",
        "$$D_{KL}(P || Q) = - \\sum_{x \\in X} P(x) \\log\\bigg(\\frac{Q(x)}{P(x)} \\bigg) .$$\n",
        "Note that $D_{KL}(P||Q) \\neq D_{KL}(Q || P)$, so KL divergence does not play the role of actual distance."
      ],
      "metadata": {
        "id": "0d008zdnN9vr"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "---\n",
        "**Example 0:**\n",
        "---\n",
        "\n",
        "\n",
        "Let us demonstrate how the KL divergence between two distributions are calculated numerically. To keep life simple, we will call the function [*losses.KLDivergence*](https://www.tensorflow.org/api_docs/python/tf/keras/losses/KLDivergence) from \"*keras*\" library in \"*Tensorflow*\", and use it to directly compute the KL divergence between two data sets.  \n",
        "\n",
        "We break this down into the following steps:"
      ],
      "metadata": {
        "id": "urYiNMKKY4Fj"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "*   Import libraries necessary to calculate KL divergence.\n",
        "\n"
      ],
      "metadata": {
        "id": "cJ9MU10ylXOu"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from tensorflow import keras\n",
        "from tensorflow.keras import layers\n",
        "import tensorflow as tf"
      ],
      "metadata": {
        "id": "JPh2SqB7lnPN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "*   Define any two distributions. I am defining target distribution $P$ and approximate distribution $Q$ as two arrays. If you want to try something different, choose any $P$ and $Q$ from the same 2D or 3D Gaussian distribution, and compute the KL divergence between them.\n",
        "\n"
      ],
      "metadata": {
        "id": "IxzUbRJMo3Yg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# True distribution P_x and approximate distribution Q_x\n",
        "P_x = [[-0.4, 1], [3, 0.2]]\n",
        "Q_x = [[0.16, 0.4], [0.14, 2.1]]"
      ],
      "metadata": {
        "id": "helIaHqhAKA4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "* Next, the function [*losses.KLDivergence*](https://www.tensorflow.org/api_docs/python/tf/keras/losses/KLDivergence) takes in two arguments - 1) reduction, and 2) name. We need to specify \"*name=kl_divergence*\", and for reduction, we want something that will sum over the batch size (to keep our ML spirit up). The choice for summing over batch size is \"*reduction=auto*\", used here."
      ],
      "metadata": {
        "id": "zb7CCDR6AStl"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "tf.keras.losses.KLDivergence(reduction=\"auto\", name=\"kl_divergence\")\n",
        "kl = tf.keras.losses.KLDivergence()\n",
        "kl(P_x, Q_x).numpy()"
      ],
      "metadata": {
        "id": "ITtJ5wXOl8XE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "\n",
        "---\n",
        "\n",
        "\n",
        "**Example 1: Let us use KL divergence loss function for CIFAR10 image classification.**\n",
        "---\n",
        "\n",
        "First, walk through the following steps to define train data and test data from CIFAR10, and then initialize a simple CNN. We will utilize Keras library in Tensorflow for this example, as its prebuilt functions are helpful for beginners.\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "P0j0hNJdoXoX"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "*   First, we import the libraries that we need.\n",
        "\n"
      ],
      "metadata": {
        "id": "pJnNb326p2nk"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import keras\n",
        "from keras.datasets import cifar10\n",
        "from keras.models import Sequential\n",
        "from keras.layers import Dense, Dropout, Flatten, Conv2D, MaxPooling2D\n",
        "from keras import backend as K"
      ],
      "metadata": {
        "id": "YNBftyn4p090"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "*   Given below are model configurations. One can choose differently, I encourage you to play with these attributes at a later time.\n",
        "\n"
      ],
      "metadata": {
        "id": "LEmXD4ivqO4-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "img_width, img_height         = 32, 32\n",
        "batch_size                    = 250\n",
        "no_epochs                     = 15\n",
        "no_classes                    = 10\n",
        "validation_split              = 0.4\n",
        "verbosity                     = 1"
      ],
      "metadata": {
        "id": "zzMMjSBjqNKL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "* Load CIFAR10 dataset as train and test data."
      ],
      "metadata": {
        "id": "GlPNVrSYqmy6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "(i_train, o_train), (i_test, o_test) = cifar10.load_data()"
      ],
      "metadata": {
        "id": "KosMUY_pqipq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "* We need to reshape data based on channels first / channels last strategy."
      ],
      "metadata": {
        "id": "8roXAtdxqfCd"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "if K.image_data_format() == 'channels_first':\n",
        "    i_train = i_train.reshape(i_train.shape[0],3, img_width, img_height)\n",
        "    i_test = i_test.reshape(i_test.shape[0], 3, img_width, img_height)\n",
        "    i_shape = (3, img_width, img_height)\n",
        "else:\n",
        "    i_train = i_train.reshape(i_train.shape[0], img_width, img_height, 3)\n",
        "    i_test = i_test.reshape(i_test.shape[0], img_width, img_height, 3)\n",
        "    i_shape = (img_width  , img_height, 3)\n"
      ],
      "metadata": {
        "id": "uSLPq487rCoW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "* We want all numbers as floats."
      ],
      "metadata": {
        "id": "iMjKVWotrb5G"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "i_train = i_train.astype('float32')\n",
        "i_test = i_test.astype('float32')"
      ],
      "metadata": {
        "id": "Kb62IZunrVqT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "* The test and train data require an appropriate normalization."
      ],
      "metadata": {
        "id": "oEPmcTXPrw06"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "i_train = i_train / 255\n",
        "i_test = i_test / 255"
      ],
      "metadata": {
        "id": "mz2e7h52rr7t"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "* Next, we need to convert target vectors to categorical targets."
      ],
      "metadata": {
        "id": "yNZINhkDry9f"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "o_train = keras.utils.to_categorical(o_train, no_classes)\n",
        "o_test = keras.utils.to_categorical(o_test, no_classes)"
      ],
      "metadata": {
        "id": "kAV1XQSjr_Gx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "* A CNN model is defined below. You can add or remove layers, or change layer sizes, or activation functions, if you want."
      ],
      "metadata": {
        "id": "Ff47tA-_tJVv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model = Sequential()\n",
        "model.add(Conv2D(32, kernel_size=(3, 3), activation='tanh', input_shape=i_shape))\n",
        "model.add(MaxPooling2D(pool_size=(2, 2)))\n",
        "model.add(Dropout(0.40))\n",
        "model.add(Conv2D(128, kernel_size=(3, 3), activation='relu'))\n",
        "model.add(MaxPooling2D(pool_size=(2, 2)))\n",
        "model.add(Dropout(0.70))\n",
        "model.add(Flatten())\n",
        "model.add(Dense(64, activation='exponential'))\n",
        "model.add(Dense(no_classes, activation='softmax'))"
      ],
      "metadata": {
        "id": "f86o38MatDpn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "* Next, it is time to train the CNN! There are two steps to this.\n",
        " * **First**, we need to \"*compile*\" the model, where we specify the **loss function** and the optimizer. These choices are used for training. We will use the function, given below, for compiling the CNN.\n",
        "```\n",
        "model.compile(loss= ,optimizer= ,metrics= )  \n",
        "```\n",
        "  Here, we will define a custom KL divergence loss function, instead of calling the inbuilt KL divergence loss ```keras.losses.kullback_leibler_divergence```. To keep life simple, we will choose ```optimizer=keras.optimizers.Adam()``` and ```metrics=['accuracy']```, respectively.\n",
        "\n",
        " * **Next**, we need to train the CNN. For this, we will use\n",
        "```\n",
        "model.fit(X, Y, epochs= , batch_size= , verbose= , validation_split= )\n",
        "```\n",
        "where ```X``` and ```Y``` are model inputs and outputs.\n",
        "\n",
        "  One can obtain training history using *.History()*. Moreover, the train and test loss and accuracies can be obtained using *.history['loss']*, *.history['val_loss']*, *.history['accuracy']*, and *.history['val_accuracy']*, respectively."
      ],
      "metadata": {
        "id": "PfNbtkpPvUZ3"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Exercise 1: Define a KL loss function to compile this CNN, and then train it. Plot the loss and accuracies for train and test data sets.**\n",
        "---"
      ],
      "metadata": {
        "id": "mTikxsQxNixZ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Forward and Reverse KL divergences:**\n",
        "---\n",
        "\n",
        "$D_{KL}(P||Q)$ and $D_{KL}(Q||P)$ are respectively known as the Forward and Reverse KL divergences for true and approximate distributions $P(x)$ and $Q(x)$.\n",
        "\n",
        "In the case of Forward KL divergence, $P(x) = 0$ does not contribute towards minimizing the loss, and $Q(x)$ is updated to approximate those data points for which $P(x)>0$. Thus, Forward KL divergence may not result in the best approximate distribution $Q(x)$, at the end of training.\n",
        "\n",
        "> Let us demonstrate this pictorially --\n",
        "![picture](https://drive.google.com/uc?export=view&id=15HnW-cd2QInRNP5LapvvPlTyGYCq-XiI)\n",
        "* P(x) is our target distribution.\n",
        "* Q(x) is the initial distribution, before we start training.\n",
        "* Using Forward KL divergence, the NN model learns to approximate for any $P(x)>0$, such approximation is bad here.\n",
        "* We obtain the following Q(x).\n",
        "![picture](https://drive.google.com/uc?export=view&id=1j-hZ-Co_R_eaECTUJPWmrtts31WjePiB)\n",
        "* Not the final approximate distribution we would like.\n",
        "\n",
        "\n",
        "On the other hand, the *Reverse KL divergence* $$D_{KL}(Q || P) = - \\sum_{x \\in X} Q(x) \\log\\bigg(\\frac{P(x)}{Q(x)} \\bigg)$$ ensures that any point $Q(x)=0$ does not contribute to optimization. Thus, $P(x)$ can not be learned starting from $Q(x)=0$; however, whenever $Q(x)>0$, reverse KL divergence contributes to learning, for both $P(x)=0$ and $P(x)>0$ cases.\n",
        "\n",
        "Finally, the choice of Forward vs. Reverse KL divergence depends entirely on true and initial approximate distributions $P(x)$ and $Q(x)$.\n"
      ],
      "metadata": {
        "id": "BmX67hP8a9ki"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "**Exercise 2: Define a Reverse KL loss, then plot and compare the train and test losses and accuracies of the same CNN using this on same test, train datasets. Any improvement?**\n",
        "---\n"
      ],
      "metadata": {
        "id": "VIgxkxw0N8Lf"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "References\n",
        "* https://github.com/christianversloot/machine-learning-articles/blob/main/how-to-use-kullback-leibler-divergence-kl-divergence-with-keras.md\n",
        "* https://keras.io/api/losses/probabilistic_losses/#kldivergence-function\n",
        "* https://www.analyticsvidhya.com/blog/2022/09/dummies-guide-to-writing-a-custom-loss-function-in-tensorflow/\n",
        "* https://dibyaghosh.com/blog/probability/kldivergence.html"
      ],
      "metadata": {
        "id": "38r0rt7_w-SE"
      }
    }
  ]
}