时不我与

V1

2022/02/27阅读:36主题:橙心

1.26

日期: 1.26 - 2.1

论文题目:

1.MLP-Mixer: An all-MLP Architecture for Vision

2.Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks

3.RepMLP: Re-parameterizing Convolutions into Fully-connected Layers for Image Recognition

4.Do You Even Need Attention? A Stack of Feed-Forward Layers Does Surprisingly Well on ImageNet

回味一下多层感知机

MLP(multi-layer perceptrons),中文就是多层感知机。刚才是接触神经网络的时候,好像就是最先接触的这个东西。就是如下图所示的这么一个东西。有输入层,隐藏层,输出层什么的,最后给个预测结果。如果学习过机器学习的一些基础课程,应该都接触过这么个东西。其实也又称之为多层全连接神经网络。也就是大量的矩阵运算balabala~~

1、MLP-Mixer

原文:MLP-Mixer: An all-MLP Architecture for Vision(arXiv:2105.01601)

论文链接: https://arxiv.org/abs/2105.01601

pytorch复现代码:https://github.com/d-li14/mlp-mixer.pytorch

好像说是谷歌是VIT团队的论文。还特意去对比了一下作者,果然好几个人都是一样的。

1.1 模型的输入,分patch

模型的整体框架如上所示,乍一看,还以为是transformer,输入都是分patch后作为输入,和transformer图很像了。当然了,后续的操作肯定还是不一样的。简单解释一下,比如一张256 \times 256的图,patch的大小是32\times 32。那么他就可以得到64个patch,每个patch的参数维度为32\times 32\times 3 = 3072,也就是说这样的输入维度为

\mathbb{R}^{b \times 64 \times 3072} \
\mathbb{R}^{b \times 64 \times 3072} \

其中,b是batchsize。输入的形式和VIT应该是一致的。

顺带提一下。对于图像的低级任务,像去雾去噪超分辨等,好像最早就有说分patch去做的,单patch的超分辨去噪似乎也没太大影响,所以早期应该是有方法ptach在MLP处理的方法吧(其实我也不确定)。至于比较高级的涉及语义的任务,分类,分割什么的就不知道有没有了~~

1.2 模型框架,MLP layers

从input进入,首先经历一个per-patch fully-connected,其实也就是个embedding的操作。然后进入一个N次的Mixer Layer。最后池化+全连接,得到后续的分类结果

文章中的MLP layers区分成了两种,channel-mixing MLPs and token-mixing MLPs

图中,紫色框框就是token-mixing MLPs,绿色框框是channel-mixing MLPs。token-mix无非就是使得不同像素间有通信,channel-mix就是使得不同的channel间有通信。作者也提到了channel-mixing MLPs类似于1x1的卷积,token-mixing MLPs类似于卷积核为n的卷积。

在MLP layers之中,先进行一次token-mixing MLP,再进行一次channel-mixing MLP 。

1.3 代码

直接看看代码吧,官方的代码是JAX/Flax框架的。

在guthub上找了一个pytorch复现的版本:https://github.com/d-li14/mlp-mixer.pytorch

import torch
import torch.nn as nn


class MlpBlock(nn.Module):
    def __init__(self, hidden_dim, mlp_dim):
        super(MlpBlock, self).__init__()
        self.mlp = nn.Sequential(
            nn.Linear(hidden_dim, mlp_dim),  
            nn.GELU(),
            nn.Linear(mlp_dim, hidden_dim)  
        )

    def forward(self, x):
        return self.mlp(x) ##每一次的全连接中间有个隐藏层,经过隐藏层再回到与输入一致的维度大小。


