updates and teaks flow matching and continuity across the whole chapter

2025-03-18 16:20:10 +01:00
1 parent 6391dbab10
commit 4589cf2860
@@ -8,7 +8,7 @@
   "source": [
    "# Denoising and Flow Matching Side-by-side\n",
    "\n",
-    "To show the capabilities of **denoising diffusion** and **flow matching**, we'll be use a learning task where we can reliably generate arbitrary amounts of ground truth data. This ensures we can quantify how well the target distribution was learned. Specifically, we'll focus on Reynolds-averaged Navier-Stokes simulations around airfoils, which have the interesting characteristic that typical solvers (such as OpenFoam) transition from steady solutions to oscillating ones for larger Reynolds numbers. This transition is exactly what we'll give as a task to diffusion models below. (Details can be found in our [diffuion-based flow prediction repository](https://github.com/tum-pbs/Diffusion-based-Flow-Prediction/).) Also, to make the notebook self-contained, we'll revisit the most important concepts from the previous section.\n",
+    "To show the capabilities of **denoising diffusion** and **flow matching**, we'll be use a learning task where we can reliably generate arbitrary amounts of ground truth data. This ensures we can quantify how well the target distribution was learned. Specifically, we'll focus on Reynolds-averaged Navier-Stokes simulations around airfoils, which have the interesting characteristic that typical solvers (such as OpenFoam) transition from steady solutions to oscillating ones for larger Reynolds numbers. This transition is exactly what we'll give as a task to diffusion models below. (Details can be found in our [diffusion-based flow prediction repository](https://github.com/tum-pbs/Diffusion-based-Flow-Prediction/).) Also, to make the notebook self-contained, we'll revisit the most important concepts from the previous section.\n",
     "[[run in colab]](https://colab.research.google.com/github/tum-pbs/pbdl-book/blob/main/probmodels-ddpm-fm.ipynb)\n",
    "\n",
    "```{note} \n",
@@ -1,12 +1,13 @@
 {
  "cells": [
    {
      "cell_type": "markdown",
      "id": "fe02764e59760330",
      "metadata": {
        "id": "fe02764e59760330"
      },
      "cell_type": "markdown",
      "source": [
-        "## Denoising\n",
+        "# Denoising\n",
        "\n",
        "From our perspective of annealed Langevin dynamics, let's consider the extreme case: let's aim for starting with _pure_ noise, and treat the annealing steps as a continuous time dimension. We define our pure noise starting time as $t=T$, and then reduce the noise to zero with a continuous _noising schedule_ at $t=0$. When we focus on Gaussian noise $\\mathcal{N}$, combining multiple instances of Gaussian noise is still Gaussian. Hence, this actually gives us a Markov chain with a **forward** chain that adds noise in a very controllable way. We'll call the forward process $q$ and denote the function for the standard deviation with $\\beta$, giving\n",
        "\n",
@@ -36,6 +37,8 @@
        "with the inverted weights $\\alpha_t = 1 - \\beta_t$ and alphas accumulated for time $t$ denoted by\n",
        "$\\overline{\\alpha}_t= \\prod_{s=1}^t \\alpha_s$.\n",
        "\n",
        "## Latent Variable Models\n",
        "\n",
        "Conceptually, this formulation gives us what's called a _latent variable model_ in the ML community. Instead of the somewhat arbitrary in between states of the Annealed Langevin Dynamics above, we now have explicitly modeled _latent_ states along the diffusion time $t$. Our distribution for targets $ x_0 \\sim q(x_0) $ is of the form  $ p_\\theta(x_0) = \\int p_\\theta(x_{0:T})dx_{1:T} , $   where  $x_1,...,x_T$ are latents with the same dimensionality as $x_0$.\n",
        "\n",
        "Above we've directly trained models with the gradient of the perturbed data to get the score. How does this work in the context of these continually changing latents?\n",
@@ -103,6 +106,8 @@
        "$x_t$ by mixing $\\sqrt{\\overline{a}_t}x_0$ with the right amount of noise\n",
        " $ \\sqrt{1-\\overline{\\alpha}_t} \\epsilon$, as outlined above in equation {eq}`ddpm-xt`.\n",
        "\n",
        "## Full Denoising Algorithm\n",
        "\n",
        "This extremely simple $L^2$ loss is really the reason for the break-through success of this approach, as it's very simple to compute, and very stable.\n",
        "This training variant often goes under the name _denoising diffusion probablistic models_ (DDPM). Unlike, e.g. GANs, for which one needs to train with a fragile balance of two networks, we only need to consider a fully supervised loss for training DDPMs. This comes at the cost of increased computational resources, though. Before addressing the computational challenges (these have also been mostly resolved by now), let's re-cap that algorithms for training and inference that we get from this denoising objective.\n",
        "\n",
@@ -167,20 +172,24 @@
        "- And in the context of physics simulations, we'd of course like to bring back our prior knowledge in the form of PDE-base constraints.\n",
        "\n",
        "Nonetheless, we've made the most important first step with DDPMs now: we have turned the deterministic networks into tools for probabilistic inference that robustly learn and reproduce complex distributions. Before starting with source code examples, let's address the computational overhead: DDPMs require a _large_ number of denoising steps. The resulting models can be orders of magnitude more expensive to run than their deterministic counterparts. Hence the central question is, how can we  reduce the number of required steps without loosing accuracy for the posterior?"
-      ],
+      ]
      "id": "fe02764e59760330"
    },
    {
      "cell_type": "markdown",
      "id": "2c66c145697a1099",
      "metadata": {
        "id": "2c66c145697a1099"
      },
      "cell_type": "markdown",
      "source": [
-        "## Train a DDPM on the Gaussian Toy Dataset"
+        "## Training with DDPM\n",
-      ],
+        "\n",
-      "id": "2c66c145697a1099"
+        "Let's see how well these ideas work for our Gaussian mixture (GM) case. To make this a standalone notebook, we first initialize the same two-peak GM case that was also used for the normalizing flows and score tests."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "eec0a45e5bc00691",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-03-17T00:09:15.870346Z",
@@ -188,13 +197,12 @@
        },
        "id": "eec0a45e5bc00691"
      },
-      "cell_type": "code",
+      "outputs": [],
      "source": [
        "import numpy as np\n",
        "\n",
        "class GaussianMixture:\n",
        "    def __init__(self, parameters):\n",
        "\n",
        "        self.parameters = parameters\n",
        "        self.distributions = [\n",
        "            {\n",
@@ -219,24 +227,25 @@
        "    {\"mean\": [0, 0], \"std\": [1, 1]},\n",
        "    {\"mean\": [3, 2], \"std\": [0.5, 0.5]}\n",
        "]\n",
        "\n",
        "mixture = GaussianMixture(parameters)"
-      ],
+      ]
      "id": "eec0a45e5bc00691",
      "outputs": [],
      "execution_count": 3
    },
    {
      "cell_type": "markdown",
      "id": "c61074cf783a4fa9",
      "metadata": {
        "id": "c61074cf783a4fa9"
      },
      "cell_type": "markdown",
      "source": [
-        "### Define the Forward Process"
+        "### Define the Forward Process\n",
-      ],
+        "\n",
-      "id": "c61074cf783a4fa9"
+        "Next, we implement the forward process for DDPM, which takes our GM distribution, and adds increasing amounts of noise until all traces of information left in the samples are removed. The visualizations below show this destructive process in action..."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "id": "59216c5f42a1a070",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-03-17T00:09:17.244046Z",
@@ -249,7 +258,18 @@
        "id": "59216c5f42a1a070",
        "outputId": "d891b00e-cd74-44a0-e202-4f29a3030603"
      },
-      "cell_type": "code",
+      "outputs": [
        {
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAABa8AAAPaCAYAAACQwqTYAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjAsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvlHJYcgAAAAlwSFlzAAAPYQAAD2EBqD+naQABAABJREFUeJzs3Xd4FNXbxvHvbiqpkJAGoUNC71VAehcpKiCKggULihWVV/Rn7w3FgkhHBBSRpvTeEjqhQyihpPfedt8/QlZCEkgwkAD357q4NDNnZ87MbHLmPHPmOQaz2WxGRERERERERERERKQMMZZ2BURERERERERERERErqTgtYiIiIiIiIiIiIiUOQpei4iIiIiIiIiIiEiZo+C1iIiIiIiIiIiIiJQ5Cl6LiIiIiIiIiIiISJmj4LWIiIiIiIiIiIiIlDkKXouIiIiIiIiIiIhImaPgtYiIiIiIiIiIiIiUOQpei4iIiIiIiIiIiEiZo+C1iIiIiAAwYsQI/P39S7sad7TY2Fhat27NO++8U+TP/Pnnn/j7+/Pnn3/euIpJoW72+T948CD+/v78/vvvN2V/IiIiIqVJwWsRERG5rfn7+1/1nwJ+UpZ89913pKWl8cwzz5TKvv39/QkICLjh++ratStdu3a94fu5HTVs2JDu3bszceJEkpOTS7s6IiIiIjeUdWlXQERERORmeO655wpcXq9evZtcE5GCXbx4kfnz5zN48GC8vLxKuzpSRD169KBJkyZ4enretH0+9dRTPPDAA8yePZunn376pu1XRERE5GZT8FpERETuCM8//3xpV0HkqubPn09WVhaDBg0q7apIMTg7O+Ps7HxT99m4cWNq1qzJ/PnzGT16NEajXqgVERGR25OC1yIiIiKXZGRkMGPGDJYuXUpISAhWVlbUrVuXhx9+mL59++Ype/78ebp168agQYN46qmnmDhxIgEBAcTGxjJz5kzmz5/P8uXLWblyJdWrV7d87vXXX+evv/6ibdu2zJw507I8KSmJNm3a0LRpU3799VcAEhMTmT9/Pps2beLMmTPExMTg5ORE06ZNeeqpp2jWrFm+Y/D396d169Z89dVXfPPNN2zatImoqCg+/PBDBg8eDMDy5cuZOnUqJ0+exNHRkQ4dOvDqq68W+3zlpn1YvHgxX3/9NatXryYuLo4qVaowbNgwRowYgcFgKPI5a9OmDSaTifnz5/PHH39w6tQpzGYztWrV4r777mPYsGEFBumCg4P55ZdfCAgIICIiAmdnZ2rUqME999zD8OHD85WdMmUK27dvJzo6GhcXF9q1a8eYMWOoWbNmnrJRUVFMnTqVdevWER4ejrW1Ne7u7jRr1owxY8ZQpUoVAMxmM3/99Rfz58/nzJkzJCcn4+bmRu3atbnvvvvyfXcKYjabWbhwIT4+PjRv3rzAMmfPnuXLL79k+/btZGZmUrdu3auOut2xYwfLly9n9+7dhIWFkZWVRdWqVenduzdPPvkkdnZ2lrJdu3blwoULADzyyCN5tnPs2DEATp8+zcKFC9m2bRsXL14kKSkJDw8POnTowJgxY/D29r7mcQYEBOTZ/uU51gcNGsQnn3xi+Xn79u388ssvBAUFkZKSQuXKlenRowdPPfVUkYPFf/75J+PHj+fjjz+mUqVKfP/99xw8eBCDwUDLli15/fXXqVWrVr7PRURE8OOPP7Jx40YiIiJwcnKiZcuWPP300zRs2LDQfeT+jgEcPXqUn3/+mX379lm24ePjQ8uWLXnttdewsbGxlM3KymL+/PksXryYkydPkp2dTY0aNbj//vsZPnx4gd/7fv368d1337F161Y6duxYpPMhIiIicqtR8FpERESEnMD1448/TmBgIDVr1mT48OGkpaWxcuVKXnrpJY4ePcrLL7+c73MhISEMGTKE6tWr079/f9LS0nBycqJdu3YsX76c7du35wleb9++HYC9e/eSnp5uCSDu3LmTrKws2rVrZykbHBzMN998Q8uWLencuTMuLi6Ehoaybt06Nm/ezI8//sjdd9+dr05xcXEMHToUBwcHevbsicFgwN3dHYAZM2bw8ccf4+LiwsCBA3F2dmbLli08+OCDODk5Xdd5GzlyJImJifTr14/MzExWrlzJhx9+yOnTp/nf//5X5HMGMG7cOJYtW4aPjw/3338/BoOBNWvW8O6777J7926+/PLLPNvasGEDL7zwAhkZGXTs2JF+/fqRkJDAsWPH+OWXX/IErzdt2sTzzz9PVlYWXbp0oWrVqoSHh7Nq1So2bNjArFmzaNCgAQCpqak8+OCDhISE0L59e7p27YrZbObixYusXbuWXr16WYLXX3/9NZMnT8bX15c+ffrg7OxMZGQkQUFBrFixokjB6xMnThAZGUm/fv0KXH/mzBmGDh1KXFwcd999N/Xq1ePs2bOMGTOm0MDllClTOH36NM2aNaNTp05kZGSwZ88evvvuOwICApgxYwZWVlZATsB67dq1BAYGMmjQICpXrpxve6tXr2bevHm0adOG5s2bY2Njw4kTJ/j9999Zv349CxcuvGa6k8qVK/Pcc89ZHtw8+uijlnWXp/CZN28e77zzDuXKlaN37964u7sTGBjIlClTWL9+Pb/99hsuLi5XP6mX2bBhA2vXrqVjx44MGzaM4OBgNm7cSFBQEMuXL8fNzc1S9ty5cwwfPpyIiAjatm1Lv379CA0NZcWKFWzYsIHvvvuOLl26XHV/R48eZciQIRgMBrp27Yqvry9JSUmEhITw22+/8eKLL1qC15mZmTz99NNs2bLF8tDFzs6OgIAA3n//ffbv38/nn3+ebx+5Dzm2bdum4LWIiIjcvswiIiIitzE/Pz+zn5+f+dtvv833b+HChZZyP/30k9nPz8/8xBNPmDMzMy3Lo6KizF26dDH7+fmZd+/ebVl+7tw5y7a//PLLfPsNCQkx+/n5mZ9//nnLsuDgYLOfn5951KhRZj8/P/O2bdss6z788EOzn5+feefOnZZlCQkJ5ujo6HzbDg0NNbdv397cu3fvQo933LhxeY4jt84NGjQwt2rVynzu3DnL8uzsbPNzzz1n+WxR5Z6XYcOGmdPT0y3LY2Njzd26dTP7+fmZAwMD8+z/auds6dKlZj8/P/PAgQPNSUlJluXJycnmQYMGmf38/MxLliyxLI+OjjY3b97c3KBBA3NAQEC+7YWGhlr+Py4uztyyZUtz69atzSdOnMhT7tixY+amTZuaBw4caFm2du1as5+fn/nDDz/Mt9309HRzYmKi5efWrVubO3bsaE5JSclXtqDrV5C5c+ea/fz8zFOnTi1wfe53ZsaMGXmWr1692nJOL/8+m80530GTyZRvW19//bXZz8/PvHz58jzLv/32W7Ofn595x44dBdYhLCwsz3XOtXnzZnPdunXNb7/99lWP8XJdunQxd+nSpcB158+fNzdo0MDcrFkz88mTJ/Os+9///mf28/MzT5gwoUj7WbhwodnPz89cr169PL9vZrPZ/MUXX5j9/PzMP//8c57ljz32mNnPz8/8ww8/5Fm+e/duc7169cytW7fO8/3M3cfl5//jjz82+/n5mVevXp2vTnFxcebs7GzLz7nn/b333jNnZWVZlmdlZZnHjx9f6HYSEhLMfn5+5vvuu69I50JERETkVqTkaCIiInJHmDRpUr5/ixYtsqxfuHAhBoOBN954A2vrf19Oc3d355lnngHg999/z7fdihUrFjgZZJUqVahcuTIBAQGYzWbg31HXY8eOxcrKyvJz7joHBweaNGliWebs7JxnRGgub29vevfuzalTp7h48WK+9TY2Nrz++ut5jgNg6dKlZGZm8vDDD+Pr62tZbjQaee211647b+4rr7yCra2t5efy5cvz7LPPAjkpFa5U2DlbuHChZXuOjo6W5Q4ODowbNw7Iew3++usvkpKSGDZsGK1bt863vcvTWPz1118kJCQwduxYateunaecn58fDzzwAIcPH+bkyZN51tnb2+fbrq2tbb5R6tbW1pZRzJcr6PoVJDQ0FAAPD49868LCwti6dSu+vr48/PDDedZ17969wGOHnO/g5Wlbco0cORKAzZs3F6luuby8vPJc51wdOnSgdu3abNmypVjbK8ySJUss39MrU3q89NJLODo6snjxYjIyMoq8zb59++Z5qwFgyJAhAAQFBVmWhYWFsWXLFipVqsQTTzyRp3zz5s3p168fcXFxrF69ukj7Lej74+rqavldM5lMzJkzBw8PD8aPH5/nO2RlZcUbb7yBwWBg6dKl+bbj7OyMnZ2d5bsjIiIicjtS2hARERG5I+Tm7S1IUlISZ8+excvLq8D8t23btgXgyJEj+dbVrVu3wIBe7ucWLlzIkSNHqF+/PgEBAXh4eNC0aVMaNGhgCV7HxMRw4sQJ2rdvnycPLsDu3buZNWsW+/btIzo6mszMzDzrw8PDqVSpUp5llStXtqQJudzhw4cBaNWqVb51VapUwcfHx5L3uKisra0LzL2dG1DN3eflCjtnhw8fxmg0FhiMbdWqFVZWVnmuwb59+wAKTJ1ypdyyR48e5bvvvsu3/syZM0BOqpbatWvTunVrvLy8+Pnnnzl06BCdOnWiefPm1KtXL1+Qun///syePZu+ffvSp08fWrVqRbNmzYo1iV9cXBxAgakwcs9hixYtCgyQt27dmsDAwHzLU1JSmDVrFqtXr7bk4s59kAI5eZ2Lw2w2s2TJEhYtWsTRo0dJSEggOzvbsv7K7+71yj3e3N+7y7m6ulK/fn127tzJqVOnqFu3bpG2eWWeagAfHx8A4uPj8+27RYsWBR5P27ZtWbJkCYcPH2bgwIGF7q9v377MmjWLMWPG0KtXL+666y6aN29O1apV85Q7ffo0cXFxVK9enR9//LHAbdnb23Pq1KkC17m6uhIdHV1oPURERERudQpei4iIyB0vKSkJKHjUK4CnpycACQkJ+dZVrFix0O22a9eOhQsXsLine truncated
            "text/plain": [
              "<Figure size 1500x1000 with 6 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "import torch\n",
        "import matplotlib.pyplot as plt\n",
@@ -315,33 +335,26 @@
        "\n",
        "ddpm = DDPMForwardProcess(num_timesteps=1000)\n",
        "ddpm.visualize_forward_process(mixture)\n"
-      ],
+      ]
      "id": "59216c5f42a1a070",
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 1500x1000 with 6 Axes>"
            ],
            "image/png": "iVBORw0KGgoAAAANSUhEUgAABa8AAAPaCAYAAACQwqTYAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjAsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvlHJYcgAAAAlwSFlzAAAPYQAAD2EBqD+naQABAABJREFUeJzs3Xd4FNXbxvHvbiqpkJAGoUNC71VAehcpKiCKggULihWVV/Rn7w3FgkhHBBSRpvTeEjqhQyihpPfedt8/QlZCEkgwkAD357q4NDNnZ87MbHLmPHPmOQaz2WxGRERERERERERERKQMMZZ2BURERERERERERERErqTgtYiIiIiIiIiIiIiUOQpei4iIiIiIiIiIiEiZo+C1iIiIiIiIiIiIiJQ5Cl6LiIiIiIiIiIiISJmj4LWIiIiIiIiIiIiIlDkKXouIiIiIiIiIiIhImaPgtYiIiIiIiIiIiIiUOQpei4iIiIiIiIiIiEiZo+C1iIiIiAAwYsQI/P39S7sad7TY2Fhat27NO++8U+TP/Pnnn/j7+/Pnn3/euIpJoW72+T948CD+/v78/vvvN2V/IiIiIqVJwWsRERG5rfn7+1/1nwJ+UpZ89913pKWl8cwzz5TKvv39/QkICLjh++ratStdu3a94fu5HTVs2JDu3bszceJEkpOTS7s6IiIiIjeUdWlXQERERORmeO655wpcXq9evZtcE5GCXbx4kfnz5zN48GC8vLxKuzpSRD169KBJkyZ4enretH0+9dRTPPDAA8yePZunn376pu1XRERE5GZT8FpERETuCM8//3xpV0HkqubPn09WVhaDBg0q7apIMTg7O+Ps7HxT99m4cWNq1qzJ/PnzGT16NEajXqgVERGR25OC1yIiIiKXZGRkMGPGDJYuXUpISAhWVlbUrVuXhx9+mL59++Ype/78ebp168agQYN46qmnmDhxIgEBAcTGxjJz5kzmz5/P8uXLWblyJdWrV7d87vXXX+evv/6ibdu2zJw507I8KSmJNm3a0LRpU3799VcAEhMTmT9/Pps2beLMmTPExMTg5ORE06ZNeeqpp2jWrFm+Y/D396d169Z89dVXfPPNN2zatImoqCg+/PBDBg8eDMDy5cuZOnUqJ0+exNHRkQ4dOvDqq68W+3zlpn1YvHgxX3/9NatXryYuLo4qVaowbNgwRowYgcFgKPI5a9OmDSaTifnz5/PHH39w6tQpzGYztWrV4r777mPYsGEFBumCg4P55ZdfCAgIICIiAmdnZ2rUqME999zD8OHD85WdMmUK27dvJzo6GhcXF9q1a8eYMWOoWbNmnrJRUVFMnTqVdevWER4ejrW1Ne7u7jRr1owxY8ZQpUoVAMxmM3/99Rfz58/nzJkzJCcn4+bmRu3atbnvvvvyfXcKYjabWbhwIT4+PjRv3rzAMmfPnuXLL79k+/btZGZmUrdu3auOut2xYwfLly9n9+7dhIWFkZWVRdWqVenduzdPPvkkdnZ2lrJdu3blwoULADzyyCN5tnPs2DEATp8+zcKFC9m2bRsXL14kKSkJDw8POnTowJgxY/D29r7mcQYEBOTZ/uU51gcNGsQnn3xi+Xn79u388ssvBAUFkZKSQuXKlenRowdPPfVUkYPFf/75J+PHj+fjjz+mUqVKfP/99xw8eBCDwUDLli15/fXXqVWrVr7PRURE8OOPP7Jx40YiIiJwcnKiZcuWPP300zRs2LDQfeT+jgEcPXqUn3/+mX379lm24ePjQ8uWLXnttdewsbGxlM3KymL+/PksXryYkydPkp2dTY0aNbj//vsZPnx4gd/7fv368d1337F161Y6duxYpPMhIiIicqtR8FpERESEnMD1448/TmBgIDVr1mT48OGkpaWxcuVKXnrpJY4ePcrLL7+c73MhISEMGTKE6tWr079/f9LS0nBycqJdu3YsX76c7du35wleb9++HYC9e/eSnp5uCSDu3LmTrKws2rVrZykbHBzMN998Q8uWLencuTMuLi6Ehoaybt06Nm/ezI8//sjdd9+dr05xcXEMHToUBwcHevbsicFgwN3dHYAZM2bw8ccf4+LiwsCBA3F2dmbLli08+OCDODk5Xdd5GzlyJImJifTr14/MzExWrlzJhx9+yOnTp/nf//5X5HMGMG7cOJYtW4aPjw/3338/BoOBNWvW8O6777J7926+/PLLPNvasGEDL7zwAhkZGXTs2JF+/fqRkJDAsWPH+OWXX/IErzdt2sTzzz9PVlYWXbp0oWrVqoSHh7Nq1So2bNjArFmzaNCgAQCpqak8+OCDhISE0L59e7p27YrZbObixYusXbuWXr16WYLXX3/9NZMnT8bX15c+ffrg7OxMZGQkQUFBrFixokjB6xMnThAZGUm/fv0KXH/mzBmGDh1KXFwcd999N/Xq1ePs2bOMGTOm0MDllClTOH36NM2aNaNTp05kZGSwZ88evvvuOwICApgxYwZWVlZATsB67dq1BAYGMmjQICpXrpxve6tXr2bevHm0adOG5s2bY2Njw4kTJ/j9999Zv349CxcuvGa6k8qVK/Pcc89ZHtw8+uijlnWXp/CZN28e77zzDuXKlaN37964u7sTGBjIlClTWL9+Pb/99hsuLi5XP6mX2bBhA2vXrqVjx44MGzaM4OBgNm7cSFBQEMuXL8fNzc1S9ty5cwwfPpyIiAjatm1Lv379CA0NZcWKFWzYsIHvvvuOLl26XHV/R48eZciQIRgMBrp27Yqvry9JSUmEhITw22+/8eKLL1qC15mZmTz99NNs2bLF8tDFzs6OgIAA3n//ffbv38/nn3+ebx+5Dzm2bdum4LWIiIjcvswiIiIitzE/Pz+zn5+f+dtvv833b+HChZZyP/30k9nPz8/8xBNPmDMzMy3Lo6KizF26dDH7+fmZd+/ebVl+7tw5y7a//PLLfPsNCQkx+/n5mZ9//nnLsuDgYLOfn5951KhRZj8/P/O2bdss6z788EOzn5+feefOnZZlCQkJ5ujo6HzbDg0NNbdv397cu3fvQo933LhxeY4jt84NGjQwt2rVynzu3DnL8uzsbPNzzz1n+WxR5Z6XYcOGmdPT0y3LY2Njzd26dTP7+fmZAwMD8+z/auds6dKlZj8/P/PAgQPNSUlJluXJycnmQYMGmf38/MxLliyxLI+OjjY3b97c3KBBA3NAQEC+7YWGhlr+Py4uztyyZUtz69atzSdOnMhT7tixY+amTZuaBw4caFm2du1as5+fn/nDDz/Mt9309HRzYmKi5efWrVubO3bsaE5JSclXtqDrV5C5c+ea/fz8zFOnTi1wfe53ZsaMGXmWr1692nJOL/8+m80530GTyZRvW19//bXZz8/PvHz58jzLv/32W7Ofn595x44dBdYhLCwsz3XOtXnzZnPdunXNb7/99lWP8XJdunQxd+nSpcB158+fNzdo0MDcrFkz88mTJ/Os+9///mf28/MzT5gwoUj7WbhwodnPz89cr169PL9vZrPZ/MUXX5j9/PzMP//8c57ljz32mNnPz8/8ww8/5Fm+e/duc7169cytW7fO8/3M3cfl5//jjz82+/n5mVevXp2vTnFxcebs7GzLz7nn/b333jNnZWVZlmdlZZnHjx9f6HYSEhLMfn5+5vvuu69I50JERETkVqTkaCIiInJHmDRpUr5/ixYtsqxfuHAhBoOBN954A2vrf19Oc3d355lnngHg999/z7fdihUrFjgZZJUqVahcuTIBAQGYzWbg31HXY8eOxcrKyvJz7joHBweaNGliWebs7JxnRGgub29vevfuzalTp7h48WK+9TY2Nrz++ut5jgNg6dKlZGZm8vDDD+Pr62tZbjQaee211647b+4rr7yCra2t5efy5cvz7LPPAjkpFa5U2DlbuHChZXuOjo6W5Q4ODowbNw7Iew3++usvkpKSGDZsGK1bt863vcvTWPz1118kJCQwduxYateunaecn58fDzzwAIcPH+bkyZN51tnb2+fbrq2tbb5R6tbW1pZRzJcr6PoVJDQ0FAAPD49868LCwti6dSu+vr48/PDDedZ17969wGOHnO/g5Wlbco0cORKAzZs3F6luuby8vPJc51wdOnSgdu3abNmypVjbK8ySJUss39MrU3q89NJLODo6snjxYjIyMoq8zb59++Z5qwFgyJAhAAQFBVmWhYWFsWXLFipVqsQTTzyRp3zz5s3p168fcXFxrF69ukj7Lej74+rqavldM5lMzJkzBw8PD8aPH5/nO2RlZcUbb7yBwWBg6dKl+bbj7OyMnZ2d5bsjIiIicjtS2hARERG5I+Tm7S1IUlISZ8+excvLq8D8t23btgXgyJEj+dbVrVu3wIBe7ucWLlzIkSNHqF+/PgEBAXh4eNC0aVMaNGhgCV7HxMRw4sQJ2rdvnycPLsDu3buZNWsW+/btIzo6mszMzDzrw8PDqVSpUp5llStXtqQJudzhw4cBaNWqVb51VapUwcfHx5L3uKisra0LzL2dG1DN3eflCjtnhw8fxmg0FhiMbdWqFVZWVnmuwb59+wAKTJ1ypdyyR48e5bvvvsu3/syZM0BOqpbatWvTunVrvLy8+Pnnnzl06BCdOnWiefPm1KtXL1+Qun///syePZu+ffvSp08fWrVqRbNmzYo1iV9cXBxAgakwcs9hixYtCgyQt27dmsDAwHzLU1JSmDVrFqtXr7bk4s59kAI5eZ2Lw2w2s2TJEhYtWsTRo0dJSEggOzvbsv7K7+71yj3e3N+7y7m6ulK/fn127tzJqVOnqFu3bpG2eWWeagAfHx8A4uPj8+27RYsWBR5P27ZtWbJkCYcPH2bgwIGF7q9v377MmjWLMWPG0KtXL+666y6aN29O1apV85Q7ffo0cXFxVK9enR9//LHAbdnb23Pq1KkC17m6uhIdHV1oPURERERudQpei4iIyB0vKSkJKHjUK4CnpycACQkJ+dZVrFix0O22a9eOhQsXsLine truncated
          },
          "metadata": {}
        }
      ],
      "execution_count": 6
    },
    {
      "cell_type": "markdown",
      "id": "b08f5dfe0fe31649",
      "metadata": {
        "id": "b08f5dfe0fe31649"
      },
      "cell_type": "markdown",
      "source": [
-        "### Setup the Dataset and Dataloader"
+        "It's visible here that the two peaks from $t=0$ start to disappear after ca. 200 steps of adding noise, and, visually, form a single Guassian peak centered around zero.\n",
-      ],
+        "\n",
-      "id": "b08f5dfe0fe31649"
+        "### Dataset and Dataloader\n",
        "\n",
        "The next code cell defines a small dataset class that we will use to train the DDPM. It precomputes the correct alpha and beta values for each $t$ so that we can quickly pull up samples with varying amounts of noise for training."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "id": "1e531e6c995b4609",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-03-17T00:02:22.356614Z",
@@ -349,7 +362,7 @@
        },
        "id": "1e531e6c995b4609"
      },
