9
$\begingroup$

Diffusion Models (http://proceedings.mlr.press/v37/sohl-dickstein15.html and https://arxiv.org/abs/2006.11239) are generative models gaining popularity in the community of deep generative modeling, beating VAEs and GANs for its sample quality and training stability.

While we can easily use the wolfram language to build VAE (https://community.wolfram.com/groups/-/m/t/1379189) and GAN (https://reference.wolfram.com/language/ref/NetGANOperator.html) models, I wonder if there is any attempt for building simple diffusion generative models using the wolfram language? An example based on MNIST would be greatly appreciated.

$\endgroup$

1 Answer 1

8
$\begingroup$

This question was posted by me three months ago. Now I have figured out how to implement the denoising diffusion probabilistic model using the handy Mathematica deep learning toolbox. This is largely motivated by the recent success of the stable diffusion model (https://github.com/CompVis/stable-diffusion). The blog from Lilian Weng and the tutorial from Calvin Luo are of crucial help through my learning.

For simplicity, we consider an example of generating handwritten digit images, learning from the MNIST dataset. First, we corrupt the target images by gradually adding Gaussian noise on them, eventually turning the original data distribution into an isotropic Gaussian distribution of equal dimension before noisification. Thereafter, we learn a hierarchy of neural nets to reverse the noisification process. Finally, starting from an isotropic Gaussian, we sequentially sample using the learned hierarchy of neural nets, and obtain novel samples of the target distribution.

Below are the implementation details. First, let's load the data. Notice that, in unconditional generating setting, we only need the 28$\times$28 dimensional handwritten digit images. The class labels for these images are therefore discarded.

data=ResourceData[ResourceObject["MNIST"]][[;;,1]]; 
RandomSample[data,20]

enter image description here We define the hyperparameters of the denoising diffusion probabilistic model:

size={28,28}; (*dimension of the target data image*)
channel=1; (*channel of the target data image, 1 for grey scale, 3 for RGB*)

T=200; (*how many steps for corrupting the image*)
c=32; (*base channel size of the UNET applied here*)
Tc=16; (*encoding of time step for corrupting the image*)
depth=3; (*depth of the UNET*)
batch=256; (*mini-batch size for training*)
b1=10^-4; (*initial variance of noise*)
bT=0.02; (*final variance of noise*)
b[t_]:=b1+(t-1)/(T-1)*(bT-b1); (*linear schedule for increasing variance of noise*)

We define a forward noising process that sequentially produces latents $X_1$ through $X_T$ by imposing Gaussian noise of variance $b_t$: $$p(X_{t+1}|X_t) = \mathcal{N}\Big{(}X_{t+1}; \sqrt{(1-b_t)}X_{t},b_t\mathbb{1}\Big{)}$$ Here, $b_t$ is scheduled to increase linearly from $b_{1}$ to $b_T$: $b_t=b_1+\frac{t-1}{T-1}(b_T-b_1)$. Given the Gaussian nature of this corruption strategy, the distribution of the corrupted image at time step t conditioned on its true value, as can be easily proved via mathematical induction, is: $$p(X_t|X_0) = \mathcal{N}\Big{(}X_{t}; \sqrt{a_t}X_0,\,\big{(}1-a_t\big{)}\mathbb{1}\Big{)}$$ where $a_t\equiv\prod_{i=1}^{t}(1-b_i)$. Therefore, we can directly obtain corrupted image samples at any time step $t$:

a[t_]:=(Table[1-b[i],{i,t}])/.List->Times;
GXT[x0_,t_,noise_]:=Sqrt[a[t]]*x0+Sqrt[1-a[t]]*noise

Below are examples of the corrupted images across time steps. As $t\rightarrow\infty$, we can always transform any original data distribution into an isotropic Gaussian distribution.

Block[{steps=Join[{0,1},Range[10,200,10]],selects=RandomSample[data,6]},
TableForm[Table[Image[GXT[ImageData[slc],t,RandomReal[NormalDistribution[0,1],size]]],{slc,selects},{t,steps}],
TableSpacing->{.5,.5},TableAlignments->Center,TableHeadings->{None,Map[Style["t="<>ToString[#],{FontFamily->"Arial",10}]&,steps]}]]

enter image description here Next, we learn a hierarchy of neural nets to sequentially reverse the noisification. Specifically, we need to learn $p_{\theta_{t}}(X_{t-1}|X_{t})$ for $t\in[1,T]$, where $p_{\theta_t}$ is a probability distribution function parameterized by a neural net with parameter $\theta_t$. The loss function for training this hierarchy of neural nets is a reformulation of the negative log likelihood (for proof, see Equation 47-58 in https://arxiv.org/pdf/2208.11970.pdf):

$$-\log p(X_0)=\sum_{t=1}^{T}\mathbb{E}_{p(X_{t}|X_0)}\big{(}D_{\text{KL}}[p(X_{t-1}|X_0,X_t)||p_{\theta_t}(X_{t-1}|X_t)]\big{)}+D_{\text{KL}}[p(X_T|X_0)||p(X_T)]$$

The reason for applying this reformulation is that, both terms here can either be analytically derived with moderate assumptions (term 1), or neglected (term 2). We can therefore apply stochastic gradient descent on $p_{\theta_{1:t}}$ to minimize the overall negative log likelihood of the observed samples: $$\begin{aligned}[t]\theta^*_{1:T}=&\underset{\theta_{1:T}}{\arg\min}\sum_{X_0}\Big{(}\sum_{t=1}^{T}\mathbb{E}_{p(X_{t}|X_0)}\big{(}D_{\text{KL}}[p(X_{t-1}|X_0,X_t)||p_{\theta_t}(X_{t-1}|X_t)]\big{)}+D_{\text{KL}}[p(X_T|X_0)||p(X_T)]\Big{)}\\=&\underset{\theta_{1:T}}{\arg\min}\sum_{X_0}\sum_{t=1}^{T}\mathbb{E}_{p(X_{t}|X_0)}\big{(}D_{\text{KL}}[p(X_{t-1}|X_0,X_t)||p_{\theta_t}(X_{t-1}|X_t)]\big{)}\end{aligned}$$ For $t=1$, we have: $$\begin{aligned}\theta^*_{1}=&\underset{\theta_{1}}{\arg\min}\sum_{X_0}\mathbb{E}_{p(X_{1}|X_0)}\big{(}D_{\text{KL}}[p(X_{0}|X_0,X_1)||p_{\theta_1}(X_{0}|X_1)]\big{)}\\=&\underset{\theta_{1}}{\arg\max}\sum_{X_0}\mathbb{E}_{p(X_{1}|X_0)}\big{(}\log p_{\theta_1}(X_{0}|X_1)\big{)}\\=&\underset{\theta_{1}}{\arg\min}\sum_{X_0}\mathbb{E}_{p(X_{1}|X_0)}\big{(}\frac{1}{2{\sigma_{\theta_1}}^2}||\mu_{\theta_1}-X_0||^2\big{)}\end{aligned}$$

For $t\in[2,T]$, we have $p(X_{t-1}|X_0,X_t)=\mathcal{N}\big{(}X_{t-1}; \mu_t,\Sigma_t\big{)}$, where: $$\mu_t=\frac{\sqrt{1-b_t}(1-a_{t-1})X_t+\sqrt{a_{t-1}}b_tX_0}{1-a_t}$$ $$\Sigma_t=\frac{b_t(1-a_{t-1})}{1-a_t}\mathbb{1}$$ We assume 1) $p_{\theta_{t}}(X_{t-1}|X_t)=\mathcal{N}(X_{t-1};\mu_{\theta_t},\sigma_{\theta_t}^2\mathbb{1})$, and 2) $\sigma_{\theta_t}^2\mathbb{1}=\Sigma_t$, i.e., $\sigma_{\theta_t}^2=\frac{b_t(1-a_{t-1})}{1-a_t}$. Given these two assumptions, we have:

$$\begin{aligned}\theta^*_{2:T}=&\underset{\theta_{2:T}}{\arg\min}\sum_{X_0}\Big{(}\sum_{t=2}^{T}\mathbb{E}_{p(X_{t}|X_0)}\big{(}D_{\text{KL}}[p(X_{t-1}|X_0,X_t)||p_{\theta_t}(X_{t-1}|X_t)]\big{)}\\=&\underset{\theta_{2:T}}{\arg\min}\sum_{X_0}\sum_{t=2}^{T}\mathbb{E}_{p(X_{t}|X_0)}\big{(}\frac{1}{2{\sigma_{\theta_t}}^2}||\mu_{\theta_t}-\mu_t||^2\big{)}\end{aligned}$$

We have the following three options for optimizing $\theta^*_{2:T}$:

  1. Approximate $\mu_t$ using $\mu_{\theta_t}$.
  2. Given $\mu_t=\frac{\sqrt{1-b_t}(1-a_{t-1})X_t+\sqrt{a_{t-1}}b_tX_0}{1-a_t}$, assume $\mu_{\theta_t}=\frac{\sqrt{1-b_t}(1-a_{t-1})X_t+\sqrt{a_{t-1}}b_t{X_0}_{\theta_t}}{1-a_t}$, approximate $X_0$ using ${X_0}_{\theta_t}$.
  3. Given $\mu_t=\frac{\sqrt{1-b_t}(1-a_{t-1})X_t+\sqrt{a_{t-1}}b_tX_0}{1-a_t}$, $X_t=\sqrt{a_t}X_0+\sqrt{1-a_t}\epsilon_t$, where $\epsilon_t\sim\mathcal{N}(0,\mathbb{1})$, we have $\mu_t=\frac{1}{1-b_{t}}X_t-\frac{b_t}{\sqrt{1-a_t}\sqrt{1-bt}}\epsilon_t$, assume $\mu_{\theta_t}=\frac{1}{1-b_{t}}X_t-\frac{b_t}{\sqrt{1-a_t}\sqrt{1-b_t}}\epsilon_{\theta_t}$, approximate $\epsilon_t$ using $\epsilon_{\theta_t}$.

Ho et al. found the third option to be closely related to score based models, and is optimal in practice (https://arxiv.org/pdf/2006.11239.pdf). Further, we can even drop all the coefficients, combine the case of $t=1$ and $t\in[2,T]$, and optimize our model using the following loss function:

$$\begin{aligned}[t]\theta^*_{1:t}=\underset{\theta_{1:t}}{\arg\min}\sum_{X_0}\sum_{t=1}^{T}\mathbb{E}_{p(X_{t}|X_0)}\big{(}||\epsilon_t-\epsilon_{\theta_t}||^2\big{)}\end{aligned}$$

In practice, we often parameterize $\epsilon_{\theta_t}$ using a UNET, which takes into input of $X_t$ and time step $t$, in order to estimate the noise $\epsilon_t$ that we impose on $X_0$ to get $X_t$. The UNET for different noisification steps share parameters, with the impact of time step accounted by a position embedding vector of the time step (https://arxiv.org/abs/1706.03762):

PositionEmbedding[t_,d_:Tc]:=Flatten[Table[{Sin[1./(10000^(2*i/d))*t],Cos[1./(100000^(2*i/d))*t]},{i,d/2}]]

Below is the UNET. Its only difference with a conventional UNET is that we have the time step embedding vector attached to each of its residual block, in order to account for the impact of denoising across time steps:

res[c_]:=NetGraph[<|"long"->Flatten[Table[{ConvolutionLayer[c,{3,3},"PaddingSize"->1],NormalizationLayer[],Ramp},2]][[1;;-2]],
      "plus"->TotalLayer[],
      "short"->ConvolutionLayer[c,{1,1}]|>,
 {NetPort["Input"]->"long"->"plus",NetPort["Input"]->"short"->"plus"}]

upres[c_,size_]:=NetGraph[<|"long"->{NormalizationLayer[],Ramp,ResizeLayer[size],ConvolutionLayer[c,{3,3},"PaddingSize"->1],                                
 NormalizationLayer[],Ramp,ConvolutionLayer[c,{3,3},"PaddingSize"->1]},
      "plus"->TotalLayer[],
      "short"->{ResizeLayer[size],ConvolutionLayer[c,{1,1}]}|>,
 {NetPort["Input"]->"long"->"plus",NetPort["Input"]->"short"->"plus"}]

contract[channel_,crop_:{{1,1},{1,1}}]:=NetGraph[{"conv"->res[channel],"pooling"->PoolingLayer[2,2,"Function"->Mean],
                               "cropping"->PartLayer[{;;,crop[[1,1]];;-crop[[1,-1]],crop[[2,1]];;-crop[[2,-1]]}]},
                     {NetPort["Input"]->"conv"->"pooling"->NetPort["Pooling"],"conv"->"cropping"->NetPort["Shortcut"]}];

expand[channel_,size_]:=NetGraph[{"deconv"->upres[channel,size],
                        "join"->CatenateLayer[],
                        "conv"->res[channel/2]},
                     {NetPort["Input"]->"deconv"->"join",
                      NetPort["Shortcut"]->"join"->"conv"}];

UNet=NetGraph[<|Table[{"contract_"<>ToString[i]->contract[c*2^(i-1)],
                  "expand_"<>ToString[i]->expand[c*2^(i-1),size/2^(i-1)],
                  "T_F_"<>ToString[i]->{LinearLayer[c*2^(i-1)],ReplicateLayer[Floor[size/2^i],2]},
                  "T_B_"<>ToString[i]->{LinearLayer[c*2^(i-1)],ReplicateLayer[Floor[size/2^i],2]},
                  "thread_F_"<>ToString[i]->ThreadingLayer[Plus],
                  "thread_B_"<>ToString[i]->ThreadingLayer[Plus]},{i,depth}],
            "T_F_0"->{LinearLayer[c/2],ReplicateLayer[size,2]},"thread_F_0"->ThreadingLayer[Plus],
            "T_B_0"->{LinearLayer[c/2],ReplicateLayer[size,2]},"thread_B_0"->ThreadingLayer[Plus],
            "preprocess"->ConvolutionLayer[c/2,{1,1}],"postprocess"->ConvolutionLayer[channel,{1,1}],
            "ubase"->res[c*(depth+1)]|>,  
 Flatten[{NetPort["Input"]->"preprocess"->"thread_F_0"->"contract_1","preprocess"->"thread_B_0",
 NetPort["T"]->"T_F_0"->"thread_F_0",NetPort["T"]->"T_B_0"->"thread_B_0",
 Table[{NetPort["contract_"<>ToString[i],"Pooling"]->"thread_F_"<>ToString[i]->
 If[i<depth,"contract_"<>ToString[i+1],"ubase"->"thread_B_"<>ToString[depth]->NetPort["expand_"<>ToString[depth],"Input"]],
 NetPort["T"]->"T_F_"<>ToString[i]->"thread_F_"<>ToString[i],
 NetPort["T"]->"T_B_"<>ToString[i]->"thread_B_"<>ToString[i],
 NetPort["contract_"<>ToString[i],"Shortcut"]->NetPort["expand_"<>ToString[i],"Shortcut"],
 NetPort["expand_"<>ToString[i],"Output"]->"thread_B_"<>ToString[i-1]->If[i>1,NetPort["expand_"<>ToString[i-1],"Input"],"postprocess"]},{i,depth}]}],
 "Input"->Prepend[size,channel],
 "T"->Tc]

Information[UNet,"SummaryGraphic"]

enter image description here

Given the trained $\epsilon^*_{t=T:1}$, we could have $\mu^*_{t=T:1}$ and $\Sigma^*_{t=T:1}$. Therefore, starting from isotropic Gaussian samples, we can sequentially sample $\mathcal{N}(\mu^*_{t},\Sigma^*_{t})$, and achieve novel samples at time step 0. Below is the code for generating samples using the trained network:

Report[net_]:=Block[{result},
 result=Table[Block[{initial,final},
initial=RandomReal[NormalDistribution[0,1],Prepend[size,channel]];
 final=Table[Block[{at=a[t],at1=a[t-1]},
  Set[initial,initial/Sqrt[1-b[t]]-b[t]/Sqrt[(1-at)]/Sqrt[1-b[t]]*net[<|"Input"->initial,"T"->PositionEmbedding[t]|>]+
  RandomReal[NormalDistribution[0,Sqrt[at1/at*b[t]]],Dimensions[initial]]];initial],{t,T,1,-1}];
  Map[NetDecoder["Image"],Append[final[[1;;-1;;8]],final[[-1]]]]],{k,10}];
 TableForm[result,TableSpacing->{.5, .5}]];

During training, we randomly select the time step from [1,T], randomly select training samples, corrupt the selected samples to selected time steps using isotropic Gaussian noise, we estimate the noise using the corrupted samples and the time step embedding. We report the generated samples every few steps:

trained=NetTrain[UNet,
 {Function[Block[{t,X0,noise,Xt,tt},
 t=Table[RandomSample[Range[T]][[1]],batch];
 X0=NetEncoder[{"Image",size,"Grayscale"}][RandomSample[data,batch]];
 noise=RandomReal[NormalDistribution[0,1],Join[{batch,channel},size]];
 Xt=Table[GXT[X0[[i]],t[[i]],noise[[i]]],{i,batch}];
 tt=Map[PositionEmbedding,t];
 <|"Input"->Xt,
   "Output"->noise,
   "T"->tt|>]],"RoundLength" ->Length[data]},
 LossFunction ->MeanSquaredLossLayer[],
 BatchSize -> batch,
 MaxTrainingRounds->10000,
 TargetDevice->"GPU",
 Method->{"ADAM","LearningRate"->10^-5},
 TrainingProgressReporting -> {{Function@Report[#Net], "Interval" -> Quantity[10, "Batches"]},"Panel"}]

Here is a snapshot of the training. As we can see, starting from isotropic Gaussian, we can get pretty sharp samples.
enter image description here The emerging of pattern from noise reminds me of the following words: enter image description here

$\endgroup$
2
  • $\begingroup$ What is this good for? When you extract from pure random noise, say, character $0$ then it does not mean that the character $0$ has any relation with the given noise. Same as you may see a face on a noisy image of moon, but when an image of better quality is taken you may find out that the purported face is just an ordinary rock with a shadow. $\endgroup$ Sep 22 at 20:19
  • $\begingroup$ Brilliant! Didn't get an answer, figured it out himself, posted it for the good of all! $\endgroup$ Sep 23 at 3:12

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.