一、论文基本信息论文题目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 2020ACL 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 distillationhidden states distillationattention maps distillationembedding distillationintermediate layer matchingBERT-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上验证。任务包括任务类型评价指标CoLA语法可接受性判断Matthews correlationMNLI自然语言推理AccuracyMRPC句子复述识别F1 / AccuracyQNLI问答式自然语言推理AccuracyQQP问题对重复识别F1 / AccuracyRTE文本蕴含AccuracySST-2情感分类AccuracySTS-B语义文本相似度Pearson / Spearman论文中明确说明SST-2、MNLI、QNLI、RTE 使用 accuracyMRPC 和 QQP 使用 F1 与 accuracySTS-B 使用 Pearson 和 SpearmanCoLA 使用 Matthews correlation。实验设定是task-specific compression。也就是说每个 GLUE 任务先 fine-tune 一个 BERT-base predecessor然后只使用该任务训练集压缩模型不额外使用大规模无标注语料。论文特别说明它没有使用外部无标注语料而是只用 GLUE 各任务训练集做压缩每个 GLUE 任务压缩过程不超过 20 GPU hours。十一、主要实验结果论文最核心的实验是把12 层 BERT-base 压缩成 6 层 BERT-of-Theseus。所有压缩模型都是6 layers66M parameters1.94× speed-up论文 Table 2 在 GLUE development set 上报告BERT-base 的 macro score 是82.5BERT-of-Theseus 达到81.2相比直接 fine-tuning 6 层模型的77.2、Vanilla KD 的78.5、BERT-PKD 的79.2、LayerDrop 的78.8BERT-of-Theseus 明显更好。开发集主要结果可以简化看成方法层数参数量加速GLUE-dev MacroBERT-base12原始1.00×82.5Fine-tuning 6-layer666M1.94×77.2Vanilla KD666M1.94×78.5BERT-PKD666M1.94×79.2LayerDrop666M1.94×78.8BERT-of-Theseus666M1.94×81.2这个结果说明直接训练一个 6 层 BERT 效果明显不够。普通蒸馏能提升一些。逐步模块替换比这些 baseline 更有效。论文还在 GLUE test server 上报告结果BERT-base macro score 为80.0BERT-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.24 层模型为77.23 层模型进一步下降。对应加速分别为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 的性能。