NASH : A Simple Unified Framework of Structured Pruning for Accelerating Encoder-Decoder Language Models (2023)

Abstract: Structured pruning methods have proven effective in reducing the model size and accelerating inference speed in various network architectures such as Transformers. Despite the versatility of encoder-decoder models in numerous NLP tasks, the structured pruning methods on such models are relatively less explored compared to encoder-only models. In this study, we investigate the behavior of the structured pruning of the encoder-decoder models in the decoupled pruning perspective of the encoder and decoder component, respectively. Our findings highlight two insights: (1) the number of decoder layers is the dominant factor of inference speed, and (2) low sparsity in the pruned encoder network enhances generation quality. Motivated by these findings, we propose a simple and effective framework, NASH, that narrows the encoder and shortens the decoder networks of encoder-decoder models. Extensive experiments on diverse generation and inference tasks validate the effectiveness of our method in both speedup and output quality.

## introduction in recent years, pre-trained language models (lms) have demonstrated their effectiveness in various downstream tasks, such as natural language understanding (nlu) and natural language generation (nlg). especially, there have been three main types of research, e.g.,encoder-only lms @xcite , decoder-only lms @xcite , and encoder-decoder lms @xcite @xcite , which aim for their specific expertise. among these various types of lms, we will focus on the widely studied and utilized encoder-decoder lms due to their flexibility in application across a range of tasks @xcite . on the other perspective of lm researches rather than performances, efficiency of lms (e.g.,computational and memory cost) have been intensively studied because of their huge computational requirements. this research direction is called model compression. among the various model compression techniques @xcite , pruning @xcite @xcite is a promising method that aims to remove redundant weights from networks, resulting in improved efficiency by saving storage capacity and enhancing inference speed. between structured pruning and unstructured pruning approaches, structured pruning is typically preferred in practice due to its relative ease of deployment on various types of hardware platforms compared to unstructured pruning @xcite . therefore, we focus on the structured pruning method specifically tailored for encoder-decoder lms. despite the remarkable advancements in encoder-decoder models, little attention has been given to structured pruning methods for encoderdecoder lms. this can be attributed to the inherent differences in the components that enhance pruning efficiency between encoder and decoder networks. consequently, traditional structured pruning methods that rely on encoder-only models may not effectively optimize encoder-decoder models. for instance, cofi @xcite , one of the sota encoder-only pruning methods, demonstrates a maximum speedup improvement of 1.53× on the cnn/dailymail @xcite dataset, with a rouge-l drop of 7.36%. this gain is considerably lower compared to the original result achieved on the encoder-only model applied to the worst qqp case in the glue @xcite , where the speedup reaches 11.0× with an accuracy drop of 1.20%. thus, it becomes crucial to investigate structured pruning methods that are specifically tailored for the encoder and decoder networks. to this end, in this paper, we pose the following question: how can we design a structured pruning method that effectively accelerates the encoderdecoder model while maintaining its performance? to the best of our knowledge, this study represents the first attempt to address this question. in order to accomplish this, we conduct systematic studies to examine the impact of structured pruning on the encoder and decoder networks, respectively. contribution. in this study, we propose an algorithm, nash, which is strongly motivated by two findings derived from our preliminary experiments. based on these findings, we propose an algorithm, illustrated in figure 1 , that consists of two parts: the encoder network, which enhances output quality by gradually reducing the width of each layer, and the decoder network, which achieves faster inference speed by uniformly selecting layers to reduce depth. we empirically evaluate the performance of nash on various nlg datasets including standard fine-tuning on a single task @xcite , multi-task learning scenarios, and recent instruction-tuning datasets @xcite . notably, in our experiments using t5-base, nash achieves a speedup of 2.5-4.2× while preserving 95% of the output quality. our experimental results show that nash can be a unified framework which is regardless of task difficulty and model type. ## preliminary transformers. we focus on the transformer network @xcite , which consists of the encoder and decoder architecture. the encoder architecture is composed of l blocks, and each block consists of a multi-head attention (mha) layer and a feed-forward (ffn) layer. an mha layer in the i-th transformer layer with n h heads outputs: and w (i,j) @xmath0, which represent the query, key, value, and output matrices, respectively. here, d and d h denote the hidden state dimension and attention head dimension, respectively. the output of the mha layer, denoted as x, is then fed into the ffn layer in the i-th transformer layer: here, the two fully-connected layers are parameterized by w 1 ∈ r d×df and w 2 ∈ r df ×d , with d f representing the dimension of the ffn layer. ## experimental motivations in this section, we separately investigate the behavior of the encoder and decoder when depth pruning is applied or not, using cofi-t5, the modified version of cofi @xcite tailored for t5 @xcite . particularly, in figure 2 , we study the results of four cases: encoder with depth-pruning (△), encoder without depth-pruning (⃝), decoder with depth-pruning (△), and decoder without depth-pruning (⃝). from these four types of cases, we aim to address the following questions: (1) does depth pruning exhibit different phenomena in each case? (2) what is the key factor for accelerating inference speed while preserving sufficient output quality? we provide detailed answers to each question by training t5-base with target sparsities of {60%, 70%, 80%, 90%, 95%} for the decoder cases, and {20%, 40%, 60%, 70%, 80%, 90%, 95%} for the encoder cases. 1 before delving into the detailed answers, we briefly address the first question: the impact of depth pruning when applied to the encoder and decoder, respectively. as depicted in figure 2 , depth pruning exhibits a significant influence on the decoder (as indicated by △ and ⃝), while the encoder shows negligible effects (as observed in △ and ⃝). consequently, the appropriate utilization of depth pruning becomes crucial. in the following paragraphs, we outline our key findings related to the second question to establish an effective structured pruning mechanism for encoder-decoder lms. finding 3.1. the number of layers in the decoder network is the dominant factor affecting the inference speed, while the decoder width does not have a significant impact. we evaluate the findings regarding the decoder network from two perspectives: (1) the effectiveness of layer-wise pruning and (2) the ineffectiveness of width pruning. firstly, as demonstrated in the second and fourth plots of figure 2 , the decoder exhibits a significant speedup with minor degradation (when comparing △ and ⃝), whereas the encoder shows no such effect (when comparing △ and ⃝). this indicates that layer-wise pruning plays a dominant role in pruning the decoder. on the other hand, when comparing the model size and speedup (as shown in the first and second plots of figure 2 with ⃝), width pruning reduces the model size but leads to both performance degradation and negligible speedup. this suggests that width pruning is not effective for the decoder. to further investigate finding 3.1, we investigate the inference speed of transformer layers, with a specific focus on understanding why width pruning is ineffective. this analysis involves two key observations: (1) finding the metric that synergizes with width pruning, and (2) identifying the component that predominantly consumes computational resources. according to figure 3 , width pruning can have a significant impact on the computational cost as the sequence length increases. however, due to the inherent nature of the autoregressive decoding process, the decoder network is constrained to a sequence length of 1. as a result, width pruning cannot effectively improve the speed of the decoder network. furthermore, as illustrated in figure 4 , layer normalization (ln) and dropout (do) collectively contribute approximately 20-25% to the overall inference time. consequently, the time allocated to these fixed operations remains constant, leading to diminished efficiency in terms of inference speed. in conclusion, width pruning is not an appropriate approach for optimizing the decoder. finding 3.2. from the perspective of encoder pruning, while achieving high-level sparsity may not be desirable, attaining low-level sparsity not only slightly accelerates inference speed but also enhances performance. by comparing the ⃝ points and ⋆ in the second and fourth plots of figure 2 , we observe that encoder pruning yields a slight speedup along with improved performance. however, when the encoder network is heavily pruned, it experiences significant performance degradation. these findings emphasize the significance of considering pruning in both the decoder and encoder networks. furthermore, they provide insights into the necessity of employing distinct pruning strategies for these two networks, considering their unique characteristics. comparison with prior observations. our key findings provide valuable insights: the appropriate strategy for encoder-decoder models involves using a small number of layers for the decoder and minimal pruning for the encoder networks. importantly, our observations offer a more generalized understanding compared to previous works @xcite . unlike prior studies that manually determined model configurations for specific tasks such as machine translation @xcite or nlu @xcite , our conclusions are derived automatically through gradual structured pruning and have been validated across both nlg and nlu tasks. furthermore, while the deepnarrow strategy proposed by @xcite demonstrates effectiveness in nlu tasks with short output sequences, it exhibits computational inefficiency when applied to nlg tasks. similarly, the contribution of processing time for encoder networks varies, necessitating the use of a narrower encoder architecture contrary to the approach proposed by @xcite . ## narrow encoder and shallow decoder based on the findings presented in section 3, we propose a structured pruning framework called nash (narrow encoder and shallow decoder) that is specifically optimized for encoder-decoder lms. our approach focuses on enhancing inference speed by utilizing uniform layer selection in the decoder network, deviating from the gradual pruning technique commonly employed in encoder-only models. additionally, we improve generation per-formance by applying gradual l0 regularization pruning specifically to the encoder network, inducing low sparsity instead of solely prioritizing inference speed improvement. ## conclusion we propose nash to address the lack of exploration in structured pruning of encoder-decoder lms. to design a structured pruning method suitable for encoder-decoder models, we first examine the behavior of pruned models with different strategies, focusing on inference speed and generation performance. our findings reveal that (1) the number of decoder network layers is the key factor in accelerating inference speed and (2) low sparsity pruning on the encoder network can enhance model performance. based on these insights, we develop nash, which constructs a narrow encoder and a shallow decoder network for encoder-decoder lms through gradual l0 regularization pruning and uniform layer selection, respectively. we demonstrate the superiority of nash in terms of speedup and output quality across various tasks. we strongly believe this work lays a strong foundation for further investigation into effective pruning approaches for encoder-decoder lm.