I would like to train a diffusion model with an additional loss on the created image. Without getting into too much details my intention is to do something like regularization, for example you may think that I want to make sure the created image is smooth, or something of the sort. My thinking was to add an additional loss during training, if the vanilla training process is: $$L = ||\epsilon-\epsilon_\theta(x_t, t)||$$ where $\epsilon_\theta$ is the model learning to predict the noise $\epsilon$ added to the original image. My suggested loss is: $$L = ||\epsilon-\epsilon_\theta(x_t, t)|| + \lambda * L'(img_\theta)$$ where $L'$ is my additional loss (which may for example induce smoothness or whatever), $img_\theta$ is the image that we get after denoising $x_t$ using the predicted noise. For standart models it is trivial that this makes sense. Due to the iterative nature of diffusion models I'm not sure if specifically for them it makes sense. I wasn't able to find any work that does something like this, would appreciate any help, does my additional loss makes sense to add?

  • $\begingroup$ Be careful not to mix different parameterizations: Here the first term is performing $\epsilon$-prediction (i.e you predict the noise part of the image), while your second term $img_\theta$, seems to be depending on a prediction of the denoised image. As @Peblo pointed out, training and sampling algorithms are typically different, and in your proposed loss you'd likely not be able to exploit the $\epsilon$-prediction reparameterization $\endgroup$
    – Ciodar
    Mar 30 at 16:28

1 Answer 1


I don't see why it should not work.

Though, the tough part is how to implement this. The algorithm requires to sample a singular time point along the timestep sequence, draw $\epsilon$, calculate $x_t$ for a batch, and evaluate the loss.

In your case you should also create the final image for the whole batch at each individual evaluation of the loss, which is extremely time consuming. But, other than that, I think it should work.

  • $\begingroup$ It's time consuming because I'm gonna have to do the whole denoising process? So, for example, it may include 50 denoising steps instead of the single step that you usually do during training, that is what you mean? Perhaps it would make sense to denoise using a single step only? then it would spare the full denoising process. Anyhow, I intend to work with low resolution (MNIST) images, so I'm not that worried about efficiency $\endgroup$ Jul 5, 2023 at 18:33
  • $\begingroup$ @GiladDeutsch yes this is what I meant with "time consuming". You may try to partially mitigate the problem by skipping a few denoising steps (look up the reply by Ciodar to this question), or by adding the penalization term at random only every N batches. Keep me posted, I am curious to know about the solution you achieve. $\endgroup$
    – Peblo
    Jul 6, 2023 at 11:46
  • $\begingroup$ are you sure that the fact that the diffusion model is recursive (i.e iteratively denoises the image) does not create an issue with this idea of learning using a loss on the created image? seems like this would create the same problems that RNNs have $\endgroup$ Jul 27, 2023 at 9:11
  • 1
    $\begingroup$ I am not sure which RNN issue you refer to, however -- regarding the recursiveness: keep in mind that the reproduction, i.e., the whole inverse process at prediction time, is indeed recursive, but the learning phase is not. During the learning phase the NN only needs to know the timestep $t$ under current evaluation, and upgrade the weights, it does not need to recur. Look at Algorithm 1 in Ho et al. 2020, step 5. $\endgroup$
    – Peblo
    Jul 31, 2023 at 12:51

You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .