BERT-of-Theseus: Compressing BERT by Progressive Module Replacing 解读

一、论文基本信息

论文题目：BERT-of-Theseus: Compressing BERT by Progressive Module Replacing

作者：Canwen Xu、Wangchunshu Zhou、Tao Ge、Furu Wei、Ming Zhou

发表会议：EMNLP 2020

方法名称：Theseus Compression / BERT-of-Theseus

官方代码：JetRunner/BERT-of-Theseus

这篇论文提出了一种非常有意思的 BERT 压缩方式：**不是直接剪掉某些层，也不是用复杂的知识蒸馏损失训练小模型，而是把原始 BERT 的模块逐步替换成更小的模块。**论文收录于 EMNLP 2020，ACL Anthology 页面显示论文页码为 7859–7869，官方摘要也明确说明，该方法先把原始 BERT 划分成多个模块，再为每个原始模块构造一个紧凑替代模块，训练过程中逐步提高替代模块出现的概率。

二、论文要解决的问题

BERT-base 有 12 层 Transformer，效果强，但推理成本高。常见压缩方法通常有几类：

第一类是直接训练一个浅层 BERT。
例如取 BERT 的前 6 层，再在下游任务上 fine-tune。这种方式简单，但精度下降明显。

第二类是知识蒸馏。
例如 DistilBERT、BERT-PKD、TinyBERT 等，用大模型作为 teacher，小模型作为 student，通过 logits、hidden states、attention maps 等多种蒸馏损失训练小模型。

第三类是剪枝或删层。
例如 LayerDrop、head pruning 等，直接删除一部分结构。

BERT-of-Theseus 认为，传统知识蒸馏虽然有效，但通常需要设计额外蒸馏损失，并且要调节多个 loss 权重；不同任务、不同模型结构下，这些超参数可能比较麻烦。论文特别强调，它的压缩过程只使用普通任务损失，不额外引入 distillation loss。(ACL 文集)

所以这篇论文想解决的问题是：

能不能不用复杂蒸馏损失，只通过训练过程中的模块替换，让一个小 BERT 逐渐接管大 BERT 的功能？

三、核心思想：逐步替换模块

BERT-of-Theseus 的核心思想可以概括为：

先让大模型和小模型混合工作，再逐步提高小模块出现的概率，最终让小模型完全替代大模型。

论文借用了“忒修斯之船”的比喻。原船的木板被一块块替换，最后所有木板都换成新的，但整体功能仍然保持。BERT-of-Theseus 中也是类似：

原始 BERT 是 predecessor model。

压缩后的 BERT 是 successor model。

原始模块是 predecessor module。

替代模块是 successor module。

训练过程中，模型不是一下子从 12 层 BERT 变成 6 层 BERT，而是每次前向传播时，随机决定某个原始模块是由大模型模块执行，还是由对应的小模型模块执行。论文中明确写到，训练时会随机用 successor module 替换 predecessor module，训练结束后把所有 successor modules 组合起来作为最终压缩模型。

这点很关键：

BERT-of-Theseus 不是直接训练一个孤立的小模型，而是让小模块在大模型的上下文中逐步学习。

四、模块如何划分

在 BERT-base 压缩实验中，原始模型是 12 层 BERT。论文的典型设置是：

把原始 BERT-base 的 12 层划分成 6 个 predecessor modules。

每个 predecessor module 包含 2 层 Transformer。

每个 successor module 使用 1 层 Transformer 来替代这 2 层。

这样最终得到的 successor model 就是一个6 层 BERT。

论文在方法部分也明确举例说明，在 BERT 压缩场景中，可以用1 个 Transformer layer 作为 2 个 Transformer layers 的 successor module。

所以它最终压缩的是深度：

12 层 BERT → 6 层 BERT

这不是 head pruning，也不是 FFN neuron pruning，而是module-level / layer-level 结构压缩。

五、训练过程：大模块和小模块混合前向

训练时，每个模块位置都有两种选择：

使用原始 predecessor module。

使用压缩 successor module。

每次前向传播时，系统根据一个替换概率随机选择。如果替换概率较低，大多数位置还是原始模块；如果替换概率较高，更多位置会使用小模块。

这种训练方式有两个作用。

第一，小模块可以在大模型提供的稳定表示环境中学习。
如果一开始就直接训练 6 层小 BERT，小模型要一次性承担全部任务，训练难度较大。Theseus 的混合前向让小模块先在部分位置替代大模块，逐步适应任务。

