hw_lab/nn-complexity.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Measuring model complexity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Intro\n",
    "In this assignment we are going to implement some complexity measures. Complexity measure is a magical blackbox function, which takes some information about a neural network, a train dataset and tells us how well our model should generalize. We have encountered several generalization bounds during the course, which can serve as complexity measures, but typically they do not work well in practice. Actually, the primary goal of this assignment is to figure how well current complexity measures work (if they work at all)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Task description\n",
    "Since we haven't polished this assignment yet, we decided to loose the requirements. We provide all the training/visualizion code (which is, actually, the most time-consuming part) and you will only have to implement the complexity measures. There are 3 complexity measures for you to implement:\n",
    "\n",
    "* [2 points] Spectral Complexity\n",
    "* [2 points] Normalized Spectral Complexity\n",
    "* [1 point] Noise Sensitivity\n",
    "\n",
    "We are going to investigate them in different setups:\n",
    "* Depending on the network width\n",
    "* Depending on the number of corrupted samples (i.e. samples with random labels)\n",
    "\n",
    "For convenience we use fully-connected NN with 1 hidden layer in all the experiments.\n",
    "\n",
    "**Important requirements**:\n",
    "* Your results *must be reproducible* (that's why we are so paranoid about fixing random seeds)\n",
    "* All the cells in your notebook must be executable in sequential order"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Spectral Complexity\n",
    "Check formula (1.2) of [this paper](https://arxiv.org/abs/1706.08498). We want you to implement this formula, setting reference matrices to be zero matrices."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Complexity Measures"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Mean Normalized Margin\n",
    "\n",
    "Mean Normalized Margin is our ad-hoc complexity measure that is build upon Spectral Complexity.\n",
    "We compute it the following way:\n",
    "* let $R_\\mathcal{A}$ be a spectral complexity of the model $\\mathcal{A}$\n",
    "* let $X$ be the training dataset.\n",
    "* let $m_i$ be a margin for a training sample $x_i$, i.e.\n",
    "$$\n",
    "m_i = l_{y_i} - \\arg\\max_{c \\neq y_i}l_c,\n",
    "$$\n",
    "where $l_c$ is the logit for class $c$, produced by our model.\n",
    "So, margin is just a difference between \"true\" logit and the logit which is closest to it, i.e. our \"second guess\".\n",
    "The idea is that when our model is very confident in its predictions, then all margins are huge.\n",
    "If our model perfectly fits training dataset then $m_i$ is positive for all training pairs $(x_i, y_i)$.\n",
    "\n",
    "We define normalized margin $\\hat{m}_i$ to be:\n",
    "\n",
    "$$\n",
    "\\hat{m}_i = m_i \\cdot \\frac{\\sqrt{n}}{R_\\mathcal{A} \\cdot \\|X\\|_F}\n",
    "$$\n",
    "\n",
    "And Mean Normalized Margin $M_\\mathcal{A}$ is just a mean value of normalized margins for our training dataset:\n",
    "\n",
    "\\begin{equation}\n",
    "\\begin{split}\n",
    "M_\\mathcal{A} &= \\frac{1}{n}\\sum_{i=1}^n \\hat{m}_i \\\\\n",
    "%&= \\frac{1}{\\sqrt{n} \\cdot R_\\mathcal{A} \\cdot \\|X\\|_2} \\cdot \\sum_{i=1}^n m_i\n",
    "\\end{split}\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Noise Sensitivity\n",
    "The idea is that if our model resistant to noise then it should generalize better.\n",
    "We check it in a very straightforward way: just add noise to the inputs and see how the outputs change.\n",
    "More precisly we compute Noise Sensitivity $N_\\mathcal{A}$ the following way:\n",
    "\n",
    "$$\n",
    "N_\\mathcal{A} = \\frac{1}{n}\\sum_{i=1}^n \\| f(x_i) - f(x_i + \\varepsilon_i) \\|_2^2,\n",
    "$$\n",
    "where:\n",
    "* $f(x_i)$ is a vector of logits\n",
    "* $\\{\\varepsilon_i\\}_{i=1}^n$ is a dataset of noise vectors sampled from $\\mathcal{N}(0, I)$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tips\n",
    "* You can delete all the provided code if it's more convenient for you to do everything from scratch\n",
    "* Check your implementation twice\n",
    "* It takes 5-10 hours for all the experiments to finish, so keep this in mind when planning your time\n",
    "* It feels like the most difficult part of this assignemnt is to get familiar with the experiments logic (and the code provided :|). It should be relatively easy to implement the measures afterwards\n",
    "* Reduce `NUM_ITERS_FOR_CONVERGENCE` value while debugging to make the experiments run faster\n",
    "* Running on a GPU will be 50-100% faster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys; sys.path.append('.')\n",
    "import random\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from utils.trainer import Trainer\n",
    "\n",
    "\n",
    "SEED = 42\n",
    "DATA_DIR = './data'\n",
    "DEVICE = 'cuda' # TIP: change to 'cuda' if you have a GPU available\n",
    "NUM_ITERS_FOR_CONVERGENCE = 20000\n",
    "\n",
    "random.seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "torch.cuda.manual_seed(SEED)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Complexity measures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import numpy as np\n",
    "from torch import Tensor\n",
    "\n",
    "\n",
    "class ComplexityMeasure:\n",
    "    def __init__(self, name:str, value:float):\n",
    "        self.name = name\n",
    "        self.value = value\n",
    "        \n",
    "    def __str__(self) -> str:\n",
    "        return f'{self.name}: {self.value}'\n",
    "    \n",
    "    def __repr__(self) -> str:\n",
    "        return f'{self.name}: {self.value}'\n",
    "\n",
    "\n",
    "def compute_accuracy(trainer:Trainer) -> ComplexityMeasure:\n",
    "    return ComplexityMeasure('Test accuracy', trainer.compute_test_accuracy())\n",
    "    \n",
    "    \n",
    "def compute_weights_norm(trainer:Trainer, p=2) -> ComplexityMeasure:\n",
    "    weights = torch.cat([p.data.cpu().view(-1) for p in trainer.model.parameters()])\n",
    "    \n",
    "    return ComplexityMeasure(f'L{p}-norm', torch.norm(weights, p).item())\n",
    "\n",
    "\n",
    "def compute_spectral_complexity(trainer:Trainer) -> ComplexityMeasure:\n",
    "    #########################################\n",
    "    ### Your code here. Difficulty: ☕☕ ###\n",
    "    ########################################\n",
    "    pass\n",
    "\n",
    "\n",
    "def compute_mean_normalized_margin(trainer:Trainer) -> ComplexityMeasure:\n",
    "    #########################################\n",
    "    ### Your code here. Difficulty: ☕☕ ###\n",
    "    ########################################\n",
    "    pass\n",
    "\n",
    "\n",
    "def compute_noise_sensitivity(trainer:Trainer, std:float=1.0, num_noised_points:int=1) -> ComplexityMeasure:\n",
    "    #######################################\n",
    "    ### Your code here. Difficulty: ☕ ###\n",
    "    ######################################\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [Part 1] Measuring complexity for a normal dataset for different widths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_config = {\n",
    "    'max_num_iters': NUM_ITERS_FOR_CONVERGENCE,\n",
    "    'model_type': 'DenseModel',\n",
    "    'batch_size': 500,\n",
    "    'device': DEVICE,\n",
    "    'data_dir': DATA_DIR,\n",
    "}\n",
    "widths = [10, 100, 250, 500, 1500, 5000]\n",
    "width_configs = [{'model_config': {'width': w}} for w in widths]\n",
    "configs = [{**base_config, **c} for c in width_configs]\n",
    "different_width_trainers = [Trainer(c) for c in configs]\n",
    "different_width_test_accs = [t.run_training(True).compute_test_accuracy() for t in different_width_trainers]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "complexity_measure_calculators = [\n",
    "    compute_accuracy,\n",
    "    compute_weights_norm,\n",
    "    compute_spectral_complexity,\n",
    "    compute_mean_normalized_margin,\n",
    "    compute_noise_sensitivity\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "different_width_complexities = [[fn(t) for t in different_width_trainers] for fn in complexity_measure_calculators]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize_complexities(x_values:List, complexities:List[List[ComplexityMeasure]], n_cols, n_rows, xlabel:str=''):\n",
    "    assert len(x_values) == len(complexities[0])\n",
    "    assert [len(set([c.name for c in cs])) == 1 for cs in complexities], \\\n",
    "        \"Each complexities row should correspond to the single complexity measure\"\n",
    "    \n",
    "    _, subplots = plt.subplots(n_rows, n_cols, figsize=(n_cols * 6, n_rows * 4))\n",
    "    if n_rows > 1: subplots = [p for row in subplots for p in row]\n",
    "    \n",
    "    for cs, subplot in zip(complexities, subplots):\n",
    "        subplot.set_title(cs[0].name)\n",
    "        subplot.plot(x_values, [c.value for c in cs])\n",
    "        subplot.set_xlabel(xlabel)\n",
    "        #subplot.set_ylabel(f'{cs[0].name}')\n",
    "        subplot.grid()\n",
    "        \n",
    "    # Hiding blank subplots for nicer visualization\n",
    "    for empty_subplot in subplots[len(complexities):]:\n",
    "        empty_subplot.set_axis_off()\n",
    "        \n",
    "    plt.subplots_adjust(hspace=0.3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize_complexities(widths, different_width_complexities, 3, 2, xlabel='Width')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Comment on plots you've obtained:**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Does accuracy reach asymptote for large width? Do complexity measures reach asymptote? Do these complexity measures increase or decrease with width? Does this seem natural?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [Part 2] Complexity measures for spoiled datasets with different amounts of spoiled samples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### [Part 2.1] Checking of complexity measures converge to some value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "max_num_iters = NUM_ITERS_FOR_CONVERGENCE\n",
    "freq_iters = 500\n",
    "num_good_points = 1000\n",
    "num_shuffled_labels_to_add = [0, 500, 2500, 5000, 50000]\n",
    "measurement_iters = [s * freq_iters for s in range(1, max_num_iters // freq_iters + 1)]\n",
    "base_config = {\n",
    "    'num_good_points': num_good_points,\n",
    "    'max_num_iters': max_num_iters,\n",
    "    'model_type': 'DenseModel',\n",
    "    'batch_size': 500,\n",
    "    'model_config': {'width': 1024},\n",
    "    'device': DEVICE,\n",
    "    'data_dir': DATA_DIR,\n",
    "}\n",
    "\n",
    "print('Constructing trainers...')\n",
    "configs = [{**base_config, **{'num_bad_points': n}} for n in num_shuffled_labels_to_add]\n",
    "trainers = [Trainer(c) for c in tqdm(configs)]\n",
    "complexity_measurements = {n:[] for n in num_shuffled_labels_to_add}\n",
    "\n",
    "def track_complexities(trainer:Trainer):\n",
    "    if trainer.num_iters_done % freq_iters != 0: return\n",
    "    cs = [fn(trainer) for fn in complexity_measure_calculators]\n",
    "    complexity_measurements[trainer.config['num_bad_points']].append(cs)\n",
    "\n",
    "for trainer in trainers:\n",
    "    trainer.on_iter_done_callbacks.append(track_complexities)\n",
    "    trainer.run_training(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_rows = 2\n",
    "n_cols = 3\n",
    "num_measures = len(complexity_measure_calculators)\n",
    "bad_points_proportions = [n / (n + num_good_points) for n in num_shuffled_labels_to_add]\n",
    "\n",
    "_, subplots = plt.subplots(n_rows, n_cols, figsize=(n_cols * 6, n_rows * 4))\n",
    "if n_rows > 1: subplots = [p for row in subplots for p in row]\n",
    "\n",
    "for c_idx, subplot in zip(range(num_measures), subplots):\n",
    "    complexities = [[compls[c_idx] for compls in complexity_measurements[n]] for n in num_shuffled_labels_to_add]\n",
    "\n",
    "    subplot.set_title(complexities[0][0].name)\n",
    "\n",
    "    for bad_points_prop, cs in zip(bad_points_proportions, complexities):\n",
    "        assert len(set(c.name for c in cs)) == 1, \"Wrong format for `complexities` argument\"\n",
    "        subplot.plot(measurement_iters, [c.value for c in cs], label=f'Bad labels proportion: {bad_points_prop:0.2f}')\n",
    "\n",
    "    subplot.set_xlabel('Iteration')\n",
    "    subplot.legend()\n",
    "    subplot.grid()\n",
    "    \n",
    "for subplot in subplots[num_measures:]:\n",
    "    subplot.set_axis_off()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Comment on plots you've obtained:**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Does accuracy reach asymptote with iterations? Which complexity measures (if any) reach asymptote?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### [Part 2.2] Final values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bad_points_proportions = [round(p, 2) for p in bad_points_proportions]\n",
    "final_measurements = [complexity_measurements[n][-1] for n in num_shuffled_labels_to_add]\n",
    "final_measurements = np.array(final_measurements).transpose().tolist()\n",
    "visualize_complexities(bad_points_proportions, final_measurements, 3, 2, xlabel='Bad points proportion')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Comment on plots you've obtained:**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Do accuracy / complexity measures increase or decrease with bad points proportion?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summarize your observations here:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Which complexity measures converge, and how are they correlated with test accuracy? Is this behavior seem desirable, or not?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}