AI/CV重磅干貨，第一時(shí)間送達(dá)

CVer

一個(gè)專注侃侃計(jì)算機(jī)視覺方向的公眾號(hào)。計(jì)算機(jī)視覺、圖像處理、機(jī)器學(xué)習(xí)、深度學(xué)習(xí)、C/C++、Python、詩和遠(yuǎn)方等。

204篇原創(chuàng)內(nèi)容

公眾號(hào)

LeViT: a Vision Transformer in ConvNet’s Clothing for Faster Inference

論文：https:///abs/2104.01136

代碼（剛剛開源）：

https://github.com/facebookresearch/LeViT

吸取CNN優(yōu)點(diǎn)！LeViT：快速推理的視覺Transformer,在速度/準(zhǔn)確性的權(quán)衡方面LeViT明顯優(yōu)于現(xiàn)有的CNN和視覺Transformer，比如ViT、DeiT等，而且top-1精度為80％的情況下LeViT比CPU上的EfficientNet快3.3倍。
作者單位：Facebook

1 簡介

本文的工作利用了基于注意力體系結(jié)構(gòu)中的最新發(fā)現(xiàn)，該體系結(jié)構(gòu)在高度并行處理硬件上具有競爭力。作者從卷積神經(jīng)網(wǎng)絡(luò)的大量文獻(xiàn)中重新評(píng)估了原理，以將其應(yīng)用于Transformer，尤其是分辨率降低的激活圖。同時(shí)作者還介紹了Attention bias，一種將位置信息集成到視覺Transformer中的新方法。

最終作者提出了LeVIT：一種用于快速推理的混合神經(jīng)網(wǎng)絡(luò)?？紤]在不同的硬件平臺(tái)上采用不同的效率衡量標(biāo)準(zhǔn)，以最好地反映各種應(yīng)用場景。作者通過廣泛的實(shí)驗(yàn)表明該方法適用于大多數(shù)體系結(jié)構(gòu)。總體而言，在速度/準(zhǔn)確性的權(quán)衡方面，LeViT明顯優(yōu)于現(xiàn)有的卷積網(wǎng)絡(luò)和視覺Transformer。例如，在ImageNet Top-1精度為80％的情況下，LeViT比CPU上的EfficientNet快3.3倍。

相同計(jì)算復(fù)雜度的情況下Transformer為什么快？

大多數(shù)硬件加速器(gpu，TPUs)被優(yōu)化以用來執(zhí)行大型矩陣乘法。在Transformer中，注意力機(jī)制和MLP塊主要依靠這些操作。相比之下，卷積需要復(fù)雜的數(shù)據(jù)訪問模式，因此它們的操作通常受io約束。這些考慮對(duì)于我們探索速度/精度的權(quán)衡是很重要的。

本文主要貢獻(xiàn)：

1. 采用注意力機(jī)制作為下采樣機(jī)制的multi-stage transformer 結(jié)構(gòu);

2. 一種計(jì)算效率高的patch descriptor，可以減少第一層特征的數(shù)量;

3. 使用Translation-invariant attention bias取代ViT中的位置嵌入;

4. 為了提高給定計(jì)算時(shí)間的網(wǎng)絡(luò)容量，作者重新設(shè)計(jì)了Attention-MLP Block。

2 LeViT的設(shè)計(jì)

2.1 LeViT設(shè)計(jì)原則

LeViT以ViT的架構(gòu)和DeiT的訓(xùn)練方法為基礎(chǔ)，合并了對(duì)卷積架構(gòu)有用的組件。第1步是獲得Compatible Representation。如果不考慮classification embedding的作用，ViT就是一個(gè)處理激活映射的Layer的堆疊。

實(shí)際上，中間“Token”嵌入可以看作是FCN體系結(jié)構(gòu)中傳統(tǒng)的C×H×W激活映射(BCHW格式)。因此，適用于激活映射(池、卷積)的操作可以應(yīng)用于DeiT的中間表征。

LeViT優(yōu)化了計(jì)算體系結(jié)構(gòu)，不一定是為了最小化參數(shù)的數(shù)量。ResNet系列比VGG更高效的設(shè)計(jì)原則之一是在其前2個(gè)階段使用相對(duì)較小的計(jì)算預(yù)算應(yīng)用strong resolution reductions。當(dāng)激活映射到達(dá)ResNet的第3階段時(shí)，其分辨率已經(jīng)縮小到足以將卷積應(yīng)用于小的激活映射，從而降低了計(jì)算成本。

2.2 LeViT組件

1、Patch embedding

初步分析表明，在transformer組的輸入上應(yīng)用一個(gè)小卷積可以提高精度。因此在LeViT中作者選擇對(duì)輸入應(yīng)用4層3×3卷積(stride2)來降低分辨率。channel的數(shù)量是C=3,32,64,128,256。

以上操作減少了對(duì)transformer下層的激活映射的輸入，同時(shí)不丟失重要信息。LeViT-256的patch extractor用184 MFLOPs將圖像形狀(3,224,224)轉(zhuǎn)換為(256,14,14)。作為比較，ResNet-18的前10層使用1042 MFLOPs執(zhí)行相同的dimensionality reduction。

