Multilinear Mixture of Experts: Scalable Expert Specialization through Factorization

Authors: james  oldfield
markos  georgopoulos
grigorios g chrysos
christos  tzelepis
yannis  panagakis
mihalis a nicolaou
jiankang  deng
ioannis  patras
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-09 18:00:47

Abstract: The Mixture of Experts (MoE) paradigm provides a powerful way to decompose dense layers into smaller, modular computations often more amenable to human interpretation, debugging, and editability. However, a major challenge lies in the computational cost of scaling the number of experts high enough to achieve fine-grained specialization. In this paper, we propose the Multilinear Mixture of Experts (μMoE) layer to address this, focusing on vision models. μMoE layers enable scalable expert specialization by performing an implicit computation on prohibitively large weight tensors entirely in factorized form. Consequently, μMoEs (1) avoid the restrictively high inference-time costs of dense MoEs, yet (2) do not inherit the training issues of the popular sparse MoEs' discrete (non-differentiable) expert routing. We present both qualitative and quantitative evidence that scaling μMoE layers when fine-tuning foundation models for vision tasks leads to more specialized experts at the class-level, further enabling manual bias correction in CelebA attribute classification. Finally, we show qualitative results demonstrating the expert specialism achieved when pre-training large GPT2 and MLP-Mixer models with parameter-matched μMoE blocks at every layer, maintaining comparable accuracy. Our code is available at: https://github.com/james-oldfield/muMoE.

Main paper: 
 introduction: 
 The Mixture of Experts (MoE) architecture [1] has reemerged as a powerful class of conditional computation, playing the pivotal role in scaling up recent large language [2, 3, 4, 5] , vision [6] , and multi-modal models [7] . MoEs apply different subsets of layers (referred to as 'experts') for each input, in contrast to the traditional approach of using the same single layer for all inputs. This provides a form of input-conditional computation [8, 9, 10, 11] that is expressive yet efficient. However, through their substantial performance gains, an important emergent property of MoEs is frequently underutilized: the innate tendency of experts to specialize in distinct subtasks. Indeed, the foundational work of Jacobs et al. [12] on MoEs describes this property, highlighting how implementing a particular function with modular building blocks (experts) often leads to subcomputations that are easier to understand individually than their dense layer counterparts-with larger expert counts allowing for more fine-grained specialization. Independent of model performance, a successful decomposition of the layer's functionality into human-comprehensible subtasks offers many significant benefits. Firstly, the mechanisms through which a network produces an output are more interpretable: the output is a sum of modular components, each contributing individual functionality. Yet, the value of interpretable computation extends beyond just transparency [13] and explainability [14] . An important corollary of successful task decomposition amongst experts is that layers are easier to debug and edit. Biased or unsafe behaviors can be better localized to specific experts' subcomputation, facilitating manual correction or surgery in a way that minimally affects the other functionality of the network. Addressing such behaviors is particularly crucial in the context of foundation models; being often fine-tuned as black boxes pre-trained on unknown, potentially imbalanced data distributions. Furthermore, there is evidence that traditional fairness techniques are less effective in large-scale models [15, 16] . However, to achieve fine-grained expert specialism at the class level (or more granular still), one needs the ability to significantly scale up the number of experts. When using only a small expert count, each expert is forced to process and generalize across multiple distinct semantic concepts, hindering specialization. Conversely, a large expert count means each can specialize to a more specific set of semantically similar inputs. Alas, the dominating 'sparse' MoE paradigm of selecting only the top-K experts [17] is not only parameter-inefficient for large expert counts, but also has several well-known issues due to its discrete expert routing-often leading to training instability and difficulties in scaling the total expert count, amongst other challenges [18, 19] .Table 1 : Benefits of the proposed µMoEs' model form over existing MoEs.Parameter-FLOPs-Differentiable efficient efficientDense MoE [1] Sparse MoE [17] µMoE (ours)In this paper, we propose the Multilinear Mixture of Experts (µMoE) layer to address these issues. µMoEs are designed to scale gracefully to dense operations involving tens of thousands of experts at once through implicit computations on a factorized form of the experts' weights. Furthermore, in contrast to the dominant sparse MoEs' [17] non-differentiable nature, µMoEs are differentiable by design, and thus do not inherit the associated training issues. We summarize the benefits of µMoEs' model form over existing MoEs in Table 1 . Crucially, we show evidence that scaling up the number of µMoE experts leads to increased expert specialism when fine-tuning foundation models for vision tasks. Our evidence is provided in three forms: (1) firstly, through the usual qualitative evaluation of inspecting inputs by their expert coefficients. Secondly (2), we further explore the causal role of each expert through counterfactual interventions [20] . Lastly, (3) we show how final-layer µMoE expert specialism facilitates the practical task of model editing-how subcomputation in specific combinations of experts biased towards demographic subpopulations can be manually corrected through straightforward guided edits.Building on these findings, we demonstrate that µMoEs offer a compelling alternative to MLPs for pre-training both vision and language models with up to 100M parameters-enabling large numbers of specialized experts while maintaining comparable performance and parameter counts to the original networks' single dense MLPs.Our contributions and core claims can be summarized as follows:• We introduce µMoE layers-a mechanism for computing vast numbers of subcomputations and efficiently fusing them conditionally on the input. • We show both qualitatively (through visualization) and quantitatively (through counterfactual intervention) that increasing the number of µMoE experts increases task modularity-learning to specialize in processing just specific input classes when fine-tuning large foundation models for vision tasks. Further, we show manual editing of µMoE expert combinations can straightforwardly mitigate demographic bias in CelebA attribute classification. • We pre-train both language (GPT2) and vision (MLP-mixer) µMoE networks, establishing experimentally that models with parameter-matched µMoE blocks are competitive with existing MLP blocks whilst facilitating expert specialism (qualitatively) throughout. 
related work: 
 Mixture of Experts Recent years have seen a resurgence of interest in the Mixture of Experts (MoE) architecture for input-conditional computation [17, 12, 21, 2] . One primary motivation for MoEs is their increased model capacity through large parameter count [17, 4, 2] . In contrast to a single dense layer, the outputs of multiple experts performing separate computations are combined (sometimes with multiple levels of hierarchy [22, 23] ). A simple approach to fusing the outputs is by taking either a convex [23] or linear [24] combination of the output of each expert. The seminal work of Shazeer et al. [17] however proposes to take a sparse combination of only the top-K most relevant experts, greatly reducing the computational costs of evaluating them all. More recent works employ a similar sparse gating function to apply just a subset of experts [2, 25] , scaling to billions [3] and trillions of parameters [4] . The discrete expert selection choice of sparse MoEs is not without its problems, however-often leading to several issues including training stability and expert under-utilization [18, 19] .Particularly relevant to this paper are works focusing on designing MoE models to give rise to more interpretable subcomputation [26, 27, 28] -hearkening back to one of the original works of Jacobs et al. [12] , where experts learned subtasks of discriminating between different lower/uppercase vowels. Indeed a common observation is that MoE experts appear to specialize in processing inputs with similar high-level features. Researchers have observed MoE experts specializing in processing specific syntax [17] and parts-of-speech [29] for language models, and foreground/background [30] and image categories (e.g. 'wheeled vehicles') [24] in vision. Evidence of shared vision-language specialism is even found in the multi-modal MoEs of Mustafa et al. [7] .Several works instead target how to make conditional computation more efficient: by sharing expert parameters across layers [31] , factorizing gating network parameters [32] , or dynamic convolution operations [33] . Relatedly, Gao et al. [34] jointly parameterize the experts' weight matrices with a Tensor-Train decomposition [35] . However, such approach still suffers from the Sparse MoE's instability and expert under-utilization issues, and stochastic masking of gradients must be performed to lead to balanced experts. Furthermore, whilst Gao et al. [34] share parameters across expert matrices, efficient implicit computation of thousands of experts simultaneously is not facilitated, in contrast to the µMoE layer.Factorized layers in the context of deep neural networks provide several important benefits. Replacing traditional operations with low-rank counterparts allows efficient fine-tuning [36] / training [37, 38] , and modeling of higher-order interactions [39, 40, 41, 42, 43] , and convolutions [44] . In addition to reducing computational costs, tensor factorization has also proven beneficial in the context of multi-task/domain learning [45, 46] through the sharing of parameters/low-rank factors across tasks. Furthermore, parameter efficiency through weight factorization often facilitates the design and efficient implementation of novel architectures such as polynomial networks [47, 48, 49] or tensor contraction layers [50] . The recent DFC layer in Babiloni et al. [51] also performs dynamic computation using the CP decomposition [52] like µMoEs. Nevertheless, the two works have very different goals and model properties due to how the weight matrices are generated. µMoEs take a sparse, convex combination of N explicit experts' latent factors. This consequently leads to specialized subcomputations in a way that facilitates the interpretability and editability presented in this paper. DFCs can be seen to apply an MLP to input vectors at this step in analogy, which does not provide the necessary model properties of interest here. 
