Approximating Two-Layer Feedforward Networks for Efficient Transformers
Authors: Róbert Csordás
Kazuki Irie
Jürgen Schmidhuber
Venue: EMNLP 2023 Findings (2023)
Keywords: transformers
moe
mixture of experts
pkm
product key memories
approximate computation
efficient transformers
language modelling
TL;DR: We present a unified view of approximating 2-layer MLPs, which enables our MoEs to match the performance of parameter-equal dense baselines.

Abstract: How to reduce compute and memory requirements of neural networks (NNs) without sacrificing performance? Many recent works use sparse Mixtures of Experts (MoEs) to build resource-efficient large language models (LMs). Here we introduce several novel perspectives on MoEs, presenting a general framework that *unifies* various methods to *approximate two-layer NNs* (e.g., feedforward blocks of Transformers), including product-key memories (PKMs). Leveraging insights from this framework, we propose methods to improve both MoEs and PKMs. Unlike prior work that compares MoEs with dense baselines under the *compute-equal* condition, our evaluation condition is *parameter-equal*, which is crucial to properly evaluate LMs. We show that our MoEs are competitive with the *dense* Transformer-XL on both the WikiText-103 and enwiki8 datasets at two different scales, while being much more resource efficient. This demonstrates that MoEs are relevant not only to extremely large LMs but also to any-scale resource-efficient LMs. Our code is public.

