ICCV2023 | 视觉 Transformer 的 Token-标签对齐-CSDN博客

本文链接：https://blog.csdn.net/qq_60090693/article/details/147699983

Token-Label Alignment for Vision Transformers

摘要-Abstract
引言-Introduction
相关工作-Related Work
提出的方法-Proposed Method
实验-Experiments
结论-Conclusion

在这里插入图片描述

本文 “Token-Label Alignment for Vision Transformers” 提出了用于视觉 Transformer（ViTs）训练的Token-Label Alignment（TL-Align）方法，旨在解决数据混合策略应用于 ViTs 时出现的 token 波动现象。通过追踪输入和转换后 token 的对应关系，为每个输出 token 获取对齐的标签，以提供更准确的训练信号。实验表明，TL-Align 在多种 ViT 架构上均提升了性能，在图像分类、语义分割、目标检测和迁移学习任务中效果显著，且具有良好的鲁棒性和泛化能力。

摘要-Abstract

Data mixing strategies (e.g., CutMix) have shown the ability to greatly improve the performance of convolutional neural networks (CNNs). They mix two images as inputs for training and assign them with a mixed label with the same ratio. While they are shown effective for vision transformers (ViTs), we identify a token fluctuation phenomenon that has suppressed the potential of data mixing strategies. We empirically observe that the contributions of input tokens fluctuate as forward propagating, which might induce a different mixing ratio in the output tokens. The training target computed by the original data mixing strategy can thus be inaccurate, resulting in less effective training. To address this, we propose a token-label alignment (TL-Align) method to trace the correspondence between transformed tokens and the original tokens to maintain a label for each token. We reuse the computed attention at each layer for efficient token-label alignment, introducing only negligible additional training costs. Extensive experiments demonstrate that our method improves the performance of ViTs on image classification, semantic segmentation, objective detection, and transfer learning tasks.

数据混合策略（例如CutMix）已被证明能够显著提升卷积神经网络（CNNs）的性能。它们将两张图像混合作为训练输入，并为其分配具有相同比例的混合标签。尽管这些策略在视觉 Transformer（ViTs）中也显示出了一定效果，但我们发现了一种 token 波动现象，这种现象抑制了数据混合策略的潜力。通过观察，我们发现输入 token 在前向传播过程中的贡献会发生波动，这可能会导致输出 token 的混合比例有所不同。因此，原始数据混合策略计算出的训练目标可能不准确，从而导致训练效果不佳。为了解决这个问题，我们提出了一种 Token-标签对齐（TL-Align）方法，用于追踪变换后的 token 与原始 token 之间的对应关系，以便为每个 token 保留一个标签。我们在每一层重用计算得到的注意力，以高效地实现 token-标签对齐，仅引入了可忽略不计的额外训练成本。大量实验表明，我们的方法在图像分类、语义分割、目标检测和迁移学习任务中均提升了视觉 Transformer 的性能。

引言-Introduction

这部分内容主要介绍了视觉Transformer（ViTs）的发展、数据混合策略的应用，指出ViTs中存在的 token 波动现象影响数据混合策略效果，并提出了 token-标签对齐（TL-Align）方法，具体内容如下：

ViTs 的发展与数据混合策略的应用：ViTs 在计算机视觉领域取得显著进展，革新了多个任务的技术水平，其成功的结构也推动了其他架构的发展。同时，数据混合策略作为一种数据增强方式，被广泛应用于现代深度架构的训练中，能有效提升模型的泛化性能，如 CutMix 通过空间域的复制粘贴操作生成混合图像。
ViTs 中 token 波动现象及影响：虽然数据混合策略在 CNNs 中研究广泛，但在 ViTs 中的兼容性研究较少。研究发现，ViTs 中的自注意力机制会破坏输入的空间结构，导致 token 与标签错位，引发 token 波动现象。这使得输出 token 的混合比例与预期不同，进而导致原始数据混合策略计算的训练目标不准确，影响训练效果。
提出 TL-Align 方法：为解决上述问题，提出了 TL-Align 方法。该方法通过追踪变换后 token 与原始 token 的对应关系，为每个 token 分配标签。具体操作是根据 token 来源为混合图像的输入 token 分配标签，再利用计算的注意力线性混合输入 token 的标签，迭代进行 token-标签对齐，从而获得更准确的训练目标。此方法仅用于训练阶段，不增加推理时的额外工作量。实验显示，TL-Align 能提升多种 ViT 模型的性能，在多个任务中验证了其鲁棒性和泛化能力。

