[论文解读] All Tokens Matter: Token Labeling for Training Better Vision Transformers

Zihang Jiang, Qibin Hou+6 more|arXiv (Cornell University)|Apr 22, 2021

Advanced Neural Network Applications参考文献 56被引用 140

一句话总结

本文提出 token labeling，即对视觉变换器的密集补充监督对象，对 ViT 变体普遍带来收益，并在较小模型上实现高准确率。它在 ImageNet 上以 26M 参数达到 84.4% Top-1，150M 参数达到 86.4%，并提升下游语义分割性能。

ABSTRACT

In this paper, we present token labeling -- a new training objective for training high-performance vision transformers (ViTs). Different from the standard training objective of ViTs that computes the classification loss on an additional trainable class token, our proposed one takes advantage of all the image patch tokens to compute the training loss in a dense manner. Specifically, token labeling reformulates the image classification problem into multiple token-level recognition problems and assigns each patch token with an individual location-specific supervision generated by a machine annotator. Experiments show that token labeling can clearly and consistently improve the performance of various ViT models across a wide spectrum. For a vision transformer with 26M learnable parameters serving as an example, with token labeling, the model can achieve 84.4% Top-1 accuracy on ImageNet. The result can be further increased to 86.4% by slightly scaling the model size up to 150M, delivering the minimal-sized model among previous models (250M+) reaching 86%. We also show that token labeling can clearly improve the generalization capability of the pre-trained models on downstream tasks with dense prediction, such as semantic segmentation. Our code and all the training details will be made publicly available at https://github.com/zihangJiang/TokenLabeling.

研究动机与目标

• 通过对每个图像补丁 token 使用密集、位置特定的监督来重新审视 ViTs 的训练方式。
• 将图像分类转化为多个 token 级识别问题，以提升对象定位与识别。
• 展示 token labeling 在不同 ViT 变体和模型规模上的通用性与有效性。

提出的方法

• 定义一个 token labeling 目标，使用一个 K 维分数图来对 N 个补丁 token 进行密集监督。
• 通过 L_total = H(X^{cls}, y^{cls}) + beta * (1/N) * sum_i H(X^i, y^i) 将图像分类 token 损失与密集 token labeling 损失结合起来。
• 使用预训练的机器标注器生成 token 级标签，以提供逐补丁的监督。
• 引入 MixToken，一种在补丁嵌入后混合 token 空间的增强方法，标签亦相应混合。
• 用一个小卷积替换 patch embedding 模块，以在 LV-ViT 变体中更好地捕捉局部信息。

实验结果

研究问题

• RQ1 dense、定位特定的 token 监督是否能超越传统的类 token 训练来提升视觉变换器？
• RQ2 token labeling 如何在 ViT 变体的不同增广和模型规模下互动？
• RQ3 token labeling 对不同标注者质量和 patch 嵌入选择是否稳健，并是否有助于下游密集预测任务？
• RQ4 在多种 ViT 架构上应用 token labeling 时，ImageNet 和 ADE20K 的性能提升是多少？

主要发现

• 在 150M 参数的 LV-ViT 上，token labeling 实现了 84.4% 的 ImageNet Top-1，在更大规模下达到 86.4%，在同等预算下优于若干基于 Transformer 的基线。
• Token labeling consistently improves seven ViT variants (DeiT, T2T-ViT, LV-ViT, etc.) with larger gains for bigger models.
• MixToken 在基于 token 的 Transformer 上优于 CutMix，并有助于在 LV-ViT-S 上结合 token labeling 与 MixToken 达到 83.3% Top-1 的 ImageNet。
• 来自 token labeling 的密集监督也提升了 ADE20K 下游语义分割的表现，在非 ImageNet-22K 预训练设置中达到 state-of-the-art（如 LV-ViT-L + UperNet 的 51.8 mIoU）。
• 该方法需要一个预训练的 annotator 来生成 token 级分数，但由于标注是预先计算的，因此额外训练成本可忽略不计。