Vision Transformer Neural Architecture Search for Out-of-Distribution Generalization: Benchmark and Insights

Summary

The paper shows that in-domain accuracy and Training-free NAS accuracy predictions correlate poorly with out-of-domain accuracy, while characteristics of the model such as Flops, Params and Embed-Dim (number of channels) correlate much better with out-of-domain accuracy. To do this, they train a supernet that contains all possible architectures (subnets), then they validate these subnets on 9 datasets ImageNet/C/D/P/O/A/R/Sketch/Stylaized.

Strengths

They explore various ways for increasing OOD accuracy, which is most important for real-world problems. They show that conventional approaches, such as Training-free NAS or relying on in-domain accuracy, do not work. While they found that Flops, Param and Embed-dim correlated well with OOD accuracy. To do this, they introduce OoD-ViT-NAS benchmark for NAS research on ViT’s OOD generalization, that includes 3000 diverse ViT architectures which are evaluated on 9 datasets. This allows us to look for new and better approaches for finding new models structures to improve OOD accuracy.

Weaknesses

You show that the embedding dimension has the highest impact among ViT architectural attributes, while network depth has a slight impact on OoD generalization. But the paper lacks an explanation that although embed-dim (number of channels) correlates with OOD accuracy better than depth, this does not necessarily mean that to increase OOD accuracy only embed-dim should be increased, since the computational cost and memory costs may be different for embed-dim and depth. Therefore, to create either the largest and most accurate model that fits in memory, or the most optimal model in terms of speed and accuracy, it may be optimal to increase the depth rather than embed-dim. There are also papers [1,2,3] that show that it is necessary to simultaneously increase several model parameters (depth, number of channels, resolution) in optimal proportions. You present charts of OOD accuracy and number of parameters (Fig. 4, K19 - K24), but it would also be great to present charts of OOD-accuracy and latency, and OOD-accuracy and GPU memory consumption. Because these 3 model characteristics (OOD-accuracy, Latency, Memory consumption) are the most important for real-world problems. 1. EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks, Mingxing Tan, Quoc V. Le, 2019 2. EfficientNetV2: Smaller Models and Faster Training, Mingxing Tan, Quoc V. Le, 2021 3. EfficientDet: Scalable and Efficient Object Detection, Mingxing Tan, Ruoming Pang, Quoc V. Le, 2019

Questions

- Have you tried plotting OOD-accuracy versus Latency or GPU memory consumption while increasing various model parameters such as: embed_dim, depth, MLP-ratio, num of heads, etc.? Or find correlations between ODE accuracy and these parameters, taking into account their impact on latency and memory consumption. - Have you tried to find optimal model scaling factors (similar to works [1, 2, 3]) to increase OOD-accuracy in optimal way with respect to latency or memory consumption? 1. EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks, Mingxing Tan, Quoc V. Le, 2019 2. EfficientNetV2: Smaller Models and Faster Training, Mingxing Tan, Quoc V. Le, 2021 3. EfficientDet: Scalable and Efficient Object Detection, Mingxing Tan, Ruoming Pang, Quoc V. Le, 2019

Rating

8

Confidence

4

Soundness

3

Presentation

3

Contribution

4

Limitations

It might be worth adding to the limitations: The results obtained on Imagenet classification datasets may not correlate with the results of real-life problems. Because Imagenet classification requires to predict the class of only one usually large object in the low resolution image without predicting its location (box, polygon or mask), whereas in real problems it is usually necessary to predict the class and position of many objects, incl. very small in high resolution image. Moreover, there is ambiguous in the Imagenet classification task, the class of which of the many objects in the image should be predicted? Thus, a continuation of the research in this direction may be the finding of correlations, NAS approaches, new network structures, their parameters and scaling factors to achieve the highest OOD accuracy on tasks close to real ones (such as Dense prediction tasks: Segmentation, Detection, etc. ) in an optimal way with respect to latency and memory consumption.

Summary

This paper introduces OoD-ViT-NAS, a benchmark designed for evaluating ViT architectures' ability to generalize under Out-of-Distribution (OoD) shifts. This paper reveals that ViT design significantly impacts OoD generalization, In-Distribution (ID) accuracy does not reliably predict OoD performance, and simpler metrics like parameter count can sometimes predict OoD accuracy better than complex training-free NAS methods.

Strengths

This paper propose a new benchmark for ViT NAS.

Weaknesses