Approximating Two-Layer Feedforward Networks for Efficient Transformers
Róbert Csordás' Kazuki Irie Jürgen Schmidhuber
The Swiss AI Lab IDSIA, USI & SUPSI "Harvard University "AI Initiative, KAUST (robert, jeuuergeeniidsiaa ch, kirie@fas harvard.
Abstract
methods are gaining popularity, principled solu- tion for resource-efficiet neural networks (NNs) remains clusive.
How to reduce compute and memory require- ments of neural networks (NNs) without sac- rificing performance? Many recent works use sparse Mixtures of Experts (MoEs) to build resource-efficient large language mod- els (LMs). Here introduce several novel perspectives on MoEs, presenting general framework that unifies various methods to ap- proximate two-layer NNs (c.g.. feedforward blocks of Transformers), including product- key memories (PKMs). Leveraging insights from this framework, propose methods to improve both MoEs and PKMs. Unlike prior work that compares MoEs with dense baselines under the compute- equal condition, our evaluation condition is parameter-equal, which is crucial properly evaluate LMs. We show that our MoEs are competitive with the dense Transformer-Xi on both the WikiText- 103 and enwiki8 datasets at two different scales, while being much more resource resource-efficient. This demonstrates that MoEs are relevant not only to extremely large LMs but also any-scale resource-efficient LMs. Our code public.
One promising approach explored by several re- cent works on extremely-large LMs is the sparse mixture of experts (MoE; Shazeer et al. (2017); Lewis et al. (2021); Lepikhin et al. (2021); Fedus et al. (2022); Clark al. (2022); Chi et al. (2022)). Unlike their dense counterparts, MoEs only com- pute subset of their activations (i.e, only few experts) at each step, offering reduced computation and memory costs. However, MoEs are not yet generally adopted as a generic/to-go approach, per- haps because of certain common beliefs on MoEs: (1) They are hard to train (involving complex en- gineering tricks to prevent collapsing), (2) they are not competitive against their dense counter- parts with the same number of parameters (in fact, prior work focuses on FLOP-equal comparison, "unfairly" comparing MoEs against dense baselines with many fewer trainable parameters), and finally, (3) they are reserved for extremely large models (they are rarely/never considered to further improve the efficiency of "small" models). Indeed, even prior works on MoF-bassed Transformer LMs only deploy MoEs in few feedforward blocks; while ideally, all such blocks should benefit from replace- ment by Molšs Here we challenge these common beliefs, and propose novel perspectives on MoEs.
1 Introduction
Despite impressive results recently achieved by large language models (LLMs; Radford al. (2019); Brown et al. (2020); Rae et al. (2021)), vast resource requirement remains their obvious limitation. In fact, most existing LLMs, such as GPT-3 (Brown et al., 2020), cannot be trained, fine- tuned or even evaluated without access to enormous compute. Many recent works strive to develop LLMs that, least, enable inference with limited resources (e.g., consumer hardware), e.g., by building "smaller" yet capable LMs (Touvron et al., 2023; Taori et al., 2023; Chiang et al., 2023) or de- veloping post-training quantization methods (Zafrir et al., 2019; Dettmers et al., 2022). While these
We present MoEs within a unified framework of methods that approximate two-layer feed- forward networks, which includes product-key memories (PKMs; Lample et al. (2019)) and top-k sparsification. This principled view not only allows us to conceptually group and compare MoEs with PKMs, it also provides insights on design choices for improving these methods. Our resulting MoE Transformer variant outperforms our improved PKMs, and performs as well as or even outperforms the dense baseline, while using a fraction of its compute for both training
Work done at IDSIA.
'https: t/gpthbb...tttttte http:///ggiiuuuucceeereerrreeeee com/roobertcsooddss/ms

---

and inference. Importantly, unlike prior work, we compare our MoEs with dense baselines with the same number of total trainable parameters, which is crucial for proper evaluation in language modeling. We conduct experiments on the standard WikiText-103 (at two different model scales) and Enwik8 datasets. We demonstrate that MoEs are not limited to extremely-large LMs, but useful as generic approach for resource-efficien NNs at any scale, and in line with the recent trend of improving "smaller" models (Touvron et al., 2023; Taori et al., 2023; Chiang et al., 2023). Finally, we release a CUDA kernel for our MoE layers which allows for achieving faster wall clock time and large memory reduction compared to the dense model.
Alternatively, this layer can be viewed as a key- value memory accessed by attention (Vaswani al. (0017),Gevv et al. (2021)), where keys and values are rows and columns of weight matrices W1 and W
k
W1
(3)
k.
1
W
01
V2
Ddr
(4)
2 Background
where k; Rdhnddd o4 &.amoo for i {1..... dn).
Transformers (Vaswani al., 2017) have two main building blocks: the self-attention layer (Parikh et al., 2016; Cheng et al., 2016; Bahdanau et al., 2015), and the two-layer feedforward, i.e. multi-layer perceptron (MLP) block. Acceleration and memory reduction of the self-attention is rather well explored (see, e.g., linear attention dating back to the unnormalised linear Transformers of 1991 (Schmidhuber, 1991; Katharopoulos al., 2020; Choromanski al., 2021; Schlag et al., 2021)), and very efficient implementations (Dao al., 2022) are also available. In constrast, resource-cfficient MLP blocks are still underexplored. This is our main focus, and it is of particular relevance today, as the proportion of the total parameter counts, compute and memory requirements due to MLP blocks in Transformers is increasing in ever-growing LLMs.
Then, the output is computed as "attention":
de y iy ReLU(k]
de
(5)
where og ReLU(k) E Rz0 the "attention weights." Note that ag u/i] where uli] e IR denotes the i-th component of Rdr in Eq. 1. Unlike the standard self-attention, the MLP block uses a ReLU activation function (instead of the softmax) without scaling.
It has been observed that, in practice only few of the factors k. x are positive (Li et al., 2023; Shen et al., 2023), making the first layer's output, i.e., u, sparse. Concretely, Shen et al. (2023) report that in Transformer with dmodel 256 and dy 1024, 10% of the channels account for 90% of the total activation mass. We confirm this trend in our own preliminary study. Fig. shows the average number of non-zero units u of size dr = 2053 in our 47M parameter dense model trained on WikiText- 103 (we refer to App. A.2 for more details). The number is below 200 for all layers. This suggests that the MLP block can be approximated without significant performance loss. Note that this is also supported by the findings Zhang et al. (2022).
Let dmodel, dr denote positive integers. Each Transformer MLP block consists of one up- projection layer with a weight matrix W1 E dd x dmmolid where typically da = 4dmodel. and one down-projection layer with parameters W2 E aud dd that projects it back to the original size. Non-linearity (typically ReLU) is applied between these two layers. That is, input Pdlma is transformed an output y Bodmo
Approximating 2-layer MLPs
= ReLU (W
(1)
Here we present a unified view on methods approximate 2-layer MLPs (Sec. 2) that includes many existing methods such as MoEs (Sec. 3.3) and PKMs (Sec. 3.2).
= Wzu
(2)
where E Rd and we omit biases (as well as batch and time dimensions) for simplicity.
"Our non-expert CUDA implementation still has much room for further optimization.
Soe the appendix "Two feedforward Layers = Attention over Parameter" in their paper version "arXiv: 1xxxi:17706.03762x33

---

Ache chames
ReLU (Wiæ) is required for Liq. 6. Going beyond this requires to also introduce some approximation to Eq. 6 as in PKMs (Sec. 3.2) and MoEs (Sec. 3.3),
200-
3.2 Product-Key Memories (PKMs)
Figure 1: Number of active channels in in our dense 47M parameter model on WikiText- 103 (out of 2053 total channels). Standard deviation over all tokens of the test and validation set.
Product- Key memories (Lample et al., 2019) con- sist of replacing W1 C RRdexdmui in Eq. 1 by two matrices Wa: Wo E B Smole slices the in- put vector x E Ridmode into two halves, La Tb E Amuital
Preliminaries. Let D E Roddd denote approx- imation of y E Rddndde in Eq. 5. Let yi I madee de- note y; 0, for E [1,...,dii] dut- The core idea is approximate the sum Eq. 5, i.e., y - Eie de by only keeping a subset C (1,..., dr} of the key-value pairs, i.c., D Lies yt- The intu- ition of this approximation is as follows. We as- sume that a good approximation ý of y is the one that minimizes their Euclidean distance e l y112 < I. which can now be expressed as Eies oguill2 where S denotes the comple- ment of i.e. 5 (1,...,d4) dp) SS Since have e e= e= Eies 0,2l S Eies a,||v,l1 (triangle in- equality; where the equality achieved when vs are orthogonal), this upper-bound Eies ax [0,ll can be minimized if each term c; = =a,1/2) 1 e IR are small. If we further assume that all value vec- tors have the same norm, the crucial factor for approximation quality is reduced to the attention weights a;, In this context, we also call a the contribution of key-value pair
R 2 so that x = Talas, where denotes con- catenation. The matrix multiplication is then per- formed on these smaller vectors: U2 = Wa and Wo Then E IRd is calculated by com- bining the elements of Uu E Rvd and us Rvār in all possible ways (i.e., Cartesian products), simi- larly the outer product, but using addition instead of multiplication, i.e., for E [1,..., du).
-==1 y/s
ux [ = un [L V dal] + uali mod v du)
(8)
In addition to applying Top-K at the output as in Sec 3.1, here Top-K can also be used to accelerate the operation above. By applying Top-K uo and u4 before combining them to compute u, only the K2 << dr components of ul]] have to be calculated, and they are guaranteed to contain the K biggest components of the full u.
In the original formulation (Lample et al., 2019), PKMs use a softmax activation function, taking in- spiration from self-attention (Vaswani et al., 2017). Instead, we'll show how non-competing activa- tion function, such as ReLU better choice (see Sec. 6.2).
Let be a positive integer. The general idea of all methods discussed this work is to keep K pairs (k; v,) whose contribution a, is the high- est, and ignore other low-contribution pairs. The goal is to find the best mechanism to select such K pairs. Here we discuss three variants: Top-K K acti- vation (Sec. 3.1), Product-Ker Key Memories (PKMs, Sec. 3.2), and Mixture of Experts (MoEs, Sec. 3.3).
3.3 Mixture of Experts (MoE)
Let Ng, G denote positive integers. MoË's par- tition di pairs of (k,, v)) (see their definition in Sec. 2) into NE groups of size each, such that
G N₁ - dn. This means that the weight matrices
W1 E Iddrdd and W2 E Rdmold xdr (Eqs. 1-2)
3.1 Top-K Activation Function
det
xdmadel
are partitioned into matrices W E I NE and W5 E R Smoes x NE for {1, {1, NE
The most straightforward implementation of the approximation described above the top-K K acti- vation function:
kJ kiGH
k k
E. arg to, K) C (1,....dh) dr)
(6)
Wi
(9)
ÿ Σ a;d;
kJ (er1
Unfortunately this only saves less than half of the entire computation: while this allows us re-
W
V.G+1
VeG+2
(e
duce computation of Eq. 2, no computation can be saved in Eq. because full computation u =
(10)

---

The output is computed as:
do not cover those that make use of reinforcement learning for expert routing) before describing our improved version Sec. 5. We'll review their expert selection function and regularization method, and highlight their key characteristics.
Ú= Σ [1 (Wf)) oC&
(11)
where sfe] E the e-th element of vector 8 E NN computed by an expert scoring func- tion sel Rdmads ++ NE (typically & = sel(x)) = softmax( with W3 E R-Nxxdwec) and & denotes a subset of indices (1,.., NE) result- ing from the Top-K operation on i.e., Ex arg ts, K). Note that in some variants, ad- ditional re-normalization is applied after Top-K, so that Exce.. s(c] 1, scc] 2 O; we define such an operation as norm topk, see its exact defini- tion in App. A.I The efficiency of MoEs comes from the fact that NE < dæ, thus calculating is cheap. Furthermore, G' and K are chosen so that G * K < dn, so the calculation performed by experts is less expensive than the dense MLP.
Sparsely Gated Mixtures of Experts. Shazeer et al. (2017) have revisited MoEs (Jacobs et al.. 1991; Ivakhnenko and Lapa, 1965) with the Top-K K operation, allowing reduction in its resource de- mands. Their method is basically the one described in Sec. 3.3 (with re-normalization after Top-K) except that they use a noisy gating function:
scl(ar) softmax(
Wyæ + (0, softplus(Waa)) (13)
where W4 R-Nxxdmuss the Gaussian noise term AO, 1) is element-wise and independent for each channel, and softplus(x) = (og(1 + e²). They use the following auxiliary regularization term for load balancing.
Given the notation above, is straightforward to see that MoEs can also be viewed as approximating 2-layer MLPs with trainable component (i.e., the selection function sel to produce s). Similarly to Eqs. and 7, Eq. 11 can be expressed as:
L= CV
Σ wcB
norm (spl(s))(z:))
(14)
aed aeciisleluacis (12)
where, compared Eqs. 5 and 7, the "contribution scores" of key-value pair (defined in Sec. 3/Pre- liminaries) have an additional factor s[c] of an expert group e to which the key-value pair belongs.
where CV(x)
d= is the coefficient of variation
and B is the set of all tokens in the batch.
Key characteristics: The scores normalized after the top-K operation (with K 1), which is equivalent applying top-K before the softmax.
The key challenge of MoEs is to learn an ex- pert selection mechanism/ffuntion sel above that assigns high scores only few experts (so that can ignore others without sacrificing performance), while avoiding a well-known issue, called expert collapsing, where only a few experts are used and the rest are never selected. To avoid this, some regu- larization is typically applied to the selection score sel(x), encouraging more uniform routing of ex- perts across the whole batch of tokens. We provide a comprehensive review of MoE variants and their details in Sec. 4 and our improved version in Sec. 5.
Switch Transformer. Fedus et al. (2022) inte- grate the MoË above into the Transformer to obtain their Switch Transformer. In terms of MoE details, one of Fedus et al. (2022)'s key claims is that top-1 routing is enough. Their selection function is sim- ply: sel(æ) softmax W3æ), but they propose hard load-balancion between experts that different hardware accelerators: At most l NN to- kens are allowed to be routed to an expert, where E Rso is the capacity factor (typically between and 1.5), defining how many times more tokens can be processed by one expert compared to the ideal case of uniform routing. Each expert is forbidden to process more than this number of tokens. For regularization, the fraction of the tokens f Ne processed by each expert, and the average selection probability E R Ne for each expert are calculated
4 Existing MoE variants
Several variations of MoEs have been proposed with many different details. Here we briefly review the most popular and representative ones (e.g., we
'In the case of the softmax(-) activation function, this cquivalent to applying Top- K to the logits before softmax.

---

(K 1; top-1 is used) as:
available at each step; Ea. aag max (sel(())) is used instead. Lewis et al. (2021) show that, while during training, the routing is enforced be com- pletely uniform, during the test time, the distribu- tion looks exponential (in fact, this is similar to the Switch Transformer but more balanced for BASE), The algorithm for solving the linear assignment problem (Eq. 19) difficult implement effi- ciently on modern accelerators. Clark et al. (2022) have proposed to use the Sinkhorn algorithm (Sinkhorn, 1964; Sinkhorn and Knopp, 1967) in- stead (resulting in a model called Sinkhorn-BASt or S- S-BASE), to approximate the solution to this problem (note that similar routing is independently discussed by Kool et al. (2021)). They report that this works well, while being simpler implement. Thus, our reimplementation of BASE S-BASE using the Sinkhorn algorithm.
fi
Σ
1i
E arg (op(((ps((()) K)} (15)
B
#CB
p- B
sel(x)
(16)
L = NRF P
(17)
where 1 denotes the indicator function (which is equal to 1 if the argument is true, and otherwise), and denotes dot product. Intuitively, this serves as an adaptive regularization that pemalizes experts that are used often with high "weights." " In addition, they use dropout with a high drop rate (10%) in the experts (but only 10% in the normal layers). Furthermore, Fedus et al. (2022) also propose to initialize the experts with 0.1 As we'll see in Sec. 5, we use a modified version of this scheme. Note that applying Top-K after softmax encour- ages collapsing: if the score of the selected ex- pert is increased, the scores of all other experts are automatically decreased. This is not the case for Shazeer al. (2017): In their method, only the se- lected experts compete with each other, so if their presence beneficial, their score can be increased. Key characteristics: Note that Top-1 I is applied after the softmax without re-normalization.
Key characteristics: During training, Sinkhorn iterations used scores obtain balanced signment. The sigmoid activation is always applied to compute the weighting score.
Overall, all load hul-bblancing methods above are rather complex. We propose simpler but effective approach for MoEs in Sec. 5.
Improving Mixture of Experts
Here we present our improved MoE variant, which we call o-MoE. We conduct thorough ablation stud- ies on our design choices in Sec. 6.
BASE layers and S-BASE. Inspired by the rout- ing strategy and the hard capacity factor of the Switch Transformer, Lewis al. (2021) propose BASE They use top-1 routing and sigmoid activation a in the selection function:
o-MoE Expert Selection Function. Our MoE make use of the top-K operation (unlike BASE). The activation we use on the selection function is sigmoid (as in Eq. 18 of BASE) instead of softmax used in Switch Transformer and Sparsely Gated Mixtures of Experts. This choice is motivated by the view of MoEs approximate 2-layer MLPs (Sec. 3). In fact, sofimax introduces competition between experts. No such competition between channels used in the regular 2-layer MLP (i.e.. there is no constraint on o₂ in Eq. 5). This suggests that, in principle, no competition is needed between terms in the sum of Eq. 12 in the MoE either, to induce sparsity. It is also well known to practitioners that softmax as regular activation negatively affects the trainability of standard MLPs. Softmax combined with top-K can also encourage expert collapsing: when the selection score of one expert increases, the score of the others automat- ically decreases. For all these reasons, we opt for
sæl(æ)) o (W3æ)
(18)
Now instead of using arg topk, they solve the fol- lowing linear assignment problem to find the index
Car C (1,..., N of the expert to which each input x E B is routed,
maximize Ne)
Σ
sc([))]cc]]
(19)
#CB
B
s.t. ViE NE), Σ (cz ==
=
N
#CB
This guarantees uniform assignment of experts, which is efficient for multi-accelerator training. The output is computed using Eq. 11 with & {ex) (a set with single element; "top-1"). How- ever, at inference time, no such balancing is pos- sible because not all tokens of the sequence are

---

sigmoid instead of softmax; we experimentally confirm that this is indeed a good choice.
simple way encourage equal expert usage within the batch and prevent unnecessary overeonfidence in selecting individual experts. Let be the set of all tokens in the batch (counting through both the batch and time dimensions). We introduce the following regularization term L:
Additionally, looking at MoEs in this framework gives us hints on combining them with Top-K K activation (Sec. 3.1) for further acceleration. We can calculate 2e = s(c] Rel. (Eq. 11) for the selected experts and perform an additional Top-K K keep the highest units among them and set the rest zero. We leave this for future work.
softmax W3x)
(20)
B
#EB
NB
o-MoE Initialization. Another design choice guided by the MLP-approximmtion view of Molis (Sec. 3) is the initialization scheme for experts. Typ- ically, experts are assumed to be independent, and the standard deviation of the initialization (Glorot and Bengio, 2010; He et al., 2015) of W is caleu- lated based on G instead of dø. Our experiments in Sec. 6.3 show that this is sub-optimal.
Σ
ple] log
(21)
Furthermore, we propose to randomly drop com- plete experts, during training; we refer to this as expert dropout. Unlike the standard dropout on the activation level, we do not apply rescaling, i.c.
(W, @m if training otherwise
In contrast, we initialize all weight matrices identically to the pre-layernorm dense baselines, not taking in account the smaller size of the indi- vidual experts, i.c., Wi ~ NOO, dmodd Wayyn
scl(w))
(22)
(oW,x)
2 and
where m E {0. [0, m ~ Bernoulli(1 - 8), where 8 is the dropout rate, and @ is the element- wise product. This prevents the dropped experts from being selected while not affecting the other ones. Intuitively, when an expert dropout removes a popular expert, it forces the less popular ones to take over. Thus, the chance of them being trained and improved increases. We experimentally show that our regularization method (Eq. 21) and expert dropout (Eq. 22) are both effective despite their simplicity.
2 where Talayer denotes the number of layers, using dmodel and dr instead of G.
W ~ NN,, V dg-nuuyess
We also take special care when initializing W3 of the selection function. We initialize to a normal distribution with the same standard devi- ation as Wf. but we also ensure that the rows of W3 have the same norm². This can be easily achieved in practice by initializing the weights to W/3 (0, 1), rescaling its rows to norm and then rescaling the whole matrix again to have the desired standard deviation. Note that each scalar score in & is the dot product of a row of W3 and æ. This initialization method ensures that only the an- gle between æ and the rows of W3 initially affects the score 8, rather than an additional random factor resulting from initialization.
Experiments
Our experimental setup is based on Dai et al. (2019)'s Transformer XL with some modifications: we use pre-layer norm and reduce the number of training steps to 100k to reduce the computational budget. Also, to match the parameter counts be- tween the baseline and MoEs, we slightly modify the hyperparameters of the baselines (Dai et al., 2019). fact, our MoE CUDA kernel can only work with dimensions divisible by 4. We round the original sizes up to the next suitable number, e.g., we change dmodel of our 47M parameter WikiText- 103 model from the original 410 to 412. Further- more, since MoEs require extra parameters for the expert selection function, we compensate for these by increasing the dr of the baseline model to match the number of parameters. Our modified baseline model on Enwik8 still has 41M parameters and performs similarly to the original Transformer XL
o-MoE Regularization. As already noted in Sec. 4, existing regularization methods for load- balancing are complex (e.g., Switch Transformers need to deal separately with the actual selection distribution and the scores, Sparsely Gated Mixture of Experts needs noise in the selection function). In contrast, we propose to simply maximize the entropy of the selection distribution p NE cal- culated across the entire batch. Intuitively, this is a
'Having rows with different norms would discourage the use of experts corresponding to rows with small norms, as thoir selection score would be low even if the angle of the selector (row of W3) fully aligns with æ.

---

(see Tab. 1). For WikiText-103, use subword units (Sennrich et al., 2016) using SentencePiece tokenizer (Kudo and Richardson, 2018) instead of the word-level vocabulary, to avoid extra tricks re- quired to reduce the parameter count and compute requirement resulting from the huge vocabulary size. On WikiText- 103, we consider two different model sizes: a 47M-parameter one (denoted by "WT-S" for "small")), and 262M-parameter one ("WT-B" for "big"). We refer Enwik8 "E8" in certain tables. For more details, see Appendix B. For all the methods considered, we use them in every MLP block of the model, which is not common practice in the literature. Typically, MoE (or other approximation methods) is used only once every n layer even only in one layer. This is not satisfactory since our goal is find generally ap- plicable method that can accelerate all layers across the whole model. Moreover, this amplifies the dif- ference between different methods, helping better illustrate effects of each of the design choices.
6.2 Product-Key Memory (PKM)
Our view of Sec. 3 suggests using a non- competitive activation such as ReLU instead of the softmax used in the original PKM (Lample et al., 2019). Our experiments confirm the benefits of this choice (Tab. 2): the performance of the ReLU variants is much closer to the dense baseline (see also related findings in Shen et al. (2023)). But even the best PKM models underperform the dense baselines, indicating the fundamental limitation of PKMs. Note that, as stated above, we conduct careful comparison between the approximation method (here, PKM) and the dense baseline using the same number of parameters. For more results and details on PKM, we refer to App. A.3.
Table 2: Performance of the parameter- parameter-matched PKM models. We provide more results Appendix/Tab. 6.
Variant
Nonlin
WT-S
WT-B
E8
Dense Baseline
ReLU
11.81
9.46
1.08
PKM
Softmax
13.96
11.10
1.16
6.1 Top-K
ReLU
12.77
9.98
1.11
We first evaluate the Top-K method (Sec. 3.1). This standalone evaluation is important as Top-K is the basis of both the PKM and the MoE approximations. Tab. shows the results. We observe that not only Top-K in the MLP blocks preserves the performance of Transformers, it even improves performance. We hypothesize that these improvements are due the reduction in feature interference as described by Elhage et al. (2022). However, we obviously can not arbitrarily reduce K: there should be trade-off between the denois- ing effect and the capacity of the network. Here, the optimal value find is K 128 K = 512.
6.3 Mixture of Experts (MoE)
Here we evaluate our a-Moßi models (Sec. 5) on Enwik8 and WikiText-103 well two additional datasets, C4 (Raffel al., 2020) and the newly proposed peS2o (Soldaini and Lo, 2023). Given the large sizes of C4 and peS2o, we cannot afford to train for a full epoch; we train for 100k steps with the same hyperparameters as for WikiText-103,
Main results. Tab. 3 shows the main results. Our o-MoE models match the performance of their parameter-equal dense baselines, while achieving significant memory and compute reduction. These models use K = 4 for N₁ 16 NE: 32, which a "moderate" level of sparsity but already offering significant compute reduction as shown in the column "% FLOPs"; concrete compute and memory reduction further shown in Fig. 2 (see Appendix A.5 for details). Naturally, there is a limit on the minimum sparsity level to preserve good performance of MoEs, which is determined by several factors. First, empirically find that experts with group size of G < 128 generally degrades performance. Second, our benchmarks with the Top-K operation (Tab. 1) and our ablations (Tab. 10 in the Appendix) show that
Table I: Effects the top-k activation function on the perplexity (WikiText-1033 and bits/character (Enwik8).
Dataset
#params
dp
bpe/perplexity
Enwik8
41M
2053
1.08
41M
2053
128
1.07
41M
2053
256
1.08
41M
2053
512
1.08
WikiText 103
47M
2053
11.81
47M
2053
64
11.86
47M
2053
128
11.74
47M
2053
256
11.74
47M
2053
512
11.68
WikiText 103
262M
4110
9.46
262M
4110
128
9.26
262M
4110
256
9.34
262M
4110
512
9.36
the minimum number of simultaneously active

---

Table 3: Performance of parameter-batchedt o-MoEs perplexity (WikiText-103, C4 and peS2o) and bits/character (Enwik8), Ours matches or surpasses the performance the dense baselines across all datasets.
Time (ms)
MLP
Memory (GB)
200
MoE
0.0
8
32
64
126
256
Number of experts (Nw)
Dataset
Model
#params
% FLOPPs
bpc/ppl
Enwik8
Dense
41M
100.0%
1.08
Figure 2: Execution time and memory usage of forward-hackward pass a single MLP and MoE layer. batch size 64 and sequence length 512, dmodel 512, K = 4, and da G NN.. Full/dashed lines show the execution time/memory consumption, respectively. As they are both linear with similar slopes, they are almost indistin- guishable. Even our sub-optimal CUDA kemne faster starting from 16 experts. Measured on an RTX 3090 with PyTorch 2.01 and CUDA 11.
a-MoE
41M
25.0%
1.08
[3] = 32768, corresponding
WikiText-103
Dense
100.0%
11.81
o-MoR
47M
25.0%
11.71
WikiText-103
Dense
262M
100.0%
9.46
o-MoRR
262M
12.5%
9.44
C4
Dense
47M
100.0%
23.76
d-MoR
47M
25.0%
23.25
C4
Dense a-MoE
262M
100.0%
17.79
262M
12.5%
17.46
peS2o
Dense
47M
100.0%
14.34
o-MoE
47M
25.0%
14.12
still remains significant (10.37 vs 10.03), unlike with the MoE with moderate sparsity. This indi- cates the importance of controlling MoE sparsity to preserve its performance against the dense baseline.
PeS2o
Dense
262M
100.0%
10.91
o-MoE
262M
12.5%
10.91
channels G K need to be above a certain critical threshold (usually around 256-512). Finally, we match the number of parameters of the bascline model; this is the last constraint. Under these constraints, we find that the performance of the dense baselines can be matched using 25% of the required FLOPs and memory for activations for our small models, and 12.5% sparsity for the big one (note that FLOPs here do not take into account the linear projection used to select the experts, which is negligible within the range of NE used here).
Comparison to Existing MoEs. We also com- pare our o-MoE to other MoE variants (Sec. 4), namely Switch Transformer (Fedus et al., 2022), S-BASE (Clark et al., 2022) and the basic softmax variant. Tab. 4 shows the results for multiple vari- ants WikiText- 103 and Enwik8. Additionally, compare o o-MoE to the most important baselines on C4 and peS2o in Tab. 5. As Switch Transformer and S S-BASE select only one single expert (K 1), we increase the expert size by a factor of 4 (instead of G = 128 in our models, we use = 512), and we decrease Ng by the same factor for fair com- parison in terms of the parameter count. Neither of them uses our proposed expert dropout. For Switch Transformer, we test variant with stan- dard dropout in the experts (see App. for details), and a version without. We also extend S S-BASE K = 4, which is similar to ours, except for the bal- ancing method. Even considering all these cases, our g-MoE outperforms Switch Transformer and S-BASE. Note that in terms of FLOPs and memory usage, all MoF. variants are equivalent given the same hyperparameters (G, dmodet. and K).
Increasing Ng and Impact of Sparsity. The results above demonstrate that our o-MoEs can be configured to match the desired performance with fewer resources. Here we conduct an extra experiment where we naively increase NE (while keeping K 1) from 16 to 128. This increases the number of parameters to 238M, while keeping the speed and memory requirements comparable the original model (column "WT-S*" in Tab. 4). This model achieves a test perplexity of 10.37, which is worse than 9.46 of the 262M dense model (see Tab. 1). Indeed, even when the parameter count is matched, there are other bottlenecks that are cru- cial, e.g., here dmmodel is much smaller (412 vs 1024). We construct another dense baseline by setting ev- ery hyperparameter like in the 47M model, except dr. which set 16480 to match the number of parameters of the Ng 128 MoF. This baseline achieves perplexity of 10.03: thus, the gap be- tween the scaled-up MoE and its dense counterpart
Ablation Studies. Finally we conduct ablation studies of individual design choices (Sec. 5). Tab. 4 shows the results. Standard dropout instead of ex- pert dropout leads to performance degradation for most of the cases, except the model with Ng = 128 experts. The softmax-based selection functions
"Unlike the original ones, our implementation does not enforce capacity factor-based hard balancing.

---

Table 4: Ablation studies. WT-S* obtained by naively
scaling NE in WT-S. More details in Sec. 6.3 & Tab. 10.
Solection proporioo
Dataset
WT-S
WT-5"
WT-B
F8
# params. (in M)
47
238
262
41
10-4
M-MoE softmax (no renorm.)
MoM
MoE softmax (renerm)
Switch Transformer no dropout
12.27
11.24
9.68
10-1
Switch Transformrr
11.88
11.10
9.77
1.10
G-128)
S-BASE (K=4, C=128)
13.01
10.96
10.50
16
82
48
64
80
96
112
1
12.32
11.31
9.77
1.32
Expert
a-MoE (K=1, G=128)
11.59
10.37
9.44
1.08
Figure 3: The total proportion of selection weights signed to given expert (x-axis; sorted by their popu- larity) on the validation of Wikitext-100 103 using the WT-S* models from Tab. 4. This is for one representa- 1.08
standard dropout
12.01
10.27
9.53
1.08
softmax (renorm.)
11.99
11.27
9.58
1.09
softmax (no renorm.)
12.05
10.54
9.62
standard init
11.00
10.59
9.67
regularization
11.83
10.41
9.51
1.08
tive layer ("Layer 5" similar plots for other layers are shown in Fig. in the appendix). The models with poor performance (Tab. 4), i..., Switch Transformer (red) and a-MoE with softmax and renormalization "softmax (renom.)" (green) be easily identified. In contrast, the fine performance differences between the rest of the models not seem be due expert collapse.
11.63
10.30
9.58
256
11.84
10.44
9.56
1.09
1,
11.90
10.83
9.58
Table 5: Perplexity of o-MoE compared to parameter- matched baselines on and peS2o datasets.
Dataset
C4
C4
PeS2o
PeS2o
dmodel
412
1024
412
1024
params
47M
262M
47M
262M
can be easily identified: they severely suffer from the expert collapse problem. The statistics are
K
Dense
128
23.76
17.79
14.34
10.91
rather similar for all other models; the fine per- formance differences among these models do not seem to be due to expert collapse. Remarkably, our entropy-regularizze models with expert dropout, especially o-MoLE, are capable of matching the ex- pert usage balancing of S-BASE without using the Sinkhorn activation function. Note that in general, we do not consider uniform expert activation to be optimal: we expect expert specialization, and thus the frequency of their usage should depend on the occurrence of the task they are performing.
o-MoE
128
4
23.25
17.46
14.12
10.91
Switch
512
24.47
18.29
14.74
11.56
S-BASE
128
35.48
18.53
16.61
11.72
(with and without (--e--oormalization)) consistently perform worse than our sigmoid one. The same is true for the standard initialization: ours is better. Interestingly, removing all regularization methods degrades performance but does not entail catas- trophic collapse even with Ng 128. We also examine the best (G, K) combinations, given a con- stant number (G K) of active pairs k; v;; we find a high K = 4 works best within this range. Further analysis of our a-MoE can be found in App. A.4.
Conclusion
Our novel view unifies methods that approximate 2-layer MLPs, such as Top-K, Mixture of Experts (MoE) and product-key memory (PKM) methods. While Top-K by itself provides limited performance improvements and speedups, further speedup requires PKM or MoE. A non-competitive activation function inspired by our unified view improves both PKM and MoE. Further novel enhancements of MoEs yield our a-MoE which outperforms existing MoEs. Importantly, our o-MoE with moderate sparsity matches the performance of parameter-equal dense baselines while being much more resource-efficient Our new insights improve the training of language models with limited hardware resources, making language modeling research more accessible.
Analyzing expert utilization. A typical failure mode of MoEs is expert collapse, where only few experts are used while others are completely ignored or underused. Here we conduct an analysis to evaluate whether various models including ours are affected by this issue. For each layer, we com- pute the proportion of the expert selection weights assigned to each expert (s))(x)) on the entire validation set of WikiText-103. We use WT-S* models from Tab. 4 with 128 experts. A represen- tative layer is shown in Fig. 3. Models with poor performance (see Tab. 4), i.e., Switch Transformer (red) and "bad" variant of o-MoE with softmax and renormalization "softmax (renom.)" (green),

