Dynamic and Efficient Inference for Text Generation via BERT Family (2023)

Abstract: Despite the excellent performance of Pre-trained Language Models on many text generation tasks, they suffer from inefficient inference on computation and memory due to their large-scale parameters and the universal autoregressive decoding paradigm. In this work, we propose a novel fine-tuning method DEER , which can make a single pre-trained model support D ynamic and E fficient inf ER ence and achieve an adaptive trade-off between model performance and latency. In particular, our critical insight is to jointly utilize the non-autoregressive (NAR) generation and dynamic parameter pruning techniques, which can flexibly control the decoding iteration steps and model sizes according to memory and latency limitations. Besides, we also explore the effectiveness of the pre-trained MLMs (i.e., the BERT family) for text generation tasks since their bidirectional attention nature is more suitable for the NAR training objective. Extensive experiments on both monolingual and multilingual pre-trained MLMs demonstrate the effectiveness of our proposed DEER method by consistently achieving (1) higher BLEU scores than the strong autoregressive Transformer model on three neural machine translation tasks with 3 \to 12 times speedup, (2) competitive performance (but with much faster inference speed) compared with the BART model on four GLGE benchmark tasks. Our code will be publicly available at GitHub https://github.com/dropreg/DEER .

## introduction large-scale pre-trained language models @xcite @xcite have shown great potential in achieving impressive performance; however, they are accompanied by substantial computational complexities and occupy significant memory space. these factors pose obstacles to their practical implementation in real-world applications. while recent studies @xcite have made attempts to address the challenges associated with compressing and accelerating inference for pre-trained transformer models, the majority of these efforts have concentrated on techniques such as knowledge distillation @xcite , quantization @xcite , and parameter pruning @xcite . the pre-trained non-autoregressive generation paradigm has received limited attention and remains relatively unexplored. to fill this blank, we first summarize two main difficulties in the deployment and application of large generative models. firstly, the prevailing generative models currently employ an autoregressive approach to generate target tokens incrementally, as seen in models like bart @xcite and t5 @xcite . while these models have gained popularity and demonstrated effectiveness, their autoregressive nature hinders efficient inference through parallelization, resulting in inefficiencies. secondly, task-specific fine-tuning is crucial when deploying pre-trained models on diverse edge devices @xcite . it is impractical to adopt a single model for all devices due to variations in memory capacity and latency constraints. consequently, multiple models with different architectural configurations need to be trained to meet these device-specific requirements, leading to additional resource consumption and increased carbon emissions. to address these challenges, we propose a novel joint training strategy called deer. this strategy offers fast inference by employing a non-autoregressive generation approach and provides flexibility in model size through the utilization of dynamic block pruning. concretely, we choose the bert family models to implement our deer method because their bidirectional attention mechanism is more suitable for non-autoregressive generation tasks. to allow encoder-based models for text generation and re-duce the error accumulation in length prediction, we combine the training objective of connectionist temporal classification @xcite (ctc) and levenshtein transformer @xcite for multi-task training. compared with previous methods, this approach has a better result than the iterative approach at the first generation step and can further improve the iteration refinement performance with the obtained good initialization. moreover, to easily adapt the bert family to non-autoregressive generation without introducing extra parameters or cumbersome post-training, we design task-specific input formats and self-attention masks @xcite . different input formats and self-attention masks can dynamically control the source and target information interaction and remedy the structural defects of the encoder-based model, making it competent for text generation. our deer also incorporates dynamic block pruning for model training and inference to make the bert family with adaptive model size. meanwhile, we use score-based parameter mask and sparsity regularization to choose and train the suitable model size for current devices, referring to movement pruning @xcite @xcite . unlike current pruning works, deer is a one-stage training method without two-stage fine-tuning for sub-models and can dynamically choose a model size instead of a fixed size. in inference, we gather the weight from the trained model for different devices when its importance score is larger than the global threshold. the sparsity regularization is also crucial, which can encourage the model to decrease the importance of weight score and control the sparsity level. we conducted extensive experiments to validate and analyze the effectiveness of our proposed deer method on both monolingual and multilingual models from the bert family. in particular, our deer method outperforms the ar model, achieving a 3× → 12× speedup on three neural machine translation tasks. additionally, deer overcomes the limitations of memory and latency, enabling support for various hardware devices without compromising the task performance of the original model. these results demonstrate the efficacy of our deer method in improving inference speed and compatibility with diverse hardware devices, while maintaining or surpassing the task performance of the original models. in a nutshell, our contributions are as follows: ## methods in this section, we first exhibit how to fine-tune the bert family model (e.g., xlm-r and roberta) as a nar text generator, which supports single-step generation ( § 3.1) and iterative-based generation ( § 3.2), as shown in figure 1 . then we introduce the dynamic block pruning for model training to reduce the computation and memory consumption in inference with dynamic model size ( § 3.3). ## experiments datasets we evaluate deer on multiple widely used text generation tasks to verify its effectiveness: 1) neural machine translation (nmt), we conduct experiments on three benchmark translation compute gradients and update weights 23: end while datasets: iwslt'14 @xmath0, wmt'16 @xmath1, and wmt'14 @xmath2. for all translation tasks, we report the results of raw (raw) data and knowledge distilled (kd) data, respectively. we use the same training/validation/test sets as in previous works and the belu score as the evaluation metric for a fair comparison. 2) monolingual text generation scenarios, we evaluate the efficacy of the proposed deer on four glge benchmarks foot_3 , including text summarization (xsum @xcite and msnews) and question generation tasks (squad 1.1 @xcite and msqg). for each dataset, we first train bart base as a teacher model and generate the distilled data as deer training data, which can reduce the multi-modality problem @xcite to facilitate the learning of nar models. the official script 6 is used for evaluation. descriptions and data statistics are shown in appendix a. ## main results in this section, we explore whether deer can provide dynamic and efficient inference on multiple tasks and datasets by evaluating its nonautoregressive capabilities and model performance with adaptive model sizes. ## conclusion in this work, we propose deer, a novel fine-tuning method that supports dynamic and efficient inference to adapt to the memory and latency limitations during deployment. our approach has achieved impressive results on multiple natural language processing tasks, including the glge benchmark and three machine translation datasets. furthermore, we have observed that the issue of length prediction consistently limits the performance of the model, especially when dealing with raw datasets. the model struggles to accurately determine the length of the target data, which somewhat affects the model evaluation. in our future work, we will prioritize addressing the challenge of length prediction, aiming to make it more convenient and applicable to a wider range of tasks and scenarios.