Title: Multilinear Mixture of Experts: Scalable Expert Specialization through Factorization
Venue: NeurIPS 2024
Accepted: yes
Keywords: interpretability

mixture of experts
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.

Multilinear Mixture of Experts: Scalable Expert Specialization through Factorization

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 finegrained specialization. In this paper, we propose the Mu ltilinear M ixture o f E xperts ( µ 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 .

1 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 topK 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].

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 100 Mparameters-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.

Weshow both qualitatively (through visualization) and quantitatively (through counterfactual intervention) that increasing the number of µ MoEexperts 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.

2 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 topK 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.

3 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 I 1 and matrices X ∈ R I 1 × I 2 in lower- and upper-case boldface latin letters respectively. Tensors X ∈ R I 1 × I 2 × ... × 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 i 1 i 2 ...i d ∈ R . Finally, we use a colon to index into all elements along a particular mode: given X ∈ R I 1 × I 2 × I 3 for example, X :: i 3 ∈ R I 1 × I 2 or equivalently X (: , : , i 3 ) ∈ R I 1 × I 2 is the matrix at index i 3 of the final mode of the tensor. We use X × n u to denote the moden (vector) product [53] of a tensor X ∈ R I 1 × I 2 × ... × I N and vector u ∈ R I n whose resulting elements are given by ( X × n u ) i 1 ...i n -1 i n +1 ...i N = ∑ I n i n =1 x i 1 i 2 ...i N u i n .

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:

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 transformations 2 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

+

=

+

Input vector

Output vector

Expert

coe

ffi

cients

Weight tensor

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).

3.1.1 Computation in factorized form

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 topK 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 write W = ∑ R r =1 u (1) r ◦ u (2) r ◦ u (3) r ∈ R N × I × O as a sum of R outer products, with factor matrices U (1) ∈ R R × N , U (2) ∈ R R × I , U (3) ∈ R R × O . This reduces the parameter count from NIO (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:

with 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 NIO , 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.

TR µ MoE Wepropose a second µ MoEvariant 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 R 1 × N × R 2 , U (2) ∈ R R 2 × I × R 3 , U (3) ∈ R R 3 × O × R 1 , where R i are the manually chosen ranks 4 . The weight tensor's elements in TR format are given by: w nio = tr ( U (1) : n : U (2) : i : U (3) : 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:

yielding a modified FLOP count of ( R 1 NR 2 + R 2 IR 3 + R 1 R 2 R 3 + R 1 OR 3 ) with just ( R 1 NR 2 + R 2 IR 3 + R 3 OR 1 ) parameters. With large N contributing to the computational cost only through R 1 NR 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.

3.1.2 µ MoEs recover dense MoEs as a special case

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 glyph[latticetop] 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].

4 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).

4.1 Expert specialism: visualization & intervention

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.

Expert 1

Expert 3

Expert 1

Expert 2

μ

MoE: 256 total experts

μ

MoE: 32 total experts

admiral

Bdmiral

admiral

eiwpe

aamira

admira

Amencan

lrnousine

American

Amencan

Mrican.c

komodo.dr

common_g

0.6

0.52

0.51

0.5

0.55

0.67

0.51

0.81

0.6

0.55

0.65

0.73

0.64

0.97

0.54

0.56

powerdr

ayn_mowe

awn_mowe

powerdri

cham_saw

204

mne

monarch

mailoox

0.5

0.71

0.61

0.52

0.51

063

0.76

0.65

0.68

0.65

0.54

0.52

0.73

290

0.72

4.1.1 Qualitative results

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.

4.1.2 Quantitative results: expert monosemanticity

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 µ MoElayer, 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 , ˆ 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 -ˆ y ( n ) c ) /y c . At the two extremes, when the full network's accuracy for

32

64

128

256

512

1024

Total # experts used in CP

MoE layer (log scale)

0.75

1.00

1.25

1.50

1.75

2.00

2.25

Mean expert polysemanticity

p

(

n

)

Experts' class-level polysemanticity (

)

Final layer

MoE (clip)

Penultimate layer

MoE (clip)

Final layer

MoE (dino)

Penultimate layer

MoE (dino)

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:

where 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.

4.2 Expert re-writing: conditional bias correction

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 ¯ a = 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:

Here, the term λ ¯ a ∈ 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

Layer 2, Expert 1

Context

Patch

Context

Patch

Layer 2, Expert 2

Layer 2, Expert 3

Context

Patch

Layer 7, Expert 1

Context

Patch

Layer 7, Expert 2

Context

Patch

Layer 7, Expert 3

1.0

10

1.0

10

1.0

1.0

1.0

1.0

10

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.

4.3 Large language/vision µ MoE networks

Finally, we train from scratch 12 layer 124 M-parameter GPT-2 [81] LLMs on OpenWebText [82] for the language domain and 8 layer S-16 variant 7 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:

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 W (1) ∈ R N × I × H , W (2) ∈ R N × H × O are the (implicit) µ MoE weight

Layer 5, Expert 8

Layer 6, Expert 1

Layer 5, Expert 37

Expert coe ffi

cients color map:

Layer 6, Expert 16

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 1 Mgenerated 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

many unfiltered experts in Figures 13 and 14). Whilst monosemanticity is not always attained, µ MoE layers nonetheless facilitate a level of specialism not facilitated by dense MLP layers.

One important result here is that µ MoE networks in this setup are significantly more parameterefficient than both dense and sparse MoEs with the same expert count, as shown in Table 4. For example, GPT-2 models with 256 sparse/dense MoE experts require a prohibitive 14 . 5 B MLP parameters alone, relative to just 57 MMLPparameters with µ MoEs of the same expert counts.

µ MoE performance Finally, we substantiate our claim that networks pre-trained and fine-tuned with parameter-matched µ MoE layers are competitive with their existing linear layer alternatives across multiple domains/machine learning tasks. We present in Table 3 the performance results for MLP-Mixer S-16 [80], NanoGPT GPT-2 [81], and (fine-tuned) CLIP ViT-B-32 [61] models on the OWT and ImageNET1k datasets. Following Section 4.1.1, we replace all linear

layers with µ MoE blocks (and a single µ MoE final layer for fine-tuning CLIP). We initialize all linear layers following the default PyTorch U [ -k, k ] initialization for a fair comparison. Please see Appendix F for experimental details and learning curves, and Appendix I for experiments with varying expert count and hierarchical µ MoEs. Crucially, whilst µ MoE layers provide additional interpretability benefits through scalable expert specialization, they do not sacrifice accuracy when parameter-matched to MLP blocks, as seen from the comparable performance.

