How diffusion fashions work: the mathematics 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 might be primarily based on diffusion fashions are GLIDE, DALLE-2, Imagen, and the total open-source steady diffusion.

However what’s the major precept behind them?

On this weblog put up, we’ll dig our manner up from the fundamental rules. There are already a bunch of various diffusion-based architectures. We’ll concentrate on essentially the most outstanding one, which is the Denoising Diffusion Probabilistic Fashions (DDPM) as initialized by Sohl-Dickstein et al after which proposed by Ho. et al 2020. Varied different approaches will likely be mentioned to a smaller extent corresponding to steady diffusion and score-based fashions.

Diffusion fashions are essentially completely different from all of the earlier generative strategies. Intuitively, they intention to decompose the picture era 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 steadily produce an excellent 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, at the very least in comparison with GANs.

Diffusion course of

The essential concept behind diffusion fashions is slightly easy. They take the enter picture $mathbf{x}_0$

Afterward, a neural community is skilled to get better the unique knowledge by reversing the noising course of. By with the ability to mannequin the reverse course of, we will generate new knowledge. That is the so-called reverse diffusion course of or, on the whole, the sampling technique of a generative mannequin.

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

Ahead diffusion

Diffusion fashions could be seen as latent variable fashions. Latent implies that we’re referring to a hidden steady characteristic house. In such a manner, they could look much like variational autoencoders (VAEs).

In follow, they’re formulated utilizing a Markov chain of $T$ steps. Right here, a Markov chain implies that every step solely will depend on the earlier one, which is a light assumption. Importantly, we’re not constrained to utilizing a particular sort 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 situation $textbf{I}$ is the identification matrix, indicating that every dimension has the identical commonplace deviation $beta_t$

Thus, we will 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 supplies a magic treatment to this.

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

If we outline $alpha_t= 1- beta_t$

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

