title denoising diffusion implicit models andrs breshttps www linkedin cominandrasberes789190210 date created 20220624 last modified 20220624 description generating images of flowers with denoising diffusion implicit models accelerator gpu introduction what are diffusion models recently denoising diffusion modelshttps arxiv orgabs2006 11239 including scorebased generative modelshttps arxiv orgabs1907 05600 gained popularity as a powerful class of generative models that can rivalhttps arxiv orgabs2105 05233 even generative adversarial networks ganshttps arxiv orgabs1406 2661 in image synthesis quality they tend to generate more diverse samples while being stable to train and easy to scale recent large diffusion models such as dalle 2https openai comdalle2 and imagenhttps imagen research google have shown incredible texttoimage generation capability one of their drawbacks is however that they are slower to sample from because they require multiple forward passes for generating an image diffusion refers to the process of turning a structured signal an image into noise stepbystep by simulating diffusion we can generate noisy images from our training images and can train a neural network to try to denoise them using the trained network we can simulate the opposite of diffusion reverse diffusion which is the process of an image emerging from noise diffusion process gifhttps i imgur comdippofa gif onesentence summary diffusion models are trained to denoise noisy images and can generate images by iteratively denoising pure noise goal of this example this code example intends to be a minimal but featurecomplete with a generation quality metric implementation of diffusion models with modest compute requirements and reasonable performance my implementation choices and hyperparameter tuning were done with these goals in mind since currently the literature of diffusion models is mathematically quite complexhttps arxiv orgabs2206 00364 with multiple theoretical frameworks score matchinghttps arxiv orgabs1907 05600 differential equationshttps arxiv orgabs2011 13456 markov chainshttps arxiv orgabs2006 11239 and sometimes even conflicting notations see appendix c 2https arxiv orgabs2010 02502 it can be daunting trying to understand them my view of these models in this example will be that they learn to separate a noisy image into its image and gaussian noise components in this example i made effort to break down all long mathematical expressions into digestible pieces and gave all variables explanatory names i also included numerous links to relevant literature to help interested readers dive deeper into the topic in the hope that this code example will become a good starting point for practitioners learning about diffusion models in the following sections we will implement a continuous time version of denoising diffusion implicit models ddimshttps arxiv orgabs2010 02502 with deterministic sampling setup hyperparameters data kid kernel inception distance see related section sampling architecture optimization data pipeline we will use the oxford flowers 102https www tensorflow orgdatasetscatalogoxfordflowers102 dataset for generating images of flowers which is a diverse natural dataset containing around 8 000 images unfortunately the official splits are imbalanced as most of the images are contained in the test split we create new splits 80 train 20 validation using the tensorflow datasets slicing apihttps www tensorflow orgdatasetssplits we apply center crops as preprocessing and repeat the dataset multiple times reason given in the next section center crop image resize and clip for image downsampling it is important to turn on antialiasing the validation dataset is shuffled as well because data order matters for the kid estimation load dataset kernel inception distance kernel inception distance kidhttps arxiv orgabs1801 01401 is an image quality metric which was proposed as a replacement for the popular frechet inception distance fidhttps arxiv orgabs1706 08500 i prefer kid to fid because it is simpler to implement can be estimated perbatch and is computationally lighter more details herehttps keras ioexamplesgenerativeganadakernelinceptiondistance in this example the images are evaluated at the minimal possible resolution of the inception network 75x75 instead of 299x299 and the metric is only measured on the validation set for computational efficiency we also limit the number of sampling steps at evaluation to 5 for the same reason since the dataset is relatively small we go over the train and validation splits multiple times per epoch because the kid estimation is noisy and computeintensive so we want to evaluate only after many iterations but for many iterations kid is estimated per batch and is averaged across batches a pretrained inceptionv3 is used without its classification layer transform the pixel values to the 0255 range then use the same preprocessing as during pretraining compute polynomial kernels using the two sets of features estimate the squared maximum mean discrepancy using the average kernel values update the average kid estimate network architecture here we specify the architecture of the neural network that we will use for denoising we build a unethttps arxiv orgabs1505 04597 with identical input and output dimensions unet is a popular semantic segmentation architecture whose main idea is that it progressively downsamples and then upsamples its input image and adds skip connections between layers having the same resolution these help with gradient flow and avoid introducing a representation bottleneck unlike usual autoencodershttps www deeplearningbook orgcontentsautoencoders html based on this one can view diffusion models as denoising autoencodershttps benanne github io20220131diffusion html without a bottleneck the network takes two inputs the noisy images and the variances of their noise components the latter is required since denoising a signal requires different operations at different levels of noise we transform the noise variances using sinusoidal embeddings similarly to positional encodings used both in transformershttps arxiv orgabs1706 03762 and nerfhttps arxiv orgabs2003 08934 this helps the network to be highly sensitivehttps arxiv orgabs2006 10739 to the noise level which is crucial for good performance we implement sinusoidal embeddings using a lambda layerhttps keras ioapilayerscorelayerslambda some other considerations we build the network using the keras functional apihttps keras ioguidesfunctionalapi and use closureshttps twitter comfcholletstatus1441927912836321280 to build blocks of layers in a consistent style diffusion modelshttps arxiv orgabs2006 11239 embed the index of the timestep of the diffusion process instead of the noise variance while scorebased models table 1https arxiv orgabs2206 00364 usually use some function of the noise level i prefer the latter so that we can change the sampling schedule at inference time without retraining the network diffusion modelshttps arxiv orgabs2006 11239 input the embedding to each convolution block separately we only input it at the start of the network for simplicity which in my experience barely decreases performance because the skip and residual connections help the information propagate through the network properly in the literature it is common to use attention layershttps keras ioapilayersattentionlayersmultiheadattention at lower resolutions for better global coherence i omitted it for simplicity we disable the learnable center and scale parameters of the batch normalization layers since the following convolution layers make them redundant we initialize the last convolution s kernel to all zeros as a good practice making the network predict only zeros after initialization which is the mean of its targets this will improve behaviour at the start of training and make the mean squared error loss start at exactly 1 this showcases the power of the functional api note how we built a relatively complex unet with skip connections residual blocks multiple inputs and sinusoidal embeddings in 80 lines of code diffusion model diffusion schedule let us say that a diffusion process starts at time 0 and ends at time 1 this variable will be called diffusion time and can be either discrete common in diffusion models or continuous common in scorebased models i choose the latter so that the number of sampling steps can be changed at inference time we need to have a function that tells us at each point in the diffusion process the noise levels and signal levels of the noisy image corresponding to the actual diffusion time this will be called the diffusion schedule see diffusionschedule this schedule outputs two quantities the noiserate and the signalrate corresponding to sqrt1 alpha and sqrtalpha in the ddim paper respectively we generate the noisy image by weighting the random noise and the training image by their corresponding rates and adding them together since the standard normal random noises and the normalized images both have zero mean and unit variance the noise rate and signal rate can be interpreted as the standard deviation of their components in the noisy image while the squares of their rates can be interpreted as their variance or their power in the signal processing sense the rates will always be set so that their squared sum is 1 meaning that the noisy images will always have unit variance just like its unscaled components we will use a simplified continuous version of the cosine schedule section 3 2https arxiv orgabs2102 09672 that is quite commonly used in the literature this schedule is symmetric slow towards the start and end of the diffusion process and it also has a nice geometric interpretation using the trigonometric properties of the unit circlehttps en wikipedia orgwikiunitcirclemediafile circletrig6 svg diffusion schedule gifhttps i imgur comjw9w0fa gif training process the training procedure see trainstep and denoise of denoising diffusion models is the following we sample random diffusion times uniformly and mix the training images with random gaussian noises at rates corresponding to the diffusion times then we train the model to separate the noisy image to its two components usually the neural network is trained to predict the unscaled noise component from which the predicted image component can be calculated using the signal and noise rates pixelwise mean squared errorhttps keras ioapilossesregressionlossesmeansquarederrorfunction should be used theoretically however i recommend using mean absolute errorhttps keras ioapilossesregressionlossesmeanabsoluteerrorfunction instead similarly to thishttps github comlucidrainsdenoisingdiffusionpytorchblobmasterdenoisingdiffusionpytorchdenoisingdiffusionpytorch pyl371 implementation which produces better results on this dataset sampling reverse diffusion when sampling see reversediffusion at each step we take the previous estimate of the noisy image and separate it into image and noise using our network then we recombine these components using the signal and noise rate of the following step though a similar view is shown in equation 12 of ddimshttps arxiv orgabs2010 02502 i believe the above explanation of the sampling equation is not widely known this example only implements the deterministic sampling procedure from ddim which corresponds to eta 0 in the paper one can also use stochastic sampling in which case the model becomes a denoising diffusion probabilistic model ddpmhttps arxiv orgabs2006 11239 where a part of the predicted noise is replaced with the same or larger amount of random noise see equation 16 and belowhttps arxiv orgabs2010 02502 stochastic sampling can be used without retraining the network since both models are trained the same way and it can improve sample quality while on the other hand requiring more sampling steps usually convert the pixel values back to 01 range diffusion times angles angles signal and noise rates note that their squared sum is always sin2x cos2x 1 the exponential moving average weights are used at evaluation predict noise component and calculate the image component using it reverse diffusion sampling important line at the first sampling step the noisy image is pure noise but its signal rate is assumed to be nonzero minsignalrate separate the current noisy image to its components network used in eval mode remix the predicted components using the next signal and noise rates this new noisy image will be used in the next step noise images denormalized images normalize images to have standard deviation of 1 like the noises sample uniform random diffusion times mix the images with noises accordingly train the network to separate noisy images to their components track the exponential moving averages of weights kid is not measured during the training phase for computational efficiency normalize images to have standard deviation of 1 like the noises sample uniform random diffusion times mix the images with noises accordingly use the network to separate noisy images to their components measure kid between real and generated images this is computationally demanding kiddiffusionsteps has to be small plot random generated images for visual evaluation of generation quality training create and compile the model below tensorflow 2 9 pip install tensorflowaddons import tensorflowaddons as tfa optimizertfa optimizers adamw pixelwise mean absolute error is used as loss save the best model based on the validation kid metric calculate mean and variance of training dataset for normalization run training and plot generated images periodically inference load the best model and generate images results by running the training for at least 50 epochs takes 2 hours on a t4 gpu and 30 minutes on an a100 gpu one can get high quality image generations using this code example the evolution of a batch of images over a 80 epoch training color artifacts are due to gif compression flowers training gifhttps i imgur comfscktzq gif images generated using between 1 and 20 sampling steps from the same initial noise flowers sampling steps gifhttps i imgur comtm5lyh3 gif interpolation spherical between initial noise samples flowers interpolation gifhttps i imgur comhk5hd5o gif deterministic sampling process noisy images on top predicted images on bottom 40 steps flowers deterministic generation gifhttps i imgur comwcvzynh gif stochastic sampling process noisy images on top predicted images on bottom 80 steps flowers stochastic generation gifhttps i imgur comkrxogzd gif lessons learned during preparation for this code example i have run numerous experiments using this repositoryhttps github comberesandrascleardiffusionkeras in this section i list the lessons learned and my recommendations in my subjective order of importance algorithmic tips min and max signal rates i found the min signal rate to be an important hyperparameter setting it too low will make the generated images oversaturated while setting it too high will make them undersaturated i recommend tuning it carefully also setting it to 0 will lead to a division by zero error the max signal rate can be set to 1 but i found that setting it lower slightly improves generation quality loss function while large models tend to use mean squared error mse loss i recommend using mean absolute error mae on this dataset in my experience mse loss generates more diverse samples it also seems to hallucinate more section 3https arxiv orgabs2111 05826 while mae loss leads to smoother images i recommend trying both weight decay i did occasionally run into diverged trainings when scaling up the model and found that weight decay helps in avoiding instabilities at a low performance cost this is why i use adamwhttps www tensorflow orgapidocspythontfkerasoptimizersexperimentaladamw instead of adamhttps keras ioapioptimizersadam in this example exponential moving average of weights this helps to reduce the variance of the kid metric and helps in averaging out shortterm changes during training image augmentations though i did not use image augmentations in this example in my experience adding horizontal flips to the training increases generation performance while random crops do not since we use a supervised denoising loss overfitting can be an issue so image augmentations might be important on small datasets one should also be careful not to use leaky augmentationshttps keras ioexamplesgenerativeganadainvertibledataaugmentation which can be done following this method end of section 5https arxiv orgabs2206 00364 for instance data normalization in the literature the pixel values of images are usually converted to the 1 to 1 range for theoretical correctness i normalize the images to have zero mean and unit variance instead exactly like the random noises noise level input i chose to input the noise variance to the network as it is symmetrical under our sampling schedule one could also input the noise rate similar performance the signal rate lower performance or even the logsignaltonoise ratio appendix b 1https arxiv orgabs2107 00630 did not try as its range is highly dependent on the min and max signal rates and would require adjusting the min embedding frequency accordingly gradient clipping using global gradient clipping with a value of 1 can help with training stability for large models but decreased performance significantly in my experience residual connection downscaling for deeper models appendix bhttps arxiv orgabs2205 11487 scaling the residual connections with 1sqrt2 can be helpful but did not help in my case learning rate for me adam optimizer shttps keras ioapioptimizersadam default learning rate of 1e3 worked very well but lower learning rates are more common in the literature tables 1113https arxiv orgabs2105 05233 architectural tips sinusoidal embedding using sinusoidal embeddings on the noise level input of the network is crucial for good performance i recommend setting the min embedding frequency to the reciprocal of the range of this input and since we use the noise variance in this example it can be left always at 1 the max embedding frequency controls the smallest change in the noise variance that the network will be sensitive to and the embedding dimensions set the number of frequency components in the embedding in my experience the performance is not too sensitive to these values skip connections using skip connections in the network architecture is absolutely critical without them the model will fail to learn to denoise at a good performance residual connections in my experience residual connections also significantly improve performance but this might be due to the fact that we only input the noise level embeddings to the first layer of the network instead of to all of them normalization when scaling up the model i did occasionally encounter diverged trainings using normalization layers helped to mitigate this issue in the literature it is common to use groupnormalizationhttps www tensorflow orgaddonsapidocspythontfalayersgroupnormalization with 8 groups for example or layernormalizationhttps keras ioapilayersnormalizationlayerslayernormalization in the network i however chose to use batchnormalizationhttps keras ioapilayersnormalizationlayersbatchnormalization as it gave similar benefits in my experiments but was computationally lighter activations the choice of activation functions had a larger effect on generation quality than i expected in my experiments using nonmonotonic activation functions outperformed monotonic ones such as reluhttps www tensorflow orgapidocspythontfkerasactivationsrelu with swishhttps www tensorflow orgapidocspythontfkerasactivationsswish performing the best this is also what imagen uses page 41https arxiv orgabs2205 11487 attention as mentioned earlier it is common in the literature to use attention layershttps keras ioapilayersattentionlayersmultiheadattention at low resolutions for better global coherence i omitted them for simplicity upsampling bilinear and nearest neighbour upsamplinghttps keras ioapilayersreshapinglayersupsampling2d in the network performed similarly however i did not try transposed convolutionshttps keras ioapilayersconvolutionlayersconvolution2dtranspose for a similar list about gans check out this keras tutorialhttps keras ioexamplesgenerativeganadagantipsandtricks what to try next if you would like to dive in deeper to the topic i recommend checking out this repositoryhttps github comberesandrascleardiffusionkeras that i created in preparation for this code example which implements a wider range of features in a similar style such as stochastic sampling secondorder sampling based on the differential equation view of ddims equation 13https arxiv orgabs2010 02502 more diffusion schedules more network output types predicting image or velocity appendix dhttps arxiv orgabs2202 00512 instead of noise more datasets related works scorebased generative modelinghttps yangsong github ioblog2021score blogpost what are diffusion models https lilianweng github ioposts20210711diffusionmodels blogpost annotated diffusion modelhttps huggingface coblogannotateddiffusion blogpost cvpr 2022 tutorial on diffusion modelshttps cvpr2022tutorialdiffusionmodels github io slides available elucidating the design space of diffusionbased generative modelshttps arxiv orgabs2206 00364 attempts unifying diffusion methods under a common framework highlevel video overviews 1https www youtube comwatch vytamrhvg1ew 2https www youtube comwatch v344w5h24h8 detailed technical videos 1https www youtube comwatch vfblgfrltngu 2https www youtube comwatch vwo7aznzbzq scorebased generative models ncsnhttps arxiv orgabs1907 05600 ncsnhttps arxiv orgabs2006 09011 ncsnhttps arxiv orgabs2011 13456 denoising diffusion models ddpmhttps arxiv orgabs2006 11239 ddimhttps arxiv orgabs2010 02502 ddpmhttps arxiv orgabs2102 09672 ddpmhttps arxiv orgabs2105 05233 large diffusion models glidehttps arxiv orgabs2112 10741 dalle 2https arxiv orgabs2204 06125 imagenhttps arxiv orgabs2205 11487 introduction what are diffusion models recently denoising diffusion models https arxiv org abs 2006 11239 including score based generative models https arxiv org abs 1907 05600 gained popularity as a powerful class of generative models that can rival https arxiv org abs 2105 05233 even generative adversarial networks gans https arxiv org abs 1406 2661 in image synthesis quality they tend to generate more diverse samples while being stable to train and easy to scale recent large diffusion models such as dall e 2 https openai com dall e 2 and imagen https imagen research google have shown incredible text to image generation capability one of their drawbacks is however that they are slower to sample from because they require multiple forward passes for generating an image diffusion refers to the process of turning a structured signal an image into noise step by step by simulating diffusion we can generate noisy images from our training images and can train a neural network to try to denoise them using the trained network we can simulate the opposite of diffusion reverse diffusion which is the process of an image emerging from noise diffusion process gif https i imgur com dippofa gif one sentence summary diffusion models are trained to denoise noisy images and can generate images by iteratively denoising pure noise goal of this example this code example intends to be a minimal but feature complete with a generation quality metric implementation of diffusion models with modest compute requirements and reasonable performance my implementation choices and hyperparameter tuning were done with these goals in mind since currently the literature of diffusion models is mathematically quite complex https arxiv org abs 2206 00364 with multiple theoretical frameworks score matching https arxiv org abs 1907 05600 differential equations https arxiv org abs 2011 13456 markov chains https arxiv org abs 2006 11239 and sometimes even conflicting notations see appendix c 2 https arxiv org abs 2010 02502 it can be daunting trying to understand them my view of these models in this example will be that they learn to separate a noisy image into its image and gaussian noise components in this example i made effort to break down all long mathematical expressions into digestible pieces and gave all variables explanatory names i also included numerous links to relevant literature to help interested readers dive deeper into the topic in the hope that this code example will become a good starting point for practitioners learning about diffusion models in the following sections we will implement a continuous time version of denoising diffusion implicit models ddims https arxiv org abs 2010 02502 with deterministic sampling setup hyperparameters data train for at least 50 epochs for good results kid kernel inception distance see related section sampling architecture optimization data pipeline we will use the oxford flowers 102 https www tensorflow org datasets catalog oxford_flowers102 dataset for generating images of flowers which is a diverse natural dataset containing around 8 000 images unfortunately the official splits are imbalanced as most of the images are contained in the test split we create new splits 80 train 20 validation using the tensorflow datasets slicing api https www tensorflow org datasets splits we apply center crops as preprocessing and repeat the dataset multiple times reason given in the next section center crop image resize and clip for image downsampling it is important to turn on antialiasing the validation dataset is shuffled as well because data order matters for the kid estimation load dataset kernel inception distance kernel inception distance kid https arxiv org abs 1801 01401 is an image quality metric which was proposed as a replacement for the popular frechet inception distance fid https arxiv org abs 1706 08500 i prefer kid to fid because it is simpler to implement can be estimated per batch and is computationally lighter more details here https keras io examples generative gan_ada kernel inception distance in this example the images are evaluated at the minimal possible resolution of the inception network 75x75 instead of 299x299 and the metric is only measured on the validation set for computational efficiency we also limit the number of sampling steps at evaluation to 5 for the same reason since the dataset is relatively small we go over the train and validation splits multiple times per epoch because the kid estimation is noisy and compute intensive so we want to evaluate only after many iterations but for many iterations kid is estimated per batch and is averaged across batches a pretrained inceptionv3 is used without its classification layer transform the pixel values to the 0 255 range then use the same preprocessing as during pretraining compute polynomial kernels using the two sets of features estimate the squared maximum mean discrepancy using the average kernel values update the average kid estimate network architecture here we specify the architecture of the neural network that we will use for denoising we build a u net https arxiv org abs 1505 04597 with identical input and output dimensions u net is a popular semantic segmentation architecture whose main idea is that it progressively downsamples and then upsamples its input image and adds skip connections between layers having the same resolution these help with gradient flow and avoid introducing a representation bottleneck unlike usual autoencoders https www deeplearningbook org contents autoencoders html based on this one can view diffusion models as denoising autoencoders https benanne github io 2022 01 31 diffusion html without a bottleneck the network takes two inputs the noisy images and the variances of their noise components the latter is required since denoising a signal requires different operations at different levels of noise we transform the noise variances using sinusoidal embeddings similarly to positional encodings used both in transformers https arxiv org abs 1706 03762 and nerf https arxiv org abs 2003 08934 this helps the network to be highly sensitive https arxiv org abs 2006 10739 to the noise level which is crucial for good performance we implement sinusoidal embeddings using a lambda layer https keras io api layers core_layers lambda some other considerations we build the network using the keras functional api https keras io guides functional_api and use closures https twitter com fchollet status 1441927912836321280 to build blocks of layers in a consistent style diffusion models https arxiv org abs 2006 11239 embed the index of the timestep of the diffusion process instead of the noise variance while score based models table 1 https arxiv org abs 2206 00364 usually use some function of the noise level i prefer the latter so that we can change the sampling schedule at inference time without retraining the network diffusion models https arxiv org abs 2006 11239 input the embedding to each convolution block separately we only input it at the start of the network for simplicity which in my experience barely decreases performance because the skip and residual connections help the information propagate through the network properly in the literature it is common to use attention layers https keras io api layers attention_layers multi_head_attention at lower resolutions for better global coherence i omitted it for simplicity we disable the learnable center and scale parameters of the batch normalization layers since the following convolution layers make them redundant we initialize the last convolution s kernel to all zeros as a good practice making the network predict only zeros after initialization which is the mean of its targets this will improve behaviour at the start of training and make the mean squared error loss start at exactly 1 this showcases the power of the functional api note how we built a relatively complex u net with skip connections residual blocks multiple inputs and sinusoidal embeddings in 80 lines of code diffusion model diffusion schedule let us say that a diffusion process starts at time 0 and ends at time 1 this variable will be called diffusion time and can be either discrete common in diffusion models or continuous common in score based models i choose the latter so that the number of sampling steps can be changed at inference time we need to have a function that tells us at each point in the diffusion process the noise levels and signal levels of the noisy image corresponding to the actual diffusion time this will be called the diffusion schedule see diffusion_schedule this schedule outputs two quantities the noise_rate and the signal_rate corresponding to sqrt 1 alpha and sqrt alpha in the ddim paper respectively we generate the noisy image by weighting the random noise and the training image by their corresponding rates and adding them together since the standard normal random noises and the normalized images both have zero mean and unit variance the noise rate and signal rate can be interpreted as the standard deviation of their components in the noisy image while the squares of their rates can be interpreted as their variance or their power in the signal processing sense the rates will always be set so that their squared sum is 1 meaning that the noisy images will always have unit variance just like its unscaled components we will use a simplified continuous version of the cosine schedule section 3 2 https arxiv org abs 2102 09672 that is quite commonly used in the literature this schedule is symmetric slow towards the start and end of the diffusion process and it also has a nice geometric interpretation using the trigonometric properties of the unit circle https en wikipedia org wiki unit_circle media file circle trig6 svg diffusion schedule gif https i imgur com jw9w0fa gif training process the training procedure see train_step and denoise of denoising diffusion models is the following we sample random diffusion times uniformly and mix the training images with random gaussian noises at rates corresponding to the diffusion times then we train the model to separate the noisy image to its two components usually the neural network is trained to predict the unscaled noise component from which the predicted image component can be calculated using the signal and noise rates pixelwise mean squared error https keras io api losses regression_losses mean_squared_error function should be used theoretically however i recommend using mean absolute error https keras io api losses regression_losses mean_absolute_error function instead similarly to this https github com lucidrains denoising diffusion pytorch blob master denoising_diffusion_pytorch denoising_diffusion_pytorch py l371 implementation which produces better results on this dataset sampling reverse diffusion when sampling see reverse_diffusion at each step we take the previous estimate of the noisy image and separate it into image and noise using our network then we recombine these components using the signal and noise rate of the following step though a similar view is shown in equation 12 of ddims https arxiv org abs 2010 02502 i believe the above explanation of the sampling equation is not widely known this example only implements the deterministic sampling procedure from ddim which corresponds to eta 0 in the paper one can also use stochastic sampling in which case the model becomes a denoising diffusion probabilistic model ddpm https arxiv org abs 2006 11239 where a part of the predicted noise is replaced with the same or larger amount of random noise see equation 16 and below https arxiv org abs 2010 02502 stochastic sampling can be used without retraining the network since both models are trained the same way and it can improve sample quality while on the other hand requiring more sampling steps usually convert the pixel values back to 0 1 range diffusion times angles angles signal and noise rates note that their squared sum is always sin 2 x cos 2 x 1 the exponential moving average weights are used at evaluation predict noise component and calculate the image component using it reverse diffusion sampling important line at the first sampling step the noisy image is pure noise but its signal rate is assumed to be nonzero min_signal_rate separate the current noisy image to its components network used in eval mode remix the predicted components using the next signal and noise rates this new noisy image will be used in the next step noise images denormalized images normalize images to have standard deviation of 1 like the noises sample uniform random diffusion times mix the images with noises accordingly train the network to separate noisy images to their components used for training only used as metric track the exponential moving averages of weights kid is not measured during the training phase for computational efficiency normalize images to have standard deviation of 1 like the noises sample uniform random diffusion times mix the images with noises accordingly use the network to separate noisy images to their components measure kid between real and generated images this is computationally demanding kid_diffusion_steps has to be small plot random generated images for visual evaluation of generation quality training create and compile the model below tensorflow 2 9 pip install tensorflow_addons import tensorflow_addons as tfa optimizer tfa optimizers adamw pixelwise mean absolute error is used as loss save the best model based on the validation kid metric calculate mean and variance of training dataset for normalization run training and plot generated images periodically inference load the best model and generate images results by running the training for at least 50 epochs takes 2 hours on a t4 gpu and 30 minutes on an a100 gpu one can get high quality image generations using this code example the evolution of a batch of images over a 80 epoch training color artifacts are due to gif compression flowers training gif https i imgur com fscktzq gif images generated using between 1 and 20 sampling steps from the same initial noise flowers sampling steps gif https i imgur com tm5lyh3 gif interpolation spherical between initial noise samples flowers interpolation gif https i imgur com hk5hd5o gif deterministic sampling process noisy images on top predicted images on bottom 40 steps flowers deterministic generation gif https i imgur com wcvzynh gif stochastic sampling process noisy images on top predicted images on bottom 80 steps flowers stochastic generation gif https i imgur com krxogzd gif lessons learned during preparation for this code example i have run numerous experiments using this repository https github com beresandras clear diffusion keras in this section i list the lessons learned and my recommendations in my subjective order of importance algorithmic tips min and max signal rates i found the min signal rate to be an important hyperparameter setting it too low will make the generated images oversaturated while setting it too high will make them undersaturated i recommend tuning it carefully also setting it to 0 will lead to a division by zero error the max signal rate can be set to 1 but i found that setting it lower slightly improves generation quality loss function while large models tend to use mean squared error mse loss i recommend using mean absolute error mae on this dataset in my experience mse loss generates more diverse samples it also seems to hallucinate more section 3 https arxiv org abs 2111 05826 while mae loss leads to smoother images i recommend trying both weight decay i did occasionally run into diverged trainings when scaling up the model and found that weight decay helps in avoiding instabilities at a low performance cost this is why i use adamw https www tensorflow org api_docs python tf keras optimizers experimental adamw instead of adam https keras io api optimizers adam in this example exponential moving average of weights this helps to reduce the variance of the kid metric and helps in averaging out short term changes during training image augmentations though i did not use image augmentations in this example in my experience adding horizontal flips to the training increases generation performance while random crops do not since we use a supervised denoising loss overfitting can be an issue so image augmentations might be important on small datasets one should also be careful not to use leaky augmentations https keras io examples generative gan_ada invertible data augmentation which can be done following this method end of section 5 https arxiv org abs 2206 00364 for instance data normalization in the literature the pixel values of images are usually converted to the 1 to 1 range for theoretical correctness i normalize the images to have zero mean and unit variance instead exactly like the random noises noise level input i chose to input the noise variance to the network as it is symmetrical under our sampling schedule one could also input the noise rate similar performance the signal rate lower performance or even the log signal to noise ratio appendix b 1 https arxiv org abs 2107 00630 did not try as its range is highly dependent on the min and max signal rates and would require adjusting the min embedding frequency accordingly gradient clipping using global gradient clipping with a value of 1 can help with training stability for large models but decreased performance significantly in my experience residual connection downscaling for deeper models appendix b https arxiv org abs 2205 11487 scaling the residual connections with 1 sqrt 2 can be helpful but did not help in my case learning rate for me adam optimizer s https keras io api optimizers adam default learning rate of 1e 3 worked very well but lower learning rates are more common in the literature tables 11 13 https arxiv org abs 2105 05233 architectural tips sinusoidal embedding using sinusoidal embeddings on the noise level input of the network is crucial for good performance i recommend setting the min embedding frequency to the reciprocal of the range of this input and since we use the noise variance in this example it can be left always at 1 the max embedding frequency controls the smallest change in the noise variance that the network will be sensitive to and the embedding dimensions set the number of frequency components in the embedding in my experience the performance is not too sensitive to these values skip connections using skip connections in the network architecture is absolutely critical without them the model will fail to learn to denoise at a good performance residual connections in my experience residual connections also significantly improve performance but this might be due to the fact that we only input the noise level embeddings to the first layer of the network instead of to all of them normalization when scaling up the model i did occasionally encounter diverged trainings using normalization layers helped to mitigate this issue in the literature it is common to use groupnormalization https www tensorflow org addons api_docs python tfa layers groupnormalization with 8 groups for example or layernormalization https keras io api layers normalization_layers layer_normalization in the network i however chose to use batchnormalization https keras io api layers normalization_layers batch_normalization as it gave similar benefits in my experiments but was computationally lighter activations the choice of activation functions had a larger effect on generation quality than i expected in my experiments using non monotonic activation functions outperformed monotonic ones such as relu https www tensorflow org api_docs python tf keras activations relu with swish https www tensorflow org api_docs python tf keras activations swish performing the best this is also what imagen uses page 41 https arxiv org abs 2205 11487 attention as mentioned earlier it is common in the literature to use attention layers https keras io api layers attention_layers multi_head_attention at low resolutions for better global coherence i omitted them for simplicity upsampling bilinear and nearest neighbour upsampling https keras io api layers reshaping_layers up_sampling2d in the network performed similarly however i did not try transposed convolutions https keras io api layers convolution_layers convolution2d_transpose for a similar list about gans check out this keras tutorial https keras io examples generative gan_ada gan tips and tricks what to try next if you would like to dive in deeper to the topic i recommend checking out this repository https github com beresandras clear diffusion keras that i created in preparation for this code example which implements a wider range of features in a similar style such as stochastic sampling second order sampling based on the differential equation view of ddims equation 13 https arxiv org abs 2010 02502 more diffusion schedules more network output types predicting image or velocity appendix d https arxiv org abs 2202 00512 instead of noise more datasets related works score based generative modeling https yang song github io blog 2021 score blogpost what are diffusion models https lilianweng github io posts 2021 07 11 diffusion models blogpost annotated diffusion model https huggingface co blog annotated diffusion blogpost cvpr 2022 tutorial on diffusion models https cvpr2022 tutorial diffusion models github io slides available elucidating the design space of diffusion based generative models https arxiv org abs 2206 00364 attempts unifying diffusion methods under a common framework high level video overviews 1 https www youtube com watch v ytamrhvg1ew 2 https www youtube com watch v 344w5h24 h8 detailed technical videos 1 https www youtube com watch v fblgfrltngu 2 https www youtube com watch v w o7aznzbzq score based generative models ncsn https arxiv org abs 1907 05600 ncsn https arxiv org abs 2006 09011 ncsn https arxiv org abs 2011 13456 denoising diffusion models ddpm https arxiv org abs 2006 11239 ddim https arxiv org abs 2010 02502 ddpm https arxiv org abs 2102 09672 ddpm https arxiv org abs 2105 05233 large diffusion models glide https arxiv org abs 2112 10741 dall e 2 https arxiv org abs 2204 06125 imagen https arxiv org abs 2205 11487
import os
os.environ["KERAS_BACKEND"] = "tensorflow"

