【寫在前面】

近段時間，Transformer-based模型在Visual Recognition領(lǐng)域取得了非常大的進(jìn)展。但是如果不借助額外的訓(xùn)練數(shù)據(jù)，Transformer-based模型離CNN-based模型還是具有一定的差距（NFNet-F5（CNN-based）：86.8%，CaiT（Transformer-based）：86.5%）。作者認(rèn)為，這是因為token embedding并沒有進(jìn)行細(xì)粒度特征表示，因此本文提出了一種新的Attention方式，通過局部信息的感知，能夠獲得更加細(xì)粒度的特征表示。

1. 論文和代碼地址

VOLO: Vision Outlooker for Visual Recognition

論文地址：https://arxiv.org/abs/2106.13112

官方代碼：https://github.com/sail-sg/volo

一看就懂的核心代碼： https://github.com/xmu-xiaoma666/External-Attention-pytorch/blob/master/attention/OutlookAttention.py

2. Motivation

前面說到了，token embedding進(jìn)行特征表示是粗粒度的。為什么是token embedding就是粗粒度的呢？因為一張圖片的size往往是比較大的(e.g., 224x224)，CNN因為只對局部的信息進(jìn)行感知（在正常圖片上進(jìn)行滑動窗口操作），所以計算復(fù)雜度和圖片大小呈線性關(guān)系。如果將圖片進(jìn)行token embedding，直接進(jìn)行Transformer的操作，那就會導(dǎo)致計算復(fù)雜度“爆炸”，因為Transformer中的Self-Attention（SA）的復(fù)雜度與輸入特征的大小是呈平方關(guān)系的。

所以Vision Transformer需要將圖片先進(jìn)行token embedding到比較小的size（e.g., 14x14），才能保證計算量在一個合理的范圍內(nèi)。但這就導(dǎo)致了一個信息損失的問題，因為token embedding到較小的size會丟失一些細(xì)粒度的信息。

“Embedding的圖片太大，會導(dǎo)致計算量過大；embedding的圖片過小，會導(dǎo)致信息損失太大”，如何在保證合理計算量的情況下，獲得更加細(xì)粒度的特征表示，這就是本文要解決的問題。

3. 方法

本文的模型可以分為兩步：

第一步，通過一系列Outlooker獲得更加細(xì)粒度的特征表示。

第二步，通過一系列的Transformer結(jié)構(gòu)聚合global的信息

3.1. Outlooker

本文提出的Outlooker其實和Transformer的Multi-head Attention非常相似。不同之處在將Multi-head Attention中的Self-Attention換成了OutlookAttention，從而能夠獲取更加細(xì)粒度的特征表示。

Outlooker的計算步驟如下：

3.2. Outlook Attention

Outlook Attention的核心思想如下：

1）每個空間位置的特征足夠豐富，用于產(chǎn)生局部聚合其相鄰特征的注意力權(quán)重；

2）深層的和局部的空間聚合可以有效地編碼細(xì)粒度的信息。

下面來看Outlook Attention的框架圖：

整個框架分為兩個分支，上面的分支用于生成attention map，下面的分支用于生成投影后的value。

3.2.1. 生成Attention Map

首先看上面的分支，特征x首先被embedding到了，其中K是局部感知的大小

然后，我們將的特征reshape成（,）。

緊接著，我們將最后一個維度進(jìn)行softmax，我們就能到了每個位置和周圍個位置的注意力權(quán)重。因此，我只需要將這個attention map和embedding后的value相乘就得到了新的feature。

3.2.2. 映射到新的特征v

下面我們來看下面的分支，由兩部分組成。

第一部分是Linear，這個操作是為了將輸入的特征x進(jìn)行一個embedding，映射到一個新的特征V。

然后就是就是unfold操作，unfold操作就是卷積中的滑動窗口操作。不同的是這里只有“卷”，沒有“積”。就是將特征中的KxK的區(qū)域給取出來。

3.2.3. 獲取weighted特征

