代码:https://github.com/Beckschen/TransMix
一、背景和动机
基于 mix-up 的数据增强方法对 ViT 这种结构很有用,因为这种结构容易产生过拟合,但是,之前的 mixup-based 方法有一个潜在的先验,那就是目标的线性插值比率和输入整张图的插值比率是相同的。这就会导致在 mixed image 里边可能没有有效的目标,但仍然会有label。
为了弥补上述现象导致的问题,作者提出了 TransMix,能够基于 attention map 来对 label 进行处理。
二、方法
2.1 Mixup
原始输入:
Mixup 使用一对图像 x A x_A xA 和 x B x_B xB,以及其对应的 label y A y_A yA 和 y B y_B yB 作为输入
输入和真值处理:
使用上述两个图像得到假的训练样本 λ x A + ( 1 − λ ) x B \lambda x_A + (1-\lambda)x_B λxA+(1−λ)xB,和真值 λ y A + ( 1 − λ ) y B \lambda y_A + (1-\lambda)y_B λyA+(1−λ)yB,此处 λ ∈ [ 0 , 1 ] \lambda \in [0,1] λ∈[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 ∈ { 0 , 1 } H W M\in\{0, 1\}^{HW} M∈{0,1}HW,是一个二值mask,决定了哪里舍弃,哪里使用
TransMix
A A 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,