Diffusion Models

Diffusion Models#

In this session we will get a short dive into Diffusion Models. We are going to see:

what can we do with them
how to train a diffusion model
an example application
a way to estimate their performance.

Setup#

from collections import defaultdict
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from matplotlib import colors
import seaborn as sns

from IPython.display import display, HTML
display(HTML("<style>.container { width:100% !important; }</style>"))

import requests

url = "https://raw.githubusercontent.com/AstroStat-Academy/assets-public/main/colab/clone_and_cd_colab.py"
colab = requests.get(url).text
exec(colab)

url = "https://raw.githubusercontent.com/AstroStat-Academy/assets-public/main/styles/plot_style.py"
style = requests.get(url).text
exec(style)

import warnings
warnings.filterwarnings(action='ignore')

# some handy functions
def add_plot_info():
    """Just add some information for the plot"""
    plt.title('Clusters')
    plt.xlabel('X')
    plt.ylabel('Y')
    plt.grid(True)

POV: You need to model a distribution!#

Let’s suppose we are given some [unknown] empirical distribution, and we need to generate mock data.

Figure 1. The sampling problem. We need to sample some data (right) based on some empirical distribution (left).

Example cases

E.g., you are:

Running an expensive simulation $\rightarrow$ cannot generate all the data analytically
Dealing with imbalanced data $\rightarrow$ need more examples of the minority class
Assessing some statistics $S$ $\rightarrow$ would like to “bootstrap” to obtain a distribution for $S$

Question

What would you do, if you had to sample pseudo-real data given a dataset?

Discuss with your teammate, then report. Think of possible pros and cons of your method.

In general, a powerful, generalizable model would:

not have parameters defined ad-hoc (parametric models are difficult to define for complex distributions)
accounting implicitly for correlations between features
adaptable to the local sub-distribution densities
unsupervised

$\rightarrow$ e.g., Generative Models

[Spoiler] (click here to expand)

E.g. KDE (Kernel Density Estimation)
sklearn.neighbors.KernelDensity

We can think of KDE as a smoothed distribution (e.g. via Gaussians) of the data.

Figure 2. An example of Kernel Density Estimation applied to a 2D color-color distribution.
(From Bonfini et al. (2021))

$\rightarrow$ From this continuous distribution we can now sample new datasets.

PROS:

Intuitive
Computationally quick

CONS:

Needs parameter tuning (kernel width)
Same global smoothing all over the feature ranges

Generative Models $vs$ Discriminative Models#

What are Generative Models? $\rightarrow$ Best understood when defined in contrast to Discriminative Models.

Figure 3. Summary comparison of Discriminative (left) $vs$ Generative (right) Models.
(From here)

A new paradigm to catalog ML models $\rightarrow$ they can be distinguished based on which probability distribution they target.

Discriminative Models

Objective $\rightarrow$ Differenciate between sample labels
Method $\rightarrow$ Learn the conditional probability $P(y | X)$
```
  "How likely is this label?"
```

Generative Models

Objective $\rightarrow$ Create new samples
Method $\rightarrow$ Learn the joint probability $P(X, y)$ $-$ or just $P(X)$ if unsupervised
```
  "How likely is this sample, or this <sample, label> pair?"
```

See this Google Developers post for a smooth intro.

Diffusion models#

What are Diffusion Models?#

Diffusion Models work by learning how to remove noise from the data.

The training consists in deteriorating the data by first adding noise …

Figure 3. Forward diffusion process.
(From here)

… and then learning how to re-construct the data by de-noising.

Figure 4. Backward diffusion process.
(From here)

$\rightarrow$ A trained model can then generate data similar to the ones used for the training (the learnt distribution).

One nifty eyey-candy example#

WARNING:

This block will:

require to download $\sim5$ GBs of pre-trained models in .cache/huggingface/, inside your home directory
require CUDA and a GPU with at least 8 GB

#%pip install -qq -U diffusers datasets transformers accelerate ftfy pyarrow==9.0.0

import torch

from huggingface_hub import login
from getpass import getpass
login(token=getpass("Paste your Hugging Face token: "))

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print('Device in use:', device)

if device.type == 'cuda':
    total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1024**3
    if total_mem_gb >= 8:
        from diffusers import StableDiffusionPipeline
        model_id = "stable-diffusion-v1-5/stable-diffusion-v1-5"
        pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to(device)
    else:
        print('Sorry peasant, not enough GPU memory!')

# Produce a sampled image! Some usage suggestions from here:
#   https://huggingface.co/blog/stable_diffusion

if device.type == 'cuda':
    total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1024**3
    if total_mem_gb >= 8:
        prompt = "A screenshot from the movie Dune 1984, zoom-in on the protagonist"
        image = pipe(prompt, num_inference_steps=60, guidance_scale=5.5, height=512, width=768).images[0]
        display(image)
    else:
        print('Sorry peasant, not enough GPU memory!')

In practice $-$ as you saw $-$ Diffusion Models are trained used conditioning.

Figure 5. Diffusion with conditioning.
(Adapted from here)

In the eye-candy example above, the conditioning is text
In Smith et al. 2022, they conditioned $-$ among others $-$ on sketches to create astro imaging (code available here)

Figure 6. Realistic galaxy imaging built with diffusion models conditioned on a sketch.
(Figure 7 from Smith et al. 2022)

This could be used e.g. to generate high-resolution imaging from cosmological simulation realizations, without passing via ray-tracing.

The model, the myth, the math#

The seminal paper is by Sohl-Dickstein et al. (2015), but the one which most determined the success of Diffusion Models is Ho, Jain & Abbeel (2020; hereafter Ho+20).

Here we will follow the method Ho+20. For a thorough mathematical explanation, follow this fantastic blog by L. Weng.

Forward process

The forward [and backward] process is a Markov chain, in the sense that the next timestep is assumed to only dependend on the previous.

In Ho+20, at each step the added noise:

is small (or else the diffusion is too fast to learn)
is sampled from a Gaussian centered at $\mu_t$ …
… with the same variance $\beta_t$ for every feature

\[ q(\mathbf{x}_t|\mathbf{x}_{t-1}) := \mathcal{N}(\mathbf{x}_t; \mathbf{\mu}_t, \mathbf{\beta}_t) \]

The choice of the $\beta_1$, $\beta_2$, …, $\beta_T$ values is called variance schedule (e.g. could be a linear or a sinusoidal progression).

We can pick $\mathbf{\mu}_t$ as we like $-$ With a smart choice (see Ho+20), we can decide that $\mathbf{\mu}_t$ is defined as a function of $\beta_t$ themselves:

\[ q(\mathbf{x}_t|\mathbf{x}_{t-1}) := \mathcal{N}(\mathbf{x}_t; \sqrt{1-\beta_t}\mathbf{x}_{t-1}, \beta_t \mathbf{I}) \]

