转载于：ChaucerG
作者：集智书童

吸取CNN优点！LeViT：快速推理的视觉Transformer,在速度/准确性的权衡方面LeViT明显优于现有的CNN和视觉Transformer，比如ViT、DeiT等，而且top-1精度为80％的情况下LeViT比CPU上的EfficientNet快3.3倍。
作者单位：Facebook

1 简介

本文的工作利用了基于注意力体系结构中的最新发现，该体系结构在高度并行处理硬件上具有竞争力。作者从卷积神经网络的大量文献中重新评估了原理，以将其应用于Transformer，尤其是分辨率降低的激活图。同时作者还介绍了Attention bias，一种将位置信息集成到视觉Transformer中的新方法。

图1 LeViT性能对比

最终作者提出了LeVIT：一种用于快速推理的混合神经网络。考虑在不同的硬件平台上采用不同的效率衡量标准，以最好地反映各种应用场景。作者通过广泛的实验表明该方法适用于大多数体系结构。总体而言，在速度/准确性的权衡方面，LeViT明显优于现有的卷积网络和视觉Transformer。例如，在ImageNet Top-1精度为80％的情况下，LeViT比CPU上的EfficientNet快3.3倍。

相同计算复杂度的情况下Transformer为什么快？

大多数硬件加速器(gpu，TPUs)被优化以用来执行大型矩阵乘法。在Transformer中，注意力机制和MLP块主要依靠这些操作。相比之下，卷积需要复杂的数据访问模式，因此它们的操作通常受io约束。这些考虑对于我们探索速度/精度的权衡是很重要的。

本文主要贡献：

采用注意力机制作为下采样机制的multi-stage transformer 结构;
一种计算效率高的patch descriptor，可以减少第一层特征的数量;
使用Translation-invariant attention bias取代ViT中的位置嵌入;
为了提高给定计算时间的网络容量，作者重新设计了Attention-MLP Block。

2 LeViT的设计

2.1 LeViT设计原则

LeViT以ViT的架构和DeiT的训练方法为基础，合并了对卷积架构有用的组件。第1步是获得Compatible Representation。如果不考虑classification embedding的作用，ViT就是一个处理激活映射的Layer的堆叠。

实际上，中间“Token”嵌入可以看作是FCN体系结构中传统的C×H×W激活映射(BCHW格式)。因此，适用于激活映射(池、卷积)的操作可以应用于DeiT的中间表征。

LeViT优化了计算体系结构，不一定是为了最小化参数的数量。ResNet系列比VGG更高效的设计原则之一是在其前2个阶段使用相对较小的计算预算应用strong resolution reductions。当激活映射到达ResNet的第3阶段时，其分辨率已经缩小到足以将卷积应用于小的激活映射，从而降低了计算成本。

2.2 LeViT组件

1、Patch embedding

初步分析表明，在transformer组的输入上应用一个小卷积可以提高精度。因此在LeViT中作者选择对输入应用4层3×3卷积(stride2)来降低分辨率。channel的数量是C=3,32,64,128,256。

以上操作减少了对transformer下层的激活映射的输入，同时不丢失重要信息。LeViT-256的patch extractor用184 MFLOPs将图像形状(3,224,224)转换为(256,14,14)。作为比较，ResNet-18的前10层使用1042 MFLOPs执行相同的dimensionality reduction。

为什么在transformer组的输入上应用一个小卷积可以提高精度？

2、No classification token

为了使用BCHW张量形式，LeViT删除了classification token。类似于卷积网络，在最后一个激活映射上使用GAP来代替，这将产生一个用于分类器的embedding。在训练中进行蒸馏，作者分别训练分类和蒸馏的Head。在测试时，平均2个Head的输出。在实践中，LeViT可以使用BNC或BCHW张量格式。

3、Normalization layers and activations

ViT架构中的FC层相当于1x1卷积。ViT在每个注意点和MLP单元之前使用层归一化。对于LeViT，每次卷积之后都要进行BN操作。然后与residual connection连接起来的每个BN权重参数初始化为零。BN可以与之前的卷积合并来进行推理，这比层归一化有运行优势(例如，在EfficientNet B0上，这种融合将GPU的推理速度提高了2倍)。而DeiT使用GELU函数，而LeViT的非线性激活都是Hardswish。

