摘要
视觉Transformer(ViTs)的架构,特别是多头注意力(MHA)机制,对硬件提出了很高的要求。在移动电话等资源受限的设备上部署ViTs,需要多种不同尺寸的模型。然而,这种方法存在局限性,例如需要分别训练和存储每个所需模型。本文介绍了一种名为HydraViT的新方法,通过堆叠注意力头来实现可扩展的ViT,从而解决了这些局限性。通过在训练过程中重复改变每一层嵌入维度的尺寸及其在MHA中对应的注意力头数量,HydraViT诱导出多个子网络。因此,HydraViT在保持性能的同时,实现了对各种硬件环境的适应性。我们的实验结果证明了HydraViT在实现具有多达10个子网络的可扩展ViT方面的有效性,涵盖了广泛的资源约束。与基线相比,HydraViT在ImageNet-1K上以相同的GMACs实现了高达5 p.p.的精度提升,以相同的吞吐量实现了高达7 p.p.的精度提升,使其成为硬件可用性多样或随时间变化的场景的有效解决方案。源代码可在https://github.com/ds-kiel/HydraViT获取。