-      "cell_type": "code",
+      "outputs": [],
      "source": [
        "from torch import nn\n",
        "from torch.optim import Adam\n",
@@ -385,30 +398,28 @@
        "        x_t = sqrt_alpha_bar * x0_sample + sqrt_one_minus_alpha_bar * eps\n",
        "\n",
        "        return x_t, torch.tensor(t, dtype=torch.float32), eps"
-      ],
+      ]
      "id": "1e531e6c995b4609",
      "outputs": [],
      "execution_count": 7
    },
    {
      "cell_type": "markdown",
      "id": "5ba5da2e4725ffb1",
      "metadata": {
        "id": "5ba5da2e4725ffb1"
      },
      "cell_type": "markdown",
      "source": [
-        "### Define the Score Network"
+        "Next, we'll define a simple network, much in line with the previous score networks. In this case, three hidden layers suffice to learn the denoising task. The main difference now is that the network receives the denoising  time `t` instead of the perturbation amount from before."
-      ],
+      ]
      "id": "5ba5da2e4725ffb1"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5e711111525c923d",
      "metadata": {
        "id": "5e711111525c923d"
      },
      "cell_type": "code",
      "outputs": [],
      "execution_count": 8,
      "source": [
-        "class ScoreNet(nn.Module):\n",
+        "class DdpmNet(nn.Module):\n",
        "    def __init__(self, in_dim=2, time_dim=1, hidden_dim=128, out_dim=2):\n",
        "        super().__init__()\n",
        "        self.net = nn.Sequential(\n",
@@ -423,20 +434,22 @@
        "        t = t.unsqueeze(1)\n",
        "        xt = torch.cat([x, t], dim=1)\n",
        "        return self.net(xt)"
-      ],
+      ]
      "id": "5e711111525c923d"
    },
    {
      "cell_type": "markdown",
      "id": "ad8de5772ec3b37b",
      "metadata": {
        "id": "ad8de5772ec3b37b"
      },
      "cell_type": "markdown",
      "source": [
        "### Initialization and Training"
-      ],
+      ]
      "id": "ad8de5772ec3b37b"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "1027e7e50cc6fe83",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
@@ -444,11 +457,10 @@
        "id": "1027e7e50cc6fe83",
        "outputId": "214a0863-f2b0-4d4e-c801-07b0ba654ce9"
      },
      "cell_type": "code",
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 1/100, Loss: 0.1529\n",
            "Epoch 2/100, Loss: 0.0871\n",
@@ -460,85 +472,7 @@
            "Epoch 8/100, Loss: 0.0742\n",
            "Epoch 9/100, Loss: 0.0706\n",
            "Epoch 10/100, Loss: 0.0704\n",
-            "Epoch 11/100, Loss: 0.0688\n",
+            "...\n",
            "Epoch 12/100, Loss: 0.0659\n",
            "Epoch 13/100, Loss: 0.0706\n",
            "Epoch 14/100, Loss: 0.0706\n",
            "Epoch 15/100, Loss: 0.0609\n",
            "Epoch 16/100, Loss: 0.0671\n",
            "Epoch 17/100, Loss: 0.0671\n",
            "Epoch 18/100, Loss: 0.0687\n",
            "Epoch 19/100, Loss: 0.0696\n",
            "Epoch 20/100, Loss: 0.0624\n",
            "Epoch 21/100, Loss: 0.0613\n",
            "Epoch 22/100, Loss: 0.0628\n",
            "Epoch 23/100, Loss: 0.0636\n",
            "Epoch 24/100, Loss: 0.0665\n",
            "Epoch 25/100, Loss: 0.0629\n",
            "Epoch 26/100, Loss: 0.0640\n",
            "Epoch 27/100, Loss: 0.0699\n",
            "Epoch 28/100, Loss: 0.0653\n",
            "Epoch 29/100, Loss: 0.0630\n",
            "Epoch 30/100, Loss: 0.0614\n",
            "Epoch 31/100, Loss: 0.0611\n",
            "Epoch 32/100, Loss: 0.0623\n",
            "Epoch 33/100, Loss: 0.0616\n",
            "Epoch 34/100, Loss: 0.0611\n",
            "Epoch 35/100, Loss: 0.0602\n",
            "Epoch 36/100, Loss: 0.0593\n",
            "Epoch 37/100, Loss: 0.0653\n",
            "Epoch 38/100, Loss: 0.0636\n",
            "Epoch 39/100, Loss: 0.0581\n",
            "Epoch 40/100, Loss: 0.0612\n",
            "Epoch 41/100, Loss: 0.0595\n",
            "Epoch 42/100, Loss: 0.0632\n",
            "Epoch 43/100, Loss: 0.0672\n",
            "Epoch 44/100, Loss: 0.0647\n",
            "Epoch 45/100, Loss: 0.0627\n",
            "Epoch 46/100, Loss: 0.0604\n",
            "Epoch 47/100, Loss: 0.0617\n",
            "Epoch 48/100, Loss: 0.0620\n",
            "Epoch 49/100, Loss: 0.0573\n",
            "Epoch 50/100, Loss: 0.0656\n",
            "Epoch 51/100, Loss: 0.0599\n",
            "Epoch 52/100, Loss: 0.0583\n",
            "Epoch 53/100, Loss: 0.0606\n",
            "Epoch 54/100, Loss: 0.0634\n",
            "Epoch 55/100, Loss: 0.0645\n",
            "Epoch 56/100, Loss: 0.0580\n",
            "Epoch 57/100, Loss: 0.0586\n",
            "Epoch 58/100, Loss: 0.0616\n",
            "Epoch 59/100, Loss: 0.0621\n",
            "Epoch 60/100, Loss: 0.0544\n",
            "Epoch 61/100, Loss: 0.0605\n",
            "Epoch 62/100, Loss: 0.0632\n",
            "Epoch 63/100, Loss: 0.0618\n",
            "Epoch 64/100, Loss: 0.0663\n",
            "Epoch 65/100, Loss: 0.0625\n",
            "Epoch 66/100, Loss: 0.0635\n",
            "Epoch 67/100, Loss: 0.0602\n",
            "Epoch 68/100, Loss: 0.0616\n",
            "Epoch 69/100, Loss: 0.0617\n",
            "Epoch 70/100, Loss: 0.0585\n",
            "Epoch 71/100, Loss: 0.0626\n",
            "Epoch 72/100, Loss: 0.0612\n",
            "Epoch 73/100, Loss: 0.0608\n",
            "Epoch 74/100, Loss: 0.0626\n",
            "Epoch 75/100, Loss: 0.0630\n",
            "Epoch 76/100, Loss: 0.0652\n",
            "Epoch 77/100, Loss: 0.0599\n",
            "Epoch 78/100, Loss: 0.0605\n",
            "Epoch 79/100, Loss: 0.0549\n",
            "Epoch 80/100, Loss: 0.0562\n",
            "Epoch 81/100, Loss: 0.0619\n",
            "Epoch 82/100, Loss: 0.0610\n",
            "Epoch 83/100, Loss: 0.0635\n",
            "Epoch 84/100, Loss: 0.0593\n",
            "Epoch 85/100, Loss: 0.0630\n",
            "Epoch 86/100, Loss: 0.0623\n",
            "Epoch 87/100, Loss: 0.0571\n",
            "Epoch 88/100, Loss: 0.0595\n",
            "Epoch 89/100, Loss: 0.0645\n",
            "Epoch 90/100, Loss: 0.0631\n",
            "Epoch 91/100, Loss: 0.0592\n",
            "Epoch 92/100, Loss: 0.0597\n",
@@ -553,7 +487,6 @@
          ]
        }
      ],
      "execution_count": 10,
      "source": [
        "num_samples = 10000\n",
        "input_dim = 2\n",
@@ -565,7 +498,7 @@
        "# %% Training the DDPM model\n",
        "time_dim = 1\n",
        "hidden_dim = 128\n",
-        "model = ScoreNet(in_dim=input_dim, time_dim=time_dim, hidden_dim=hidden_dim, out_dim=input_dim)\n",
+        "model = DdpmNet(in_dim=input_dim, time_dim=time_dim, hidden_dim=hidden_dim, out_dim=input_dim)\n",
        "optimizer = Adam(model.parameters(), lr=1e-3)\n",
        "loss_fn = nn.MSELoss()\n",
        "\n",
@@ -583,25 +516,27 @@
        "        optimizer.step()\n",
        "        epoch_loss += loss.item()\n",
        "    avg_loss = epoch_loss / len(data_loader)\n",
-        "    print(f\"Epoch {epoch}/{epochs}, Loss: {avg_loss:.4f}\")\n",
+        "    print(f\"Epoch {epoch}/{epochs}, Loss: {avg_loss:.4f}\")\n"
-        "\n",
+      ]
        "betas = 0.02 * torch.ones(T)\n",
        "alphas = 1 - betas\n",
        "alpha_bar = torch.cumprod(alphas, dim=0)"
      ],
      "id": "1027e7e50cc6fe83"
    },
    {
      "cell_type": "markdown",
      "id": "24d5227d55d6d5e4",
      "metadata": {
        "id": "24d5227d55d6d5e4"
      },
      "cell_type": "markdown",
      "source": [
-        "### Inference with the Reverse Process"
+        "### Inference with the Reverse Process\n",
-      ],
+        "\n",
-      "id": "24d5227d55d6d5e4"
+        "Above, it was enough to have single steps of the denoising process for training the network. At inference time, we now need to perform the whole denoising process: starting from an initial pure noise sample, updating it over the course of 1000 iterations for denoising, to finally obtain a sample from the target distribution at the end.\n",
        "\n",
        "The code below implements this process, and records different snapshots of the denoising process for visualization."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "bb1b68299909aa1e",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-03-17T00:02:25.210393Z",
@@ -609,10 +544,14 @@
        },
        "id": "bb1b68299909aa1e"
      },
-      "cell_type": "code",
+      "outputs": [],
      "source": [
        "n_infer = 1000\n",
        "\n",
        "betas = 0.02 * torch.ones(T)\n",
        "alphas = 1 - betas\n",
        "alpha_bar = torch.cumprod(alphas, dim=0)\n",
        "\n",
        "with torch.no_grad():\n",
        "    x = torch.randn(n_infer, input_dim)\n",
        "\n",
@@ -637,22 +576,24 @@
        "\n",
        "        if t in record_steps:\n",
        "            snapshots[t] = x.clone().cpu().numpy()"
-      ],
+      ]
      "id": "bb1b68299909aa1e",
      "outputs": [],
      "execution_count": 11
    },
    {
      "cell_type": "markdown",
      "id": "df3ce91df721e024",
      "metadata": {
        "id": "df3ce91df721e024"
      },
      "cell_type": "markdown",
      "source": [
-        "### Visualize the Sampled Trajectories"
+        "### Visualize the Sampled Trajectories\n",
-      ],
+        "\n",
-      "id": "df3ce91df721e024"
+        "Next, we visualize how the initial distribution slowly turns into the target distribution from our GM case."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "id": "cafff2d56a61be60",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-03-17T00:02:26.532752Z",
@@ -660,7 +601,7 @@
        },
        "id": "cafff2d56a61be60"
      },
-      "cell_type": "code",
+      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "\n",
@@ -723,12 +664,12 @@
        "    plt.suptitle(\"Diffusion Model Reverse Process\")\n",
        "    plt.tight_layout()\n",
        "    plt.show()\n"
