Efficient Latent Space Compression for Lightning-Fast Fine-Tuning and Inference of Transformer-Based Models ^†

Abstract

This paper presents a technique to reduce the number of parameters in a transformer-based encoder–decoder architecture by incorporating autoencoders. To discover the optimal compression, we trained different autoencoders on the embedding space (encoder’s output) of several pre-trained models. The experiments reveal that reducing the embedding size has the potential to dramatically decrease the GPU memory usage while speeding up the inference process. The proposed architecture was included in the BART model and tested for summarization, translation, and classification tasks. The summarization results show that a 60% decoder size reduction (from 96 M to 40 M parameters) will make the inference twice as fast and use less than half of GPU memory during fine-tuning process with only a 4.5% drop in R-1 score. The same trend is visible for translation and partially for classification tasks. Our approach reduces the GPU memory usage and processing time of large-scale sequence-to-sequence models for fine-tuning and inference. The implementation and checkpoints are available on GitHub.

1. Introduction

This paper is an extended version of the work published in [1] to explore the proposed architecture in more detail to further study its strength and shortcomings. The primary focus of this paper is to create a more efficient fine-tuning paradigm for pre-trained transformer-based [2] models. It is no secret that the current landscape in natural language processing (NLP) portrays a direct correlation between the model’s scale and capabilities. Admittedly, the trend is more visible in decoder-only large language models like GPT-family [3,4,5] and LLaMA [6]. However, other task-specific models can also experience a boost in performance with a higher number of parameters. This tendency has serious environmental [7] and usability impacts.

Exploring the new possibilities to enhance the accessibility of large-scale models is a crucial research area. The performance of an architecture, no matter how capable it may be, becomes irrelevant if researchers cannot make further improvements or practitioners cannot build upon it. Even the process of loading some of the large-scale models, like the encoder-only switch transformer [8], necessitates a substantial amount of engineering effort. The impact of this trend on the sequence-to-sequence (seq2seq) architecture is even more pronounced due to the presence of both an encoder and a decoder component. Without adequate resources, it is challenging to effectively utilize well-known seq2seq models such as BART [9] and Pegasus [10], which consist of 400 M and 570 M parameters, respectively. It is worth noting that these numbers are considered relatively small compared to the current standards.

This research examines how the dimensionality reduction capability of autoencoders [11] can be used to compress the latent space of an encoder–decoder architecture. Reducing the size of the embedding space leads to a smaller decoder, as it reduces the input dimension of the decoder. Both the fine-tuning and inference processes are expected to benefit from this approach. However, we anticipate a significant improvement during inference, as the decoder plays a crucial role in this stage. The main contributions of this paper are as follows:

• The best algorithm and architecture to perform an effective embedding space compression were found.
• An extensive number of experiments were conducted to find the perfect balance between the autoencoder’s compression rate and the model’s capability.
• The proposed architecture underwent testing across various datasets and was evaluated for summarization, translation, and classification tasks.
• The autoencoder’s capacity to effectively generalize unseen datasets and diverse tasks was explored.

The rest of the paper is organized as follows. In the next section, we provide an overview of the prior literature on the topic of model compression. Section 3 will delve into a comprehensive explanation of the proposed architecture and its individual components, offering a more detailed perspective. Next, we discuss the experiments that were carried out to both validate the approach and assess its effectiveness across various tasks. Subsequently, we present the outcomes of the experiments, concluding with a comprehensive analysis of the findings. You can find the implementation details and pre-trained weights in the provided GitHub repository: https://github.com/AlaFalaki/compressedBART (accessed on 13 July 2023).

2. Background

The transformer architecture has established its dominance across various NLP tasks. The BART [9] model reigned supreme on the summarization task leaderboard for a long time and demonstrated proficiency in translation as well. A number of studies extend it for summarization like the current state-of-the-art approach, BRIO [12], which generates a large pool of summaries and select the best candidate, or the HydraSum [13] architecture, with multiple decoders for capturing different styles.