Now, turns out that the Markov chain is actually recursive, so we can get directly from $\mathbf{x}_0$ to $\mathbf{x}_t$ as:

\[ q(\mathbf{x}_t|\mathbf{x}_{0}) := \mathcal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t}\mathbf{x}_{0}, (1-\bar{\alpha}_t) \mathbf{I}) \]

where $\bar{\alpha}_t := \prod_{s=1}^{t}\alpha_s$, and $\alpha_t:=1-\beta_t$.

$\rightarrow$ The forward process is trivial: no problems here, gg.

Backward process

Here we would like to know $q(\mathbf{x}_{t-1}|\mathbf{x}_{t})$, but guess what $-$ it’s untractable!

Intuitively $-$ there are too many possible reconstructions, starting from pure noise.

Good news: we can know $q(\mathbf{x}_{t-1}|\mathbf{x}_{t}, \mathbf{x}_{0})$!

That is, if we condition on the input, we can de-noise from timestep $t$ to $t-1$.

Intuitively $-$ we are adding information on where the starting noise came from.

\[ q(\mathbf{x}_{t-1}|\mathbf{x}_t, \mathbf{x}_0) := \mathcal{N}(\mathbf{x}_{t-1}; \mathbf{\tilde{\mu}_t}(\mathbf{x}_t, \mathbf{x}_0), \mathbf{\tilde{\beta}}_t) \]

were:

$ \mathbf{\tilde{\mu}_t}(\mathbf{x}_t, \mathbf{x}_0) := {\sqrt{\bar{\alpha}_t-1}\beta_t \over 1-\bar{\alpha}_t} \mathbf{x}_0 + {\sqrt{\bar{\alpha}_t}(1-\bar{\alpha}_{t-1}) \over 1-\bar{\alpha}_t} \mathbf{x}_t $

$\mathbf{\tilde{\beta}}_t := {1-\bar{\alpha}_{t-1} \over 1-\bar{\alpha}_t}\beta_t$

IMPORTANT: Here we assumed that the reverse process is also Gaussian. This is true as long as the diffusion is “slow” (many small timesteps).

The caveat

“Ok, so $q(\mathbf{x}_t|\mathbf{x}_{0})$ is trivial, and I can calculate analytically $\mathbf{\tilde{\mu}_t}$ and $\mathbf{\tilde{\beta}}_t$, so I have $q(\mathbf{x}_{t-1}|\mathbf{x}_t, \mathbf{x}_0)$ as well.”

“What do I have to model? Why I am still here, anyways?”

Chill. Even if we have $q(\mathbf{x}_{t-1}|\mathbf{x}_t, \mathbf{x}_0)$, that is only meaningful for the training set $\rightarrow$ we need to generalize to unseen data!

Intuitively $-$ we would be overfitting.

$\rightarrow$ We need to model the reverse process for any input $-$ we will use a Neural Network!

The model

What do we need to predict?

Basically, we want to approximate:

\[ q(\mathbf{x}_{t-1}|\mathbf{x}_t, \mathbf{x}_0)\]

with a model (see Figure 2.2):

\[ p_\theta(\mathbf{x}_{t-1} | \mathbf{x}_t)\]

$\rightarrow$ Now $ q(\mathbf{x}_{t-1}|\mathbf{x}_t, \mathbf{x}_0)$ (known) becomes our training label that we need to predict!

If we can predict the Gaussian noise, we can subtract it and reconstruct the image, right?

the variance schedule is fixed by construction
our NN only needs to predict $\mathbf{\tilde{\mu}_t}(\mathbf{x}_t, \mathbf{x}_0)$ !

\[ p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t) \rightarrow \mathbf{\mu_\theta}(\mathbf{x}_t, t) \]

TL;DR

We need to train a NN that minimizes the following Loss Function:

\[ L = E_q \big[ {1 \over 2\beta_t} || \mathbf{\tilde{\mu}_t}(\mathbf{x}_t, \mathbf{x}_0) - \mathbf{\mu_\theta}(\mathbf{x}_t, t) ||^2 \big] \]

Thanks to some smart parametrization of $\mu_t$ (again, see the blog by L. Weng), the loss can be simplified to:

\[ L = E_q \big[ || \mathbf{\epsilon} - \mathbf{\epsilon_\theta}(\sqrt{\bar{\alpha}_t}\mathbf{x}_{0} + \sqrt{1 - \bar{\alpha}_t}\mathbf{\epsilon}, t) ||^2 \big] ~~~~~ (1) \]

where:

\[ \mathbf{\epsilon} \sim \mathcal{N(0, \mathbf{I})} \]

(i.e., is a value randomly drawn from a Normal distribution) and $\epsilon_\theta$ is our NN that predicts noise given a noisy image and $t$.

We can interpret Equation 1 as:

“Predict me how much noise has been added to this image, at timestep t.”

Algorithm

The training algorithm is now extremely simple $\rightarrow$ we don’t even need the forward pass!

Figure 7. Diffusion Model algorithm from Ho+20.
(From here)

The sampling algorithm is not conditioned on $\mathbf{x}_0$, as desired $\rightarrow$ $\mathbf{x}_0$ used only for training

My first diffusion model!#

import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import torch
import os
import seaborn as sns

# Global Seaborn setup for the reminder of the notebook:
sns.set(font_scale=1.2)

from src.diffusion_utils import extract, make_beta_schedule, p_sample_loop

path_results = "results"
if not os.path.exists(path_results):
    os.makedirs(path_results)

force_cpu = False
# to force-use CPU (debug purposes or peasants without a proper GPU)
if force_cpu:
    device = torch.device('cpu')
else:
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # Automatically move all tensors to the device:
    torch.set_default_tensor_type(torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor)

print('Device in use [cpu/cuda]:', device)

The following example has been coded based on this work by J. Rafid Siddiqui.

Generating some mock data#

from sklearn.datasets import make_swiss_roll

data, _ = make_swiss_roll(1000, noise=0.2)
data    = data[:, [0, 2]]
# dropping 3rd dimension, only working in 2D

features = ['X1', 'X2']
# arbitrary feature names

df_data = pd.DataFrame(data, columns=features)
display(df_data.head(3))

Visualizing data#

def plot_data(df_data, features, display=True, ax=None, color=None):
    
    color_codes_wanted = ['grey', 'green', 'purple']

    if color is None:
        palette = ["#4c7d99"]
        color_codes = ['thiscolor']
    else:
        palette = [color]
        color_codes = ['thiscolor']
    c = lambda x: palette[color_codes.index(x)]

    
    if display:
        plt.figure(figsize=(8, 6))

    if ax is not None:
        ax = sns.scatterplot(data=df_data, x=features[0], y=features[1], ax=ax, color=c('thiscolor'))
    else:
        ax = sns.scatterplot(data=df_data, x=features[0], y=features[1], color=c('thiscolor'))
            
    if ax is not None:
        ax.set_ylim(-12.5,15)
        ax.set_xlim(-12.5,15)
    else:
        plt.ylim(-12.5,15)
        plt.xlim(-12.5,15)
    
    if display:
        plt.show()