-      ],
+      ]
      "id": "cafff2d56a61be60",
      "outputs": [],
      "execution_count": 12
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "id": "5cd9e3d82a956ac6",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-03-17T00:02:33.517119Z",
@@ -741,27 +682,42 @@
        "id": "5cd9e3d82a956ac6",
        "outputId": "15be7112-1732-406f-ad94-bbd35441215b"
      },
      "cell_type": "code",
      "source": [
        "visualize_snapshots(snapshots)"
      ],
      "id": "5cd9e3d82a956ac6",
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAABcwAAAMSCAYAAABXqYdyAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjAsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvlHJYcgAAAAlwSFlzAAAPYQAAD2EBqD+naQABAABJREFUeJzs3Xd8FNX6x/HP1vTeCaFDQu9NKYJUAWmCitg7ol7s5dqv3atXsTfsnSIiTao06Z3QCZBAeq+b3Z3fH5FoTBCChBh/37cvXrgzZ2bOOWxydp858xyTYRgGIiIiIiIiIiIiIiL/z5lruwIiIiIiIiIiIiIiIn8HCpiLiIiIiIiIiIiIiKCAuYiIiIiIiIiIiIgIoIC5iIiIiIiIiIiIiAiggLmIiIiIiIiIiIiICKCAuYiIiIiIiIiIiIgIoIC5iIiIiIiIiIiIiAiggLmIiIiIiIiIiIiICKCAuYiIiIiIiIiIiIgIoIC5iIiIyFlx5ZVXEhsbW2l7aWkpr732GoMGDaJNmzbExsayaNGi8v2ffPIJF110Ee3atSM2NpaPPvqoxuvav39/+vfvX+PXqcsSExOJjY3lgQce+EvnWbt2LbGxsUydOvUs1UxERERERGqStbYrICIiIvJ38ceAt81mw9fXl6ioKFq1asWgQYPo1asXFovltM85bdo03njjDbp27crQoUOxWq00btwYgB9//JGnn36aVq1acfXVV2O32+nQocPZbFKd0b9/f5KSkgD46KOP6NmzZ5XlHnzwQWbMmAHA5MmTuf32289ZHc+FtWvXctVVV1XYZrVaCQoKokOHDlx99dV07dq1lmpXtyUmJnLhhRdW2GaxWAgMDKRNmzZcccUV9O3bt5ZqJyIiIiJ/FwqYi4iIiPzB5MmTAXC5XOTl5bFv3z6+//57vvvuO9q0acNLL71UHvQ+4fnnn6eoqKjSuZYuXYq3tzcffvghdru90j6At99+m4iIiBpqTWXnYhb7mbJarXz33XdVBszz8/OZN28eVqsVp9NZC7U7d6Kjoxk9ejQARUVF7Ny5k59++olFixbxyiuvMHTo0FquYd3l5+fH1VdfDUBJSQm7d+9m+fLlLF++nIcffrjSDQsRERER+f9FAXMRERGRP6hq1nJ6ejpPPfUU8+fP59prr2X69OmEhISU769Xr16V50pNTSUoKKhSsPzEPuCcBssBGjRocE6vVx0XXHABCxcuJCsri6CgoAr7Zs+eTVFREQMHDuSnn36qpRqeG9HR0ZXeh++++y7//e9/efHFFxUw/wv8/f0r9e306dN56KGHeOWVVxg3bhxeXl61VDsRERERqW3KYS4iIiJyGkJDQ3nllVfo1q0bx48f5+23366w/485zB944AFiY2NJTEwkKSmJ2NhYYmNj6d+/P1OnTiU2Npa1a9cClO87cfyp8mdXlS/dMAxmzpzJZZddRo8ePWjbti19+/bl+uuvZ+7cuRXKniyHucPh4N1332XEiBG0b9+eTp06MWHChErH/7GOiYmJTJkyhe7du9O2bVvGjBlTPnu+usaPH4/D4eD777+vtO/bb78lKiqK3r17n/T41NRUnnjiCfr370+bNm3o0aMHkydPZseOHVWWz8/P59lnn6VPnz60bduWIUOGMG3aNAzDOOk1ioqKeOeddxg5ciQdOnSgY8eOXHrppcyZM6f6Da6GSy65BICkpCQyMzMr7Z8zZw5XXnklXbp0oW3btgwdOpQ333wTh8NRXiYlJYWWLVsyatSok17nhhtuIDY2lr1791bYvnXrVu644w7OP/982rRpQ9++fXn00UdJSUmpdI4T71GHw8Hrr7/O4MGDadOmTfl72uFw8MknnzB69Gi6du1K+/bt6d+/P7feeiurV6+udL4DBw7wwAMP0LdvX9q0acN5553H3XffzcGDB0+r705lzJgxeHt7U1hYyL59+4DffoaPHj3Kp59+yogRI2jXrh1XXnll+XEJCQncd9999O7dmzZt2tCrVy/uu+8+EhISqryOy+Xiyy+/5LLLLqNz5860a9eOgQMH8vDDD1c6xul08vnnnzN+/Hg6depE+/btGTVqFJ999hlut7vSuRcvXszVV19Nr169yusyceJEPv/88wrljh49yiOPPMLAgQNp164d3bp1Y8SIETz66KNkZWX9tY4UERER+QfQDHMRERGR02Q2m5k0aRLr1q3jxx9/5KGHHsJkMlVZdsCAAURHR/Pxxx8DlKeA8PPzo2XLlkyePJmZM2eSlJRUngLmr3jllVd45513qF+/PkOHDsXPz4+0tDS2b9/O/Pnzueiii/70eIfDwfXXX8+6deto0qQJEyZMoLi4mAULFjBlyhR2797NXXfdVem4pKQkxo0bR0xMDCNHjiQnJ4e5c+cyadIkpk2bRo8eParVjvPOO4/o6Gi+++47rrnmmvLtO3bsYNeuXUyePBmzueo5H0ePHmXChAmkpqbSo0cPhg0bxvHjx5k/fz7Lli1j6tSp9OvXr0Kbr7nmGrZv305cXBwjRowgLy+PN998k3Xr1lV5jdzcXK6++mp27dpF69atGTt2LG63m5UrV3L33Xezb98+pkyZUq02nwmbzVbh9Ync7pGRkQwaNAh/f3+2bNnCq6++ypo1a5g2bRpWq5WIiAjOO+88Vq5cyZ49eyrdeElNTWX16tW0bt2aFi1alG//7rvvePTRR7Hb7fTv35/IyEgOHz7Mt99+y5IlS/jmm2+qfMrijjvuYPv27fTp04cBAwaUP5Xx4IMPMmfOHFq0aMHIkSPx9PQkNTWVjRs3smLFCs4777zyc/z888/cfvvtOJ1O+vXrR4MGDUhJSWHhwoUsW7aMTz75hNatW//lPj1xk+SPP9NPP/00GzZsoG/fvvTt27d8DYNt27Zx7bXXUlBQQP/+/WnWrBkHDx5k9uzZLF68mGnTptGuXbvy8zgcDm655RZWrVpFVFQUw4cPx9fXl6SkJBYtWkTnzp1p1KgRULZY8C233MLKlStp3Lgxw4cPx8PDg7Vr1/LUU0+xdetWXnzxxfJzf/311zz66KOEhYXRr18/goKCyMjIYM+ePcyYMYMrrrgCKPv3veSSS8jPz6dPnz4MGjSIkpISEhMTmT17NhMnTqz0ZIeIiIjI/zcKmIuIiIhUQ+fOnbFarWRkZJCYmEhMTEyV5QYMGMCAAQOYOXMmUDnNS/fu3Vm3bh1JSUlnZeHKr7/+moiICObMmVMpnURVs5H/aNq0aaxbt44+ffrw1ltvYbWWfUycPHky48aN45133uGCCy6gU6dOFY5bt24dt99+e4Wg//Dhw7nhhhv44IMPqh0wN5lMXHLJJbz66qts3ryZjh07AvDNN99gNpsZO3Ysq1atqvLYxx9/nNTUVP71r39x6623lm+fMGECEydO5IEHHmDJkiX4+PgA8OGHH7J9+3YGDRrEq6++Wh6Iv/HGGxk7dmyV13jmmWfYtWsX99xzDzfeeGP59pKSEiZNmsQ777zDkCFDaNmyZbXafTq++uorAFq0aIGfn1/59hkzZjBjxgwGDhzISy+9hKenZ/m+qVOn8vrrr/P555+X37QZPXo0K1euZNasWdx///0VrjF79mxcLld5/nSAQ4cO8fjjjxMdHc1nn31WIYXQmjVruO6663j66ad54403KtU5KSmJH374geDg4PJteXl5/Pjjj7Ru3Zpvv/220iK6v5/lnJOTw913342npyeff/45zZo1K9+3d+9eLr30Uv7973+X/5ydqenTp1NUVIS3tzfNmzevsG/nzp3MnDmzws+6YRjcf//95Ofn8+KLL3LxxReX75s7dy5TpkzhvvvuY+7cueXvq9dff51Vq1bRr18/XnvttQppmhwOB/n5+eWv3377bVauXMnEiRN56KGHyvvI5XLxyCOPMH36dAYPHsyAAQOAsp9/m83G999/XyFVFFT8+V+wYAHZ2dk89NBD5e+HEwoLC096M0pERETk/xN9IhIRERGpBrvdTmBgIMDfLn2B1WqtFHwEKgQrT2b69OmYTCYeeOCB8mA5QEhISHnw+dtvv610XHR0dIXgNEDv3r2pV68e27Ztq24TABg7diwWi4VvvvkGKAvkzZkzh169ep00V3xycjIrV66kXr163HDDDRX2derUiWHDhpGdnV0h9/mMGTMwm83ce++9FQKFMTExFdJunJCVlcXs2bNp06ZNhWA5gIeHB/feey+GYfDDDz+cUbt/LykpialTpzJ16lReeOEFrrrqKl599VV8fX154oknKpT95JNPsFqtPPPMMxWC5QCTJk0iMDCwQp0GDBiAn58fP/zwAy6Xq0L5WbNmYbPZGDZsWPm2L7/8ktLSUh5++OFK+fZ79uxJ//79Wbp0aYWA7wl33nlnpfefyWTCMAzsdnuVAdrfz3CeNWsWubm53HHHHRWC5VB242DcuHHs2rWL/fv3VzrPyeTm5pb37UsvvcSNN97Iww8/DMCUKVMq9eENN9xQ6cbYpk2bOHjwIB07dqwQLAe46KKL6Ny5M4cOHWLjxo1AWaD7iy++wNPTkyeeeKLSmgZ2u728n9xuN5999hlhYWE8+OCDFX6mLRYLDzzwACaTqdL7zGq1VvjZPaGqn/8/thHA29u7yu0iIiIi/99ohrmIiIhINf1ZfuvaMmLECD799FMuuugihg4dSteuXenYsWOFmcgnk5+fz+HDh4mIiKBp06aV9p+YJR4fH19pX1xcXJVB+sjISLZs2VL9hlC2CGqfPn2YP38+Dz/8MPPmzaOgoIDx48ef9Jhdu3YBZU8A/DFdyYk2zJ49m127djFq1KjyNkdFRVW5CGq3bt0qbdu+fTsulwuTycTUqVMr7Xc6nQBnJa92UlISr7/+eoVtAQEBfPzxxxVmrxcVFbF7926CgoLK0//8kd1u58CBA+WvPT09GTp0KN988w0rV66kb9++QFnam3379jFw4MAKQdYT/47r1q1j+/btlLine truncated
            "text/plain": [
              "<Figure size 1500x800 with 6 Axes>"
-            ],
+            ]
            "image/png": "iVBORw0KGgoAAAANSUhEUgAABcwAAAMSCAYAAABXqYdyAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjAsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvlHJYcgAAAAlwSFlzAAAPYQAAD2EBqD+naQABAABJREFUeJzs3Xd8FNX6x/HP1vTeCaFDQu9NKYJUAWmCitg7ol7s5dqv3atXsTfsnSIiTao06Z3QCZBAeq+b3Z3fH5FoTBCChBh/37cvXrgzZ2bOOWxydp858xyTYRgGIiIiIiIiIiIiIiL/z5lruwIiIiIiIiIiIiIiIn8HCpiLiIiIiIiIiIiIiKCAuYiIiIiIiIiIiIgIoIC5iIiIiIiIiIiIiAiggLmIiIiIiIiIiIiICKCAuYiIiIiIiIiIiIgIoIC5iIiIiIiIiIiIiAiggLmIiIiIiIiIiIiICKCAuYiIiIiIiIiIiIgIoIC5iIiIyFlx5ZVXEhsbW2l7aWkpr732GoMGDaJNmzbExsayaNGi8v2ffPIJF110Ee3atSM2NpaPPvqoxuvav39/+vfvX+PXqcsSExOJjY3lgQce+EvnWbt2LbGxsUydOvUs1UxERERERGqStbYrICIiIvJ38ceAt81mw9fXl6ioKFq1asWgQYPo1asXFovltM85bdo03njjDbp27crQoUOxWq00btwYgB9//JGnn36aVq1acfXVV2O32+nQocPZbFKd0b9/f5KSkgD46KOP6NmzZ5XlHnzwQWbMmAHA5MmTuf32289ZHc+FtWvXctVVV1XYZrVaCQoKokOHDlx99dV07dq1lmpXtyUmJnLhhRdW2GaxWAgMDKRNmzZcccUV9O3bt5ZqJyIiIiJ/FwqYi4iIiPzB5MmTAXC5XOTl5bFv3z6+//57vvvuO9q0acNLL71UHvQ+4fnnn6eoqKjSuZYuXYq3tzcffvghdru90j6At99+m4iIiBpqTWXnYhb7mbJarXz33XdVBszz8/OZN28eVqsVp9NZC7U7d6Kjoxk9ejQARUVF7Ny5k59++olFixbxyiuvMHTo0FquYd3l5+fH1VdfDUBJSQm7d+9m+fLlLF++nIcffrjSDQsRERER+f9FAXMRERGRP6hq1nJ6ejpPPfUU8+fP59prr2X69OmEhISU769Xr16V50pNTSUoKKhSsPzEPuCcBssBGjRocE6vVx0XXHABCxcuJCsri6CgoAr7Zs+eTVFREQMHDuSnn36qpRqeG9HR0ZXeh++++y7//e9/efHFFxUw/wv8/f0r9e306dN56KGHeOWVVxg3bhxeXl61VDsRERERqW3KYS4iIiJyGkJDQ3nllVfo1q0bx48f5+23366w/485zB944AFiY2NJTEwkKSmJ2NhYYmNj6d+/P1OnTiU2Npa1a9cClO87cfyp8mdXlS/dMAxmzpzJZZddRo8ePWjbti19+/bl+uuvZ+7cuRXKniyHucPh4N1332XEiBG0b9+eTp06MWHChErH/7GOiYmJTJkyhe7du9O2bVvGjBlTPnu+usaPH4/D4eD777+vtO/bb78lKiqK3r17n/T41NRUnnjiCfr370+bNm3o0aMHkydPZseOHVWWz8/P59lnn6VPnz60bduWIUOGMG3aNAzDOOk1ioqKeOeddxg5ciQdOnSgY8eOXHrppcyZM6f6Da6GSy65BICkpCQyMzMr7Z8zZw5XXnklXbp0oW3btgwdOpQ333wTh8NRXiYlJYWWLVsyatSok17nhhtuIDY2lr1791bYvnXrVu644w7OP/982rRpQ9++fXn00UdJSUmpdI4T71GHw8Hrr7/O4MGDadOmTfl72uFw8MknnzB69Gi6du1K+/bt6d+/P7feeiurV6+udL4DBw7wwAMP0LdvX9q0acN5553H3XffzcGDB0+r705lzJgxeHt7U1hYyL59+4DffoaPHj3Kp59+yogRI2jXrh1XXnll+XEJCQncd9999O7dmzZt2tCrVy/uu+8+EhISqryOy+Xiyy+/5LLLLqNz5860a9eOgQMH8vDDD1c6xul08vnnnzN+/Hg6depE+/btGTVqFJ999hlut7vSuRcvXszVV19Nr169yusyceJEPv/88wrljh49yiOPPMLAgQNp164d3bp1Y8SIETz66KNkZWX9tY4UERER+QfQDHMRERGR02Q2m5k0aRLr1q3jxx9/5KGHHsJkMlVZdsCAAURHR/Pxxx8DlKeA8PPzo2XLlkyePJmZM2eSlJRUngLmr3jllVd45513qF+/PkOHDsXPz4+0tDS2b9/O/Pnzueiii/70eIfDwfXXX8+6deto0qQJEyZMoLi4mAULFjBlyhR2797NXXfdVem4pKQkxo0bR0xMDCNHjiQnJ4e5c+cyadIkpk2bRo8eParVjvPOO4/o6Gi+++47rrnmmvLtO3bsYNeuXUyePBmzueo5H0ePHmXChAmkpqbSo0cPhg0bxvHjx5k/fz7Lli1j6tSp9OvXr0Kbr7nmGrZv305cXBwjRowgLy+PN998k3Xr1lV5jdzcXK6++mp27dpF69atGTt2LG63m5UrV3L33Xezb98+pkyZUq02nwmbzVbh9Ync7pGRkQwaNAh/f3+2bNnCq6++ypo1a5g2bRpWq5WIiAjOO+88Vq5cyZ49eyrdeElNTWX16tW0bt2aFi1alG//7rvvePTRR7Hb7fTv35/IyEgOHz7Mt99+y5IlS/jmm2+qfMrijjvuYPv27fTp04cBAwaUP5Xx4IMPMmfOHFq0aMHIkSPx9PQkNTWVjRs3smLFCs4777zyc/z888/cfvvtOJ1O+vXrR4MGDUhJSWHhwoUsW7aMTz75hNatW//lPj1xk+SPP9NPP/00GzZsoG/fvvTt27d8DYNt27Zx7bXXUlBQQP/+/WnWrBkHDx5k9uzZLF68mGnTptGuXbvy8zgcDm655RZWrVpFVFQUw4cPx9fXl6SkJBYtWkTnzp1p1KgRULZY8C233MLKlStp3Lgxw4cPx8PDg7Vr1/LUU0+xdetWXnzxxfJzf/311zz66KOEhYXRr18/goKCyMjIYM+ePcyYMYMrrrgCKPv3veSSS8jPz6dPnz4MGjSIkpISEhMTmT17NhMnTqz0ZIeIiIjI/zcKmIuIiIhUQ+fOnbFarWRkZJCYmEhMTEyV5QYMGMCAAQOYOXMmUDnNS/fu3Vm3bh1JSUlnZeHKr7/+moiICObMmVMpnURVs5H/aNq0aaxbt44+ffrw1ltvYbWWfUycPHky48aN45133uGCCy6gU6dOFY5bt24dt99+e4Wg//Dhw7nhhhv44IMPqh0wN5lMXHLJJbz66qts3ryZjh07AvDNN99gNpsZO3Ysq1atqvLYxx9/nNTUVP71r39x6623lm+fMGECEydO5IEHHmDJkiX4+PgA8OGHH7J9+3YGDRrEq6++Wh6Iv/HGGxk7dmyV13jmmWfYtWsX99xzDzfeeGP59pKSEiZNmsQ777zDkCFDaNmyZbXafTq++uorAFq0aIGfn1/59hkzZjBjxgwGDhzISy+9hKenZ/m+qVOn8vrrr/P555+X37QZPXo0K1euZNasWdx///0VrjF79mxcLld5/nSAQ4cO8fjjjxMdHc1nn31WIYXQmjVruO6663j66ad54403KtU5KSmJH374geDg4PJteXl5/Pjjj7Ru3Zpvv/220iK6v5/lnJOTw913342npyeff/45zZo1K9+3d+9eLr30Uv7973+X/5ydqenTp1NUVIS3tzfNmzevsG/nzp3MnDmzws+6YRjcf//95Ofn8+KLL3LxxReX75s7dy5TpkzhvvvuY+7cueXvq9dff51Vq1bRr18/XnvttQppmhwOB/n5+eWv3377bVauXMnEiRN56KGHyvvI5XLxyCOPMH36dAYPHsyAAQOAsp9/m83G999/XyFVFFT8+V+wYAHZ2dk89NBD5e+HEwoLC096M0pERETk/xN9IhIRERGpBrvdTmBgIMDfLn2B1WqtFHwEKgQrT2b69OmYTCYeeOCB8mA5QEhISHnw+dtvv610XHR0dIXgNEDv3r2pV68e27Ztq24TABg7diwWi4VvvvkGKAvkzZkzh169ep00V3xycjIrV66kXr163HDDDRX2derUiWHDhpGdnV0h9/mMGTMwm83ce++9FQKFMTExFdJunJCVlcXs2bNp06ZNhWA5gIeHB/feey+GYfDDDz+cUbt/LykpialTpzJ16lReeOEFrrrqKl599VV8fX154oknKpT95JNPsFqtPPPMMxWC5QCTJk0iMDCwQp0GDBiAn58fP/zwAy6Xq0L5WbNmYbPZGDZsWPm2L7/8ktLSUh5++OFK+fZ79uxJ//79Wbp0aYWA7wl33nlnpfefyWTCMAzsdnuVAdrfz3CeNWsWubm53HHHHRWC5VB242DcuHHs2rWL/fv3VzrPyeTm5pb37UsvvcSNN97Iww8/DMCUKVMq9eENN9xQ6cbYpk2bOHjwIB07dqwQLAe46KKL6Ny5M4cOHWLjxo1AWaD7iy++wNPTkyeeeKLSmgZ2u728n9xuN5999hlhYWE8+OCDFX6mLRYLDzzwACaTqdL7zGq1VvjZPaGqn/8/thHA29u7yu0iIiIi/99ohrmIiIhINf1ZfuvaMmLECD799FMuuugihg4dSteuXenYsWOFmcgnk5+fz+HDh4mIiKBp06aV9p+YJR4fH19pX1xcXJVB+sjISLZs2VL9hlC2CGqfPn2YP38+Dz/8MPPmzaOgoIDx48ef9Jhdu3YBZU8A/DFdyYk2zJ49m127djFq1KjyNkdFRVW5CGq3bt0qbdu+fTsulwuTycTUqVMr7Xc6nQBnJa92UlISr7/+eoVtAQEBfPzxxxVmrxcVFbF7926CgoLK0//8kd1u58CBA+WvPT09GTp0KN988w0rV66kb9++QFnam3379jFw4MAKQdYT/47r1q1j+/btlLine truncated
          },
-          "metadata": {}
+          "metadata": {},
          "output_type": "display_data"
        }
      ],
-      "execution_count": 13
+      "source": [
        "visualize_snapshots(snapshots)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7b73e6a1",
      "metadata": {},
      "source": [
        "Qualitatively, the DDPM network manages to accurate capture the two peaks in our distribution, and - best of all - it does so with a very simple and stable training process. All we need are the initial samples, and controlled perturbations with the noising schedule added on top!\n",
        "\n",
        "Looking ahead, three important extensions of the DDPM algorithm still need to be addressed:\n",
        "- Faster inference speed, primarily getting away with fewer network calls.\n",
        "- Conditioning on external parameters and observations.\n",
        "- And in the context of physics simulations, we'd of course like to bring back our prior knowledge in the form of PDE-base constraints.\n",
        "\n",
        "Nonetheless, we've made the most important first step with DDPMs now: we have turned the deterministic networks into tools for probabilistic inference that robustly learn and reproduce complex distributions. So we've arrived at a very capable and stable method to learn distributions from data. The only caveat left is the computational cost. The probabilistic DDPM network require many sequential evaluations to produce a sample. These _neural function evaluations_ (NFEs) are of course costly, and while many practical cases can do with fewer than the 1000 we've used above, even 10 mean that the network is 10x slower than the \"regular\" deterministic version. This brings us to the last remaining step in this _generative AI journey_: how can we make the inference process faster, without negatively affecting the accuracy of our target, the inferred posterior distribution?\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "diffusion",
      "language": "python",
@@ -778,11 +734,8 @@
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython2",
      "version": "2.7.6"
    },
    "colab": {
      "provenance": []
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}
@@ -1,21 +1,24 @@
 {
  "cells": [
    {
      "cell_type": "markdown",
      "id": "18dcb0d9d85f5323",
      "metadata": {
        "id": "18dcb0d9d85f5323"
      },
      "cell_type": "markdown",
      "source": [
        "\n",
-        "## Flow Matching\n",
+        "# Flow Matching\n",
        "\n",
-        "To motivate the transition from _denoising_ to _flow matching_, the score formulation from {eq}`score-perturbed` comes in handy: as a reminder, there we were integrating the gradient $\\nabla_x$ of the log likelihoods to transform noise into a target distribution over the course of a continuous virtual time axis. To briefly re-cap, we integrated $\\nabla_{\\tilde{x}} \\log p_\\sigma(\\tilde{x})$ over time with Langevin dynamics\n",
+        "To reduce the many function evaluations of DDPM, we'll turn to _flow matching_ {cite}`lipman2022flow`. To motivate the transition from denoising to flow matching, the score formulation from {eq}`score-perturbed` comes in handy: as a reminder, there we were integrating the gradient $\\nabla_x$ of the log likelihoods to transform noise into a target distribution over the course of a continuous virtual time axis. To briefly re-cap, we integrated $\\nabla_{\\tilde{x}} \\log p_\\sigma(\\tilde{x})$ over time with Langevin dynamics\n",
        "$dx/dt = \\epsilon \\nabla_x \\log p(x) + \\sqrt{2 \\epsilon} z_i $, and used a training objective of $\\arg \\min_\\theta \\mathbb{E} \\left[ || s_\\theta(\\tilde{x}) - \\nabla_{\\tilde{x}} \\log p_\\sigma(\\tilde{x}) ||^2 \\right]$ to learn the score $s_\\theta(\\tilde{x})$. Score modeling leveraged the mathematical equivalence between the probability flows of diffusion processes and the trajectories of probability densities in the form of ordinary differential equations (ODEs).\n",
        "\n",
-        "For _flow matching_, our goal is to construct a _flow_ in a velocity field, and\n",
+        "## Learning Flows with Velocities\n",
        "the gradients in score matching can actually readily be interpreted as a _velocity_: neglecting the stochastic terms of the Langevin time evolution above, we're left with the velocity $dx/dt = \\nabla_x \\log p(x)$ of $x$ being the score. The main challenge we were fighting with in the context of score matching was how to obtain a good reference velocity, so that we cuold train a neural network.\n",
        "\n",
-        "Now, with the powerful machinery of denoising at hand, we can revisit the  concept of _flows and velocities_: we aim for learning a continuous-time flow, for which we prescribe as-simple-as-possible reference velocities. As before, the velocity viewpoint frees us from any constraints on the network architecture, and the simple velocities will make training and inference substantially faster than before.\n",
+        "For _flow matching_, our goal is to construct a _flow_ in a velocity field, and\n",
        "the gradients in score matching can actually readily be interpreted as a _velocity_: neglecting the stochastic terms of the Langevin time evolution above, we're left with the velocity $dx/dt = \\nabla_x \\log p(x)$ of $x$ being the score. The main challenge we were fighting with in the context of score matching was how to obtain a good reference velocity, so that we could train a neural network.\n",
        "\n",
        "Now, with the powerful machinery of _denoising_ at hand, we can revisit the  concept of _flows and velocities_: we aim for learning a continuous-time flow, for which we prescribe as-simple-as-possible reference velocities. As before, the velocity viewpoint frees us from any constraints on the network architecture, and the simple velocities will make training and inference substantially faster than before.\n",
        "\n",
        "Let's formalize this: methods such as _flow matching_ are typically categorized as \"continuous-time flow models\", and transform samples $x$ from a sampling distribution $p_0$ to samples of a target or posterior distribution $p_1$. This mapping is expressed via the ODE\n",
        "\n",
@@ -25,7 +28,6 @@
        "\n",
        "where $v_\\theta(x,t)$ represents a neural network with parameters $\\theta$. In flow matching, the network $v_\\theta$ is trained by regressing a vector field that generates probability paths that map from $p_0$ to $p_1$.\n",
        "\n",
        "\n",
        "We say that a smooth vector field $u : [0,1] \\times \\mathbb{R}^d \\to \\mathbb{R}^d$, called _velocity_, generates the probability paths $p_t$, if it satisfies the continuity equation $\\partial p / \\partial t = - \\nabla \\cdot (p_t u_t)$.\n",
        "Informally, this means that we can sample from the distribution $p_t$ by sampling $x_0 \\sim p_0$ and then solving the ODE $dx = u(x,t)dt$ with initial condition $x_0$. In the following, we will denote $u(x,t)$ by $u_t(x)$.\n",
        "To regress the velocity field, we define the _flow matching_ objective\n",
@@ -53,9 +55,10 @@
        "\\end{align}\n",
        "$$ (conditional-flow-matching)\n",
        "\n",
-        "This version is tractable and can be used for actual training runs, in contrast to the un-conditional objective from {eq}`eq-flow-matching`.\n",
+        "This version is tractable and can be used for actual training runs, in contrast to the un-conditional objective from equation {eq}`eq-flow-matching`.\n",
        "This means that we can train $v_\\theta(x,t)$ to regress $u_t(x)$ generating the mapping from $p_0$ to the target distribution $p_1$.\n",
        "\n",
        "## Mappings and Conditioning\n",
        "\n",
        "Especially important: we have a lot of freedom when specifying the mapping from $p_0$ to $p_1$ via the conditioning variable $z$ and the conditional likelihoods $p_t$ in this formulation.\n",
        "So how do we best make use of this freedom?\n",
@@ -84,10 +87,22 @@
        "```\n",
        "\n",
        "In addition, a fundamental advantage of flow matching over denoising diffusion models comes from a smart choice of the generating velocity paths: aiming for a linear velocity towards the target means we would obtain a straight path from an original sample towards our goal point. **Linear** here means that we could take a **single** Euler step to compute it from any starting point! As illustrated in the picture right above, this stands in huge contrast to the hundreds of steps we might have to make to follow one of the curved (unconstrained) denoising paths. It turns out we don't always obtain perfectly straight paths, so a single step can be sub-optmial (or might require further fine tuning / distillation of the flows). Nonetheless, flow matching typically works with a substantially reduced number of iterations compared to denoising. As inference time is directly proportional to the number of steps, this will result in a corresponding speed up when computing a sample from our posterior distribution. And, most importantly, we've not sacrificied any of the original goals: a convenient sampling procedure and an accurate coverage of the posterior just from data.\n"
-      ],
+      ]
      "id": "18dcb0d9d85f5323"
    },
    {
      "cell_type": "markdown",
      "id": "69695675",
      "metadata": {},
      "source": [
        "## Implementing Flow Matching\n",
        "\n",
        "Just like for DDPM, we first need to initialize our standard Gaussian Mixture (GM) case:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c041087b8b9191f0",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-03-16T22:58:52.827556Z",
@@ -95,13 +110,12 @@
        },
        "id": "c041087b8b9191f0"
      },
-      "cell_type": "code",
+      "outputs": [],
      "source": [
        "import numpy as np\n",
        "\n",
        "class GaussianMixture:\n",
        "    def __init__(self, parameters):\n",
        "\n",
        "        self.parameters = parameters\n",
        "        self.distributions = [\n",
        "            {\n",
@@ -126,24 +140,23 @@
        "    {\"mean\": [0, 0], \"std\": [1, 1]},\n",
        "    {\"mean\": [3, 2], \"std\": [0.5, 0.5]}\n",
        "]\n",
        "\n",
        "mixture = GaussianMixture(parameters)"
-      ],
+      ]
      "id": "c041087b8b9191f0",
      "outputs": [],
      "execution_count": 1
    },
    {
      "cell_type": "markdown",
      "id": "7f9d408886c23020",
      "metadata": {
        "id": "7f9d408886c23020"
      },
      "cell_type": "markdown",
      "source": [
-        "### Setup Dataset and Dataloader"
+        "Setting up the training dataset is even simpler than for the denoising task. We don't need to include a noising schedule, but rather, we simply compute a straight velocity `u_t` for each requested sample after selecting a random time `t`."
-      ],
+      ]
      "id": "7f9d408886c23020"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8ac07b3409739df9",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-03-16T23:32:38.653547Z",
@@ -151,7 +164,7 @@
        },
        "id": "8ac07b3409739df9"
      },
-      "cell_type": "code",
+      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
@@ -160,7 +173,6 @@
        "\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "\n",
        "class FlowMatchingDataset(Dataset):\n",
        "    def __init__(self, mixture, n_samples=1000, sigma_min=1e-4):\n",
        "        super().__init__()\n",
@@ -185,22 +197,22 @@
        "        u_t = torch.tensor(u_t, dtype=torch.float32)\n",
        "\n",
        "        return x_t, t, u_t\n"
-      ],
+      ]
      "id": "8ac07b3409739df9",
      "outputs": [],
      "execution_count": 2
    },
    {
      "cell_type": "markdown",
      "id": "471468fdc74f2af2",
      "metadata": {
        "id": "471468fdc74f2af2"
      },
      "cell_type": "markdown",
      "source": [
-        "### Small MLP for Learning the Velocity"
+        "Similar to before, we'll use a simple neural network with three layers as \"backbone\" for the flow matching task:"
-      ],
+      ]
      "id": "471468fdc74f2af2"
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "id": "e298ca1a4bee5d7f",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-03-16T23:32:39.251649Z",
@@ -208,7 +220,7 @@
        },
        "id": "e298ca1a4bee5d7f"
      },
-      "cell_type": "code",
+      "outputs": [],
      "source": [
        "class VelocityNet(nn.Module):\n",
        "    def __init__(self, in_dim=2, time_dim=1, hidden_dim=128, out_dim=2):\n",
@@ -224,22 +236,24 @@
        "    def forward(self, x, t):\n",
        "        xt = torch.cat([x, t], dim=1)\n",
        "        return self.net(xt)\n"
-      ],
+      ]
      "id": "e298ca1a4bee5d7f",
      "outputs": [],
      "execution_count": 3
    },
    {
      "cell_type": "markdown",
      "id": "832f39392dda6db4",
      "metadata": {
        "id": "832f39392dda6db4"
      },
      "cell_type": "markdown",
      "source": [
-        "### Train a Velocity Model"
+        "### Training a Velocity Model\n",
-      ],
+        "\n",
-      "id": "832f39392dda6db4"
+        "Now we can start training the velocity model. This means with simply let the network predict the \"supervised\" ground truth velocity we receive from the dataset class. Just like for denoising, we have a very well-defined and stable learning task. The code cell below runs this training for 50 epochs:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "id": "853faf1062eab3a8",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-03-16T23:33:04.270517Z",
@@ -251,41 +265,10 @@
        "id": "853faf1062eab3a8",
        "outputId": "b9f4b482-5adb-4071-b5ee-8dced95c5227"
      },
      "cell_type": "code",
      "source": [
        "n_samples = 10000\n",
        "batch_size = 128\n",
        "n_epochs = 50\n",
        "learning_rate = 1e-3\n",
        "\n",
        "dataset = FlowMatchingDataset(mixture, n_samples=n_samples)\n",
        "dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)\n",
        "\n",
        "model = VelocityNet().to(device)\n",
        "optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n",
        "criterion = nn.MSELoss()\n",
        "\n",
        "for epoch in range(n_epochs):\n",
        "    running_loss = 0.0\n",
        "    for x_t, t, u_t in dataloader:\n",
        "        x_t = x_t.to(device)\n",
        "        t = t.to(device)\n",
        "        u_t = u_t.to(device)\n",
        "        optimizer.zero_grad()\n",
        "        pred_v = model(x_t, t)\n",
        "        loss = criterion(pred_v, u_t)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        running_loss += loss.item() * x_t.size(0)\n",
        "    running_loss /= len(dataset)\n",
        "    print(f\"Epoch {epoch + 1}/{n_epochs}, Loss: {running_loss:.4f}\")\n"
      ],
      "id": "853faf1062eab3a8",
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 1/50, Loss: 2.8302\n",
            "Epoch 2/50, Loss: 2.2797\n",
@@ -340,35 +323,76 @@
          ]
        }
      ],
