[Paper Review] Accelerating BERT Inference for Sequence Labeling via Early-Exit

Xiaonan Li, Yunfan Shao+4 more|arXiv (Cornell University)|May 28, 2021

Topic Modeling39 references4 citations

TL;DR

This paper proposes SENTEE and TOKEE, two early-exit mechanisms to accelerate BERT inference for sequence labeling tasks. TOKEE enables token-level early exits using a window-based uncertainty criterion and self-sampling fine-tuning, achieving up to 75% inference cost reduction with minimal performance drop, outperforming DistilBERT at equivalent speed-up ratios.

ABSTRACT

Both performance and efficiency are crucial factors for sequence labeling tasks in many real-world scenarios. Although the pre-trained models (PTMs) have significantly improved the performance of various sequence labeling tasks, their computational cost is expensive. To alleviate this problem, we extend the recent successful early-exit mechanism to accelerate the inference of PTMs for sequence labeling tasks. However, existing early-exit mechanisms are specifically designed for sequence-level tasks, rather than sequence labeling. In this paper, we first propose a simple extension of sentence-level early-exit for sequence labeling tasks. To further reduce the computational cost, we also propose a token-level early-exit mechanism that allows partial tokens to exit early at different layers. Considering the local dependency inherent in sequence labeling, we employed a window-based criterion to decide for a token whether or not to exit. The token-level early-exit brings the gap between training and inference, so we introduce an extra self-sampling fine-tuning stage to alleviate it. The extensive experiments on three popular sequence labeling tasks show that our approach can save up to 66%-75% inference cost with minimal performance degradation. Compared with competitive compressed models such as DistilBERT, our approach can achieve better performance under the same speed-up ratios of 2X, 3X, and 4X.

Motivation & Objective

• Address the high computational cost of BERT in sequence labeling tasks, especially in real-time or mobile scenarios.
• Extend early-exit mechanisms—previously used for sequence-level tasks like text classification—to token-level sequence labeling, where per-token confidence is required.
• Design a fine-grained, dynamic computational cost allocation strategy that allows individual tokens to exit early based on local context.
• Mitigate the train-inference discrepancy introduced by halting and copying representations of exited tokens during inference.
• Improve efficiency without sacrificing performance, outperforming distilled models like DistilBERT under equivalent speed-up ratios.

Proposed method

• Propose SENTEE, a sentence-level early-exit mechanism that allows all tokens in a sequence to exit together when the maximum token uncertainty falls below a threshold.
• Introduce TOKEE, a token-level early-exit mechanism that enables individual tokens to exit at different layers based on their confidence, improving computational efficiency.
• Apply a window-based uncertainty criterion to consider local context when deciding whether a token can exit, capturing sequential dependencies.
• Implement a self-sampling fine-tuning stage that simulates early exits during training by sampling halting layers based on uncertainty, reducing the gap between training and inference.
• Use a halting-and-copy mechanism to propagate representations of exited tokens to higher layers without re-computation, preserving model structure.
• Integrate the early-exit mechanism into BERT’s standard encoder-decoder framework, maintaining compatibility with standard sequence labeling heads.

Experimental results

Research questions

• RQ1 Can early-exit mechanisms designed for sequence-level tasks be effectively adapted to sequence labeling, where predictions are per-token and dependencies exist across tokens?
• RQ2 Does token-level early-exit with a window-based uncertainty criterion lead to better computational efficiency than sentence-level early-exit in long sequences?
• RQ3 How can the train-inference discrepancy caused by halting and copying representations be effectively mitigated during fine-tuning?
• RQ4 What is the performance-cost trade-off of TOKEE compared to distilled models like DistilBERT under varying speed-up ratios?
• RQ5 Can fine-grained, dynamic computational allocation via TOKEE maintain high accuracy while significantly reducing inference cost?

Key findings

• TOKEE achieves up to 75% reduction in inference cost on sequence labeling tasks with minimal performance degradation, outperforming DistilBERT at equivalent speed-up ratios.
• SENTEE achieves up to 66% inference cost reduction, but its speedup ratio diminishes with increasing sequence length due to the need for all tokens to exit together.
• TOKEE maintains higher performance at high speed-up ratios (e.g., 4×) compared to SENTEE, which reduces its effective depth to only 2.5 layers on average under 4× speed-up.
• Self-sampling fine-tuning significantly reduces the performance drop at high speed-up ratios, outperforming random sampling and no fine-tuning.
• The layer distribution in TOKEE shows that more tokens exit earlier as speed-up increases, with a few still using deeper layers, enabling fine-grained cost allocation.
• The window-based uncertainty criterion improves robustness by considering local context, leading to more reliable early-exit decisions than relying solely on individual token confidence.