By pre-training on a multilingual dataset, the mBART [14] model has achieved an enhanced performance in machine translation task. Various approaches have been investigated to enhance translation quality, including the utilization of back translation for monolingual models [15] and parameter sharing across layers to reduce the number of parameters and computational time [16].

For several classification tasks, the pre-trained XLNet [17] model demonstrated superior performance compared to BERT. LadaBERT [18] presents a novel approach that combines various dimensionality reduction techniques, including low rank decomposition and pruning. This approach leads to a compact and highly efficient model that can seamlessly integrate into diverse applications. The ERNIE-Doc [19] model introduces a recursive architecture that offers an improved efficiency in handling large sequences. It surpasses competing models in multiple tasks, including classification.

Additional researchers are actively investigating the pre-training of large models to achieve high performance across multiple downstream tasks. The T5 [20] architecture is one of the well-known examples that resulted in in a pre-trained model to handle translation, summarization, question answering, and classification with high accuracy. These advancements are scaling up the parameter count and slowing the inference process.

Multiple research studies have already been conducted to tackle the growing size of neural network models. Quantization [21] 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 [22] to further close the gap in evaluation results. Recent works experimented with the effect of the knowledge distillation [23] method to transfer information from a larger network to a smaller one without a significant loss in accuracy. There are multiple papers that present different combinations of fine-tuning and distillation on top of BERT [24,25]. However, DistilBERT [26] obtained the best results by training the smaller student model on BERT and then fine-tuning it for downstream tasks, resulting in a more generalized pre-trained model. Lastly, the DistilBART [27] model aimed to create a smaller student network for summarization.

The pruning [28] method’s influence on transfer learning has 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 [29] or components such as the number of self-attention heads [30] and layers [31]. This technique was also used in the lottery ticket hypothesis [32,33] to uncover subnetworks performing on par with the full model. This method focusing on linear layers is adopted for the seq2seq architectures like BART and Pegasus to speed up the inference process [34].

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 [35] and Longformer [36] papers experiment on different attention patterns to reduce connections and result in fewer computations. Wang et al. presented the Linformer [37] network that projects the self-attention vectors to lower dimensions. Reformer [38] paper studies the idea of grouping key and query vectors based on locality-sensitive hashing (LSH) to reduce the computations needed to find similar vectors.

It is important to highlight several studies that have effectively combined two or more methods to construct smaller models, while maintaining a high level of accuracy without significant compromise. The literature contains numerous experiments that explore the combination of distillation with pruning [39] or quantization [40], among others. Furthermore, Thierry et al. made the EdgeBERT [41] 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 uses a pre-trained autoencoder connecting the network’s encoder ($T_{enc}$) to its decoder ($T_{dec}$) (Figure 1). The mentioned approach will result in a smaller $T_{dec}$ and reduce the overall number of trainable parameters. The modules are described in the following sub-sections.

3.1. Transformer Encoder ($T_{\mathit{enc}}$)

The encoder ($T_{enc}$) 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, DistilBERT, ${\mathrm{BART}}_{\mathrm{enc}}$ (only encoder), and a custom transformer model as a baseline comparison. The custom transformer model is used to evaluate the effectiveness of our approach on a small pre-trained model with only six encoder layers. Subsequently, its scores are used as a baseline and are not supposed to be competitive with the ones of the other approaches. This component is frozen during the training of the summarizer model.

3.2. Autoencoder ($\mathit{AE}$)

The AE (Figure 2) purpose is to reduce the $T_{enc}$’s output size to a smaller latent space using the following equation:

$$\begin{gathered} {X^{\prime }}_{AE}=A{E}_{dec}\left(A{E}_{enc}\left(X_{AE}\right)\right) \\ Z= A{E}_{enc}\left(X_{AE}\right) \end{gathered}$$

(1)