前面在3.2.2中，我們已經(jīng)獲得了每個位置的attention map，因此，我們在這一步中只需要做矩陣相乘，就可以得到Outlook Attention之后的結(jié)果。最后我們通過Fold函數(shù)，就能將feature map還原到輸入的大小。

從下面這張圖可以看出，在Outlook Attention中，每一個中心點(diǎn)的位置都要周圍kxk個位置進(jìn)行attention操作，這個步驟就有點(diǎn)類似卷積。

Outlook Attention的偽代碼如下：

pytorch可執(zhí)行的代碼：

class OutlookAttention(nn.Module):

    def __init__(self,dim,num_heads=1,kernel_size=3,padding=1,stride=1,qkv_bias=False,
                    attn_drop=0.1):
        super().__init__()
        self.dim=dim
        self.num_heads=num_heads
        self.head_dim=dim//num_heads
        self.kernel_size=kernel_size
        self.padding=padding
        self.stride=stride
        self.scale=self.head_dim**(-0.5)

        self.v_pj=nn.Linear(dim,dim,bias=qkv_bias)
        self.attn=nn.Linear(dim,kernel_size**4*num_heads)

        self.attn_drop=nn.Dropout(attn_drop)
        self.proj=nn.Linear(dim,dim)
        self.proj_drop=nn.Dropout(attn_drop)

        self.unflod=nn.Unfold(kernel_size,padding,stride) #手動卷積
        self.pool=nn.AvgPool2d(kernel_size=stride,stride=stride,ceil_mode=True) 

    def forward(self, x) :
        B,H,W,C=x.shape

        #映射到新的特征v
        v=self.v_pj(x).permute(0,3,1,2) #B,C,H,W
        h,w=math.ceil(H/self.stride),math.ceil(W/self.stride)
        v=self.unflod(v).reshape(B,self.num_heads,self.head_dim,self.kernel_size*self.kernel_size,h*w).permute(0,1,4,3,2) #B,num_head,H*W,kxk,head_dim

        #生成Attention Map
        attn=self.pool(x.permute(0,3,1,2)).permute(0,2,3,1) #B,H,W,C
        attn=self.attn(attn).reshape(B,h*w,self.num_heads,self.kernel_size*self.kernel_size \
                    ,self.kernel_size*self.kernel_size).permute(0,2,1,3,4) #B，num_head，H*W,kxk,kxk
        attn=self.scale*attn
        attn=attn.softmax(-1)
        attn=self.attn_drop(attn)

        #獲取weighted特征
        out=(attn @ v).permute(0,1,4,3,2).reshape(B,C*self.kernel_size*self.kernel_size,h*w) #B,dimxkxk,H*W
        out=F.fold(out,output_size=(H,W),kernel_size=self.kernel_size,
                    padding=self.padding,stride=self.stride) #B,C,H,W
        out=self.proj(out.permute(0,2,3,1)) #B,H,W,C
        out=self.proj_drop(out)

        return out

3.3. VOLO的不同變種

作者在文中給出了不同變種的VOLO，其中Stage 1為Outlook Attention+MLP,Stage 2為Self-Attention+MLP。

3.4. 復(fù)雜度分析

文中給出了Self-Attention、Local Self-Attention，和Outlook Attention的復(fù)雜度。通常，所以O(shè)A的復(fù)雜度是低于SA和LSA的。

4. 實驗

不同VOLO變種在Imagenet上的準(zhǔn)確率對比

用不同模塊進(jìn)行局部感知的結(jié)果（說明OA確實比LSA和Conv更強(qiáng)一些）

在分割任務(wù)上的performance

5. 總結(jié)

這篇論文其實也是在為了解決Self-Attention只有全局信息感知的缺點(diǎn)，Outlook Attention相當(dāng)于是加入了一個局部感知的功能，使得模型能夠感知局部特征，從而使最終提取的特征更加細(xì)粒度。這篇工作也是將ImageNet數(shù)據(jù)集上Performance刷到了一個新高度。

- END -

歡迎加入GiantPandaCV交流群，可以添加號主微信

OutLook Attention：具有局部信息感知能力的ViT