plot_data(df_data, features)

Splitting#

We are not going to use the validation set, but we leave this block here for good habit in case a real application is more sophisticated.

from sklearn.model_selection import train_test_split
from prettytable import PrettyTable

df_X_train, df_X_valid, df_y_train, df_y_valid = \
    train_test_split(df_data[features],
                     _, test_size=0.3, random_state=42)

table = PrettyTable()
table.title = str('Data shape')
table.field_names = ['set', 'X', 'y']
table.add_row(['train',   np.shape(df_X_train), np.shape(df_y_train)])
table.add_row(['valid', np.shape(df_X_valid), np.shape(df_y_valid)])
print(table)

Scaling#

# Scaling:
from sklearn.preprocessing import StandardScaler

scaler = StandardScaler()

df_X_train_n = df_X_train.copy()
df_X_valid_n = df_X_valid.copy()

df_X_train_n[features] = scaler.fit_transform(df_X_train_n[features])
df_X_valid_n[features] = scaler.transform(df_X_valid_n[features])

display(df_X_train.head(3))
display(df_X_train_n.head(3))

print("Scaled data (obviously they don't fit in the plot, but we don't care)")
plot_data(df_X_train_n, features)

# Converting to torch tensors:
X_train = torch.tensor(df_X_train_n[features].values).to(torch.float16)
X_valid = torch.tensor(df_X_valid_n[features].values).to(torch.float16)

n_features = np.shape(X_train)[1]
print('Number of features:', n_features)

Diffusion Parameters#

number of timesteps

i.e., for how many steps we add noise

$\mathbf{\beta}$ schedule (and therefore $\mathbf{\alpha} = \mathbf{\alpha}(\mathbf{\beta})$):

$ q(\mathbf{x}_t|\mathbf{x}_{0}) := \mathcal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t}\mathbf{x}_{0}, (1-\bar{\alpha}_t) \mathbf{I}) $

i.e., how intense is the noise, at each time $t$

n_timesteps = 100
# number of timesteps of the diffusion procedure
# WARNING: Do not increase much beyond 100 or it will saturate the RAM

# Beta schedule:
betas = make_beta_schedule(schedule='sigmoid', n_timesteps=n_timesteps, start=1e-3, end=3e-2)

alphas = 1 - betas
alphas_prod = torch.cumprod(alphas, 0)
alphas_prod_p = torch.cat([torch.tensor([1]).float(), alphas_prod[:-1]], 0)
alphas_bar_sqrt = torch.sqrt(alphas_prod)
one_minus_alphas_bar_log = torch.log(1 - alphas_prod)
one_minus_alphas_bar_sqrt = torch.sqrt(1 - alphas_prod)

from matplotlib import pyplot as plt

alphas_bar_sqrt_ = alphas_bar_sqrt.detach().cpu().numpy()

plt.plot(alphas_bar_sqrt_, label=r"$\mu_t$ multiplier: ${\sqrt{\bar{\alpha}_t}}$")
plt.plot((1 - alphas_bar_sqrt_**2), label=r"$\sigma_t^2$: $(1 - \bar{\alpha}_t)$")
plt.xlabel('Timestep')
plt.legend(fontsize="large", loc='center left', bbox_to_anchor=(1, 0.5))
plt.show()

Forward Diffusion#

NOTE: We actually only use this only for showcasing the diffusion process …

For any given time $t$, we need to sample $\mathbf{x}_t$ from the distribution $q(\mathbf{x}_t|\mathbf{x}_{0})$:

\[\mathbf{x}_t \sim q(\mathbf{x}_t|\mathbf{x}_{0}) := \mathcal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t}\mathbf{x}_{0}, (1-\bar{\alpha}_t) \mathbf{I})\]

… but with a simple re-parametrization trick we can take the mean and just sample around it:

\[\mathbf{x}_t = \sqrt{\bar{\alpha}_t}\mathbf{x}_{0} + (1-\bar{\alpha}_t) \epsilon\]

where: $\epsilon \sim \mathcal{N}(0, \mathbf{I})$

def q_x(x_0, t, epsilon=None):
    '''
    q(x_t | x_0)
    '''
    if epsilon is None:
        epsilon = torch.randn_like(x_0)
    alphas_t = extract(alphas_bar_sqrt, t, x_0)
    alphas_1_m_t = extract(one_minus_alphas_bar_sqrt, t, x_0)
    return (alphas_t * x_0 + alphas_1_m_t * epsilon)

n_plots = 10
fig, axs = plt.subplots(1, n_plots, figsize=(28, 3))

for i, t in enumerate(np.linspace(0, n_timesteps-1, n_plots)):
    t = torch.tensor([int(t)])
    q_t = q_x(X_train, t)
    axs[i].scatter(q_t[:, 0].cpu(), q_t[:, 1].cpu(),color='white',edgecolor='gray', s=5);
    axs[i].set_axis_off(); axs[i].set_title('$q(\mathbf{x}_'+str(t.cpu().numpy())+')$') 

Reverse Diffusion/Reconstruction#

Unlike the forward diffusion, the reverse diffusion process requires training the neural network model to predict the noise
The model will be a simple fully-connected NN with a few layers

Training $\rightarrow$ boils down to following Algorithm 1, and calculating the correct Loss

Sampling $\rightarrow$ Trivial once the NN is trained, follows Algorithm 2 (we will give it for granted, but see code in diffusion_utils.py if you are curious)

batch_size = 128

n_epochs = 1000
# training epochs
# NOTE: At every batch of every epoch, a different time step is sampled

import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

class DenseWithTime(nn.Module):
    '''A modification of a basic dense layer, where the X input is actually
    replaced by the X input multiplied by the time embedding'''
    def __init__(self, num_in, num_out, n_timesteps):
        super(DenseWithTime, self).__init__()
        self.num_out = num_out
        self.lin = nn.Linear(num_in, num_out)
        self.embed = nn.Embedding(n_timesteps, num_out)
        # definition of the model (nn.Embedding) to embed the time
        self.embed.weight.data.uniform_()

    def forward(self, x, t):
        '''This gets executed when `DenseWithTime` is called.'''
        data = self.lin(x)
        t_embedded = self.embed(t)
        # embedding the specific time <t>
        out = t_embedded.view(-1, self.num_out) * data
        # reshaping time embedding and multiplying it by X
        return out
        
