[论文解读] BERT-of-Theseus: Compressing BERT by Progressive Module Replacing

Canwen Xu, Wangchunshu Zhou+3 more|arXiv (Cornell University)|Feb 7, 2020

Topic Modeling参考文献 44被引用 34

一句话总结

BERT-of-Theseus 在训练过程中逐步用更小的替代模块替换 BERT 模块，在保持 GLUE 上对 BERT-base 性能超过 98% 的前提下实现约 1.94x 的推理加速，同时无需额外蒸馏损失。

ABSTRACT

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.

研究动机与目标

• 推动在不依赖额外蒸馏损失的情况下，减小大型 Transformer 模型如 BERT 的尺寸与计算量。
• 引入一个渐进式模块替换（Theseus Compression）框架，在训练过程中在前身模块和后继模块之间交替。
• 证明所得的紧凑模型在 GLUE 上维持近似原始性能同时提供加速。
• 显示基于课程学习的替换调度比恒定替换率得到更好的结果。
• 提供对哪些层次和替换策略对性能影响最大的分析。

提出的方法

• 将原始 BERT 划分为模块并为每个模块定义紧凑的替代模块。
• 在训练过程中，以概率 p 将前身模块替换为后继模块，在一个前向传播中混合两者。
• 仅优化任务相关损失（如交叉熵），同时冻结前身的嵌入和输出，以实现跨模块的梯度流。
• 在收敛后，将所有后继模块组装成完整的后继模型，并使用相同的损失进行微调。
• 应用一个由课程学习驱动的替换调度器，在训练步骤中逐步增加替换概率。
• 提供一个简单的线性调度器 p_d = min(1, kt + b) 来控制动态替换率并热身学习。

实验结果

研究问题

• RQ1 Can progressively replacing modules within a large pretrained model yield effective compression without extra distillation objectives?
• RQ2 Does curriculum-based scheduling improve the trade-off between compression and accuracy compared to constant replacement rates?
• RQ3 How does Theseus Compression compare to KD-based baselines on GLUE in terms of performance, speed, and model size?
• RQ4 Is the approach model-agnostic to Transformer-based architectures and potential to apply to other domains?

主要发现

方法	CoLA	MNLI	MRPC	QNLI	QQP	RTE	SST-2	STS-B	宏平均
BERT-base	54.3	83.5	89.5	91.2	89.8	71.1	91.5	88.9	82.5
DistilBERT	43.6	79.0	87.5	85.3	84.9	59.9	90.7	81.2	76.5
Vanilla KD	45.1	80.1	86.2	88.0	88.1	64.9	90.5	84.9	78.5
BERT-PKD	45.5	81.3	85.7	88.4	88.4	66.5	91.3	86.2	79.2
BERT-of-Theseus	51.1	82.3	89.0	89.5	89.6	68.2	91.5	88.7	81.2

• BERT-of-Theseus 在 6 层紧凑模型下实现 1.94x 的推理加速，同时在 GLUE 上保留了 BERT-base 的 98.4%（开发集）和 98.3%（测试集）性能。
• 该方法在 GLUE 的大多数任务上优于原生 KD 和 PKD 基线。
• 课程替换调度器在性能上始终优于恒定替换和反课程策略。
• 替换早期 Transformer 层往往比替换后期层对性能的负 impact 更大，表明早期层对语言特征贡献更大。
• 中间任务转移（以 MNLI 作为预训练任务）在若干任务上与 DistilBERT 和 PD-BERT 相比表现具有竞争力或更优。