From Sparse to Soft Mixtures of Experts

This work focuses on a number of issues in vanilla Sparse MoEs, and proposes Soft MoE. Instead of hard assignment of tokens, Soft MoE “softs” the tokens by weighted averaging them before assigned to experts. Experiments with ViT models and CV tasks demonstrate that Soft MoE outperforms both Sparse MoEs and dense models.

Instead of sending each expert a bag of vectors from the input the vectors entering each expert are composed of a weighted sum of all input vectors. My understanding of the implementation was assisted by Phil Wang's replication in PyTorch, which I have shortened below:

def forward(self, x):
        """
        einstein notation
        b - batch
        n - sequence length
        e - number of experts
        s - number of slots per expert
        d - feature dimension
        """
        slot_embeds = self.slot_embeds

        logits = einsum('b n d, e s d -> b n e s', x, slot_embeds)

        # get dispatch and combine weights (softmax across right dimensions)

        dispatch_weights = logits.softmax(dim = 1)

        combine_weights = rearrange(logits, 'b n e s -> b n (e s)')
        combine_weights = combine_weights.softmax(dim = -1)

        # derive slots by weighted average of input tokens using the dispatch weights from above

        slots = einsum('b n d, b n e s -> b e s d', x, dispatch_weights)

        # route the slots per expert to each expert

        out = self.experts(slots) # vectors are mapped across dim e to each expert

        # combine back out

        out = rearrange(out, ' b e s d -> b (e s) d')
        out = einsum('b s d, b n s -> b n d', out, combine_weights)

The contributions of this paper are mainly the experimental results, the experiments aim to demonstrate that a soft-MoE layer allows faster inference and shorter training time on some large scale training benchmarks. Specifically:

• Soft MoE L/16 vs ViT H/14:
  • Half the training time
  • 2x faster inference
• Soft MoE B/16 vs ViT/14:
  • 5.7x faster at inference time

These results are demonstrated comprehensively with ablation experiments investigating the different assumptions, such as whether the weighted sum is necessary, whether the weighted sum can be uniform and whether the weighted sum at the output is necessary; finding that the decisions all improve performance.

The results are demonstrated on expensive large scale benchmarks, including training many ViT style models on JFT-4B and further experiments testing these models on a contrastive benchmark using WebLI. The results presented support the contributions of the paper, demonstrating improved performance across the experiments compared to other MoE methods.

The authors also note that the proposed method would violate the causal assumption if used in an autoregressive model, meaning it can't be used as defined above for autoregressive language modelling.

This paper introduces a novel soft MOE method that keeps the benefits of sparse MOE (processing a subset of tokens for lower cost), and also has full differentiability and better performance. The proposed method let the MOE experts to process a number of slots, where each slot is a combination of all tokens. The idea of processing a number of slots is how the inference cost is controlled. It outperforms dense (non MOE) networks and also typical sparse MOE networks: with similar inference cost, the accuracy metric is much better.

The following paper proposes Soft Mixture-of-Experts (MoE), a fully-differentiable sparse transformers capable of passing different weighted combinations of all input tokens to each expert. Compared to other MoEs, experts in Soft MoE only process a subset of (combined) tokens, resulting in larger model capacity (and performance) with lower inference cost. Thanks to this, Soft MoE are able to alleviate training instability, token dropping, scalability issues, and ineffective finetuning process that is apparent in previous type of MoE. Experiments on visual recognition tasks in Section 3 and 4 show the superiority of Soft MoE over dense Vision Transformers (ViT) and popular MoEs not limited to Tokens Choice and Experts Choice. In addition to that, Soft MoE provably scales well; since while the Huge/14 version of it has $40\times$ more parameters than ViT Huge/14, it bring substantial improvement in terms of classification accuracy with only 2% gain in inference.

We are glad of the mostly positive comments and scores from the reviewers, and we are very thankful for their very useful comments to help to improve our submission.

• All reviewers commented on the lack of natural language processing results and/or the challenges of applying Soft MoEs in this domain.

Indeed, NLP and text auto-regressive tasks in particular, are a highly important application in current ML trends. We want to emphasize that we are actively working on this, but we decided to limit the scope of the paper to vision tasks for a few reasons:

