Masked Diffusion as Self-supervised Representation Learner

We would like to thank you for your positive evaluation of our work and providing valuable and insightful comments. Please find our responses to your comments in what follows:

Q1: "This paper decomposes the interrelation between the generative capability and representation learning ability inherent in diffusion models." - I am not really sure what is meant by this or why it is significant.

We apologize for the confusion.

Recent works [1, 2, 3] which use diffusion models as representation learners usually assume that the superior generative quality of diffusion models supports them as meaningful discriminative learning methods. However, in this paper, we analyze the diffusion models from the self-supervised learning perspective and extend denoising diffusion models into a self-supervised pre-training algorithm, using denoising as a preliminary task. Then, we do not need to stick to denoising, as we did in our paper, if our objective is representation learning rather than image synthesis. This conclusion is significant since it allows researchers not to focus too much on the generative performance of diffusion models if their goal is to acquire meaningful representations. Instead, future works can explore various corruptions and adapt them to the diffusion framework for various data and downstream tasks. In practice, many works [1, 4, 5] which use extracted diffusion representations for downstream segmentation tasks can be potentially improved.

Q2: "a scalable self-supervised representation learner..." - What is meant or referred to with scalable here? It does not appear to be mentioned elsewhere.

We are sorry not making this information clear enough. The scalibility of MDM has been discussed in Appendix B of the original paper. We find that the representations learned by MDM continue to improve with longer training (Figure 7), which suggests that MDM is a scalable approach.

Q3: While the prediction task demands a focus on high-level, low-frequency structural aspects of images," - I would have liked this to be better explored or more well supported to make statements such as this.

Thanks for pointing this out. We revised this statement in the manuscript to make it clearer. We made the statement in the context of the few-shot scenarios where a lightweight network is used to avoid over-fitting. Due to the limited model capacity and labels, the model is not able to learn a good dense prediction if it has to process low-level, high-frequency details of images (e.g., texture). Therefore, it is essential that MDM learns high-level, low-frequency structural information in pre-training and provides the learned meaningful representations during fine-tuning.

Q4: "the representation ability of diffusion models does not originate from their generative power." - I don't see where this is investigsted in the proposed manuscript. E.g. the generative power of the proposed model is not shown.

We did not show the generative power of MDM because MDM is not designed to be a generative model. As shown in Section 3.1, the generative power of a denoising diffusion model is based on two prerequisites: (1) The forward process should be Gaussian so that Equation 2 and the whole derivation hold. (2) $q(x_{T})$ should be close to $N(0, I)$ so that we can sample from a standard Gaussian distribution and iteratively apply Equation 2 to generate images. MDM apparently does not fulfill these prerequisites (prerequisite 1 does not hold since MDM does not use Gaussian noise; prerequisite 2 does not hold since $q(x_{T})$ in MDM is not $N(0, I)$) and therefore should not be expected to have a good generative capability. Despite that, as discussed in Q1, MDM shows robust representation learning ability. That's why we claim that the representation ability of diffusion models does not originate from their generative power.

Q5: Will trained model weights also be released on acceptance?

Yes, we will release our code with checkpoints on acceptance.

Thank you for your careful reading of our paper and your valuable comments, which have helped us to improve the presentation of our manuscript. Please find our answers to your remarkable comments in the following:

Q1: Benchmarking for Segmentation Tasks.

We totally agree that these experiments are important and meaningful. In our manuscript, we focus on the semantic segmentation under a label-efficient setting. We plan to conduct a direct comparison with methods such as Leopart on COCO and PVOC in our future work.

Q2: General SSL Representation Benchmarking.

Thanks so much for the constructive suggestion.

In the revised manuscript, we constrain our objective to representation learning specifically for segmentation, rather than general SSL representation learning.

Nevertheless, we evaluate our method on other downstream tasks and get promising results compared to other prevailing SSL techniques. In particular, we conduct the MDM pre-training on CIFAR-10 and then use linear probing to test how well MDM performs on classification task. The experimental details and results are shown in Appendix C. We provide the results below:

The table above shows that MDM achieves much better linear evaluation accuracy than DDPM which has the same learning type and network and competitive performance compared to the contrastive learning methods.

