NeurIPS 2021 ｜ Twins：重新思考高效的视觉注意力模型设计

Twins 是美团和阿德莱德大学协作提出的视觉留意力模型，相关论文已被 NeurIPS 2021 会议接纳。本文首要讲述 Twins 处理的难点、规划和完成思路，以及在美团场景的探究落地，希望能对从事视觉算法研发的同学有所帮助和启发。

概要

Twins ^[1] 是美团和阿德莱德大学协作提出的视觉留意力模型，相关论文已被 NeurIPS 2021 会议接纳，代码也已在GitHub上进行开源。NeurIPS（Conference on Neural Information Processing Systems）是机器学习和核算神经科学相关的学术会议，也是人工智能方向的世界尖端会议。

Twins 提出了两类结构，分别是 Twins-PCPVT 和 Twins-SVT：

• Twins-PCPVT 将金字塔 Transformer 模型 PVT ^[2] 中的固定方位编码（Positional Encoding）更改为团队在 CPVT ^[3] 中提出的条件式方位编码 （Coditional Position Encoding, CPE），然后使得模型具有平移等变性（即输入图画发生平移后，输出一起相应发生改变），能够灵敏处理来自不同空间标准的特征，然后能够广泛运用于图画切割、检测等变长输入的场景。
• Twins-SVT 提出了空间可分离自留意力机制（Spatially Separable Self-Attention，SSSA）来对图画特征的空间维度进行分组，分别核算各部分空间的自留意力，再利用大局自留意力机制对其进行交融。这种机制在核算上更高效，功能更优。

Twins 系列模型完成简略，布置友好，在 ImageNet 分类、ADE20K 语义切割、COCO 方针检测等多个经典视觉使命中均获得了业界抢先的成果。

布景

2020 年 9 月，谷歌的视觉留意力模型 （Vision Transformer， ViT） ^[4] 成功将原本用于自然语言处理的 Transformer ^[5] 运用到视觉的分类使命中。ViT 将一幅输入图画切分为若干个图画块（Patch），并把一个图画块类比为一个文字（Word）作为 Transformer 编码器的输入（如图 1 所示），经过 L 层的编码器处理后运用一般的多层感知机（Multilayer Perceptron, MLP）映射到类别空间。ViT 的模型功能大幅超过了卷积神经网络，尔后迅速发展成为了当时视觉领域研讨的首要热点。

Transformer 编码器中多头留意力（Multi-head attention）的基本核算办法由下式给定，其间 Q、K、V 分别为 Query（查询）、Key（键）、Value（值） 的缩写，d 为编码维度，softmax 为归一化函数，留意力机制能够理解为对输入按照相关性加权的进程。

原生的视觉留意力模型做主干网络并不能很好地适配方针检测、语义切割等常用的稠密预测使命。此外，比较于卷积神经网络，ViT 核算量通常要更大，推理速度变慢，不利于在实践事务中运用。因而规划更高效的视觉留意力模型，并更好地适配下流使命成为了当下研讨的要点。香港大学、商汤联合提出的金字塔视觉留意力模型 PVT ^[2] 借鉴了卷积神经网络中的图画金字塔范式来生成多标准的特征，这种结构能够和用于稠密使命的现有后端直接结合，支撑多种下流使命，如图 2（c）所示。但因为 PVT 运用了静态且定长的方位编码，经过插值办法来习惯变长输入，不能针对性依据输入特征来编码，因而功能受到了约束。别的，PVT 沿用了 ViT 的大局自留意力机制，核算量仍然较大。

微软亚研院提出的 Swin ^[6] 复用了 PVT 的金字塔结构。在核算自留意力时，运用了对特征进行窗口分组的办法（如图 3 所示），将留意力机制限定在一个个小的窗口（赤色格子），然后经过对窗口进行错位使不同组的信息发生交互。这样能够避免核算大局自留意力而削减核算量，其缺点是丢失了大局的留意力，一起因为窗口错位发生的信息交互才能相对较弱，必定程度上影响了功能。

视觉留意力模型规划的难点

简略总结一下，当时视觉留意力模型规划中需求处理的难点在于：

• 高功率的核算：缩小和卷积神经网络在运算功率上的差距，促进实践事务运用；
• 灵敏的留意力机制：即能够具备卷积的部分感触野和自留意力的大局感触野才能，兼二者之长；
• 利于下流使命：支撑检测、切割等下流使命，尤其是输入标准改变的场景。

