self-attention 也许并不是 VIT 取得成功的关键？ShiftViT：极简 Shift 操作顶替上岗...

作者 | 科技猛兽  编辑 | 极市平台

点击下方卡片，关注“自动驾驶之心”公众号

ADAS巨卷干货，即可获取

点击进入→自动驾驶之心技术交流群

后台回复【数据集下载】获取计算机视觉近30种数据集！

导读

 

本文提出了一个类似 ViT 的基础模型 ShiftViT，其中 Self-Attention 机制被一个极其简单的移位操作所取代。ShiftViT 可以获得比 Swin Transformer 更好的性能。

本文目录

1 ShiftViT：当移位操作遇到视觉 Transformer
(来自中国科学技术大学，MSRA)
1 ShiftViT 论文解读
1.1 背景和动机
1.2 Shift 操作
1.3 Shift 块
1.4 ShiftViT 架构
1.5 实验结果

1 ShiftViT：当移位操作遇到视觉 Transformer

论文名称：When Shift Operation Meets Vision Transformer: An Extremely Simple Alternative to Attention Mechanism (AAAI 2022)

论文地址：

https://arxiv.org/pdf/2201.10801.pdf

1.1 背景和动机

注意力机制 (Self-Attention) 被广泛认为是视觉 Transformer 模型成功的关键，因为它提供了一种灵活而强大的空间关系建模方法。然而，注意力机制真的是 ViT 模型不可或缺的一部分吗？它能被其他替代品取代吗？

尽管最近 ViT 变体模型的表现令人印象深刻，但仍然不清楚是什么使 ViT 模型真正 work 的。一些传统观点倾向于将这种 work 的机制归功于 Self-Attention，因为它提供了一种灵活而强大的方式来建模空间关系。与 CNN 中的卷积运算相比，Self-Attention 有两个显著的优点。其一，长短距离建模： 这种机制提供了同时捕获短距离依赖和长距离依赖的可能性，并摆脱了卷积的局部性限制。其二，动态权重： 两个空间位置之间的动态注意力值依赖于它们自身的特征，而不是一个固定，可学习的卷积核。由于这些良好的特性，很长一段时间人们认为是 Self-Attention 机制给 ViT 模型带来了强大的表征能力。

然而，这两个优势真的是成功的关键吗？答案很可能是否定的。

其一，长短距离建模的必要性： 一些现有的工作证明，即使没有这些特性，ViT 变体仍然可以很好地工作。对于第一个问题，完全的全局依赖关系可能不是不可避免的。越来越多的工作像早期的 Swin Transformer 通过引入局部注意机制获得了更好的性能。

其二，动态权重的必要性： MLP-Mixer 提出用线性投影层代替注意层，其中线性权重不是动态生成的，而是固定，可学习的。在这种情况下，它仍然可以在 ImageNet 上达到不错的性能。

根据以上论据，这两个优势都不是成功的关键。那么 ViT 成功的根本原因是什么呢？为了解决这个问题，为了揭开注意力机制的神秘角色，本文将 Self-Attention 机制简化为一个极其简单的例子：没有全局感受野，没有动态权重，没有 FLOPs 甚至没有任何参数量的移位操作 (Shift Operation)，想知道 ViT 在这种极端的情况下能否保持良好的性能。

1.2 Shift 操作

移位操作在计算机视觉中并不新鲜。早在2017年，就有人提出它是空间卷积运算的一种有效替代方法[1]。具体而言，它使用一个三明治式的体系结构，两个  卷积和一个 Shift 操作，来近似一个  卷积。

ShiftViT 还可以归入纯 MLP 类结构，其中移位操作被视为一种极其特殊的 Token Mixer。与现有的 MLP 工作相比，Shift 操作甚至简单得多，因为它不仅没有全局感受野，没有动态权重，甚至没有计算量和参数量。此外，由于权重是固定大小，MLP-Mixer 无法处理可变的输入大小。然而 ShiftViT 可以，因此使基础模型适用于更多的视觉任务，如物体检测和语义分割。

