E - LANG : Energy-Based Joint Inferencing of Super and Swift Language Models (2022)

Abstract: Building huge and highly capable language models has been a trend in the past years. Despite their great performance, they incur high computational cost. A common solution is to apply model compression or choose light-weight architectures, which often need a separate fixed-size model for each desirable computational budget, and may lose performance in case of heavy compression. This paper proposes an effective dynamic inference approach, called E-LANG, which distributes the inference between large accurate Super-models and light-weight Swift models. To this end, a decision making module routes the inputs to Super or Swift models based on the energy characteristics of the representations in the latent space. This method is easily adoptable and architecture agnostic. As such, it can be applied to black-box pre-trained models without a need for architectural manipulations, reassembling of modules, or re-training. Unlike existing methods that are only applicable to encoder-only backbones and classification tasks, our method also works for encoder-decoder structures and sequence-to-sequence tasks such as translation. The E-LANG performance is verified through a set of experiments with T5 and BERT backbones on GLUE, SuperGLUE, and WMT. In particular, we outperform T5-11B with an average computations speed-up of 3.3X on GLUE and 2.9X on SuperGLUE. We also achieve BERT-based SOTA on GLUE with 3.2X less computations. Code and demo are available in supplementary materials.

## introduction with the introduction of influential language models such as bert @xcite , a trend in natural language processing (nlp) research has been to develop high capacity models and push their performance to new levels. consequently, state-of-the-art (sota) results were achieved on various benchmarks using these models; gpt-3 @xcite , xlnet @xcite , roberta @xcite , t5 @xcite , electra @xcite , and deberta @xcite to name a few. a potential down-side, however, is that the number of parameters or floating point operations (flops) for these models can get extremely large. for example, gshard @xcite comes with 600b parameters with an enormous amount of computation. this in turn results in a higher inference latency, which is not desirable for latency-sensitive applications. a common solution to speed-up the large language models is to apply model compression @xcite . although generally successful, compression does come with a trade-off on accuracy, and may lose performance if compression is heavy. in addition, these methods usually compress a model to a fixed smaller size, where a separate model is required for each possible computational budget. an alternative approach explored in the literature is to leverage dynamic inferencing in a way that examples may be routed to different (potentially lower cost) paths throughout the network. for example, a temporal early-exit model @xcite terminates the procedure of reading the input sequence when sufficient evidence has been found for accurate predictions. instance-wise early-exiting @xcite is another technique, which allows a sample to adaptively choose from multiple available exit nodes if some conditions are met. consequently, earlier exists require less computation and lead to a lower latency. adjusting the size of the model at the inference time by choosing adaptive width and depth is also another approach employed for dynamic inference @xcite . there is a variety of adaptive/dynamic inference approaches proposed, however, a general down-side for many of these methods is that often times they require a careful architecture design, manipulation of network modules, or even re-training. in this paper, we propose a simple but rather effective approach of dynamically distributing the inference between the original large model (called the super model) and a light-weight (e.g., compressed) model referred to as the swift model. to this end, we design an energy-based decision making module that routes examples to the appropriate model based on the negative free energy of the latent space representations, such that the swift model attains a high accuracy on the examples sent to it. the remaining samples are then forwarded to the super model that is supposed to have a good performance on all examples. since the swift model can make highly accurate predictions over the majority of the samples, e-lang significantly reduces the overall computational cost, while maintains the high accuracy of the super model. although simple, this strategy achieves sota results on multiple structures (e.g., t5 and bert) and benchmarks (e.g., glue and superglue). due to its desirable practical characteristics, this method is a strong candidate for the practical application of super models. the main contributions of the paper are as follows: ## related works as mentioned, compression is a widely used strategy to speed-up the large language models @xcite . this involves incorporating techniques such as quantization of weights and activations @xcite @xcite @xcite , knowledge distillation (kd) @xcite @xcite , pruning/sharing @xcite , multi-device distribution @xcite , or a combination of these techniques @xcite . among all the compression techniques, creating a fixed-size small version of large models along with distillation has been popular in the recent years. @xcite introduced distillbert, which was a smaller version of bert trained with distillation for general purposes. another compact variant of bert was proposed by mobile-bert @xcite in which inverted bottleneck structures and progressive knowledge transfer were used. tinybert @xcite also presented a novel two-stage transformer distillation for both pre-training and task-specific fine-tuning. in @xcite , the usage of grouped convolutions was studied to design squeezebert. elm @xcite , a layer mapping search framework, was also proposed for improving downstream bert distillation. a recent method, ghost-bert @xcite , employed softmaxnormalized 1d convolutions as ghost modules to generate more features with cheap operations. although compression techniques in general are effective, they come with a trade-off on accuracy, and may lose performance in case of high ratio compression. in addition, an individual fixed-size model is required for each possible computational budget. as stated in the introduction, the alternative solution is dynamic inference, which can be achieved with either early-exit or length/depthadaptive models. one of the first temporal earlyexit strategies was proposed by reasonet @xcite , which stops its reading procedure when sufficient evidence has been found for answering a question. similarly, in @xcite , an early stopping method applicable to classification tasks was presented. deebert @xcite also proposed an instance-wise multi-exit method via the entropy of the output probability distribution to speed-up bert inference. as a length-adaptive method, @xcite introduced a dynamic inference framework with one-shot training of transformers for both sequence-and token-level classification. also, in @xcite , an architecture named dyn-abert was proposed for adaptively adjusting the computations by choosing sub-networks of different widths and depths. both length-adaptive and dynabert utilized knowledge distillation and data augmentation to improve their performance. although early-exit and adaptive methods have made significant progress and work well in practice, they often require architectural manipulation and re-training. in addition, they are only applicable to encoder-only backbones and classification tasks. in contrast, our method can work with out-of-the-box pre-trained models without a need for re-training and are also applicable for encoder-decoder structures and sequence-to-sequence tasks. ## proposed method we propose a new energy-based joint inference method called e-lang, where a large/accurate language model (super) is jointly employed with a small/fast one (swift) to achieve efficient inference without sacrificing the accuracy. to this end, inspired by the method in @xcite , a routing mechanism empowered by energy-based models (ebm) is introduced to dynamically distribute the input samples between the super and swift models. similar to the out-of-distribution (ood) detection problem, our goal is to identify the ood samples that are hard to handle for the swift and forward them to the super model. on the other hand, we have the in-distribution data for which the swift can make highly reliable and accurate predictions. in other words, the routing mechanism needs to detect whether or not the input data fits in the swift's distribution (i.e., the one the swift has been trained with). inspired by the success of ebms in dealing with ood detection problems @xcite , the energy characteristics of data samples for an efficient and effective routing are investigated in our work. the overall framework of e-lang is shown in figure 1 . ## experimental results in this section, the performance of e-lang on different architectures such as t5 and bert; and benchmarks such as glue @xcite , superglue @xcite , and wmt @xcite is evaluated and compared with the super models and previous works. ## conclusion in this paper, we introduced e-lang, an energybased joint inference approach, which integrates super and swift language models for achieving efficient inference without sacrificing the accuracy. our method can work with both encoder-only (e.g., bert) and encoder-decoder (e.g., t5) architectures, and is also applicable for text classification and sequence-to-sequence problems. the proposed joint inference strategy was theoretically and experimentally analyzed with an extensive set of experiments and ablation studies. our results showed that e-lang outperforms sota in both fixed-size and dynamic inference over different benchmarks such as glue and superglue. one future direction to this work is to apply e-lang to multiple super and swift models with different sizes.