Original: T
Revised: T

Original: OWARDS
Revised: OWARDS

Original: A U NIFIED
Revised: A U NIFIED

Original: V IEW OF S PARSE
Revised: V IEW OF S PARSE

Original: F EEDF ORWARD N ETWORK IN T RANSFORMER
Revised: F EED - F ORWARD N ETWORK IN T RANSFORMER

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Large and sparse feed-forward networks (S-FFN) such as Mixture-of-Experts (MoE) have demonstrated to be an efﬁcient approach for scaling up Transformers model size for pretraining .
Revised: Large and sparse feed-forward networks (S-FFN) such as Mixture-of-Experts (MoE) have demonstrated to be an efﬁcient approach for scaling up Transformers model size for pretraining .

Original: By only activating part of the FFN parameters conditioning on input, S-FFN improves generalization performance while keeping training and inference cost (in FLOPs) ﬁxed.
Revised: By only activating part of the FFN parameters conditioning on input, S-FFN improves generalization performance while keeping training and inference cost (in FLOPs) ﬁxed.

Original: A growing body of work has been focusing on improving the S-FFN design, including routing and load balancing methods in the context of MoEs.
Revised: A growing body of work has been focusing on improving the S-FFN design, including routing and load balancing methods in the context of MoEs.

Original: Previously, another line of work motivates from a neural memory perspective and develops sparse neural memory techniques for
Revised: Previously, another line of work motivates from a neural memory perspective and develops sparse neural memory techniques for

Original: S-FFN.
Revised: S-FFN.

Original: This work merges the two seemingly different lines of work.
Revised: This work merges the two seemingly different lines of work.

Original: We present a uniﬁed framework to categorize design choices along two axes: memory block size and memory block selection method.
Revised: We present a uniﬁed framework to categorize design choices along two axes: memory block size and memory block selection method.

Original: Using this uniﬁed framework, we compare several S-FFN architectures for language modeling and provide insights into their relative efﬁcacy and efﬁciency.
Revised: Using this uniﬁed framework, we compare several S-FFN architectures for language modeling and provide insights into their relative efﬁcacy and efﬁciency.

Original: We show that a smaller memory block size leads to lower perplexity.
Revised: We show that a smaller memory block size leads to lower perplexity.

Original: Additionally, we ﬁnd that selection through a gate, in general, improves the perplexity-FLOPs trade-off but has worse perplexity than selection using hidden states without a gate.
Revised: Additionally, we ﬁnd that selection through a gate, in general, improves the perplexity-FLOPs trade-off but has worse perplexity than selection using hidden states without a gate.

Original: Based on these insights, we propose a new selection method — Avg-K that selects blocks through their mean aggregated hidden states.
Revised: Based our analysis, we propose a new selection method — Avg-K that selects blocks through their mean aggregated hidden states.

Original: With 1% additional FLOPs, Avg-K achieves 2.16 lower perplexity than a vanilla transformer (16.96), outperforming  Switch Transformer (16.45).
Revised: Not only outperforming all MoE models considered, Avg-K is also the ﬁrst to performs well without load balancing while Switch Transformer may degenerate.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Large-scale pretrained language models (LLMs) achieve remarkable performance and generalization ability for NLP tasks (Radford & Narasimhan, 2018;
Revised: Large-scale pretrained language models (LLMs) achieve remarkable performance and generalization ability for NLP tasks (Radford & Narasimhan, 2018;

Original: Devlin et al., 2019; Liu et al., 2019; Radford et al., 2019; Brown et al., 2020; Raffel et al., 2020; Chowdhery et al., 2022).
Revised: Devlin et al., 2019; Liu et al., 2019; Radford et al., 2019; Brown et al., 2020; Raffel et al., 2020; Chowdhery et al., 2022).

Original: Scaling up the model size (the number of parameters) has been shown as a reliable recipe for better generalization, unlocking new capabilities, while the performance has not shown signs of plateauing (Kaplan et al., 2020;
Revised: Scaling up the model size (the number of parameters) has been shown as a reliable recipe for better generalization, unlocking new capabilities, while the performance has not shown signs of plateauing (Kaplan et al., 2020;

Original: Zhang et al., 2022a; Chowdhery et al., 2022; Hoffmann et al., 2022;
Revised: Zhang et al., 2022a; Chowdhery et al., 2022; Hoffmann et al., 2022;

Original: Wei et al., 2022).
Revised: Wei et al., 2022).

Original: However, the computational resources required to train larger language models are formidable, calling for more efﬁcient training and inference solutions of LLMs (Borgeaud et al., 2022; Schwartz et al., 2020; Tay et al., 2020).
Revised: However, the computational resources required to train larger language models are formidable, calling for more efﬁcient training and inference solutions of LLMs (Borgeaud et al., 2022; Schwartz et al., 2020; Tay et al., 2020).

Original: One promising direction is sparse scaling which increases the number of parameters while keeping the training and inference cost (in FLOPs) ﬁxed.
Revised: One promising direction is sparse scaling which increases the number of parameters while keeping the training and inference cost (in FLOPs) ﬁxed.

Original: Recent work focuses on scaling up a transformer’s feed-forward network (FFN) with sparsely activated parameters, resulting in a scaled and sparse FFN (S-FFN).
Revised: Recent work focuses on scaling up a transformer’s feed-forward network (FFN) with sparsely activated parameters, resulting in a scaled and sparse FFN (S-FFN).

Original: There have been two major approaches to achieve S-FFN.
Revised: There have been two major approaches to achieve S-FFN.

Original: One treats S-FFN as a neural memory (Sukhbaatar et al., 2015a) where a sparse memory retrieves and activates only parts of the memory cells (Lample et al., 2019).
Revised: One treats S-FFN as a neural memory (Sukhbaatar et al., 2015a) where a sparse memory retrieves and activates only parts of the memory cells (Lample et al., 2019).

Original: The other adopts Mixture-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 a few among many experts for a particular input.
Revised: The other adopts Mixture-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 a few among many experts for a particular input.

Original: While both memory and MoE models achieve S-FFN, they have been considered two completely different approaches.
Revised: While both memory and MoE models achieve S-FFN, they have been considered two completely different approaches.

Original: We aim to draw the connections between these two classes of S-FFN:
Revised: We aim to draw the connections between these two classes of S-FFN:

Original: What critical design choices do they have in common?
Revised: What critical design choices do they have in common?

Original: Which design choices are essential for their modeling capability and computation efﬁciency?
Revised: Which design choices are essential for their modeling capability and computation efﬁciency?

Original: Can the effective ingredients of each method be transferred and combined to improve performance further?
Revised: Can the effective ingredients of each method be transferred and combined to improve performance further?

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: Add
Revised: Add

Original: [Figure]
Revised: [Figure]

Original: Add + Normalize y
Revised: Add + Normalize y

Original: Down-projection
Revised: Down-projection

Original: Key Table
Revised: Key Table

Original: Value Table
Revised: Value Table

Original: Transformer Layer Option 1: Sparse Mixture-of-Expert
Revised: Transformer Layer Option 1: Sparse Mixture-of-Expert

Original: Option 2: Sparse Neural Memory
Revised: Option 2: Sparse Neural Memory

