Multilinear Mixture of Experts: Scalable Expert Specialization through Factorization
Authors: James Oldfield
Markos Georgopoulos
Grigorios Chrysos
Christos Tzelepis
Yannis Panagakis
Mihalis Nicolaou
Jiankang Deng
Ioannis Patras
Venue: NeurIPS 2024 poster (2024)
Keywords: interpretability
mixture of experts
Primary area: interpretability_and_explainability
TL;DR: muMoE layers perform factorized (fully-differentiable) computation that efficiently scales MoE's expert count, leading to increasingly fine-grained expert specialization

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
James Oldfield"* Markos Georgopoulos Grigorios G. Chrysos
Christos Tzelepis
Yannis Panagakist.0 Mihalis A. Nicolaou
Jiankang Deng'
Toannis Patras'
'Queen Mary University of London University of Wisconsin-Madison CCity University of
London Natiooall and Kapodistrian University of Athens AArchmeedss AI, Athena RC TThe
Cyprus Institute
'Imperial College London
Abstract
The Mixture of Experts (MoE) paradigm provides powerful way to decompose dense layers into smaller, modular computations often more amenable to human interpretation, debugging, and editability. However, 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 (JMMoE) layer to address this, focusing on vision models. fMME layers enable scal- able expert specialization by performing large weight tensors entirely in factorized form. Consequently. jMoEs (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 special- ized experts 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 yMoE blocks at every layer, maintaining comparable accuracy. Our code is available at: https: / //gi thub com/ james- oldf eldd/uMot.
implicit computation on prohibitively
Introduction
The Mixture of Experts (MoE) architecture [1] has reemerged as 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 as 'experts") ) for each input, in contrast to the traditional approach of using the same single layer for inputs. This pro- vides form of input-conditional computation [8, 9, 10, 11] that is expressive yet efficient. However, through their substantial performance gains, an important emergent property MoEs is frequently underutilized: the innate tendency of experts to specialize in distinct subtasks. Indeed, the founda- tional work of Jacobs et al. [12] on MoEs describes this property, highlighting how implementing particular function with modular building blocks (experts) often leads to subcomputations that are easier to understand individually than their dense layer counterparts 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 com- ponents, each contributing individual functionality. Yet, the value of interpretable computation
"Corresponding author: j. j.a. oldd aldfield@qmul.ee ac. uc.uk
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

extends beyond just transparency [13] and explainability [14]. An important corollary of suc- cessful task decomposition amongst experts that layers are easier to debug and edit. Biased or unsafe behaviors can be better localized to specific experts" subcomputation, facilitating man- ual 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 of- ten 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 spe- cialize to a more specific set of semantically similar inputs. Alas, the dominating "sparse" MoE paradigm of selecting only the top-K experts [17] is not only parameter-ineffcieet for large expert counts, but also has several well-known issues due to its discrete expert routing- -oeeennooteennn 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 designed to scale gracefully to dense operations in- volving tens of thousands of experts at once through implicit computations on factorized form of the ex- perts' weights. Furthermore, in contrast to the dominant sparse MoEs' [17] non-differentiable nature, jMoEs are differentiable by design, and thus do not inherit the associated training issues. We summarize the benefits of pMoEs' model form over existing MoEs in Table 1. Crucially, we show evidence that scaling up the number of μMoß 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 μMоε expert specialism facilitates the practi- cal task of model editing-hore how subcomputation in specific combinations of experts biased towards demographic subpopulations can be manually corrected through straightforward guided edits.
Table 1: Benefits of the proposed pMoEs' model form over existing MoEs.
Parameter-
FLOPs-
Differentiable
efficient
efficient
Dense MoE [1] Spurse MoEl [17]
AMoE (ours)
Building on these findings, we demonstrate that μMoEs offer compelling alternative to MLPs for pre-training both vision and language models with up to 100M parameters-enabling large numbers of specialized experts while maintaining comparable performance and parameter counts the original networks' single dense MLPs.
Our contributions and core claims can be summarized as follows:
We introduce μMoË layers-a mechanism for computing vast numbers of subcomputations and efficiently fusing them conditionally on the input.
We show both qualitatively (through visualization) and quantitatively (through counterfactual intervention) that increasing the number of μMMo experts increases task modulariiy-leann 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-mixe) μMМE networks, establishing experimentally that models with parameter-man matched μMMσ 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 appu--conitionall 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
2

---

