告别注意力机制：MLP-Mixer如何用全MLP架构颠覆视觉识别？

最新推荐文章于 2026-04-24 10:00:47 发布
原创 最新推荐文章于 2026-04-24 10:00:47 发布 · 960 阅读 ·
CC 4.0 BY-SA版权
版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。

告别注意力机制：MLP-Mixer如何用全MLP架构颠覆视觉识别？

【免费下载链接】vision_transformer 【免费下载链接】vision_transformer 项目地址: https://gitcode.com/gh_mirrors/vi/vision_transformer

你是否还在为Transformer模型的复杂注意力机制感到困扰？是否想过用简单的多层感知机（MLP）就能实现媲美Transformer的视觉识别效果？本文将带你深入探索MLP-Mixer架构的工作原理，通过具体代码示例和配置参数，展示如何从零开始构建这一革命性的视觉模型。读完本文，你将能够理解MLP-Mixer的核心设计思想，掌握其实现细节，并了解如何在实际项目中应用这一架构。

MLP-Mixer架构概述

MLP-Mixer是一种完全基于多层感知机（MLP）的视觉识别架构，它摒弃了传统CNN中的卷积操作和Transformer中的自注意力机制，仅通过简单的矩阵变换和非线性激活函数实现视觉特征学习。这种极简设计不仅降低了模型复杂度，还在多个视觉任务上取得了与Transformer相当的性能。

MLP-Mixer架构示意图

MLP-Mixer的核心组件包括：

  • Patch Embedding（补丁嵌入）：将输入图像分割为固定大小的补丁，并将每个补丁映射到低维向量空间
  • Mixer Block（混合器块）：由Token Mixing和Channel Mixing两个MLP层组成，分别对空间维度和通道维度进行特征混合
  • Classification Head（分类头）：对最终的特征向量进行分类预测

核心组件详解

Patch Embedding：图像分块与嵌入

Patch Embedding是MLP-Mixer的第一步，它将输入图像分割为多个固定大小的补丁，并通过线性变换将每个补丁映射到低维向量空间。这一过程类似于Transformer中的Patch Embedding，但实现更加简单直接。

# 代码来源：[vit_jax/models_mixer.py](https://link.gitcode.com/i/f595f35e359725391afd64471f913c79)
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')

在这段代码中，我们使用卷积操作实现Patch Embedding：

  • self.patches.size指定了补丁的大小（如16x16）
  • 卷积核大小和步长均设置为补丁大小，确保每个补丁被映射为一个特征向量
  • einops.rearrange将输出特征图的空间维度展平，得到形状为(batch_size, num_patches, hidden_dim)的特征矩阵

Mixer Block：特征混合的核心单元

Mixer Block是MLP-Mixer的核心组件，它由Token Mixing和Channel Mixing两个子模块组成，分别负责对空间维度和通道维度的特征进行混合。

# 代码来源：[vit_jax/models_mixer.py](https://link.gitcode.com/i/f595f35e359725391afd64471f913c79)
class MixerBlock(nn.Module):
  """Mixer block layer."""
  tokens_mlp_dim: int
  channels_mlp_dim: int

  @nn.compact
  def __call__(self, x):
    # Token Mixing
    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
    
    # Channel Mixing
    y = nn.LayerNorm()(x)
    return x + MlpBlock(self.channels_mlp_dim, name='channel_mixing')(y)

Token Mixing和Channel Mixing的工作原理：

  1. Token Mixing：通过转置操作将通道维度和空间维度交换，然后使用MLP对空间维度进行特征混合
  2. Channel Mixing：直接对通道维度进行MLP特征混合
  3. 两者均采用残差连接，有助于缓解深层网络的梯度消失问题

MLP Block：简单而强大的特征变换

MlpBlock是Mixer Block中的基本单元，它由两个全连接层和GELU激活函数组成，实现对输入特征的非线性变换。

# 代码来源：[vit_jax/models_mixer.py](https://link.gitcode.com/i/f595f35e359725391afd64471f913c79)
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)

