Towards A Unified View of Sparse Feed-Forward Network in Pretraining Large Language Model
Authors: Zeyu Liu
Tim Dettmers
Xi Victoria Lin
Veselin Stoyanov
Xian Li
Venue: EMNLP 2023 Main (2023)
Keywords: Mixture-of-Experts
Efficiency
Transformer
Architecture
Pretraining
LLM
TL;DR: After a systematic analysis from a key-value memory view, we found a better way to parameterize the gate in mixture of expert architecture.

Abstract: Large and sparse feed-forward layers (S-FFN) such as Mixture-of-Experts (MoE) have proven effective in scaling up Transformers model size for pretraining large language models. By only activating part of the FFN parameters conditioning on input, S-FFN improves generalization performance while keeping training and inference costs (in FLOPs) fixed. In this work, we analyzed two major design choices of S-FFN: the memory block (a.k.a. expert) size and the memory block selection method under a general conceptual framework of sparse neural memory. Using this unified framework, we compare several S-FFN architectures for language modeling and provide insights into their relative efficacy and efficiency. We found a simpler selection method — Avg-K that selects blocks through their mean aggregated hidden states, achieving lower perplexity in language model pretraining compared to existing MoE architectures including Switch Transformer (Fedus et al., 2021) and HashLayer (Roller et al., 2021).

Towards A Unified View of Sparse Feed-Forward Network in Pretraining Large Language Model
Leo Z. Liu 9
Tim Dettmers
Xi Victoria Lin University of Washington, Seattle
Veselin Stoyanov Meta Al.
Xian Li
Abstract
(Borgeaud al., 2022; Schwartz al., 2020; Tay et al., 2020).
Large and sparse feed-forward layers (S-FFN) such as Mixture-of-Expeee (MoB) have proven effective in scaling up Transformers model size for pretraining large language models. By only activating part the FFN parameters condi- tioning input, S-FFN improves generaliza- tion performance while keeping training and inference costs (in FLOPs) fixed. this work, we analyzed two major design choices of S- FFN: the memory block (a.k.a. expert) size and the memory block selection method under general conceptual framework of sparse neu- ral memory. Using this unified framework, we compare several S-FFN architectures for lan- guage modeling and provide insights into their relative efficacy and efficiency. We found simpler selection method Avg-K that selects blocks through their mean aggregated hidden states, achieving lower perplexity in language model pretraining compared to existing MoE ar- chitectures including Switch Transformer (Fe- dus al., 2021) and HashLayer (Roller et al., 2021).
One promising direction is sparse scaling which increases the number of parameters while keeping the training and inference cost (in FLOPs) fixed. Recent work focuses scaling up transformer's feed-forward network (FFN) with sparsely acti- vated parameters, resulting in scaled and sparse FFN (S-FFN). There have been two major ap- proaches to achieve S-FFN. One treats S-FFN a neural memory (Sukhbaatar et al., 2015a) where a sparse memory retrieves and activates only parts of the memory cells (Lample et al., 2019). The other adopts Mixture-of ofx-pet-of-Expert Network (MoE) (Lepikhin et al., 2021; Fedus et al., 2021: Du et al., 2021; Roller et al., 2021; Lewis et al., 2021; Chi et al., 2022) that replaces a single FFN module with multiple equal-sized ones (called "experts") and only activates few among many experts for a particular input.
While both memory and MoE models achieve S-FFN, they have been considered two distinct ap- proaches. We aim to draw the connections between these two classes of S-FFN: What critical design choices do they have in common? Which design choices are essential for their modeling capabil- ity and computation efficiency? Can the effective ingredients of each method be transferred and com- bined to improve performance further?
1 Introduction
Large-scale pretrained language models (LLMs) achieve remarkable performance and general- ization ability for NLP tasks (Radford and Narasimhan, 201 8; Devlin et al., 2019; Liu et al., 2019; Radford et al., 2019; Brown et al., 2020; Raf- fel et al., 2020; Chowdhery et al., 2022). Scaling up the model size (the number of parameters) has been shown as reliable recipe for better generalization, unlocking new capabilities, while the performance has not shown signs of platcauing (Kaplan et al., 2020; Zhang et al., 2022a; Chowdhery et al., 2022: Hoffmman et al., 2022; Wei et al., 2022). However, the computational resources required to train larger language models are formidable, calling for more efficient training and inference solutions of LLMSs
In order to answer these questions, we start from the neural memory view of FFN (Sukhbaatar et al., 2015a) ($2.1) 1) and reduce all S-FFN's 's to the same mathematical form ($3.1). Then, characterize these methods along two dimensions memory block size (e.g. expert size) and memory block selection method (e.g. gating) ($3.2). Using this framework, we made the following contributions:
We study wider range of memory block sizes besides the commonly used block size MoE architectures (Fedus et al., 2022) and show that reducing the block size keeps improving
Work done as a Meta Al Resident, Correspondence to: Lco Zeyu Liu < iu2ecs. washington.ele edu;>

---