1. By narrowing the scope of domains, we can do a more in-depth analysis of Soft MoEs in vision tasks (see all the detailed investigation in the appendix).
2. Limiting the scope of tasks also makes the presentation of the experiments, and the overall reading and understanding of the model clearer (in our opinion), and allows one to focus on the fundamental aspects of the proposed method (we are glad that most reviewers value the quality of the presentation very highly).
3. We want to present our current (very good) results to the ML community, and encourage it to explore and improve the proposed approach for NLP and many other modalities and tasks. Certainly, tens or hundreds of researchers can make faster progress than just a handful, we believe.

• Reviewer HCxh commented:

Unfortunately, the dataset that is being used for training is not publicly available; making it hard to be used for benchmarking with other papers.

We have used the rebuttal period to run additional experiments pre-training on the publicly available LAION-400M dataset (using SigLIP pre-training, as described in [1]). The results are consistent with those obtained on JFT. Other researchers will be able to replicate our results, since our code is open-sourced and we will release the configuration file to run the LAION pre-training.

Here we show different downstream results obtained from LAION-400M pre-trained models: https://drive.google.com/file/d/16EN8jDbXxLC4-CGxWpBeP0EuU4dBIIva/view

As the plot shows, Soft MoE has better quality than both ViT and Experts Choice, when matching the number of steps (which also roughly matches training cost). We will include these (and more detailed) results obtained on this publicly available dataset in the final version of the paper, most likely in the appendix due to space limitations.

• Reviewer pKE3 said:

When assessing the training and inference speed, the authors chose dense ViT as the baseline. It would be better to include the results for Sparse MoEs additionally.

We compare quality vs. comparing training and inference speed of Sparse MoEs and Dense in Figure 3 (main paper) and Table 9 (appendix J). For long training runs, we compare Soft MoE only with the already published results of ViT in [2], since these runs are very costly and we couldn’t afford to train all methods for millions of steps.

• Reviewer wHYN made some comments on the presentation.

Regarding Figure 2: we tried to imply that the dispatch and combine weights were used to map input tokens to slots (dispatch) and slots to output tokens (combine). We’ll try to find a better representation of this, or annotate the image for further clarification.

Regarding time complexity: indeed Figure 6 shows the fact that the time complexity of Sparse MoE is higher than Soft MoEs, due to the sorting operations: when the number of experts grows a lot, the sorting operations start to dominate the runtime and the throughput decreases a lot for Experts Choice and Tokens Choice. It’s a bit hard to visualize the time complexity in a single plot, since it depends on multiple factors such as the number of experts, the group size, and the model dimension. We’ll add additional plots in the appendix trying to visualize these dependencies.

We’ll also add additional references to where to find the supporting results when claims are made.

• Reviewer fyPo commented:

I hope the authors can release as much details as possible of the implementation details. Looking forward to the adaptation into different fields.

We have already open-sourced our implementation, including details about the training hyperparameters used in this paper. We have excluded the link to the code in the draft to avoid breaking double-blind review process, but this will be part of the final version. We will also include the configuration files to replicate the LAION experiments mentioned above.

We hope that we have addressed your questions and concerns. Let us know if any further clarification is needed.

[1] Sigmoid Loss for Language Image Pre-Training, https://arxiv.org/abs/2303.15343

[2] Scaling Vision Transformers, https://arxiv.org/abs/2106.04560

The paper introduces Soft Mixture-of-Experts (MoE), a novel approach that combines the benefits of sparse MoEs with full differentiability, addressing issues such as training instability, token dropping, scalability, and finetuning challenges seen in previous MoE models. Soft MoE allows experts to process a subset of tokens, reducing inference costs while maintaining model capacity and performance. Experimental results show that Soft MoE significantly outperforms dense networks and conventional sparse MoE models in terms of classification accuracy while offering faster inference and shorter training times, particularly in comparison to ViT models.

The reviewers were concerned lack of natural language processing applications, lack of results on public benchmark datasets, and choice of baseline.

Despite the concerns, all the reviewers are positive. Most of the reviewers' concerns and questions have been well-addressed by the authors through the rebuttal.

The highlight of this work is that the authors demonstrate a simple design choice can help effectively tackle many limitations of vanilla Sparse MoEs.

Accept (spotlight)