為什么在transformer組的輸入上應(yīng)用一個(gè)小卷積可以提高精度？

2、No classification token

為了使用BCHW張量形式，LeViT刪除了classification token。類似于卷積網(wǎng)絡(luò)，在最后一個(gè)激活映射上使用GAP來代替，這將產(chǎn)生一個(gè)用于分類器的embedding。在訓(xùn)練中進(jìn)行蒸餾，作者分別訓(xùn)練分類和蒸餾的Head。在測試時(shí)，平均2個(gè)Head的輸出。在實(shí)踐中，LeViT可以使用BNC或BCHW張量格式。

3、Normalization layers and activations

ViT架構(gòu)中的FC層相當(dāng)于1x1卷積。ViT在每個(gè)注意點(diǎn)和MLP單元之前使用層歸一化。對(duì)于LeViT，每次卷積之后都要進(jìn)行BN操作。然后與residual connection連接起來的每個(gè)BN權(quán)重參數(shù)初始化為零。BN可以與之前的卷積合并來進(jìn)行推理，這比層歸一化有運(yùn)行優(yōu)勢(例如，在EfficientNet B0上，這種融合將GPU的推理速度提高了2倍)。而DeiT使用GELU函數(shù)，而LeViT的非線性激活都是Hardswish。

4、Multi-resolution pyramid

LeViT在transformer架構(gòu)中集成了ResNet stage。在各個(gè)stage中，該體系結(jié)構(gòu)類似于一個(gè)visual transformer:一個(gè)帶有交替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 * (resolution**4) * 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之間，一個(gè)縮小的注意塊減少了激活映射的大小:在Q轉(zhuǎn)換之前應(yīng)用一個(gè)subsampling，然后傳播到soft activation的輸出。這將一個(gè)大小為的輸入張量映射到一個(gè)大小為的輸出張量。由于尺度的變化這個(gè)注意塊的使用沒有殘差連接。同時(shí)為了防止信息丟失，這里將注意力頭的數(shù)量設(shè)為。

6、Attention bias instead of a positional embedding

在transformer架構(gòu)中的位置嵌入是一個(gè)位置依賴可訓(xùn)練的向量，在將token嵌入輸入到transformer塊之前，將其添加到token嵌入。如果沒有它，轉(zhuǎn)換器輸出將獨(dú)立于輸入標(biāo)記的排列。位置嵌入的Ablations會(huì)導(dǎo)致分類精度的急劇下降。

然而，位置嵌入只包含在注意塊序列的輸入上。因此，由于位置編碼對(duì)higher layer也很重要，所以它很可能仍然處于中間表示中。

因此，LeViT在每個(gè)注意塊中提供位置信息，并在注意機(jī)制中明確地注入相對(duì)位置信息:只是在注意力圖中添加了注意偏向。對(duì)于每個(gè)head ，每2個(gè)像素和之間的標(biāo)量值計(jì)算方式為：

第一項(xiàng)是經(jīng)典的注意力。第二個(gè)是translation-invariant attention bias。每個(gè)Head有H×W參數(shù)對(duì)應(yīng)不同的像素偏移量。對(duì)稱差異和鼓勵(lì)用 flip invariance進(jìn)行訓(xùn)練。

self.attention_biases = torch.nn.Parameter(
            torch.zeros(num_heads, len(attention_offsets)))

7、Smaller keys

由于translation-invariant attention bias偏置項(xiàng)減少了key對(duì)位置信息編碼的壓力，因此LeViT減少了key矩陣相對(duì)于V矩陣的大小。如果key大小為, V則有2D通道。key的大小可以減少計(jì)算key product 所需的時(shí)間。

對(duì)于沒有殘差連接的下采樣層，將V的維數(shù)設(shè)置為4D，以防止信息丟失。

8、Attention activation

在使用常規(guī)線性投影組合不同Heads的輸出之前，對(duì)product 應(yīng)用Hardswish激活。這類似于ResNet bottleneck residual block，V是一個(gè)1×1卷積的輸出，對(duì)應(yīng)一個(gè)spatial卷積，projection是另一個(gè)1×1卷積。

9、Reducing the MLP blocks

在ViT中，MLP residual塊是一個(gè)線性層，它將嵌入維數(shù)增加了4倍，然后用一個(gè)非線性將其減小到原來的嵌入維數(shù)。但是對(duì)于視覺架構(gòu)，MLP通常在運(yùn)行時(shí)間和參數(shù)方面比注意Block更昂貴。

對(duì)于LeViT, MLP是1x1卷積，然后是通常的BN。為了減少計(jì)算開銷，將卷積的展開因子從4降低到2。一個(gè)設(shè)計(jì)目標(biāo)是注意力和MLP塊消耗大約相同數(shù)量的FLOPs。

2.3 LeViT家族

3 實(shí)驗(yàn)

3.1 速度對(duì)比

ResNet50的精度，但是是起飛的速度。

3.2 SOTA對(duì)比