-      "execution_count": 4
+      "source": [
        "n_samples = 10000\n",
        "batch_size = 128\n",
        "n_epochs = 50\n",
        "learning_rate = 1e-3\n",
        "\n",
        "dataset = FlowMatchingDataset(mixture, n_samples=n_samples)\n",
        "dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)\n",
        "\n",
        "model = VelocityNet().to(device)\n",
        "optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n",
        "criterion = nn.MSELoss()\n",
        "\n",
        "for epoch in range(n_epochs):\n",
        "    running_loss = 0.0\n",
        "    for x_t, t, u_t in dataloader:\n",
        "        x_t = x_t.to(device)\n",
        "        t = t.to(device)\n",
        "        u_t = u_t.to(device)\n",
        "        optimizer.zero_grad()\n",
        "        pred_v = model(x_t, t)\n",
        "        loss = criterion(pred_v, u_t)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        running_loss += loss.item() * x_t.size(0)\n",
        "    running_loss /= len(dataset)\n",
        "    print(f\"Epoch {epoch + 1}/{n_epochs}, Loss: {running_loss:.4f}\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "4e03ffe8ace08e19",
      "metadata": {
        "id": "4e03ffe8ace08e19"
      },
      "cell_type": "markdown",
      "source": [
-        "### Inference via Solving the ODE"
+        "### Inference via Solving the ODE\n",
-      ],
+        "\n",
-      "id": "4e03ffe8ace08e19"
+        "At inference time, we actually only have to integrate the velocities produced by the network starting from a random sample. While this could be done with simple Euler steps, we'll again use the ODE solvers from the `torchdiffeq` package for additional flexibility. (Feel free to experiment with other integration schemes below.)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "IU2P4X9RexTO",
      "metadata": {
        "id": "IU2P4X9RexTO"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "    import google.colab  # only to ensure that we are inside colab\n",
        "    %pip install --quiet torchdiffeq\n",
        "except ImportError:\n",
        "    print(\"This notebook is running locally, please install torchdiffeq manually.\")"
-      ],
+      ]
      "metadata": {
        "id": "IU2P4X9RexTO"
      },
      "id": "IU2P4X9RexTO",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "id": "8afac219",
      "metadata": {},
      "source": [
        "The integration itself is done in `odeint` using the trained network in the `ode_func` wrapper. By default, we're integrating the trajectories over 100 steps, but in practice, much fewer should still give good results. Nonetheless, this is also an order of magnitude less that what we used for the DDPM version."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "691f3460aa86e1a1",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-03-17T00:06:20.371503Z",
@@ -376,12 +400,12 @@
        },
        "id": "691f3460aa86e1a1"
      },
-      "cell_type": "code",
+      "outputs": [],
      "source": [
        "import torch\n",
        "from torchdiffeq import odeint\n",
        "\n",
-        "def integrate_flow(model, x0, t_span=(0.0, 1.0), n_steps=100, method='dopri5'):\n",
+        "def integrate_flow(model, x0, t_span=(0.0, 1.0), n_steps, method='dopri5'):\n",
        "\n",
        "    t = torch.linspace(t_span[0], t_span[1], n_steps).to(x0.device)\n",
        "\n",
@@ -400,22 +424,24 @@
        "trajectory, time_points = integrate_flow(model, x0_gen, n_steps=100)\n",
        "\n",
        "x_gen_final = trajectory[-1].detach().cpu().numpy()"
-      ],
+      ]
      "id": "691f3460aa86e1a1",
      "outputs": [],
      "execution_count": 7
    },
    {
      "cell_type": "markdown",
      "id": "ae7e42318e2841b2",
      "metadata": {
        "id": "ae7e42318e2841b2"
      },
      "cell_type": "markdown",
      "source": [
-        "### Visualize the Sampling Trajectories"
+        "### Visualize the Sampling Trajectories\n",
-      ],
+        "\n",
-      "id": "ae7e42318e2841b2"
+        "We start with points from a standard gaussian, and color them based on their initial position. The plots below again show states at six different times of the  trajectories."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5d381b5a8e3514e1",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-03-17T00:06:56.685929Z",
@@ -428,7 +454,18 @@
        "id": "5d381b5a8e3514e1",
        "outputId": "20fb889d-a765-4787-91bb-8bf4f4b96612"
      },
-      "cell_type": "code",
+      "outputs": [
        {
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAABcwAAAOHCAYAAAAaAgCCAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjAsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvlHJYcgAAAAlwSFlzAAAPYQAAD2EBqD+naQABAABJREFUeJzs3XeYHWXd//H3nH7O9t5r6qY3kpBACIQEQq+CUiwgYBd9VHzsyk957IoKFhQQ6Z0QSggpJKT33jabzfZeTz8zvz9WIjEddpNs+Ly4cu3FOTP33HPv7Pme+c5dDMuyLEREREREREREREREPuJsp7oCIiIiIiIiIiIiIiKnAyXMRURERERERERERERQwlxEREREREREREREBFDCXEREREREREREREQEUMJcRERERERERERERARQwlxEREREREREREREBFDCXEREREREREREREQEUMJcRERERERERERERARQwlxEREREREREREREBFDCXEREREREREREREQEUMJc5LS3du1a7r//fjo6Ok7J8Ts6Ovje977H5MmTGTNmDLfccgtbtmw57v337NnDbbfdxtixY5k4cSLf+MY3aGlpOWQ70zT561//ygUXXMDIkSO5/PLLmTNnTm+eioiISJ/rr3HbNE2ef/557rrrLs477zzGjBnDZZddxp/+9CdCodBh93nmmWeYPXs2I0eOZNasWfzzn//s7dMRERHpU/01bv+3SCTCJZdcwpAhQ3jooYcOeV/32yInRglzkdPcunXr+MMf/nBKArhpmtxxxx3MmTOHm2+++UCy+5ZbbqGiouKY+9fV1XHTTTdRWVnJ3XffzWc+8xkWLVrEpz/9acLh8EHb/uY3v+GXv/wlU6dO5Xvf+x65ubl8/etf59VXX+2jsxMREel9/TVuBwIBvv3tb9Pa2sqNN97I//7v/zJy5Ejuv/9+br/9dizLOmj7J598ku9+97sMGjSI733ve4wZM4Z7772Xv/zlL314hiIiIr2rv8bt//bYY49RW1t7xPd1vy1yYhynugIicvp6/fXXWbduHb/73e+4+OKLAZg9ezYXXXQR999/P7/61a+Ouv+DDz5IIBDg+eefJzc3F4BRo0bx6U9/mhdeeIEbbrgBgPr6ev7xj39w00038f3vfx+A66+/nptvvpmf//znXHzxxdjt9j48UxERkf7vw8Rtp9PJE088wbhx4w689rGPfYy8vDzuv/9+li1bxpQpUwAIBoP85je/Yfr06fz+978/sK1pmjzwwAPccMMNJCUl9eGZioiI9H8f9n77Pc3Nzfzxj3/k9ttvPxCX30/32yInTj3MRU5j999/Pz//+c8BmDFjBkOGDGHIkCFUVVWdlOO/8cYbpKenM2vWrAOvpaamMnv2bObPn39IL/H/9uabbzJ9+vQDyXKAKVOmUFxczGuvvXbgtbfeeotIJMInPvGJA68ZhsHHP/5x6urqWLduXS+elYiISN/oz3Hb5XIdlCx/z8yZM4GeKdbes2LFCtra2g6K2wA33XQTfr+fhQsXfsgzERER6Xv9OW6/3y9/+UtKSkq44oorDvu+7rdFTpx6mIucxmbOnElFRQVz5szh29/+NikpKUBPED2SQCBAIBA4Ztl2u/2Yvb+2bdvGsGHDsNkOfrY2cuRInnrqKfbu3cuQIUMOu299fT3Nzc2MGDHikPdGjRrF4sWLDzqOz+djwIABh2z33vsTJkw45jmJiIicSv05bh9JU1MTwIFzAdi6dSvAITF++PDh2Gw2tm3bxpVXXnlCxxERETnZzoS4vXHjRl588UUef/xxDMM44nF0vy1yYpQwFzmNDR06lGHDhjFnzhwuvPBC8vPzj7nP3/72N/7whz8cc7u8vDzefvvto27T2Nh42MCZmZkJQENDwxEDeENDAwAZGRmHvJeRkUFbWxvhcBiXy0VjYyNpaWmHBPj39n2vLBERkdNZf47bR6tffHw806ZNO+g4drudtLS0g7Z1uVwkJycrbouISL/Q3+O2ZVn85Cc/4ZJLLmHs2LFH7Bmv+22RE6eEucgZ5qqrrmL8+PHH3M7tdh9zm2AwiMvlOuT1914LhUJH3Pe99w63/3vHfq/8Ix3n/duJiIiciU6XuH04Dz74IO+++y4/+MEPSExMPOg4TqfziPVU3BYRkTPV6RS3n3/+eXbu3HnYecuP5zi63xY5MiXMRc4wBQUFFBQU9EpZHo/nsPOmvffa0b4EvPfe4fZ/L/B7PJ6jHue/txMRETnTnC5x+7/NnTuX3/72t1x33XWHzFXu8XiIRCKH3S8UCilui4jIGet0idtdXV38+te/5rbbbiMnJ+cDHUf32yJHpoS5yBmmu7sbv99/zO3sdvtR52aDniFajY2Nh7z+3pCt94aKHc577x1u/8bGRpKTkw885c7IyGDFihVYlnXQMLH39j3acURERPqz0yVuv9/SpUv55je/yfTp0/nRj3502OPEYjGam5sPmpYlHA7T1tamuC0iImes0yVuP/TQQ0QiES655JIDU7HU1dUB0NHRQVVVFZmZmbhcLt1vi3wASpiLnOaOtHDHkfz973/vtTnVhg4dypo1azBN86CFSDZu3IjX66WkpOSI+2ZlZZGamsrmzZsPeW/jxo0MHTr0wP+XlZXxzDPPsGfPHgYOHHjg9Q0bNhx4X0REpD/or3H7PRs2bOCLX/wiI0aM4Le//S0Ox6G3C+/F5c2bN3PeeecdeH3z5s2YpnlQjBcRETmd9de4XVtbS3t7O5deeukh7z344IM8+OCDvPjii5SVlel+W+QDUMJc5DTn9XoB6OzsPK7te3NOtYsvvpg33niDN998k4svvhiAlpYWXn/9dc4///yD5kGrrKwEoLCw8MBrs2bN4sUXX6S2tvbAMLFly5ZRUVHBpz71qQPbzZgxg5/97Gc8/vjjfP/73wd6FjB58sknycrKYuzYscc+cRERkdNAf47be/bs4Y477iAvL48///nPRxyiPXnyZJKTk3niiScOSpg/8cQTeL1epk+ffsy6ioiInA76a9y+5ZZbuPDCCw8qr7m5me9///tcc801zJgx48AiprrfFjlxSpiLnOaGDx8OwG9+8xsuueQSnE4n559/Pj6f77Db9+acahdddBFjxozh29/+Nrt37yYlJYUnnniCWCzGl770pYO2fS8B/v6n6HfddRevv/46t956K7feeit+v5+HHnqIwYMHc+211x7YLjs7m1tvvZWHHnqIaDTKyJEjeeutt1i9ejW//OUvsdvtvXI+IiIifa2/xu2uri5uu+02Ojo6uO2221i4cOFB2xcWFh64ofZ4PHz5y1/mxz/+MV/+8pc599xzWb16NS+//DJ33303ycnJvXI+IiIifa2/xu3hw4cfqPt73puaZeDAgQcl03W/LXLiDMuyrFNdCRE5uj/96U88+eSTNDY2Ypom8+fPP/C0uK+1t7fz85//nLfeeotQKMTIkSP55je/yciRIw/a7oILLgA4ZNjZrl27uO+++1izZg1Op5PzzjuPe+65h/T09IO2M02Tv/71rzz11FM0NDRQXFzMHXfcwRVXXNG3JygiItLL+mPcrqqqYsaMGUcs9+qrr+a+++476LWnn36av//971RVVZGTk8NNN93EJz/5yRMe3i4iInIq9ce4fTjvxfJvfvOb3HbbbQe9p/ttkROjhLmIiIiIiIiIiIiICGA79iYiIiIiIiIiIiIiImc+JcxFRERERERERERERFDCXEREREREREREREQEUMJcRERERERERERERARQwlxEREREREREREREBADHqa7ABzFhwgTC4TAZGRmnuioiIiK9prGxEZfLxerVq091VXqV4raIiJyJFLdFRET6jxOJ2/2yh3koFCIajZ7qahw3y7IIhUJYlnWqq3JGUvv2LbVv31L79q3+1r7RaJRQKHSqq9HrFLfl/dS+fUvt27fUvn2rv7Wv4vbpob9dN/2N2rdvqX37ltq3b/W39j2RuN0ve5hnZmYCMH/+/FNck+Pj9/vZtm0bZWVl+Hy+U12dM47at2+pffuW2rdv9bf2nTFjxqmuQp9Q3Jb3U/v2LbVv31L79q3+1r6K26eH/nbd9Ddq376l9u1bat++1d/a90Tidr/sYS4iIiIiIiIiIiIi0tuUMBcRERERERERERERQQlzERERERERERERERFACXMREREREREREREREUAJcxERERERERERERERQAlzERERERERERERERFACXMREREREREREREREUAJcxERERERERERERERQAlzERERERERERERERFACXMREREREREREREREUAJcxERERERERERERERQAlzERERERERERERERFACXMREREREREREREREUAJcxERERERERERERERQAlzERERERERERERERFACXMREREREREREREREUAJcxERERERERERERERQAlzERERERERERERERFACXMREREREREREREREUAJcxERERERERERERERQAlzERERERERERERERFACXMREREREREREREREUAJcxERERERERERERERQAlzERERERERERERERFACXMREREREREREREREUAJcxERERERERERERERQAlzERERERERERERERFACXMREREREREREREREUAJcxERERERERERERERQAlzERERERERERERERFACXMREREREREREREREUAJcxERERERERERERER4BQkzLu7u5k2bRpDhgxh06ZNJ/vwIiIicgIUt0VERPoPxW0REZEP76QnzP/0pz8Ri8VO9mFFRETkA1DcFhER6T8Ut0VERD68k5ow37NnD48//jhf+tKXTuZhRURE5ANQ3BYREek/FLdFRER6x0lNmN97773ceOONlJSUnMzDioiIyAeguC0iItJ/KG6LiIj0DsfJOtDrr7/Ozp07uf/++9myZcsxt58xY8YR36utrSU7Oxu/39+bVewzgUDgoLine truncated
            "text/plain": [
              "<Figure size 1500x1000 with 6 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
@@ -457,7 +494,6 @@
        "\n",
        "# Determine indices in the trajectory corresponding to these times.\n",
        "# We assume time_points is a 1D tensor of size n_steps.\n",
        "# Convert time_points to numpy for easy indexing.\n",
        "time_np = time_points.detach().cpu().numpy()\n",
        "n_steps = len(time_np)\n",
        "\n",
@@ -504,38 +540,29 @@
        "plt.tight_layout(rect=[0, 0.03, 1, 0.95])\n",
        "plt.show()\n",
        "\n"
-      ],
+      ]
      "id": "5d381b5a8e3514e1",
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 1500x1000 with 6 Axes>"
            ],
            "image/png": "iVBORw0KGgoAAAANSUhEUgAABcwAAAOHCAYAAAAaAgCCAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjAsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvlHJYcgAAAAlwSFlzAAAPYQAAD2EBqD+naQABAABJREFUeJzs3XeYHWXd//H3nH7O9t5r6qY3kpBACIQEQq+CUiwgYBd9VHzsyk957IoKFhQQ6Z0QSggpJKT33jabzfZeTz8zvz9WIjEddpNs+Ly4cu3FOTP33HPv7Pme+c5dDMuyLEREREREREREREREPuJsp7oCIiIiIiIiIiIiIiKnAyXMRURERERERERERERQwlxEREREREREREREBFDCXEREREREREREREQEUMJcRERERERERERERARQwlxEREREREREREREBFDCXEREREREREREREQEUMJcRERERERERERERARQwlxEREREREREREREBFDCXEREREREREREREQEUMJc5LS3du1a7r//fjo6Ok7J8Ts6Ovje977H5MmTGTNmDLfccgtbtmw57v337NnDbbfdxtixY5k4cSLf+MY3aGlpOWQ70zT561//ygUXXMDIkSO5/PLLmTNnTm+eioiISJ/rr3HbNE2ef/557rrrLs477zzGjBnDZZddxp/+9CdCodBh93nmmWeYPXs2I0eOZNasWfzzn//s7dMRERHpU/01bv+3SCTCJZdcwpAhQ3jooYcOeV/32yInRglzkdPcunXr+MMf/nBKArhpmtxxxx3MmTOHm2+++UCy+5ZbbqGiouKY+9fV1XHTTTdRWVnJ3XffzWc+8xkWLVrEpz/9acLh8EHb/uY3v+GXv/wlU6dO5Xvf+x65ubl8/etf59VXX+2jsxMREel9/TVuBwIBvv3tb9Pa2sqNN97I//7v/zJy5Ejuv/9+br/9dizLOmj7J598ku9+97sMGjSI733ve4wZM4Z7772Xv/zlL314hiIiIr2rv8bt//bYY49RW1t7xPd1vy1yYhynugIicvp6/fXXWbduHb/73e+4+OKLAZg9ezYXXXQR999/P7/61a+Ouv+DDz5IIBDg+eefJzc3F4BRo0bx6U9/mhdeeIEbbrgBgPr6ev7xj39w00038f3vfx+A66+/nptvvpmf//znXHzxxdjt9j48UxERkf7vw8Rtp9PJE088wbhx4w689rGPfYy8vDzuv/9+li1bxpQpUwAIBoP85je/Yfr06fz+978/sK1pmjzwwAPccMMNJCUl9eGZioiI9H8f9n77Pc3Nzfzxj3/k9ttvPxCX30/32yInTj3MRU5j999/Pz//+c8BmDFjBkOGDGHIkCFUVVWdlOO/8cYbpKenM2vWrAOvpaamMnv2bObPn39IL/H/9uabbzJ9+vQDyXKAKVOmUFxczGuvvXbgtbfeeotIJMInPvGJA68ZhsHHP/5x6urqWLduXS+elYiISN/oz3Hb5XIdlCx/z8yZM4GeKdbes2LFCtra2g6K2wA33XQTfr+fhQsXfsgzERER6Xv9OW6/3y9/+UtKSkq44oorDvu+7rdFTpx6mIucxmbOnElFRQVz5szh29/+NikpKUBPED2SQCBAIBA4Ztl2u/2Yvb+2bdvGsGHDsNkOfrY2cuRInnrqKfbu3cuQIUMOu299fT3Nzc2MGDHikPdGjRrF4sWLDzqOz+djwIABh2z33vsTJkw45jmJiIicSv05bh9JU1MTwIFzAdi6dSvAITF++PDh2Gw2tm3bxpVXXnlCxxERETnZzoS4vXHjRl588UUef/xxDMM44nF0vy1yYpQwFzmNDR06lGHDhjFnzhwuvPBC8vPzj7nP3/72N/7whz8cc7u8vDzefvvto27T2Nh42MCZmZkJQENDwxEDeENDAwAZGRmHvJeRkUFbWxvhcBiXy0VjYyNpaWmHBPj39n2vLBERkdNZf47bR6tffHw806ZNO+g4drudtLS0g7Z1uVwkJycrbouISL/Q3+O2ZVn85Cc/4ZJLLmHs2LFH7Bmv+22RE6eEucgZ5qqrrmL8+PHH3M7tdh9zm2AwiMvlOuT1914LhUJH3Pe99w63/3vHfq/8Ix3n/duJiIiciU6XuH04Dz74IO+++y4/+MEPSExMPOg4TqfziPVU3BYRkTPV6RS3n3/+eXbu3HnYecuP5zi63xY5MiXMRc4wBQUFFBQU9EpZHo/nsPOmvffa0b4EvPfe4fZ/L/B7PJ6jHue/txMRETnTnC5x+7/NnTuX3/72t1x33XWHzFXu8XiIRCKH3S8UCilui4jIGet0idtdXV38+te/5rbbbiMnJ+cDHUf32yJHpoS5yBmmu7sbv99/zO3sdvtR52aDniFajY2Nh7z+3pCt94aKHc577x1u/8bGRpKTkw885c7IyGDFihVYlnXQMLH39j3acURERPqz0yVuv9/SpUv55je/yfTp0/nRj3502OPEYjGam5sPmpYlHA7T1tamuC0iImes0yVuP/TQQ0QiES655JIDU7HU1dUB0NHRQVVVFZmZmbhcLt1vi3wASpiLnOaOtHDHkfz973/vtTnVhg4dypo1azBN86CFSDZu3IjX66WkpOSI+2ZlZZGamsrmzZsPeW/jxo0MHTr0wP+XlZXxzDPPsGfPHgYOHHjg9Q0bNhx4X0REpD/or3H7PRs2bOCLX/wiI0aM4Le//S0Ox6G3C+/F5c2bN3PeeecdeH3z5s2YpnlQjBcRETmd9de4XVtbS3t7O5deeukh7z344IM8+OCDvPjii5SVlel+W+QDUMJc5DTn9XoB6OzsPK7te3NOtYsvvpg33niDN998k4svvhiAlpYWXn/9dc4///yD5kGrrKwEoLCw8MBrs2bN4sUXX6S2tvbAMLFly5ZRUVHBpz71qQPbzZgxg5/97Gc8/vjjfP/73wd6FjB58sknycrKYuzYscc+cRERkdNAf47be/bs4Y477iAvL48///nPRxyiPXnyZJKTk3niiScOSpg/8cQTeL1epk+ffsy6ioiInA76a9y+5ZZbuPDCCw8qr7m5me9///tcc801zJgx48AiprrfFjlxSpiLnOaGDx8OwG9+8xsuueQSnE4n559/Pj6f77Db9+acahdddBFjxozh29/+Nrt37yYlJYUnnniCWCzGl770pYO2fS8B/v6n6HfddRevv/46t956K7feeit+v5+HHnqIwYMHc+211x7YLjs7m1tvvZWHHnqIaDTKyJEjeeutt1i9ejW//OUvsdvtvXI+IiIifa2/xu2uri5uu+02Ojo6uO2221i4cOFB2xcWFh64ofZ4PHz5y1/mxz/+MV/+8pc599xzWb16NS+//DJ33303ycnJvXI+IiIifa2/xu3hw4cfqPt73puaZeDAgQcl03W/LXLiDMuyrFNdCRE5uj/96U88+eSTNDY2Ypom8+fPP/C0uK+1t7fz85//nLfeeotQKMTIkSP55je/yciRIw/a7oILLgA4ZNjZrl27uO+++1izZg1Op5PzzjuPe+65h/T09IO2M02Tv/71rzz11FM0NDRQXFzMHXfcwRVXXNG3JygiItLL+mPcrqqqYsaMGUcs9+qrr+a+++476LWnn36av//971RVVZGTk8NNN93EJz/5yRMe3i4iInIq9ce4fTjvxfJvfvOb3HbbbQe9p/ttkROjhLmIiIiIiIiIiIiICGA79iYiIiIiIiIiIiIiImc+JcxFRERERERERERERFDCXEREREREREREREQEUMJcRERERERERERERARQwlxEREREREREREREBADHqa7ABzFhwgTC4TAZGRmnuioiIiK9prGxEZfLxerVq091VXqV4raIiJyJFLdFRET6jxOJ2/2yh3koFCIajZ7qahw3y7IIhUJYlnWqq3JGUvv2LbVv31L79q3+1r7RaJRQKHSqq9HrFLfl/dS+fUvt27fUvn2rv7Wv4vbpob9dN/2N2rdvqX37ltq3b/W39j2RuN0ve5hnZmYCMH/+/FNck+Pj9/vZtm0bZWVl+Hy+U12dM47at2+pffuW2rdv9bf2nTFjxqmuQp9Q3Jb3U/v2LbVv31L79q3+1r6K26eH/nbd9Ddq376l9u1bat++1d/a90Tidr/sYS4iIiIiIiIiIiIi0tuUMBcRERERERERERERQQlzERERERERERERERFACXMREREREREREREREUAJcxERERERERERERERQAlzERERERERERERERFACXMREREREREREREREUAJcxERERERERERERERQAlzERERERERERERERFACXMREREREREREREREUAJcxERERERERERERERQAlzERERERERERERERFACXMREREREREREREREUAJcxERERERERERERERQAlzERERERERERERERFACXMREREREREREREREUAJcxERERERERERERERQAlzERERERERERERERFACXMREREREREREREREUAJcxERERERERERERERQAlzERERERERERERERFACXMREREREREREREREUAJcxERERERERERERERQAlzERERERERERERERFACXMREREREREREREREUAJcxERERERERERERERQAlzERERERERERERERFACXMREREREREREREREUAJcxERERERERERERERQAlzERERERERERERERFACXMREREREREREREREUAJcxERERERERERERER4BQkzLu7u5k2bRpDhgxh06ZNJ/vwIiIicgIUt0VERPoPxW0REZEP76QnzP/0pz8Ri8VO9mFFRETkA1DcFhER6T8Ut0VERD68k5ow37NnD48//jhf+tKXTuZhRURE5ANQ3BYREek/FLdFRER6x0lNmN97773ceOONlJSUnMzDioiIyAeguC0iItJ/KG6LiIj0DsfJOtDrr7/Ozp07uf/++9myZcsxt58xY8YR36utrSU7Oxu/39+bVewzgUDgoLine truncated
          },
          "metadata": {}
        }
      ],
      "execution_count": 8
    },
    {
      "cell_type": "markdown",
      "id": "4255a24220b673d4",
      "metadata": {
        "id": "4255a24220b673d4"
      },
      "cell_type": "markdown",
      "source": [
        "This looks just like it should - both density peaks are well represented, and with _flow matching_, we now have a method at hand that yields accurate samples from the posterior with a (relatively) mild computational overhead in comparison to a deterministic model. 👍\n",
        "\n",
        "## Summary\n",
        "\n",
-        "The flow matching algorithm is an important milestone in the field of diffusion mdoels, and it concludes our trip through the history of generative modeling approaches in deep learning.\n",
+        "The **flow matching** algorithm is an important milestone in the field of diffusion models, and it concludes our trip through the history of generative modeling approaches in deep learning. Interestingly, after they were proposed, denoising, flow matching, and other variants were shown to be more similar than one would think based on the derivation. Nonetheless, they're easier to understand via their respective views on the problem, rather than from a more generic mathematical framework. Nonetheless, this field is a very exciting. We can recommend browsing recent developments online!\n",
        "\n",
-        "As such, it's timem to show what these methods can do with specific examples. The following notebook will do just that. We'll start by comparing denoising diffusion models with flow matching models for a RANS airfoil task with uncertainty."
+        "Nonetheless, it's time to show what these methods can do with specific and more complex examples than our GM toy problem. The following notebook will do just that. We'll start by comparing denoising diffusion models with flow matching models for a RANS airfoil task with uncertainty."
-      ],
+      ]
      "id": "4255a24220b673d4"
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "diffusion",
      "language": "python",
@@ -552,11 +579,8 @@
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython2",
      "version": "2.7.6"
    },
    "colab": {
      "provenance": []
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}
@@ -750,14 +750,21 @@
        "\n",
        "The way in which we can achieve this are _Neural ODEs_{cite}`chen2019node`.\n",
        "They're especially interesting in the context of physics simulations. Mathematically speaking, they replace the mapping $g_i$ with a learned velocity predictor $g_\\theta$. This means\n",
