字节跳动提出让所有 token 都参与 loss，提高 transformer 的表现

19 All Tokens Matter: Token Labeling for Training Better Vision Transformers

Score /5: ⭐️⭐️⭐️⭐️

论文地址：

本想偷个懒，看看有没有人写这篇文章的阅读笔记，结果没找到，只能亲自写了。

字节跳动的文章，目的是提高准确率，采用的方法主要是将图片 embedding 后，让所有的 token 都参与最后 loss 的计算，而不是只用 classification token。

其中还借鉴了一些 CutMix 之类的 data augmentation 的方法。

其实我想看一下，如果只用 all token，但不用 cutmix 的效果会是怎样，不过 ablation study 里没有提到

摘要

• 提出 token labeling，新的 training objective。
• 利用了图片中的所有 token，而并非像以往的 ViT 一样，只利用了 classification token 来计算 loss。
• 相当于将一个图像分类问题，转换成了多 token 尺度的识别问题，且给每一个 patch 都分配了一个独立的，与位置相关联的，由机器生成的监督（supervision）。

简介

• 普通的 ViT 只用了 classification token 计算全局信息中的损失，这样的话其它一些 patch 中的丰富的局部信息就被忽视了。
• 作者提出将一个由机器生成的 K 维的分数图，来监督所有的 token，其中 K 是数据中类的个数。
• 这样，每个 patch 就和一个单独的与位置相关的监督联系在了一起，这个监督表示目标物体是否在对应的 image patch 中存在。

Token Labeling 方法

token labeling

以往的 transformer 计算 loss 都是用以下的公式：

$$L_{cls}=H\left(X^{cls},y^{cls}\right)$$

但作者提出的方法将计算所有 token 的 loss，且每个 token 应该有一个相对应的独立的，与位置相关的标签。因此，ground truth 不再是一个简单的 K 维的向量，而是一个 $K\times N$ 的矩阵，或者叫 K 维分数图，表示为 $\left[y^1,\dots ,y^N\right]$，其中 N 是输入 patch token 的个数。

这个新的 token labeling loss 由以下的公式来计算（相当于是在所有 token 上取了个平均）

$$L_{tl}=\frac{1}{N}\sum _{i=1}^{N}H\left(X^i,y^i\right)$$

于是，最后的总 loss 由以下的公式计算：

$$\begin{array}{rl}{L}_{\text{total }}& =H\left(X^{cls},y^{cls}\right)+\beta \cdot L_{tl}\\ & =H\left(X^{cls},y^{cls}\right)+\beta \cdot \frac{1}{N}\sum _{i=1}^{N}H\left(X^i,y^i\right)\end{array}$$

$\beta$ 是一个可调超参，实验中取值 0.5

优点

• 不像其它蒸馏方法一样需要一个教师模型，且计算量比较低
• dense score map 可以由预训练模型直接生成，比如 EfficientNet, NFNet 等等，训练中，我们只需要将生成的分数图截取出来再插值，使之与截出来的图在位置上对应好。

Token Labeling with MixToken

• 之前的研究表明，在 Transformer 中使用 augmentation 能起到很好的效果
• 如果直接在原图上使用 CutMix，那在经过 embedding 后生成的 token 里，有的 token 里会有来自附近几个 patch 的信息，很难给出一个干净且正确的标签。
• 作者因此在 patch embedding 之后再使用类似 CutMix 的方法，来避免上述问题。
• 具体地说，两张图片 $I_1,I_2$，对应的标签为 $Y_1=\left[y_1^1,\dots ,y_1^N\right],Y_2=\left[y_2^1,\dots ,y_2^N\right]$，那首先经过 patch embedding，将它们变成 token： $T_1=\left[t_1^1,\dots ,t_1^N\right]\text{ and }{T}_2=\left[t_2^1,\dots ,t_2^N\right]$，再用下面的公式生成新的 token：

$$\hat{T}=T_1\odot M+T_2\odot (1-M)$$

其中 M 是一个二值掩码， $\odot$ 是点积，M 是根据论文“Regularization strategy to train strong classifiers with localizable features” 中的方法生成的。

接着用下式对标签进行相同的操作：

$$\hat{Y}=Y_1\odot M+Y_2\odot (1-M)$$

这样，类的标签就可以写成以下形式：

$$y^{\widehat{cls}}=\bar{M}{y}_1^{cls}+(1-\bar{M})y_2^{cls}$$

其中 $\bar{M}$ 是 $M$ 中所有元素的平均值。

实验及其它

这些就不写了，感兴趣的可以看一下，截几张图放在这里。