Multilinear Mixture of Experts: Scalable Expert Specialization through Factorization

We thank Reviewer 1sqG for their positive assessment of the paper; for praising the “innovative” quantitative metrics, the “comprehensive evaluation”, and “competitive performance with added benefits”. We address the stated weaknesses below:

[W1] Limited model/dataset scope

The reviewer is correct that we focus primarily on vision models and ImageNET1k. However, we highlight that our experiments also demonstrate the benefits of the layer on multi-modal vision-language models (CLIP and DINO in Figures 3, 14-17, and 21), and LLMs using the large and diverse Open WebText natural language dataset [1].

[W2] Robustness analysis

Our quantitative evaluation of vision models is only on the test set and does not include out-of-distribution (OOD) data, a limitation we have acknowledged in the conclusion.

However, in our experiments on language models (as shown in Figures 5, 12, and 13), we use sequences generated by large language models (LLMs) as input instead of the training set data. Additionally, we conducted an extra experiment for Reviewer-qC51 (in the rebuttal PDF), where the input prompt is free-form text, demonstrating initial evidence of the experts' generalizability beyond the training and test data.

We thank the reviewer for their suggestion to focus explicitly on robustness. We will accordingly include this explicitly in the limitations as a possible avenue for future work.

[W3 + Q1] Need for implementation details

We highlight that we include the full pseudocode of all muMoE layers in Appendix B, and an anonymous link in the paper's abstract links to our full model source code implementing the layers (please see MuMoE.py in particular).

For full details about the computational requirements, hardware, and configurations used for the experiments, please see Table 6 in the appendix.

[Q2] Large MoE models

We do not anticipate any particular difficulties applying muMoEs to even larger networks.

MuMoEs are indeed a uniquely attractive alternative at very large scales, unlocking (through their parameter efficiency) a route towards achieving fine-grained specialization that is not possible with Sparse/Soft MoEs due to their prohibitively large parameter counts.

However, we unfortunately do not have access to the computational resources to verify this experimentally. As shown in Table 6, we have access to a maximum of 4 A100 GPUs, which limits us to pre-training only the smallest GPT-2 model with 124M parameters. We have added the following to the “limitations” section to acknowledge this:

Despite demonstrating the benefits of scalable expert specialization with muMoE layers for various architectures with different modalities, we have not conducted experiments on the largest scales of commercial LLMs. Future work should explore this where resources permit.

---

• [1] Aaron Gokaslan and Vanya Cohen. Openwebtext corpus, 2019.

We thank Reviewer YTQ3 for engaging thoroughly with the paper. However, we respectfully argue below that the reviewer overlooks our stated goals and claims in the paper, and the benefits to scalable expert specialization the methodology brings that are otherwise infeasible:

In particular, muMoE layers are a parameter-efficient method for scalable expert specialization. MuMoEs aim is to provide interpretability and controllability benefits to dense layers without increasing the parameter or FLOP counts of the original networks. We respectfully highlight that these goals and scope are outlined in the paper’s claims (e.g. [L17, L81, L353]), and that two reviewers praise this explicitly:

• Reviewer XreH: "specialization of experts is one of the most important goals in MoE".
• Reviewer Su68: "µMoE models demonstrate a more parameter-efficient alternative to traditional MoEs, an important consideration for practical applications."

Concretely, as detailed from [L344], our experiments comparing model accuracy in Section 4.3 serve only to verify that muMoE layers do not trade-off accuracy relative to existing dense layers of both equivalent parameters and FLOP counts. Sparse MoEs introduce significantly more parameters; thus, there is no possible fair comparison in this setting.

In the paper, we make no claims about competing with Sparse MoEs’ performance—only that our model form avoids Sparse MoEs’ non-differentiability by design. Rather than exploring conditional computation for increasing performance and network capacity, the paper instead has the orthogonal goal of facilitating more interpretable subcomputations (e.g. allowing downstream tasks/properties of manual bias correction, interpretability, or network editing/steering as shown in the rebuttal PDF), without introducing any additional parameters over MLPs.