where $X_{AE}^{S\times D}$ is the input, $A{E}_{enc}$ indicates the encoder responsible for compressing the input to latent space $Z^{S\times C}$ and a decoder $A{E}_{dec}$ generating the output $X{\prime }_{AE}^{S\times D}$, trying to reconstruct the input $X_{AE}^{S\times D}$ during the training process. Variables S and D denote the sequence length and input’s embedding dimension, respectively. The values depend on the chosen pre-trained model’s configuration and are set to 512 and 768 in the experiments. However, $C$, representing the compressed latent space size, will vary to find out the optimal latent space size. A six-layer linear AE (three for encoder and three for decoder) is selected after comparing its reconstruction ability to the same architecture with long short-term memory (LSTM) [42] or convolutional neural network (CNN) [43] building blocks (see Section 5.1).

The final AE architecture with six linear layers is independently trained for each selected pre-trained encoder $T_{enc}$. It attempts to reconstruct the output of $T_{enc}$ using a smaller representation $Z$. The frozen $A{E}_{enc}$ is then used in our summarizer architectures to pass a compressed representation to the decoder $T_{dec}$ (from size $D$ to $C$). Refer to Appendix A for more information about the autoencoder architecture details.

3.3. Transformer Decoder ($T_{\mathit{dec}}$)

The decoder component of the architecture ($T_{enc}$) is an original transformer decoder with three 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’s overall number of trainable parameters.

4. Experiments

We performed two experiments to: (i) evaluate the effectiveness of the proposed architecture as described in Section 3 and (ii) showcase its effectiveness by integrating it with BART-base model on summarization, translation, and classification tasks.

The first study (we call it “validation”) starts by justifying the use of the linear layer autoencoder by comparing its performance to compress and reconstruct the BERT’s output with a different number of layers and building units (LSTM and CNN). Then, we investigated the advantage of using “frozen autoencoders” with several studies listed below:

• AE (presented architecture): A pre-trained and frozen autoencoder that connects the encoder to a 3-layer decoder.
• AE-S: The same architecture without pre-training the autoencoders. They will be trained jointly with the decoder from scratch.
• LL: A small 1-layer learnable linear model to lower the encoder’s output dimensionality from $D$ to $C$.
• PCA: The classical dimensionality reduction algorithm, incremental PCA [44], trained to project the outputs of the encoder to the 458 first principal components to preserve more than 90% of variances and use them as the decoder input.

The summaries are generated using three different decoding strategies. The fastest and easiest method is selecting 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 chooses a token from the top K probable outputs proportionally to their respective probabilities.

Following the “Validation” phase, we thoroughly explore the potential of the proposed architecture by evaluating it on different tasks. The compression rates are slightly different in these experiments because the compression sizes (latent space dimension) must be dividable by the number of the decoder’s attention heads (12).

Evaluation: The summarization and translation experiments use ROUGE [45] and BLEU [46] scores to measure the model’s generation quality. The ROUGE score is the standard metric to calculate how close a generated summary is, compared to the target. It will measure the overlapping n-grams between the two (R-1 for unigrams, R-2 for bigrams, etc.) and the longest common subsequence n-grams (R-L). The classification task is measured using the accuracy metric. The BLEU score compares the machine-generated translation with one or more reference translations. It counts how many words and phrases from the machine translation are also present in the references and assigns a score based on this overlap. The sentient analysis study will use the BART architecture with a classification head.

Hardware and Hyperparameters: The effects of the proposed architecture will be analyzed in subsections aimed at the GPU memory consumption and fine-tuning/inference computational time. The experiments aim for maximum performance using Nvidia V100 GPU with 32 GB memory. Therefore, each task uses the largest possible batch size with respect to hardware constraints. Moreover, the datasets are vastly varied in size and characteristics. Thus, it is not possible to compare the tasks against each other.

Lastly, the autoencoder models are trained with fixed hyperparameters to make them comparable. They are trained using the 1cycle [47] training policy that increases the learning rate from 2 × 10⁻⁵ to a maximum of 5 × 10⁻⁴, while decreasing the momentum beta from 0.95 to 0.85 for a faster convergence. We also used the Adam [48] optimizer and the label smoothing CrossEntropy [49] loss function for both autoencoder and summarizer models.