Original: Replace Replace
Revised: Replace Replace

Original: Top-
Revised: Top-

Original: [Figure]
Revised: [Figure]

Original: GShard, Switch Transformer, HashLayer, BASELayer...
Revised: GShard, Switch Transformer, HashLayer, BASELayer...

Original: [Figure]
Revised: [Figure]

Original: In order to answer these questions, we start from the neural memory view of FFN (Sukhbaatar et al., 2015a) (§2.1) and reduce all S-FFN to the same mathematical form (§3.1).
Revised: In order to answer these questions, we start from the neural memory view of FFN (Sukhbaatar et al., 2015a) (§2.1) and reduce all S-FFN to the same mathematical form (§3.1).

Original: Then, we characterize various S-FFN methods along two dimensions — memory block size (e.g. expert size) and memory block selection method (e.g. gating) (§3.2).
Revised: Then, we characterize various S-FFN methods along two dimensions — memory block size (e.g. expert size) and memory block selection method (e.g. gating) (§3.2).

Original: Using this framework, we made the following contributions:
Revised: Using this framework, we made the following contributions:

Original: • We study a wide range of memory block sizes besides common block size in MoEs (Fedus et al., 2022) and show that, compared with a larger one, a smaller block size can keep improving the perplexity with little incurred extra FLOPs (§5.1 §5.2), leading to better perplexity/computation trade-offs.
Revised: • We study a wide range of memory block sizes besides common block size in MoEs (Fedus et al., 2022) and show that, compared with a larger one, a smaller block size can keep improving the perplexity with little incurred extra FLOPs (§5.1 §5.2), leading to better perplexity/computation trade-offs.

Original: • We conduct a systematic exploration of block selection methods to quantify their relative efﬁcacy and efﬁciency (§5.2).
Revised: • We conduct a systematic exploration of block selection methods to quantify their relative efﬁcacy and efﬁciency (§5.2).

Original: Speciﬁcally, we ﬁnd that the selection method through a gating function, in general, improves the FLOPs-Perplexity trade-off.
Revised: Speciﬁcally, we ﬁnd that the selection method through a gating function, in general, improves the FLOPs-Perplexity trade-off.

Original: However, the parameterization of the current gating function has worse perplexity than a selection method that uses the FFN hidden states.
Revised: However, the parameterization of the current gating function has worse perplexity than a selection method that uses the FFN hidden states.

Original: • Drawing on insights above, we propose a simple gate for S-FFN— Avg-K (§6), which efﬁciently selects memory blocks based on the mean aggregated hidden states of each block.
Revised: • Drawing on insights above, we propose a simple gate for S-FFN— Avg-K (§6), which efﬁciently selects memory blocks based on the mean aggregated hidden states of each block.

Original: With 1% additional FLOPs, Avg-K achieves 2.16 lower perplexity than a vanilla transformer
Revised: With 1% additional FLOPs, Avg-K achieves 2.16 lower perplexity than a vanilla transformer

Original: (16.96), outperforming Switch Transformer (16.45).
Revised: (16.96), outperforming Switch Transformer (16.45).

Original: As an additional bonus, we show that it is effective without load balancing constraints, while the conventional one degrades (Shazeer et al., 2017; Eigen et al., 2014).
Revised: Avg-K is the ﬁrst MoE model performs well without load balancing constraint, as suggested by previous work (Shazeer et al., 2017; Eigen et al., 2014).

Original:  B ACKGROUND 2. F EED -F ORWARD N ETWORK
Revised:  B ACKGROUND 2. F EED -F ORWARD N ETWORK

Original: A transformer layer (Vaswani et al., 2017) consists of a self-attention block and a Feed-Forward
Revised: A transformer layer (Vaswani et al., 2017) consists of a self-attention block and a Feed-Forward

Original: Network (FFN) block.
Revised: Network (FFN) block.

Original: FFN receives an input vector x 2 R d from the self-attention block, multiplies it with K 2 R d m ⇥ d , applies an non-linear function to obtain the hidden states m 2 R d m and applies another afﬁne transformation V 2 R d m ⇥ d to produce a d -dimensional output (Eq. 1).
Revised: FFN receives an input vector x 2 R d from the self-attention block, multiplies it with K 2 R d m ⇥ d , applies an non-linear function to obtain the hidden states m 2 R d m and applies another afﬁne transformation V 2 R d m ⇥ d to produce a d -dimensional output (Eq. 1).

Original: [Equation]
Revised: [Equation]

Original: FFN is a multi-layer perceptron, but it can also be described as a neural memory (Sukhbaatar et al., 2015b; 2019; Geva et al., 2021)(Eq. 2).
Revised: FFN is a multi-layer perceptron, but it can also be described as a neural memory (Sukhbaatar et al., 2015b; 2019; Geva et al., 2021)(Eq. 2).

Original: In this view, FFN consists of d m key-value pairs, called memory cells .
Revised: In this view, FFN consists of d m key-value pairs, called memory cells .

Original: Each key is represented by a d -dimensional k
Revised: Each key is represented by a d -dimensional k

Original: i 2 R d , and together form the key table ; likewise, a value table V 2 R d m ⇥ d .
Revised: i 2 R d , and together form the key table ; likewise, a value table V 2 R d m ⇥ d .

Original: When the memory receives the query input x 2 R d , it multiplies x with every k i ; followed by the non-linear function, it produces memory coefﬁcient
Revised: When the memory receives the query input x 2 R d , it multiplies x with every k i ; followed by the non-linear function, it produces memory coefﬁcient

Original: m i =
Revised: m i =

Original: f
Revised: f

Original: ( x · k i ) for the i -th memory cell.
Revised: ( x · k i ) for the i -th memory cell.

Original: Finally, the output of FFN is the sum of its values v i weighted by their corresponding memory coefﬁcient m i .
Revised: Finally, the output of FFN is the sum of its values v i weighted by their corresponding memory coefﬁcient m i .

Original: Conventionally, the size of FFN — d m — is set to be 4 · d .
Revised: Conventionally, the size of FFN — d m — is set to be 4 · d .

Original: 2. S CALING UP FFN
Revised: 2. S CALING UP FFN

Original: As discussed in §1, scaling up the number of parameters in FFN serves as a lever to improve transformer performance.
Revised: As discussed in §1, scaling up the number of parameters in FFN serves as a lever to improve transformer performance.

Original: Since a conventional FFN already takes about two-thirds of a transformer layer’s parameters (Geva et al., 2021), scaling up FFN will greatly affect the parameter size of a transformer model.
Revised: Since a conventional FFN already takes about two-thirds of a transformer layer’s parameters (Geva et al., 2021), scaling up FFN will greatly affect the parameter size of a transformer model.

Original: However, one could sparsely activate the parameter to control the required compute.
Revised: However, one could sparsely activate the parameter to control the required compute.

Original: In this section, we will review two lines of work to achieve a scaled and sparse FFN (S-FFN).
Revised: In this section, we will review two lines of work to achieve a scaled and sparse FFN (S-FFN).