`class Linear_BN(torch.nn.Sequential):
    def __init__(self, a, b, bn_weight_init=1, resolution=-100000):
        super().__init__()
        self.add_module('c', torch.nn.Linear(a, b, bias=False))
        bn = torch.nn.BatchNorm1d(b)
        torch.nn.init.constant_(bn.weight, bn_weight_init)
        torch.nn.init.constant_(bn.bias, 0)
        self.add_module('bn', bn)

        global FLOPS_COUNTER
        output_points = resolution**2
        FLOPS_COUNTER += a  b  output_points

    @torch.no_grad()
    def fuse(self):
        l, bn = self._modules.values()
        w = bn.weight / (bn.running_var + bn.eps)**0.5
        w = l.weight * w[:, None]
        b = bn.bias - bn.running_mean * bn.weight / \
            (bn.running_var + bn.eps)**0.5
        m = torch.nn.Linear(w.size(1), w.size(0))
        m.weight.data.copy_(w)
        m.bias.data.copy_(b)
        return m

    def forward(self, x):
        l, bn = self._modules.values()
        x = l(x)
        return bn(x.flatten(0, 1)).reshape_as(x)
`

4、Multi-resolution pyramid

LeViT在transformer架构中集成了ResNet stage。在各个stage中，该体系结构类似于一个visual transformer:一个带有交替MLP和激活块的残差模块。下面是注意块的修改。

`class Attention(torch.nn.Module):
    def __init__(self, dim, key_dim, num_heads=8,
                 attn_ratio=4,
                 activation=None,
                 resolution=14):
        super().__init__()
        self.num_heads = num_heads
        self.scale = key_dim ** -0.5
        self.key_dim = key_dim
        self.nh_kd = nh_kd = key_dim * num_heads
        self.d = int(attn_ratio * key_dim)
        self.dh = int(attn_ratio  key_dim)  num_heads
        self.attn_ratio = attn_ratio
        h = self.dh + nh_kd * 2
        self.qkv = Linear_BN(dim, h, resolution=resolution)
        self.proj = torch.nn.Sequential(activation(), Linear_BN(
            self.dh, dim, bn_weight_init=0, resolution=resolution))

        points = list(itertools.product(range(resolution), range(resolution)))
        N = len(points)
        attention_offsets = {}
        idxs = []
        for p1 in points:
            for p2 in points:
                offset = (abs(p1[0] - p2[0]), abs(p1[1] - p2[1]))
                if offset not in attention_offsets:
                    attention_offsets[offset] = len(attention_offsets)
                idxs.append(attention_offsets[offset])
        self.attention_biases = torch.nn.Parameter(
            torch.zeros(num_heads, len(attention_offsets)))
        self.register_buffer('attention_bias_idxs',
                             torch.LongTensor(idxs).view(N, N))

        global FLOPS_COUNTER
        #queries * keys
        FLOPS_COUNTER += num_heads  (resolution4)  key_dim
        # softmax
        FLOPS_COUNTER += num_heads  (resolution*4)
        #attention * v
        FLOPS_COUNTER += num_heads  self.d  (resolution**4)

    @torch.no_grad()
    def train(self, mode=True):
        super().train(mode)
        if mode and hasattr(self, 'ab'):
            del self.ab
        else:
            self.ab = self.attention_biases[:, self.attention_bias_idxs]

    def forward(self, x):  # x (B,N,C)
        B, N, C = x.shape
        qkv = self.qkv(x)
        q, k, v = qkv.view(B, N, self.num_heads, -
                           1).split([self.key_dim, self.key_dim, self.d], dim=3)
        q = q.permute(0, 2, 1, 3)
        k = k.permute(0, 2, 1, 3)
        v = v.permute(0, 2, 1, 3)

        attn = (
            (q @ k.transpose(-2, -1)) * self.scale
            +
            (self.attention_biases[:, self.attention_bias_idxs]
             if self.training else self.ab)
        )
        attn = attn.softmax(dim=-1)
        x = (attn @ v).transpose(1, 2).reshape(B, N, self.dh)
        x = self.proj(x)
        return x
`

5、Downsampling

在LeViT stage之间，一个缩小的注意块减少了激活映射的大小:在Q转换之前应用一个subsampling，然后传播到soft activation的输出。这将一个大小为的输入张量映射到一个大小为的输出张量。由于尺度的变化这个注意块的使用没有残差连接。同时为了防止信息丢失，这里将注意力头的数量设为。

`class Subsample(torch.nn.Module):
    def __init__(self, stride, resolution):
        super().__init__()
        self.stride = stride
        self.resolution = resolution

    def forward(self, x):
        B, N, C = x.shape
        x = x.view(B, self.resolution, self.resolution, C)[
            :, ::self.stride, ::self.stride].reshape(B, -1, C)
        return x
`