Twins 模型规划

从这些难点问题动身，基于对当时视觉留意力模型的详尽剖析，美团视觉智能部重新考虑了自留意力机制的规划思路，提出了针对性的处理方案。首先将 PVT ^[2] 和 CPVT ^[4] 相结合，构成 Twins-PCPVT 来支撑标准改变场景的下流使命。再从自留意机制的功率和感触野角度动身，规划了兼容部分和大局感触野的新型自留意力，叫做空间可分离自留意力 （Spatially Separable Self-Attention，SSSA）， 构成了 Twins-SVT。

Twins-PCPVT

Twins-PCPVT 经过将 PVT 中的方位编码（和 DeiT ^[7] 相同固定长度、可学习的方位编码）替换为 CPVT ^[4] 中的条件方位编码 （Conditional Positional Encodings，CPE）。生成 CPE 的模块叫做方位编码器（Positional Encoding Generator， PEG），PEG 在 Twins 模型中的具体方位是在每个阶段的第 1 个 Transformer Encoder 之后，如下图 4 所示：

条件方位编码

下图 5 展现了团队在 CPVT ^[4] 中提出的条件方位编码器的编码进程。首先将 $N\ast d$ 的输入序列转为 $H\ast W\ast d$ 的输入特征，再用 $F$ 依据输入进行条件式的方位编码，而且输出标准和输入特征相同，因而能够转为 $N\ast d$ 序列和输入特征进行逐元素的加法交融。

其间，编码函数 $F$ 能够由简略的深度可分离卷积完成或者其他模块完成，PEG 部分的简化代码如下。其间输入 feat_token 为形状为 $B\ast N\ast d$ 的张量，$B$ 为 batch，$N$ 为 token 个数，$C$ 为编码维度（同图 5 中 $d$）。将 feat_token 转化为 $B\ast d\ast H\ast W$ 的张量 cnn_feat 后，经过深度可分离卷积 (PEG) 运算，生成和输入 feat_token 相同形状的张量，即条件式的方位编码。

class PEG(nn.Module):
    def __init__(self, in_chans, embed_dim):
        super(PEG, self).__init__()
        self.peg = nn.Conv2d(in_chans, embed_dim, 3, 1, 1, bias=True, groups=embed_dim)
    def forward(self, feat_token, H, W):
        B, N, C = feat_token.shape
        cnn_feat = feat_token.transpose(1, 2).view(B, C, H, W) 
        x = self.peg(cnn_feat) + cnn_feat
        x = x.flatten(2).transpose(1, 2)
        return x

因为条件方位编码 CPE 是依据输入生成，支撑可变长输入，使得 Twins 能够灵敏处理来自不同空间标准的特征。别的 PEG 选用卷积完成，因而 Twins 一起保留了其平移等变性，这个性质对于图画使命非常重要，如检测使命中方针发生偏移，检测框需随之偏移。试验标明 Twins-PCPVT 系列模型在分类和下流使命，尤其是在稠密使命上能够直接获得功能提高。该架构说明 PVT 在只是经过 CPVT 的条件方位编码增强后就能够获得很不错的功能，由此说明 PVT 运用的方位编码约束了其功能发挥。

Twins-SVT

Twins-SVT （如下图 6 所示）对大局留意力策略进行了优化改进。大局留意力策略的核算量会跟着图画的分辨率成二次方增加，因而如安在不明显丢失功能的情况下降低核算量也是一个研讨热点。Twins-SVT 提出新的交融了部分-大局留意力的机制，能够类比于卷积神经网络中的深度可分离卷积 （Depthwise Separable Convolution），并因而命名为空间可分离自留意力（Spatially Separable Self-Attention，SSSA）。与深度可分离卷积不同的是，Twins-SVT 提出的空间可分离自留意力（如下图 7 所示）是对特征的空间维度进行分组，并核算各组内的自留意力，再从大局对分组留意力成果进行交融。

空间可分离自留意力选用部分-大局自留意力（LSA-GSA）相互替换的机制，分组核算的部分留意力能够高效地传导到大局。LSA 能够大幅降低核算成本，复杂度从输入的平方 O(H2W2d)O(H^2W^2d)O(H2W2d) 降为线性的 $O(mnHWd)$。其间分组部分留意力 LSA 要害完成（初始化函数略）如下：

