Vision MLP（MLP-Mixer，RepMLP，ResMLP，gMLP，aMLP）

MLP–>CNN–>Transformer–>MLP
天道好轮回。CNN家族和Vision Transformer博主已经整理过，不再赘述，本期博文主要整理Vision MLP范式的文章们。

Do You Even Need Attention? A Stack of Feed-Forward Layers Does Surprisingly Well on ImageNet
这篇文章想要大家首先需要思考的问题是，当Transformer开始流行于视觉领域时，其注意力机制真的是保障图像分类任务性能的关键吗？

• Transformer于CNN的区别和优势。Attention为网络引入了全局感受野（global receptive field），而注意力模块的效率及有效性与网络是至关重要的。

但局部感知和平移不变不就是前CNN的优点吗？？所以反过来看，其实单纯的“全局性”正是普通神经网络MLP的重要能力。所以这篇文章，在不改动注意力分类网络其他结构和参数设置的情况下，单纯将Transformer里面每个注意力网络层替换为一个简单的前传网络层FFN（feed-forward layer），从而变成一个纯多层感知机的网络。模型架构如上图，由Feed-Forward Layer（Features），Transpose，Feed-Forward Layer（Patches），Transpose组成。看后面的代码会非常清楚具体的做法。

上图右侧的训练结果可以看到，Tiny 网络中 FF only 模型表现一般，但当模型增大为 Base 与 Large 时，FF only 表现相当给力，这从侧面也说明了MLP 模型在较复杂较大的网络中有不俗的表征能力 （无限逼近理论要来了吗？）。

这或许能够证明出色的架构，以及大量的数据学习优异的嵌入是十分重要的，model architecture和loss具有一定的等价性。

class LinearBlock(nn.Module):

    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, num_tokens=197):
        super().__init__()

        # First stage
        self.mlp1 = Mlp(in_features=dim, hidden_features=int(dim * mlp_ratio), act_layer=act_layer, drop=drop)
        self.norm1 = norm_layer(dim)

        # Second stage
        self.mlp2 = Mlp(in_features=num_tokens, hidden_features=int(
            num_tokens * mlp_ratio), act_layer=act_layer, drop=drop)
        self.norm2 = norm_layer(num_tokens)

        # Dropout (or a variant)
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()

    def forward(self, x):
        x = x + self.drop_path(self.mlp1(self.norm1(x))) #mlp
        x = x.transpose(-2, -1) #transpose
        x = x + self.drop_path(self.mlp2(self.norm2(x))) #mlp
        x = x.transpose(-2, -1) #transpose
        return x 

paper：https://arxiv.org/abs/2105.02723
code：https://github.com/lukemelas/do-you-even-need-attention

MLP-Mixer: An all-MLP Architecture for Vision
所以这篇备受关注的谷歌MLP-Mixer文章，就直接尝试将Vision Transformer架构中的Attention全部变为MLP，即其只基于多层感知机结构，只依赖基础的矩阵相乘，重复地在空间特征或者通道特征上计算抽取。完整架构如上图：

• 输入的处理和Vision Transformer一致，切成Patch再展平，然后通过Per-patch Fully-connected将每个patch转换成feature embedding。值得注意的是，后面会做token-mixing MLP，所以不需要加入position embedding也能感知位置。
• X然后经过多个 Mixer 层，其包含两类 MLP 计算层：token-mixing MLPs（计算不同方块张量之间的联系） 和 channel-mixing MLPs（计算通道之间的特征）。先token再channel，这两种类型的层交替执行以促进两个维度间的信息交互。其中每个MLP由两层fully-connected和一个GELU构成。公式为：$$U=X+W_2\sigma (W_{1}LN(X))$$$$Y=U+W_4\sigma (W_{3}LN(U))$$
• 然后pooling之后接全连接用于分类。

在极端的情况下，Mixer 架构可以看做是一个特殊的 CNN，使用 1×1 卷积进行 channel mixing，同时全感受野和参数共享的的单通道深度卷积进行 token mixing。

class MlpBlock(nn.Module):
  mlp_dim: int

  @nn.compact
  def __call__(self, x): #这部分和上图的右侧MLP一致
    y = nn.Dense(self.mlp_dim)(x) #MLP+GELU+MLP
    y = nn.gelu(y)
    return nn.Dense(x.shape[-1])(y)


