通过堆叠注意力头实现可扩展视觉变压器,适配不同硬件需求。
HydraViT: Stacking Heads for a Scalable ViT
- 在训练中动态调整每层嵌入维度与注意力头数,生成多个子网络。
- 支持最多10个子网络,在相同计算量下比基线高5个百分点准确率。
- 适合部署环境多变或需长期维护的设备,如手机端视觉任务。
视觉变换器(ViT)的架构,尤其是多头注意力(MHA)机制,对硬件资源要求较高。在手机等资源受限设备上部署时,通常需要训练和存储多个不同尺寸的模型,但该方法存在训练和存储成本高的问题。本文提出HydraViT,通过在训练过程中反复改变每一层的嵌入维度及其对应的注意力头数量,实现注意力头的堆叠,从而诱导出多个子网络。该方法使模型可在广泛硬件环境下自适应运行,同时保持高性能。实验表明,HydraViT可实现最多10个子网络,覆盖多种资源约束场景;在ImageNet-1K上,相同GMACs下比基线最高提升5个百分点准确率,相同吞吐量下提升7个百分点,是一种应对硬件多样性或动态变化的有效方案。源码已公开于https://github.com/ds-kiel/HydraViT。
原文摘要 · Abstract (English)
The architecture of Vision Transformers (ViTs), particularly the Multi-head Attention (MHA) mechanism, imposes substantial hardware demands. Deploying ViTs on devices with varying constraints, such as mobile phones, requires multiple models of different sizes. However, this approach has limitations, such as training and storing each required model separately. This paper introduces HydraViT, a novel approach that addresses these limitations by stacking attention heads to achieve a scalable ViT. By repeatedly changing the size of the embedded dimensions throughout each layer and their corresponding number of attention heads in MHA during training, HydraViT induces multiple subnetworks. Thereby, HydraViT achieves adaptability across a wide spectrum of hardware environments while maintaining performance. Our experimental results demonstrate the efficacy of HydraViT in achieving a scalable ViT with up to 10 subnetworks, covering a wide range of resource constraints. HydraViT achieves up to 5 p.p. more accuracy with the same GMACs and up to 7 p.p. more accuracy with the same throughput on ImageNet-1K compared to the baselines, making it an effective solution for scenarios where hardware availability is diverse or varies over time. Source code available at https://github.com/ds-kiel/HydraViT.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。