We would be grateful if the reviewer could re-assess the work with the paper’s stated goals and claims in mind, and to judge the method through the novelty it brings to scalable expert specialization at no additional computation cost--along with its downstream benefits shown in the paper.

Sparse/Soft MoE comparison results

We emphasize the significance of muMoE’s parameter-efficient expert specialization: even at the moderately large expert counts we explore in this paper, Sparse MoEs quickly approach a prohibitively large number of parameters. Even the GPT2-small model with a 2048 expert Sparse MoE has on the order of 100 billion MLP parameters--compared to just ~100 million parameters with CPmuMoEs. Meanwhile, the GPT2-XL model with 2048 expert Sparse MoE layers has on the order of a trillion MLP parameters alone. For even larger dense models whose base parameter counts are already in the billions, the large expert counts necessary for fine-grained specialization are infeasible with Sparse MoEs.

We accordingly focus on an experimental setup in which it is possible to use large numbers of experts in both Sparse and Soft MoEs--without an infeasible number of parameters or FLOPs. In this setting, the computational requirements of all MoEs can be matched fairly. Please find in the global rebuttal PDF (at the top of openreview) a comparison of muMoEs to both Soft and Sparse MoEs for fine-tuning CLIP on ImageNET1k: whilst it is not the goal of the layer, both variants of muMoEs provide Pareto-optimal improvements to the parameter count vs accuracy frontier.

Furthermore, as outlined above, the paper’s aim is not to increase accuracy or replace existing MoEs — but rather to replace parameter- and FLOPs-matched dense layers with interpretable subcomputations. This is why we compare dense layers as a baseline and not transformer-specific MoE layers such as the suggested [1].

What’s more, we highlight that the MoE and routing strategy in “From Sparse to Soft Mixtures of Experts” [1] is tailored to transformer/token-based architectures, and the method and paper aim at improving transformer model performance. In contrast, the muMoE focuses instead on scalable expert specialization, and in dense layers more generally. As demonstrated in the paper, our focus is on facilitating applications such as last-layer fine-tuning models with MoEs (and subsequently removing model bias in Table 2). These tasks addressed in the paper are not possible with the token-specific methodology of [1].

---

• [1]: Puigcerver, Joan et al. “From Sparse to Soft Mixtures of Experts.” ICLR 2024.

Low rankness at train-time

The reviewer is correct that Appendix C.2 explores the minimal impact of applying truncated SVD on pre-trained weights. However, all the results in Section 4.3 of the main paper do validate the ability of the muMoE layer to retain performance with rank constraints at train time. As stated on [L769], this is our main evidence to support the claims of the paper.

Benchmarking of speed & memory

Please find (in the rebuttal PDF) a summary of latency and peak memory usage for various layers. We use the same number of experts for all MoE variants and set the rank of muMoEs to match the parameter count of a linear layer mapping from 768 to 1000 dimensions. Whilst muMoEs do not offer significant speed gains over alternative MoEs, they have an order of magnitude less peak memory usage. We will include a discussion of these results in the revised main paper with the additional page count.

“Dense MoEs” suggestion

We thank the reviewer for the helpful suggestion to use “Dense MoE” to refer to the baseline methodology. Not only does such a name accurately describe the properties of this model's forward pass, but it indeed helps distinguish it from the recent ICLR paper. We will gladly adopt this name in the paper.

Dear Reviewer-YTQ3,

We are grateful for the thorough and thoughtful review of our submission.

The reviewer's primary concern is missing comparisons to Sparse MoE's capabilities. We kindly draw our attention to rebuttal, where we provide new requested experiments showing that the proposed muMoE layer offers Pareto-optimal improvements to accuracy over Sparse and Soft MoE. We provide further requested experiments to benchmark speed and peak memory usage and discuss all remaining concerns raised.