---

Limitations
ful for hardware donations from NVIDIA and IBM. The resources used for this work were partially pro- vided by Swiss National Supercomputing Centre (CSCS) project s 154, s s1205 and d123. Finally, we would like to thank Dániel Berényi for his support with the CUDA kernel development.
Our experiments show that if we naively increase the number of experts, the performance gap be- tween MoE models and their dense counterparts in- creases. This indicates the need for careful control of sparsity and challenge for MoEs.
which remains
Our CUDA kernel is sub-optimal and I/O limited. However, even in its current form, already yields significant performance boosts and memory reduction. We expect that an expert CUDA programmer could improve the speed of our kerne by at least factor of 2.
References
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Ben- gio. 2015. Neural machine translation by jointly learning to align and translate. In Int. Conf. on Learn- ing Representations (ICLR), San Diego, CA, USA.
Tom B Brown et al. 2020. Language models are few- shot learners. In Proc. Advances in Neural Informa- rion Processing Systems (NeurIPS), Virtual only.
We do not consider load balancing between hard- ware accelerators is done in Switch Transformers and S-BASE. Our goal is make larger model fit single accelerator, or multiple accelerators in the standard data-parallel training. Our preliminary experiments suggest that such balancing entails a performance hit.
Jianpeng Cheng, Li Dong, and Mirella Lapata, 2016. Long short-term memory-n for machine read- ing. In Proc. Conf. on Empirical Methods in Natu- ral Language Processing (EMNLP), pages 551-561, Austin, TX, USA.
We could not reproduce the 277M Enwik8 model of Dai et al. (2019), because we could not fit the beaseline model any of our machines. We tried to use rotary positional encodings with PyTorch 2.0's memory-efficiiee attention reduce it's memory consumption; however, this resulted in a significant performance degradation (even for the smaller models).
Zewen Chi, Li Dong, Shaohan Huang. Damai Dai. Shuming Ma, Barun Patra, Saksham Singhal, Payal Bajaj, Xia Song. Xian-Ling Mao, Heyan Huang, and Furu Wei. 2022. On the representation collapse of sparse mixture of experts. In Proc. Advances in Neu- ral Information Processing Systems (NeurIPS), New Orleans, Louisiana, USA.
Wei-L.in Chiang, Zhuohan Li, 7i Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang. Yonghao Zhuang, Joseph E. Gonzalez, lon Stoica, and Eric P. Xing. 2023. Vicuna: An open- source chatbot impressing gpt-4 with 90%* chatgpt quality.
Our study focuses on end-to-end trainable MoEs. Other MoE methods (Irie et al., 2018; Li et al., 2022) that pre-train LMs on disjoint data, to recombine them later into a single model, are out-of-scope.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamás Sarlós, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Be- langer. Lucy J. Colwell, and Adrian Weller. 2021. Rethinking attention with performers. Ini. Conf. on Learning Representations (ICLR), Virtual only.
Our study only considers standard Transform- ers; however, similar acceleration methods are of utmost importance for shared- layer Transformers, such as Universal Transformers (Dehghani et al., 2019) and NDRs (Csordás et al., 2022). In fact, layer sharing dramatically reduces the number of parameters. Compensating for this by naively in- creasing dmooel or du results in prohibitively high memory overhead and slow execution. In contrast, MoEs allow increasing the number of parameters without such dramatic drawbacks. We leave shared- layer MoEs for future work.
Aidan Clark, Diego Las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini, Jordan Hoffmann, Bog- dan Damoc, Blake A. Hechtman, Trevor Cai, Se- bastian Borgeaud, George van den Driessche, Eliza Rutherford, Tom Hennigan, Matthew J. Johnson, Albin Cassirer, Chris Jones, Elena Buchatskaya, David Budden, Laurent Sifre, Simon Osindero, Oriol Vinyals, Mare' Aurelio Ranzato, Jack W. Rae, Frich Elsen, Koray Kavukcuoglu, and Karen Simonyan. 2022. Unified scaling laws for routed language mod- cls. In Proc. Int. Conf. on Machine Learning (ICML).
A Acknowledgements
This research was partially funded by ERC Ad- vanced grant no: 742870, project AlgoRNN, and by Swiss National Science Foundation grant no: 200021 200021_192356 192356, project NEUSYM. We are thank-
Róbert Csordás, Kazuki Irie, and Jürgen Schmidhuber. 2022. The neural data router: Adaptive control flow in transformers improves systematic generalization. In Int. Conf. on Learning Representations (ICLR). Virtual only.

