Residual Mixture of Experts

TL;DR

This paper tackles the high training cost of large Mixture-of-Experts (MoE) vision transformers by introducing Residual Mixture of Experts (RMoE), which factorizes MoE weights into an input-independent core learned from non-MoE pretraining and an input-dependent residual learned during downstream finetuning. RMoE adopts a bilevel optimization framework and practical design choices (weight inheritance, stop-gradient alignment, and Firefly-based layer selection) to achieve MoE-like performance with substantially lower training cost. Across segmentation and object detection tasks, and on multiple backbones (Swin-T/CvT-13/Swin-L/BeiT-L), RMoE delivers notable accuracy gains with less than 3% extra training expense, and maintains favorable inference costs. The work provides actionable guidance (Every-Last layer placement, 8 experts, top-1 gate) for practitioners aiming to scale vision transformers efficiently in downstream scenarios.

Abstract

Mixture of Experts (MoE) is able to scale up vision transformers effectively. However, it requires prohibiting computation resources to train a large MoE transformer. In this paper, we propose Residual Mixture of Experts (RMoE), an efficient training pipeline for MoE vision transformers on downstream tasks, such as segmentation and detection. RMoE achieves comparable results with the upper-bound MoE training, while only introducing minor additional training cost than the lower-bound non-MoE training pipelines. The efficiency is supported by our key observation: the weights of an MoE transformer can be factored into an input-independent core and an input-dependent residual. Compared with the weight core, the weight residual can be efficiently trained with much less computation resource, e.g., finetuning on the downstream data. We show that, compared with the current MoE training pipeline, we get comparable results while saving over 30% training cost. When compared with state-of-the-art non- MoE transformers, such as Swin-T / CvT-13 / Swin-L, we get +1.1 / 0.9 / 1.0 mIoU gain on ADE20K segmentation and +1.4 / 1.6 / 0.6 AP gain on MS-COCO object detection task with less than 3% additional training cost.

Figures (6)

• Figure 1: RMoE: balancing performance improvement and additional training cost between upper-bound MoE and lower-bound non-MoE training pipelines. RMoE-D and RMoE-I are two variants of RMoE.
• Figure 2: Training pipelines for (a) non-MoE, (b)(c) RMoE and (d) MoE. Here we use 3 experts as an example to illustrate how RMoE works in the intermediate and downstream finetune stages. Compared with MoE training, we inherit the non-MoE transformer after the upstream pretraining to save the cost. In this figure, we simplify the general transformer block design by only showing self-attention (SA) and multilayer perceptron (MLP) modules. Norm and add operations are not shown here.
• Figure 3: Visualization of experts weights evolution during an MoE training process. Here we show the visualization of the experts in the last layer of an 8-experts MoE Swin-T. We project the weights to a 2D plane using principle components analysis. Each point is the weight of an expert at a training epoch. The X and Y axis are the first and second principle components, respectively: (a) in the training process, experts weights at different epochs are clustered and the variance of each weight cluster (points of the same color) is much smaller than the variance of cluster centers (blue solid line) during the training. Motivated by this pattern, our RMoE on (b) first learns the centers of experts' weights (red solid line) and then learns their residual weights (dash blue line).
• Figure 4: Balance loss on the ADE20K segmentation tasks with different loss weights. Weight=0 denotes removing the balance loss. Both RMoE-I and MoE can get a good balance under a small weight, but RMoE-D can only get balanced when a large balance loss weight is used.
• Figure 5: Visualization of experts specialization of RMoE and MoE training. A darker color in a cell indicates that the expert is more specialized to tackle the images in that class of ImageNet. We separately reorder the ImageNet classes for RMoE and MoE to get a better visualization.