class MixerBlock(nn.Module):
  """Mixer block layer."""
  tokens_mlp_dim: int
  channels_mlp_dim: int

  @nn.compact
  def __call__(self, x): #这一部分和上面的公式中一致
    y = nn.LayerNorm()(x)
    y = jnp.swapaxes(y, 1, 2)
    y = MlpBlock(self.tokens_mlp_dim, name='token_mixing')(y) #token MLP
    y = jnp.swapaxes(y, 1, 2)
    x = x + y #残差部分
    y = nn.LayerNorm()(x)
    return x + MlpBlock(self.channels_mlp_dim, name='channel_mixing')(y) #channel MLP


class MlpMixer(nn.Module):
  """Mixer architecture."""
  patches: Any
  num_classes: int
  num_blocks: int
  hidden_dim: int
  tokens_mlp_dim: int
  channels_mlp_dim: int

  @nn.compact
  def __call__(self, inputs, *, train): #把各组件搭起来
    del train
    x = nn.Conv(self.hidden_dim, self.patches.size,
                strides=self.patches.size, name='stem')(inputs)
    x = einops.rearrange(x, 'n h w c -> n (h w) c')
    for _ in range(self.num_blocks):
      x = MixerBlock(self.tokens_mlp_dim, self.channels_mlp_dim)(x)
    x = nn.LayerNorm(name='pre_head_layer_norm')(x)
    x = jnp.mean(x, axis=1)
    return nn.Dense(self.num_classes, kernel_init=nn.initializers.zeros,
                    name='head')(x)

其效果基本现有SOTA模型的表现持平，同时看作者的前几个token-mixing MLP的weights可视化，可以发现其和CNN逐层由粗到细、由轮廓到慢慢具象化的演变很像。

• paper：https://arxiv.org/pdf/2105.01601.pdf

• code：https://github.com/google-research/vision_transformer/tree/linen

MLP-Mixer相对Transformer的优势？
MLP-Mixer用Mixer的MLP来替代ViT的Transformer，减少了特征提取的自由度，并且巧妙的可以交替进行patch间信息交流和patch内信息交流。

MLP-Mixer相对CNN的优势？
CNN的特点是inductive bias，即局部假设这一设计模式。但当算了提高了，大量数据是可以用来直接学习的，如ViT的范式。而MLP-Mixer则说明了大量数据+算力是可以直接堆砌学习的。

MLP-Mixer和Conv1x1的不同之处在哪里？
1x1卷积可以结合不同channels的信息，但无法结合不同空间位置的信息。具体来说，Channel-mixing MLPs相当于1x1卷积，而Token-mixing MLPs相当于广义的depth-wise卷积。

那么真的是MLP吗？
其实这一点也遭到了LeCun的批评，如果真的是标准的MLP，那应该要将输入展平为一个一维向量，然后再接变换矩阵。现在的做法更像是一维卷积。

RepMLP: Re-parameterizing Convolutions into Fully-connected Layers for Image Recognition
CNN和MLP是否能够结合呢？

• CNN擅长捕捉局部的特征或模式识别，即归纳偏置或局部先验（local prior）。
• MLP更加擅长于建立特征的长依赖/全局关系与空间关系（所以 ViT 需要更大的训练集或数据扩增来训练模型）

那么两者可以结合吗？于是这篇清华的文章提出了 RepMLP，巧妙利用“重参数”（re-parameterization）的方法，将局部的先验信息加进了全连接层。

如上图，RepMLP由三个部分组成，分别是Global Perceptron, Partition Perceptron and Local Perceptron。具体来说，其同时利用多层神经网络提取长期依赖关系与空间信息，并且利用结构化的重参数化，在网络训练时候将多个卷积模块与全连接并行，用其抽取对应的局部先验知识并最后进行信息融合汇总。