5 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.

References

Robert A Jacobs, Michael I Jordan, Steven J Nowlan, and Geoffrey E Hinton. Adaptive mixtures of local experts. Neural computation , 3(1):79-87, 1991.

Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, Gianna Lengyel, Guillaume Bour, Guillaume Lample, Lélio Renard Lavaud, Lucile Saulnier, Marie-Anne Lachaux, Pierre Stock, Sandeep Subramanian, Sophia Yang, Szymon Antoniak, Teven Le Scao, Théophile Gervet, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mixtral of experts, 2024.

Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. GShard: Scaling giant models with conditional computation and automatic sharding. In Int. Conf. Learn. Represent. (ICLR) , 2021.

William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. The Journal of Machine Learning Research , 23(1):5232-5270, 2022.

Trevor Gale, Deepak Narayanan, Cliff Young, and Matei Zaharia. Megablocks: Efficient sparse training with mixture-of-experts. Proceedings of Machine Learning and Systems , 5, 2023.

Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. Adv. Neural Inform. Process. Syst. (NeurIPS) , 34:8583-8595, 2021.

Basil Mustafa, Carlos Riquelme Ruiz, Joan Puigcerver, Rodolphe Jenatton, and Neil Houlsby. Multimodal contrastive learning with LIMoe: the language-image mixture of experts. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Adv. Neural Inform. Process. Syst. (NeurIPS) , 2022.

David Ha, Andrew M. Dai, and Quoc V. Le. Hypernetworks. In Int. Conf. Learn. Represent. (ICLR) , 2017.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Adv. Neural Inform. Process. Syst. (NeurIPS) , 30, 2017.

Yizeng Han, Gao Huang, Shiji Song, Le Yang, Honghui Wang, and Yulin Wang. Dynamic neural networks: A survey. IEEE Trans. Pattern Anal. Mach. Intell. (TPAMI) , 44(11):7436-7456, 2021.

Yinpeng Chen, Xiyang Dai, Mengchen Liu, Dongdong Chen, Lu Yuan, and Zicheng Liu. Dynamic convolution: Attention over convolution kernels. In IEEE Conf. Comput. Vis. Pattern Recog. (CVPR) , pages 11030-11039, 2020.

Robert A Jacobs, Michael I Jordan, and Andrew G Barto. Task decomposition through competition in a modular connectionist architecture: The what and where vision tasks. Cognitive science , 15(2):219-250, 1991.

Zachary C. Lipton. The mythos of model interpretability. Communications of the ACM , 61(10): 36-43, September 2018. ISSN 1557-7317.

Marco Tulio Ribeiro, Sameer Singh, and Carlos Guestrin. " why should i trust you?" explaining the predictions of any classifier. In Proceedings of the 22nd ACM SIGKDD international conference on knowledge discovery and data mining , pages 1135-1144, 2016.

Yuzhen Mao, Zhun Deng, Huaxiu Yao, Ting Ye, Kenji Kawaguchi, and James Zou. Last-layer fairness fine-tuning is simple and effective for neural networks. In Proceedings of the 2nd Workshop on Spurious Correlations, Invariance and Stability at the International Conference on Machine Learning (ICML 2023) , 2023.

Valeriia Cherepanova, Vedant Nanda, Micah Goldblum, John P Dickerson, and Tom Goldstein. Technical challenges for training fair neural networks. arXiv preprint arXiv:2102.06764 , 2021.

Noam Shazeer, *Azalia Mirhoseini, *Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-ofexperts layer. In Int. Conf. Learn. Represent. (ICLR) , 2017.

Muqeeth Mohammed, Haokun Liu, and Colin Raffel. Models with conditional computation learn suboptimal solutions. In I Can't Believe It's Not Better Workshop: Understanding Deep Learning Through Empirical Falsification , 2022.

Joan Puigcerver, Carlos Riquelme, Basil Mustafa, and Neil Houlsby. From sparse to soft mixtures of experts. In Int. Conf. Learn. Represent. (ICLR) , 2024.

Yanai Elazar, Shauli Ravfogel, Alon Jacovi, and Yoav Goldberg. Amnesic probing: Behavioral explanation with amnesic counterfactuals. Transactions of the Association for Computational Linguistics , 9:160-175, 2021.

Emmanuel Bengio, Pierre-Luc Bacon, Joelle Pineau, and Doina Precup. Conditional computation in neural networks for faster models. In Int. Conf. Mach. Learn. Worksh. (ICMLW) , 2015.

M.I. Jordan and R.A. Jacobs. Hierarchical mixtures of experts and the em algorithm. In Proceedings of 1993 International Conference on Neural Networks (IJCNN-93-Nagoya, Japan) , volume 2, pages 1339-1344 vol.2, 1993. doi: 10.1109/IJCNN.1993.716791.

David Eigen, Marc'Aurelio Ranzato, and Ilya Sutskever. Learning factored representations in a deep mixture of experts. In Int. Conf. Mach. Learn. Worksh. (ICMLW) , volume abs/1312.4314, 2013.

Brandon Yang, Gabriel Bender, Quoc V Le, and Jiquan Ngiam. Condconv: Conditionally parameterized convolutions for efficient inference. Adv. Neural Inform. Process. Syst. (NeurIPS) , 32, 2019.

Nan Du, Yanping Huang, Andrew M Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, et al. Glam: Efficient scaling of language models with mixture-of-experts. In Int. Conf. Mach. Learn. (ICML) , pages 5547-5569. PMLR, 2022.

Shashank Gupta, Subhabrata Mukherjee, Krishan Subudhi, Eduardo Gonzalez, Damien Jose, Ahmed H Awadallah, and Jianfeng Gao. Sparsely activated mixture-of-experts are robust multi-task learners. arXiv preprint arXiv:2204.07689 , 2022.

Suchin Gururangan, Mike Lewis, Ari Holtzman, Noah Smith, and Luke Zettlemoyer. Demix layers: Disentangling domains for modular language modeling. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies . Association for Computational Linguistics, 2022. doi: 10. 18653/v1/2022.naacl-main.407.