Additionally, we provide clarifications about the goal and claims of the paper as providing scalable expert specialization to dense layers without increasing either the FLOPs or parameter counts (orthogonal to the aim of pushing capabilities through more parameters/FLOPs).

If there are any other questions or concerns that we can address to enhance your support for our paper, we would greatly appreciate the opportunity to respond.

Sincerely, Submission #3407 Authors.

Despite demonstrating the benefits of scalable expert specialization with muMoE layers for various architectures with different modalities, we have not conducted experiments on the largest scales of commercial LLMs. Furthermore, our experimental results demonstrated comparable accuracies of muMoE networks for models only up to parameter counts on the order of 100 million. Where resources permit, future work should explore the scalability of expert specialization and performance of muMoEs in even larger-scale LLMs.

We thank Reviewer-XreH for praising the well-motivated problem studied in the paper and the clear presentation of the factorization methodology.

More results on language models

Whilst we do not have the computational resources to pre-train new additional large language models during the rebuttal, we do provide more LLM results on the GPT-2 language models already trained for steering the LLM outputs using the experts, inspired by the reviewer’s comments.

In particular, we use a GPT-2 model trained with muMoE layers at each MLP layer, using 2048 experts (at every layer). By modifying the forward pass of the trained model—specifically, adding selected expert cluster center vectors to each token's input latent activation vector before applying the muMoE layer—we can consistently control the model to generate outputs aligned with specific themes. Illustrations of this approach, using 3 different experts, are included in the global rebuttal PDF.

The selected experts guide the language model's outputs toward discussing topics such as climate change, police brutality, or foreign politics. We suggest that these findings further demonstrate the effectiveness of the muMoE layer in facilitating controllable generation of language model outputs—all whilst requiring no additional training.

[Q1] muMoE vs Soft MoE

Soft MoEs perform a dense forward pass taking a linear combination of all $N$ experts’ outputs simultaneously. With large expert counts, this is problematic due to both high parameter counts and restrictively high FLOP costs.

As detailed in Section 3.1.2, the existing soft MoE is a special case of muMoE when (a) the weight tensors are explicitly stored without decomposition and (b) no implicit fast factorized computation is performed. The key technical contribution of the muMoE layer is to make a similar full dense forward pass (using all $N$ experts simultaneously) feasible at high expert counts. In contrast, muMoEs elegantly skirt both issues with Soft MoE: the former (a) through factorization, and the latter issue (b) through the derivations of the fast implicit forward passes.

[Q2] How to find top-activating experts

The reviewer is correct that we apply no top-k routing at training time, nor at test time. After the model is trained however, we can still rank inputs by their expert coefficients for the qualitative experiments. Concretely, we pass a dataset of inputs through the network and rank images/text-token inputs based on the magnitude of the corresponding expert coefficient (the strength with which each expert is activate). It is through this standard technique that the top-activating images are ranked and displayed in Figures 4 and 5.

As shown in the “Expert coefficients color map” in the top-right of Figure 5, the color indicates the soft expert coefficient weight (on the interval [0,1]) with which this particular expert processes each text token.

We're pleased to hear that the reviewer has raised their score following the additional experiments and rebuttal. We wish to thank Reviewer XreH for their time and efforts in engaging with our response.

Are there any additional questions about our work we can answer that would further increase the reviewer's rating of the paper? Please do not hesitate to ask us any further questions you may have.

Best,

Submission #3407 Authors.

We thank Reviewer-qC51 for their detailed review; for praising the ease with which the methodology can be implemented and the benefits of downstream applications in bias mitigation. We address the stated weaknesses below and draw attention to the existing ablation studies requested. Additional experiments requested to further show the benefits of the interpretability gained are also provided through steering LLMs’ outputs using the experts.

Comparisons to Sparse/Soft MoE

