LECO: Improving Early Exiting via Learned Exits and Comparison-based Exiting Mechanism
Jingfan Zhang, Ming Ren Tan, Pengyu Dai, Wei Zhu · 2023
Recently, dynamic early exiting has attracted much attention since it can accelerate the inference speed of pre-trained models (PTMs).However, previous work on early exiting has neglected the intermediate exits' architectural designs.In this work, we propose a novel framework, Learned Exits and COmparison-based early exiting (LECO) to improve PTMs' early exiting performances.First, to fully uncover the potentials of multi-exit BERT, we design a novel search space for intermediate exits and employ the idea of differentiable neural architecture search (DNAS) to design proper exit architectures for different intermediate layers automatically.Second, we propose a simpleyet-effective comparison-based early exiting mechanism (COBEE), which can help PTMs achieve better performance and speedup tradeoffs.Extensive experiments show that our LECO achieves the SOTA performances for multi-exit BERT training and dynamic early exiting.