[论文解读] 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 Science被引用 10

一句话总结

Auto-Prox 自动通过一个 joint correlation metric 的进化搜索发现 Vision Transformers 的零成本代理，使无需训练的 ViT 架构搜索成为可能，并在多个数据集上实现强排序和准确性。

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.

研究动机与目标

• 通过消除手工代理设计并提升跨领域泛化能力来解决无训练 ViT 架构搜索的挑战。
• 创建 ViT-Bench-101，以在多个数据集上对照实际性能基准零成本代理。
• 开发一个自动代理搜索空间和有效的进化搜索策略，以发现通用且高质量的代理。
• 提出一个联合相关性度量，以在多样化数据集上优化代理的泛化能力。
• 证明所发现的代理能够以极低的搜索成本实现有竞争力或最先进的排序与最终 ViT 性能。

提出的方法

• 使用多样化的 ViT 架构及其在若干数据集上的真实性能来构建 ViT-Bench-101。
• 使用 ViT 统计数据（激活、权重、梯度）及一组丰富的一元和二元运算来定义零成本代理搜索空间，以形成计算图。
• 将代理表示为计算图（每一层的两输入表达树），并在一个由联合相关性度量（Joint Correlation Metric，JCM）引导的保持精英的进化搜索中对它们进行进化。
• 引入保持精英策略，以在世代之间保持高性能代理并防止退化。
• 在独立的 ViT 搜索空间中计算代理 AutoProxA 和 AutoProxP，并分析它们的公式以理解与 ViT 性能的相关性。
• 使用发现的代理通过对候选 ViTs 进行评分并按照 Auto-Prox 分数选择最佳，从而执行无需训练的 ViT 架构搜索。

实验结果

研究问题

• RQ1 完全自动的代理发现过程是否能够产生在多个 ViT 设计空间和数据集上具有泛化能力的零成本代理？
• RQ2 在多个数据集上汇集性能相关性的联合相关性度量，是否比单一数据集代理提高了代理的泛化能力？
• RQ3 保持精英的进化搜索如何影响发现代理的效率和质量？

主要发现

• Auto-Prox 在多个数据集上的排序相关性与实际 ViT 性能相比，超过了现有的零成本代理。
• 联合相关性度量提升了跨数据集的代理泛化能力，使得在多样化任务和数据条件下能够稳定发现代理。
• 保持精英策略提高搜索效率并防止代理在进化过程中的退化。
• 使用 Auto-Prox 进行的无训练 ViT 搜索以极低的搜索成本（在单个 NVIDIA A40 上为 0.1 GPU 天）获得有竞争力或最先进的结果。
• Auto-Prox 发现的代理在 ViT 排序任务和若干 ViT-Bench-101 数据集上的最终蒸馏/准确性方面，超过了手工设计的代理，如 TF-TAS。
• 通过无训练搜索由 Auto-Prox 选择的 ViT 在 ImageNet-1K 上的评估空间中可以获得较强的性能。