Q3: Visualization of the dynamic masking strategy.

In the revised manuscript, we have added the visualization of the dynamic masking strategy in Figure 1.

Q4: Demonstration of Versatility.

We totally agree that adapting the method to other data domains can be interesting. However, many extra modifications need to be made due to the domain gap. We think it will be an interesting and meaningful direction of further investigation.

Q5: Robustness Against Noisy Data to Measure The Method's Real-world Applicability.

Thanks for giving this constructive suggestion. We agree with the reviewer that a robustness analysis is helpful.

We added robustness analysis in Appendix B. We apply 15 types of corruption from [1] to the test data (Gaussian Noise, Shot Noise, Impulse Noise, Defocus Blur, Frosted Glass Blur, Frost, Brightness, Contrast, Elastic, Pixelate, JPEG, Speckle, Gaussian Blur, Spatter and Saturat). As shown in Figure 8, MDM demonstrates strong robustness against various types of corruption, even without employing any data augmentation related to the applied corruption (e.g., noise) during both pre-training and downstream segmentation training.

Q6: Discussion on Limitations.

We added the discussion of the method's limitations and future work in section 6.

[1] D. Hendrycks and T. G. Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.

We would like to thank you for reviewing our paper and providing valuable and helpful comments. Please find our responses to your comments in what follows:

Q1: The contribution of this paper.

The reviewer proposes that the contribution is limited because the technical implementation is simple. However, we respectfully disagree with this assessment as Reviewer Sbaz and Reviewer caGX agree unanimously that the proposed method is indeed novel. From the self-supervised learning perspective, even the famous denoising diffusion probability models (DDPM) can also be treated as denoising noisy image with different levels of noise during training, and the technical implementation is really simple as well. The simplicity of implementation does not diminish the effectiveness of MDM as a robust self-supervised representation learning method, much like DDPM's simplicity doesn't detract from its prowess as an impressive image synthesizer. To alleviate the reviewer's concern, we summarize our contributions below:

• In this paper, we analyze the diffusion models from the self-supervised learning perspective and extend denoising diffusion models into a self-supervised pre-training algorithm, using denoising as a preliminary task. We then empirically verify that the representation ability of diffusion models does not originate from their generative power by proposing our MDM which does not have generative power but still possesses powerful representation learning ability. This conclusion is significant since it allows researchers not to focus too much on the generative performance of diffusion models if their goal is to acquire meaningful representations. In practice, many works [1, 2, 3] which use extracted diffusion representations for downstream segmentation tasks can be potentially improved.
• Our MDM offers a different way to approach representation learning. We conduct extensive experiments to validate the effectiveness of our method. The state-of-the-art segmentation performance and competitive classification performance show the superiority of the proposed MDM.
• The label-efficient attribute of MDM has a wide range of applications. In particular, when developing medical AI, it is usually easy to collect large number of images but not feasible to make the corresponding labels for all the images. The proposed method can be a perfect solution for such scenarios.

Q2: The fair comparison with MAE should be added and the results on FFHQ-34 implies the improvement brought by MDM cannot generalize.

We appreciate this constructive suggestion. We agree that MDM w/ SSIM loss should be compared with MAE w/ SSIM loss for a more fair comparison. We updated our new results in Table 4:

We can draw two conclusions from this table:

• SSIM loss is not suitable for MAE. From our empirical results, SSIM loss can encourage the model to learn the structure information of the whole image because it takes into account structure while MSE loss assumes pixel-wise independence, which is beneficial for downstream dense prediction tasks. However, MAE calculates the loss per patch and only on invisible parts. Using SSIM loss to substitute MSE loss therefore does not make sense in MAE. Similarly, it is also not reasonable to use SSIM loss for DDPM when its reconstruction objective is the added noise.

• When adopting each corresponding suitable loss function, MDM achieves better performance than MAE and DDPM on both GlaS and FFHQ-34.

Q3: Questions regarding the layers of decoder representation.

We are sorry not making this information clear enough.

