[Paper Review] Auto-Prox: Training-Free Vision Transformer Architecture Search via Automatic Proxy Discovery

Zimian Wei, Lujun Li+7 more|arXiv (Cornell University)|Dec 14, 2023

Machine Learning in Materials Science10 citations

TL;DR

Auto-Prox automatically discovers zero-cost proxies for Vision Transformers using an evolutionary search with a joint correlation metric, enabling training-free ViT architecture search and achieving strong ranking and accuracy across multiple datasets.

ABSTRACT

The substantial success of Vision Transformer (ViT) in computer vision tasks is largely attributed to the architecture design. This underscores the necessity of efficient architecture search for designing better ViTs automatically. As training-based architecture search methods are computationally intensive, there is a growing interest in training-free methods that use zero-cost proxies to score ViTs. However, existing training-free approaches require expert knowledge to manually design specific zero-cost proxies. Moreover, these zero-cost proxies exhibit limitations to generalize across diverse domains. In this paper, we introduce Auto-Prox, an automatic proxy discovery framework, to address the problem. First, we build the ViT-Bench-101, which involves different ViT candidates and their actual performance on multiple datasets. Utilizing ViT-Bench-101, we can evaluate zero-cost proxies based on their score-accuracy correlation. Then, we represent zero-cost proxies with computation graphs and organize the zero-cost proxy search space with ViT statistics and primitive operations. To discover generic zero-cost proxies, we propose a joint correlation metric to evolve and mutate different zero-cost proxy candidates. We introduce an elitism-preserve strategy for search efficiency to achieve a better trade-off between exploitation and exploration. Based on the discovered zero-cost proxy, we conduct a ViT architecture search in a training-free manner. Extensive experiments demonstrate that our method generalizes well to different datasets and achieves state-of-the-art results both in ranking correlation and final accuracy. Codes can be found at https://github.com/lilujunai/Auto-Prox-AAAI24.

Motivation & Objective

• Address the challenge of training-free ViT architecture search by eliminating manual proxy design and enhancing cross-domain generalization.
• Create ViT-Bench-101 to benchmark zero-cost proxies against actual performance across multiple datasets.
• Develop an automatic proxy search space and an effective evolutionary search strategy to discover generic, high-quality proxies.
• Propose a joint correlation metric to optimize proxy generalization across diverse datasets.
• Demonstrate that the discovered proxies enable competitive or state-of-the-art ranking and final ViT performance with minimal search cost.

Proposed method

• Construct ViT-Bench-101 with diverse ViT architectures and their true performance on several datasets.
• Define a zero-cost proxy search space using ViT statistics (activations, weights, gradients) and a rich set of unary and binary operations to form computation graphs.
• Represent proxies as computation graphs (two-input expression trees per layer) and evolve them using an elitism-preserve evolutionary search guided by a Joint Correlation Metric (JCM).
• Introduce an elitism-preserve strategy to maintain high-performing proxies across generations and prevent deterioration.
• Compute proxies AutoProxA and AutoProxP in separate ViT search spaces and analyze their formulas to understand what correlates with ViT performance.
• Use the discovered proxy to perform training-free ViT architecture search by scoring candidate ViTs and selecting the best according to the Auto-Prox score.

Experimental results

Research questions

• RQ1 Can a fully automatic proxy discovery process yield zero-cost proxies that generalize across multiple ViT design spaces and datasets?
• RQ2 Does a joint correlation metric that aggregates performance correlations across several datasets improve proxy generalization compared to single-dataset proxies?
• RQ3 How does an elitism-preserve evolutionary search affect the efficiency and quality of discovered proxies?

Key findings

• Auto-Prox achieves higher ranking correlation with actual ViT performance across multiple datasets than existing zero-cost proxies.
• The joint correlation metric improves cross-dataset proxy generalization, enabling stable proxy discovery across diverse tasks and data regimes.
• Elitism-preserve strategy enhances search efficiency and prevents deterioration of proxies during evolution.
• Training-free ViT search using Auto-Prox yields competitive or state-of-the-art results with very low search cost (0.1 GPU days on a single NVIDIA A40).
• Auto-Prox-discovered proxies outperform hand-crafted proxies like TF-TAS on ViT ranking tasks and final distillation/accuracy on several ViT-Bench-101 datasets.
• ViTs selected by Auto-Prox via training-free search can achieve strong performance on ImageNet-1K when evaluated in the benchmarked spaces.