Cross-token Modeling with Conditional Computation
This work addresses efficiency and accuracy issues in vision models for researchers and practitioners, representing an incremental improvement over existing methods like MLP-Mixer.
The paper tackles the challenge of scaling cross-token modules like self-attention in transformers by proposing Sparse-MLP, an all-MLP model that uses sparsely-activated MLPs for cross-token modeling, achieving a 2.5% improvement in ImageNet Top-1 accuracy over MLP-Mixer with fewer parameters and computational cost.
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.