import math
import matplotlib.pyplot as plt
import tensorflow as tf
import tensorflow_datasets as tfds

import keras
from keras import layers
from keras import ops




dataset_name = "oxford_flowers102"
dataset_repetitions = 5
num_epochs = 1  
image_size = 64

kid_image_size = 75
kid_diffusion_steps = 5
plot_diffusion_steps = 20


min_signal_rate = 0.02
max_signal_rate = 0.95


embedding_dims = 32
embedding_max_frequency = 1000.0
widths = [32, 64, 96, 128]
block_depth = 2


batch_size = 64
ema = 0.999
learning_rate = 1e-3
weight_decay = 1e-4




def preprocess_image(data):
    
    height = ops.shape(data["image"])[0]
    width = ops.shape(data["image"])[1]
    crop_size = ops.minimum(height, width)
    image = tf.image.crop_to_bounding_box(
        data["image"],
        (height - crop_size) // 2,
        (width - crop_size) // 2,
        crop_size,
        crop_size,
    )

    
    
    image = tf.image.resize(
        image, size=[image_size, image_size], antialias=True
    )
    return ops.clip(image / 255.0, 0.0, 1.0)


def prepare_dataset(split):
    
    
    return (
        tfds.load(dataset_name, split=split, shuffle_files=True)
        .map(preprocess_image, num_parallel_calls=tf.data.AUTOTUNE)
        .cache()
        .repeat(dataset_repetitions)
        .shuffle(10 * batch_size)
        .batch(batch_size, drop_remainder=True)
        .prefetch(buffer_size=tf.data.AUTOTUNE)
    )