class ConditionalModel(nn.Module):
    def __init__(self, n_features, n_timesteps):
        '''Model layers definition.'''
        super(ConditionalModel, self).__init__()
        self.lin1 = DenseWithTime(n_features, 128, n_timesteps)
        self.lin2 = DenseWithTime(128, 128, n_timesteps)
        self.lin3 = DenseWithTime(128, 128, n_timesteps)
        self.lin4 = DenseWithTime(128, 128, n_timesteps)
        self.out = nn.Linear(128, n_features)
    
    def forward(self, x, t):
        '''This gets executed when `model` is called.'''
        x = F.relu(self.lin1(x, t))
        x = F.relu(self.lin2(x, t))
        x = F.relu(self.lin3(x, t))
        x = F.relu(self.lin4(x, t))
        return self.out(x)
    
model = ConditionalModel(X_train.shape[1], n_timesteps)
optimizer = optim.Adam(model.parameters(), lr=1e-3)    

%%time

losses = {'train':[], 'valid':[]}

for epoch in range(n_epochs):
    permutation = torch.randperm(X_train.size()[0])
    # random permutations of the batch indexes
    epoch_loss_cum = 0
    # cumulative epoch loss over batches
    
    # Iterating over batches:
    for i in range(0, X_train.size()[0], batch_size):
        # Retrieving current batch:
        indices = permutation[i:i+batch_size]
        batch_x = X_train[indices]
        batch_size_ = len(batch_x)
        # calculating actual size since the last batch can be smaller than <batch_size>
        
        # Algorithm 1 line 3: selecting a random timestep for each sample in the batch:
        t = torch.randint(0, n_timesteps, size=(batch_size_,))
        
        # Sampling noise (first term in Loss Equation) -----------------------
        e = torch.randn_like(batch_x)
        #---------------------------------------------------------------------

        # Predicting noise using NN (second term in Loss Equation) -----------
        # - The "extract" calls are just convenience functions which pick a
        #    value of the alphas at time <t> and reshape it to the batch size
        # - The core of the prediction is where "model" is called 
        #
        # Multiplier of x_0 in the Equation:
        alpha_bar_sqrt = extract(alphas_bar_sqrt, t, batch_x)
        # Multiplier of epsilon in the Equation:
        one_minus_alpha_bar_sqrt = extract(one_minus_alphas_bar_sqrt, t, batch_x)
        # Model input:
        input_e = batch_x * alpha_bar_sqrt + one_minus_alpha_bar_sqrt*e
        # Model prediction:
        e_theta = model(input_e, t)
        #---------------------------------------------------------------------
        
        # Loss ---------------------------------------------------------------
        # Loss Equation:
        loss = (e - e_theta).square().mean()
        # NOTE: average loss across samples and features, for this epoch, this batch
        epoch_loss_cum += loss.detach().cpu().numpy()
        #---------------------------------------------------------------------

        # Optimizing the NN --------------------------------------------------
        # Before the backward pass, zero-ing all the network gradients:
        optimizer.zero_grad()
        # Backward pass: computing gradient of the loss with respect to parameters
        loss.backward()
        # Performing gradient clipping
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.)
        # Calling the step function to update the parameters:
        optimizer.step()
        #---------------------------------------------------------------------

    ###############################################
    # Done! The rest is just to display stuff ... #
    ###############################################

    # Storing loss:
    n_batches = np.ceil(X_train.size()[0]/batch_size)
    losses['train'].append(epoch_loss_cum/n_batches)
    # averaging over batches

    # Printing loss (every 1/10th of the total epochs):
    if (epoch % int(n_epochs/10) == 0):       
        print('Epoch: %-3s [/%3s] | Loss: %s' % (epoch, n_epochs, loss.detach().cpu().numpy()))

    # Plotting data (every 1/10th of the total epochs):
    if (epoch % int(n_epochs/10) == 0):       
        
        # Generating a whole NEW time sequence by drawing from the learnt distribution:
        x_seq_gen = p_sample_loop(model, X_train.shape, n_timesteps,
                                  alphas,betas, one_minus_alphas_bar_sqrt)
        
        # Looping through all timesteps for this epoch, and plotting a few of them:
        n_plots = 8
        fig, axs = plt.subplots(1, n_plots, figsize=(28, 3))
        for j, t in enumerate(np.linspace(0, n_timesteps-1, n_plots)):
            t = int(t)
            # Picking dataset generated for timestep <t>:
            x_t_gen = x_seq_gen[t]
            if device != 'cpu': x_t_gen = x_t_gen.detach().cpu().numpy()            
            # Reverting normalization:
            df_data_gen = pd.DataFrame(data=x_t_gen, columns=features)
            df_data_gen[features] = scaler.inverse_transform(df_data_gen[features])
            # Plotting:
            plot_data(df_data_gen, features, display=False, ax=axs[j])
            axs[j].xaxis.label.set_visible(False) 
            axs[j].yaxis.label.set_visible(False) 
            axs[j].set_title('$q(\mathbf{x}_{'+str(t)+'})$')
        plt.tight_layout()
        plt.show()

# Plotting the Loss:
import matplotlib.pyplot as plt

if device != 'cpu':
    losses_train = [loss for loss in losses['train']]
    losses_valid = [loss for loss in losses['valid']]

fig, axes = plt.subplots(1, 1, figsize=(6, 4))
ax = axes
ax.plot(losses_train, label='train')
ax.plot(losses_valid, label='valid')

ax.legend()
ax.set_xlabel('Epoch')
ax.set_ylabel('Loss')

plt.tight_layout()
plt.show() 

Animation#

You don’t need to care about the following code, it is only for presentation purposes!

We will create an image sequence that resembles the evolution of the data along the timesteps, both in the forward (diffusion) and in the backward (denoising) pass.

In the GIF:

we first diffuse the original data
we then start from random noise and converge towards the sampled distribution

IMPORTANT: Notice that those are two different datasets! The first is the training dataset being diffused, the second is a new dataset randomly sampled (from noise) and inverse-diffused!

I.e., the start of the inverse-diffusion is not the end-point of the diffusion! This is the whole point of learning the distribution $-$ we want to sample new data!

They sequences are just collated together to display the intuition behind the process.

import matplotlib
from matplotlib.colors import to_hex
import io
from PIL import Image
from tqdm import tqdm

cmap = matplotlib.colors.LinearSegmentedColormap.from_list("", ["tomato","teal","steelblue"])

Generating the diffusion image sequence#

%%time
    
imgs_diffusion = []

