cloneofsimo / d3pm

Minimal Implementation of a D3PM in pytorch

Minimal Implementation of a D3PM (Structured Denoising Diffusion Models in Discrete State-Spaces), in pytorch

Special thanks to fal.ai for the compute resources for this project.

This is minimal (400 LOC), but fully faithful implementation of a D3PM Structured Denoising Diffusion Models in Discrete State-Spaces. in pytorch.

I have tried to keep the code as simple as possible with much comments and explanation that is somewhat lacking on the original jax implementation, so that it is easy to understand. As far as I know, this is the first, faithful reimplementation of D3PM in pytorch. (Please correct me if I am wrong). Of course, this implementation was heavily based on the official implementation.

Difference between this implementation and the official implementation:

• This one has conditional sampling, so as you can see, generations are class-conditioned.
• This one uses rather different/simple model architecture.
• This one simplfies the official implementation very very much, so it is 400 LOC.
• This one does not use truncated logistic reparameterization, but you can use that if you wish.
• Only has uniform sample with inverse-linear beta scheudule, but you can change that with couple loc as well.

Usage

Following is completely self-contained example.

python d3pm_runner.py

Following uses dit.py, for CIFAR-10 dataset.

python d3pm_runner_cifar.py

Requirements

Install torch, torchvision, pillow, tqdm

pip install torch torchvision pillow tqdm

Citation

This implementation:

@misc{d3pm_pytorch,
  author={Simo Ryu},
  title={Minimal Implementation of a D3PM (Structured Denoising Diffusion Models in Discrete State-Spaces), in pytorch},
  year={2024},
  howpublished={\url{https://github.com/cloneofsimo/d3pm}}
}

Original Paper:

@article{austin2021structured,
  title={Structured denoising diffusion models in discrete state-spaces},
  author={Austin, Jacob and Johnson, Daniel D and Ho, Jonathan and Tarlow, Daniel and Van Den Berg, Rianne},
  journal={Advances in Neural Information Processing Systems},
  volume={34},
  pages={17981--17993},
  year={2021}
}

d3pm's Issues

Question about loss function in D3PM

Hello,

I hope this message finds you well. I have a question regarding the implementation of the loss function in D3PM. Specifically, in the forward method, there is a line of code located here:

predicted_x0_logits = self.model_predict(x, t, cond)

I was wondering why it isn't self.model_predict(x_t, t, cond) instead, predicting x_0 through x_t?

Thank you for your attention to this matter.

Absorbing State Diffusion

Could you please share the code for Absorbing State Forward Diffusion, where all states moved to [MASKED] state eventually?

A question about "a[t - 1, x, :]"

Thank you very much for posting the pytorch implementation of D3PM. I have questions about the following function:

*def _at(self, a, t, x):
# t is 1-d, x is integer value of 0 to num_classes - 1
bs = t.shape[0]
t = t.reshape((bs, [1] * (x.dim() - 1)))
# out[i, j, k, l, m] = a[t[i, j, k, l], x[i, j, k, l], m]
return a[t - 1, x, :]

This function seems to convert x0 to xt based on accumulated Qt. However, in the original paper, the conversion of x0 to xt is done by Qt multiplied by x0 (Eq 3). But this function does not seem to express this meaning. This is just selecting some values from Qt, and there no exist any calculation relationship between x0 and Qt (in the original D3PM code, the conversion from x0 to xt is also achieved in this way). At the same time, I also noticed that the xt obtained in this way was directly sent to the network. What is the way to convert from x0 to xt? Or is there something wrong with my understanding of this code? Looking forward to your reply.

Concerns Regarding `q_posterior_logits`

Hi Simo,

Thank you very much for your work. I have a question about one part: why is it t-2 here?

qmats2 = self.q_mats[t - 2].to(dtype=softmaxed.dtype)