Diffusion Decoder based Compression

I'm always super interested in expanding the frontier of learned compression. The current methods we use are derived from VAEs (equivalence shown in [3]) and use GANs to improve perceptual quality but in this blog I explore how diffusion models could be used for compression.

The feasibility of any method we explore in a learned compression pipeline hinges on the complexity of that method. Since decode and sometimes encode are performed on battery constrained edge devices, the scope of a new method to add compute is very limited, often 1-2KMACs/px. As a result training time improvements are favoured over encode/decode. In this blog though, we'll go against this advice and augment our decoder with a denoising diffusion model.

Intro to Diffusion Models

Before we dive into diffusion models, let's summarise the main methods for generating from a distribution currently available to us, and see how diffusion models [1, 2] look compared to them. This is shown in Figure 1.

Figure 1: Summary of various generative models

The forward process

The forward step in a diffusion model consists of adding small amounts of Gaussian noise to our image (which is sampled from our data distribution), until it looks like a Gaussian sample from $N (0, I)$ . We define a number of steps we want to achieve this transition in, and at each step, we add noise with a variance we pre-set (we call this the variance schedule). Since each transition depends only on the present state, the forward process is Markovian and can be written as:

q (x_{t} | x_{t - 1}) = N (x_{t}; \sqrt{1 - β_{t}} x_{t - 1}, β_{t} I)

q (x_{1 : T} | x_{0}) = \prod_{t = 1}^{T} q (x_{t} | x_{t - 1})

Here, $β_{t}$ is our variance at time step $t$ . As well as adding noise, we also scale down the previous sample by $\sqrt{1 - β_{t}}$ which is required if we are to bound the variance of $x_{t}$ in the limit of a standard normal. This will be graphically shown later.

Now, let's see what this means in practice. Let's take an image from the Kodak dataset, and define a variance schedule over $T = 1000$ steps, starting at $β = 10^{- 4}$ and ending at $β = 0.02$ . Figures 2 and 3 below show a grey scale image and its grey value pixel distribution.

Figure 2: Greyscale image forward noising process up to T=300

Figure 3: Histogram of normalised pixel values in forward noising process up to T=300

Let's have a closer look at what is actually happening in the forward noising process and some of its useful properties. Essentially, at each step $t$ we are adding small amounts of Gaussian noise with a particular variance to the sample from the previous step, $t - 1$ , while scaling it down with $\sqrt{1 - β_{t}}$ . To help with notation, let's write $β_{t} = 1 - α_{t}$ . Figure 4 shows the forward step diagrammatically.

Figure 4: Forward process, as an addition of Gaussians, step by step!

Figure 4 shows that at any point in the forward process, we can break down the current sample $x_{t}$ into two terms:

The initial data sample we are noising ( $x_{0}$ from $p_{x}$ ) and a cumulative product of scaling coefficients for $x_{0}$ , square root of ${\bar{α}}_{t} = \prod_{s = 1}^{t} α_{s}$ (which depends on the time step we are on and the self-defined variance schedule)
A set of mean-zero Gaussians with different scales where the scales depend on the current time step and self-defined variance schedule. Using the properties of Gaussians, we can combine all these zero-mean Gaussians into one Gaussian with a particular variance, which is just the sum of the variances of all the (scaled) Gaussians up to that point and takes the form $1 - {\bar{α}}_{t}$ . A simple derivation of this is shown later.

Scaling x

Why do we need to scale down the previous input $x_{t - 1}$ at each step by $\sqrt{1 - β_{t}}$ ? Well, let's have a look at what happens if we don't. Figure 5 shows the same process as Figure 4, but where the $x_{t - 1}$ is not scaled down by $\sqrt{1 - β_{t}}$ . We observe that the variance of the distribution continues to grow, so it keeps getting wider, which means it is not an $N (0, I)$ .

Figure 5: Distribution of forward process without scaling down previous input

Sampling $x_{t}$ at an arbitrary time step $t$

Since we are adding a lot of Gaussians to a sample, this affords us the nice properties Gaussians bring with them. One useful property we can extract from this is the ability to sample a noisy $x_{t}$ at any time step (without having to go through the entire forward process up to that point), given an $x_{0}$ . This is because we define the variance schedule and therefore, at any time steps, we know all the Gaussian samples that were added to $x_{0}$ to get that particular noisy $x_{t}$ . Figure 4 already shows this, but here we provide more detail.