train_dataset = prepare_dataset("train[:80%]+validation[:80%]+test[:80%]")
val_dataset = prepare_dataset("train[80%:]+validation[80%:]+test[80%:]")




@keras.saving.register_keras_serializable()
class KID(keras.metrics.Metric):
    def __init__(self, name, **kwargs):
        super().__init__(name=name, **kwargs)

        
        self.kid_tracker = keras.metrics.Mean(name="kid_tracker")

        
        
        
        self.encoder = keras.Sequential(
            [
                keras.Input(shape=(image_size, image_size, 3)),
                layers.Rescaling(255.0),
                layers.Resizing(height=kid_image_size, width=kid_image_size),
                layers.Lambda(keras.applications.inception_v3.preprocess_input),
                keras.applications.InceptionV3(
                    include_top=False,
                    input_shape=(kid_image_size, kid_image_size, 3),
                    weights="imagenet",
                ),
                layers.GlobalAveragePooling2D(),
            ],
            name="inception_encoder",
        )

    def polynomial_kernel(self, features_1, features_2):
        feature_dimensions = ops.cast(ops.shape(features_1)[1], dtype="float32")
        return (
            features_1 @ ops.transpose(features_2) / feature_dimensions + 1.0
        ) ** 3.0

    def update_state(self, real_images, generated_images, sample_weight=None):
        real_features = self.encoder(real_images, training=False)
        generated_features = self.encoder(generated_images, training=False)

        
        kernel_real = self.polynomial_kernel(real_features, real_features)
        kernel_generated = self.polynomial_kernel(
            generated_features, generated_features
        )
        kernel_cross = self.polynomial_kernel(real_features, generated_features)

        
        batch_size = real_features.shape[0]
        batch_size_f = ops.cast(batch_size, dtype="float32")
        mean_kernel_real = ops.sum(
            kernel_real * (1.0 - ops.eye(batch_size))
        ) / (batch_size_f * (batch_size_f - 1.0))
        mean_kernel_generated = ops.sum(
            kernel_generated * (1.0 - ops.eye(batch_size))
        ) / (batch_size_f * (batch_size_f - 1.0))
        mean_kernel_cross = ops.mean(kernel_cross)
        kid = mean_kernel_real + mean_kernel_generated - 2.0 * mean_kernel_cross

        
        self.kid_tracker.update_state(kid)

    def result(self):
        return self.kid_tracker.result()

    def reset_state(self):
        self.kid_tracker.reset_state()