在这里插入图片描述
图1. 所提出的 TL-Align 方法概述。(a) 类似 CutMix 的方法在模型训练中被广泛使用，这些方法在输入空间中对 token 及其标签进行空间混合。(b) 它们最初是为卷积神经网络（CNNs）设计的，并且假设处理后的 token 在空间上与输入 token 对齐。我们发现，由于视觉 Transformer（ViTs）具有全局感受野和自适应权重，这一假设并不适用于 ViTs。(c) 与现有方法相比，我们的方法无需预训练的教师网络，就能有效且高效地对齐 token 和标签。

提出的方法-Proposed Method

预备知识-Preliminaries

这部分内容主要介绍了卷积神经网络（CNN）、视觉 Transformer（ViT）的基本架构，以及数据混合策略中的 CutMix 方法，为后续阐述研究问题和提出解决方法做铺垫，具体如下：

CNN 与 ViT 架构：在深度学习时代，CNN 是计算机视觉领域的主流架构，对众多任务性能提升显著。但近年来，ViT 的出现对其主导地位构成挑战。ViT 将图像 “分块” 成 token，通过交替的自注意力（SA）和多层感知器（MLP）进行处理。
数据混合策略之 CutMix：训练策略对模型性能影响重大，数据混合是 CNNs 和 ViTs 训练中重要的数据增强手段，可提升模型泛化能力。CutMix 是常用的数据混合策略，旨在从给定训练样本 $(X, y)$ 创建虚拟训练样本。它从一个输入图像 $X_{1}$ 中随机选择局部区域，替换另一个输入图像 $x_{2}$ 相同区域的像素，生成新样本 $\tilde{x}$ 。新样本标签 $\bar{y}$ 是原始标签 $y_{1}$ 和 $y_{2}$ 的组合，计算公式为 $\begin{cases}\tilde{X} =M \odot X_{1}+(1 - M) \odot X_{2}\\\tilde{y} =\lambda y_{1}+(1-\lambda) y_{2}\end{cases}$ 。其中， $M$ 是表示像素所属图像的二进制掩码， $\lambda$ 反映两个标签的混合比例，由从 $x_{1}$ 裁剪区域的像素占比确定。

Token 波动现象-The Token Fluctuation Phenomenon

该部分内容主要阐述了在视觉 Transformer（ViT）中发现的 token 波动现象，分析了其产生的原因及对训练造成的影响，具体如下：

现象发现：CutMix 原本为 CNNs 设计，其假设特征提取过程不会改变混合比例。但研究发现，ViT 中的自注意力机制与 CNNs 不同，会导致部分 token 出现波动。
理论分析
- 用 $z_{i}$ 表示图像 $Z$ 的一个 token，经过空间操作后的第 $i$ 个变换 token $\hat{z}_{i}$ 可表示为 $\hat{z}_{i}=\sum_{j = 1}^{N}w_{i, j}^{s}z_{j}$ ，其中 $w_{i, j}^{s}$ 是空间混合矩阵 $w^{s}(z)$ 的元素。在此基础上，定义原始 token $z_{i}$ 对混合 token $\hat{z}_{j}$ 的贡献 $c(z_{i}, \hat{z}_{j})=\frac{|w_{i, j}^{s}|}{\sum_{k = 1}^{N}|w_{k, j}^{s}|}$ ，token $z_{i}$ 在所有混合图像 token 中的存在性 $p(z_{i})=\sum_{j = 1}^{N}c(z_{i}, \hat{z}_{j})$ 。
- 对于非步长深度卷积，由于平移不变性， $\sum_{l = 1}^{N}|w_{i, l}^{s}|=\sum_{j = 1}^{N}|w_{k, j}^{s}|=\sum_{k = 1, l = 1}^{M}|K_{k, l}|$ ，可推出内部 token 的 $p(z_{i}) = 1$ ，即卷积过程中内部 token 的效果不变。然而，ViT 中的自注意力机制不存在平移不变性，上述等式不成立，且自注意力引起的空间混合矩阵 $w^{s}(z)$ 的输入依赖性会进一步放大 $p (z)$ 的波动，极端情况下某些 token 的 $p (z)$ 可能接近 0 或 1。
影响阐述：token 的波动会改变混合比例（即 $\lambda$ ），使网络可能完全忽略其中一个混合图像，导致处理后 token 的实际标签偏离按原公式计算的混合标签，进而降低训练效果。

