论文链接：https://arxiv.org/pdf/2111.09833.pdf
代码链接：https://github.com/Beckschen/TransMix
1. 动机
最近很多研究已经发现，基于ViT的网络很难优化，如果训练数据不足，很容易过拟合。快速解决这一问题的方法是在训练过程中应用数据增强和正则化技术。其中，Mixup和CutMix等基于Mix的方法被证明对基于ViT的网络泛化特别有帮助。但是，以往基于Mixup的方法有一个潜在的先验知识，即假设目标的线性插值比率应该与输入插值中提出的比率保持一致。这可能会导致一种奇怪的现象，有时由于增广的随机过程，Mixup图像中没有有效的对象，但标签空间中仍然有响应。作者认为上述假设并不完全正确的，因为并非所有像素都是相等的。如上图1所示，背景中的像素对标签空间的贡献不像显著区域中的像素那么大。虽然也有工作发现了这个问题，但是那些方法仅在输入层面上Mix最有描述性的部分来解决它，这对输入的操作可能会缩小增强的空间（因为它们往往较少考虑将背景图像放入Mix图像中），同时还需要更多的参数或训练吞吐量来提取输入显著区域。
2. 贡献
本文没有研究如何在输入层面更好地Mix图像，而是更多地关注如何通过学习标签分配来缓和输入与标签空间之间的差距。具体贡献如下：
- 为了弥补输入空间和标签空间之间的差距，本文提出了TransMix，它基于Vision transformer的attention map来mix labels。attention map对相应输入图像的加权越高，标签的置信度越大。TransMix非常简单，只需几行代码就可以实现，无需向基于ViT的模型引入任何额外的参数和flop。
- 实验结果表明，在ImageNet分类上，该方法能在不同尺度上持续改进各种基于ViT的模型。在ImageNet上使用TransMix进行预处理后，基于ViT-C的模型在语义分割、目标检测和实例分割方面也表现出了较好的可移植性。在评估4个不同的基准时，TransMix也表现得更加健壮。
3. 方法
3.1 回顾: CutMix数据增强
CutMix是一种简单的数据增强技术，结合2个输入标签对 ( x A , y A ) (x_A, y_A) (xA,yA)和 ( x B , y B ) (x_B, y_B) (xB,yB)来增强一个新的训练样本 ( x ~ , y ~ ) (\tilde{x}, \tilde{y}) (x~,y~)。公式如下：
其中 M ∈ { 0 , 1 } H W M \in \{0, 1\}^{HW} M∈{
0,1}HW为二进制掩码，表示从两幅图像中剔除和填充的位置，1为二进制掩码， ⊙ \odot ⊙为逐元乘法， λ \lambda λ是 y A y_A yA在混合标签中的比例。
在增强过程中，在 x B x_B xB中移除一个随机采样区域，并用从 x A x_A xA的 A A A中裁剪的patch填充，其中patch的边界坐标均匀采样为 ( r x , r y , r w , r h ) (r_x, r_y, r_w, r_h) (rx,ry,rw,rh),混合目标分配因子 λ \lambda λ等于裁剪面积比 r w r h W H \frac{r_wr_h}{WH} WHrwrh。
3.2 回顾 ：self-attention机制
self-attention操作在一个输入矩阵 x ∈ R N × d x \in \mathbb{R}^{N \times d} x∈RN×d，其中 N N N是token的数目， d d d是每个token的维数。输入 x x x被线性映射为query，key和value，并通过使用权重矩阵 w q ∈ R d × d q w_q \in \mathbb{R}^{d \times d_q} wq∈Rd×dq, w k ∈ R d × d k w_k \in \mathbb{R}^{d \times d_k} wk∈Rd×dk和 w v ∈ R d × d v w_v \in \mathbb{R}^{d \times d_v} wv∈Rd×dv，即 q = x w q q=x_wq q=xwq, k = x w k k=x_wk k=xwk和 v = x w v v=x_wv v=xwv，其中 d q = d k d_q = d_k dq=dk。通过query和key来计算注意力图 A ( q , k ) = S o f t m a x ( q k > k ) ∈ R N × N \mathcal{A}(q, k) = Softmax(q_k>\sqrt{k}) \in R^{N×N} A(q,k)=Softmax(qk>k)∈RN×N，Self-Attention操作的输出定义为 v v v中 N N N个token特征的加权和，其权重与注意力图相对应:
通过将query、key和value分别用不同的、学习过的线性投影 g g g次投影到 d k d_k dk、 d k d_k dk和 d v d_v dv维上，可以将单头自我注意扩展到多头自我注意。
3.3 TransMix
作者提出TransMix在attention map的指导下分配mixup labels，attention map被具体定义为multi-head class attention A，它是作为自我注意的一部分计算的。在分类任务中，class token是一个查询 q q q，其对应的key k k k是所有输入token，class attention A是从类token到输入token的注意映射，总结出哪些输入token对最终分类器最有用，然后提出使用class attention A来mix labels。
-
Multi-head Class Attention
Vision transformer (ViTs)将图像 x ∈ R 3 × H × W x \in \mathbb{R}^{3 \times H \times W} x∈R3×H×W分割并嵌入到$p $个patch token x p a t c h e s ∈ R p × d x_{patches} \in \mathbb{R}^{p \times d} xpatches∈Rp×d中，通过class token x c l s ∈ R 1 × d x_{cls} \in \mathbb{R}^{1 \times d} xcls∈R1×d聚合全局信息，其中 d d d为嵌入的维数。ViT作用于patch嵌入 z = [ x c l s , x p a t c h e s ] ∈ R ( 1 + p ) × d z = [x_{cls}, x_{patches}] \in \mathbb{R}^{(1+p) \times d} z=[xcls,xpatches]∈R(1+p)×d。
给定一个具有 g g g个注意头和输入patch嵌入 z z z的Transformer，利用映射矩阵 w q w_q wq, w k ∈ R d × d w_k \in \mathbb{R}^{d \times d} wk∈Rd×d对多注意头进行参数化，每个头的class attention可表示为:
其中 q ⋅ k T ∈ R 1 × ( 1 + p ) q \cdot k^T \in \mathbb{R}^{1 \times (1+p)} q⋅kT∈R1×(1+p)表示类token是一个query，其对应的key值是所有输入tokens， A ∈ [ 0 , 1 ] p A \in [0, 1]^p A∈[0,1]p是class token到图像patch token的注意映射，概括了哪些patch对最终分类器最有用。当注意中有多个头时，我们只需对所有的注意力头求平均，得到 A ∈ [ 0 , 1 ] p A \in [0,1]^p A∈[0,1]p。在实现中，在公式（6）中的A可以作为最后一个Transformer块的中间输出而不需要进行架构修改。 -
Mixing labels with the attention map A
作者遵循CutMix中提出的输入混合过程，该过程在Eqn (1)中定义。然后在注意图 A A A的引导下重新计算 λ \lambda λ ( y A y_A yA在Eqn(2)中的比例):
其中 ↓ ( ⋅ ) \downarrow(\cdot) ↓(⋅)表示将原始 M M M从 H W HW HW变换为 p p p个像素的最近邻插值下采样。注意，为了简化忽略了Eqn (7)中的维解压缩。通过这种方式，网络可以学会根据每个数据点在注意图中的响应动态地重新分配标签的权重。被注意力映射更好聚焦的输入将在混合标签中分配一个更高的值。
3.4 伪代码实现
4. 部分实验结果
4.1 三个任务结果对比
- ImageNet Classification
- 迁移到语义分割任务
- 迁移到目标检测与实例分割任务
4.2 鲁棒分析
-
遮挡
-
对空间结构变换的敏感性
-
自然对抗与Out-of-distribution检测
4.3 泛化性研究
4.4 与最先进的Mixup变体比较
TransMix明显优于所有其他Mixup变体。与传统的CutMix相比，基于显著性的方法(如SaliencyMix和Puzzle-Mix)并没有显示出视觉Transformer的优势。我们分析了这些方法调整繁琐，在新体系结构中应用困难。例如，Attentive-CutMix不仅会带来额外的时间，还会带来参数开销，因为它引入了一个外部模型来提取显著性图。Puzzle-Mix执行的速度最低，因为它在一次训练迭代中向前和向后两次。相比之下，TransMix获得了显著的2.1%的性能提升，具有最高的训练吞吐量，而且没有参数开销。
4.5 TransMix可视化
第一行说明旧的基于区域的标签分配是反直觉的，因为图像A的前景被图像B的patch遮挡，TransMix通过Transformer的注意纠正了标签分配。如果具有鉴别性的细粒度属性出现，TransMix能够提升标签权重(例如第二行中出现了波美拉尼亚犬的脸颊和眼睛)。
4.6 消融实验
5. 结论
在本文中，作者提出TransMix，一种简单而有效的数据增强技术，它为Vision transformer分配了带有注意力引导的Mixup标签。TransMix自然地利用Transformer的注意力图来分配混合目标的可信度，并将DeiT-S和大型变型XCiT-L在ImageNet上的准确度提高了0.9%。在总共10个基准上进行了大量的实验，验证了TransMix的有效性、可移植性、鲁棒性和通用性。
由于这是第一个将基于mixup的方法向增强视觉transformer的方向推进的工作，因此存在以下限制:
- TransMix不能很好地处理那些没有class token的主干网络，因为它强烈依赖于类注意力。这个限制可以以架构修改为代价来减轻。
- TransMix要求注意图与输入在空间上对齐，这表明它可能不兼容基于变形的Transformer(如PSViT， DeformDETR)。未来，可以利用变形偏移网格将注意力图标定到输入空间位置，从而解决这一问题。
- 由于裁剪后的带有尖锐矩形边界的patch与背景的区别非常明显，Transformer可能会很自然地对裁剪后的patch产生好奇，然后对patch进行关注，因此无论patch是否包含有用的信息，都会得到一个基本的注意权值。这种现象也会出现在以往的基于显著性的方法中，因为裁剪后的patch边缘会增强一阶/二阶特征统计量。