COST - EFF : Collaborative Optimization of Spatial and Temporal Efficiency with Slenderized Multi-exit Language Models (2022)

Abstract: Transformer-based pre-trained language models (PLMs) mostly suffer from excessive overhead despite their advanced capacity. For resource-constrained devices, there is an urgent need for a spatially and temporally efficient model which retains the major capacity of PLMs. However, existing statically compressed models are unaware of the diverse complexities between input instances, potentially resulting in redundancy and inadequacy for simple and complex inputs. Also, miniature models with early exiting encounter challenges in the trade-off between making predictions and serving the deeper layers. Motivated by such considerations, we propose a collaborative optimization for PLMs that integrates static model compression and dynamic inference acceleration. Specifically, the PLM is slenderized in width while the depth remains intact, complementing layer-wise early exiting to speed up inference dynamically. To address the trade-off of early exiting, we propose a joint training approach that calibrates slenderization and preserves contributive structures to each exit instead of only the final layer. Experiments are conducted on GLUE benchmark and the results verify the Pareto optimality of our approach at high compression and acceleration rate with 1/8 parameters and 1/19 FLOPs of BERT.

## introduction pre-training generalized language models and finetuning them on specific downstream tasks has become the dominant paradigm in natural language processing (nlp) since the advent of transformers @xcite and bert @xcite . however, pre-trained language models (plms) are predominantly designed to be vast in the pursuit of model capacity and generalization. with such a concern, the model storage and inference time of plms are usually high, making them challenging to be deployed on resourceconstrained devices @xcite . recent studies indicate that transformer-based plms bear redundancy spatially and temporally which comes from the excessive width and depth of the model @xcite . with static compression methods including network pruning @xcite and knowledge distillation @xcite , spatial overheads of plms (i.e., model parameters) can be reduced to a fixed setting. from the perspective of input instances rather than the model, early exiting without passing all the model layers enables the dynamic acceleration at inference time and diminishes the temporal overheads @xcite . however, static compression can hardly find an optimal setting that is both efficient on simple input instances and accurate on complex ones, while early exiting cannot diminish the redundancy in model width and is impotent for reducing the actual volume of model. further, interpretability studies indicate that the attention and semantic features across layers are different in bert @xcite . hence, deriving a multi-exit model from a pre-trained single-exit model like bert incurs inconsistency in the training objective, where each layer is simultaneously making predictions and serving the deeper layers @xcite . empirically, we find that the uncompressed bert is not severely influenced by such inconsistency, whereas small capacity models are not capable of balancing shallow and deep layers. plugging in exits after compression will lead to severe performance degradation, which hinders the complementation of the two optimizations. to fully exploit the efficiency of plms and mitigate the above-mentioned issues, we design a slenderized multi-exit model and propose a collaborative optimization approach of spatial and temporal efficiency (cost-eff) as depicted in figure 1. unlike previous works, e.g., dynabert @xcite and cofi @xcite , which obtain a squat model, we keep the depth intact while slenderizing the plm. superiority of slender architectures over squat ones is supported by @xcite and @xcite in generic machine learning and plm design. to address the inconsistency in compressed multi-exit model, we first distill an multi-exit bert from the original plm as both the teaching assistant (ta) and the slenderization backbone, which is more effective in balancing the trade-off between layers than compressed models. then, we propose a collaborative approach that slenderizes the backbone with the calibration of exits. such a slenderization diminishes less contributive structures to each exit as well as the redundancy in width. after the slenderization, task-specific knowledge distillation is conducted with the objectives of hidden representations and predictions of each layer as recovery. specifically, the contributions of this paper are as follows. the compression and acceleration of plms were recently investigated to neutralize the overhead of large models by various means. the structured pruning objects include, from small to large, hidden dimensions @xcite , attention heads @xcite , multihead attention (mha) and feed-forward network (ffn) modules @xcite and entire transformer layers @xcite . considering the benefit of the overall structure, we keep all the modules while reducing their sizes. besides pruning out structures, a fine-grained approach is unstructured pruning which prunes out weights. unstructured pruning can achieve high sparsity of 97% @xcite but is not yet adaptable to general computing platforms and hardware. during the recovery training of compressed models, knowledge distillation objectives include predictions of classifiers @xcite , features of intermediate representations @xcite and relations between samples @xcite . also, the occasion of distillation varies from general pre-training and task-specific fine-tuning @xcite . distillation enables the training without ground-truth labels complementing data augmentation. in this paper, data augmentation is not leveraged as it requires a long training time, but our approach is well adaptable to it if better performance is to be pursued. dynamic early exits come from branchynet @xcite , which introduces exit branches after specific convolution layers of the cnn model. the idea is adopted to plms as transformer layer-wise early exiting @xcite @xcite . however, early exiting only accelerates inference but does not reduce the model size and the redundancy in width. furthermore, owing to the inconsistency between shallow and deep layers, it is hard to achieve high speedup using early exiting alone. the prevailing plms, e.g., roberta @xcite and xlnet @xcite are variants of transformer with similar overall structures, well-adaptable to the optimizations that we propose. apart from plms with increasing size, albert @xcite ) is distinctive with a small volume of 18m (million) parameters obtained from weight sharing of transformer layers. weight sharing allows the model to store the parameters only once, greatly reducing the storage overhead. however, the shared weights have no contribution to inference speedup. instead, the time required for albert to achieve bert-like accuracy increases. ## methodology in this section, we analyze the major structures of transformer-based plms and devise corresponding optimizations. the proposed cost-eff has three key properties, namely static slenderization, dynamic acceleration and collaborative training. ## conclusion in this paper, we statically slenderize and dynamically accelerate plms in the pursuit of inference efficiency as well as preserving the capacity. to integrate the perspectives, we propose a collaborative optimization approach that achieves a mutual gain of static slenderization and dynamic acceleration. specifically, the size of plm is reduced in model width, and the inference is adaptable to the complexity of inputs without introducing redundancy for simple inputs and inadequacy for hard inputs. comparative experiments are conducted on glue benchmark and verify the pareto optimality of our approach at high compression and acceleration rate.