\begin{aligned} t = 0 : & x_{0} \sim p_{x} (x) \\ t = 1 : & x_{1} = \sqrt{1 - β_{1}} x_{0} + N (0, β_{1} I) \\ = \sqrt{1 - β_{1}} x_{0} + 0 + \sqrt{β_{1}} z_{1} \\ t = 2 : & x_{2} = \sqrt{1 - β_{2}} x_{1} + N (0, β_{2} I) \\ = \sqrt{1 - β_{2}} (\sqrt{1 - β_{1}} x_{0} + \sqrt{β_{1}} z_{1}) + \sqrt{β_{2}} z_{2} \\ = \sqrt{α_{2}} (\sqrt{α_{1}} x_{0} + \sqrt{1 - α_{1}} z_{1}) + \sqrt{1 - α_{2}} z_{2} \\ = \sqrt{α_{2} α_{1}} x_{0} + \sqrt{(1 - α_{1}) α_{2} + (1 - α_{2})} \bar{z} \\ = \sqrt{{\bar{α}}_{2}} x_{0} + \sqrt{(1 - {\bar{α}}_{2})} \bar{z} \end{aligned}

For any $t$ :

x_{t} = \sqrt{{\bar{α}}_{t}} x_{0} + \sqrt{(1 - {\bar{α}}_{t})} \bar{z}

Here, $\bar{z}$ is a sample from a standard normal. The figure below demonstrates how we can jump from $x_{0}$ directly to $x_{50}$ and $x_{300}$ through combining all the variances of the Gaussians in the chain up to those points. This is an extremely useful property and will help us in training as we shall see later.

Figure 6: Sampling at an arbitrary step T

The reverse process

The learned reverse process is written as:

p_{θ} (x_{t - 1} | x_{t}) = N (x_{t - 1}; μ_{θ} (x_{t}, t), Σ_{θ} (x_{t}, t) I) .

Like variational inference, $p_{θ}$ is an approximate distribution for the reverse of our forward process, $q (x_{t - 1} ∣ x_{t})$ , and we close the gap between the two using KL-divergence.

$q (x_{t - 1} ∣ x_{t})$ is generally intractable; however, it can be shown to be tractable when conditioned on $x_{0}$ . This results in the following formulation of what we call the forward process posterior:

q (x_{t - 1} | x_{t}, x_{0}) = N (μ_{post}, {\tilde{β}}_{t} I)

where the posterior mean is:

μ_{post} = \frac{\sqrt{{\bar{α}}_{t - 1}} β_{t}}{1 - {\bar{α}}_{t}} x_{0} + \frac{\sqrt{α_{t}} (1 - {\bar{α}}_{t - 1})}{1 - {\bar{α}}_{t}} x_{t}

{\tilde{β}}_{t} = \frac{β_{t} (1 - {\bar{α}}_{t - 1})}{1 - {\bar{α}}_{t}}

This can be derived with (2.116) in Bishop's book. Therefore, if we can predict our posterior mean $μ_{post}$ and our variance (which can easily be computed as all terms are always known), we should be able to sample from our reverse posterior $q (x_{t - 1} | x_{t}, x_{0})$ .

If we have a look at this mean, during the reverse/sampling process, we know $\bar{α}$ and $β$ as these are self-defined. $x_{t}$ is the current step we are on, which we know too (since we start with a sample from $N (0, I)$ ). The only term we do not know at sampling time is $x_{0}$ . This is what we need a neural network to help us predict. We can train a network $f_{θ} (x_{t}, t)$ to directly predict this, given $x_{t}$ and $t$ as inputs, however, this is empirically shown not to work too well [1], and a better method is to do the following:

We know from the forward process that:

x_{t} = \sqrt{{\bar{α}}_{t}} x_{0} + \sqrt{(1 - {\bar{α}}_{t})} ε

where $ε \sim N (0, I)$ which we can rearrange to get $x_{0}$ . Once $ε$ is predicted, this gives an approximation of $x_{0}$ , say ${\tilde{x}}_{0}$ :

{\tilde{x}}_{0} = \frac{1}{\sqrt{{\bar{α}}_{t}}} (x_{t} - \sqrt{1 - {\bar{α}}_{t}} ε)

We can plug this formula in for $x_{0}$ in the forward process posterior definition to obtain the reverse sampling step:

x_{t - 1} = \frac{1}{\sqrt{α_{t}}} (x_{t} - \frac{β_{t}}{\sqrt{1 - {\bar{α}}_{t}}} ε) + \sqrt{{\tilde{β}}_{t}} z, z \sim N (0, I)

This leaves us with the unknown noise sample $ε$ from a standard normal which we need to predict; all other terms are known. We can define a neural network to predict this term:

ε \approx ε_{θ} (x_{t}, t)