Datasets: The CNN/Dailymail (300 K samples) [50] and the Newsroom (1.3 M samples) [51] datasets were used in the following experiments. We combined the training set of mentioned datasets and randomly selected 60% of the samples to train the autoencoders. This was performed to hold onto unseen data to evaluate the generalization ability of the summarizer model. The summarizer models are trained using the CNN/Dailymail training set and evaluated on the pre-defined test set.

The translation task utilizes the WMT16 dataset’s Romanian-to-English subset [52]. It has 610 K samples in the training set and 2 K for each validation and test set. The training set is used to create a Romanian language tokenizer with the same vocabulary size as BART’s pre-trained tokenizer.

The Sentiment140 [53] dataset is used for the classification experiment. It is a sentiment analysis dataset consisting of twitter data with corresponding annotations (0 for negative and 4 representing positive). There are 1.6 M data points with random 80–10–10 splits for training, validation, and testing, respectively. A simple preprocessing step was performed to remove unnecessary characters (e.g., @ signs, URLs, or HTML codes).

5. Results

In the upcoming subsections, we will initially conduct a comprehensive examination of the architecture’s viability and subsequently assess its performance in summarization, translation, and classification tasks.

5.1. Validation

To assess their performance, we conducted experiments involving various types and depths of AE layers. Specifically, we trained autoencoder architectures using linear, long short-term memory (LSTM), and convolutional neural network (CNN) layer types, with the mean square error (MSE) loss, over a span of seven epochs. A latent space representation of size 64 was selected for this comparison benchmark (

Table 1. The comparison of MSE loss between linear, LSTM, and CNN blocks to train an autoencoder with a 64 compression size.

Number of Layers Types	4	6	8
Linear	0.0813	0.0766	0.0776
LSTM	0.0863	0.0810	0.0849
CNN	0.2666	0.2659	0.2750

). The linear autoencoder demonstrates a superior performance compared to both LSTM and CNN architectures (refer to Appendix A,

Table A1. The MSE loss value of the selected 3 network types (LSTM, linear, CNN) with a different number of layers.

Type	Compression Rate	Number of Layers	MSE Loss
Linear	32	6	0.0930
	64	4	0.0752
		6	0.0766
		8	0.0775
	128	6	0.0637
	256	6	0.0453
	384	6	0.0278
	448	6	0.0239
	512	6	0.0200
LSTM	32	6	0.0905
	64	2	0.1176
		4	0.0863
		6	0.0810
		8	0.0849
		10	0.1043
	128	6	0.0670
	256	6	0.0543
	384	6	0.0462
	448	6	0.0427
	512	6	0.0400
CNN	64	4	0.2666
		6	0.2659
		8	0.2750

for the complete list of all compression rates). Also, the six-layer design choice results in a better score in all experiments.

Table 2. Comparing the decoder’s number of parameters in a 3-layer decoder network with a 768 input size to the number of parameters of the same decoder after applying the autoencoder to reduce the encoder’s output dimension.

Decoder Input Dimension	AE Parameters Count	Decoder Parameters Count (by Encoder Type)
		(Distil)BERT/ Transformer	Reduction (%)	BARTenc	Reduction (%)
768 (Default)	-	32 M	-	48 M	-
C = 512	2.3 M	21 M	26.26	32 M	28.48
C = 384	1.8 M	16 M	44.44	24 M	46.18
C = 128	1.1 M	5 M	79.84	8 M	80.91
C = 32	1 M	1 M	92.47	2 M	93.5

and

Table 3. Comparing the total number of parameters in a 3-layer decoder network with a 768 input size to the number of parameters of the same decoder after applying the autoencoder to reduce the encoder’s output dimension.

Decoder Input Dimension	AE Parameters Count	Network’s Total Parameters Count (by Encoder Type)
		Transformer	BART	BERTenc	DistilBERT
768 (Default)	-	70 M	188 M	142 M	98 M
C = 512	2.3 M	61 M	174 M	134 M	90 M
C = 384	1.8 M	55 M	165 M	128 M	84 M
C = 128	1.1 M	44 M	149 M	116	75 M
C = 32	1 M	40 M	143 M	112 M	68 M

