MobileViT v2（Apple）论文与代码解析

paper：Separable Self-attention for Mobile Vision Transformers

official implementation：https://github.com/apple/ml-cvnets

third-party implementation：https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/mobilevit.py

存在的问题

MobileViT在多个移动视觉任务中取得了SOTA的性能，虽然其参数较少，但和基于卷积的网络相比，延迟较高。MobileViT的主要效率瓶颈在于Transformer中的multi-head self-attention，它需要 \(O(k^2)\) 的时间复杂度（\(k\) 是token或patch的数量）。此外，MHA需要昂贵的操作比如batch-wise matrix mulplication来计算self-attention，从而影响资源受限设备上的延迟。

创新点

针对多头注意力高延迟和高计算成本的问题，本文提出了一种新的分离自注意力机制，通过逐元素操作计算自注意力，大大减少了计算复杂度和成本，将计算复杂度降至 \(O(k)\)，使其更适合在资源受限的设备上运行。

在MobileViT的基础上，结合本文提出的separable self-attention，作者设计了一个新的网络MobileViT v2。通过实验，作者证明了MobileViT v2在多个任务（如ImageNet分类、MS-COCO目标检测、PASCAL VOC分割）上比MobileViT有显著的性能提升。 例如，MobileViTv2在ImageNet数据集上取得了75.6%的Top-1准确率，比MobileViT高出约1%，且在移动设备上的运行速度快了3.2倍。

方法介绍

分离自注意力和MHA类似，输入 \(\mathbf{x}\) 经过三个分支处理，即input \(\mathcal{I}\)、key \(\mathcal{K}\)、value \(\mathcal{V}\)。输入分支 \(\mathcal{I}\) 用一个权重为 \(W_{\mathbf{I}}\in \mathbb{R}^d\) 的线性层将 \(\mathbf{x}\) 中的每个维度为 \(d\) 的token映射为一个标量。权重 \(W_{\mathbf{I}}\) 作为图4b中的潜在节点latent node \(L\)。这个线性映射是一个内积运算，它计算了潜在节点 \(L\) 和 \(\mathbf{x}\) 之间的距离，得到一个 \(k\) 维的向量（这里的\(k\)是token的数量）。然后这个 \(k\) 维向量通过softmax得到一个上下文得分context score \(\mathbf{c}_{\mathbf{s}}\in \mathbb{R}^k\)。

和transformer计算每个token和所有 \(k\) 个tokens之间的attention（或context）得分不同，本文提出的方法只计算和潜在token \(L\) 之间的context score。这将计算attention（或context）得分的成本从 \(O(k^2)\) 降到了 \(O(k)\)。

context score \(\mathbf{c}_{\mathbf{s}}\) 用来计算一个context向量 \(\mathbf{c}_{\mathbf{v}}\)。具体来说，输入 \(\mathbf{x}\) 首先通过权重为 \(\mathbf{W}_{\mathbf{K}}\in \mathbb{R}^{d\times d}\) 的key分支 \(\mathcal{K}\) 线性映射到一个 \(d\) 维空间得到输出 \(\mathbf{x}_{\mathbf{K}}\in \mathbb{R}^{k\times d}\)。然后用 \(\mathbf{c}_{\mathbf{s}}\) 对 \(\mathbf{x}_{\mathbf{K}}\) 加权求和得到context向量 \(\mathbf{c}_{\mathbf{v}}\in \mathbb{R}^d\)

  

\(\mathbf{x}\) 中的所有token共享 \(\mathbf{c}_{\mathbf{v}}\) 中编码的上下文信息。输入 \(
\mathbf{x}\) 首先通过权重为 \(\mathbf{W}_{\mathbf{V}}\in \mathbb{R}^{d\times d}\) 的value分支 \(\mathcal{V}\) 线性映射到一个 \(d\) 维空间，然后接上一个ReLU激活得到输出 \(\mathbf{x}_{\mathbf{V}}\in \mathbb{R}^{k\times d}\)。\(\mathbf{c}_{\mathbf{v}}\) 中的上下文信息通过broadcasted element-wise相乘操作传播到 \(\mathbf{x}_{\mathbf{V}}\) 中。得到的输出经过另外一个权重为 \(\mathbf{W}_{\mathbf{O}}\in \mathbb{R}^{d\times d}\) 的线性层得到最终的输出 \(\mathbf{y}\in \mathbb{R}^{k\times d}\)。分离自注意力的完整公式如下