Aya Abdelsalam Ismail, Sercan O Arik, Jinsung Yoon, Ankur Taly, Soheil Feizi, and Tomas Pfister. Interpretable mixture of experts. Transactions on Machine Learning Research , 2023. ISSN 2835-8856.

Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. Base layers: Simplifying training of large, sparse models. In Int. Conf. Mach. Learn. (ICML) , 2021.

Lemeng Wu, Mengchen Liu, Yinpeng Chen, Dongdong Chen, Xiyang Dai, and Lu Yuan. Residual mixture of experts, 2022.

Fuzhao Xue, Ziji Shi, Futao Wei, Yuxuan Lou, Yong Liu, and Yang You. Go wider instead of deeper. In Conf. on Artifi. Intel. (AAAI) , volume 36, pages 8779-8787, 2022.

Andrew Davis and Itamar Arel. Low-rank approximations for conditional feedforward computation in deep neural networks. arXiv preprint arXiv:1312.4461 , 2013.

Yunsheng Li, Yinpeng Chen, Xiyang Dai, mengchen liu, Dongdong Chen, Ye Yu, Lu Yuan, Zicheng Liu, Mei Chen, and Nuno Vasconcelos. Revisiting dynamic convolution via matrix decomposition. In Int. Conf. Learn. Represent. (ICLR) , 2021.

Ze-Feng Gao, Peiyu Liu, Wayne Xin Zhao, Zhong-Yi Lu, and Ji-Rong Wen. Parameter-efficient mixture-of-experts architecture for pre-trained language models. In Proceedings of the 29th International Conference on Computational Linguistics , Gyeongju, Republic of Korea, October 2022. International Committee on Computational Linguistics.

I. Oseledets. Tensor-train decomposition. SIAM J. Sci. Comput. , 33:2295-2317, 2011.

Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.09685 , 2021.

Alexander Novikov, Dmitrii Podoprikhin, Anton Osokin, and Dmitry P Vetrov. Tensorizing neural networks. Adv. Neural Inform. Process. Syst. (NeurIPS) , 28, 2015.

Timur Garipov, Dmitry Podoprikhin, Alexander Novikov, and Dmitry Vetrov. Ultimate tensorization: compressing convolutional and fc layers alike. arXiv preprint arXiv:1611.03214 , 2016.

Alexander Novikov, Mikhail Trofimov, and Ivan Oseledets. Exponential machines. In Int. Conf. Learn. Represent. Worksh. , 2017.

Markos Georgopoulos, James Oldfield, Mihalis A Nicolaou, Yannis Panagakis, and Maja Pantic. Mitigating demographic bias in facial datasets with style-based multi-attribute transfer. Int. J. Comput. Vis. (IJCV) , 129(7):2288-2307, 2021.

Francesca Babiloni, Ioannis Marras, Gregory Slabaugh, and Stefanos Zafeiriou. Tesa: Tensor element self-attention via matricization. In IEEE Conf. Comput. Vis. Pattern Recog. (CVPR) , pages 13945-13954, 2020.

Markos Georgopoulos, Grigorios Chrysos, Maja Pantic, and Yannis Panagakis. Multilinear latent conditioning for generating unseen attribute combinations. In Int. Conf. Mach. Learn. (ICML) , 2020.

Yixin Cheng, Grigorios G. Chrysos, Markos Georgopoulos, and Volkan Cevher. Multilinear operator networks, 2024.

Jean Kossaifi, Antoine Toisoul, Adrian Bulat, Yannis Panagakis, Timothy M. Hospedales, and Maja Pantic. Factorized higher-order cnns with an application to spatio-temporal emotion estimation. In IEEE Conf. Comput. Vis. Pattern Recog. (CVPR) . IEEE, June 2020.

Adrian Bulat, Jean Kossaifi, Georgios Tzimiropoulos, and Maja Pantic. Incremental multidomain learning with network latent tensor factorization. In Conf. on Artifi. Intel. (AAAI) , volume 34, pages 10470-10477, 2020.

Yongxin Yang and Timothy M. Hospedales. Deep multi-task representation learning: A tensor factorisation approach. In Int. Conf. Learn. Represent. (ICLR) , 2017.

Grigorios G Chrysos, Stylianos Moschoglou, Giorgos Bouritsas, Yannis Panagakis, Jiankang Deng, and Stefanos Zafeiriou. P-nets: Deep polynomial neural networks. In IEEE Conf. Comput. Vis. Pattern Recog. (CVPR) , pages 7325-7335, 2020.

Grigorios G. Chrysos, Stylianos Moschoglou, Giorgos Bouritsas, Jiankang Deng, Yannis Panagakis, and Stefanos P Zafeiriou. Deep polynomial neural networks. IEEE Trans. Pattern Anal. Mach. Intell. (TPAMI) , page 1-1, 2021. ISSN 1939-3539.

Francesca Babiloni, Ioannis Marras, Filippos Kokkinos, Jiankang Deng, Grigorios Chrysos, and Stefanos Zafeiriou. Poly-nl: Linear complexity non-local layers with 3rd order polynomials. In Int. Conf. Comput. Vis. (ICCV) , pages 10518-10528, 2021.

Jean Kossaifi, Aran Khanna, Zachary Lipton, Tommaso Furlanello, and Anima Anandkumar. Tensor contraction layers for parsimonious deep nets. In IEEE Conf. Comput. Vis. Pattern Recog. Worksh. (CVPRW) , pages 26-32, 2017.

Francesca Babiloni, Thomas Tanay, Jiankang Deng, Matteo Maggioni, and Stefanos Zafeiriou. Factorized dynamic fully-connected layers for neural networks. In Int. Conf. Comput. Vis. Worksh. (ICCVW) , pages 1374-1383, October 2023.

Frank Lauren Hitchcock. The expression of a tensor or a polyadic as a sum of products. Journal of Mathematics and Physics , 6:164-189, 1927.

Tamara G. Kolda and Brett W. Bader. Tensor decompositions and applications. SIAM Review , 51(3):455-500, 2009. doi: 10.1137/07070111X.

Ben Peters, Vlad Niculae, and André F. T. Martins. Sparse sequence-to-sequence models. In Anna Korhonen, David Traum, and Lluís Màrquez, editors, Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics , pages 1504-1519, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.18653/v1/P19-1146.

Gonçalo M. Correia, Vlad Niculae, and André F. T. Martins. Adaptively sparse transformers. In Kentaro Inui, Jing Jiang, Vincent Ng, and Xiaojun Wan, editors, Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP) , pages 2174-2184, Hong