class LSA(nn.Module):
    def forward(self, x, H, W):
        B, N, C = x.shape
        h_group, w_group = H // self.ws, W // self.ws # 依据窗口大小核算长（H）和宽（W）维度的分组个数
        total_groups = h_group * w_group
        x = x.reshape(B, h_group, self.ws, w_group, self.ws, C).transpose(2, 3) # 将输入依据窗口进行分组 B* h_group * ws * w_group * ws * C
        qkv = self.qkv(x).reshape(B, total_groups, -1, 3, self.num_heads, C // self.num_heads).permute(3, 0, 1, 4, 2, 5) # 核算各组的 q， k， v 
        q, k, v = qkv[0], qkv[1], qkv[2]
        attn = (q @ k.transpose(-2, -1)) * self.scale # 核算各组的留意力
        attn = attn.softmax(dim=-1) # 留意力归一化
        attn = self.attn_drop(attn) # 留意力 Dropout 层
        attn = (attn @ v).transpose(2, 3).reshape(B, h_group, w_group, self.ws, self.ws, C) # 用各组内的部分自留意力给 v 进行加权
        x = attn.transpose(2, 3).reshape(B, N, C)
        x = self.proj(x) # MLP 层
        x = self.proj_drop(x) # Dropout 层
        return x

高效交融 LSA 留意力的 GSA 要害完成（初始化函数略）如下。比较于 ViT 原始的大局自留意力，GSA 的 K、V 是在缩小特征的基础上核算的，但 Q 是大局的，因而留意力仍然能够恢复到大局。这种做法明显削减了核算量。

class GSA(nn.Module):
    def forward(self, x, H, W):
        B, N, C = x.shape
        q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) # 依据输入特征 x 核算查询张量 q
        x_ = x.permute(0, 2, 1).reshape(B, C, H, W)
        x_ = self.sr(x_).reshape(B, C, -1).permute(0, 2, 1) # 缩小输入特征的标准得到 x_
        x_ = self.norm(x_) # 层归一化 LayerNorm
        kv = self.kv(x_).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) # 依据缩小标准后的特征后 x_，核算 k, v 
        k, v = kv[0], kv[1]
        attn = (q @ k.transpose(-2, -1)) * self.scale # 核算大局自留意力
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)
        x = (attn @ v).transpose(1, 2).reshape(B, N, C) # 依据大局自留意力对 v 加权
        x = self.proj(x)
        x = self.proj_drop(x)
        return x

从上述代码中能够看出，SVT 系列在完成上选用现有主流深度学习结构中的已有操作，不需求额定的底层适配，因而布置起来比较方便。

试验

ImageNet-1k 分类

Twins-PCPVT 和 Twins-SVT 在 ImageNet-1k 分类使命上，比较同等量级模型均获得 SOTA 成果，吞吐率占优。别的，Twins 支撑 TensorRT 布置，Twins-SVT-S 模型运用 NVIDIA TensorRT 7.0 推理能够有 1.6 倍的加速，吞吐率能够从 PyTorch 完成的 1059（images/s）提高到 1732。

ADE20K 切割

在语义切割使命 ADE20K 上，Twins 模型做主干网分别运用 FPN 和 Upernet 后端，比较 PVT 和 Swin 也达到了更好成果，见下表 2：

COCO 方针检测（RetinaNet 结构）

在经典的 COCO 方针检测使命中，运用 RetinaNet 结构，Twins 模型大幅优于 PVT。而且 Twins-PCPVT 系列证明 PVT 在经过 CPVT 的编码办法增强之后，能够媲美 Swin 同量级模型，见下表 3：

COCO 方针检测（Mask-RCNN 结构）

在 Mask-RCNN 结构下，Twins 模型在 COCO 上也有很好的功能优势，且在更长时间训练（3x）时得以保持，见下表 4：

在高精地图多要素语义切割场景的运用

高精地图是主动驾驶中的要害组成部分，在美团无人配送、网约车等事务承担着非常重要的效果。路途场景要害要素的语义提取作为高精建图的前序流程，对建图的质量有直接的影响。多要素语义切割是语义提取的重要一环，业界一般选用经典的语义切割算法来完成。

