How diffusion models work: the math from scratch

Diffusion fashions are a brand new class of state-of-the-art generative fashions that generate numerous high-resolution photos. They’ve already attracted plenty of consideration after OpenAI, Nvidia and Google managed to coach large-scale fashions. Instance architectures which can be based mostly on diffusion fashions are GLIDE, DALLE-2, Imagen, and the complete open-source secure diffusion.

However what’s the foremost precept behind them?

On this weblog put up, we’ll dig our approach up from the fundamental rules. There are already a bunch of various diffusion-based architectures. We’ll deal with probably the most distinguished one, which is the Denoising Diffusion Probabilistic Fashions (DDPM) as initialized by Sohl-Dickstein et al after which proposed by Ho. et al 2020. Numerous different approaches will probably be mentioned to a smaller extent equivalent to secure diffusion and score-based fashions.

Diffusion fashions are basically completely different from all of the earlier generative strategies. Intuitively, they intention to decompose the picture technology course of (sampling) in lots of small “denoising” steps.

The instinct behind that is that the mannequin can right itself over these small steps and regularly produce a very good pattern. To some extent, this concept of refining the illustration has already been utilized in fashions like alphafold. However hey, nothing comes at zero-cost. This iterative course of makes them sluggish at sampling, not less than in comparison with GANs.

Diffusion course of

The fundamental thought behind diffusion fashions is fairly easy. They take the enter picture $mathbf{x}_0$

Afterward, a neural community is skilled to recuperate the unique knowledge by reversing the noising course of. By having the ability to mannequin the reverse course of, we are able to generate new knowledge. That is the so-called reverse diffusion course of or, generally, the sampling strategy of a generative mannequin.

How? Let’s dive into the mathematics to make it crystal clear.

Ahead diffusion

Diffusion fashions may be seen as latent variable fashions. Latent signifies that we’re referring to a hidden steady function house. In such a approach, they could look just like variational autoencoders (VAEs).

In observe, they’re formulated utilizing a Markov chain of $T$ steps. Right here, a Markov chain signifies that every step solely relies on the earlier one, which is a light assumption. Importantly, we aren’t constrained to utilizing a selected kind of neural community, in contrast to flow-based fashions.

Given a data-point $textbf{x}_0$

Ahead diffusion course of. Picture modified by Ho et al. 2020

Since we’re within the multi-dimensional state of affairs $textbf{I}$ is the id matrix, indicating that every dimension has the identical normal deviation $beta_t$

Thus, we are able to go in a closed type from the enter knowledge $mathbf{x}_0$

The image $:$ in $q(mathbf{x}_{1:T})$

To date, so good? Effectively, nah! For timestep $t=500 < T$

The reparametrization trick offers a magic treatment to this.

The reparameterization trick: tractable closed-form sampling at any timestep

If we outline $alpha_t= 1- beta_t$

Be aware: Since all timestep have the identical Gaussian noise we’ll solely use the image $boldsymbol{epsilon}$ any more.

Thus to provide a pattern $mathbf{x}_t$

Since $beta_t$

Variance schedule

The variance parameter $beta_t$

Latent samples from linear (high) and cosine (backside)
schedules respectively. Supply: Nichol & Dhariwal 2021

Reverse diffusion

As $T to infty$

The query is how we are able to mannequin the reverse diffusion course of.

Approximating the reverse course of with a neural community

In sensible phrases, we do not know $q(mathbf{x}_{t-1} vert mathbf{x}_{t})$

As a substitute, we approximate $q(mathbf{x}_{t-1} vert mathbf{x}_{t})$

Reverse diffusion course of. Picture modified by Ho et al. 2020