Kong, China, November 2019. Association for Computational Linguistics. doi: 10.18653/v1/ D19-1223.

J. Douglas Carroll and Jih Jie Chang. Analysis of individual differences in multidimensional scaling via an n-way generalization of 'eckart-young' decomposition. Psychometrika , 35: 283-319, 1970.

fvcore: Flop counter for pytorch models. https://github.com/facebookresearch/ fvcore . Accessed: 2024-05-16.

Qibin Zhao, Guoxu Zhou, Shengli Xie, Liqing Zhang, and Andrzej Cichocki. Tensor ring decomposition. ArXiv , abs/1606.05535, 2016.

Lee Sharkey. A technical note on bilinear layers for interpretability. arXiv preprint arXiv:2305.03452 , 2023.

Michael T. Pearce, Thomas Dooms, and Alice Rigg. Weight-based decomposition: A case for bilinear MLPs, 2024.

Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In Int. Conf. Mach. Learn. (ICML) , 2021.

Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In Int. Conf. Comput. Vis. (ICCV) , 2021.

Gabriel Ilharco, Mitchell Wortsman, Samir Yitzhak Gadre, Shuran Song, Hannaneh Hajishirzi, Simon Kornblith, Ali Farhadi, and Ludwig Schmidt. Patching open-vocabulary models by interpolating weights. Adv. Neural Inform. Process. Syst. (NeurIPS) , 35:29262-29277, 2022.

Gabriel Ilharco, Marco Tulio Ribeiro, Mitchell Wortsman, Ludwig Schmidt, Hannaneh Hajishirzi, and Ali Farhadi. Editing models with task arithmetic. In Int. Conf. Learn. Represent. (ICLR) , 2023.

Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravec, Zac Hatfield-Dodds, Robert Lasenby, Dawn Drain, Carol Chen, et al. Toy models of superposition. arXiv preprint arXiv:2209.10652 , 2022.

Tilman Räuker, Anson Ho, Stephen Casper, and Dylan Hadfield-Menell. Toward transparent ai: A survey on interpreting the inner structures of deep neural networks. In 2023 IEEE Conference on Secure and Trustworthy Machine Learning (SaTML) , pages 464-483. IEEE, 2023.

Shauli Ravfogel, Grusha Prasad, Tal Linzen, and Yoav Goldberg. Counterfactual interventions reveal the causal effect of relative clause representations on agreement prediction. In Arianna Bisazza and Omri Abend, editors, Proceedings of the 25th Conference on Computational Natural Language Learning , pages 194-209, Online, November 2021. Association for Computational Linguistics.

Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in gpt. Adv. Neural Inform. Process. Syst. (NeurIPS) , 35:17359-17372, 2022.

Cynthia Rudin. Stop explaining black box machine learning models for high stakes decisions and use interpretable models instead. Nature machine intelligence , 1(5):206-215, 2019.

Stephen Casper. Broad critiques of interpretability research. 2023. URL https://www. alignmentforum.org/s/a6ne2ve5uturEEQK7/p/gwG9uqw255gafjYN4 .

Shlomi Hod, Daniel Filan, Stephen Casper, Andrew Critch, and Stuart Russell. Quantifying local specialization in deep neural networks. arXiv preprint arXiv:2110.08058 , 2021.

Joy Buolamwini and Timnit Gebru. Gender shades: Intersectional accuracy disparities in commercial gender classification. In Conference on fairness, accountability and transparency , pages 77-91. PMLR, 2018.

Timnit Gebru, Jamie Morgenstern, Briana Vecchione, Jennifer Wortman Vaughan, Hanna Wallach, Hal Daumé Iii, and Kate Crawford. Datasheets for datasets. Communications of the ACM , 64(12):86-92, 2021.

Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Int. Conf. Comput. Vis. (ICCV) , December 2015.

Saachi Jain, Hannah Lawrence, Ankur Moitra, and Aleksander Madry. Distilling model failures as directions in latent space. In Int. Conf. Learn. Represent. (ICLR) , 2023.

Moritz Hardt, Eric Price, and Nati Srebro. Equality of opportunity in supervised learning. In Adv. Neural Inform. Process. Syst. (NeurIPS) , 2016.

Mei Wang and Weihong Deng. Mitigating bias in face recognition using skewness-aware reinforcement learning. In IEEE Conf. Comput. Vis. Pattern Recog. (CVPR) , pages 9322-9331, 2020.

Preethi Lahoti, Alex Beutel, Jilin Chen, Kang Lee, Flavien Prost, Nithum Thain, Xuezhi Wang, and Ed Chi. Fairness without demographics through adversarially reweighted learning. Adv. Neural Inform. Process. Syst. (NeurIPS) , 33:728-740, 2020.

Mohsan Alvi, Andrew Zisserman, and Christoffer Nellåker. Turning a blind eye: Explicit removal of biases and variation from deep neural network embeddings. In Proceedings of the European Conference on Computer Vision (ECCV) Workshops , 2018.

Ilya O Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, et al. MLPmixer: An all-MLP architecture for vision. Adv. Neural Inform. Process. Syst. (NeurIPS) , 34: 24261-24272, 2021.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. OpenAI Blog , 2019. URL https://cdn.openai.com/better-language-models/language_models_ are_unsupervised_multitask_learners.pdf .

Aaron Gokaslan and Vanya Cohen. Openwebtext corpus. http://Skylion007.github.io/ OpenWebTextCorpus , 2019.

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A largescale hierarchical image database. In IEEE Conf. Comput. Vis. Pattern Recog. (CVPR) , pages 248-255, 2009.

Tolga Bolukbasi, Adam Pearce, Ann Yuan, Andy Coenen, Emily Reif, Fernanda Viégas, and Martin Wattenberg. An interpretability illusion for bert. arXiv preprint arXiv:2104.07143 , 2021.

Alex Rogozhnikov. Einops: Clear and reliable tensor manipulations with einstein-like notation. In Int. Conf. Learn. Represent. (ICLR) , 2022.

Ledyard R. Tucker. Some mathematical notes on three-mode factor analysis. Psychometrika , 31:279-311, 1966.

Carl Eckart and Gale Young. The approximation of one matrix by another of lower rank. Psychometrika , 1(3):211-218, 1936.

