Paper: Masked Distillation Advances Self-Supervised Transformer Architecture Search (ICLR-2024)

Reviewer comment (Experiments):
The authors mentioned that experiments were conducted on CIFAR-10 and CIFAR-100 for assessing the transferability of self-supervised architecture search methods. However, it should be noted that the paper does not provide the corresponding experimental results.

Author response:
We apologize for missing the results on CIFAR-10 and CIFAR-100 in the manuscript. We have provided the corresponding results below in Table 1 and added these results to the current version of our paper. For the purpose of comparison, we also present the results achieved by ViT-B and AutoFormer-S pre-trained on ImageNet-1k in Table 1. It can be observed that the proposed MaskTAS gets consistent improvement on the four datasets compared with ViT-B and AutoFormer-S. We believe this improvement benefits from the self-supervised architecture search framework. Compared with the supervised counterparts, the searched architectures of MaskTAS are not limited by the range and quality of data annotation. Thus, the proposed MaskTAS is capable of producing transformer architectures that can better transfer to other data domains.

Method | CIFAR-10 | CIFAR-100 | PETS | Flowers  
-------------------|------------------|------------------|------------------|------------------
ViT-B | 98.1\% | 87.1\% | 93.8\% | 89.5\%  
AutoFormer-S | 98.3\%  | 88.1\%  | 94.1\%  | 97.8\%
MaskTAS-small | **99.2\%** | **90.3\%** | **94.8\%** | **98.6\%**

**Table 1** Transferability comparison on different data domains.

Reviewer comment (Reproducibility):
I was not able to find the architecture of the final searched model in the paper. I would like to see the entire architecture of searched models, rather than a single number that indicating the final performance. This is particularly important in the NAS literature.

Author response:
**Q1: The architecture of the final searched model.**  

Thanks for the valuable suggestion. We provide the searched architectures on ImageNet-1k as follows. In order to present the network architecture more clearly, we will present them in the form of diagrams and add them to the supplementary material.  

MaskTAS-tiny: \{*layer num*: 12, *embed dim*: 192, *mlp ratio*: [4.0, 4.0, 4.0, 4.0, 4.0, 3.5, 3.5, 4.0, 4.0, 3.5, 4.0, 4.0], *head num*: [4, 4, 4, 4, 3, 3, 3, 3, 3, 4, 3, 4]\}  
MaskTAS-small: \{*layer num*: 13, *embed dim*: 384, *mlp ratio*: [4.0, 4.0, 3.0, 4.0, 3.0, 4.0, 3.0, 3.0, 4.0, 3.0, 3.0, 4.0, 3.0]; *head num*: [7, 7, 7, 7, 5, 7, 7, 5, 6, 5, 6, 6, 5]\}  
MaskTAS-base: \{*layer num*: 14, *embed dim*: 576, *mlp ratio*: [3.5, 4.0, 3.0, 4.0, 4.0, 3.5, 3.0, 3.5, 3.0, 3.0, 3.0, 3.0, 3.0, 3.0]; *head num*: [10, 9, 10, 10, 10, 10, 9, 10, 9, 9, 10, 9, 9, 10]\}

Reviewer comment (Experiments):
Typically, in a NAS paper, the final output is the architecture, which should be trained again from scratch to test the "true" performance. In this work, the weights that are distilled during training are preserved and used for fine-tuning in the downstream tasks. It would be better if the authors could provide an explanation on this.

Author response:
**Q2: The architectures should be trained again from scratch to test the true performance.**  

For one-shot NAS, pre-training a high-quality supernet is essential for candidate architectures to inherit weights directly [1,2]. Thus, in addition to training from scratch, inheriting the pre-trained weights for fine-tuning is also a popular way of architecture retraining. For better comparison, we provide the results of both fine-tuning and training from scratch in Table 1. The results on all three specifications of model size, *i.e.*, tiny, small and base, show that our MaskTAS consistently outperforms AutoFormer in both fine-tuning and training from scratch settings.

Method | Model Size | Finetune | Train from scratch
-------------------|------------------|------------------|------------------
AutoFormer-T | 5.7M | 74.9\% | 74.9\%
MaskTAS-tiny | 5.5M | 75.6\% | 75.3\%
AutoFormer-S | 22.9M | 81.8\% | 81.7\%
MaskTAS-small | 22.1M | 82.5\% | 82.0\%
AutoFormer-B | 53.7M | 82.6\% | 82.6\%
MaskTAS-base | 52.0M | 83.8\% | 82.6\%

**Table 1** Top-1 accuracy comparison achieved by fine-tuning and training from scratch.

Reviewer comment (Reproducibility):
Some of the experimental details are missing, for instance, where do all the numbers in Table 1 come from?

Author response:
**Q3: Some of the experimental details are missing, for instance, where do all the numbers in Table 1 come from?**  

We are sorry for the missing experimental details. In Table 1 of our paper, we compare the performance of the proposed MaskTAS with existing visual models. Specifically, the performance comparison with existing transformer models utilizes the results reported by AutoFormer [1] and ViTAS [3]. The results of CNN models in the table are mostly inherited from the corresponding paper. In addition, we will carefully review the experimental part and try to complete the experimental details.

