Masked Diffusion Models are Fast Distribution Learners

1. Masked image pre-training is a popular stragedy to provide the ViT with image prior. Observed from Algorithm 1, there is few novelty compared to such kind of works.
2. Though authors present 2D Swiss roll as an intuitive example, theoretical support for the proposed method of approximating the primer distribution should be improved.
3. Experiments on deploying the proposed pre-trained method onto more existing diffusion methods are desired to prove the characteristics of plug-and-play.
4. Please clarify the fairness of comparisons with other diffusion-based methods, e.g., backbone.

1. The authors state that “our masked pre-training technique can be universally applied to various diffusion models that directly generate images in the pixel space, aiding in the learning of pre-trained models with superior generalizability.” The proposed Masked Diffusion Models are well suited to use VIT as the backbone. However, the current mainstream diffusion models use CNNs as the backbone. There is still insufficient evidence to determine whether Transformers contribute to the performance of diffusion-based image generation tasks. Therefore, this article lacks experimental or theoretical analysis on combining the proposed framework with previous CNN-based Diffusion models.

2. In Figure 4, the authors discussed the training efficiency of MaskDM compared to baseline DDPM. However, this baseline DDPM is implemented based on VIT. The reviewer believes it is also necessary to compare it with the CNN-based DDPM model to illustrate the training efficiency of the proposed model.

3. Some strategies to reduce the training overhead of diffusion models, such as Latent Diffusion Models and Cascaded Diffusion Models [1] are similar to the proposed method. The authors did not discuss the differences between these methods and the advantages of the proposed method in this article.

   • [1] Ho, Jonathan, et al. "Cascaded diffusion models for high fidelity image generation." The Journal of Machine Learning Research 23.1 (2022): 2249-2281.

4. Similarly, the authors did not report the training efficiency of other strategies in Fig. 4, which is not enough to demonstrate the efficiency of the proposed model.

5. In Table 2, the FID of DDIM is 3.26. This FID corresponds to the “DDPM ($\eta=1$ and $\hat{\sigma}$)” situation of the original article of DDIM. Therefore, labeling it as DDIM here will cause misunderstanding.

1. There are ambiguities in the manuscript. E.g. what is the point of learning a swiss roll distribution? This is mentioned as an example of learning the distribution, but it is not clear how it relates to the image generation task. Also, the paper's main claim in the title is on learning distributions, but there are no experiments supporting this claim. The paper only provides experiments on image generation quality based on FID and qualitative comparison. The distribution learning claim needs to be backed by performing more downstream tasks using the proposed model.

2. There has been a couple of works combining masked autoencoders and specifically ViTs with diffusion models. The authors only mention two of these, but do not contrast their methodology in detail against those, and do not compare their model's performance against them. Here are more papers proposing a similar idea. In general, this makes the novelty of the work limited. [a] Wei, Chen, et al. "Diffusion Models as Masked Autoencoders." arXiv (2023). [b] Pan, Zixuan, Jianxu Chen, and Yiyu Shi. "Masked Diffusion as Self-supervised Representation Learner." arXiv (2023).

3. The method is only evaluated using FID. There are much more metrics commonly used for image generation such as KID, IS, CLIP Score, Precision / Recall.

4. The performance gain is either marginal or worse e.g. compared to LDM.

Minor:

Fig. 5 on page viii does not have a caption.

The authors mention that they achieved a record value of 6.73 in FID. However, this value is never used in any table or text.

1. The paper only shows results on CelebA and CelebA-HQ which is not sufficient. More results on different other datasets need to be present to further demonstrate the effectiveness of the proposed method.
2. Results on CelebA 64x64 and 128x128 from Figure 4 (a) and (b) did not show that the proposed method has significant advantages over the baseline.
3. In Table 4, the results for the baseline which is U-ViT are missing.

Inconsistent/missing baselines: The paper presents varying baseline models in Tables 2-4, all of which are generative models that can be conceptually reused in these experiments. Alternating between baselines in different experiments can be perplexing and lead to potential misinterpretation.

Vanilla U-ViT: It is important to report the results of a vanilla U-ViT architecture without the masked pre-training, as this would help discern whether the performance improvement stems from the pre-training algorithm or the architectural modifications. What's even good is to include different transformer based diffusion model results. Note that I'm not demanding these experiments, I just state this as a weakness of the work.

Redundant/insufficient experiment: The evaluation on both CelebA and CelebA-HQ datasets may yield similar insights, given their substantial similarity. In addition to these, it is customary for generative models to be evaluated on larger-scale datasets such as ImageNet and LSUN. These experiments can be down at lower resolutions (64x64) to manage computational costs effectively.

Presentation: IMHO, the story of the primer distribution and the marginal distribution, while potentially insightful, appears to complicate understanding. Given the straightforward nature of the work and its simple implementation, streamlining the presentation and perhaps relocating the swiss roll toy example to a separate section could enhance readability.

Visual comparison: A qualitative comparison between this work and other existing methods is notably absent, making it challenging to directly appreciate the advantages of this approach.

Inconsistencies regarding VGGFace2: The paper frequently references VGGFace2, particularly in the abstract, where it claims a 46% quality improvement through fine-tuning on just 10% data from a different dataset. However, detailed results and discussions related to VGGFace2 are conspicuously absent in the main paper. To substantiate these claims, it is crucial to provide pertinent results and discussions within the main body of the paper.

Incomplete related work:

• MAE for other vision tasks or even broader topics (nlp, robotics) should be mentioned.
• Pre-training methods for generative models (GANs, VAEs) should be mentioned as well. One example is "Contrastive Learning for Unpaired Image-to-Image Translation, Park et al., ECCV'20".
• MDT and MaskDiT are mentioned in related work but not compared to in experiments. What's the reason?

1. In experiments it doesn't detail the time advantage of pre-training.
2. In experiments no detail on amount of fine-tuning and ablations on sensitivity to it
3. Overall the masking seems very hyper parameter sensitive as detailed in 4.2. This is further detailed in Appendix