摘要

在视觉变换器中，注意力机制是稀疏的。我们观察到，最终预测仅基于最具信息量的一小部分标记（tokens），这对于准确的图像识别已经足够。基于这一观察结果，我们提出了一种动态标记稀疏化框架，该框架可以根据输入逐步且动态地剪枝冗余标记。具体而言，我们设计了一个轻量级预测模块，用于根据当前特征估计每个标记的重要性得分。该模块被添加到不同的层中，以分层次地剪枝冗余标记。为了以端到端的方式优化预测模块，我们提出了一种注意力掩码策略，通过阻止标记与其他标记之间的交互来实现可微分的剪枝。得益于自注意力机制的特性，非结构化的稀疏标记仍然对硬件友好，这使得我们的框架易于实现实际加速。通过分层次地剪枝66%的输入标记，我们的方法显著减少了31%~37%的浮点运算次数（FLOPs），并将吞吐量提高了40%以上，而各种视觉变换器的精度下降不超过0.5%。配备动态标记稀疏化框架后，DynamicViT模型在ImageNet数据集上可以实现与最先进的卷积神经网络（CNNs）和视觉变换器相比非常有竞争力的复杂度/精度权衡。代码可在 https://github.com/raoyongming/DynamicViT 获取。