@keras.saving.register_keras_serializable()
def sinusoidal_embedding(x):
    embedding_min_frequency = 1.0
    frequencies = ops.exp(
        ops.linspace(
            ops.log(embedding_min_frequency),
            ops.log(embedding_max_frequency),
            embedding_dims // 2,
        )
    )
    angular_speeds = ops.cast(2.0 * math.pi * frequencies, "float32")
    embeddings = ops.concatenate(
        [ops.sin(angular_speeds * x), ops.cos(angular_speeds * x)], axis=3
    )
    return embeddings


def ResidualBlock(width):
    def apply(x):
        input_width = x.shape[3]
        if input_width == width:
            residual = x
        else:
            residual = layers.Conv2D(width, kernel_size=1)(x)
        x = layers.BatchNormalization(center=False, scale=False)(x)
        x = layers.Conv2D(
            width, kernel_size=3, padding="same", activation="swish"
        )(x)
        x = layers.Conv2D(width, kernel_size=3, padding="same")(x)
        x = layers.Add()([x, residual])
        return x

    return apply


def DownBlock(width, block_depth):
    def apply(x):
        x, skips = x
        for _ in range(block_depth):
            x = ResidualBlock(width)(x)
            skips.append(x)
        x = layers.AveragePooling2D(pool_size=2)(x)
        return x

    return apply


