SHViT：单头视觉Transformer搭配内存高效的巨型设计

最近，高效视觉Transformer在资源受限设备上展现出卓越的性能。传统方法通常采用4x4分块嵌入和4级宏观结构，同时在微观层面使用多头注意力机制。本文旨在以内存高效的方式解决所有设计层面的计算冗余问题。我们发现，使用更大步幅的分块嵌入不仅能降低内存访问成本，还能通过利用早期阶段减少空间冗余的token表示，实现具竞争力的性能。进一步的初步分析表明，早期注意力层可以被卷积替代，而后期几个注意力头可能具有计算冗余。为此，我们引入单头注意力模块，既能内在地防止头冗余，又能通过并行结合全局和局部信息来提升准确率。基于我们的解决方案，我们提出SHViT——一个单头视觉Transformer，实现了最佳的速度与精度权衡。例如，在ImageNet-1k上，SHViT-S4在GPU、CPU和iPhone12移动设备上的速度分别比MobileViTv2 x1.0快3.3倍、8.1倍和2.4倍，同时准确率提高1.3个百分点。在MS COCO上进行对象检测和实例分割任务中使用Mask-RCNN头部，本模型的性能与FastViT-SA12相当，但在GPU和移动设备上的 backbone latency分别低3.8倍和2.0倍。