pbar = tqdm(range(n_timesteps))
for t in pbar:
    plt.clf()
    q_x_t = q_x(X_train, torch.tensor([t]))
    if device != 'cpu': q_x_t = q_x_t.detach().cpu().numpy()

    # Indexing colors using first-half palette (the rest is used for the de-noising!):
    color = to_hex(cmap(t/n_timesteps/2))

    # Reverting normalization:
    df_data_noise = pd.DataFrame(data=q_x_t, columns=features)
    df_data_noise[features] = scaler.inverse_transform(df_data_noise[features])
    # Plotting:
    plt.title('Diffusing training set')
    plot_data(df_data_noise, features, display=False, color=color)

    img_buf = io.BytesIO()
    plt.savefig(img_buf, format='png')
    img = Image.open(img_buf)
    imgs_diffusion.append(img)

Generating the de-noising [reverse] image sequence#

Let’s generate a brand new dataset composed of n_samples_gen samples!

%%time

imgs_inv_diffusion = []

# Generating a whole NEW time sequence by drawing from the learnt distribution:
n_samples_gen = 2000
X_gen_shape = (n_samples_gen, X_train.shape[1])
x_seq_gen = p_sample_loop(model, X_gen_shape, n_timesteps,
                          alphas,betas, one_minus_alphas_bar_sqrt)

pbar = tqdm(range(n_timesteps))
for t in pbar: 
    plt.clf()
    # Picking dataset generated for timestep <t>:
    x_t_gen = x_seq_gen[t]
    if device != 'cpu': x_t_gen = x_t_gen.detach().cpu().numpy() 

    # Indexing colors using second-half palette (the rest is used for the noising!):
    color = to_hex(cmap(t/n_timesteps/2+0.5))

    # Reverting normalization:
    df_data_gen = pd.DataFrame(data=x_t_gen, columns=features)
    df_data_gen[features] = scaler.inverse_transform(df_data_gen[features])
    # Plotting:
    plt.title('Generating a new dataset')
    plot_data(df_data_gen, features, display=False, color=color)
    
    img_buf = io.BytesIO()
    plt.savefig(img_buf, format='png')
    img = Image.open(img_buf)
    imgs_inv_diffusion.append(img)    

# Pasting the 2 seqeunces into a sequential GIF animation:
imgs_timeline = imgs_diffusion + imgs_inv_diffusion

imgs_timeline[0].save(path_results+'/diffusion_roll.gif', format='GIF', append_images=imgs_timeline,
             save_all=True, duration=50, loop=0)

import IPython

def show_image(filename):
    with open(filename, 'rb') as f:
        img = f.read()
    IPython.display.display(IPython.display.Image(data=img, format='gif'))

show_image(path_results+'/diffusion_roll.gif')

In-class Exercise: diffuse a BPT#

The “Baldwin, Phillips & Terlevich” (BPT) diagrams are used to distinguish sources based on specific spectral emission lines. The strengths of these lines depend on the heating source (see for more here).

BPTs allow to distinguish:

    AGNs (Seyfert)
    LINERs
    Star-forming galaxies
    Composite objects

Figure 8. Example of classification via BPT diagram. Theoretical or observationally-calibrated curves allow to distinguish the different sub-populations. (From Parra et al. (2010), ApJ, 720, 555.

Dataset: We will use a dataset of SDSS galaxies prepared by Daoutis Charalampos, carrying information about multiple emission lines and galaxy classification (paper in preparation).

The class is only used for displaying purposes, and discarded for the actual analysis

NOTE: Don’t be intimidated! $-$ The following code is 99% copy-pasted from above, with minor differences in the train/test splitting and plotting.

import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import numpy as np
import torch
import os

from src.diffusion_utils import extract, make_beta_schedule, p_sample_loop

path_results = "results"
if not os.path.exists(path_results):
    os.makedirs(path_results)

force_cpu = False
# to force-use CPU (debug purposes)
if force_cpu:
    device = torch.device('cpu')
else:
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # Automatically move all tensors to the device:
    torch.set_default_tensor_type(torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor)

print('Device in use [cpu/cuda]:', device)

Importing Data#

!pip install wget --quiet

import os
import wget
import pandas as pd

path_data = './data'
# path of data folder

path_to_data = path_data+"/BPT_emlines.csv"
# path/to/file to be created

# > BPT_emlines.csv
#link_data = "https://github.com/astrostatistics-in-crete/2023_summer_school/blob/ac06b58182bcfb9e1db0826147b1bf771069a1a7/DL_Diffusion/data/BPT_emlines.csv"

if not os.path.exists(path_data):
    os.makedirs(path_data)

#if not os.path.exists(path_to_data):
#    response = wget.download(link_data, path_to_data)

df_data = pd.read_csv(path_to_data)

# Limiting data:
frac_sampling = 0.02
df_data = df_data.groupby('Target_class', group_keys=False).apply(lambda x: x.sample(frac=frac_sampling))
# stratified sampling

display(df_data.head(5))
print(df_data.shape)

features = ['log(NII_6584/H_ALPHA)', 'log(OIII_5007/H_BETA)']
# x-axis and y-axis features in the BPT diagram
print('Features used in this exercise:', features)

Visualizing data#

import seaborn as sns
from matplotlib import pyplot as plt

# Global Seaborn setup for the reminder of the notebook:
sns.set(font_scale=1.2)

def plot_data(df_data, features, show_classes=False, display=True, ax=None, color=None):
    
    color_codes_wanted = ['grey', 'green', 'purple']

    if color is None:
        palette = ["#4c7d99"]
        color_codes = ['thiscolor']
    else:
        palette = [color]
        color_codes = ['thiscolor']
    c = lambda x: palette[color_codes.index(x)]

    
    if display:
        plt.figure(figsize=(8, 6))

    if show_classes:
        if ax is not None:
            ax = sns.scatterplot(data=df_data, x=features[0], y=features[1],
                        hue='Target_class', palette='tab10', ax=ax)  
        else:
            ax = sns.scatterplot(data=df_data, x=features[0], y=features[1],
                        hue='Target_class', palette='tab10')
    else:
        if ax is not None:
            ax = sns.scatterplot(data=df_data, x=features[0], y=features[1], ax=ax, color=c('thiscolor'))
        else:
            ax = sns.scatterplot(data=df_data, x=features[0], y=features[1], color=c('thiscolor'))
            
    # Kewley+01:
    xx_Ke = np.linspace(-2.5, 0.3, 1000)
    yy_Ke = (0.61/( xx_Ke  - 0.47  )) + 1.19

    # Schawinski+07:
    xx_Sc = np.linspace(-0.180, 1.5, 1000)
    yy_Sc = 1.05*xx_Sc + 0.45

    # Kauffmann+03:
    xx_Ka = np.linspace(-1.3, 0., 1000)
    yy_Ka = 0.61/(xx_Ka -0.05) + 1.3

    if ax is not None:
        ax.plot(xx_Ke, yy_Ke, '-' , color='#505050') # Kewley+01
        ax.plot(xx_Sc, yy_Sc, '-',  color='#505050') # Schawinski+07
        ax.plot(xx_Ka, yy_Ka, '--', color='#505050') # Kauffmann+03     

        ax.set_ylim(-1.5,1.25)
        ax.set_xlim(-1.75,0.5)
    else:
        plt.plot(xx_Ke, yy_Ke, '-' , color='#505050') # Kewley+01
        plt.plot(xx_Sc, yy_Sc, '-',  color='#505050') # Schawinski+07
        plt.plot(xx_Ka, yy_Ka, '--', color='#505050') # Kauffmann+03     
        
        plt.ylim(-1.5,1.25)
        plt.xlim(-1.75,0.5)
    
    if display:
        plt.show()
    
plot_data(df_data, features, show_classes=True)
plot_data(df_data, features, show_classes=False)

Splitting#

You know the drill …

from sklearn.model_selection import train_test_split
from prettytable import PrettyTable

df_X_train, df_X_valid, df_y_train, df_y_valid = \
    train_test_split(df_data[['log(NII_6584/H_ALPHA)', 'log(OIII_5007/H_BETA)']],
                     df_data['Target_class'], test_size=0.3, random_state=42)

table = PrettyTable()
table.title = str('Data shape')
table.field_names = ['set', 'X', 'y']
table.add_row(['train',   np.shape(df_X_train), np.shape(df_y_train)])
table.add_row(['valid', np.shape(df_X_valid), np.shape(df_y_valid)])
print(table)

Scaling#

# Scaling:
from sklearn.preprocessing import StandardScaler

scaler = StandardScaler()

df_X_train_n = df_X_train.copy()
df_X_valid_n = df_X_valid.copy()

df_X_train_n[features] = scaler.fit_transform(df_X_train_n[features])
df_X_valid_n[features] = scaler.transform(df_X_valid_n[features])

display(df_X_train.head(3))
display(df_X_train_n.head(3))

print("Scaled data (obviously they don't fit in the plot, but we don't care)")
plot_data(df_X_train_n, features)

# Converting to torch tensors:
X_train = torch.tensor(df_X_train_n[features].values).to(torch.float16)
X_valid = torch.tensor(df_X_valid_n[features].values).to(torch.float16)

n_features = np.shape(X_train)[1]
print('Number of features:', n_features)

Your turn!

Objective: Learn a model which can sample from a heuristic BPT distribution.

Task:

First, just copy-paste the code from above
- do not change a thing!
- just train the model!
Think of how you would change the code if you were to calculate the loss also for the validation set - don't do it, just think!
Try to improve the model by changing network and diffusion parameters

Try to generate a distribution which is as close as possible to the original one. Observe how the model convergence changes with the modification of the hyperparameters, and note down your observations for later discussion.

Try one/two things at a time, and re-run every time to see what happens. You can experiment with:

Different $\beta$ schedules (use: 'linear', 'quad', 'sigmoid', change start/end values of betas - use the multiplier vs. noise plot)
Different network architectures
- number of layers
- number of neurons
- activation functions
Number of training epochs and batch sizes

Neo flexes the confidence typical of a student at the last day of the AstroStatistics School.
Spoiler $-$ he gets his ass handed over by Morpheus.

# Neo, try your code here..

[Solution below]

Diffusion Parameters#

number of timesteps

i.e., for how many steps we add noise

$\mathbf{\beta}$ schedule (and therefore $\mathbf{\alpha} = \mathbf{\alpha}(\mathbf{\beta})$):

$ q(\mathbf{x}_t|\mathbf{x}_{0}) := \mathcal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t}\mathbf{x}_{0}, (1-\bar{\alpha}_t) \mathbf{I}) $