第二，大模型和小模型之间存在更深层次交互。
普通蒸馏通常是 teacher 输出目标，student 单独学习；BERT-of-Theseus 则让 predecessor modules 和 successor modules 在同一次前向传播中共同组成一个混合模型，梯度可以穿过两类模块形成更直接的训练互动。论文也强调，这种方式让原模型和压缩模型在训练阶段共同工作，形成比普通蒸馏更深的交互。

六、predecessor 是否更新？

训练时，predecessor modules 是冻结的。也就是说，大 BERT 的参数不更新，只作为稳定参考结构存在。

论文明确说明，在反向传播时，所有 predecessor modules 的权重被冻结；embedding layer 和 output layer 也直接采用 predecessor model 中的参数，并在该训练阶段冻结。(ACL 文集)

这带来一个好处：

训练目标更明确：让 successor modules 学会接替 predecessor modules，而不是让大模型和小模型同时乱动。

如果大模型模块也更新，整个混合模型的分布会不断变化，小模块更难稳定学习。

七、只使用任务损失，不使用额外蒸馏损失

这是这篇论文最重要的特点之一。

很多 BERT 压缩方法会设计复杂蒸馏损失，例如：

logits distillation

hidden states distillation

attention maps distillation

embedding distillation

intermediate layer matching

BERT-of-Theseus 不这么做。它训练时只使用普通下游任务损失，例如分类任务中的交叉熵。论文特别强调，与 KD 方法相比，Theseus Compression 不需要额外设计相似性损失，也不需要平衡多个损失权重。

所以它的训练逻辑更简单：

把大模型的一部分模块随机换成小模块。

用原任务标签训练整个混合网络。

小模块在混合网络中逐步学会替代大模块。

这也是它和 TinyBERT、BERT-PKD 这类蒸馏方法最大的区别。

八、Curriculum Replacement：为什么要逐步提高替换概率

如果替换概率一开始就很高，小模块还没学好，就要承担大量计算，训练容易不稳定。

如果替换概率一直很低，小模块获得训练机会不足，最终接管大模型时会表现不好。

所以论文提出Curriculum Replacement，也就是课程式替换。

它的思想是：

训练初期，让 predecessor modules 出现更多，小模块在大模型帮助下学习。

训练后期，逐步提高 successor modules 的替换概率，让小模块承担更多责任。

最后，当替换概率接近 1 时，模型基本完全由 successor modules 组成，也就过渡成真正的压缩模型。论文第 3.3 节提出了 curriculum learning driven replacement scheduler，并使用线性 scheduler 动态调整 replacement rate。

这个设计很自然：

先易后难。

先依赖大模型。

再逐步摆脱大模型。

论文实验也表明，使用 curriculum scheduler 的模型稳定优于固定替换率，而 anti-curriculum 会导致明显性能下降。

九、Successor Fine-tuning：最后为什么还要微调

当模块替换训练完成后，所有 successor modules 被组合成最终压缩模型。

但是训练时，successor modules 经常和 predecessor modules 混合工作；推理时，所有位置都换成 successor modules。训练和推理之间仍然存在一点分布差异。

因此论文又加入了successor fine-tuning阶段。

也就是说：

先通过模块替换让 successor modules 学会接替 predecessor modules。

再把所有 successor modules 组合成完整小模型。

最后对这个完整小模型进行一轮 fine-tuning。

论文在方法部分明确说明，在 module replacing compression 收敛后，会收集所有 successor modules，并进一步做 successor fine-tuning，使训练和推理过程更接近。

这一步的意义是：

让小模型内部所有模块真正适应彼此，而不是只适应和大模型模块混合工作。

十、实验设置

论文主要在GLUE benchmark上验证。任务包括：

论文中明确说明，SST-2、MNLI、QNLI、RTE 使用 accuracy；MRPC 和 QQP 使用 F1 与 accuracy；STS-B 使用 Pearson 和 Spearman；CoLA 使用 Matthews correlation。

实验设定是task-specific compression。也就是说，每个 GLUE 任务先 fine-tune 一个 BERT-base predecessor，然后只使用该任务训练集压缩模型，不额外使用大规模无标注语料。论文特别说明，它没有使用外部无标注语料，而是只用 GLUE 各任务训练集做压缩；每个 GLUE 任务压缩过程不超过 20 GPU hours。

