从零复现MLP-Mixer：纯MLP架构的PyTorch实现与调参指南

1. 从卷积与注意力之外看MLP-Mixer的设计动机第一次看到MLP-Mixer这个结构的时候我的反应和大多数人一样这不就是把Transformer里的注意力换成MLP吗能有多复杂但真正动手复现之后才发现这个看似简单的结构背后藏着不少值得琢磨的设计取舍。MLP-Mixer是2021年那篇同名论文提出的架构核心思路非常直接——既然卷积神经网络靠局部窗口提取特征Vision Transformer靠全局注意力建模关系那能不能用纯MLP同时完成这两件事答案是能而且效果还不差。这篇文章适合谁看如果你已经了解基本的PyTorch操作知道什么是张量的shape变换也大致明白Transformer的基本结构那接下来的内容会让你很顺畅地理解MLP-Mixer的每一个设计细节。如果你对CNN和ViT都有一定认知那更好因为MLP-Mixer本质上是在这两条路线之间找了一条中间道路。我打算分几个层面来拆解先讲清楚MLP-Mixer的整体架构和设计逻辑然后逐模块分析核心组件的实现细节接着给出完整的PyTorch代码和逐行注释最后分享一些我在复现过程中踩过的坑和调参经验。整篇内容会围绕一个目标让你看完之后能自己从零写出一份可运行的MLP-Mixer实现并且理解每一行代码为什么这么写。2. MLP-Mixer整体架构拆解2.1 核心设计思路两条MLP各管一件事MLP-Mixer的核心思想可以用一句话概括用两种不同类型的MLP层交替堆叠分别负责“空间维度上的信息混合”和“通道维度上的信息混合”。这个思路听起来抽象但用生活化的类比就很好理解。想象你手里有一叠扑克牌每张牌代表一个图像patch的特征向量。你有两种操作方式第一种是把同一张牌上的不同花色信息混合起来通道混合第二种是把不同牌之间相同位置的信息混合起来空间混合。MLP-Mixer做的事情就是交替执行这两种操作让信息在空间和通道两个维度上充分流动。具体来说输入图像首先被切成若干个不重叠的patch每个patch展平后通过一个线性层映射到指定维度。假设输入是224×224的图像patch大小设为16×16那就会得到(224/16)×(224/16)196个patch。每个patch被映射成比如512维的向量最终得到一个196×512的矩阵。这个矩阵就是MLP-Mixer的基本工作单元。接下来就是交替执行两种MLPToken-mixing MLP在196这个维度上做MLP也就是让不同patch之间交换信息。操作方式是先转置矩阵变成512×196然后对196这个维度做全连接再转置回来。Channel-mixing MLP在512这个维度上做MLP也就是让同一个patch内部的不同通道交换信息。操作方式是直接对512维度做全连接。这两种操作交替堆叠N次最后通过全局平均池化得到分类结果。整个结构没有卷积、没有注意力只有矩阵乘法、转置和激活函数。2.2 为什么这样设计能work你可能会问这么简单的结构凭什么能和CNN、ViT掰手腕我个人的理解是关键在于信息混合的充分性。CNN通过堆叠卷积层逐步扩大感受野ViT通过自注意力一次性建模全局关系而MLP-Mixer通过交替的token-mixing和channel-mixing让每个位置的信息都有机会和其他所有位置的信息交互。Token-mixing MLP的权重矩阵形状是(196, 196)这意味着每个patch的输出是所有196个patch输入的加权组合。虽然参数量比注意力机制小得多注意力需要计算196×196的注意力矩阵而这里只是一个固定的权重矩阵但效果上确实实现了全局信息混合。另一个关键点是参数共享策略。MLP-Mixer在不同层之间不共享权重但在同一层的token-mixing和channel-mixing内部所有patch共享同一组MLP权重。这种设计既控制了参数量又保证了操作的对称性。2.3 和CNN、ViT的对比为了更直观地理解MLP-Mixer的定位我整理了一个对比表格特性CNNViTMLP-Mixer局部性归纳偏置强弱无全局信息交互逐层扩大单层全局单层全局参数量Base规模中等大中等计算复杂度O(N)O(N²)O(N²)实现难度低中低对数据量的需求低高高从表格可以看出MLP-Mixer在实现难度上比ViT低不少但在数据需求上和大规模ViT类似都需要大量数据才能发挥优势。这也是为什么原论文要在JFT-300M这种规模的数据集上预训练才能达到SOTA效果。3. 核心模块的PyTorch实现细节3.1 Patch Embedding层把图像变成序列Patch Embedding是整个模型的入口作用是把H×W×C的图像转换成N×D的序列。N是patch数量D是隐藏维度。实现方式有两种一种是用卷积kernel_sizestridepatch_size另一种是用 unfold 线性层。我推荐用卷积实现因为代码更简洁而且可以利用cuDNN的优化。import torch import torch.nn as nn class PatchEmbedding(nn.Module): def __init__(self, img_size224, patch_size16, in_channels3, hidden_dim512): super().__init__() self.num_patches (img_size // patch_size) ** 2 self.proj nn.Conv2d( in_channels, hidden_dim, kernel_sizepatch_size, stridepatch_size ) def forward(self, x): # x: (B, C, H, W) x self.proj(x) # (B, D, H/P, W/P) x x.flatten(2) # (B, D, N) x x.transpose(1, 2) # (B, N, D) return x这里有个细节需要注意卷积层的输出通道数就是hidden_dim卷积核大小和步长都等于patch_size。这样每个patch区域被独立映射成一个D维向量且patch之间没有重叠。flatten(2)把空间维度展平transpose把通道维度换到最后得到(B, N, D)的标准序列格式。注意如果你的输入图像尺寸不是patch_size的整数倍卷积操作会自动丢弃边缘部分。建议在数据预处理阶段就把图像resize到合适的大小避免信息浪费。3.2 Token-Mixing MLP跨patch信息交互Token-mixing MLP是MLP-Mixer区别于普通MLP的关键模块之一。它的作用是在patch维度上做全连接让不同空间位置的信息相互混合。实现上需要先转置做完MLP后再转置回来。class TokenMixer(nn.Module): def __init__(self, num_patches, hidden_dim, token_dim, dropout0.0): super().__init__() self.norm nn.LayerNorm(hidden_dim) self.mlp nn.Sequential( nn.Linear(num_patches, token_dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(token_dim, num_patches), nn.Dropout(dropout) ) def forward(self, x): # x: (B, N, D) residual x x self.norm(x) x x.transpose(1, 2) # (B, D, N) x self.mlp(x) # (B, D, N) x x.transpose(1, 2) # (B, N, D) return x residual这里的token_dim通常比num_patches小很多比如num_patches196时token_dim可能设为128或256。这个瓶颈结构的作用是压缩信息减少参数量同时引入非线性变换。LayerNorm放在MLP之前是Pre-Norm的常见做法训练更稳定。实操心得token_dim的选择很关键。设得太小会限制信息混合能力设得太大参数量会爆炸。我的经验是token_dim取num_patches的1/2到1/4之间比较合适。比如196个patchtoken_dim取64到128之间。3.3 Channel-Mixing MLP通道维度信息交互Channel-mixing MLP和标准的Transformer FFN几乎一样就是在通道维度上做两层全连接中间加激活函数。区别在于MLP-Mixer里这个操作是独立于token-mixing的两者交替执行。class ChannelMixer(nn.Module): def __init__(self, hidden_dim, channel_dim, dropout0.0): super().__init__() self.norm nn.LayerNorm(hidden_dim) self.mlp nn.Sequential( nn.Linear(hidden_dim, channel_dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(channel_dim, hidden_dim), nn.Dropout(dropout) ) def forward(self, x): residual x x self.norm(x) x self.mlp(x) return x residualchannel_dim通常是hidden_dim的2到4倍这和Transformer FFN的扩展比例一致。比如hidden_dim512时channel_dim可以设为2048或1024。这个扩展-压缩的结构让模型有足够的容量学习复杂的通道交互模式。3.4 Mixer Block交替堆叠的基本单元一个Mixer Block就是Token-Mixing MLP和Channel-Mixing MLP的串联。注意顺序先做token-mixing再做channel-mixing这个顺序在论文里是固定的但理论上可以交换。我试过交换顺序效果差异不大但保持和论文一致更稳妥。class MixerBlock(nn.Module): def __init__(self, num_patches, hidden_dim, token_dim, channel_dim, dropout0.0): super().__init__() self.token_mixer TokenMixer(num_patches, hidden_dim, token_dim, dropout) self.channel_mixer ChannelMixer(hidden_dim, channel_dim, dropout) def forward(self, x): x self.token_mixer(x) x self.channel_mixer(x) return x每个Mixer Block内部都有残差连接这是深层网络能训练起来的关键。残差连接让梯度可以直接回传避免了梯度消失问题。LayerNorm放在每个子模块之前Pre-Norm这也是当前主流做法。4. 完整模型组装与训练配置4.1 完整MLP-Mixer模型代码把前面所有模块组装起来就得到了完整的MLP-Mixerclass MLPMixer(nn.Module): def __init__( self, img_size224, patch_size16, in_channels3, num_classes1000, hidden_dim512, num_blocks8, token_dim256, channel_dim2048, dropout0.0 ): super().__init__() self.patch_embed PatchEmbedding( img_size, patch_size, in_channels, hidden_dim ) num_patches self.patch_embed.num_patches self.blocks nn.Sequential(*[ MixerBlock(num_patches, hidden_dim, token_dim, channel_dim, dropout) for _ in range(num_blocks) ]) self.norm nn.LayerNorm(hidden_dim) self.head nn.Linear(hidden_dim, num_classes) # 参数初始化 self.apply(self._init_weights) def _init_weights(self, m): if isinstance(m, nn.Linear): nn.init.trunc_normal_(m.weight, std0.02) if m.bias is not None: nn.init.zeros_(m.bias) elif isinstance(m, nn.LayerNorm): nn.init.zeros_(m.bias) nn.init.ones_(m.weight) def forward(self, x): x self.patch_embed(x) x self.blocks(x) x self.norm(x) x x.mean(dim1) # 全局平均池化 x self.head(x) return x这个实现对应的是MLP-Mixer的Base版本配置8个blockhidden_dim512token_dim256channel_dim2048。原论文还有Small、Large、Huge等不同规模的配置主要区别在block数量、隐藏维度和MLP维度上。4.2 不同规模配置的参数对照配置BlocksHidden DimToken DimChannel Dim参数量Small85122562048~18MBase127683843072~59MLarge2410245124096~207MHuge3212806405120~431M参数量计算方式每个Mixer Block的参数量约为2×num_patches×token_dim 2×hidden_dim×channel_dim。以Base配置为例num_patches196token_dim384hidden_dim768channel_dim3072单个block参数量约为2×196×384 2×768×3072 ≈ 150K 4.7M ≈ 4.85M12个block约58M加上embedding和分类头约59M。4.3 训练配置与超参数选择原论文用的是AdamW优化器学习率带余弦退火权重衰减0.05batch size 4096大规模训练。如果我们在小数据集上训练需要调整这些超参数。# 训练配置示例 config { optimizer: AdamW, lr: 1e-3, weight_decay: 0.05, batch_size: 256, epochs: 100, warmup_epochs: 5, scheduler: cosine, mixup: 0.2, cutmix: 1.0, label_smoothing: 0.1, dropout: 0.1 }在小数据集上训练时我建议把dropout调大一些0.1到0.3加上Mixup和CutMix数据增强学习率用1e-3到5e-4之间。如果过拟合严重可以进一步增大weight_decay或减小模型规模。常见问题很多人复现时发现loss不下降或者震荡严重。最常见的原因是学习率太大或者warmup不够。MLP-Mixer对学习率比较敏感建议先用小学习率试跑几个epoch确认loss在稳定下降后再调大。5. 实操过程中的问题排查与调参经验5.1 训练不收敛的排查思路我在第一次复现MLP-Mixer时遇到了loss完全不下降的问题。排查过程如下第一步检查数据预处理。确认图像归一化参数是否正确ImageNet的mean和stdpatch切分有没有问题。我当时的bug是忘了做归一化导致输入值域不对模型完全学不动。第二步检查参数初始化。MLP-Mixer对初始化比较敏感特别是token-mixing MLP的权重。如果初始化方差太大深层网络容易梯度爆炸。用trunc_normal_(std0.02)是比较稳妥的选择。第三步检查学习率和warmup。MLP-Mixer没有卷积的平移不变性归纳偏置训练初期比较脆弱。warmup阶段让学习率从很小的值慢慢升上去给模型一个适应的过程。第四步检查LayerNorm的位置。Pre-Norm比Post-Norm更容易训练如果用了Post-Norm深层网络可能完全训不起来。5.2 过拟合与欠拟合的应对策略MLP-Mixer参数量不小在小数据集上很容易过拟合。我整理了一个问题速查表现象可能原因解决方案训练loss低验证loss高过拟合增大dropout、weight_decay加数据增强训练loss和验证loss都高欠拟合增大模型规模减小正则化检查学习率训练loss震荡学习率太大减小学习率增加warmup验证loss突然飙升梯度爆炸加梯度裁剪检查初始化训练极慢batch size太小或模型太大增大batch size用混合精度5.3 显存优化技巧MLP-Mixer的token-mixing操作需要转置和全连接显存占用比同等参数量的CNN要高。我试过几个有效的优化方法混合精度训练用torch.cuda.amp自动混合精度显存占用可以减少30%到50%速度也有提升。梯度累积如果显存不够放大的batch size可以用梯度累积模拟大batch效果。激活检查点对Mixer Block使用gradient checkpointing用计算换显存。减小token_dimtoken-mixing的参数量和计算量都和token_dim直接相关适当减小可以显著降低显存。# 混合精度训练示例 from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for images, labels in dataloader: optimizer.zero_grad() with autocast(): outputs model(images) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()5.4 学习率调度与warmup的实操细节学习率调度对MLP-Mixer的训练效果影响很大。我用过的最稳的方案是线性warmup加余弦退火def get_lr(epoch, warmup_epochs, total_epochs, base_lr, min_lr1e-6): if epoch warmup_epochs: # 线性warmup return base_lr * (epoch 1) / warmup_epochs else: # 余弦退火 progress (epoch - warmup_epochs) / (total_epochs - warmup_epochs) return min_lr 0.5 * (base_lr - min_lr) * (1 math.cos(math.pi * progress))warmup_epochs一般设为总epoch的5%到10%。比如训练100个epochwarmup设5到10个epoch。base_lr根据batch size调整batch size 256时用1e-3左右batch size 1024时可以用3e-3到5e-3。实操心得如果训练过程中发现验证集准确率波动很大可以在余弦退火的基础上加一点周期性重启或者用ReduceLROnPlateau在验证指标不提升时自动降低学习率。我试过在CIFAR-100上训练加warmup比不加warmup最终准确率高2到3个百分点。6. 从复现到理解的几个关键认知6.1 为什么token-mixing用转置而不是直接操作这个问题我琢磨了很久。Token-mixing需要在patch维度上做全连接但PyTorch的nn.Linear默认在最后一维上操作。输入是(B, N, D)如果直接对N做全连接需要把N换到最后一维所以先transpose(1,2)变成(B, D, N)做完MLP后再transpose回来。这个转置操作在实现上是必要的但也会带来额外的内存拷贝开销。有没有办法避免转置可以用einsum或者手动实现矩阵乘法但可读性和灵活性都不如直接转置。实际测试下来转置带来的开销在整体计算中占比很小不用太担心。6.2 全局平均池化 vs 分类tokenViT用一个额外的分类token来聚合全局信息MLP-Mixer用的是全局平均池化。这两种方式各有优劣分类token更灵活可以学习到更复杂的聚合方式全局平均池化更简单没有额外参数也不需要考虑位置编码。我试过在MLP-Mixer上加分类token效果和全局平均池化差不多但参数量多了一点。如果任务需要更复杂的全局聚合可以试试分类token的方案。6.3 位置编码的必要性MLP-Mixer没有显式的位置编码因为token-mixing MLP的权重矩阵本身就编码了位置信息。每个patch的输出是所有patch的加权和权重矩阵的每一行对应一个输出位置每一列对应一个输入位置。这种全连接的方式天然包含了位置信息不需要额外加位置编码。但这也意味着MLP-Mixer对输入patch的顺序是敏感的。如果你打乱patch顺序输出会完全不同。这和ViT不同ViT如果不加位置编码打乱patch顺序输出不变。6.4 后续扩展方向MLP-Mixer的后续工作有不少有意思的改进。比如有的工作把token-mixing换成稀疏的或局部的降低计算复杂度有的工作把MLP-Mixer和卷积结合引入局部归纳偏置还有的工作把MLP-Mixer用到分割、检测等下游任务上。如果你想在这个方向上做实验我建议先从这几个角度入手一是改进token-mixing的效率比如用低秩分解或稀疏连接二是引入多尺度特征类似CNN的金字塔结构三是探索不同的混合顺序和组合方式。我个人在实际操作中的体会是MLP-Mixer最大的价值不在于它本身有多强而在于它证明了一个事情没有卷积、没有注意力纯MLP也能做视觉任务。这个结论打开了架构设计的新思路后续很多工作都在这个方向上继续探索。如果你正在学习视觉模型架构MLP-Mixer是一个很好的切入点代码量不大结构清晰适合拿来练手和做实验。