In fact, we discussed the block choosing strategy in Appendix B of our original paper. Also, we visualized the clustering of features extracted from different levels of the decoder for MDM and DDPM in Figure 6. The features from the deeper layers (blocks with smaller values) exhibit coarse semantic masks, while those from shallower layers (blocks with larger values) reveal finer details yet lack the same level of semantic coherence for coarse segmentation. Therefore, we choose the middle blocks B = {8, 9, 10, 11, 12} and B = {5, 6, 7, 8, 9, 10, 11, 12} among the 18 decoder blocks for two medical image datasets and two natural image datasets, respectively.

For a fair comparison, MDM and DDPM use the same blocks for feature extraction. Due to the fact that MAE uses a different architecture (ViT) from MDM and DDPM (U-Net), we extract feature maps from the deepest 12 ViT-L blocks for MAE, following [1].

We would like to thank you for reviewing our paper and providing valuable and helpful comments. Please find our responses to your comments in what follows:

Q1: The statistical results shown in Table 1 look promising if all the methods are compared fairly.

• Thanks for the positive comment. We tried our best to make sure all the experiments being fair:
  • We use the same settings for all methods unless it is infeasible.
  • We use each method's official implementation and checkpoints (if released).
  • We conduct all the experiments 10 times using 10 random seeds and report the standard deviation and mean.

Q2: The writing of this paper needs to be improved.

• We have revised the manuscript for most parts which may cause confusion and reorganized the manuscript. The new manuscript is much clearer and the claims are well-supported. Regarding the specific claim you mentioned, we say that ''such efforts risk deviating from the theoretical underpinnings of diffusion'' for the reason that: As shown in Section 3.1, the whole theoretical derivation of denoising diffusion model is based on the assumption that the tractable posterior used to generate images from a standard Gaussian distribution is Gaussian, which is fulfilled in denoising diffusion by adding Gaussian noise in the forward process. However, recent methods such as Cold Diffusion [1] have replaced the Gaussian noise with other corruptions which can not be described by Gaussian distributions, making them deviate from the theoretical underpinnings of diffusion and suffer from worse sample quality.

Q3: How well the method performs on other downstream tasks?

• To evaluate the performance of our method on other downstream tasks, we conduct the MDM pre-training on CIFAR-10 and then use linear probing to test how well MDM performs on a classification task. The experimental details and results are shown in Section C of the Appendix. We provide the results below:

    The table above shows that MDM achieves much better linear evaluation accuracy than DDPM which has the same learning type and network and competitive performance compared to the contrastive learning methods.

Q4: Justify the choice of SSIM over MSE. Is it a trade-off between learning a general representation than tuning toward specific downstream tasks?

• From the experimental results, we observe a phenomenon that even the MSE loss is quite low during pre-training (which indicates that the image restoration is good based on MSE loss), the segmentation performance of the downstream task using the learned representations is not always good. The reason could be that MSE loss assumes pixel-wise independence. As another commonly used loss function in image restoration task, SSIM loss can encourage the model to learn the structure information of the whole image because it takes into account structure, which is beneficial for downstream dense prediction task.

  Regarding the trade-off between learning a general representation than tuning toward specific downstream tasks, we believe that better loss functions than SSIM loss must exist for specific datasets and downstream tasks. We do not intend to claim SSIM loss is only choice.

  However, we do not think using SSIM loss in our MDM is a trade-off for the reason that our MDM with SSIM loss not only achieves state-of-the-art semantic segmentation performance on all four datasets we used, but also demonstrates competitive classification performance on CIFAR-10, as discussed above.

Q5: Figure 2 and Figure 3 are quite repetitive and it is very hard to conclude which methods are better. Reporting Dice scores along each figure will be helpful.

• Thanks for pointing this out. It is indeed hard to conclude which methods are better from Figure 2 and Figure 3 in the original manuscript. In the revised manuscript, we report the dice score of each prediction in the corner for Figure 3 and Figure 4 (Figure 2 and Figure 3 in previous version).

[1] A. Bansal, E. Borgnia, H.-M. Chu, J. S. Li, H. Kazemi, F. Huang, M. Goldblum, J. Geiping, and T. Goldstein. Cold diffusion: Inverting arbitrary image transforms without noise, 2022.

Dear Reviewers,

As the deadline of ICLR rebuttal period is approaching, we look forward to hearing your feedback on our responses. We would be happy to address any remaining concerns that you may still have.

Thanks,

Authors