MlpBlock的设计非常简洁：

  • 第一层将输入特征映射到更高维度空间（mlp_dim）
  • GELU激活函数引入非线性变换
  • 第二层将特征映射回原始维度

完整架构实现

将上述组件组合起来，我们就得到了完整的MLP-Mixer架构：

# 代码来源：[vit_jax/models_mixer.py](https://link.gitcode.com/i/f595f35e359725391afd64471f913c79)
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
    # Patch Embedding
    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')
    
    # Mixer Blocks
    for _ in range(self.num_blocks):
      x = MixerBlock(self.tokens_mlp_dim, self.channels_mlp_dim)(x)
    
    # Classification Head
    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

配置参数解析

MLP-Mixer的性能很大程度上取决于超参数的选择。下面是一个在CIFAR-10数据集上训练Mixer-B/16模型的配置示例：

# 代码来源：[vit_jax/configs/mixer_base16_cifar10.py](https://link.gitcode.com/i/85d0fd3a6f0222ecd41436f98548ca4a)
def get_config():
  """Returns config for training Mixer-B/16 on cifar10."""
  config = common.get_config()
  config.model_type = 'Mixer'
  config.model = models.get_mixer_b16_config()
  config.dataset = 'cifar10'
  config.total_steps = 10_000
  config.pp = ml_collections.ConfigDict(
      {'train': 'train[:98%]', 'test': 'test', 'crop': 224})
  return config

关键参数解析：

  • 模型架构：Mixer-B/16表示基础版模型，补丁大小为16x16
  • 训练步数：总共训练10,000步
  • 数据预处理：训练集使用98%的数据，测试集使用全部数据，图像裁剪大小为224x224

Mixer-B/16的具体架构参数：

  • num_blocks=12：12个Mixer Block
  • hidden_dim=768：隐藏层维度为768
  • tokens_mlp_dim=384：Token Mixing MLP的中间层维度
  • channels_mlp_dim=3072：Channel Mixing MLP的中间层维度

性能评估与应用场景

MLP-Mixer在多个视觉任务上表现出色，特别是在图像分类任务上，与同等规模的Transformer模型性能相当，但计算复杂度更低。根据model_cards/lit.md中的评估结果，类似架构的模型在ImageNet数据集上的零样本分类准确率可达75.7%。

Vision Transformer架构示意图

MLP-Mixer的优势场景：

  • 资源受限设备：由于没有复杂的注意力机制，更适合在计算资源有限的设备上部署
  • 大规模数据训练：简单的架构更容易扩展到更大规模的数据集
  • 迁移学习任务：预训练的MLP-Mixer模型可以很容易地迁移到其他视觉任务

快速上手指南

要在自己的项目中使用MLP-Mixer，只需按照以下步骤操作：

  1. 克隆项目仓库：

    git clone https://gitcode.com/gh_mirrors/vi/vision_transformer

  2. 查看完整实现代码：

  3. 参考Jupyter Notebook示例：

总结与展望

MLP-Mixer架构展示了一种极简主义的视觉识别方法，仅使用MLP就能实现与Transformer相当的性能。这种设计不仅降低了模型复杂度，还为视觉识别提供了新的思路：特征混合不一定需要复杂的注意力机制，简单的MLP同样可以实现强大的特征学习能力。

随着研究的深入，MLP-Mixer的变种模型不断涌现，如ResMLP、gMLP等，进一步推动了纯MLP架构在视觉领域的应用。未来，我们有理由相信，这种简洁高效的设计思想将在更多领域得到应用。

如果你对MLP-Mixer感兴趣，可以通过以下资源深入学习：

希望本文能帮助你理解MLP-Mixer的核心思想和实现细节，激发你在视觉识别领域的创新思考！

【免费下载链接】vision_transformer 【免费下载链接】vision_transformer 项目地址: https://gitcode.com/gh_mirrors/vi/vision_transformer

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

实付元
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明：

1.余额是钱包充值的虚拟货币，按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载，可以购买VIP、付费专栏及课程。

余额充值