Implementation (in pytorch) of the D3PM algorithm (diffusion with categorical variable) : https://arxiv.org/pdf/2107.03006.pdf
In this repository we will implement the D3PM algorithm on discrete MNIST data.
in the forward discrete we have something like that :
Training is currently being done.
Result after training on 10% of the data points :
We train using nb_temporal_step = 254, lambda = 0.01 and hidden_dim = 32 (no need to have huge number of latent images).
Also we only have 3 bins (categorical variables) in the above exemple images.
Also as the base model we use ContextUnet architecture coming from the deeplearning.ai course on diffusion process : https://deeplearning.ai/short-courses/how-diffusion-models-work/
The modification is concern the passage here :
I want to point out that there is currently some (tiny) differences with the paper implementation :
In our implementation, we currently directly parametrize 
as being the neural network and not 
A second interesting point : in the current code we only use the Uniform Noise and not the absorbing noise setup which seems to give better result in the paper.
[!CAUTION] This was a repo for an old discrete diffusion method. Today i would use a simpler discrete diffusion set like the one used for diffusion llm like https://arxiv.org/pdf/2508.02193