On the effectiveness of MoE-enhanced transformer for accurate and generalizable mask-based semantic segmentation
摘要
The Mixture-of-Experts (MoE) paradigm is a promising approach to efficiently scale neural networks, especially for large-scale Transformers. This paper investigates integrating MoE into the state-of-the-art mask-based semantic segmentation models, MaskFormer and Mask2Former. We hypothesize that routing inputs to specialized experts will enhance performance in terms of accuracy and domain generalization. We present two novel architectures, Mask-MoE and Mask2-MoE, which embed the MoE paradigm within the transformer decoder to foster specialization among experts for different semantic concepts. Our experiments show that these architectures significantly improve in both in-domain accuracy and cross-domain generalization performance. On the Cityscapes benchmark, our models achieve a better accuracy-efficiency trade-off by reducing confusion between visually similar classes. Furthermore, in a sim-to-real zero-shot domain generalization experiments, our MoE-enhanced models outperform baselines. Our in-depth investigation shows that this improvement stems from a learned, domain-agnostic routing policy in which the experts consistently divide their labor across domains.