Paper: Masked Diffusion as Self-supervised Representation Learner (ICLR-2024)

Reviewer comment (Writing):
"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.

Author response:
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.

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

Author response:
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.

Reviewer comment (Evaluation):
"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.

Author response:
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.

Reviewer comment (Experiments):
"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.

Author response:
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.

Reviewer comment (Reproducibility):
Will trained model weights also be released on acceptance?

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

Reviewer comment (Experiments):
I am concerned by the apparent lack of a validation set. There seems to be a number of hyper parameters involved such as the choice of timesteps, blocks, patch sizes and training schedules covered in section B. Other choices such as loss function (SSIM or MSE) were also made. It is unclear if test sets were involved in these choices and if so can we expect the reported results to generalize?

Author response:
Thanks for the comment.

We did not use a validation set in our paper for two main reasons: (1) We follow the configuration of GlaS and MoNuSeg in [6] and the configuration of FFHQ-34 and CelebA-19 in [1] for a fair comparison. (2) We evaluate the label-efficient attribute of our method in a few-shot setting where the training data is limited and valuable.
    
However, the reported results can still generalize for the following reasons: (1) We use the same data split setting to train and evaluate other comparison methods and do early stopping only based on training loss. Therefore, it is still a fair comparison. (2) We adopt an unsupervised clustering algorithm to choose the blocks (see Figure 6), which does not use any label information in test sets. (3) We intentionally do not tune the hyper parameters for MoNuSeg and CelebA-19, and MDM still achieves state-of-the-art performance. (4) As shown in Table 6 (timesteps) and Table 7 (patch size), even with sub-optimal hyper parameter settings, MDM can still achieve state-of-the-art performance. This implies that MDM is robust to the hyper parameters. (5) For all experiments, we run 10 times with 10 random seeds to make sure that the results are reliable.

Reviewer comment (Experiments):
See the question in the above related to the use of validation sets.

Author response:
Thanks for the comment.

We did not use a validation set in our paper for two main reasons: (1) We follow the configuration of GlaS and MoNuSeg in [6] and the configuration of FFHQ-34 and CelebA-19 in [1] for a fair comparison. (2) We evaluate the label-efficient attribute of our method in a few-shot setting where the training data is limited and valuable.
    
However, the reported results can still generalize for the following reasons: (1) We use the same data split setting to train and evaluate other comparison methods and do early stopping only based on training loss. Therefore, it is still a fair comparison. (2) We adopt an unsupervised clustering algorithm to choose the blocks (see Figure 6), which does not use any label information in test sets. (3) We intentionally do not tune the hyper parameters for MoNuSeg and CelebA-19, and MDM still achieves state-of-the-art performance. (4) As shown in Table 6 (timesteps) and Table 7 (patch size), even with sub-optimal hyper parameter settings, MDM can still achieve state-of-the-art performance. This implies that MDM is robust to the hyper parameters. (5) For all experiments, we run 10 times with 10 random seeds to make sure that the results are reliable.

Reviewer comment (Experiments):
**Benchmarking for Segmentation Tasks**: The primary focus on segmentation necessitates benchmarking against specialized self-supervised learning (SSL) methods designed for this task, both at the instance-level and pixel/patch-level. A direct comparison with methods such as Leopart, IIC, MaskContrast, DenseCL, MoCoV2, and DINO on standard datasets like COCO and PVOC would provide a holistic evaluation. Refer to paper: Self-Supervised Learning of Object Parts for Semantic Segmentation for methods specific to semantic segmentation.

Author response:
**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.

Reviewer comment (Experiments):
**General SSL Representation Benchmarking**: If general SSL representation learning is the core objective, the method should be contrasted with prevailing SSL techniques (e.g., DINO, MoCoV2) across a spectrum of downstream tasks, including classification, object detection, and linear probing.

Author response:
**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:
| Learning Type      | Method   | Network    | Accuracy (%) |
|------------------- |----------|------------|--------------|
| Contrastive        | MoCov1   | ResNet-50  | 85.0         |
| Contrastive        | MoCov2   | ResNet-50  | 93.4         |
| Contrastive        | SimCLR   | ResNet-50  | 94.0         |
| Contrastive        | SimCLRv2 | ResNet-101 | 96.4         |
|                   |          |            |              |
| Generative         | DDPM     | U-Net      | 90.1         |
| Generative         | MDM      | U-Net      | 94.8         |

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.

