[论文解读] Sparse-MLP: A Fully-MLP Architecture with Conditional Computation.

Yuxuan Lou, Fuzhao Xue+2 more|arXiv (Cornell University)|Sep 5, 2021

Domain Adaptation and Few-Shot Learning参考文献 41被引用 5

一句话总结

Sparse-MLP 提出了一种完全基于 MLP 的架构，将稀疏的专家混合（MoE）机制集成到 MLP-Mixer 中，通过两阶段 MoE 层替换部分密集 MLP 块，实现跨维度的特征混合。通过引入 Re-represent 层以增强专家容量和路由效率，该模型在参数更少、计算成本更低的情况下，ImageNet Top-1 准确率比密集 MLP 基线高出 2.5%。

ABSTRACT

Mixture-of-Experts (MoE) with sparse conditional computation has been proved an effective architecture for scaling attention-based models to more parameters with comparable computation cost. In this paper, we propose Sparse-MLP, scaling the recent MLP-Mixer model with sparse MoE layers, to achieve a more computation-efficient architecture. We replace a subset of dense MLP blocks in the MLP-Mixer model with Sparse blocks. In each Sparse block, we apply two stages of MoE layers: one with MLP experts mixing information within channels along image patch dimension, one with MLP experts mixing information within patches along the channel dimension. Besides, to reduce computational cost in routing and improve expert capacity, we design Re-represent layers in each Sparse block. These layers are to re-scale image representations by two simple but effective linear transformations. When pre-training on ImageNet-1k with MoCo v3 algorithm, our models can outperform dense MLP models by 2.5\% on ImageNet Top-1 accuracy with fewer parameters and computational cost. On small-scale downstream image classification tasks, i.e. Cifar10 and Cifar100, our Sparse-MLP can still achieve better performance than baselines.

研究动机与目标

• 通过在 Mixture-of-Experts（MoE）中引入稀疏条件计算，提升 MLP-Mixer 的效率与可扩展性。
• 解决视觉 Transformer 和 MLP-Mixer 架构中密集 MLP 块存在的计算效率低下和容量受限问题。
• 通过稀疏块中的两阶段 MoE 层，实现通道与 patch 维度之间的更有效特征混合。
• 通过引入 Re-represent 层对表示进行重缩放，以减少路由开销并提高专家利用率。
• 在图像分类基准上实现 SOTA 性能，同时参数量和 FLOPs 均低于密集 MLP 模型。

提出的方法

• 将 MLP-Mixer 中的部分密集 MLP 块替换为实现两阶段 MoE 的稀疏块：一阶段用于跨 patch 的通道间混合，另一阶段用于跨通道的 patch 间混合。
• 引入 Re-represent 层——在 MoE 路由前后应用两个简单的线性变换，以重缩放特征表示，提升专家容量和路由效率。
• 在 ImageNet-1k 上使用 MoCo v3 进行预训练，以确保在稀疏 MoE 设置下的鲁棒自监督表示学习。
• 采用路由器机制，为每个 token 在每个 MoE 层中动态选择 top-k 专家，实现条件计算并降低 FLOPs。
• 设计稀疏块以在保持计算效率的同时，通过专家的稀疏激活实现高容量特征学习。
• 通过在专家选择前对输入表示进行稳定和优化，优化路由过程。

实验结果

研究问题

• RQ1 稀疏 MoE 机制能否被有效集成到像 MLP-Mixer 这类完全基于 MLP 的架构中，以提升效率和性能？
• RQ2 两阶段 MoE 设计（分别处理通道和 patch 维度）如何影响特征表示和模型准确率？
• RQ3 Re-represent 层在稀疏 MLP 架构中在多大程度上提升了专家利用率并减少了路由开销？
• RQ4 稀疏 MLP 架构是否能在 ImageNet-1k 和小规模基准上，以更低的参数量和 FLOPs 超越密集 MLP 基线？
• RQ5 所提出的 Sparse-MLP 在 CIFAR-10 和 CIFAR-100 等下游任务上是否能保持强泛化能力，且仅需极少的架构修改？

主要发现

• 在使用 MoCo v3 在 ImageNet-1k 上预训练后，Sparse-MLP 的 ImageNet Top-1 准确率比密集 MLP-Mixer 基线高出 2.5%。
• 与密集 MLP 模型相比，该模型在参数更少、计算成本更低的情况下实现了性能提升。
• 在 CIFAR-10 和 CIFAR-100 上，尽管数据集规模较小，Sparse-MLP 仍优于现有基线模型。
• Re-represent 层有效提升了专家容量，并减少了 MoE 机制中的路由相关计算开销。
• 两阶段 MoE 设计实现了高效的跨维度特征混合，在几乎不增加 FLOPs 的情况下增强了表征学习能力。
• 稀疏架构在多个基准上保持了强大的泛化能力，展现出可扩展性和效率优势。