Modular Transformers: Compressing Transformers into Modularized Layers for Flexible Efficient Inference (2023)

Abstract: Pre-trained Transformer models like T5 and BART have advanced the state of the art on a wide range of text generation tasks. Compressing these models into smaller ones has become critically important for practical use. Common neural network compression techniques such as knowledge distillation or quantization are limited to static compression where the compression ratio is fixed. In this paper, we introduce Modular Transformers, a modularized encoder-decoder framework for flexible sequence-to-sequence model compression. Modular Transformers trains modularized layers that have the same function of two or more consecutive layers in the original model via module replacing and knowledge distillation. After training, the modularized layers can be flexibly assembled into sequence-to-sequence models that meet different performance-efficiency trade-offs. Experimental results show that after a single training phase, by simply varying the assemble strategy, Modular Transformers can achieve flexible compression ratios from 1.1x to 6x with little to moderate relative performance drop.

## introduction the ever increasing size of pre-trained sequence-tosequence (seq2seq) models @xcite @xcite @xcite has supported advances in the state of the art in a wide range of natural language processing (nlp) tasks. for example, bart @xcite has 400 million parameters while t5 @xcite pushes this number to 11 billion. this makes large pretrained seq2seq models hard to deploy and prone to negative environmental impacts @xcite @xcite , motivating researchers to investigate methods to compress large pre-trained models into smaller, faster ones that retain strong performance. previous work * work done while interning at the allen institute for ai has shown that bert @xcite , a popular encoder-only pre-trained transformer @xcite , can be effectively compressed and accelerated via different neural network compression techniques @xcite @xcite @xcite . however, there is limited research on methods to compress pre-trained seq2seq models @xcite . on the other hand, pre-trained seq2seq models are generally more space and timeconsuming compared to their encoder-only counterparts since they require storing the decoder as well and generally decode in an auto-regressive fashion. to satisfy ever-changing resource constraints varying in different applications and over time, existing seq2seq compression techniques must separately train and store many compact models with different compression ratios, which is computationally inefficient and may also violate space constraints. moreover, as suggested by @xcite , as opposed to encoder-only models, it is nontrivial to find proper sizes for seq2seq models with a given resource constraint as the depth for the encoder and the decoder must be jointly tuned. as such, searching for compact seq2seq models meeting different resource constraints can be very costly. this motivates us to investigate the problem of flexible seq2seq compression that can dynamically adjust compression ratio to meet varying resource constraints without training multiple compact models. in this work, we present modular transformers, a framework to compress transformers into modularized layers for flexible efficient inference. with the guidance from the original model, modular transformers trains modularized layers that have the same function of different numbers of consecutive layers in the original model via multigrained module replacing and knowledge distillation. specifically, we first map each of the modularized transformer layers to a sub-module (i.e., a number of consecutive layers) of the encoder or decoder of the original model. we then train the modularized layers by randomly assembling a model for each training step (replacing submodules of the original model with their corresponding modularized layers), while keeping the original model parameters frozen. we propose a curriculum-replacing strategy to train the modularized layers in a fine-to-coarse fashion. we also use attention and representation distillation to further encourage the modularized layers to behave similarly to the original sub-modules. after training, the modularized layers can be flexibly assembled into seq2seq transformers that meet different performance-efficiency trade-offs. the compression ratio for encoders and decoders can also be seamlessly adjusted to find the optimal parameterization within a fixed computation budget without any additional training. in addition, we introduce two deterministic assembling strategies tailored for smaller sizes and lower latency. both of them replace the original model following a decoder-to-encoder, top-to-bottom, and fine-tocoarse fashion. we empirically verify the effectiveness of modular transformers by compressing t5 @xcite , a popular pre-trained seq2seq model, on representative text generation tasks including text summarization, machine translation, and question generation. empirical results show that our approach consistently outperforms prior seq2seq compression techniques across different datasets while also enabling users to flexibly adjust the efficiencyperformance trade-off of the model to meet different requirements for deployment. ## related work pre-trained model compression prior work has shown that bert @xcite , a popular encoder-only pre-trained transformer @xcite , can be effectively compressed and accelerated. as summarized in @xcite and @xcite , popular bert compression techniques include knowledge distillation @xcite @xcite @xcite which trains a compact student network to mimic the behavior of the original teacher model, pruning @xcite @xcite @xcite which prunes redundant neurons or structures in the original model, module replacing @xcite which trains compact successor sub-modules to replace that in the original model, and quantization @xcite that compresses a neural network by reducing the number of bits used to represent its parameters. in addition, a number of work also investigated efficient inference with bert-like models with dynamic inference methods including early exit @xcite @xcite @xcite or adaptive computation time @xcite . these approaches only reduce inference latency while keep the model size unchanged. modular transformers focuses on the first category which reduces the size of the model while also accelerating inference. sequence-to-sequence compression compared to conventional neural network compression, seq2seq compression is relatively less explored. the de-facto method for compressing seq2seq model is sequence-level knowledge distillation @xcite , which uses the original teacher model to generate pseudo-labels with beam search on the train set and train the compact student model with the input-pseuodo label pairs. a number of recent studies (junczys-dowmunt, 2019; @xcite have verified the effectiveness of the pseudo labeling method for distilling pre-trained transformer-based seq2seq models. however, @xcite shows that a simple shrink and fine-tune method performs competitively with pseudo labeling while requiring fewer computations. another work related to modular transformers is layerdrop @xcite , which randomly dropout transformer layers when pre-training a seq2seq model and then prunes certain layers during inference for improved efficiency. layerdrop differs from our method for two main reasons. first, layerdrop must be applied during pre-training while modular transformers can be applied to any pre-trained models. this makes the scope of application of our method much broader because most existing models are not pre-trained with layerdrop. second, during inference, layer-drop prunes layers while modular transformers replaces sub-networks with compact modules with similar functionality. this hopefully reduces the performance drop, especially when the target compression ratio is relatively large. ## methodology in this section, we describe the proposed modular transformers framework in detail. we first recap module replacing in §3.1. we then describe the architecture design and training method for modular transformers in §3.2 and §3.3, respectively. then, we present the idea of dynamic assembling in §3.4. finally, we discuss the relationship between modular transformers and vanilla module replacing in §3.5. ## experiments in this section, we empirically verify the effectiveness of modular transformers by using it to compress t5 @xcite , a popular pre-trained seq2seq transformer, on a number of representative natural language generation tasks and datasets. ## conclusion we introduce modular transformers, a framework for flexible and accurate seq2seq compression that adapts to different performance-efficiency tradeoffs, which has important practical implications such as reducing the computation cost and environmental impact of deployed systems. our method defines a set of modularized layers with different granularity and trains them to share the same functionality as the sub-networks of the original model they replace. we combine multi-grained module replacing and knowledge distillation, and design two flexible assembling strategies for size-first and speed-first inference. empirical results on various text generation tasks show that our approach consistently outperforms previous methods while alleviating the need to re-train the compact model in order to adapt to new memory and inference time constraints.