[论文解读] RomeBERT: Robust Training of Multi-Exit BERT

Shijie Geng, Peng Gao+2 more|arXiv (Cornell University)|Jan 24, 2021

Topic Modeling参考文献 27被引用 6

一句话总结

RomeBERT 提出了一种针对多出口 BERT 的鲁棒训练框架，通过梯度正则化自蒸馏缓解早期出口与晚期出口之间的性能不平衡。通过在单阶段过程中联合训练 BERT 主干网络和多个分类器，RomeBERT 实现了更优的速度-性能权衡，在 GLUE 基准测试中优于 DeeBERT 和 FastBERT，推理时间最高减少 65.2%，同时保持接近基线模型的准确率。

ABSTRACT

BERT has achieved superior performances on Natural Language Understanding (NLU) tasks. However, BERT possesses a large number of parameters and demands certain resources to deploy. For acceleration, Dynamic Early Exiting for BERT (DeeBERT) has been proposed recently, which incorporates multiple exits and adopts a dynamic early-exit mechanism to ensure efficient inference. While obtaining an efficiency-performance tradeoff, the performances of early exits in multi-exit BERT are significantly worse than late exits. In this paper, we leverage gradient regularized self-distillation for RObust training of Multi-Exit BERT (RomeBERT), which can effectively solve the performance imbalance problem between early and late exits. Moreover, the proposed RomeBERT adopts a one-stage joint training strategy for multi-exits and the BERT backbone while DeeBERT needs two stages that require more training time. Extensive experiments on GLUE datasets are performed to demonstrate the superiority of our approach. Our code is available at https://github.com/romebert/RomeBERT.

研究动机与目标

• 为解决多出口 BERT 中早期出口显著劣于晚期出口的性能不平衡问题。
• 通过实现 BERT 主干网络与多出口分类器的端到端联合训练，消除两阶段训练的需求。
• 通过稳健的训练技术提升早期出口性能，同时不损害最终出口的准确率。
• 实现优于现有方法（如 DeeBERT 和 FastBERT）的效率-性能权衡。
• 实现无需冻结 BERT 主干网络的一阶段联合训练，从而保留模型表达能力。

提出的方法

• 引入受 GEM 和梯度手术启发的梯度正则化（GR），以平衡早期与晚期出口之间的梯度，减少训练冲突。
• 采用自蒸馏（SD）将最终出口的软标签知识传递至早期出口，促进一致性并提升早期出口性能。
• 执行 BERT 主干网络与所有多出口分类器的一阶段联合训练，避免 DeeBERT 的两阶段微调过程。
• 对最终出口使用交叉熵损失，对早期出口使用蒸馏损失，同时应用 GR 以稳定梯度更新。
• 从最终分类器向所有中间分类器进行软标签蒸馏，以对齐各出口层的预测结果。
• 提出一种新颖的梯度正则化技术，在反向传播过程中约束冲突梯度，尤其在多出口设置中至关重要。

实验结果

研究问题

• RQ1 自蒸馏是否能在不损害最终出口准确率的前提下提升多出口 BERT 中早期出口的性能？
• RQ2 BERT 主干网络与多出口分类器的一阶段联合训练是否优于冻结 BERT 的两阶段微调？
• RQ3 梯度正则化是否能有效减少梯度冲突并提升多出口 BERT 的训练稳定性？
• RQ4 在多样化的自然语言处理任务中，RomeBERT 与 DeeBERT 和 FastBERT 相比，在推理效率与准确率权衡方面表现如何？
• RQ5 与先前方法相比，RomeBERT 中出口层的分布向更早层偏移的程度如何？

主要发现

• 在 RTE 数据集上，RomeBERT 达到 69.5% 的准确率——比 BERT-base 高 0.1%，同时将期望运行时间减少 10.6%，在准确率与效率方面均优于 DeeBERT 和 DeeBERT+SD。
• 在 QNLI 上，RomeBERT 实现 88.7% 的准确率，仅需 34.1% 的期望运行时间，而 DeeBERT 的时间成本高出 15.2%，准确率为 87.1%。
• 在 QQP 上，RomeBERT 的 F1 得分为 70.8%，相比 BERT-base 节省 65.2% 时间，而 DeeBERT 需要多花费 20.1% 时间才能达到相同性能。
• 当运行时间限制在 34.1% 时，RomeBERT 维持 88.7% 的准确率，而 DeeBERT 下降至 82.4%，DeeBERT+SD 降至 81.7%。
• RomeBERT 的出口层分布显著更倾向于早期层（例如在 SST-2 上 60% 的情况下为第一层），表明其早期出口能力明显优于 DeeBERT 和 DeeBERT+SD。
• 消融实验表明，梯度正则化在所有层上均提升性能，尤其在 RTE 和 QNLI 上效果显著，对 QQP 和 MNLI 的早期层性能最高提升达 20%。