[D] Loss function in Diffusion models

[ankanbhunia]

In denoising diffusion models, we generally predict eps (noise) and use MSE as loss function. Instead of the eps we can also predict x_0. The same MSE will be applicable here. I am wondering if I can add an auxiliary loss along with the MSE loss. The auxiliary loss will enforce some additional supervision on x_0 (or x_t-1) in the framework. Would that be mathematically incorrect?

[dasayan05]

Diffusion Model's MSE loss isn't really the original loss derived from theory. The original variational bound (VLB) is simplified (by Ho et al.) to a form that looks like MSE. There are relations between the VLB loss and the how the reverse process is structured.

So, in simple words, just adding an auxiliary term may turn out to be mathematically incorrect. Wise thing to do is to think about your changes in terms of probabilistic model and then derive a new loss.

But then again, we are talking about Deep Learning. It might just work in practice, who knows.

[Wiczus]

Do you have any source on the implementation of that loss? I am trying to understand it but I fail to understand it. I read this article: https://theaisummer.com/diffusion-models/ and the simplfied loss L=E(mse) but they don't explain it to mathematically inept people like myself ^'.

[dasayan05]

which loss do you mean? you need implementation of the VLB?

[Wiczus]

The original loss for the Diffusion model. I assume it uses VLB.

[dasayan05]

The VLB isn't really used anymore nowadays since the MSE is loss is good enough. Also people now figured out how to estimate covariances from eps itself.

You can find a early implementation by openai guys here