abstract：

我们提出MLP-Mixer架构(或简称“Mixer”)，这是一个具有竞争力但在概念和技术上都很简单的替代方案，它不使用卷积或自关注。相反，Mixer的架构完全基于多层感知器(mlp)，这些感知器可以在空间位置或特征通道上重复应用。Mixer仅依赖于基本的矩阵乘法例程、对数据布局的更改（重塑和换位）以及标量非线性。

intro：

图1描述了Mixer的宏观结构。它接受一系列线性投影图像补丁(也称为令牌)，形状为“patches x channels”表，作为输入，并保持该维度。Mixer使用两种类型的MLP层:通道混合MLP和令牌混合MLP。信道混合mlp允许在不同信道之间进行通信;它们独立地操作每个令牌，并将表中的各个行作为输入。令牌混合mlp允许在不同的空间位置(令牌)之间进行通信;它们独立地在每个通道上操作，并将表中的各个列作为输入。这两种类型的层相互交织，以支持两个输入维度的交互。

在极端情况下，我们的架构可以看作是一个非常特殊的CNN，它使用1×1卷积进行通道混合，使用单通道深度卷积的完整接受场和参数共享进行令牌混合。然而，相反的情况并不成立，因为典型的cnn并不是Mixer的特例。此外，卷积比mlp中的普通矩阵乘法更复杂，因为它需要对矩阵乘法进行额外的昂贵简化和/或专门的实现。

mixer architecture：

现代深度视觉架构由混合特征的层组成(i)在给定的空间位置，(ii)在不同的空间位置之间，或同时混合特征。

在cnn中，(ii)是用N × N个卷积(对于N > 1)和池化实现的。更深层的神经元有更大的接受野[1,29]。同时，1×1卷积也执行(i)，更大的核同时执行(i)和(ii)。

在Vision transformer和其他基于注意力的架构中，自注意力层允许(i)和（ii），而mlp块执行(i)。Mixer架构背后的思想是清楚地分离每个位置（通道混合）操作(i)和跨位置（令牌混合）操作（ii）。

这两个操作都是通过mlp实现的。图1总结了该体系结构。

代码：

class MlpBlock(nn.Module):
  mlp_dim: int

  @nn.compact
  def __call__(self, x):
    y = nn.Dense(self.mlp_dim)(x)
    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)
    y = jnp.swapaxes(y, 1, 2)
    x = x + y
    y = nn.LayerNorm()(x)
    return x + MlpBlock(self.channels_mlp_dim, name='channel_mixing')(y)


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
  model_name: Optional[str] = None

  @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)
    if self.num_classes:
      x = nn.Dense(self.num_classes, kernel_init=nn.initializers.zeros,
                   name='head')(x)
    return x

mlpblock是一个简单的mlp块，包含两个全连接层和一个GELU激活函数。第一个全连接层增加维度，第二个全连接层减少维度回到原始输入的维度。

mixerblock是核心，首先对输入进行层归一化，然后应用token mixing，即对mlpblock交换维度，接着添加残差连接。之后再次对结果进行层归一化，并应用channel mixing，最后再次添加残差连接。

mlpmixer是完整的MLP-Mixer模型，首先使用一个卷积来处理输入图像，然后将图像patches展平。接着，通过多个mixerblock来处理patches。在所有MixerBlock之后，它应用层归一化并进行平均池化。如果指定了num_classes，它将添加一个全连接层作为分类头。

可以将公式总结为：

这里σ是单元非线性（GELU[16]）。Ds和Dc分别是令牌混合和通道混合mlp中可调的隐藏宽度。注意，Ds的选择与输入补丁的数量无关。因此，网络的计算复杂度在输入patch的数量上是线性的，不像ViT的复杂度是二次的。由于Dc与patch大小无关，因此整体复杂度在图像的像素数上是线性的

如上所述，将相同的通道混合MLP(令牌混合MLP)应用于x的每一行(列)。将通道混合MLP的参数(在每一层内)捆绑在一起是一种自然的选择——它提供了位置不变性，这是卷积的一个突出特征。但是，跨通道绑定参数的情况要少见得多。例如，在一些cnn中使用的可分离卷积[9,40]，独立于其他通道对每个通道应用卷积。然而，在可分离卷积中，不同的卷积核应用于每个通道，而不像Mixer中的令牌混合mlp，它为所有通道共享相同的核(完全接受场)。

当增加隐藏维度C或序列长度S时，参数绑定可以防止体系结构增长过快，并节省大量内存。

Mixer中的每一层（除了初始patch投影层）都有相同大小的输入。这种“各向同性”的设计最类似于变压器，或者其他领域的深度rnn，它们也使用固定的宽度。这与大多数具有金字塔结构的cnn不同：较深的层具有较低的分辨率输入，但有更多的通道。请注意，虽然这些是典型的设计，但也存在其他组合，如各向同性ResNets[38]和金字塔状vit[52]。

除了MLP层，Mixer还使用其他标准的体系结构组件:跳过连接[15]和层规范化[2]。与vit不同，Mixer不使用位置嵌入，因为令牌混合mlp对输入令牌的顺序很敏感。最后，Mixer使用具有全局平均池化层的标准分类头，然后是线性分类器。

注意，在多模态中，视频-文本，音频-文本应用了以上结构：

信息可以在不同的模态和时间序列之间流动。每个区块由两个MLP层和一个GELU激活功能（描述为𝛷). 此外，在每个块中也应用了跳连接。假设 是一个输入特征，其中 𝑡 是时间序列的长度，并且 𝑑 是模态的数量。在每一层中，MLP-communicator模块可表示如下：

其中，𝑖 从 1 到 𝑑 表示行数，𝑗 从 1 到 𝑡 表示列数。Norm() 表示 LayerNorm （层归一化），W 表示每个块中线性层的权重。输入特征 𝑋 首先通过 Time-Mixing MLP 过程，并通过跳连接生成 𝑍，这一步骤允许水平对应特征之间的通信。然后，𝑍 经过 Modality-Mixing MLP 过程生成 𝑌，特征在纵向方向上进行融合。最终得到的特征 𝑌 融合了两个方向的特征信息。该结构允许输入特征中的每个元素可以沿着两个维度与其他特征进行互动。

        解释：在Time-Mixing的过程中，假设我们有三个模态，X j,i 就是表示在第j个时间点，第i个模态的特征值。在这个公式中，我们将所有时间点下不同模态的特征进行了整合，然后Norm(X *,i)是对这个特征值进行正规化，使其分布在一个固定的区间内，通常是0到1或者-1到1。这个正规化的过程可以使得不同的特征值有可比性，且防止部分大数值特征对整个模型训练结果的过度影响。

        接着W_1和W_2就是两个权重矩阵，它们和正规化后的特征值进行乘积操作。φ是激活函数，例如本模型中的GELU激活函数，这个函数起到的作用类似一个开关，可以根据函数结果决定神经元是否被激活，即该特征是否被放入下一层中去。

        然后，我们再看第二个公式：

        这个公式和前一个公式的结构是一样的，都是先进行一次正规化，然后使用两个权重矩阵和输入的特征进行乘积操作，并通过激活函数处理。

        在这个公式中，我们将所有模态在同一个时间点的特征进行了整合，即在时间 j 的所有模态的特征都放入了这个公式中。这也就体现了这个模型中，在不同模态和时间系列进行信息流动和整合的目的。