class MixerBlock(nn.Module):
    def __init__(self, num_tokens, hidden_dim, tokens_mlp_dim, channels_mlp_dim):
        super(MixerBlock, self).__init__()
        self.ln_token = nn.LayerNorm(hidden_dim)
        self.token_mix = MlpBlock(num_tokens, tokens_mlp_dim) 
        self.ln_channel = nn.LayerNorm(hidden_dim)
        self.channel_mix = MlpBlock(hidden_dim, channels_mlp_dim)

    def forward(self, x):
        out = self.ln_token(x).transpose(1, 2)
        x = x + self.token_mix(out).transpose(1, 2) ##先进行一次token-mixing MLP
        out = self.ln_channel(x)
        x = x + self.channel_mix(out) ##再进行一次channel-mixing MLP
        return x


class MlpMixer(nn.Module):
    ##描述整体的MLP-Mixer框架
    def __init__(self, num_classes, num_blocks, patch_size, hidden_dim, tokens_mlp_dim, channels_mlp_dim, image_size=224):
        super(MlpMixer, self).__init__()
        num_tokens = (image_size // patch_size)**2

        self.patch_emb = nn.Conv2d(3, hidden_dim, kernel_size=patch_size, stride=patch_size, bias=False)
        self.mlp = nn.Sequential(*[MixerBlock(num_tokens, hidden_dim, tokens_mlp_dim, channels_mlp_dim) for _ in range(num_blocks)])
        self.ln = nn.LayerNorm(hidden_dim)
        self.fc = nn.Linear(hidden_dim, num_classes)

    def forward(self, x):
        x = self.patch_emb(x)
        x = x.flatten(2).transpose(1, 2) ##制造生成分patch的input
        x = self.mlp(x) #进行n次mlp layer
        x = self.ln(x)
        x = x.mean(dim=1) #average pooling的操作
        x = self.fc(x)  #全连接至分类数
        return x


def mixer_s32(num_classes=1000, **kwargs):
    return MlpMixer(num_classes, 8, 32, 512, 256, 2048, **kwargs)

2、External Attention

原文:Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks(arXiv:2105.02358)

论文链接: https://arxiv.org/abs/2105.02358

官方Jittor和torch代码:https://github.com/MenghaoGuo/-EANet

清华计图团队的工作,所以代码也是有一份jittor框架的代码。

这篇文章提出了一种新型的attention机制,可能是因为这种attention机制的线性性质,所以会在MLP这个板块中出现。感觉还是蛮有意思的。思想也不复杂。用文章摘要的一句话就是:

This paper proposes a novel attention mechanism which we call external attention, based on two external, small, learnable, and shared memories, which can be implemented easily by simply using two cascaded linear layers and two normalization layers

两个外部的、小的、可学习的和共享的存储器,存储器?是个啥?

2.1 self-attentio的缺陷

从摘要到introduction,作者一直都有说self-attention的缺陷在于:1.计算量大,2.他只局限于当前样本,没有关注数据集中的其他样本,缺乏信息的交互。比如说,在分割任务中,在不同样本中也存在着同一类别的特征,其他样本的特征,能否辅助该样本的分割呢。

2.2 self-attention vs External Attention

文章通过对比self-attention来引出external attention。

对于上图a就是经典的self-attention。通过先将feature map投影到QKV,通过Q,K计算注意力权重,最后分配到V上:

\begin{aligned} A =(\alpha)_{i, j}&=\operatorname{softmax}\left(Q K^{T}\right) \ F_{\text {out }} &=A V \end{aligned} \
\

b展示的是a的简化版,即无需投影,所有操作对feature map直接操作:

\begin{aligned} A &=\operatorname{softmax}\left(F F^{T}\right) \ F_{\text {out }} &=A F \end{aligned} \
\

c展示也就是文章所提出的external attention,文章中给出的公式是:

\begin{aligned} A &=\operatorname{Norm}\left(F M_{k}^{T}\right) \ F_{\text {out }} &=A M_{v} \end{aligned} \
\

其中这个A表示了注意力图。M_{k}M_{v}是两个不同的记忆单元。M \in \mathbb{R}^{S \times d}

文章也给出了一个EA操作用于语义分割的框架图,可以看出,在backbone之后引入该模块。EA操作的前后还是给出一个shortcut,类似resnet,使用了sumation将前后直接相加起来。这个后面带个MLP操作不知道是个什么意思,文章里貌似也没看到什么说法。代码中,好像也没个,最后也就是卷积加个上采样的操作。

2.3 double-normalization

EA操作中,norm操作使用的是double-normalization。文章中说是为了避免输入特征的比例敏感,相当于做的两次norm。

\begin{aligned} (\tilde{\alpha})_{i, j} &=F M_{k}^{T} \ \alpha_{i, j} &=\frac{\exp \left(\tilde{\alpha}_{i, j}\right)}{\sum_{k} \exp \left(\tilde{\alpha}_{k, j}\right)} \ \alpha_{i, j} &=\frac{\alpha_{i, j}}{\sum_{k} \alpha_{i, k}} \end{aligned} \
\

2.4 实验

文章的实验做的很多,包括图像分类、语义分割、图像生成、点云分类和点云分割等任务,应该还是想验证该注意力对高级任务的有效性,就像前面提到的分割的例子一样。EA操作,又节省复杂度,又能使得sample之间也获得关联。

2.5 代码

官方Jittor和torch代码:https://github.com/MenghaoGuo/-EANet

External_attention的代码:

class External_attention(nn.Module):
    '''
    Arguments:
        c (int): The input and output channel number.
    '''
    def __init__(self, c):
        super(External_attention, self).__init__()
        
        self.conv1 = nn.Conv2d(c, c, 1)
        self.k = 64
        self.linear_0 = nn.Conv1d(c, self.k, 1, bias=False)
        self.linear_1 = nn.Conv1d(self.k, c, 1, bias=False)
        self.linear_1.weight.data = self.linear_0.weight.data.permute(1, 0, 2)        
        self.conv2 = nn.Sequential(
            nn.Conv2d(c, c, 1, bias=False),
            norm_layer(c))        
        
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.Conv1d):
                n = m.kernel_size[0] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, _BatchNorm):
                m.weight.data.fill_(1)
                if m.bias is not None:
                    m.bias.data.zero_()
 

    def forward(self, x):
        idn = x
        ##刚进EA模块的时候是一个正常的四维的特征图。比如1*512*16*16
        x = self.conv1(x)

        b, c, h, w = x.size()  ## 1 512 16 16
        n = h*w  ## n=16*16=256
        x = x.view(b, c, h*w)   # b * c * n   1*512*256

        attn = self.linear_0(x) # b, k, n    用kernel为1的1d卷积,得到 1*64*256
        #实际上这一步的1x1卷积就是将输入乘上一个(512x64)的可学习的卷积。
        attn = F.softmax(attn, dim=-1) # b, k, n

        attn = attn / (1e-9 + attn.sum(dim=1, keepdim=True)) #  # b, k, n
        ##进行两次norm的操作,即double-normalization
        
        x = self.linear_1(attn) # b, c, n
        #再进行第二次的矩阵乘法 从1*64*256重新变回1*512*256

        x = x.view(b, c, h, w)#变回四维 1*512*16*16
        x = self.conv2(x) #卷积 1*512*16*16 -> 1*512*16*16
        x = x + idn  #与输入直接值相加  1*512*16*16 -> 1*512*16*16
        x = F.relu(x)
        return x