As the reviewer correctly states, our sole interest is the independent benefits of scalable expert specialization (these are the claims and paper's goals as stated in [L17, L81, L353]). muMoEs are parameter-efficient layers for providing interpretability and controllability benefits without increasing the parameter or FLOPs count of the original networks. We highlight that two reviewers praise this explicitly:

• Reviewer XreH: "specialization of experts is one of the most important goals in MoE".
• Reviewer Su68: "µMoE models demonstrate a more parameter-efficient alternative to traditional MoEs, an important consideration for practical applications."

Consequently, we argue that the fair and relevant comparison to validate the paper’s claims is with a fixed computational budget: between the original networks with MLP layers and parameter- and FLOPs-matched muMoE layers. This serves to verify that we do not trade-off accuracy when introducing the many benefits of muMoE layers. We do not claim (nor desire) to increase parameter counts / capacity / performance. Sparse and Soft MoEs significantly increase the number of parameters and the number of FLOPs, respectively, making a fair comparison impossible.

In contrast, Soft and Sparse MoEs mostly aim to increase performance through increased parameter counts and/or FLOPs. Crucially, fine-grained expert specialization in large models is impractical with both Soft MoEs and Sparse MoEs. For Soft MoEs, a restrictively high FLOP count is required. For Sparse MoEs, the parameter count alone is problematic: even the smallest GPT2 variant with 2048 expert Sparse MoE layers has on the order of 100 billion MLP parameters--compared to just ~100 million parameters with muMoEs. For even larger dense models whose base parameter counts are already in the billions, large expert counts with Sparse MoE are infeasible.

Requested comparisons

Inspired by the comments of the reviewer, we've included a new detailed comparison of muMoEs with both soft and sparse MoEs for CLIP fine-tuning (in the rebuttal PDF). This single-layer experimental setup makes it possible to use an otherwise prohibitively large number of experts in Sparse/Soft MoEs, all while maintaining a feasible number of parameters and FLOPs. Although optimizing the parameter count / FLOPs VS accuracy trade-off is not the aim of the paper, both muMoE variants achieve Pareto-optimal improvements in this aspect for last-layer fine-tuning of CLIP on ImageNET1k.

Ablations

We kindly draw the reviewer’s attention to the following figures where the requested ablations can already be found in the submission:

• Performance increase as a function of expert count: please see Figure 21 of the appendix ([L933]) where performance increases as expert count is increased.
• Placement of muMoE layers: please see Figures 3, 16, 17 for final and penultimate layer muMoEs configurations. We note that Section 4.3’s results are with networks with muMoEs at every layer, and results from all layers are shown in Figure 11 of the appendix.
• Different model sizes: experiments are already performed throughout on 4 network architectures of various sizes (CLIP, DINO, GPT2, and MLP Mixers).
• Gating function: An ablation study is already presented on [L888] between activation functions.

Additional use cases of experts

As requested, to further demonstrate the usefulness of the interpretability and experts, we show how in GPT2 models, the experts can be used to steer the model’s output to particular themes of interest.

Concretely, we take a GPT-2 model trained with CPmuMoE layers at every MLP layer, each with 2048 experts. We find that intervening in the trained model’s forward pass (concretely: adding particular expert cluster center vectors to every token’s input latent activation vector before applying the muMoE layer), reliably steers the model to produce outputs related to specific themes. Examples are shown for 3 experts in the rebuttal PDF at the top of the page.

Chosen experts steer the LLM outputs toward talking about themes such as climate change, police brutality, or programming. We suggest these results constitute further evidence of the usefulness of the muMoE layer in providing controllable generation of LLM outputs (without any additional training needed).

Figure 2

Our experiments are separated into qualitative studies (Section 4.1.1.) and quantitative studies (Sections 4.1.2 and 4.2). The sole purpose of Figure 2 is a qualitative study to visually motivate the benefits of increasing the expert count. Separate quantitative results are found throughout the latter two subsections. Furthermore, the distribution of expert assignments is already visualized in Figure 20 of the appendix.

Dear Reviewer-qC51,

Thank you for your thorough review of our submission.

The reviewer's key issues were about comparisons to Soft/Sparse MoE and missing ablation studies. In our rebuttal, we highlight where each of the requested ablations can be found throughout the main paper and the appendix. Furthermore, we address the reviewer's primary issue by including new comparisons to both Soft and Sparse MoEs, where the muMoE layer strictly outperforms the two existing alternatives.

If there are any other questions or concerns that we can address to further strengthen your support for our paper, we would be eager to respond.

Sincerely, Submission #3407 Authors.

Thanks to Reviewer-qC51 for their response.

We are glad to hear that our rebuttal has answered most of your concerns. We thank the reviewer for their further insights and chance to answer the remaining questions below:

1. Varying parameter counts of baselines

As described in the Caption of Figure 1, we increase the Soft/Sparse MoEs' parameter counts by increasing the expert counts.

We also use a linear transformation (to match the functional form of the baseline layers), but instead compose the weight matrix as a product of two tall matrices $\mathbf{W}_1$ and $\mathbf{W}_2$, whose number of rows are varied to increase the linear layer's parameter count. Concretely we compute: $(\mathbf{W}_1^\top\mathbf{W}_2)\mathbf{x}+\mathbf{b}$.

We are thankful to the reviewer for the attentive study of our responses. We will include these details in the revised version.

2. CP vs TR

Figure 21 of the appendix demonstrates that parameter-matched muMoEs strictly outperform linear layers in this experiment, for all possible choice of expert counts. Even with equal parameter counts, we see performance gains of up to 0.7% above the linear layer (which is non-trivial on ImageNET).

However, to parameter-match the single linear layers, we must decrease the muMoE layer rank upon increasing the expert count. Crucially, as we see from Figure 7 of the appendix, TRMuMoEs are often much more parameter-efficient at high expert counts (a result we discuss in depth in Section C.1).

As a consequence of this, TRMuMoEs' resulting expert matrix ranks are increasingly larger than that of CPMuMoEs, even when the layers are parameter-matched. For example, the parameter-matched layers with 512 experts in Figure 21 have an expert matrix rank of 165 for the CPMuMoE compared to a much larger 208 for the TRMuMoE.

We attribute TRMuMoE's even greater performance gains over CPMuMoEs to the more favorable relationship between tensor rank and expert matrix rank (a larger weight matrix rank meaning the resulting layers' activations live in a larger dimensional subspace).