结构重参数的意思是：训练时的结构对应一组参数，推理时则对应另一组参数，而只要能把前者的参数等价转换为后者，就可以把前者的结果也等价转换。文章中公式太多了，核心就是如何把Conv+BN -> FC。
$$W=RS（Conv（M，F，p），（Chw，Ohw）{）}^T$$直接看代码吧：

    def _convert_conv_to_fc(self, conv_kernel, conv_bias):
        I = torch.eye(self.C * self.h * self.w #self.fc3_groups).repeat(1, self.fc3_groups).reshape(self.C * self.h * self.w // self.fc3_groups, self.C, self.h, self.w).to(conv_kernel.device)
        fc_k = F.conv2d(I, conv_kernel, padding=conv_kernel.size(2)#2, groups=self.fc3_groups)
        fc_k = fc_k.reshape(self.O * self.h * self.w # self.fc3_groups, self.C * self.h * self.w).t()
        fc_bias = conv_bias.repeat_interleave(self.h * self.w)
        return fc_k, fc_bias

• paper：https://arxiv.org/pdf/2105.03404.pdf
• code：https://github.com/DingXiaoH/RepMLP

ResMLP: Feedforward networks for image classification with data-efficient training
这篇和前面的一些文章都很像，主要构建了一个残差架构，其残差块只由一个隐藏层的前馈网络和一个线性patch交互层组成。模型图如上，有两部分：

• 线性层Linear，其中图像 patches在通道之间独立且相同地交互
• 两层前馈网络，其中通道中的每个 patch独立地相互作用

公式为：
$$Z=X+Aff((AAff(X{)}^T{)}^T)$$$$Y=Z+Aff(CGELU(BAff(Z)))$$
其中Aff是类似LN的东西，GELU是激活函数。不过因为在深度学习中，层数的加深一般来说总是有好处的。所以这一启发性的行为，也对后续的工作很重要。

• paper：https://arxiv.org/abs/2105.03404

Pay Attention to MLPs
MLP-Mixer的增强版，带gating的MLP。有两个版本，分别是gMLP和aMLP。

• gMLP的g是“gate”的意思，简单来说gMLP就是将MLP-Mixer跟门控机制结合起来。
  $$\begin{gathered} [X_1,X_2]=X \\ Y=W{X}_2+b \\ O=X_1\otimes Y \end{gathered}$$即将输入沿着特征维度分为两半，然后将其中一半传入MLP-Mixer，作为另一半的gate。

• aMLP的a是“attention”的意思，它将一个简单的单头Self Attention结合进去作为gate：
  $$\begin{gathered} [X_1,X_2]=X \\ {Y}_1=W{X}_2+b \\ {Y}_2=SelfAttention(X) \\ O=X_1\otimes (Y_1+Y_2) \end{gathered}$$从实验结果来看，gMLP略差于标准的Self Attention，而aMLP则是普遍优于Self Attention，这进一步肯定了门控机制的价值。

paper：https://arxiv.org/abs/2105.08050

---

下一篇继续整理最新的论文：Vision MLP（CycleMLP，Swin Transformer，ConvMixer）

---

---

彩蛋
最后彩蛋一篇暂时不知道放在哪里的文章，一篇号称在Graph届中无需message passing，只要简单的MLP就可超越GNN的工作（毕竟某种程度上Transformer也算是全连接的Graph了）。

Graph-MLP：Node Classification without Message Passing in Graph
无需message passing，只要简单的MLP就可超越GNN？模型如上图：

• 上半图是普通使用MP的GNN。用邻接矩阵A来指导消息传播和聚合过程，后面接一个任务loss。
• 下部分是作者说的MLP-based的结构。直接将节点属性通过MLP来映射为其表示，而邻接矩阵A仅仅用来作为loss指导优化过程，不会参与计算。

主要的优化方式其实是利用和Neighboring Contrastive Loss。即认为距离比较近的节点特征相似的，距离比较远的节点特征应该不相似。所以直接用对比学习将相似节点和不相似节点表示分别拉近和推开:
$$L=-log\frac{{\sum }_{j=1}^{B}exp(sim(z_i,z_j)/\tau )}{{\sum }_{k=1}^{B}exp(sim(z_i,z_k)/\tau )}$$其实博主想吐槽的是早在普通的Graph Embedding时期（如LINE，Node2vec）就有这种类似做法了，但是能同时搭上MLP和Contrastive的车也是真的强。

paper：https://arxiv.org/abs/2106.04051