Token-标签对齐-Token-Label Alignment

这部分主要介绍了 Token-标签对齐（TL-Align）方法，通过一系列具体操作，解决了视觉 Transformer（ViT）中 token 与标签不匹配的问题，为模型训练提供更准确的标签，具体内容如下：

方法提出背景：ViT 中自注意力机制使处理后的 token 与初始 token 不匹配，为解决该问题，提出 TL-Align 方法，通过追踪输入和转换后 token 的对应关系，获取对齐的标签。
具体操作流程
- 初始 token 及标签处理：将 CutMix 后的混合输入 $\tilde{x}$ 分割并展平为原始图像 token $\{\tilde{x}_{1}, \tilde{x}_{2}, \cdots, \tilde{x}_{N}\}$ ，投影并添加位置嵌入得到 $Z^{0}$ 。同时，为每个 token $z_{i}$ 分配标签嵌入 $y_{i}$ ，形成初始标签嵌入 $Y^{0}$ ，初始化方式依据不同数据混合范式而定。
- 基于注意力的标签更新：ViT 通常采用多头自注意力（MSA），在计算中，先计算 $\cdot W_{Q}$ 、 $\cdot W_{K}$ 、 $\cdot W_{V}$ ，得到注意力矩阵 $\mathcal{A}(Q, K) = Softmax\left(Q \cdot K^{T} / \sqrt{d}\right)$ ，进而计算 $\hat{Z} = SA(Z)=\mathcal{A}(Q, K) \cdot V$ .
  为对齐标签，使用相同注意力矩阵 $\mathcal{A}(Q, K)$ 更新标签嵌入 $Y$ ，即 $\hat{Y}=\mathcal{A}(Q, K) \cdot Y$ 。对于 MSA，通过对所有注意力矩阵求平均来对齐标签，公式为 $\hat{Y}= \text{TL-Align-S}(Z, Y):=\frac{1}{H} \sum_{i = 1}^{H} \mathcal{A}_{i}(Q, K) \cdot Y$ ，其中 $H$ 为头数， $\mathcal{A}_{i}$ 是第 $i$ 个头的注意力矩阵。
- 逐层对齐与最终标签确定：每个 Transformer 块通过空间和通道混合处理 token，TL-Align 以类似方式对齐标签嵌入。对于层次化视觉 Transformer 如 Swin，合并 token 时采用添加标签嵌入并归一化的方式获取对齐标签。经过逐层对齐，得到最终的对齐 token $Z^{L}$ 和标签 $Y^{L}$ ，根据模型不同，图像的最终对齐标签 $y_{align}$ 为类 token 的标签 $y_{cls}^{L}$ 或所有空间 token 标签的平均值 $\frac{1}{N} \sum_{i = 1}^{N} y_{i}^{L}$ ，并用于训练网络。
  
  图2. 所提出的 TL-Align 方法示意图。我们追踪输入 token 和转换后 token 之间的对应关系，并据此对齐标签。我们重用计算得到的注意力，对输入 token 的标签进行线性混合，以获得转换后 token 的标签。通过迭代进行 token-标签对齐，为每个输出 token 获取一个标签。
方法优势：TL-Align 作为一个即插即用模块，在训练过程中自适应调整每个 token 的标签，在保持 token 和标签对齐的同时，仅引入可忽略的训练成本，且在推理时不增加额外计算成本。

实验-Experiments

这部分主要通过多组实验，全面评估了 TL-Align 方法的性能，验证了其有效性、鲁棒性和泛化性，具体内容如下：