-        "$$\\frac{\\partial z(t)}{\\partial t} = g_\\theta(z(t), t).$$\n",
+        "\n",
        "$$\n",
        "\\frac{\\partial z(t)}{\\partial t} = g_\\theta(z(t), t).\n",
        "$$\n",
        "\n",
        "Then the sequence of $g_i$ steps that made up $g$  can be replaced by integrating the velocity, which gives a simple ODE integration.\n",
        "The continuous time axis is introduced, with $t=0$ starting at a normal Gaussian distribution, to $t=1$ for the target distribution.\n",
        "The ODE is solved along this timeline, e.g. to transform the base distribution $p_Z$ into the target $p_Y$ by querying $g$ for a velocity at each time point along the way.\n",
        "Even better, for an ODE solve there's an analytic formulation for the gradient of the backpropagation path. This is a neat example of a differentiable physics solver (the ODE solve), providing an efficient way to compute a gradient, and aligns with the topics discussed in {doc}`physgrad`.\n",
        "\n",
        "The change in probabilities over time can also be computed conveniently via the trace of the learned function:\n",
-        "$$\\frac{\\partial \\log p(z(t))}{\\partial t} = - \\mathrm{Tr} \\left( \\frac{\\partial g_\\theta}{\\partial z(t)} \\right).$$\n",
+        "\n",
        "$$\n",
        "\\frac{\\partial \\log p(z(t))}{\\partial t} = - \\mathrm{Tr} \\left( \\frac{\\partial g_\\theta}{\\partial z(t)} \\right).\n",
        "$$\n",
        "\n",
        "Compared to the Normalizing Flows above, an important difference in the NeuralODE picture is that now we have a single function $g_\\theta(\\cdot,t)$ that is repeatedly evaluated at different points in the time interval $t \\in [0,1]$. This might seem like a trivial change at first, but it's a crucial step towards more powerful probabilistic models such as diffusion models.\n",
        "It turns out to be important that we can re-use a learned function, instead of having to manually construct many layers with large numbers of trainable parameters.\n",
