Diffusion probabilistic model in deep generative modeling

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.

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]


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}],


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./(100000^(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"}]

"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]},
"preprocess"->ConvolutionLayer[c/2,{1,1}],"postprocess"->ConvolutionLayer[channel,{1,1}],
"ubase"->res[c*(depth+1)]|>,
NetPort["contract_"<>ToString[i],"Shortcut"]->NetPort["expand_"<>ToString[i],"Shortcut"],
"Input"->Prepend[size,channel],
"T"->Tc]

Information[UNet,"SummaryGraphic"]


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",
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.
The emerging of pattern from noise reminds me of the following words:

• 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. Commented Sep 22, 2022 at 20:19
• Brilliant! Didn't get an answer, figured it out himself, posted it for the good of all! Commented Sep 23, 2022 at 3:12

In the interest of presenting a simpler, more pedagogical introduction to diffusion model concepts, inspired by by François Flueret's post on a toy diffusion model, and corresponding python code, let us examine a simple 1D and 2D example (the code is in fact general). The text below is adapted from my blog.

Toy Problems and Sample Data

We'll begin by defining some sample distributions that we will try to learn; they also admit higher-dimensional analogues if desired. For generality, we'll make the input always a vector. This gets a bit awkward if it is just a vector of length one, but no matter...

sampleGaussianMixture[size_, p_ : 0.3, std_ : 0.2] := With[
{r = RandomVariate[NormalDistribution[0, std], size]},
{#} & /@ (r + Sign[RandomReal[{0, 1}, size] - p]/2)]

sampleRamp[size_List] := MapThread[Min, RandomReal[{0, 1}, Prepend[2]@size]]
sampleRamp[size_Integer] := sampleRamp[{size}]

sampleSpiral[size_] := With[
{u = RandomReal[{0, 1}, size]},
With[
{rho = u*0.65 + 0.25 + 0.15*RandomReal[{0, 1}, size],
theta = 3*Pi*u},
Transpose[{Cos[theta]*rho, Sin[theta]*rho}]]]

(*take a look*)
trainingSample1D = Standardize@sampleGaussianMixture[10^4];
trainingSample2D = Standardize@sampleSpiral[10^4];
GraphicsRow@{Histogram@Flatten@trainingSample1D, Histogram3D@trainingSample2D}


Now we will try to build up the diffusion model. We'll begin by defining functions that define some relevant constants ; in practice we'll just include the lists that get generated into the data.

Next, we will implement the analytic expression for the diffusion process at time t:

Clear[alpha, beta, alphaBar, sigma, xt];
beta[tMax_Integer] := Subdivide[10^-4, 0.2, tMax - 1]
alpha[tMax_Integer] := 1. - beta[tMax]
alphaBar[tMax_Integer] := Exp@Accumulate@Log[alpha[tMax]]
sigma[tMax_Integer] := Sqrt@beta[tMax]

xt[tMax_Integer] := With[
{a = Sqrt[alphaBar[tMax]],
b = Sqrt[1. - alphaBar[tMax]]},
FunctionLayer[
a[[#t]]*#x0 + b[[#t]]*#eps &,
"t" -> "Integer"]]


To try running the diffusion, we'll need a source of numbers drawn from the standard normal (Gaussian) distribution that we provide to the #eps input of the neural network. We'll also provide a time #t and a sample #x0. We will ultimately define that as part of the neural network below, but for now we'll just compute these numbers in a batch and perform a calculation with them. The function that gets returned by xt can be used as "just a function" or as part of a neural network, which gives us some freedom.

diffuse[xt_, sample_List, t_Integer] := With[
{eps = RandomVariate[NormalDistribution[], Dimensions[sample]]},
xt[<|"t" -> t, "x0" -> sample, "eps" -> eps|>]]

(*demo*)
diffuse[xt[200], trainingSample1D, 200]; // AbsoluteTiming

(*{0.288201, Null}*)


Let's use this to show how our starting samples get diffused into the Gaussian. Notice that our diffusion process, xt, is applied to each individual point--they don't know anything about each other, but by 200 steps of diffusion, the ensemble recovers a normal distribution. Neat!

With[
{process = xt[200]},
GraphicsRow@Table[
Histogram@Flatten@diffuse[process, trainingSample1D, t],
{t, {1, 5, 10, 50, 100, 200}}]]


Building a network to learn the inverse process

Having constructed the diffusion process, our next step is to develop a model that can undo the diffusion. It needs to take the time and the "diffused" x as inputs. We'll rescale the time so that it is between -1/2 and +1/2, and just run this through a few fully-connected layers. There's no rhyme or reason to this particular network architecture, and I'm sure it is totally overkill for the problem. After defining the model, we'll use a NetGraph to put it together with the other pieces. We'll need a set of parameters to define the model, so we'll save those in an association for convenience.

(*convenience function for defining parameters*)
parameters[nDim_, tMax_, nH_] := <|"nDim" -> nDim, "tMax" -> tMax, "nHidden" -> nH|>

(*allow our xt to accept the parameter association*)
xt[param_Association] := xt[param["tMax"]]

(*define a network for generating random normal numbers*)
eps[param_Association] := With[{nDim = param["nDim"]},
RandomArrayLayer[NormalDistribution, "Output" -> {nDim}]]

(*the "undiffusion" model that we will learn; appends a time *)
model[param_Association] := With[
{tMax = param["tMax"], nDim = param["nDim"], nHidden = param["nHidden"]},
NetGraph[
<|"tScale" -> FunctionLayer[(# - 1)/(tMax - 1) - 0.5 &, Input -> "Integer"],
"model" -> NetChain[
{AppendLayer["Input" -> nDim, "Element" -> "Real"],
LinearLayer[nHidden], Ramp,
LinearLayer[nHidden], Ramp,
LinearLayer[nHidden], Ramp,
LinearLayer[nDim]}]|>,
{NetPort["t"] -> "tScale",
{NetPort["xt"], "tScale"} -> "model"}]]

(*construct the network for learning the diffusion undoer*)
trainingNet[param_Association] := NetGraph[
<|"eps" -> eps[param],
"xt" -> xt[param],
"model" -> model[param],
"loss" -> MeanSquaredLossLayer[]|>,
{ {NetPort["t"], NetPort["x0"], "eps"} -> "xt",
{NetPort["t"], "xt"} -> "model",
{"model", "eps"} -> "loss"},
"t" -> "Integer", "x0" -> {param["nDim"]}]

(*pull a sample for SGD*)
sampler[param_Association, trainingSample_List] := Function[
<|"x0" -> RandomSample[trainingSample, #BatchSize],
"t" -> RandomInteger[param["tMax"], #BatchSize]|>];


Let's try it on our 1D example. First, let's get a sense of the model architecture:

trainingNet@parameters[1, 200, 128]


Now let's perform the calculation. Our training data consists of samples drawn from the training data and random times (using the sampler function defined above) to generate these data on demand. (Maybe there's a more efficient way to do this by building a network that generates t and feeds it in, so we can just pass the entire collection of datapoints in? For simplicity we'll ignore that.)

How long should you train for? I've found that the loss function is not a great indicator of training quality (more on this below), so I'm just going to set a boundary of 2 minutes and stop there. The hyperparameters defining the network are also chosen pretty much at random; of course you can do better:

param = parameters[1, 200, 128]
trainedNet1D = NetTrain[trainingNet[param], sampler[param, trainingSample1D],
TimeGoal -> 120] (*train for 2 minutes on CPU*)

(*<|"nDim" -> 1, "tMax" -> 200, "nHidden" -> 128|>*)


Generating samples

Now we've got to extract the trained model and use it to generate new samples. We'll start by defining the functions that let us do this. It is more natural to apply a Fold operation


z[nSamples_, t_Integer] := RandomVariate[NormalDistribution[], nSamples]

z[nSamples_, 1] := ConstantArray[0., nSamples]

timeStep[alpha_List, alphaBar_List, sigma_List, trainedModel_NetGraph][x_, t_Integer] := With[
<|"xt" -> x, "t" -> ConstantArray[t, Length[x]]|>,
zVec = z[Dimensions[x], t]},
1/Sqrt[alpha[[t]]]*(x - (1 - alpha[[t]])/Sqrt[1 - alphaBar[[t]]]*output) + sigma[[t]]*zVec
]

generate[trainedNet_NetGraph, param_Association, nSamples_Integer] :=
With[
{trainedModel = NetExtract[trainedNet, "model"],
a = alpha@param["tMax"],
ab = alphaBar@param["tMax"],
sig = sigma@param["tMax"]},
Fold[
timeStep[a, ab, sig, trainedModel],
RandomVariate[NormalDistribution[], {nSamples, param["nDim"]}], (*x*)
Reverse@Range[param["tMax"]] (*tValues*)]]


Now apply it to the model we trained above:

Histogram@Flatten@generate[trainedNet1D, param, 10^4]


This almost exactly reproduces the training data (which we've seen before)

As an aside: It's also fun to compare to with Mathematica's built in distribution learning methods--this defaults to performing a kernel density estimate:

With[
{auto = LearnDistribution[Flatten[trainingSample1D]]},
Histogram[RandomVariate[auto, 10^4], {0.1}, "PDF"]]


Our little diffusion model has done quite well for itself (of course, it is a bit of overkill for such a simple problem), but our goal is to illustrate the point.

Application to the 2D Problem

Let's apply it to our spiral 2D problem. We'll just use a humongous network with an arbitrarily chosen and almost certainly too large number of hidden layers, and see how well we do with 200 time steps:

param = parameters[2, 200, 256]
trainedNet = NetTrain[
trainingNet[param], sampler[param, trainingSample2D],
Histogram3D@generate[%, param, 10^4]

(*<|"nDim" -> 2, "tMax" -> 200, "nHidden" -> 256|>*)


Not bad! We're certainly capture some of the aspects of the spiral, it's not as clean as we might like.

Let's train the heck out of it (we'll just start from where we left off. This will take ~6 minutes on CPU, provide you don't overhead your poor laptop). Again, this is probably too much, but maybe it is useful. For each of our samples we have 200 times to consider, so if we perform

trainedNetv2 = NetTrain[
trainedNet, sampler[param, trainingSample2D],
Method -> "ADAM", MaxTrainingRounds -> 10^5]

Histogram3D@generate[%, param, 10^4]


Perhaps it is easier to compare these as a DensityHistogram:

DensityHistogram /@ {trainingSample2D,
generate[trainedNet, param, 10^4],
generate[trainedNetv2, param, 10^4]} // GraphicsRow


This suggests that you really do need to run long training rounds like 10^5. But now we see a good reproduction of the underlying probability distribution. Other parameters might include changing the number of time steps that are taken; using 1000 steps (instead of 200) might allow for finer resolution. This is left as an exercise for the reader...