i.e., how intense is the noise, at each time $t$

n_timesteps = 100
# number of timesteps of the diffusion procedure
# WARNING: Do not increase much beyond 100 or it will saturate the RAM

# Beta schedule:
betas = make_beta_schedule(schedule='sigmoid', n_timesteps=n_timesteps, start=1e-5, end=1e-1)

alphas = 1 - betas
alphas_prod = torch.cumprod(alphas, 0)
alphas_prod_p = torch.cat([torch.tensor([1]).float(), alphas_prod[:-1]], 0)
alphas_bar_sqrt = torch.sqrt(alphas_prod)
one_minus_alphas_bar_log = torch.log(1 - alphas_prod)
one_minus_alphas_bar_sqrt = torch.sqrt(1 - alphas_prod)

from matplotlib import pyplot as plt

alphas_bar_sqrt_ = alphas_bar_sqrt.detach().cpu().numpy()

plt.plot(alphas_bar_sqrt_, label=r"$\mu_t$ multiplier: ${\sqrt{\bar{\alpha}_t}}$")
plt.plot((1 - alphas_bar_sqrt_**2), label=r"$\sigma_t^2$: $(1 - \bar{\alpha}_t)$")
plt.xlabel('Timestep')
plt.legend(fontsize="large", loc='center left', bbox_to_anchor=(1, 0.5))
plt.show()

Forward Diffusion#

NOTE: We actually only use this only for showcasing the diffusion process …

For any given time $t$, we need to sample $\mathbf{x}_t$ from the distribution $q(\mathbf{x}_t|\mathbf{x}_{0})$:

\[\mathbf{x}_t \sim q(\mathbf{x}_t|\mathbf{x}_{0}) := \mathcal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t}\mathbf{x}_{0}, (1-\bar{\alpha}_t) \mathbf{I})\]

… but with a simple re-parametrization trick we can take the mean and just sample around it:

\[\mathbf{x}_t = \sqrt{\bar{\alpha}_t}\mathbf{x}_{0} + (1-\bar{\alpha}_t) \epsilon\]

where: $\epsilon \sim \mathcal{N}(0, \mathbf{I})$

def q_x(x_0, t, noise=None):
    '''
    q(x_t | x_0)
    '''
    if noise is None:
        noise = torch.randn_like(x_0)
    alphas_t = extract(alphas_bar_sqrt, t, x_0)
    alphas_1_m_t = extract(one_minus_alphas_bar_sqrt, t, x_0)
    return (alphas_t * x_0 + alphas_1_m_t * noise)

n_plots = 8
fig, axs = plt.subplots(1, n_plots, figsize=(28, 3))

for i, t in enumerate(np.linspace(0, n_timesteps-1, n_plots)):
    t = torch.tensor([int(t)])
    q_t = q_x(X_train, t)
    axs[i].scatter(q_t[:, 0].cpu(), q_t[:, 1].cpu(),color='white',edgecolor='gray', s=5);
    axs[i].set_axis_off(); axs[i].set_title('$q(\mathbf{x}_'+str(t.cpu().numpy())+')$') 