present the comparisons of model sizes in terms of the number of parameters. The first step is calculating the decoder’s size without using the dimensionality reduction method. As previously mentioned, the default encoder’s representation dimension ($d_{model}$) is 768 (based on the choice of the encoder), which results in a decoder with either 48 M (for ${\mathrm{BART}}_{\mathrm{enc}}$) or 33 M (other options) parameters. We utilize this value as a reference number to calculate 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. We have decoders with different parameter counts while using the same hyper-parameters, because we use each model’s pre-trained tokenizers with different vocabulary sizes, which affects the decoder’s last layer output dimension. BERT, DistilBERT, and transformer models use BERT’s pre-trained tokenizer, and BART uses its own.

Table 4. The ROUGE score for using a pre-trained autoencoder on top of pre-trained transformer-based encoders with different compression sizes. The summaries are generated using greedy decoding method. The percentages indicate the relative change in the ROUGE score compared to the baselines.

Models	R-1	R-2	R-L
Transformer	0.346	0.143	0.312
$+\mathrm{AE} (C=512$)	0.368 (106%)	0.157 (109%)	0.325 (104%)
$+\mathrm{AE} (C=384$)	0.363 (104%)	0.154 (107%)	0.322 (103%)
$+\mathrm{AE} (C=128$)	0.308 (89%)	0.114 (79%)	0.286 (91%)
$+\mathrm{AE} (C=32$)	0.156 (45%)	0.019 (13%)	0.174 (55%)
${\mathrm{BART}}_{\mathrm{enc}}$	0.355	0.142	0.310
$+\mathrm{AE} (C=512$)	0.341 (96%)	0.128 (90%)	0.298 (96%)
$+\mathrm{AE} (C=384$)	0.332 (93%)	0.121 (85%)	0.291 (93%)
$+\mathrm{AE} (C=128$)	0.257 (72%)	0.063 (44%)	0.239 (77%)
$+\mathrm{AE} (C=32$)	0.145 (40%)	0.014 (9%)	0.168 (54%)
BERT	0.349	0.133	0.306
$+\mathrm{AE} (C=512$)	0.339 (97%)	0.123 (92%)	0.298 (97%)
$+\mathrm{AE} (C=384$)	0.332 (95%)	0.119 (89%)	0.294 (96%)
$+\mathrm{AE} (C=128$)	0.278 (79%)	0.074 (55%)	0.256 (83%)
$+\mathrm{AE} (C=32$)	0.168 (48%)	0.021 (15%)	0.187 (61%)
DistilBERT	0.317	0.124	0.283
$+\mathrm{AE} (C=512$)	0.333 (105%)	0.123 (99%)	0.298 (105%)
$+\mathrm{AE} (C=384$)	0.334 (105%)	0.122 (98%)	0.297 (104%)
$+\mathrm{AE} (C=128$)	0.287 (90%)	0.083 (66%)	0.265 (93%)
$+\mathrm{AE} (C=32$)	0.161 (50%)	0.020 (16%)	0.180 (63%)

displays the results of the proposed method on the text summarization task, with the utilization of the classic ROUGE (R-1, R-2, and R-L) metric, employing the greedy decoding strategy (refer to Appendix B for the weighted random sampling and beam search scores). Autoencoders are trained with four (32, 128, 384, 512) different latent space dimensions ($C$) to analyze the effect of these compression rates. The models are also evaluated using the BERTScore [54] metric, and the results follow the same trend (refer to Appendix C).

In addition to reducing the number of parameters in the network by decreasing the decoder size, our configurations exhibit the capability to enhance the model’s summarization ability, resulting in higher scores compared to the original setup for both the transformer and DistilBERT models. The experiments show that adding an AE with $C=512$ outperforms the same vanilla encoder–decoder network. Both ${\mathrm{BART}}_{\mathrm{enc}}$ and BERT experiments with the same compression rate outperformed their vanilla model in several metrics by only using the beam search method. At the same time, the rest of the scores with a smaller decoder are still competitive. An even more interesting result is that a smaller dedicated transformer with the proposed method and $C=384$ performed better than both ${\mathrm{BART}}_{\mathrm{enc}}$ and BERT in the summarization benchmark.