@@ -7,12 +7,14 @@
        "id": "718ff8c0ea48fa63"
      },
      "source": [
-        "## Score Matching\n",
+        "# Score Matching\n",
        "\n",
        "A first important step is to realize that there's a convenient alternative to probability densities that let's us work with unnormalized functions: the so-called _score_. This is the name that was established for the gradient of the log likelihood function: $\\nabla_x \\log p(x)$. When we learn this function with a neural network, which we'll call $s_\\theta(x)$ to indicate this is a learned _score_, the great thing is that we don't need to worry about normalization anymore. When integrating we'd need to have the right constant offset, but for the \"local\" gradients at $x$, all that matters is the derivative at this point.\n",
        "\n",
        "\n",
-        "### Gaussian Toy Dataset with Analytic Scores"
+        "## Gaussian Toy Dataset with Analytic Scores\n",
        "\n",
        "We'll re-use the same Gaussian mixture case from the previous sections, but for experiments with the score function we'll expand the base class with a `score()` function so that we can check how well learned approximations fare. Luckily, we can compute the reference score directly from the Gaussian mixture model."
      ]
    },
    {
@@ -106,12 +108,14 @@
        "id": "6273c841be6767d5"
      },
      "source": [
-        "### Visualize Samples, Likelihoods and Scores"
+        "### Visualizing Samples, Likelihoods and Scores\n",
        "\n",
        "We'll also directly define a helper function that visualizes samples, the score field and the likelihood in a single graph:"
      ]
    },
    {
      "cell_type": "code",
-      "execution_count": 5,
+      "execution_count": null,
      "id": "1e646ecf577abbf2",
      "metadata": {
        "ExecuteTime": {
@@ -127,14 +131,14 @@
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAABcwAAAHkCAYAAAAD/WxfAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjAsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvlHJYcgAAAAlwSFlzAAAPYQAAD2EBqD+naQABAABJREFUeJzs3XV0FOfbxvHvxgMhQPAAwYK7OxQvXtzd3YtbcSlQ3N3dtbgUD65BAgQnBALx7L5/5J35ZaMzCwXa3p9zcg47u8+zs5YN99xzPQaTyWRCCCGEEEIIIYQQQgghhPiPs/reOyCEEEIIIYQQQgghhBBC/AikYC6EEEIIIYQQQgghhBBCIAVzIYQQQgghhBBCCCGEEAKQgrkQQgghhBBCCCGEEEIIAUjBXAghhBBCCCGEEEIIIYQApGAuhBBCCCGEEEIIIYQQQgBSMBdCCCGEEEIIIYQQQgghACmYCyGEEEIIIYQQQgghhBCAFMyFEEIIIYQQQgghhBBCCEAK5kKIOGTNmpVZs2Z9790QQggh/jW2bt1K1qxZefbsme6xs2bNImvWrJpuK9/hQgghhBBC6GfzvXdAiP+Ku3fvMmfOHK5fv87bt29JlCgR7u7ulC9fnhYtWnzv3RNCCCG+q3/T9+SgQYPYtm1btNctWrToG++NEEIIIYQQQg8pmAvxDVy+fJmWLVvi6upKgwYNSJYsGS9evODq1ausXLnyH1cIEEIIIb6mf+P3pJ2dHWPHjo2yPVu2bCRJkoTq1atjZ2f3HfZMCCGEEEIIERspmAvxDcyfP58ECRKwefNmnJ2dza579+7dd9orIYQQ4sfwI31PBgQE4Ojo+MXz2NjYULt27Rivt7a2/uL7EEIIIYQQQnx9kmEuxDfw5MkT3N3doxQBAJIkSaL+e8uWLbRs2ZLixYuTK1cuqlWrxtq1a6OMKV++PJ06deLcuXPUrVuXPHnyULNmTc6dOwfAwYMHqVmzJrlz56Zu3brcunXLbPygQYPInz8/T58+pV27duTLl49SpUoxe/ZsTCZTnI/n1atXDB48mBIlSpArVy6qV6/O5s2bo9xu1apVVK9enbx581K4cGHq1q3Lrl274pxfCCHEf4vW70nFjh07qF+/vvr90qxZM06dOmV2mzVr1lC9enVy5cpFqVKlGD16NB8/fjS7TYsWLahRowY3btygWbNm5M2bl2nTpgEQHBzMzJkzqVSpErly5aJs2bJMnjyZ4ODgL368MWWYHz9+nKZNm5IvXz7y589Px44duX//fpzzBQcHM378eIoVK0b+/Pnp3LkzL1++/OL9FEIIIYQQ4r9IOsyF+AZSp06Nh4cH9+7dI0uWLDHebt26dWTOnJny5ctjY2PD0aNHGT16NCaTiWbNmpnd1svLi379+tG4cWNq1arF0qVL6dy5M6NHj2b69Ok0adIEgIULF9K7d2/279+PldX/jpGFhYXRvn178ubNy4ABAzh58iSzZs0iLCyMXr16xbiPb9++pWHDhhgMBpo1a4aLiwsnTpxg6NChfPr0idatWwOwceNGxo4dS5UqVWjZsiVBQUHcvXuXq1evUrNmzS94NoUQQvzbaP2eBJg9ezazZs0if/789OzZE1tbW65evcrZs2cpVaoUEL4w5uzZsylRogRNmjTh0aNHrFu3juvXr7Nu3TpsbW3V+Xx9fenQoQPVq1enVq1aJEmSBKPRSJcuXbh06RINGzYkU6ZM3Lt3jxUrVvD48WPmzp2r6XH5+PiYXba1tSVBggTR3nb79u0MGjSIUqVK0b9/fwICAli3bh1NmzZl27ZtpEmTJsb7GTp0KDt37qRGjRoUKFCAs2fP0rFjR037KIQQQgghhDAnBXMhvoG2bdvSoUMHfvnlF/LkyUPBggUpXrw4RYsWNftP++rVq3FwcFAvN2/enHbt2rFs2bIoBfNHjx6xfv168ufPD4C7uzvt2rVj+PDh7Nu3D1dXVwASJkzIiBEjuHDhAkWLFlXHBwUFUbp0aYYNGwZA06ZN6dy5M4sWLaJFixa4uLhE+1imT59OWFgYu3btInHixAA0adKEvn37Mnv2bBo3boyDgwPHjh0jc+bMzJw58ys8g0IIIf7NtH5Penl5MWfOHCpVqsTMmTPNDgQrZ0j5+PiwYMECSpUqxaJFi9TbZMyYkd9++42dO3dSr149ddybN28YPXo0jRs3Vrft2LGDM2fOsGrVKgoVKqRuz5w5MyNHjuTy5csUKFAg1sfk7+9P8eLFzbYVKVKEVatWRbnt58+fGTduHA0aNGDMmDHq9jp16vDzzz+zYMECs+0R3blzh507d9K0aVNGjhwJQLNmzejXrx93796NdR+FEEIIIYQQUUkkixDfQMmSJVm/fj3ly5fnzp07LF68mHbt2lGmTBkOHz6s3i5isdzPzw8fHx+KFCnC06dP8fPzM5vT3d1dLZYD5M2bF4BixYqpxfKI258+fRplvyIW4ZWO8ZCQEP76669oH4fJZOLgwYOUL18ek8mEj4+P+lOqVCn8/Py4efMmAM7Ozrx8+ZJr165pfp6EEEL8N2n9nvzzzz8xGo1069bNrFgO4d9jAGfOnCEkJISWLVua3aZBgwY4OTlx/Phxs3F2dnbUrVvXbNv+/fvJlCkTGTNmNPuuK1asGIAagRYbe3t7li1bZvYzcODAaG975swZPn78SPXq1c3uz8rKirx588Z6f8rjibwwaqtWreLcRyGEEEIIIURU0mEuxDeSJ08eZs+eTXBwMHfu3OHPP/9k+fLl9OrVi+3bt+Pu7s6lS5eYNWsWV65cISAgwGy8n5+f2WncqVKlMrteuS5lypRm252cnACi5LZaWVmRNm1as20ZMmQAwNvbO9rH4OPjw8ePH9mwYQMbNmyI8TYAHTp04MyZMzRo0IB06dJRsmRJatSoQcGCBaMdJ4QQ4r9Ny/fkkydPsLKyIlOmTDHO8/z5cyC8ozwiOzs70qZNG+U7LkWKFNjZ2Zlt8/Ly4sGDB1E6xBVaFiK1tramRIkScd4O4PHjx0DMRW7luzw63t7eWFlZ4ebmZrY98uMXQgghhBBCaCMFcyG+MTs7O/LkyUOePHlInz49gwcPZv/+/dSqVYvWrVuTMWNGBg0aRKpUqbC1teX48eMsX74co9FoNo+1tXW088e0XctinnFR9qFWrVrUqVMn2ttkzZoVgEyZMrF//36OHTvGyZMnOXjwIGvXrqVbt2707Nnzi/dFCCHEv1NM35Pdu3f/W+4v4tldCqPRSJYsWRg8eHC0YyIfnP5Synf05MmTSZYsWZTrY/puF0IIIYQQQnx9UjAX4jvKlSsXAK9fv+bIkSMEBwczb948s0gVLad9W8JoNPL06VO1qxzCc9EhfPG16Li4uBA/fnyMRqOmrrl48eJRrVo1qlWrRnBwMD169GD+/Pl06tQJe3v7r/NAhBBC/GtF/J4EcHNzw2g08uDBA7Jnzx7tGOU79OHDh2ZnUgUHB/Ps2TNN319ubm7cuXOH4sWLq1EvfydlP5MkSaK5K12ROnVqjEYjT548Mesqf/jw4VfdRyGEEEIIIf4rJMNciG/g7Nmz0XZ4K7mjGTNmVLvHIt7Oz8+PLVu2/G37tWbNGvXfJpOJNWvWYGtrG+Mp6NbW1lSpUoUDBw5w7969KNcrcSwA79+/N7vOzs6OTJkyYTKZCAkJ+UqPQAghxL+Blu9JgIoVK2JlZcWcOXOinHmljC9RogS2trasWrXKbM7Nmzfj5+dH2bJl49yfqlWr8urVKzZu3BjlusDAQPz9/bU/OA1Kly6Nk5MTCxYsiPY7MuL3a2RlypQBiLKY6IoVK77qPgohhBBCCPFfIR3mQnwDY8eOJSAggEqVKpExY0ZCQkK4fPky+/btI3Xq1NStW5e3b99ia2tL586dady4MZ8/f2bTpk0kSZKEN2/efPV9sre35+TJkwwcOJA8efJw8uRJjh07RufOnXFxcYlxXL9+/Th37hwNGzakQYMGuLu78+HDB27evMlff/3F+fPnAWjXrh1JkyalQIECJEmShIcPH7J69WrKli0baxarEEKI/x4t35MA6dKlo3PnzsydO5emTZtSuXJl7OzsuH79OsmTJ6dfv364uLjQqVMnZs+eTfv27SlfvjyPHj1i7dq15M6dm1q1asW5P7Vr12bfvn2MHDmSc+fOUaBAAcLCwnj48CH79+9n8eLF5M6d+6s9ficnJ0aNGsWvv/5K3bp1qVatGi4uLjx//pzjx49ToEABRowYEe3Y7NmzU6NGDdauXYufnx/58+fn7NmzeHl5fbX9E0IIIYQQ4r9ECuZCfAO//vor+/fv5/jx42zYsIGQkBBcXV1p2rQpXbp0wdnZGWdnZ2bOnMmMGTOYNGkSSZMmpUmTJri4uDBkyJCvvk/W1tYsXryYUaNGMWXKFOLHj0/37t3p1q1brOOSJk3Kpk2bmDNnDocOHWLdunUkSpQId3d3+vfvr96uUaNG7Nq1i2XLluHv70/KlClp0aIFXbt2/eqPRQghxD+blu9JRa9evUiTJg2rV69m+vTpODo6kjVrVmrXrq3epkePHri4uLB69WomTJhAwoQJadiwIX379sXW1jbO/VG62JcvX86OHTs4dOgQjo6OpEmThhYtWpjFmX0tNWvWJHny5CxcuJAlS5YQHBxMihQpKFSokHrAICbjx48nceLE7Nq1i8OHD1O0aFEWLlyoqZteCCGEEEIIYc5g+horAQoh/lEGDRrEgQMH8PDw+N67IoQQQgghhBBCCCHED0MyzIUQQgghhBBCCCGEEEIIpGAuhBBCCCGEEEIIIYQQQgBSMBdCCCGEEEIIIYQQQgghAMkwF0IIIYQQQgghhBBCCCEA6TAXQgghhBBCCCGEEEIIIQApmAshhBBCCCGEEOI/7tmzZ2TNmpWtW7eq2wYNGkT+/Pn/tvlnzZpF1qxZzW6XNWtWfvvtt69yn1/DuXPnyJo1K+fOnfveuLine truncated
            "text/plain": [
              "<Figure size 1500x600 with 4 Axes>"
-            ],
+            ]
            "image/png": "iVBORw0KGgoAAAANSUhEUgAABcwAAAHkCAYAAAAD/WxfAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjAsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvlHJYcgAAAAlwSFlzAAAPYQAAD2EBqD+naQABAABJREFUeJzs3XV0FOfbxvHvxgMhQPAAwYK7OxQvXtzd3YtbcSlQ3N3dtbgUD65BAgQnBALx7L5/5J35ZaMzCwXa3p9zcg47u8+zs5YN99xzPQaTyWRCCCGEEEIIIYQQQgghhPiPs/reOyCEEEIIIYQQQgghhBBC/AikYC6EEEIIIYQQQgghhBBCIAVzIYQQQgghhBBCCCGEEAKQgrkQQgghhBBCCCGEEEIIAUjBXAghhBBCCCGEEEIIIYQApGAuhBBCCCGEEEIIIYQQQgBSMBdCCCGEEEIIIYQQQgghACmYCyGEEEIIIYQQQgghhBCAFMyFEEIIIYQQQgghhBBCCEAK5kKIOGTNmpVZs2Z9790QQggh/jW2bt1K1qxZefbsme6xs2bNImvWrJpuK9/hQgghhBBC6GfzvXdAiP+Ku3fvMmfOHK5fv87bt29JlCgR7u7ulC9fnhYtWnzv3RNCCCG+q3/T9+SgQYPYtm1btNctWrToG++NEEIIIYQQQg8pmAvxDVy+fJmWLVvi6upKgwYNSJYsGS9evODq1ausXLnyH1cIEEIIIb6mf+P3pJ2dHWPHjo2yPVu2bCRJkoTq1atjZ2f3HfZMCCGEEEIIERspmAvxDcyfP58ECRKwefNmnJ2dza579+7dd9orIYQQ4sfwI31PBgQE4Ojo+MXz2NjYULt27Rivt7a2/uL7EEIIIYQQQnx9kmEuxDfw5MkT3N3doxQBAJIkSaL+e8uWLbRs2ZLixYuTK1cuqlWrxtq1a6OMKV++PJ06deLcuXPUrVuXPHnyULNmTc6dOwfAwYMHqVmzJrlz56Zu3brcunXLbPygQYPInz8/T58+pV27duTLl49SpUoxe/ZsTCZTnI/n1atXDB48mBIlSpArVy6qV6/O5s2bo9xu1apVVK9enbx581K4cGHq1q3Lrl274pxfCCHEf4vW70nFjh07qF+/vvr90qxZM06dOmV2mzVr1lC9enVy5cpFqVKlGD16NB8/fjS7TYsWLahRowY3btygWbNm5M2bl2nTpgEQHBzMzJkzqVSpErly5aJs2bJMnjyZ4ODgL368MWWYHz9+nKZNm5IvXz7y589Px44duX//fpzzBQcHM378eIoVK0b+/Pnp3LkzL1++/OL9FEIIIYQQ4r9IOsyF+AZSp06Nh4cH9+7dI0uWLDHebt26dWTOnJny5ctjY2PD0aNHGT16NCaTiWbNmpnd1svLi379+tG4cWNq1arF0qVL6dy5M6NHj2b69Ok0adIEgIULF9K7d2/279+PldX/jpGFhYXRvn178ubNy4ABAzh58iSzZs0iLCyMXr16xbiPb9++pWHDhhgMBpo1a4aLiwsnTpxg6NChfPr0idatWwOwceNGxo4dS5UqVWjZsiVBQUHcvXuXq1evUrNmzS94NoUQQvzbaP2eBJg9ezazZs0if/789OzZE1tbW65evcrZs2cpVaoUEL4w5uzZsylRogRNmjTh0aNHrFu3juvXr7Nu3TpsbW3V+Xx9fenQoQPVq1enVq1aJEmSBKPRSJcuXbh06RINGzYkU6ZM3Lt3jxUrVvD48WPmzp2r6XH5+PiYXba1tSVBggTR3nb79u0MGjSIUqVK0b9/fwICAli3bh1NmzZl27ZtpEmTJsb7GTp0KDt37qRGjRoUKFCAs2fP0rFjR037KIQQQgghhDAnBXMhvoG2bdvSoUMHfvnlF/LkyUPBggUpXrw4RYsWNftP++rVq3FwcFAvN2/enHbt2rFs2bIoBfNHjx6xfv168ufPD4C7uzvt2rVj+PDh7Nu3D1dXVwASJkzIiBEjuHDhAkWLFlXHBwUFUbp0aYYNGwZA06ZN6dy5M4sWLaJFixa4uLhE+1imT59OWFgYu3btInHixAA0adKEvn37Mnv2bBo3boyDgwPHjh0jc+bMzJw58ys8g0IIIf7NtH5Penl5MWfOHCpVqsTMmTPNDgQrZ0j5+PiwYMECSpUqxaJFi9TbZMyYkd9++42dO3dSr149ddybN28YPXo0jRs3Vrft2LGDM2fOsGrVKgoVKqRuz5w5MyNHjuTy5csUKFAg1sfk7+9P8eLFzbYVKVKEVatWRbnt58+fGTduHA0aNGDMmDHq9jp16vDzzz+zYMECs+0R3blzh507d9K0aVNGjhwJQLNmzejXrx93796NdR+FEEIIIYQQUUkkixDfQMmSJVm/fj3ly5fnzp07LF68mHbt2lGmTBkOHz6s3i5isdzPzw8fHx+KFCnC06dP8fPzM5vT3d1dLZYD5M2bF4BixYqpxfKI258+fRplvyIW4ZWO8ZCQEP76669oH4fJZOLgwYOUL18ek8mEj4+P+lOqVCn8/Py4efMmAM7Ozrx8+ZJr165pfp6EEEL8N2n9nvzzzz8xGo1069bNrFgO4d9jAGfOnCEkJISWLVua3aZBgwY4OTlx/Phxs3F2dnbUrVvXbNv+/fvJlCkTGTNmNPuuK1asGIAagRYbe3t7li1bZvYzcODAaG975swZPn78SPXq1c3uz8rKirx588Z6f8rjibwwaqtWreLcRyGEEEIIIURU0mEuxDeSJ08eZs+eTXBwMHfu3OHPP/9k+fLl9OrVi+3bt+Pu7s6lS5eYNWsWV65cISAgwGy8n5+f2WncqVKlMrteuS5lypRm252cnACi5LZaWVmRNm1as20ZMmQAwNvbO9rH4OPjw8ePH9mwYQMbNmyI8TYAHTp04MyZMzRo0IB06dJRsmRJatSoQcGCBaMdJ4QQ4r9Ny/fkkydPsLKyIlOmTDHO8/z5cyC8ozwiOzs70qZNG+U7LkWKFNjZ2Zlt8/Ly4sGDB1E6xBVaFiK1tramRIkScd4O4PHjx0DMRW7luzw63t7eWFlZ4ebmZrY98uMXQgghhBBCaCMFcyG+MTs7O/LkyUOePHlInz49gwcPZv/+/dSqVYvWrVuTMWNGBg0aRKpUqbC1teX48eMsX74co9FoNo+1tXW088e0XctinnFR9qFWrVrUqVMn2ttkzZoVgEyZMrF//36OHTvGyZMnOXjwIGvXrqVbt2707Nnzi/dFCCHEv1NM35Pdu3f/W+4v4tldCqPRSJYsWRg8eHC0YyIfnP5Synf05MmTSZYsWZTrY/puF0IIIYQQQnx9UjAX4jvKlSsXAK9fv+bIkSMEBwczb948s0gVLad9W8JoNPL06VO1qxzCc9EhfPG16Li4uBA/fnyMRqOmrrl48eJRrVo1qlWrRnBwMD169GD+/Pl06tQJe3v7r/NAhBBC/GtF/J4EcHNzw2g08uDBA7Jnzx7tGOU79OHDh2ZnUgUHB/Ps2TNN319ubm7cuXOH4sWLq1EvfydlP5MkSaK5K12ROnVqjEYjT548Mesqf/jw4VfdRyGEEEIIIf4rJMNciG/g7Nmz0XZ4K7mjGTNmVLvHIt7Oz8+PLVu2/G37tWbNGvXfJpOJNWvWYGtrG+Mp6NbW1lSpUoUDBw5w7969KNcrcSwA79+/N7vOzs6OTJkyYTKZCAkJ+UqPQAghxL+Blu9JgIoVK2JlZcWcOXOinHmljC9RogS2trasWrXKbM7Nmzfj5+dH2bJl49yfqlWr8urVKzZu3BjlusDAQPz9/bU/OA1Kly6Nk5MTCxYsiPY7MuL3a2RlypQBiLKY6IoVK77qPgohhBBCCPFfIR3mQnwDY8eOJSAggEqVKpExY0ZCQkK4fPky+/btI3Xq1NStW5e3b99ia2tL586dady4MZ8/f2bTpk0kSZKEN2/efPV9sre35+TJkwwcOJA8efJw8uRJjh07RufOnXFxcYlxXL9+/Th37hwNGzakQYMGuLu78+HDB27evMlff/3F+fPnAWjXrh1JkyalQIECJEmShIcPH7J69WrKli0baxarEEKI/x4t35MA6dKlo3PnzsydO5emTZtSuXJl7OzsuH79OsmTJ6dfv364uLjQqVMnZs+eTfv27SlfvjyPHj1i7dq15M6dm1q1asW5P7Vr12bfvn2MHDmSc+fOUaBAAcLCwnj48CH79+9n8eLF5M6d+6s9ficnJ0aNGsWvv/5K3bp1qVatGi4uLjx//pzjx49ToEABRowYEe3Y7NmzU6NGDdauXYufnx/58+fn7NmzeHl5fbX9E0IIIYQQ4r9ECuZCfAO//vor+/fv5/jx42zYsIGQkBBcXV1p2rQpXbp0wdnZGWdnZ2bOnMmMGTOYNGkSSZMmpUmTJri4uDBkyJCvvk/W1tYsXryYUaNGMWXKFOLHj0/37t3p1q1brOOSJk3Kpk2bmDNnDocOHWLdunUkSpQId3d3+vfvr96uUaNG7Nq1i2XLluHv70/KlClp0aIFXbt2/eqPRQghxD+blu9JRa9evUiTJg2rV69m+vTpODo6kjVrVmrXrq3epkePHri4uLB69WomTJhAwoQJadiwIX379sXW1jbO/VG62JcvX86OHTs4dOgQjo6OpEmThhYtWpjFmX0tNWvWJHny5CxcuJAlS5YQHBxMihQpKFSokHrAICbjx48nceLE7Nq1i8OHD1O0aFEWLlyoqZteCCGEEEIIYc5g+horAQoh/lEGDRrEgQMH8PDw+N67IoQQQgghhBBCCCHED0MyzIUQQgghhBBCCCGEEEIIpGAuhBBCCCGEEEIIIYQQQgBSMBdCCCGEEEIIIYQQQgghAMkwF0IIIYQQQgghhBBCCCEA6TAXQgghhBBCCCGEEEIIIQApmAshhBBCCCGEEOI/7tmzZ2TNmpWtW7eq2wYNGkT+/Pn/tvlnzZpF1qxZzW6XNWtWfvvtt69yn1/DuXPnyJo1K+fOnfveuLine truncated
          },
-          "metadata": {}
+          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
@@ -214,12 +218,12 @@
        "\n",
        "With normalizing flows, this is always satisfied by the way the model and networks are constructed. However, as we have seen in the previous section, normalizing flows do not scale well to high dimensional data and are computationally expensive.\n",
        "\n",
-        "For alternative methods, making sure that the model $q_\\theta$ integrates to 1 becomes a big difficulty.\n",
+        "For general and flexible methods, making sure that the model $q_\\theta$ integrates to 1 becomes extremely difficult.\n",
        "What makes the score so interesting is that it depends on local information only. That means, because it considers the *gradient* of the log likelihood, its value only depends on the local values of the likelihood. This makes it much easier for a neural network to learn the score.\n",
        "\n",
        "We will first explore how to learn the score from samples of the target distribution using a neural network. Then we will introduce a first method how to use the score to sample from the target distribution.\n",
        "\n",
-        "### Learning the Score"
+        "## Learning the Score"
      ]
    },
    {
@@ -229,7 +233,7 @@
        "id": "6f44ec41228f5c32"
      },
      "source": [
-        "We could directly learn this gradient with an $L^2$ term, which however, in this context is known as the \"Fisher divergence\": $\\mathbb{E}_{x \\sim p(x)} [||\\nabla_x \\log p(x) - s_\\theta(x)||^2]$. However, this would require having access to the ground truth gradients $\\nabla_x \\log p(x)$?, which unfortunately is not the case for relevant settings. The trick established in ML to compute targets for learning is to slightly perturb the dataset with Gaussian noise {cite}`vincent2011dsm`. This turns it from a collection of point-wise samples into a continuous function of which we can compute the gradients.\n",
+        "We could directly learn $\\nabla_x \\log p(x)$, the gradient of the log likelihood function, with an $L^2$ term. In this context, $L^2$ is  known as the \"Fisher divergence\": $\\mathbb{E}_{x \\sim p(x)} [||\\nabla_x \\log p(x) - s_\\theta(x)||^2]$. However, this would require having access to the ground truth gradients $\\nabla_x \\log p(x)$?, which unfortunately is not the case for relevant settings. The trick established in ML to compute targets for learning is to slightly perturb the dataset with Gaussian noise {cite}`vincent2011dsm`. This turns it from a collection of point-wise samples into a continuous function of which we can compute the gradients.\n",
        "\n",
        "So, for the dataset $\\{ x_1, ..., x_n\\}$ we consider the perturbed dataset $\\{ \\tilde{x}_1, ..., \\tilde{x}_n\\}$ by adding Gaussian noise $\\tilde{x} = x + \\sigma z$ with $z \\sim \\mathcal{N}(0, I)$.\n",
        "For now we keep the noise level $\\sigma > 0 $ fixed at a certain value, but this is a crucial parameter that we'll revisit soon.\n",
@@ -270,7 +274,9 @@
        "id": "19630c5d2c4b42f9"
      },
      "source": [
-        "### Define Score Network and Dataset"
+        "### Define Score Network and Dataset\n",
        "\n",
        "To implement these ideas for our Gaussian mixture problem, we first define a simple score network with four hidden layers, and a dataset class that perturbs our samples as described above. The `DenoisingScoreMatchingDataset` class has a parameter `sigma` to control the amount of perturbation."
      ]
    },
    {
@@ -329,7 +335,9 @@
        "id": "be4ba82aeabf26e6"
      },
      "source": [
-        "### Training"
+        "### Training\n",
        "\n",
        "Now we can train the score network to estimate the gradients of the perturbed dataset, which are given by the perturbation in the dataset (`noise`) below, divided by $\\sigma^2$."
      ]
    },
    {
@@ -349,8 +357,8 @@
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 10/100, Loss: 192.8589\n",
            "Epoch 20/100, Loss: 197.8407\n",
@@ -417,7 +425,9 @@
        "id": "2645142b87a208f4"
      },
      "source": [
-        "### Compare Learned and Reference Scores"
+        "### Compare Learned and Reference Scores\n",
        "\n",
        "Now that we have a trained network, let's qualitatively evaluate whether the predicted score is accurate or not. As shown above, we have a ground truth vector field for the score. The next cell plots it next to the one predicted by our trained network."
      ]
    },
    {
@@ -438,14 +448,14 @@
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAA9gAAAH7CAYAAAA3jTfAAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjAsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvlHJYcgAAAAlwSFlzAAAPYQAAD2EBqD+naQABAABJREFUeJzsnXVUVN33xp8ZWgERFXUwUBELG7sLWxRRMRAVTAzs7sDXTmzFwBZbwQALOxG7G0RCpWPO74/5nfPOEErce9H3ez5ruZbOjHdP3nv22c9+towQQsDhcDgcDofD4XA4HA4nR8hz+wlwOBwOh8PhcDgcDofzX4An2BwOh8PhcDgcDofD4QgAT7A5HA6Hw+FwOBwOh8MRAJ5gczgcDofD4XA4HA6HIwA8weZwOBwOh8PhcDgcDkcAeILN4XA4HA6Hw+FwOByOAPAEm8PhcDgcDofD4XA4HAHgCTaHw+FwOBwOh8PhcDgCwBNsDofD4XA4HA6Hw+FwBIAn2BwOh8PhcDgczl+Kk5MTnJyccvtpcDic/0c7t58Ah/Nfply5cpl63I4dO1CnTh2Rn03GKJVKHDt2DN7e3nj37h2SkpJgZmaGqlWrolevXqhWrVquPTcOh8PhcH6Hj48PJk+ejIMHD6Jy5cq5/XT+SBITE7Fnzx4cPnwY79+/h1wuR+HChVGjRg3069cPZcqUye2nyOH8J+AJNocjIosWLdL499GjRxEYGJjm9ty+qM2bNw/e3t5o0aIFOnbsCC0tLbx58waXL19G8eLFeYLN4XA4HM5fzsiRI3Hp0iW0b98e3bp1Q3JyMl6/fo0LFy6gevXqub4W4XD+K/AEm8MRETs7O41/P3jwAIGBgWluT01cXBwMDAzEfGqMb9++Yffu3ejevTvmzp2rcR8hBBEREZI8DwBITk6GUqmErq6uZDE5HA6HwxGKP/U6FhQUhICAAIwePRpDhgzRuC8lJQU/fvyQ7LkkJCRAR0cHcjnvVOX8N+HfbA4nl3FyckKHDh0QHByM3r17o2rVqli2bBkAlcR89erVaf5P8+bNMWnSJI3bfvz4gfnz56NJkyawtrZGq1atsHHjRiiVyl/G//jxIwghqFGjRpr7ZDIZChQokCbOggUL0Lx5c1hbW6Nx48aYMGGCRiIeHh6OKVOmoH79+qhcuTI6deqEw4cPp4lbrlw5bNmyBV5eXmjZsiUqV66MV69eAQBevXqFkSNHonbt2qhcuTLs7e1x/vx5jWMkJSVhzZo1sLW1ReXKlVGnTh307NkTgYGBv3zNHA6Hw/nfJDQ0FJMnT0b9+vVhbW2N9u3b4+DBgxqPSUxMxMqVK2Fvb4+aNWuiWrVq6NWrF65fv67xuF9dx1avXo1y5crh3bt3mDRpEmxsbFCzZk1MnjwZcXFxaZ7X0aNHYW9vjypVqqB27doYPXo0vnz5kuZx+/btQ8uWLVGlShU4ODjg9u3bmXrdHz58AIB0r/VaWlrInz9/mvdpypQpaNiwIaytrdG8eXPMnDkTiYmJGsek1+mqVauie/fuuHDhgsZxbty4gXLlyuHkyZNYvnw5GjVqhKpVqyI6OhqAqvDg4uKCmjVromrVqujTpw/u3LmjcYzo6GjMnz+frTvq1auH/v3749GjR5l67RyO1PAKNofzBxAVFYWBAweiffv26NSpU5qk9nfExcWhT58+CA0NhaOjI4oWLYp79+5h2bJlCAsLw9SpUzP8vwqFAgDg6+uLNm3a/LJyHhMTg969e+PVq1fo2rUrKlasiMjISPj7+yM0NBSmpqaIj4+Hk5MT3r9/j969e6NYsWLw9fXFpEmT8OPHDzg7O2sc08fHBwkJCejevTt0dXWRL18+vHjxAj179kThwoUxcOBA5MmTB6dPn4abmxtWr16NVq1aAQDWrFmDDRs2oFu3bqhSpQqio6MRHByMR48eoUGDBll6DzkcDofz3+bbt2/o3r07ZDIZevfuDVNTU1y6dAlTp05FdHQ0+vXrB0CV0B04cAAdOnRAt27dEBMTg4MHD8LV1RUHDhxAhQoVNI6b3nWM4u7ujmLFimHMmDF4/PgxDhw4AFNTU4wfP549Zt26dVi5ciXatm0LBwcHREREYNeuXejduzeOHDkCY2NjAMCBAwcwY8YMVK9eHc7Ozvjw4QOGDh2KfPnyoWjRor987fRaf/z4cdSoUQPa2hmnAKGhoXBwcMDPnz/RvXt3lC5dGqGhofDz80N8fDx0dXXx7ds3ODo6Ii4uDk5OTsifPz8OHz6MoUOHYtWqVew6TfH09ISOjg5cXFyQmJgIHR0dXLt2DQMHDoS1tTWGDx8OmUwGHx8fODs7Y/fu3ahSpQoAYObMmfDz80OfPn1QpkwZREVF4c6dO3j16hUqVar0m0+dw8kFCIfDkYzZs2cTKysrjdv69OlDrKysyJ49e9I83srKiqxatSrN7c2aNSMTJ05k/167di2pVq0aefPmjcbjlixZQipUqEA+f/78y+c1YcIEYmVlRWrVqkXc3NzIli1byMuXL9M8buXKlcTKyoqcOXMmzX1KpZIQQoiXlxexsrIiR48eZfclJiaSHj16kGrVqpGfP38SQgj58OEDsbKyIjVq1CDh4eEax3J2diYdOnQgCQkJGsfv0aMHsbW1Zbd16tSJDBo06JevjcPhcDj/fQ4dOkSsrKxIUFBQho+ZMmUKadCgAYmIiNC4ffTo0aRmzZokLi6OEEJIcnKyxvWHEEK+f/9O6tevTyZPnsxu+9V1bNWqVcTKykrj8YQQ4ubmRmrXrs3+/fHjR1KhQgWybt06jcc9e/aMVKxYkd2emJhI6tWrR+zs7DSe2759+4iVlRXp06dPhq+bENU1lK436tevT8aMGUN27dpFPn36lOaxEyZMIOXLl0/3vaTX+vnz5xMrKyty69Ytdl90dDRp3rw5adasGUlJSSGEEHL9+nViZWVFWrRowd5fehxbW1syYMAAdkxCCImLiyPNmzcn/fv3Z7fVrFmTzJ49+5evj8P5k+AScQ7nD0BXVxf29vbZ/v++vr6oWbMmjI2NERERwf7Ur18fKSkpuHXr1i//v4eHB2bMmIFixYrh7Nmz+Oeff9CuXTs4OzsjNDSUPe7MmTMoX758mp1pQCUnB4BLly6hUKFC6NChA7tPR0cHTk5OiI2NTfNcbG1tYWpqyv4dFRWF69evo23btoiOjmavJTIyEg0bNsTbt2/ZczI2NsaLFy/w9u3bLL9nHA6Hw/nfgRCCM2fOoHnz5sxfhP5p2LAhfv78ySTHWlparIdaqVQiKioKycnJsLa2xuPHj9McO/V1TB1HR0eNf9vY2CAqKopJpM+ePQulUom2bdtqPKeCBQuiZMmSuHHjBgAgODgY4eHhcHR01Ojv7tKlC4yMjH77+mUyGbZs2QJ3d3cYGxvjxIkTmDNnDpo1awZ3d3fWg61UKnHu3Dk0a9YsXTd2eq2/ePEiqlSpAhsbG3Zf3rx50aNHD3z69AkvX77U+H+dO3eGvr4++/eTJ0/w9u1bdOzYEZGRkex1x8bGol69erh16xZrcTM2NsaDBw801iMczp8Ml4hzOH8AhQsXzpEhyrt37/Ds2TPUq1cv3ft/Z1Qml8vRu3dv9O7dG5GRkbh79y727t2LS5cuYfTo0di9ezcA4P3797C1tf3lsT59+oSSJUumMS+h7qSfP3/WuL1YsWIa/37//j0IIVi5ciVWrlyZbozw8HAULlwYI0eOxLBhw9C6dWtYWVmhYcOGsLOzQ/ny5X/5HDkcDofzv0VERAR+/PiBffv2Yd++fRk+hnL48GFs3boVb968QVJSErs99TUro9soVJpNoXLv79+/w9DQEG/fvgUhJMNrK5Vy02tnyZIlNe7X0dFB8eLFM4yvjq6uLoYOHYqhQ4fi69evuHXrFnbs2IHTp09DW1sbS5YsQUREBKKjo1G2bNlfHuvz58+oWrVqmttLly7N7reysmK3p36P6Mb4xIkTM4zx8+dP5MuXD+PGjcOkSZPQtGlTVKpUCU2aNEHnzp0z/bo5HKnhCTaH8wegvqubGVJSUjT+rVQq0aBBA7i6uqb7eAsLi0wfO3/+/GjRogVatGgBJycn3Lx5E58+fYK5uXmWnmNmSf3a6Y71gAED0KhRo3T/T4kSJQAAtWrVwtmzZ3H+/HkEBgbi4MGD2L59O2bPno1u3bqJ8nw5HA6H8/dBry2dOnVCly5d0n1MuXLlAKgMxyZNmoSWLVvCxcUFBQoUgJaWFjZs2MDMwtT51TU8I6dsQgh7XjKZDJs2bYKWllaax+XJk+fXLyybmJmZoX379rC1tUWHDh3g6+uLhQsXihILSPse0dc/YcKEND3tFPra27VrBxsbG5w9exaBgYHYsmULNm3ahNWrV6NJkyaiPWcOJ7vwBJvD+YPJly9fmtEZiYmJCAsL07itRIkSiI2NRf369QWNb21tjZs3byIsLAzm5uYoUaIEXrx48cv/Y25ujmfPnkGpVGosLF6/fg0g7W5+auiOtI6OTqZej4mJCbp27YquXbsiJiYGffr0werVq3mCzeFwOByGqakp8ubNC6VS+dtri5+fH4oXL441a9YwSTQArFq1SvDnVaJECRBCUKxYMZQqVSrDx9Fr57t37zTUaklJSfj48WO2lVs6OjooV64c3r59i8jISBQoUACGhoa/vdYrFAq8efMmze1ZvdYbGhpm6lpvZmbGlHbh4eHo0qUL1q9fzxNszh8J78HmcP5gihcvnmYEx/79+9NUsNu2bYt79+7h8uXLaY7x48cPJCcnZxgjLCwsTa8UoErkr127BrlczirGtra2ePr0Kc6ePZvm8XQ3unHjxggLC8OpU6fYfcnJydi5cyfy5MmDWrVq/eIVAwUKFEDt2rWxb98+fP36Nc396hK+yMhIjfvy5s2LEiVKaIwR4XA4HA5HS0sLrVu3hp+fH54/f57mfvVrC60k0+saoBondf/+fcGflLine truncated
            "text/plain": [
              "<Figure size 1000x500 with 2 Axes>"
-            ],
+            ]
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAA9gAAAH7CAYAAAA3jTfAAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjAsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvlHJYcgAAAAlwSFlzAAAPYQAAD2EBqD+naQABAABJREFUeJzsnXVUVN33xp8ZWgERFXUwUBELG7sLWxRRMRAVTAzs7sDXTmzFwBZbwQALOxG7G0RCpWPO74/5nfPOEErce9H3ez5ruZbOjHdP3nv22c9+towQQsDhcDgcDofD4XA4HA4nR8hz+wlwOBwOh8PhcDgcDofzX4An2BwOh8PhcDgcDofD4QgAT7A5HA6Hw+FwOBwOh8MRAJ5gczgcDofD4XA4HA6HIwA8weZwOBwOh8PhcDgcDkcAeILN4XA4HA6Hw+FwOByOAPAEm8PhcDgcDofD4XA4HAHgCTaHw+FwOBwOh8PhcDgCwBNsDofD4XA4HA6Hw+FwBIAn2BwOh8PhcDgczl+Kk5MTnJyccvtpcDic/0c7t58Ah/Nfply5cpl63I4dO1CnTh2Rn03GKJVKHDt2DN7e3nj37h2SkpJgZmaGqlWrolevXqhWrVquPTcOh8PhcH6Hj48PJk+ejIMHD6Jy5cq5/XT+SBITE7Fnzx4cPnwY79+/h1wuR+HChVGjRg3069cPZcqUye2nyOH8J+AJNocjIosWLdL499GjRxEYGJjm9ty+qM2bNw/e3t5o0aIFOnbsCC0tLbx58waXL19G8eLFeYLN4XA4HM5fzsiRI3Hp0iW0b98e3bp1Q3JyMl6/fo0LFy6gevXqub4W4XD+K/AEm8MRETs7O41/P3jwAIGBgWluT01cXBwMDAzEfGqMb9++Yffu3ejevTvmzp2rcR8hBBEREZI8DwBITk6GUqmErq6uZDE5HA6HwxGKP/U6FhQUhICAAIwePRpDhgzRuC8lJQU/fvyQ7LkkJCRAR0cHcjnvVOX8N+HfbA4nl3FyckKHDh0QHByM3r17o2rVqli2bBkAlcR89erVaf5P8+bNMWnSJI3bfvz4gfnz56NJkyawtrZGq1atsHHjRiiVyl/G//jxIwghqFGjRpr7ZDIZChQokCbOggUL0Lx5c1hbW6Nx48aYMGGCRiIeHh6OKVOmoH79+qhcuTI6deqEw4cPp4lbrlw5bNmyBV5eXmjZsiUqV66MV69eAQBevXqFkSNHonbt2qhcuTLs7e1x/vx5jWMkJSVhzZo1sLW1ReXKlVGnTh307NkTgYGBv3zNHA6Hw/nfJDQ0FJMnT0b9+vVhbW2N9u3b4+DBgxqPSUxMxMqVK2Fvb4+aNWuiWrVq6NWrF65fv67xuF9dx1avXo1y5crh3bt3mDRpEmxsbFCzZk1MnjwZcXFxaZ7X0aNHYW9vjypVqqB27doYPXo0vnz5kuZx+/btQ8uWLVGlShU4ODjg9u3bmXrdHz58AIB0r/VaWlrInz9/mvdpypQpaNiwIaytrdG8eXPMnDkTiYmJGsek1+mqVauie/fuuHDhgsZxbty4gXLlyuHkyZNYvnw5GjVqhKpVqyI6OhqAqvDg4uKCmjVromrVqujTpw/u3LmjcYzo6GjMnz+frTvq1auH/v3749GjR5l67RyO1PAKNofzBxAVFYWBAweiffv26NSpU5qk9nfExcWhT58+CA0NhaOjI4oWLYp79+5h2bJlCAsLw9SpUzP8vwqFAgDg6+uLNm3a/LJyHhMTg969e+PVq1fo2rUrKlasiMjISPj7+yM0NBSmpqaIj4+Hk5MT3r9/j969e6NYsWLw9fXFpEmT8OPHDzg7O2sc08fHBwkJCejevTt0dXWRL18+vHjxAj179kThwoUxcOBA5MmTB6dPn4abmxtWr16NVq1aAQDWrFmDDRs2oFu3bqhSpQqio6MRHByMR48eoUGDBll6DzkcDofz3+bbt2/o3r07ZDIZevfuDVNTU1y6dAlTp05FdHQ0+vXrB0CV0B04cAAdOnRAt27dEBMTg4MHD8LV1RUHDhxAhQoVNI6b3nWM4u7ujmLFimHMmDF4/PgxDhw4AFNTU4wfP549Zt26dVi5ciXatm0LBwcHREREYNeuXejduzeOHDkCY2NjAMCBAwcwY8YMVK9eHc7Ozvjw4QOGDh2KfPnyoWjRor987fRaf/z4cdSoUQPa2hmnAKGhoXBwcMDPnz/RvXt3lC5dGqGhofDz80N8fDx0dXXx7ds3ODo6Ii4uDk5OTsifPz8OHz6MoUOHYtWqVew6TfH09ISOjg5cXFyQmJgIHR0dXLt2DQMHDoS1tTWGDx8OmUwGHx8fODs7Y/fu3ahSpQoAYObMmfDz80OfPn1QpkwZREVF4c6dO3j16hUqVar0m0+dw8kFCIfDkYzZs2cTKysrjdv69OlDrKysyJ49e9I83srKiqxatSrN7c2aNSMTJ05k/167di2pVq0aefPmjcbjlixZQipUqEA+f/78y+c1YcIEYmVlRWrVqkXc3NzIli1byMuXL9M8buXKlcTKyoqcOXMmzX1KpZIQQoiXlxexsrIiR48eZfclJiaSHj16kGrVqpGfP38SQgj58OEDsbKyIjVq1CDh4eEax3J2diYdOnQgCQkJGsfv0aMHsbW1Zbd16tSJDBo06JevjcPhcDj/fQ4dOkSsrKxIUFBQho+ZMmUKadCgAYmIiNC4ffTo0aRmzZokLi6OEEJIcnKyxvWHEEK+f/9O6tevTyZPnsxu+9V1bNWqVcTKykrj8YQQ4ubmRmrXrs3+/fHjR1KhQgWybt06jcc9e/aMVKxYkd2emJhI6tWrR+zs7DSe2759+4iVlRXp06dPhq+bENU1lK436tevT8aMGUN27dpFPn36lOaxEyZMIOXLl0/3vaTX+vnz5xMrKyty69Ytdl90dDRp3rw5adasGUlJSSGEEHL9+nViZWVFWrRowd5fehxbW1syYMAAdkxCCImLiyPNmzcn/fv3Z7fVrFmTzJ49+5evj8P5k+AScQ7nD0BXVxf29vbZ/v++vr6oWbMmjI2NERERwf7Ur18fKSkpuHXr1i//v4eHB2bMmIFixYrh7Nmz+Oeff9CuXTs4OzsjNDSUPe7MmTMoX758mp1pQCUnB4BLly6hUKFC6NChA7tPR0cHTk5OiI2NTfNcbG1tYWpqyv4dFRWF69evo23btoiOjmavJTIyEg0bNsTbt2/ZczI2NsaLFy/w9u3bLL9nHA6Hw/nfgRCCM2fOoHnz5sxfhP5p2LAhfv78ySTHWlparIdaqVQiKioKycnJsLa2xuPHj9McO/V1TB1HR0eNf9vY2CAqKopJpM+ePQulUom2bdtqPKeCBQuiZMmSuHHjBgAgODgY4eHhcHR01Ojv7tKlC4yMjH77+mUyGbZs2QJ3d3cYGxvjxIkTmDNnDpo1awZ3d3fWg61UKnHu3Dk0a9YsXTd2eq2/ePEiqlSpAhsbG3Zf3rx50aNHD3z69AkvX77U+H+dO3eGvr4++/eTJ0/w9u1bdOzYEZGRkex1x8bGol69erh16xZrcTM2NsaDBw801iMczp8Ml4hzOH8AhQsXzpEhyrt37/Ds2TPUq1cv3ft/Z1Qml8vRu3dv9O7dG5GRkbh79y727t2LS5cuYfTo0di9ezcA4P3797C1tf3lsT59+oSSJUumMS+h7qSfP3/WuL1YsWIa/37//j0IIVi5ciVWrlyZbozw8HAULlwYI0eOxLBhw9C6dWtYWVmhYcOGsLOzQ/ny5X/5HDkcDofzv0VERAR+/PiBffv2Yd++fRk+hnL48GFs3boVb968QVJSErs99TUro9soVJpNoXLv79+/w9DQEG/fvgUhJMNrK5Vy02tnyZIlNe7X0dFB8eLFM4yvjq6uLoYOHYqhQ4fi69evuHXrFnbs2IHTp09DW1sbS5YsQUREBKKjo1G2bNlfHuvz58+oWrVqmttLly7N7reysmK3p36P6Mb4xIkTM4zx8+dP5MuXD+PGjcOkSZPQtGlTVKpUCU2aNEHnzp0z/bo5HKnhCTaH8wegvqubGVJSUjT+rVQq0aBBA7i6uqb7eAsLi0wfO3/+/GjRogVatGgBJycn3Lx5E58+fYK5uXmWnmNmSf3a6Y71gAED0KhRo3T/T4kSJQAAtWrVwtmzZ3H+/HkEBgbi4MGD2L59O2bPno1u3bqJ8nw5HA6H8/dBry2dOnVCly5d0n1MuXLlAKgMxyZNmoSWLVvCxcUFBQoUgJaWFjZs2MDMwtT51TU8I6dsQgh7XjKZDJs2bYKWllaax+XJk+fXLyybmJmZoX379rC1tUWHDh3g6+uLhQsXihILSPse0dc/YcKEND3tFPra27VrBxsbG5w9exaBgYHYsmULNm3ahNWrV6NJkyaiPWcOJ7vwBJvD+YPJly9fmtEZiYmJCAsL07itRIkSiI2NRf369QWNb21tjZs3byIsLAzm5uYoUaIEXrx48cv/Y25ujmfPnkGpVGosLF6/fg0g7W5+auiOtI6OTqZej4mJCbp27YquXbsiJiYGffr0werVq3mCzeFwOByGqakp8ubNC6VS+dtri5+fH4oXL441a9YwSTQArFq1SvDnVaJECRBCUKxYMZQqVSrDx9Fr57t37zTUaklJSfj48WO2lVs6OjooV64c3r59i8jISBQoUACGhoa/vdYrFAq8efMmze1ZvdYbGhpm6lpvZmbGlHbh4eHo0qUL1q9fzxNszh8J78HmcP5gihcvnmYEx/79+9NUsNu2bYt79+7h8uXLaY7x48cPJCcnZxgjLCwsTa8UoErkr127BrlczirGtra2ePr0Kc6ePZvm8XQ3unHjxggLC8OpU6fYfcnJydi5cyfy5MmDWrVq/eIVAwUKFEDt2rWxb98+fP36Nc396hK+yMhIjfvy5s2LEiVKaIwR4XA4HA5HS0sLrVu3hp+fH54/f57mfvVrC60k0+saoBondf/+fcGflLine truncated
          },
-          "metadata": {}
+          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
@@ -512,7 +522,7 @@
        "This motivates the non-trivial construction via the perturbed, conditional density: all steps required in the second equation can be computed efficiently. Hence, this gives a practical method for learning score functions, provided that we know a suitable value for $\\sigma$.\n",
        "\n",
        "This leads us to the next step, let's assume we have successfully trained  our network, how can we use $\\nabla_{\\tilde{x}} \\log p_\\sigma(\\tilde{x})$\n",
