cloneofsimo/d3pm

Minimal Implementation of a D3PM in pytorch

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)

Error in hybrid loss

The hybrid loss from eq. 5 is defined (and also accordingly implemented in the official implementation) as:

losses = vb_losses + self.hybrid_coeff * ce_losses

Here it is implemented as follows though:

d3pm/d3pm_runner.py

Line 294 in 3ceb637

return self.hybrid_loss_coeff * vb_loss + ce_loss, {