Approximating Two-Layer Feedforward Networks for Efficient Transformers (2023)

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

## introduction despite impressive results recently achieved by large language models (llms; @xcite ; @xcite ; @xcite ), vast resource requirement remains their obvious limitation. in fact, most existing llms, such as @xcite , cannot be trained, finetuned or even evaluated without access to enormous compute. many recent works strive to develop llms that, at least, enable inference with limited resources (e.g., on consumer hardware), e.g., by building "smaller" yet capable lms @xcite @xcite or developing post-training quantization methods @xcite . while these methods are gaining popularity, a principled solution for resource-efficient neural networks (nns) remains elusive. one promising approach explored by several recent works on extremely-large lms is the sparse mixture of experts (moe; @xcite ; @xcite ; @xcite ; @xcite ; @xcite ; @xcite ). unlike their dense counterparts, moes only compute a subset of their activations (i.e, only a few experts) at each step, offering reduced computation and memory costs. however, moes are not yet generally adopted as a generic/to-go approach, perhaps because of certain common beliefs on moes: we present moes within a unified framework of methods that approximate two-layer feedforward networks, which includes product-key memories (pkms; @xcite ) and top-k sparsification. this principled view not only allows us to conceptually group and compare moes with pkms, it also provides insights on design choices for improving these methods. our resulting moe transformer variant outperforms our improved pkms, and performs as well as or even outperforms the dense baseline, while using a fraction of its compute for both training and inference. importantly, unlike prior work, we compare our moes with dense baselines with the same number of total trainable parameters, which is crucial for proper evaluation in language modeling. we conduct experiments on the standard wikitext-103 (at two different model scales) and enwik8 datasets. we demonstrate that moes are not limited to extremely-large lms, but useful as a generic approach for resource-efficient nns at any scale, and in line with the recent trend of improving "smaller" models @xcite @xcite . finally, we release a cuda kernel for our moe layers which allows for achieving faster wall clock time and large memory reduction compared to the dense model. 2 ## background transformers @xcite have two main building blocks: the self-attention layer @xcite @xcite , and the two-layer feedforward, i.e, multi-layer perceptron (mlp) block. acceleration and memory reduction of the self-attention is rather well explored (see, e.g., linear attention dating back to the unnormalised linear transformers of 1991 @xcite @xcite ), and very efficient implementations @xcite are also available. in constrast, resource-efficient mlp blocks are still underexplored. this is our main focus, and it is of particular relevance today, as the proportion of the total parameter counts, compute and memory requirements due to mlp blocks in transformers is increasing in ever-growing llms. let d model , d ff denote positive integers. each transformer mlp block consists of one upprojection layer with a weight matrix w 1 ∈ r d ff ×d model where typically d @xmath0, and one down-projection layer with parameters w 2 ∈ r d model ×d ff that projects it back to the original size. non-linearity (typically relu) is applied between these two layers. that is, an input @xmath1 alternatively, this layer can be viewed as a keyvalue memory accessed by attention @xcite al. (2017) 3 ,geva et al. (2021) ), where keys and values are rows and columns of weight matrices w 1 and w 2 : then, the output is computed as "attention": unlike the standard self-attention, the mlp block uses a relu activation function (instead of the softmax) without scaling. it has been observed that, in practice, only a few of the factors k ⊺ i x are positive @xcite , making the first layer's output, i.e., u, sparse. concretely, @xcite report that in a transformer with d @xmath2, 10% of the channels account for 90% of the total activation mass. we confirm this trend in our own preliminary study. fig. 1 shows the average number of non-zero units in u of size d @xmath3. a.2 for more details). the number is below 200 for all layers. this suggests that the mlp block can be approximated without a significant performance loss. note that this is also supported by the findings of @xcite . ## approximating 2-layer mlps here we present a unified view on methods to approximate 2-layer mlps (sec. 2) that includes many existing methods such as moes (sec. 3.3) and pkms (sec. 3.2). the core idea is to approximate the sum in eq. 5, i.e., @xmath4, ..., d ff } of the key-value pairs, i.e., ŷ = @xmath5. the intuition of this approximation is as follows. we assume that a good approximation ŷ of y is the one that minimizes their euclidean distance @xmath6, which can now be expressed as let k be a positive integer. the general idea of all methods discussed in this work is to keep k pairs (k i , v i ) whose contribution α i is the highest, and ignore other low-contribution pairs. the goal is to find the best mechanism to select such k pairs. here we discuss three variants: top-k activation (sec. 3.1), product-key memories (pkms, sec. 3.2), and mixture of experts (moes, sec. 3.3). ## existing moe variants several variations of moes have been proposed with many different details. here we briefly review the most popular and representative ones (e.g., we do not cover those that make use of reinforcement learning for expert routing) before describing our improved version in sec. 5. we'll review their expert selection function and regularization method, and highlight their key characteristics. sparsely gated mixtures of experts. @xcite have revisited moes @xcite with the top-k operation, allowing a reduction in its resource demands. their method is basically the one described in sec. 3.3 (with re-normalization after top-k) except that they use a noisy gating function: ## improving mixture of experts here we present our improved moe variant, which we call σ-moe. we conduct thorough ablation studies on our design choices in sec. 6. σ-moe expert selection function. our moe make use of the top-k operation (unlike base). the activation we use on the selection function is sigmoid (as in eq. 18 of base) instead of softmax used in switch transformer and sparsely gated mixtures of experts. this choice is motivated by the view of moes as approximate 2-layer mlps (sec. 3). in fact, softmax introduces competition between experts. no such competition between channels is used in the regular 2-layer mlp (i.e., there is no constraint on α i in eq. 5). this suggests that, in principle, no competition is needed between terms in the sum of eq. 12 in the moe either, to induce sparsity. it is also well known to practitioners that softmax as regular activation negatively affects the trainability of standard mlps. softmax combined with top-k can also encourage expert collapsing: when the selection score of one expert increases, the score of the others automatically decreases. for all these reasons, we opt for sigmoid instead of softmax; we experimentally confirm that this is indeed a good choice. additionally, looking at moes in this framework gives us hints on combining them with top-k activation (sec. 3.1) for further acceleration. we can calculate u @xmath7. 11) for the selected experts and perform an additional top-k to keep the highest units among them and set the rest to zero. we leave this for future work. in contrast, we initialize all weight matrices identically to the pre-layernorm dense baselines, not taking in account the smaller size of the individual experts, i.e., w e 1 ∼ n (0, we also take special care when initializing w 3 of the selection function. we initialize it to a normal distribution with the same standard deviation as w e furthermore, we propose to randomly drop complete experts, during training; we refer to this as expert dropout. unlike the standard dropout on the activation level, we do not apply rescaling, i.e., ## experiments our experimental setup is based on @xcite 's transformer xl with some modifications: we use pre-layer norm and reduce the number of training steps to 100k to reduce the computational budget. also, to match the parameter counts between the baseline and moes, we slightly modify the hyperparameters of the baselines @xcite . in fact, our moe cuda kernel can only work with dimensions divisible by 4. we round the original sizes up to the next suitable number, e.g., we change d model of our 47m-parameter wikitext-103 model from the original 410 to 412. furthermore, since moes require extra parameters for the expert selection function, we compensate for these by increasing the d ff of the baseline model to match the number of parameters. our modified baseline model on enwik8 still has 41m parameters and performs similarly to the original transformer xl (see tab. 1 ). for wikitext-103, we use subword units @xcite using sentencepiece tokenizer @xcite instead of the word-level vocabulary, to avoid extra tricks required to reduce the parameter count and compute requirement resulting from the huge vocabulary size. on wikitext-103, we consider two different model sizes: a 47m-parameter one (denoted by "wt-s" for "small"), and a 262m-parameter one ("wt-b" for "big"). we refer to enwik8 as "e8" in certain tables. for more details, see appendix b. for all the methods considered, we use them in every mlp block of the model, which is not a common practice in the literature. typically, moe (or other approximation methods) is used only once every n th layer or even only in one layer. this is not satisfactory since our goal is to find a generally applicable method that can accelerate all layers across the whole model. moreover, this amplifies the difference between different methods, helping better illustrate effects of each of the design choices. ## conclusion our novel view unifies methods that approximate 2-layer mlps, such as top-k, mixture of experts (moe) and product-key memory (pkm) methods. while top-k by itself provides limited performance improvements and speedups, further speedup requires pkm or moe. a non-competitive activation function inspired by our unified view improves both pkm and moe. further novel enhancements of moes yield our σ-moe which outperforms existing moes. importantly, our σ-moe with moderate sparsity matches the performance of parameter-equal dense baselines while being much more resource-efficient. our new insights improve the training of language models with limited hardware resources, making language modeling research more accessible.