[论文解读] Auto-scaling Vision Transformers without Training

Wuyang Chen, Wei Huang+4 more|arXiv (Cornell University)|Feb 24, 2022

Advanced Neural Network Applications被引用 7

一句话总结

As-ViT 提出了一种无需训练的自动化框架，用于设计和扩展视觉变换器（ViTs）。该框架首先通过基于复杂度的搜索发现一个最优的“种子”架构，然后系统性地将其扩展为深度与宽度均衡的更大变体。该方法仅使用 12 小时 V100 GPU 时间，便实现了 SOTA 性能——在 ImageNet-1k 上达到 83.5% 的 top-1 准确率，在 COCO 检测任务上达到 52.7% 的 mAP，同时引入了一种新型渐进式重令牌化策略，将训练 FLOPs 最多降低 56.2%。

ABSTRACT

This work targets automated designing and scaling of Vision Transformers (ViTs). The motivation comes from two pain spots: 1) the lack of efficient and principled methods for designing and scaling ViTs; 2) the tremendous computational cost of training ViT that is much heavier than its convolution counterpart. To tackle these issues, we propose As-ViT, an auto-scaling framework for ViTs without training, which automatically discovers and scales up ViTs in an efficient and principled manner. Specifically, we first design a "seed" ViT topology by leveraging a training-free search process. This extremely fast search is fulfilled by a comprehensive study of ViT's network complexity, yielding a strong Kendall-tau correlation with ground-truth accuracies. Second, starting from the "seed" topology, we automate the scaling rule for ViTs by growing widths/depths to different ViT layers. This results in a series of architectures with different numbers of parameters in a single run. Finally, based on the observation that ViTs can tolerate coarse tokenization in early training stages, we propose a progressive tokenization strategy to train ViTs faster and cheaper. As a unified framework, As-ViT achieves strong performance on classification (83.5% top1 on ImageNet-1k) and detection (52.7% mAP on COCO) without any manual crafting nor scaling of ViT architectures: the end-to-end model design and scaling process cost only 12 hours on one V100 GPU. Our code is available at https://github.com/VITA-Group/AsViT.

研究动机与目标

• 为解决缺乏系统性、高效的无需昂贵训练的视觉变换器（ViTs）设计与扩展方法的问题。
• 降低 ViT 训练带来的高计算成本，其训练开销显著高于卷积神经网络（CNNs）。
• 在统一的端到端框架中自动化架构发现与扩展过程。
• 实现在无需手动调整超参数或进行大规模超参数搜索的前提下，获得高性能的 ViTs。
• 开发一种可扩展的、无需训练的搜索方法，其与实际模型准确率具有强相关性。

提出的方法

• 采用无需训练的架构搜索，利用网络复杂度度量（尤其是期望长度失真）作为模型准确率的代理指标，从而实现对 ViT 拓扑结构的快速评估。
• 从基于先前人工设计的 ViT 拓扑的宽松搜索空间中发现‘种子’ViT 拓扑结构，通过基于复杂度的排序替代昂贵的训练过程。
• 提出一种系统性的扩展规则，通过逐层逐步增加深度与宽度，确保每一步都保持网络复杂度的均衡。
• 渐进式重令牌化在训练过程中动态调整令牌化方式（通过扩张和步长），减少令牌数量，从而加速训练。
• 该框架可在单次运行中生成一系列参数量不同的 ViT 变体，起点为单一种子架构。
• 该方法利用了 ViTs 在训练初期可容忍粗粒度令牌化的观察，从而实现高效的渐进式优化。

实验结果

研究问题

• RQ1 基于网络复杂度度量的无需训练搜索是否能可靠预测 ViT 性能，而无需实际训练？
• RQ2 何种最优的扩展策略能平衡 ViT 的深度与宽度，同时保持性能？
• RQ3 如何在不牺牲准确率的前提下，加速 ViT 训练，特别是针对大模型？
• RQ4 能否从单一种子架构系统性地扩展为多个变体，且附加成本极低？
• RQ5 渐进式重令牌化是否能提升训练效率，同时保持模型准确率？

主要发现

• 期望长度失真度量在计算成本与实际 ViT 准确率之间的 Kendall-tau 相关系数之间实现了最佳平衡，使其成为无需训练搜索的理想选择。
• As-ViT 框架仅用七小时 V100 GPU 时间便发现‘种子’ViT 拓扑结构，随后在五小时内完成扩展，总计 12 小时完成完整设计与扩展流程。
• As-ViT 在 ImageNet-1k 上达到 83.5% 的 top-1 准确率，在 COCO 目标检测任务上达到 52.7% 的 mAP，优于现有手工设计的 ViTs。
• 渐进式重令牌化策略相比固定粗粒度令牌化，将训练 FLOPs 最多降低 56.2%，显著加速训练过程。
• 该框架可从单一种子架构生成一整套不同参数量的 ViT 变体，支持在不同硬件约束下的高效部署。
• 该方法表明，通过基于复杂度的度量可有效扩展 ViTs，从而在搜索与设计阶段绕过昂贵的训练循环。