Table 5. The comparison between using the pre-trained autoencoder (AE), training the autoencoder’s encoder jointly with the network from scratch (AE S), using a simple linear layer model for the projection (LL), and PCA to perform the dimensionality reduction. The percentages indicate the relative change in the ROUGE score compared to the baseline (BERT).

Models	R-1	R-2	R-L
BERT	0.349	0.133	0.306
$+\mathrm{AE} (C=512$)	0.339 (97%)	0.123 (92%)	0.298 (97%)
$+\mathrm{AE} \mathrm{S} (C=512$)	0.289 (82%)	0.079 (59%)	0.262 (85%)
$+\mathrm{LL} (C=512$)	0.277 (79%)	0.083 (62%)	0.260 (84%)
$+\mathrm{PCA} (C=458$)	0.143 (40%)	0.016 (12%)	0.156 (50%)

presents the results of our ablation study, illustrating the impact of the autoencoder pre-training step on the final score. It surpasses both training the $A{E}_{enc}$ 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 as well. The results also show that using an autoencoder with a latent space size of 384 ($C=384$) results in a ROUGE score close to the vanilla model. It reduces the decoder size by 46% from 48 M to 24 M for the ${\mathrm{BART}}_{\mathrm{enc}}$ model and 44% from 33 M to 16 M parameters for the other models. The critical point is that the combination of this configuration in association with the greedy decoding algorithm shows no noticeable degrading quality in the generated summary (refer to Appendix D).

Our findings demonstrate that by reducing the encoder’s output dimension from 768 to 384, we achieve a 44% reduction in decoder size while preserving 95% of the R-1 score in the worst-case scenario using BERT as the encoder. (refer to Appendix E for additional experiments on intermediate compression rates) Alternatively, with DistilBERT, we observe a potential increase of up to 105% in the R-1 score. Moreover, even with an almost 80% reduction in the decoder size, we can still achieve 90% of the ROUGE score with our dedicated transformer and DistilBERT architectures. This method will be integrated into a full pre-trained sequence-to-sequence model to measure its footprint in the next section.

5.2. Summarization

This section’s experiment utilizes the proposed architecture in a sequence-to-sequence pre-trained model (BART) to demonstrate its abilities. We fine-tune (i) vanilla BART-base and (ii) BART-base with autoencoder (various compression rates), while both BART’s encoder and the $A{E}_{enc}$ are frozen. The ROUGE score for these experiments (using beam search decoding) and the number of decoder parameters are reported in

Table 6. The ROUGE score results from integrating BART-base with our proposed architecture for summarization. The autoencoders are frozen, and “fe” stands for “frozen encoder”. The last column indicates the number of parameters in the decoder part of the network. The summaries are generated using the beam search decoding method.

Models		R-1	R-2	R-L	# Dec Params
BART		0.419	0.198	0.393	96 M
BARTfe	+AE (C = 504)	0.401	0.182	0.375	56 M
	+AE (C = 384)	0.400	0.181	0.374	40 M
	+AE (C = 120)	0.351	0.133	0.331	21 M
	+AE (C = 24)	0.182	0.026	0.170	11 M

. The decoder component is responsible for generating tokens one at a time during inference, so compressing it will largely optimize the generation process.

Our proposed architecture with a 384 compression size shows a 4.5% drop in the R-1 score while reducing the model’s decoder size by nearly 60%. The slight decrease in the ROUGE score does not seem to translate to a lower-quality summary (refer to Appendix F). Moreover, the network size decrease will compensate for a few lost overlapping n-grams in summaries by reducing the computational time and memory usage by a large margin at training and inference time. The following experiments are performed using a batch size of 16 and 32 for training and inference, respectively.

5.2.1. Computational Time

As stated in

Table 7. Comparing vanilla BART and the proposed architecture on the minutes it takes to fine-tune the model for 1 epoch and performing inference on the full test set (summarization task).

Model	Computational Time (Minutes)
	Fine-Tuning	Inference
BART	145	309
BART + AE (C = 384)	79	140