3、RepMLP

原文:RepMLP: Re-parameterizing Convolutions into Fully-connected Layers for Image Recognition(arXiv:2105.02358)

论文链接: https://arxiv.org/abs/2105.01883

官方pytorch代码:https://github.com/DingXiaoH/RepMLP

清华丁霄汉的工作,最早关注他的工作是ICCV2019的ACNet(1908.03930)(用三个不同的卷积核的三通路卷积来训练,用一个3X3卷积核在测试)。他后续的工作也是主要是在结构重参数化。今年的CVPR2021中,也 有两篇相关的工作,分别是RepVGG和DiverseBranchBlock。对结构重参数化感兴趣的,可以关注一下这个大佬。

结构重参数化,简而言之。就是结构(或者说是模型)A对应一组参数X,结构B对应一组参数Y,如果我们能将X等价转换为Y,就能将结构A等价转换为B。在实际的使用中,使用一个较为复杂的网络进行训练,再想办法将这个复杂的网络简化,最终得到的效果是一致的。也就是说,对于同样一个输入,这个复杂的网络和简化的网络能得到完全一致的输出。

这篇RepMLP也是一篇相关的工作,正好MLP是这突如其来的热点,他就是其中之一。

3.1 卷积和MLP处理图像的优缺点

这篇文章很巧妙的从深度学习处理图像中几个重要的性质开始说起。即长距离建模(或者说是全局信息),位置信息以及局部先验。