Option 1: Mixturre-of-Exert
Transformer Layer
Option Sparse Memory.
Add Normalto lbe
V
PPN
FFN
FFN
FFN4
Vinu Thtw
Top-
Top
Tode
Add
GATE
Load Balancing Constraint
projection
GShard, Switch HashLayer, BASELayer...
'Thee
LowRank Product
Figure Sparse Mixture-of -ixture-of-Expert and Sparse Neural Memory as two different methods.
the perplexity with little incurred extra FLOPs ($5.1 (5.2), leading to better perplexity/com- putation trade-offs.
Background
2.1 Feed-Forward Network
transformer layer (Vaswani et al., 2017) consists of a self-attention block and a Fieed-Forwar Net- work block (FFN). FFN receives an input vector E Rd from the self-attention block, multiplies it with parameter matrix K C Rdm.dd applies a non-linear function f to obtain the hidden states E IR and applies another affine transforma- tion V Rddxdd produce a d-dimensional out- put y. This multi-layer perceptron architecture can be expressed as:
We conduct systematic exploration of block selection methods quantify their relative ef- ficacy and efficiency ($5.2). Specifically, we find that the selection method through a gat- ing function, general, improves the FLOOP- Perplexity trade-off. However, the parame- terization of the current MoE's (Fedus et al., 2021; Lewis et al., 2021) gating function multiplying token representation with sepa- rately learned matrix has worse perplexity than using the FFN hidden states - multiply- ing token representation with FFN's matrix (82.1).
= FFN(x) - f(x K V - m V. (1)
Additionally, we could also view it as neural mem- ory (Sukhbaatar et al., 2015b, 2019; Geva et al., 2021) (Bq. 2).
dm -1
dm
y f(x.k;) V
(2)
Drawing on the insights above, we propose simple gate for S-FFN- Avg-K ($3.3) as hybrid design choice between sparse neu- ral memory and mixture of experts. It effi- ciently selects memory blocks based on the mean aggregated hidden states of each block. With 1 % additional FLOPs, Avg-K achieves lower perplexity ( (14.80) on the validation set of The Pile (Gao et al., 2020) than both Switch Transformer (16.45) (Fedus et al., 2021) and HashLayer (15.75) (Roller et al., 2021). More- over, Avg-K is the first MoE model that per- forms well without load balancing constraint, while conventional MoE transformers like Switch Transformer will degenerate (Fedus et al., 2021; Shazeer et al., 2017; Eigen et al., 2014).
In this view, FFN consists of d,,, key- pairs, known as memory cells. Each key is represented by d-dimensional k; E Rd. and together form the key table; likewise, value vectors V; E Rd constitutes a value table V Rdd xd The memory multiplies the query input x E Rdd with every k; followed by the non-linear function, it produces memory coefficient
f(x k;) for the i-th memory cell. Finally, the output of FFN is the sum of its values v, weighted by their corresponding memory coefficient my Conventionally, the size of FFN dm is set be d.
2.2 Scaling up FFN
As discussed in $1, scaling up the number of pa- rameters in FFN serves as a lever to improve trans-

---

former performance. Since standard FFN ac- counts for about two-thirds of transformer layer's parameters (Geva et al., 2021), scaling up FFN will greatly affect the parameter size of a trans- former model. However, one could sparsely acti- vate the parameter to control the required compute. In this section, we and discuss two approaches to achieve scaled and sparse FFN (S-FFNN, One has mixture-of-expert model activate few experts ($2.2.1), and the other specifies memory model sparsify ($2.2.2).
Static gate, in contrast a learnable gate, does not have any differentiable parameters. Instead, it uses a static mapping that encodes load-balancing constraints to route input (Roller et al., 2021; Guru- rangan et al., 2021). For example, RandHash from HashLayer (Roller et al., 2021) uses hash table that maps from token type randomly selected expert(s). DEMix (Gururangan et al., 2021) en- sures each expert only sees data from pre-defined domain.
2.2.2 Sparse Neural Memory
2.2.1 Mixture of Experts (MoE)
The other line of work follows the memory view of FFN (Eq. 2). straightforward to increase the memory size dm to a much larger value dm >>>>> 4- d. By only using the top-k entries in the memory co- efficient m = x K (Eq. 2), one could sparsely activate the value table, resulting in vanilla sparse memory (VanillaM). However, the naive imple- mentation of this approach results in computation cost proportional linearly to the memory size. Lam- ple et al. (2019) explored the following two tech- niques in this direction to scale computation sub- linearly.
Mixture of experts (MoE; Jacobs et al. (1991)) con- sists of set of expert models {f) and gating function g : I + R to estimates the rele- vance of each expert. Finally, the output is the sum of experts' output weighted by the gate's weight estimation for that particular expert.
MoE(xx- =
Σ iEE= iEE=
gx(x) (xx)
(3)
Recent work (Du et al., 2021; Lepikhin et al., 2021; Roller et al., 2021; Lewis et al., 2021; Zhou et al., 2022) have applied this approach to transformer by dissecting the FFN into multiple expert blocks and sparse activation of the MoE (SMoE). In SMoE, the gating function (or "router") routes input token to subset' (e.g. 2) of experts, (xbs)t(g(x)). Previous work mainly adopts two types gates.
Low-Rank Key Memory (LoRKM) A straight- forward technique is to assume that the full key table K _ddxdm is composed of and approxi- mated by downward projection D Raxd and low-rank key table K Rdm xdg.
K
D.
where de < d. LoRKM produces memory coeffi- cients by m, f((x D) k.) = f(t k.).
Learned gate is parameterized by set of learnable expert embeddings 0 = (eu: ;eg-i] c TR Bxd where each embedding corresponds to expert. The relevance of the i-th expert is obtained by
Product Key Memory (PKM) Building upon LoRKM, PKM further decomposes the low-rank key table by assuming different low-rank keys have structured sharing with each other. See Appendix B for more technical details. Due to such factoriza- tion, PKM has negligible key table K D- (e.g.. < 0.3%) relative to the parameters in the value table.
exp (e, x)
g (x)
E, exp (ej x)
To enforce load balancing when routing, previ- work have employed an additional auxiliary loss (Lepikhin et al., 2021; Fedus et al., 2021; Artetxe al., 2021; Du et al., 2021; Chi et al.. 2022) or framed expert utilization as a constrained optimization problem (Lewis et al., 2021: Zhou et al., 2022).
3 A Unified View of Sparse FFNs
We show the connections between MoE and neural memory despite their different formulations the surface. We first derive variant form of MoE establish its connection with sparse memory ($3.1). Then, we propose a unified framework for S-FFN ($3.2).
'Conventionally, SMoE implements load balancing con- straints to prevent the overuse of certain experts and the under- utilization of others, and avoid convergence to local optima (Shuzcer et al., 2017; Eigen et al., 2014).

---

3.1 A Neural Memory View of MoE
the core of S-FFN computation, the sparse memory directly uses the memory parameters for selection - the dot product between input token x and key vectors k, whereas SMoE depends on a separately parameterized gate g-
MoEs use a gating function to estimate the im- portance of all experts and combine each expert's output through linear combination. Here, inspired by the memory view on FFNs ($2), we could view MoE a wide neural memory chunked into B FFNs:
3.2 The Unified Framework
We propose general framework that unifies the two different approaches to achieve S-FFN. We view both as instances of a memory with large key and value table - K E Rd xdz V E Rdmxd where dm >> 4 d. We distinguish the different methods along two dimensions illustrated below and summarized in Table 1:
B-1 MoB(x)= = Σ
(x)
FFN (x)
1=0
B-1
dm
Σ (x)
(i)
(
m
i=0
j=0
8 1 dm 1
Bi(x)) - m())) .v
Memory block size specifies how many memory cells share the same relevance weight at selection time, and thus together treated as a memory block. We use g to denote the size of one block. In other words, we split the K, V along the -n-dimmnnion into g-size blocks. Therefore, memory consists of B = dm/g blocks in total. Formally, we write
1=0
j=0
Bdd
Σ
[i(x)-mg) (0) v,
Bddm-
m Vln
(4)
where FFN() denotes the i-th FFN.
K = ()) KO .... (((1))) C R Rdn
Based on linear algebra, we are able to write the standard MoF formulation ($2.2.1) in a similar summation form as that of neural memory (Eq 4; 2.1). MoE in its neural memory form has value ta- ble V [V()).... v(B-1)] E BB-dm xd the concatenation of value tables vii) E Rdmxd
V V v().... v(B-1) Rduxdee
For example, sparse memory has block size g =
1 - trivially treating memory cell as a "block"; and SMoE has the block size 9 = 1 ($3. I). Current approaches generally use fixed block sizes, but this is mostly artifact of how the methods were derived rather than mathematical constraint. For example, we can design SMoE versions instead of 1 expert of size 4 d. or uses experts of size d. We can similarly chunk memory coefficients m into blocks of size in sparse memories.
from all FFN( In this view, Vi with / = i dm +j corresponds to the j-th value vector in the i-th chunk of value table v(i Thus, its correspond- ing memory coefficient my = (x( - m is pro- duced by weighting the j-th memory coefficient (2)
of FFN(), m; (2) x (a) by the relevance score of (FN(), g (x). Through this memory view, one could see that sparse MoE(x) is sparse mem- ory operating in terms of B memory blocks; it uses the gate g(x) to narrow down the calculation over the stacked value tables V to the value tables from FFN()) for i E subset(g( gxx) (i.e. sparsify).
Memory block selection method is the specific function that compute the relevance of each mem- ory blocks for selection. Since SMoE is also a type of sparse memory. we distinguish the selec- tion method by a new criterion - whether one allows input x to directly interact with the key ta- ble K", As discussed in $3.1, SMoE uses the es- timation from an individually parameterized gate to select, while sparse memory solely and directly uses a key table. Thus, current SMoE is a type of indirect selection method, and sparse memory a direct one. Various SMoEs are further character- ized by whether their gating function has learned parameters or consists of a static mapping ($2.2. 1).
Comparison with Sparse Memory Both SMoE and sparse neural memory are neural memory, but there are several differences: 1) whether memory cells share the same relevance weight: in sparse neural memory, each memory cell receives an in- dividual weight mi. In contrast, in SMoE, each group of memory cells shares the same rele- vance weight gx().. 2) memory selection criterion: if we center the key-value computation (Eq. 2) at

---

Memory block
Memory block selection method
Model Name
size (g)
Full Fall-parameter Key
VanillaM
Direct
Low-rank Key
LoRKM, PKM (Lample et al., 2019)
Switch Transformer(Fddus al., 2021),
GShard (Lepikhin al., 2021),
Learned gate
GLaM (Du al., 2021).
Indirect
BASELayer (Lewis et al., 2021),
X-MoE (Chi al.. 2022)
HushLayer(Roller al., 2021).
Static gate
DEMix (Gururangan et al., 2021)
Table I: S-FEN methods decomposed along the defined design dimensions.
Meanwhile, sparse memory is characterized by how much factorization the key table uses ($2.2.2).
Due to the linearity of averaging, the operation e; x is equivalent to calculating the average of dot products within block without GeLU. Since all tokens share the of averaged key vectors, our method is efficient. In summary, Avg-K mi- grates MoE's advantages: a gate design, and a full pall-araameteieed key table in each expert. As an advantage of sparse memory, it uses the average key vectors of each expert to route tokens and thus increases the memory selection method's depen- dency on the key table. We provide more rationale for our choice of average function in Appendix D.I.
3.3 A New Selection Method - Avg-K
The gate design in the Mixture-of offuxert methods ensures that all experts are activated for rout- ing tokens. Without the gate design, conventional sparse neural memory methods ($2.2.2) uses the full key table before sparsifying the computation in the value tables, which explains the increasing com- putation cost for scaling up sparse neural memory, As later shown in Fig. 3 ($5.1), this additional com- putation doesn't 't bring proportional improvement its performance, compared with MoE.
Experiment Setup
However, there is still merit in sparse neural memory that MoE could acquire. Based on the contrastive analysis from $5.2, we found that when the model makes more use of each expert's 's key table for routing tokens, the more its performance will increase. This is precisely the deficiency in the conventional Mixture of Experts (MoE) model, it relies separately learned parameter static gating mechanism ($2.2.1).
4.1 Models
We choose Dense Baseline using transformer chitectures used in GPT-3 models (Brown et al., 2020), which has 24 transformer layers, with 1024, f = GeLU activation function, and with a memory size (or FFN hidden size) to be d. This model is also referred to as the "base model" be- cause S-FFN's size is chosen based on the configu- ration of this model. Our choice of architecture size leads base model with 355M parameters (See Appendix A). Our reason for the chosen model size two-fold: 1) this size is similar to the community- estimated size of OpenAI text-ada ada-001; ; 2) in- dicated in Clark al. (2022), 355M the smallest size that separates the performances of different architecture designs.
To this end, we propose new routing method Avg-K- as a hybrid design choice between spanse neural memory and MoE methods. Avg-K repre- sents each block with the average of its key table (i) E poxd along g-dimension:
1 9
Σ 0=
k i
Avg (1), dim=0)
e:
Then, we use the dot product between x and the averages to select the top-b selected block and route the token there for memory calculation (Eq. 2):
S-FFN Given model above, we replace some of its FFNs with S-FFN. Similar to (Lepikhin al., 2021), we replace the FFN at every 6 layers (layer 5, 11, 17, 23, indexed from 0), leading to 4 S-FFNs in total across 24 layers. We use denote the number of memory blocks used and control how
[1 {
E {top-b of (eg. , x]
g: (x)
otherwise

---

activated the S-FFN is. We use the formulation of dm E- (4- d) to control the size of S-FFN,
indirect block selection methods lead to lower per- plexity than a standard FFN, but which type of method has better PLOPs-Perplexity trade-off and what are the relative efficacy and efficiency of dif- ferent methods require further study ($5.2).
so the S-FFN will activate b = &o out of B
dm memory blocks. In Table 2, we list all S-FFN models used for analysis in $5. We count FLOPs analytically following Narayanan et al. (2021) and do not account if a worker finishes computation before another (when using model parallelism). We use the number of learnable parameters to consider whether two models equally expressive. In Table 2, we list all S-FFNN models used for analysis in $5. For experimentation of our Avg-K, we don't t enforce load balancing for simplicity.
5.1 Memory block size
Since block size is natural number, we aim to answer a straightforward question - given a fixed number of active memory cells k, does smaller memory block size lead to lower per- plexity? We use simple and robust selection meth- ods to disentangle the impact of hyperparameter choices. Specifically, we use random hash as rec- ommended in HashLayer (Roller al., 2021) (de- noted RandHash) for indirect block selection and exact top-k memory block (denoted VanillaM) for direct block selection. For all experiments, we use E 16.
PKM-FEN Since the factorized key table in PKM has little (< 0.3%) learnable parameter rel- ative to the value table, we propose an indirect variant called PKM-FFN to match the number of parameters of other models like RandHash. This variant has memory block size = 1 and the same key-value table as RandHash. PKM-FFNN has a gate whose g x)) is the same as the m from PKM and gi m, and no load-balancing is enforced.
RandHash randomly selects k/g unique memory blocks among all B dm/g blocks - -essentially sampling unique values from Uniform( ... B - 1].. Originally, with block size 9 = 4096, RandHash assigns token to 4096/ 4096 = 1 block; with block size 2048, 1096/2018 2 blocks.
4.2 Language Modeling
Pretraining Data We pretrain all S-FFN models on a total of 453GB text with 1 112B tokens from a union of six English-only datasets, including En- glish subset of CC100 and the five datasets used to pretrain RoBERTa (Liu al., 2019) specif- ically BookCorpus, English Wikipedia, CC-News, OpenWebText, CC-Stories (details in Appendix A.3). We adopt the same Byte-Pair Encoding as GPT-2 (Radford et al., 2019) and RoBERTa (Liu et al., 2019) with a vocabulary of 50K subword units. All models are trained for 60B tokens.
VanillaM originally has a block size g - and selects top-k scalars in memory coefficients m = GeLU(x T We made a minimal change to extend it to larger block size g: given m, we chunk it into B blocks then, we select the top-h blocks using the average of each block:
m [m:").... (B-1)];
Evaluation settings We evaluate our models' ability to predict the next token in sequence as measured by perplexity. We report both in-domain and out-of-domain perplexity to indicate general- ization ability. For out-of-domain, we use valida- tion data from The Pile (Gao et al., 2020), public dataset that combines data from 22 diverse sources.
Avg(GeLU(x: (K())) dim=0)²
In Fig. 2, we observe that smaller block size leads to an improvement of 0.4(15.75 + 15.35) perplexity for RandHash and improvement of 0.87(15.56 + 14.69) for VanillaM.
In Appendix C.2, we provide theoretical justi- fications for this observation which shows that a smaller block size improves model capacity by in- cluding more combinations of memory cells. For example, with g/2, half memory cells of expert-1 could be activated together with half of the expert- 2; however, this combination is impossible with a larger block size.
5 Analysis Results
In this section, we use the proposed unified view to systematically study the design choice of S-FFN. Specifically, (I). we study wide range of block sizes other than the incidental choice used in ex- isting work and investigate its impact on language modeling perplexity 1). (2). Both direct and

---

Selection
Method name
g
method lype
VanillaM
(1,64,266, 1024, 2048, 1096) ($5.1)
Direct
LoRKM
[1]
PKM (Lample al., 2019)
[1}
{4, 16,
{4096,
RandHash (Roller al., 2021)
{1, 64, 256, 1024, 2048, 4096) ($5.1)
(32)*)
(8192)")
Indirect
Switch (Fedus et al., 2021)
{4096)
PKM-FFN ($4.1)
(1}
Table All the S-FEN models used in experiments and analysis in g is the number of memory cells grouped in a memory block, k is the active memory cells, and E control the sizes of memory dm. E - (4 - d). Some settings(**) are only used for PKM.
VanillaM
RandomHash
Domain = Outmof-Domain
Domain = In-Domain Val.
Domain = In-Domain Train
17
Dense Baseline
17
Dense Baseline
Dense Baseline
Perplexity( 4)
19
16
16
15
15
64
256
1024
2048
4096
64
256
1024
2048
4096
64
256
1024
2048
4096
Figure 2: Perplexity (lower the better) consistently improves as memory block size decreases for both direct (VanillaM) and indirect (RandHash) selection method in S-FFN models. Ranking individual out-of-domain test sets generally follows the ranking by average perplexity (c.g. 20 out of 22).
5.2 Memory block selection method
Effect of gating function We start with trastive comparisons among PKM-FFN E=16 3= RandHash with memory block size and 4096 active memory blocks. From the three parameter-matched models, we can learn important lessons to improve the design of the gate:
Next, we investigate the impact of the selection method, specifically, the FL.Pstperllexity trade- off for direct and indirect methods determine the overall usefulness of each S-FFN method.
Preoseepyyyyxxtty trade-off We study the effi- ciency of direct and indirect selection methods in S-FFN models characterized by FLOPS- trade-off. We conduct experiments across differ- ent scales of the memory by varying E E (1,16) additionally, we run E 32 for PKM.
1. Comparing with PKM-FFN PKM R=32 essentially moves the parameters from a full- parameter key table to double the size of the value table.
2. PKM-FFNE=16 and RandHash E=16 have the same (in size) key and value tables. But the former uses gate jointly learned with key table, while the latter uses learning free gate.
In Fig. 3, we marginalize different factors used in the two selection methods - i.c. types of gates, factorization techniques on key table, etc. and consider each type of selection method whole. When change different marginalized factors, observe that indirect methods tend improve more as we use more FLOPs (with larger memory sizes controlled by E). Thus, the indirect method has a better Prrppeeetyyy trade-off.
As shown in Table 3, on tit-of-dommin, PKM- FFN outperforms PKM by 0.87 perplexity and slightly outperform RandHash E=16 by 0.16. Therefore, it is essential to have full- parameter, and thus expressive enough, key ta- ble to produce memory coefficients.
Table 3 shows the improvement E=11 PKM-FFN RandHash 1=16
*Avg(-) performs better than other simple aggregators Min(.), Max(-), and Avg(l []: see ablations in Table 7.

---

Domain Out-of-Domain
Domain In- Domain Val.
Selection method type
17.5
Direct
17.0
Indirect
17.0
Memory size controlled by E
16.5
Perplexity(
16.5
16
32
16.0
15.5
15.5
15.0
15,0
0.20
0.25
0.30
0.20
0.25
0.30
Total Train FLOPs (ZFLOPs)
Total Train FLOPs (ZFLOPs)
Figure 3: Pr.Pexperllexity trade-off of indirect block selection is better than that of direct block selection. Indirect methods (orange cross) have more perplexity improvement relative to increases FLOPs than direct methods (blue dots). See more detailed legend (e.g. include methods like LoRKM) in Fig. 5.
Active
Selection
#Parameters
Train
Out-of- Mikoo--Dmmiin Avg. (1)
In-Domain (1)
Name
memory cells (A:)
method type
(Entire Model)
ZFLOPs
(See Table for each domain)
Train
Val.
Dense Baseline
4096
354.7M
0.212
16.96
19.60
17.16
8-M
4096
590.2M
0.205
16.66
19.45
16.87
P=M
4096
858.7M
0.205
16.06
Direct
EKM
8192
858.7M
0.213
16.16
19.05
16.45
VanillaMF= F=
4096
858.3M
0.333
14.69
17.48
14.90
PKM-FFN 16
4096
858.9M
0,213
15.19
17.82
15.48
Indirect
RandHash 1=
4096
858.3M
0.212
15.35
17.88
15.45
Table 3: List of experiments for contrastively comparing designs. This table assume each memory cell is memory block, i.e. g = The top two best performing models (bolded) have full-parameter key table depend more on dot product to activate parameters. Ranking individual out-of-domain test set generally follows the ranking by average perplexity (e.g. 21 out of 22).
over Dense Baseline (16.96) are 2.27, 1.77, and 1.61 respectively on out-of-domain. They only differ by how much they depend on the key table for selection VanillaM directly uses it, PKM-FFN indirectly gains information from it, and RandHash completely ignores it. Due to the consistent observation across out-of-domain test sets, we conclude that the more dependent on the key table the selection method is, the better language model it will lead to; and indirect usage (PKM-FFNN is not enough.
Language Modeling Pretraining. Table shows that the proposed S-FFN design outperforms all other indirect methods. With < 1% additional FLOPs, Avg-K achieves 2.16 lower perplexity than Dense Baseline (16.96), outperform Fedus et al. (2021) by 1.65 and Roller et al. (2021) by 0.5.
Analysis Fig 6, although VanillaM increases its performance from block size = 4096 to 256, the improvement is less significant than that of Avg-K. The comparison suggests that, with larger block size, GeLU activation protects the average operation in VanillaM (after GeLU) af- fected by (potentially many) large negatives be- cause lim, so GeLU((z)
5.3 Performance of Avg-K
In contrast, this "negative value" problem is mitigated by using a smaller block size, due to more blocks available for
We benchmark our proposed Avg-K approach ($3.3) in this section.

---

#Parameters
Train 7.FLOPs
Out-of-Domain Avg. (L)
In- In-Domain (})
Selection method
(Entire Model)
(See Table 10 for details)
Train
Val.
354.7M
Dense Baseline
0.212
16.96
19.60
17.16
RandHash
4096
0.212
15.75
18.78
16.26
(Roller al., 2021)
0.212
15.35
17.88
15.45
Switch
4096
0.212
16.45
18.20
16.00
858.3M
(Fedus et al., 2021)
PKM-FIN
0.213
15.19
17.82
15.48
4096
0.212
16.44
19.04
16.59
Avg-K
256
0.213
14.91
17.57
15.19
64
0.214
14.80
17.51
15.11
Table 4: Avg-K out-performs other indirect block selection methods. Switch transformer is trained with the load balancing loss to prevent model degradation (Shazeer et al., 2017; Eigen et al., 2014). Ranking on individual out-of-domain test set mostly follows the ranking by average perplexity (e.g. 21 22).
selection. Since negative dot products affect Avg-K more, prefers blocks with more or very posi- tive dot products, whereas VanillaM is not shielded from extreme negatives, so might fail to detect those blocks. Therefore, Avg-K could achieve an even slightly better perplexity than VanillaM for block size g < 256. See full discussions in Ap- pendix D.2.
efficiency. For example, one could process the un- factorized key table by ANN methods like FAISS (Johnson et al., 2021) and ScaNN (Guo et al., 2020). One successful example is applying vanilla Locality-Sensitiin Hashing to Reformer (Kitaev et al., 2020). However, in our preliminary study, found that perplexity is greatly affected by the search quality, and building data structure after every update is expensive and hard avoid. We leave detailed discussion Appendix F.2.
In Figure 7, we also include load balancing analysis of Avg-K. To our surprise, the mode col- lapse (i.e., imbalanced usage of memory block) issue in Avg-K still exists as load balancing loss is not enforced. Given its superior performance, this suggests that Avg-K learned good representation for each memory block despite the disadvantage of load imbalance.
Conclusion
We provide a unified framework for designing sparse FFN in transformers and analyze existing S-FFN methods such as MoEs in the language mod- eling task. Using this framework based on sparse neural memory, we found that smaller memory block (c.g. expert) size improves perplexity at the cost of slightly higher computation cost. Selection methods with gates have better Perpl-terllexity trade-offs than without, while the gating function in current MoEs is sub-optimal. This framework en- ables to instantiate simpler S-FFN architecture that outperforms MoEs while still being efficient in training and inference.
6 Related Work
Excluded S-FFN Terraformer's Jaszezur et al. (2021) technique on FFN is closest to our PKМ-FFN because there is a low-rank learned gate to operate each memory cells for selec- tion. However, we exclude this method because framework uses all memory cells in each block, but Terraformer selects cell in each memory block (see our study Appendix I). In finetuning sce- narios, Zhang et al. (2022b) studies the connection berween FFN and SMoE by turning trained FFN into experts and separately learning gate. In con- trast, we focus on pretraining from scratch.
Limitations
Limitations of a smaller block size 9 With model parallelism (Lepikhin et al., 2021), multiple GPUs contains different memory block and paral- lelize the calculations. If with block size & d, token only routed memory block on one device, each device doubles its chance to receive more tokens with block size g = d. Therefore. each GPU processes more tokens and requires more
Approximate Nearest Neighbour (ANN) search One might wonder whether ANN techniques could help to search for the best key in VanillaM rather than trade the expressiveness of the key table for

---

compulation time, but we didn't measure the wall time difference in our work. Better implementa- tions could be developed to make a smaller block size for practical usage. In addition, since each memory block has its representation stored in the gating function, the smaller block will lead to more block representation stored in the gate, e.g,, more learned parameters in the learned gate and larger table for the static gate. Although RandHash with memory block size g 4- d cost essentially the same with memory block size g = 1, computing g(x) for learned gates requires more cost (details in Appendix C.3.1).
Ethics Statements
Due to the nature of pretraining, the carbon foot- print of our work is large estimated by the amount of FLOPs and GPUs reported in the paper. We did make the effort to minimize the cost at design stage of the project. In our preliminary study, we ask for recommendation from one of the authors of Artetxe et al. (2021) to choose and verify the minimal model size and amount of tokens that suf- ficiently differentiate different design choices.
Another ethical concern of the paper is from the pretraining data we use. As we used the same data source as Artetxe et al. (2021), we refer the reader to the ethics statements in Artetxe al. (2021) for how much trained model absorbs bias and toxicity from training data.
As discussed in Appendix C.3.2, smaller mem- ory block size will induce higher communication cost given the current all_to_an to_all-baasee imple- mentation framework (e.g. Switch Transformer). We think reducing memory block size to 1 is too extreme to be practical; and there should be sweet spot between and 4096 (or the chosen expert size) allowed by the implementation and hardware sta- tus.
A Acknowledgement
We would like to thank (in random order) helpful feedbacks from Suchin Gururangan, Xiaochuang Han, Luke Zettlemoyer, Noah Smith, pre-doctoral members of Noah's Ark, and anonymous reviewers. In developing the experimentation of this work, great help has been received from Jingfei Du, Susan Zhang, and researchers at Meta AI. We would also want to thank Zhaoheng Billy Li for his help in developing analytical formula in Appendix C.2.
Limitations of the unified framework Since our method Avg-K essentially applies an average pool- ing to the key table K", a better alternative may exist. Our method also heavily depends on dot product information, but this might not be the best information be used. Due to the curse of di- mensionality, future work might want to focus on finding a better metric than dot product and other aggregation methods than average to measure dis- tance between high-dimensional vectors.
References
Mikel Artetxe, Shruti Bhosale, Naman Goyal, Todor Mihaylov, Myle Ott. Sam Shleifer, Xi Victoria Lin, Jingfei Du, Srinivasan Iyer, Ramakanth Pa- sunuru, Giri Anamtharaman, Xian Li, Shuohui Chen, Halil Akin, Mandeep Baines, Louis Martin, Xing Zhou, Punit Singh Koura, Brian O'Horo, Jeff Wang. Luke Zettlemoyer. Mona Diab, Zomitsa Kozareva, and Ves Stoyanov. 2021. Efficient large scale lan- guage modeling with mixtures of experts. CoRR. abs/2112 12.
Also, we didn't 't train Avg-K with load-balancii balaneing due to our current limit in budget, but we in- clude our rationale in Appendix D.3 for why Avg-K should work with load bald-baanning..
Additionally, in large-scale SMoE training, the speed is limited by the most heavy-loaded GPU when model parallelism is used. Therefore, load balancing is essential. We also note that our scale is relatively small and does not use model paral- lelism, so the problem is not pronounced for us. Future follow-up should look at how to incorpo- rate load balancing into the unified framework and inspire better actionable design choice. We think such unification requires more advanced theoreti- cal connection with memory block and block selec- tion method, which likely involves consideration of training procedure.
Sebastian Borgeaud, Arthur Mensch, Jordan Hoff- mann. Trevor Cai, Eliza Rutherford, Katie Milli- can, George Bm Van Den Driessche, Jean- -aptiste Lespiau, Bogdan Damoc, Aidan Clark, et al. 2022. Improving language models by retrieving from tril- lions of tokens. In International Conference on Ma- chine Learning, pages 1206-2240 PMLR.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss. Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu. Clemens Winter, Christopher Hesse, Mark Chen, Eric

---

Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. 2020. Language models are few-shot learners. In Ad- vances in Neural Information Processing Systems 33: Annual Conference on Neural Information Process- ing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual.
Barret Zoph, Liam Fedus, Maarten Bosma, Zong- wei Zhou, Tao Wang, Yu Emma Wang, Kellie Web- ster, Marie Pellat, Kevin Robinson, Kathy Meier- Hellstern, Toju Duke, Lucas Dixon, Kun Zhang, Quoe V. Le, Yonghui Wu, Zhifeng Chen, and Claire Cui. 2021. Glam: Efficient scaling of language mod- els with mixture-of-experts, CoRR, abs/2 12.0695..
David Eigen, Marc' Aurelio Ranzato, and Ilya Sutskever. 2014. Learning factored representations in deep mixture of experts. In 2nd International Conference on Learning Representations. ICLR 2014, Banff. AB, Canada, April 14-16, 2014, Workshop Track Proceed- ings.
Zewen Chi, Li Dong, Shaohan Huang, Damai Dai, Shuming Ma. Barun Patra, Saksham Singhal, Payal Bajaj. Xia Song, and Furu Wei. 2022. On the rep- resentation collapse of sparse mixture of experts. CoRR, abs/2204.09179.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts. Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vin- odkumar Prabhakaran, Emily Reif, Nan Du Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengeheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, An- drew M. Dai, Thanumalayan Sankaranarayana Pil- lai, Marie Pellat, Aitor Lewkowycz, Erica Morcira, Rewon Child, Oleksandr Polozov. Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. 2022. Palm: Scaling language mod- cling with pathways. CORR, abs/2204.02311.
William Fedus, Jeff Dean, and Barret Zoph. 2022. A review of sparse expert models in deep learning.
William Fedus, Barret Zoph, and Noam M. Shazeer. 2021. Switch transformers: Scaling to trillion param- cter models with simple and efficient sparsity. ArXiv, 0as/2101.0396.
Leo Gao, Stella Biderman, Sid Black, Laurence Gold- ing, Travis Hoppe, Charles Foster, Jason Phang. Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. 2020. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00022
Mor Geva, Roci Schuster, Jonathan Berant, and Omer Levy. 2021. Transformer feed forward layers are key- value memories. In Proceedings of the 2021 Confer- ence Empirical Methods in Natural Language Pro- cessing, EMNLP 2021, Virtual Event Punta Cana, Dominican Republic, 7-11 November, 2021, pages 5484-5495, Association for Computational Linguis- tics.
Aaron Gokaslan and Vanya Cohen. 2019. Openwehtext corpus. http://west archive, sttt/p///tttp:////// skylion807 - git.....uu...... too/ppeneebTextCorps..
Aidan Clark, Diego de Las Casas, Aurelia Guy. Arthur Mensch, Michela Paganini, Jordan Hoff- mann, Bogdan Damoc, Blake A. Hechtman, Trevor Sai,Sebbstian Borgeaud, George van den Driessche, Eliza Rutherford, Tom Hennigan, Matthew Johnson, Katie Millican, Albin Cassirer, Chris Jones, Elena Buchatskaya, David Budden, Laurent Sifre, Simon Osindero, Oriol Vinyals, Jack W. Rac, Erich Elsen, Koray Kavukcuoglu, and Karen Simonyan. 2022. Unified scaling laws for routed language models. CoRR, abs/2202.01 169.
Ruiqi Guo, Philip Sun, Erik Lindgren, Quan Geng. David Simcha, Felix Chem, and Sanjiv Kumar. 2020. Accelerating large-scale inference with anisotropic vector quantization. In International Conference on Machine Learning.
Suchin Gururangan, Mike Lewis, Ari Holtzman, Noah A. Smith, and Luke Zettlemoyer. 2021. Demix layers: Disentangling domains for modular language modeling. CoRR, abs/2108.05036.
Jacob Devlin, Ming- Hei Chang, Kenton Lee, and Kristina Toutanova. 2019. BERT: pre-training of deep bidirectional transformers for language under- standing. In Proceedingy of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Tech- nologies, NAACL-HLT 2019. Minneapolis, MN, USA. June 2-7, 2019, Volume I (Long and Short Pupers), pages 4171-4186. Association for Computational Linguistics.
Jordan Hoffmann, Sehastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford. Diego de Las Casas, Lisa Anne Hendricks. Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George Damoc, Aurelia Guy, Simon Osindero, Karen Si- monyan, Erich Elsen, Jack W. Rac, Oriol Vinyals, and Laurent Sifre. 2022. Training compute-optimal large language models. CoRR, abs/2203. 15556.
den Driessche, Bogdan
Nan Du, Yanping Huang, Andrew M. Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat,
Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan, and Geoffrey E. Hinton. 1991. Adaptive mixtures of local experts. Neural Computation, 3:79-87.

---

Sebastian Jaszezur, Aakanksha Chowdhery, Afroz Mo- hiuddin, Lukasz Kaiser, Wojciech Gajewski, Henryk Michalewski, and Jonni Kanerva. 2021. Sparse is enough in scaling transformers. In Advances in Neu- ral Information Processing Systems 34: Annual Con- ference on Neural Information Processing Systems 2021, NewrIPS 2021, December 6- 6-14, 2021, virtual, pages 9995-9907.
megatron Imm In SC '21: The International Con- ference for High Performance Computing, Network- ing. Storage and Analysis, St. Louis, Missouri, USA, November 14 19, 2021, pages 58:1 58:1-58:15. ACM.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf. Edward Z. Yang. 7achary DeVito, Martin Raison, Alykhan Te- jani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. 2019. Pytorch: An imperative style, high-performance deep learning li- brary- In Advances Neural Information Processing Systems 32: Annual Conference on Neural Informa- tion Processing Systems 2019, NewrIPS 2019, De- cember 8-14, 2019, Vancouver, BC, Canada, pages 8024-8033 8035.
Jeff Johnson. Matthijs Douze, and Hervé Jégou. 2021. Billion-scale similarity search with gpus. IEEE Trans. Big Data, 733):535-547.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. 2020. Scaling laws for neural language models. CoRR, 0os/2001.08361
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. 2020. Reformer: The efficient transformer. In 8th International Conference Learning Representa- tions, ICLR 2020, Addis Ababa, Ethiopia, April 26- 30, 2020. OpenReview.net.
Alec Radford and Karthik Narasimhan. 2018. Im- proving language understanding by generative pre- training.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Nlya Sutskever. 2019. Language models are unsupervised multitask learners.
Guillaume Lample, Alexandre Sablayrolles, Mare' Aurelio Ranzato, Ludovic Denoyer, and Hervé Jégou. 2019. Large memory layers with product keys. In Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, December 8- 8-14, 2019, Vancouver, BC, Canada, pages 8546-8557.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. 2020. Exploring the limits of transfer learning with unified text-to-text trans- former. J. Mach. Learn. Res., 21:140:1-140:5:
Stephen Roller, Sainbayar Sukhbaatar, Arthur Szlam, and Jason Weston. 2021. Hash layers for large sparse models. CoRR, abs/2106.044.
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam M. Shazeer, and Z. Chen. 2021. Gshard: Scaling giant models with conditional computation and automatic sharding. ArXiv, abs/2006.1666 16668.
Roy Schwartz, Jesse Dodge, Noah Smith, and Oren Etzioni. 2020. Green ai. Communications of the ACM, 63:54-63 63.
Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zeutlemoyer. 2021. BASE layers: Simplifying training of large, sparse models. In Pro- ceedings of the 38th International Conference on Machine Learning. ICML 2021, 18-24 July 2021, Vir- tual Event, volume 139 of Proceedings of Machine Learning Research, pages 6265 6274. PMLR
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe V. Le, Geoffrey E. Hinton, and Jeff Dean. 2017. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In Sth International Conference on Learning Representa- tions, ICLR 2017, Toulon, France, April 24-26. 2017, Conference Track Proceedings. OpenReview.net.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Man- dar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. 2019. RoBERTa: A Robustly Optimized BERT Pretrain- ing Approach.
Sainbayar Sukhbaatar, Edouard Grave, Guillaume Lam- ple, Herve Jegou, and Armand Joulin. 2019. Aug- menting self-attention with persistent memory. arXiv preprint arXiv: 1907:09007
Sainbayar Sukhbaatar, Arthur D. Szlam, Jason Weston, and Rob Fergus. 2015a. End-to-end memory net- works. In NIPS.
Sebastian Nagel. 2016.
Co-news. http:
wtt.:r/r//sveetttp/////meeeee / org/2016/10//nes--dateeo petttaset-vvailabie
Sainbayar Sukhbaatar, Jason Weston, Rob Fergus, et al. 2015b. End-to-end memory networks. Advances neural information processing systems, 28.
Deepak Narayanan, Mohammad Shocybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Kor- thikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, Amar Phanishayee, and Matei Zaharia. 2021. Efficient large-scale language model training on GPU clusters using
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Met- zler. 2020. Efficient transformers: survey. ACM Computing Surveys (CSUR).

---

Tricu Trinh and Quoe 2018. A simple method for commonsense reasoning. ArXiv, abs/1806.02847,
with different random seeds and found the results are almost identical. Therefore, in the interest of time and carbon cost, we didn't 't run with different random seeds for the experiments in this paper.
Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention all you need. ArXiv, abs/ 170.03376..
A.2 Infrastructure
Jason Wei. Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph. Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, Chi, Tatsunori Hashimoto, Oriol Vinyals, Percy Liang, Jeff Dean, and William Fedus. 2022. Emer- gent abilities of large language models. CoRR, abs/2206.07682.
We used fairseq periments. For training. all the training are done on V100 GPUs. Each training takes 32 GPUs; rare circumstances, the number GPU is adjusted for training speed.
codebase to conduct ex-
A.3 Data
Guillaume Wenzek, Maree-Anne Lachaux, Alexis Con- neau, Vishrav Chaudhary, Francisco Guzm' an, Ar- mand Joulin, and Edouard Grave. 2020. Cenet: Ex- tracting high quality monolingual datasets from web crawl data. In LREC.
Here is detailed description of our pretraining corpus.
BookCorpus (Zhu al., 2015) consists of more than 10K unpublished books (4GB);
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher De- wan, Mona Diab, Xian Li, Xi Victoria Lin, Todor Mi- haylov, Myle Ott, Sam Shleifer. Kurt Shuster, Daniel Simig, Punit Singh Koura, Anjali Sridhar, Tianlu Wang, and Luke Zettlemoyer. 2022a. OPT: open pre-trained transformer language models. CoRR, Abs/2205.01066
English Wikipedia, excluding lists, tables and headers (12GB);
CC-News (Nagel, 2016) contains 63 mil- lions English news articles crawled between September 2016 and February 2019 (76GB);
Zhengyan Zhang, Yankai Lin, Zhiyuan Liu, Peng Li, Maosong Sun, and Jie Zhou. 2022b. MoEfication: Transformer feed-forward layers are mixtures of ex- perts. In Findings of the Association for Compu- kational Linguistics: ACI. 2022, pages 877-890, Dublin, Ireland. Association for Computational Lin- guistics.
OpenWeb (Gokaslan and Cohen, 2019). an open source recreation of the WebText dataset used to train GPT-2 (38GB);
CC-Stories (Trinh and Le, 2018) contains a subset of CommonCrawl data filtered to match the story-like style of Winograd schemas (31GB);
Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Y. Zhao, Andrew M. Dai, Zhifeng Chen, Quoe Le, and James Laudon. 2022. Mixture- of-experts with expert choice routing. CoRR, abs/2202.09368..
English CC100 (Wenzek al., 2020), dataset extracted from CommonCrawl snap- shots between January 2018 and December 2018, filtered to match the style of Wikipedia (292GB).
Yukun Zhu, Ryan Kiros, Richard S. Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. 2015. Aligning books and movies: Towards story-like like visual explanations by watching movies and reading books. 2015 IEEE International Conference on Computer Vision (ICCV), pages 19- 27.
B Product Key Memory
As mentioned in $2.2.2, LoRKM assumes that the full key table K E I dxdd is composed of and ap- proximated by downward projection D E Rdddd and a low-rank key table K E Rdmxdd
A Experimental details
A.1 Hyperparameters
Table 5 specifies shared hyperparameters across all experiments, in which Table 5a contains ones for training data, optimizer, and efficient infrastructure techniques; and Table 5b for architecture. Then, Table describes the hyperparameters specifically for Switch, Table 6b for LoRKM, Table 6c for PKM. In preliminary study, we train baselines
D DT
Lample al. (2019) further approximated the low-rank key table K. assumes that each low- rank key vector k, E Te is created from the con- catenation from two sub-key vectors c, c' E IR ;

---

(a) Shared configuration for data, optimizer, and efficient infrastructure
Name
Values
#Tokens for training
60c9
#Tokens for warm-up
375 1021
#Tokens per batch
0.5 10212
#Tokens per sample
2048
#GPU
32
GPU
NVIDIA Tesla V100 32GB
Optimizer
Adam [8s (0.9, 0.98), e le 8)
Weight Decay
0.01
Peak Learning Rate
3e-4
Learning Rate Scheduler
polynomial_decaa
clip_norm
0.0
DistributedDataParrllll backend
FullyShardedDDttaaaaa
memory-efficient-ffple
True
fp16-init-scale
4
checkpoint-actiivttios
True
(b) Shared configuration for architecture.
Name
Values
Objective
Causal Language Model (CLM)
Activation function(()
GeLU
Model dimension (d)
1024
dn of non-S-FFiN
4 1024
#Attention Head
16
#Layer
24
Dropout Rate
0.0
Attention Dropout Rate
0.0
True
Table 5: Shared configuration. This is also used for training base model.
(a) Switch
Name
Values
moe-gating-use--p33
True
0.01
moe-gate-loss-wt
i.e. CLM loss + 0.01- auxiliary loss (Fedus et al., 2021)
Divide expert gradients by
/# Expert B
(c) PKM
(b) LoRKM
Name
Values
Name
Values
de # key table ($2.2.2)
128
de
128
BatchNorm after x D
False
BatchNorm after x D
True
Table 6: Specific architecture configuration

---

and the sub-key vectors is from smaller and non-overlapped sub-key tables C, C' C Rvvmx x 2 Unlike LoRKM where key vectors are indepen- dent of each other, key vectors PKM have some overlaps with each other and have a structured shar- ing,
where Zy is the indices of selected memory block for token position æ, and similarly for Zy. RandHash, though, is an exception where uni- form sampling is used. Therefore, Eir] could also be analytically calculated for various 9. when as- suming tokens also uniformly distributed."
k, = [ & (mod vdm)
|
RRdd
1 LSS-FNN
[[
LS.FEN Elre]
One could exploit this structure to efficiently com- pute the memory coefficient One first calculates the dot product between sub-key tables and down- projected input individually and combines them with negligible cost to form the full dot product m E Id
x
No. of such
block assignments
m,
5 [i/v +s' (mod vdm)
i-1
1-1 B- b x
II 1=0
dt
B ー k
where
C RRaa
k=O
2
[e
Probability overlaps
Probability of non-overlaps
s' = t
(C)T BRV...
2 :
t= DE RETT
i cells
Block size
overlap
C.1 VanillaM with block size g 1
For VanillMM with block size g > 1, we also tried three other simple aggregation function, but they all under-perform Average. We show their results in Table 7.
In Fig. 4a, 4b, we evaluate our model with E 16 validation subset and calculate the estimations across various g. It is observed that less sharing happens as block size decreases. However, the empirical estimation for RandHash are relatively constant across granularity. We sus- pect this is due to the Zipf's law of tokens. Also, we note that the magnitude of E[r] are different for different methods. We defer the reason of this phenomena to future work.
C.2 Analysis of smaller block sizes
We first quantify the intuition -usage of model memory is more spread out" by number of activated memory cells shared between two random tokens E[r]. We define this quantity to be average of every S-FFN layer, reflect the overall behavior
C.3 Cost of smaller block sizes
[r
1
Σ E(re]
LS.FTN
C.3.1 Cost of gate
RandHash efficient for computation because a hash table thcoretically has time complexity O(1). In contrast, conventional learned gate 2.2.1 has an d-dimensional embedding for each memory block. Therefore, with total of B memory blocks, it has the time complexity of O(d B). In Table 8 we show how the FLOPs percentage of learned gate in single newward-bacwwand computation changes with respect to the change in memory block size, where we assume setup $4 is adopted.
where LS-FEN the number of S-FFN. Because block selection usually depends on a contextualized token embedding, it's hard to draw tokens an i.i.d. fashion. Therefore, we estimate the the quantity by evaluating the model on a validation set. We sample N token pairs from each sequence for estimation:
E(rr]
Σ
vall
N
sCval
N-1
Σ
II n nIy
IIn the formulae, we use - and x interchangeably for hetter presentation
(20(

---

#arameeters
Train
Out-of-Domain
In-Domain
Selection method
Aggregator
(Entire Model)
ZFLOPs
(22 domains; Avg. + Std.)
Train
Val.
Dense Bascline
354.7M
0.212
16.96 5.20
19.60
17.16
Avg()
15.56 4.62
18.33
15.87
Avg(
15.67 4.66
15.94
VanillaM
4096
858.3M
0.333
Max(-)
16.11 + 4.86
16.33
Min(-)
94.86 = 57.63
17.08
Table 7: VanillaM with different simple aggregators
Memory block size (g)
TFLOPs of
4096
2048
1024
512
256
128
64
32
1
4 learned gates
0.275
0.552
1.10
2.20
4.40
8.80
17.6
35.2
1124
(across 24 layers)
Entire model
1850
1850
1860
1860
1860
1860
1870
1890
2980
%
0.0149
0.0298
0.0595
0.
0.237
0.473
0.941
1.86
37.718
Table 8: FLOPs percentage of learned gate increases when memory block size decreases
C.3.2 Cost communication
negative value.
The conventional training framework of MoE (Fe- dus et al., 2021) depends on all_to_all all opera- tions (Paszke al., 2019) to route tokens to dif- ferent devices. One might expect the communica- tion cost remains the same if the number of device doesn't 't change. However, this assumes the tokens are identified by their type. In fact, the training framework further identify the routed tokens type by the experts it routed to. Therefore, the com- munication cost scales linearly with respect to the change in the number of memory block.
Also, we consider the experiment with max- pooled hidden states (i.e., Max(-)). This experiment shows that memory block hardly has single key- value cell that dominates underperforms Avg(.) and Avg( [). What makes worse, the max operation will overlook lots of hidden states at selection, but the overlooked hid- den states still contribute to the computation. contrast, the performance increases when we con- sider the average (or average of the absolute values) where every hidden state contributes to the decision. Although the situation is slightly different in Avg-K, the "max- mmax-pooledd" version of Avg-K will only over- estimate the hidden states information and the aggregated value won't be indicative of the hidden states used for computation.
others since Max(-)
D Avg-K
D.1 Rationale to use Avg in Avg-K
more,
We heavily base our choice on experiments with ag- gregators in VanillaM (in Table 7). From the exper- iments with average absolute value (after GeLU), we hypothesized that a positive feature is good at predicting the value of a label/token against all oth- ers. In contrast, negative value is good at negating the prediction of single token. As such, positive features are more predictive than negative ones. Al- though the situation might be different for Avg-K (before GeLU), we expect the selection will only be affected more because of the larger impact of
The last consideration we have is that the average function is linear. When we select experts, we use the dot product between input and averaged keys. Due the linearity, this value is equivalent taking the dot product between the input and every key and taking the average (See Appendix D.2). Thus, using this design choice saves great amount of computation compared with VanillaM, while keeping the neural memory analogy.

---

Layer
Layer
Layer
1400
1600
1200
1400
Etr
1200
1500
800
1000
1000
600
256
1024
2048
4096
Layer 23
Layer Average
1600
1100
1400
1000
EDJ
1200
VanillaM
900
1000
800
800
1024
2048
40916
64
256
1024
2048
4096
(a) VanillaM: empirical estimation from sampling
Hasd
Layer =
Layer
Layer
300
300
300
200
200
200
100
100
100
0-
128
256
512
1024
2048.
4095
Awerage
300
200
Analytical
Empirical
100
16
32
64
1024
2040
4096
128
236
512
1024
2048
4096
(b) RandHash: both analytical value from close form calculation and empirical value from sampling
Figure 4: Expected shared memory cells across various block size g
D.2 Avg-K analysis through comparison with VanillaM
select blocks with the average of dot products:
9-1
e; x
1=
Avg-K essentially applies an average pooling the unfactorized K2 create representation of each block. Due to the linearity of averaging, the oper- ation e; x is equivalent to calculate the average of dot products within a block before GeLU and
k
xx
g j=0
(K(²) T
Avg
x
_dim=o)
(Avg-K)

---

Selection method
Domain Out-of-Domain
Domain In- Domain Val.
PKM
17.5
VanillaM
17.0
Dense Beseline
Dense Baseline
RandomHashig
17.0
LoRKM
16.5
16.5
PKM-FFN
Perplexity
Switch
16.0
16.0
RandomHash(g 4
15.5
Memory size controlled by E
15.5
15.0
15.0
16
14.5
Dense Besaline
[Dense Baseline
32
Selection method type
0.20
0.25
0.30
0.20
0.30 Total Train FLOPs (ZFLOPs)
0.25
Total Train FLOPs (ZFLOPs)
Direct Indirect
Figure 5: PeePex-Pryleextty trade-off different models where direct/indirect methods are further distinguished by
model name.
In contrast, VanillaM uses GeLUU(5.11):
average
after
the same block size (Fig. 6). This means the model will accommodate small design changes.
Additionally, the requirement of a load balanc- ing loss is determined by the sparsity of gradients. If one "expert" gets updated with gradients while other "experts" are starved then the expert will be selected all of the time leading a form of mode collapse. fact, with the same memory block size (g = 4096), we are surprised to observe Avg-K (w/o load balancing loss; row 6 in Table 4) could still perform on par with Switch (w/ load balancing; row in Table 4). As our load balancing analysis suggests in Appendix D.4, when the number of ex- perts is small, the mode collapse issue in Avg-K is severe. This makes us more confident that Avg-K will perform better with standard load balancing loss added.
9-
GeLU kg) xx
(VanillaM)
Because GeLU is a non-linear function, average from Avg-K could be shared across tokens. In con- trast, VanillaM can't, and thus making Avg-K effi- cient.
In Fig we experiment both methods with var- ious g. We observe when g decreases from 1096, the perplexity of Avg-K drops more drastically than VanillaM. We believe this observation highlights the impact of GeLU, Because lim GeLU(x:)
2---0
0, it protects the average in VanillaM from some very negative values. Thus, Avg-K with larger in- cluded more and potentially very negative values to average over, and thus leads worse choices than ones made by VanillaM. On the other hand, when g decrcases, this "negative value" problem is mit- igated. When there are more blocks available for selection (smaller g), because negative dot products affects Avg-K more, it prefers blocks with more or very positive dot products; whereas, VanillaM is protected from negative value it fails to detect those blocks. Therefore, Avg-K with g S 256 could achieve an even better perplexity.
Therefore, we helieve that with the loss added, the model will accommodate and could still per- form competitively.
D.4 Avg-K load balancing analysis
On the same validation set as used in $C, we also conduct a load balancing analysis of mem- ory blocks. Fig. shows that Avg-K and VanillaM disproportionally used some memory blocks.
E Preliminary study for related work
D.3
Why Avg-K with load balancing should work?
F.1 Terraformer analysis
Comparing VanillaM and Avg-K, one would expect Avg-Kt be greatly affected by extremely negative hidden states (before GeLU). Yet, the final model with A could even outperform VanillaM with
Controller in Terraformer Jaszezur et al. (2021) uses a controller to score all memory cells and pre- select a subsets Controller(x) - for computa-

---

Avg-K
VanillaM
Domain Out-of-Domain
Domain = In-Domain Val.
Domain = In-Domain Train
17
Dense Baseline
Dense Baseline
Dense Baseline
19
Perplexity(
16
16
18
15
15
64
256
1024
2048
4096
64
256
1024
2048
4096
256
1024
2049
4096
Figure 6: Perplexity performance (lower the better) of Avg-K and VanillaM across various g. We observe large drop in perplexity when g decreases Avg-K and less so in VanillaM; and Avg-K slightly outperform VanillaM with g 256.
tion.
as a gating function is justified ($5.2), the decision choice of "chunking memory but only select the best memory cells" seems unmotivated. Thus, we exclude this design setup from our analysis.
Σ
f(x ki)
y
iCController(x)
This is closest to our PKM-FFN, since their con- troller is essentially gate with low-rank key ta- ble in LoRKM g(x) = (x D) I where DE Rdxde K E Rd.. xdz, and de < d. The differ- ence is that they additionally assume the estimation from gate (and memory) could be seen as chun- ked into blocks and only select top-1 memory cell scored by the controller from each blocks:
E.2 ANN
Since ANN is an approximation exact search, we propose randomly sabotage VanillaM, which uses the exact search. Given a k, we randomly swap n% of the top-k of memory coefficient m (exact search results) with non-top-k values (during training and validation). and has accuracy (100 n)% We call it Naive-ANN. This is meant to set up random baseline for ANN, because different ANN techniques might make systematic mistakes, rather than a random one. However, we believe this could still serve as a proxy and shed light on how affects performance. As we see in Fig. 9, the model quality is sensitive to the quality of ANN.
B
y 8(x)
Σ
(1) f(x k..) v
where j" = arg max, B(x) Therefore, their number of active memory cells is equal to dm/ / Similar to our contrastive pair of PKM-FFN and VanillaM, hypothesize "vanilla" version of their methods. Memory is chunked into blocks of size g K [ ( (B-1)) and sim- ilarly for V9, Then, one chooses the top-1 with x (K()) T. We call it VoniolaCon
In our preliminary study, we found building data structure after every update is expensive. This leads
to some critical drawback when we apply the tech- niques to model parameter. Although one could amortize the cost by periodically building, the out- dated data structure will lead to lower accuracy. If one chooses a hyperparameter that leads to higher quality, the cost of preprocessing and the corre- sponding search will be even higher. What makes worse, the current ANN methods' search either don't 't support speedup by using GPU, or is not very well-integrated with GPUs slower than calcu- lating the exact dot product with CUDA kernl..
B-1 y Σ [(x k(9) v/
where j* arg max,x (K(1) In Fig. 8.
compare VanillaController to VanillaM with g because the actual section is the level of g 1. We set in VanillaM the one determined by equation above. We observe VanillaM outperforms VanillaController.. Although the controller design

---

York aM
Fandliann
27
6.8-
0.3
3.0-
0.5
644
633
CA
C2
0.2
Sorted
Seroed
Served Indin
Sartod Indoo
Saroed
g
1096
Vori Ia
Tarciluu
0.33
0.4
6.02
0.2%
0.33
0.20
0.10-
0.35
0.2
6.16-
0.00
0.30
0.00
0.3
0.22
0.00
So 1042
Sanod r00
Sortac Indse
Sortac Index
Scried Indox
(b) 2048
Jayer
Layer- 31
Layer- 17
Layer 22
Lager- Avarage
2.00
0.325
0.35
0.200
3.00
0.30
0.05
3.0h
6.15)
0.59
ELEC2
(c) g
= 1024
Vork sM
Aarcliian
K
Layer
Lever
Liver
Laper Aetuve
0.06
0.35
0.35
0.00
0.02
0.04-
0.31
0.00
6.01
0.03
0.33
0.00
0.00
0.13
0.02
0.02
0.0%
0.02
20 noc
Sanad Toe
Sortec Indox
Sortec Index
Seroed Index
(d) .
256
York AM
Tarcltir
Kgg
0,5357
0.035
0,7359
5.050
0.6333
0.0022
0.3100
6.0000
6.0000
0.3322
0.390
60000
60053
0.0002
03322
03259
0,000
0,003
Sarad reeK
Scrid TCo
Sortac Irdee
Sotte tndex
Scried Indux
(e)
64
Figure 7: Load balaneing of VanillaM, RandHash, Avg-K. The height of bar represents the proportion of memory block usage with which the memory block are sorted (in descending order).

---

VaniliaM
VanilaContrallar
Domain Oul-of-Domain
Domain In- In-Domain Val.
Domain In-Domain Train
17
Denii Basaine
Diris Basaline
Derise Basalina
Perplexityy (3)
18
18
15
16
15
84
1024
4096
84
1024
4096
64
1024
4096
k
k
k
Average perplexity performance (lower the better)
ABDV
DU Mabaaaaa
Devai Erex,
255
170
11.0
155
150
Darann Sunsiaa
Denia Fe
Dacan 10
Dacin dudnber 20
215
Depir Hodortkes
220
Draan EaPoe
210
Poppodg11
130
255
215
20
160
Desia Steeriemee mnd7
Dence iPiPaam
150
Dare -Paa0 Milli
Davin Madil Crins
Due Aorrta
150
150
145
120
Due en U
Drsan - RC
Drean WMaged ce
Dernn YrAde
Dasu Dater
125
1075
150
155
155
132
495
4x
1004
KX
(b) Performance on individual domain in PILE perplexity (lower the better)
Figure 8: Perplexity performance (lower the better) of VanillaM (g=1) and VanillaControlleer with E 16.

---

Naive-ANe
Domain Dut-o--Domain
Domain In-Domain Val.
Domain In-Domain Train
Perplexty(
17.0
Danse Bascline
17.0
Basalra
19.5
Dansa Baseline
16.5
19.0
8
8
12
86B
8 Accuracy(% T
8
Accuracy(% )
Accuracy(%
Average perplexity performance (lower the better)
17.5
1T
1650
1912
17.2
11.8
1930
1838
5055
1T
Treday
Domah Stnut
240
Dets Oesein
445
2.
224
Deman OpenWet
Coraln PhiPoour
Camah Cets
StacCactunge
Delo Bicura
Dew BHRn
1515
TA43
137
1879
29
122
184
Canals UpPPO
Uburkun2
Wilppds.nn
Oovoi-orrah
A v1N
17:5
1025
1100
Nema Reston
1630
508
1.5
360
NGG
250
250
Aareromy 11
Amcers 1)
Amany/h 11
Ar T1
Amazayo 1:
(b) Performance individual domain perplexity in PILE (lower the better)
Figure 9: Perplexity performance (lower the better) of Naive-ANN accuracy with

---

Selection method
Direct
Indirect
Type
Dense
Selection method
PKM
VanillaM
PKM-FFN RandHash
Baseline
E
16
32
32
16
16
16
4096
4096
4096
8192
4096
4096
4096
ArXiv
12.39
12.20
11.82
11.89
10.75
11.05
11.22
Bibliotik
21.17
20.82
20.15
20.25
18.49
19.11
19.33
BookCorpus
18.90
18.61
18.09
18.19
16.80
17.26
17.49
CommonCrawl
18.98
18.68
18.09
18.20
16.68
17.23
17.40
DM_ MM_aamthaasics
10.27
10.34
10.05
10.28
9.70
9.72
9.91
Muron__mmais
17.51
17.23
16.67
16.68
15.55
15.90
16.18
EuroParl
18.35
17.79
17.01
17.05
14.48
15.50
15.03
FreeLaw
13.62
13.34
12.84
12.93
11.70
12.11
12.29
Github
7.01
6.91
6.67
6.68
6.08
6.25
6.37
Gutenberg
PG-19
23.14
22.61
21.83
22.03
19.88
20.74
21.07
Out-of-Domain
HackerNews
23.58
23.35
22.44
22.62
20.71
21.36
21.76
Out-of-
WIH_ExPrtter
21.87
21.45
20.59
20.77
18.81
19.48
19.69
OpenSubtitles
18.03
17.99
17.46
17.44
16.48
16.84
17.10
Open WebText2
16.45
16.15
15.60
15.68
14.19
14.73
14.74
PhilPapers
30.63
30.02
28.50
28.74
25.14
26.44
26.60
PubMed_Abstrand
18.88
18.50
17.71
17.92
16.11
16.75
16.90
PubMed Central
10.57
10.40
10.08
10.14
9.37
9.66
9.71
StackExchange
14.10
13.85
13.37
13.45
12.05
12.46
12.72
USPTO
14.28
14.07
13.58
13.68
12.55
12.96
13.09
Ubuntu_IRa
16.14
15.40
14.95
15.08
14.14
14.35
14.62
Wikipedia_er en
16.26
16.03
15.33
15.48
14.07
14.51
14.59
YoutubeSubtilles 10.98
10.86
10.41
10.41
9.48
9.87
9.77
Average
16.96
16.66
16.06 16.16
14.69
15.19
15.35
Table Detailed out-of-domain perplexity for Table 3. Best performance each domain is in bold. Relative ranking on each domain generally follows the relative ranking by averaged performance (i.c. last row).

---

Dense Baseline
Selection method
RandHash
Switch
PKM-FFFN
Avg-K
4096
4096
4096
256
64
ArXiv
12.39
11.42
11.22
12.32
11.05
12.09
10.99
10.85
Bibliotik
21.17
19.84
19.33
19.87
19.11
20.49
18.70
18.61
BookCorpus
18.90
17.94
17.49
17.85
17.26
18.33
16.86
16.82
CommonCrawl
18.98
17.84
17.40
17.70
17.23
18.42
16.89
16.81
DM_ MM__Maaematiiss
10.27
10.22
9.91
10.63
9.72
10.25
9.62
9.61
Enron_ Emails
17.51
16.70
16.18
17.36
15.90
17.20
15.65
15.55
EuroParl
18.35
15.55
15.03
19.63
15,50
17.38
15.32
15.13
FreeLaw
13.62
12.56
12.29
12.80
12.11
13.20
11.84
11.77
Github
7.01
6.51
6.37
6.96
6.25
6.89
6.18
6.12
Gutenberg_PG-11 P9
23.14
21.55
21.07
21.81
20.74
22.36
20.07
19.98
Out-of-Domain
HackerNews
23.58
22.30
21.76
22.59
21.36
22.95
20.82
20.65
NIH_ExPorter
21.87
20.22
19.69
20.99
19.48
21.09
19.09
19.01
OpenSubtitles
18.03
17.33
17.10
17.31
16.84
17.74
16.62
16.66
OeenWe2Text2
16.45
15.16
14.74
15.51
14.73
15.87
14.48
14.39
PhilPapers
30.63
27.53
26.60
30.84
26.44
29.51
25.90
25.72
Pustec__bstracts
18.88
17.39
16.90
18.36
16.75
18.24
16.34
16.33
PubMed_Cenm
10.57
9.94
9.71
10.28
9.66
10.30
9.50
9.45
StackExchange
14.10
13.04
12.72
13.78
12.46
13.73
12.23
12.13
USPTO
14.28
13.41
13.09
13.54
12.96
13.89
12.73
12.62
Ubuntu IRC
16.14
14.95
14.62
14.78
14.35
15.50
14.26
13.59
Wikipedia_en
16.26
14.98
14.59
15.67
14.51
15.73
14.23
14.12
YoutuheSubritles
10.98
10.06
9.77
11.25
9.87
10.59
9.73
9.67
Average
16.96
15.75
15.35
16.45
15.19
16.44
14.91
14.80
Table 10: Detailed out-of-domain perplexity for Table 4. Best performance on each domain is in bold. Relative ranking on each domain generally follows the relative ranking by averaged performance (i.e. last row).


Reviews:
[{"review_id": "JlCwLVIgeV", "reviewer": "Reviewer_M6uy", "summary": "", "questions": "-  Does the proposed Analysising framework cover some new variations or combinations of architectural choices? ", "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 appreciate M6uy recognizing our S-FFN framework as an innovative contribution and highlighting our clear writing. We provide further clarifications and context to address M6uy’s concerns about the scope and impact of our work. \n\n\n**Q1: Scope of Analysis**: The paper primarily focuses on the impact of memory block size and selection [...]. However, it **might not cover all potential variations or combinations** of architectural choices. Especially, the research landscape in natural language processing evolves rapidly. The paper's findings and conclusions **might become dated if not updated or considered in the context of subsequent advancements** in the field.\n\n**A1:** \n\nWith finite resources, it is not feasible for us to **cover all potential variations or combinations** of architectural choices. Therefore, we chose to focus on two key design choices of S-FFN, namely memory block size and selection method, which are widely used in sparse Transformers but yet lack thorough understanding. We conduct comprehensive ablations covering commonly used hyperparameters specific to these two design choices as well as their combinations. The empirical results are complementary to existing work on Mixture-of-Experts (MoEs) which usually only focus on one specific selection method. \n\nRegarding your concern that this work may be dated, we want to emphasize that focusing on understanding the vanilla FFN layer instead of a special instantiation makes our work stay relevant to all Transformer architectures which incorporate sparsity in the FFN layers. Given that S-FFN has been actively studied and its potential for scaling [3], we believe this work is a timely contribution to the field that provides complementary knowledge to existing work on S-FFN. \n\n**Q2: Scalability**: The paper **does not [...] explore [...] larger models or more complex language tasks**. [...] the scalability of the \\textbf{\\texttt{Avg-K}} method might need to be scrutinized further.\n\n**A2:**\n\n> Larger models\n\nOur experiments cover a representative range of scales — from 355M to 859M\n by parameter size and from 0.212 to 0.333 by ZFLOPs (Table 3). While we agree that this experiment cannot be extrapolated to models as large as Switch [2] or GShard [4], we think the scaling curve we provide gives the initial evidence that S-FFN with gate design achieves much better FLOPs-perplexity trade-off than S-FFN without gate (Line 429-443, Figure 3&5).\n\nAdditionally, the reason for the chosen base model size is two-fold:\n1. At the time of development, we followed  [1, 2] to determine our base model size, **the smallest standard size for base architecture in [1, 2] is about 222M**. When designing our experiment, we chose the base size of 355M, similar to the community-estimated size of OpenAI text-ada-001.\n2. Also informed by Fig 2 in [3], 355M is the smallest size that separates the performances of different architecture designs. \n\nWe certainly hope there will be more future work on the scaling behavior of our unified framework. Thanks for pointing out the inclarity on choice of size, we will clarify this in the camera-ready version.\n\n> More complex language tasks\n\nFor the evaluation of models, we found it **conventional to primarily use perplexity (i.e. auto-regressive language modeling task)** when developing or exploring model architecture design [1, 2, 3, 4, 5, 6, 7]. To gain more insight into the generalization capability, we also included the PILE datasets [9] to demonstrate the model's performance on text from 22 diverse domains. Additionally, considering the strong capabilities of auto-regressive/decoder models like GPT-3, “complex” tasks like question answering, coreference resolution, etc. could be unified by properly prompting the language model. This approach exhibits a strong correlation with the model's perplexity performance (e.g. Fig 2&4 in [8]). Therefore, as it's impractical to run every evaluation metrics, we believe our evaluation setup is sufficient for the purpose of our work.\n\n---\n\n> Does the proposed Analysis framework **cover some new variations or combinations** of architectural choices?\n\nYes, we included a new MoE design called Avg-K (details in Sec 3.3 and 5.3) with practical potential. This can be seen as a variation of regular MoE models where the routing mechanism is changed. In addition, we also included a new design called PKM-FFN (see more details starting Line 345, and Line 444) as a new variation only intended for analysis purposes.\n\n\n**References:**\n\n[1] Roller, Stephen, Sainbayar Sukhbaatar, and Jason Weston. \"Hash layers for large sparse models.\" https://arxiv.org/pdf/2106.04426.pdf\n\n[2] William Fedus, Barret Zoph, and Noam M. Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. https://jmlr.org/papers/volume23/21-0998/21-0998.pdf\n\n[3] Aidan Clark, Diego de las Casas, Aurelia Guy, Arthur Mensch, et. al. Unified Scaling Laws for Routed Language Models. https://arxiv.org/abs/2202.01169 \n\n[4] Dmitry Lepikhin et. al. GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding https://arxiv.org/abs/2006.16668\n\n[5] Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, Luke Zettlemoyer. BASE Layers: Simplifying Training of Large, Sparse Model\n\n[6] Suchin Gururangan, Mike Lewis, Ari Holtzman, Noah A. Smith, Luke Zettlemoyer. DEMix Layers: Disentangling Domains for Modular Language Modeling\n\n[7] Ofir Press, Noah A. Smith, Mike Lewis Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation\n\n[8] Mikel Artetxe, et al. Efficient Large Scale Language Modeling with Mixtures of Experts\n\n[9] Leo Gao et. al.The Pile: An 800GB Dataset of Diverse Text for Language Modeling", "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": "- Conceptual Framework: The paper introduces a unified framework of sparse neural memory, which provides an  innovative insight for analyzing the impact of memory block size and selection methods in S-FFN architectures. \n\n- The paper is generally clearly-written. ", "reasons_to_reject": "- Scope of Analysis: The paper primarily focuses on the impact of memory block size and selection methods within the context of S-FFN architectures, particularly Mixture-of-Experts (MoE). However, it might not cover all potential variations or combinations of architectural choices. Espeecially,  the research landscape in natural language processing evolves rapidly. The paper's findings and conclusions might become dated if not updated or considered in the context of subsequent advancements in the field.\n\n- Scalability: The paper does not extensively explore the scalability of the proposed approach to even larger models or more complex language tasks. As model sizes and task complexities continue to grow, the scalability of the \\textbf{\\texttt{Avg-K}} method might need to be scrutinized further.\n\n", "questions_to_authors": "-  Does the proposed Analysising framework cover some new variations or combinations of architectural choices? ", "paper_topic_and_main_contributions": "The reviewed paper focuses on the effectiveness and efficiency of large and sparse feed-forward networks (S-FFN), specifically Mixture-of-Experts (MoE). The study examines the two critical aspects of S-FFN: memory block size and memory block selection method, within a comprehensive framework of sparse neural memory.\nThe researchers compare various S-FFN architectures for language modeling, shedding light on their relative effectiveness and efficiency. They propose a unified framework to analyze the impact of memory block size and selection method. The central idea revolves around the notion of sparse neural memory, emphasizing the activation of specific FFN parameters based on input conditions.\nThe paper delves into the intricate design choices of S-FFN architectures, particularly MoE, showcasing how memory block size and selection methods influence model performance and efficiency. ", "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": "hWTo7cLYbW", "reviewer": "Reviewer_mMBN", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 3, "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": "**Q1:** For a comprehensive analysis of MoE, it is very important to explore the **communication overhead among the experts (all to all)**.\n\n**A1:**\n\nSince we have a large variety of implementations it is difficult to make a fair comparison between communication overhead between models since the overhead is mostly dependent on how optimized the implementation is. For example, we thought about blocking that is partitioned on a single node or GPU, such that there could be a “local routing” program to send tokens to different memory blocks efficiently which is much more efficient than all-to-all communication. However, doing the full optimizations for all implementations and block sizes is out of the scope of our work.\n\nWe agree this is an important limitation of our work as we discussed in the limitation section(Line 552-572). We want to additionally comment that, historically, the proof-of-concept of MoE in neural networks started with [1] and the engineering implementation that makes MoE more accessible and practical (directly) started with [2]. We think similar things happen here — analyzing the communication cost and proposing solutions to the problem might each deserve the work of a paper. \n\n**Q2:** For the new selection method (Avg-K), do you have **any insights that guide you** to determine this approach? It will be better if there is a deeper exploration.\n\n**A2:**\n\nYes. We summarize our insights in Line 300-308. Specifically, we found that the gate design of MoE provides a better performance-trainFLOPs trade-off (see Fig 3), so we need a gate. In sec 5.2, the contrastive comparison of different gate designs controls how “directly” the input vector and key table parameter can see each other’s information — from direct dot product (most direct) to hash table (least direct). In Line 473-474, we wrote “the more dependent on the key table the selection method is, the better”. And we also found that having a full-parameter key table is important (Line 450-463). \n\nTherefore, \n1. Avg-K has a gate design;\n2. We use aggregated key vector to dot product with input vector to make it as direct as possible;\n3. We didn’t apply any factorization trick to the experts’ key tables.\n\nWe will include these details in the updated draft.\n\n\n**References:**\n\n[1] Noam Shazeer et. al. Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer https://arxiv.org/abs/1701.06538\n\n[2] Dmitry Lepikhin et. al. GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding https://arxiv.org/abs/2006.16668", "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 authors propose an interesting research problem on the key design choices in different S-FFN approaches.\n\n2. The paper presents a comprehensive analysis of different S-FFN approaches and a unified framework for the key design choices.\n\n3. The authors propose a new selection method (Avg-K) and demonstrate its performance benefits through extensive experiments.", "reasons_to_reject": "1. For a comprehensive analysis of MoE, it is very important to explore the communication overhead among the experts (all to all).\n\n2. For the new selection method (Avg-K), do you have any insights that guide you the determine this approach? It will be better if there is a deeper exploration.", "questions_to_authors": "", "paper_topic_and_main_contributions": "In this paper, the authors propose an interesting question that lies in the essence of the S-FFN approaches. The authors establish a unified framework to characterize various S-FFN approaches, which feature two key dimensions - memory block size and memory block selection method. With the elaborate analysis, the authors demonstrate a single gate for S-FFN - Avg-K in both sparse neural memory and MoE. The extensive evaluation shows the proposed Avg-K outperforms the state-of-the-art approaches.", "excitement": 4, "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": "0jPql233SV", "reviewer": "Reviewer_mxPW", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 3, "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": "**Q1:** While [...] better FLOPs-perplexity trade-off, it is **unclear how Avg-K can provide better actual training speed** perplexity trade-off due to increased communication cost.\n\n**A1:**\n\nAvg-K on its own does not increase communication cost by much from Switch Transformer. During training, the increased communication cost of Avg-K comes from sending the averaged key of each expert to all other GPUs. In our experiments, this means, for each batch, the additional communication cost is #num_GPU x 1024 x #experts floating points. Usually this is less than one million FLOPs. According to Table 8 in appendix, this cost is negligible. During inference, when computing token-by-token generations, our selection method is cheap because the average can be pre-computed and thus reused.\n\nSo in short, Avg-K only has slightly higher communication cost than baselines and is very efficient during inference.\n\n\n**Q2:** It will be better to **compare [...] with a more recent MoE routing approach such as X-MoE**.\n\n**A2:**\n\nWe thank the reviewer for raising this concern. We indeed noticed X-MoE during the project, but there are a few reasons that we are not able to compare:\n1. X-MoE is a contemporary work towards the end of this project. Nevertheless, the even harder part for adding X-MoE to our experiment is that they didn’t release their code. Therefore, unfortunately, we are not able to compare shoulder-to-shoulder with them.\n2. In general, there are many S-FFN variants and thus it is hard to cover them all by nature. Even for [1] where they have much more GPU than us, the analysis focuses on 4 variants. In our work, we tried our best to broaden the coverage and covered 6 variants (of various #expert, see Table 2) for analysis. In addition, we also include some side experiments when we explore related work (see discussion in Sec 6 and details in Appendix D).\n\n**References:**\n\n[1] Aidan Clark, Diego de las Casas, Aurelia Guy, Arthur Mensch, et. al. Unified Scaling Laws for Routed Language Models. https://arxiv.org/abs/2202.01169", "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. Strong perplexity improvement over Switch Transformer on out-of-domain evaluation of language modeling.\n2. The paper is well-written and easy to follow. \n3. Comprehensive analyses for the hyper-parameters and the baselines. The analyses deepen the understanding of current MoE methods.", "reasons_to_reject": "1. While the method provides better FLOPs-perplexity trade-off, it is unclear how Avg-K can provide better actual training speed perplexity trade-off due to increased communication cost.\n2. It will be better to compare the performance of the proposed method with more recent MoE routing approach such as X-MoE.", "questions_to_authors": "", "paper_topic_and_main_contributions": "This paper provides a unified mathematical view of sparse mixture-of-experts and sparse neural memory. Based on this view and empirical observations, the author proposes a new gating method Avg-k to select memory blocks based on the averaged hidden states of each block. Avg-k achieves substantially better perplexity compared with a strong baseline (Switch Transformer).", "excitement": 4, "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": "lvjOfNg28S", "reviewer": "Reviewer_J9vn", "summary": "", "questions": "", "limitations": "", "rating": 4, "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": "**Q1:** This paper [...] from a purely theoretical and model design perspective, but **lacks integration with engineering implementation**. Given [...] a complex system architecture, it is suggested to further explore the potential of the Avg-K [...] **under the premise of model and data parallelism**.\n\n**A1:**\n\nWe extend our gratitude to the reviewer for addressing concerns related to engineering implementation. Nevertheless, we wish to clarify that **the scope of our work unfortunately does not include the communication cost.**\n\nDuring the initial project scoping, we observed historical parallels where the start of the proof-of-concept for MoE in neural networks was established in [1], followed by the engineering implementation for practical efficiency in [2]. Drawing a parallel, our work follows a similar trajectory — we introduce and demonstrate the potential of a (theoretically) unified framework to shape S-FFN design. Consequently, future research could concentrate on bridging the gaps (practically) to escalate the impact. This would entail delving into the analysis of communication costs and proposing corresponding solutions, both of which hold the potential for dedicated papers. \n\nHere are some of our considerations:\n\nEach MoE variant may require distinct implementation considerations to obtain near-optimal runtime. For instance, we thought of an optimization strategy that becomes relevant as memory block sizes decrease, leading to increased communication costs. In such a context, where multiple small memory blocks coexist within a single GPU, an intuitively efficient \"local routing\" program could be employed to cleverly allocate tokens among these local blocks. In contrast, the currently standard all-to-all implementation could inadvertently lead to a large number of duplicate tokens being dispatched, thereby accounting for the elevated communication cost. Implementing such strategies could potentially involve refining the underlying data structure or even creating a specialized CUDA kernel.\n\nWe leave the above to future work due to access to limited resources.\n\n**References:**\n\n[1] Noam Shazeer et. al. Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer https://arxiv.org/abs/1701.06538\n\n[2] Dmitry Lepikhin et. al. GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding https://arxiv.org/abs/2006.16668", "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": "This paper systematically analyzes the role of S-FFN in large-scale language models and summarizes it into a unified design framework. After thorough experimental analysis, the authors propose a superior Avg-K method that surpasses common methods such as MoE. This provides insightful considerations for the larger parameterization and training of LLMs.", "reasons_to_reject": "This paper provides a unified S-FFN framework from a purely theoretical and model design perspective, but lacks integration with engineering implementation. Given that large-scale pretraining is a complex system architecture, it is suggested to further explore the potential of the Avg-K method to improve efficiency and effectiveness under the premise of model and data parallelism.", "questions_to_authors": "", "paper_topic_and_main_contributions": "Topic: Large and sparse feed-forward networks S-FFN for pretraining large language models.\n\nMain contribution: This paper analyzes the design choices of S-FFN, including the memory block size and selection method, under a general conceptual framework of sparse neural memory. The authors compare several S-FFN architectures for language modeling and provide insights into their relative efficacy and efficiency. They propose a simpler selection method, Avg-K, that selects blocks through their mean aggregated hidden states, which achieves lower perplexity in language modeling pretraining compared to existing MoE architectures. This work provides a better understanding of S-FFN and its potential for scaling up Transformers model size for pretraining large language models.", "excitement": 4, "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": {}}]

Author rebuttal:

Meta-review: This paper analyzes design decisions for sparse feedforward layers in large language models. It presents a unified framework that connects sparse mixtures of experts with sparse neural memory and empirically compares the two methods along with new variants that take elements of each. Inspired by their analysis, the authors present a new routing method called avg-k that outperforms previous methods. 

Reviewers generally found the paper clear and appreciated the strong empirical results, comprehensive analysis, and insights from the new conceptual framework. However, a central concern with the paper brought up by all reviews is potential lack of scalability to settings where model parallelism is necessary (models in the paper are up to 850M parameters and are not trained with model parallelism):
* The finding that a small block size (essentially activating a larger number of smaller experts) works better given the same compute budget may not hold because communication costs will increase. 
* The lack of load balancing in Avg-K will become more of a problem because speed will be limited by the devices receiving more tokens. While the paper does argue that load balancing should combine well with avg-k, there aren’t experiments (even at smaller scale) to justify this.
* Possibly due to the lack of load balancing, avg-k works poorly compared to baselines when using a larger block size, but that might be the more realistic setting when thinking about a model-parallel version.

In the rebuttal, the authors argue that while they believe there exist implementation details or method variants that could reduce communication overhead, the engineering required to develop an efficient model-parallel implementation would be very high. Instead, they see the value of the work in presenting their unified framework and in providing findings (e.g that Avg-K with a small block size is FLOP-efficient) that could be built upon in future engineering-focused research (e.g. adding load balancing and reducing communication overhead). While I do think the limitations may diminish the impact of the paper, I don’t think we should limit research on language model sparsity only to researchers with huge engineering budgets, and I believe the findings still could be of interest to the community.
Decision: Accept-Main