Denoising Diffusion Probabilistic Models (DDPMs) are a fascinating way to do generative modeling. In essence, the idea is to construct a forward Markov chain that gradually adds noise starting with a clean sample of your data, and going all the way to pure noise. The reverse process recovers successively denoised versions of the data starting with pure noise, and ending with a generated data sample that is from the same distribution as the data the model was trained on. This simple-sounding method showed striking reuslts on generating natural images with this 2020 paper, although the idea of using diffusion processes in machine learning is from way back in this 2015 paper. In the last 5 years, the use of DDPMs has grown well beyond computer vision – it is a central workhorse in generative modeling of molecules for de novo protein based drug design.

I wrote a simple minDDPM model to:

  • Generate a scalar random variable from a mixture of 3 scalar Gaussians. The distritbution of generated values looks very close to the distribution of training data. The denoising network here was a simple MLP.
  • Generate a random sinewave of random amplitude and phase. A bit more complicated than generaing scalars. The generated sine-waves look okayish, not perfect. The denoising network here was a simple CNN.

Note that at generation time, the same initial pure noise starting point can provide many different samples, i.e. the generation process is not deterministic. This is because at each denoising step, we sample from a conditional distribution rather than taking a single deterministic update. More concretely, in a DDPM we model a reverse transition of the form:

\[p_{\theta}(x_{t-1}|x_t)\]

which is typically a Gaussian (in continuous diffusion). Even though the model predicts parameters of this distribution (e.g., mean and variance, or logits), we draw a random sample from it at every step.

So during generation one starts from the same \(x_T\) (pure noise), but at each step \(x_t \rightarrow x_{t-1}\), one injects randomness by sampling. Small differences introduced early get amplified over subsequent steps, leading to different final samples. In other words, the reverse process is a stochastic Markov chain, and the randomness at each step is what enables the model to generate diverse outputs rather than always collapsing to a single mode.

As a side note, if you remove this sampling and always take the mean of the distribution, the process becomes deterministic – which leads to a variant called DDIM (I = implicit) which allows for faster sampling (fewer denoising steps) at the cost of just one deterministic sampling for any given initial noise input \(x_T\). During the reverse sampling process starting with \(x_T\), crude details of the final sample \(x_0\) take shape first, and then finer details emerge.

A large part of the art in diffusion modeling is designing the denoising network. That is, the network that learns to predict the noise that was added to a clearn data sample to create the corrupted data sample, given as inputs: the corrupted sample and the noise level. Typically, this network ends up corresponding to the type of data modality and the inductive biases needed to capture structure in the data. For example, images typically employ some type of convolutional architecture, for e.g. a U-net.

While diffusion modeling is commonly done on continuous data, it is also possible to construct diffusion on discrete spaces – for example, a sequence of tokens as in natural laungauge or protein sequences. This formulation is called Discrete DDPM, or D3PM. I wrote a minD3PM to learn to generate sequences of ordered integers of a fixed sequence length with wrap around and random phase. For example, using a sequence length of 7 and vocab size of 10 (0, 1, .. 9), one would generate sequences like [4,5,6,7,8,9,0], [6,7,8,9,0,1,2], and so on. In the discrete diffusion set up, adding noise can mean either randomly flipping a token to a different token, or masking a token. In my model, I chose masking only. A simple transformer-based denoising network is able to learn the distribution of training data and generate very realistic looking sequences.

This looks a lot like training a BERT model because the architecture uses transformer layers, and because we are trying to predict the identity of masked tokens. The difference of course is that BERT just uses a single noise level (= masking probability) and of course, BERT is not a generative model, i.e. it is meant for representation learning. the minD3PM in my simple project can generate a sequence from the source distribution starting with a sequence of all mask tokens.