We will include a table in the revised version comparing the matrix ranks to the expert counts in this experiment, and include this discussion in the revised supplementary material.

3. Implications of layer placement

Our most noteworthy observation about the implications of layer placement is in vision models. As shown in Figure 4 and discussed in [L326], we find that MuMoEs at early layers exhibit specialism to more coarse-grained patterns (e.g. texture and colors), whilst later layers exhibit more fine-grained expert specialisms (to classes and objects). Please also observe this in Figure 11 of the appendix for both MuMoE variants. To our knowledge, there is no significant trade-off when varying placements of the layer in terms of training resources or final performance.

---

We thank the reviewer again for the insightful discussion. Might they have any further questions we could answer to increase their score?

Sincerely,

Submission #3407 Authors.

We thank Reviewer-Su68 for their detailed assessment of the paper: for praising the “novel[ty]” of the layer and the “technically sound” methodology. Furthermore, we are glad the reviewer appreciates the “extensive qualitative and quantitative evidence to support the claims” of the paper.

All stated weaknesses are addressed below, pointing out where the requested details can be found:

[W1]: Factorization challenges

Low-rank factorization is indeed a well-known technique in deep neural networks, frequently employed to reduce the number of parameters. However, as far as we are aware, our work is the first to introduce a novel technical link between hierarchical MoEs and multilinear models. Instead of utilizing factorization to simply reduce parameter counts, our work provides a way of designing trainable layers with prohibitively large numbers of experts in the first place. This is only possible through the paper’s extra derivations and formulations of MuMoE layers’ fast forward passes (beyond simple factorization).