Original: One has a mixture-of-expert model activate a few experts (§2.2.1), and the other speciﬁes a memory model (§2.2.2).
Revised: One has a mixture-of-expert model activate a few experts (§2.2.1), and the other speciﬁes a memory model (§2.2.2).

Original: 2.2. M IXTURE OF E XPERTS (M O E)
Revised: 2.2. M IXTURE OF E XPERTS (M O E)

Original: Mixture of experts (MoE; Jacobs et al.
Revised: Mixture of experts (MoE; Jacobs et al.

Original: (1991)) consists of a set of expert models { f i ( x ) } B  1 i =0 and a gating function g : R d !
Revised: (1991)) consists of a set of expert models { f i ( x ) } B  1 i =0 and a gating function g : R d !

Original: R B to estimates the importance of each expert.
Revised: R B to estimates the importance of each expert.

Original: Finally, the output is the sum of experts’ output weighted by the gate’s weight estimation for that particular expert.
Revised: Finally, the output is the sum of experts’ output weighted by the gate’s weight estimation for that particular expert.

Original: [Equation]
Revised: [Equation]

Original: Recent work (Du et al., 2021;
Revised: Recent work (Du et al., 2021;

Original: Lepikhin et al., 2021; Roller et al., 2021; Lewis et al., 2021;
Revised: Lepikhin et al., 2021; Roller et al., 2021; Lewis et al., 2021;

Original: Zhou et al., 2022) adopt it to transformer by treating an FFN as one expert and S parsely activating the MoE ( SMoE ).
Revised: Zhou et al., 2022) adopt it to transformer by treating an FFN as one expert and S parsely activating the MoE ( SMoE ).

Original: In SMoE, the gating function (or “router”) routes an input token x to a subset (e.g. 1 or 2) of experts — E = subset ( g ( x )) .
Revised: In SMoE, the gating function (or “router”) routes an input token x to a subset (e.g. 1 or 2) of experts — E = subset ( g ( x )) .

Original: Conventionally, SMoE enforces load balancing constraints to avoid overly using a few experts while under-utilizing others, and converging to local optima (Shazeer et al., 2017; Eigen et al., 2014).
Revised: Conventionally, SMoE enforces load balancing constraints to avoid overly using a few experts while under-utilizing others, and converging to local optima (Shazeer et al., 2017; Eigen et al., 2014).

Original: Various SMoEs mainly have two types of gates.
Revised: Various SMoEs mainly have two types of gates.

Original: Learned gate is parameterized by a set of learnable expert embeddings ✓ =
Revised: Learned gate is parameterized by a set of learnable expert embeddings ✓ =

Original: [ e 0 ; · · · ; e B  1 ] 2 R B ⇥ d , corresponding to each experts.
Revised: [ e 0 ; · · · ; e B  1 ] 2 R B ⇥ d , corresponding to each experts.

Original: The importance of the i -th expert is obtained by g i ( x ) =
Revised: The importance of the i -th expert is obtained by g i ( x ) =

Original: exp( e i · x )
Revised: exp( e i · x )

Original: P j exp( e j · x ) .
Revised: P j exp( e j · x ) .

Original: To enforce load balancing when routing, SMoEs use an additional auxiliary loss (Lepikhin et al., 2021; Fedus et al., 2021; Artetxe et al., 2021; Du et al., 2021; Chi et al., 2022) or frame expert utilization as a constrained optimization problem (Lewis et al., 2021;
Revised: To enforce load balancing when routing, SMoEs use an additional auxiliary loss (Lepikhin et al., 2021; Fedus et al., 2021; Artetxe et al., 2021; Du et al., 2021; Chi et al., 2022) or frame expert utilization as a constrained optimization problem (Lewis et al., 2021;

Original: Zhou et al., 2022).
Revised: Zhou et al., 2022).

Original: Static gate , in contrast to a learnable gate, does not have any differentiable parameters.
Revised: Static gate , in contrast to a learnable gate, does not have any differentiable parameters.

Original: Instead, it uses a static mapping that encodes load-balancing constraints to route input (Roller et al., 2021;
Revised: Instead, it uses a static mapping that encodes load-balancing constraints to route input (Roller et al., 2021;

Original: Gururangan et al., 2021).
Revised: Gururangan et al., 2021).

Original: For example, RandHash from HashLayer (Roller et al., 2021) uses a hash table that maps from token type to randomly selected expert(s) and enforces the load balancing with random selection.
Revised: For example, RandHash from HashLayer (Roller et al., 2021) uses a hash table that maps from token type to randomly selected expert(s) and enforces the load balancing with random selection.

Original: DEMix (Gururangan et al., 2021) ensures i -th expert only sees data from i -th domain and trains each expert for equal steps to ensure equal utilization.
Revised: DEMix (Gururangan et al., 2021) ensures i -th expert only sees data from i -th domain and trains each expert for equal steps to ensure equal utilization.

Original: 2.2. S PARSE N EURAL M EMORY
Revised: 2.2. S PARSE N EURAL M EMORY

Original: The other line of work follows the memory view of FFN (Eq. 2).
Revised: The other line of work follows the memory view of FFN (Eq. 2).

Original: It is straightforward to increase the memory size d m to a much larger value d m  4 · d .
Revised: It is straightforward to increase the memory size d m to a much larger value d m  4 · d .

Original: By only using the top- k entries in m , one could sparsely activate the value table, resulting in a vanilla sparse memory ( VanillaM ).
Revised: By only using the top- k entries in m , one could sparsely activate the value table, resulting in a vanilla sparse memory ( VanillaM ).

Original: The straightforward application, however, uses the entire key table to produce all the memory coefﬁcients m = x · K > , resulting in dense computation proportional to the memory size.
Revised: The straightforward application, however, uses the entire key table to produce all the memory coefﬁcients m = x · K > , resulting in dense computation proportional to the memory size.

Original: Memory models need to use models that do not look at the full key table to scale computation sublinearly to the number of memory cells and take advantage of sparse computation.
Revised: Memory models need to use models that do not look at the full key table to scale computation sublinearly to the number of memory cells and take advantage of sparse computation.

Original: Lample et al.
Revised: Lample et al.

Original: (2019) explored the following two techniques in this direction.
Revised: (2019) explored the following two techniques in this direction.

Original: Low-Rank Key Memory (LoRKM)
Revised: Low-Rank Key Memory (LoRKM)

Original: A straightforward technique is to assume that the key table is composed of and approximated by a down-projection D 2 R d ⇥ d ` and a low-rank key table ˜ K
Revised: A straightforward technique is to assume that the key table is composed of and approximated by a down-projection D 2 R d ⇥ d ` and a low-rank key table ˜ K

Original: R d m ⇥ d ` , where d ` ⌧ d .
Revised: R d m ⇥ d ` , where d ` ⌧ d .

Original: LoRKM produces memory coefﬁcients by m i =
Revised: LoRKM produces memory coefﬁcients by m i =

Original: f (( x · D ) ·

Original: ˜ k

Original: [Equation]
Revised: [Equation]

Original: ) = f ( t · ˜ k i ) . Product Key Memory (PKM) Building upon LoRKM, PKM assumes that the low-rank key table is further approximated by the outer product , w.r.t.
Revised: Product Key Memory (PKM) Building upon LoRKM, PKM assumes that the low-rank key table is further approximated by the outer product , w.r.t.

