Towards A Unified View of Sparse Feed-Forward Network in Pretraining Large Language Model (2023)

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

## introduction large-scale pretrained language models (llms) achieve remarkable performance and generalization ability for nlp tasks @xcite @xcite @xcite @xcite . 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 @xcite @xcite @xcite . however, the computational resources required to train larger language models are formidable, calling for more efficient training and inference solutions of llms @xcite @xcite . 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 on scaling up a transformer's feed-forward network (ffn) with sparsely activated parameters, resulting in a scaled and sparse ffn (s-ffn). there have been two major approaches to achieve s-ffn. one treats s-ffn as a neural memory @xcite where a sparse memory retrieves and activates only parts of the memory cells @xcite . the other adopts mixture-of-expert network (moe) @xcite @xcite @xcite 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. while both memory and moe models achieve s-ffn, they have been considered two distinct approaches. 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 capability and computation efficiency? can the effective ingredients of each method be transferred and combined to improve performance further? in order to answer these questions, we start from the neural memory view of ffn @xcite ( §2.1) and reduce all s-ffn's to the same mathematical form ( §3.1). then, we 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: transformer layer option 1: sparse mixture-of-expert option 2: sparse neural memory replace replace top-1 gshard, switch transformer, hashlayer, baselayer... lowrank key, product key + load balancing constraint gate top-6 ffn 2 ffn 1 ffn 3 ffn 4 ffn figure 1: sparse mixture-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/computation trade-offs. ## a unified view of sparse ffns we show the connections between moe and neural memory despite their different formulations on the surface. we first derive a variant form of moe to establish its connection with sparse memory ( §3.1). then, we propose a unified framework for s-ffn ( §3.2). ## analysis results in this section, we use the proposed unified view to systematically study the design choice of s-ffn. specifically, (1). 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). (2). both direct and indirect block selection methods lead to lower perplexity than a standard ffn, but which type of method has better flops-perplexity trade-off and what are the relative efficacy and efficiency of different methods require further study ( §5.2). ## conclusion we provide a unified framework for designing sparse ffn in transformers and analyze existing s-ffn methods such as moes in the language modeling task. using this framework based on sparse neural memory, we found that smaller memory block (e.g. expert) size improves perplexity at the cost of slightly higher computation cost. selection methods with gates have better flops-perplexity trade-offs than without, while the gating function in current moes is sub-optimal. this framework enables us to instantiate a simpler s-ffn architecture that outperforms moes while still being efficient in training and inference.