seminal work of Shazeer et al. [17] however proposes to take a sparse combination of only the top- most relevant experts, greatly reducing the computational costs of evaluating them all. More recent works employ a similar sparse gating function to apply just subset of experts [2. 25], scaling to billions [3] and trillions of parameters [4]. The discrete expert selection choice of sparse MoEs 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, 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 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/hackrroud [30] and image categories (e.g. 'wheeled vehicles") [24] in vision. Evidence of shared vision- vision-language specialism is even found in the multi-modal MoEs of Mustafa et [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-Trai 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 al. [34] share parameters across expert matrices, efficient implicit computation of thousands of experts simultaneously is not facilitated, in contrast to the μMoF. layer.
Factorized layers in the context of deep neural networks provide several important benefits. Re- placing 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 tanlimmask//ommin learning [45, 46] through the sharing of parameters/low-rn -aak factors across tasks. Furthermore, parameter efficiency through weight factorization often facilitates the design and efficient implementation of novel architectures such polynomial networks [47, 48, 49] or tensor contraction layers [50]. The recent DFC layer in Babiloni et al. [51] also performs dynamic computa- tion using the CP decomposition [52] like yMoEs. 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 specialized subcomputations in way that facilitates the interpretability and editability presented in this paper. DFCs can be seen to apply an MLP input vectors 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 unique resource-efficient models and forward passes in Section 3.1.1. Finally, we show in Section 3.1.2 how jMoEs recover linear MoBs as a special case.
Notation We denote scalars x E with lower-case letters, and vectors x E R I and matrices X E R x12 in lower- and upper-case boldface latin letters respectively. Tensors X' E IR h ggx..xx...Il X of order d are denoted with calligraphic letters. We refer to the (11, (2, ith)-th element of this tensor with both X(i,, .2, id) E R and Fiiig...aa E R. Finally, we use a colon to index into all elements along a particular mode: given X' E Rlixlills for example, X E RR/x11 or equivalently X((;;;,,) i2) 1 is the matrix at index ig of the final mode of the tensor. We use X X, u to denote the mode-n (vector) product [53] of tensor X E ER/1x /2/xx2xx...2 IN and vector E R whose resulting elements are given by (X' Xx u) u)jj...n-
-in = E/n
Dinmt Tirig...n Us
3.1 The 5MoE layer
MMoEs provide a scalable way to execute and fuse large numbers of operations on an in- put vector by formalizing conditional computation through resource-efficient multilinear oper-
3

---

ations. A μMoE layer comprised of N many experts (and single level of expert hierar- chy) is parameterized by weight tensor W RNx1x0 and expert gating parameter RIXN. Given an input vector 2 E R (denoting the hidden representation of an individual to- ken, for example), its forward pass can be expressed through the series of tensor contractions:
(G 2) ERRY W X1 X2 N
(CT22)
Weight tensor W
Expert confficients
X De PEEmE
X1a
Wn E RO
(1)
(estpst dis
where a is the vector of expert coefficients and is the entmax activation [54, 55]. The μMoE layer can be understood as taking sparse, con- vex combination of N many affine transforma- tions? of input vector z, weighted by the efficients in a. The first tensor contraction in the forward pass (2; W. (xx) matrix- multiplies the input vector with every expert's weight matrix. The following tensor contraction with expert coefficients a takes a linear combi- nation of the results, yielding the output vector. The forward pass can be visualized intuitively as multiplying and summing over the modes 3D tensor, which we illustrate in Figure I. Furthermore, μMoEs readily generalize to hierarchical conditional computations by introducing additional modes the weight tensor and corresponding vectors of expert coefficients (see Appendix E).
7 Lirgut dis
Output vector
Input vector X2Z
Figure 1: μМОЕ layer series of tensor contractions: the experts' weight matrices (yellow 2D slices) matrix-multiplies with the input vector and summed (weighted by the red expert coefficients).
The forward pass of an (unfactorized)
3.1.1
Computation in factorized form
Our key insight is that the dense μMME forward pass over all N experts simultancously can be com- puted entirely in factorized form, needing never materialize prohibitively large weight tensors. This allows μMoEs' computations to scale gracefully to many thousands of experts simultane- ously, without the problematic top-K gating [17]. To achieve this, (1) first parameterize the experts' weights W E R TxO with a tensor factorization and (2) re-derive fast forward passes of Equation (1) operate solely in factorized form.
In the context of a MMME layer, the various choices of tensor factorizations make different trade-offs regarding parameter/FL.00 FOO 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μMоE Imposing CP structure [52, 56] of rank R the weight tensor, (1)
write W up ur ur E NxIxO as a sum of It outer products, with factor matrices U(1) This reduces the parameter count from NTO (such as with sparse/dense MoEs and regular jMMoss just R(N 7 + O). Crucially,
ΣR R U E RRRI R
U33) E RRO
further rewrite the
CPμMoF. layer's forward pass entirely factorized form without ever materializing the full tensor (plugging the CP-composed tensor into Equation (I)) as:
N y ΣΣ Σ
R
(1) u o u
(2)
u
(3)
2g0n =
Σ (U (U (3) E R
(2)
with Equation (2) being analogous to the fast computation in Babiloni et al. [51], only here the operations combining the weights and producing the outputs can be expressed single step. Whilst the original naive CPμMoE forward pass has a FLOP count" of NIO, the fast computation
IIncrementing the dimension of the second 'input" mode of the weight tensor W (X--1)) and
appending a 1 the input vector z E 4 folds a per-expert bias term into the computation.
We adopt the convention of counting fused multiply-adds as one operation [57]. Note that the small additional expert coefficients cost is constant across models and thus ignored in comparisons.

---

above has just R(N + + 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 pMoEs.
TRuMoE We propose second μMoo variant based on the Tensor Ring [58] (TR) factorization that can offer even better efficiency for large values of N. In TR format, W E R NxXx has three factor tensors: 2()) E (RRxNNX, R2, (((2) RR2xXxx Rs, 2 E RRxOXx, where R, are the manually
chosen ranks". The weight tensor's elements in TR format are given by: Wnio tr( U ( ! [58]. TRaMoE'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:
R2 Σ
(2 X2 () X2 z)
(3) mirg Tyan
E ERO,
(3)
[Rg Ral
yielding modified FLOP count of N R2TR3 Rí RR + R, with just (R, N R2 R + R2ORt, parameters. With large N contributing to the computational cost only through Rí N. R2, the TRuMoE can prove even more resource-efficient than CPyMoES by choosing small values of R, Rg. We refer readers to Appendix D for a further discussion of decomposition choice, derivations of how tensor rank translates expert matrix rank, and FLOPs comparisons.
3.1.2 MoMEs recover dense MoEs as a special case
Finally, we note how unfactorized μMoE layers with 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 yo = a W W... This highlights an interesting technical connection between neural network layers: dense MoE layers in this tensor
formulation can be seen to share a similar functional form to bilinear layers, which have also found applications in interpretability [59, 60].
Experiments
We start in Section 4.1 by presenting both qualitative and quantitative experiments validating that the experts learn to specialize in processing different semantic clusters of the input data. In Section 4.2 we demonstrate one practical benefit of the learned specialism-showe 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 for experiments with hierarchical j.Mos.
how expert-conditional re-writing
Implementation details Before applying the activation function to the expert coefficients we apply batch- and layer-normalization to YMоσ layers in vision and language models respectively (see Appendix H.3 for an ablation). Interestingly, we do not find the need for any load- bala-cinniing losses. We fix the TRaMoE ranks be R Ra throughout (see Appendix D.1.2).
4.1 Expert specialism: visualization & intervention
Our first objective is to show that scaling M.Mo''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 μMoË 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 ImageNETIk (following the fine-tuning protocol in Ilharco et al. [63, 64]). Whilst fine-tuning large foundation models is an important application of μMoΕ layers in its own right (c.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.
'Setting R 1 recovers Tensor Train [35] μMoE.
5

---

Eanee E
1ppin'
IMoE: 256 total experts
:MOE: 32 total experts
Figure 2: Specialization in 256 vs 32 total expert CPμMME layers (fine-tuned on CLIP ViT-B-32) Each row displays randomly selected images processed (with coefficient 2 0.5) by the first few experts for the two models. The more scale the expert count, the greater the apparent expert specialism (to single visual themes or image categories).
4.1.1 Qualitative results
We first show random examples in Figure 2 of images processed (with expert coefficient 2 0.5) by the experts by each CPμMMo layer (the class labels and expert coefficients overlaid in white and green text respectively). Using only a modest number of experts (c.g. 32) appears to lead to some "polysemanticity" [65] in experts-with some processing unrelated classes of images (c.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 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 pMoEs. 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- undaanndddrrtt patterns of input features 169, 70J. 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 spe- cialism μМоB expert exhibits the ImageNETik dataset as an (imperfect) proxy [71].
Following the causal intervention protocol of Elazar et al. [20], ask the specific counter- factual question about solely each expert n in a EMoE layer in turn: "had expert 's weight ma- trix W. not contributed its computation, would the network's 's test-sel accuracy for class have dropped?" Practically speaking, given a net- work fine-tuned with μMMF layer, we achieve this by intervening in the forward pass by zero- ing the th expert's weight matrix W,, := O, leaving every other aspect of the forward pass completely untouched. Let the elements of y, y(n) E C denote the test set accuracy for the C = 1000 ImageNETIk classes, pre- and cost-intervention of expert n respectively. We collect the normalized difference to per-class ac- curacy in the vector (m), whose elements are
Experts' class-level polysemanticity
Final layer MMME (clip)
2.00
Panultimate layor (MooE-ciipp)
1.75
Final layer μMoE (dino)
Penultimate layer pMoE (dino)
1.50
Mein extert
1.00
0.75
32
256
512
Total W experts used in CPyMMoE layer (log scale!
Figure 3: Higher expert counts lead to more monosemantic experts: mean expert class-level polysemanticity of Equation (4) (L) as a function of the total number of experts. Results are shown for both CLIP ViT-B-32 and DINO models fine-tuned on ImageNETIk with CPyMoE layers.
given by d(m) (ye- g At the two extremes, when the full network's 'ss accuracy for class drops completely from Ye to O upon manually excluding expert n's computation get
d( = 1, whilst d!m) = means the absence of the subcomputation did not change class e's test set accuracy all. We thus estimate the 'elass-level polyscmanticity" of expert n as the distance between

---

Table 2: Fairness metrics for baseline models and after applying standard fairness techniques, for the two experiments on CelebA. A - e128 model is used as the final layer.
Eai Blus uncarth "Old feveddee fir 'Ac'' prediction bead
(2) Bias lowaeds 'Blasi msllo for 'Blosd Harr prodiction hoad
Tages
Equalix of ope.
Salpore
Jos. 521
Taggs
Equaatt of
5TD
Salpop
Ics.20 pc:
Hirtdari lncar
0.113
1.24
0.136
3.513
15.33 333
0.11%
11:30g
(1340
0.311
05:11
427%
+ aahe dihiss [2S]
0-414
1.224 01977
0.776 0J866
01424
17:165
0.01
01633 01081
03374
0.191
01.011
179K
esper theech (eurs)
8,866
0.786
$1.000
0.47
0J98
0.866
98.995
175K
the difference vector and the one-hot vector:
p(n) = I|d 160112
(4)
where index argmax, (d.")) of (mn) has a value of 1 (and values of 0 everywhere else). This encodes the signature of perfectly class-level monosemantic expert, for which all accuracy for single class alone is lost in the counterfactual scenario in which the expert n did not contribute. We plot in Figure the average expert polysemanticity p(*) for all experts with non-zero difference vectors", observing a steady drop in its value 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 final classification layers and as penultimate layers (followed by a ReLU activation and linear classification layer), and for multiple pre- pre-trained foundation models. We further refer readers to the bar plots of the values of (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 simultancously 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- 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 MMoE final layer models for the 40 binary attributes in CelebA, jointly optimizing 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 set collecting the expert coefficients a E R from forward passes of the training images belonging the target subpopulation. We evaluate the subpopulation's 's mean expert coefficients á = 1/|C1 aec a E R.Y, proposing to manually re-write the output of this expert combination. We modify the layer's forward pass for the tt output head for attribute of interest (e.g. 'blond hair') as:
Vo =
2 + T a.
(5)
Here, the term Aā E R specifies, for each expert, how much to increase/decrease the logits for attribute with À being a scaling hyperparameter". Taking the dot product with an input image's expert coefficients a applies the relevant experts' correction terms (in the same way selects a subset of the most relevant experts' weight matrices). We report range of standard faimess metrics for both the model rewriting and networks trained with existing techniques (that aim to mitigate demographic
'Le. we include only experts that, when ablated in isolation, alter the class accuracy; please see the Appendix for discussion on expert load.
'We set À := N for all experiments for simplicity, but we note that its value could require tuning in different experimental setups. The sign of À is chosen to correct the bias in the target direction (whether to move the logits positivelylnegatively towards CelebA's e.g. young/old binary age labels respectively).
7

---

Layer 2, Expert 1
Context Patch
Layer 2, Expert 2
Context Patch
Layer 2, Expert 3
Context Patch
Layer 7, Expert
Context Patch
Layer 7, Expert
Context Patch
Layer 7, Expert
Context Patch
Figure Top-activating patches (top rows) and their full images (second rows) for the first experts across 2 CP CPpMoE-es e64 layers in μMMoE MLP-mixe [80] MMo blocks exhibit coarse- coarse-grained specialism (e.g. texture) earlier and more fine-grained specialism (e.g. objects) deeper in the network.
bias without requiring images' sensitive attribute value at test time). These are shown in Table for the two different experiments CelebA. where the proposed intervention outperforms bascline alternative methods in the majority of settings. Please see Appendix for details about the baseline methods and faimess metrics used, and further discussion of results.
4.3 Large language/vision μMMoE networks
Finally, we train from scratch 12 layer 124M- parammeeeeeter GPT-2 [81] LLMs on Open WebTex [82] for the language domain and 8 layer S- 1-16 variant" MLP-Mixers [80] on ImageNETIk [83] for vision. We replace every MLP block's 's linear layers with 2 μMoE layers. Each token t's input vector E R is therefore transformed with μMoE blocks of the form:
y:
( bh
GELU
ΣΣ
(1) Pridln
a; === @ 2z),
m==1==1
where a, E N the expert coefficients for each specific token and block, H is the dimension of the block's 's hidden layer, and ((1) E NRXII m2(2 E N are the (implicit) μMoE weight
'The S-16 model is the largest configuration that fits into 4x80GB A100 GPUs using the original paper's batch size of 4096.

---

Repert confficients nolor map:
81
14
0.
Layer 5, Expert
Layer 5, Expert 37
#s: Hortiea 1800 tys The the rebordd BAY the
ti lock pomoilla Pervend rudrtty Oe
19
chocker "Mutiecar
14. tull be doeir the Ell ieguaeee such an cara wiihh somvet usdd
b Ietum CADEl 0.55 1000 1 Indlanctonn3 C:3
pain my comasion
reshee thet
berptt
begs
partur cborget tiee teach
(HB5 Al, casesc
allow
P3X: Parini 1855
#13: frn itter
EP ent cealton povermmet poing dediicn
#33: te (ir) 1-
exst
accuudt
012:
centre hod priiinnn stgrtturt
have pobdecss
211: p1r ios Iely
usd Hnl andera compared
#15:
racic sotar aouuut
P33: Xill auurags par
brcher as08 take compleddy difirent
037:
418
aweraned funres weer hoand sking taos
018: 200r whht [P (Pouia suum TM0
P21:
und tex disuas that had boes die drccla colondaain
P21: enomaass naricll calk, The dffermnce aggitl cart Higgu
P22 K51 dccchar
#31
aaee saner waster kncreye less
234:
Coaventent
P23: Liet and
crargiaa spplectte
223: pariive wnlp 64%
hore
022:
panglot tas buuft Aepwyy pass they
#28:1 7 ES5 H HUl
233.
PXL: Bush daci 2000 mar vartary atsevittio
2010
PX1: 2) Ne (ar)
(ee] (ve lrigh #32: 111 www.eindeeettiiieecooo Image magac
#32: warking croepe Sdaodll paing
Layer 6, Expert
Layer 6, Expert 16
#1: Eer our is con
grotal book incoeneendd RMM
iini-oto koays diamte taddcte vo the mud adtunitingg ringg soout
#2: yours,
fhattitthe los Does diletta potcchiit
and
Saan inndenen vnttase terettolteee
ontir
iwow ploal Tedate Then Lee cianne etteeceti disemmtte aiteett
G.L 12
Dea
USA TOBAT
verdeed
#8: Trey Bie mont highty riparied and scoem
the wert with
P8: 1a plon tiigg Lifiioloo Lay dietts ei seenoeouti-eee many Amariias are buyy with High
atficr thene
#11 eety
inpertart dklle waw artteer chaat scr.
ecoweate sodkin speodd text
#12: pactators. this the crouced anctomat wartingg ng lathonn world.
#12: malistaa Chy crea via votame get nuthil enddimeet--1e the
sp lumoo
mo
Follon
The Dally Calter Nowo
soodis I9 cameat
#15: Polon dhett der (olaaen)
The
P18: ruteett ofser, volaan pange luu part of the your.
bin youc or dimmms Mschacd ccmpptt ticolass ppecettedd--
#13 topelight damig eoemaal craclal reamants Cancars ahore
for
222 Bhunch Mit fox
222 KDB
cowter excane calor towws
#25 /ini one prmpaneaa the better
do boki
#33: led tor the sother mumsse lemthrrs.
424
038
aarlortrrtee US and medile (an
P2R Vieil thea vas als souineed saciar ro China 2015 Karce Ben-Caronn
228
aree right #32 upe on penats do mo sne hugehunn hunt lap ohehe der, with shho
Hell
#33:V
#32: 11h low Hortegttn
Figure 5: Top-activating generated tokens for 4 manually selected experts for GPT-2 trained with CPμMoE blocks at every layer (each token is highlighted by the coefficient of the expert in question), exhibiting specializations to concepts including compound adjectives and equality operators.
tensors for each of the two layers. We manually set the sMoE 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 μMME block performs computations with N experts yet has the same parameter counts and FLOPs as a single, dense MLP block.
MoME-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 experts across all 8 blocks where such scale-specific specialism is apparent across the entire network.
J.MoE-PP22 For LLMs, we see promising qualitative evidence of experts specializing throughout a corpus of 1 1M generated 100-token sequences. At layer 5, for example, the generated tokens that use expert 8 with the highest coefficient are compound adjectives (Figure 5), whilst expert 37 most highly activates for equality and comparison operators in code and scientific text (please examples of

---

Table Comparison of μMoEs and dense MLPs across different models and tasks. We use 64
μMoE experts for the two vision tasks and N 256 for GPT2. MLP mixers and GPT2s are pre-trained for 300 epochs and 100k iterations respectively, whilst CLIP fine-tuned for epochs.
MLP-mixxer S- 16 (ImageNETIk)
GPT-2 NanoGPT (OWT)
CLIP B- (ImageNETIk)
Val. acc. (1)
#params
Val. loss (J)
#params
Val. acc. (1)
#params
MLPs
70.31
18.5M
2.876
124M
77.99
769K
TRIMoES
71.26
18.3M
2.886
124M
78.71
771K
CPμMoEs
71.29
18.6M
2.893
124M
78.07
769K
many unfiltered experts in Figures 13 and 14). Whilst monosemanticity is not always attained, μMoE layers nonetheless facilitate level of specialism not facilitated by dense MLP layers.
One important result here is that μMoE networks in this setup are significantly more parameter- efficient than both dense and sparse MoEs with the same expert count, as shown Table 4. For example, GPT-2 models with 256 sparse/dense MoE experts require a prohibitive 14.5B MLP parameters alone, relative to just 57M MLP parameters with μMoEs of the same expert counts.
μMоE performance Finally, 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 Ta- ble 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 ImageNETIk datasets. Following Section 1. we replace all linear
Table 4: MLP parameters required for net- works with the same expert counts.
NanoGPT (xpt2)
MLP-Mixer (8-16)
Model
256
Dense/Sparse MoE
14.5B
1,13B
CP,Mo TRyMMoF
57.0M
17.7M
57.4M
17.4M
layers with μMoE blocks (and single μМоΕ final layer for fine-tuning CLIP). We initialize all linear layers following the default Py PyTorch U[-k, kj] 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 pMoEs. Crucially, whilst μMoE layers provide additional interpretability benefits through scalable expert specialization, they do not sacrifice accuracy when parameter-matee matched to MLP blocks, as seen from the comparable performance.
5
Conclusion
In this paper, we introduced the Multilinear Mixture of Experts layer (jMME). We demonstrated that larger expert counts lead to increased specialization, and how jMME layers make this computationally tractable through factorized forward passes. μMMEs scale large expert counts much more gracefully than existing MoEs, yet avoid the issues from popular gating mechanisms. As further practical example of μMME'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 faimnes/interprttabliity for accuracy.
Limitations Firstly, it is important to state again that our quantitative evaluation only captures data [70, 84]. Furthermore, expert specialism expert behavior on the test set, not qualitatively (through the expert coefficients) due to the absence in large models only demonstrated
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 EMoE 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 jMMos in even larger-scale LLMs.
10

---

References
[1] Robert A Jacobs, Michael Jordan, Steven J Nowlan, and Geoffrey E Hinton. Adaptive mixtures of local experts, Neural computation. 3(1):79-87, 1991.
[2] 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.
[3] Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. GShard: Scaling giant models with condi- tional computation and automatic sharding. In Int. Conf. Learn. Represent. (ICLR), 2021.
[4] 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.
[5] Trevor Gale, Deepak Narayanan, Cliff Young, and Matei Zaharia. Megablocks: Efficient sparse training with
Proceedings of Machine Learning and Systems, 5, 2023.
[6] 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.
[7] 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.
[8] David Ha, Andrew M. Dai, and Quoc V. Le. Hypernetworks, In Int. Conf. Learn. Represent. (ICLR), 2017.
[9] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Adv. Neural Inform. Process. Syst. (NeurIPS), 30, 2017.
[10] Yizeng Han, Gao Huang, Shiji Song, Le Yang, Honghui Wang, and Yulin Wang. Dynamic neural networks: A survey. IEEE Trans. Puttern Anal. Mach. Intell. (TPAMI), 44(11):443-7455 2021.
[11] Yinpeng Chen, Xiyang Dai, Mengehen Liu, Dongdong Chen, Lu Yuan, and Zicheng Liu. Dynamic convolution: Attention over convolution kernels. In IEEE Conf. Comput. Vis. Pattern Recog. (CVPR), pages 11000-10039 2020.
[12] Robert A Jacobs, Michael I Jordan, and Andrew G Barto. Task decomposition through com- petition in a modular connectionist architecture: The what and where vision tasks. Cognitive science, 15(2):219-25)
1991.
[13] Zachary C. Lipton. The mythos of model interpretabillity. Communications of the ACM, 61(10): 36-43, September 2018. ISSN 1557-7317.
[14] Marco Tulio Ribeiro, Sameer Singh, and Carlos Guestrin. " why should 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.
[15] 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.
[16] Valeriaa Cherepanova, Vedant Nanda, Micah Goldblum, John Dickerson, and Tom Goldstein. Technical challenges for training fair neural networks. arXiv preprini
2021.
[17] Noam Shazeer, AAzalia Mirhoseini, *Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of- experts layer. In Int. Conf. Learn. Represent. (ICLR), 2017.
11

---

[18]
Mugeeth Mohammed, Haokun Liu, and Colin Raffel. Models with conditional computation learn suboptimal solutions. In Can't Believe It's Not Better Workshop: Understanding Deep Learning Through Empirical Falsification, 2022.
[19]
Joan Puigcerver, Carlos Riquelme, Basil Mustafa, and Neil Houlsby. From sparse to soft mixtures of experts. In Int. Conf. Learn. Represent. (ICLR), 2024.
[20] Yanai Elazar, Shauli Ravfogel, Alon Jacovi, and Yoav Goldberg. Amnesic probing: Behavioral explanation with amnesic counterfactuals. Transactions of the Association for Computational Linguistics, 9: 9:160-175, 2021.
[21] Emmanuel Bengio, Puerre-Luc Bacon, Joelle Pineau, and Doina Precup. Conditional com- putation in neural networks for faster models. In Int. Conf. Mach. Learn. Worksh. (ICMLW), 2015.
[22] M.I. Jordan and R.A. Jacobs. Hierarchical mixtures of experts and the em algorithm. In Proceedings of 1993 International Conference on Neural Networks JJCNN-93-NNgoyy Japan), volume 2. pages 1339-1344 vol.2, 1993. doi: 10. (119///CNNN11933.1677777777
[23] David Eigen, Mare' Aurelio Ranzato, and Ilya Sutskever. Learning factored representations in a deep mixture of experts. In Int. Conf. Mach. Learn. Worksh. (ICMLW), volume abs/1312.4311, 2013.
[24] 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.
[25] 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 Inl. Conf. Mach. Learn, (ICML), pages 5547-5569,
PMLR, 2022.
[26] 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.07688 2022.
[27] 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.
[28] Aya Abdelsalam Ismail, Serean O Arik, Jinsung Yoon, Ankur Taly, Soheil Feizi, and Tomas Pfister. Interpretable mixture of experts. Transactions Machine Learning Research, 2023. ISSN 2835-8856.
[29] 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.
[30] Lemeng Wu, Mengchen Liu, Yinpeng Chen, Dongdong Chen, Xiyang Dai, and Lu Yuan. Residual mixture of experts, 2022.
[31] 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.
[32] Andrew Davis and Itamar Arel. Low-rank approximations for conditional feedforward compu- tation in deep neural networks. arXiv preprint arXiv:
2013.
[33] Yunsheng Li. Yinpeng Chen, Xiyang Dai, mengehen 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.
[34] Ze-Feng Gao, Peiyu Liu, Wayne Xin Zhao, Zhong-Yi 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.
[35] I. Oseledets. Tensor-train decomposition. SIAM J. Sei. Comput., 33:2295-2317, 2011.
[36] 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: 216.05.066 2021.
12

---

[37] Alexander Novikov, Dmitrii Podoprikhin, Anton Osokin, and Dmitry P Vetrov. Tensorizing neural networks. Adv. Neural Inform. Process. Syst. (NeurIPS), 28, 2015.
[38] Timur Garipov, Dmitry Podoprikhin, Alexander Novikov, and Dmitry Vetrov. Ultimate ten- sorization: compressing convolutional and fc layers alike. arXiv preprint arXiv: 1.1122214 2016.
[39] Alexander Novikov, Mikhail Trofimov, and Ivan Oseledets. Exponential machines. In Int. Conf.
Learn. Represent. Worksh., 2017.
[40] Markos Georgopoulos, James Oldfield, Mihalis A Nicolaou, Yannis Panagakis, and Maja Pantic. Mitigating demographic bias in facial datasets with style-based Comput. Vis. (IJCV), 129(7))2288-225) -20077 2021.
transfer. Int. J.
[41] Francesca Babiloni, loannis Marras, Gregory Slabaugh, and Stefanos Zafeiriou. Tesa: Tensor element self-attention via matricization, In IEEE Conf. Comput. Vis. Pattern Recog. (CVPR), pages 13995-13954, 2020.
[42] Markos Georgopoulos, Grigorios Chrysos, Maja Pantic, and Yannis Panagakis. Multilinear latent conditioning for generating unseen attribute combinations. In Int. Conf. Mach. Learn. (ICML), 2020.
[43] Yixin Cheng, Grigorios G. Chrysos, Markos Georgopoulos, and Volkan Cevher. Multilinear operator networks, 2024.
[44] 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.
[45] Adrian Bulat, Jean Kossaifi, Georgios Tzimiropoulos, and Maja Pantic. Incremental multi- domain learning with network latent tensor factorization. In Conf. on Artifi. Intel. (AAAI), volume 34, pages 10470-10477) 2020.
[46] Yongxin Yang and Timothy M. Hospedales. Deep multi-task representation learning: A tensor factorisation approach. In Int. Conf. Learn. Represent. (ICLR), 2017.
[47] 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.
[48] Grigorios G. Chrysos, Stylianos Moschoglou, Giorgos Bouritsas, Jiankang Deng, Yannis Panagakis, and Stefanos Zafeiriou. Deep polynomial neural networks. IEEE Trans. Pattern Anal. Mach. Intell. (TPAMI), page 1-1, 2021. ISSN 1939-3539.
[49] 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-10522 2021.
[50] 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.
[51] Francesca Babiloni, Thomas Tanay, Jiankang Deng, Matteo Maggioni, and Stefanos Zafeiriou. Factorized dynamic fully-conneeted layers for neural networks. In Int. Conf. Comput. Vis. Worksh. (ICCVW). pages 1374-1388 October 2023.
[52] Frank Lauren Hitchcock. The expression of a tensor or a polyadic as a sum of products. Journal of Mathematics and Physics, 6: 1644188, 1927.
[53] Tamara G. Kolda and Brett W. Bader. Tensor decompositions and applications. SIAM Review,
5 51(3):555-500 2009. doi: 10.1137/070701111
[54] Ben Peters, Vlad Niculae, and André F. T. Martins. Sparse sequence-to-sequemme 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 15041519, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.18653/v1/P19-11
[55] 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
13

---

Kong, China, November 2019. Association for Computational Linguistics. doi: 10.18653/v1/ 1223.
[56] 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.
[57] fvcore: Flop counter for pytorch models. https:// http://www.....cccccooeeeeeeeaaa fvcore. Accessed: 2024-05-16.
[58] Qibin Zhao, Guoxu Zhou, Shengli Xie, Liqing Zhang, and Andrzej Cichocki. Tensor ring decomposition. ArXiv, abs/1606.05535, 2016.
[59] Lee Sharkey. A technical note bilinear layers for interpretability. arXiv preprini arXiv:2305.0 03452, 2023.
[60] Michael T. Pearce, Thomas Dooms, and Alice Rigg. Weight-based dccomposition: A case for bilinear MLPs, 2024.
[61] 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.
[62] Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. Int. Conf. Comput. Vis. (ICCV), 2021.
[63] Gabriel Ilbarco, 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-29925 2022.
[64] Gabriel Ilharco, Marco Tulio Ribeiro, Mitchell Wortsman, Ludwig Schmidt, Hannaneh Ha- jishirzi, and Ali Farhadi. Editing models with task arithmetic. Int. Conf. Learn. Represent. (ICLR), 2023.
[65] Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravec, Zac Hatfield-Dodi Dodds, Robert Lasenby, Dawn Drain. Carol Chen, al. Toy models of superposition. arXiv preprint arXiv:2209. 10652, 2022.
[66] Tilman Räuker, Anson Ho, Stephen Casper, and Dylan Hadfield-Menell. Toward transparent ai: 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.
[67] 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.
[68] 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.
[69] 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.
[70] Stephen Casper. Broad critiques of interpretability research. 2023. URL https://www.
al alignmentforum.. org/ s/ /ttaa//vveeuturrrr//////// http://///////////gg/gee//e/////
[71] Shlomi Hod, Daniel Filan, Stephen Casper, Andrew Critch, and Stuart Russell. Quantifying local specialization in deep neural networks. arXiv preprint arXiv:2 10.00058, 2021.
[72] 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
[73] Timnit Gebru, Jamie Morgenstern, Briana Vecchione, Jennifer Wortman Vaughan, Hanna Wallach, Hal Daumé Ii,, and Kate Crawford. Datasheets for datasets. Communications of the ACM, 64(12):86 92066692277 2021.
[74] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Int. Conf. Comput. Vis. (ICCV), December 2015.
14

---

[75] Saachi Jain, Hannah Lawrence, Ankur Moitra, and Aleksander Madry. Distilling model failures as directions latent space. In Int. Conf. Learn. Represent. (ICLR), 2023.
[76] Moritz Hardt, Eric Price, and Nati Srebro. Equality of opportunity in supervised learning. In Adv. Neural Inform. Process. Syst. (NewrIPS), 2016.
[77] Mei Wang and Weihong Deng. Mitigating bias in face recognition using skewness aware reinforcement learning. In IEEE Conf. Comput. Vis. Pattern Recog. (CVPR), pages 1322-9331,
2020.
[78] 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.
[79] Mohsan Alvi, Andrew Zisserman, and Christoffer Nellåker. Turning 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.
[80] Iyy O Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, et al. MLP- mixer: An alI-MLP architecture for vision. Adv. Neural Inform. Process. Syst. (NeurIPS), 34: 24261-24272, 2021.
[81] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and llya Sutskever. Language models are unsupervised multitask learners. OpenAl Blog, 2019. URL https://cdn. http://///ddooooaii.e.eamaa are ure_unsupervised saaiedddmlltttttttttaeeee mlliitaak__earrerrss learners pdf.
[82] Aaron Gokaslan and Vanya Cohen. Openwebtext corpus. http://Skylionn007, 1atp:/11iioo00.ithuuu... io/ OpenWebbeetCCorpss, 2019.
[83] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In IEEE Conf. Comput. Vis. Pattern Recog. (CVPR), pages 248-255, 2009.
[84] Tolga Bolukbasi, Adam Pearce, Ann Yuan, Andy Coenen, Emily Reif, Fernanda Viégas, and Martin Wattenberg. An interpretability illusion for bert. arXiv preprint arX1v:2104.07143, 2021.
[85] Alex Rogozhnikov. Einops: Clear and reliable tensor manipulations with einstein-like notation. In Int. Conf. Learn. Represent. (ICLR), 2022.
[86]
Ledyard Tucker. Some mathematical notes on three-mode factor analysis. Psychometrika,
3 31:279-30 331, 1966.
[87] Carl Eckart and Gale Young. The approximation of one matrix by another of lower rank. Psychometrika, 1(3):211-21 218, 1936.
[88] Pratyusha Sharma, Jordan T. Ash, and Dipendra Misra. The truth is in there: Improving reasoning language models with layer-selective rank reduction. In Int. Conf. Learn. Represent (ICLR), 2024.
[89] 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 JEEE Conf. Comput. Vis. Pattern Recog. (CVPR), pages 7959-7971, 2022.
[90] 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 JEEE Conf. Comput. Vis. Pattern Recog, (CVPR), pages 8919-8928, 2020.
15

---

Appendix
Table of Contents
A Broader impact
16
B Fast μMoË implementations
17
B.1 CPμMoE cinsum implementation
17
B.2 TRIMoE einsum implementation
17
MoE forward pass visualization
Decomposition choice, matrix rank, computational cost
D.1 Tensor ranks to matrix rank
D.2 Why low-rankness reasonable assumption?
20
D.3
MoE/μMMB parameter count comparisons
20
E Hierarchical μMoo. model derivations
21
E.1 Hierarchical CPμMoE...
22
E.2
Hierarchical TRuMoE
22
F
Experimental details
23
F.1
Network configurations and hyperparamters
23
F2 Weight initialization
................................
23
Expert specialism: additional results
23
Large scale models
................................
23
G.2
LLM steering
................................
28
G.3 CLIP ViT-B-32
28
H
Ablation studies
28
H.1
Entmax softmax
28
H.2
Fast forward pass computation speedups
34
H.3
Batch normalization
34
H.4
Expert load
35
Additional performance results
35
L.1
CL.P ViT-B-32 ImageNETIk ablations
35
1.2 Hierarchical μMоEs
36
1.3 Comparisons dense/sparse MoEs
36
J Fairness baselines metric details
37
K Fairness: additional results
38
K.1 Model re-writing
38
L NeurIPS Paper Checklist
40
A Broader impact
This paper presents work whose goal advance the field of interpretable machine learning. Our goal is not improve model capabilities but rather orthogonal one of designing architectures
16

---

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 implement the fast forward passes of the μMoB models in batch-wise manner, where each mini-batch element is a 2D matrix of shape Z RT'xC' (with 'token' and 'channel' dimensions) with Py Porohh and einops' [85] einsum:
B.1 CPμMoE einsum implementation
The CPμMoE forward pass can be implemented with:
CPmuMoE (r=CP rank, b-batch_ dim, t=tokens, =experteet_coee coefs, =*=expert_ dims) = einaum (G3, a[0] @G1. .T, z0G2. T, 'r btr, t -> t o²)
i=input_dine dim, p=output_dime
And two-level hierarchical CPμMoE with additional factor matrix
CPmuMoE (r=CP rank, b=batch_ dim, t=tokens,
i-input dim, p=output_dim, exx]=exp coefs, =*=expert_ dims)
****
A 2-level hierarchical CPmuMoE, assuming Gi's of appropriate shape y einsum (G4, a[0]0G1.1 e[1] 20G3 T, o bob, btt. btr -> b oto
B.2 TRuMoE einsum implementation
TRuMoEs can be implemented with:
TRmuMoE (r**TR ranks, b-batch. dim, t-tokens,
s-input_dim, o=output, dim, a[e]=expert_coofa coefs, n#=xxpertddims)
batched mode-2 tensor vector products
f1 Sinsum(a[O) G1, "b t ni, ri mi r2 -> b r1 r2')
f2 = einsum (z, 62, "b i, r2 r3 -> bttre r3))
batch mutch-yultiply f10f2
fout einsum (f1, f2, t r1 r2, bt rtr2 12 r3 <> b r1 x3*)
contract with final TR core y einsum (G3, fout, Fr ri, b btri
b bbto)
And two-level hierarchical version with an additional TR-core as:
TRmuMoE (r*=TR ranks, b-batch. dim t=tokens, i-input_ddie _imm o=output_ dim, [[]]=exper__com coefs,, n*=expert_dimg
A 2-level hiererchical TRmuNoE, assuming additional TR cores Gi
f1 einsum (a[0], G1, bb t n1, ri ni r2 -> b r1 r2')
f2 = sinsum [1]), G2. bb t n2, r2 n2 12 r3 >> r2 r3))
f3 einsum G3, "b i, r3 - r4
batch batch-multiply f1ef2013
fout cinsum(f1, 12, "b t r1 r2, bt > b t r1 r3*)
fout einsum (fout, f3, bb rl r3, b 14 > b r1 r4))
contract with final core einsum (G4, fout, 'I o ri, btrir -bbo

---

C μMMe forward pass visualization
For intuition, we provide a visualization in Figure 6 of the step-by-step series of tensor contractions W x X2
E IR that the μMoB computes (in non-factorized form).
(G')
Weight teroor W. [1/5]
(GG
Weight censor W
[2/5]
(G
[3/5]
Expert conficents
X1a
X18
(oople die
Loopes Jun
Outgun vector
Interndiatt aunget
NF 96
14032E costraction
X
Contrnct reight (aleng
seput vector
Y R
Concrace waight tupes vactor
2nd Soda)
(c''s)
[4/5]
[6/5]
opperte
Expart
coolicienta
X1a
X
IMoE's forward pass Visualized step-by-step
Tecrpst etc.
ovez the pertial cetputs
Bac over octpoca cORESIOLNES)
Figure 6: An intuitive visualization of the μMoE (unfactorized) forward pass, visualized (as series of tensor contractions) in steps. Each step contributes to producing the output vector y E Ro either by contracting with the expert coefficients appropriate mode of the collective weight tensor W E NNxxx
RN with the input vector R along the
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.
Table A computational comparison of decomposition choice for /MoE layers and existing MoEs.
Param-efliciet
Param
Param-eflicient efficient
(medium N)
(large N)
# Parameters
Estimated FLOPS
Max. expert materx rank
min(I.O] min
Sparse MoE
NIO R(N +
X1O
CPaMaK
RIN +
min[T.O.
TRuMoE
R.NRs + RoIR, + R,0R, Ra + R.N., + R.R.R, ROR min [Rs min Ro.10
D.1 Tensor ranks to matrix rank
One important consideration how the chosen tensor ranks bound the resulting experts' matrix rank in pMoE layers. Here, we derive the matrix ranks as a function of tensor ranks for each model in turn.
D.1.1 CPpMMEs:: rank analysis
CPμMoES are parameterized by factor matrices U(1) RxN, UT2) RR (3() E RxO for chosen CP-rank R. Following Section 3 of Kolda and Bader [53] which provides the matricization/un- folding of CP tensors, we can write expert n's weight matrix as
W = U2) (USH @ U
E I'0
(6)
18

---

where @ is the Khatri-Rao product [53], and U.! E RRx1 is the column of the factor matrix associated with expert n (including a singleton dimension for the Khatri-Rao product to be well- defined). Through the linear algebra rank inequality for matrix products, have
rank( (W,) rank (V U
T o (UO))
= min rank U rank
( + U(3)
RXT
OXR
Therefore single CPuMoE's rith expert's matrix rank is bounded by min
D.1.2 TRuMoEs: rank analysis
We now turn our attention to TRpMoES, where we will see that the TR ranks R1, R2, Rs translate very favorably into matrix rank smaller computational cost than with CPyMoEs. First recall that TRyMoES are parameterized instead by core tensors 4 R XNxR2 u((2) RR RIxXR
u(3) E R with chosen ranks R1 R2, R3. We can derive expression to materialize expert n's 's matrix through the sum of matrix products of the TR cores as:
Rg Σ
W =
(3)
(2)
E RRix
(8)
PaRt
ting
1==1
Ox R2
Rs X Ra
R
The matrix product rank inequality applies each T x O matrix summand, whilst the matrix rank inequality applies to the outer matrix sum:
Ra
rank W.) = rank
Σ 1==1
(U+) ( Fast
(9)
rank ((U()) (
(10)
Past
min rank (U() (gn) rank (un)( rank (unk( iirs)
(ank(U(3)
(11)
Consequently, expert n's materialized weight matrix in TRyMoES has more generous upper bound of min { RR min R,
Through this analysis, we observe that one can choose large values of R3 yet small R4, R2 to yield high expert matrix rank with few parameters, justifying the choice of Ri Ro 4 in the main paper.
D.1.3 TuckeryiMoEs: rank analysis
One popular alternative decomposition is the Tucker decomposition [86]. Here we derive the resulting matrix rank of this alternative μMME variant and detail why it's not desirable as the proposed μМоE variants.
A TuckeryMoE composes an (MoE weight tensor through the series of mode-n products [53]: W = Z x1 ( X2 U X3 U33) where Z Ry R2 Ro is the so-called 'core tensor' and U E NxRx U2 E RI U, E Ro are the 'factor matrices' for the tensor's three modes.
Again following Kolda and Bader [53] a single expert n's weight matrix can rewritten through the matricization involving the Kronecker product @ as:
W, ((2
(U@) & U E Tx0
(12)
'Regardless of how large R3 is. the rank of the matrix cannot exceed min 0).
19

---

where Z(2) E R x (Ro RRx is the so-called mode-2 (matrix) unfolding of the core tensor [53].
Consequently, the same rank inequality applies:
rank w.) rank
(22
( @ U )
(13)
min rank U rank (22) rank U U
(14)
R2
(Rx Ox(Bo-RN)
Where we see the much more restrictive matrix rank upper bound applies: min {min(1, I, Ry), min( R1, Ro R min(O, Ro)). Thus in practice, both R1, Ro need to be large to yield large matrix rank, which in conflict with the goal of maintaining moderate number of parameters.
D.2
Why is low-rankness reasonable assumption?
Given we' ve seen that parameter-efficient μMoE layers lead to low-rank expert weight matrices, natural ques- tion is whether or not low-rankness in MLP linear layers' weight matrices
MLP-Mixer ImageNET1K accuracy with Truncated SVD MLP weights
reasonable assumption or constraint.
70
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 no signif- icant drop in accuracy from their linear layer counterparts (see also Appendix for many more results).
50
40
ImsooMETIs
30
20
To investigate this further we perform a rank ablation on trained MIP--Mxer model with the original linear layers' weights. Concretely, we compute the truncated SVD of each MLP block's s 2 linear layer weight matrices, We explore the impact on the model's ImageNETII Ik val- idation set accuracy when using only the top-k singular vectors/values (the best rank-k approximation [87]). The validation set accuracy using truncated SVD weights in every mixer block is plotted in Figure 7-we see here that discarding as many as half the total number of (bottom) singular vectors/values to approximate the original weights leads to negligible difference to the validation set accuracy. In other words, low-rank approximations of MIP 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.
10
20
40
60
80
100
Percentage
Figure 7: Val. accuracy for S-16 MLP-mixer when performing truncated SVD on all MLP's linear layers' weight; model accuracy is closely retained even with half the singular vectors.
D.3 MoE/μMoË parameter count comparisons
We plot in Figure 8 the parameter counts for μMoF layers function of the expert counts (sweeping from N = 2 experts through to N = 16, 384), relative to dense/sparse MoEs (with rank R4 = Ri = 4 TRuMoEs), for the first layer in a MLP-mixes channel-mixing block [80]. As can be seen, both μMMB variants are vastly more parameter-efficien than dense/sparse MoEs.
Given TRaMoEs 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 MMMF, linear layers, and alternative MoEs shown in Table 6, where the μMoEs perform favorably.
20

---

MMME layer parameter counts with dimensions = 512, 2048
CPMMoE (R 512)
1020
Number of parameters (log scale)
TRMMoE (R, 128)
Soft/Sparse MoE
10a
10a
107
10c
64
512
4096
Number experts (log scale)
Figure 8: yMoF layer parameter count function of expert count.
Table 6: Comparison of different layers' peak memory usage and latency (per single input). We 128 experts in each MoE layer, and set the rank of the μMoEs to parameter-match that of the linear layer.
Layer type
Peak memory usage (MB)
Latency per single input (ms)
Linear layer
12.07
0,01
Dense MoE (N = 128)
390.17
1.17
Sparse MoE (N 128)
765. 16
0.80
MRuMoR (N 128)
15.87
0.94
CPμMoE (N 128)
14.02
1.05
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 multiple levels of expert one hierarchy-ooe one simply increments the number of modes the weight tensor and includes another tensor contraction with 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 small parameter cost. Here, derive the fast forward passes for jMoE layers in their most general form with E' levels of expert hierarchy. For intuition, we first further visualize μMоE layers with 2 levels of hierarchy in Figure 9-note how we have an extra mode the weight tensor, and an extra contraction over the new expert mode to combine its outputs.
Given that hierarchical MooEs involve very high-order tensors, adopt the popular mode-n product [53] to express the forward passes in as readable a way as possible. The mode-n (vector) product of tensor X e IR X and vector R I is denoted by X' x,, u [53], with its elements given by:
In
(X xn dh
Fe Uin
We first introduce the formulation of an E-level hierarchical μMoE layer from Equation (1) in the main paper: given input z E / the most general form μMoE layer is parameterized by weight
tensor W Nr X.....xNxxxxxx and E many expert gating parameters {Ge I [x E
Je=1- The
21

---

2nd level's expert
BI= e ())
X2 22 confficintss
exeperrta_21
1st level's expert coefficients X1 a.1
[a.iopute]
W WER WER NIXNaXXX
Output
Weight tensor
Input vector X3
Figure 9: Illustration of two-hierarchy /MoE layer's (unfactorized) forward pass series of tensor contractions. The N1 N2₂ many experts' weight matrices are visualized as 2D horizontal slices in yellow, which are (1) matrix-multiplies with the input vector, (2) summed over the first expert mode (weighted by the first expert coefficients a₁ in red), and (3) summed over the second expert mode (weighted by the second expert mode's coefficients a2 in dark green).
explicit, unfactorized forward pass is given by:
a (GG z) N: Ve E (I....:E),
= W X1 X X ag X
01m1 GENg
W
mgg 2) RR.
(15)
OXi
where Equation (15) is expressed as sums over the E-many expert modes to make it clear that hierarchical μMoßs take convex combinations of II Ne many experts outputs (given there are Ne experts at each level of hierarchy). With expert coefficients (a, E R.N. J. the factorized forward passes of the most general hierarchical yMoE layers are given for the two variants below.
E.1 Hierarchical CPμMоE
The full CPμMoE model of rank an implicit weight tensor W = D uf ou o of o...o (E+2) E Nixxx Nex with factor matrices U(!) RRX, ..... U RR (F+1)
RxT U(F+2) E RxO The implicit, factorized forward pass is given by:
R
(2) o
u. (E+2)
X1 a1 X2 X ag X
(E+2)
uyn Oln1
u aEp thi 24)
(E)
(E-1)
Σ (E+2) (U()))) ... (U(() (U(E+1))) 2), ERO
(16)
E.2 Hierarchical TRuMoE
In TR format, We W NgNxx x.xXxxxNXxxxx has E + 2 factor tensors: 4() R RR.MM (F)
R R R (R-11) E R Rg (FF-2) E Rg #2 where R, are the manually
chosen ranks. The weight tensor's elements given by:
Rgio = tr(( ml
( U
(E+1)
1o:
22

---

We derive the fast factorized forward in terms of series of mode- products:
Σ Σ W(n ,,, 2,), a5 (n
(17)
(B+2) Tg-230
(U X2 a1) (u))) X2 a aE) X2 mPR-=
ERO.
(18)
R4 Re12
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 lattps://githe
meriiee--pytcccce the open-source code for NanoGPT www.wddiiinnmmlpppp.... (https:/ gi github. com hum/karpathy//aann
for GPT2 [81]. and the robust fine-tuning protocol of [89] for CLIP [61]. These values summarized in Table 7. We plot the learning curves for the training of both models Figures 10 and 11.
Table 7: Experimental configuration and settings for the results reported in the main paper in
Section 4.3.
Learning
Batch size
Weight docny
Warmup steps
Training duntion
Stochastic depth
RandAugment strength
Mixup strongth
Mixed precision
Random send
rate
Dropont
Handware
MLP Mixer
1e.3
4096
1e.4
10k
300
True
15
0
0.5
bt16 fp16
4xA S0GB
NTnoGPT
iter.
10 epochs
fp16
90GB
Rank choices Throughout all experiments in the main paper, we fix the TRpMoE ranks for the first modes to be R Rs 4. This way, can maximize the effective expert matrix ranks low parameter cost, as shown in Appendix D.1.2. The final TR rank R3 is varied to parameter-match the networks in question. For CPμMoEs, we set the single CP rank R to parameter-man match the baselines.
Training times Each MLP mixer model takes just under 3 days to train on 4xA 100 80GB GPUs. The NanoGPT models take 2-3 days train for 100k iterations, with the same resources.
F22 Weight initialization
We initialize each element of the factor matrices/tensors for the input and output modes from a Ul-VR, vkj distribution (following PyTorch's linear layers' initialization strategy), for features, where in feafurees is the dimension of the input 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 (1, distribution. For TRoMMES, the weight matrices can instead be replicated along the expert mode by initializing each slice (e.g. 91 (1 i,:)) as a diagonal matrix with elements sampled from (11 rr. 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 MI.P-mixers trained with both CPμMoE and TRuMoE blocks. Examples shown for the first two experts as they appear numerically for
23

---

Large scale MLP mixer results: ImageNET1k validation accuracy
Large scale MLP mixer results: ImageNET1k training loss
70
MLP
CPuMoE (ours)
TRoMoE (ours)
Validation set accurasy
50
Training
30
20
MLP
2:39
CPuMoE (ours)
2:44
TRuMMoE (ours)
50
100
150
200
250
300
0
50
100
150
200
250
300
Epochs
Epochs
Figure 10: Training loss and validation accuracy for the MLP-mixers models for 300 epochs.
Large (NanoGPT) results:
Large (NanoGPT) results:
OpenWebText2 training loss
OpenWebText2 validation loss
11
11
CPμMoE (ours)
CPμMoE (ours)
TRuMoE (ours)
TRuMoE (ours)
Validation loss
Trainin loss
B6B
2 2,886 2,876
2.862
20000
4D000
60000
80800
100000
20000
40000
60000
80000
100000
Iterations
Iterations
Figure 11: Training and validation loss for the GPT-2 models for 100k iterations.
each of the 8 layers, where we observe the same phenomenon earlier blocks specializing to textures, and later blocks to higher-level abstract concepts/objects.
Secondly, Figure 13 we show the top 32 activating tokens for the first 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).
24

---

Expert
Expert 2
Layer
Layer 2
Layer 3
Layer
Laye 5
Laver
Layer
Layer
(a) MP block MLP-Mixers: top-activating tokens.
Expert
Expert
Layer
Layer
Layer
Layer
Layer 5
Layer
Layer
Layer
(b) TR,MoE block MiP-mixers: top-activating tokens.
Figure 12: Top-activating patches (and their surrounding image context) for the first experts two blocks in ML--mixer models. μMoE blocks (with 64) exhibit coarse-grained specialism (e.g.. texture) earlier and more fine- fine-grained specialism (e.g., object category) deeper in the network.
25

---

Eepart color mape
Layer 5, Expert
Layer 5, Expert 2
kid
the yearts aoul wmaine
wasoatatientiio oodic the bega Sevie Uniin, iito
the souded torkon had
sabees securty
Congrast
liTte rocent coim
trp
the Moull thehh Unitud Stateo two
sagt
Moderis
pasttmmnt
Machine Cow
Good waiched
very voii lary the total communier and though itt
panterr merutt
the 1.800 Sovect Uittss huvt hoor monden juce os
the Sivier Uoei the
dacorbe an coed angine Wel.
vilil sidir the kamyy aherr the berth well verrs Germans. The Commane came
Pestores teem
workd
tacover'd
Insemar
#23 (oa: 2016 coye nent and talted
The fondly stores tarms thene the Hepps withh brsued statement
they witch
caur the
cor the aesati and cesturianiis the
222. whdch
your uquaa cupudar Foycho hos
pallad the
2010
fi berlin Wall, Derfi also
#30: wis The
by
shoce 2932 Homee" min
tuat
Layer 5, Expert 3
Layer 5, Expert 4
for than
said inppetart know
syewe
bastclly gres the Stane Drpart the
#10 prosecs
peaceiul,
non nut
+tRm aspeating to as
010 Unitr sartheeeneed
#1L
the
F11
exn't the
premitsee
#14 the
ocd to fon
616 trt "yxnssss sui (Wrr' to dong foror the
*15
015
#17
rit that worther trat
K18
FL8
#20
420.
#21
fou thet
#21.
422
924
t0
#20
only thaggs my
120
Soatinn but
427
ro who
420
(dbrrt time what tii
the
428
*XL. nerk they aolad ine it thes gor they and
F2L SoctoN the tech lata
ane sould
#22 Na Elyyer now ane Argies cuppersice for mrathh and
F22
ph cytter, wetcn the called's eystem on againe the puslic
Layer 5, Expert 5
Layer 5, Expert 6
#1:
in
#1.
e1: srink
dangercall worts is abd changee
Suiit tae pat be ariee Chhici ciici
recsiing hare rog tracit a cassa. thingryee
kerr bos Inotived heraaaana anl wodted hard atudeetour
the cosute
The
rew cy
the
essemant faet fass-
parforrid mudiit and #11: gratt ane spreas working ing
#11: foder povermenn monttrr the vorcon on. and the Auciic vn- atsta have deneloped
importandd things the hespored pas years
dully cumctal wilti
and mollee wark
athertarter.
hostial
has then burnec the tut ond ceti nulk on frr The Oburn comiitrration huo refued
aficted Toy Siiie
wup
and symana di pclss in psice and ruper tal
sueg the years
tout and adviited tot
than millien
"en cintttann sos the repert to with devoo
chiigg and
#25: b
tos
such es
430:
can et beat and the voics mondte restaite and improvdd oo
435.
to ru. s0B anboayy Hivarr wown
wird and unirr the bocker of Coudtt Casleed and Custone
was "Temsibe" bn tua
#52
432)
Figure 13: Top-activating 32 tokens for the first unfiltered experts 1-6 (as ordered numerically) at layer 5 in the CP/MMOE GPT2 model (Please find the next 6 experts in Figure 14).
26

---

contictente
Layer 5, Expert 7
Layer 5, Expert
PL light suggart EArgENE
20t lng ie actuet
press
-3: Mie (We manothe Mertiopon Thu the number people empleyed in vor dladt
#3 lust hou cun by how
icsolpiicl right
P5. srarig and
*21
hank
reat BE Extow
K14
houn compagg
*18 2055 Erng fov preideenee aaee socowecledd
428
apponer licass pan
your inmadient
433
439
Rubry
#24
ware ponbrrasponn 35
138 16. 22 delae Annie the nee
Raail teetre
#31
Layer 5. Expert 9
Layer 5, Expert 10
#2: Pr
weree thei'v
#2
avrr with tack Hends and femily the reed.
cloal
A4:
Fush
asic tnsugh Mne
City when you tak the exolidor
most
dering trad ty Seusigo
#15 to Coople preaee
617
F18
by
420
822
tkis
736
inee ane
X77
best ldacified
213 ialail
Layer 5, Expert 11
Layer 5, Expert 12
expecialy
tyrih hi
very
Agan
you
2055
the
PAG: tegge have toes
033 T14 co cand
traamtaameenn sia varry of Mincrdee
A12
moced endinction
#15 you spas with
K18
Ehammanttiatt trass searg
in posrion
+55
2013
senchiiee
faamavook lorite picbal adacotion
#20 dic medk undi chenth nod bo
undc saae Synt, Shordh
922
-22
P15: eogacs and
tvensoee the weteranne 13
+25
437
*.8
resocaation
#06 in having hamt time soditi mover
433 the Year Thes 20u of acrd fod mower theimmeette
an
undi Hemds kimb
412
Figure 14: Top-activating 32 tokens for the unfiltered experts 7-12 (as ordered numerically) at layer in the MPμMMoE GPT2 model.
27

---

G.2 LLM steering
Here we provide additional evidence that the experts' specialization is mechanistically relevant to the functionality of the network, in the sense that use them to steer the LLM's output.
In particular, we use 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 spaciiilill adding selected expert cluster center vectors each token's input latent activation vector before applying the μMoE layer-we we can consistently control the model to generate outputs aligned with specific themes. Illustrations of this approach, using 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 topies 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 2 0.5) of the first few experts as they are ordered numerically. Furthermore, when we use even larger number of experts (i.e. 2048) we observe select few experts developing what appear to be very fine-gre seeegraandd 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 cating 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 shown in Figure 18 and Figure 19 when using a CPμMM as a final and penultimate layer respectively. As can be seen, often observe that 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 chosen in numerical order and only those yielding 2 0.5 total accuracy change any class upon counterfactual ablation.
H
Ablation studies
H.1 Entmax 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 cPpMooE-5512 final layer models fine-tuned with various numbers experts. As can seen, the entmax function consistently produces more monosemantic experts for larger total expert counts. We attribute this to the sparsity in entmax's 's post-activation distribution (whereas the softmax function can just readily output uniform distribution over all expert coefficients).
28

---

dodiine dus Raw model outputs
Sarov2 3imple
The begest isue of today's world the presiurc on planet's fragle bioophere do
makins aboriatie Bi ppurees the chart di exllatiion hut there no sudeece for them soye the EXFT report
laaget that an cast majorit of medis don't wart sae so it neceaing dificult
One
The bagest isue of todlays world internet wilch only Inportin tant aspects our lives out also has impo tant benchts that geadly boschtt Rebody wares to be charged with crime, and nobody
the
The Nagpet lsus of boday's world not climate charge but the fyluue make the unakling the right thing and awold a massive catsstrophe Wact noted United
The ligpest lous of today's werl this As the world continue and mone into the
poss tho nut thet And the pada domand food
harder
soms key nestons
on The liayest of bodayy's Fow to manage growing food production system, though we must grea chillonge but als strtogic apporturity That's bocause city doell
manage a
Steering LLM outputs +
The Nageet loun of todby's wordd the comtniing reed for moss nowe-umaan immigrition
lica aporoved and the rest I1.5 millon) will at bert, temporary Ther's because of
the fational. world of today euperincing erfourdd change. We citaittrrdd the
of toriay's
problams
proslame of today
cooted a sabuaion
Intervening to increase:
Intervening to increase: "l "programming" " expert
/"eather"//"clinnnn /"climmte" expert
Syrph? Sandee
The biggee: baus of today's world is chot of dimats changs, and many people are the morsoon tiat hreses into March. Bet that could the Hages lsuu lacing lidia
The kigpet leue of today's world is that of the foewl valse with nealt from
Fdoruary
whic a lamboda, will be bis that is now Value
today's is world is the esk of giart iais sepecialy the Briis lole, high terrpe tempetatur expected to mach on lighed maxtmum of X.
"The higgne:
The Haueet iess of todan's world is the lack of, constant: DeoA -> (TTTT)] -t
volue, but
lhety hour.
again
per
typocck imp emonts
The Bauest isse of today's world is the accomulation not ooosl or typ rad
with the best have to sperochh this problem caininuely So for to good, anyany! So, here
bigane:
today's is world the worn samperature that you can expect equator But
gold,
The fiest type is and the second.
inbatre
with
ligpett lown of todoy's world not schv
be able cvatte And most we
it. We lot more water and warmer
loe more sursiine But the climate warmed, now there area restivelyy low semperaturs is gasting
world not burn
Thea types comernct a member
Any
be created anyohers the set
The Happet less today's world
reash all We will ses miliors bilken,, and billors of world in idontly some kidd
The big bigges: bave today's world is this As gummer aporoaches, mane fammer grow Ther''s wothing vnnry shoue, epaeislly the Souther Califonis mountains If
things And the way we thin
The
type (higher type) ovedladed oyee. life fanct on My type is resut, in like
result
The hig higgee hoss of today's world a lack of snndine oo which makes the hest perfect storm sunahine hare, at high pemsure and it's going to be diody for
after run) 1) and type
The Higgest sloays number and the aim of all the seter vrlume alays number. Bue the [bigee] type always
steanee of clow. "A
The bigane haue today's world is lack of rainfal and ectreme semperatures so piono to the high comperalure but those regions should be day, and the temocratare
ahotss
will
The Hlgeett issue soday's world ceceation is net a auperdaas,, and a static class carnect be. la a single consiruct and any
"Tee higgne today's is world is chat, aveeage the atmorphere too hot be cool, a though the wammer ar metting some parts the Avatic continen
haze windishear
Intervening to increase: "pllice violence" " expert
Intervening to increase: ll "foreign politics" " expert
Sayppi 301243
Niggett
is that don't have eoough pol officers
The Hagpett leoe today/s world is live
Americane must leaders said.
to
have many immigrants from lante Statee the helping the "We hive a drgg proslem,
oeia sheye vary the American politee Kowa, the North Kireats the US and Great Britin
depestation
infigners with drago. We've
get
The ligent isus of today's world the miltarization of polioe, sapecially in the US, fos
The biapest isss of today's world is coming the ratioaaioo that since the cold era, the Stares
Ferguson: and the ail muhible cxllann.
medi ability get The US milltar's lack
whut's hagpening
Spoee Union They sucrebly the histanc dea in 1963, the Stwiess all stop the Berlin Talran.
its
for resolztion
accountabiity puts
world aecumolation of yesourcos thet being
baaest
and
movel, and the eflects of that 2016 Derard Trmm, Tease police oftiwes States,
the every seose of the word They one to other tne for The regior is the protlem: the
fog
sees no trouble
The bigget imus of bocla's is world not immigaation. but the war as drups The war
The bigget ieus of today's world is has to scho the iaue the Black Sea when througe Washingttnn 's represettal ve from the time of Napoleen until the Empin''s resegné
country's x efforts keep us sone. The War on Drugs has been an unjind and
bigestt
world this As United States
The blakest issue today's is world is beth sides alse pursuing plan avoid permenent at-atoo atallic allince out ther cermot reach for tos is now calling its long. especiily in Wastington, does not mant
lagest in Chufle spoll achrrotttoell polis death "dtwing" villnne and ether
getting
US hitss //t
The Higgest todly's world that consuncly the same side and
gre cortarrlyy
of Pal inanss howove mich we are there as well" Elmon by or And alo viclent mony
The liaestt isyv of tyday's world a courcee's atgleiic response to liae' inegtion burched musise oper stion to fndd Baghhas acled neighbote and captiine cccupy the city
before people
seek
that, polce getting kiled of repert the polre, other tncluding the of
hodoy's
way
The bigges isse of today's world is fight between the two leades over their mass the Eor Oitoman Rabid munial pupirat tied naciee And dificabiss left
drones. That's where thin,, a police officer it's araater
were
with
broken immigration who boing contenally killed by a pesson who owns scficle They have lilos five offices pesi wodh.
birgest bacly'' world that we
We
war-oeary Romaan Prendent
have economic systom, We
police office
The isuu tolly's
conscquantaal tha coonomic poditors are bring keps for falf year Donald Tramp the bet dnys were also
evcry love bry and over Russa
Figure 15: Steering LLM outputs by forcefully activating experts: adding specific manually
chosen expert's 's cluster centers to GPT-2's activation vectors at particular layers reliably steer the LLM generations towards specific themes, based on the learned expert specialism. For example, we see an expert that steers discussion towards police violence, or about the climate. The initial prompt in every instance is the text: 'The biggest issue of today's 's world is"
29

---

:MME: 256 total experts
IMME: 32 total experts
Figure 16: High vs low total expert count: Randomly selected training set images with expert coefficient 2 0.5 for the first 10 numerical experts (of those processing any images with coefficient 2 0.5). Results are with CP-r512 j.MoE layers with 256 (left) and 32 (right) total experts respectively. We highlight the apparent specialism of the experts when higher total number is used. (Please zoom for detail)
30

---

Expert 834
Expert 1203
Expert 316
Expert 206
10
Expert 203
Expert 1602
CPmuMoE-r512 2048 total experts
Figure 17: Fine-grained expert specialisms: Manually selected experts (and images ranked by highest expert coefficients) processing what appears to be very fine- tme-grained categories (e.g. animals with footballs, trolleys in water, etc.). Model fine-tuned on ImageNET1k with high number of 2048 experts and CP-r512 μMoE final CLIP layer. (Please zoom for detail)
31

---

HANNOOI MOWT NONT
32 total experts
CUATIN
64 total exparts
120 total exparts
256 total experts
wws.eesee siktin expent
Inanurtie eletin exet
Harrederr st urig uppert
Elivedaen de
512 total emparte
1024 total exparts
Figure 18: Penultimate layer CPμMoE: Percentage of per-class test set accuracy lost when inter- vening and ablating particular experts (along the columns). In general, the more total experts (rows), the more class-level monosemantic the experts are as indicated by the mass centred on fewer classes, and with higher magnitude. Shown are the first experts in each model (row) change 0.5 of any class" accuracy when counterfactually ablated.
32

---

32 tdtal erportss
onth
total emparts
120 Local exQRiES
ververier
Inenvedtan
250 Dotal empurts
Weanerbee eveen
512 total expartes
Oew
repart
Iiteraabe abliling maac 231
1024 total uxperts
Figure 19: Final layer CPμMoE: Percentage of per-class test set accuracy lost when intervening and ablating particular experts (along the columns). In general, the more total experts (rows), the more class-level monosemantic the experts are as indicated by the mass centred on fewer classes, and with higher magnitude. Shown are the first experts in each model (row) to change 2 0.5 of any class" accuracy when counterfactually ablated.
33

---

Experts' class-level polysemanticity (+)
Experts' class-level polysemanticity (4)
1.4
1.3
Entmax (dino)
Entmax (clip)
Mean paysensntiity
1.3
(dino)
poysementiity
1.2
(clip)
1.2
1.1
1.0
1.0
0.9
0.9
0.8
0.8
32
64
128
256
512
1024
32
64
128
256
512
1024
experts CPyMoE Blog scale)
used CPyMPE layer flog scale)
DINO backbone
(b) CLIP backbone
Figure 20: Softmmx vs Entmax ablation CPpMoE-r512 final layers trained on ImageNET, and the resulting class-level polysemanticity. For large values of experts, the entmax activation produces more specialized experts.
H.2 Fast forward pass computation speedups
We next report Table 8 the actual num- ber of FLOPs (as reported by https:// www.etrron2..aaaheeeeeeooo wttp:/iiiieeeeneeeeeeeeeeeeeeeee mooulee/ffrcore//eee
Table 8: Original μMoE layers' FLOPs vs the fast einsum forward passes in Appendix (for
count. html) when executing PyTorch μMoE layers using the naive forward pass relative 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).
N = 512 experts with 768-dimensional input and output dimensions).
CPμMoE
TR
Original FLOPs 155.1B
622.8B
Fast model FLOPs
1.4M
3.5M
H.3 Batch normalization
We next perform an ablation study for the use of batch normalization (BN) before the activa- tion function for the expert coefficients. We study CPμMMoE final layer layers with CLIP ViT- B-32, quantifying BN's effect on expert class- monosemanticity
polysemanticity Batchnorm ablation
Mean expert expert polsemamictty gift
1.3
-Xe
Final layer JMME (CLIP) (Batchnarm)
Firal LMoE (CLIP) batchnorm)
1.2
function of the expert count. Concretely, we perform the same class- level polysemanticity experiments in the main paper, with and without batch normalization in Figure 21. As can be seen clearly, the batch nor- malization models lead to individual experts that increasingly class- monosemantic as desired (as function of the total expert count).
1.1
1.0
0.9
0.8
32
128
Total experts used CPyMeE layer Olog scale)
Figure 21: Ablation study: batch normalization leads to more class-level monosemantic experts.
34

---

Expert CPyMoE, total exparts
17500
15000
12500
7500
Lunboe imiges
5090
2500
100
200
303
400
500
Exportinaxx (a) 512 total experts
Cimber image snligned with conti-teet
Expert CPμMoE, experts
8000
6000
4000
Enpertinaux
(b) 768 total experts
Figure 22: Expert load: Number of training set images with expert coefficient d, 2 0.5 for CPpMoE models fine-tuned on ImageNETIk. Bars are drawn with 3x width and colored sequentially repeating order of distinct colors to help visually distinguish between neighbors.
H.4 Expert load
Here, plot the expert load in Figure 22 to give a visual indication of how many images processed by each expert with a, 2 0.5 for CPuMoE final layers fine-tuned on ImageNETIk with CLIP backbone. Whilst clearly, not all experts have images with coefficient of least 0.5, relatively uniform spread over all experts. Furthermore, we note the cost from 'dead' experts is not particularly troublesome in an μMoB given its factorized form form-speaking informally, 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 ImageNETIk ablations
Here, we compare the performance of parameter-matcher μMoE final layers (for varying expert counts N) to linear layers for fine-tuning large vision-language models (CLIP ViT-B-32) on ImageNETI k. Following the robust fine-tuning protocol of [89], we use largest possible batch size (to fit on one A100 GPU) of 4096, and the same learning rate of 3e 05.
For μлОΕ 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 ImageNETIk validation loss after 10 epochs of training, where all expert counts out-perform the linear layers initialized the same default way with elements from k k). However, to marameter-match single dense linear layers, we must decrease the μMMo layer rank upon increasing the expert count. This is a concrete example of where the extra parameter efficiency of TRuMoEs can come in useful (as discussed in Appendix D.1.2). Consequently, TRuMoES resulting expert matrix ranks are increasingly larger than that of CPμMoES in the parameter-matched setting. For example, the parameter-ne parameter-matched layers with 512 experts in Figure 23a have a max expert matrix rank of 165 for the CPyMMEE compared to a much larger 208 for the TRuMoE.
35

---

CLIP fine-tuning on ImageNET1k: Rank comparison
CLIP ViT-B-32 fine-tuning accuracy on ImageNET1k (linear vs j.MoE)
78.7
Linear
ImageNETI, validation set accuracy
CPμМоЕ
400
78.6
TRuMoE
78.5
350
Resulting matrix
78.4
300
78.3
78.2
200
CPpMoE
TRUMoE
128 256 512 Number experts
8
16
32 64 128 256 512
Number of experts
(a) Accuracy comparison (μMoE Linear)
(b) Rank comparison (CPμMoE TRμMoE)
Figure 23; Comparative analysis of fine-tuning CLIP ViT-B-32 with pMoE layers using different configurations. All experiments have the same number of parameters.
Table 9: Hierarchical S-16 MRixeoEmixerss and CPyaMoE-nixers: ImageNETIk val. accuracy at 300 epochs pre-training; N1 = 64, N2 = 2 experts).
Model
Val. acc. (t)
Experts per block
Params
MLP
70.31
n/a
18.5M
CPμMoE (hierarchy=l)
71.29
G4
18.6M
MRAMoE (hierarchy=l)
71.26
64
18.3M
MPp (hierarchy=2)
64-2
19.5M
TRμMoE (hierarchy=2)
71.56
64-2
18.7M
We attribute TRJMoE'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 larger dimensional subspace) (see Figure 23b).
I.2 Hierarchical MMoEs
Hierarchical μMME 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 two-level hierarchical μМoE with N1 64 experts for the first level and N₂ = 2 experts for the second layer (128 total effective experts). As shown through the results in Table the hierarchical pMoE's also perform well against the MLP alternatives, whilst providing even better parameter-efficienn
Hierarchical μMoE fine-tuning layers We also perform additional experiments with hierarchical μMoEs used to fine-tune CLIP ViT-B-32 models ImageNETIk. Here the experimental setup in [63, 64], training each model for a single epoch with the specified learning rate of - 05. We fine-tune hierarchical MMME CLIP models with up to 4 levels of hierarchy shown in Table 10, where the best-performing models (averaged over 5 runs) are found with levels of hierarchy.
I.3 Comparisons to dense/sparse MoEs
The goal of the μMoE layer to facilitate more interpretable subcomputations with a similar number of parameters and FLOOPs 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 the scalable expert specialization provided,
36

---

Table 10: Hierarchical μMoEs:: Mean validation-se aceuracy with a CLIP ViT-B-32 fine-tuned with hierarchical μMoE final layers on ImageNETIk. Shown are the number of parameters as the number of total experts increases to 8192 with levels of hierarchy, and the corresponding number of parameters needed for each expert total using a hierarchy 1 μMoF, and regular MoE. Results are the average over 5 runs with different seeds. Additional expert modes for TRIMMoE have the additional ranks set equal to the corresponding number of experts at the new mode(s) (c.g. and 4).
(a) Hierarchical CP Es (R 512) fine-tuning CLIP ViT-B- -B- B-32 on ImageNETIK.
Hierachhy
Val ace
Wright tensor shape
Total P expents
Parans
4 Paraas needed (wr) hicrarchy (MoEE A Params needed (wr regular Mols)
73.78 0.07
W RiRKZZO
128
1,069,565
1069,568
98.432,000
glaszarso
1,233,408
0.14
W giza-2iietto
512
1.561,88
395,7280001
70.82 1.06
W glaektaaaaeeeo
1024
1,077,245
2,210,448
787,4450000
2
73.89
0.30
W & tis
512
1,074,685
1.561,008
393,728,000
73.85
0.08
W gu Riat-toadtto
2048
1,079,00
3,527,168
1,574,9122 20.000
0.00
W Rianeeeeeeeo
11.391,488
6,299,665.0000
(b) Hierarchical TRJMoES (R3 512) fine-tuning CLIP ViT-B-32 on ImageNETIK.
Hierechy
aal
Wright lensor shape
experts
Parmss
Piraens newded hieranchy pMoF)
Params needed (wi repulor MolF)
74.00
1.:0
W. gioor
128
3,723,264
3,723,264
38.4822000
2
74.72 0.08
W glasiano
256
3,724.832
3.823,616
196.864.00
3
74.75 0.14
W
512
3,756 400
4,024,320
393,728.000
74.76 0.11
W 10
1004
3,727.968
8561,456
787. 287.456,000
74,82 0.11
W glasstrr
512
3,726,400
4.024.320
393,728.000
74.77
3,729.536
5,228.544
02,0,9100000
74.73 0.11
W Ritodeeeeee
8192
3,752,672
10,045,440
6,299 000
MoB model accurary ImageNETik
Mof model accuracy ImagelEE
EPyMee
73
TRuMME
Unnar
Sp3rsoftcE SoftHoE
72
Accursc
CPyMeE
Unear
SparseMot
69
89
SolMoo
10F
107
10P
10P
104 PLOPS
107
Parameter count
Figure 24: Results fine-tuning CLIP V- final layers only on ImageNETIk for epoch. For μMoE layers, we increase parameter counts by varying the ranks for fixed 64 experts. For dense ("Soft") and sparse MoEs, we increase the parameters through increased expert counts.
the μMoEs also perform very favorably against the alternative MoE models when fine-tuning CLIP on ImageNETIk.
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 subpop- P(ý
ulations to be equal, defined in Hardt et al. [76] as P(ý = 1|4 = 0,
1|A = 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.
37

---

Standard deviation bias computes the standard deviation of the accuracy for the different subpopulations [77]. Intuitively. small STD bias indicates similar performance across groups.
Max-Min Fairness quantifies the worst-case performance for the different demographic subpopulations [78], with max m P(Y y a, Y = y). We compute this as the minimum the test-set accuracy for the subpopulations in each experiment.
Baselines
Oversample we oversample the low-support subpopulation 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' match the number 'blond females' for the first experiment, and oversample the number of 'old females' to match the number of 'old males' the second.
Blind thresholding implemented by unconditionally increasing/deereasing the logits in the target direction for all outputs. Concretely, the results in the main paper are achieved by
setting = 2.5 and to 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- re-writing [76].
Adversarial debiasing we observe in Table 2 the same poor performance for the adversarial debiasing technique 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 the final representation, information about correlated attributes (e.g. wearing makeup) are likely still present. This makes it fundamentally challenging to apply fairness taimnss--through-ee throughhunawwrrness techniques in the CelebA multi-class setting.
K Fairness: additional results
K.1
Model re-writing
The full per-subpopulation test set accuracies 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.
38

---

Target attribute: "Blond Hair". Target subpopulation: "Blond Hair"+
CypMoE--5122e128
Before intervention
Vale Fena
96
89
96
88
100
41
100
43
Oner
Siende
Obher
Stende
Blind
Expert thresholding
Blind threshalding
Expert thresholding
threshalding
After intervention
Temuie
Na Nemak
86
98
90
86
98
95
90
97
84
96
85
97
84
96
87
Other
Binds
Other
Blande
Other
Bonde
Other
Blondo
Blind
Expert thresholding
Blind thresholding
Expert thresholding
thresholding
Before- refore-aftel difft
Forale
Kanal.
0.51
Hewe Femal
9.9
Hale Forrae
0.76
1.6
Other
Other
Blonne
Ronte
Blonde
(a) 'Young blond" intervention for Blond hair attribute prediction head
Target attribute: "Young". Target subpopulation: "Old" "Old"*"Femmi""" "Female"
152MeE-512-e128
TR1MoE-5522e112
intervention
Before
Vaee Fenaee
55
97
Ma Femm
56
97
75
94
75
94
old
Yourg
06
Young
Blind thresholding
Expert thresholding
Blind threshalding
Expert thresholding
After Intervention
Temsie
Remake
84
87
87
83
8.5
87
92
74
93
70
Male
77
93
Mile
93
70
Hole
77
92
od
Young
OU
od
Young
OL
Yourg
Blind thresholding
Expert thresholding
Blind thresholding
Expert thresholding
Before- after diff.
Funale
Feraie
Fomale
36
23
24
24
Wais
1.6
(b) 'Old female' intervention for age attrihute prediction head
Figure 25: CelebA Subpopulation accuracies before (first rows) and after intervention (second rows), followed by their absolute difference (third rows). Green rectangles denote the target subpopulation each experiment (subfigure).
39

---

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 fine- tuning 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 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 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 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 ex- perimental results of the paper to the extent that affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: Full experiment setingg/conggllyyepprammeee are provided in Table and the supporting code (https://gitae com/james oldfi wwddelldoooe provides even more explicit experimental instructions. Learning curves are also plotted Figures 10 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 instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: Model code for μMMEs and the experiments in the paper are found at:htpp:: //github com/ james comes-oldfill
oldd e www.dfeeeeellffooooooooooooo Id/muMoE.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, 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, 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?
40

---

Answer: [No]
Justification: We do include mean (and STD) of the results over multiple fine-tuning models, but we only have single single runs of large models, we always set all random seeds to for reproducibility.
the large models due resource constraints. For these
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [Yes]
Justification: Details are provided Appendix
9. Code Of Ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: /ttpp:::///uurss.... c//ccc//p/// wc/ppblicc/thiccsuueeee.
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]
The paper proposed 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.
Justification:
11. Safeguards
Question: Does the paper deseribe 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 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 the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Yes, the open-source codebases on which we base our code are explicitly referenced.
Justification:
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.
Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
41

---

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


Reviews:
[{"review_id": "BoqrZnofJF", "reviewer": "Reviewer_Su68", "summary": "The paper presents an interesting and novel approach with the µMoE layer, offering a method for scalable expert specialization. Extensive experiments in multi-dimensional have been conducted to show the effectiveness of the proposed method. However, the marginal improvements in results, coupled with several other weaknesses, including experimental variability, novelty of factorization, clarity of metrics, and scalability, limit the overall impact of the work.", "questions": "N/A", "limitations": "Yes, the authors addressed limitations and the direction for future research.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. 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.", "weaknesses": "1. The idea of factorization is not new, and it seems that the authors simply apply factorization in the MoE context. The authors should specify and address any new challenges that arise when applying factorization in the MoE context.\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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank **Reviewer-Su68** for their detailed assessment of the paper: for praising the “novel[ty]” of the layer and the “technically sound” methodology. Furthermore, we are glad the reviewer appreciates the “extensive qualitative and quantitative evidence to support the claims” of the paper.\n\nAll stated weaknesses are addressed below, pointing out where the requested details can be found:\n\n## [W1]: Factorization challenges\n\nLow-rank factorization is indeed a well-known technique in deep neural networks, frequently employed to reduce the number of parameters. However, as far as we are aware, our work is the first to introduce a novel technical link between hierarchical MoEs and multilinear models. Instead of utilizing factorization to simply reduce parameter counts, our work provides a way of designing trainable layers with prohibitively large numbers of experts in the first place. **This is only possible through the paper’s extra derivations and formulations of MuMoE layers’ fast forward passes** (beyond simple factorization).\n\nWe have an in-depth discussion in Appendix C [L722] of how the factorization choice impacts MoEs. For example, we derive the theoretical relationship between tensor and subsequent expert matrix rank, and the relative strengths of different factorizations in this applied context. Furthermore, in Appendix C.2 we also discuss the impact of low-rankness in the specific context of the models studied in this paper. We believe this further detailed explanation and exploration of the methodology also addresses the reviewers’ second comment about further enhancing comprehension.\n\n\n## [W2]: Marginal Improvements and lack of multiple runs\n\nAs stated on [L293], the **results in Table 2 are already the mean over 10 runs**. We have added this to the caption to make it more clear. Furthermore, an improvement of (for example) 5% to the min-max fairness metric (for “age”) and half the equality of opportunity (for “blond hair”) are significant improvements over existing techniques.  \n\n## [W3]: Metrics not clearly explained\n\n**The metrics used in the paper are** **described in detail** **in Appendix I** ([L958]), as written on [L307] of the main paper. Furthermore, we describe in depth their interpretations in the context of the celebA bias problem studied, and the implementation details of the baselines. We will also explicitly reference this appendix discussing the metrics in Table 2’s caption.\n\n## [W4]: Evaluation on OOD data for vision models\n\nIndeed, our quantitative evaluation of vision models is limited to the test set (and not OOD data), which we already noted as a limitation in the conclusion.\n\nDespite this, our experiments on language models in Figures 5, 12, and 13 use *generated* LLM sequences as their input (rather than that of the training set). Furthermore, we also perform an additional experiment on steering LLMs for **Reviewer-qC51**, where the input prompt is free-form text. This provides initial evidence of the generalizability of the experts outside of the training/test data. \n\n## [W5]: Lacking quantitative measures of interpretability\n\nWe kindly draw attention to both Section 4.1.2 and Section 4.2. In the former, a new quantitative metric capturing expert specialization is introduced. In the latter, interpretability is further measured quantitatively through its downstream ability to edit models compared to existing methods. We highlight that **Reviewer-1sqG** praises the “innovative metric\", stating that it `\"provide[s] a clear measure of the impact of increasing the number of experts\"`. If the reviewer has any other valid metric in mind, we will be glad to consider it in our evaluation.\n\n\n## [W6]: Different model architectures and sizes\n\nWe kindly note that we already perform experiments across **four different architectures**: CLIP, DINO, MLP-Mixer, and GPT-2 (in Figure 3 and Table 4). These models are each of vastly different sizes and configurations (and operate on various modalities including language and vision), already demonstrating the wide applicability of the layer.\n\nThat said, we have access to a maximum of 4 A100 GPUs, limiting the size of the models we can pre-train to GPT2-sized models. We have added a limitation to the paper’s conclusion to note explicitly how we have not tested the layer at the largest scales of commercial LLMs:\n\n\n    Despite demonstrating the benefits of muMoE layers for various architectures with different modalities, we have not conducted experiments on the largest scales of commercial LLMs. Future work should explore this where resources permit.\n\nFor the camera ready, we will include additional experiments on a GPT-2 model trained from scratch with 2048 expert CPMuMoEs at each layer, to provide further evidence of the layer’s benefit at very large expert counts.  \n\n## “Ease of implementation” and visual illustration\n\nWe would like to draw attention to Appendix B, which contains pseudocode implementations of all muMoE variants. Furthermore, we note that the `MuMoE.py` file in the anonymous code linked to in the abstract provides a full implementation of the layers in PyTorch.\n\nTo aid with visual illustration, we plan to include an annotated variant of Figure 1 in the appendix that shows a step-by-step explanation of each stage of the computation of the forward pass for clarity.  \n\nWe thank the reviewer for their suggestion to improve the ease of implementation of the methods and to include additional visual aids to improve comprehension of the method.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6xMh3vuhHv", "reviewer": "Reviewer_qC51", "summary": "The paper proposes a new way of computation in a mixture of expert (MoE) layer. Unlike SparseMoE, instead of using the top-K operator and then doing full computation only for that expert, this paper do not do topK, but does a linear combination of all experts, but each expert is now a low rank matrix, which saves compute/parameters. The authors do show some properties that performing MoE computation in such a way can help in interpretability of the experts, their task specialization and so on. This can further help in certain cases like to mitigate biases.", "questions": "- 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.", "limitations": "The authors have identified the limitations of the work.", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- 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.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank **Reviewer-qC51** for their detailed review; for praising the ease with which the methodology can be implemented and the benefits of downstream applications in bias mitigation. We address the stated weaknesses below and draw attention to the existing ablation studies requested. Additional experiments requested to further show the benefits of the interpretability gained are also provided through steering LLMs’ outputs using the experts.\n\n## Comparisons to Sparse/Soft MoE\n\nAs the reviewer correctly states, our sole interest is the independent benefits of scalable expert specialization (these are the claims and paper's goals as stated in [L17, L81, L353]). **muMoEs are parameter-efficient layers for providing interpretability and controllability benefits without increasing the parameter or FLOPs count** of the original networks. We highlight that two reviewers praise this explicitly: \n\n- **Reviewer XreH**: `\"specialization of experts is one of the most important goals in MoE\".`\n- **Reviewer Su68:** `\"µMoE models demonstrate a more parameter-efficient alternative to traditional MoEs, an important consideration for practical applications.\"`\n\nConsequently, we argue that the fair and relevant comparison to validate the paper’s claims is **with a fixed computational budget**: between the original networks with MLP layers and parameter- and FLOPs-matched muMoE layers. This serves to verify that we do not trade-off accuracy when introducing the many benefits of muMoE layers. We do not claim (nor desire) to *increase* parameter counts / capacity / performance. Sparse and Soft MoEs significantly increase the number of parameters and the number of FLOPs, respectively, making a fair comparison impossible.\n\nIn contrast, Soft and Sparse MoEs mostly aim to increase *performance* through increased parameter counts and/or FLOPs. Crucially, fine-grained expert specialization in large models is impractical with both Soft MoEs and Sparse MoEs. For Soft MoEs, a restrictively high FLOP count is required. For Sparse MoEs, the parameter count alone is problematic: even the smallest GPT2 variant with 2048 expert Sparse MoE layers has on the order of 100 *billion* MLP parameters--compared to just ~100 *million* parameters with muMoEs. **For even larger dense models whose base parameter counts are already in the billions, large expert counts with Sparse MoE are infeasible.**\n\n\n## Requested comparisons\n\nInspired by the comments of the reviewer, we've included a new detailed comparison of muMoEs with both soft and sparse MoEs for CLIP fine-tuning (in the rebuttal PDF). This single-layer experimental setup makes it possible to use an otherwise prohibitively large number of experts in Sparse/Soft MoEs, all while maintaining a feasible number of parameters and FLOPs.\nAlthough optimizing the parameter count / FLOPs VS accuracy trade-off is not the aim of the paper, both muMoE variants achieve Pareto-optimal improvements in this aspect for last-layer fine-tuning of CLIP on ImageNET1k.\n\n## Ablations\n\nWe kindly draw the reviewer’s attention to the following figures where **the** **requested ablations** **can** **already** **be** **found in the submission**:\n\n\n- **Performance increase as a function of expert count**: please see Figure 21 of the appendix ([L933]) where performance increases as expert count is increased.\n- **Placement of muMoE layers**: please see Figures 3, 16, 17 for final and penultimate layer muMoEs configurations. We note that Section 4.3’s results are with networks with muMoEs at *every* layer, and results from *all* layers are shown in Figure 11 of the appendix.\n- **Different model sizes**: experiments are already performed throughout on 4 network architectures of various sizes (CLIP, DINO, GPT2, and MLP Mixers).\n- **Gating function**: An ablation study is already presented on [L888] between activation functions.\n\n\n## Additional use cases of experts\n\nAs requested, to further demonstrate the usefulness of the interpretability and experts, we show how in GPT2 models, the experts can be used to **steer the model’s output to particular themes of interest**.\n\nConcretely, we take a GPT-2 model trained with CPmuMoE layers at every MLP layer, each with 2048 experts. We find that intervening in the trained model’s forward pass (concretely: adding particular expert cluster center vectors to every token’s input latent activation vector before applying the muMoE layer), reliably steers the model to produce outputs related to specific themes. Examples are shown for 3 experts in the rebuttal PDF at the top of the page. \n\nChosen experts steer the LLM outputs toward talking about themes such as climate change, police brutality, or programming. We suggest these results constitute further evidence of the usefulness of the muMoE layer in providing controllable generation of LLM outputs (without any additional training needed).\n\n\n## Figure 2\n\nOur experiments are separated into qualitative studies (Section 4.1.1.) and quantitative studies (Sections 4.1.2 and 4.2). The sole purpose of Figure 2 is a qualitative study to visually motivate the benefits of increasing the expert count. Separate quantitative results are found throughout the latter two subsections. Furthermore, the distribution of expert assignments is already visualized in Figure 20 of the appendix.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vcED87vQuJ", "reviewer": "Reviewer_XreH", "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.", "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?\n* 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.", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "* The problem studied in this paper is well-motivated, specialization of experts is one of the most important goals in MoE.\n* The formulated illustration of factorization is clear.\n* Figures and tables are clear and easy to read.", "weaknesses": "* It would be better if more results on language models could be included.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank **Reviewer-XreH** for praising the well-motivated problem studied in the paper and the clear presentation of the factorization methodology.\n\n\n## More results on language models\n\nWhilst we do not have the computational resources to pre-train new additional large language models during the rebuttal, we do provide more LLM results on the GPT-2 language models already trained for **steering the LLM outputs using the experts,** inspired by the reviewer’s comments.\n\nIn particular, we use a GPT-2 model trained with muMoE layers at each MLP layer, using 2048 experts (at every layer). By modifying the forward pass of the trained model—specifically, adding selected expert cluster center vectors to each token's input latent activation vector before applying the muMoE layer—we can consistently control the model to generate outputs aligned with specific themes. Illustrations of this approach, using 3 different experts, are included in the global rebuttal PDF.\n\nThe selected experts guide the language model's outputs toward discussing topics such as climate change, police brutality, or foreign politics. We suggest that these findings further demonstrate the effectiveness of the muMoE layer in facilitating controllable generation of language model outputs—all whilst requiring no additional training.\n\n\n## [Q1] muMoE vs Soft MoE\n\nSoft MoEs perform a dense forward pass taking a linear combination of all $N$ experts’ outputs simultaneously. With large expert counts, this is problematic due to both high parameter counts and restrictively high FLOP costs.\n\nAs detailed in Section 3.1.2, the existing soft MoE is a special case of muMoE when (a) the weight tensors are explicitly stored without decomposition and (b) no implicit fast factorized computation is performed. The key technical contribution of the muMoE layer is to make a similar full dense forward pass (using *all* $N$ experts simultaneously) feasible at high expert counts. In contrast, muMoEs elegantly skirt both issues with Soft MoE: the former (a) through factorization, and the latter issue (b) through the derivations of the fast implicit forward passes. \n\n\n## [Q2] How to find top-activating experts\n\nThe reviewer is correct that we apply no top-k routing at training time, nor at test time. After the model is trained however, we can still rank inputs by their expert coefficients for the qualitative experiments. Concretely, we pass a dataset of inputs through the network and rank images/text-token inputs based on the magnitude of the corresponding expert coefficient (the strength with which each expert is activate). It is through this standard technique that the top-activating images are ranked and displayed in Figures 4 and 5.\n\nAs shown in the “Expert coefficients color map” in the top-right of Figure 5, the color indicates the soft expert coefficient weight (on the interval [0,1]) with which this particular expert processes each text token.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dXRywLOH8P", "reviewer": "Reviewer_YTQ3", "summary": "The paper proposes the use of (factorized) Multilinear Mixture of Experts as an alternative to Sparse MoEs, which are prone to training issues due to the sparse top-k activation. Several factorization options are described and compared, which result in models of better accuracy in vision tasks, when matching the parameter count, compared to models using vanilla dense MLPs. The paper also includes experiments on a language task, where the proposed method is almost as good as the baseline.", "questions": "This is a suggestion, not a question: The term \"Soft MoE\" used to refer to the seminal work of Robert A. Jacobs et al. from 1991 can be confusing with the recent work [\"From Sparse to Soft Mixture of Experts\"](https://arxiv.org/abs/2308.00951) by Puigcerver et al, from last year. Since the term Soft is only used once in the 1991 work, I would suggest using \"Dense MoE\" to refer to it, which is also more opposed to \"Sparse MoE\".", "limitations": "No negative societal impacts are specific to this work, in my opinion.", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 4, "contribution": 1, "strengths": "- 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.", "weaknesses": "- 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\".", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer YTQ3 for engaging thoroughly with the paper. However, we respectfully argue below that the reviewer overlooks our stated goals and claims in the paper, and the benefits to scalable expert specialization the methodology brings that are otherwise infeasible: \n\nIn particular, muMoE layers are a parameter-efficient method for scalable expert specialization. **MuMoEs aim is to provide interpretability and controllability benefits to dense layers without increasing the parameter or FLOP counts** of the original networks. We respectfully highlight that these goals and scope are outlined in the paper’s claims (e.g. [L17, L81, L353]), and that two reviewers praise this explicitly: \n\n- **Reviewer XreH**: `\"specialization of experts is one of the most important goals in MoE\".`\n- **Reviewer Su68:** `\"µMoE models demonstrate a more parameter-efficient alternative to traditional MoEs, an important consideration for practical applications.\"`\n\nConcretely, as detailed from [L344], our experiments comparing model accuracy in Section 4.3 serve only to verify that muMoE layers do not trade-off accuracy relative to existing dense layers of both equivalent parameters and FLOP counts. **Sparse MoEs introduce significantly more parameters; thus, there is no possible fair comparison in this setting.**\n \nIn the paper, we make no claims about competing with Sparse MoEs’ performance—only that our model form avoids Sparse MoEs’ non-differentiability by design. Rather than exploring conditional computation for increasing performance and network capacity, the paper instead has the orthogonal goal of facilitating more interpretable subcomputations (e.g. allowing downstream tasks/properties of manual bias correction, interpretability, or network editing/steering as shown in the rebuttal PDF), without introducing any additional parameters over MLPs.\n\nWe would be grateful if the reviewer could re-assess the work with the paper’s stated goals and claims in mind, and to judge the method through the novelty it brings to scalable expert specialization at no additional computation cost--along with its downstream benefits shown in the paper.\n\n\n## Sparse/Soft MoE comparison results\n\n**We emphasize the significance of muMoE’s parameter-efficient expert specialization:** even at the moderately large expert counts we explore in this paper, Sparse MoEs quickly approach a prohibitively large number of parameters. Even the GPT2-small model with a 2048 expert Sparse MoE has on the order of 100 billion MLP parameters--compared to just ~100 million parameters with CPmuMoEs. Meanwhile, the GPT2-XL model with 2048 expert Sparse MoE layers has on the order of a trillion MLP parameters alone. For even larger dense models whose base parameter counts are already in the billions, the large expert counts necessary for fine-grained specialization are infeasible with Sparse MoEs.\n\nWe accordingly focus on an experimental setup in which it *is* possible to use large numbers of experts in both Sparse and Soft MoEs--without an infeasible number of parameters or FLOPs. In this setting, the computational requirements of all MoEs can be matched fairly. Please find in the global rebuttal PDF (at the top of openreview) a comparison of muMoEs to both Soft and Sparse MoEs for fine-tuning CLIP on ImageNET1k: **whilst it is not the goal of the layer, both variants of muMoEs provide Pareto-optimal improvements to the parameter count vs accuracy frontier.**\n\nFurthermore, as outlined above, the paper’s aim is not to increase accuracy or replace existing MoEs — but rather to replace parameter- and FLOPs-matched dense layers with interpretable subcomputations. This is why we compare dense layers as a baseline and not transformer-specific MoE layers such as the suggested [1].\n\nWhat’s more, we highlight that the MoE and routing strategy in “From Sparse to Soft Mixtures of Experts” [1] is tailored to transformer/token-based architectures, and the method and paper aim at improving transformer model performance. In contrast, the muMoE focuses instead on scalable expert specialization, and in **dense layers** more generally. As demonstrated in the paper, our focus is on facilitating applications such as last-layer fine-tuning models with MoEs (and subsequently removing model bias in Table 2). These tasks addressed in the paper are not possible with the token-specific methodology of [1].\n\n\n----------\n- [1]: Puigcerver, Joan et al. “From Sparse to Soft Mixtures of Experts.” ICLR 2024.\n\n\n## Low rankness at train-time\n\nThe reviewer is correct that Appendix C.2 explores the minimal impact of applying truncated SVD on *pre-trained* weights. However, all the results in Section 4.3 of the main paper **do** validate the ability of the muMoE layer to retain performance with rank constraints **at train time**. As stated on [L769], this is our main evidence to support the claims of the paper.\n\n\n## Benchmarking of speed & memory\n\nPlease find (in the rebuttal PDF) a summary of latency and peak memory usage for various layers. We use the same number of experts for all MoE variants and set the rank of muMoEs to match the parameter count of a linear layer mapping from 768 to 1000 dimensions. Whilst muMoEs do not offer significant speed gains over alternative MoEs, they have an order of magnitude less peak memory usage. We will include a discussion of these results in the revised main paper with the additional page count.\n\n## “Dense MoEs” suggestion\n\nWe thank the reviewer for the helpful suggestion to use “Dense MoE” to refer to the baseline methodology. Not only does such a name accurately describe the properties of this model's forward pass, but it indeed helps distinguish it from the recent ICLR paper. We will gladly adopt this name in the paper.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zIqeuPWkCy", "reviewer": "Reviewer_1sqG", "summary": "In the paper, the author introduces $\\mu$MoE layers, which factorize large weight tensors to facilitate implicit computation, thereby accelerating the computation process. By increasing the number of experts, the model enhances its specialization in vision tasks. Further, the author conducted experiments with the GPT2 and MLP-Mixer models, demonstrating that these models maintain high accuracy and exhibit increased specialization. Both qualitative and quantitative evaluations conducted by the author were commendably thorough and well-executed. The experiment is comprehensive.", "questions": "1. 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?", "limitations": "The authors adequately discuss the limitations and do not have the potential negative societal impact.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. **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.", "weaknesses": "1. **Limited Scope:** The research primarily focuses on vision models and specific datasets, which may limit the generalizability of the findings. Expanding the validation to more diverse models (LLaMA, gemma, and $e.t.c$) and datasets could provide a more comprehensive view of the $\\mu$MoE layers' effectiveness.\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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer 1sqG for their positive assessment of the paper; for praising the “innovative” quantitative metrics, the “comprehensive evaluation”, and “competitive performance with added benefits”. We address the stated weaknesses below:\n\n\n## [W1] Limited model/dataset scope\n\nThe reviewer is correct that we focus primarily on vision models and ImageNET1k. However, we highlight that our experiments also demonstrate the benefits of the layer on multi-modal vision-language models (CLIP and DINO in Figures 3, 14-17, and 21), and LLMs using the large and diverse Open WebText natural language dataset [1].\n\n\n## [W2] Robustness analysis\n\nOur quantitative evaluation of vision models is only on the test set and does not include out-of-distribution (OOD) data, a limitation we have acknowledged in the conclusion.\n\nHowever, in our experiments on language models (as shown in Figures 5, 12, and 13), we use sequences generated by large language models (LLMs) as input instead of the training set data. Additionally, we conducted an extra experiment for Reviewer-qC51 (in the rebuttal PDF), where the input prompt is free-form text, demonstrating initial evidence of the experts' generalizability beyond the training and test data.\n\nWe thank the reviewer for their suggestion to focus explicitly on robustness. We will accordingly include this explicitly in the limitations as a possible avenue for future work.\n\n\n## [W3 + Q1] Need for implementation details\n\nWe highlight that we include the full pseudocode of all muMoE layers in Appendix B, and an anonymous link in the paper's abstract links to our full model source code implementing the layers (please see `MuMoE.py` in particular). \n\nFor full details about the computational requirements, hardware, and configurations used for the experiments, please see Table 6 in the appendix.\n\n\n## [Q2] Large MoE models\n\nWe do not anticipate any particular difficulties applying muMoEs to even larger networks.\n\nMuMoEs are indeed a uniquely attractive alternative at very large scales, unlocking (through their parameter efficiency) a route towards achieving fine-grained specialization that is not possible with Sparse/Soft MoEs due to their prohibitively large parameter counts.\n\nHowever, we unfortunately do not have access to the computational resources to verify this experimentally. As shown in Table 6, we have access to a maximum of 4 A100 GPUs, which limits us to pre-training only the smallest GPT-2 model with 124M parameters. We have added the following to the “limitations” section to acknowledge this:\n\n\n    Despite demonstrating the benefits of scalable expert specialization with muMoE layers for various architectures with different modalities, we have not conducted experiments on the largest scales of commercial LLMs. Future work should explore this where resources permit.\n\n---\n\n- [1] Aaron Gokaslan and Vanya Cohen. Openwebtext corpus, 2019.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We are grateful to each of the 5 reviewers for their thorough comments; for their positive assessments of both the paper and the novelty of the layer in facilitating parameter-efficient, scalable expert specialization:


- **Reviewer Su68** praises the “thorough” mathematical formulation, “technically sound” factorization approach, and `"extensive qualitative and quantitative evidence to support the claims"`. They also highlight the importance of the parameter-efficiency of muMoEs: `"a more parameter-efficient alternative to traditional MoEs, an important consideration for practical applications".`
- **Reviewer qC51** highlights the ease with which the method can be implemented in existing MoE setups and improvements to mitigating bias. 
- **Reviewer XreH** also praises the paper’s motivation (`"specialization of experts is one of the most important goals in MoE".`), and states that the technical formulations of the layer are clear.
- **Reviewer YTQ3** highlights the model performance over MLPs and the detailed appendix and technical details. They also summarize the main paper as thorough and well-written.
- **Reviewer 1sqG** praises the “innovative metric” introduced to quantify the expert specialization, the “comprehensive evaluation”, and the benefit of the layer as providing `"Competitive Performance with Added Benefits"`.

Our initial response addresses all of the stated weaknesses and questions. In the instances where reviewers have asked for experiments that already exist, we note explicitly in which sections they are located.

We highlight **three** new sets of experiments requested by various reviewers (which can be found in the rebuttal PDF):

1. **Comparisons of the muMoE layer to both Soft and Sparse MoEs** for CLIP fine-tuning: the muMoE offers Pareto-optimal improvements to accuracy over both MoE alternatives and dense layers.
2. **Specialized experts’ additional ability to steer the generation of GPT-2 muMoEs**: adding specific experts’ gating vectors to the activations reliably causes GPT-2 models with muMoEs to produce text of specific themes (such as about politics, or the weather).
3. Peak memory usage (PMU) and latency metrics comparisons between MoE layers, where muMoEs have an order of magnitude less PMU than Sparse MoEs.

We are happy to elaborate further on any remarks made in the rebuttal.

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

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

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