-        "to obtain a generative model for $p(x)$? That is, how can we produce actual samples $x$ with the score?"
+        "to obtain a generative model for $p(x)$? That is, how can we produce actual samples $x$ with the score? "
      ]
    },
    {
@@ -524,7 +534,7 @@
      "source": [
        "## Langevin Dynamics\n",
        "\n",
-        "_Langevin Dynamics_ are traditionally used for molecular systems with deterministic and stochastic forces.\n",
+        "To arrive at a practical method, we turn to _Langevin Dynamics_. They're traditionally used for molecular systems with deterministic and stochastic forces.\n",
        "It was shown there that iteration steps along the score, with just the right amount of additional perturbation at each step actually converge to samples from $p(x)$.\n",
        "\n",
        "Specifically, for the update\n",
@@ -535,12 +545,8 @@
        "\n",
        "with  $i = 0,1,..., K$  and  $z_i \\sim \\mathcal{N}(0, I)$,\n",
        "the iterate $x_K$ converges to a sample from $p(x)$ as  $K \\to \\infty$  and $\\epsilon \\to 0$  under a set of regularity conditions.\n",
-        "For the theory, we should also consider how  $x_0$ was sampled from a distribution for  initialization, typically called$\\pi(x)$.\n",
+        "For the theory, we should also consider how  $x_0$ was sampled from a distribution for  initialization, typically called $\\pi(x)$.\n",
-        "For updating $x_i$, we can of course make use of the trained network $s_\\theta$ to obtain the score $\\nabla_x \\log p(x)$ in each step of the iteration.\n",
+        "For updating $x_i$, we can of course make use of the trained network $s_\\theta$ to obtain the score $\\nabla_x \\log p(x)$ in each step of the iteration.\n"
        "\n",
        "-overview image?-\n",
        "\n",
        "\n"
      ]
    },
    {
@@ -550,7 +556,9 @@
        "id": "8c477e0570efd8c1"
      },
      "source": [
-        "### Langevin Dynamics Algorithm"
+        "### Langevin Dynamics Algorithm\n",
        "\n",
        "Let's give this a try with our score function! The `langevin_dynamics` function below implements the integration steps for a varying number of steps (0 to 500) to show how it converges. We're starting with points on a dense regular grid, to show whether (and how) points from all over domain move towards the densities of the underlying distribution via the score gradients."
      ]
    },
    {
@@ -570,11 +578,10 @@
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "output_type": "stream",
          "text": [
-            "/usr/local/lib/python3.11/dist-packages/torch/functional.py:539: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at /pytorch/aten/src/ATen/native/TensorShape.cpp:3637.)\n",
+            " \n"
            "  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]\n"
          ]
        }
      ],
@@ -623,7 +630,9 @@
        "id": "5cb44219fecf2b59"
      },
      "source": [
-        "### Visualize the Trajectories"
+        "### Visualize the Trajectories\n",
        "\n",
        "Having the sample points ready in the `samples` array, we can visualize them as colored dots over the score contours. We'll color the initial grid setup with a smooth RGB gradient, so that the colors of dots later on show where initial positions have ended up across the target density landscape."
      ]
    },
    {
@@ -644,14 +653,14 @@
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAABtgAAAPXCAYAAAChK0/mAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjAsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvlHJYcgAAAAlwSFlzAAAPYQAAD2EBqD+naQABAABJREFUeJzs3XV8FEcbwPHf7t7l4i4kWHB3d2sphQqFlgotpdBSfVvqRt3dlbpRqAAVChSXAi3uroG4y+Vkd94/Eo4EsncJJEg73/fzft53N7M7c8PlbrLPzDOKEEIgSZIkSZIkSZIkSZIkSZIkSZIkSVKlqGe6AZIkSZIkSZIkSZIkSZIkSZIkSZJ0LpEBNkmSJEmSJEmSJEmSJEmSJEmSJEmqAhlgkyRJkiRJkiRJkiRJkiRJkiRJkqQqkAE2SZIkSZIkSZIkSZIkSZIkSZIkSaoCGWCTJEmSJEmSJEmSJEmSJEmSJEmSpCqQATZJkiRJkiRJkiRJkiRJkiRJkiRJqgIZYJMkSZIkSZIkSZIkSZIkSZIkSZKkKpABNkmSJEmSJEmSJEmSJEmSJEmSJEmqAhlgkyRJkiRJkiRJkiRJkiRJkiRJkqQqkAE2SZIkSZIkSZIkSZIkSZIkSZIkSaoCGWCTpLPYjh07uPPOOxkwYABt2rShT58+3HDDDXz99dflyn344YfMmzfvDLWyhNPp5JVXXqF37960bduWK664guXLl5/RNkmSJEmSJJ1J58pYrrCwkLfffpvx48fTtWtXmjVrxs8//2xafs+ePYwfP54OHTrQtWtX7r//frKysk4oZxgGkydPZuDAgbRp04aLL76Y3377rSZfiiRJkiRJ0ml3roz5Nm7cyNNPP82wYcNo3749/fv356677mLfvn0VlpdjPknyTRFCiDPdCEmSTrR27VrGjBlDQkICw4cPJyYmhuTkZDZs2MDBgwf5888/PWU7dOjABRdcwIsvvnjG2nvPPfcwZ84cxowZQ2JiItOnT2fTpk18+eWXdO7c+Yy1S5IkSZIk6Uw4l8ZySUlJDBo0iISEBOrUqcPff//NCy+8wIgRI04om5KSwvDhwwkJCeG6666jqKiIzz77jPj4eH744Qf8/Pw8ZV977TU+/vhjRo0aRZs2bZg/fz6LFi3i9ddfZ9iwYafzJUqSJEmSJNWIc2nMd+edd7J27VqGDBlCs2bNSE9P59tvv6WoqIipU6fStGlTT1k55pOkyrGc6QZIklSxDz/8kJCQEH788UdCQ0PL/SwzM/MMtapiGzdu5Pfff+eBBx5g/PjxAAwfPpyLLrqIV199le+///4Mt1CSJEmSJOn0OpfGcrGxsSxbtoyYmBg2bdrE5Zdfblr2ww8/xG638/PPP5OQkABA27ZtueGGG5g+fTpXXnklAKmpqXz++eeMHj2axx9/HIArrriCa6+9lpdffpkhQ4agaVrNvzhJkiRJkqQadC6N+caOHcurr75aLjg2dOhQLr74Yj7++GNeffVVz3k55pOkypEpIiXpLHXw4EEaN258wpczQFRUlOf/N2vWjKKiIqZPn06zZs1o1qwZDz30kOfnqampPPzww/Ts2ZPWrVszbNgwfvzxx3L3W7VqFc2aNWPWrFm8/vrr9OrVi/bt23PLLbeQnJzss62zZ89G0zTPlyuAzWbj8ssvZ926dZW6hyRJkiRJ0r/JuTSW8/PzIyYmplKva+7cufTv39/zoAWgZ8+eJCYm8scff3jOzZs3D5fLxTXXXOM5pygKV199NSkpKaxbt65S9UmSJEmSJJ3NzqUxX8eOHcsF1wASExNp0qQJe/fuLXdejvkkqXLkCjZJOkvVrl2bdevWsXPnznJLtI/38ssvM2nSJNq2bcuoUaMAqFevHgAZGRmMGjUKRVEYPXo0kZGRLFmyhEcffZSCggLGjh1b7l4ffPABiqJw0003kZmZyZdffsnYsWOZOXMm/v7+pm3Ytm0biYmJBAcHlzvftm1bz8/j4+NPphskSZIkSZLOSefSWK6yUlNTyczMpHXr1if8rG3btixZssRzvG3bNgIDA2nUqNEJ5Y7+XKYRlyRJkiTpXHeuj/mEEGRkZNCkSRPPOTnmk6TKkwE2STpLjRs3jptuuonhw4fTtm1bOnXqRI8ePejWrRtWq9VT7tJLL+XJJ5+kbt26XHrppeXu8cYbb6DrOr/++isREREAXH311dxzzz28++67XHXVVeW+eHNzc5k1a5YnUNayZUsmTpzItGnTGDNmjGlb09PTK5z1fPRcWlrayXeEJEmSJEnSOehcGstV1tExndm4LycnB6fTiZ+fH+np6URFRaEoygnlyt5LkiRJkiTpXHauj/l++eUXUlNTufPOOz3n5JhPkipPpoiUpLNUr169+P777xk4cCDbt2/nk08+Yfz48fTt25f58+f7vF4Iwdy5cxk4cCBCCLKysjz/7d27N/n5+WzZsqXcNcOHDy+3Cm3IkCHExMSwePFir3UVFxefsMQcStJEHv25JEmSJEnSf8m5NJarLIfDAVCpcZ8cH0qSJEmS9F9wLo/59uzZw9NPP02HDh247LLLPOflmE+SKk+uYJOks1jbtm159913cTqdbN++nXnz5vHFF19w1113MWPGDBo3bmx6bVZWFnl5eUydOpWpU6ealimrfv365Y4VRaF+/focPnzYazv9/f1xOp0nnD/6hVwdKYkkSZIkSZLONefKWK6yjj4oqcy4T44PJUmSJEn6rzgXx3zp6encfPPNhISE8NZbb6FpmudncswnSZUnA2ySdA7w8/Ojbdu2tG3blsTERB5++GFmz57NHXfcYXqNYRgAXHLJJeVmoZTVrFmzamlfTEwMqampJ5xPT08HIDY2tlrqkSRJkiRJOhed7WO5yjo6pjs6xisrPT2d8PBwzwzmmJgYVq1ahRCiXMogOT6UJEmSJOnf6lwZ8+Xn53PTTTeRn5/Pt99+S1xcXLmfyzGfJFWeDLBJ0jnm6AajvnIYR0ZGEhQUhGEY9OzZs1L3PnDgQLljIQQHDhzw+UXevHlzVq1aRUFBQbkl6hs2bACgRYsWlapfkiRJkiTp3+5sHMtVVlxcHJGRkWzevPmEn23cuJHmzZt7jlu0aMEPP/zAnj17ys3aluNDSZIkSZL+C87WMZ/D4eCWW25h//79fP755xWurpNjPkmqPLkHmySdpVauXIkQ4oTzR/MpN2zY0HMuMDCQvLy8cuU0TeOCCy5gzpw57Ny584T7HL+8HGDGjBkUFBR4jmfPnk16ejp9+/b12tYhQ4ag63q5pexOp5Off/6Zdu3aER8f7/V6SZIkSZKkf5tzaSxXFYMHD2bRokUkJyd7zq1YsYL9+/czZMgQz7lBgwZhtVr57rvvPOeEEHz//ffExcXRoUOHamuTJEmSJEnSmXIujfl0XWfixImsX7+et956y+t4TI75JKly5Ao2STpLPfvss9jtds4//3waNmyIy+Vi7dq1/PHHH9SuXZsRI0Z4yrZq1YoVK1bw+eefExsbS506dWjXrh333nsvq1atYtSoUVxxxRU0btyY3NxctmzZwooVK/j777/L1RkWFsY111zDiBEjyMzM5Msvv6R+/fqMGjXKa1vbtWvHkCFDeP3118nMzKR+/fpMnz6dw4cP89xzz9VI/0iSJEmSJJ3NzqWxHMA333xDXl6eZ5b1woULSUlJAeC6664jJCQEgFtuuYXZs2czZswYxowZQ1FREZ9++ilNmzZl5MiRnvvVqlWLMWPG8Omnn+J2u2nTpg3z5s1j9erVvPrqq+X2+ZAkSZIkSTpXnUtjvhdffJEFCxYwYMAAcnJymDlzZrmfX3rppZ7/L8d8klQ5iqgoxC5J0hm3ZMkSZs+ezbp160hJScHlcpGQkEDfvn259dZbiYqK8pTdu3cvjz/+OJs2baK4uJjLLruMF198EYDMzEzee+89FixYQEZGBuHh4TRu3JihQ4d6vnhXrVrFmDFjeP3119mxYwc//vgjhYWFdO/enSeeeIKEhASf7XU4HLz55pv8+uuv5Obm0qxZM+666y769OlTMx0kSZIkSZJ0FjvXxnIDBw7k8OHDFf5s/vz51KlTx3O8a9cuXnzxRdasWYPVaqVfv3489NBDREdHl7vOMAwmT57M1KlTSUtLIzExkQkTJnDJJZdUuT8lSZIkSZLORufSmO+66647IVhX1o4dO8odyzGfJPkmA2ySJHm+oN96661yy7wlSZIkSZKks58cy0mSJEmSJP37yTGfJJ195B5skiRJkiRJkiRJkiRJkiRJkiRJklQFMsAmSZIkSZIkSZIkSZIkSZIkSZIkSVUgA2ySJEmSJEmSJEmSJEmSJEmSJEmSVAVyDzZJkiRJkiRJkiRJkiRJkiRJkiRJqgK5gk2SJEmSJEmSJEmSJEmSJEmSJEmSqsByphtwMjp37ozT6SQmJuZMN0WSJEmSSE9Px8/Pj9WrV5/ppkjSv4oc80mSJElnEznmk6SaIcd8kiRJ0tmkKmO+czLA5nA4cOgO9roPkWipjWqyEC/byCVH5CEoyYKpoBCuhBKhhlVYXsfggJ5U7pyKSqJWx7QtaUYGhRRSNs9mrBJNsBJUYfliHBwxUsqdsyl+1FbiTetIEodx4ip3rraSgA2/CsvnU0C6SPccKygEK8HEEG1ax372Y1A+W2gDElFQKiyfTRa55Hr6FhQiiCCc8ArL6+gcEgfKnVPRqKfUN21TukilSBRBmXbFKHEEmvWtKCbVOFLunL/iT5yaYFpHsp6ESzjLnYvX6mBVKu7bQiOfLL183wapIURo5n172LWvTD+VqG1tYNq3uXoWBXoulOnbMC2SYM3kfSvcpDoPljunYSHOVs+0TVnOVBxGYblzEX618FcDKyzvNOxkOZLLnfNTA4i0mb9vM+yH0EX5921UQF0sirXC8nZXHvnOTMq+7gBrCCF+5n2bXrAPjuvbmOCGpuULHJkUO/PK1RFkiyTAzLine truncated
            "text/plain": [
              "<Figure size 2000x1000 with 6 Axes>"
-            ],
+            ]
            "image/png": "iVBORw0KGgoAAAANSUhEUgAABtgAAAPXCAYAAAChK0/mAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjAsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvlHJYcgAAAAlwSFlzAAAPYQAAD2EBqD+naQABAABJREFUeJzs3XV8FEcbwPHf7t7l4i4kWHB3d2sphQqFlgotpdBSfVvqRt3dlbpRqAAVChSXAi3uroG4y+Vkd94/Eo4EsncJJEg73/fzft53N7M7c8PlbrLPzDOKEEIgSZIkSZIkSZIkSZIkSZIkSZIkSVKlqGe6AZIkSZIkSZIkSZIkSZIkSZIkSZJ0LpEBNkmSJEmSJEmSJEmSJEmSJEmSJEmqAhlgkyRJkiRJkiRJkiRJkiRJkiRJkqQqkAE2SZIkSZIkSZIkSZIkSZIkSZIkSaoCGWCTJEmSJEmSJEmSJEmSJEmSJEmSpCqQATZJkiRJkiRJkiRJkiRJkiRJkiRJqgIZYJMkSZIkSZIkSZIkSZIkSZIkSZKkKpABNkmSJEmSJEmSJEmSJEmSJEmSJEmqAhlgkyRJkiRJkiRJkiRJkiRJkiRJkqQqkAE2SZIkSZIkSZIkSZIkSZIkSZIkSaoCGWCTpLPYjh07uPPOOxkwYABt2rShT58+3HDDDXz99dflyn344YfMmzfvDLWyhNPp5JVXXqF37960bduWK664guXLl5/RNkmSJEmSJJ1J58pYrrCwkLfffpvx48fTtWtXmjVrxs8//2xafs+ePYwfP54OHTrQtWtX7r//frKysk4oZxgGkydPZuDAgbRp04aLL76Y3377rSZfiiRJkiRJ0ml3roz5Nm7cyNNPP82wYcNo3749/fv356677mLfvn0VlpdjPknyTRFCiDPdCEmSTrR27VrGjBlDQkICw4cPJyYmhuTkZDZs2MDBgwf5888/PWU7dOjABRdcwIsvvnjG2nvPPfcwZ84cxowZQ2JiItOnT2fTpk18+eWXdO7c+Yy1S5IkSZIk6Uw4l8ZySUlJDBo0iISEBOrUqcPff//NCy+8wIgRI04om5KSwvDhwwkJCeG6666jqKiIzz77jPj4eH744Qf8/Pw8ZV977TU+/vhjRo0aRZs2bZg/fz6LFi3i9ddfZ9iwYafzJUqSJEmSJNWIc2nMd+edd7J27VqGDBlCs2bNSE9P59tvv6WoqIipU6fStGlTT1k55pOkyrGc6QZIklSxDz/8kJCQEH788UdCQ0PL/SwzM/MMtapiGzdu5Pfff+eBBx5g/PjxAAwfPpyLLrqIV199le+///4Mt1CSJEmSJOn0OpfGcrGxsSxbtoyYmBg2bdrE5Zdfblr2ww8/xG638/PPP5OQkABA27ZtueGGG5g+fTpXXnklAKmpqXz++eeMHj2axx9/HIArrriCa6+9lpdffpkhQ4agaVrNvzhJkiRJkqQadC6N+caOHcurr75aLjg2dOhQLr74Yj7++GNeffVVz3k55pOkypEpIiXpLHXw4EEaN258wpczQFRUlOf/N2vWjKKiIqZPn06zZs1o1qwZDz30kOfnqampPPzww/Ts2ZPWrVszbNgwfvzxx3L3W7VqFc2aNWPWrFm8/vrr9OrVi/bt23PLLbeQnJzss62zZ89G0zTPlyuAzWbj8ssvZ926dZW6hyRJkiRJ0r/JuTSW8/PzIyYmplKva+7cufTv39/zoAWgZ8+eJCYm8scff3jOzZs3D5fLxTXXXOM5pygKV199NSkpKaxbt65S9UmSJEmSJJ3NzqUxX8eOHcsF1wASExNp0qQJe/fuLXdejvkkqXLkCjZJOkvVrl2bdevWsXPnznJLtI/38ssvM2nSJNq2bcuoUaMAqFevHgAZGRmMGjUKRVEYPXo0kZGRLFmyhEcffZSCggLGjh1b7l4ffPABiqJw0003kZmZyZdffsnYsWOZOXMm/v7+pm3Ytm0biYmJBAcHlzvftm1bz8/j4+NPphskSZIkSZLOSefSWK6yUlNTyczMpHXr1if8rG3btixZssRzvG3bNgIDA2nUqNEJ5Y7+XKYRlyRJkiTpXHeuj/mEEGRkZNCkSRPPOTnmk6TKkwE2STpLjRs3jptuuonhw4fTtm1bOnXqRI8ePejWrRtWq9VT7tJLL+XJJ5+kbt26XHrppeXu8cYbb6DrOr/++isREREAXH311dxzzz28++67XHXVVeW+eHNzc5k1a5YnUNayZUsmTpzItGnTGDNmjGlb09PTK5z1fPRcWlrayXeEJEmSJEnSOehcGstV1tExndm4LycnB6fTiZ+fH+np6URFRaEoygnlyt5LkiRJkiTpXHauj/l++eUXUlNTufPOOz3n5JhPkipPpoiUpLNUr169+P777xk4cCDbt2/nk08+Yfz48fTt25f58+f7vF4Iwdy5cxk4cCBCCLKysjz/7d27N/n5+WzZsqXcNcOHDy+3Cm3IkCHExMSwePFir3UVFxefsMQcStJEHv25JEmSJEnSf8m5NJarLIfDAVCpcZ8cH0qSJEmS9F9wLo/59uzZw9NPP02HDh247LLLPOflmE+SKk+uYJOks1jbtm159913cTqdbN++nXnz5vHFF19w1113MWPGDBo3bmx6bVZWFnl5eUydOpWpU6ealimrfv365Y4VRaF+/focPnzYazv9/f1xOp0nnD/6hVwdKYkkSZIkSZLONefKWK6yjj4oqcy4T44PJUmSJEn6rzgXx3zp6encfPPNhISE8NZbb6FpmudncswnSZUnA2ySdA7w8/Ojbdu2tG3blsTERB5++GFmz57NHXfcYXqNYRgAXHLJJeVmoZTVrFmzamlfTEwMqampJ5xPT08HIDY2tlrqkSRJkiRJOhed7WO5yjo6pjs6xisrPT2d8PBwzwzmmJgYVq1ahRCiXMogOT6UJEmSJOnf6lwZ8+Xn53PTTTeRn5/Pt99+S1xcXLmfyzGfJFWeDLBJ0jnm6AajvnIYR0ZGEhQUhGEY9OzZs1L3PnDgQLljIQQHDhzw+UXevHlzVq1aRUFBQbkl6hs2bACgRYsWlapfkiRJkiTp3+5sHMtVVlxcHJGRkWzevPmEn23cuJHmzZt7jlu0aMEPP/zAnj17ys3aluNDSZIkSZL+C87WMZ/D4eCWW25h//79fP755xWurpNjPkmqPLkHmySdpVauXIkQ4oTzR/MpN2zY0HMuMDCQvLy8cuU0TeOCCy5gzpw57Ny584T7HL+8HGDGjBkUFBR4jmfPnk16ejp9+/b12tYhQ4ag63q5pexOp5Off/6Zdu3aER8f7/V6SZIkSZKkf5tzaSxXFYMHD2bRokUkJyd7zq1YsYL9+/czZMgQz7lBgwZhtVr57rvvPOeEEHz//ffExcXRoUOHamuTJEmSJEnSmXIujfl0XWfixImsX7+et956y+t4TI75JKly5Ao2STpLPfvss9jtds4//3waNmyIy+Vi7dq1/PHHH9SuXZsRI0Z4yrZq1YoVK1bw+eefExsbS506dWjXrh333nsvq1atYtSoUVxxxRU0btyY3NxctmzZwooVK/j777/L1RkWFsY111zDiBEjyMzM5Msvv6R+/fqMGjXKa1vbtWvHkCFDeP3118nMzKR+/fpMnz6dw4cP89xzz9VI/0iSJEmSJJ3NzqWxHMA333xDXl6eZ5b1woULSUlJAeC6664jJCQEgFtuuYXZs2czZswYxowZQ1FREZ9++ilNmzZl5MiRnvvVqlWLMWPG8Omnn+J2u2nTpg3z5s1j9erVvPrqq+X2+ZAkSZIkSTpXnUtjvhdffJEFCxYwYMAAcnJymDlzZrmfX3rppZ7/L8d8klQ5iqgoxC5J0hm3ZMkSZs+ezbp160hJScHlcpGQkEDfvn259dZbiYqK8pTdu3cvjz/+OJs2baK4uJjLLruMF198EYDMzEzee+89FixYQEZGBuHh4TRu3JihQ4d6vnhXrVrFmDFjeP3119mxYwc//vgjhYWFdO/enSeeeIKEhASf7XU4HLz55pv8+uuv5Obm0qxZM+666y769OlTMx0kSZIkSZJ0FjvXxnIDBw7k8OHDFf5s/vz51KlTx3O8a9cuXnzxRdasWYPVaqVfv3489NBDREdHl7vOMAwmT57M1KlTSUtLIzExkQkTJnDJJZdUuT8lSZIkSZLORufSmO+66647IVhX1o4dO8odyzGfJPkmA2ySJHm+oN96661yy7wlSZIkSZKks58cy0mSJEmSJP37yTGfJJ195B5skiRJkiRJkiRJkiRJkiRJkiRJklQFMsAmSZIkSZIkSZIkSZIkSZIkSZIkSVUgA2ySJEmSJEmSJEmSJEmSJEmSJEmSVAVyDzZJkiRJkiRJkiRJkiRJkiRJkiRJqgK5gk2SJEmSJEmSJEmSJEmSJEmSJEmSqsByphtwMjp37ozT6SQmJuZMN0WSJEmSSE9Px8/Pj9WrV5/ppkjSv4oc80mSJElnEznmk6SaIcd8kiRJ0tmkKmO+czLA5nA4cOgO9roPkWipjWqyEC/byCVH5CEoyYKpoBCuhBKhhlVYXsfggJ5U7pyKSqJWx7QtaUYGhRRSNs9mrBJNsBJUYfliHBwxUsqdsyl+1FbiTetIEodx4ip3rraSgA2/CsvnU0C6SPccKygEK8HEEG1ax372Y1A+W2gDElFQKiyfTRa55Hr6FhQiiCCc8ArL6+gcEgfKnVPRqKfUN21TukilSBRBmXbFKHEEmvWtKCbVOFLunL/iT5yaYFpHsp6ESzjLnYvX6mBVKu7bQiOfLL183wapIURo5n172LWvTD+VqG1tYNq3uXoWBXoulOnbMC2SYM3kfSvcpDoPljunYSHOVs+0TVnOVBxGYblzEX618FcDKyzvNOxkOZLLnfNTA4i0mb9vM+yH0EX5921UQF0sirXC8nZXHvnOTMq+7gBrCCF+5n2bXrAPjuvbmOCGpuULHJkUO/PK1RFkiyTAzLine truncated
          },
