Figure 1. We propose OoD-ViT-NAS, the first comprehensive benchmark for NAS on OoD generalization of ViT architectures. The heatmap shows Kendall ฯ ranking correlations between OoD accuracy and (a) ID accuracy, (b) 9 Training-free NAS proxies, and (c) ViT architectural attributes. Key finding: embedding dimension consistently has the highest correlation with OoD accuracy, while existing Training-free NAS methods largely fail at predicting OoD accuracy despite excelling at ID.
Abstract
While Vision Transformers have achieved success across various machine learning tasks, deploying them in real-world scenarios faces a critical challenge: generalizing under Out-of-Distribution (OoD) shifts. A crucial research gap remains in understanding how to design ViT architectures, both manually and automatically, to excel in OoD generalization.
To address this gap, we introduce OoD-ViT-NAS, the first systematic benchmark for ViT Neural Architecture Search focused on OoD generalization. This benchmark includes 3,000 ViT architectures of varying computational budgets evaluated on 8 common large-scale OoD datasets. Our analysis uncovers that ViT architecture designs have a considerable impact on OoD accuracy (up to 11.85%); that ID accuracy is often a poor indicator of OoD accuracy; that existing Training-free NAS methods are largely ineffective at predicting OoD accuracy; and that simple proxies like #Param or #Flops surprisingly outperform more complex methods. Finally, we discover that increasing embedding dimensions generally enhances OoD performance, a finding traceable to improved learning of high-frequency components, outperforming SOTA domain-invariant training methods under comparable settings.
Key Findings
- Architecture choice materially affects OoD generalization. The observed accuracy range reaches 11.85% on some shifts.
- ID accuracy is a poor OoD proxy. Architectures optimized for ID performance can be suboptimal under distribution shift.
- Existing training-free NAS proxies do not reliably predict OoD accuracy. Simple parameter and FLOP counts can outperform specialized proxies.
- Embedding dimension is the strongest architectural attribute. It has the highest correlation with OoD accuracy and improves high-frequency learning.
- The benchmark spans Autoformer Tiny, Small, and Base search spaces. It varies embedding, Q-K-V dimensions, heads, MLP ratio, and depth.
- One-shot NAS makes the benchmark tractable. We sample 3,000 subnets that inherit pre-trained supernet weights.
- Eight datasets cover algorithmic, natural, and generative shifts. Evaluation includes ImageNet-C/P, ImageNet-A/O/R/Sketch, Stylized ImageNet, and ImageNet-D.
Embedding Dimension Drives OoD Robustness
Among all ViT structural attributes, embedding dimension has by far the highest correlation with OoD accuracy (ฯ โ 0.45 avg). Depth shows only a slight impact (ฯ โ 0.19), while MLP ratio and #heads exhibit very low correlation (<0.10).
Our frequency analysis reveals why: increasing embedding dimension helps ViTs learn more High-Frequency Components (HFC). Models learning more HFC achieve better OoD generalization. This insight drives a practical design rule โ by only increasing the embedding dimension of ViT-B-32 (768 โ 840), our architecture outperforms compound-scaled ViT-L-32 in OoD accuracy with far fewer parameters (98.6M vs 305.6M), achieving IN-R OoD Acc of 48.28% vs 44.33%.
Citation
@article{ho2024vision,
title={Vision Transformer Neural Architecture Search for Out-of-Distribution Generalization: Benchmark and Insights},
author={Ho, Sy-Tuyen and Vo, Tuan Van and Ebrahimkhani, Somayeh and Cheung, Ngai-Man Man},
journal={Advances in Neural Information Processing Systems},
volume={37},
pages={84197--84245},
year={2024}
}