©PaperWeekly 原创 · 作者 | 国雍
单位 | 马普所
研究方向 | 计算机视觉
视觉 Transformer（Vision Transformer，ViT）在图像分类等任务中表现出色，但在面对常见的图像扰动（如噪声或模糊）时，其性能会显著下降。为此，来自马克思普朗克研究所的研究者对 ViT 的关键组成部分——自注意力机制进行研究分析，发现当前的视觉 transformer 模型在自注意力机制中存 “token overfocusing” 的问题，即注意力机制过度依赖于少数重要 token。然而这些 token 对图像扰动非常敏感。 为提高模型的鲁棒性，研究人员提出了两种通用技术：Token-aware Average Pooling (TAP) 和 Attention Diversification Loss (ADL)。所提出的方法可以应用于大多数 Vit 架构之上，在几乎不增加训练开销的情况下大大提升了分类精度和鲁棒性。此外，这种改进也可以很好地推广到其他下游任务，例如语义分割。本工作发表在 ICCV 2023，为提升 ViT 的鲁棒性提供了有效途径。 论文标题：
Robustifying Token Attention for Vision Transformers
https://arxiv.org/abs/2303.11126
注意力机制中的Token Overfocusing现象 以最近的 FAN 架构为例，本文将 FAN 最后一层的注意力可视化为注意力矩阵。第行表示第个输出 token “关注”哪些输入 token（列）——红色越深表示注意力得分越高。对于多头自注意力则通过对注意力头进行平均来可视化该矩阵。 由下图可观察到，列方向上的注意力通常非常稀疏，这意味着 大多数输入 token 没有被关注到，并且每个输出 token 只关注相同的少数重要输入 token 。 本文将这种现象称为 Token Overfocusing 。然而，当面对高斯噪声等干扰时，会导致模型注意力完全转移到其他不同的 token（图 1，第二列）。 这可以理解为原始 token 没有捕捉到稳定的信息。进一步定量分析发现，同一幅图像的干净版本和加入高斯噪声的版本之间的注意力图的余弦相似度非常低，即同一图像输入，注意力变化非常剧烈。这说明标准的自注意力机制对输入扰动极为不稳定。作者发现这种现象存在于各种架构中，包括 DeiT 和 RVT，并且还出现在语义分割等模型中。 ▲ 图1 通过可视化角度观察模型对抗扰动图像的稳定性（左图为 196×196 的注意力矩阵可视化，右图为干净示例与扰动示例之间的注意力余弦相似度可视化）
解决思路与技术方案 为解决 Token Overfocusing 问题，本文提出了两个通用技术用于提高注意力机制的稳定性和鲁棒性：1） 提出 Token-aware Average Pooling（TAP） ，通过学习每个 token 的 pooling 区域，让更多 token 参与注意力计算；2） 提出 Attention Diversification Loss（ADL） ，最大化不同 token 间的注意力向量差异，提高多样性。 1. Token-aware Average Pooling鼓励更多token参与
该方法试图鼓励更多的输入 token 参与自注意力机制，即在注意力图中获得更多具有高得分的列。为此，本文鼓励每个输入 token 从其局部邻域显式地聚合有用的信息，以防该 token 本身不包含重要信息。 如下表所示，在自注意力之前引入任何局部聚合都能一致地提高鲁棒性（参见表 1 最后一列）。这些方法对所有 token 应用了固定的卷积核或池化区域。然而，token 通常彼此不同，每个 token 都应该需要一个特定的局部聚合策略。因此本文采用 自适应方式选择正确的邻域大小和聚合策略。
▲ 表1 基于FAN-B-Hybrid的局部聚合方法比较
基于上述想法，本文提出 Token-aware Average Pooling（TAP）模块，使每个 token 能够选择适当的区域进行局部聚合。具体而言，TAP 对每个 token 进行平均池化，并自适应地调整池化区域。 如下图所示，TAP 利用多分支结构，在多个分支上进行加权求和，每个分支具有特定的池化区域。与简单地改变卷积核大小类似，TAP 通过改变扩张率来调整池化区域。这背后的主要观察是，没有扩张的大卷积核的平均池化将导致相邻池化区域之间的重叠非常大，从而导致输出 token 中的严重冗余。例如，在表 1 中可以看到，AvgPool5x5 会导致准确率约 1.2% 的大幅下降。
▲ 图2 提出的 Token-aware Average Pooling 模块（左）和整体架构（右）
基于这些观察，本文基于具有不同扩张率的平均池化构建了 TAP：给定 个分支，考虑范围 内的扩张率。其中， 表示恒等映射，即没有任何计算，即没有局部聚合。由 K 确定的最大扩张率为超参数。TAP 包含一个轻量级的扩张率预测器，用于预测在允许的扩张率范围内要使用哪种扩张率（即图2中的哪个分支）。该预测器非常高效，由于它将特征维度从 C 降低到图2中的 ，因此只增加了最小的计算开销和模型参数。
2. Attention Diversification Loss提高注意力多样性
该方法旨在改善输出 token 之间的注意力多样性，即鼓励图 1 中的不同行对应不同的输入 token。基于该目标，本文提出了一种注意力多样化损失（ADL），以减小不同输出 token（行）之间的注意力的余弦相似度。 例如，当两行（即输出 token）具有非常不相交的注意力模式时，期望其余弦相似度接近 0，即较低的相似度。然而，即使对于没有被关注的 token，注意力得分也不会为零。对于较大的 N，计算点积并将这些值相加往往会导致余弦相似度显著大于零。 设置阈值过滤注意力小值，关注最重要的值，使相似度计算更准确： 设 为指示函数， 为第 层中第 个 token（行）的注意力向量。引入一个依赖于 token 数量 的阈值 ，即 。因此，经阈值处理后的注意力变为
通 过近似计算相似度来降低计算量：为避免计算 N 个行之间的配对余弦相似度的二次复杂度，本文通过计算每个单独的注意力向 量 与平均注意力 ≔ 之间的余弦相似度来近似计算。 在考虑具有 L 层的模型时，本文通过以下方式对所有层的 ADL 损失进行平均：
在实践中，本文将 ADL 与标准的交叉熵损失相结合，并引入一个超参数 来控制 ADL 的重要性：
实验评估 本文在图像分类和语义分割任务上验证了所提出的方法的有效性。 本文基于两个最先进的鲁棒性架构 RVT 和 FAN（使用‘’Base‘’模型大小，即 RVT-B 和 FAN-B-Hybrid）构建了本文的方法，与多个鲁棒性 baseline 在 ImageNet 图像分类任务上进行了评估。实验证明，无论是单独使用 TAP 或 ADL 还是联合使用两者，本文方法都能有效提高模型的鲁棒性，联合使用两者则可以产生更佳效果。
▲ 表2 将 TAP 和 ADL 应用于不同架构上的实验结果 （1）在 ImageNet 上的对比： 相较于 RVT 和 FAN 两个强 baseline，本文方法在 ImageNet-C 上的鲁棒性提高 >2.1%，并且在其他鲁棒性 Benchmark（包括 ImageNet-A/R/SK）上也有类似的提高。在 ImageNet 上的分类准确率同样得到了提高。 此外，本文还基于 FAN-B-Hybrid 模型分析了本文方法对 ImageNet-C 上不同类型扰动的鲁棒性。从下表可以看出，本文最佳模型（+TAP&DAL）对抗各种噪声扰动尤其有效，特别是在高斯噪声上，其 mCE 下降了 6.25%。这充分证明了方法通过提高注意力机制的鲁棒性，可持续增强不同 Vit 对各类扰动的对抗能力。
▲表3 基于 FAN-B-Hybrid 模型在 ImageNet-C 各个 corruption 类型上的 corruption error 比较
注意力稳定性分析： 下图可视化了当面临图像扰动（例如高斯噪声）时注意力的变化。在各示例中，基准模型存在严重的 token overfocusing 问题，并且在面临扰动时产生了明显的注意力转移。应用 TAP 后，模型将注意力分配给了周围更多的 token，在一定程度上缓解了 token 过度聚焦问题。然而，仍然可观察到干净示例和扰动示例之间的注意力转移。 应用 ADL 训练模型时，注意力呈现对角线模式，类似残差结构，令 token 既保留自身信息又聚合周围信息，大幅提升了相似图像间注意力图的稳定性。当将 TAP 和 ADL 结合在一起时，模型进一步鼓励对角线模式在局部区域内扩展，使得 token 将更多地关注自身之外的邻域，从而获得更强的特征。本文通过计算整个 ImageNet 上干净示例和扰动示例之间的注意力余弦相似度来定量评估注意力的稳定性。结果显示，加入本文方法的模型大幅提升了相似度分数，表明本文方法可有效提高注意力的稳定性。
每个 head 的注意力和注意力多样性分析： 下图展示了每个 head 的注意力图，可见，基准模型各 head 之间注意力高度相似，多样性极低。相比之下，本文方法各 head 可以产生明显不同的注意力模式，只有 2 个 head 呈现对角线模式，其他 head 则具有全局范围的注意力。这种结合局部和全局过滤器的设计提高了 head 间的注意力多样性。定量结果也证实，本文方法可以大幅降低不同 head 间的相似度，即增加多样性。
本文接着验证了提出的方法在语义分割上的泛化能力。在 Cityscapes 数据集上训练模型，并在包含不同扰动类型的 Cityscapes-C 和包含多种不利条件下街景图像的 ACDC 上评估鲁棒性。实验采用与 SegFormer 相似的设置。结果表明，本文的 TAP 和 ADL 技术可以很好的泛化至分割任务，并显著提升了鲁棒性。
定量对比： 如下表所示，在 Cityscapes 和 ACDC 数据集上，单独使用 TAP 或 ADL 即可获得显著提升，联合使用两者效果更佳，在 Cityscapes-C 和 ACDC 上 mIoU 分别提升 2.4% 和 3.1%。同时在干净 Cityscapes 图像上 mIoU 也获得 0.5% 的改进。与主流分割模型相比，本文的最佳模型也展现出了明显优势。这些结果表明，所提出的两种技术不仅适用于图像分类，而且在语义分割任务中也具有很好的泛化性能。
▲ 表4 各语义分割模型在 Cityscapes 验证集、Cityscapes-C 和 ACDC 测试集上的分割 mIoU 对比
视觉对比： 在存在雪天扰动时，基准模型无法检测到部分道路区域。而在夜间条件下，基准模型将汽车的一部分识别为骑车者，并且预测结果存在许多伪影。相比之下，本文方法在这些情况下都展现出更强的鲁棒性，可准确检测关键目标。
分析与讨论 应用于不同架构的有效性： 除了 RVT 和 FAN 之外，本文方法对其他 Transformer 架构也具有很好的适用性。在 DeiT 和 Swin Transformer 上应用本文方法，同样可以获得显著的鲁棒性提升和准确率改进。
▲ 表5 将TAP和ADL应用于不同架构上的实验结果 超参数 K 和 的影响：TAP 模块中的分支数 K 决定了每个 toke n 可以融合信息的池化区域大小种类，K 值越大则可以融合更多样化的局部信息来提高鲁棒性，但过大会增加模型负担。较大的权重 可以增强 ADL 的作用，即更强力地鼓励模型进一步提高注意力多样化，但是 太大会抑制标准训练损失，导致模型训练效果下降。 如下图所示，当只引入 TAP 时，随着 K 值的增加，本文模型始终优于基准模型，并在 时取得最佳结果；当只使用 ADL 时，可以观察到太小或太大的 会降低本文方法的效果。经验证明，设置 和 可以在保证鲁棒性的同时，使额外计算和内存成本最小。
结语 本文针对视觉 Transformer 中的 token overfocusing 问题进行了研究，即其自注意力机制过度依赖少数重要 token，在存在扰动的情况下往往会产生高度分散的注意力模式。为此，本文提出两种通用技术：1）Token-aware Average Pooling 模块通过为每个 token 学习自适应的平均池化方案，使局部邻域 token 参与到自注意力机制中；2）Attention Diversification Loss 明确减小不同 token 间注意力的余弦相似度。
实验表明，本文方法应用到各种网络上，在不同 benchmark 和学习任务上获得了显著的鲁棒性提升。
如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。
总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。
PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读 ，也可以是学术热点剖析 、科研心得 或竞赛经验讲解 等。我们的目的只有一个，让知识真正流动起来。
📝 稿件基本要求：
• 文章确系个人原创作品 ，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注
• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题
• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬 ，具体依据文章阅读量和文章质量阶梯制结算
📬 投稿通道：
• 投稿邮箱： [email protected]
• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者
• 您也可以直接添加小编微信（pwbot02 ）快速投稿，备注：姓名-投稿
△长按添加PaperWeekly小编
🔍
现在，在「知乎」 也能找到我们了
进入知乎首页搜索「PaperWeekly」
点击「关注」 订阅我们的专栏吧