其中 \(*\) 和 \(\sum \) 分别是broadcastable element-wise的相乘和相加操作。

图1将分离自注意力与Transformer和Linformer进行了比较。由于self-attention的时间复杂度没有考虑用于实现这些方法的操作的成本，一些操作可能成为资源受限设备的瓶颈。为了整理理解，除了理论指标外，还测了单个CPU核心上不同 \(k\) 值的模块级延迟。与Transformer和Linformer中的MHA相比，本文提出的分离自注意力快速又高效。

实验结果

作者用分离自注意力替换MobileViT中的MHA得到了MobileViT v2，此外作者还去掉了MobileViT block中的skip-connection和fusion block因为它们对性能的提升很小。和MobileViT设计了XXS、XS、S三种架构不同 ，MobileViT v2采用了一个宽度缩放因子 \(\alpha\in \{0.5,2.0\}\) 来对模型均匀的缩放得到一系列的网络。 

和其它模型在ImageNet上的分类结果如表2所示，可以看到相比于MobileViT v1，v2在相似的参数和下延迟大大降低，例如MobileViTv2-0.5和MobileViT-XXS相比，延迟下降了2/3，精度提高了1.2%。

代码解析

这里是timm中的实现。输入大小为(1, 3, 224, 224)，模型选择"mobilevitv2_100"即mobilevitv2-1.0。分离自注意力的实现在类LinearSelfAttention中，我们来看forward函数，代码如下。

输入shape=(1, 128, 4, 256)，分别表示batch_size、channel、patch_area、patch_num（token_num）。这里和v1中一样采用2x2的patch大小。接下来是self.qkv_proj就是上面提到的 \(\mathcal{I},\mathcal{K},\mathcal{V}\) 分支，输出通道数是2*embed_dim+1，split后分别得到query、key、value，对query取softmax就得到了 \(\mathbf{c}_{\mathbf{s}}\)，key和value分别对应上文的 \(\mathbf{x}_{\mathbf{K}},\mathbf{x}_{\mathbf{V}}\)。24行的加权求和对应上文的式(2)得到 \(\mathbf{c}_{\mathbf{v}}\)，最后对 \(\mathbf{x}_{\mathbf{V}}\) 取ReLU并与 \(\mathbf{c}_{\mathbf{v}}\) 乘对应式(3)得到最终的输出。

self.qkv_proj = nn.Conv2d(
    in_channels=embed_dim,
    out_channels=1 + (2 * embed_dim),
    bias=bias,
    kernel_size=1,
)
def _forward_self_attn(self, x: torch.Tensor) -> torch.Tensor:
    # (1,128,4,256)
    # [B, C, P, N] --> [B, h + 2d, P, N]
    qkv = self.qkv_proj(x)  # Conv2d(128, 257, kernel_size=(1, 1), stride=(1, 1)), (1,257,4,256)

    # Project x into query, key and value
    # Query --> [B, 1, P, N]
    # value, key --> [B, d, P, N]
    query, key, value = qkv.split([1, self.embed_dim, self.embed_dim], dim=1)
    # (1,1,4,256),(1,128,4,256),(1,128,4,256)

    # apply softmax along N dimension
    context_scores = F.softmax(query, dim=-1)  # (1,1,4,256)
    context_scores = self.attn_drop(context_scores)

    # Compute context vector
    # [B, d, P, N] x [B, 1, P, N] -> [B, d, P, N] --> [B, d, P, 1]
    context_vector = (key * context_scores).sum(dim=-1, keepdim=True)  # (1,128,4,256)->(1,128,4,1)

    # combine context vector with values
    # [B, d, P, N] * [B, d, P, 1] --> [B, d, P, N]
    out = F.relu(value) * context_vector.expand_as(value)  # (1,128,4,256)
    out = self.out_proj(out)  # Conv2d(128, 128, kernel_size=(1, 1), stride=(1, 1)), (1,128,4,256)
    out = self.out_drop(out)
    return out