Pratyusha Sharma, Jordan T. Ash, and Dipendra Misra. The truth is in there: Improving reasoning in language models with layer-selective rank reduction. In Int. Conf. Learn. Represent. (ICLR) , 2024.

Mitchell Wortsman, Gabriel Ilharco, Jong Wook Kim, Mike Li, Simon Kornblith, Rebecca Roelofs, Raphael Gontijo Lopes, Hannaneh Hajishirzi, Ali Farhadi, Hongseok Namkoong, et al. Robust fine-tuning of zero-shot models. In IEEE Conf. Comput. Vis. Pattern Recog. (CVPR) , pages 7959-7971, 2022.

Zeyu Wang, Klint Qinami, Ioannis Christos Karakozis, Kyle Genova, Prem Nair, Kenji Hata, and Olga Russakovsky. Towards fairness in visual recognition: Effective strategies for bias mitigation. In IEEE Conf. Comput. Vis. Pattern Recog. (CVPR) , pages 8919-8928, 2020.

Appendix

Table of Contents

A Broader impact

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.

B Fast µ MoE implementations

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:

B.1 CP µ MoE einsum implementation

The CP µ MoE forward pass can be implemented with:

And a two-level hierarchical CP µ MoE with an additional factor matrix as:

B.2 TR µ MoE einsum implementation

TR µ MoEs can be implemented with:

And a two-level hierarchical version with an additional TR-core as:

C µ MoE forward pass visualization

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).

+

=

+

Output vector

Expert

coe

ffi

cients

Weight tensor

MoE as

tensor contraction

[1/5]

+

=

+

Output vector

Expert

coe

ffi

cients

Weight tensor

Contract weight tensor

with the input vector

(along 2nd mode)

[2/5]

+

+

Contract weight tensor

with the input vector

(along 2nd mode)

=

Output vector

Expert

coe

ffi

cients

Intermediate output:

[3/5]

Input vector

+

=

+

Output vector

Expert

coe

ffi

cients

Sum over the N partial outputs

(tensor contraction with expert coefficients)

[4/5]

Input vector

+

=

+

Output vector

Expert

coe

ffi

cients

Sum over the N partial outputs

(tensor contraction with expert coefficients)

[5/5]

μ MoE's forward pass Visualized step-by-step

D Decomposition choice, matrix rank, and computational cost

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.

D.1 Tensor ranks to matrix rank

One 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.

D.1.1 CP µ MoEs: rank analysis

CP µ MoEs are parameterized by factor matrices U (1) ∈ R R × N , U (2) ∈ R R × I , U (3) ∈ 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 as

where glyph[circledot] is the Khatri-Rao product [53], and U (1) : 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

Therefore a single CP µ MoE's n th expert's matrix rank is bounded by min { I, O, R } .

D.1.2 TR µ MoEs: rank analysis

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 tensors U (1) ∈ R R 1 × N × R 2 , U (2) ∈ R R 2 × I × R 3 , U (3) ∈ R R 3 × O × R 1 , 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:

The matrix product rank inequality applies to each I × O matrix summand, whilst the matrix sum rank inequality applies to the outer matrix sum:

Consequently, expert n 's materialized weight matrix in TR µ MoEs has a more generous upper bound of min { R 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.

D.1.3 Tucker µ MoEs: rank analysis

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 moden products [53]: W = Z × 1 U (1) × 2 U (2) × 3 U (3) , where Z ∈ R R N × R I × R O is the so-called 'core tensor' and U 1 ∈ R N × R N , U 2 ∈ R I × R I , U 3 ∈ R O × R O are 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:

where Z (2) ∈ R R I × ( R O · R N ) is the so-called mode2 (matrix) unfolding of the core tensor [53]. Consequently, the same rank inequality applies:

Where 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 ) } . Thus in practice, both R I , R O need to be large to yield a large matrix rank, which is in conflict with the goal of maintaining a moderate number of parameters.