We have an in-depth discussion in Appendix C [L722] of how the factorization choice impacts MoEs. For example, we derive the theoretical relationship between tensor and subsequent expert matrix rank, and the relative strengths of different factorizations in this applied context. Furthermore, in Appendix C.2 we also discuss the impact of low-rankness in the specific context of the models studied in this paper. We believe this further detailed explanation and exploration of the methodology also addresses the reviewers’ second comment about further enhancing comprehension.

[W2]: Marginal Improvements and lack of multiple runs

As stated on [L293], the results in Table 2 are already the mean over 10 runs. We have added this to the caption to make it more clear. Furthermore, an improvement of (for example) 5% to the min-max fairness metric (for “age”) and half the equality of opportunity (for “blond hair”) are significant improvements over existing techniques.

[W3]: Metrics not clearly explained

The metrics used in the paper are described in detail in Appendix I ([L958]), as written on [L307] of the main paper. Furthermore, we describe in depth their interpretations in the context of the celebA bias problem studied, and the implementation details of the baselines. We will also explicitly reference this appendix discussing the metrics in Table 2’s caption.

[W4]: Evaluation on OOD data for vision models

Indeed, our quantitative evaluation of vision models is limited to the test set (and not OOD data), which we already noted as a limitation in the conclusion.

Despite this, our experiments on language models in Figures 5, 12, and 13 use generated LLM sequences as their input (rather than that of the training set). Furthermore, we also perform an additional experiment on steering LLMs for Reviewer-qC51, where the input prompt is free-form text. This provides initial evidence of the generalizability of the experts outside of the training/test data.

[W5]: Lacking quantitative measures of interpretability

We kindly draw attention to both Section 4.1.2 and Section 4.2. In the former, a new quantitative metric capturing expert specialization is introduced. In the latter, interpretability is further measured quantitatively through its downstream ability to edit models compared to existing methods. We highlight that Reviewer-1sqG praises the “innovative metric", stating that it "provide[s] a clear measure of the impact of increasing the number of experts". If the reviewer has any other valid metric in mind, we will be glad to consider it in our evaluation.

[W6]: Different model architectures and sizes

We kindly note that we already perform experiments across four different architectures: CLIP, DINO, MLP-Mixer, and GPT-2 (in Figure 3 and Table 4). These models are each of vastly different sizes and configurations (and operate on various modalities including language and vision), already demonstrating the wide applicability of the layer.

That said, we have access to a maximum of 4 A100 GPUs, limiting the size of the models we can pre-train to GPT2-sized models. We have added a limitation to the paper’s conclusion to note explicitly how we have not tested the layer at the largest scales of commercial LLMs:

Despite demonstrating the benefits of muMoE layers for various architectures with different modalities, we have not conducted experiments on the largest scales of commercial LLMs. Future work should explore this where resources permit.

For the camera ready, we will include additional experiments on a GPT-2 model trained from scratch with 2048 expert CPMuMoEs at each layer, to provide further evidence of the layer’s benefit at very large expert counts.

“Ease of implementation” and visual illustration

We would like to draw attention to Appendix B, which contains pseudocode implementations of all muMoE variants. Furthermore, we note that the MuMoE.py file in the anonymous code linked to in the abstract provides a full implementation of the layers in PyTorch.

To aid with visual illustration, we plan to include an annotated variant of Figure 1 in the appendix that shows a step-by-step explanation of each stage of the computation of the forward pass for clarity.

We thank the reviewer for their suggestion to improve the ease of implementation of the methods and to include additional visual aids to improve comprehension of the method.

Dear Reviewer-Su68,

Thank you for your time and effort in providing the review of our submission.

In our rebuttal, we clarify the reviewer's key concerns: that the paper's fairness results are indeed over repeated runs, and draw attention to where explanations of the metrics can be found. We also provide additional clarifications and responses to all further questions raised.

We would like to ask if the reviewer has any additional questions we are able to answer that would increase their support in the paper?

Sincerely, Submission #3407 Authors.