局部先验:由于conv层只处理局部邻域,图像的局部性(即一个像素与其邻居的关系比远处的像素更密切)使得卷积神经网络在图像识别中取得了成功。这是全连接操作不会有的,因为全连接操作是全局的,每个点之间都能产生关系,就不存在什么局部的说法。这也说明了卷积操作的合理性和必要性。

The locality of images (i.e., a pixel is more related to its neighbors than the distant pixels) makes Convolutional Neural Network (ConvNet) successful in image recognition, as a conv layer only processes a local neighborhood. In this paper, we refer to this inductive bias as the local prior.

全局信息:传统的纯卷积网络当中,我们会通过不断的卷积或者encoder减小尺寸,扩大感受野来获得全局信息。事实上在transformer的工作上就是一个有丰富全局信息的方法。同样的,对比卷积来说,他缺乏局部先验,所以可能需要大量的数据来进行预训练。同样的,全连接操作是天然获取全局信息的。

位置信息:很多图像是有位置先验的(比如说一个人的面部,眼睛肯定是在鼻子上面的),但是卷积操作是无法利用这些位置信息的。而全连接也是天然有位置信息的,他的数据分布是排序的。

综上所述,全连接天然拥有全局信息和位置信息,但是我也想要有局部先验呀。怎么办,如果是我的话,那就直接并行两条支路呗。但是,作者的野心不止于此,他想要一个纯的MLP网络。于是就是作者的老本行了,将卷积操作重参数化为MLP。

3.2 RepMLP模块

在RepVGG中,训练时是多分支结构,推理时是单分支结构。

那在RepMLP中呢?

看上图,还是颇为复杂的。左边的这个训练阶段简单拆分一下,分为三部分Global Perceptron、Partition Perceptron、Local Perceptron。

Global Perceptron

顾名思义,Global Perceptron就是用于提取全局信息的模块。输入进入之后,分成了两条线。蓝线的操作就是一个reshape感觉,本质上就是将图片分patch,然后叠加在bathsize上面。但是直接这么做似乎缺乏了patch之间的相关性。于是就有了绿色的这条支路,将原图进行pooling,使得大小为patch的数量,也就是一个patch对应1个点。然后进行一系列操作。最后叠加到蓝线的输出中。

Partition Perceptron

从global出来之后,通过各种reshape的操作,将图像的特征图形式的样式序列化。在进行全连接的操作。因为全连接可以天然的有位置信息,于是这个模块是可以获取位置信息的。

Local Perceptron

局部信息也就是前面提到的,卷积可以获得但全连接不能获得的信息。在训练模块中,依然使用的是卷积操作。对进来的特征图,分别进行卷积核为1,3,5,7的分组卷积,得到相对于的特征图,最后合并起来。并与Partition Perceptron的输出相加。