, the BART + AE ($C=384$) architecture can be fine-tuned 45% faster. It is a one-hour reduction (or 45%) compared to the BART’s 2:40 h for one epoch. Moreover, the inference (beam search) timing is a more important factor since the models are usually fine-tuned once to be used in numerous applications in the future. The 54% decrease in inference time on the test set (the 5:09 compared to 2:20 h) could greatly impact the responsiveness of downstream applications.

The block pruning technique [34] reported a 1.35 speedup by reducing the “linear” parameter count from 99 M to 23 M in the BART-base model. The smallest DistilBART [27] model showed a 1.66 faster inference rate with a 100 M decoder size. Our approach resulted in a much faster inference time with a 2.2 increase in speed rate.

5.2.2. GPU Memory Usage

The comparison of GPU memory consumption is shown in

Table 8. Comparing vanilla BART and the proposed architecture on the GPU memory (MBs) used to fine-tune the model for 1 epoch and performing inference on the full test set (summarization task).

Model	GPU Memory Usage (MB)
	Fine-Tuning	Inference
BART	28,384	30,624
BART + AE (C = 384)	10,240	23,680

. The BART+AE ($C=384$) experiment shows a 57% reduction in GPU memory usage while fine-tuning and a 20% reduction during inference. The massive drop in GPU utilization while fine-tuning makes it easier for both researchers and practitioners to try larger models in their experiments with minimal sacrifice in the quality.

5.3. Translation

The BART-base model is fine-tuned for translation to set a baseline score. Additionally, the proposed architecture is fine-tuned on the WMT16 dataset, as described in Section 3. The translations are generated using beam search (K = 4), and the BLEU score is used to measure the translation’s quality, as presented in

Table 9. The BLEU score results from integrating BART-base with our proposed architecture for translation. The autoencoders are frozen, and “fe” stands for “frozen encoder”. The last column indicates the number of parameters in the decoder part of the network. The translations are generated using the beam search decoding method.

Models		BLEU	# Dec Params
BART		21.05	96 M
BARTfe	+AE (C = 504)	18.93	56 M
	+AE (C = 384)	18.63	40 M
	+AE (C = 120)	13.95	21 M
	+AE (C = 24)	1.27	11 M

. The proposed architecture with a 50% compression rate (from 768 to 384) shows a ~10% decrease in the BLEU score. Meanwhile, the generated translations are close in meaning, even identical in some cases (refer to Appendix G). This experiment also reveals that an autoencoder pre-trained on one dataset can perform well on others without additional fine-tuning. The following experiments are conducted using a batch size of 32 and 64 for training and inference, respectively.

5.3.1. Computational Time

The translation task shows excellent improvement in training speed. As stated in

Table 10. Comparing vanilla BART and proposed architecture on the minutes required to fine-tune the model for 1 epoch and performing inference on the full test set (translation task).

Model	Computational Time (Minutes)
	Fine-Tuning	Inference
BART	247	23
BART + AE (C = 384)	110	21

, the fine-tuning process is 55% faster, and the inference speed increases by 9%. Our architecture’s novelty is in reducing the decoder’s size to maximize efficiency that primarily targets inference. However, the WMT16 dataset, on average, consists of a 50% lesser word count in target sequences (25 words, compared to CNN/DM’s 50). It will reduce the decoder’s workload, which might appear to be an unsatisfactory speed-up rate. This result shows that the more extended the sequences are, the more impactful our approach will be on inference speed.

5.3.2. GPU Memory Usage

As presented in

Table 11. Comparing vanilla BART and the proposed architecture on the GPU memory (MBs)used to fine-tune the model for 1 epoch and performing inference on the full test set (translation task).

Model	GPU Memory Usage (MB)
	Fine-Tuning	Inference
BART	25,024	31,680
BART + AE (C = 384)	13,120	27,200

, with only 10% translation quality loss, the proposed architecture uses 37% and 14% less GPU memory while training and inference, respectively. The translation dataset’s short source and target sequences require less GPU memory to process and store the weights. Nevertheless, there is a meaningful decrease in memory usage during the model’s life cycle.