1. While the paper introduces a new benchmark, it doesn't sufficiently discuss the practical implications of the findings or potential applications. Including a section on how these insights can influence future ViT designs would add significant value. Furthermore, assessing whether ViTs designed based on these insights outperform those designed by humans would be meaningful to include in experiments and results. 2. The analysis of the experimental results is somewhat superficial. It presents observations without delving into deeper discussions. There should be more profound insights and detailed discussions on why certain ViT architectural attributes perform better under OoD conditions. This could involve providing theoretical justifications or proposing hypotheses that could be explored in future research. 3. The paper states that simple proxies like #Param and #Flop outperform more complex Training-free NAS methods but does not provide enough insight into why this might be the case. A deeper analysis of these findings, including possible reasons and implications, would strengthen the argument.

Questions

Please address the weaknesses mentioned above.

Rating

6

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

No potential negative societal impact.

Summary

This paper studies the out of distribution generalization of vision transformer architecture designs. Specifically the paper studies how in-distribution accuracies relate to OOD accuracies for a set pf 3000 architectures. In addition the paper also studies the correlation between different zero-cost proxies and OOD accuracies (on different variants of imagenet), observing that simply parameters or FLOPs are good proxies. Furthermore, the paper also studies the impact of different architectural decisions on the OOD accuracies, discovering that embedding dimension is of utmost importance for OOD generalization. The authors release code and the raw results to reproduce their experiments on AutoFormer-T,B and S space.

Strengths

- The paper is the first one to study out-of-ditribution generalization of architectures discovered by NAS - The paper also studies zero cost proxies on the ViT space showing ineffectiveness of modern/newer ZCPs - The paper is written in a clear and easy to understand manner - Experiments on ZCPs are thorough and the results are insigntful

Weaknesses

- Currently I am missing the potential use-cases of the benchmark. The paper is more of a study and not a method (and fits better in the datasets and benchmark track of NeurIPS). Also though the paper is posed as a benchmark,I miss the intended use-case of the benchmark. Do the authors intend on releasing surrogates based on proxies, to allow API based querying these metrics for any possible architecture? - The authors only study 3000 architectures, which are too few to lead to generalizable insights. - Analysis of architecture importance: The search space of AutoFormer is a bit biased in the sense that there is a larger (factor two) variabilities in embedding dimension and a less variability for number of layers. I wonder if this is the reason embedding dimension has a higher correlation with OOD accuracy compared to other architecture dimensions like number of heads, embedding dimensions? - Multi-Objective Search: Since inheriting weights of a supernet and evaluation is fast, generally one would perform multi-objective random or evolutionary search directly on the pretrained supernet. This would lead to different Pareto-Fronts for different dataset distributions. Could the authors also present the results of multi-objective search directly on the OOD datasets? - What is the cost/search time different in computing Pareto-Fronts based on MO-search v/s Zero-Cost-Proxies (a lot of these proxies are actually not zero cost)?

Questions

- Do the authors train their own supernet, or train a one-shot supernet from scratch? - How does this benchmark contribute to aiding development of future NAS techniques for OOD generalization? - Check weakness

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

The authors discuss limitations and broader impact sufficiently.

Summary

This work presented OoD-ViT-NAS, a comprehensive NAS benchmark with a focus on out-of-distribution generalisation of vision Transformer architectures. The authors created a benchmark with 3000 diverse ViT architectures evaluated on 8 common large-scale OoD datasets, and provided the first comprehensive investigation of how ViT architectures affect OoD generalisation. They also conducted the first study exploring NAS for ViT’s OoD, analysing the impact of individual ViT architectural attributes on OoD generalisation.

Strengths

This study introduces the first large-scale benchmark for ViT OoD generalisation, which could be a valuable resource to the NAS field. The authors conducted a thorough investigation on the impact of ViT architectures in relation to OoD generalisation, covering multiple aspects such as architectural attributes, correlation between ID and OoD performance and the effectiveness of training-free NAS methods. This study reported several findings, such as the low correlation between ID and OoD accuracy, the superiority of simple proxies, i.e., model params and FLOPs, over more computationally complex training-free metrics. Although the latter has been revealed in other work on computer vision tasks, the authors further validate that point on the OoD tasks. The reproducibility of this work is good as the authors provided a substantial amount of details of their methodology, hyper-parameters and computational resources.

Weaknesses

This study is limited to the Autoformer architecture space, which may limit the generalisability of the findings to other ViT architectures. Also it comprises the novelty of this work. Although the empirical studies are comprehensive, there are not many theoretical explanations for the observed phenomena. More analysis and insights should be provided to support this work. The presentation can be improved. For example, Fig 1 appears on Page 2, but is first mentioned at Line 226, Page 7.

