daviddoji/pbdl-book
1
0 Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

prob models introduction with code from Benjamin

Browse Source
This commit is contained in:
NT committed 2025-03-17 15:32:01 +01:00
1 parent 3b73717017
commit fb5229a105
7 files changed
+3786 -1
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

_toc.yml
+4
View File
@@ -35,6 +35,10 @@ parts:
- caption: Probabilistic Learning
chapters:
- file: probmodels-intro.md
- file: probmodels-normflow.ipynb
- file: probmodels-score.ipynb
- file: probmodels-diffusion.ipynb
- file: probmodels-flowmatching.ipynb
- file: probmodels-ddpm-fm.ipynb
- file: probmodels-phys.md
- file: probmodels-sbisim.ipynb
probmodels-diffusion.ipynb
+788
View File
@@ -0,0 +1,788 @@
{
"cells": [
{
"metadata": {
"id": "fe02764e59760330"
},
"cell_type": "markdown",
"source": [
"## 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",
"$$\n",
"q(x_{1:T}|x_0) := \\prod_{t=1}^T q(x_t|x_{t-1})\n",
"\\text{ with }\n",
"q(x_t|x_{t-1}) = \\mathcal{N} \\big( \\sqrt{1-\\beta_t} x_{t-1}, \\beta_t I \\big) ~.\n",
"$$\n",
"\n",
"The process that we've been considering in the previous sections above (the \"Langevin dynamics\") then represents the **reverse** process. We'll directly denote it with $p_\\theta$ as this is what we want to learn later on. The reverse process over the diffusion time $t$ is given by:\n",
"\n",
"$$\n",
"p_\\theta(x_{0:T}) := p(x_T) \\prod_{t=1}^T p_\\theta(x_{t-1}|x_t)\n",
"\\text{ with }\n",
"p_\\theta(x_{t-1}|x_t) = \\mathcal{N} \\big( \\mu_\\theta(x_t, t), \\beta_t I \\big)\n",
"$$\n",
"\n",
"which looks quite simple at first, but we've simply hidden all the difficulties in $\\mu_\\theta$, the learned mean of our distributions, giving the key building block of $p_\\theta$.\n",
"\n",
"Thanks to our construction with the Gaussians having standard deviation $\\beta_t$, we have an analytic form for every step along the way:\n",
"Given a data point $x_0$, we can sample the noisy latent state $x_t$ from the forward Markov chain via\n",
"\n",
"$$\n",
" q(x_t|x_0) = \\mathcal{N}(x_t, \\sqrt{\\overline{\\alpha}_t}x_0, (1-\\overline{\\alpha}_t)I)) ,\n",
"$$\n",
"\n",
"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",
"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",
"The exact marginal likelihood of the diffusion process involves an intractable integral, but the step-wise accumulation of Gaussian noise represents a variational lower bound that we can use instead. Via _Jensen's inequality_ for the log marginal likelihood, we obtain an expression for the _evidence lower bound_ (ELBO)\n",
"\n",
"$$\n",
"\\log p_\\theta(x_0) \\geq \\mathbb{E}_{q(x_1, \\ldots, x_T | x_0)} \\left[ \\log \\frac{p_\\theta(x_0, x_1, \\ldots, x_T)}{q(x_1, \\ldots, x_T | x_0)} \\right].\n",
"$$\n",
"\n",
"Expanding and simplifying the terms we get a lower bound for the maximum likelihood objective:\n",
"$$\n",
"\\mathbb{E}[- \\log p_\\theta(x_0)] \\leq \\mathbb{E} \\left[ - \\log p(x_T) - \\sum_{t \\geq 1} \\log \\frac{p_\\theta(x_{t-1}|x_t)} {q(x_t|x_{t-1})} \\right]\n",
"$$\n",
"\n",
"Note that we've also made use of the Bayes conditioning trick from the top here. This training objective can be reformulated in terms of a KL divergence, to give three terms:\n",
"\n",
"$$\n",
"\\begin{aligned}\n",
"\\mathbb{E} & \\left[ L_T + L_{t-1} + L_1 \\right] = \\\\\n",
"\\mathbb{E} & \\left[ \\mathrm{KL}(q(x_T|x_0) ||p(x_T)) + \\sum_{t > 1} \\mathrm{KL}(q(x_{t-1}|x_t, x_0)|| p_\\theta(x_{t-1}|x_t)) - \\log p_\\theta(x_0|x_1) \\right]\n",
"\\end{aligned}\n",
"$$\n",
"\n",
"This looks complicated, but the three terms have distinct meanings:\n",
"- the first one, $L_T$ does not depend on $\\theta$;\n",
"- the last one, $L_1$, is easy to train with;\n",
"- and the middle $L_{t-1}$ term is just the KL divergence between two known Gaussian distributions.\n",
"\n",
"Hence, the middle term $L_{t-1}$ is key, and we actually have the Gaussians under control with the $\\alpha$ and $\\beta$ terms above. For the forward process the corresponding expressions are a bit lengthy, but we can write them out. This gives the conditional Gaussian\n",
"$q(x_{t-1}|x_t,x_0) = \\mathcal{N} \\big( \\tilde{\\mu}_t(x_t, x_0), \\tilde{\\beta}_tI \\big)$, with mean and standard deviation being\n",
"\n",
"$$\n",
"\\tilde{\\mu}_t(x_t, x_0) = \\frac{\\sqrt{\\overline{\\alpha}_{t-1}}\\beta_t}{1-\\overline{\\alpha}_t} x_0 + \\frac{\\sqrt{\\alpha_t}(1-\\overline{\\alpha}_{t-1})}{1-\\overline{\\alpha}_t}x_t\n",
"\\text{ and }\n",
"\\tilde{\\beta}_t = \\frac{1 - \\overline{\\alpha}_{t-1}}{1- \\overline{\\alpha}_t} \\beta_t\n",
"$$\n",
"\n",
"This actually gives us a feasible training objective that has been used in a number of papers, but we don't need to stop there. Rather we can further simplify this construction by realizing that we're dealing with noise $\\epsilon$ being added to our means ($x+\\epsilon$). Instead of predicting the means (the signal $x$ itself), predicting the noise $\\epsilon$ on top of it is easier, and by subtracting it we likewise get $x$. In the terminology above, we get\n",
"\n",
"$$\n",
"x_t(x_0, \\epsilon) = \\sqrt{\\overline{a}_t}x_0 + \\sqrt{1-\\overline{\\alpha}_t} \\epsilon \\text{ for } \\epsilon \\sim \\mathcal{N}(0,I) ,\n",
"$$ (ddpm-xt)\n",
"\n",
"and instead of predicting the mean $\\mu_\\theta(x_t, t)$ we predict the noise $\\epsilon_\\theta(x_t,t)$ from which we can compute the mean via\n",
"\n",
"$$\n",
"\\mu_\\theta(x_t, t) = \\frac{1}{\\sqrt{\\alpha_t}} \\left( x_t - \\frac{1-\\alpha_t}{\\sqrt{1-\\overline{\\alpha}_t}} \\epsilon_\\theta(x_t, t)\\right) .\n",
"$$\n",
"\n",
"This simplifies the loss term $L_{t-1}$ from above to\n",
"\n",
"$$\n",
"\\mathbb{E}_{x_0, \\epsilon} \\left[ \\frac{\\beta_t^2}{2 \\beta_t \\alpha_t (1-\\overline{\\alpha}_t)} || \\epsilon - \\epsilon_\\theta(\\sqrt{\\overline{\\alpha}}_t x_0 + \\sqrt{1-\\overline{\\alpha}_t} \\epsilon, t) ||^2 \\right] + C ,\n",
"$$\n",
"\n",
"and even better, it turns out we can remove the first $\\beta^2$ factor. People have noticed that it actually provides a sub-optimal scaling, and removing it increases the weighting for samples with small $\\beta$ values at the end of the reverse chain. This gives a simplified denosing loss of\n",
"\n",
"$$\n",
"L_\\mathrm{DM}(\\theta) = \\mathbb{E}_{t, x_0, \\epsilon} \\left[ || \\epsilon - \\epsilon_\\theta \\big (\n",
" x_t(x_0, \\epsilon), t\n",
" \\big) ||^2 \\right] ,\n",
"$$\n",
"\n",
"where can easily compute correct values for\n",
"$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",
"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",
"```{figure} resources/prob02-ddpm.jpg\n",
"---\n",
"height: 240px\n",
"name: probmod-ddpm\n",
"---\n",
"DDPM process (in contrast to CNFs): denoising the simple base distribution (left) over a large number of denoising steps to arrive at a sample from the target distribution (right).\n",
"```\n",
"\n",
"\n",
"The training procedure can be summarized in pseudo-code with:\n",
"\n",
"```{prf:algorithm} DDPM Training\n",
":label: ddpm-training\n",
"\n",
"**Inputs** Target distribution $q$, noise schedule $\\overline{\\alpha}$\n",
"\n",
"**Output** Trained noise estimator NN $\\epsilon_\\theta$\n",
"\n",
"1. repeat until converged:\n",
"\n",
" 2. $x_0 \\sim q(x_0)$\n",
" 3. $t \\sim \\text{Uniform}(1,\\dots,T)$\n",
" 4. $\\epsilon \\sim \\mathcal N(0,I)$\n",
" 5. $x_t = \\sqrt{\\overline{a}_t}x_0 + \\sqrt{1-\\overline{\\alpha}_t} \\epsilon$\n",
" 6. optimizer_step($ \\nabla_\\theta ||\\epsilon - \\epsilon_\\theta(x_t,t)||^2 $)\n",
"\n",
"7. return $\\epsilon_\\theta$\n",
"\n",
"```\n",
"\n",
"We uniformly sample an integer time $t$, and compute the noise sample $x_t$. This is quite simple, just blend a sample from step 0 of our forward process, $x_0$ with uniform noise according to the weight $\\overline{\\alpha}_t$. Let the network try to estimate the noise given $x_t$ and the time $t$, and then do a step with your favorite optimizer (usually Adam) along the gradient of the $L^2$ loss. Note that the noise schedule $\\overline{\\alpha}_t$ and the corresponding coefficients can be precomputed for a chosen number of steps $T$. This reduces the calculation to a simple lookup at training time. The loss is particularly simple as we can omit the weighting of the noise schedule as explained above.\n",
"\n",
"For inference, we need to properly evaluate the weighting terms, which makes the process slightly more complicated. The weighting can be precomputed just like for training, but the expensive bit is that we now need to go through all $T$ steps to fully denoise the sample:\n",
"\n",
"```{prf:algorithm} DDPM Inference\n",
":label: ddpm-inference\n",
"\n",
"**Inputs** Noise estimator $\\epsilon_\\theta$, noise schedule $\\alpha,\\overline{\\alpha},\\sigma$\n",
"\n",
"**Output** Sample from distribution $p$\n",
"\n",
"1. $x \\sim \\mathcal N(0,I)$\n",
"2. for $t = T,\\dots,1$:\n",
"\n",
" 3. if $t>1$: $z \\sim \\mathcal N(0,I)$\n",
" 4. else: $z = 0$\n",
" 5. $x = \\frac{1}{\\sqrt{\\alpha_t}} \\left( x - \\frac{1-\\alpha_t}{\\sqrt{1-\\overline{\\alpha}_t}} \\epsilon_\\theta(x, t)\\right) + \\sigma_t z$\n",
"\n",
"6. return $x$\n",
"```\n",
"\n",
"This means we have to sequentially evaluate the trained neural network $\\epsilon_\\theta$ for $T$ times, and the original DDPM papers actually used $T=1000$ to obtain very high quality samples. In practice, $T$ can be reduced significantly, by ca. one order of magnitude, without while still giving excellent results, but this nonetheless means we have to evaluate the network around 100 times.\n",
"\n",
"While this is clearly more expensive than the deterministic networks we have used in previous chapters, which produced the mean output in a single pass, the powerful, fundamental change is that the DDPM network can learn and reliably reproduce distributions. I.e., if your data for a single flow input contains 20% solutions that spin left, and 80% that turn right, the trained DDPM network will reproduce these solutions with the right probability. When you repeatedly run inference with different initial noise values for $x$, you will see left spins with a probability of 0.2, and right spinning solutions with 0.8. The network itself is of course still deterministic, for a constant initial $x$ in line 1 of the inference algorithm, it will always produce the same output.\n",
"\n",
"Another powerful aspect that we'll get back to is also that these networks can be reliable conditioned to steer the generated samples. This is highly important for our initial goals from simulation-based inference: we like to learn the posterior distribution $p(x|y)$ that produces the correct outputs for a given observation $y$. Looking ahead, three important extensions of the algorithm so far are 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. 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"
},
{
"metadata": {
"id": "2c66c145697a1099"
},
"cell_type": "markdown",
"source": [
"## Train a DDPM on the Gaussian Toy Dataset"
],
"id": "2c66c145697a1099"
},
{
"metadata": {
"ExecuteTime": {
"end_time": "2025-03-17T00:09:15.870346Z",
"start_time": "2025-03-17T00:09:15.866991Z"
},
"id": "eec0a45e5bc00691"
},
"cell_type": "code",
"source": [
"import numpy as np\n",
"\n",
"class GaussianMixture:\n",
" def __init__(self, parameters):\n",
"\n",
" self.parameters = parameters\n",
" self.distributions = [\n",
" {\n",
" 'mean': np.array(dist['mean']),\n",
" 'std': np.array(dist['std']),\n",
" 'cov': np.diag(np.array(dist['std']) ** 2)\n",
" }\n",
" for dist in parameters\n",
" ]\n",
"\n",
" def sample(self, num_samples):\n",
" samples = []\n",
" num_distributions = len(self.distributions)\n",
" for _ in range(num_samples):\n",
" idx = np.random.randint(num_distributions) # Choose a random Gaussian\n",
" dist = self.distributions[idx]\n",
" sample = np.random.multivariate_normal(mean=dist['mean'], cov=dist['cov'])\n",
" samples.append(sample)\n",
" return np.array(samples)\n",
"\n",
"parameters = [\n",
" {\"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
},
{
"metadata": {
"id": "c61074cf783a4fa9"
},
"cell_type": "markdown",
"source": [
"### Define the Forward Process"
],
"id": "c61074cf783a4fa9"
},
{
"metadata": {
"ExecuteTime": {
"end_time": "2025-03-17T00:09:17.244046Z",
"start_time": "2025-03-17T00:09:16.341782Z"
},
"colab": {
"base_uri": "https://localhost:8080/",
"height": 1000
},
"id": "59216c5f42a1a070",
"outputId": "d891b00e-cd74-44a0-e202-4f29a3030603"
},
"cell_type": "code",
"source": [
"import torch\n",
"import matplotlib.pyplot as plt\n",
"import seaborn as sns\n",
"sns.set_theme(style=\"ticks\")\n",
"\n",
"class DDPMForwardProcess:\n",
" def __init__(self, num_timesteps=1000):\n",
" self.num_timesteps = num_timesteps\n",
"\n",
" # Standard beta schedule from DDPM paper\n",
" scale = 1000 / num_timesteps\n",
" beta_start = scale * 0.0001\n",
" beta_end = scale * 0.02\n",
" self.betas = torch.linspace(beta_start, beta_end, num_timesteps)\n",
"\n",
" self.alphas = 1 - self.betas\n",
" self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)\n",
" self.sqrt_betas = torch.sqrt(self.betas)\n",
" self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)\n",
" self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1 - self.alphas_cumprod)\n",
"\n",
" def forward_process(self, x_0, t):\n",
" epsilon = torch.randn_like(x_0)\n",
" mean = self.sqrt_alphas_cumprod[t] * x_0\n",
" std = self.sqrt_one_minus_alphas_cumprod[t]\n",
" return mean + std * epsilon, epsilon\n",
"\n",
" def visualize_forward_process(self, mixture, num_samples=1000):\n",
"\n",
" timesteps_to_show=[0, 200, 400, 600, 800, 999]\n",
" x_0 = torch.FloatTensor(mixture.sample(num_samples))\n",
"\n",
" fig, axes = plt.subplots(2, 3, figsize=(15, 10))\n",
" axes = axes.ravel()\n",
"\n",
" xx, yy = np.mgrid[-5:5:100j, -5:5:100j]\n",
" positions = np.vstack([xx.ravel(), yy.ravel()])\n",
"\n",
" for idx, t in enumerate(timesteps_to_show):\n",
" x_t, _ = self.forward_process(x_0, t)\n",
" samples = x_t.numpy()\n",
"\n",
" from scipy.stats import gaussian_kde\n",
" kernel = gaussian_kde(samples.T)\n",
" density = np.reshape(kernel(positions).T, xx.shape)\n",
"\n",
" ax = axes[idx]\n",
"\n",
" ax.contour(xx, yy, density, levels=10, alpha=1.0)\n",
"\n",
" ax.scatter(samples[:, 0], samples[:, 1], alpha=1.0, s=4)\n",
"\n",
" ax.set_xlim(-5, 5)\n",
" ax.set_ylim(-5, 5)\n",
" ax.set_aspect('equal')\n",
" ax.set_title(f't = {t}')\n",
"\n",
" plt.suptitle(\"Forward process (data to noise)\")\n",
"\n",
" plt.tight_layout()\n",
" plt.show()\n",
"\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
},
{
"metadata": {
"id": "b08f5dfe0fe31649"
},
"cell_type": "markdown",
"source": [
"### Setup the Dataset and Dataloader"
],
"id": "b08f5dfe0fe31649"
},
{
"metadata": {
"ExecuteTime": {
"end_time": "2025-03-17T00:02:22.356614Z",
"start_time": "2025-03-17T00:01:52.958922Z"
},
"id": "1e531e6c995b4609"
},
"cell_type": "code",
"source": [
"from torch import nn\n",
"from torch.optim import Adam\n",
"from torch.utils.data import Dataset, DataLoader\n",
"import numpy as np\n",
"\n",
"class ForwardDiffusionDataset(Dataset):\n",
" def __init__(self, num_samples, input_dim,\n",
" T, mixture, beta_value=0.02):\n",
" super().__init__()\n",
" self.num_samples = num_samples\n",
" self.input_dim = input_dim\n",
" self.T = T\n",
" self.beta_value = beta_value\n",
" self.gm = mixture\n",
"\n",
" self.betas = beta_value * torch.ones(T)\n",
" self.alphas = 1 - self.betas\n",
" self.alpha_bar = torch.cumprod(self.alphas, dim=0)\n",
"\n",
" self.x0 = torch.FloatTensor(self.gm.sample(num_samples))\n",
"\n",
" def __len__(self):\n",
" return self.num_samples\n",
"\n",
" def __getitem__(self, idx):\n",
" x0_sample = self.x0[idx]\n",
" t = torch.randint(low=1, high=self.T + 1, size=(1,)).item()\n",
" eps = torch.randn(self.input_dim)\n",
" sqrt_alpha_bar = torch.sqrt(self.alpha_bar[t - 1])\n",
" sqrt_one_minus_alpha_bar = torch.sqrt(1 - self.alpha_bar[t - 1])\n",
"\n",
" 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
},
{
"metadata": {
"id": "5ba5da2e4725ffb1"
},
"cell_type": "markdown",
"source": [
"### Define the Score Network"
],
"id": "5ba5da2e4725ffb1"
},
{
"metadata": {
"id": "5e711111525c923d"
},
"cell_type": "code",
"outputs": [],
"execution_count": 8,
"source": [
"class ScoreNet(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",
" nn.Linear(in_dim + time_dim, hidden_dim),\n",
" nn.ReLU(),\n",
" nn.Linear(hidden_dim, hidden_dim),\n",
" nn.ReLU(),\n",
" nn.Linear(hidden_dim, out_dim)\n",
" )\n",
"\n",
" def forward(self, x, t):\n",
" t = t.unsqueeze(1)\n",
" xt = torch.cat([x, t], dim=1)\n",
" return self.net(xt)"
],
"id": "5e711111525c923d"
},
{
"metadata": {
"id": "ad8de5772ec3b37b"
},
"cell_type": "markdown",
"source": [
"### Initialization and Training"
],
"id": "ad8de5772ec3b37b"
},
{
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "1027e7e50cc6fe83",
"outputId": "214a0863-f2b0-4d4e-c801-07b0ba654ce9"
},
"cell_type": "code",
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"Epoch 1/100, Loss: 0.1529\n",
"Epoch 2/100, Loss: 0.0871\n",
"Epoch 3/100, Loss: 0.0863\n",
"Epoch 4/100, Loss: 0.0747\n",
"Epoch 5/100, Loss: 0.0759\n",
"Epoch 6/100, Loss: 0.0737\n",
"Epoch 7/100, Loss: 0.0715\n",
"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",
"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",
"Epoch 93/100, Loss: 0.0595\n",
"Epoch 94/100, Loss: 0.0570\n",
"Epoch 95/100, Loss: 0.0632\n",
"Epoch 96/100, Loss: 0.0609\n",
"Epoch 97/100, Loss: 0.0570\n",
"Epoch 98/100, Loss: 0.0580\n",
"Epoch 99/100, Loss: 0.0606\n",
"Epoch 100/100, Loss: 0.0624\n"
]
}
],
"execution_count": 10,
"source": [
"num_samples = 10000\n",
"input_dim = 2\n",
"T = 1000\n",
"dataset = ForwardDiffusionDataset(num_samples=num_samples, input_dim=input_dim, T=T, mixture=mixture, beta_value=0.02)\n",
"batch_size = 64\n",
"data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)\n",
"\n",
"# %% 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",
"optimizer = Adam(model.parameters(), lr=1e-3)\n",
"loss_fn = nn.MSELoss()\n",
"\n",
"epochs = 100\n",
"model.train()\n",
"for epoch in range(1, epochs + 1):\n",
" epoch_loss = 0\n",
" for x_t_batch, t_batch, eps_batch in data_loader:\n",
" optimizer.zero_grad()\n",
" # Normalize time t to range [0,1]\n",
" t_norm = t_batch / T\n",
" eps_pred = model(x_t_batch, t_norm)\n",
" loss = loss_fn(eps_pred, eps_batch)\n",
" loss.backward()\n",
" 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",
"\n",
"betas = 0.02 * torch.ones(T)\n",
"alphas = 1 - betas\n",
"alpha_bar = torch.cumprod(alphas, dim=0)"
],
"id": "1027e7e50cc6fe83"
},
{
"metadata": {
"id": "24d5227d55d6d5e4"
},
"cell_type": "markdown",
"source": [
"### Inference with the Reverse Process"
],
"id": "24d5227d55d6d5e4"
},
{
"metadata": {
"ExecuteTime": {
"end_time": "2025-03-17T00:02:25.210393Z",
"start_time": "2025-03-17T00:02:24.998920Z"
},
"id": "bb1b68299909aa1e"
},
"cell_type": "code",
"source": [
"n_infer = 1000\n",
"\n",
"with torch.no_grad():\n",
" x = torch.randn(n_infer, input_dim)\n",
"\n",
" record_steps = {999, 500, 200, 100, 50, 1}\n",
" snapshots = {}\n",
"\n",
" for t in range(T, 0, -1):\n",
"\n",
" t_tensor = (torch.ones(n_infer) * t / T).float()\n",
" eps_theta = model(x, t_tensor)\n",
"\n",
" beta_t = betas[t - 1]\n",
" alpha_t = alphas[t - 1]\n",
" alpha_bar_t = alpha_bar[t - 1]\n",
"\n",
" coef = beta_t / torch.sqrt(1 - alpha_bar_t)\n",
" x = (1 / torch.sqrt(alpha_t)) * (x - coef * eps_theta)\n",
"\n",
" if t > 1:\n",
" sigma_t = torch.sqrt(beta_t)\n",
" x = x + sigma_t * torch.randn_like(x)\n",
"\n",
" if t in record_steps:\n",
" snapshots[t] = x.clone().cpu().numpy()"
],
"id": "bb1b68299909aa1e",
"outputs": [],
"execution_count": 11
},
{
"metadata": {
"id": "df3ce91df721e024"
},
"cell_type": "markdown",
"source": [
"### Visualize the Sampled Trajectories"
],
"id": "df3ce91df721e024"
},
{
"metadata": {
"ExecuteTime": {
"end_time": "2025-03-17T00:02:26.532752Z",
"start_time": "2025-03-17T00:02:26.526617Z"
},
"id": "cafff2d56a61be60"
},
"cell_type": "code",
"source": [
"import matplotlib.pyplot as plt\n",
"\n",
"def get_angle_colors(positions):\n",
" angles = np.arctan2(positions[:, 1], positions[:, 0])\n",
" angles_deg = (np.degrees(angles) + 360) % 360\n",
" colors = np.zeros((len(positions), 3))\n",
" for i, angle in enumerate(angles_deg):\n",
" segment = int(angle / 120)\n",
" local_angle = angle - segment * 120\n",
" if segment == 0: # 0 degrees to 120 degrees (R->G)\n",
" colors[i] = [1 - local_angle/120, local_angle/120, 0]\n",
" elif segment == 1: # 120 degrees to 240 degrees (G->B)\n",
" colors[i] = [0, 1 - local_angle/120, local_angle/120]\n",
" else: # 240 degrees to 360° (B->R)\n",
" colors[i] = [local_angle/120, 0, 1 - local_angle/120]\n",
"\n",
" return colors\n",
"\n",
"def visualize_snapshots(snapshots):\n",
"\n",
" sorted_steps = sorted(snapshots.keys(), reverse=True)\n",
" n_plots = len(sorted_steps)\n",
"\n",
" fig, axes = plt.subplots(2, 3, figsize=(5 * 3, 8))\n",
" axes = axes.ravel()\n",
"\n",
" if n_plots == 1:\n",
" axes = [axes]\n",
"\n",
" # Create grid for density visualization\n",
" xx, yy = np.mgrid[-5:5:100j, -5:5:100j]\n",
" positions = np.vstack([xx.ravel(), yy.ravel()])\n",
"\n",
" idx = 0\n",
" for ax, t in zip(axes, sorted_steps):\n",
"\n",
" samples = snapshots[t]\n",
"\n",
" if idx == 0:\n",
" c = get_angle_colors(samples)\n",
" idx += 1\n",
"\n",
" x_0 = torch.FloatTensor(mixture.sample(num_samples))\n",
" x_t, _ = ddpm.forward_process(x_0, t)\n",
" samples_ = x_t.numpy()\n",
"\n",
" from scipy.stats import gaussian_kde\n",
" kernel = gaussian_kde(samples_.T)\n",
" density = np.reshape(kernel(positions).T, xx.shape)\n",
"\n",
" ax.contour(xx, yy, density, levels=10, alpha=0.8)\n",
"\n",
" ax.scatter(samples[:, 0], samples[:, 1], alpha=1.0, s=5, color=c)\n",
" ax.set_title(f\"t = {t}\")\n",
" ax.set_xlabel(\"\")\n",
" ax.set_ylabel(\"\")\n",
" ax.grid(True)\n",
"\n",
" plt.suptitle(\"Diffusion Model Reverse Process\")\n",
" plt.tight_layout()\n",
" plt.show()\n"
],
"id": "cafff2d56a61be60",
"outputs": [],
"execution_count": 12
},
{
"metadata": {
"ExecuteTime": {
"end_time": "2025-03-17T00:02:33.517119Z",
"start_time": "2025-03-17T00:02:27.762131Z"
},
"colab": {
"base_uri": "https://localhost:8080/",
"height": 824
},
"id": "5cd9e3d82a956ac6",
"outputId": "15be7112-1732-406f-ad94-bbd35441215b"
},
"cell_type": "code",
"source": [
"visualize_snapshots(snapshots)"
],
"id": "5cd9e3d82a956ac6",
"outputs": [
{
"output_type": "display_data",
"data": {
"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": {}
}
],
"execution_count": 13
}
],
"metadata": {
"kernelspec": {
"display_name": "diffusion",
"language": "python",
"name": "diffusion"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 2
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython2",
"version": "2.7.6"
},
"colab": {
"provenance": []
}
},
"nbformat": 4,
"nbformat_minor": 5
}
probmodels-flowmatching.ipynb
+562
View File
@@ -0,0 +1,562 @@
{
"cells": [
{
"metadata": {
"id": "18dcb0d9d85f5323"
},
"cell_type": "markdown",
"source": [
"\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",
"$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",
"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",
"\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",
"$$\n",
" dx/dt = v_\\theta(x,t) ,\n",
"$$\n",
"\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",
"\n",
"$$\n",
"\\begin{align}\n",
" \\mathcal{L}_\\mathrm{FM}(\\theta) = \\mathbb{E}_{t\\sim \\mathcal{U}(0,1), x \\sim p_t(x)} ~||~ v_\\theta(x,t) - u_t(x)\n",
" ~||^2.\n",
"\\end{align}\n",
"$$ (eq-flow-matching)\n",
"\n",
"In order to evaluate this loss, we need to sample from the probability distribution $p_t(x)$ and we need to know the velocity $u_t(x)$.\n",
"In contrast to the perturbed data that we used for score matching, we'll solve this problem differently for flow matching:\n",
"we apply a trick by introducing a latent variable $z$ distributed according to $q(z)$ and define the conditional likelihoods $p_t(x|z)$ that depend on the latent variable so that $p_t(x) = \\int p_t(x|z)q(z)dz$. Interestingly, if the conditional likelihoods are generated by the velocities $u_t(x|z)$, then the velocity $u_t(x)$\n",
"can be expressed in terms of $u_t(x|z)$ and $p_t(x|z)$ as the following expectation:\n",
"$u_t(x) = \\mathbb{E}_{q(z)}[u_t(x|z)p_t(x|z)/p_t(x)]$.\n",
"Here we have the freedom to choose paths $p_t(x|z)$ that are easy to sample from and for which we know the generating velocities $u_t(x|z)$.\n",
"So, as long as we can generate enough velocities conditioned on $z$, we can obtain the expectation over the course of producing all these samples, and obtain the right generating velocity for training our network.\n",
"\n",
"Flow networks can then be trained with the _conditional flow matching_ loss\n",
"\n",
"$$\n",
"\\begin{align}\n",
" \\mathcal{L}_\\mathrm{CFM}(\\phi) = \\mathbb{E}_{q(z,t),p_t(x|z)} ~||~ v_\\phi(x,t)-u_t(x|z) ~||^2.\n",
"\\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 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",
"\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",
"It turns out, the \"simplest possibility\" of aiming for straight, non-crossing paths is a great choice. The goal is a velocity from the current state $x$ directly towards the target $x_1$ at time $t$ with\n",
"$(x_1-x) / (1-t)$. The only slight complication is that we need to consider a minimal amount of noise $\\sigma_\\mathrm{min} > 0$ at time $t=1$ to ensure that despite having discrete samples, we keep a continuous distribution. This aspect is in line with the perturbed samples from score matching above. Hence $\\sigma_\\mathrm{min}$ influences the end points of the prescribed paths and its probability distribution.\n",
"\n",
"Plugging it into the equations, we consider the coupling $q(z) = p_1(x)$ together with conditional probability and generating velocity\n",
"\n",
"$$\n",
"\\begin{align}\n",
" p_t(x|x_1) &= \\mathcal{N}( t x_1, (1-(1-\\sigma_\\mathrm{min})t)I) \\\\\n",
" u_t(x|x_1) &= \\frac{x_1-(1-\\sigma_\\mathrm{min})x}{1-(1-\\sigma_\\mathrm{min})t} ~.\n",
"\\end{align}\n",
"$$ (conditional_ot_prob)\n",
"\n",
"Conditioned on $x_1$, this coupling transports a point $x_0 \\sim \\mathcal{N}(0,I)$ from the sampling distribution to the posterior distribution on the linear trajectory $t x_1$ ending in $x_1$. At the same time it decreases the standard deviation from $1$ to a smoothing constant $\\sigma_\\mathrm{min}$. In this case, the transport path coincides with the optimal transport between two Gaussian distributions.\n",
"\n",
"So it might seem - on first sight - that we've made a huge step back: we're back to transforming distributions, a concept for which we argued above in the section on _normalizing flows_ that it was a bad idea. However, with the methodology from score matching and denoising we've arrived at a fundamentally different (and more powerful) way to transform distributions. E.g., we're freed from constraints to preserve densities, and have a very tractable and convenient learning objective. In contrast to _Neural-ODEs_ we can train with single steps, instead of having to backpropagate through the full chain from end to start.\n",
"\n",
"```{figure} resources/prob03-fm.jpg\n",
"---\n",
"height: 240px\n",
"name: probmod-fm\n",
"---\n",
"Flow matching (in contrast to CNFs and DDPMs): the simple base distribution (left) is transfored by ODE integration of the inferred flow velocity to obtain a sample from the target distribution (right). Due to the linear paths, only very few integration steps are necessary.\n",
"```\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"
},
{
"metadata": {
"ExecuteTime": {
"end_time": "2025-03-16T22:58:52.827556Z",
"start_time": "2025-03-16T22:58:52.822630Z"
},
"id": "c041087b8b9191f0"
},
"cell_type": "code",
"source": [
"import numpy as np\n",
"\n",
"class GaussianMixture:\n",
" def __init__(self, parameters):\n",
"\n",
" self.parameters = parameters\n",
" self.distributions = [\n",
" {\n",
" 'mean': np.array(dist['mean']),\n",
" 'std': np.array(dist['std']),\n",
" 'cov': np.diag(np.array(dist['std']) ** 2)\n",
" }\n",
" for dist in parameters\n",
" ]\n",
"\n",
" def sample(self, num_samples):\n",
" samples = []\n",
" num_distributions = len(self.distributions)\n",
" for _ in range(num_samples):\n",
" idx = np.random.randint(num_distributions) # Choose a random Gaussian\n",
" dist = self.distributions[idx]\n",
" sample = np.random.multivariate_normal(mean=dist['mean'], cov=dist['cov'])\n",
" samples.append(sample)\n",
" return np.array(samples)\n",
"\n",
"parameters = [\n",
" {\"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
},
{
"metadata": {
"id": "7f9d408886c23020"
},
"cell_type": "markdown",
"source": [
"### Setup Dataset and Dataloader"
],
"id": "7f9d408886c23020"
},
{
"metadata": {
"ExecuteTime": {
"end_time": "2025-03-16T23:32:38.653547Z",
"start_time": "2025-03-16T23:32:38.648261Z"
},
"id": "8ac07b3409739df9"
},
"cell_type": "code",
"source": [
"import torch\n",
"import torch.nn as nn\n",
"import torch.optim as optim\n",
"from torch.utils.data import Dataset, DataLoader\n",
"\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",
" self.n_samples = n_samples\n",
" self.mixture = mixture\n",
" self.sigma_min = sigma_min\n",
" self.x1 = mixture.sample(n_samples)\n",
"\n",
" def __len__(self):\n",
" return self.n_samples\n",
"\n",
" def __getitem__(self, idx):\n",
"\n",
" x0 = np.random.multivariate_normal([0.0, 0.0], np.eye(2), 1)[0]\n",
" t = np.random.rand() # scalar in [0,1]\n",
" x1 = self.x1[idx]\n",
"\n",
" x_t = (1 - ( 1 - self.sigma_min) * t) * x0 + t * x1\n",
" u_t = (x1 - x0)\n",
" x_t = torch.tensor(x_t, dtype=torch.float32)\n",
" t = torch.tensor([t], dtype=torch.float32)\n",
" u_t = torch.tensor(u_t, dtype=torch.float32)\n",
"\n",
" return x_t, t, u_t\n"
],
"id": "8ac07b3409739df9",
"outputs": [],
"execution_count": 2
},
{
"metadata": {
"id": "471468fdc74f2af2"
},
"cell_type": "markdown",
"source": [
"### Small MLP for Learning the Velocity"
],
"id": "471468fdc74f2af2"
},
{
"metadata": {
"ExecuteTime": {
"end_time": "2025-03-16T23:32:39.251649Z",
"start_time": "2025-03-16T23:32:39.246074Z"
},
"id": "e298ca1a4bee5d7f"
},
"cell_type": "code",
"source": [
"class VelocityNet(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",
" nn.Linear(in_dim + time_dim, hidden_dim),\n",
" nn.ReLU(),\n",
" nn.Linear(hidden_dim, hidden_dim),\n",
" nn.ReLU(),\n",
" nn.Linear(hidden_dim, out_dim)\n",
" )\n",
"\n",
" 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
},
{
"metadata": {
"id": "832f39392dda6db4"
},
"cell_type": "markdown",
"source": [
"### Train a Velocity Model"
],
"id": "832f39392dda6db4"
},
{
"metadata": {
"ExecuteTime": {
"end_time": "2025-03-16T23:33:04.270517Z",
"start_time": "2025-03-16T23:32:40.419723Z"
},
"colab": {
"base_uri": "https://localhost:8080/"
},
"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",
"text": [
"Epoch 1/50, Loss: 2.8302\n",
"Epoch 2/50, Loss: 2.2797\n",
"Epoch 3/50, Loss: 2.2231\n",
"Epoch 4/50, Loss: 2.1215\n",
"Epoch 5/50, Loss: 2.1258\n",
"Epoch 6/50, Loss: 2.0737\n",
"Epoch 7/50, Loss: 2.1223\n",
"Epoch 8/50, Loss: 2.0625\n",
"Epoch 9/50, Loss: 2.0656\n",
"Epoch 10/50, Loss: 2.0527\n",
"Epoch 11/50, Loss: 2.0776\n",
"Epoch 12/50, Loss: 2.0550\n",
"Epoch 13/50, Loss: 2.0521\n",
"Epoch 14/50, Loss: 2.0700\n",
"Epoch 15/50, Loss: 2.0398\n",
"Epoch 16/50, Loss: 2.0406\n",
"Epoch 17/50, Loss: 2.0620\n",
"Epoch 18/50, Loss: 2.0879\n",
"Epoch 19/50, Loss: 2.0883\n",
"Epoch 20/50, Loss: 2.0175\n",
"Epoch 21/50, Loss: 2.0247\n",
"Epoch 22/50, Loss: 2.0447\n",
"Epoch 23/50, Loss: 2.0342\n",
"Epoch 24/50, Loss: 2.0436\n",
"Epoch 25/50, Loss: 2.0334\n",
"Epoch 26/50, Loss: 2.0474\n",
"Epoch 27/50, Loss: 2.0547\n",
"Epoch 28/50, Loss: 2.0368\n",
"Epoch 29/50, Loss: 2.0394\n",
"Epoch 30/50, Loss: 2.0229\n",
"Epoch 31/50, Loss: 2.0394\n",
"Epoch 32/50, Loss: 2.0581\n",
"Epoch 33/50, Loss: 1.9796\n",
"Epoch 34/50, Loss: 2.0157\n",
"Epoch 35/50, Loss: 2.0325\n",
"Epoch 36/50, Loss: 2.0319\n",
"Epoch 37/50, Loss: 2.0511\n",
"Epoch 38/50, Loss: 2.0593\n",
"Epoch 39/50, Loss: 1.9919\n",
"Epoch 40/50, Loss: 2.0522\n",
"Epoch 41/50, Loss: 2.0709\n",
"Epoch 42/50, Loss: 2.0460\n",
"Epoch 43/50, Loss: 2.0688\n",
"Epoch 44/50, Loss: 2.0273\n",
"Epoch 45/50, Loss: 2.0711\n",
"Epoch 46/50, Loss: 2.0120\n",
"Epoch 47/50, Loss: 2.0344\n",
"Epoch 48/50, Loss: 2.0347\n",
"Epoch 49/50, Loss: 2.0537\n",
"Epoch 50/50, Loss: 2.0440\n"
]
}
],
"execution_count": 4
},
{
"metadata": {
"id": "4e03ffe8ace08e19"
},
"cell_type": "markdown",
"source": [
"### Inference via Solving the ODE"
],
"id": "4e03ffe8ace08e19"
},
{
"cell_type": "code",
"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": []
},
{
"metadata": {
"ExecuteTime": {
"end_time": "2025-03-17T00:06:20.371503Z",
"start_time": "2025-03-17T00:06:20.321439Z"
},
"id": "691f3460aa86e1a1"
},
"cell_type": "code",
"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",
"\n",
" t = torch.linspace(t_span[0], t_span[1], n_steps).to(x0.device)\n",
"\n",
" def ode_func(t, x):\n",
" t_tensor = t.expand(x.shape[0], 1)\n",
" return model(x, t_tensor)\n",
"\n",
" trajectory = odeint(ode_func, x0, t, method=method,\n",
" atol=1e-5, rtol=1e-5)\n",
"\n",
" return trajectory, t\n",
"\n",
"n_gen = 500\n",
"\n",
"x0_gen = torch.randn(n_gen, 2).to(device)\n",
"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
},
{
"metadata": {
"id": "ae7e42318e2841b2"
},
"cell_type": "markdown",
"source": [
"### Visualize the Sampling Trajectories"
],
"id": "ae7e42318e2841b2"
},
{
"metadata": {
"ExecuteTime": {
"end_time": "2025-03-17T00:06:56.685929Z",
"start_time": "2025-03-17T00:06:53.477554Z"
},
"colab": {
"base_uri": "https://localhost:8080/",
"height": 940
},
"id": "5d381b5a8e3514e1",
"outputId": "20fb889d-a765-4787-91bb-8bf4f4b96612"
},
"cell_type": "code",
"source": [
"import numpy as np\n",
"import matplotlib.pyplot as plt\n",
"\n",
"import seaborn as sns\n",
"sns.set_theme(style=\"ticks\", palette=\"pastel\")\n",
"\n",
"def get_angle_colors(positions):\n",
" angles = np.arctan2(positions[:, 1], positions[:, 0])\n",
" angles_deg = (np.degrees(angles) + 360) % 360\n",
" colors = np.zeros((len(positions), 3))\n",
" for i, angle in enumerate(angles_deg):\n",
" segment = int(angle / 120)\n",
" local_angle = angle - segment * 120\n",
" if segment == 0: # 0 degrees to 120 degrees (R->G)\n",
" colors[i] = [1 - local_angle/120, local_angle/120, 0]\n",
" elif segment == 1: # 120 degrees to 240 degrees (G->B)\n",
" colors[i] = [0, 1 - local_angle/120, local_angle/120]\n",
" else: # 240 degrees to 360° (B->R)\n",
" colors[i] = [local_angle/120, 0, 1 - local_angle/120]\n",
"\n",
" return colors\n",
"\n",
"# Desired time points to visualize\n",
"desired_times = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0]\n",
"\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",
"# Get the index for each desired time: here we choose the index with minimum absolute difference.\n",
"indices = [np.argmin(np.abs(time_np - t_val)) for t_val in desired_times]\n",
"\n",
"# Create subplots: We'll use 2 rows and 3 columns, one subplot per time.\n",
"fig, axes = plt.subplots(nrows=2, ncols=3, figsize=(15, 10))\n",
"axes = axes.ravel() # flatten the 2D array for easier indexing\n",
"\n",
"# Create grid for density visualization\n",
"xx, yy = np.mgrid[-5:5:100j, -5:5:100j]\n",
"positions = np.vstack([xx.ravel(), yy.ravel()])\n",
"\n",
"for i, idx in enumerate(indices):\n",
" ax = axes[i]\n",
" # Get the samples at this time point: shape (batch, 2)\n",
" x_t = trajectory[idx].detach().cpu().numpy()\n",
"\n",
" if i == 0:\n",
" c = get_angle_colors(x_t)\n",
"\n",
" x_0 = mixture.sample(5000)\n",
" t = time_np[idx]\n",
" eps = np.random.randn(5000, 2)\n",
" x_t_forward = t * x_0 + (1-t) * eps\n",
" samples_ = x_t_forward\n",
"\n",
" # Compute KDE for visualization\n",
" from scipy.stats import gaussian_kde\n",
" kernel = gaussian_kde(samples_.T)\n",
" density = np.reshape(kernel(positions).T, xx.shape)\n",
"\n",
" # Plot density contours\n",
" ax.contour(xx, yy, density, levels=10, alpha=0.8)\n",
"\n",
" # Scatter plot: plot each sample as a point\n",
" ax.scatter(x_t[:, 0], x_t[:, 1], alpha=0.5, s=10, color=c)\n",
" ax.set_title(f\"t = {time_np[idx]:.2f}\")\n",
" ax.set_xlabel(\"\")\n",
" ax.set_ylabel(\"\")\n",
" ax.grid(True)\n",
"\n",
"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
},
{
"metadata": {
"id": "4255a24220b673d4"
},
"cell_type": "markdown",
"source": [
"## 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",
"\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."
],
"id": "4255a24220b673d4"
}
],
"metadata": {
"kernelspec": {
"display_name": "diffusion",
"language": "python",
"name": "diffusion"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 2
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython2",
"version": "2.7.6"
},
"colab": {
"provenance": []
}
},
"nbformat": 4,
"nbformat_minor": 5
}
probmodels-intro.md
+117
View File
@@ -0,0 +1,117 @@
Introduction to Probabilistic Learning
=======================
So far we've treated the target function $f(x)=y$ as being deterministic, with a unique solution $y$ for every input. That's certainly a massive simplification: in practice, solutions can be ambiguous, our learned model might mix things up, and both effects could show up in combination. This all calls for moving towards a probabilistic setting, which we'll address here. The machinery from previous sections will come in handy, as the probabilistic viewpoint essentially introduces another dimension for the problem. Instead of a single $y$, we now have a multitude of solutions drawn from a distribution $Y$, each with a probability $p_Y(y)$, often shortened to $p(y)$.
Samples $y \sim p(y)$ drawn from the distribution should follow this probability, so that we can distinguish rare and and frequent cases.
To summarize, instead of individual solutions $y$ we're facing a large number of samples $y \sim p(y)$.
![Divider](resources/divider5.jpg)
## Uncertainty
All measurements, models, and discretizations that we are working with exhibit uncertainties. For measurements and observations, they typically appear in the form of measurement errors. Model equations, on the other hand, usually encompass only parts of a system we're interested in (leaving the remainder as an uncertainty), while for numerical simulations we inevitably introduce discretization errors. In the context of machine learning, we additionally have errors introduced by the trained model. All these errors and unclear aspects make up the uncertainties of the predicted outcomes, the _predictive uncertainty_. For practical applications it's crucial to have means for quantifying this uncertainty. This is a central motivation for working with probabilistic models, and for adjacent fields such as in "uncertainty quantification" (UQ).
```{note} Aleatoric vs. Epistemic Uncertainty.
The predictive uncertainty in many cases can
be distinguished in terms of two types of uncertainty:
- _Aleatoric_ uncertainty denotes uncertainty within the data, e.g., noise in measurements.
- _Epistemic_ uncertainty, on the other hand, describes uncertainties within a model such as a trained neural network.
A word of caution is important here:
while this distinction seems clear cut, both effects overlay and can be difficult to tell apart. E.g., when facing discretization errors, uncertain outcomes could be caused by unknown ambiguities in the data, or by a suboptimal discrete representation.
These aspects can be very difficult to disentangle in practice.
```
Closely aligned, albeit taking a slightly different perspective, are so-called _simulation-based inference_ (SBI) methods. Here the main motivation is to estimate likelihoods in computer-based simulations, so that reliable probability distributions for the solutions can be obtained. The SBI viewpoint provides a methodological approach for working with computer simulations and uncertainties, and will provide a red thread for the following sections.
## Forward or Backward?
At this point it's important to revisit the central distinction between forward and inverse ("backward") problems: most classic numerical methods target ➡️ **forward** ➡️ problems to compute solutions for steady-state or future states of a system.
Forward problems arise in many settings, but across the board, at least as many problems are ⬅️ **inverse** ⬅️ problems, where a forward simulation plays a central role, but the main question is not a state that it generates, but rather the value of parameter of simulator to explain a certain measurement or observation. To formalize this, our simulator $f$ is parametrized by a set of inputs $\nu$, e.g., a viscosity, and takes states $x$ to produce a modified state $y$. We have an observation $\tilde{y}$ and are interested in the value of $\nu$ to produce the observation. In the easiest case this inverse problem can tackled as a minimization problem
$\text{arg min}_{\nu} | f(x;\nu) - \tilde{y} |_2^2$. Solving it would tell us the viscosity of an observed material, and similar problems arise in pretty much all fields, from material science to cosmology. To simplify the notation, we'll merge $\nu$ into $x$, and minimize for $x$ correspondingly, but it's important to keep in mind that $x$ can encompass any set of parameters or state samples that we'd like to solve for with our inverse problem.
In the following, we will focus on inverse problems, as these best illustrate the capabilities of the probabilistic modeling, but the algorithms discussed are not exclusively applicable to inverse problems (an example will follow).
## Simulation-based Inference
For inverse problems, it is in practice not sufficient to match a single observation $\tilde{y}$. Rather, we'd like to ensure that the parameter we obtain explains a wide range of observations, and we might be interested in the possibility of multiple values explaining our observations. Similarly, quantifying the uncertainty of the estimate is important in real world settings: is the observation explained by only a very narrow range of parameters, or could the parameter vary by orders of magnitude without really influencing the observation? These questions require a statistical analysis, typically called _inference_, to draw conclusions about the results obtained from the inverse problem solve. To connect this viewpoint with the distinction regarding epistemic and aleatoric uncertainties above, we're primarily addressing the latter here: which uncertainties lie in our observations, given a scientific hypothesis in the form of a simulator.
To formalize these inverse problems let's consider
a vector-valued input$x$ that can contain states and / or
the aforementioned parameters (like $\nu$).
We also have a
distribution of latent variables 
$z \sim p(z|x)$ that describes the unknown part of our system.
Examples for z are unobservable and stochastic variables , intermediate simulation steps, or the control flow of simulator.
```{note} Bayes theorem is fundamental for all of the following. For completeness, here it is: $p(x|y)~p(y) = p(y|x)~p(x)$. And it's worth keeping in mind that both sides are equivalent to the joint probabilities, i.e. $... = p(x,y) = p(y,x)$.
```
For $x$ there is a prior distribution X with a probability density $p(x)$for the inputs,
and the simulator produces an observation or output $y \sim p(y | x, z)$. Thus, $x$ can take different values, maybe it contains some noise, and the $z$ is out of our control, and can likewise influence the $y$ that are produced.
The function for the conditional probability $p(y|x)$ is called the **likelihood** function, and is a crucial value in the following. Note that it does not depend on $z$, as these latent states are out of our control.
So we actually need to
compute the marginal likelihood $p(y|x) = \int p(y, z | x) dz$ by integrating over all possible $z$.
This is necessary because the likelihood function shouldn't depend on $z$, otherwise we'd need to know the exact values of $z$ before being able to calculate the likelihood.
Unfortunately, this is often intractable, as $z$ can be difficult to sample, and in some case we can't even control it in a reasonable way.
Some algorithms have been proposed to compute likelihoods, one popular one is Approximate Bayesian Computation (ABC), but all approaches are highly expensive and require a lot of expert knowledge to set up. They suffer from the _curse of dimensionality_, i.e. become very expensive when facing larger numbers of degrees of freedom. Thus,
obtaining good approximations of the likelihood will be a topic that we'll revisit below.
With a function for the likelihood we can compute the
**distribution of the posterior**, the main quantity we're after,
in the following way:
$p(x|y) = \frac{p(y|x)p(x)}{\int p(y|x') p(x') dx'}$,
where the denominator
$\int p(y|x') p(x') dx'$ is called the _evidence_.
The evidence is just $p(y)$, which shows
that the equation for the posterior follows directly from Bayes' theorem $p(x|y) = p(y|x) p(x) / p(y)$.
The evidence can be computed with stochastic methods such as Markov Chain Monte Carlo (MCMC).
It primarily "normalizes" our posterior distribution and is typically easier to obtain than the likelihood, but nonetheless still a challenging term.
```{admonition} Leveraging Deep Learning
:class: tip
This is were deep learning turns out to be extremely useful: we can use it to train a conditional density estimator $q_\theta(x|y)$ for the posterior $p(x|y)$ that allows sampling, and can be trained from simulations $y \sim p(y|x)$ alone.
```
Deep learning has been instrumental to provide new ways of addressing the classic challenges of obtaining accurate estimates of posterior distributions, and this is what we'll focus on in this chapter. Previously, we called our neural networks $f_\theta$, but in the following we'll use $q_\theta = f_\theta$ to make clear we're dealing with a learned probability. Specifically, we'll target neural networks that learn a probability density, i.e. $\int q_\theta(x) dx = 1$.
We'll often first target unconditional densities, and then show how they can be modified to learn conditional versions $q_\theta(x|y)$.
Looking ahead, the learned SBI methods, i.e. approaches for computing posterior distributions, have the following properties:
✅ Pro:
* Fast inference (once trained)
* Less affected by curse of dimensionality
* Can represent arbitrary priors
❌ Con:
* Require costly upfront training
* Lacks rigorous theoretical guarantees
In the following we'll explain how to obtain and derive a very popular and powerful family of methods that can be summarized as **diffusion models**. We could simply provide the final algorithm (which will turn out to be surprisingly simple), but it's actually very interesting to see where it all comes from.
We'll focus on the basics, and leave the _physics-based extensions_ (i.e. including differentiable simulators) for a later section. The path towards diffusion models also introduces a few highly interesting concepts from machine learning along the way, and provides a nice "red thread" for discussing seminal papers from the past few years. Here we go...
<br>
![Divider](resources/divider6.jpg)
```{note} Historic Alternative: Bayesian Neural Networks
A classic variant that should be mentioned here are "Bayesian Neural Networks". They
follow Bayes more closely, and pre-scribe a prior distribution on the neural network
parameters to learn the posterior distribution. Every weight and bias in the NN are assumed to be Gaussian with an own mean and variance, which are adjusted at training time. For inference, we can then "sample" a network, and use it like any regular NN.
Despite being a very good idea on paper, this method turned out to have problems with learning complex distributions, and requires careful tuning of the hyperparameters involved. Hence, these days, it's strongly recommended to use flow matching (or at least a diffusion model) instead.
If you're interested in details, BNNs with a code example can be found, e.g., in v0.3 of PBDL: https://arxiv.org/abs/2109.05237v3 .
```
probmodels-normflow.ipynb
+1269
View File
File diff suppressed because it is too large. Load diff
probmodels-score.ipynb
+1035
View File
File diff suppressed because it is too large. Load diff
references.bib
+11 -1
View File
@@ -15,7 +15,7 @@
@inproceedings{lino2025dgn,
title={Diffusion Graph Networks},
title={Learning Distributions of Complex Fluid Simulations with Diffusion Graph Networks},
author={Mario Lino and Tobias Pfaff and Nils Thuerey},
booktitle={International Conference on Learning Representations},
year={2025}
@@ -34,6 +34,16 @@
booktitle={International Conference on Learning Representations},
year={2025}
}
@inproceedings{koehler2024ape,
title={APEBench: A Benchmark for Autoregressive Neural Emulators of PDEs},
author={Felix Koehler and Simon Niedermayr and Ruediger Westermann and Nils Thuerey},
journal={Advances in Neural Information Processing Systems (NeurIPS)},
year={2024}
}
@inproceedings{shehata2025trunc,
title={Truncation Is All You Need: Improved Sampling Of Diffusion Models For Physics-Based Simulations},