M ini ALBERT : Model Distillation via Parameter-Efficient Recursive Transformers (2023)

Abstract: Pre-trained Language Models (LMs) have become an integral part of Natural Language Processing (NLP) in recent years, due to their superior performance in downstream applications. In spite of this resounding success, the usability of LMs is constrained by computational and time complexity, along with their increasing size; an issue that has been referred to as overparameterisation. Different strategies have been proposed in the literature to alleviate these problems, with the aim to create effective compact models that nearly match the performance of their bloated counterparts with negligible performance losses. One of the most popular techniques in this area of research is model distillation. Another potent but underutilised technique is cross-layer parameter sharing. In this work, we combine these two strategies and present MiniALBERT, a technique for converting the knowledge of fully parameterised LMs (such as BERT) into a compact recursive student. In addition, we investigate the application of bottleneck adapters for layer-wise adaptation of our recursive student, and also explore the efficacy of adapter tuning for fine-tuning of compact models. We test our proposed models on a number of general and biomedical NLP tasks to demonstrate their viability and compare them with the state-of-the-art and other existing compact models. All the codes used in the experiments and the pre-trained compact models will be made publicly available.

## introduction following the introduction of bert @xcite , generic pre-trained language models (lms) have started to dominate the field of nlp. virtually all state-of-the-art nlp models are built on top of some large pre-trained transformer as a backbone and are subsequently fine-tuned on their target dataset. while this pre-train and fine-tune approach has resulted in significant improvements across a wide range of nlp tasks, the widespread use of resource-exhaustive and overparameterised transformers has also raised concerns among researchers about their energy consumption, environmental impact, and ethical implications @xcite . as a response to this, different approaches have appeared with the aim to make large lms more efficient, accessible, and environmentally friendly. model compression is a line of research that has recently received considerable attention. it involves encoding a larger and slower but more performant model into a smaller and faster one with the aim to retain much of the former's performance capability @xcite . knowledge distillation @xcite , quantisation @xcite , and pruning @xcite are three examples of such methods. adapter modules @xcite are recently introduced as an effective mechanism to address the parameter inefficiency of large pre-trained models. in this approach, several 'bottleneck adapters' @xcite are embedded inside different locations within the original network. during fine-tuning, the parameters of the original model are kept fixed, and for each new task only the adapters are fine-tuned. this only adds a small number of parameters to the overall architecture and allows for a much faster and more efficient fine-tuning on different downstream tasks. another approach to improve efficiency of lmbased transformers is shared parameterisation, which was popularised by albert @xcite . while the original formulation of transformers @xcite employs full parameterisation wherein each model parameter is independent of other modules and used only once in the forward pass, shared parameterisation allows different modules of the network to share parameters, resulting in a more efficient use of resources given the same parameterisation budget. however, a common downside of this approach is slower inference time and reduced performance. @xcite posits two different parameterisation methods as an attempt to address the compute and memory challenges of transformer models and explores layer-wise adaptation in an encoder-decoder architecture. these methods exploit cross-layer parameter sharing in a way that would allow for the model to be utilised on mobile devices with strict memory constraints while achieving state-of-theart results on two seq2seq tasks for english. in this work, we exploit some of the above approaches to create a number of compact and efficient encoder-only models distilled from much larger language models. the contributions of this work are as follows: ## methods in this work, we introduce a method to distil the knowledge of a fully parameterised transformer into a lightweight efficient recursive transformer via layer-to-layer distillation. we also experiment with layer-wise adaption of the recursive transformer via bottleneck adapters and factorise the embedding layer for extra parameter saving. in this section, we explain each component of our compact models in detail. ## experiments and results we evaluate our general models on the widely used glue benchmark @xcite . all the models were pre-trained on four tesla v100 32gb gpus with a total batch size of 192 (48 each) and fine-tuning was done using only one tesla v100. a random seed of 42 was used consistently throughout training for fair comparison. for all of the datasets, in order to do full fine-tuning, we use a learning rate from the set {5e-5, 3e-5, 1e-5}. for large datasets (mnli, qqp, qnli, and sst-2), we train models for a maximum of 5 epochs, and up to ten epochs on other datasets. the hyperparameters used for full-finetuning are listed in for adapter-tuning, the learning rate was selected from the set {5e-5, 5e-4, 1e-3}. models were trained for a maximum of 10 epochs for large datasets and up to 15 epochs for other datasets. the biomedical models are evaluated on the task of named entity recognition (ner), which is one of the most prominent tasks in biomedical nlp. we use four well-known datasets, namely, ncbidisease @xcite , bc5cdr-disease @xcite , bc5cdr-chem @xcite , and bc2gm @xcite . we generally follow the same pre-processing pipeline as @xcite . for biomedical ner, we use pre-processed datasets from @xcite . we fine-tune the models using a learning rate from the set {5e-5, 3e-5, 1e-5} and perform adapter-tuning with a learning rate from {5e-5, 5e-4, 1e-3}. the hyperparameters for the biomedical datasets are presented in tables 9 and 10 . the results of the baselines and our biomedical models are available in ## discussion we trained our recursive students both with and without adapters and found that generally having adapters would increase the learning capacity of the model since each iteration provides an extra degree of freedom to the model in order to capture a specific type of feature. as shown in tables 1 and 2, in virtually all of the studies, models with adapters outperformed models of comparable size without adapters. in general, both models with and without adapter have outperformed their baseline by a significant margin (as shown in tables 2 and 1 ), demonstrating the effectiveness of the proposed layer-to-layer distillation loss for constructing powerful compact recursive models. comparison between the attention maps of our trained student and teacher models (figure 3 in appendix) suggests the specific recursive architecture we have introduced in this work is indeed capable of mimicking the components of a larger model. in addition, our experiments revealed that utilising adapters in the pre-training stage enables the model to effectively use adapter-tuning with only minor performance drops. however, adaptertuning in models that have not used adapters in the pre-training stage causes major performance drops. we also found that adapter tuning is nearly 30% faster than full fine-tuning. therefore, with adapter-tuning, the models can be trained for a larger number of epochs given the same training time, which can potentially increase the performance of the model. in general, a higher learning rate was required for adapter-tuning than for full fine-tuning. for both general and biomedical tasks we found that a learning rate of 5e -4 results in the best performance, however, in some cases, a higher or lower learning rate such as 5e -5 or 1e -3 was deemed better. another method explored in this work for parameter saving is embedding factorisation. in our experiments, we observed that regardless of drastic parameter reduction, models using this approach are still able to perform on par or even in some cases better than models with full-rank embedding. ## conclusion and future works in this work, we explored distilling the knowledge of fully-parameterised language models into recursive students with cross-layer parameter sharing. we used a layer-to-layer distillation technique to observe the learning capacity of the recursive students compared to their fully-parameterised teachers. we used bottleneck adapters for improving the performance of our recursive students and also assessed the benefits of adapter-tuning at the finetuning stage. furthermore, an embedding factorisation technique was used for additional parameter reduction, which was evaluated with and without an extra distillation loss to match the student's embeddings with the teacher. finally, by integrating all of the strategies outlined above, we were able to train compact recursive students with no more than 12m parameters, yielding competitive performance on both general and biomedical nlp. in the future, we hope to investigate various parameter-sharing and embedding factorisation strategies, as well as other layer-wise adaption techniques such as prompttuning. we would also like to train recursive students with larger hidden sizes and more training iterations to compress massive fully-parameterised models with minor performance drops. loss. despite the fact that the recursive student has only one layer, it has been able to perfectly mimic its teacher in some of the heads which shows the efficiency of the layer-to-layer distillation loss.