D.2 Why is low-rankness a reasonable assumption?

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 topk singular vectors/values (the best rankk 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

20

40

60

80

100

Percentage of singular vectors kept

0

10

20

30

40

50

60

70

ImageNET1k test set accuracy

68.50 69.41 69.71 70.02 70.11 70.29

65.28

26.61

9.22

0.44

MLP-Mixer ImageNET1K accuracy with

Truncated SVD MLP 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.

D.3 MoE/ µ MoE parameter count comparisons

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.

8

64

512

4096

Number of experts (log scale)

10 6

10 7

10 8

10 9

10 10

Number of parameters (log scale)

MoE layer parameter counts

with dimensions

I

=512,

O

=2048

CPMMoE (

R

=512)

TRMMoE (

R

3 =128)

Soft/Sparse MoE

E Hierarchical µ MoE model derivations

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 moden product [53] to express the forward passes in as readable a way as possible. The moden (vector) product of a tensor X ∈ R I 1 × I 2 × ... × I N and vector u ∈ R I n is denoted by X × n u [53], with its elements given by:

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 N 1 × ... × N E × I × O and E many expert gating parameters { G e ∈ R I × N e } E e =1 . The

[

n_outputs]

+

=

+

[

[

n_inputs]

Input vector

Output vector

+

+

Weight tensor

[

n_experts_2]

1st level's expert

coe

ffi

cients

2nd level's expert

coe

ffi

cients

explicit, unfactorized forward pass is given by:

where 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 N e } E e =1 , the factorized forward passes of the most general hierarchical µ MoE layers are given for the two variants below.

E.1 Hierarchical CP µ MoE

The full CP µ MoE model of rank R has an implicit weight tensor W = ∑ R r =1 u (1) r ◦ u (2) r ◦ u (3) r ◦· · · ◦ u ( E +2) r ∈ R N 1 ×···× N E × I × O , with factor matrices U (1) ∈ R R × N 1 , . . . , 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:

E.2 Hierarchical TR µ MoE

In TR format, W ∈ R N 1 ×···× N E × I × O has E +2 factor tensors: U (1) ∈ R R 1 × N 1 × R 2 , . . . , 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 × R 1 , where R i are the manually chosen ranks. The weight tensor's elements are given by:

We derive the fast factorized forward pass in terms of a series of mode2 products:

F Experimental details

F.1 Network configurations and hyperparamters

Here 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 100 k iterations, with the same resources.

F.2 Weight initialization

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.

G Expert specialism: additional results

G.1 Large scale models

We first show in Figure 12 the top-activating examples for MLP-mixers trained with both CP µ MoE and TR µ MoE blocks. Examples are shown for the first two experts as they appear numerically for

0

50

100

150

200

250

300

Epochs

0

10

20

30

40

50

60

70

Validation set accuracy

70.31%

71.29%

71.26%

Large scale MLP mixer results:

ImageNET1k validation accuracy

MLP

CP MoE (ours)

TR MMoE (ours)

0

50

100

150

200

250

300

Epochs

2

3

4

5

6

7

Training loss

2.44

2.39

2.39

Large scale MLP mixer results:

ImageNET1k training loss

MLP

CP MoE (ours)

0

20000

40000

60000

80000

100000

Iterations

3

4

5

6

7

8

9

10

11

Training loss

2.862

2.880

2.868

Large scale GPT-2 (NanoGPT) results:

OpenWebText2 training loss

MLP

CP MoE (ours)

0

20000

40000

60000

80000

100000

Iterations

3

4

5

6

7

8

9

10

11

Validation loss

2.876

2.893

2.886

Large scale GPT-2 (NanoGPT) results:

OpenWebText2 validation loss

MLP

CP MoE (ours)

each of the 8 layers, where we observe the same phenomenon of earlier blocks specializing to textures, and later blocks to higher-level abstract concepts/objects.

Secondly, in Figure 13 we show the top 32 activating tokens for the first 6 experts (as they appear numerically) for layer 5 in GPT2 models trained with CP µ MoEs replacing every MLP block. Whilst there are clear coherent themes amongst the top-activating tokens, we do see some examples of multiple themes being processed with high coefficients by the same experts (e.g. example #20 in expert 2's top-activating examples appears unrelated to the context of the other top-activating tokens) indicating a certain degree of expert polysemanticity (as expected in the large open domain of web text).

Context

Patch

Context

Patch

Context

Patch

Context

Patch

Context

Patch

Context

Patch

Context

Patch

Context

Patch

Expert 1

Expert 2

Layer 1

Layer 2

Layer 3

Layer 4

Layer 5

Layer 6

Layer 7

Layer 8

Expert 1

Expert 2

Layer 1

Layer 2

Layer 3

Layer 4

Layer 5

Layer 6

Layer 7

Layer 8

Patch

Context

Patch

Context

Patch

Context

Patch

Context

Patch

Context

Patch

Context

Patch

Context

Patch

Context

Layer 5, Expert 1

Layer 5, Expert 3

Layer 5, Expert 5

Layer 5, Expert 2

Expert coe ffi cients color map:

Layer 5, Expert 4

Layer 5, Expert 6

Layer 5, Expert 7

Layer 5, Expert 9

Layer 5, Expert 11

Layer 5, Expert 8

Layer 5, Expert 10

Layer 5, Expert 12

Expert coe ffi

cients color map:

G.2 LLMsteering

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.

G.3 CLIP ViT-B-32

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.

H Ablation studies

H.1 Entmax vs softmax

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 climate crisis is

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 di

ffi

cult

to communicate what is called "the truth" with the public.        In fact, in some

The biggest issue of today's world is the internet, which not only protects the most

important aspects of our lives, but also has important bene

fi

ts that greatly bene

fi

t

other people. We can bypass the police, governments, banks and everyone else.

Nobody wants to be charged with a crime, and nobody

The 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

fi

nding out that food is

becoming harder to

fi

nd. 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 dwell

The biggest issue of today's world is the continuing need for mass-human

immigration. Of the 3 million new visas issued in

fi

scal 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>] , it

The 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

fi

rst 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 .        If we

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 o

ffi

cers, 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 got

The 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 puts

The biggest issue of today's world is the accumulation of resources that are being

moved, and the e

ff

ects of that. In 2016, Donald Trump, Texas police o

ffi

cers are

killed in the street.        'He has insulted the Mexican

fl

ag, the

fl

ag 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 e

ff

orts to keep us safe.        The War on Drugs has been an unjust and

unw

The 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 o

ffi

cer, 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 o

ffi

cer or a policeman who is

unarmed and is being continually killed by a person who owns a vehicle. They have

killed

fi

ve o

ffi

cers 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-pressure

The 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, here

The 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

warmer

The 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. If

The 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 a

The 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 100

The 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 a

The 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 Britain

The 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' di

ffi

culties 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 18

The 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 long-

term 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

fi

nd Baghdad's exiled neighbour and to capture or

occupy the city.

The biggest issue of today's world is a

fi

ght 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 di

ffi

culties 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

Raw model outputs

Sample 1

Sample 2

...

Sample 8

Intervening to increase:

"police violence" expert

Intervening to increase:

"weather"/"climate" expert

Intervening to increase:

"programming" expert

Intervening to increase:

"foreign politics" expert

Sample 1

Sample 2

...

Sample 8

Sample 1

Sample 2

...

Sample 8

Steering LLM outputs ↓

the last days were also a

μ

MoE: 256 total experts

μ

MoE: 32 total experts

Expert 1

Expert 3

Expert 5

Expert 6

Expert 7

Expert 8

Expert 9

Expert 10

Expert 1

Expert 2

Expert 2

Expert 4

Expert 5

Expert 6

Expert 7

Expert 8

Expert 834

Expert 1203

Expert 316

Expert 206

Expert 203

Expert 1602

0.99

0.76

coccerba.

nous

erep

shopping_.

hammermea:

53

0.41

ciopping.

shoppuig

0.34

0.29

Wred

Anerican_

aunffi

shopping

Fddou

ZET

0.22

0.21

coral_ree

singrey

sandbar

Anican

-021

oralree

stingray

0.18

platypus

LtO

andb

CPmuMoE-r512: 2048 total experts

32 total experts

64 total experts

128 total experts

256 total experts

512 total experts

1024 total experts

32 total experts

64 total experts

128 total experts

256 total experts

512 total experts

1024 total experts

32

64

128

256

512

1024

Total # experts used in CP

MoE layer (log scale)

0.8

0.9

1.0

1.1

1.2

1.3

1.4

Mean expert polysemanticity

p

(

n

)

Experts' class-level polysemanticity (

)

Entmax (dino)

Softmax (dino)

(a) DINO backbone

32

64

128

256

512

1024

Total # experts used in CP

MoE layer (log scale)

0.8

0.9

1.0

1.1

1.2

1.3

Mean expert polysemanticity

p

(

n

)

Experts' class-level polysemanticity (

)