If we apply the reverse formulation for all timesteps ($p_theta(mathbf{x}_{0:T})$

By moreover conditioning the mannequin on timestep $t$, it’s going to study to foretell the Gaussian parameters (which means the imply $boldsymbol{mu}_theta(mathbf{x}_t, t)$

However how can we practice such a mannequin?

Coaching a diffusion mannequin

If we take a step again, we are able to discover that the mixture of $q$ and $p$ is similar to a variational autoencoder (VAE). Thus, we are able to practice it by optimizing the unfavourable log-likelihood of the coaching knowledge. After a sequence of calculations, which we cannot analyze right here, we are able to write the proof decrease certain (ELBO) as follows:

Let’s analyze these phrases:

1. The $mathbb{E}_{q(x_1 vert x_0)} [log p_{theta} (mathbf{x}_0 vert mathbf{x}_1)]$

2. $D_{KL}(q(mathbf{x}_T vert mathbf{x}_0) vertvert p(mathbf{x}_T))$

3. The third time period $sum_{t=2}^T L_{t-1}$

It’s evident that by means of the ELBO, maximizing the probability boils all the way down to studying the denoising steps $L_t$

Essential notice: Though $q(mathbf{x}_{t-1} vert mathbf{x}_{t})$

Intuitively, a painter (our generative mannequin) wants a reference picture ($textbf{x}_0$

In different phrases, we are able to pattern $textbf{x}_t$

Be aware that $alpha_t$

This little trick offers us with a completely tractable ELBO. The above property has yet one more essential facet impact, as we already noticed within the reparameterization trick, we are able to signify $mathbf{x}_0$

the place $boldsymbol{epsilon} sim mathcal{N}(textbf{0},mathbf{I})$

By combining the final two equations, every timestep will now have a imply $tilde{boldsymbol{mu}}_t$

Subsequently we are able to use a neural community $epsilon_{theta}(mathbf{x}_t,t)$

Thus, the loss perform (the denoising time period within the ELBO) may be expressed as:

This successfully reveals us that as an alternative of predicting the imply of the distribution, the mannequin will predict the noise $boldsymbol{epsilon}$ at every timestep $t$.

Ho et.al 2020 made a number of simplifications to the precise loss time period as they ignore a weighting time period. The simplified model outperforms the complete goal:

The authors discovered that optimizing the above goal works higher than optimizing the unique ELBO. The proof for each equations may be discovered on this glorious put up by Lillian Weng or in Luo et al. 2022.

Moreover, Ho et. al 2020 resolve to maintain the variance mounted and have the community study solely the imply. This was later improved by Nichol et al. 2021, who resolve to let the community study the covariance matrix $(boldsymbol{Sigma})$ as properly (by modifying $L_t^textual content{easy}$

Coaching and sampling algorithms of DDPMs. Supply: Ho et al. 2020

Structure

One factor that we’ve not talked about to this point is what the mannequin’s structure seems like. Discover that the mannequin’s enter and output must be of the identical dimension.

To this finish, Ho et al. employed a U-Internet. In case you are unfamiliar with U-Nets, be at liberty to take a look at our previous article on the most important U-Internet architectures. In a number of phrases, a U-Internet is a symmetric structure with enter and output of the identical spatial dimension that makes use of skip connections between encoder and decoder blocks of corresponding function dimension. Often, the enter picture is first downsampled after which upsampled till reaching its preliminary dimension.

Within the authentic implementation of DDPMs, the U-Internet consists of Large ResNet blocks, group normalization in addition to self-attention blocks.

The diffusion timestep $t$ is specified by including a sinusoidal place embedding into every residual block. For extra particulars, be at liberty to go to the official GitHub repository. For an in depth implementation of the diffusion mannequin, take a look at this superior put up by Hugging Face.

The U-Internet structure. Supply: Ronneberger et al.

Conditional Picture Technology: Guided Diffusion

An important side of picture technology is conditioning the sampling course of to govern the generated samples. Right here, that is additionally known as guided diffusion.

There have even been strategies that incorporate picture embeddings into the diffusion with a view to “information” the technology. Mathematically, steerage refers to conditioning a previous knowledge distribution $p(textbf{x})$ with a situation $y$, i.e. the category label or a picture/textual content embedding, leading to $p(textbf{x}|y)$.

To show a diffusion mannequin $p_theta$

The truth that the conditioning is being seen at every timestep could also be a very good justification for the superb samples from a textual content immediate.

Basically, guided diffusion fashions intention to study $nabla log p_theta( mathbf{x}_t vert y)$

$p_theta(y)$

And by including a steerage scalar time period $s$, we’ve got:

Utilizing this formulation, let’s make a distinction between classifier and classifier-free steerage. Subsequent, we’ll current two household of strategies aiming at injecting label data.

Classifier steerage

Sohl-Dickstein et al. and later Dhariwal and Nichol confirmed that we are able to use a second mannequin, a classifier $f_phi(y vert mathbf{x}_t, t)$

We will construct a class-conditional diffusion mannequin with imply $mu_theta(mathbf{x}_t|y)$

Since $p_theta sim mathcal{N}(mu_{theta}, Sigma_{theta})$

Within the well-known GLIDE paper by Nichol et al, the authors expanded on this concept and use CLIP embeddings to information the diffusion. CLIP as proposed by Saharia et al., consists of a picture encoder $g$ and a textual content encoder $h$. It produces a picture and textual content embeddings $g(mathbf{x}_t)$

Subsequently, we are able to perturb the gradients with their dot product:

Because of this, they handle to “steer” the technology course of towards a user-defined textual content caption.

Algorithm of classifier guided diffusion sampling. Supply: Dhariwal & Nichol 2021

Classifier-free steerage

Utilizing the identical formulation as earlier than we are able to outline a classifier-free guided diffusion mannequin as:

Steerage may be achieved with out a second classifier mannequin as proposed by Ho & Salimans. As a substitute of coaching a separate classifier, the authors skilled a conditional diffusion mannequin $boldsymbol{epsilon}_theta (mathbf{x}_t|y)$

Be aware that this may also be used to “inject” textual content embeddings as we confirmed in classifier steerage.

This admittedly “bizarre” course of has two main benefits:

• It makes use of solely a single mannequin to information the diffusion.

• It simplifies steerage when conditioning on data that’s tough to foretell with a classifier (equivalent to textual content embeddings).

Imagen as proposed by Saharia et al. depends closely on classifier-free steerage, as they discover that it’s a key contributor to producing samples with sturdy image-text alignment. For more information on the method of Imagen take a look at this video from AI Espresso Break with Letitia:

Scaling up diffusion fashions

You is likely to be asking what’s the drawback with these fashions. Effectively, it is computationally very costly to scale these U-nets into high-resolution photos. This brings us to 2 strategies for scaling up diffusion fashions to greater resolutions: cascade diffusion fashions and latent diffusion fashions.

Cascade diffusion fashions

Ho et al. 2021 launched cascade diffusion fashions in an effort to provide high-fidelity photos. A cascade diffusion mannequin consists of a pipeline of many sequential diffusion fashions that generate photos of accelerating decision. Every mannequin generates a pattern with superior high quality than the earlier one by successively upsampling the picture and including greater decision particulars. To generate a picture, we pattern sequentially from every diffusion mannequin.

Cascade diffusion mannequin pipeline. Supply: Ho & Saharia et al.

To amass good outcomes with cascaded architectures, sturdy knowledge augmentations on the enter of every super-resolution mannequin are essential. Why? As a result of it alleviates compounding error from the earlier cascaded fashions, in addition to because of a train-test mismatch.

It was discovered that gaussian blurring is a important transformation towards attaining excessive constancy. They discuss with this system as conditioning augmentation.

Secure diffusion: Latent diffusion fashions

Latent diffusion fashions are based mostly on a fairly easy thought: as an alternative of making use of the diffusion course of straight on a high-dimensional enter, we venture the enter right into a smaller latent house and apply the diffusion there.

In additional element, Rombach et al. proposed to make use of an encoder community to encode the enter right into a latent illustration i.e. $mathbf{z}_t = g(mathbf{x}_t)$

If the loss for a typical diffusion mannequin (DM) is formulated as:

then given an encoder $mathcal{E}$ and a latent illustration $z$, the loss for a latent diffusion mannequin (LDM) is:

Latent diffusion fashions. Supply: Rombach et al

For extra data take a look at this video:

Rating-based generative fashions

Across the similar time because the DDPM paper, Tune and Ermon proposed a unique kind of generative mannequin that seems to have many similarities with diffusion fashions. Rating-based fashions sort out generative studying utilizing rating matching and Langevin dynamics.

Rating-matching refers back to the strategy of modeling the gradient of the log likelihood density perform, also referred to as the rating perform. Langevin dynamics is an iterative course of that may draw samples from a distribution utilizing solely its rating perform.

the place $delta$ is the step dimension.

Suppose that we’ve got a likelihood density $p(x)$ and that we outline the rating perform to be $nabla_x log p(x)$

Then through the use of Langevin dynamics, we are able to straight pattern from $p(x)$ utilizing the approximated rating perform.

In case you missed it, guided diffusion fashions use this formulation of score-based fashions as they study straight $nabla_x log p(x)$

Including noise to score-based fashions: Noise Conditional Rating Networks (NCSN)

The issue to this point: the estimated rating features are normally inaccurate in low-density areas, the place few knowledge factors can be found. Because of this, the standard of knowledge sampled utilizing Langevin dynamics is not good.

Their resolution was to perturb the information factors with noise and practice score-based fashions on the noisy knowledge factors as an alternative. As a matter of reality, they used a number of scales of Gaussian noise perturbations.

Thus, including noise is the important thing to make each DDPM and rating based mostly fashions work.

Rating-based generative modeling with rating matching + Langevin dynamics. Supply: Generative Modeling by Estimating Gradients of the Information Distribution

Mathematically, given the information distribution $p(x)$, we perturb with Gaussian noise $mathcal{N}(textbf{0}, sigma_i^2 I)$

Then we practice a community $s_theta(mathbf{x},i)$

Rating-based generative modeling by means of stochastic differential equations (SDE)

Tune et al. 2021 explored the connection of score-based fashions with diffusion fashions. In an effort to encapsulate each NSCNs and DDPMs beneath the identical umbrella, they proposed the next:

As a substitute of perturbing knowledge with a finite variety of noise distributions, we use a continuum of distributions that evolve over time in response to a diffusion course of. This course of is modeled by a prescribed stochastic differential equation (SDE) that doesn’t depend upon the information and has no trainable parameters. By reversing the method, we are able to generate new samples.

Rating-based generative modeling by means of stochastic differential equations (SDE). Supply: Tune et al. 2021

We will outline the diffusion course of ${ mathbf{x}(t) }_{tin [0, T]}$

the place $mathbf{w}$ is the Wiener course of (a.okay.a., Brownian movement), $mathbf{f}(cdot, t)$ is a vector-valued perform referred to as the drift coefficient of $mathbf{x}(t)$, and $g(cdot)$ is a scalar perform generally known as the diffusion coefficient of $mathbf{x}(t)$. Be aware that the SDE usually has a novel sturdy resolution.

To make sense of why we use an SDE, here’s a tip: the SDE is impressed by the Brownian movement, through which quite a lot of particles transfer randomly inside a medium. This randomness of the particles’ movement fashions the continual noise perturbations on the information.

After perturbing the unique knowledge distribution for a sufficiently very long time, the perturbed distribution turns into near a tractable noise distribution.

To generate new samples, we have to reverse the diffusion course of. The SDE was chosen to have a corresponding reverse SDE in closed type:

To compute the reverse SDE, we have to estimate the rating perform $nabla_mathbf{x} log p_t(mathbf{x})$

the place $mathcal{U}(0, T)$ denotes a uniform distribution over the time interval, and $lambda$ is a constructive weighting perform. As soon as we’ve got the rating perform, we are able to plug it into the reverse SDE and remedy it with a view to pattern $mathbf{x}(0)$ from the unique knowledge distribution $p_0(mathbf{x})$

There are a selection of choices to unravel the reverse SDE which we cannot analyze right here. Be sure that to verify the unique paper or this glorious weblog put up by the creator.

Overview of score-based generative modeling by means of SDEs. Supply: Tune et al. 2021

Abstract

Let’s do a fast sum-up of the details we discovered on this blogpost:

• Diffusion fashions work by regularly including gaussian noise by means of a sequence of $T$ steps into the unique picture, a course of generally known as diffusion.

• To pattern new knowledge, we approximate the reverse diffusion course of utilizing a neural community.

• The coaching of the mannequin relies on maximizing the proof decrease certain (ELBO).

• We will situation the diffusion fashions on picture labels or textual content embeddings with a view to “information” the diffusion course of.

• Cascade and Latent diffusion are two approaches to scale up fashions to high-resolutions.

• Cascade diffusion fashions are sequential diffusion fashions that generate photos of accelerating decision.

• Latent diffusion fashions (like secure diffusion) apply the diffusion course of on a smaller latent house for computational effectivity utilizing a variational autoencoder for the up and downsampling.

• Rating-based fashions additionally apply a sequence of noise perturbations to the unique picture. However they’re skilled utilizing score-matching and Langevin dynamics. Nonetheless, they find yourself in the same goal.

• The diffusion course of may be formulated as an SDE. Fixing the reverse SDE permits us to generate new samples.

Lastly, for extra associations between diffusion fashions and VAE or AE take a look at these very nice blogs.

Cite as

@article{karagiannakos2022diffusionmodels,
    title   = "Diffusion fashions: towards state-of-the-art picture technology",
    creator  = "Karagiannakos, Sergios, Adaloglou, Nikolaos",
    journal = "https://theaisummer.com/",
    12 months    = "2022",
    howpublished = {https://theaisummer.com/diffusion-fashions/},
  }

References

[1] Sohl-Dickstein, Jascha, et al. Deep Unsupervised Studying Utilizing Nonequilibrium Thermodynamics. arXiv:1503.03585, arXiv, 18 Nov. 2015

[2] Ho, Jonathan, et al. Denoising Diffusion Probabilistic Fashions. arXiv:2006.11239, arXiv, 16 Dec. 2020

[3] Nichol, Alex, and Prafulla Dhariwal. Improved Denoising Diffusion Probabilistic Fashions. arXiv:2102.09672, arXiv, 18 Feb. 2021

[4] Dhariwal, Prafulla, and Alex Nichol. Diffusion Fashions Beat GANs on Picture Synthesis. arXiv:2105.05233, arXiv, 1 June 2021

[5] Nichol, Alex, et al. GLIDE: In the direction of Photorealistic Picture Technology and Enhancing with Textual content-Guided Diffusion Fashions. arXiv:2112.10741, arXiv, 8 Mar. 2022

[6] Ho, Jonathan, and Tim Salimans. Classifier-Free Diffusion Steerage. 2021. openreview.internet

[7] Ramesh, Aditya, et al. Hierarchical Textual content-Conditional Picture Technology with CLIP Latents. arXiv:2204.06125, arXiv, 12 Apr. 2022

[8] Saharia, Chitwan, et al. Photorealistic Textual content-to-Picture Diffusion Fashions with Deep Language Understanding. arXiv:2205.11487, arXiv, 23 Could 2022

[9] Rombach, Robin, et al. Excessive-Decision Picture Synthesis with Latent Diffusion Fashions. arXiv:2112.10752, arXiv, 13 Apr. 2022

[10] Ho, Jonathan, et al. Cascaded Diffusion Fashions for Excessive Constancy Picture Technology. arXiv:2106.15282, arXiv, 17 Dec. 2021

[11] Weng, Lilian. What Are Diffusion Fashions? 11 July 2021

[12] O’Connor, Ryan. Introduction to Diffusion Fashions for Machine Studying AssemblyAI Weblog, 12 Could 2022

[13] Rogge, Niels and Rasul, Kashif. The Annotated Diffusion Mannequin . Hugging Face Weblog, 7 June 2022

[14] Das, Ayan. “An Introduction to Diffusion Probabilistic Fashions.” Ayan Das, 4 Dec. 2021

[15] Tune, Yang, and Stefano Ermon. Generative Modeling by Estimating Gradients of the Information Distribution. arXiv:1907.05600, arXiv, 10 Oct. 2020

[16] Tune, Yang, and Stefano Ermon. Improved Methods for Coaching Rating-Primarily based Generative Fashions. arXiv:2006.09011, arXiv, 23 Oct. 2020

[17] Tune, Yang, et al. Rating-Primarily based Generative Modeling by means of Stochastic Differential Equations. arXiv:2011.13456, arXiv, 10 Feb. 2021

[18] Tune, Yang. Generative Modeling by Estimating Gradients of the Information Distribution, 5 Could 2021

[19] Luo, Calvin. Understanding Diffusion Fashions: A Unified Perspective. 25 Aug. 2022

* Disclosure: Please notice that a few of the hyperlinks above is likely to be affiliate hyperlinks, and at no extra price to you, we’ll earn a fee should you resolve to make a purchase order after clicking by means of.