1.3 Shift 块

如下图 1(b) 所示为一个基本的 Shift 块架构。可以看到 Shift 块由3个按顺序堆叠的组件构成：Shift 操作，Layer Norm，MLP 块。

图1：ShiftViT 和 Shift 块的架构

给定一个输入张量，它的一部分通道将沿左、右、上、下4个空间方向移动，其余通道保持不变。移动后，超出范围的像素被丢弃，剩下空的像素被填充为零。其中 Shift 操作的步长设置为1个像素。

具体而言, 假设输入特征  的形状是 , 其中  是通道数,  和  为高度和宽度。输出特征  具与输入特征  的形状相同，可以写成：

其中，  是一个比率因子，控制多少百分比的 channel 被移位，其值通常被设置为 1/121/12 。

值得注意的是，移位操作不包含任何参数和计算，唯一需要的操作是内存的复制，因此移位操作的效率很高，非常容易实现，对 TensorRT 这样的深度学习推理库更友好。其 PyTorch 的伪代码如下所示。

def shift(feat, gamma=1/12):
    # feat is a tensor with a shape of 
    # [Batch, Channel, Height, Width]
    B, C, H, W = feat.shape
    g = int(gamma * C)
    out = zeros_like(feat)
    # spatially shift
    out[:, 0*g:1*g, :, :-1] = x[:, 0*g:1*g, :, 1:]  
    out[:, 1*g:2*g, :, 1:] = x[:, 1*g:2*g, :, :-1]
    out[:, 2*g:3*g, :-1, :] = x[:, 2*g:3*g, 1:, :]
    out[:, 3*g:4*g, 1:, :] = x[:, 3*g:4*g, :-1, :]
    # remaining channels
    out[:, 4*g:, :, :] = x[:, 4*g:, :, :]
    return out

1.4 ShiftViT 架构

为了与基线 Swin Transformer 进行公平的比较，作者设计了不同参数数量和计算复杂度的模型。基本通道数 CC 和不同 Stage 的 Block 数的详细配置如下：

• Shift-T :

• Shift-S :

• Shift-B :

在 Shift-ViT 中，几乎所有的参数都集中在 MLP 的部分。因此可以通过控制 MLP 的 Expansion Ratio 改变模型的深浅：Expansion Ratio 设置大一点，模型就得变浅，Expansion Ratio 设置小一点，模型可以适当加深。如果未指定，则将 Expansion Ratio 设置为2。

1.5 实验结果

ImageNet-1K 图像分类

图2：ShiftViT 与 Swin Transformer 的性能对比

对于图像分类任务，作者完全按照 Swin Transformer 中的设置。在骨干网络之后附加了平均池化层 (average pooling layer) 和分类头。所有参数都被随机初始化，并使用 AdamW 优化器训练300个 Epoch。学习率从0.001开始，随着余弦计划逐渐衰减到0。所有数据增强和正则化技巧都遵循 Swin Transformer，Batch Size 大小设置为1024。

因为这项工作的目标是探索注意力机制是否可以被一个极其简单的 Shift 操作所取代，所以作者将 Swin Transformer 作为基线模型，并从构建 ShiftViT 的轻量级版本开始。ShiftViT/light 与 Swin Transformer 几乎相同，只是注意层被 Shift 操作所取代。用 Shift 操作代替注意力会导致参数和 FLOPs 的减少。实验结果如下图2所示，结果表明，移位操作不包含任何可学习的参数和算术计算，因此比注意机制弱。例如，Shift-T/light 只有 20M 参数和 3.0 FLOPs，比 Swin-T 精度低了近 33%。但值得注意的是，轻量级 Shift-T/light 的绝对精度并不差。与典型的ResNet-50 骨干网相比，Shift-T/light 功能更强大，效率更高。