Original: vector concatenation operator, of two smaller sub-key tables C , C 0 2 R p d m ⇥ d ` 2 .
Revised: vector concatenation operator, of two smaller sub-key tables C , C 0 2 R p d m ⇥ d ` 2 .

Original: Unlike LoRKM where key vectors are independent of each other,
Revised: Unlike LoRKM where key vectors are independent of each other,

Original: PKM has them entangled with one another, i.e. ˜ k
Revised: PKM has them entangled with one another, i.e. ˜ k

Original: i =
Revised: i =

Original: [Equation]
Revised: [Equation]

Original: i
Revised: b i/

Original: [Equation]
Revised: [Equation]

Original: mod
Revised: mod

Original: [Equation]
Revised: [Equation]

Original: i
Revised: i

Original: [Equation]
Revised: [Equation]

Original: d ` .
Revised: d ` .

Original: One could exploit this structure to efﬁciently compute the memory coefﬁcient m .
Revised: One could exploit this structure to efﬁciently compute the memory coefﬁcient m .

Original: One ﬁrst calculates the dot product between sub-key tables and down-projected input t individually and combines them with a negligible cost to form the full dot product m 2 R d m :
Revised: One ﬁrst calculates the dot product between sub-key tables and down-projected input t individually and combines them with a negligible cost to form the full dot product m 2 R d m :

Original: [Equation]
Revised: [Equation]

Original: ⌘ , where s = t
Revised: ⌘ , where s = t

Original: [Equation]
Revised: [Equation]

