1. 从Token-Mixing到Channel-MixingMLP-Mixer的核心设计哲学很多朋友第一次看到MLP-Mixer的论文时第一反应是这也太简单了吧——没有卷积、没有注意力整个网络就是一堆全连接层堆起来的。但真正动手复现的时候才会发现简单的外表下藏着不少设计上的取舍。这一篇接着上一篇的内容重点聊Token-Mixing MLP和Channel-Mixing MLP的具体实现以及我在PyTorch复现过程中踩过的那些坑。先回顾一下MLP-Mixer的整体思路。它把输入图像切成若干个Patch每个Patch经过线性映射变成一个Token所有Token排成一个二维矩阵形状是(num_tokens, hidden_dim)。然后交替做两种操作一种是跨Token方向的全连接让不同空间位置的信息互相交流这叫Token-Mixing MLP另一种是跨Channel方向的全连接让同一个Token内部不同特征维度之间融合这叫Channel-Mixing MLP。两种操作交替堆叠最后接一个全局池化和分类头。这个设计的核心洞察在于空间信息混合和通道信息混合是可以解耦的。传统卷积把这两件事绑在一起做注意力机制也是同时处理空间和通道关系。MLP-Mixer把它们拆开各自用独立的MLP处理参数量反而更可控。理解这一点后面的代码实现就顺理成章了。2. Token-Mixing MLP的实现细节与转置陷阱2.1 Token-Mixing到底在做什么Token-Mixing MLP的作用是让不同空间位置的Token互相交换信息。假设输入张量形状是(batch, num_tokens, hidden_dim)Token-Mixing要做的是对每一个通道维度把所有Token的值拿出来组成一个向量过一个全连接层再放回去。这里有个关键点全连接层的输入维度是num_tokens输出维度也是num_tokens。也就是说它是在Token数量这个维度上做线性变换。如果你用的是16x16的Patch切分那num_tokens 196这个全连接层的权重矩阵就是196 x 196参数量大约3.8万。看起来不大但如果输入分辨率提高Token数量会平方级增长这也是MLP-Mixer在高分辨率下比较吃力的原因之一。2.2 转置操作最容易写错的地方在PyTorch里实现Token-Mixing最自然的写法是先转置再过全连接。我见过不少初版实现是这样的class TokenMixer(nn.Module): def __init__(self, num_tokens, hidden_dim): super().__init__() self.mlp nn.Linear(num_tokens, num_tokens) def forward(self, x): # x: (batch, num_tokens, hidden_dim) x x.transpose(1, 2) # (batch, hidden_dim, num_tokens) x self.mlp(x) # 在最后一维做变换 x x.transpose(1, 2) # 转回来 return x这段代码逻辑上没问题但性能上有个隐患连续的transpose操作会导致张量在内存中变得不连续后续的矩阵乘法可能触发额外的内存拷贝。我在实测中发现对于num_tokens196, hidden_dim512这种规模加上.contiguous()之后反而比不加还慢因为PyTorch的底层BLAS库对非连续张量的处理已经足够优化了。提示不要盲目加.contiguous()先用torch.cuda.synchronize()配合计时器实测一下不同PyTorch版本和硬件上的表现可能不一样。2.3 权重初始化对训练稳定性的影响Token-Mixing层的权重初始化很关键。如果直接用PyTorch默认的初始化在深层网络中容易出现梯度爆炸。我试过几种方案最后发现比较稳的做法是对Token-Mixing和Channel-Mixing的全连接层都使用截断正态分布初始化标准差设为sqrt(2 / (fan_in fan_out))同时把偏置项初始化为零。这个初始化策略的依据来自论文的附录部分本质上是一种Xavier初始化的变体。实测下来用这个初始化配合适当的LayerNorm训练过程中loss曲线明显更平滑不会出现前几百步loss剧烈震荡的情况。3. Channel-Mixing MLP的GELU选择与维度扩展比3.1 为什么需要维度扩展Channel-Mixing MLP的结构和标准Transformer里的FFN几乎一样先升维过激活函数再降维。论文里默认的扩展比是4也就是说如果hidden_dim512中间层就是2048。这个4倍扩展比不是随便定的。从信息论的角度看升维是为了给非线性变换提供更大的表达空间降维则是为了压缩信息、提取精华。我在实验中发现如果把扩展比降到2模型在中小规模数据集上的表现会明显下降但如果提到8参数量翻倍的同时收益却很有限。所以4是一个比较平衡的选择。3.2 GELU vs ReLU实测对比论文用的是GELU激活函数。我一开始图省事用了ReLU结果发现收敛速度慢了大约15%左右。后来换成GELU训练曲线明显更陡。GELU的公式是x * Φ(x)其中Φ(x)是标准正态分布的累积分布函数。PyTorch里可以直接用nn.GELU()也可以手动实现近似版本。class ChannelMixer(nn.Module): def __init__(self, hidden_dim, expansion_ratio4): super().__init__() mid_dim hidden_dim * expansion_ratio self.fc1 nn.Linear(hidden_dim, mid_dim) self.act nn.GELU() self.fc2 nn.Linear(mid_dim, hidden_dim) def forward(self, x): return self.fc2(self.act(self.fc1(x)))这段代码看起来简单但有几个细节值得注意。第一fc1和fc2之间不要加LayerNormLayerNorm应该放在整个Mixer Block的输入端。第二GELU的计算量比ReLU大如果对推理速度有极致要求可以考虑用nn.ReLU()替代但训练阶段建议还是用GELU。3.3 Dropout的位置与比例Dropout加在哪里很有讲究。我试过三种方案加在fc1之后、加在fc2之后、两边都加。实测下来只在fc1之后加Dropout效果最好比例设在0.1到0.3之间。两边都加会导致欠拟合尤其是当训练数据不够大的时候。4. Mixer Block的残差连接与LayerNorm前置设计4.1 Pre-Norm还是Post-NormMLP-Mixer用的是Pre-Norm结构也就是LayerNorm放在残差分支的输入端而不是输出端。这个选择和Transformer的原始设计不同但和后来的很多改进版本一致。Pre-Norm的好处是训练更稳定尤其是在深层网络中。我做过一个对比实验12层的Mixer用Post-Norm的时候训练到第6层左右梯度就开始异常换成Pre-Norm之后梯度范数一直保持在合理范围内。所以如果你要复现MLP-Mixer强烈建议用Pre-Norm。4.2 残差连接的必要性残差连接在MLP-Mixer里不是可选项而是必需品。没有残差连接的话即使只有几层模型也几乎训不起来。原因很简单Token-Mixing和Channel-Mixing都是线性变换加非线性激活如果没有残差信息在多层传递后会严重衰减。class MixerBlock(nn.Module): def __init__(self, num_tokens, hidden_dim, expansion_ratio4, dropout0.1): super().__init__() self.norm1 nn.LayerNorm(hidden_dim) self.token_mixer TokenMixer(num_tokens, hidden_dim) self.norm2 nn.LayerNorm(hidden_dim) self.channel_mixer ChannelMixer(hidden_dim, expansion_ratio) self.dropout nn.Dropout(dropout) def forward(self, x): # Token-Mixing分支 residual x x self.norm1(x) x self.token_mixer(x) x self.dropout(x) x residual x # Channel-Mixing分支 residual x x self.norm2(x) x self.channel_mixer(x) x self.dropout(x) x residual x return x这段代码里有个容易忽略的点两个分支的Dropout是共享的还是一个实例我建议用两个独立的Dropout实例因为两个分支的随机性应该分开。虽然共享一个实例在代码上更简洁但实测效果略差。4.3 LayerNorm的维度选择LayerNorm的归一化维度是hidden_dim也就是对每个Token内部的特征做归一化。这一点和Token-Mixing的维度正好相反不要搞混。如果你不小心把LayerNorm的维度设成了num_tokens模型也能跑但效果会差很多因为那样归一化的是不同空间位置的信息破坏了Token之间的独立性。5. 整体网络组装与分类头的实现5.1 Patch Embedding的实现方式Patch Embedding可以用两种方式实现一种是用nn.Conv2d配合合适的kernel_size和stride另一种是用nn.Unfold加线性层。我推荐用nn.Conv2d因为它在GPU上的效率更高而且代码更简洁。class PatchEmbedding(nn.Module): def __init__(self, img_size224, patch_size16, in_channels3, hidden_dim512): super().__init__() num_patches (img_size // patch_size) ** 2 self.proj nn.Conv2d(in_channels, hidden_dim, kernel_sizepatch_size, stridepatch_size) self.num_patches num_patches def forward(self, x): # x: (batch, 3, 224, 224) x self.proj(x) # (batch, hidden_dim, 14, 14) x x.flatten(2) # (batch, hidden_dim, 196) x x.transpose(1, 2) # (batch, 196, hidden_dim) return x这里flatten(2)把空间维度展平然后转置把Token维度放到中间。注意num_patches的计算(224 // 16) ** 2 196这个值后面Token-Mixing层要用到。5.2 分类头的设计分类头很简单先对所有Token做全局平均池化然后过一个LayerNorm最后接一个全连接层输出类别数。class MLPMixer(nn.Module): def __init__(self, img_size224, patch_size16, in_channels3, num_classes1000, hidden_dim512, num_layers8, expansion_ratio4, dropout0.1): super().__init__() self.patch_embed PatchEmbedding(img_size, patch_size, in_channels, hidden_dim) num_tokens self.patch_embed.num_patches self.blocks nn.ModuleList([ MixerBlock(num_tokens, hidden_dim, expansion_ratio, dropout) for _ in range(num_layers) ]) self.norm nn.LayerNorm(hidden_dim) self.head nn.Linear(hidden_dim, num_classes) def forward(self, x): x self.patch_embed(x) for block in self.blocks: x block(x) x self.norm(x) x x.mean(dim1) # 全局平均池化 x self.head(x) return x这个结构就是MLP-Mixer的完整实现。看起来确实不复杂但每个组件的设计都有其道理。6. 训练调参中的实际经验与常见问题6.1 学习率与优化器的选择MLP-Mixer对学习率比较敏感。我用AdamW优化器学习率设在1e-3到3e-4之间配合余弦退火调度。如果学习率超过1e-2loss会直接飞掉如果低于1e-4收敛速度慢得让人抓狂。权重衰减建议设在0.05左右这个值比Transformer常用的0.01要大一些。原因是MLP-Mixer的全连接层参数量大需要更强的正则化来防止过拟合。6.2 批量大小与梯度累积在单卡显存有限的情况下批量大小可能只能设到32或64。这时候可以用梯度累积来模拟更大的批量。我的经验是累积步数设为4等效批量大小128训练效果和直接用128的批量差不多。注意用梯度累积的时候LayerNorm的统计量是在小批量上计算的这可能会引入一些噪声。如果发现训练不稳定可以试试用SyncBatchNorm或者干脆减小累积步数。6.3 常见报错与排查报错一维度不匹配。最常见的是Token-Mixing层的输入维度设错了。检查num_tokens是否和Patch Embedding的输出一致。报错二显存溢出。MLP-Mixer的参数量随num_tokens平方增长如果输入分辨率是448x448num_tokens会变成784Token-Mixing层的参数量直接翻四倍。解决办法是减小hidden_dim或者减少层数。报错三训练loss不下降。先检查LayerNorm的位置是否正确再检查初始化是否合理。如果都没问题试试把学习率调小一个数量级。6.4 推理阶段的优化推理阶段可以把Dropout关掉这是基本操作。另外如果部署环境对延迟敏感可以考虑把Token-Mixing和Channel-Mixing的权重合并减少矩阵乘法的次数。不过这个优化需要手动推导实现起来比较麻烦收益也有限。我在实际项目里还发现一个小技巧把hidden_dim设成8的倍数GPU的利用率会更高。这是因为很多GPU的Tensor Core对维度有对齐要求虽然PyTorch会自动处理但对齐之后确实能快几个百分点。7. 关于MLP-Mixer适用场景的一些个人判断MLP-Mixer刚出来的时候很多人觉得它能取代卷积和注意力。但实际用下来我的判断是它在中等规模数据集上表现不错尤其是当数据量足够大、分辨率不太高的时候。但如果你的任务需要处理高分辨率输入或者对局部细节非常敏感卷积网络可能还是更合适的选择。另外MLP-Mixer的参数量对num_tokens非常敏感这意味着它不太适合输入尺寸变化很大的场景。如果你要做多尺度训练可能需要额外处理Token数量的对齐问题。我在一个小规模图像分类任务上对比过MLP-Mixer和ResNet在数据量充足的情况下MLP-Mixer的top-1准确率能比ResNet高1到2个百分点但训练时间多了大约30%。所以选型的时候要权衡精度和效率。最后分享一个调试技巧如果你不确定Token-Mixing和Channel-Mixing是否真的在起作用可以把其中一个的权重冻结观察性能变化。我试过冻结Token-Mixing准确率直接掉了8个百分点说明空间信息混合对这个任务确实关键。这种消融实验虽然简单但能帮你快速理解每个组件的重要性。
