Improving the Performance of Compressed Transformers with Post-Hoc Embedding Alignment
Organizations: Indian Institute of Technology Guwahati, Guwahati, Assam, India · B.K. Birla Institute of Engineering & Technology, Pilani, Rajasthan, India
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.
Figures & tables
| Model | Params | Size | Compression method | ||
|---|---|---|---|---|---|
| BERT-base ( Devlin et al., 2019 ) | 12 | 768 | 110M | 100% | Reference model |
| DistilBERT ( Sanh et al., 2019 ) | 6 | 768 | 66M | 60% | Distilled from BERT-base |
| ASC ( Rai et al., 2024 ) | 4 | 768 | 52M | 48% | Layers pruned from BERT-base |
| BERT-Mini ( Turc et al., 2019 ) | 4 | 256 | 11M | 10% | Pretrained at reduced scale |
| BERT-Tiny ( Turc et al., 2019 ) | 2 | 128 | 4.4M | 4% | Pretrained at reduced scale |
| Model | Variant | POS | NER | EQA | Size | Speedup |
|---|---|---|---|---|---|---|
| DistilBERT | fine-tuned | 100 | 88 | 89 | 60.00 | 1.60 |
| + Linear | 100 | 102 | 91 | 61.48 | 1.58 | |
| + AlAd | 101 | 106 | 93 | 65.89 | 1.54 | |
| ASC | fine-tuned | 96 | 53 | 68 | 48.00 | 2.12 |
| + Linear | 99 | 75 | 70 | 49.01 | 2.10 | |
| + AlAd | 100 | 80 | 74 | 53.89 | 2.01 |
| Model | Variant | POS | NER | EQA | Size | Speedup |
|---|---|---|---|---|---|---|
| DistilBERT | + LoRA | 73 | 6 | 26 | 60.30 | 1.55 |
| + LoRA + AlAd | 96 | 90 | 58 | 71.08 | 1.47 | |
| ASC | + LoRA | 82 | 24 | 12 | 48.30 | 2.04 |
| + LoRA + AlAd | 98 | 76 | 69 | 59.08 | 1.89 | |
| BERT-M | + LoRA | 72 | 10 | 3 | 10.25 | 2.48 |
| + LoRA + AlAd | 92 | 71 | 34 | 15.03 | 2.28 |
| Model | Variant | POS | NER | EQA | Size | Speedup |
|---|---|---|---|---|---|---|
| DistilBERT | fine-tuned | 100 | 88 | 89 | 60.00 | 1.60 |
| frozen + AlAd | 95 | 94 | 7 | 65.89 | 1.54 | |
| ASC | fine-tuned | 96 | 53 | 68 | 48.00 | 2.12 |
| frozen + AlAd | 98 | 75 | 34 | 53.89 | 2.01 | |
| BERT-M | fine-tuned | 92 | 27 | 61 | 10.20 | 2.50 |
| frozen + AlAd | 92 | 71 | 3 | 12.68 | 2.41 |
| Model | Adapter | POS | NER | EQA | Sz | Spd |
|---|---|---|---|---|---|---|
| DistilBERT | Linear | 100 | 102 | 91 | 61.48 | 1.58 |
| W-1 | 101 | 104 | 91 | 62.28 | 1.55 | |
| W-3 | 101 | 106 | 92 | 64.09 | 1.54 | |
| W-5 | 101 | 106 | 93 | 65.89 | 1.54 | |
| ASC | Linear | 99 | 75 | 70 | 49.01 | 2.10 |
| W-1 | 99 | 76 | 70 | 50.28 | 2.03 |
| Model | Performance improvement | Size increase | Speedup loss | ||
| POS | NER | EQA | |||
| DistilBERT | 9.8% | 3.8% | |||
| ASC | 12.3% | 5.2% | |||
| BERT-M | 24.3% | 3.6% | |||
| BERT-T | 42.1% | 18.9% | |||
| With AlAd | Against | POS | NER | EQA | Size | Speedup |
|---|---|---|---|---|---|---|
| DistilBERT | BERT-base | 101 | 106 | 93 | 66 | 1.54 |
| ASC | DistilBERT | 100 | 91 | 83 | 90 | 1.26 |
| BERT-M | ASC | 102 | 152 | 96 | 26 | 1.14 |
| BERT-T | BERT-M | 99 | 214 | 58 | 56 | 0.91 |
| Paradigm | Compressed model | Weight correspondence | Optimisation target | Module at inference |
|---|---|---|---|---|
| Knowledge Distillation | Produced by the method | No | Teacher outputs or features | Discarded |
| PEFT | Given | No | Task loss | Retained or merged |
| Post-training repair | Given | Yes | Original layer output | None |
| Representation alignment | Given | No | Task loss of stitched model | Not deployed |
| AlAd | Given | No | Reference model embeddings | Retained |
Adapter Tuning
Adapter tuning adapts pretrained models with small, task-specific modules instead of updating all model parameters.
See all 76 papers