Original: Since PKM has a negligible factorized key table ˜ K · D > ( < 0 .
Revised: d `

Original: 3% ), one could have multiple randomly initialized key tables to process batched input together and encourage load balancing — uniform usage of V .
Revised: [Equation]

Original:  A U NIFIED V IEW OF S PARSE FFN S
Revised:  A U NIFIED V IEW OF S PARSE FFN S

Original: We show the connections between MoE and neural memory despite their different surface forms.
Revised: We show the connections between MoE and neural memory despite their different surface forms.

Original: We ﬁrst derive an equivalent form of MoE to establish its connection with sparse memory (§3.1).
Revised: We ﬁrst derive an equivalent form of MoE to establish its connection with sparse memory (§3.1).

Original: Then, we propose a uniﬁed framework for S-FFN (§3.2).
Revised: Then, we propose a uniﬁed framework for S-FFN (§3.2).

Original: 3. A C LOSER LOOK AT M O E
Revised: 3. A C LOSER LOOK AT M O E

Original: MoEs use a gating function to estimate the importance of all experts and combine each expert’s output through linear combination.
Revised: MoEs use a gating function to estimate the importance of all experts and combine each expert’s output through linear combination.

Original: Here, inspired by the memory view on FFNs (§2), we derive a formulation of MoE models (Eq. 3) as seen from the memory view:
Revised: Here, inspired by the memory view on FFNs (§2), we derive a formulation of MoE models (Eq. 3) as seen from the memory view:

Original: [Equation]
Revised: [Equation]

Original: Apply the memory view of FFN 
Revised: expert is [Equation] FFN )

Original: [Equation]
Revised: [Equation]

Original: Eq. 2) ) 
Revised: A ( Apply the memory view of FFN (Eq. 2) ) [Equation] Distributive property of multiplication ) [Equation]

Original: [Equation]
Revised: [Equation]

Original: In Eq 5, v l is the l -th row of the stack of d m -size value tables from B FFN experts; and ml =
Revised: In Eq 5, v l is the l -th row of the stack of d m -size value tables from B FFN experts; and m l

Original:  m i · d m + j is its memory coefﬁcient obtained from g i ( x ) · m ( i ) j .
Revised: = m i · d m + j is its memory coefﬁcient obtained from g i ( x ) · m ( i ) j .

Original: For example, m d m is is obtained by weighting m (1)0 — the 0 -th memory coefﬁcient in the 1 -th FFN expert — with g 0 ( x ) — the estimation of the 1 -th FFN from gate.
Revised: For example, m d m is is obtained by weighting m (1)0 — the 0 -th memory coefﬁcient in the 1 -th FFN expert — with g 0 ( x ) — the estimation of the 1 -th FFN from gate.

Original: Building upon such memory view, one could see that is a sparse memory operating in terms block and uses its gate to narrow down the summation over the stacked value tables V to ones in FFN ( i ) , for i 2 subset ( g ( x )) =
Revised: Building upon such memory view, one could see that is a sparse memory operating in terms block and uses its gate to narrow down the summation over the stacked value tables V to ones in FFN ( i ) , for i 2 subset ( g ( x )) =

Original: E .
Revised: E .

Original: Comparison with Sparse Memory Both SMoE and Sparse Memory are neural memory, but there are several differences: 1) Memory cells share the same importance weight: in sparse memory, each memory cell receives an individual weight.
Revised: Comparison with Sparse Memory Both SMoE and Sparse Memory are neural memory, but there are several differences: 1) Memory cells share the same importance weight: in sparse memory, each memory cell receives an individual weight.

Original: In contrast, in SMoE groups of 4 · d share the same importance weight g i ( x ) .
Revised: In contrast, in SMoE groups of 4 · d share the same importance weight g i ( x ) .

Original: Single selection and importance weight: Sparse memory uses a single dot product between input token and any key vector x · k for both, whereas SMoE depends on a separately parameterized gate g in addition to the dot product.
Revised: Single selection and importance weight: Sparse memory uses a single dot product between input token and any key vector x · k for both, whereas SMoE depends on a separately parameterized gate g in addition to the dot product.

Original: 3. T HE U NIFIED F RAMEWORK
Revised: 3. T HE U NIFIED F RAMEWORK

Original: We propose a general framework that uniﬁes the two different approaches to achieve S-FFN.
Revised: We propose a general framework that uniﬁes the two different approaches to achieve S-FFN.

Original: We cast both as instances of a memory with large key and value table — K , V 2 R d m ⇥ d , where d m  4 · d .
Revised: We cast both as instances of a memory with large key and value table — K , V 2 R d m ⇥ d , where d m  4 · d .

Original: We distinguish the different methods along two dimensions illustrated below and summarized in
Revised: We distinguish the different methods along two dimensions illustrated below and summarized in

Original: Table 1:
Revised: Table 1:

Original: Memory block size speciﬁes how many memory cells share the same importance weight at selection time, and thus together treated as a memory block .
Revised: Memory block size speciﬁes how many memory cells share the same importance weight at selection time, and thus together treated as a memory block .

Original: We use g to denote the size of one block.
Revised: We use g to denote the size of one block.

Original: In other words, we split the K , V along the d m -dimension into g -size blocks.
Revised: In other words, we split the K , V along the d m -dimension into g -size blocks.

Original: Therefore, a memory consists of B = d m /g blocks in total
Revised: Therefore, a memory consists of B = d m /g blocks in total.

Original: Formally, we write
Revised: Formally, we write

Original: [Equation]
Revised: [Equation]

Original: d v [Table] For example, sparse memory has block size g = 1 — trivially treating 1 memory cell as a “block”; and SMoE has the block size g = 4 · d (§3.1).
Revised: For example, sparse memory has block size g = 1 — trivially treating 1 memory cell as a “block”; and SMoE has the block size g = 4 · d (§3.1).

Original: Current approaches generally use ﬁxed block sizes, but this is mostly an artifact of how the methods were derived rather than a mathematical constraint.
Revised: Current approaches generally use ﬁxed block sizes, but this is mostly an artifact of how the methods were derived rather than a mathematical constraint.

Original: For example, we can design SMoE versions instead of 1 expert of size 4 · d , or uses 2 experts of size 2 · d .
Revised: For example, we can design SMoE versions instead of 1 expert of size 4 · d , or uses 2 experts of size 2 · d .

Original: We can similarly chunk memory coefﬁcients m into blocks of size g when dealing with sparse memories. 
Revised: We can similarly chunk memory coefﬁcients m into blocks of size g in sparse memories. [Table]

Original: Memory block selection method is the speciﬁc function that compute the importance of each memory blocks for selection.
Revised: Memory block selection method is the speciﬁc function that compute the importance of each memory blocks for selection.

Original: Since SMoE is also a type of sparse memory, we distinguish the selection method by a new criterion — whether one allows input x to directly interact with the key table K g .
Revised: Since SMoE is also a type of sparse memory, we distinguish the selection method by a new criterion — whether one allows input x to directly interact with the key table K g .

Original: As discussed in §3.1, SMoE uses the estimation from an individually parameterized gate to select, while sparse memory solely and directly uses a key table.
Revised: As discussed in §3.1, SMoE uses the estimation from an individually parameterized gate to select, while sparse memory solely and directly uses a key table.

Original: Thus, current SMoE is a type of indirect selection method, and sparse memory a direct one.
Revised: Thus, current SMoE is a type of indirect selection method, and sparse memory a direct one.

Original: Various SMoEs are further characterized by whether their gating function has learned parameters or consists of a static mapping (§2.2.1).
Revised: Various SMoEs are further characterized by whether their gating function has learned parameters or consists of a static mapping (§2.2.1).

Original: On the other hand, sparse memory is characterized by how much factorizationis applied to the key table  (§2.2.2).
Revised: Meanwhile, sparse memory is characterized by how much factorization the key table uses (§2.2.2).

Original:  E XPERIMENT S ETUP 4. M ODELS
Revised:  E XPERIMENT S ETUP 4. M ODELS

Original: We choose Dense Baseline using transformer architectures used in GPT-3 models (Brown et al., 2020), which has 24 transformer layers, with d = 1024 , GeLU activation functions and with a memory size (or FFN hidden size) to be 4 · d .
Revised: We choose Dense Baseline using transformer architectures used in GPT-3 models (Brown et al., 2020), which has 24 transformer layers, with d = 1024 , GeLU activation functions and with a memory size (or FFN hidden size) to be 4 · d .

Original: S-FFN Given a model above, we replace some of its FFNs with an S-FFN.
Revised: S-FFN Given a model above, we replace some of its FFNs with an S-FFN.

Original: Similar to (Lepikhin et al., 2021), we replace the FFN at every 6 layers, leading to 4 S-FFNs in total across 24 layers.
Revised: Similar to (Lepikhin et al., 2021), we replace the FFN at every 6 layers, leading to 4 S-FFNs in total across 24 layers.

Original: Since memory block size g is a perception imposed upon memory cells, we use k to denote the number of active memory cells and control how activated the S-FFN is.
Revised: Since memory block size g is a perception imposed upon memory cells, we use k to denote the number of active memory cells and control how activated the S-FFN is.

Original: We use the formulation of d m = E · (4 · d ) to control the size of S-FFN, so the S-FFN will activate b
Revised: We use the formulation of d m = E · (4 · d ) to control the size of S-FFN, so the S-FFN will activate b

Original: = kg out of B = d m g memory blocks.
Revised: = kg out of B = d m g memory blocks.

Original: In Table 2, we list all S-FFN models used for analysis in §5.
Revised: In Table 2, we list all S-FFN models used for analysis in §5.

Original: We count FLOPs analytically following Narayanan et al.
Revised: We count FLOPs analytically following Narayanan et al.

Original: and do not account if a worker ﬁnishes computation before another (when using model parallelism).
Revised: and do not account if a worker ﬁnishes computation before another (when using model parallelism).

Original: 4. L ANGUAGE M ODELING
Revised: We use the number of learnable parameter to consider whether two models are equally expressive. PKM-FFN Since the factorized key table in PKM has little ( < 0 . 3% ) learnable parameter relative to the value table, we propose a variant called PKM-FFN to match the number of parameter of other models like RandHash. This variant has memory block size g = 1 and the same key-value table as

Revised: RandHash.

Revised: PKM-FFN has a gate whose g ( x ) is the same as the m from a PKM and g i

Original:  Pretraining Data We pretrain all S-FFN models on a total of 453GB text with 112B tokens from a union of six English-only datasets, including English subset of CC100 (Wenzek et al., 2020) and the ﬁve datasets used to pretrain RoBERTa
Revised: = m i . 4. L ANGUAGE M ODELING Pretraining Data We pretrain all S-FFN models on a total of 453GB text with 112B tokens from a union of six English-only datasets, including English subset of CC100 (Wenzek et al., 2020) and the ﬁve datasets used to pretrain RoBERTa

Original: (Liu et al., 2019) — speciﬁcally BookCorpus (Zhu et al., 2015), English Wikipedia, CC-News (Nagel, 2016), OpenWebText (Gokaslan & Cohen, 2019),
Revised: (Liu et al., 2019) — speciﬁcally BookCorpus (Zhu et al., 2015), English Wikipedia, CC-News (Nagel, 2016), OpenWebText (Gokaslan & Cohen, 2019),

Original: CC-Stories (Trinh & Le, 2018)
Revised: CC-Stories (Trinh & Le, 2018)

Original: (details in Appendix A.2)
Revised: (details in Appendix A.2)

Original: 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.
Revised: 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.

Original: All models are trained for 60B tokens for convergence.
Revised: All models are trained for 60B tokens for convergence.

Original: See Appendix A.1 for hyperparameters.
Revised: See Appendix A.1 for hyperparameters.

Original: [Table]
Revised: [Table]

Original: Evaluation settings We evaluate our models’ ability to predict the next token in a sequence as measured by perplexity.
Revised: Evaluation settings We evaluate our models’ ability to predict the next token in a sequence as measured by perplexity.

Original: We report both in-domain and out-of-domain perplexity to indicate generalization ability.
Revised: We report both in-domain and out-of-domain perplexity to indicate generalization ability.

Original: For out-of-domain, we use data from The Pile (Gao et al., 2020), a public dataset that combines data from 22 diverse sources.
Revised: For out-of-domain, we use data from The Pile (Gao et al., 2020), a public dataset that combines data from 22 diverse sources.

Original:  A NALYSIS UNDER THE UNIFIED VIEW
Revised:  A NALYSIS UNDER THE UNIFIED VIEW

Original: In this section
Revised: In this section

Original: [Table]
Revised: [Table]

Original: we use the proposed uniﬁed view to systematically study the design choice of S-FFN.
Revised: we use the proposed uniﬁed view to systematically study the design choice of S-FFN.

Original: Speciﬁcally, (1) .
Revised: Speciﬁcally, (1) .

Original: we study a wide range of block sizes other than the incidental choice used in existing work and investigate its impact on language modeling perplexity (§5.1).
Revised: we study a wide range of block sizes other than the incidental choice used in existing work and investigate its impact on language modeling perplexity (§5.1).

Original: Both direct and indirect block selection methods lead to lower perplexity than a standard FFN
Revised: Both direct and indirect block selection methods lead to lower perplexity than a standard FFN

Original: [Table]
Revised: [Table]

Original: but which type of method has better FLOPs-Perplexity trade-off and what are the relative efﬁcacy and efﬁciency of different methods require further study (§5.2)
Revised: but which type of method has better FLOPs-Perplexity trade-off and what are the relative efﬁcacy and efﬁciency of different methods require further study (§5.2)

Original: [Table]
Revised: [Table]

Original: 5. M EMORY BLOCK SIZE
Revised: 5. M EMORY BLOCK SIZE

Original: Since block size is a natural number, we aim to answer a straightforward question — given a ﬁxed k , does smaller memory block size lead to lower perplexity?
Revised: Since block size is a natural number, we aim to answer a straightforward question — given a ﬁxed k , does smaller memory block size lead to lower perplexity?

Original: We use simple and robust selection methods to disentangle the impact of hyperparameter choices.
Revised: We use simple and robust selection methods to disentangle the impact of hyperparameter choices.

Original: Speciﬁcally, we use random hash in
Revised: Speciﬁcally, we use random hash in

Original: HashLayer (Roller et al., 2021) (denoted RandHash ) for indirect block selection and exact top-k memory block (denoted VanillaM ) for direct block selection.
Revised: HashLayer (Roller et al., 2021) (denoted RandHash ) for indirect block selection and exact top-k memory block (denoted VanillaM ) for direct block selection.

Original: For all experiments, we use E = 16 .
Revised: For all experiments, we use E = 16 .

Original: RandHash randomly selects b = kg unique memory blocks among all B =
Revised: RandHash randomly selects b = kg unique memory blocks among all B =

Original: d m g blocks — essentially sampling b unique values from Uniform( [0 , · · · , B  1] ).
Revised: d m g blocks — essentially sampling b unique values from Uniform( [0 , · · · , B  1] ).

Original: Originally, with block size g = 4096 , a
Revised: Originally, with block size g = 4096 , a

Original: RandHash assigns a token to 40964096 = 1 block; with block size g = 2048 , 40962048 = 2 blocks.
Revised: RandHash assigns a token to 40964096 = 1 block; with block size g = 2048 , 40962048 = 2 blocks.

Original: VanillaM originally has a block size g = 1 and selects top-
Revised: VanillaM originally has a block size g = 1 and selects top-

Original: k scalars in memory coefﬁcients m =
Revised: k scalars in memory coefﬁcients m =

Original: GeLU ( x · K > ) .
Revised: GeLU ( x · K > ) .

Original: We made a minimal change to extend it to larger block size g : given m , we chunk it into B blocks — m g =
Revised: We made a minimal change to extend it to larger block size g : given m , we chunk it into B blocks — m g =

Original: [ m (0) ; · · · ; m ( B  1) ] ; then, we select the top- b blocks using the average of each block — Avg ( GeLU ( x · ( K ( i ) ) > ) ,dim=0 ) .
Revised: [ m (0) ; · · · ; m ( B  1) ] ; then, we select the top- b blocks using the average of each block — Avg ( GeLU ( x · ( K ( i ) ) > ) ,dim=0 ) .

Original: In Fig.
Revised: In Fig.

Original: 2, we observe that smaller block size leads to an improvement of 0 .
Revised: 2, we observe that smaller block size leads to an improvement of 0 .

Original: 15 . 35) perplexity for RandHash and an improvement of 0 .
Revised: 15 . 35) perplexity for RandHash and an improvement of 0 .

Original: 14 . 69) for VanillaM.
Revised: 14 . 69) for VanillaM.

Original: This observation holdsTable no matter how one selects: 1) with indirect (RandHash) or direct (VanillaM) selection and 2) with random (RandHash) or similarity-based (VanillaM) block selection.
Revised: This observation holds no matter how one selects: 1) with indirect (RandHash) or direct (VanillaM) selection and 2) with random (RandHash) or similarity-based (VanillaM) block selection.

Original: In Appendix B.2, we provide theoretical justiﬁcations for this observation which shows that a smaller block size improves model capacity by including more combinations of memory cells.
Revised: In Appendix B.2, we provide theoretical justiﬁcations for this observation which shows that a smaller block size improves model capacity by including more combinations of memory cells.

Original: 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 larger block size.
Revised: 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 larger block size.

Original: Costs for a smaller block size g
Revised: Costs for a smaller block size g

Original: With model parallelism (Lepikhin et al., 2021), multiple GPUs contains different memory block and parallelize the calculations.
Revised: With model parallelism (Lepikhin et al., 2021), multiple GPUs contains different memory block and parallelize the calculations.

Original: If with block size g = 4 · d , a token is only routed to 1 memory block on one device, each device doubles its chance to receive more tokens with block size g = 2 · d .
Revised: If with block size g = 4 · d , a token is only routed to 1 memory block on one device, each device doubles its chance to receive more tokens with block size g = 2 · d .

Original: Therefore, each GPU processes more tokens and requires more computation time.
Revised: Therefore, each GPU processes more tokens and requires more computation time.

Original: Better implementations could be developed to make a smaller block size for practical usage.
Revised: Better implementations could be developed to make a smaller block size for practical usage.

Original: 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 a larger table for the static gate.
Revised: 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 a larger table for the static gate.

Original: 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 B.3.1).
Revised: 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 B.3.1).

Original: 5. B LOCK SELECTION METHOD
Revised: 5. B LOCK SELECTION METHOD

Original: Next, we investigate the impact of the selection method, speciﬁcally, the FLOPs-perplexity trade-off for direct and indirect methods to determine the overall usefulness of each S-FFN method.
Revised: Next, we investigate the impact of the selection method, speciﬁcally, the FLOPs-perplexity trade-off for direct and indirect methods to determine the overall usefulness of each S-FFN method.

Original: FLOPs-perplexity trade-off
Revised: FLOPs-perplexity trade-off

Original: We study the efﬁciency of direct and indirect selection methods in S-FFN models characterized by FLOPS-perplexity trade-off.
Revised: We study the efﬁciency of direct and indirect selection methods in S-FFN models characterized by FLOPS-perplexity trade-off.

Original: We conduct experiments across different scales of the memory by varying E 2 { 4 , 16 } ; additionally, we run E = 32 for PKM.a more detailed
Revised: We conduct experiments across different scales of the memory by varying E 2 { 4 , 16 } ; additionally, we run E = 32 for PKM.

Original: Indirect methods (orange

Original: ) have more perplexity improvement relative increases

Original: FLOPs direct methods

Original: In Fig.
Revised: In Fig.

Original: 3, we marginalize different factors used in the two selection methods — i.e. types of gates, factorization techniques on key table,or memory sizes controlled by E , etc. — and consider each type of selection method as a whole.
Revised: 3, we marginalize different factors used in the two selection methods — i.e. types of gates, factorization techniques on key table, etc. — and consider each type of selection method as a whole.

Original: When we change different marginalized factors, we observe that indirect methods tend to improve more as we use more FLOPs.
Revised: When we change different marginalized factors, we observe that indirect methods tend to improve more as we use more FLOPs (with larger memory sizes controlled by E ).

Original: Thus, the indirect method has a better FLOPs-perplexity trade-off.
Revised: Thus, the indirect method has a better FLOPs-perplexity trade-off.

Original: Effect of gating function We start with  comparisons among PKM E =16 ,PKM E =32 ,
Revised: Effect of gating function We start with contrastive comparisons among PKM-FFN E =16 ,

Original:  RandHash E
Revised: PKM E =32 , RandHash E

Original: =16 with memory block size g = 1 and 4096 active memory blocks from which we can learn important lessons to improve the design of a S-FFN architecture.
Revised: =16 with memory block size g = 1 and 4096 active memory blocks. From the three parameter-matched models, we can learn important lessons to improve the design of gate:

Original: We ﬁrst make several observations about the “edit distance” between them: 1) RandHash E

Original:  =16 and PKM E =16 differ by inserting an full-parameter key table K and substituting hash table, which haslearnable parameters, for a factorized key table, which has minuscule learnable parameters ( < 0 .
Revised: Comparing with PKM-FFN E =16 , PKM E =32 essentially moves the parameters from a fullparamter key table to double the size of value table.

Original: 3% ; §2.2.2); 2) other than selection methods, RandHash E =16 and PKM E=32 differ by substituting the fullparameter key table K for doubling the value table V . For further analysis, we introduce PKM-FFN whose output g ( x ) is the same as the m from a PKM and g i
Revised: PKM-FFN E =16 and RandHash E

Original: = m i . With PKM-FFN added, the comparisons above become more contrastive:
Revised: =16 have the same key and value tables; but the former uses gate jointly learned with key table, while the later uses a learning-free gate.

Original: [Table]
Revised: [Table]

Original: PKM-FFN E =16 uses learned gate, while RandHash E =16 uses static gate 2. Different from PKM E =16 , PKM-FFN E =16 add the additional full-parameter table to be the key table, while PKM E =32 add it to double the size d m of value table.
Revised: As shown in Table

Original: As shown in Table 3, on out-of-domain, PKM-FFN E =16 outperforms PKM E =32 (16.06) by 0 .perplexity and slightly outperform RandHash E =16 by 0.16.
Revised: [Table] on out-of-domain, PKM-FFN E =16 outperforms PKM E =32 (16.06) by 0 .perplexity and slightly outperform RandHash E =16 by 0.16.

Original: Therefore, it is essential to have a full-parameter, and thus expressive enough, key table to produce memory coefﬁcients .
Revised: Therefore, it is essential to have a full-parameter [Table] and thus expressive enough, key table to produce memory coefﬁcients [Table]

Original: Table 3 shows the improvement of VanillaM E =16 , PKM-FFN E =16 , RandHash E
Revised: Table 3 shows the improvement of VanillaM E =16 , PKM-FFN E =16 , RandHash E

Original: =16 over
Revised: =16 over

Original: Dense Baseline (16.96) are 2.27, 1.77, and 1.61 respectively on out-of-domain.
Revised: Dense Baseline (16.96) are 2.27, 1.77, and 1.61 respectively on out-of-domain.

Original: They only differ by how much the dot product information is used for selection — VanillaM directly uses it, PKM-FFN indirectly uses it when the gate is jointly trained, and RandHash completely ignores it.
Revised: They only differ by how much they depends on key table for selection — VanillaM directly uses it, PKM-FFN indirectly gain information from it, and RandHash completely ignores it.

Original: Therefore, we conclude that the more dependent on dot product the selection method is, the better language model it will lead to; and indirect usage (PKM-FFN) is not enough.
Revised: Therefore, we conclude that the more dependent on key table the selection method is, the better language model it will lead to; and indirect usage (PKM-FFN) is not enough.

Original:  A NEW ROUTING METHOD — A VG -K
Revised:  A NEW ROUTING METHOD — A VG -K

Revised: Shown by Fig.

Original: Based on the insights from  the analysis above, two factors characterize abetter gating function: (1) having full-parameter key and value tables K g , V g , it improves expressiveness with direct operations between x and K g such as using dot product, and (2) is  compute-efﬁcient with smaller g .
Revised: 3, a gated MoE has clear advantage over sparse memory. Basing on the contrastive analysis from §5.2, we believe the current MoE already has a good design choice — a full-parameter key table. However, an important improvement is to make more use of each expert’s key table for routing tokens. Additionally, it needs to be compute-efﬁcient with smaller g (§5.1).

Original: To this end, we propose a new routing method — Avg-K .
Revised: To this end, we propose a new routing method — Avg-K .

Original: We represent each block with the average of its key table K ( i ) 2 R g ⇥ d along g -dimension — e i = 1 g
Revised: We represent each block with the average of its key table K ( i ) 2 R g ⇥ d along g -dimension — e i = 1 g

Original: P g  1 j =0 k
Revised: P g  1 j =0 k

Original: ( i )
Revised: ( i )

Original: j
Revised: j

Original: = Avg ( K ( i ) ,dim=0 ) .
Revised: = Avg ( K ( i ) ,dim=0 ) .

Original: 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):
Revised: 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):

Original: [Equation]
Revised: [Equation]

Revised: [Table]

Original: Due to the linearity of averaging, the operation e i · x is equivalent to calculating the average of dot products within a block without GeLU.
Revised: Due to the linearity of averaging, the operation e i · x is equivalent to calculating the average of dot products within a block without GeLU.

Original: Since all tokens share the averages, our method is efﬁcient.
Revised: Since all tokens share the averages, our method is efﬁcient.

Revised: We provide more rationale for our choice of average function in Appendix C.1.

Original: Table 4 shows that the proposed S-FFN design outperforms all other indirect methods.
Revised: Table 4 shows that the proposed S-FFN design outperforms all other indirect methods.

Original: 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.
Revised: 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.

Original: Due to the similar form as VanillaM, in Appendix C.1, we include  contrastive analysis  to understand the drastic improvement from block size g = 4096 to 256 .
Revised: In Appendix C.2, we include a contrastive analysis with VanillaM to understand the drastic improvement from block size g = 4096 to 256 .

Original: Load balancing We provide a detailed analysis of Avg-K ’s load balancing behavior in Appendix
Revised: Load balancing Previously, load balancing has two-fold implication: 1) to improve empirical training time Fedus et al.

Original: C.2.
Revised: (2021); 2) more importantly, to prevent the expert representations from degeneration — i.e. routing most tokens to a single expert (Shazeer et al., 2017; Eigen et al., 2014).

Original: Without load balancing constraints, SMoE with learned gate often degenerate to using only  a handful of memory blocks (§2.2.1) while Avg-K does not suffer training instability. We found [Table]
Revised: Based on our observation, we hypothesize that load balancing is only optional for learning a good expert representation (i.e., the second implication) since our experiment with Switch enforces load balancing and experiment with Avg-K doesn’t.

Original:  Avg-Kuses blocks ina non-uniform fashion and varies across layers. Overall, Avg-K still matches
Revised: We provide analysis for block usage for Avg-K in

Original: Switch’s performance with memory block size g = 4096 or out-performs all other SMoE with memory block size g 

Original: [Table]
Revised: Appendix C.3.

Original: Interpretability Geva et al.
Revised: Interpretability Geva et al.

Original: (2021) shows that rows of K  corresponds (i [Table] e. by dot products) to different —both surface- and semantic-level — textual features.
Revised: (2021) shows that rows of K of an end-to-end trained model strongly corresponds (i.e. by dot product) to human-interpretable — surface- and semantic-level — textual features.

Original: Since Avg-K heavily depends on dot products, their analysis could be easily extended for interpreting our model.
Revised: Since Avg-K heavily depends on dot products, their analysis could be easily extended for interpreting our model.

Original: In comparison, RandHash is a lot harder to interpret due to random selection, and the relationship between e i and K ( i ) remains unclear for Switch or SMoE with conventional learned gate.
Revised: In comparison, RandHash is a lot harder to interpret due to random selection, and the relationship between e i and K ( i ) remains unclear for Switch or SMoE with conventional learned gate.

Original: We defer such analysis to future work since this is beyond the scope of our paper.
Revised: We defer such analysis to future work since this is beyond the scope of our paper.

Original: Limitation Since Avg-K essentially applies an average pooling to the key table K g , a better alternative may exist.
Revised: Limitation Since Avg-K essentially applies an average pooling to the key table K g , a better alternative may exist.

Original: Our method also heavily depends on dot product information, but this might not be the best information to be used.
Revised: Our method also heavily depends on dot product information, but this might not be the best information to be used.

Original: Additionally, in large-scale SMoE training, the speed is limited by the most heavy-loaded GPU when model parallelism is used.
Revised: Additionally, in large-scale SMoE training, the speed is limited by the most heavy-loaded GPU when model parallelism is used.

Original: Therefore, load balancing is essential.
Revised: Therefore, load balancing is essential.

Original: We note that our scale is relatively small and does not use model parallelism, so the problem is not pronounced for us.
Revised: We note that our scale is relatively small and does not use model parallelism, so the problem is not pronounced for us.

Original: Future follow-up should look at how to incorporate load balancing into the uniﬁed framework and how Avg-K corporate with load balancing techniques.
Revised: Future follow-up should look at how to incorporate load balancing into the uniﬁed framework. Such uniﬁcation requires more advanced theoretical connection with memory block and block selection method, which likely involves consideration of training procedure.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Excluded S-FFN Terraformer’s Jaszczur et al.
Revised: Excluded S-FFN Terraformer’s Jaszczur et al.

Original: (2021) technique on FFN is closest to our PKM-FFN because there is a low-rank learned gate to operate on each memory cells for selection.
Revised: (2021) technique on FFN is closest to our PKM-FFN because there is a low-rank learned gate to operate on each memory cells for selection.

Original: However, we exclude this method because our framework uses all memory cells in each block, but Terraformer selects 1 cell in each memory block. In Appendix D.1, we studied and found their “one-per-block” design limits the model’s performance.
Revised: However, we exclude this method because our framework uses all memory cells in each block, but Terraformer selects 1 cell in each memory block (see our study in Appendix D.1).

Original: In ﬁnetuning scenario, Zhang et al.
Revised: In ﬁnetuning scenario, Zhang et al.

Original: (2022b)also studies the connection between FFN and SMoE by splitting trained FFN into memory blocks of various sizes and separately learning a gate.
Revised: (2022b) studies the connection between FFN and SMoE by turning trained FFN into experts and separately learning a gate.

Original: In contrast, we focus on pretraining from scratch .
Revised: In contrast, we focus on pretraining from scratch .

Original: 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 efﬁciency.
Revised: 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 efﬁciency.

Original: For example, one could process the unfactorized key table by ANN methods like FAISS (Johnson et al., 2021) and ScaNN (Guo et al., 2020).
Revised: For example, one could process the unfactorized key table by ANN methods like FAISS (Johnson et al., 2021) and ScaNN (Guo et al., 2020).

Original: One successful example is applying vanilla
Revised: One successful example is applying vanilla

Original: Locality-Sensitive Hashing to Reformer (Kitaev et al., 2020).
Revised: Locality-Sensitive Hashing to Reformer (Kitaev et al., 2020).

Original: However, in our preliminary study, we found that perplexity is greatly affected by the search quality, and building a data structure after every update is expensive and hard to avoid.
Revised: However, in our preliminary study, we found that perplexity is greatly affected by the search quality, and building a data structure after every update is expensive and hard to avoid.

Original: We leave detailed discussion and analysis to Appendix D.2.
Revised: We leave detailed discussion and analysis to Appendix D.2.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We provide a uniﬁed framework for designing sparse FFN in transformers and analyze existing
Revised: We provide a uniﬁed framework for designing sparse FFN in transformers and analyze existing

Original: S-FFN methods such as SMoEs in the language modeling task.
Revised: S-FFN methods such as SMoEs in the language modeling task.

Original: Using this framework, we found that smaller memory block (e.g. expert) size improves perplexity at the cost of slightly higher computation cost.
Revised: Using this framework, we found that smaller memory block (e.g. expert) size improves perplexity at the cost of slightly higher computation cost.

Original: Selection methods with gates have better FLOPs-Perplexity trade-offs than without, while the gating function in current SMoEs is sub-optimal.
Revised: Selection methods with gates have better FLOPs-Perplexity trade-offs than without, while the gating function in current SMoEs is sub-optimal.

Original: This framework enables us to instantiate a simpler
Revised: This framework enables us to instantiate a simpler

Original: S-FFN architecture that outperforms SMoEs while still being efﬁcient in training and inference.
Revised: S-FFN architecture that outperforms SMoEs while still being efﬁcient in training and inference.

Revised: A T RAINING S ETTING A. H YPERPARAMETERS

Revised: Table 6 speciﬁes shared hyperparameters across all experiments, in which Table 5a contains ones for training data, optimizer, and efﬁcient infrastructure techniques; and Table 5b for architecture.

Revised: Then,

Revised: Table 6a describes the hyperparameters speciﬁcally for Switch, Table 6b for LoRKM, Table 6c for

Revised: PKM,

Revised: A.

Revised: D ATA

Revised: Here is a detailed description of our pretraining corpus.

Revised: • BookCorpus (Zhu et al., 2015) consists of more than 10K unpublished books (4GB); • English Wikipedia , excluding lists, tables and headers (12GB); • CC-News (Nagel, 2016) contains 63 millions English news articles crawled between September 2016 and February 2019 (76GB); • OpenWebText (Gokaslan & Cohen, 2019), an open source recreation of the WebText dataset used to train GPT-2 (38GB);