Source-linked AI summary
BERT-of-Theseus: Compressing BERT by Progressive Module Replacing
Canwen Xu, Wangchunshu Zhou, Tao Ge, Furu Wei, Ming Zhou
TL;DR
Large pretrained models impose memory and latency costs, while existing BERT compression methods often rely on explicit distillation losses. BERT-of-Theseus progressively replaces predecessor modules with compact successors using task-specific training, and retains over 98% of BERT-base performance while improving speed on GLUE.
Problem
Pretrained models contain millions or billions of parameters, making memory consumption and latency costly for production applications; existing knowledge-distillation methods explicitly use distillation losses.
Method
BERT-of-Theseus specifies compact successor modules, randomly replaces predecessor modules with them during training, progressively increases replacement, and uses only a task-specific loss.
Results
98.4% and 98.3% of BERT-base performance are retained on the GLUE development and test sets, respectively, while the compressed model is 1.94× faster and outperforms other knowledge-distillation baselines.
Takeaways & Limitations
Theseus Compression provides a model-compression pathway that transfers knowledge through progressive module replacement without an additional loss function.
Takeaways & Limitations
The reported approach compresses BERT, while applying Theseus Compression to heterogeneous modules and other neural models remains future work.
Abstract
from arXiv · showhide
In this paper, we propose a novel model compression approach to effectively compress BERT by progressive module replacing. Our approach first divides the original BERT into several modules and builds their compact substitutes. Then, we randomly replace the original modules with their substitutes to train the compact modules to mimic the behavior of the original modules. We progressively increase the probability of replacement through the training. In this way, our approach brings a deeper level of interaction between the original and compact models. Compared to the previous knowledge distillation approaches for BERT compression, our approach does not introduce any additional loss function. Our approach outperforms existing knowledge distillation approaches on GLUE benchmark, showing a new perspective of model compression.
1 Introduction
Large pretrained language models are computationally expensive because of their size and latency, motivating compression methods that preserve performance. BERT-of-Theseus progressively replaces BERT modules with smaller substitutes and reports faster inference with high retained performance.
- Millions or billions of parameters make pretrained models costly in memory and computation, while high latency hinders production deployment.
- Knowledge distillation compresses models by training a compact student to mimic a larger teacher, but its retained performance depends on the distillation setup.
- BERT-of-Theseus progressively substitutes BERT modules with fewer-parameter successor modules, randomly replacing predecessor modules during training before combining successors for inference.
- Theseus Compression uses the task-specific loss alone, avoiding the additional distillation losses and loss-weight tuning used by knowledge-distillation methods.
- 1.94× faster inference retains more than 98% of the original BERT's performance and outperforms other knowledge-distillation compression baselines.
2 Related Work
Prior model-compression work reduces cost through quantization, pruning, distillation, architectural changes, and dynamic acceleration. BERT-of-Theseus instead progressively replaces modules while retaining the predecessor–successor workflow shown in Figure 1.
- Model Compression: Model compression targets smaller size and lower computational cost while retaining as much performance as possible.
- Model Compression: Quantization reduces numerical representation bits, whereas weight pruning removes selected weights from the model.
- Model Compression: Knowledge distillation trains compact models to behave like original models, with prior variants using hints, intermediate representations, ensembles, multilingual training, or augmented data.
- Faster BERT: Other faster-BERT approaches prune attention heads, quantize to 2-bit, drop structures, share parameters, or use modified architectures.
- Faster BERT: BERT-of-Theseus compresses a 6-layer predecessor into a 3-layer successor by probabilistically replacing corresponding two-layer predecessor modules with one-layer successors during training.
3 BERT-of-Theseus
BERT-of-Theseus compresses BERT by randomly replacing predecessor modules with compact successor modules during training, then fine-tuning the assembled successor. A curriculum scheduler progressively increases replacement to transition from predecessor-guided learning toward successor inference.
- 3.1 Module Replacing: BERT-of-Theseus specifies a compact successor module for each predecessor module and randomly replaces predecessor modules with successors during training.The hybrid model trains predecessor and successor modules together while successor modules learn to reproduce the predecessor's behavior.
- 3.1 Module Replacing: Random module replacement adds training noise that regularizes successor learning, similarly to Dropout.Each replacement decision is an independent Bernoulli variable, and the resulting hybrid model changes across training passes.
- 3.1 Module Replacing: Training uses only a regular task-specific loss, such as cross-entropy, while predecessor-module weights remain frozen.Gradients pass through both predecessor and successor modules; the predecessor embedding and output layers are also frozen and reused.
- 3.2 Successor Fine-tuning and Inference: After replacement converges, all successor modules are combined and fine-tuned together before inference.This postreplacement phase makes training closer to inference, and the smaller successor modules yield a smaller successor model.
- 3.3 Curriculum Replacement: A curriculum-driven scheduler dynamically increases the replacement rate, unifying module replacement with an easy-to-hard transition toward successor fine-tuning.The scheduler begins with more predecessor guidance and later reduces that guidance; it also provides a warm-up mechanism for Transformer training.
- 3.3 Curriculum Replacement: Figure 2 contrasts a constant replacement rate with the scheduler and distinguishes module replacing from successor fine-tuning.The two phases are marked using different gray shades.
4 Experiments
The experiments evaluate task-specific compression of BERT into a six-layer model on GLUE, using task training data without external unlabeled corpora. BERT-of-Theseus retains most predecessor performance, exceeds several compression baselines, and supports intermediate-task transfer.
- 4.2 Experimental Settings: The task-specific setting uses only each GLUE task’s training set, without external unlabeled corpora, to test the compression approach directly.The authors report that task-specific compression takes no longer than 20 GPU hours for any GLUE task.
- 4.3 Training Details: The predecessor is fine-tuned BERT-base, while successor training initializes from BERT-base’s first six layers to aid convergence on small datasets.Predecessor fine-tuning uses batch size 32, learning rate 2×10−5, and four epochs.
- 4.5 Experimental Results: 98.4% and 98.3% of BERT-base performance are retained on GLUE development and test sets, respectively.The model outperforms the fine-tuning, vanilla KD, and PKD baselines, while remaining slightly below PD-BERT on MNLI.
- 4.5 Experimental Results: BERT-of-Theseus achieves nearly identical performance to BERT-base on RTE and exceeds BERT-base on QQP.The authors report good performance on both datasets with more than 350K samples and datasets with fewer than 4K samples.
- 4.6 Intermediate-Task Transfer Learning: MNLI intermediate-task transfer matches DistilBERT on MRPC and outperforms it on the other sentence-level tasks.The transferred model also outperforms PD-BERT on three tasks, according to the authors.
5 Analysis
The analyses examine module location, replacing rates, and replacement scheduling, showing that compression behavior depends on where and how replacement is applied. Curriculum scheduling and moderate replacing rates consistently yield stronger results than their relevant alternatives.
- Impact of Module Replacement: Replacing the first module significantly harms performance, whereas replacing the last two modules has limited influence.The analysis attributes the first-module sensitivity to reduced representation capability becoming a bottleneck for subsequent layers.
- Impact of Replacing Rate: Equivalent-learning-rate adjustments do not explain the performance differences between replacing-rate settings.The learning rate is adjusted as lr = lr′/p when the equivalent learning rate is fixed, and the resulting curves remain trivially different.
- Impact of Replacing Rate: A replacing rate between 0.5 and 0.7 consistently gives satisfying performance across GLUE tasks, while p = 0.1 causes significant drops.The best replacing rate still differs across tasks.
- Impact of Replacement Scheduler: Curriculum scheduling consistently outperforms constant replacing rates, while anti-curriculum scheduling produces a substantial performance drop.The comparison evaluates constant, curriculum, and anti-curriculum schedulers on GLUE-dev.
- Different Compression Settings: BERT-of-Theseus consistently outperforms fine-tuned truncated BERT baselines when different numbers of layers are replaced with one layer.The setting includes replacing 3/4 layers with one Transformer layer, producing a 4/3-layer BERT model.
6 Discussion
The discussion presents Theseus Compression as a new model-compression pathway and identifies broader architectures as future application targets. It specifically highlights heterogeneous modules and neural models beyond BERT as directions for further exploration.
- Contributions: Theseus Compression is presented as a novel approach that compresses BERT more effectively than other Knowledge Distillation-compressed models.The authors characterize it as a new genre and path toward model compression.
- Future Work: Future work will explore applying Theseus Compression to heterogeneous network modules with different architectures or hidden-space sizes.Examples include ShuffleNet units replacing ResBlocks and Reformer layers replacing Transformer layers, with feature mapping enabling different input and output sizes.
- Future Work: The approach is proposed for investigation on other neural models, including ResNet, Convolutional Neural Networks, and Graph Neural Networks.The paper also proposes combining compression with dynamic acceleration methods to improve pretrained-language-model efficiency.