methodology: 
 We first formulate the proposed µMoE layer in Section 3.1, introducing 2 unique resource-efficient models and forward passes in Section 3.1.1. Finally, we show in Section 3.1.2 how µMoEs recover linear MoEs as a special case.Notation We denote scalars x ∈ R with lower-case letters, and vectors x ∈ R I1 and matrices X ∈ R I1×I2 in lower-and upper-case boldface latin letters respectively. Tensors X ∈ R I1×I2×...×I d of order d are denoted with calligraphic letters. We refer to the (i 1 , i 2 , . . . , i d )-th element of this tensor with both X (i 1 , i 2 , . . . , i d ) ∈ R and x i1i2...i d ∈ R. Finally, we use a colon to index into all elements along a particular mode: given X ∈ R I1×I2×I3 for example, X ::i3 ∈ R I1×I2 or equivalently X (:, :, i 3 ) ∈ R I1×I2 is the matrix at index i 3 of the final mode of the tensor. We use X × n u to denote the mode-n (vector) product [53] of a tensor X ∈ R I1×I2×...×I N and vector u ∈ R In whose resulting elements are given by (X× n u) i1...in-1in+1...i N = In in=1 x i1i2...i N u in .3.1 The µMoE layer µMoEs provide a scalable way to execute and fuse large numbers of operations on an input vector by formalizing conditional computation through resource-efficient multilinear oper-ations. A µMoE layer comprised of N many experts (and a single level of expert hierarchy) is parameterized by weight tensor W ∈ R N ×I×O and expert gating parameter G ∈ R I×N . Given an input vector z ∈ R I (denoting the hidden representation of an individual token, for example), its forward pass can be expressed through the series of tensor contractions:+ = + Input vectorExpert coefficientsFigure 1 : The forward pass of an (unfactorized) µMoE layer as a series of tensor contractions: the experts' weight matrices (yellow 2D slices) are matrix-multiplied with the input vector and summed (weighted by the red expert coefficients).EQUATION)where a is the vector of expert coefficients and φ is the entmax activation [54, 55] . The µMoE layer can be understood as taking a sparse, convex combination of N many affine transformations2 of input vector z, weighted by the coefficients in a. The first tensor contraction in the forward pass ( i W :i: z i ∈ R N ×O ) matrixmultiplies the input vector with every expert's weight matrix. The following tensor contraction with expert coefficients a takes a linear combination of the results, yielding the output vector.The forward pass can be visualized intuitively as multiplying and summing over the modes in a 3D tensor, which we illustrate in Figure 1 . Furthermore, µMoEs readily generalize to hierarchical conditional computations by introducing additional modes to the weight tensor and corresponding vectors of expert coefficients (see Appendix E).Our key insight is that the dense µMoE forward pass over all N experts simultaneously can be computed entirely in factorized form, needing never materialize prohibitively large weight tensors. This allows µMoEs' computations to scale gracefully to many thousands of experts simultaneously, without the problematic top-K gating [17] . To achieve this, we (1) first parameterize the experts' weights W ∈ R N ×I×O with a tensor factorization and (2) re-derive fast forward passes of Equation (1) to operate solely in factorized form.In the context of a µMoE layer, the various choices of tensor factorizations make different trade-offs regarding parameter/FLOP counts and rank constraints. We derive two unique resource-efficient µMoE variants to suit different computational budgets and choices of expert counts. We now present the derivations of the forward passes of the factorized µMoE models (with einsum pseudocode implementations in Appendix B):CPµMoE Imposing CP structure [52, 56] of rank R on the weight tensor, we can writeEQUATION∈ R R×O . This reduces the parameter count from N IO (such as with sparse/dense MoEs and regular µMoEs) to just R(N + I + O). Crucially, we can further rewrite the CPµMoE layer's forward pass entirely in factorized form without ever materializing the full tensor (plugging the CP-composed tensor into Equation ( 1)) as:EQUATIONwith Equation (2) being analogous to the fast computation in Babiloni et al. [51] , only here the operations of combining the weights and producing the outputs can be expressed in a single step. Whilst the original naive CPµMoE forward pass has a FLOP count 3 of N IO, the fast computation above has just R(N + I + O) (the same number of factorized layer parameters). With moderate values of both R and N , the layer becomes significantly more resource-efficient than vanilla µMoEs.We propose a second µMoE variant based on the Tensor Ring [58] (TR) factorization that can offer even better efficiency for large values of N . In TR format, W ∈ R N ×I×O has three factor tensors:U (1) ∈ R R1×N ×R2 , U (2) ∈ R R2×I×R3 , U (3) ∈ R R3×O×R1, where R i are the manually chosen ranks 4 . The weight tensor's elements in TR format are given by: w nio = tr U:n: U:i: U:o: [58] . TRµMoE's forward passes can be computed efficiently by contracting the first two factor tensors with the input/expert coefficients vectors and then combining the results:EQUATIONyielding a modified FLOP count of (R1 N R 2 + R 2 IR 3 + R 1 R 2 R 3 + R 1 OR 3 ) with just (R 1 N R 2 + R 2 IR 3 + R 3 OR 1 )parameters. With large N contributing to the computational cost only through R 1 N R 2 , the TRµMoE can prove even more resource-efficient than CPµMoEs by choosing small values of R 1 , R 2 . We refer readers to Appendix D for a further discussion of decomposition choice, derivations of how tensor rank translates to expert matrix rank, and FLOPs comparisons.Finally, we note how unfactorized µMoE layers with a single level of expert hierarchy recover dense MoE layers [17, 11] as a special case. When computing Equation (1) over the full materialized weight tensor, one can alternatively write the output element-wise as y o = a W ::o z. This highlights an interesting technical connection between neural network layers: dense MoE layers in this tensor formulation can be seen to share a similar functional form to bilinear layers, which have also found applications in interpretability [59, 60] . 
