# RomeBERT: Robust Training of Multi-Exit BERT

arXiv: https://arxiv.org/abs/2101.09755 · HTML: https://mlsanity.com/paper/2101.09755/

## Key facts

- Cited by 19 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 28 references that are in the map.
- First posted on arXiv 24 Jan 2021.

## Builds on (most cited first, 26 of 26)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [Non-local Neural Networks](https://mlsanity.com/paper/1711.07971/) (2017)
- [Gradient Episodic Memory for Continual Learning](https://mlsanity.com/paper/1706.08840/) (2017)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [LXMERT: Learning Cross-Modality Encoder Representations from Transformers](https://mlsanity.com/paper/1908.07490/) (2019)
- [Self-training with Noisy Student improves ImageNet classification](https://mlsanity.com/paper/1911.04252/) (2019)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [Deep Networks with Stochastic Depth](https://mlsanity.com/paper/1603.09382/) (2016)
- [Born Again Neural Networks](https://mlsanity.com/paper/1805.04770/) (2018)
- [Adaptive Computation Time for Recurrent Neural Networks](https://mlsanity.com/paper/1603.08983/) (2016)
- [DeeBERT: Dynamic Early Exiting for Accelerating BERT Inference](https://mlsanity.com/paper/2004.12993/) (2020)
- [BERT Loses Patience: Fast and Robust Inference with Early Exit](https://mlsanity.com/paper/2006.04152/) (2020)
- [FastBERT: a Self-distilling BERT with Adaptive Inference Time](https://mlsanity.com/paper/2004.02178/) (2020)
- [Fast Convergence of DETR with Spatially Modulated Co-Attention](https://mlsanity.com/paper/2101.07448/) (2021)
- [Fast Transformers with Clustered Attention](https://mlsanity.com/paper/2007.04825/) (2020)
- [Spatially Adaptive Computation Time for Residual Networks](https://mlsanity.com/paper/1612.02297/) (2016)
- [End-to-End Object Detection with Adaptive Clustering Transformer](https://mlsanity.com/paper/2011.09315/) (2020)
- [Knowledge Distillation in Generations: More Tolerant Teachers Educate Better Students](https://mlsanity.com/paper/1805.05551/) (2018)
- [Dynamic Graph Representation Learning for Video Dialog via Multi-Modal Shuffled Transformers](https://mlsanity.com/paper/2007.03848/) (2020)
- [Multi-Pass Transformer for Machine Translation](https://mlsanity.com/paper/2009.11382/) (2020)
- [Character Matters: Video Story Understanding with Character-Aware Relations](https://mlsanity.com/paper/2005.08646/) (2020)

## Built on by (most cited first, 19 of 19)

- [AutoMix: Automatically Mixing Language Models](https://mlsanity.com/paper/2310.12963/) (2023)
- [Let's Sample Step by Step: Adaptive-Consistency for Efficient Reasoning and Coding with LLMs](https://mlsanity.com/paper/2305.11860/) (2023)
- [Consistent Accelerated Inference via Confident Adaptive Transformers](https://mlsanity.com/paper/2104.08803/) (2021)
- [Model Compression and Efficient Inference for Large Language Models: A Survey](https://mlsanity.com/paper/2402.09748/) (2024)
- [A Survey on Green Deep Learning](https://mlsanity.com/paper/2111.05193/) (2021)
- [A Survey on Model Compression and Acceleration for Pretrained Language Models](https://mlsanity.com/paper/2202.07105/) (2022)
- [Uncertainty in Natural Language Processing: Sources, Quantification, and Applications](https://mlsanity.com/paper/2306.04459/) (2023)
- [Tracing Representation Progression: Analyzing and Enhancing Layer-Wise Similarity](https://mlsanity.com/paper/2406.14479/) (2024)
- [A Survey on Dynamic Neural Networks for Natural Language Processing](https://mlsanity.com/paper/2202.07101/) (2022)
- [Agreement-Based Cascading for Efficient Inference](https://mlsanity.com/paper/2407.02348/) (2024)
- [ConsistentEE: A Consistent and Hardness-Guided Early Exiting Method for Accelerating Language Models Inference](https://mlsanity.com/paper/2312.11882/) (2023)
- [DEED: Dynamic Early Exit on Decoder for Accelerating Encoder-Decoder Transformer Models](https://mlsanity.com/paper/2311.08623/) (2023)
- [BEEM: Boosting Performance of Early Exit DNNs using Multi-Exit Classifiers as Experts](https://mlsanity.com/paper/2502.00745/) (2025)
- [Training dynamic models using early exits for automatic speech recognition on resource-constrained devices](https://mlsanity.com/paper/2309.09546/) (2023)
- [Scalable Transformers for Neural Machine Translation](https://mlsanity.com/paper/2106.02242/) (2021)
- [FLOWGEN: Fast and slow graph generation](https://mlsanity.com/paper/2207.07656/) (2022)
- [Multimodal Adaptive Inference for Document Image Classification with Anytime Early Exiting](https://mlsanity.com/paper/2405.12705/) (2024)
- [Dynamic Self-Distillation via Previous Mini-batches for Fine-tuning Small Language Models](https://mlsanity.com/paper/2411.16991/) (2024)
- [Two-Stage Voting for Robust and Efficient Suicide Risk Detection on Social Media](https://mlsanity.com/paper/2510.08365/) (2025)

## Abstract (arXiv)

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.