十一、主要实验结果

论文最核心的实验是把12 层 BERT-base 压缩成 6 层 BERT-of-Theseus。

所有压缩模型都是：

6 layers

66M parameters

1.94× speed-up

论文 Table 2 在 GLUE development set 上报告，BERT-base 的 macro score 是82.5，BERT-of-Theseus 达到81.2；相比直接 fine-tuning 6 层模型的77.2、Vanilla KD 的78.5、BERT-PKD 的79.2、LayerDrop 的78.8，BERT-of-Theseus 明显更好。

开发集主要结果可以简化看成：

这个结果说明：

直接训练一个 6 层 BERT 效果明显不够。

普通蒸馏能提升一些。

逐步模块替换比这些 baseline 更有效。

论文还在 GLUE test server 上报告结果：BERT-base macro score 为80.0，BERT-of-Theseus 为78.6，优于 Fine-tuning、Vanilla KD 和 BERT-PKD 等对比方法。

十二、为什么它效果好？

这篇论文真正有价值的地方，不只是“把 12 层变成 6 层”，而是它提出了一种更平滑的压缩路径。

直接训练小模型时，小模型从一开始就必须完成全部任务，学习压力大。

普通蒸馏时，小模型单独学习 teacher 的输出或中间表示，teacher 和 student 之间主要是 loss-level 联系。

BERT-of-Theseus 不同：

小模块被直接插入大模型内部。

它必须在真实前向路径中接替大模块的功能。

它前后可能连接 predecessor modules，也可能连接 successor modules。

这种混合组合让小模块在更接近真实推理的环境中学习。

可以理解为：

不是让学生远远看老师做题，而是让学生逐步接手老师的一部分工作。

这比普通 teacher-student 蒸馏更像“岗位交接”。

十三、模块替换的位置有什么影响？

论文还分析了不同模块被替换时的影响。

结果发现，替换靠前模块带来的性能下降更大，替换后面模块影响较小。例如在 QNLI、MNLI、QQP 上，替换第一个 predecessor module 造成明显下降，而替换第 5、6 个模块影响小得多。论文解释说，前几层主要负责抽取语言特征，如果前面表示能力下降，会成为后续层的瓶颈。

这个结论很重要：

BERT 的不同层并不等价。

前层模块更敏感，后层模块更容易替换。

这也说明，虽然 BERT-of-Theseus 最终采用统一的 2 层替换 1 层，但如果进一步优化，完全可以考虑不同层采用不同压缩比例。

十四、替换概率怎么选？

论文实验发现，固定替换率在0.5 到 0.7范围内通常能取得较好效果；如果替换率太小，例如 0.1，性能会明显下降。原因很直观：替换率太小，successor modules 出场机会太少，学不到足够能力。

不过最终更推荐使用 curriculum scheduler，而不是固定替换率。因为固定替换率很难同时兼顾前期稳定和后期接管，而课程式替换可以自然完成从大模型到小模型的过渡。

所以替换概率的核心经验是：

不能太低，否则小模块训练不足。

不能一开始太高，否则训练不稳定。

逐步提高替换概率最合理。

十五、能不能压缩到更小？

论文还测试了更激进的压缩，把更多 predecessor layers 替换为更少 successor layers。

在 Table 6 中，6 层 BERT-of-Theseus 的 GLUE-dev macro score 为81.2，4 层模型为77.2，3 层模型进一步下降。对应加速分别为1.94×、2.82×、3.66×。

这说明：

BERT-of-Theseus 可以支持更高压缩率，但压缩越激进，精度下降越明显。

6 层是一个比较合理的折中点：速度接近 2 倍，精度保留较好。

4 层、3 层虽然更快，但表示能力下降更明显，尤其在 CoLA、RTE 等小数据或敏感任务上会更脆弱。

十六、和知识蒸馏的区别

BERT-of-Theseus 和知识蒸馏有相似目标：都希望小模型模仿大模型。

但方法不同。

知识蒸馏是 teacher 和 student 分开运行。
student 通过 logits、hidden states、attention maps 等损失模仿 teacher。

BERT-of-Theseus 是 predecessor 和 successor 混合运行。
successor modules 被直接插入 predecessor model 的前向路径中，逐步替换原模块。

所以它的特点是：