此处，咱们以 DeepLab 系列 ^[8] 为代表做介绍，切割模型通常分为编码和解码两个阶段，运用卷积神经网络提取特征，并选用空间金字塔迟化（Spatial Pyramid Pooling），以及不同标准空洞卷积（Atrous Conv）操作（如下图 8a 所示）来增加大局感触野。这种规划一方面受限于卷积神经网络的特征提取才能，另一方面临大局联系的建模才能有限，导致在切割使命上对细节重视不行，边缘往往不行清晰。

Twins 虽然大幅提高了视觉留意力模型的功率和功能，但为了保持和卷积神经网络的挨近的推理功率，咱们仍需求对模型的后端结构作进一步的优化。不同于论文中为了与其他办法做公正对比而运用的较重的 FPN ^[9] 或 UperNet ^[10] 后端，咱们规划了如下图 9 所示的简略轻量的后端，并在事务数据集上的功能和推理速度之间获得了很好的平衡。这个后端是依据 Twins 的特性而规划，因为 Twins 兼顾了大局和部分两种留意力，因尔后端无需选用复杂的规划来增大感触野，只经过各标准特征的线性改变和缩放，就直接恢复到相同标准并进行拼接（Concat），简略维度变换后就能够输出切割成果。

从下图 10 的切割成果对比看，Twins 为主干网的模型能够提获得到更精密的图画边缘，如隔离带、路途标牌、路灯杆等要害路途要素和标注真值（Ground Truth）之间的差异更小。

总结

视觉留意力模型是当时视觉领域的研讨要点，而且已经在各类视觉使命上展现了比较经典卷积神经网络的优越性，但在功率上仍需精心优化，在效果上也需求继续提高。探究规划更高功率的留意力模型并促进前沿的视觉研讨转向工业落地，对美团事务也具有重要的含义。

此次，美团和阿德莱德大学协作规划的 Twins 系列模型架构，有效降低了核算成本，提高了模型功能，更好地支撑了如检测和切割等稠密使命。此外，咱们将 Twins 运用在美团高精地图的要素语义切割场景中，带来了更精密的切割成果，提高了高精地图的建图质量。后续，视觉团队将继续探究高效的视觉留意力模型规划，并期望在美团更广泛的事务场景中得到实践和运用。

参考文献

• [1] Twins: Revisiting the Design of Spatial Attention in Vision Transformers
• [2] Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions
• [3] Conditional Positional Encodings for Vision Transformers
• [4] An image is worth 16×16 words: Transformers for image recognition at scale
• [5] Attention Is All You Need
• [6] Swin Transformer: Hierarchical Vision Transformer using Shifted Windows
• [7] Training data-efficient image transformers & distillation through attention
• [8] Encoder-Decoder with Atrous Separable Convolution for Semantic Image Segmentation
• [9] Panoptic Feature Pyramid Networks
• [10] Unified Perceptual Parsing for Scene Understanding

作者简介

祥祥、田值、张勃、晓林在视觉智能部，海兵、华夏在主动车配送部。

团队简介及招聘信息

美团视觉智能部 AutoML 算法团队旨在经过 AutoML 及前沿的视觉技能赋能公司各项事务、加速算法落地，涵盖 AutoML、切割、检测（2D、3D）、Self-training 等技能方向。欢迎感兴趣的校招和社招同学发送简历至：chuxiangxiang@meituan.com，实习、正式均可。

美团主动车配送部高精地图团队是高精地图技能研发团队，咱们的职责是为美团主动驾驶供给高精度、高鲜度、大范围的高精地图服务。高精地图是涉及多种学科的归纳技能，不仅需求依托 SLAM、地舆测绘、深度学习、多传感器定位等算法进行高精地图的构建，而且还需求利用大数据技能、高功能核算、高并发服务等进行大规模高精地图处理、存储和查询服务。高精地图团队长期招聘核算机视觉、SLAM、系统开发等专家，感兴趣的同学能够将简历发送至：tech@meituan.com（邮件主题：美团高精地图）。

阅读美团技能团队更多技能文章合集

前端 | 算法 | 后端 | 数据 | 安全 | 运维 | iOS | Android | 测试

| 在大众号菜单栏对话框回复【2021年货】、【2020年货】、【2019年货】、【2018年货】、【2017年货】等要害词，可查看美团技能团队历年技能文章合集。

| 本文系美团技能团队出品，著作权归属美团。欢迎出于分享和交流等非商业意图转载或运用本文内容，敬请注明“内容转载自美团技能团队”。本文未经许可，不得进行商业性转载或者运用。任何商用行为，请发送邮件至tech@meituan.com请求授权。