HydraViT：堆叠注意力头实现可扩展的 ViT
计算机视觉与模式识别
2024-12-06 v2 人工智能
机器学习
摘要
Vision Transformer（ViT）的架构，特别是多头注意力（MHA）机制，对硬件提出了巨大的需求。将 ViT 部署在具有不同约束的设备（如手机）需要多个不同大小的模型。然而，这种方法存在局限性，如单独训练和存储每个所需模型。本文引入了 HydraViT，这是一种新方法，通过堆叠注意力头实现可扩展的 ViT。通过在训练期间不断改变每个层中嵌入维度的大小及其相应注意力头的数量，HydraViT 诱导多个子网络。从而，HydraViT 在保持性能的同时实现了对广泛硬件环境的适应性。我们的实验结果表明，HydraViT 在实现覆盖广泛资源约束的 10 个子网络方面具有有效性。与基线方法相比，HydraViT 在 ImageNet-1K 上以相同的 GMACs 实现最高可达 5 个百分点的额外准确率，或以相同的吞吐量实现最高可达 7 个百分点的额外准确率，使其成为硬件可用性多样或随时间变化的场景中的有效解决方案。源代码可在 https://github.com/ds-kiel/HydraViT 获取。
引用
@article{arxiv.2409.17978,
title = {HydraViT: Stacking Heads for a Scalable ViT},
author = {Janek Haberer and Ali Hojjat and Olaf Landsiedel},
journal= {arXiv preprint arXiv:2409.17978},
year = {2024}
}
备注
Accepted at NeurIPS'24, please cite the conference version