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*d$ 的输入序列转为 $H*W*d$ 的输入特征，再用 $F$ 根据输入进行条件式的位置编码，而且输出尺寸和输入特征相同，因此可以转为 $N*d$ 序列和输入特征进行逐元素的加法融合。

其中，编码函数 $F$ 可以由简单的深度可分离卷积实现或者其他模块实现，PEG 部分的简化代码如下。其中输入 feat_token 为形状为 $B*N*d$ 的张量，$B$ 为 batch，$N$ 为 token 个数，$C$ 为编码维度（同图 5 中 $d$）。将 feat_token 转化为 $B*d*H*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(H^2W^2d)$ 降为线性的 $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)