Thus to supply 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 will 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 method for all timesteps ($p_theta(mathbf{x}_{0:T})$

By moreover conditioning the mannequin on timestep $t$, it is going to be taught to foretell the Gaussian parameters (that means the imply $boldsymbol{mu}_theta(mathbf{x}_t, t)$

However how will we practice such a mannequin?

Coaching a diffusion mannequin

If we take a step again, we will discover that the mixture of $q$ and $p$ is similar to a variational autoencoder (VAE). Thus, we will practice it by optimizing the destructive log-likelihood of the coaching knowledge. After a sequence of calculations, which we can’t analyze right here, we will write the proof decrease sure (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 via the ELBO, maximizing the probability boils all the way down to studying the denoising steps $L_t$

Necessary notice: Although $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 will pattern $textbf{x}_t$

Observe that $alpha_t$

This little trick supplies us with a totally tractable ELBO. The above property has another essential aspect impact, as we already noticed within the reparameterization trick, we will 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 will use a neural community $epsilon_{theta}(mathbf{x}_t,t)$

Thus, the loss operate (the denoising time period within the ELBO) could 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 total goal:

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

Moreover, Ho et. al 2020 determine to maintain the variance mounted and have the community be taught solely the imply. This was later improved by Nichol et al. 2021, who determine to let the community be taught the covariance matrix $(boldsymbol{Sigma})$ as nicely (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 up to now is what the mannequin’s structure appears like. Discover that the mannequin’s enter and output must be of the identical dimension.

To this finish, Ho et al. employed a U-Web. If you’re unfamiliar with U-Nets, be at liberty to take a look at our previous article on the main U-Web architectures. In a number of phrases, a U-Web 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 characteristic dimension. Often, the enter picture is first downsampled after which upsampled till reaching its preliminary dimension.

Within the authentic implementation of DDPMs, the U-Web consists of Vast 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, try this superior put up by Hugging Face.

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

Conditional Picture Technology: Guided Diffusion

A vital facet of picture era is conditioning the sampling course of to control the generated samples. Right here, that is additionally known as guided diffusion.

There have even been strategies that incorporate picture embeddings into the diffusion to be able to “information” the era. 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 an excellent justification for the superb samples from a textual content immediate.

Typically, guided diffusion fashions intention to be taught $nabla log p_theta( mathbf{x}_t vert y)$

$p_theta(y)$

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

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 info.

Classifier steerage

Sohl-Dickstein et al. and later Dhariwal and Nichol confirmed that we will 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 will perturb the gradients with their dot product:

Because of this, they handle to “steer” the era 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 will outline a classifier-free guided diffusion mannequin as:

Steerage could be achieved with no 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)$

Observe 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 info that’s tough to foretell with a classifier (corresponding 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 robust image-text alignment. For more information on the strategy of Imagen try this video from AI Espresso Break with Letitia:

Scaling up diffusion fashions

You is perhaps 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 increased resolutions: cascade diffusion fashions and latent diffusion fashions.

Cascade diffusion fashions

Ho et al. 2021 launched cascade diffusion fashions in an effort to supply 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 increased decision particulars. To generate a picture, we pattern sequentially from every diffusion mannequin.

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

To accumulate good outcomes with cascaded architectures, robust 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 as a result of a train-test mismatch.

It was discovered that gaussian blurring is a important transformation towards reaching excessive constancy. They seek advice from this method as conditioning augmentation.

Steady diffusion: Latent diffusion fashions

Latent diffusion fashions are primarily based on a slightly easy concept: as an alternative of making use of the diffusion course of straight on a high-dimensional enter, we undertaking 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 info try this video:

Rating-based generative fashions

Across the identical time because the DDPM paper, Tune and Ermon proposed a distinct sort 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 technique of modeling the gradient of the log chance density operate, often known as the rating operate. Langevin dynamics is an iterative course of that may draw samples from a distribution utilizing solely its rating operate.

the place $delta$ is the step dimension.

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

Then through the use of Langevin dynamics, we will straight pattern from $p(x)$ utilizing the approximated rating operate.

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

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

The issue up to now: the estimated rating capabilities 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 truth, they used a number of scales of Gaussian noise perturbations.

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

Rating-based generative modeling with rating matching + Langevin dynamics. Supply: Generative Modeling by Estimating Gradients of the Knowledge 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 via 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 accordance with a diffusion course of. This course of is modeled by a prescribed stochastic differential equation (SDE) that doesn’t rely upon the information and has no trainable parameters. By reversing the method, we will generate new samples.

Rating-based generative modeling via 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.ok.a., Brownian movement), $mathbf{f}(cdot, t)$ is a vector-valued operate referred to as the drift coefficient of $mathbf{x}(t)$, and $g(cdot)$ is a scalar operate often known as the diffusion coefficient of $mathbf{x}(t)$. Observe that the SDE usually has a novel robust resolution.

To make sense of why we use an SDE, here’s a tip: the SDE is impressed by the Brownian movement, during which numerous 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 operate $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 optimistic weighting operate. As soon as we’ve the rating operate, we will plug it into the reverse SDE and resolve it to be able to pattern $mathbf{x}(0)$ from the unique knowledge distribution $p_0(mathbf{x})$

There are a selection of choices to resolve the reverse SDE which we can’t analyze right here. Be certain to test the unique paper or this glorious weblog put up by the writer.

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

Abstract

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

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

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

• The coaching of the mannequin is predicated on maximizing the proof decrease sure (ELBO).

• We will situation the diffusion fashions on picture labels or textual content embeddings to be able 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 steady 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 could 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 try these very nice blogs.

Cite as

@article{karagiannakos2022diffusionmodels,
    title   = "Diffusion fashions: towards state-of-the-art picture era",
    writer  = "Karagiannakos, Sergios, Adaloglou, Nikolaos",
    journal = "https://theaisummer.com/",
    yr    = "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 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.web

[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 Knowledge Distribution. arXiv:1907.05600, arXiv, 10 Oct. 2020

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

[17] Tune, Yang, et al. Rating-Based mostly Generative Modeling via Stochastic Differential Equations. arXiv:2011.13456, arXiv, 10 Feb. 2021

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

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

* Disclosure: Please notice that among the hyperlinks above is perhaps affiliate hyperlinks, and at no further price to you, we’ll earn a fee in the event you determine to make a purchase order after clicking via.