5.4. Classification

The BART-base model with a classification head was used for the classification task. Compared to previous sections, the main difference is using a feedforward layer instead of a decoder. There are three experiments presented in

Table 12. The accuracy results from integrating BART-base with our proposed architecture for classification. The autoencoders are frozen, and “fe” stands for “frozen encoder”. The last column indicates the number of parameters in the classification head part of the network.

Models		Accuracy (%)	# Classifier Head Params
BART		86.73	592 K
BARTfe		80.19	592 K
BARTfe	+AE (C = 504)	78.12	263 K
	+AE (C = 384)	75.94	148 K
	+AE (C = 120)	70.04	16 K
	+AE (C = 24)	59.43	1 K

. First, the baseline accuracy for fine-tuning the model. Second, freezing the BART encoder to fine-tune the head component that reduces the accuracy by 6%. Third, the addition of the pre-trained autoencoder with various compression sizes. The 35% compression rate (C = 504) reduces the head size by 55% while showing an 8% reduction in accuracy. The following experiments are performed using a batch size of 64 and 128 for training and inference, respectively.

The proposed architecture does not make a significant contribution to the classification task. It is better to freeze the encoder and sacrifice fewer accuracy points. It is worth noting that the addition of the autoencoder results in a behavior as seen in previous experiments regarding accuracy drop. However, it does not improve efficiency by a large margin since the classifier head size is minimal, and there is little room for improvement. Also, it should be noted that the autoencoder adds about 1.5 M parameters (depending on compression rate) to the architecture.

5.4.1. Computational Time

The computational time results are displayed in

Table 13. Comparing vanilla BART and the proposed architecture on the minutes it takes to fine-tune the model for 1 epoch and performing inference on the full test set (classification task).

Model	Computational Time (Minutes)
	Fine-Tuning	Inference
BART	1200	83
BARTfe	450	83
BART + AE (C = 384)	480	85

. There is no noticeable change in the prediction duration. However, it will result in a more efficient training process since we are fine-tuning a small percentage of the network while the rest is frozen. The addition of an autoencoder performs marginally worse.

5.4.2. GPU Memory Usage

The

Table 14. Comparing vanilla BART and the proposed architecture on the GPU memory (MBs) used to fine-tune the model for 1 epoch and performing inference on the full test set (classification task).

Model	GPU Memory Usage (MB)
	Fine-Tuning	Inference
BART	27,264	22,831
BARTfe	3244	22,831
BART + AE (C = 384)	3264	22,966

presents the GPU usage results. It follows the same trend as the computational time benchmarks and shows no advantage in using the proposed architecture for the classification task.

6. Conclusions

This paper presents a method to shrink the decoder size (number of parameters) in a sequence-to-sequence setting. It has been performed by incorporating an autoencoder as a dimensionality reduction method between the encoder and decoder. The proposed method was validated using a custom summarizer network with various pre-trained encoders. The main idea is that decreasing the size of the encoder’s output leads to a dramatic decrease in the decoder size. It results in an optimized network in terms of GPU memory consumption and computational time.

Later, we thoroughly studied the strengths and shortcomings of the proposed architecture on the BART-base model. It is shown that the mentioned architecture on summarization can reduce the decoder’s size by 60% (from 96 M to 40 M) while maintaining 95.5% (in terms of R1 score) of the text generation capability. This will significantly impact speeding up both (i) fine-tuning by 45% while using 57% less GPU memory and (ii) the inference process by 54% while utilizing 20% less GPU memory. It will make the large pre-trained models more accessible to anyone without high-end hardware with negligible quality.

The automatic translation experiment follows the same trend as summarization. However, the binary classification task did not improve the model’s efficiency while showing the same trade-off between compression rate and accuracy loss. Furthermore, the additional experiments revealed the autoencoder’s impressive generalization capability on different datasets.

Our method can be directly used with other approaches, such as distillation, pruning and quantization, to reduce the network size further. One of our future research projects will investigate which combination could lead to the best compromise between the size and the model’s accuracy.