Reviewer comment (Experiments):
**Demonstration of Versatility**: To set this work apart, the authors can also consider to showcase results across diverse data domains. Presenting outcomes on datasets related to audio, text, or time-series would emphasize the method's adaptability.

Author response:
**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.

Reviewer comment (Experiments):
**Robustness Against Noisy Data**: The paper could benefit from an evaluation of the model's robustness against noisy or corrupted data, providing a measure of its real-world applicability.

Author response:
**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. 

[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.

Reviewer comment (Presentation):
**Interpretability and Visualization**: Including a section on interpretability, with visualizations illustrating the diffusion process or the dynamic masking strategy, would help readers better grasp the underlying mechanisms.

Author response:
**Q3:**
Visualization of the dynamic masking strategy.

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

Reviewer comment (Presentation):
**Discussion on Limitations**: A more explicit section discussing the method's limitations, potential pitfalls, or scenarios where it might not be the best fit would offer a balanced perspective.

Author response:
**Q6:**
Discussion on Limitations.

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

Reviewer comment (Novelty):
My main concern is the novelty of the method. The paper mentions that with the fixed t, the method degrades to a vanilla masked autoencoder with SSIM loss. This basically means that the only contribution of the paper is masking the image with a dynamic masking ratio during training, which concerns me regarding the contribution of the paper.

Author response:
**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.

Reviewer comment (Experiments):
Table 4 shows that using SSIM improves the performance significantly both for DDPM and MDM while MDM is still being better. I think a more fair comparison of MDM w/ SSIM loss would have been with MAE w/ SSIM loss, since it would show the effect of changing t dynamically.

Author response:
**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:

| Method | Loss Type | GlaS 10% (8) | FFHQ-34 |
|--------|-----------|--------------|---------|
| DDPM | MSE | 82.32±0.77 | 58.75±0.16 |
| | | | |
| DDPM* | MSE | 78.77±1.01 | 51.63±0.16 |
| | **SSIM** | **81.79±1.13** | **56.97±0.18** |
| | | | |
| MAE | **MSE** | **79.47±1.08** | **57.06±0.20** |
| | SSIM | 76.72±1.75 | 56.55±0.13 |
| | | | |
| MDM | MSE | 82.70±0.79 | 55.06±0.21 |
| | **SSIM** | **84.51±1.15** | **59.18±0.11** |

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.

Reviewer comment (Experiments):
It is not very clear from the paper which decoder level representation is given as input to the segmentation network. Is it only the last layer or some combinations of the last few layers? I also think it would be interesting to know the performance change as a function of using different decoder level features. Additionally, for fair comparison, the same level decoder features should be used for also DDPM and MAE. I didn't see a clear statement about this in the paper.

Author response:
**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].

Reviewer comment (Writing):
The writing of this paper needs to be improved. Many claims in the introduction section are not very well-supported (e.g. "such efforts risk deviating from the theoretical underpinnings of diffusions") and are not very well organized.

Author response:
- 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.

Reviewer comment (Evaluation):
The choice of SSIM over MSE is rather empirical and not well justified.

Author response:
- 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.

Reviewer comment (Experiments):
Since the goal of the proposed method is to learn a better representation, how well the method performs on other downstream tasks?

Author response:
- 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:

| Learning Type      | Method   | Network    | Accuracy (%) |
|------------------- |----------|------------|--------------|
| Contrastive        | MoCov1   | ResNet-50  | 85.0         |
| Contrastive        | MoCov2   | ResNet-50  | 93.4         |
| Contrastive        | SimCLR   | ResNet-50  | 94.0         |
| Contrastive        | SimCLRv2 | ResNet-101 | 96.4         |
|                   |          |            |              |
| Generative         | DDPM     | U-Net      | 90.1         |
| Generative         | MDM      | U-Net      | 94.8         |


   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.

Reviewer comment (Presentation):
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.

Author response:
- 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).