Masked Distillation Advances Self-Supervised Transformer Architecture Search

We thank the reviewer for the valuable feedback and thoughtful comments. We appreciate the reviewer's acknowledgment that our work mitigates the high costs associated with supervised learning, formulates a noteworthy proposition of extracting knowledge from pre-trained networks, and the paper is written in a clear and accessible manner. In the following, we provide detailed explanations to address the concerns.

Q1: The experiments conducted on CIFAR-10 and CIFAR-100 for assessing the transferability of self-supervised architecture search methods.

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.

Table 1 Transferability comparison on different data domains.

Q2: What is the rationale behind the author's decision to employ the Transformer architecture?

From Table 1 of the paper, we can observe that MaskTAS is competitive compared to vanilla CNN models. Moreover, our MaskTAS also performs better than the manually-designed ResNet [1], ResNeXt [2] and DenseNet [3], demonstrating the potentials of pure transformer models for visual representation. However, transformer-based vision models now are still inferior to the lightweight CNN models based on inverted residual blocks, such as MobileNetV3 [4] and EfficientNet [5]. The reason is that inverted residuals are optimized for edge devices, so the model sizes and FLOPs are much smaller than vision transformers. Motivated by this, we plan to investigate the efficient search and deployment of ViT architectures on mobile devices in forthcoming research. This requires to optimize the design choices of ViT architecture to adapt to mobile devices, and develop new transformer architecture search method with low latency and high parameter efficiency. Your feedback has been valuable in shaping the direction of our future investigations.

References:

[1] He K, Zhang X, Ren S, et al. Deep residual learning for image recognition[C]//Proceedings of the IEEE conference on computer vision and pattern recognition. 2016: 770-778.
[2] Xie S, Girshick R, Dollár P, et al. Aggregated residual transformations for deep neural networks[C]//Proceedings of the IEEE conference on computer vision and pattern recognition. 2017: 1492-1500.
[3] Huang G, Liu Z, Van Der Maaten L, et al. Densely connected convolutional networks[C]//Proceedings of the IEEE conference on computer vision and pattern recognition. 2017: 4700-4708.
[4] Howard A, Sandler M, Chu G, et al. Searching for mobilenetv3[C]//Proceedings of the IEEE/CVF international conference on computer vision. 2019: 1314-1324.
[5] Tan M, Le Q. Efficientnet: Rethinking model scaling for convolutional neural networks[C]//International conference on machine learning. PMLR, 2019: 6105-6114.

We are grateful for your insightful comments and constructive suggestions. We provide a thorough response to each point raised to ensure clarity and a comprehensive understanding of our approach.

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]}

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.

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

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.

References:

[1] Chen M, Peng H, Fu J, et al. Autoformer: Searching transformers for visual recognition[C]//ICCV 2021: 12270-12280.
[2] Yu J, Jin P, Liu H, et al. Bignas: Scaling up neural architecture search with big single-stage models[C]//ECCV 2020: 702-717.
[3] Su X, You S, Xie J, et al. ViTAS: Vision transformer architecture search[C]//ECCV 2022: 139-157.

We thank the reviewer for the valuable feedback and thoughtful comments. We appreciate that the reviewer acknowledges the writing is good, the motivation is clear, and the method is sound. We will address your concerns in detail and we hope our response can further demonstrate the strengthens of our method.

Q1: It is unclear whether the performance gain is coming from knowledge distillation or training with the reconstruction objective.

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

Table 1 Ablation study on the impact of self-supervised learning and teacher-student distillation.

From Table 1, 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.

Q2: The experimental results using only masked image modeling on ImageNet & Cifar & ADE20K.

Thanks for the suggestion. The results using only masked image modeling on ImageNet have been provided in Table 1. We further report the results on CIFAR & ADE20K in Table 2, where ''w/o $KD_{100}$'' and ''w/o $KD_{1600}$'' pre-train the model using only masked image modeling with 100 and 1600 epochs respectively.

Table 2 Performance using only MIM on CIFAR and ADE20K.

Notably, by increasing the pre-training epochs from 100 to 1600, ''MaskTAS-small w/o KD'' can achieve comparable or even better performance than the complete ''MaskTAS-small'' model, as shown in Table 2. It indicates the great potential of masked image modelling in self-supervised architecture search.

After carefully reading the authors' responses and other reviewers' comments, I raised my rating to 6. Although I maintain some reservations about the experimental approach, particularly as the new experiments suggest that the performance gains are primarily attributable to knowledge distillation, I think this paper meets the bar of ICLR.

We greatly appreciate your careful reading and constructive comments. We have carefully considered each comment and made revisions accordingly.

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.

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.

Table 1 Transferability comparison on different data domains.

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

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.

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.

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.

References:

[1] He K, Chen X, Xie S, et al. Masked autoencoders are scalable vision learners[C]//CVPR 2022.
[2] Dong X, Bao J, Zhang T, et al. Bootstrapped masked autoencoders for vision BERT pretraining[C]//ECCV 2022.
[3] Feichtenhofer C, Li Y, He K. Masked autoencoders as spatiotemporal learners[J]//NeurIPS 2022.