这里还是简单的介绍了一下几个模块大概的一个情况,很多细节和具体的公式还是需要在论文中去看。这里也推荐一篇讲解(https://mp.weixin.qq.com/s/FwITC1JEG1vr2Y1ePzSvuw),将数据流都解释的很详细了,所以也不多花篇幅在这个上面了。

3.3 将卷积变成全连接?怎么变?

文章定义FC kernel为\mathrm{W}^{(1)}(\mathrm{Ohw}, C h w),卷积的conv kernel为\mathrm{F}(O, C, K, K),padding为p

我们现在想要将FC操作和CONV操作给相加起来。

\operatorname{MMUL}\left(\mathrm{M}^{(\text {in })}, \mathrm{W}^{\prime}\right)=\operatorname{MMUL}\left(\mathrm{M}^{(\mathrm{in})}, \mathrm{W}^{(1)}\right)+\mathrm{CONV}\left(\mathrm{M}^{(\mathrm{in})}, \mathrm{F}, p\right) \
\operatorname{MMUL}\left(\mathrm{M}^{(\text {in })}, \mathrm{W}^{\prime}\right)=\operatorname{MMUL}\left(\mathrm{M}^{(\mathrm{in})}, \mathrm{W}^{(1)}\right)+\mathrm{CONV}\left(\mathrm{M}^{(\mathrm{in})}, \mathrm{F}, p\right) \

希望获得这么一个结果。即,一次fn操作和一次conv操作可以用一个fn操作就给它代替掉了。

我们知道,两个相同尺度的全连接操作是可以直接相加的,如果想要全连接和卷积能够相加的话,就要把卷积操作,转换成一个等价的全连接操作。也就是希望找到一个W满足下面这样的关系。也就是把卷积核的权重参数,想办法变成等价的全连接的权重参数(这是一定存在的,因为卷积本身就是一种特殊的全连接)

\operatorname{MMUL}\left(\mathrm{M}^{(\mathrm{in})}, \mathrm{W}^{(\mathrm{F}, p)}\right)=\mathrm{CONV}\left(\mathrm{M}^{(\mathrm{in})}, \mathrm{F}, p\right) \
\operatorname{MMUL}\left(\mathrm{M}^{(\mathrm{in})}, \mathrm{W}^{(\mathrm{F}, p)}\right)=\mathrm{CONV}\left(\mathrm{M}^{(\mathrm{in})}, \mathrm{F}, p\right) \

对于一个输入M^{(in)},他要得到输出M^{(out)}。可以通过卷积实现,也可以通过全连接实现。如下式:

\mathrm{M}^{(\text {out })}=\operatorname{CONV}\left(\mathrm{M}^{(\text {in })}, \mathrm{F}, p\right)=\operatorname{MMUL}\left(\mathrm{M}^{(\mathrm{in})}, \mathrm{W}^{(\mathrm{F}, p)}\right) \
\mathrm{M}^{(\text {out })}=\operatorname{CONV}\left(\mathrm{M}^{(\text {in })}, \mathrm{F}, p\right)=\operatorname{MMUL}\left(\mathrm{M}^{(\mathrm{in})}, \mathrm{W}^{(\mathrm{F}, p)}\right) \

下面正式看看这么做:

首先定义一下全连接操作,其实也就是矩阵乘法

\mathrm{V}^{(\text {out })}=\operatorname{MMUL}\left(\mathrm{V}^{(\text {in })}, \mathrm{W}\right)=\mathrm{V}^{(\mathrm{in})} \cdot \mathrm{W}^{\top} \
\mathrm{V}^{(\text {out })}=\operatorname{MMUL}\left(\mathrm{V}^{(\text {in })}, \mathrm{W}\right)=\mathrm{V}^{(\mathrm{in})} \cdot \mathrm{W}^{\top} \

结合前面给的定义,也就是说我现在要寻找一个\mathrm{W}^{(\mathrm{F}, p) \top},输入乘上这样一个矩阵,它能起到的作用等同于F为卷积核,padding p的卷积操作。下面这个\mathrm{W}^{(\mathrm{F}, p) \top}是我们要去寻找到哦。为了方便理解,这里也判断一下各个张量的维度,{V}^{(\text {out })} \in (Ohw){V}^{(\text {in })} \in (Chw)\mathrm{W}^{(\mathrm{F}, p) \top} \in (Chw,Ohw)

\mathrm{V}^{(\text {out })}=\mathrm{V}^{(\mathrm{in})} \cdot \mathrm{W}^{(\mathrm{F}, p) \top} \
\mathrm{V}^{(\text {out })}=\mathrm{V}^{(\mathrm{in})} \cdot \mathrm{W}^{(\mathrm{F}, p) \top} \

先插入一个单位矩阵I(Chw,Chw),不影响运算,值和维度都不会发生改变:

\mathrm{V}^{(\mathrm{out})}=\mathrm{V}^{(\mathrm{in})} \cdot\left(\mathrm{I} \cdot \mathrm{W}^{(\mathrm{F}, p) \mathrm{T}}\right) \
\mathrm{V}^{(\mathrm{out})}=\mathrm{V}^{(\mathrm{in})} \cdot\left(\mathrm{I} \cdot \mathrm{W}^{(\mathrm{F}, p) \mathrm{T}}\right) \

\mathrm{W}^{(\mathrm{F}, p) \top}终究还是需要源自于卷积核F的(不管他是怎么变过来的),上面(\mathrm{I} \cdot \mathrm{W}^{(\mathrm{F}, p) \mathrm{T}})这个式子可以表示的是对\mathrm{M}^{(\mathrm{I})}=reshape(I)进行F卷积操作。\mathrm{M}^{(\mathrm{I})} \in (C h w, C, h, w),\mathrm{F} \in (O, C, K, K),\operatorname{CONV}\left(\mathrm{M}^{(\mathrm{I})}, \mathrm{F}, p\right)的结果的维度应该是(Chw,O,h,w),再经过下面第三个公式的RS后有变成了二维的(Chw,Ohw),再让输入乘以它。得到的结果刚刚好是输出想有的尺寸(Ohw)

\begin{array}{c} \mathrm{M}^{(\mathrm{I})}=\operatorname{RS}(\mathrm{I},(C h w, C, h, w)) \ \mathrm{I} \cdot \mathrm{W}^{(\mathrm{F}, p) \top}=\operatorname{CONV}\left(\mathrm{M}^{(\mathrm{I})}, \mathrm{F}, p\right) \ \mathrm{V}^{(\mathrm{out})}=\mathrm{V}^{(\mathrm{in})} \cdot \mathrm{RS}\left(\mathrm{I} \cdot \mathrm{W}^{(\mathrm{F}, p) \top},(C h w, O h w)\right) \end{array} \
\

第一个公式意思就是将矩阵I变成4维的形式,第二个公式表示的是对I乘以这个矩阵需要等同于对的I的reshape做卷积操作。第三个式子表示,将卷积得到的重新reshape成矩阵,并让输入乘以他。

前面搞了这么久,绕过来绕过去。但是我们要明确的还是,我们已知的是卷积,要求的是FC的权重剧中。结合前面的分析可以得到:

\mathrm{W}^{(\mathrm{F}, p)}=\mathrm{RS}\left(\mathrm{CONV}\left(\mathrm{M}^{(\mathrm{I})}, \mathrm{F}, p\right),(C h w, O h w)\right)^{\top} \
\mathrm{W}^{(\mathrm{F}, p)}=\mathrm{RS}\left(\mathrm{CONV}\left(\mathrm{M}^{(\mathrm{I})}, \mathrm{F}, p\right),(C h w, O h w)\right)^{\top} \

文章最后也给了伪代码以供参考:

Algorithm 1 PyTorch code for converting groupwsie conv into FC.
Input: C, h, w, g, O, conv kernel, conv bias
I = torch.eye(C * h * w // g).repeat(1, g).reshape(C * h * w // g, C, h, w)
fc kernel = F.conv2d(I, conv kernel, padding=conv kernel.size(2)//2, groups=g)
fc kernel = fc kernel.reshape(O * h * w // g, C * h * w).t() # Note the transpose
fc bias = conv bias.repeat interleave(h * w)
return: fc kernel, fc bias

3.4 总结和思考

其实,我想了很久没想明白。后来,也有一点点自己的理解吧。首先明确任务,就是要将一个卷积核参数转化成全连接权重参数,他们的维度不一样。卷积是一种特殊的全连接,或者说是稀疏的全连接。如果现在拿到手一个卷积操作,不管通过手工还好,拼凑还好,肯定能变换成一个合适的全连接权重,是满足参数要求的。

那么作者妙在哪里呢?他想了一个办法,使得这个过程能够很轻松的通过普通运算就能搞出来。这个转换有效,过程可微,满足了我们去训练网络非常重要的一些性质。

那么是方法是什么呢?其实前面很多公式已经给出了具体过程了。其实,在我理解看来,就是利用了单位矩阵的良好性质,文章中称之为identity matrix(自身的矩阵)。在线性代数中,和单位矩阵进行乘法运算都会成为自己本身。同样的,在这里是不是可以理解为,对单位矩阵进行卷积,相当于是在保留这个卷积的运算,后续的reshape成二维也好,怎么样也好,他保存的计算属性应该是一致的,所以可以去等价的替换。

不管是ACNet还是RepVgg,作者都是对卷积操作做替换。这次直接卷积到全连接。其实思路可能是好想的,因为卷积源于全连接,他们之间有着千丝万缕的联系。但是难就难在如何巧妙的转换,如何让这个过程可微又方便,所以真的很妙~

参考

1.https://mp.weixin.qq.com/s/FwITC1JEG1vr2Y1ePzSvuw

2.https://zhuanlan.zhihu.com/p/369970953

4、Do You Even Need Attention?

原文:Do You Even Need Attention? A Stack of Feed-Forward Layers Does Surprisingly Well on ImageNet(arXiv:2105.02723)

论文链接: https://arxiv.org/abs/2105.02723

pytorch代码:https://github.com/lukemelas/do-you-even-need-attention

牛津大学的工作(居然只有一个作者),这是一篇很短的文章,仅仅只有三页。这篇短文更像一篇实验报告,可能是想让狂热的transformer稍微冷静一下下吧。

4.1 将attention层换成了feed-forward层

文章开篇质疑了attention的作用,认为VIT起作用不一定是因为attention。作者验证的方式就是将vit中的attention层换成了feed-forward层进行实验。

实验结果表明VIT的强大性能可能更多地归因于其他因素,而不是注意力机制,例如由patch embedding和训练策略产生的影响。

4.2 具体的代码实现

from torch import nn
class LinearBlock(nn.Module):
 def __init__(self, dim, mlp_ratio=4., drop=0., drop_path=0., act=nn.GELU,
        norm=nn.LayerNorm, n_tokens=197): # 197 = 16**2 + 1
        super().__init__()
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        # FF over features
        self.mlp1 = Mlp(in_features=dim, hidden_features=int(dim*mlp_ratio), act=act, drop=drop)
        self.norm1 = norm(dim)
        # FF over patches
        self.mlp2 = Mlp(in_features=n_tokens, hidden_features=int(n_tokens*mlp_ratio), act=act, drop=drop)
        self.norm2 = norm(n_tokens)
 def forward(self, x):
        x = x + self.drop_path(self.mlp1(self.norm1(x)))
        x = x.transpose(-2, -1)
        x = x + self.drop_path(self.mlp2(self.norm2(x)))
        x = x.transpose(-2, -1)
        return x
    
class Mlp(nn.Module):
 def __init__(self, in_features, hidden_features, act_layer=nn.GELU, drop=0.):
        super().__init__()
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, in_features)
        self.drop = nn.Dropout(drop)
 def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x

代码也比较好看, 就是相比vit换attention为MLP

分类:

后端

标签:

后端

作者介绍

时不我与
V1