[论文解读] Accelerating BERT Inference for Sequence Labeling via Early-Exit

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

Topic Modeling参考文献 39被引用 4

一句话总结

本文提出 SENTEE 和 TOKEE 两种早期退出机制，用于加速序列标注任务中的 BERT 推理。TOKEE 通过基于窗口的不确定性准则与自采样微调，实现基于 token 的早期退出，最高可将推理成本降低 75%，性能损失极小，在同等加速比下优于 DistilBERT。

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.

研究动机与目标

• 解决 BERT 在序列标注任务中计算成本过高的问题，尤其是在实时或移动场景下。
• 将早期退出机制从此前用于序列级任务（如文本分类）扩展至基于 token 的序列标注任务，以满足对每个 token 置信度的需求。
• 设计一种细粒度、动态的计算成本分配策略，使每个 token 可根据局部上下文决定是否提前退出。
• 缓解推理过程中因停止并复制已退出 token 表示而引入的训练-推理差异。
• 在不牺牲性能的前提下提升效率，在同等加速比下优于如 DistilBERT 等蒸馏模型。

提出的方法

• 提出 SENTEE，一种句子级早期退出机制，当最大 token 不确定性低于阈值时，允许序列中所有 token 同时退出。
• 引入 TOKEE，一种基于 token 的早期退出机制，使不同 token 可在不同层基于其置信度提前退出，提升计算效率。
• 采用基于窗口的不确定性准则，在决定是否退出时考虑局部上下文，捕捉序列依赖关系。
• 实施自采样微调阶段，通过基于不确定性的采样方式模拟训练过程中的早期退出，缩小训练与推理之间的差距。
• 使用停止并复制机制，将已退出 token 的表示传播至高层，无需重新计算，同时保持模型结构。
• 将早期退出机制集成到 BERT 的标准编码器-解码器框架中，保持与标准序列标注头的兼容性。

实验结果

研究问题

• RQ1 专为序列级任务设计的早期退出机制能否有效适配到序列标注任务中？此类任务要求逐 token 预测且存在跨 token 依赖关系。
• RQ2 在长序列中，基于窗口的不确定性准则的 token 级早期退出是否比句子级早期退出具有更高的计算效率？
• RQ3 在微调过程中，如何有效缓解因停止并复制表示而引入的训练-推理差异？
• RQ4 在不同加速比下，TOKEE 与 DistilBERT 等蒸馏模型相比，其性能与成本之间的权衡如何？
• RQ5 通过 TOKEE 实现的细粒度、动态计算分配，能否在显著降低推理成本的同时保持高准确率？

主要发现

• TOKEE 在序列标注任务中最高可将推理成本降低 75%，性能损失极小，在同等加速比下优于 DistilBERT。
• SENTEE 最高可实现 66% 的推理成本降低，但随着序列长度增加，其加速比逐渐减弱，原因在于所有 token 必须同时退出。
• 在高加速比（如 4×）下，TOKEE 的性能显著优于 SENTEE，后者在 4× 加速比下平均有效深度仅剩 2.5 层。
• 自采样微调显著降低了高加速比下的性能下降，优于随机采样与无微调方法。
• TOKEE 的层分布显示，随着加速比提升，更多 token 选择在早期退出，少数仍使用深层网络，实现细粒度的成本分配。
• 基于窗口的不确定性准则通过考虑局部上下文，提升了鲁棒性，使早期退出决策比仅依赖单个 token 置信度更可靠。