def UpBlock(width, block_depth):
    def apply(x):
        x, skips = x
        x = layers.UpSampling2D(size=2, interpolation="bilinear")(x)
        for _ in range(block_depth):
            x = layers.Concatenate()([x, skips.pop()])
            x = ResidualBlock(width)(x)
        return x

    return apply


def get_network(image_size, widths, block_depth):
    noisy_images = keras.Input(shape=(image_size, image_size, 3))
    noise_variances = keras.Input(shape=(1, 1, 1))

    e = layers.Lambda(sinusoidal_embedding, output_shape=(1, 1, 32))(
        noise_variances
    )
    e = layers.UpSampling2D(size=image_size, interpolation="nearest")(e)

    x = layers.Conv2D(widths[0], kernel_size=1)(noisy_images)
    x = layers.Concatenate()([x, e])

    skips = []
    for width in widths[:-1]:
        x = DownBlock(width, block_depth)([x, skips])

    for _ in range(block_depth):
        x = ResidualBlock(widths[-1])(x)

    for width in reversed(widths[:-1]):
        x = UpBlock(width, block_depth)([x, skips])

    x = layers.Conv2D(3, kernel_size=1, kernel_initializer="zeros")(x)

    return keras.Model([noisy_images, noise_variances], x, name="residual_unet")







