Compressing Transformer-Based Sequence to Sequence Models With Pre-trained Autoencoders for Text Summarization

1 INTRODUCTION . It is safe to say that the combination of Transformer ( Vaswani et al. , 2017 ) architecture , and transfer learning concept dramatically modified the landscape of Natural Language Processing ( NLP ) . Introduction of large-scale pre-trained language models like BERT ( Devlin et al. , 2018 ) , GPT-2 ( Radford et al. , 2019 ) , MegatronLM ( Shoeybi et al. , 2019 ) , BART ( Lewis et al. , 2019 ) , and GPT-3 ( Brown et al. , 2020 ) keeps on improving state-of-the-art results by fine-tuning them for downstream tasks such as Sentiment Analysis , Question Answering , and Summarization . However , the upward trend of the network size in mentioned models raises serious environmental ( Strubell et al. , 2019 ) and usability issues . The number of parameters in Transformers-based models is constantly rising . As a result , the hardware requirement for fine-tuning or inference drastically increased in the past couple of years . It is challenging for both researchers and developers to use these models and build on top of them without sufficient resources . To put it in perspective , the earliest and smallest BERT had 110M parameters ( Devlin et al. , 2018 ) , and the latest and largest Switch Transformer ( Fedus et al. , 2021 ) model was introduced with 1.5 trillion parameters which makes it accessible only using high-end servers . This issue is even more consequential for tasks such as automatic text summarization and machine translation that incorporate sequence-to-sequence architecture . This architecture consists of an encoder ( encoding the input sequences ) paired with a decoder ( generates tokens conditioned on the encoder ’ s representation ) . It means even more parameters will be added to the model for these tasks which is the min focus of this paper . It is important to discover techniques to reduce the overall network parameters while maintaining the quality of the generated text . In this paper , autoencoders ’ ( AE ) ( Liu et al. , 2019 ) property of dimensionality reduction is evaluated in a setting with sequence-to-sequence architecture and pretrained encoders . The autoencoder will act as an intermediate model to compress the encoder ’ s final representation and decoder will use this compressed latent representation to generate summaries with minimal information loss . The idea is to find the ideal trade-off between the compression ratio and model ’ s text generation capability . 2 BACKGROUND . Multiple different approaches have already been explored to tackle the problem of neural network models growing size . Quantization ( Gupta et al. , 2015 ) is one of the first experiments which apply to any deep learning model by using a half-precision ( 16-bit ) floating point to greatly reduce the network size and memory usage . Micikevicius et al . proposed a mixed precision algorithm in ( Micikevicius et al. , 2017 ) to further close the gap in evaluation results . Recent works experimented the effect of knowledge distillation ( Bucila et al. , 2006 ) method to transfer information from a larger network to a smaller one without significant loss in accuracy . There are multiple papers that present different combinations of fine-tuning and distillation on top of BERT . ( Chatterjee , 2019 ; Turc et al. , 2019 ) However , DistilBERT ( Sanh et al. , 2019 ) obtained the best results with training the smaller student model on BERT and then fine-tuning it for downstream tasks that resulted in a more generalized pre-trained model . Their approach led to similar implementations on other classic transformer models such as DistilGPT2 with 33 % less parameters ( two times faster ) which resulted in 21.1 perplexity score comparing to the GPT-2 ’ s 16.3 , and a 35 % smaller RoBERTa ( Liu et al. , 2019 ) model while maintaining 95 % of the accuracy on GLUE , named DistilRoBERTa.1 The pruning ( LeCun et al. , 1990 ) method ’ s influence on transfer learning have recently gained attention from researchers . It refers to determining the parts in the network that have the weaker effect on the model accuracy and removing them without compromising the model ’ s accuracy on downstream tasks . The main ideas are to either focus on finding the less important weights ( Gordon et al. , 2020 ) , or components such as number of self-attention heads ( Michel et al. , 2019 ) and layers ( Sajjad et al. , 2020 ) . This technique was also used in the Lottery Ticket hypothesis ( Frankle & Carbin , 2018 ; Prasanna et al. , 2020 ) to uncover subnetworks performing on par with the full model . The latest research area focuses on rethinking the self-attention mechanism to eliminate its quadratic memory usage connection with respect to the input sequence length . The goal is to find the best trade-off between performance and memory usage . Big Bird ( Zaheer et al. , 2020 ) and Longformer ( Beltagy et al. , 2020 ) papers experiment on different attention patterns to reduce connections and result in fewer computations . Wang et al . presented the Linformer ( Wang et al. , 2020 ) network that projects the self-attention vectors to lower dimensions . Reformer ( Kitaev et al. , 2020 ) paper studies the idea of grouping key and query vectors based on the locality sensitive hashing ( LSH ) to reduce the computations needed to find similar vectors . It is worth noting several studies that combined two or more methods to build even smaller models without significantly compromising the accuracy . Several such experiments are present in the literature , namely , a combination of distillation with pruning ( Hou et al. , 2020 ) , or quantization ( Sun et al. , 2020 ) . Also , Tabmbe et al . made the EdgeBERT ( Tambe et al. , 2020 ) model by leveraging both pruning and quantization along with other methods . These techniques are not exclusive , and it is possible to study them independently without testing all possible combinations . This is why we can focus our study only on one reduction method and only a few important models . 3 PROPOSED METHOD . Our proposed architecture is a sequence-to-sequence Transformer ( T ) model that includes a pretrained autoencoder ( AE ) connecting the network ’ s encoder ( Tenc ) to its decoder ( Tdec ) . ( Fig . 1 ) The mentioned approach will result in a smaller Tdec and reduce the overall number of trainable parameters . The modules are described in the sub-sections . 3.1 TRANSFORMER ENCODER ( Tenc ) The encoder ( Tenc ) part of the proposed architecture is a pre-trained transformer-based model . We have selected a set of models that include some of the best models for text summarization ; BERT , 1The results related to DistilRoBERTa model is available at https : //github.com/huggingface/ transformers/tree/master/examples/research_projects/distillation DistilBERT , BART ’ s encoder ( all base versions ) , and a custom transformer model . The custom transformer model is used to evaluate the effectiveness of our approach on a small pre-trained model with only 6 encoder layers and subsequently its scores are used as baseline and not supposed to be competitive with the ones of the other approaches . These models are frozen during the training processes to reduce the number of influential factors . 3.2 AUTOENCODER ( AE ) . The AE ( Fig . 1 ) purpose is to reduce the Tenc ’ s output size to a smaller latent space using the following equation : X ′AE = AEdec ( AEenc ( XAE ) ) Z = AEenc ( XAE ) Where XS×DAE is the input , AEenc indicates the encoder responsible for compressing the input to latent space ZS×C and a decoder AEdec generating the output X ′ S×D AE trying to reconstruct the input XS×DAE during the training process . Variables S and D denote the sequence length and input ’ s embedding dimension respectively . The values are dependent to the chosen pre-trained Tenc model ’ s configuration and are set to 512 and 768 in this paper . However , C that represents the compressed latent space size , will vary to find out the optimal latent space size . A 6-layer linear AE ( 3 for encoder and 3 for decoder ) was selected after comparing its reconstruction ability to the same architecture with Long Short-Term Memory ( LSTM ) ( Hochreiter & Schmidhuber , 1997 ) , or Convolutional Neural Network ( CNN ) ( LeCun et al. , 1998 ) building blocks . The final AE architecture with 6 linear layers was independently trained for each selected pre-trained encoder Tenc . It attempts to reconstruct the output of Tenc using a smaller representation Z . The frozen AEenc is then used in our summarizer architectures to pass a compressed representation to the decoder Tdec . ( From size D to C ) Refer to appendix A.1 for more information about the hidden layers ’ sizes . 3.3 TRANSFORMER DECODER ( Tdec ) The decoder component of the architecture ( Tenc ) is an original transformer decoder with 3-layers and a linear head on top in all the experiments . It is the only piece of the network that is not frozen after the AE has been trained . Its embedding dimension ties to the AE ’ s latent space size ( C ) that can drastically alter the architecture overall number of trainable parameters . 3.4 DATASETS . A combination of the CNN/Dailymail ( 300K samples ) ( Hermann et al. , 2015 ) and the Newsroom ( 1.3M samples ) ( Grusky et al. , 2018 ) datasets used for training all the summarizer models ; and the pre-defined test set of CNN/Dailymail dataset utilized to evaluate them . However , we randomly selected only 60 % of these combinations to train each autoencoder models to hold on unseen data to evaluate the generalization ability of the summarizer model . 3.5 EXPERIMENTS . We performed several experiments to evaluate the effectiveness of the proposed method . First ( AE ) , we combined different pre-trained encoders with several autoencoder compressed latent space sizes paired with a 3-layers decoder component in each instance . Second ( AE-S ) , we used the same autoencoders ( without the pre-training step ) and trained them jointly with the decoder from scratch to measure the effect of the pre-training step on the autoencoder . Third ( LL ) , using a small 1-layer learnable linear model to lower the encoder ’ s output dimensionality from D to C. Lastly ( PCA ) , we used the classical dimensionality reduction algorithm , incremental PCA ( Ross et al. , 2008 ) trained to project the outputs of the encoder to the 458 first principal components to preserve more than 90 % of variances and used them as the decoder input . All these models are trained with fixed hyperparameters to make them comparable . They are trained using the 1cycle ( Smith & Topin , 2019 ) training policy that increases the learning rate from 2e-5 to a maximum of 5e-4 while decreasing the momentum beta from 0.95 to 0.85 for faster convergence . We also used the Adam optimizer and the label smoothing CrossEntropy ( Pereyra et al. , 2017 ) loss function . The decoder ’ s feedforward layers and the eight attention heads ’ inner dimensions are 1024 and 32 , respectively , and their input sizes are dependent on the dimensionality reduction rate . 4 RESULTS The summaries are generated using three different decoding strategies . The fastest and easiest method to develop is picking the most probable output at each timestep , known as the Greedy algorithm . A more extensive approach is to use the Beam Search algorithm that develops K paths in parallel using the top K tokens with the highest scores . Finally , we used the Weighted Random Sampling algorithm that randomly chose a token from the top K probable outputs proportionally to their respective probabilities . In this paper , we used K value equals 5 and 2 for the beam search and weighted random sampling , respectively , based on the previous experiments . We experimented with different AE layer types and numbers . Linear , Long Short-Term Memory ( LSTM ) , and Convolutional Neural Network ( CNN ) AE architectures were trained with the Mean Square Error ( MSE ) loss for seven epochs to measure their performance . A latent space representation of size 64 was selected for this comparison benchmark ( Table 1 ) . The linear autoencoder outperforms both LSTM and CNN in all experiments . ( Refer to Appendix A.1 , Table 6 for the full list of all compression rates ) Also , the 6-layer design choice results in a better score in all experiments . The model size comparisons ( number of parameters ) are presented in Table 2 . The first step is to calculate the decoder ’ s size without using the dimensionality reduction method . As previously mentioned , based on the choice of the pre-trained encoder models , the default encoder ’ s representation dimension is 768 which results in a decoder with either 48M ( for BART encoder ) or 33M ( other options ) parameters . We use this value as the reference number to measure the reduction percentage . It is important to keep in mind that even though there are parameters being added to the model from the autoencoder that affects the percentage , the number of trainable parameters in all the experiments is equal to the decoder size since the rest of the components ( encoder , and autoencoder ) are frozen during the tests . The reason why we have decoders with different parameter counts while using the same hyper-parameters is because we use each model ’ s pre-trained tokenizers which have different vocabulary size , which affects the decoder ’ s last layer output dimension . BERT , DistilBERT , and Transformer models use the BERT ’ s pre-trained tokenizer , and BART uses its own . The proposed method results on the text summarization task are shown in Table 3 where the classic ROUGE ( R-1 , R-2 , R-3 , R-L , and R-W ) metric is reported using the greedy , weighted random sampling , and beam search decoding strategies . Autoencoders are trained with four ( 32 , 128 , 384 , 512 ) different latent space dimensions ( C ) to analyze the effect of these dimensionality reductions . The models are also evaluated using the BERTScore ( Zhang et al. , 2019 ) metric and the results follow the same direction . ( Refer to Appendix A.4 ) Our configurations not only reduce the number of parameters in the network by reducing the decoder size , but also shows the ability to increase the model ’ s summarization ability with a higher score compared to the original setup for the Transformer and DistilBERT models . The experiments show that adding an AE with C = 512 outperforms the same vanilla encoder-decoder network . Both BART and BERT experiments with the same AE size outperformed their vanilla model in several metrics only using the beam search method ; however , the rest of the scores with a smaller decoder are still competitive . An even more surprising result is that a smaller dedicated Transformer with the proposed method and C = 384 performed better than BART and BERT in all the summarization benchmarks . It also worth noting that the greedy inference method constantly results in better scores , with the weighted random sampling method following closely . The fact that beam search algorithm leans towards shorter sequences ( Wu et al. , 2016 ) reduces the ROUGE scores since there are fewer matching N-grams in the generated and target summaries . It does not mean that the sentences structure/quality are flawless using greedy/weighted random sampling , or poor using beam search , the results just reflect what the ROUGE score is measuring : an N-gram overlap between the generated and the target sequences . Our ablation results show ( Table 4 ) the impact of autoencoder pre-training step on the final score . It surpasses both training the AEenc jointly with the network from scratch ( AE S ) and using a simple learnable linear layer for projecting ( LL ) . Lastly , the PCA dimensionality reduction technique does not produce desirable results . The results also show that using an autoencoder with a latent space size of 384 ( C = 384 ) generates ROUGE score close to the vanilla model . It reduces the decoder size by 46 % from 48M to 24M for the BART model and 44 % from 33M to 16M parameters for the other models . The critical point is that the combination of this configuration associated with greedy decoding algorithm shows no noticeable degrading quality in the generated summary . ( Table 5 )

Summary: The paper proposes a new autoencoder-based seq2seq model for text summarization tasks. The paper tries to find the best trade-off between compression ratio and model performance. The paper conducts extensive experiments by evaluating the loss of accuracy with ROUGE.