Entmax (clip)

Softmax (clip)

(b) CLIP backbone

H.2 Fast forward pass computation speedups

We next report in Table 8 the actual number of FLOPs (as reported by https:// detectron2.readthedocs.io/en/latest/ _modules/fvcore/nn/flop_count.html ) when executing PyTorch µ MoE layers using the naive forward pass relative to the cost when using the fast einsum computation derived in Appendix B-the fast computation is many orders of magnitude less expensive (using one A100 GPU).

H.3 Batch normalization

We next perform an ablation study for the use of batch normalization (BN) before the activation function for the expert coefficients. We study CP µ MoEfinal layer layers with CLIP ViTB-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).

32

64

128

256

512

Total # experts used in CP

MoE layer (log scale)

0.8

0.9

1.0

1.1

1.2

1.3

Mean expert polysemanticity

p

(

n

)

Experts' class-level polysemanticity (

)

Batchnorm ablation

Final layer

MoE (CLIP) (Batchnorm)

Final layer

MoE (CLIP) (No batchnorm)

0

100

200

300

400

500

Expert index

0

2500

5000

7500

10000

12500

15000

17500

Number of images assigned with coefficient > 0.5

Expert load: CP

MoE, 512 total experts

(a)

512

total experts

0

100

200

300

400

500

600

700

Expert index

0

2000

4000

6000

8000

10000

Number of images assigned with coefficient > 0.5

Expert load: CP

MoE, 768 total experts

(b)

768

total experts

H.4 Expert load

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.

I Additional performance results

I.1 CLIP ViT-B-32 ImageNET1k ablations

Here, we compare the performance of parameter-matched µ MoEfinal 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 3 e -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.

1

2

4

8

16

32

64

128 256 512

Number of experts

78.0

78.1

78.2

78.3

78.4

78.5

78.6

78.7

ImageNET1k validation set accuracy

CLIP ViT-B-32 fine-tuning accuracy on ImageNET1k (linear vs  MoE)

Linear

CP MoE

TR MoE

(a) Accuracy comparison ( µ MoE vs Linear)

1

2

4

8

16

32

64

128 256 512

Number of experts

200

250

300

350

400

Resulting matrix rank

CLIP fine-tuning on ImageNET1k: Rank comparison

CP MoE

TR MoE

(b) Rank comparison (CP µ MoE vs TR µ MoE)

We attribute TR µ MoE's even greater performance gains over CP µ MoEs here 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) (see Figure 23b).

I.2 Hierarchical µ MoEs

Hierarchical µ MoE Mixers We train from scratch two hierarchical µ MoE MLP-mixer S-16 models for 300 epochs on ImageNET following the same configuration as in Section 4.3 of the main paper. Concretely, we use a two-level hierarchical µ MoE with N 1 = 64 experts for the first level and N 2 = 2 experts for the second layer ( 128 total effective experts). As shown through the results in Table 9, the hierarchical µ MoE's also perform well against the MLP alternatives, whilst providing even better parameter-efficiency.

Hierarchical µ MoE fine-tuning layers We also perform additional experiments with hierarchical µ MoEs used to fine-tune CLIP ViT-B-32 models on ImageNET1k. Here we use the experimental setup in [63, 64], training each model for a single epoch with the specified learning rate of 1 e -05 . We fine-tune hierarchical µ MoE CLIP models with up to 4 levels of hierarchy as shown in Table 10, where the best-performing models (averaged over 5 runs) are found with 2 levels of hierarchy.

I.3 Comparisons to dense/sparse MoEs

The goal of the µ MoE layer is to facilitate more interpretable subcomputations with a similar number of parameters and FLOPs to regular dense layers. Whilst the layer does not aim to improve on the capabilities of existing MoE layers, we nonetheless provide an initial comparison study here in Figure 24 for completeness. As can be seen, in addition to the scalable expert specialization provided,

10 6

10 7

10 8

10 9

Parameter count

69

70

71

72

73

MoE model val. accuracy on ImageNET1k

CP MoE

TR MoE

Linear

SoftMoE

10 6

10 7

FLOPS

69

70

71

72

73

MoE model val. accuracy on ImageNET1k

CP MoE

TR MoE

Linear

SoftMoE

the µ MoEs also perform very favorably against the alternative MoE models when fine-tuning CLIP on ImageNET1k.

J Fairness baselines & metric details

Here we present more details about the fairness comparisons and metrics used in the main paper.

Metrics

Equality of opportunity requires the true positive rates for the sensitive attribute subpopulations to be equal, defined in Hardt et al. [76] as P ( ˆ Y = 1 | A = 0 , Y = 1) = P ( ˆ Y = 1 | A = 1 , Y = 1) for sensitive attribute A , target attribute Y , and predictor ˆ Y . 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 maxmin y ∈Y ,a ∈A P ( ˆ Y = y | A = a, Y = y ) . We compute this as the minimum of the test-set accuracy for the 4 subpopulations in each experiment.

Baselines

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 ¯ a 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 results

K.1 Model re-writing

The 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.

Target attribute: "Blond_Hair". Target subpopulation: "Blond_Hair"+"Male"

Other

Blonde

CP MoE-r512-e128

Female

Male

Before

intervention

96

89

100

41

Other

Blonde

Blind

thresholding

Female

Male

After

intervention

86

98

97

84

Other

Blonde

Expert

thresholding

Female

Male

95

90

96

85

Other

Blonde

Blind

thresholding

Female

Male

Before-after

diff.

9.8

9.8

2.5

43

Other

Blonde

Expert

thresholding

Female

Male

0.51

1.1

3.2

44

Other

Blonde

TR MoE-r512-e128

Female

Male

96

88

100

43

Other

Blonde

Blind

thresholding

Female

Male

86

98

97

84

Other

Blonde

Expert

thresholding

Female

Male

95

90

96

87

Other

Blonde

Blind

thresholding

Female

Male

9.5

9.9

2.5

40

Other

Blonde

Expert

thresholding

Female

Male

0.76

1.6

3.8

44

(a) 'Young blond' intervention for Blond hair attribute prediction head

Old

Young

CP MoE-r512-e128

Female

Male

Before

intervention

55

97

75

94

Old

Young

Blind

thresholding

Female

Male

After

intervention

84

87

93

70

Old

Young

Expert

thresholding

Female

Male

87

83

77

93

Old

Young

Blind

thresholding

Female

Male

