What is Sampling?

Sampling has already been explained in other sections in this guide. To recap:

#What: sampling means to pick data out according to a data distribution \(p_\text{data},\) which assigns probabilities to data. For example, a cat distribution \(p_\text{Cat}\) would assign high probability to cats, and everything else like dogs, humans, scenery, ...
#How: For diffusion models, solve the corresponding ODE/SDE which describe paths that take Gaussian noise and gradually transform it into realistic images.

What is a Sampler?

While not all of them were conceived this way, all diffusion samplers can be viewed as SDE/ODE solvers - or in other words, numerical integration methods.

[Deep Dive] Sampling Algorithm

Using the Euler method, one can sample from a diffusion model as follows:

Algorithm 1: Diffusion Sampling (Euler-Maruyama Method)

Require: diffusion model \(u_\theta,\) step sizes \([h_0,h_1,...,h_{n-1}]\), diffusion coefficient \(\eta_t\)

Set \(t = 0\)
Draw a sample \(x_0\sim p_{\text{init}}\)
for \(i=0,...,n-1\) do
Draw random noise \(\epsilon \sim \mathcal N(0,1)\)
\(x_{t+h}=x_t+h_iu_\theta(x_t,t)+\eta_t\sqrt{h_i}\epsilon\)
Update \(t\gets t+h_i\)
end for (\(t\) should be \(1\) after the loop)
return \(x_1\)

Where \(\mathcal N(0,1)\) is the Standard Gaussian distribution.

A few notes:

We recover flow matching / ODE sampling by setting \(\eta_t=0.\)
One can adapt the above algorithm to other samplers by changing line 5 to using said samplers' update rules rather than Euler's.

To adapt the above into k-diffusion, and by extension popular UIs like forge and comfy, a few terminologies need to shift. Mainly:

k-diffusion works with a data prediction model \(x_\theta\) that would predict the clean image directly, rather than \(u_\theta\) which predicts the change needed to get there.
k-diffusion directly works in a noise schedule \(\sigma_t\) rather than step sizes.

These aren't huge issues, as one can be translated into another without much trouble.

Algorithm 2 mostly follows the conventions in Algorithm 1, and thus you may find it easier to follow the changes; Algorithm 2.1 is a more direct inscription of the actual code.

Algorithm 2: k-diffusion Sampling (Euler-Maruyama Method)

Require: data prediction model \(x_\theta,\) noise schedule \([\sigma_0,\sigma_1,...,\sigma_n]\), diffusion coefficient \(\eta\)

Draw a sample \(x\sim\mathcal N(0, \sigma_0^2)\)
for \(i=0,...,n-1\) do
Draw random noise \(\epsilon\sim\mathcal N(0,1)\)
Set \(h_\text{down},h_\text{up}\gets\text{find\_step\_size}(\sigma_i,\sigma_{i+1},\eta)\)
Set \(u_\theta\gets (x-x_\theta(x,\sigma_i))/\sigma_i\)
Update \(x\gets x+h_\text{down}u_\theta+\eta h_\text{up}\epsilon\)
end for
return \(x\)

Algorithm 2.1: k-diffusion Sampling True-to-the-Code (Euler-Maruyama Method)

Require: data prediction model \(x_\theta,\) noise schedule \([\sigma_0,\sigma_1,...,\sigma_n]\), diffusion coefficient \(\eta\)

Draw a sample \(x\sim\mathcal N(0, \sigma_0^2)\)
for \(i=0,...,n-1\) do
Draw random noise \(\epsilon\sim\mathcal N(0,1)\)
Set \(\sigma_\text{down},\sigma_\text{up}\gets\text{get\_ancestral\_step}(\sigma_i,\sigma_{i+1},\eta)\)
Set \(d\gets (x-x_\theta(x,\sigma_i))/\sigma_i\)
Set \(dt\gets \sigma_\text{down}-\sigma_i\)
Update \(x\gets x+d\cdot dt+\sigma_\text{up}\epsilon\)
end for
return \(x\)