[Paper Review] Cross-token Modeling with Conditional Computation

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

Advanced Neural Network Applications4 citations

TL;DR

This paper proposes Sparse-MLP, an all-MLP vision model that uses two-stage Mixture-of-Experts (MoE) for cross-token modeling: token-mixing MoE for spatial context and channel-mixing MoE for channel-wise information flow. It achieves 79.2% ImageNet Top-1 accuracy with 62.8% fewer parameters and 85.5% less pretraining cost than ViT-L/16, outperforming MLP-Mixer by 2.5% with lower computational cost.

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.

Motivation & Objective

• To address the instability of training MoE-based self-attention in vision transformers by replacing cross-token attention with conditional computation.
• To design a scalable, efficient all-MLP architecture that applies MoE in both spatial (token) and channel dimensions for improved modeling capacity.
• To improve MoE efficiency through importance-score routing and representation shape redesign for reduced routing computation.
• To achieve state-of-the-art performance on ImageNet and downstream vision benchmarks with lower parameter and computational costs than ViT and MLP-Mixer.

Proposed method

• Introduces token-mixing MoE (MoE_S), which applies sparse experts across the patch dimension within each channel to model spatial context.
• Introduces channel-mixing MoE (MoE_C), which applies sparse experts across channels within each patch to model channel-wise interactions.
• Proposes an importance-score routing strategy that ranks tokens or channels by their importance scores to reduce routing computation.
• Redesigns the image representation shape in Sparse blocks to improve the efficiency of the token-mixing MoE gating network.
• Employs rescale sub-layers in Sparse blocks to reduce computational cost and improve training stability.
• Uses MoCo v3 for self-supervised pretraining on ImageNet-1k to evaluate downstream performance.

Experimental results

Research questions

• RQ1 Can Mixture-of-Experts be effectively and stably applied to cross-token modeling in vision transformers, replacing self-attention?
• RQ2 Does applying MoE in both spatial and channel dimensions improve model capacity and efficiency in an all-MLP architecture?
• RQ3 Can importance-score routing reduce MoE routing computation without sacrificing model performance?
• RQ4 Does redesigning the image representation shape enhance the efficiency of token-mixing MoE?
• RQ5 Can an all-MLP model with two-stage MoE outperform ViT and MLP-Mixer in accuracy and efficiency?

Key findings

• Sparse-B achieves 77.9% ImageNet Top-1 accuracy, 2.0% higher than MLP-Mixer-B/16 with comparable computational cost and 1.2% higher than ViT-B/16 with less cost.
• Sparse-L achieves 79.2% Top-1 accuracy, outperforming MLP-Mixer-L/16 by 2.5% with 62.8% fewer parameters and 85.5% less pretraining cost.
• Sparse-L outperforms ViT-L/16 by 1.8% with less than half the pretraining cost.
• The ablation study shows that rescale sub-layers reduce pretraining cost by 30.5% and improve accuracy from 76.9% to 77.9%.
• Importance-score routing reduces MoE routing computation while maintaining model capacity, improving efficiency.
• The two-stage MoE design with spatial and channel-wise conditional computation proves effective, enabling an all-MLP model to surpass both ViT and MLP-Mixer in performance and efficiency.