A Survey on Dynamic Neural Networks for Natural Language Processing (2023)

Abstract: Effectively scaling large Transformer models is a main driver of recent advances in natural language processing. Dynamic neural networks, as an emerging research direction, are capable of scaling up neural networks with sub-linear increases in computation and time by dynamically adjusting their computational path based on the input. Dynamic neural networks could be a promising solution to the growing parameter numbers of pretrained language models, allowing both model pretraining with trillions of parameters and faster inference on mobile devices. In this survey, we summarize the progress of three types of dynamic neural networks in NLP: skimming, mixture of experts, and early exit. We also highlight current challenges in dynamic neural networks and directions for future research.

## introduction scaling up model capacity is an obvious yet effective approach for better performance in natural language processing (nlp) tasks @xcite @xcite . however, the resulting increase in computational complexity and memory consumption becomes a bottleneck for scaling, making these models hard to train and use. on the other hand, it is not necessary to allocate the same amount of computation to all instances. for example, categorizing "i love you" as a positive sentence does not require a model containing dozens of transformer layers. to resolve the aforementioned problems, dynamic neural networks have been a significant thrust of recent research in nlp. dynamic networks can adjust their computational path based on the input for better efficiency, making it possible to train models with trillions of parameters and accelerate models in a low-resource setting. in this survey, we review the latest state of research on three types of dynamic neural networks that have been adopted in nlp: skimming, mixtures of experts (moe), and early exit, as illustrated in figure 1 . these three types of techniques share a common idea of dynamically adjusting computation with respect to input, to save computation through bypassing unnecessary modules in a large neural network. however, they implement the goal via different approaches. skimming was well-researched in the era of recurrent neural networks (rnn). skimming models save computation timewise by dynamically allocating computation to different time steps, based on the input tokens. since rnn models process the input sequence recurrently, it allows skimming models to achieve a substantial acceleration, especially when the sequence is long @xcite . different from rnn, recent works on transformers skip tokens between layers instead of time steps. for transformer models @xcite @xcite , the input tokens are fed into the model in parallel, while models have dozens of transformer layers. this motivates the development of moe and early exit. moe horizontally extends a feedforward neural network (ffnn) with multiple sub-networks. during inference, only one or a few of these sub-networks will be activated for computation, thus can save widthwise computation. early exit, on the other hand, terminates inference at an early layer, without exhausting full computational capacity, thus saves depthwise computation. early exit techniques often insert a series of lightweight classifiers which help decide when to exit, based on an exit strategy. note that this stream of works is distinct from static model acceleration, which is often referred to as model compression, including knowledge distillation, weight sharing, pruning and quantization @xcite @xcite @xcite (etc., see another survey @xcite ). the major difference is that the computational path in a statically compressed model does not condition on the input and is invariable for all examples in inference. these two streams of research are in fact orthogonal and recent works @xcite , @xcite and @xcite have shown that static and dynamic approaches can be combined for even faster inference and better performance. to summarize, our contribution is two-fold: (1) we review the latest studies on the topic of dynamic neural networks for nlp by providing a comprehensive comparison and organize them with a new taxonomy, as shown in figure 1 . ( 2) we analyze current challenges in dynamic neural networks and point out directions for future research. ## skimming skimming techniques, as summarized in skipping and early stopping skipping and early stopping aim to improve efficiency for a long sequence by skipping some tokens or stopping reading early. lstm-jump @xcite is a skipping mechanism to ignore irrelevant information for natural language understanding (nlu). at each step, the current states are used to compute a "jumping softmax", which decides how many steps to jump forward and whether to stop reading. lstm-jump employs policy gradient to train the model to make non-differentiable discrete jumping decisions. the reward is a binary function which rewards a correct prediction and penalizes an incorrect prediction of the label. compared to a standard lstm, lstm-jump achieves better accuracy with up to 6× speed-up. skip rnn @xcite introduces a binary gate to learn whether to skip a state update. if the gate decides to skip a time step, the hidden states will be directly copied without any update. to stop reading early as needed, rea-sonet @xcite introduces a terminal state which decides whether to terminate early for machine reading comprehension on each time step at the token level. jumper @xcite first splits a paragraph to several sub-sentences and encodes them into sentence embeddings. they then apply early stopping at a sentence when the policy network decides to stop reading. @xcite use eye-tracking devices and confirm that skipping and early stopping are common when humans read text. they propose reading inspired model to mimic the behaviors of humans, which allows the model to decide whether to skip a single time step or stop reading early. @xcite add a rereading operation to lstm-jump @xcite which allows the model to stay on the current time step, allocating more computation to important information. the aforementioned techniques can only go forward, which makes it impossible to regret if hav- ## mixture of experts increasing the number of parameters in a model often leads to increased computation and memory consumption. to take the advantages of parameter scaling without a proportional increase in computation, mixture of experts (moe) @xcite is introduced to large language models, as summarized in the key element of moe methods is the routing mechanism. the routing mechanism has to be lightweight, not to significantly slower the speed of the model. we categorize moe methods into two groups: learned routing and unlearnable routing. learned routing often requires some load balancing mechanisms to ensure that all experts are trained with enough examples thus are useful during inference. unlearnable routing usually slightly underperforms learned routing but does not require complicated load balancing. moe layers with learned routing a straightforward idea to implement moe is to learn a router to allocate inputs to experts. sparsely-gated moe layer @xcite ) contains up to thousands of feed-forward sub-networks with a trainable gating network which determines a sparse combination of these experts to use for each example. there are two major challenges to address: (1) sparsity. the gating network predicts a softmax weight for the experts based on the input. the gating network is trained by simple back-propagation, together with other parts of the model. then, only the top-k experts in the layer will be activated based on the softmax prediction of the gating network. they insert one moe layer between stacked lstm layers and achieve improvement on language modeling and machine translation tasks. (2) load balancing. @xcite observe a self-reinforcing phenomenon that the gating network tends to converge to a state where it always produces large weights for the same few experts. they resolve the problem by defining the importance of an expert relative to a batch of training examples to be batch-wise sum of the gate values for that expert. then, they introduce an additional loss, the square of the coefficient of variation of the set of importance values, to encourage a more balanced update during training. besides encouraging a balanced update, the authors also introduce a loss function with a smooth estimator that estimate the number of examples assigned to each expert for a batch of inputs, to encourage experts to receive roughly equal numbers of training examples. gshard @xcite enables scaling up multilingual neural machine translation transformer beyond 600 billion parameters. it adapts sparsely-gated moe @xcite to transformer @xcite by replacing every other feed forward layer with an moe layer, which routes to top-2 experts. when scaling to multiple devices, the moe layer is sharded across devices, i.e., each device has different allocated experts, while all other layers are replicated. to achieve workload balance, gshard employs a threshold, namely expert capacity, to limit the maximum number of tokens processed by one single expert. they also introduce a local group dispatching mechanism, which partitions all tokens in a training batch evenly into groups to be processed independently in parallel, to balance the overall workload. following @xcite , they use an additional loss to enforce even ## early exit early exit techniques aim to terminate model inference in early layers, to save computation and sometimes improve performance by resolving the overthinking problem @xcite , i.e., possible performance degradation at a later layer. it can be useful especially in the era of pretrained language models (plm), since increasing the size of plms can often lead to better performance, although a smaller model can already predict most examples (i.e., "easy examples") correctly. the main idea of early exit is to exit inference at an earlier layer, rather than the last layer. early exit often involves a series of internal classifiers inserted into a large network, providing signals for early exiting. the core of early exit methods is the exit criterion. based on their exit strategies, we categorize the early exit methods into three classes: confidence-based, ensemble-based and learning-based, as listed in ## challenges and future directions evaluation evaluating dynamic neural networks can be difficult since we cannot pre-define a few break points to compare different methods at the exact same amount of computation or time. elue score @xcite ) may be a promising solution to this problem by considering both computation and performance, depicting the pareto front. besides, different works have different calculation for speed-up ratio. for example, some works use the ratio of layers involved in computation to estimate speed-up ratio @xcite @xcite . this can be misleading since internal classifiers introduce extra computational costs, especially when more complicated mechanism introduced, e.g., futurelayer imitation @xcite . also, the reported speed of moe models, greatly differs on different hardware and distribution settings, making it hard to compare across papers. data parallelism one drawback of dynamic neural networks is their inefficiency on data parallelism. to be specific, moe methods introduce extra communication costs for dynamic routing and could be a bottleneck for efficiency. skimming and early exit methods often employ an "online inference" setting where the batch size is fixed to 1, to achieve maximum acceleration. however, for batched inference, the efficiency of these methods will drastically degrade, since the already-exited instances will have to wait all instances to exit, which causes a low parallelism and low utilization of gpu. optimized runtime since dynamic neural networks are an emerging type of neural networks, most hardware and libraries are not welloptimized for these models. for example, sparse matrix multiplication in moe needs specialized hardware and software support to achieve its theoretical efficiency. also, current dynamic neural networks are often implemented in eager execution, which prevents them from lowlevel optimization of graph execution. there have been works exploring optimized runtime for moe @xcite @xcite and early exit @xcite while more to be done in the future.