Once we have estimated $ε$ , we can compute ${\tilde{x}}_{0}$ , after which we can predict $μ_{post}$ and sample the $x_{t - 1}$ . Then we can repeat this process, until we get to $x_{0}$ which is our sample from the distribution. Figure 7 shows the samples this process passes through, from $x_{T}$ on the right back to $x_{0}$ on the left, for example from $x_{4}$ to $x_{3}$ and then to $x_{2}$ .

Figure 7: The reverse process from x_4 to x_3 and then from x_3 to x_2

What does ${\tilde{x}}_{0}$ look like at each time step? That depends on where you are; if you are quite early in the chain, it looks like noise with some structure. If you are close to the end of the chain, it almost looks like an image. An example is shown in Figure 8, taken from the DDPM paper.

Figure 8: Approximated x̃_0 given a sample x_t, where T → t → 0 from left to right

Training

The loss function is a maximisation problem of the variational lower bound on the log-likelihood of the data:

E_{q (x_{0})} [\log p_{θ} (x_{0})] \geq - E_{q (x_{0 : T})} [\log q (x_{1 : T} | x_{0}) - \log p_{θ} (x_{0 : T})] = - L

This can be reduced to a sum of KL-divergences between $q (x_{t - 1} ∣ x_{t}, x_{0})$ and $p_{θ} (x_{t - 1} ∣ x_{t})$ . Since both are Gaussian, this is just a KL between two Gaussians. Working through this, we can end up with a simplified loss term (we need to drop some scaling constants that arise when we do the KL between the Gaussians), corresponding to a weighted variational lower bound:

L_{simple} = E_{t, x_{0}, ε} (‖ ε - ε_{θ} (x_{t}, t) ‖_{2}^{2})

This is the simplified objective we can use to train our denoising function $ε_{θ}$ , which is generally selected to be some sort of U-Net architecture. What this denoising function learns during training is the noise that was added to samples from the dataset to noise them. Once it is fully trained, it can then be used to remove that noise. The training process is then extremely simple and is performed as below:

We sample an image $x_{0}$ from our dataset and select a random time step $t$
We sample an instance of noise: $ε \sim N (0, I)$
We generate the noisy sample $x_{t}$ at this time step using the sampled noise $ε$ and variance schedule (see forward process section)
We feed this $t$ and $x_{t}$ to our denoising function $ε_{θ} (x_{t}, t)$ to approximate $ε$
We use an L2-metric to compute a loss between $ε$ and $ε_{θ}$

Conditional compression denoising decoder

Is there a way we could use diffusion models in our compression pipeline? Can they be used for explicit likelihood or implicit distribution matching?

An immediate idea would be to replace our decoder with a conditional diffusion model where we condition the diffusion process on the quantised latent space. This would enable conditional image generation, trained with the rate as usual and the denoising loss in place of the distortion. As the denoising loss is a bound on the likelihood, this is the explicit likelihood route.

There are a few changes that we would need to consider for the pipeline. Let's break them down into architecture, training and inference.

Architecture

Until the decoder, our architecture can remain the same, however, we would then need to upsample our quantised latent space to the image scale (diffusion models work on the image scale).

We need to define a function $ε_{θ}$ , which empirically is shown to work best as a U-Net architecture. The input of the U-Net will need to be the noisy image $x_{t}$ concatenated with the upsampled latent space, ${\hat{y}}_{u s}$ , along with the time step $t$ .

The training

The training function we minimise now becomes:

L = E_{x_{0} \sim p (x_{0})} (E_{t, ε \sim N (0, I)} (‖ ε - ε_{θ} (x_{t}, t, {\hat{y}}_{u s}) ‖_{2}^{2}) + λ R)

where $ε_{θ}$ has an additional input to force it to be conditional, $R$ is the rate of $\hat{y}$ and $λ$ trades the rate off against the denoising loss. Note that the decoder then produces a plausible image given $\hat{y}$ rather than the one closest to $x_{0}$ , so fidelity (e.g. PSNR) is not directly controlled.

The inference

The encoding works exactly the same, but the decoding differs. After the $\hat{y}$ is produced, we sample $x_{T} \sim N (0, I)$ , and perform the computation of the reverse process (Figure 7) for each of the $T$ steps to get our final output $\hat{x}$ .

Figure 9: Simple architecture of diffusion based compression pipeline. Here, the noise profile is denoted with z instead of with ε.

I ran the above model and it produced promising results on the first iteration. I have high hopes for the compression performance, but the largest concern is with execution time: decoding runs the U-Net once per step, $T$ times, far beyond the 1-2KMACs/px budget mentioned above. We'd need to see a lot more algorithmic progress in diffusion decoders to enable efficient decoding.