ImageNet分类实验
- 实验设置：基于 PyTorch 和 timm 库，在多种 ViT 架构（DeiT、PVT、Swin的不同变体）上进行实验。小模型从 scratch 训练 300 轮，使用 CutMix 且保持其他设置不变，仅将 CutMix 的混合目标替换为 TL-Align 得到的标签；大模型则微调预训练模型。
- 实验结果：TL-Align 提升了不同架构模型的性能，如 DeiT-T、DeiT-S、DeiT-B 的 Top-1 准确率分别提升 1.0%、0.8%、0.5%. 与其他训练策略相比，在保持参数数量和训练速度的同时，TL-Align 性能更优。
  表1. ImageNet 图像分类任务的结果。我们比较了不同视觉 Transformer 骨干网络在使用和不使用我们的 TL-Align 方法时的参数数量、浮点运算次数（FLOPs）和准确率。
  
  表2. 在 ImageNet 上我们的 TL-Align 方法与其他训练策略的比较。
下游任务实验
- 语义分割：在 ADE20K 数据集上，以 DeiT-S 和 Swin 系列为骨干网络，TL-Align 提高了不同模型尺度的分割性能。
  表3. ADE20K 语义分割任务的结果。
- 目标检测和实例分割：在 COCO 2017 数据集上，将 TL-Align 应用于 Swin，采用 Cascade Mask-RCNN 框架和 3x 训练策略，所有 Swin 变体模型性能均有提升。
  表4. 在 COCO 数据集上进行目标检测和实例分割的实验结果。
- 迁移学习：在 CIFAR-10、CIFAR-100、Flowers 和 Cars 数据集上，对 ImageNet 预训练模型微调，TL-Align 使 DeiT 不同变体模型性能显著提升。
  表5. 不同迁移学习数据集上的结果。
性能分析和可视化实验
- token-标签对齐有效性：计算原始目标和对齐标签的均方根误差（RMSE），发现模型越大，RMSE 越小，Swin 的 RMSE 低于相似大小的 DeiT，表明小模型和 DeiT 类骨干网络受 token 波动影响大，TL-Align 对其提升更显著。
  
  图3. 原始 CutMix 目标与通过 TL-Align 获得的标签之间的均方根误差（RMSE）。
- 波动 token 混合比例可视化：计算基于相似性的“真实”混合比例，与其他方法对比，发现 CutMix 等方法假设输出 token 与输入 token 空间对应，计算固定混合比例，而TL-Align 通过逐层对齐为 token 分配动态标签，更准确。
  
  图4. 不同层波动 token 的混合比例 $λ$ 可视化。我们将 TL-Align 的结果与 CutMix、token 相似度、TransMix 和 TokenLabeling 的结果进行了比较。
- 鲁棒性和泛化性评估：在多个损坏和分布外数据集上评估，采用 AutoAttack 评估对抗鲁棒性，结果表明 TL-Align 提高了模型的鲁棒性和泛化性。
  表7. 使用不同数据混合策略的消融实验。
- 不同数据混合策略的消融实验：将 TL-Align 应用于多种数据混合策略，结果显示均有性能提升，验证了其泛化性。
  表7. 使用不同数据混合策略的消融实验。
- 不同标签对齐操作的消融实验：对 DeiT-S 进行实验，发现仅用第 12 层注意力图对齐或部分层对齐且禁用归一化，性能提升有限或下降，证明了 TL-Align 中利用注意力和归一化逐层对齐的重要性。
  表8. 不同 TL-Align 操作的消融实验。
- 对齐标签可视化：可视化 DeiT-S 和 Swin-S 的对齐标签，发现其与原始标签不同，使用原始比例作为训练目标可能产生错误信号，TL-Align 可纠正标签。
  
  图5. 在 DeiT-S 和 Swin-S 上的结果可视化。我们对输入图像、混合图像、原始标签嵌入以及经过 token-标签对齐后的标签嵌入进行了可视化展示。

结论-Conclusion

这部分内容总结了 TL-Align 方法的研究成果，并对未来研究方向进行了展望，具体如下：

研究成果总结：本文提出的 Token-标签对齐（TL-Align）方法有效解决了视觉 Transformer（ViTs）训练中由于 token 波动导致的训练信号不准确问题。数据混合策略虽能提升 CNNs 和 ViTs 性能，但 ViTs 中的 token 波动现象限制了其效果。TL-Align 通过追踪 token 对应关系获取准确训练信号，实验表明该方法能持续提升多种 ViT 模型在不同任务上的性能。
未来研究方向展望：TL-Align 在其他架构（如 MLP-like 模型）上的泛化性能尚不清楚，这为后续研究提供了一个具有潜力的探索方向。