Reviewer comment (Experiments):
It's not clear the impact of the self-supervised learning and teach-student distillation framework individually. Is it possible to add some ablation studies?

Author response:
**Q2: The impact of the self-supervised learning and distillation individually.** 

Thanks for the valuable suggestion. We have conducted ablation studies to explore the impact of self-supervised learning and teacher-student distillation individually. The results are reported in Table 3.
Method | w/ $SSL$&$KD$ | w/o $SSL$ | w/o $KD_{100}$ | w/o $KD_{1600}$  
-------------------|------------------|------------------|------------------|------------------
MaskTAS-tiny | 75.6\% | 75.0\% | 73.3\% | 75.5\%
MaskTAS-small | 82.5\% | 81.8\% | 80.6\% | 82.5\%
MaskTAS-base | 83.8\% | 82.7\% | 81.2\% | 84.0\%

**Table 3**  Ablation study on the impact of self-supervised learning and knowledge distillation.  

From Table 3, we can make the following observations. 1) The model equipped with both the two components, *i.e.*, ''w/ $SSL$&$KD$'', surpasses that with only distillation, *i.e.*, ''w/o $SSL$''. We think this is because of the capacity gap between teacher and student, thus the distillation process may cause information losses. 2) For better comparison, we adopt two configurations for the model without distillation, where ''w/o $KD_{100}$'' and ''w/o $KD_{1600}$'' refer to the model with 100 and 1600 pre-training epochs respectively. Although ''w/o $KD_{100}$'' suffers from a significant performance degradation, ''w/o $KD_{1600}$'' achieves comparable performance as ''w/ $SSL$&$KD$'' by increasing the number of pre-training epochs. 3) In conclusion, self-supervised learning is beneficial for performance improvement, while distillation helps speed up training. Thus, our MaskTAS employs both of the two components to achieve a tradeoff between performance and overhead.

Reviewer comment (Evaluation):
It's not clear on the overhead and tradeoff of training another teacher encoder-decoder network.

Author response:
**Q3: The overhead and tradeoff of training another teacher encoder-decoder network.**   

Due to the MIM-based self-supervised learning strategy, pre-training another teacher encoder-decoder network can inevitably bring a high computational overhead, *e.g.*, MAE [1] and its variants [2-3]. In our work, we directly employ the MIM pre-trained models released from the official MAE implementations as our teacher model to avoid the high overhead of teacher pre-training. Surprisingly, we find by building upon the teacher pre-trained in MAE, the whole distillation framework can efficiently yield high-performance student models with performance and overhead trade-offs.

Reviewer comment (Reproducibility):
Could you clarify how does MaskTAS "which completely avoids the expensive costs of data labeling inherited from supervised learning"? How does MaskTAS perform classification tasks? Will labels be need for fine-tuning the classification heads?

Author response:
**Q1-1: MaskTAS avoids the expensive costs of data labeling.**  

As an architecture search method for vision transformers, MaskTAS is mainly designed to avoid the data labeling cost for architecture search. By developing self-supervised architecture search framework, MaskTAS can discover optimal transformer architectures without using manual labels. Then, these searched architectures can be adapted to downstream tasks through a fine-tuning or re-training process, where different tasks may have different requirements for data labels. In terms of the standard classification task, labels are needed for fine-tuning the classification heads.

Reviewer comment (Experiments):
Is there any evidence to support MaskTAS can "generalize well to various data domains and tasks by searching specialized transformer architectures in self-supervised manner"?

Author response:
**Q1-2: MaskTAS generalizes well to various data domains and tasks.**  

We thank the reviewer for the valuable comment. We provide evidences to support the generalization ability of MaskTAS. 1) In addition to ImageNet-1k, we also perform architecture search on other data domains, including CIFAR-10, CIFAR-100, PETS and Flowers. The classification results are reported in Table 1. 2) We also evaluate the generalization ability of MaskTAS by transferring the searched architectures to semantic segmentation task. The experimental results of searched transformer models in small and base size are reported in Table 2. We have added these results and corresponding analyses to our paper.  
Method | CIFAR-10 | CIFAR-100 | PETS | Flowers  
-------------------|------------------|------------------|------------------|------------------
ViT-B | 98.1\% | 87.1\% | 93.8\% | 89.5\%  
AutoFormer-S | 98.3\%  | 88.1\%  | 94.1\%  | 97.8\%
MaskTAS-small | **99.2\%** | **90.3\%** | **94.8\%** | **98.6\%**

**Table 1** Transferability comparison on different data domains.

Method | Pre-train data | Small | Base  
-------------------|------------------|------------------|------------------
Naive supervised | IN1K w/ labels | - | 47.4\%
Twins-PCPVT | IN1K w/ labels | 46.2\% | 47.1\%
Twin-SVT | IN1K w/ labels | 45.9\% | 47.7\%
ViTAS-Twins | IN1K w/ labels | 47.9\% | 50.2\%
MaskTAS | IN1K w/o labels | **48.3\%** | **51.6\%**

**Table 2** Performance comparison (mIoU) on ADE20K semantic segmentation.