为了弥补 Shift 操作和注意力机制之间的复杂性差距，作者采用更多的构建块，以确保它具有与 Swin 基线相似的参数数量。在这样的公平比较中，ShiftViT 取得了比 Swin Transformer 更好的结果。虽然 Shift 操作在空间建模中弱于注意力机制，但其简单的架构使网络能够向纵深发展。这样就大大缓解了换挡操作的弱点。在相同的计算预算下，ShiftViT 的整体性能与基于注意力的 Swin Transformer 相当。这些实验证明，注意力机制可能不是 ViT work 的关键，即使是极其简单的操作也能达到类似的效果。

COCO 目标检测和 ADE20K 语义分割

实验结果如上图2所示，结果表明，在小模型中，ShiftViT 骨干网络在 COCO 上的 mAP 为 45.4%，在 ADE20K 上的 mIoU 为 46.3%，明显优于 Swin-T 骨干网络 。对于大模型，ShiftViT 骨干网络的性能似乎已经饱和，但仍然与 Swin 基线相当。

虽然 Shift 操作在空间建模中弱于注意力机制，但是好处是没有参数和计算量，使得模型可以堆得更深。在相似参数和计算量的前提下，ShiftViT 的整体性能与基于注意力的 Swin Transformer 相当。这些实验证明，注意力机制对于 ViT 来说可能是不必要的。即使是极其简单的操作也能达到类似的效果。

与同期骨干模型性能对比

如下图3所示，为了进一步证明有效性，作者将 ShiftViT 骨干模型与现有的几个模型进行了比较。对于 ImageNet-1K 上的图像分类任务，与 CNN、ViT 和 MLP 三种不同类型的模型进行比较。结果详见图4。总的来说，ShiftViT 可以达到与最先进的性能相当的性能。另一件有趣的事情是比较两个并行工作 S2-MLP 和 AS-MLP。这两件作品在移位操作上有相似的想法，但他们在构建块中引入了一些辅助模块，例如投影前层和投影后层。实验表明，ShiftViT 的性能略好于这两个工作，证明了仅用简单的 Shift 操作构建骨干架构就足够好了。

图3：ShiftViT 与其他 MLP 类架构在 ImageNet 上的精度对比

总结

本文提出了一个类似 ViT 的基础模型 ShiftViT，其中 Self-Attention 机制被一个极其简单的移位操作所取代。ShiftViT 可以获得比 Swin Transformer 更好的性能。此外，作者还分析了 ViT 成功的原因，并暗示注意力机制可能不是 ViT 模型在多数视觉任务上真正好用的关键因素。也许在未来 ViT 的有关研究中，研究者更应该重视 ViT 剩余的部分。

参考

1. ^Shift: A Zero FLOP, Zero Parameter Alternative to Spatial Convolutions

往期回顾

史上最全综述 | 3D目标检测算法汇总！（单目/双目/LiDAR/多模态/时序/半弱自监督）

国内首个自动驾驶学习社区

近1000人的交流社区，和20+自动驾驶技术栈学习路线，想要了解更多自动驾驶感知（分类、检测、分割、关键点、车道线、3D目标检测、多传感器融合、目标跟踪、光流估计、轨迹预测）、自动驾驶定位建图（SLAM、高精地图）、自动驾驶规划控制、领域技术方案、AI模型部署落地实战、行业动态、岗位发布，欢迎扫描下方二维码，加入自动驾驶之心知识星球，这是一个真正有干货的地方，与领域大佬交流入门、学习、工作、跳槽上的各类难题，日常分享论文+代码+视频，期待交流！

【自动驾驶之心】全栈技术交流群

自动驾驶之心是首个自动驾驶开发者社区，聚焦目标检测、语义分割、全景分割、实例分割、关键点检测、车道线、目标跟踪、3D目标检测、BEV感知、多传感器融合、SLAM、光流估计、深度估计、轨迹预测、高精地图、NeRF、规划控制、模型部署落地、自动驾驶仿真测试、产品经理、硬件配置、AI求职交流等方向；

添加汽车人助理微信邀请入群

备注：学校/公司+方向+昵称