[论文解读] Cross-token Modeling with Conditional Computation

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

Advanced Neural Network Applications被引用 4

一句话总结

该论文提出 Sparse-MLP，一种基于两阶段专家混合（MoE）的全 MLP 视觉模型，用于跨标记建模：通过标记混合 MoE 建模空间上下文，通过通道混合 MoE 实现通道间信息流动。该模型在参数量减少 62.8%、预训练成本降低 85.5% 的情况下，实现了 79.2% 的 ImageNet Top-1 准确率，优于 MLP-Mixer 2.5%，且计算成本更低。

ABSTRACT

Mixture-of-Experts (MoE), a conditional computation architecture, achieved promising performance by scaling local module (i.e. feed-forward network) of transformer. However, scaling the cross-token module (i.e. self-attention) is challenging due to the unstable training. This work proposes Sparse-MLP, an all-MLP model which applies sparsely-activated MLPs to cross-token modeling. Specifically, in each Sparse block of our all-MLP model, we apply two stages of MoE layers: one with MLP experts mixing information within channels along image patch dimension, the other with MLP experts mixing information within patches along the channel dimension. In addition, by proposing importance-score routing strategy for MoE and redesigning the image representation shape, we further improve our model's computational efficiency. Experimentally, we are more computation-efficient than Vision Transformers with comparable accuracy. Also, our models can outperform MLP-Mixer by 2.5\% on ImageNet Top-1 accuracy with fewer parameters and computational cost. On downstream tasks, i.e. Cifar10 and Cifar100, our models can still achieve better performance than baselines.

研究动机与目标

• 通过用条件计算替代跨标记注意力，解决基于 MoE 的视觉 Transformer 训练不稳定性问题。
• 设计一种可扩展、高效的全 MLP 架构，在空间（标记）和通道维度均应用 MoE，以提升建模能力。
• 通过重要性分数路由和表示形状重构，提升 MoE 效率，降低路由计算开销。
• 在 ImageNet 和下游视觉基准上实现 SOTA 性能，同时参数量和计算成本低于 ViT 和 MLP-Mixer。

提出的方法

• 提出标记混合 MoE（MoE_S），在每个通道内对补丁维度应用稀疏专家，以建模空间上下文。
• 提出通道混合 MoE（MoE_C），在每个补丁内对通道维度应用稀疏专家，以建模通道间交互。
• 提出一种重要性分数路由策略，通过重要性分数对标记或通道进行排序，以减少路由计算。
• 对稀疏模块中的图像表示形状进行重构，以提升标记混合 MoE 门控网络的效率。
• 在稀疏模块中引入缩放子层，以降低计算成本并提升训练稳定性。
• 使用 MoCo v3 在 ImageNet-1k 上进行自监督预训练，以评估下游性能。

实验结果

研究问题

• RQ1 能否在视觉 Transformer 中有效且稳定地应用专家混合（MoE）进行跨标记建模，以替代自注意力机制？
• RQ2 在全 MLP 架构中同时在空间和通道维度应用 MoE，是否能提升模型容量与效率？
• RQ3 重要性分数路由能否在不损失模型性能的前提下减少 MoE 的路由计算？
• RQ4 重构图像表示形状是否能提升标记混合 MoE 的效率？
• RQ5 采用两阶段 MoE 的全 MLP 模型能否在准确率和效率上超越 ViT 和 MLP-Mixer？

主要发现

• Sparse-B 在 ImageNet 上实现 77.9% 的 Top-1 准确率，比 MLP-Mixer-B/16 高 2.0%，计算成本相当；比 ViT-B/16 高 1.2%，且成本更低。
• Sparse-L 实现 79.2% 的 Top-1 准确率，比 MLP-Mixer-L/16 高 2.5%，参数量减少 62.8%，预训练成本降低 85.5%。
• Sparse-L 比 ViT-L/16 高 1.8%，且预训练成本不足其一半。
• 消融实验表明，缩放子层将预训练成本降低 30.5%，并将准确率从 76.9% 提升至 77.9%。
• 重要性分数路由在保持模型容量的同时减少 MoE 路由计算，显著提升效率。
• 两阶段 MoE 设计（空间与通道的条件计算）被证明有效，使全 MLP 模型在性能与效率上均超越 ViT 和 MLP-Mixer。