experiments: 
 We start in Section 4.1 by presenting both qualitative and quantitative experiments validating that the experts learn to specialize in processing different semantic clusters of the input data. In Section 4.2 we demonstrate one practical benefit of the learned specialism-showing how expert-conditional re-writing can correct for specific demographic bias in CelebA attribute classification. Finally, in Section 4.3 we train both large language and large vision models with µMoE layers throughout-providing qualitative evidence of expert specialism and model performance competitive with networks using MLP blocks. Please see Appendix H for detailed ablation studies, and Appendix I for experiments with hierarchical µMoEs.Implementation details Before applying the activation function to the expert coefficients we apply batch-and layer-normalization to µMoE layers in vision and language models respectively (see Appendix H.3 for an ablation). Interestingly, we do not find the need for any load-balancing losses.We fix the TRµMoE ranks to be R 1 = R 2 = 4 throughout (see Appendix D.1.2).Our first objective is to show that scaling µMoE's expert count leads to more specialized experts.We provide evidence of this effect both qualitatively (through visualization) and quantitatively (through intervention).To isolate the impact of µMoE layers and varying expert counts, we first explore the controlled setting of fine-tuning large foundation models CLIP [61] ViT-B-32 and DINO [62] on ImageNET1k (following the fine-tuning protocol in Ilharco et al. [63, 64] ). Whilst fine-tuning large foundation models is an important application of µMoE layers in its own right (e.g. as explored later in Section 4.2 for fairer models), the ability to cheaply train many models with different µMoE layer configurations forms an ideal setting in which to study their properties. Each row displays randomly selected images processed (with coefficient ≥ 0.5) by the first few experts for the two models. The more we scale the expert count, the greater the apparent expert specialism (to single visual themes or image categories).We first show random examples in Figure 2 of images processed (with expert coefficient ≥ 0.5) by the experts by each CPµMoE layer (the class labels and expert coefficients are overlaid in white and green text respectively). Using only a modest number of experts (e.g. 32) appears to lead to some 'polysemanticity' [65] in experts-with some processing unrelated classes of images (e.g. 'gators', 'limos', and a 'quilt' for Expert 1 on the right). On the other hand, using a much larger number of total experts appears to yield more specialization, with many experts contributing their computation to only images of the same single class label or broader semantic category. Please see Figure 16 in the Appendix for many more random images for the first 10 experts per model to observe this same trend more generally, and Figure 17 for even finer-grained specialism with 2048-expert µMoE layers.The qualitative evidence above hints at the potential of a prominent benefit to scaling up the number of experts with µMoEs. Such subjective interpretations alone about expect specialism are hypotheses, rather than conclusions however [66] . Similarities in images processed by the same expert give us an intuitive explanation of its function but do not show the expert's computation contributes causally [20, 67, 68] to the subtask of processing specific human-understandable patterns of input features [69, 70] . However, the absence of ground-truth labels for interpretable features of the input one may be interested in (e.g. specific types of textures in images, or words related to 'Harry Potter') makes this difficult to quantify in any objective or systematic manner.Despite the absence of fine-grained labels, we can quantify and compare the class-level specialism a µMoE expert exhibits on the ImageNET1k dataset as an (imperfect) proxy [71] . Following the causal intervention protocol of Elazar et al. [20] , we ask the specific counterfactual question about solely each expert n in a µMoE layer in turn: "had expert n's weight matrix W n not contributed its computation, would the network's test-set accuracy for class c have dropped?" Practically speaking, given a network fine-tuned with an µMoE layer, we achieve this by intervening in the forward pass by zeroing the n th expert's weight matrix W n := 0, leaving every other aspect of the forward pass completely untouched. Let the elements of y, ŷ(n) ∈ R C denote the test set accuracy for the C = 1000 ImageNET1k classes, pre-and post-intervention of expert n respectively. We collect the normalized difference to per-class accuracy in the vector d (n) , whose elements are given by d(n) c = (y c - ŷ(n) c )/y c .At the two extremes, when the full network's accuracy for class c drops completely from y c to 0 upon manually excluding expert n's computation we get d(n) c = 1, whilst d (n) c= 0 means the absence of the subcomputation did not change class c's test set accuracy at all. We thus estimate the 'class-level polysemanticity' of expert n as the distance between the difference vector and the one-hot vector:EQUATIONwhere index argmax c (d(n) c ) of 1 (n)has a value of 1 (and values of 0 everywhere else). This encodes the signature of a perfectly class-level monosemantic expert, for which all accuracy for a single class alone is lost in the counterfactual scenario in which the expert n did not contribute. We plot in Figure 3 the average expert polysemanticity p (n) for all experts with non-zero difference vectors 5 , observing a steady drop in its value as N increases from 32 to 1024 total experts. In other words, increasing N leads to individual experts increasingly responsible for a single subtask: classifying all inputs of just one class. As shown in Figure 3 we observe this trend both when µMoEs are used as final classification layers and as penultimate layers (followed by a ReLU activation and linear classification layer), and for multiple pre-trained foundation models. We further refer readers to the bar plots of the values of d (n) (the per-class accuracy changes) in Figures 18 and 19 , where this trend is observable through mass concentrated on increasingly fewer class labels as the number of experts increases.We further validate the modular expert hypothesis of µMoEs and simultaneously provide a concrete example of its usefulness by correcting demographic bias in attribute classification. Classifiers trained to minimize the standard binary cross-entropy loss often exhibit poor performance for demographic subpopulations with low support [72, 73] . By identifying which combination of experts is responsible for processing target subpopulations, we show how one can straightforwardly manually correct mispredictions in a targeted way-without any re-training.We focus on mitigating bias towards two low-support subpopulations in models with µMoE final layers fine-tuned on CelebA [74] : (a) bias towards images labeled as 'old females' for age prediction [75] , and (b) bias towards images labeled as 'blond males' for blond hair prediction [15] . Concretely, we train N = 128 multi-label µMoE final layer models for the 40 binary attributes in CelebA, jointly optimizing a pre-trained CLIP ViT-B-32 model [61] backbone, again following the fine-tuning setup in Ilharco et al. [63, 64] . All results presented in this section are the average of 10 runs with different random seeds.Experimental setup Let C be a set collecting the expert coefficients a ∈ R N from forward passes of the training images belonging to the target subpopulation. We evaluate the subpopulation's mean expert coefficients ā = 1/|C| a∈C a ∈ R N , proposing to manually re-write the output of this expert combination. We modify the layer's forward pass for the o th output head for attribute of interest (e.g. 'blond hair') as:EQUATIONHere, the term λā ∈ R N specifies, for each expert, how much to increase/decrease the logits for attribute o, with λ being a scaling hyperparameter 6 . Taking the dot product with an input image's expert coefficients a applies the relevant experts' correction terms (in the same way it selects a subset of the most relevant experts' weight matrices). We report a range of standard fairness metrics for both the model rewriting and networks trained with existing techniques (that aim to mitigate demographic bias without requiring images' sensitive attribute value at test time). These are shown in Table 2 for the two different experiments on CelebA, where the proposed intervention outperforms baseline alternative methods in the majority of settings. Please see Appendix J for details about the baseline methods and fairness metrics used, and further discussion of results.Finally, we train from scratch 12 layer 124M-parameter GPT-2 [81] LLMs on OpenWebText [82] for the language domain and 8 layer S-16 variant7 MLP-Mixers [80] on ImageNET1k [83] for vision. We replace every MLP block's 2 linear layers with 2 µMoE layers. Each token t's input vector z t ∈ R I is therefore transformed with µMoE blocks of the form:y t = N n2=1 H h=1 w (2) n2h: GELU N n1=1 I i=1 w (1) n1i: z ti a tn1 h a tn2 , a t = φ(G z t ),where a t ∈ R N are the expert coefficients for each specific token and block, H is the dimension of the block's hidden layer, and tensors for each of the two layers. We manually set the µMoE ranks to parameter-match each original network and set the number of experts (per block) to N = 64 for vision models and N = 256 for LLMs. Consequently, with this configuration, each layer's µMoE block performs computations with N experts yet has the same parameter counts and FLOPs as a single, dense MLP block. µMoE-Mixer For vision, our key findings are that earlier µMoE channel-mixing blocks' experts appear (qualitatively) to exhibit specialisms to colors, shapes, and textures, whilst later layers exhibit more object-specific specialization. We plot the patches from the training set for which each expert most contributes its computation in Figure 4 for both a shallow and deep layer to illustrate this-earlier layers' experts contribute strongly to the processing of similar patches (top rows, e.g. specific edges) whilst later layers' experts process tokens based more on the similarity of their surrounding semantic context (bottom rows, e.g. images of animals). We further show in Figure 12 results for the first 2 experts across all 8 blocks where such scale-specific specialism is apparent across the entire network. µMoE-GPT2 For LLMs, we see promising qualitative evidence of experts specializing throughout a corpus of 1M generated 100-token sequences. At layer 5, for example, the generated tokens that use expert 8 with the highest coefficient are compound adjectives (Figure 5 ), whilst expert 37 most highly activates for equality and comparison operators in code and scientific text (please see examples of Table 3 : Comparison of µMoEs and dense MLPs across different models and tasks. We use N = 64 µMoE experts for the two vision tasks and N = 256 for GPT2. MLP mixers and GPT2s are pre-trained for 300 epochs and 100k iterations respectively, whilst CLIP is fine-tuned for 10 epochs. W (1) ∈ R N ×I×H , W (2) ∈ R N ×H×O are 