不需要设计蒸馏损失。

不需要调多个 loss 权重。

大模型和小模块在同一个前向图中协同工作。

压缩过程更像结构替换，而不是单纯输出模仿。

十七、它是不是剪枝？

严格来说，BERT-of-Theseus不是传统意义上的剪枝。

它不是先训练一个大模型，然后判断哪些 heads、neurons 或 layers 冗余，再把它们删掉。

它做的是：

为大模型模块构造紧凑替代模块。

通过训练让替代模块逐步接管原模块。

最后丢弃原模块，只保留替代模块。

所以它更准确地说是：

结构化模型压缩。

模块级替换压缩。

深度压缩。

如果一定要放到剪枝语境里，可以说它和 layer pruning 的目标相似，都是减少 Transformer 层数；但实现方式不是“删层”，而是“用小模块替换大模块”。

十八、方法优点

18.1 训练目标简单

它不需要额外蒸馏损失，只使用普通任务损失。

这让方法更简洁，也减少了 loss 权重调参问题。

18.2 压缩过程平滑

逐步替换比直接训练小模型更稳定。

小模块不是一下子承担全部任务，而是在大模型帮助下逐渐接管。

18.3 模块级设计比较通用

论文强调，这种压缩方式不依赖 BERT 特有结构，理论上也可以扩展到其他模块化神经网络。ACL 页面摘要也说明，该方法相比一些 BERT 专用蒸馏方法，展示了一种新的模型压缩视角。

18.4 GLUE 上效果强

6 层 BERT-of-Theseus 在 GLUE-dev 上保留了 BERT-base 大部分性能，并明显优于直接 fine-tuning、Vanilla KD、BERT-PKD 和 LayerDrop。

18.5 适合任务级压缩

它只使用下游任务数据即可压缩，不需要额外无标注预训练语料。论文明确采用 task-specific compression setting，并说明每个 GLUE 任务压缩不超过 20 GPU hours。

十九、方法局限

19.1 需要先得到 fine-tuned predecessor

每个任务都要先 fine-tune 一个完整 BERT-base predecessor，然后再做 Theseus Compression。

这意味着它不是直接给一个通用小 BERT，而是更偏任务级压缩。

19.2 主要压缩深度，宽度不变

它从 12 层变成 6 层，但 hidden size、attention head 数量、FFN hidden size 等基本保持 BERT-base 结构。

因此，它不像 DynaBERT 那样同时压缩宽度和深度。

19.3 替代模块结构比较简单

典型设置是 2 层 Transformer 替换成 1 层 Transformer。

这种设计直接有效，但不一定是最优。不同层可能需要不同替代结构，前层可能更敏感，后层可能更容易压缩。

19.4 仍然需要任务训练数据

它不是纯后处理压缩。每个任务都要进行模块替换训练和 successor fine-tuning。

19.5 更激进压缩时精度下降明显

6 层效果较好，但压到 4 层、3 层时 GLUE macro score 明显下降。说明模块替换并不能无限压缩，模型容量仍然是硬约束。

二十、整体评价

BERT-of-Theseus 的核心价值在于，它把 BERT 压缩从“训练一个小模型去模仿大模型”，变成了“让小模块逐步接管大模块”。

它不依赖复杂的蒸馏损失，也不需要设计 attention map matching、hidden state matching 这类任务。它真正的创新点在训练路径：

先混合使用大模块和小模块。

再逐步提高小模块出现概率。

最后完全切换到小模型。

这使压缩过程更加平滑，也让 successor modules 在真实前向路径中学习如何替代 predecessor modules。

它不是最细粒度的剪枝方法，也不是最极致的压缩方法，但它提供了一个非常清晰的新思路：

模型压缩可以通过“结构替换过程”完成，而不一定非要依赖显式蒸馏损失或重要性剪枝。

二十一、一句话总结

《BERT-of-Theseus: Compressing BERT by Progressive Module Replacing》提出 Theseus Compression，将 12 层 BERT-base 划分为多个 predecessor modules，并用更小的 successor modules 在训练过程中逐步替换它们；通过课程式提高替换概率，小模块逐渐接管大模块，最终形成 6 层压缩 BERT。它的核心贡献不是剪掉某些 head 或 neuron，而是用“渐进式模块替换”实现深度压缩，并且只依赖普通任务损失就能在 GLUE 上取得接近 BERT-base 的性能。