---

Zihang Dai, Zhilin Yang. Yiming Yang, Jaime Car- bonell, Quoe Le, and Ruslan Salakhutdinov. 2019. -xa Attentive language models beyond a fixed-length context. In Proc. Association for Com- putational Linguistics (ACL), pages 2978-29888 Flo- rence, Italy.
Kazuki Irie, Shankar Kumar, Michael Nirschl, and Hank Liao. 2018. RADMM: Recurrent adaptive mixture model with applications to domain robust language modeling. In Proc. IEEE Int. Conf. on Acoustics, Speech and Signal Processing (ICASSP), pages 6079- 6083, Calgary, Canada.
Tri Dao. Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. 2022. Flashattention: Fast and memory-efficient exact attention with io-awareness. In Proc. Advances in Neural Information Processing Systems (NeurIPS), New Orleans, Louisiana, USA.
Aleksei Grigorievitch Ivakhnenko and Valentin Grig- orévich Lapa. 1965. Cybernetic Predicting Devices. CCM Information Corporation.
Robert A. Jacobs. Michael I. Jordan, Steven J. Nowlan, and Geoffrey E. Hinton. 1991. Adaptive mixtures of local experts. Neural Compututaion, 3(1):79-87.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals. Jakob Uszkoreit, and Lukasz Kaiser. 2019. Universal Transformers. In Int. Conf. on Learning Representu- tions (ICLR), New Orleans, LA, USA.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pap- pas, and François Fleuret. 2020. Transformers are mns: Fast autoregressive transformers with linear attention. In Proc. Int. Conf. on Machine Learning (ICML), volume 119, pages 5156-5165, Virtual Only.
Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zeulemoyer. 2022. Llm.int8(: 8-bit matrix multipli- cation for transformers at scale. In Proc. Advances in Neural Information Processing Systems (NeurIPS), New Orleans, Louisiana, USA.
Diederik P. Kingma and Jimmy Ba. 2015. Adam: A method for stochastic optimization. Int. Conf. on Learning Representations (ICLR), San Diego, CA. USA.
Wouter Kool, Chris J Maddison, and Andriy Mnih. 2021. Nelson Elhage. Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravec, 7ac Hadfield-Dodds, Robert Lasenby, Dawn Drain, Carol Chen, Roger Grosse, Sam McCandlish, Jared Kaplan, Dario Amodei, Martin Wattenberg, and Christopher Olah. 2022. Toy models of superpo- sition. Transformer Circuits Thread.
Unbiased gradient estimation with balanced assign- ments for mixtures of experts. In I (Still) Can Believe It's Not Better Workshop, NeurlPS, Virtual Only.
William Fedus, Barret Zoph. and Noam Shazeer. 2022. Switch transformers: Scaling trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research (JMLR), 23(1):5232- 5270.
Taku Kudo and John Richardson. 2018. Sentencepiece: A simple and language independent subword tok- enizer and detokenizer for neural text processing, In Proc. Conf. on Empirical Methods in Natural Lan- guage Processing (EMNLP), pages 66-71, Brussels, Belgium.
Jerry& A Fodor and Zenon W Pylyshyn. 1988. Connec- tionism and cognitive architecture: A critical analysis. Cognition, 28( 1812331377
Guillaume Lample, Alexandre Sablayrolles, Aureloo Ranzato, Ludovic Denoyer, and Hervé Jégou. 2019. Large memory layers with prod- uct keys. In Proc. Advances in Neural Information Processing Systems (NeurIPS), pages 2546-8557, Vancouver, Canada.
Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. 2021. Transformer feed-forward layers key-value memories. In Proc. Conf. on Empirical Methods in Natural Language Processing (EMNLP). pages 5484-5495, Punta Cana, Dominican Republic.
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yamping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. 2021. Gshard: Scaling giant models with conditional com- putation and automatic sharding. Int. Conf. on Learning Representations (ICLR), Virtual only.
Xavier Glorot and Yoshua Bengio. 2010. Understanding the difficulty of training deep feedforward neural networks. In Proc. Int. Conf. on Artificial Intelligence and Statistics (AISTATS), pages 449-256, Sardinia, Italy.
Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. 2021. BASE layers: Simplifying training of large, sparse models. In Proc. Int. Conf. on Machine Learning (ICML), volume 139, pages 6265-6274, Virtual only.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. 2015. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proc. IEEE Int. Conf. on Computer Vision (ICCV), pages 1026- 1034, Santiago. Chile.
Margaret Li, Suchin Gururangan, Tim Dettmers, Mike Lewis, Tim Althoff, Noah A Smith, and Luke 7etle- moyer. 2022. Branch-train-merg parallel training of expert language models. Preprint arXiv:2208.03300
Dieuwke Hupkes, Verna Dankers, Mathijs Mul, and Eliia Bruni. 2020. Compositionality decomposed: How do neural networks generalise? Journal of Artificial Intelligence Rexearch, pages 757-795.
Embarrassingly

---

Zonglin Li. Chong You, Srinadh Bhojanapalli, Daliang Li, Ankit Singh Rawat, Sashank J. Reddi, Ke Ye, Felix Chern, Felix Yu, Ruiqi Guo, and Sanjiv Kumar. 2023. The lazy neuron phenomenon: On emergence of activation sparsity in transformers. In Int. Conf. on Learning Representations (ICLR), Kigali, Rwanda.
Zhou, Wei Li, and Peter J. Liu. 2020. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research (JMLR), 21:140:
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. 2021. Linear transformers are secretly fast weight programmers. In Proc. Int. Conf. on Machine Learn- ing (ICML), volume 139, pages 9355-9366, Virtual only.
Peter Pagin and Dag Westersthhl. 2010. Compositional- ity I: Definitions and variants. Philosophy Compass, 5(3):250-264.
Ankur P. Parikh, Oscar Täckström, Dipanjan Das, and Jakob Uszkoreit. 2016. A decomposable attention model for natural language inference. In Proc. Conf. on Empirical Methods in Natural Language Process- ing (EMNLP), pages 2249-2255, Austin, TX, USA.
Jürgen Schmidhuber. 1991. Learning to control fast- weight memories: An alternative to recurrent nets, Technical Report FKI-147-91, Institut für Informatik, Technische Universität München.
Rico Sennrich, Barry Haddow, and Alexandra Birch. 2016. Neural machine translation of rare words with subword units. In Proc. Association for Computa- tional Linguistics (ACL.), pages 1715-1725, Berlin, Germany.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf. Edward Yang. Zachary Vito, Martin Raison, Alykhan Te- jani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. 2019. Pytorch: An imperative style, uigh-eerformance deep learning library. In Proc. Advances in Neural Information Pro- cessing Systems (NeurIPS), pages 8024-8035, Van- couver, Canada.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quae Le, Geoffrey Hinton, and Jeff Dean. 2017. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In Int. Conf. on Learning Representations (ICLR), Toulon, France.
Kai Shen, Junliang Guo, Xu Tan, Siliang Tang, Rui Wang, and Jiang Bian. 2023. A study on relu and softmax in transformer. Preprint (tti:2302.004661
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. 2019. Language models are unsupervised multitask learners.
Richard Sinkhorn. 1964. A relationship between arbi- trary positive matrices and doubly stochastic matri- ces. The annals of mathematical statistics, 35(2):876 879.
Jack W. Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, H. Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susan- nah Young. Eliza Rutherford, Tom Hennigan, Ja- cob Menick, Albin Cassirer, Richard Powell, George van den Driessche, Lisa Anne Hendricks, Mari- beth Rauh. Po-Sen Huang, Amelia Glaese, Jo- hannes Welbl, Sumanth Dathathri, Saffron Huang. Jonathan Uesato, John Mellor, Irina Higgins, Antonia Creswell, Nat McAleese, Amy Wu, Erich Elsen, Sid- dhant M. Jayakumar, Elena Buchatskaya, David Bud- den, Esme Sutherland, Karen Simonyan, Michela Pa- ganini, Laurent Sifre. Lena Martens, Xiang Lorraine Li, Adhiguna Kuncoro, Aida Nematzadeh, Elena Gribovskaya, Domenic Donato, Angeliki Lazaridou, Arthur Mensch, Jean- Baptiste Lespiau, Maria Tsim- poukelli, Nikolai Grigorev, Doug Fritz, Thibault Sot- tiaux, Mantas Pajarskas, Toby Pohlen, Zhitao Gong, Daniel Toyama, Cyprien de Masson d' Autume, Yujia Li. Tayfun Terzi, Vladimir Mikulik, Igor Babuschkin, Aidan Clark, Diego de Las Casas, Aurelia Guy, Chris Jones, James Bradbury, Matthew J. Johnson, Blake A. Hechtman. Laura Weidinger, lason Gabriel, William Isaac. Edward Lockhart, Simon Osindero, Laura Rimell, Chris Dyer. Oriol Vinyals. Kareem Ayoub, Jeff Stanway, Lorrayne Bennett, Demis Hassabis, Ko- ray Kavukcuoglu, and Geoffrey Irving. 2021. Scaling language models: Methods, analysis & insights from training gopher. Preprint arXiv:21
Richard Sinkhorn and Paul Knopp. 1967. Concerning nonnegative matrices and doubly stochastic matrices. Pucific Journal of Mathematics, 21(3)3433-33 348.
Luca Soldaini and Kyle Lo. 2023. peS2o (Pretrain- ing Tefficienlly on S2ORC) Dataset. Technical re- port, Allen Institute for AI. https://github. com/ allenai/pes2o.
Rohan Taori, Ishaan Guirajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. 2023. Stanford alpaca: An cearuuctio-followiii llama model. https:/ gi thub. wtt../.cooaaaaaa coon/tatsu-lall wttpsssuuaaa//staaaaaaaa stanford_alpaca.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière. Naman Goyal, Eric Hambro, Faisal Azhar, Aurélien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. 2023. Llama: Open and efficient foundation language models. Preprint arXiv:2302.13977.
Ashish Vaswani. Neam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez. Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Proc. Advances in Neural Information Processing Systems (NIPS), pages 5998-6098999 66088, Long Beach, CA, USA.
Colin Raffel, Noam Shazeer, Adam Roberts, Kather- ine Lee, Sharan Narang, Michael Matena, Yanqi

---

Ofrr Zafrir, Guy Boudoukh. Peter Izsak, and Moshe Wasserblat. 2019. QSBERT: quantized 8bit BERT In Workshop on Energy Efficient Machine Learn- ing and Cognitive Computing NeurlPS, Vancouver, Canada.
Activ
+30
Zhengyan Zhang, Yankai Lin, Zhiyuan Liu, Peng Li, Maosong Sun, and Jie Zhou. 2022. MoFfication: Transformer feed-forward layers are mixtures of ex- perts. In Proc. Findings of the Association for Com- putational Linguistics (ACL), pages 877-890, Dublin, Ireland.
Figure 4: Number of active channels in u in our dense 262M parameter model Wikitext- 103. du 4110 for this model, so the sparsity is below ~ 5%. Standard deviation over all tokens of the test and validation set.
A Further details and analyses
A.1 Definition of normalised Top-K
200
Ad Adow-crarees
Using the setting of Sec. 3.3, we define the normal-
630
ized top-K K operation as follows:
En arg (o K)
(23)
s(i] if E En
Uxper
(op](s)[[]
(24)
0
otherwise
Figure 5: Number of active channels in u in our dense
41M parameter model on Enwik8. dg = 2053 for this model, thus the sparsity is below 15%. Standard deviation over all tokens of the test and validation set.
(s)
norm (8p = E; [t
(25)
A.2 Measuring the Number of Active Channels in u
divide the input vectors into two sub-keys without an additional projection. Finally, unlike Lample et al. (2019), we use the same learning rate for all parts of the network.
In order to explore whether (k; u.) sparsity- based approach is feasible, measure the number of nonzero entries in the up-projected vector in our baseline models (which, because of the ReLU activation function, is the same as the positive en- tries). We show the results of our 47M model in Fig. 1. Note that dn 2053 (See Tab. 8) for the same model, which means that on average only 1-10% of the channels are active. We show the same analysis for the 262M model in Fig. 4. Interestingly, the counts remain the same, even though dæ 4110 for this model. The 41M parameter model En- wik8 shows a stark difference in the distribution of the channels between layers; see Fig. This suggests that the key factor determining the count distribution is the dataset, and the size of the model plays only secondary role. Fortunately, the spar- sity is very high for all models considered.
In addition to the parameter-equal comparison of Sec. 6.2, there is another possibly "fair" way of setting the size of the PKM-based model: match the number of values (this would result in fewer pa- rameters because of the key approximation), even though Elhage et al. (2022) suggest that the keys typically play a vital role, and reducing their capac- ity will cause performance loss. See Tab. 6 for the corresponding results. Note that, for Enwik8 and Wikitext-103 small, the parameter-equal set- ting increases the number of sub-keys from 46 to 62 (2166 vs. 3844 values). This helps significantly.
A.4 Further Analyses of Our a-MoE
We also examine the best (G, K) given a constant number (G K) of active pairs k; v; In this setting, reducing K by factor of m (K' = K) involves increasing G (G' mG), which, for a constant number of parameters, reduces NE to No Ng.
A.3 More Details and Results on PKM
Our PKM (Sec. 3.2) is based on Lample et al. (2019) with the following basic modifications. First, we do not use batch normalization (BN). As Lample et al. (2019) shows that BN is only benefi- cial for models with very large memory size, we remove as it simplifies inference where the effec- tive batch size varies over time. Also, we directly
The results can be seen in the 2nd block of Tab. 10. We find that higher K is beneficial. Given this, we ask the question how the selection distribution of the models with K > 1 is different from selecting the same experts together and acting as a larger

---

Table 6: The performance of the PKM model variants, Both value-count and parameter matched variants are shown. Additionally, show the effect of the initialization inspired by unified view, which is marginal for PKMs.
Variant
Setting
Noulinearity
WT-S
WT-M
E8
Dense Baseline
ReLU
11.81
9.46
1.08
PKM
value-count
Softmax
14.11
11.29
1.20
PKM
value-count
ReLU
13.32
10.16
1.12
PKM
total params.
Softmax
13.96
11.10
1.16
PKM
total params.
ReLU
12.77
9.98
1.11
PKM init
total params.
ReLU
12.75
9.96
1.11
expert. Are these models combining experts in more meaningful ways? To test this, measure the distribution of experts that are used together on Wikitext- 103 with our 47M MoE model with = The result can be seen Fig. 6: the network combines experts a rich way, further supporting the use of K > 1. Note that, it remains an open question whether such "compositions" may help the generalization and compositional behavior of the network (Fodor and Pylyshyn, 1988; Pagin and Westerstähl, 2010; Hupkes et al.. 2020).
is negligible), dmodel G .g. However, both the memory usage and the execution time of the MoE. are almost independent of NE. except for a small linear factor due the selection network (see Fig. 2). Figures 2, and 9 show the actual measured execution time and memory usage on RTX 3090 GPU.
Note that there is significant difference in terms of speed and memory usage between differ- ent MoE variants given the same dmodel, G, and K. This is because they only differ the selection mechanism and regularization, and not in the way the experts are executed. Since all methods are configured to have the same number of parameters as the dense baselines, and K experts used in parallel, the factor of reduction in both FLOPs and memory usage is given by Ng We show this factor for all models in Tab. 7.
0.20
5
0.15
Exper ID
9
0.10
0.05
B Implementation details
0.00
5
13
We train all of our models for 100k steps with co- sine learning rate decay, starting from the initial learning rate of 0,00025 and decaying to O. We use the Adam optimizer (Kingma and Ba, 2015) with default PyTorrc parameters (Paszke al., 2019). We use gradient clipping with a max gradient norm of 0.25. We show the other hyperparameters of our dense models in Tab. 8. We train our models with XL memory of the same size the context size. However, following Dai al. (2019), eval- the models using longer memory. Unlike the typerparameter-tunn
Expert ID
Figure 6: Expert co-occurrence a o-MoE model with NE 16 experts and K 4. Each shows the distribution of experts used together with the one corre- sponding to the row. Measured on the validation of Wikitext- 103 in the 3rd layer of our 47M a o-MoE model. The other layers and models behave qualitatively the same.
Detailed Usage Count Analysis. We show the relative proportion of experts selected for all layers in Fig. 7. For more details, please refer to Sec. 6.3.
memory sizes Trans- former XL, we use times the context size (this approximates the size of the memory by Dai et al. (2019), while being simple).
A.5 More on Resource Efficiency
For execution time and memory usage, both the dense MLP and the MoE layers are linear in dmodel (Fig. 9), the MLP is linear in dg, and MoE is linear (Fig. 8) and K. For the same number of parameters (except for the selection network, which
The hyperparameters of the MoE models match those of their dense counterparts with the same number of parameters, except for the MoF-specific ones, which are shown Tab. denotes the expert dropout and y denotes the regularization

---

Table 7: The relative amount FLOPs and memory used by the feedforward block of the MoE transformer compared to its dense counterpart. The same configurations are shown in Tab. 10.
Datasel
Wikitexi 103
Wikitext 103
Wikitext 103
Enwik8
drudkl
412
412
1024
512
# params
47M
237M
262M
41M
G
K
a-MoF (ours)
25.0%
3.
12.5%
25.0%
standard dropout
25.0%
3.1%
12.5%
25.0%
softmax (after top-k)
25.0%
3.1%
12.5%
25.0%
softmax (before top-k)
128
4
25.0%
3.1%
12.5%
25.0%
standard init
128
25.0%
3.1%
12.5%
25.0%
no reg (7 = 0,8 = 0)
128
25.0%
3.1%
12.5%
25.0%
= 8, C =
64
25.0%
3.1%
12.5%
25.0%
= 2, C = 256
256
2
25.0%
3.1%
12.5%
25.0%
K = 1, 512
512
25.0%
3.1%
12.5%
25.0%
N' 2Ng: G 64
64
4
12.5%
1.6%
12.5%
K
128
6.2%
0.8%
6.2%
K = 2
128
2
12.5%
1.6%
12.5%
K =
128
8
50.0%
6.2%
50.0%
Switch, K = 1, 512
$12
25.0%
3.1%
12.5 %
25.0%
no dropout K = 4, C = 128
512
25.0%
3.1%
12.5%
25.0%
128
4
25.0%
3.1%
25.0%
K = 1, G 128
128
6.2%
0.8%
6.2%
no dropout
128
6.2%
0.8%
6.2%
S-BASE
25.0%
3.1%
12.5
25.0%
K 1, 1 512
512
25.0%
3.1%
12.5%
25.0%
strength used for the loss L (See Eq. 21). For the non-MoE layers, the same dropout used as for the baselines. For Switch Transformers, we use ? 0.01 with regularization of the form presented in Eq. 17, following Fedus et al. (2022). The other variants, including S-BASE, use the regularizer proposed by (Eq. 21).
-0,...K-1]] 1.,.K--I,,.CMM( M) E Nx
CVMM(V M,M)
(26)
M
Vin, m] M S[m]. m,
Our CUDA kernel is based the blog post developing matrix multiplication kernel by Si- mon Boehm (https://sibolt 22/CUDA-MMM), However, there are major differ- ences: unlike standard matrix multiplication, in our case, different matrices could be used for different batch elements of the input. In order to be able to reuse matrices fetched from the global memory of the GPU, first do a preprocessing step: we sort the selection indices, and obtain reordering vector. This gives us an ordering of the input and output batch elements, such that the consecutive indices are multiplied by the same matrix with high probability. Fortunately, multiple channels have to be fetched/written out at once, so this reordering has minimal overhead. Our kernel has an addi- tional grid dimension compared to standard ma- trix multiplication, iterating over the matrix index, (0. K 1}. We find that skipping matrices that do have any corresponding inputs has min- imal overhead. To avoid checking all elements the reordering vector, we precompute their offsets.
Our small PKM models use 46 subkeys resulting in 462 2116 values for the mn-matched case and 62 subkeys (3844 values) for the parameter- matched case. The PKM equivalent of the 262M parameter model on Wikitext-109 103 has 64 subkeys (4096 values) for the dg-matched and 89 subkeys (7921 values) for the parameter-matchie case. The PKM models do not use dropout in the PKM layers, and have heads.
com/ com/articles/
B.1 A Few Words on the CUDA Kernel
We call the key operation for our MoE layers con- ditional vector-matrix multiplication, or CVMM, and we define it as follows. Given batch of vec- tors, E NNM where is the batch size and M is the number channels, a set of K matri- ces M. E Mx and selection indices E

---

Table 8: Hyperparameters of dense baselines and their MoE counterparts. For the MoE-specific hyperparameters. please refer Tab. 9. "SetencePiece" tokenization used for Wikitext- 103e, C4 and PES2O datasets, and "Character" for Enwik8.
Tokenization
#params
Amodel
du
Flayen
Theads
head
batch dropout
Ir warmup
SentencePiece
47M
412
2053
16
10
41
256
64
0.1
SentencePicce
238M
412
16480
16
10
41
256
64
0.1
SentencePiece
262M
1024
4110
18
16
64
512
64
0.2
4000
Character 41M
512
2053
12
64
512
32
0.1
Table 9: MoF-speciiicc hyperparameters for different model variants. denotes the scaler for the load balancing term the loss and is the probability of the expert dropout. The standard, transformer-specific hyperparameters are the same as for the baselines. Please refer to Tab. 8. "SetencePiece" tokenization used for Wikitext- 103, C4 and PES20 datasets, and "Character" for Enwik8.
Tokenization
#params
dmadd
Ng
C
K
7
SentencePiece
47M
412
16
128
0.001
SentencePiece
237M
412
128
0.05
0.001
SentencePiece
262M
1024
32
128
0.2
0.001
Character
4IM
512
16
128
0.05
0.0001
Our kernel uses shared memory and register caching: however, it does not use asynchronous loads, which makes I/O bound. It also does not support tensor cores and mixed precision. The pre-processing step uses the radix sort from the CUB library. However, computing the offsets re- quires counting the number of vectors assigned to single matrix. This information, well as the offset, which is their sum, are freely available as sub-results that the radix sort computes anyways; however, we found no way of extracting it from the CUB implementation. We estimate that by implementing a more efficient preprocessing step. asynchronous loads, and tensor core support, our kernel can be further accelerated by a factor of two.
B.2 Additional Results on MoEs
Additional results of different MoE variants with model details are shown in Tab. 10. We repeat the entries from Tab. for easier comparison.

---

Table 10: Detailed ablation results. WT-S* is obtained by naively scaling Ng in WT-S. More details in Sec. 6.3. We do not evaluate all versions of the 262M Wikitext-105 model due to its long training time. However, we aim to include what we believe are the most interesting variants. 7 = means no regularization applied the selection scores (See Eq. 21), 8 denotes expert dropout.
Variant
WT-S
WT-5*
WT-B
dmedel
412
412
1024
512
params
47M
237M
262M
41M
-MoE (ours)
11.59
10.37
standard dropout
12.01
softmax (after top-k)
11.89
11.27
softmax (before top-k)
12.05
10.54
9.62
standard
11.80
10.59
9.67
no reg (n = 0,
128
11.83
10.41
9.51
1.08
K =, 64
64
11.63
10.30
9.58
1.08
K ==2,6 256
256
11.84
10.44
9.56
1.09
K I,G = 512
512
11.90
10.83
9.58
1.09
N = 2N,, 64
64
11.81
10.53
1.08
K 1
128
12.26
11.30
1.09
K =
128
11.90
10.66
1.09
==8
128
11.58
10.22
1.08
Switch, 1, = 512
512
12.27
11.24
9.68
1.08
no dropout
11.88
11.10
9.77
K 4, = 128
128
12.05
11.37
1.10
K 1,G = 128
128
12.61
11.89
1.11
no dropout
128
12.35
11.78
1.10
S-BASE, K 4, 128
13.01
10.50
K =1 512
12.32
11.31
9.77

---

Losr
16"h
16.*
P-A.G
SOnG4 9-A 1-1
S0ASE J-A. G-
Pai 0x07
-o6 Hare 20 (o
Even Tendorer
Even Tendoreu
16"
Layer
COASE You
SCACE C-A-C-05 128
12-
voliee
SOACE YoA C-I
SOKCE Brik C-12
Figure 7: Total proportions of selection weights assigned to given expert (x-axis) the validation set of Wikitext- 103 using the WT-S* models from Tah. 4. Experts are sorted by their popularity. All layers are shown ("Layer 5 is also shown Fig. 3 the main text). The models with poor performance can be distinguished easily (Switch Transformer and M-MoE with softmax and renormalization. "softmax (renom.)"). Their poor performance may be partially explained by expert collapse. In contrast, the fine performance differences between the rest of the models do not seem to be due the expert collapse phenomenon.

---

MLP
Memory (GB)
Time (ms)
MoE
200
100
0-
256
512
1024
Expert size (G)
Figure 8: Mcasured execution time and memory us- age forward-backward pass of a single MLP and MoE layer. [B] 32768, corresponding the realistic scenario of batch size 64 and sequence length 512, Ng 32 and dr G Ng. Full lines show the execution time, and dashed ones the memory consumption. Because they are both linear with similar slopes, they almost indistinguishable. Even with our suboptimal CUDA kernel, the wall-clock time is faster starting from experts.
dmodel 512, K
150
Memory (GB)
Time (mS)
MoE
0.6
50
0.2
128
512
1024
2048
Awodd
Figure 9: Measured execution time and memory us- age forward-backwer pass of a single MLP and MoB layer. [B] 32768, corresponding to the real- istic scenario of a batch size 64 and sequence length 512, K 4, Nę 32, G 128 and dæ G Ng. Full lines show the execution time, and dashed ones the memory consumption. Even with our suboptimal CUDA kernel, the wall-clock time is faster starting from 16 experts.


Reviews:
[{"review_id": "tS0GnXilHq", "reviewer": "Reviewer_d9Y9", "summary": "", "questions": "1. It would be valuable to include additional experiments on diverse text datasets such as PTB and C4. Expanding the experimental evaluation beyond Enwik8 and WikiText-103 can provide a more comprehensive understanding of the proposed approach's performance across various text domains and scales.\n2. Could you provide further clarification and elaboration on the design principles underlying the $\\sigma$-MoE model? ", "limitations": "", "rating": 3, "confidence": 3, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable time spent on reviewing our work.\n\n> *“The empirical investigations reveal that the proposed MoEs perform competitively when compared to the dense Transformer-XL, underscoring the applicability of MoEs to language models of varying scales.”*\n\nThank you very much for acknowledging this. We would like to emphasize that this is under the *parameter-equivalent condition*, which is a unique and *novel* aspect of our study, challenging the common belief that MoEs underperform their dense counterparts under this condition. To the best of our knowledge, no prior work has shown that MoEs are competitive in this setting.\n\n> *“1. The proposed sigma-MoE framework requires further justification. ... The paper should delve into the rationale behind these design choices and their benefits, particularly in terms of their impact on the performance of the MoE model.”*\n> (Question 2) *”Could you provide further clarification and elaboration on the design principles underlying the sigma-MoE model?”*\n\nWe are sorry for the confusion. We will emphasize these points better in the final version of our paper. The $\\sigma$-MoE consists of the following key architectural differences compared to previous work: non-competitive selection function (sigmoid), global entropy regularization, optional expert dropout, and the initialization of the selection mechanism. Here is a short summary of the rationale for each of them:\n\n- Sigmoid activation function: Section 3 highlights the similarity between the MoEs and the top-k approximation of the feedforward layers; we observe that one of the key differences\nbetween them is the use of a competitive (softmax in MoEs) vs. non-competitive activation function (in standard FF layers).\nIn Sec 6.1 we empirically show that the top-k approximation of the FF layers performs well.\nThis success has motivated us to also use a non-competitive activation function (sigmoid) in our MoE. \nThis explanation is given in line 403.\n- Global entropy regularization: We found that the existing regularization methods were complex and somewhat arbitrary (L451). Our goal was to have each of the experts used for some processing step, with minimal restrictions. A natural way to achieve this is to regularize the entropy of the distribution of the selection scores (computed within the batch, in practice). This results in Eqs. 20 and 21.\n- Expert dropout (L465-476): Our primary goal was to increase expert utilization.\nRandomly dropping experts help us prevent the collapsing phenomena where only a few (or single) experts are consistently used and trained. \n- Initialization of the selection mechanism: The rationale is briefly described in line 445. We normalize the length of the row vectors of the projection matrix. This is helpful because, given that each row is responsible for computing the score of the corresponding expert by taking the dot product with the input, if any of the rows had a smaller norm than the others, the corresponding expert would be selected more rarely, which would encourage the collapsing (assuming a uniform distribution of the input vectors).\n\nFollowing the reviewer’s suggestion:\n\n> *”3. To enhance the clarity of the novelty introduced beyond the MoE variant design, it's advisable to provide further elaboration and illustration.”*\n\nWe will add a figure summarizing these modifications in the final version. If the reviewer has any better/concrete suggestions, we’ll be also happy to take them.\n\n> *“2. it would be prudent to expand the experimental scope to include other datasets.”* \n> (Question 1) *”It would be valuable to include additional experiments on diverse text datasets such as PTB and C4.”*\n\nWe would like to thank the reviewer for this suggestion. We conducted extra experiments using the same models on C4 and the newly proposed peS2o dataset (https://github.com/allenai/peS2o, a collection of cleaned, filtered text from academic papers). We chose peS2o instead of PTB because PTB is tiny and is known to be all about \"regularization\" (see for example Merity et al. 2017: Regularizing and Optimizing LSTM Language Models). We used identical hyperparameters as in the Wikitext 103 experiments in our paper, including the same number of iterations. This means that we cannot do a full epoch on the training data, which would be prohibitively expensive with the compute budget we have. Instead, we train it for the same number of tokens as for Wikitext 103. The results are as follows:\n\n|Dataset | | | C4 | C4 | peS2o | peS2o |\n| --- | --- | --- | ---:   | ---:   | ---:   | ---:   |\n| $d_{model}$ | | | 412 | 1024 | 412 | 1024 |\n| # params | | | 47M |  262M |  47M |  262M |\n| | G | K  |  |  |  |  |\n|  |  |  |    |    |    |    |\n| Dense | 128 | 1 | 23.76 | 17.79 | 14.34 | 10.91 |\n| $\\sigma$-MoE (ours) | 128 | 4 | 23.25 | 17.46 | 14.12 | 10.91 |\n| Switch, $K=1, G=512$ | 512 | 1 | 24.47 | 18.29 | 14.74 | 11.56 |\n| S-BASE | 128 | 4 | 35.48 | 18.53 | 16.61 | 11.72 |\n\nThis confirms the results we reported in the paper on two more datasets.\nWe hope that this helps to clear the reviewer's doubts about the generality of our 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": "1. The exploration of approximating two-layer feedforward networks (FFNs) is both novel and captivating. Notably, the authors provide a comprehensive perspective that encompasses various methods, and they also conduct a comprehensive comparison of the most notable MoE variants.\n2. The motivation behind the paper is distinctly articulated and substantiated. The seamless progression from motivation to theoretical analysis, and ultimately to the proposal of the method, is presented in a coherent and natural manner.\n3. The proposed approach is effectively supported by empirical experiments, demonstrating its ability to attain performance on par with dense networks.", "reasons_to_reject": "1. The proposed $\\sigma$-MoE framework requires further justification. The specific design approach for $W_3$ and the uniform initialization of weight matrices are notable contributions. The paper should delve into the rationale behind these design choices and their benefits, particularly in terms of their impact on the performance of the MoE model.\n\n2. The empirical investigations provided appear somewhat limited in scope. Given that all experiments are exclusively conducted on WikiText-103 and Enwik8, which share similar data distributions, it would be prudent to expand the experimental scope to include other datasets. This would provide additional support for the performance claims of the $\\sigma$-MoE model.\n\n3. While the paper introduces novel elements through the MoE variant design, the novelty level might be constrained. To enhance the clarity of the novelty introduced beyond the MoE variant design, it's advisable to provide further elaboration and illustration.\n\n", "questions_to_authors": "1. It would be valuable to include additional experiments on diverse text datasets such as PTB and C4. Expanding the experimental evaluation beyond Enwik8 and WikiText-103 can provide a more comprehensive understanding of the proposed approach's performance across various text domains and scales.\n2. Could you provide further clarification and elaboration on the design principles underlying the $\\sigma$-MoE model? ", "paper_topic_and_main_contributions": "This paper explores diverse techniques for approximating two-layer neural networks (NNs). To begin with, the authors introduce a comprehensive framework that unifies the Top-K activation function, Mixtures of Experts (MoEs), and product-key memories (PKMs). By thoroughly analyzing their approach, they subsequently present enhancements for both MoEs and PKMs. The empirical investigations reveal that the proposed MoEs perform competitively when compared to the dense Transformer-XL, underscoring the applicability of MoEs to language models of varying scales.", "excitement": 3, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "UkdVUPycze", "reviewer": "Reviewer_JndA", "summary": "", "questions": "1. How does the approximation of the two-layer feedforward network affect the overall performance of the transformer model? \n\n2. Have you conducted any experiments to evaluate the trade-off between efficiency and performance?", "limitations": "", "rating": 3, "confidence": 2, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable time spent on reviewing our work.\n\n> *”1. The paper lacks a clear motivation for why approximating two-layer feedforward networks is important for efficient transformers. The authors should provide a more compelling argument for why this is a necessary step toward creating more efficient transformers.”*\n\nWe’d like to draw the reviewer’s attention to the fact that every Transformer “layer” consists of one self-attention layer and one two-layer feedforward block. In typical large language models popular today, the feedforward block is configured to be very large, becoming an important computational bottleneck for efficient large transformer language models. This has motivated not only our work but also many prior works to accelerate it (for example, please see references for baseline models presented in Table 4).\n\n> *”2. The paper does not provide a thorough comparison of their proposed method with existing methods for reducing compute and memory requirements of neural networks. The authors should include a more comprehensive analysis of how their approach compares to other methods in terms of performance and efficiency.”*\n> *”3. Figure 2 shows the execution time and memory usage of a forward-backward pass of a single MLP and MoE layer, however, the authors should provide more information on the computational complexity of their approach and how it compares to other MoE methods.”* (We grouped Question 2 and 3 because they are directly related.)\n\nWe first would like to point out that Tables 3 and 4 compare the performance of our models to different baselines, and Table 3 also compares their efficiency with their dense counterparts.\n\nThat said, we agree with the reviewer that the comparison in terms of efficiency can still be potentially confusing in the current presentation.\nIn fact, given the same hyperparameters ($d_{model}$, $G$ and $K$), all MoE approaches have an almost identical number of flops and memory usage (namely, relative to their dense counterparts, both their FLOP and memory reduction factor are given by $\\frac{G}{N_E}$). This is simply because they differ only in the way experts are selected and the selection mechanism only accounts for a negligible proportion of the total resources. This is the reason why we originally did not report FLOPS and the memory usage of other MoE baselines.\n\nIn the final version, we will not only add the corresponding explanation, but also a new table displaying FLOPs/memory reduction compared to the dense model for all models/hyper-parameters configurations we evaluated:\n\n\n| Dataset | | | Wikitext 103 | Wikitext 103 | Wikitext 103 | Enwik8 |\n| --- | --- | --- | ---:   | ---:   | ---:   | ---:   |\n| $d_{model}$ | | | 412 | 412 | 1024 | 512 |\n| # params | | | 47M | 237M | 262M | 41M |\n| | G | K  |  |  |  |  |\n|  |  |  |    |    |    |    |\n| $\\sigma$-MoE (ours) | 128 | 4 | 25.0\\% | 3.1\\% | 12.5\\% | 25.0\\% |\n| $~~~$ standard dropout | 128 | 4 | 25.0\\% | 3.1\\% | 12.5\\% | 25.0\\% |\n| $~~~$ softmax (after top-k) | 128 | 4 | 25.0\\% | 3.1\\% | 12.5\\% | 25.0\\% |\n| $~~~$ softmax (before top-k) | 128 | 4 | 25.0\\% | 3.1\\% | 12.5\\% | 25.0\\% |\n| $~~~$ standard init | 128 | 4 | 25.0\\% | 3.1\\% | 12.5\\% | 25.0\\% |\n| $~~~$ no reg ($\\gamma=0,\\delta=0$) | 128 | 4 | 25.0\\% | 3.1\\% | 12.5\\% | 25.0\\% |\n| $~~~$ $K=8, G=64$ | 64 | 8 | 25.0\\% | 3.1\\% | 12.5\\% | 25.0\\% |\n| $~~~$ $K=2, G=256$ | 256 | 2 | 25.0\\% | 3.1\\% | 12.5\\% | 25.0\\% |\n| $~~~$ $K=1, G=512$ | 512 | 1 | 25.0\\% | 3.1\\% | 12.5\\% | 25.0\\% |\n| $~~~$ $N_E' = 2N_E, G=64$ | 64 | 4 | 12.5\\% | 1.6\\% | -| 12.5\\% |\n| $~~~$ $K=1$ | 128 | 1 | 6.2\\% | 0.8\\% | - | 6.2\\% |\n| $~~~$ $K=2$ | 128 | 2 | 12.5\\% | 1.6\\% | - | 12.5\\% |\n| $~~~$ $K=8$ | 128 | 8 | 50.0\\% | 6.2\\% | - | 50.0\\% |\n|  |  |  |    |    |    |    |\n| Switch, $K=1, G=512$ | 512 | 1 | 25.0\\% | 3.1\\% | 12.5\\% | 25.0\\% |\n| $~~~$ no dropout | 512 | 1 | 25.0\\% | 3.1\\% | 12.5\\% | 25.0\\% |\n| $~~~$ $K=4, G=128$ | 128 | 4 | 25.0\\% | 3.1\\% | - | 25.0\\% |\n| $~~~$ $K=1, G=128$ | 128 | 1 | 6.2\\% | 0.8\\% | - | 6.2\\% |\n| $~~~$$~~~$ no dropout | 128 | 1 | 6.2\\% | 0.8\\% | - | 6.2\\% |\n|  |  |  |    |    |    |    |\n| S-BASE | 128 | 4 | 25.0\\% | 3.1\\% | 12.5\\% | 25.0\\% |\n| $~~~$ $K=1, G=512$ | 512 | 1 | 25.0\\% | 3.1\\% | 12.5\\% | 25.0\\% |\n\n\n\nAdditionally, here are some execution time and memory usage benchmarks for the different MoE models given different expert sizes and $d_\\text{model}=512$, $K=4$, $N_{E}=32$ on a 2080Ti GPU (which empirically confirm our statement above about comparable time/memory usage among all the MoE methods):\n\n|  | 64 | 128 | 256 | 512 | 1024 |\n| --- |  ---: | ---: | ---: | ---: | ---: |\n | $\\sigma$-MoE | 17.8ms | 24.7ms | 43.3ms | 80.5ms | 155.5ms |\n | Switch | 16.1ms | 24.9ms | 43.4ms | 80.5ms | 155.5ms |\n | S-BASE | 16.6ms | 25.6ms | 44.1ms | 81.4ms | 156.7ms |\n\n |  | 64 | 128 | 256 | 512 | 1024 |\n| --- |  ---: | ---: | ---: | ---: | ---: |\n | $\\sigma$-MoE | 135MB | 199MB | 327MB | 583MB | 1095MB |\n | Switch | 135MB | 199MB | 327MB | 583MB | 1095MB |\n | S-BASE | 135MB | 199MB | 327MB | 583MB | 1095MB |\n\nWe will add these tables to the final version.\nThank you for drawing our attention to this potentially confusing aspect. \n\n> *“4. While the algorithmic design of the proposed method appears intuitive, the authors would benefit from a more detailed theoretical or analytical analysis of the existing content of the paper, which is currently lacking in detail. The authors should provide more information on the theoretical underpinnings of their approach and how it relates to existing research in the field.”*\n\nWe are not fully sure to understand exactly what kind of *“more detailed theoretical or analytical analysis of the existing content of the paper”* the reviewer is requesting. Regarding *“theoretical underpinnings of their approach and how it relates to existing research in the field”*, we’d like to claim that our current Section 2, 3 and 4 provide a rather comprehensive summary of conceptual comparisons between different approaches. Sec 2 starts with a general framework (theoretical view) of approximate two-layer feedforward networks. Sec 3 then introduces top-K activations (an essential component of any MoEs) and connects PKMs and MoEs. Finally, Sec 4 reviews and conceptually compares existing MoE methods, before describing our method in Sec 5. If the reviewer still thinks that this is not enough, please specifically tell us what is missing. Thank you.\n\n> *“How does the approximation of the two-layer feedforward network affect the overall performance of the transformer model?”*\n\nEssentially, all our experimental results (Tables 1, 2, 3, 4) show *“how approximation of the two-layer feedforward network affects the overall performance of the transformer model”*. These tables evaluate *different* types of approximate methods for the two-layer feedforward networks.\n\n> *”Have you conducted any experiments to evaluate the trade-off between efficiency and performance?”*\n\nYes, Table 3 compares both FLOPs and the performance (as measured by perplexity). Using our sigma-MoE, we obtain 75\\% - 87.5\\% speed up in FLOPs without losing performance. We also reported additional numbers related to this in the new tables above.\n\nWe hope our response brings clarifications and improves the reviewer’s view on our contributions.", "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": "1. This paper proposes a new perspective on MoEs and a general framework that unifies various methods to approximate two-layer neural networks, which demonstrates the effectiveness of Mixtures of Experts (MoEs) in creating resource-efficient large language models (LMs).\n\n2. This paper studies an important problem of reducing the computing and memory requirements of neural networks while maintaining performance. The paper proposes a novel method of using sparse Mixtures of Experts (MoEs) to create resource-efficient LMs and presents a general framework that unifies various methods to approximate two-layer neural networks. The paper's proposed methods to improve both MoEs and PKMs could also be useful for researchers and practitioners working on creating resource-efficient LMs.\n ", "reasons_to_reject": "1. The paper lacks a clear motivation for why approximating two-layer feedforward networks is important for efficient transformers. The authors should provide a more compelling argument for why this is a necessary step toward creating more efficient transformers.\n2. The paper does not provide a thorough comparison of their proposed method with existing methods for reducing compute and memory requirements of neural networks. The authors should include a more comprehensive analysis of how their approach compares to other methods in terms of performance and efficiency.\n3. The proposed method has not been subject to a comprehensive complexity analysis or a comparative analysis of training time. Figure 2 shows the execution time and memory usage of a forward-backward pass of a single MLP and MoE layer, however, the authors should provide more information on the computational complexity of their approach and how it compares to other MoE methods.\n4. While the algorithmic design of the proposed method appears intuitive, the authors would benefit from a more detailed theoretical or analytical analysis of the existing content of the paper, which is currently lacking in detail. The authors should provide more information on the theoretical underpinnings of their approach and how it relates to existing research in the field.", "questions_to_authors": "1. How does the approximation of the two-layer feedforward network affect the overall performance of the transformer model? \n\n2. Have you conducted any experiments to evaluate the trade-off between efficiency and performance?", "paper_topic_and_main_contributions": "The paper studies an important problem, which is about reducing the compute and memory requirements of neural networks (NNs) without sacrificing performance. Specifically, the paper proposes using sparse Mixtures of Experts (MoEs) to create resource-efficient large language models (LMs). The paper presents a general framework that unifies various methods to approximate two-layer neural networks, including feedforward blocks of Transformers, and proposes methods to improve both MoEs and product-key memories (PKMs). ", "excitement": 3, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Missing_References": "None", "Typos_Grammar_Style_and_Presentation_Improvements": "None"}, "extra_text": {}}, {"review_id": "O45fV309cW", "reviewer": "Reviewer_QrbU", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable time spent on reviewing our work.\n\n> *“Some architectural choices are not entirely motivated”*\n\nWe would have appreciated it a lot if the reviewer had specified *exactly* which “architectural choices” s/he is referring to. A general response is that all the introduced choices are motivated by the ultimate goal of making the approximate feedforward layers perform as good as the non-approximate one. We will be happy to clarify any explanations in our text that the reviewer still finds confusing. \n\nAlso, some other reviewers (Reviewer d9Y9 in particular) also asked for clarification on certain design choices. Our response to them (we typically provide specific pointers to the part of the text that explains them) might also help the reviewer.\n\n> *”Like all MoE approaches, this method is relatively complicated to implement compared to vanilla Transformers, requiring a custom CUDA kernel for efficient implementation.”*\n\nAs the reviewer wrote, this is true for all other MoE baselines or sparse acceleration methods in general. In fact, this is even true for many other well-known methods (e.g., even LSTM) before they became standard today. A custom CUDA kernel is ultimately required for any efficient implementation. Importantly (as the reviewer acknowledges), we will publicly release our kernel upon acceptance of our paper.\n\nFinally, we would like to emphasize one important aspect/novelty of our work which is omitted in the review. We demonstrate that MoEs are competitive to their dense counterparts even under the *parameter-equivalent condition*. This challenges the common belief that MoEs underperform their dense counterparts under this condition. To the best of our knowledge, no prior work has shown such a result (prior work only focused on the FLOP-equivalent comparison).\n\nWe hope that our response above brings clarifications to the concerns raised by the reviewer.\nIf the reviewer has any specific questions regarding the first point, we will be happy to answer them during the author-review discussion period.", "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": "- Interesting approach\n\n- Compelling results\n\n- The authors promise to release the code", "reasons_to_reject": "- Some architectural choices are not entirely motivated\n\n- Like all MoE approaches, this method is relatively complicated to implement compared to vanilla Transformers, requiring a custom CUDA kernel for efficient implementation.", "questions_to_authors": "", "paper_topic_and_main_contributions": "The paper presents a novel Mixture-of-Experts (MoE) for Transformer models.\n\nThe paper starts with a review of existing MoE approaches, characterizing them as approximations of the 2-layer MLP of the transformer block, comparing it with other approximation schemes (top-k and Product-Key memories). After reviewing the different characteristics of these approaches, the paper proposes a novel MoE which uses a non-competitive expert selection function (sigmoid) followed by top-k selection, with a specific initialization scheme, expert dropout and a load-balancing regularization objective.\n\nThey evaluate the method on language modelling tasks, on standard datasets, and report that their approach provides comparable quality of parameter-equivalent dense models while being more compute-efficient, and also outperform other MoEs and approximation schemes. Ablations studies are also reported. The authors conclude that their MoE is beneficial even for small models.", "excitement": 3, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "qaw1JaOPaF", "reviewer": "Reviewer_FXRA", "summary": "", "questions": "The author emphasises that the \\sigma MoE is applied to each MLP block of the model. What if the method is also applied in the usual way (e.g. once in every nth layer or even only in a single layer), how does the effectiveness change? Furthermore, including the number of layers applied, what is the main controllable factor in your approach that reduces FLOPs? In the \\sigma n MoE initialisation? Why is the social care initialisation W_3 required? Does this effectiveness still validate on several downstream benchamark performance? What if the given reduced FLOPs are more than 50% (e.g. 75%) so that the Top-K activation function can be compared, can the proposed method outperform the Tok-K approach?", "limitations": "", "rating": 2, "confidence": 2, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable time spent on reviewing our work.\n\n> *”Evaluated under parameter-equivalent conditions, the empirical results underscore the competitive performance of MoEs compared to dense baselines.”*\n\nThank you very much for acknowledging this. In fact, evaluating MoEs under the “parameter-equivalent” condition is a unique and *novel* aspect of our study, challenging the common belief that MoEs underperform their dense counterparts under this condition. To the best of our knowledge, no prior work has shown MoEs are competitive under this setting.\n\n> *“Lack of FLOPs information on other approaches' tables leads to unclear comparisons on how this approach is attractive in terms of FLOPs and bits/characters”*\n\nThank you for drawing our attention to this potentially confusing aspect.\nIn fact, given the same hyperparameters ($d_{model}$, $G$ and $K$), all these MoE approaches have an almost identical number of flops and memory usage (namely, relative to their dense counterparts, both their FLOP and memory reduction factor are given by $\\frac{G}{N_E}$). This is simply because they differ only in the way experts are selected and the selection mechanism only accounts for a negligible proportion of the total resources. This is the reason why we originally did not report FLOPS and the memory usage of other MoE baselines. That said, we fully agree with the reviewer that this was confusing as this explanation was missing in the current text. We will not only add the corresponding explanation, but also a new table displaying FLOPs/memory reduction compared to the dense model for all models/hyper-parameters configurations we evaluated:\n\n\n| Dataset | | | Wikitext 103 | Wikitext 103 | Wikitext 103 | Enwik8 |\n| --- | --- | --- | ---:   | ---:   | ---:   | ---:   |\n| $d_{model}$ | | | 412 | 412 | 1024 | 512 |\n| # params | | | 47M | 237M | 262M | 41M |\n| | G | K  |  |  |  |  |\n|  |  |  |    |    |    |    |\n| $\\sigma$-MoE (ours) | 128 | 4 | 25.0\\% | 3.1\\% | 12.5\\% | 25.0\\% |\n| $~~~$ standard dropout | 128 | 4 | 25.0\\% | 3.1\\% | 12.5\\% | 25.0\\% |\n| $~~~$ softmax (after top-k) | 128 | 4 | 25.0\\% | 3.1\\% | 12.5\\% | 25.0\\% |\n| $~~~$ softmax (before top-k) | 128 | 4 | 25.0\\% | 3.1\\% | 12.5\\% | 25.0\\% |\n| $~~~$ standard init | 128 | 4 | 25.0\\% | 3.1\\% | 12.5\\% | 25.0\\% |\n| $~~~$ no reg ($\\gamma=0,\\delta=0$) | 128 | 4 | 25.0\\% | 3.1\\% | 12.5\\% | 25.0\\% |\n| $~~~$ $K=8, G=64$ | 64 | 8 | 25.0\\% | 3.1\\% | 12.5\\% | 25.0\\% |\n| $~~~$ $K=2, G=256$ | 256 | 2 | 25.0\\% | 3.1\\% | 12.5\\% | 25.0\\% |\n| $~~~$ $K=1, G=512$ | 512 | 1 | 25.0\\% | 3.1\\% | 12.5\\% | 25.0\\% |\n| $~~~$ $N_E' = 2N_E, G=64$ | 64 | 4 | 12.5\\% | 1.6\\% | -| 12.5\\% |\n| $~~~$ $K=1$ | 128 | 1 | 6.2\\% | 0.8\\% | - | 6.2\\% |\n| $~~~$ $K=2$ | 128 | 2 | 12.5\\% | 1.6\\% | - | 12.5\\% |\n| $~~~$ $K=8$ | 128 | 8 | 50.0\\% | 6.2\\% | - | 50.0\\% |\n|  |  |  |    |    |    |    |\n| Switch, $K=1, G=512$ | 512 | 1 | 25.0\\% | 3.1\\% | 12.5\\% | 25.0\\% |\n| $~~~$ no dropout | 512 | 1 | 25.0\\% | 3.1\\% | 12.5\\% | 25.0\\% |\n| $~~~$ $K=4, G=128$ | 128 | 4 | 25.0\\% | 3.1\\% | - | 25.0\\% |\n| $~~~$ $K=1, G=128$ | 128 | 1 | 6.2\\% | 0.8\\% | - | 6.2\\% |\n| $~~~$$~~~$ no dropout | 128 | 1 | 6.2\\% | 0.8\\% | - | 6.2\\% |\n|  |  |  |    |    |    |    |\n| S-BASE | 128 | 4 | 25.0\\% | 3.1\\% | 12.5\\% | 25.0\\% |\n| $~~~$ $K=1, G=512$ | 512 | 1 | 25.0\\% | 3.1\\% | 12.5\\% | 25.0\\% |\n\n\n\nAdditionally, here are some execution time and memory usage benchmarks for the different MoE models given different expert sizes and $d_\\text{model}=512$, $K=4$, $N_{E}=32$ on a 2080Ti GPU (which empirically confirm our statement above about comparable time/memory usage among all the MoE methods):\n\n|  | 64 | 128 | 256 | 512 | 1024 |\n| --- |  ---: | ---: | ---: | ---: | ---: |\n | $\\sigma$-MoE | 17.8ms | 24.7ms | 43.3ms | 80.5ms | 155.5ms |\n | Switch | 16.1ms | 24.9ms | 43.4ms | 80.5ms | 155.5ms |\n | S-BASE | 16.6ms | 25.6ms | 44.1ms | 81.4ms | 156.7ms |\n\n |  | 64 | 128 | 256 | 512 | 1024 |\n| --- |  ---: | ---: | ---: | ---: | ---: |\n | $\\sigma$-MoE | 135MB | 199MB | 327MB | 583MB | 1095MB |\n | Switch | 135MB | 199MB | 327MB | 583MB | 1095MB |\n | S-BASE | 135MB | 199MB | 327MB | 583MB | 1095MB |\n\n> *“What if the method is also applied in the usual way (e.g. once in every nth layer or even only in a single layer), how does the effectiveness change?”*\n\nIntroducing more dense components will directly slow down the model. Since we show that our $\\sigma$-MoE layer does not deteriorate the performance compared to the dense counterpart, using it only every n-th layer will only cause harm (we are not fully sure to understand what exactly the reviewer refers to as “effectiveness“).\n\n\n> *”Furthermore, including the number of layers applied, what is the main controllable factor in your approach that reduces FLOPs? “*\n\nIn addition to the standard Transformer parameters such $d_{\\text{model}}$, the size of the expert ($G$) and the number of active experts ($K$) also influence the FLOPs (we are not fully sure to understand the reviewer’s question, please let us know if this is effectively what the reviewer wished to know).\n\n> *”In the $\\sigma$-MoE initialisation? Why is the social care initialisation $W_3$ required?”*\n\nThe corresponding explanation is provided in Line 445. Essentially, each element of s, i.e, score $s\\_k$ for expert selection, is a dot product between the input x and a row of $W\\_3$. This initialization simply ensures that all rows of $W\\_3$ have the same norm, such that scores initially only depend on the angles between x and the rows of $W\\_3$ (otherwise there will be a bias towards rows with bigger norms to be selected, which may encourage collapsing).\n\n> *“Does this effectiveness still validate on several downstream benchamark performance?“*\n\nWhile it is not entirely to us what the reviewer refers to as *“effectiveness”*, the speed gain should remain effective for downstream tasks. Regarding the performance, we can not tell it based on our experiments. Here our scope is pure language modeling, and therefore, we evaluate them based on perplexity. The question of how the performance of MoE models in general transfers to downstream tasks compared to their dense counterpart is an interesting research direction for future work.\n\n> *”What if the given reduced FLOPs are more than 50\\% (e.g. 75\\%) so that the Top-K activation function can be compared, can the proposed method outperform the Tok-K approach?”*\n\nWe are not fully sure to understand the reviewer’s question. Which model does *“if the given reduced FLOPs are more than 50\\% (e.g. 75\\%)“* refer to? We confirm that the FLOPs reduction is 75\\% for the small MoE models and 87.5\\% for the big ones (as shown in Table 3). The performance results for the top-K activation are already reported in Table 1. Comparing these to those in Table 4, we observe that the best sigma-MoE outperforms the top-K on the Wikitext-103/47M setting, (11.59 v. 11.68), underperforms it on Wikitext-103/262M (9.44 v. 9.26) and performs similarly on Enwik8 (1.08 v. 1.07). However, we would like to note that, unlike the MoE approaches, the acceleration obtainable by the top-K method alone is limited, as discussed in L199.\n\nWe hope our response above brings clarifications to all the concerns raised by the reviewer.\nIf you find our response useful, please consider increasing the scores. Thank you.", "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": "Clear comparisons to existing methods and evaluation direction. Upon analysis, the steps for the proposed method are straightforward. The effectiveness of the method is validated and ablation studies on different hyperparameter settings and components are provided.", "reasons_to_reject": "Extending and combining the existing technique lacks novelty and, as pointed out in limitations, involves considerable sensitive engineering tricks. Lack of FLOPs information on other approaches' tables leads to unclear comparisons on how this approach is attractive in terms of FLOPs and bits/characters.", "questions_to_authors": "The author emphasises that the \\sigma MoE is applied to each MLP block of the model. What if the method is also applied in the usual way (e.g. once in every nth layer or even only in a single layer), how does the effectiveness change? Furthermore, including the number of layers applied, what is the main controllable factor in your approach that reduces FLOPs? In the \\sigma n MoE initialisation? Why is the social care initialisation W_3 required? Does this effectiveness still validate on several downstream benchamark performance? What if the given reduced FLOPs are more than 50% (e.g. 75%) so that the Top-K activation function can be compared, can the proposed method outperform the Tok-K approach?", "paper_topic_and_main_contributions": "The paper presents a comprehensive investigation of the use of sparse Mixtures of Experts (MoEs) and product-key-memories (PKMs) to improve the efficiency of any-scale-resource-efficient LMs in terms of computational and memory requirements. Evaluated under parameter-equivalent conditions, the empirical results underscore the competitive performance of MoEs compared to dense baselines.", "excitement": 2, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The manuscript underwent a comprehensive evaluation by four reviewers, revealing both strengths and weaknesses in its content.

Pros

1. Novelty and Contribution: The reviewers largely concur that the paper provides a novel investigation into the approximation of two-layer feedforward networks.
2. Empirical Robustness: The empirical outcomes are broadly considered to be compelling, particularly highlighting the MoEs' competitive edge over dense Transformer architectures.
3. Community Relevance: By focusing on diminishing computational and memory overhead without compromising model efficacy, the paper holds significant appeal for both researchers and practitioners in the field.


Cons

1. Lack of Rigorous Justification: The design choices within the MoE framework warrant further elucidation. The reviewers note that these choices are not sufficiently justified, especially regarding their impact on model performance, and the authors were unable to clarify their motivations convincingly in their rebuttal.
2. Need for Novelty Elaboration: Though the paper incorporates novel components, some reviewers question the extent of this novelty and advocate for a more detailed exposition to distinguish the proposed method from existing MoEs.
3. Comparative Evaluation: Reviewers express the need for a more systematic comparison with established methodologies, focusing on aspects such as computational complexity, performance metrics, and overall efficiency. The rebuttal partially resolved this issue.
4. Theoretical Depth: A deeper theoretical or analytical discourse is recommended by the reviewers to both validate the design decisions and to situate the work more solidly within the existing body of literature.
Decision: Accept-Findings