一、RepMLP

原文链接:https://arxiv.org/pdf/2105.01883.pdf

RepMLP(re-parameterized MLP)是考虑到FC层比起卷积来说不擅长捕捉局部信息的特点而提出的。其训练和推断是不同的。

训练阶段由全局感知器,partition感知器和局部感知器组成。

全局感知器

图中N为batch size

        将feature map分割为partition。为了捕捉partition之间的交互,使用平均池化处理每个partition,输入到BN和2层的MLP,然后reshape,加到partition map上。

partition感知器

     由一个FC层和BN层组成,以partition map作为输入。FC层为类似group卷积的group FC以减少参数。

group FC在编程时采用group 1\times1卷积实现,步骤如下:(1)将V^{(\textup{in})}reshape为空间维度为1\times1的feature map;(2)使用g分组的1\times1卷积处理;(3)将处理后的feature map reshape为V^{(\textup{out})}。即:

{M}'=\textup{RS}(V^{(\textup{in})},(\textup{N},P,1,1)),{F}'=\textup{RS}(W,(Q,\frac{P}{g},1,1))

\textup{gMMUL}(V^{(\textup{in})},W,g)=\textup{RS}(\textup{gConv}({M}',{F}',g,0),(\textup{N},Q))

RS表示reshape,gMMUL表示group FC。W为group FC的权重矩阵(大小应为(Q,P/g)),{F}'表示W转换的group卷积核(Q个核,每个核大小为P/g),PQ分别为FC层输入和输出的维度。且应有\textup{N}=NHW/hwP=ChwQ=Ohw

局部感知器

        将partition map通过多个并行的卷积层(保持分辨率大小和输入一致,后接BN),卷积组数g应和Partition感知器中相同

       最后所有卷积输出和Partition感知器的输出相加,还原形状,得到最终输出。

推断阶段会把RepMLP转化为3个FC层。

   关键是两步:

        1.将BN merge到之前的卷积中:

{F}'_{i.:,:,:}=\frac{\gamma_i}{\sigma_i}F_{i.:,:,:},{b}'=-\frac{\mu_i\gamma_i}{\sigma_i}+\beta_i

    2.将卷积转化为FC层(IChw维的单位矩阵):

M^{(I)}=\textup{RS}(I,(Chw,C,h,w)),W^{(F,p)}=\textup{RS}(\textup{Conv}(M^{(I)},F,p),(Chw,Ohw))^T

其中p为padding,F为卷积核,W^{(F,p)}为FC层的权重。

        这样可以把FC3和局部感知器的卷积合并。

二、ResMLP

原文链接:https://arxiv.org/pdf/2105.03404.pdf

Res表示residual。

模型结构

        首先将原图分割为N\times N个patch,然后通过线性层得到d维特征,输入到ResMLP中。图中的A为按列的仿射变换;T为转置。

Residual Multi-Perceptron层

        线性层+前馈层。没有使用LN,而采用对每一列进行仿射变换:

\textup{Aff}_{\alpha,\beta}(x)=\textup{Diag}(\alpha)x+\beta

        这个变换在每个残差块进行两次(两次分别称为为pre和post)。它们在推断时会整合到线性层。

        前馈网络和Transformer一样,为双层MLP,激活函数变为GELU。

Z=X+\textup{Aff}\left ((A\: \textup{Aff}(X)^T)^T \right ),Y=Z+\textup{Aff}\left ( C\: \textup{GELU}(B\: \textup{Aff(Z)}) \right )

其中A,B,C为线性层权重,AN^2\times N^2维,B4d\times d维,Cd\times 4d维。

三、S²-MLPv2

原文链接:https://arxiv.org/pdf/2108.01072.pdf

S²-MLP

        patch embedding层+数个S²-MLP块+分类头

patch embedding层将图像分割成p\times p的patch,然后通过FC得到d维向量。

S²-MLP块

        4个作用于通道维度的MLP+空间移位层。

空间移位:将X沿通道方向分成4份,然后沿长和宽的正负方向分别平移1个单位。

split attention

        设K个大小为n\times c的feature map\{X_k\}_{k=1}^K;其中n为patch数,c为通道数。沿空间维度进行求和得到c维向量a

a=\sum_{k=1}^K\textbf{1}X_k

其中\textbf{1}\in \mathbb{R}^n为全为1的行向量。

        然后a通过MLP把维度变为Kc\hat{a}=\sigma(aW_1)W_2\sigma为GELU),再reshape为K\times c矩阵\hat{A},沿第一维度softmax后得到\bar{A},按照下面的公式生成新的feature map\hat{X}

\hat{X}[i,:]=\sum_{k=1}^KX_k[i,:]\odot \bar{A}[k,:]

其中\odot为按元素乘法。

S²-MLPv2:patch embedding层+数个S²-MLPv2块+分类头

Y=\textup{S}^2\textup{-MLPv2}(\textup{LN}(X))+X,Z=\textup{CM-MLP}(\textup{LN}(Y))+Y

S²-MLPv2块包含S²-MLPv2 component和channel-mixing MLP (CM-MLP)。CM-MLP结构和前面MLP-Mixer(见7种视觉MLP整理(上)中的二)中的结构相同。

S²-MLPv2块结构

\textrm{MLP}_1X的维度变为3c。然后分为3个维度为c的feature map,其中两个按下图做空间移位,另一个不变。 

然后将3个feature map(h\times w\times c)reshape为hw\times c的矩阵,通过split attention后再通过\mathrm{MLP}_2