Improving the Performance of Compressed Transformers with Post-Hoc Embedding Alignment

Abstract

Model compression is necessary for deployment of transformer models in resource-constrained environments. However, compression reduces the performance of transformer models on downstream tasks. Existing solutions such as Knowledge Distillation (KD) and Parameter Efficient Fine-Tuning (PEFT) have their own limitations. KD trains the student under the teacher's supervision, but needs the teacher's training pipeline. PEFT methods such as LoRA adapt a model to a new task, but they do not recover the loss caused by compression. Our intuition is that a compressed model can be improved by aligning its token embeddings with those of a better performing reference model. We learn this alignment with a small feed-forward Alignment Adapter (AlAd) placed after the final layer of the compressed model. AlAd reads a window of neighbouring tokens, so it can use local context. AlAd is trained post-hoc, after the compressed model has been produced. The compressed model can either be kept frozen or jointly fine-tuned with AlAd. We study the effect of AlAd with BERT-base as the reference model and four compressed variants of the BERT family. We evaluate AlAd on three token-level tasks: POS tagging, NER and Extractive Question Answering. AlAd improves performance of compressed models on all three tasks. Joint fine-tuning gives a larger improvement in performance. The performance gain is the largest for the smallest models. AlAd provides a favorable trade-off between the efficiency and the performance of compressed models.