Questions

Line 119 shows that the work is based on Autoformer. What is the reason behind this choice? What is new? Line 145, OoD Classification Accuracy, what is the difference to the ID Acc in Line 142? Maybe a formula could clearly explain that. Line 154, we make use of the One-Shot NAS approach. Why not use zero shot NAS instead? What are the benefits that one short NAS can bring but not zero shot NAS? Table 2 shows correlations of some zero proxies. Can SWAP be applied here as it demonstrates a much higher correlation than those on the table? SWAP-NAS: Sample-Wise Activation Patterns for Ultra-fast NAS. Peng et. al. ICLR'24.

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors discussed the limitations of this work, but in the appendix instead of the main text.

Thank you for your response. One remaining concern is regarding #Param and #Flop in A4.

Thank you for your responses and the latest results provided in A2. A2 and A3 have largely addressed my concerns. A1 has partially addressed my concerns. Although I think the author's response in A1 is not perfect, I acknowledge that this is an open issue worthy of further discussion, so what the authors have done so far is acceptable. One remaining concern is regarding #Param and #Flop in A4. From the results in Fig. 4 of the paper, it appears that the model's performance is generally monotonically related to #Param and #Flop. Does this mean that for this task, fine-grained architectural design might not be necessary, and that simply increasing the model size could improve performance? Besides, the conclusion that larger models perform better seems very intuitive, but aside from demonstrating that the architecture scales well, what further conclusions can be drawn from this? Additionally, this correlation trend is overall, but within certain #Param and #Flop ranges, there is some variance. So, when model size and computational cost are in certain ranges, does this #Param and #Flop-based approach become ineffective? Is it difficult to more accurately assess the performance differences between two models of similar size?

I leave the same rating: Strong Accept.

Thanks to the authors for additional comparison results that add confidence to the conclusions of the paper. It can be seen that increasing embed-dim increases OOD accuracy in an optimal way relative to increasing the inference latency. Two minor remarks to the rebuttal: 1. It is better in tables R.1 and R.2 to use GPU memory consumption instead of the number of parameters, which do not always correlate well with GPU memory consumption . 2. In the Figure R.2, it is better to use larger ranges for MLP-Ratio and #Heads, so that curves are visible instead of points in one place.

Response to rebuttal

I thank the authors for their response. I am increasing my score to 6.

Summary for AC and Reviewer discussion

Thank you AC for handling our paper. We thank the reviewers for all of their insightful comments and for taking the time to thoroughly review our paper. We would like to briefly summary our work and Author-Reviewer discussion. **Main contribution.** In this work, we conduct an in-depth investigation into the impact of ViT architecture on OoD generalization. To achieve this, we first construct a comprehensive OoD-ViT-NAS benchmark covering 8 large-scale OoD datasets and 3,000 ViT architectures. This benchmark allows us to thoroughly study the impact of ViT architecture on OoD generalization, revealing that ViT architectures have a considerable impact on OoD generalization. Then, we conduct the ViT architectural studies on OoD generalization through the lens of (Training-Free) NAS and ViT structural attributes. Importantly, we discover that, among ViT structural attributes, increasing embedding dimensions improves OoD generalization leading to a simple yet effective method that can produce high OoD performance ViTs architectures. **Reproducibility.** Our submission includes the program code, OoD-ViT-NAS benchmark, and detailed instructions for reproducibility purposes. All of these will be released to the public to facilitate future OoD research. **Rebuttal.** In total we receive four reviews for our work. We are particularly grateful for the reviewers' constructive feedback to further improve our work. These have led to: - Additional experiment to validate finding on embedding dimension's effectiveness for established human design ViT, e.g., ViT [1] and Swin [3] - Additional High-Frequency analysis to support our embedding dimension finding - Additional latency and model size analysis - Additional experiment on the comparison of Training-Free NAS and Evolutionary Search Importantly, additional results consistently support our findings in the main paper. **Final Reviewers’ Decision.** After we submitted our rebuttal, we received positive feedback from all reviewers. In the end, we received the above acceptance threshold across the board (i.e., 6, 6, 7, 8) including a “Strong Accept” rating (Reviewer f7Da) and a feedback on the moderate-to-high impact on the field (Reviewer Lm5n). This encouraging feedback motivates us to further explore this research direction. One again, we sincerely thank all reviewers and AC again for their valuable time on this work. Best Regards.