sail-sg / MDT

Masked Diffusion Transformer is the SOTA for image synthesis. (ICCV 2023)
Apache License 2.0
500 stars 35 forks source link

Questions about MAE Loss #35

Open ALEX13679173326 opened 5 months ago

ALEX13679173326 commented 5 months ago

Great works! However, I also wonder that the reconstruction loss (MSE Loss) in MAE is lost in the code. I couldn't find it. Does the MDT computes the noising prediction MSE loss twice in training process (one for complete images and the other for masked image)?

gasvn commented 5 months ago

https://github.com/sail-sg/MDT/blob/4b8bcbb656750cb8b763c309711234652b95b1c5/masked_diffusion/train_util.py#L227 Check this function.