-          "metadata": {}
+          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
@@ -711,10 +720,36 @@
      "metadata": {
        "id": "1e25c16fdfac1c16"
      },
      "source": [
        "This seems to work well! The regularly spaced points correctly move towards the higher density regions, and the reddish points on the right primarily end up in the high-density cluster on the right side, while other points (blue, green and red all mixed) end up in the lower density peak in the center."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "81f4777c",
      "metadata": {},
      "source": [
        "## Annealed Langevin Dynamics\n",
        "\n",
-        "### Network and Dataset Definition with Noise Level"
+        "While this provably converges in the limit, it causes some difficulties in practice: the approach requires a fairly large amount of noise $\\sigma$ to ensure the gradients actually \"lead\" the samples towards the correct targets. If $\\sigma$ is too small, we won't have a score (no direction from the gradient), and if it's too large, the overlapping Gaussians will reduce the quality of our data samples (drowning important details in noise).\n",
        "\n",
        "Considering the space of all $x$ we get from $\\pi(x)$, increasing noise levels $\\sigma$ of the perturbed dataset will cover larger regions of the perturbed data space. This will give gradients wherever we start out, \n",
        "and is important in practice: we shouldn't have to make a lot of assumptions to find good initial points for our iterations. This would only shift the problem from generating the outputs to generating suitable inputs for the algorithm. This is way many classic works need to carefully specify prior distributions to make sure the algorithms converge. It's a delicate balance: working with large $\\sigma$ is important for practical applications, \n",
        "but if the perturbation is too large we have $p(x) \\not\\approx p_\\theta(x) $\n",
        "and the samples we produce will be useless in the worst case.\n",
        "\n",
        "This poses the question: why should we be restricted to a single $\\sigma$?\n",
        "It turns out that we can resolve the problem of _no score VS bad quality_ by going from large to small amounts of noise.\n",
        "Let's consider multiple noise scales $0 < \\sigma_1 < \\sigma_2 < ... < \\sigma_L$ , giving what's known as _Annealed Langevin Dynamics_.\n",
        "(Looking ahead, this also bring us closer to a central topic in the area of diffusion models: the different and changing levels of noise.)\n",
        "\n",
        "We repeatedly apply Langevin Dynamics for each noise scale, starting from the largest noise $\\sigma_L$ until the smallest noise $\\sigma_1$, and\n",
        "train a network $s_\\theta(x, \\sigma_j)$ with the noise scale $\\sigma_j$ as an additional input. \n",
        "Now we can iterate over $x_i$, and while updating it we also use smaller and smaller $\\sigma_j$ to make sure we converge to an accurate sample $x \\sim p_\\theta(\\tilde{x}) $.\n",
        "\n",
        "### Network and Dataset Definition with Noise Level\n",
        "\n",
        "Let's implement this idea. First, we need to modify our score network to be aware of the noise level. For this, it will receive an additional parameter $\\sigma$ as input. Likewise, our dataset is extended to cover different amounts of noise."
      ]
    },
    {
@@ -782,7 +817,9 @@
        "id": "41110367f065c0f8"
      },
      "source": [
-        "### Training"
+        "### Training\n",
        "\n",
        "With these two ingredients, training proceeds just like before..."
      ]
    },
    {
@@ -802,8 +839,8 @@
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 10/100, Loss: 3306.6889\n",
            "Epoch 20/100, Loss: 3374.6170\n",
@@ -869,7 +906,7 @@
        "id": "679141fd87497e8b"
      },
      "source": [
-        "### Sampling with Langevin Dynamics and Noise Levels"
+        "For sampling with Langevin Dynamics and varying noise levels, we modify the sampling process to reduce the $\\sigma$ values step by step, and then leverage the noise-aware score network to provide the right gradients."
      ]
    },
    {
@@ -929,7 +966,9 @@
        "id": "1b51131422b627ae"
      },
      "source": [
-        "### Visualization of Trajectories"
+        "### Visualization of Trajectories\n",
        "\n",
        "Now we can finally visualize the trajectories. In addition to the moving sample locations, the contours of the underlying score function will now change correspondingly with the varying $\\sigma$ perturbations."
      ]
    },
    {
@@ -950,14 +989,14 @@
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAABcIAAAPgCAYAAAAROg/TAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjAsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvlHJYcgAAAAlwSFlzAAAPYQAAD2EBqD+naQABAABJREFUeJzs3XV4FFcXwOHf7G427kIESNCggSBFirsWa/EPipQWt1KKS1taoBSXUlqgpZTiVLDi7u4SQrC468rM90fKtimWhI0A9+2Tp+zuzNw7u8nOzJl7z5EURVEQBEEQBEEQBEEQBEEQBEEQhNeUKq87IAiCIAiCIAiCIAiCIAiCIAg5SQTCBUEQBEEQBEEQBEEQBEEQhNeaCIQLgiAIgiAIgiAIgiAIgiAIrzURCBcEQRAEQRAEQRAEQRAEQRBeayIQLgiCIAiCIAiCIAiCIAiCILzWRCBcEARBEARBEARBEARBEARBeK2JQLggCIIgCIIgCIIgCIIgCILwWhOBcEEQBEEQBEEQBEEQBEEQBOG1JgLhgiAIgiAIgiAIgiAIgiAIwmtNBMIF4Q3l7+/P/PnzM7VsgwYN+PTTT7Pcxv379/H392fjxo1ZXlcQBEEQBEEQBEEQBEEQzEUEwgXhFbZx40b8/f25ePHiS2/rzJkzzJ8/n/j4eDP0TBAEQRCE3HD06FHGjBlD06ZNqVChAg0bNmTcuHGEh4dnehthYWEMHTqUKlWqUKlSJfr378+9e/eeuuy6deto3rw55cuXp0mTJvz000/m2hVBEARBeO2Fh4fz9ddf87///Y/AwED8/f05fvx4lrYhjtuCkH2avO6AIAh548KFC6jVatPjs2fPsmDBAtq1a4eDg0OGZbdv344kSbndRUEQBEEQXmDmzJnExcXRrFkz/Pz8uHfvHqtWrWLfvn1s3rwZd3f3566flJREjx49SEhI4MMPP8TCwoIVK1bQvXt3Nm/ejLOzs2nZNWvWMGnSJJo2bUqvXr04deoUn3/+OSkpKfTr1y+nd1UQBEEQXnl37tzhu+++w8/PD39/f86ePZul9cVxWxBejgiEC8IbytLSMtPLarXaHOyJIAiCIAjZNWbMGCpXroxK9c9Ez9q1a9O9e3dWrVrF8OHDn7v+6tWrCQ4OZt26dQQEBJjWb926NcuXL2fEiBEApKamMnv2bOrVq8e8efMA6NixI7Iss3jxYjp16oSjo2MO7aUgCIIgvB7Kli3L8ePHcXJyYvv27VkOhIvjtiC8HJEaRRBeI59++imBgYGEhYUxYMAAAgMDqV69OtOnT8doNGZY9t85wufPn8+MGTMAaNiwIf7+/vj7+3P//n3gyRzhsbGxTJ8+ndatWxMYGEilSpXo27cv165dy6U9FQRBEIS8tXv3bjp27EjFihVNx81//+SWqlWrZgiCP37OycmJoKCgF66/Y8cOypcvb7qYBihWrBg1atRg27ZtpueOHz9ObGwsXbt2zbB+t27dSE5OZt++fS+3I4IgCIKQg/LLcdvOzg4nJ6dsry+O24LwcsSIcEF4zRiNRvr06UNAQACffPIJR48e5YcffqBQoUJPHAQfa9y4McHBwfzxxx+MGTPGNJ3KxcXlqcvfu3ePXbt20axZMwoWLEhkZCS//vor3bt3588//6RAgQI5tn+CIAiCkNeOHj3KwIEDKVCggGla8rp16wgODqZMmTJ4eXk9d/2UlBRSUlJe2I5arc7WaK2kpCSSkpIyTI9+GlmWuX79Oh06dHjitfLly3Po0CESExOxs7PjypUrAJQrVy7DcmXLlkWlUnH16lXatGmT5b4KgiAIQk7L78ftzBLHbUF4eSIQLgivmbS0NJo3b87AgQMB6NKlC+3atWP9+vXPDISXKlWKMmXK8Mcff9CoUSMKFiz43Db8/f3ZsWNHhhFobdq0oXnz5qxfv97UtiAIgiC8jhYtWoRWq+WXX37B29sbgPbt29O4cWMKFSpkmoL8LMuWLWPBggUvbMfHx4c9e/ZkuX8rV65Er9fTvHnz5y4XGxuLTqd7ah7xx8+Fh4djZ2dHREQEarUaV1fXDMtptVqcnJyyVJxTEARBEHJTfj9uZ5Y4bgvCyxOBcEF4DXXp0iXD48qVK/Pbb7+Zbfv/zhluNBqJj4/HxsaGIkWKmO48C4IgCMLrSKfTcfbsWZo0aWK6mIb0WVStWrVi3bp16HS659bXaNu2LZUrV35hW1mp5/HYyZMnWbhwIc2bN6dGjRrPXTYtLQ14ei2Qx20/XiY1NRULC4tn9jM1NTXLfRUEQRCEnJbfj9tZIY7bgvDyRCBcEF4zlpaWT6Q0cXR0JC4uzmxtyLLMjz/+yOrVq7l//36G/OMvk+9MEARBEPK7mJgY9Ho9RYoUeeK1YsWKYTQaefjwIX5+fs/cRqFChShUqJDZ+3b79m0GDRpEiRIl+Pzzz1+4/OOLZp1O98Rrjy+kHy9jZWWFXq9/6nbS0tKwsrLKbrcFQRAEIcfk5+N2VonjtiC8PBEIF4TXjFqtzvE2lixZwty5c+nQoQNDhw7F0dERlUrFtGnTUBQlx9sXBEEQhLzy+Dj73yLUkH6jGMBgMDx3G0lJSSQnJ2eqrWfV6/ivR48e0adPH+zs7Fi6dCl2dnYvXMfJyQmtVktERMQTrz1+zsPDA0ifcm00GomKisowzVqn0xEbG2taThAEQRDyk/x63M4OcdwWhJcnAuGCIAAgSVKml92xYwfVqlVj2rRpGZ6Pj49/YWEuQRAEQXiVubi4YGtrS1BQ0BOvBQUFoVKpXlh064cffjBrrtGYmBh69+6NTqdj9erVmb64ValUlCxZkkuXLj3x2oULFyhUqJApoF66dGkALl26RN26dU3LXbp0CVmWKVWqVKbaFARBEITclB+P29kljtuC8PJEIFwQBACsra0BSEhIeOGyarX6iZHf27ZtIywsDF9f3xzpnyAIgiDkByqVinr16rFz507u3btnmiodHh7OH3/8QdWqVbG1tX3uNsyZazQ5OZl+/foRFhbGjz/++Nyp3Q8fPiQlJYVixYqZnmvatCmzZs3i4sWLlC9fHkgPDBw7dozevXublqtevTpOTk788ssvGS6of/nlF6ytralXr94L+yoIgiAIuS2/HbezQhy3BcH8RCBcEAQAypYtC8Ds2bNp0aIFFhYW1K9fHxsbmyeWrVevHgsXLmTMmDEEBgZy48YNfv/993yRN00QBEEQctrQoUM5cOAAXbt2pWvXrmg0GtasWYPBYGDUqFEvXN+cuUY//vhjLly4QIcOHbh9+za3b982vWZra0ujRo1Mj0ePHs2JEye4fv266bmuXbuybt06PvzwQ3r37o1Go2HFihW4urpmuKC2srJiyJAhTJ06lSFDhlC7dm1OnTrFb7/9xvDhw0WNEEEQBCHfyk/HbYBFixYBcOvWLQC2bNnC6dOnARgwYIBpOXHcFgTzE4FwQRAACAgIYOjQoaxZs4aDBw8iyzK7d+9+aiD8o48+IiUlhd9//52tW7dSpkwZvv32W2bNmpUHPRcEQRCE3OXr68uqVav4+uuvWbp0KYqiEBAQwKxZs0yjs3LLtWvXANiwYQMbNmzI8JqPj0+GQPjT2NnZ8dNPPzFt2jQWL16MLMtUq1aNMWPGPJHntFu3blhYWPDDDz+wZ88evLy8GDNmDD179jTvTgmCIAiCGeWn4zbA3LlzMzz+9/H734HwpxHHbUF4OZIiKtsJgiAIgiAIgiAIgiAIgiAIrzFVXndAEARBEARBEARBEARBEARBEHKSCIQLgiAIgiAIgiAIgiAIgiAIrzURCBcEQRAEQRAEQRAEQRAEQRBeayIQLgiCIAiCIAiCIAiCIAiCILzWRCBcEARBEARBEARBEARBEARBeK2JQLggCIIgCIIgCIIgCIIgCILwWtPkdQeyo0qVKuh0Otzd3fO6K4IgCILwQhEREWi1Wk6dOpXXXckT4rgtCIIgvErEcVsctwVBEIRXR1aO269kIDwtLQ2j0fjM1+Pj40lOTsbJyQkrK6tc6VNkZCRGoxEXFxcsLCxypU2AxMRE4uLicHR0xM7OLtfazSxFUYiPjycxMRELCwvc3d2RJCmvu2VWKSkpxMbGAuDh4YFarc7bDmWTXq8nNjYWnU6Hra0tTk5Oed0lQcggNTWVqKgoACRJQqvVYm1tjZWVVb7/uzMYDCiKktfdyDMvOm4LuUNnNPIoIQGjoqBWqfC2s8Mik387iTodEcnJKIqCtYUFnnZ25PTR/F5KJAbFyOO/HA+tI3aa9PM6Y5qBlMgkFFnB0skarb1lpreblJREbGwsWq32mUGee9zDgMH0uCAFsSD3zu9ellHWE518L8NzLjaFUKsy7kNwMCjy3w8kcHICZ2fz9kWWZZKTk0lKSsJgMKBSqfD09HztzgfzAz16EpRE1JIKBxyQcvyv1MwUBcUoI2nSv5cSjXHEG6IACVBw0rhjo7bP1KZSEo3EReoBsHfWYOuYtctecdwWx20hZyUZUwlLi8XaYMQhNe1fr0jYevjBq/b9lcN0qfGkJkXy+PvQxt4TjdbG7O0oKKTKyQBYqWwwGFJJ0SeQQCo6tQpFAgkJXwsvVGZIMJGik3kUnf75Syg4qqPRSDoANBbW2Dh6wb17oE//PkeSwN4e3Nxeuu2nURRISVaQJAlr87+9ABjjopETYgEJJAnJ3YtHKanoZRmNSoWXvT0Wqpd7b1N08SSlRYHpLFrCzb7IS/Y88yJ0D9ArOlP7rhZeWKqsc7TNrBy3X8lAuIeHBwC7d+9+4rWQkBA6duyItbU106dPp1atWjnen19++YVZs2bh4ODA6tWrKViwYI63CbB582Y+//xz7OzsGDJkCN26dcuVdjPr8uXLTJw4kfj4eOzs7HjnnXf4+OOPsbHJoW+UXJaQkMCMGTPYtm0bNjY2lC9fnmnTpuHl5ZXXXcuShIQEFi5cyIYNG9BqtTg5OdG/f386d+4sLLine truncated
            "text/plain": [
              "<Figure size 1500x1000 with 6 Axes>"
-            ],
+            ]
            "image/png": "iVBORw0KGgoAAAANSUhEUgAABcIAAAPgCAYAAAAROg/TAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjAsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvlHJYcgAAAAlwSFlzAAAPYQAAD2EBqD+naQABAABJREFUeJzs3XV4FFcXwOHf7G427kIESNCggSBFirsWa/EPipQWt1KKS1taoBSXUlqgpZTiVLDi7u4SQrC468rM90fKtimWhI0A9+2Tp+zuzNw7u8nOzJl7z5EURVEQBEEQBEEQBEEQBEEQBEEQhNeUKq87IAiCIAiCIAiCIAiCIAiCIAg5SQTCBUEQBEEQBEEQBEEQBEEQhNeaCIQLgiAIgiAIgiAIgiAIgiAIrzURCBcEQRAEQRAEQRAEQRAEQRBeayIQLgiCIAiCIAiCIAiCIAiCILzWRCBcEARBEARBEARBEARBEARBeK2JQLggCIIgCIIgCIIgCIIgCILwWhOBcEEQBEEQBEEQBEEQBEEQBOG1JgLhgiAIgiAIgiAIgiAIgiAIwmtNBMIF4Q3l7+/P/PnzM7VsgwYN+PTTT7Pcxv379/H392fjxo1ZXlcQBEEQBEEQBEEQBEEQzEUEwgXhFbZx40b8/f25ePHiS2/rzJkzzJ8/n/j4eDP0TBAEQRCE3HD06FHGjBlD06ZNqVChAg0bNmTcuHGEh4dnehthYWEMHTqUKlWqUKlSJfr378+9e/eeuuy6deto3rw55cuXp0mTJvz000/m2hVBEARBeO2Fh4fz9ddf87///Y/AwED8/f05fvx4lrYhjtuCkH2avO6AIAh548KFC6jVatPjs2fPsmDBAtq1a4eDg0OGZbdv344kSbndRUEQBEEQXmDmzJnExcXRrFkz/Pz8uHfvHqtWrWLfvn1s3rwZd3f3566flJREjx49SEhI4MMPP8TCwoIVK1bQvXt3Nm/ejLOzs2nZNWvWMGnSJJo2bUqvXr04deoUn3/+OSkpKfTr1y+nd1UQBEEQXnl37tzhu+++w8/PD39/f86ePZul9cVxWxBejgiEC8IbytLSMtPLarXaHOyJIAiCIAjZNWbMGCpXroxK9c9Ez9q1a9O9e3dWrVrF8OHDn7v+6tWrCQ4OZt26dQQEBJjWb926NcuXL2fEiBEApKamMnv2bOrVq8e8efMA6NixI7Iss3jxYjp16oSjo2MO7aUgCIIgvB7Kli3L8ePHcXJyYvv27VkOhIvjtiC8HJEaRRBeI59++imBgYGEhYUxYMAAAgMDqV69OtOnT8doNGZY9t85wufPn8+MGTMAaNiwIf7+/vj7+3P//n3gyRzhsbGxTJ8+ndatWxMYGEilSpXo27cv165dy6U9FQRBEIS8tXv3bjp27EjFihVNx81//+SWqlWrZgiCP37OycmJoKCgF66/Y8cOypcvb7qYBihWrBg1atRg27ZtpueOHz9ObGwsXbt2zbB+t27dSE5OZt++fS+3I4IgCIKQg/LLcdvOzg4nJ6dsry+O24LwcsSIcEF4zRiNRvr06UNAQACffPIJR48e5YcffqBQoUJPHAQfa9y4McHBwfzxxx+MGTPGNJ3KxcXlqcvfu3ePXbt20axZMwoWLEhkZCS//vor3bt3588//6RAgQI5tn+CIAiCkNeOHj3KwIEDKVCggGla8rp16wgODqZMmTJ4eXk9d/2UlBRSUlJe2I5arc7WaK2kpCSSkpIyTI9+GlmWuX79Oh06dHjitfLly3Po0CESExOxs7PjypUrAJQrVy7DcmXLlkWlUnH16lXatGmT5b4KgiAIQk7L78ftzBLHbUF4eSIQLgivmbS0NJo3b87AgQMB6NKlC+3atWP9+vXPDISXKlWKMmXK8Mcff9CoUSMKFiz43Db8/f3ZsWNHhhFobdq0oXnz5qxfv97UtiAIgiC8jhYtWoRWq+WXX37B29sbgPbt29O4cWMKFSpkmoL8LMuWLWPBggUvbMfHx4c9e/ZkuX8rV65Er9fTvHnz5y4XGxuLTqd7ah7xx8+Fh4djZ2dHREQEarUaV1fXDMtptVqcnJyyVJxTEARBEHJTfj9uZ5Y4bgvCyxOBcEF4DXXp0iXD48qVK/Pbb7+Zbfv/zhluNBqJj4/HxsaGIkWKmO48C4IgCMLrSKfTcfbsWZo0aWK6mIb0WVStWrVi3bp16HS659bXaNu2LZUrV35hW1mp5/HYyZMnWbhwIc2bN6dGjRrPXTYtLQ14ei2Qx20/XiY1NRULC4tn9jM1NTXLfRUEQRCEnJbfj9tZIY7bgvDyRCBcEF4zlpaWT6Q0cXR0JC4uzmxtyLLMjz/+yOrVq7l//36G/OMvk+9MEARBEPK7mJgY9Ho9RYoUeeK1YsWKYTQaefjwIX5+fs/cRqFChShUqJDZ+3b79m0GDRpEiRIl+Pzzz1+4/OOLZp1O98Rrjy+kHy9jZWWFXq9/6nbS0tKwsrLKbrcFQRAEIcfk5+N2VonjtiC8PBEIF4TXjFqtzvE2lixZwty5c+nQoQNDhw7F0dERlUrFtGnTUBQlx9sXBEEQhLzy+Dj73yLUkH6jGMBgMDx3G0lJSSQnJ2eqrWfV6/ivR48e0adPH+zs7Fi6dCl2dnYvXMfJyQmtVktERMQTrz1+zsPDA0ifcm00GomKisowzVqn0xEbG2taThAEQRDyk/x63M4OcdwWhJcnAuGCIAAgSVKml92xYwfVqlVj2rRpGZ6Pj49/YWEuQRAEQXiVubi4YGtrS1BQ0BOvBQUFoVKpXlh064cffjBrrtGYmBh69+6NTqdj9erVmb64ValUlCxZkkuXLj3x2oULFyhUqJApoF66dGkALl26RN26dU3LXbp0CVmWKVWqVKbaFARBEITclB+P29kljtuC8PJEIFwQBACsra0BSEhIeOGyarX6iZHf27ZtIywsDF9f3xzpnyAIgiDkByqVinr16rFz507u3btnmiodHh7OH3/8QdWqVbG1tX3uNsyZazQ5OZl+/foRFhbGjz/++Nyp3Q8fPiQlJYVixYqZnmvatCmzZs3i4sWLlC9fHkgPDBw7dozevXublqtevTpOTk788ssvGS6of/nlF6ytralXr94L+yoIgiAIuS2/HbezQhy3BcH8RCBcEAQAypYtC8Ds2bNp0aIFFhYW1K9fHxsbmyeWrVevHgsXLmTMmDEEBgZy48YNfv/993yRN00QBEEQctrQoUM5cOAAXbt2pWvXrmg0GtasWYPBYGDUqFEvXN+cuUY//vhjLly4QIcOHbh9+za3b982vWZra0ujRo1Mj0ePHs2JEye4fv266bmuXbuybt06PvzwQ3r37o1Go2HFihW4urpmuKC2srJiyJAhTJ06lSFDhlC7dm1OnTrFb7/9xvDhw0WNEEEQBCHfyk/HbYBFixYBcOvWLQC2bNnC6dOnARgwYIBpOXHcFgTzE4FwQRAACAgIYOjQoaxZs4aDBw8iyzK7d+9+aiD8o48+IiUlhd9//52tW7dSpkwZvv32W2bNmpUHPRcEQRCE3OXr68uqVav4+uuvWbp0KYqiEBAQwKxZs0yjs3LLtWvXANiwYQMbNmzI8JqPj0+GQPjT2NnZ8dNPPzFt2jQWL16MLMtUq1aNMWPGPJHntFu3blhYWPDDDz+wZ88evLy8GDNmDD179jTvTgmCIAiCGeWn4zbA3LlzMzz+9/H734HwpxHHbUF4OZIiKtsJgiAIgiAIgiAIgiAIgiAIrzFVXndAEARBEARBEARBEARBEARBEHKSCIQLgiAIgiAIgiAIgiAIgiAIrzURCBcEQRAEQRAEQRAEQRAEQRBeayIQLgiCIAiCIAiCIAiCIAiCILzWRCBcEARBEARBEARBEARBEARBeK2JQLggCIIgCIIgCIIgCIIgCILwWtPkdQeyo0qVKuh0Otzd3fO6K4IgCILwQhEREWi1Wk6dOpXXXckT4rgtCIIgvErEcVsctwVBEIRXR1aO269kIDwtLQ2j0fjM1+Pj40lOTsbJyQkrK6tc6VNkZCRGoxEXFxcsLCxypU2AxMRE4uLicHR0xM7OLtfazSxFUYiPjycxMRELCwvc3d2RJCmvu2VWKSkpxMbGAuDh4YFarc7bDmWTXq8nNjYWnU6Hra0tTk5Oed0lQcggNTWVqKgoACRJQqvVYm1tjZWVVb7/uzMYDCiKktfdyDMvOm4LuUNnNPIoIQGjoqBWqfC2s8Mik387iTodEcnJKIqCtYUFnnZ25PTR/F5KJAbFyOO/HA+tI3aa9PM6Y5qBlMgkFFnB0skarb1lpreblJREbGwsWq32mUGee9zDgMH0uCAFsSD3zu9ellHWE518L8NzLjaFUKsy7kNwMCjy3w8kcHICZ2fz9kWWZZKTk0lKSsJgMKBSqfD09HztzgfzAz16EpRE1JIKBxyQcvyv1MwUBcUoI2nSv5cSjXHEG6IACVBw0rhjo7bP1KZSEo3EReoBsHfWYOuYtctecdwWx20hZyUZUwlLi8XaYMQhNe1fr0jYevjBq/b9lcN0qfGkJkXy+PvQxt4TjdbG7O0oKKTKyQBYqWwwGFJJ0SeQQCo6tQpFAgkJXwsvVGZIMJGik3kUnf75Syg4qqPRSDoANBbW2Dh6wb17oE//PkeSwN4e3Nxeuu2nURRISVaQJAlr87+9ABjjopETYgEJJAnJ3YtHKanoZRmNSoWXvT0Wqpd7b1N08SSlRYHpLFrCzb7IS/Y88yJ0D9ArOlP7rhZeWKqsc7TNrBy3X8lAuIeHBwC7d+9+4rWQkBA6duyItbU106dPp1atWjnen19++YVZs2bh4ODA6tWrKViwYI63CbB582Y+//xz7OzsGDJkCN26dcuVdjPr8uXLTJw4kfj4eOzs7HjnnXf4+OOPsbHJoW+UXJaQkMCMGTPYtm0bNjY2lC9fnmnTpuHl5ZXXXcuShIQEFi5cyIYNG9BqtTg5OdG/f386d+4sLLine truncated
          },
-          "metadata": {}
+          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
@@ -1004,6 +1043,13 @@
        "id": "377fcbf4bd708f79"
      },
      "source": [
        "\n",
        "These plots nicely show how we start with a very smoothly varying score function, as shown by the wide contour lines, towards the sharper peaks of the target distribution.\n",
        "The annealing successfully removes the need for a manually prescribed perturbation.\n",
        "\n",
        "\n",
        "## Score Summary\n",
        "\n",
        "Now we're almost at the goal, but before making the last step to derive the famous \"denoising\" task, let's summarize the ground we've covered so far:\n",
        "* Denoising score matching works well even for high-dimensional data such as images.\n",
        "* There's no need to backpropagate gradients through many steps, i.e. the method is much more scalable than CNFs or NeuralODEs.\n",
@@ -1019,17 +1065,17 @@
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "colab": {
      "provenance": []
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}