[Paper Review] RomeBERT: Robust Training of Multi-Exit BERT

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

Topic Modeling27 references6 citations

TL;DR

RomeBERT proposes a robust training framework for multi-exit BERT using gradient regularized self-distillation to alleviate performance imbalance between early and late exits. By jointly training the BERT backbone and multiple classifiers in a one-stage process, RomeBERT achieves superior speed-performance tradeoffs, outperforming DeeBERT and FastBERT across GLUE benchmarks with up to 65.2% inference time reduction while maintaining near-base model accuracy.

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.

Motivation & Objective

• To address the performance imbalance in multi-exit BERT, where early exits significantly underperform late exits.
• To eliminate the need for two-stage training by enabling joint end-to-end training of BERT backbone and multi-exit classifiers.
• To improve early exit performance without degrading final exit accuracy through robust training techniques.
• To achieve better efficiency-performance tradeoffs than existing methods like DeeBERT and FastBERT.
• To enable one-stage, joint training that avoids freezing the BERT backbone, thus preserving model expressiveness.

Proposed method

• Introduces gradient regularization (GR) inspired by GEM and Gradient Surgery to balance gradients across early and late exits and reduce training conflicts.
• Employs self-distillation (SD) to transfer soft label knowledge from the final exit to earlier exits, promoting consistency and improving early exit performance.
• Performs one-stage joint training of the BERT backbone and all multi-exit classifiers, avoiding the two-stage fine-tuning of DeeBERT.
• Uses cross-entropy loss for final exit and distillation loss for early exits, with GR applied to stabilize gradient updates.
• Applies soft label distillation from the final classifier to all intermediate classifiers to align predictions across exit layers.
• Employs a novel gradient regularization technique that constrains conflicting gradients during backpropagation, especially critical in multi-exit settings.

Experimental results

Research questions

• RQ1 Can self-distillation improve the performance of early exits in multi-exit BERT without degrading final exit accuracy?
• RQ2 Does one-stage joint training of the BERT backbone and multi-exit classifiers outperform two-stage fine-tuning with frozen BERT?
• RQ3 Can gradient regularization effectively reduce gradient conflicts and improve training stability in multi-exit BERT?
• RQ4 How does RomeBERT compare to DeeBERT and FastBERT in terms of inference efficiency and accuracy tradeoffs across diverse NLP tasks?
• RQ5 To what extent does the exit-layer distribution shift toward earlier layers in RomeBERT compared to prior methods?

Key findings

• On the RTE dataset, RomeBERT achieves 69.5% accuracy—0.1% better than BERT-base—while reducing expected running time by 10.6%, outperforming DeeBERT and DeeBERT+SD in both accuracy and efficiency.
• On QNLI, RomeBERT attains 88.7% accuracy with only 34.1% of the expected running time, compared to DeeBERT’s 15.2% higher time cost and 87.1% accuracy.
• On QQP, RomeBERT achieves 70.8% F1-score with 65.2% time saving compared to BERT-base, while DeeBERT requires 20.1% more time to match the same performance.
• When constrained to 34.1% running time, RomeBERT maintains 88.7% accuracy, whereas DeeBERT drops to 82.4% and DeeBERT+SD to 81.7%.
• RomeBERT’s exit-layer distribution is significantly more skewed toward earlier layers (e.g., 1st layer in 60% of cases on SST-2), indicating stronger early exit capability than DeeBERT and DeeBERT+SD.
• Ablation studies confirm that gradient regularization improves performance across all layers, especially on RTE and QNLI, with up to 20% performance gain on early layers for QQP and MNLI.