@keras.saving.register_keras_serializable()
class DiffusionModel(keras.Model):
    def __init__(self, image_size, widths, block_depth):
        super().__init__()

        self.normalizer = layers.Normalization()
        self.network = get_network(image_size, widths, block_depth)
        self.ema_network = keras.models.clone_model(self.network)

    def compile(self, **kwargs):
        super().compile(**kwargs)

        self.noise_loss_tracker = keras.metrics.Mean(name="n_loss")
        self.image_loss_tracker = keras.metrics.Mean(name="i_loss")
        self.kid = KID(name="kid")

    @property
    def metrics(self):
        return [self.noise_loss_tracker, self.image_loss_tracker, self.kid]

    def denormalize(self, images):
        
        images = self.normalizer.mean + images * self.normalizer.variance**0.5
        return ops.clip(images, 0.0, 1.0)

    def diffusion_schedule(self, diffusion_times):
        
        start_angle = ops.cast(ops.arccos(max_signal_rate), "float32")
        end_angle = ops.cast(ops.arccos(min_signal_rate), "float32")

        diffusion_angles = start_angle + diffusion_times * (
            end_angle - start_angle
        )

        
        signal_rates = ops.cos(diffusion_angles)
        noise_rates = ops.sin(diffusion_angles)
        

        return noise_rates, signal_rates

    def denoise(self, noisy_images, noise_rates, signal_rates, training):
        
        if training:
            network = self.network
        else:
            network = self.ema_network

        
        pred_noises = network(
            [noisy_images, noise_rates**2], training=training
        )
        pred_images = (noisy_images - noise_rates * pred_noises) / signal_rates

        return pred_noises, pred_images

    def reverse_diffusion(self, initial_noise, diffusion_steps):
        
        num_images = initial_noise.shape[0]
        step_size = 1.0 / diffusion_steps

        
        
        
        next_noisy_images = initial_noise
        for step in range(diffusion_steps):
            noisy_images = next_noisy_images

            
            diffusion_times = ops.ones((num_images, 1, 1, 1)) - step * step_size
            noise_rates, signal_rates = self.diffusion_schedule(diffusion_times)
            pred_noises, pred_images = self.denoise(
                noisy_images, noise_rates, signal_rates, training=False
            )
            

            
            next_diffusion_times = diffusion_times - step_size
            next_noise_rates, next_signal_rates = self.diffusion_schedule(
                next_diffusion_times
            )
            next_noisy_images = (
                next_signal_rates * pred_images + next_noise_rates * pred_noises
            )
            

        return pred_images

    def generate(self, num_images, diffusion_steps):
        
        initial_noise = keras.random.normal(
            shape=(num_images, image_size, image_size, 3)
        )
        generated_images = self.reverse_diffusion(
            initial_noise, diffusion_steps
        )
        generated_images = self.denormalize(generated_images)
        return generated_images

    def train_step(self, images):
        
        images = self.normalizer(images, training=True)
        noises = keras.random.normal(
            shape=(batch_size, image_size, image_size, 3)
        )

        
        diffusion_times = keras.random.uniform(
            shape=(batch_size, 1, 1, 1), minval=0.0, maxval=1.0
        )
        noise_rates, signal_rates = self.diffusion_schedule(diffusion_times)
        
        noisy_images = signal_rates * images + noise_rates * noises

        with tf.GradientTape() as tape:
            
            pred_noises, pred_images = self.denoise(
                noisy_images, noise_rates, signal_rates, training=True
            )

            noise_loss = self.loss(noises, pred_noises)  
            image_loss = self.loss(images, pred_images)  

        gradients = tape.gradient(noise_loss, self.network.trainable_weights)
        self.optimizer.apply_gradients(
            zip(gradients, self.network.trainable_weights)
        )

        self.noise_loss_tracker.update_state(noise_loss)
        self.image_loss_tracker.update_state(image_loss)

        
        for weight, ema_weight in zip(
            self.network.weights, self.ema_network.weights
        ):
            ema_weight.assign(ema * ema_weight + (1 - ema) * weight)

        
        return {m.name: m.result() for m in self.metrics[:-1]}

    def test_step(self, images):
        
        images = self.normalizer(images, training=False)
        noises = keras.random.normal(
            shape=(batch_size, image_size, image_size, 3)
        )

        
        diffusion_times = keras.random.uniform(
            shape=(batch_size, 1, 1, 1), minval=0.0, maxval=1.0
        )
        noise_rates, signal_rates = self.diffusion_schedule(diffusion_times)
        
        noisy_images = signal_rates * images + noise_rates * noises

        
        pred_noises, pred_images = self.denoise(
            noisy_images, noise_rates, signal_rates, training=False
        )

        noise_loss = self.loss(noises, pred_noises)
        image_loss = self.loss(images, pred_images)

        self.image_loss_tracker.update_state(image_loss)
        self.noise_loss_tracker.update_state(noise_loss)

        
        
        images = self.denormalize(images)
        generated_images = self.generate(
            num_images=batch_size, diffusion_steps=kid_diffusion_steps
        )
        self.kid.update_state(images, generated_images)

        return {m.name: m.result() for m in self.metrics}

    def plot_images(self, epoch=None, logs=None, num_rows=3, num_cols=6):
        
        generated_images = self.generate(
            num_images=num_rows * num_cols,
            diffusion_steps=plot_diffusion_steps,
        )

        plt.figure(figsize=(num_cols * 2.0, num_rows * 2.0))
        for row in range(num_rows):
            for col in range(num_cols):
                index = row * num_cols + col
                plt.subplot(num_rows, num_cols, index + 1)
                plt.imshow(generated_images[index])
                plt.axis("off")
        plt.tight_layout()
        plt.show()
        plt.close()





model = DiffusionModel(image_size, widths, block_depth)




model.compile(
    optimizer=keras.optimizers.AdamW(
        learning_rate=learning_rate, weight_decay=weight_decay
    ),
    loss=keras.losses.mean_absolute_error,
)



checkpoint_path = "checkpoints/diffusion_model.weights.h5"
checkpoint_callback = keras.callbacks.ModelCheckpoint(
    filepath=checkpoint_path,
    save_weights_only=True,
    monitor="val_kid",
    mode="min",
    save_best_only=True,
)


model.normalizer.adapt(train_dataset)


model.fit(
    train_dataset,
    epochs=num_epochs,
    validation_data=val_dataset,
    callbacks=[
        keras.callbacks.LambdaCallback(on_epoch_end=model.plot_images),
        checkpoint_callback,
    ],
)




model.load_weights(checkpoint_path)
model.plot_images()