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

TL;DR

This paper introduces OoD-ViT-NAS, a comprehensive benchmark for evaluating Vision Transformer architectures on out-of-distribution generalization, revealing key factors influencing robustness and the limited effectiveness of existing NAS methods.

Contribution

It presents the first systematic benchmark for ViT NAS focused on OoD generalization, analyzes factors affecting robustness, and evaluates the performance of training-free NAS methods.

Findings

ViT architecture design significantly impacts OoD generalization.

ID accuracy is a poor predictor of OoD accuracy.

Simple proxies like parameter count outperform complex NAS methods in predicting OoD robustness.

Abstract

While ViTs have achieved across machine learning tasks, deploying them in real-world scenarios faces a critical challenge: generalizing under OoD shifts. A crucial research gap exists in understanding how to design ViT architectures, both manually and automatically, for better OoD generalization. To this end, we introduce OoD-ViT-NAS, the first systematic benchmark for ViTs NAS focused on OoD generalization. This benchmark includes 3000 ViT architectures of varying computational budgets evaluated on 8 common OoD datasets. Using this benchmark, we analyze factors contributing to OoD generalization. Our findings reveal key insights. First, ViT architecture designs significantly affect OoD generalization. Second, ID accuracy is often a poor indicator of OoD accuracy, highlighting the risk of optimizing ViT architectures solely for ID performance. Third, we perform the first study of NAS for ViTs OoD robustness, analyzing 9 Training-free NAS methods. We find that existing Training-free NAS methods are largely ineffective in predicting OoD accuracy despite excelling at ID accuracy. Simple proxies like Param or Flop surprisingly outperform complex Training-free NAS methods in predicting OoD accuracy. Finally, we study how ViT architectural attributes impact OoD generalization and discover that increasing embedding dimensions generally enhances performance. Our benchmark shows that ViT architectures exhibit a wide range of OoD accuracy, with up to 11.85% improvement for some OoD shifts. This underscores the importance of studying ViT architecture design for OoD. We believe OoD-ViT-NAS can catalyze further research into how ViT designs influence OoD generalization.