用于高效推理的多尾视觉 Transformer
计算机视觉与模式识别
2024-03-19 v3
摘要
近来，Vision Transformer (ViT) 在图像识别中取得了令人瞩目的性能，并逐步成为各类视觉任务中强有力的骨干网络。为满足 Transformer 的序列输入需求，ViT 的尾部首先将每幅图像切分为定长的视觉 token 序列，随后的自注意力层构建 token 间的全局关系，为下游任务提供有用表征。经验上，用更多 token 表示图像可带来更好性能，但自注意力层关于 token 数量的二次计算复杂度会严重影响 ViT 推理效率。为降低计算量，一些剪枝方法在 Transformer 编码器中逐步剪除无信息 token，却未改变 Transformer 之前的 token 数量。实际上，减少作为 Transformer 编码器输入的 token 数可直接降低后续计算开销。基于此，本文提出多尾视觉 Transformer (MT-ViT)。MT-ViT 采用多个尾部为后续 Transformer 编码器生成不同长度的视觉序列，并引入尾预测器决定哪条尾对该图像最高效且能给出准确预测。两模块均以端到端方式借助 Gumbel-Softmax 技巧优化。在 ImageNet-1K 上的实验表明，MT-ViT 能在精度不降的前提下显著减少 FLOPs，并在精度与 FLOPs 上均优于其他对比方法。
引用
@article{arxiv.2203.01587,
title = {Multi-Tailed Vision Transformer for Efficient Inference},
author = {Yunke Wang and Bo Du and Wenyuan Wang and Chang Xu},
journal= {arXiv preprint arXiv:2203.01587},
year = {2024}
}