Before-after

diff.

29

10

17

24

Old

Young

Expert

thresholding

Female

Male

31

15

1.7

1.2

Old

Young

TR MoE-r512-e128

Female

Male

56

97

75

94

Old

Young

Blind

thresholding

Female

Male

85

87

93

70

Old

Young

Expert

thresholding

Female

Male

92

74

77

92

Old

Young

Blind

thresholding

Female

Male

29

10

17

24

Old

Young

Expert

thresholding

Female

Male

36

23

2

1.6

Target attribute: "Young". Target subpopulation: "Old"+"Female"

(b) 'Old female' intervention for age attribute prediction head

L NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: Claims regarding both qualitative and quantitative expert specialism for finetuning large foundation models are demonstrated in Section 4.1, where the benefits of scaling the expert counts are also substantiated both qualitatively and quantitatively. Claims regarding bias mitigation are substantiated in Section 4.2. Qualitative expert specialism is provided for large models (along with their performance) in Section 4.3.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: The limitations clearly state the lack of evaluation for out-of-domain data for vision, and the difficulties in further evaluating expert specialism quantitatively in large models (given the lack of ground-truth).

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [NA]

Justification: Technical derivations of models are made throughout (and further basic derivations of expert matrix rank), but no novel theoretical results are presented.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Full experiment settings/config/hyperparameters are provided in Table 7, and the supporting code ( https://github.com/james-oldfield/muMoE ) provides even more explicit experimental instructions. Learning curves are also plotted in Figures 10 and 11 for additional transparency. Pseudocode implementations are also given in Appendix B.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: Model code for µ MoEs and the experiments in the paper are found at: https: //github.com/james-oldfield/muMoE .

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: As found in Table 7, where we state we follow these choices based on the default parameters of the original papers introducing the models, or the default configurations used by the open-source maintainer for GPT2.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: We do include mean (and STD) of the results over multiple fine-tuning models, but we only have single runs over the large models due to resource constraints. For these single runs of large models, we always set all random seeds to 0 for reproducibility.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: Details are provided in Appendix F.

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: No ethical concerns to note.

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: The paper proposed a layer that provides more transparent, explainable, and editable networks. We discuss positive social impacts throughout the paper, but also acknowledge and discuss the potential negative impacts in Appendix A.

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: No models posing a high risk of misuse are to be released.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: Yes, the open-source codebases on which we base our code are explicitly referenced.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: None introduced.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: No human subjects involved.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\n1. **Innovative Metric:** The introduction of a quantitative metric for \"class-level polysemanticity\" is a novel approach to evaluating the specialization of experts within the model, providing a clear measure of the impact of increasing the number of experts.\n   \n2. **Comprehensive Evaluation:** The paper undertakes a broad range of experiments to validate the effectiveness of $\\mu$MoE layers. The inclusion of both qualitative and quantitative analyses strengthens the argument for the utility of µMoE layers.\n   \n3. **Competitive Performance with Added Benefits:** The ability of $\\mu$MoE layers to compete with larger models like GPT-2 and MLP-Mixer, coupled with the advantages of additional interpretability, is highlighted as a significant strength of the study.\n\nWEAKNESSES:\n1. **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.\n   \n2. **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.\n   \n3. **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.\n\nQUESTIONS:\n1. Could you specify the computational hardware and configuration utilized for these experiments?\n2. 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?\n\nLIMITATIONS:\nThe authors adequately discuss the limitations and do not have the potential negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- The proposed approach (using either the CP structure, or Tensor Ring structure, or both) achieves better accuracy than MLPs on image tasks, on two different architectures: MLP-Mixer S-16 and CLIP B-32, with models of similar parameter count.\n- The discussion about different factorization strategies for the MLPs in the appendix is a delight. In fact, all the subsections in the appendix contain many useful details. I enjoyed reading several of them.\n- The main paper contains a thorough and well written description of the proposed method, as well as the experiments. Any of the details and ablations that are missing from the main paper can be found in the appendix.\n\nWEAKNESSES:\n- 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.\n- 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.\n- 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.\n- 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.\n\n-----------\nDuring 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\".\n\nQUESTIONS:\nThis 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\".\n\nLIMITATIONS:\nNo negative societal impacts are specific to this work, in my opinion.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-19"}

comments: 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.

STRENGTHS:
* The problem studied in this paper is well-motivated, specialization of experts is one of the most important goals in MoE.
* The formulated illustration of factorization is clear.
* Figures and tables are clear and easy to read.

WEAKNESSES:
* It would be better if more results on language models could be included.

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:
Yes.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-07-23

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- The idea seems to be simple enough to be easily implemented in existing MoE setups. \n- Anecdotal and some statistical results suggests that it does help in  mitigating bias in certain cases in a constrained setting.\n\nWEAKNESSES:\n- 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.\n\n- Again, for Figure 3, there should be a way to compare it at least with sparse MoE. \n\n- 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. \n\n- 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.\n\nQUESTIONS:\n- As mentioned above, I think it is quite important to show how \\muMoE perform compared to SoftMoE and SparseMoE. \n\n- Making Figure 2 more statistically significant rather than anecdotes. \n\n- 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. \n\nSuggestions:\n- I think the notation used in Eqn 1 is a bit difficult to follow, specifically with the colons used without brackets.\n\nLIMITATIONS:\nThe authors have identified the limitations of the work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. Introducing the Multilinear Mixture of Experts (µMoE) layer is novel and addresses the computational cost challenges associated with scaling the number of experts for fine-grained specialization. The mathematical formulation of µMoE layers is thorough, and the factorization approach is well-explained and technically sound.\n2. Practicality: The proposed µMoE models demonstrate a more parameter-efficient alternative to traditional MoEs, an important consideration for practical applications.\n3. The paper includes extensive qualitative and quantitative evidence to support the claims. Experiments on vision and language models illustrate the versatility and effectiveness of the proposed method.\n4. This paper is well-written and easy to follow.\n\nWEAKNESSES:\n1. 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.\n2. 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.\n3. 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.\n2. 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.\n4. While the paper qualitatively demonstrates interpretability, it lacks objective, quantitative measures of the interpretability gains compared to other methods.\n8. 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.\n\nMore Comments:\n1. 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.\n2. 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.\n\nQUESTIONS:\nN/A\n\nLIMITATIONS:\nYes, the authors addressed limitations and the direction for future research.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-29"}