conclusion: 
 In this paper, we introduced the Multilinear Mixture of Experts layer (µMoE). We demonstrated that larger expert counts lead to increased specialization, and how µMoE layers make this computationally tractable through factorized forward passes. µMoEs scale to large expert counts much more gracefully than existing MoEs, yet avoid the issues from popular gating mechanisms. As a further practical example of µMoE's task decomposition, we illustrated how manual guided edits can be made to correct bias towards demographic subpopulations in fine-tuned foundation models. Having also shown matching performance in addition to expert specialism in both large vision and language models, we believe µMoE layers constitute an important step towards facilitating increasingly performant models that do not trade off fairness/interpretability for accuracy.Limitations Firstly, it is important to state again that our quantitative evaluation only captures expert behavior on the test set, not out-of-distribution data [70, 84] . Furthermore, expert specialism in large models is only demonstrated qualitatively (through the expert coefficients) due to the absence of fine-grained labels. Developing ways of quantifying fine-grained expert specialism is an important direction for future research. Finally, our experimental results demonstrated comparable accuracies of µMoE networks only for models with parameter counts on the order of 100 million. Where resources permit, future work should explore the scalability of expert specialization and performance of µMoEs in even larger-scale LLMs.  
 
 Appendix: This paper presents work whose goal is to advance the field of interpretable machine learning. Our goal is not to improve model capabilities but rather an orthogonal one of designing architectures more interpretable and controllable. As with many work with an interpretability focus, however, the µMoE layer could nonetheless facilitate the further development of SOTA models through its more expressive computation. We thus encourage the development of further guardrails against potentially harmful dual-uses of such technology. We release our code upon acceptance to facilitate further research along such lines.We here detail how to implement the fast forward passes of the µMoE models in a batch-wise manner, where each mini-batch element is a 2D matrix of shape Z ∈ R T ×C (with 'token' and 'channel' dimensions) with PyTorch and einops' [85] einsum:The CPµMoE forward pass can be implemented with:# CPmuMoE ( r = CP rank , b = batch_dim , t = tokens , # i = input_dim , o = output_dim , a [ e ]= expert_coefs , n *= expert_dims ) y = einsum ( G3 , a [ 0 ] @G1 .T , z@G2 .T , 'r o , b t r , b t r -> b t o ')And a two-level hierarchical CPµMoE with an additional factor matrix as:# CPmuMoE ( r = CP rank , b = batch_dim , t = tokens , # i = input_dim , o = output_dim , a [ e ]= expert_coefs , n *= expert_dims ) # ################ # A 2 -level hierarchical CPmuMoE , assuming Gi 's of appropriate shape y = einsum ( G4 , a [ 0 ] @G1 .T , a [ 1 ] @G2 .T , z@G3 .T , 'r o , b t r , b t r , b t r -> b t o ')TRµMoEs can be implemented with: For intuition, we provide a visualization in Figure 6 of the step-by-step series of tensor contractions W × 1 a × 2 z ∈ R O that the µMoE computes (in non-factorized form). In this section we present a further detailed discussion of decomposition choice, validating our choices and comparing alternative options. The computational costs of each fast µMoE forward pass and tensor-matrix rank relationships implications derived in this section are summarized in Table 5 . + I + O) R(N + I + O) min{I, O, R} TRµMoE R1N R2 + R2IR3 + R3OR1 R2IR3 + R1N R2 + R1R2R3 + R1OR3 min R3 • min{R1, R2}, I, OOne important consideration is how the chosen tensor ranks bound the resulting experts' matrix rank in µMoE layers. Here, we derive the matrix ranks as a function of tensor ranks for each model in turn.CPµMoEs are parameterized by factor matricesEQUATION∈ R R×O for chosen CP-rank R. Following Section 3 of Kolda and Bader [53] which provides the matricization/unfolding of CP tensors, we can write expert n's weight matrix asEQUATION:nEQUATIONwhere is the Khatri-Rao product [53] , and U:n ∈ R R×1 is the column of the factor matrix associated with expert n (including a singleton dimension for the Khatri-Rao product to be welldefined). Through the linear algebra rank inequality for matrix products, we have rank(W n ) = rank U (2) U (1) :n U (3) ≤ min rank(U (2)), rank(U (1) :nU (3)) .(EQUATIONTherefore a single CPµMoE's nth expert's matrix rank is bounded by min{I, O, R}.We now turn our attention to TRµMoEs, where we will see that the TR ranks R 1 , R 2 , R 3 translate very favorably into matrix rank at smaller computational cost than with CPµMoEs. First recall that TRµMoEs are parameterized instead by core tensorsU (1) ∈ R R1×N ×R2 , U (2) ∈ R R2×I×R3 , U (3) ∈ R R3×O×R1 , with chosen ranks R 1 , R 2 , R 3 .We can derive an expression to materialize expert n's matrix through the sum of matrix products of the TR cores as: 3) r3::EQUATIONEQUATION:n:EQUATION::r3EQUATIONThe matrix product rank inequality applies to each I × O matrix summand, whilst the matrix sum rank inequality applies to the outer matrix sum:EQUATION::r3EQUATIONrank U (3) r3:: U (1) :n: U (2) ::r3EQUATIONmin rank U (3) r3:: , rank U (1) :n: , rank U (2) ::r3 , .Consequently, expert n's materialized weight matrix in TRµMoEs has a more generous upper bound of minR 3 • min{R 1 , R 2 }, I, O 8 .Through this analysis, we observe that one can choose large values of R 3 yet small R 1 , R 2 to yield a high expert matrix rank with few parameters, justifying the choice of R 1 = R 2 = 4 in the main paper.One popular alternative decomposition is the Tucker decomposition [86] . Here we derive the resulting matrix rank of this alternative µMoE variant and detail why it's not as desirable as the proposed µMoE variants.A TuckerµMoE composes an µMoE weight tensor through the series of mode-n products [53] : 3) , where Z ∈ R R N ×R I ×R O is the so-called 'core tensor' andEQUATIONU 1 ∈ R N ×R N , U 2 ∈ R I×R I , U 3 ∈ R O×R Oare the 'factor matrices' for the tensor's three modes.Again following Kolda and Bader [53] a single expert n's weight matrix can be rewritten through the matricization involving the Kronecker product ⊗ as:EQUATIONwhereZ (2) ∈ R R I ×(R O •R N )is the so-called mode-2 (matrix) unfolding of the core tensor [53] . Consequently, the same rank inequality applies:EQUATION≤ min rank(U (2) I×R I), rank(Z (2)R I ×(R O •R N )), rank(U (1) n ⊗ U (3)EQUATIONWhere we see the much more restrictive matrix rank upper bound applies: min {min(I, R I ), min(R I , R O • R N ), min(O, R O )}.Truncated SVD MLP weights Given we've seen that parameter-efficient µMoE layers lead to low-rank expert weight matrices, a natural question is whether or not low-rankness in MLP linear layers' weight matrices is a reasonable assumption or constraint.Our strongest piece of evidence supporting the claim is experimental in nature: we've seen from the results in Section 4.3 that using all parameter-matched µMoE layers for both MLP mixers and GPT-2 models leads to no significant drop in accuracy from their linear layer counterparts (see also Appendix I for many more results).To investigate this further we perform a rank ablation on our trained MLP-Mixer model with the original linear layers' weights. Concretely, we compute the truncated SVD of each MLP block's 2 linear layer weight matrices. We explore the impact on the model's ImageNET1k validation set accuracy when using only the top-k singular vectors/values (the best rank-k approximation [87] ). The validation set accuracy using truncated SVD weights in every mixer block is plotted in Figure 7 -we see here that discarding as many as half the total number of (bottom) singular vectors/values to approximate the original weights leads to negligible difference to the validation set accuracy. In other words, low-rank approximations of MLP Mixers' weights retain their representational power sufficiently well to produce nearly the same validation set accuracy as the original model. Such findings are consistent with results in recent work in the language domain [88] , where low-rank approximations of MLP layers can even sometimes boost original performance. The accuracy retained by MLP Mixers here even after such aggressive rank reduction constitutes further evidence that full-rank weights are not always necessary.We plot in Figure 8 the parameter counts for µMoE layers as a function of the expert counts (sweeping from N = 2 experts through to N = 16, 384), relative to dense/sparse MoEs (with rank R 1 = R 2 = 4 TRµMoEs), for the first layer in a MLP-mixer channel-mixing block [80] . As can be seen, both µMoE variants are vastly more parameter-efficient than dense/sparse MoEs.Given TRµMoEs offer even better parameter efficiency for larger numbers of experts, we suggest opting for CPµMoEs when using expert counts less than ∼ 128, and considering TRµMoEs for higher values.Latency and memory usage comparisons between the µMoE, linear layers, and alternative MoEs are shown in Table 6 , where the µMoEs perform favorably. In the main paper, the fast forward passes are derived for a single level of expert hierarchy. One additional attractive property of µMoEs is their straightforward extension to multiple levels of expert hierarchy-one simply increments the number of modes of the weight tensor and includes another tensor contraction with new expert coefficients. Hierarchical µMoEs intuitively implement "and" operators in expert selection at each level, and further provide a mechanism through which to increase the total expert count at a small parameter cost. Here, we derive the fast forward passes for µMoE layers in their most general form with E levels of expert hierarchy. For intuition, we first further visualize µMoE layers with 2 levels of hierarchy in Figure 9 -note how we have an extra mode to the weight tensor, and an extra contraction over the new expert mode to combine its outputs.Given that hierarchical µMoEs involve very high-order tensors, we adopt the popular mode-n product [53] to express the forward passes in as readable a way as possible. The mode-n (vector) product of a tensor X ∈ R I1×I2×...×I N and vector u ∈ R In is denoted by X × n u [53] , with its elements given by:(X × n u) i1...in-1in+1...i N = In in=1 x i1i2...i N u in .We first introduce the formulation of an E-level hierarchical µMoE layer from Equation (1) in the main paper: given input z ∈ R I , the most general form of µMoE layer is parameterized by weight tensor W ∈ R N1×...×N E ×I×O and E many expert gating parameters {G e ∈ R I×Ne } E e=1 . The explicit, unfactorized forward pass is given by:EQUATIONwhere Equation ( 15) is expressed as sums over the E-many expert modes to make it clear that hierarchical µMoEs take convex combinations of E e=1 N e many experts' outputs (given there are N e experts at each level of hierarchy). With expert coefficients {a e ∈ R Ne } E e=1 , the factorized forward passes of the most general hierarchical µMoE layers are given for the two variants below.The full CPµMoE model of rank R has an implicit weight tensorW = R r=1 u (1) r • u (2) r • u (3) r • • • • • u (E+2) r ∈ R N1×•••×N E ×I×O , with factor matrices U (1) ∈ R R×N1 , . . . , U (E) ∈ R R×N E , U (E+1) ∈ R R×I , U (E+2) ∈ R R×O .The implicit, factorized forward pass is given by:EQUATION)In TR format, W ∈ R N1×•••×N E ×I×O has E + 2 factor tensors:U (1) ∈ R R1×N1×R2 , . . . , U (E) ∈ R R E ×N E ×R E+1 , U (E+1) ∈ R R E+1 ×I×R E+2 , U (E+2) ∈ R R E+2 ×O×R1, where R i are the manually chosen ranks. The weight tensor's elements are given by:w n1...n E io = tr U (1) :n1: • • • U (E):n E : U (E+1) :i:U (E+2):o:. We derive the fast factorized forward pass in terms of a series of mode-2 products:EQUATION)F Experimental detailsHere we provide the full experimental details and setups to reproduce the performance results in the paper for each of the networks. We further include the per-epoch accuracy plots for additional transparency into the training processes.The experimental configurations used to reproduce the performance results in the main paper follow as closely as possible those specified in the main paper of MLP-mixer [80] and open-source code (https://github.com/lucidrains/mlp-mixer-pytorch), the open-source code for NanoGPT (https://github.com/karpathy/nanoGPT) for GPT2 [81] , and the robust fine-tuning protocol of [89] for CLIP [61] . These values are summarized in Table 7 . We plot the learning curves for the training of both models in Figures 10 and 11 . Rank choices Throughout all experiments in the main paper, we fix the TRµMoE ranks for the first two modes to be R 1 = R 2 = 4. This way, we can maximize the effective expert matrix ranks at a low parameter cost, as shown in Appendix D.1.2. The final TR rank R 3 is varied to parameter-match the networks in question. For CPµMoEs, we set the single CP rank R to parameter-match the baselines.Training times Each MLP mixer model takes just under 3 days to train on 4xA100 80GB GPUs. The NanoGPT models take 2-3 days to train for 100k iterations, with the same resources.We initialize each element of the factor matrices/tensors for the input and output modes from a U [-√ k, √ k] distribution (following PyTorch's linear layers' initialization strategy), for k = 1/in_features, where in_features is the dimension of the input to each factor matrix/tensor during the factorized forward passes.Factor matrices for the expert modes are initialized to replicate the weight matrices along the expert mode (plus optional noise). For CPµMoEs, this corresponds to sampling the factor matrices' elements from a N (1, σ) distribution. For TRµMoEs, the weight matrices can instead be replicated along the expert mode by initializing each slice (e.g. G 1 (:, i, :)) as a diagonal matrix with its elements sampled from N (1, σ). In all our experiments we set σ := 1 to introduce noise along the first expert mode, and σ := 0 for additional expert modes.We first show in Figure 12 Expert coefficients color map:Figure 13 : Top-activating 32 tokens for the first unfiltered experts 1-6 (as ordered numerically) at layer 5 in the CPµMoE GPT2 model (Please find the next 6 experts in Figure 14 ).Here we provide additional evidence that the experts' specialization is mechanistically relevant to the functionality of the network, in the sense that we use them to steer the LLM's output.In particular, we use a larger GPT-2 model trained from scratch with µMoE layers at each MLP layer, using 2048 experts at every layer, following the setup in Section 4.3. 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 µMoE layer-we can consistently control the model to generate outputs aligned with specific themes. Illustrations of this approach, using 4 different manually chosen experts (with their first 8 generated samples) are shown in Figure 15 . 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 µMoE layer in facilitating controllable generation of language model outputs.However, we note that these initial results are hand-selected examples of some of the experts which do exhibit sensible specialization. We find many experts, when activated, do not steer the generations in such an interpretable high-level manner.Qualitative visualization Additional results to further substantiate the claims in the main paper about expert class-modularity are presented here. Firstly in Figure 16 are many more random images (of those with expert coefficient ≥ 0.5) of the first few experts as they are ordered numerically. Furthermore, when we use an even larger number of experts (i.e. 2048) we observe a select few experts developing what appear to be very fine-grained specialisms, as shown in Figure 17 . For example, images with large coefficients for #203 are often animals on top of laptops, whilst images with high coefficients for #1203 are animals eating corn.Counterfactual intervention barplots Next, we show barplots of the class labels whose test set accuracies are most changed under the counterfactual question in the main paper: "had (expert n) not contributed its weight, how would the class predictions have changed?". These are shown in Figure 18 and Figure 19 when using a CPµMoE as a final and penultimate layer respectively. As can be seen, we often observe that a higher number of experts (the final rows in brown color) lead to experts that, upon ablation, cause the model to lose almost all its accuracy for fewer classes. Experts here are chosen in numerical order and only those yielding ≥ 0.5 total accuracy change to any class upon counterfactual ablation.We find the use of the entmax activation function [54, 55] to produce more monosemantic experts, as quantified by the measure of polysemanticity used in the main paper. We show in Figure 20 the mean expert polysemanticity (of those experts that affect the class accuracy upon ablation) for CPµMoE-r512 final layer models fine-tuned with various numbers of experts. As can be seen, the entmax function consistently produces more monosemantic experts for larger total expert counts. We attribute this to the sparsity in entmax's post-activation distribution (whereas the softmax function can just as readily output a uniform distribution over all expert coefficients).The biggest issue of today's world is the pressure on the planet's fragile biosphere to become more carbon-rich. People have been making alternative energy sources since the dawn of civilisation, but there is no evidence for them, says the UNEP report.The biggest issue of today's world is the lack of unity. One of the biggest problems is that as the vast majority of the media don't want to say so, it is increasingly difficult to communicate what is called "the truth" with the public.In fact, in someThe biggest issue of today's world is the internet, which not only protects the most important aspects of our lives, but also has important benefits that greatly benefit other people. We can bypass the police, governments, banks and everyone else. Nobody wants to be charged with a crime, and nobodyThe biggest issue of today's world is not climate change, but the failure to make the right choices about who to send to war," he said. "The danger is that our leaders are unwilling to do the right thing and avoid a massive catastrophe." Watt noted that the United The biggest issue of today's world is this:As the world continues to get more and more food into the hands of the poor, we are also finding out that food is becoming harder to find. And the pace of demand for food is quickly in decline due to some key reasons The biggest issue of today's world is not how to manage a growing population (although we must do that), but how to manage a growing food production system. Growing food in urban areas is a challenge, but it also is a strategic opportunity. That's because city dwellThe biggest issue of today's world is the continuing need for mass-human immigration. Of the 3 million new visas issued in fiscal year 2010, 2.4 million have been approved, and the rest (1.6 million) will be, at best, temporary. That's because of The biggest issue of today's world is not the problems of the last century; it's the problems of today.The world of today is experiencing a profound change. We are less and less rational, more and more embittered, and the world will inevitably evolve to a situation The biggest issue of today's world is that of the 'new' value: 'a' with `result', from which all operations are overloaded.The other class that lists a type is `f`, which is a lambda, will be a type that is a new Value.The biggest issue of today's world is the lack of, as a constant:()<A> -> A<B> . And the get() type only makes the function return value, but, from type<A> implements A<T<T<T>] , itThe biggest issue of today's world is the accumulation of non-ditional , or __init__ , in every object.Here, we have an instance of , and an are, with type F , and . The first type is , and the second,The biggest issue of today's world is not to solve, but to be able to create. And it is: Any new construct cannot be created anywhere in the world (except at most one). These types of construct is a member of the set , where .The biggest issue of today's world is this: As a result, the world will grow one world after another, after all.We will see millions, billions, and billions of things. And the way we think of the world is in the identity of some kind of The biggest issue of today's world is that it's a problem of type (higher type) in an overloaded type, like a function. My type is a result, in a type like a, you get a result, after run(T), and the type t is a result is The biggest issue of today's world is the absence of a class."A big number is always a number, and the sum of all the other values is always a number. But the (bigest) type is always a, and the same is always a. But the The biggest issue of today's world is that, on every stream, no matter of size, exception is not a superclass, and a static class cannot be. It is a single construct and void is a new int. But a non-class is not any of its complement.The biggest issue of today's world is that we don't have enough police officers, and we have too many immigrants from the United States who aren't helping with the deportation of illegal immigrants."We have a drug problem, we've got a problem with drugs. We've gotThe biggest issue of today's world is the militarization of police, especially in the US. It also has to do with the media's ability to get to the bottom of what's happening in Ferguson, and the killing of multiple civilians.The US military's lack of accountability putsThe biggest issue of today's world is the accumulation of resources that are being moved, and the effects of that. In 2016, Donald Trump, Texas police officers are killed in the street."He has insulted the Mexican flag, the flag of the United States,"The biggest issue of today's world is not immigration, but the war on drugs. The war on drugs is a racist, violent criminal regime that is in the process of dismantling our country's efforts to keep us safe.The War on Drugs has been an unjust and unwThe biggest issue of today's world is this:As the United States says, the #8 target in Charlottesville #Charlottesville police death is "driving" violence, and other #2's getting killed in the US. https://t.w/r #2's killed are The biggest issue of today's world is that we are constantly on the same side and we are constantly on the side of the Palestinians, however much we are there as well," Clinton said in July. "And we also have to be here before the people that are so violent, so many The biggest issue of today's world is a lack of respect for the police, and the way law enverses that, police can be getting killed in a lot of other ways, including the use of drones. That's where I think, as a police officer, it's a greater The biggest issue of today's world is that we have a broken immigration system. We have a broken economic system. We have a police officer or a policeman who is unarmed and is being continually killed by a person who owns a vehicle. They have killed five officers this past week.The biggest issue of today's world is that of climate change, and many people are already building their homes on land from Mount Meru in the month-long dry monsoon that freezes into March. But that could be the biggest issue facing India; even in February, a low-pressureThe biggest issue of today's world is the risk of giant rain, especially in the British Isles, with the high temperature expected to reach an highest maximum of 2C. It's likely to become a drier pattern again today with winds of up to 10 miles per hour.The biggest issue of today's world is the warm temperatures that you can expect to see today, with the heat centering around the equator. But, it is not too cold, so we have to approach this problem cautiously. So far so good, anyway! So, hereThe biggest issue of today's world is not to burn down the planet. But it is to freeze it. We need a lot more water ice, and we need a lot more sunshine. But the climate has warmed, and now there is an area of relatively low temperature that is getting warmerThe biggest issue of today's world is this: As summer approaches, many farmers will grow one-to-two tons of wine every year.But, most will have some water. There's nothing to worry about, especially in the Southern California mountains. IfThe biggest issue of today's world is a lack of sunshine, which makes the heat evaporate away. The heat in the day is a much-remarensed mist. We're in a perfect storm of sunshine here, at high pressure, and it's going to be cloudy for aThe biggest issue of today's world is a lack of rainfall and extreme temperatures, so the coastal area will be dry to some extent. That means the rest of the area will be prone to the high temperature, but those regions should be dry, and the temperature could be as high as 100The biggest issue of today's world is that, on average, the atmosphere is too hot to be able to cool, although the warmer air is melting to some parts of the Arctic continent. [3] The area is covered with a haze of wind-shear (twice), with aThe biggest issue of today's world is that the Americans must implement a plan to end the war with Iraq and to improve postwar relations, American leaders said. However, the Soviet Union was always wary of the American policies.Korea, the North Koreans, the US and Great BritainThe biggest issue of today's world is coming to the realization that since the cold war era, the United States has turned its back on the Soviet Union. They secretly pushed through the historic deal in 1983, and the Soviets were all too eager for a resolution to stop the Berlin and Tehran.The biggest issue of today's world is the United States' difficulties with the East, and the US, in every sense of the world. They see no other use for the region is the obvious obvious of a new problem: the problem of Western imperialism. Also, the US sees no trouble The biggest issue of today's world is how to solve the issue of the Black Sea when it is. On a modern scale, the US Embassy in Moscow had changed the status quo through Washington's representative.From the time of Napoleon until the Russian Empire's recognition of Russia in 18The biggest issue of today's world is this, both sides are also pursuing a plan to avoid a permanent strategic alliance but they cannot reach for this is now calling its longterm friendship could be achieved with the current.The US government, especially in Washington, does not want to The biggest issue of today's world is a country's strategic response to Iraq's invasion and the rest of Iraq ruled in the 1950s.A week after the start of the war, it launched a massive operation to find Baghdad's exiled neighbour and to capture or occupy the city.The biggest issue of today's world is a fight between the two leaders over their mutual aspirations. Those were long-lasting issues over the Berlin-Ottoman-Rabid government has now tried to resolve. And those difficulties have left the country with a war-weary Russian President and The biggest issue of today's world is surely, on a broad level, some of the most consequential economic positions are being kept for half a year. Donald Trump and Vladimir Putin have been at every level to try to end a crisis over Russia -but the last days were also a We next perform an ablation study for the use of batch normalization (BN) before the activation function for the expert coefficients. We study CPµMoE final layer layers with CLIP ViT-B-32, quantifying BN's effect on expert classmonosemanticity as a function of the expert count. Concretely, we perform the same classlevel polysemanticity experiments as in the main paper, with and without batch normalization in Figure 21 . As can be seen clearly, the batch normalization models lead to individual experts that are increasingly class-monosemantic as desired (as a function of the total expert count). Here, we plot the expert load in Figure 22 to give a visual indication of how many images are processed by each expert with a e ≥ 0.5 for CPµMoE final layers fine-tuned on ImageNET1k with a CLIP backbone. Whilst clearly, not all experts have images with a coefficient of at least 0.5, we see a relatively uniform spread over all experts. Furthermore, we note the cost from 'dead' experts is not particularly troublesome in an µMoE given its factorized form-speaking informally, we would rather have too many experts than too few, so long as there exist select individual experts conducting the subcomputations of interest.Here, we compare the performance of parameter-matched µMoE final layers (for varying expert counts N ) to linear layers for fine-tuning large vision-language models (CLIP ViT-B-32) on ImageNET1k. Following the robust fine-tuning protocol of [89] , we use the largest possible batch size (to fit on one A100 GPU) of 4096, and the same learning rate of 3e -05.For µMoE layers, we reduce the layer ranks to parameter match single linear layers for each value of total expert count. We plot in Figure 23a the ImageNET1k validation loss after 10 epochs of training, where all expert counts out-perform the linear layers initialized the same default way with elements from U [-k, k]. However, to parameter-match single dense linear layers, we must decrease the µMoE layer rank upon increasing the expert count. This is a concrete example of where the extra parameter efficiency of TRµMoEs can come in useful (as discussed in Appendix D.1.2). Consequently, TRµMoEs' resulting expert matrix ranks are increasingly larger than that of CPµMoEs in the parameter-matched setting. For example, the parameter-matched layers with 512 experts in Figure 23a have a max expert matrix rank of 165 for the CPµMoE compared to a much larger 208 for the TRµMoE. Table 10 : Hierarchical µMoEs: Mean validation-set accuracy with a CLIP ViT-B-32 fine-tuned with hierarchical µMoE final layers on ImageNET1k. Shown are the number of parameters as the number of total experts increases to 8192 with 4 levels of hierarchy, and the corresponding number of parameters needed for each expert total using a hierarchy 1 µMoE, and regular MoE. Results are the average over 5 runs with different seeds. Additional expert modes for TRµMoEs have the additional ranks set equal to the corresponding number of experts at the new mode(s) (e.g. 2 and 4). the µMoEs also perform very favorably against the alternative MoE models when fine-tuning CLIP on ImageNET1k.Here we present more details about the fairness comparisons and metrics used in the main paper.• Equality of opportunity requires the true positive rates for the sensitive attribute subpopulations to be equal, defined in Hardt et al. [76] as P ( Ŷ = 1|A = 0, Y = 1) = P ( Ŷ = 1|A = 1, Y = 1) for sensitive attribute A, target attribute Y , and predictor Ŷ . In the first of our CelebA experiments we measure the absolute difference of the true positive rates between the 'blond female' and 'blond male' subpopulations for the 'blond hair' target attribute. For the second we measure the difference between that of the 'old female' and 'old male' subpopulations, taking the 'old' label as the true target attribute.• Standard deviation bias computes the standard deviation of the accuracy for the different subpopulations [77] . Intuitively, a small STD bias indicates similar performance across groups. • Max-Min Fairness quantifies the worst-case performance for the different demographic subpopulations [78] , with max min y∈Y,a∈A P ( Ŷ = y|A = a, Y = y). We compute this as the minimum of the test-set accuracy for the 4 subpopulations in each experiment.• Oversample we oversample the low-support subpopulation to balance the number of input images that have the sensitive attribute for the value of the target attribute wherein bias occurs. For example, we oversample the 'blond males' to match the number of 'blond females' for the first experiment, and oversample the number of 'old females' to match the number of 'old males' for the second. • Blind thresholding is implemented by unconditionally increasing/decreasing the logits in the target direction for all outputs. Concretely, the results in the main paper are achieved by setting λ := 2.5 and ā to a vector of ones in Equation ( 5) for all experiments. We find this value of λ to give us the best results for the attribute-blind re-writing [76] . • Adversarial debiasing we observe in Table 2 the same poor performance for the adversarial debiasing technique as is reported in Wang et al. [90] . We hypothesize that the same issues face the technique in our experimental setup. In particular, even in the absence of discriminative information for the 'gender' label in the final representation, information about correlated attributes (e.g. wearing makeup) are likely still present. This makes it fundamentally challenging to apply fairness-through-unawareness techniques in the CelebA multi-class setting.K Fairness: additional resultsThe full per-subpopulation test set accuracies are shown in Figure 25 for the two experiments in the main paper. The first rows show the accuracies before layer re-write, the second rows after re-write, and the third rows the absolute difference between the two. As can be seen in the 'before-after difference' final rows of Figure 25 , the proposed expert-conditional re-write provides much more precision in changing only the computation for the target populations.Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained? Answer: [NA] Justification: No human subjects involved.Incrementing the dimension of the second 'input' mode of the weight tensor W ∈ R N ×(I+1)×O and appending a 1 to the input vector z ∈ R I+1 folds a per-expert bias term into the computation.We adopt the convention of counting fused multiply-adds as one operation[57]. Note that the small additional expert coefficients cost is constant across models and thus ignored in comparisons.Setting R1 = 1 recovers a Tensor Train[35] µMoE.I.e. we include only experts that, when ablated in isolation, alter the class accuracy; please see the Appendix for discussion on expert load.We set λ := N for all experiments for simplicity, but we note that its value could require tuning in different experimental setups. The sign of λ is chosen to correct the bias in the target direction (whether to move the logits positively/negatively towards CelebA's e.g. young/old binary age labels respectively).The S-16 model is the largest configuration that fits into 4x80GB A100 GPUs using the original paper's batch size of 4096.Regardless of how large R3 is, the rank of the matrix cannot exceed min{I, O}. 


Decision: The authors proposed a multilinear Mixture of Experts Layer. This method, unlike existing MoE approaches, scales to a large number of experts, and the authors show improved interpretability, specialisation and accuracy. The proposed method is differentiable as it is a linear combination of the different experts, and scalable by using low-rank factorisations of these experts.

During the rebuttal, the authors also clarified that although they propose an MoE architecture, the motivation is to improve the parameterisation of standard dense models, which is why the parameter-counts of their models are not as high as Soft- or Sparse-MoE baselines.

The authors also performed numerous experiments throughout the rebuttal which addressed most of the reviewers concerns. Authors should incorporate all of these changes into the final camera ready version of the paper.

Reviews: review: paper_summary: paper_summary: The paper presents an interesting and novel approach with the µMoE layer, offering a method for scalable expert specialization. Extensive experiments in multi-dimensional have been conducted to show the effectiveness of the proposed method. However, the marginal improvements in results, coupled with several other weaknesses, including experimental variability, novelty of factorization, clarity of metrics, and scalability, limit the overall impact of the work.
strength_weakness: strength_weakness: 1. The idea of factorization is not new, and it seems that the authors simply apply factorization in the MoE context. The authors should specify and address any new challenges that arise when applying factorization in the MoE context.
2. Marginal Improvements and lack of multiple runs: The experimental results in Table 2 show only marginal improvements over existing methods, which might not be compelling enough to justify the additional complexity introduced by the µMoE layers. All experiments were only run once without reporting mean and standard deviation. It is recommended that the authors add results from at least three runs to provide a better understanding of the variability and reliability of their results.
3. The metrics used in the paper are not clearly explained. The authors should include a paragraph either in the experiments section or in the appendix to explain the metrics used.
2. The evaluation captures expert behavior only on the test set, not on out-of-distribution data, which is crucial for understanding the robustness and generalizability of the model.
4. While the paper qualitatively demonstrates interpretability, it lacks objective, quantitative measures of the interpretability gains compared to other methods.
8. It is unclear how the proposed methods perform across different model architectures and sizes. The authors should discuss whether the method scales well and provide evidence to support their claims.

More Comments:
1. The paper is generally well-written and clearly presents the proposed methodology and results. However, some sections, particularly those dealing with mathematical formulations, could benefit from additional explanations or visual aids to enhance comprehension.
2. Addressing the limitations and adding some discussions, such as out-of-distribution performance and quantitative measures for interpretability, would significantly strengthen the paper. Further research should also focus on demonstrating more significant performance gains and improving the ease of implementation to enhance the practical utility of the proposed method.
questions: questions: N/A
limitations: limitations: Yes, the authors addressed limitations and the direction for future research.
paper_summary: paper_summary: The paper proposes a new way of computation in a mixture of expert (MoE) layer. Unlike SparseMoE, instead of using the top-K operator and then doing full computation only for that expert, this paper do not do topK, but does a linear combination of all experts, but each expert is now a low rank matrix, which saves compute/parameters. The authors do show some properties that performing MoE computation in such a way can help in interpretability of the experts, their task specialization and so on. This can further help in certain cases like to mitigate biases.
strength_weakness: strength_weakness: - Figure 2 can be made much stronger by showing statistical differences between the two different cases -- like entropy of experts, overlap among experts, distribution, etc, and also comparing it with similar parameter and/or compute sized sparse/soft MoEs.

- Again, for Figure 3, there should be a way to compare it at least with sparse MoE. 

- Comparison with Sparse and SoftMoE in vision and LLMs are missing -- keeping the same parameter size and/or compute fixed across algorithms, as well as how does performance scale with more parameters in the MoE part. These are essential, as increasing the capacity of the network by increasing parameter count, while preserving the same compute is one of the main motivations behind MoEs, other than the interpretability lens which the works focuses mostly on. 

- Several ablations seems to be missing -- performance improvements when you increase number of experts, placement of the MoE layers, experiments on different model sizes, larger datasets, effect of different non-linearities in the gating function, etc.
questions: questions: - As mentioned above, I think it is quite important to show how \muMoE perform compared to SoftMoE and SparseMoE. 

- Making Figure 2 more statistically significant rather than anecdotes. 

- More experiments needs to be conducted to showcase the efficacy of the interpretability of the experts and its use cases. And how are those benefits compared to Soft and SparseMoE. 

Suggestions:
- I think the notation used in Eqn 1 is a bit difficult to follow, specifically with the colons used without brackets.
limitations: limitations: The authors have identified the limitations of the work.
paper_summary: paper_summary: This paper proposes the Multilinear Mixture of Experts (µMoE) layer for scalable expert specialization by performing an implicit computation on prohibitively large weight tensors entirely in factorized form. Both qualitative and quantitative results show that scaling µMoE layers when fine-tuning foundation models for vision tasks leads to specialized experts at the class level.
strength_weakness: strength_weakness: * It would be better if more results on language models could be included.
questions: questions: * The description of the dense µMoE layer in Section 3.1 is a little bit confusing to me, for example, what is the difference between it and the soft MoE layer?
* According to the methodology description, e.g. in Section 3.1.1, there is no “top-k” routing in the µMoE layer. Nevertheless, how are the “top-activating” experts selected in Figures. 4 & 5?
limitations: limitations: Yes.
paper_summary: paper_summary: The paper proposes the use of (factorized) Multilinear Mixture of Experts as an alternative to Sparse MoEs, which are prone to training issues due to the sparse top-k activation. Several factorization options are described and compared, which result in models of better accuracy in vision tasks, when matching the parameter count, compared to models using vanilla dense MLPs. The paper also includes experiments on a language task, where the proposed method is almost as good as the baseline.
strength_weakness: strength_weakness: - The results on GPT-2 are close to the dense MLPs, but not better than them. This is problematic, since this architecture is the biggest one trained (with 124M params) and probably the one trained on the richest data source. This begs the question of whether the proposed approach can scale well to bigger backbone architectures and datasets.
- Appendix C.2 discusses why the low-rank assumption in MLPs is reasonable. In particular, a trained MLP-Mixer is taken and the rank of the matrices is reduced with truncated SVD. This shows that the low rank approximation is reasonable once the model is trained, but completely ignores the effect of this approximation _during training_, especially when the architecture size and training data scales.
- Most importantly: there is not quality comparison between the proposed approach and Sparse MoEs. The recipes for training Sparse MoEs have improved significantly, and they are present in most state-of-the-art language models, and multi-modal models (for instance, take a look at the [Mixtral report](https://arxiv.org/abs/2401.04088), [DeepSeekMoE](https://arxiv.org/abs/2401.06066), [Gemini 1.5 report](https://arxiv.org/abs/2403.05530), [Databrick's DBRX](https://www.databricks.com/blog/introducing-dbrx-new-state-art-open-llm) which is available on HuggingFace, and many others). In vision tasks, Soft MoEs in ["From Sparse to Soft Mixture of Experts"](https://arxiv.org/abs/2308.00951) have shown also much better results than vanilla dense transformers and Sparse MoEs. Thus, either of these should be considered as a stronger baseline than the use of vanilla MLPs.
- A discussion about the speed of the different methods is also missing from the main paper, only parameter count is considered, which may not be the only (or main) limiting factor.

-----------
During the rebuttal the authors satisfactorily addressed two of my weaknesses. They also pointed out that the method is not an alternative to Sparse MoEs, but to provide a better parametrization of standard (dense) models. I still believe that Sparse MoEs would be a better baseline than Dense MoEs (formerly referred to as Soft MoEs in the paper), as I mentioned in my comment to the author's rebuttal. I also find concerning the fact that when the method is applied on relatively large models, it provides little benefit, although I acknowledge that it may be more interpretable. Given this, I'm slightly increasing my initial rating from "Reject" to "Borderline accept".
questions: questions: This is a suggestion, not a question: The term "Soft MoE" used to refer to the seminal work of Robert A. Jacobs et al. from 1991 can be confusing with the recent work ["From Sparse to Soft Mixture of Experts"](https://arxiv.org/abs/2308.00951) by Puigcerver et al, from last year. Since the term Soft is only used once in the 1991 work, I would suggest using "Dense MoE" to refer to it, which is also more opposed to "Sparse MoE".
limitations: limitations: No negative societal impacts are specific to this work, in my opinion.
paper_summary: paper_summary: In the paper, the author introduces $\mu$MoE layers, which factorize large weight tensors to facilitate implicit computation, thereby accelerating the computation process. By increasing the number of experts, the model enhances its specialization in vision tasks. Further, the author conducted experiments with the GPT2 and MLP-Mixer models, demonstrating that these models maintain high accuracy and exhibit increased specialization. Both qualitative and quantitative evaluations conducted by the author were commendably thorough and well-executed. The experiment is comprehensive.
strength_weakness: strength_weakness: 1. **Limited Scope:** The research primarily focuses on vision models and specific datasets, which may limit the generalizability of the findings. Expanding the validation to more diverse models (LLaMA, gemma, and $e.t.c$) and datasets could provide a more comprehensive view of the $\mu$MoE layers' effectiveness.
   
2. **Lack of Robustness Analysis:** The paper would benefit from including experiments that assess the robustness and performance of µMoE layers on out-of-distribution data, offering insights into the model's reliability under varied conditions.
   
3. **Need for Implementation Details:** There is a notable gap in the disclosure of practical implementation details and computational requirements for µMoE layers. Addressing this could aid in the reproducibility and application of the findings.
questions: questions: 1. Could you specify the computational hardware and configuration utilized for these experiments?
2. Is it feasible to apply this methodology to large-scale MoE Language Models with a substantial number of experts, similar to the extensive parameter count found in Mixtral 8$\times$7B?
limitations: limitations: The authors adequately discuss the limitations and do not have the potential negative societal impact.

Comments: title: Discussion summary
Thanks to Reviewer YTQ3; responses and suggested additions to paper's "limitations"
Might Reviewer-Su68 have any remaining questions?
Thanks to the reviewer
Responses to additional questions
Clarifying questions
Thanks to Reviewer XreH
comment: Dear Area Chair(s) and Reviewers,

We thank the AC kindly for their effort in handling our submission, and for the help facilitating the discussion with the reviewers.
Throughout the discussions, reviewers praised the paper for the well-motivated approach to the important goal of scalable expert specialization (**Reviewer XreH**), its technical soundness and `"extensive qualitative and quantitative evidence to support the claims"` (**Reviewer Su68**), and the innovative new metric introduced to quantify expert specialization (**Reviewer 1sqG**).

**All three of the reviewers to respond to our rebuttals (**qC51**, **XreH**, and **YTQ3**) raised their score following our new experimental results**: **Reviewer qC51** increased their score to a 6 following the new comparisons to Soft and Sparse MoEs, where muMoEs demonstrated Pareto-optimal improvements to accuracy for CLIP fine-tuning. **Reviewer YTQ3** also increased their score after these experiments and additional clarifications about the paper's aims. We also performed new experiments on steering GPT-2 models for **Reviewer XreH**, who raised their support for the paper too.

Two reviewers are yet to respond to our rebuttals. However, **our rebuttals address all concerns of the remaining reviewers**, which we summarize below (and describe how the key concerns are addressed):

- **Reviewer-1sqG**: the reviewer already leaves a favorable score of 7, and we propose to acknowledge additional limitations in the camera-ready (regarding robustness to OOD data and experiments on very large models).
- **Reviewer-Su68**: we clarify in our rebuttal that **the fairness results are indeed already over multiple runs**, and highlight where detail of the metrics can be located in the appendix.

Given the lack of responses from the remaining reviewers, we kindly request that you consider our rebuttal's comprehensive resolution of the remaining concerns when making your decision about the paper.

We will incorporate all changes suggested by the reviewers and clarifications that emerged from the insightful technical discussions into our camera-ready version. We are grateful to the reviewers for helping improve the paper significantly.

Sincerely,

Submission #3407 Authors.
Dear Reviewer YTQ3,

Thank you kindly for your time and effort in engaging with our rebuttal.

We are delighted to hear that we have addressed two of your primary concerns and that the goals of the paper have been well-clarified. We are grateful for the increase to the paper's score.

Your comments have been invaluable in helping improve the clarity of the paper's aims and objectives in the revised paper. We will make it clear throughout the revised text that the layer does not have the goal of outperforming the *accuracy/loss* of either Soft or Sparse MoEs. Regardless, we will include all comparisons to both Soft and Sparse MoEs and supporting discussions from this thread in the camera-ready version.

Regarding the final concern, the reviewer is absolutely correct to note that the original GPT-2 model indeed has a slightly lower validation loss (e.g. 0.01 below the TRMuMoE GPT-2 model), despite not providing any interpretable expert specialization of the muMoEs. Furthermore, whilst we don't anticipate any problems unique to muMoEs when scaling to even larger models (due to the layers' parameter-efficiency), we acknowledge that we have not validated this experimentally.

Inspired by the reviewer's comments, we propose to include the following limitation to the revised paper to acknowledge this explicitly (jointly in response to **Reviewer 1sqG**):

```
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 the reviewer once again for their careful and detailed assessment of the paper. We are more than happy to listen to any other suggestions they might have for the paper.

Sincerely,

Submission #3407 Authors.
I appreciate very much the author's response to my initial review.

I do understand better now the focus of the paper. I still believe believe that Sparse MoEs are a valid baseline: if Dense MoEs (or Soft MoEs as called in the initial version of the paper) are a valid baseline, which have a huge memory and runtime footprint, why not using instead Sparse MoEs which have the same memory footprint but far less runtime cost?

My concern about the scalability of the approach (i.e. not so good results on GPT-2, where models are bigger) was not satisfactorily addressed in your rebuttal.

Nevertheless, the other two weaknesses that I pointed out have been clarified. I also have to admit that, perhaps, I initially scored the paper too harshly, since I misjudged the focus of the paper. I apologise for any distress that this may have caused to the authors, and will increase my score accordingly.
Dear Reviewer-Su68,

We are keenly awaiting your feedback on our rebuttal, where we clarify your key concerns: that the fairness results are indeed based on repeated runs and highlight where the metric explanations are detailed.

As the author-reviewer discussion is almost finished, we would appreciate an opportunity following our rebuttal to answer any additional questions you have that would increase your score of the paper.

Sincerely,

Submission #3407 Authors.
Dear **Reviewer YTQ3**,

We eagerly await your response to our rebuttal where we address your primary concerns including the requested comparisons to sparse and soft MoEs and the requested benchmarking for speed.

We highlight that **Reviewer qC51** had a similar question about comparisons to Sparse MoEs in their initial review, and has since firmly increased their score after our rebuttal and experimental comparisons. In addition to this, **Reviewer XreH** also increased their score following our rebuttal and new experiments results. 

We very much appreciate the insights and attention to detail provided in your initial review, and we would be very grateful for the chance to respond to any remaining concerns or questions you have that would increase your support for the paper.

Sincerely,

Submission #3407 Authors.
Dear Reviewer qC51,

We thank you for the **strong support** and confidence in our paper.

We will include the linear layer experimental setup details and discussions about the 2 further insightful questions asked by the reviewer (regarding CP vs TR, and the significance of muMoE layer placement) in the revised main paper. Furthermore, we are open to any additional suggestions the reviewer might have to improve the paper. 

Sincerely,
Submission #3407 Authors.
Thanks to the authors for the answers. 
I have increased my score to 6. 
I strongly suggest the authors to include all the details/discussions presented here, to be included in the main paper.
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.
Dear Reviewers

This is another reminder to engage with the authors in this phase of the rebuttal. The deadline to respond to authors is EOD Anywhere on Earth timezone today.
Thanks to the authors for the extensive answers to the comments. While it answered most of my concerns, I have a few clarifying questions. 

1. Figure 1 (right) of rebuttal document, how do you get mutliple points for Linear, SparseMoE and SoftMoE? Is the dimension of MLP changed? or applied on multiple layers?
2. In Figure 21 of Appendix, while there is a slight increase in performance (max diff of +0.3%) for TR\mu MoE, there is hardly any difference for CP\mu MoE. Is there any explanation behind it?
3. While I understand the authors did experiments, visualizations with \mu MoE in different configuration - final layer, penultimate layer or all layers. Is there any conclusion that one can draw about the placement of expert layers from interpretability or performance and the trade-off if any?
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.
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.
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.
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.
Thank you for the detailed response to the reviews and additional experimental results. After reading all the reviews and responses, I will raise my score to 5.