# How to train Diffusion model with additional loss?

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?

• 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 Mar 30 at 16:28

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.
• 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. Jul 31, 2023 at 12:51