Reverse Diffusion/Reconstruction#

Unlike the forward diffusion, the reverse diffusion process requires training the neural network model to predict the noise
The model will be a simple fully-connected NN with a few layers

Training $\rightarrow$ boils down to following Algorithm 1, and calculating the correct Loss

Sampling $\rightarrow$ Trivial once the NN is trained, follows Algorithm 2 (we will give it for granted, but see code in diffusion_utils.py if you are curious)

batch_size = 128

n_epochs = 200
# training epochs
# NOTE: At every batch of every epoch, a different time step is sampled

import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

class DenseWithTime(nn.Module):
    '''A modification of a basic dense layer, where the X input is actually
    replaced by the X input multiplied by the time embedding'''
    def __init__(self, num_in, num_out, n_timesteps):
        super(DenseWithTime, self).__init__()
        self.num_out = num_out
        self.lin = nn.Linear(num_in, num_out)
        self.embed = nn.Embedding(n_timesteps, num_out)
        # definition of the model (nn.Embedding) to embed the time
        self.embed.weight.data.uniform_()

    def forward(self, x, t):
        '''This gets executed when `DenseWithTime` is called.'''
        data = self.lin(x)
        t_embedded = self.embed(t)
        # embedding the specific time <t>
        out = t_embedded.view(-1, self.num_out) * data
        # reshaping time embedding and multiplying it by X
        return out
        
class ConditionalModel(nn.Module):
    def __init__(self, n_features, n_timesteps):
        '''Model layers definition.'''
        super(ConditionalModel, self).__init__()
        self.lin1 = DenseWithTime(n_features, 128, n_timesteps)
        self.lin2 = DenseWithTime(128, 128, n_timesteps)
        self.lin3 = DenseWithTime(128, 128, n_timesteps)
        self.out = nn.Linear(128, n_features)
    
    def forward(self, x, t):
        '''This gets executed when `model` is called.'''
        x = F.relu(self.lin1(x, t))
        x = F.relu(self.lin2(x, t))
        x = F.relu(self.lin3(x, t))
        return self.out(x)
    
model = ConditionalModel(X_train.shape[1], n_timesteps)
optimizer = optim.Adam(model.parameters(), lr=1e-3)    

%%time

losses = {'train':[], 'valid':[]}

for epoch in range(n_epochs):
    permutation = torch.randperm(X_train.size()[0])
    # random permutations of the batch indexes
    epoch_loss_cum = 0
    # cumulative epoch loss over batches
    
    # Iterating over batches:
    for i in range(0, X_train.size()[0], batch_size):
        # Retrieving current batch:
        indices = permutation[i:i+batch_size]
        batch_x = X_train[indices]
        batch_size_ = len(batch_x)
        # calculating actual size since the last batch can be smaller than <batch_size>
        
        # Algorithm 1 line 3: selecting a random timestep for each sample in the batch:
        t = torch.randint(0, n_timesteps, size=(batch_size_,))
        
        # Sampling noise (first term in Loss) --------------------------------
        e = torch.randn_like(batch_x)
        #---------------------------------------------------------------------

        # Predicting noise using NN (second term in Loss Equation) -----------
        # - The "extract" calls are just convenience functions which pick a
        #    value of the alphas at time <t> and reshape it to the batch size
        # - The core of the prediction is where "model" is called 
        #
        # Multiplier of x_0 in the Equation:
        alpha_bar_sqrt = extract(alphas_bar_sqrt, t, batch_x)
        # Multiplier of epsilon in the Equation:
        one_minus_alpha_bar_sqrt = extract(one_minus_alphas_bar_sqrt, t, batch_x)
        # Model input:
        input_e = batch_x * alpha_bar_sqrt + one_minus_alpha_bar_sqrt*e
        # Model prediction:
        e_theta = model(input_e, t)
        #---------------------------------------------------------------------
        
        # Loss ---------------------------------------------------------------
        # Loss Equation:
        loss = (e - e_theta).square().mean()
        # NOTE: average loss across samples and features, for this epoch, this batch
        epoch_loss_cum += loss.detach().cpu().numpy()
        #---------------------------------------------------------------------

        # Optimizing the NN --------------------------------------------------
        # Before the backward pass, zero-ing all the network gradients:
        optimizer.zero_grad()
        # Backward pass: computing gradient of the loss with respect to parameters
        loss.backward()
        # Performing gradient clipping
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.)
        # Calling the step function to update the parameters:
        optimizer.step()
        #---------------------------------------------------------------------

    ###############################################
    # Done! The rest is just to display stuff ... #
    ###############################################

    # Storing loss:
    n_batches = np.ceil(X_train.size()[0]/batch_size)
    losses['train'].append(epoch_loss_cum/n_batches)
    # averaging over batches

    # Printing loss (every 1/10th of the total epochs):
    if (epoch % int(n_epochs/10) == 0):       
        print('Epoch: %-3s [/%3s] | Loss: %s' % (epoch, n_epochs, loss.detach().cpu().numpy()))

    # Plotting data (every 1/10th of the total epochs):
    if (epoch % int(n_epochs/10) == 0):       
        
        # Generating a whole NEW time sequence by drawing from the learnt distribution:
        x_seq_gen = p_sample_loop(model, X_train.shape, n_timesteps,
                                  alphas,betas, one_minus_alphas_bar_sqrt)
        
        # Looping through all timesteps for this epoch, and plotting a few of them:
        n_plots = 8
        fig, axs = plt.subplots(1, n_plots, figsize=(28, 3))
        for j, t in enumerate(np.linspace(0, n_timesteps-1, n_plots)):
            t = int(t)
            # Picking dataset generated for timestep <t>:
            x_t_gen = x_seq_gen[t]
            if device != 'cpu': x_t_gen = x_t_gen.detach().cpu().numpy()            
            # Reverting normalization:
            df_data_gen = pd.DataFrame(data=x_t_gen, columns=features)
            df_data_gen[features] = scaler.inverse_transform(df_data_gen[features])
            # Plotting:
            plot_data(df_data_gen, features, display=False, ax=axs[j])
            axs[j].xaxis.label.set_visible(False) 
            axs[j].yaxis.label.set_visible(False) 
            axs[j].set_title('$q(\mathbf{x}_{'+str(t)+'})$')
        plt.tight_layout()
        plt.show()

# Plotting the Loss:
import matplotlib.pyplot as plt

if device != 'cpu':
    losses_train = [loss for loss in losses['train']]
    losses_valid = [loss for loss in losses['valid']]

fig, axes = plt.subplots(1, 1, figsize=(6, 4))
ax = axes
ax.plot(losses_train, label='train')
ax.plot(losses_valid, label='valid')
# NOTE: Plotting every few epochs to avoid cluttering

ax.legend()
ax.set_xlabel('Epoch')
ax.set_ylabel('Loss')

plt.tight_layout()
plt.show() 

Animation#

You don’t need to care about the following code, it is only for presentation purposes!

We will create an image sequence that resembles the evolution of the data along the timesteps, both in the forward (diffusion) and in the backward (denoising) pass.

In the GIF:

we first diffuse the original data
we then start from random noise and converge towards the sampled distribution

IMPORTANT: Notice that those are two different datasets! The first is the training dataset being diffused, the second is a new dataset randomly sampled (from noise) and inverse-diffused!

I.e., the start of the inverse-diffusion is not the end-point of the diffusion! This is the whole point of learning the distribution $-$ we want to sample new data!

They sequences are just collated together to display the intuition behind the process.

import matplotlib
from matplotlib.colors import to_hex
import io
from PIL import Image
from tqdm import tqdm

cmap = matplotlib.colors.LinearSegmentedColormap.from_list("", ["tomato","teal","steelblue"])

Generating the diffusion image sequence#

%%time
    
imgs_diffusion = []

pbar = tqdm(range(n_timesteps))
for t in pbar:
    plt.clf()
    q_x_t = q_x(X_train, torch.tensor([t]))
    if device != 'cpu': q_x_t = q_x_t.detach().cpu().numpy()

    # Indexing colors using first-half palette (the rest is used for the de-noising!):
    color = to_hex(cmap(t/n_timesteps/2))

    # Reverting normalization:
    df_data_noise = pd.DataFrame(data=q_x_t, columns=features)
    df_data_noise[features] = scaler.inverse_transform(df_data_noise[features])
    # Plotting:
    plt.title('Diffusing training set')
    plot_data(df_data_noise, features, display=False, color=color)

    img_buf = io.BytesIO()
    plt.savefig(img_buf, format='png')
    img = Image.open(img_buf)
    imgs_diffusion.append(img)

Generating the de-noising [reverse] image sequence#

Let’s generate a brand new dataset composed of n_samples_gen samples!

%%time

imgs_inv_diffusion = []

# Generating a whole NEW time sequence by drawing from the learnt distribution:
n_samples_gen = 2000
X_gen_shape = (n_samples_gen, X_train.shape[1])
x_seq_gen = p_sample_loop(model, X_gen_shape, n_timesteps,
                          alphas,betas, one_minus_alphas_bar_sqrt)

pbar = tqdm(range(n_timesteps))
for t in pbar: 
    plt.clf()
    # Picking dataset generated for timestep <t>:
    x_t_gen = x_seq_gen[t]
    if device != 'cpu': x_t_gen = x_t_gen.detach().cpu().numpy() 

    # Indexing colors using second-half palette (the rest is used for the noising!):
    color = to_hex(cmap(t/n_timesteps/2+0.5))

    # Reverting normalization:
    df_data_gen = pd.DataFrame(data=x_t_gen, columns=features)
    df_data_gen[features] = scaler.inverse_transform(df_data_gen[features])
    # Plotting:
    plt.title('Generating a new dataset')
    plot_data(df_data_gen, features, display=False, color=color)
    
    img_buf = io.BytesIO()
    plt.savefig(img_buf, format='png')
    img = Image.open(img_buf)
    imgs_inv_diffusion.append(img)    

# Pasting the 2 sequences into a sequential GIF animation:
imgs_timeline = imgs_diffusion + imgs_inv_diffusion

imgs_timeline[0].save(path_results+'/diffusion_BPT.gif', format='GIF', append_images=imgs_timeline,
             save_all=True, duration=50, loop=0)

import IPython

def show_image(filename):
    with open(filename, 'rb') as f:
        img = f.read()
    IPython.display.display(IPython.display.Image(data=img, format='gif'))

show_image(path_results+'/diffusion_BPT.gif')

[End of solution]

Question

What did you observe?

[Spoiler]

The $\beta$ schedule/epochs should be such to guarantee complete noise diffusion

Corollary: Because $\beta_t$ is the same for every feature, it is important that the data are normalized
A simpler model suffices for the BTP: the roll is a more complicated dataset!

Corollary: As usual, a validation set is actually needed to avoid overfitting!

Measuring the difference via Kullback–Leibler (KL) divergence#

How can we quantitatively assess the difference between the original and a sampled distribution?

Or, how do we know that our model produces trustworthy samples?

X_gen = x_seq_gen[-1][:X_train.size()[0]]
# extracting the last generated dataset, i.e. the one at time t

if device != 'cpu':
    X_gen_   = X_gen.detach().cpu().numpy()
    X_train_ = X_train.detach().cpu().numpy()
else:
    X_gen_   = X_gen.detach().numpy()
    X_train_ = X_train.detach().numpy()
    
plt.scatter(X_train_[:,0], X_train_[:,1], alpha=0.5, label='original')
plt.scatter(X_gen_[:,0],   X_gen_[:,1],   alpha=0.5, label='sampled')
plt.legend()
plt.show()

Kullback-Leibler Divergence

\[D_{KL} (~p(x)||q(x)~) = \sum_{x \in X} p(x)~ln{p(x) \over q(x)}\]

Intuitively: Information represented by the approximated distribution $q(x)$ when the original distribution is $p(x)$ $-$ see e.g. here.

The video below (from here) shows the evolution of the K-L divergence as $q$ varies $\rightarrow$ Notice that $D_{KL} = 0$ $\Leftrightarrow$ $p(x)$ = $q(x)$

from IPython.display import Video
Video("images/KL_Divergence.mp4", embed=True, width=600)

from src.KLdivergence import KLdivergence

print('KL divergence between training data and itself: %.2f' %\
     KLdivergence(X_train_, X_train_+1e-10))
# ... plus a tiny number to avoid division by 0
print('KL divergence between generated data and training data: %.2f' %\
     KLdivergence(X_gen_, X_train_))

###EOF

Diffusion Models

Contents

Diffusion Models#

Setup#

POV: You need to model a distribution!#

Generative Models \(vs\) Discriminative Models#

Diffusion models#

What are Diffusion Models?#

One nifty eyey-candy example#

The model, the myth, the math#

My first diffusion model!#

Generating some mock data#

Visualizing data#

Splitting#

Scaling#

Diffusion Parameters#

Forward Diffusion#

Reverse Diffusion/Reconstruction#

Animation#

Generating the diffusion image sequence#

Generating the de-noising [reverse] image sequence#

In-class Exercise: diffuse a BPT#

Importing Data#

Visualizing data#

Splitting#

Scaling#

Diffusion Parameters#

Forward Diffusion#

Reverse Diffusion/Reconstruction#

Animation#

Generating the diffusion image sequence#

Generating the de-noising [reverse] image sequence#

Measuring the difference via Kullback–Leibler (KL) divergence#