【Transformer】TransMix: Attend to Mix for Vision Transformers

代码：https://github.com/Beckschen/TransMix

一、背景和动机

基于 mix-up 的数据增强方法对 ViT 这种结构很有用，因为这种结构容易产生过拟合，但是，之前的 mixup-based 方法有一个潜在的先验，那就是目标的线性插值比率和输入整张图的插值比率是相同的。这就会导致在 mixed image 里边可能没有有效的目标，但仍然会有label。

为了弥补上述现象导致的问题，作者提出了 TransMix，能够基于 attention map 来对 label 进行处理。

二、方法

2.1 Mixup

原始输入：

Mixup 使用一对图像 $x_A$ 和 $x_B$，以及其对应的 label $y_A$ 和 $y_B$ 作为输入

输入和真值处理：

使用上述两个图像得到假的训练样本 $\lambda x_A+(1-\lambda )x_B$，和真值 $\lambda y_A+(1-\lambda )y_B$，此处 $\lambda \in [0,1]$ 是一个从 Beta 分布拿到的随机数。

如图 1 所示，背景像素是没有办法和前景对 label 起到相同的作用，也就是并非所有像素对 label 的贡献是相同的。

所以本文聚焦于如何使用可学习的方法来实现输入和label空间的统一。

作者发现，vision transformer 产生的 attention map 可以较好的使用于该任务。

如图 1 所示，作者使用 attention map 的权重作为 $\lambda$ 的值，label 能够被 re-weighted，每个像素点的权重都是不同的，所以不会使用相同的值对图像中的所有像素进行组合。而且由于使用的是 attention map，所以该方法能够适用于任何 ViT-based 的方法，并且没有额外参数。

2.2 TransMix

CutMix 数据增强：

CutMix 是一个简单的增强方式，将两个 label 结合起来，生成一个新 label：

• $M\in \{0,1{\}}^{HW}$，是一个二值mask，决定了哪里舍弃，哪里使用

TransMix

$A$ 是从 cls token 到输入图像的 token 的 attention map，表示了每个 patch 对最终的分类结果的重要性。对于多头的 attention，作者使用了平均的方法。

使用 attention map 对 label 进行处理：

下箭头表示最近邻插值，可以把 M 从 HW 大小下采样成 p 个像素。

这样的话，网络就可以给 label 的每个点基于 attention map 来动态分配权重。

伪代码：

三、效果

TransMix 的可视化如下：

第一行展示了area-based 的 label assignment ，把 image A 的一块粘贴到 B 上，TransMix 能够使用 attention map 对 label 进行修正，能够提升突变区域 label 的 weight，