[論文レビュー] 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

ひとこと要約

本稿では、序列ラベル付けタスクにおけるBERT推論を高速化する2つの早期終了機構であるSENTEEとTOKEEを提案する。TOKEEは、窓ベースの不確実性基準と自己サンプリング微調整を用いて、トークンレベルの早期終了を可能にし、最小限の性能低下で最大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の高い計算コストを軽減すること、特にリアルタイムまたはモバイル環境での応用を想定すること。
• 従来、文レベルのタスク（例：テキスト分類）に用いられていた早期終了機構を、個々のトークンに対する予測とトークン間の依存関係を必要とする、トークンレベルの序列ラベル付けに拡張すること。
• 局所的文脈に基づいて個々のトークンが異なる層で早期に終了できる、細粒度で動的な計算コスト割り当て戦略を設計すること。
• 推論時に終了したトークンの表現をコピーするプロセスによって生じるトレーニングと推論の乖離を緩和すること。
• 性能を犠牲にすることなく効率性を向上させ、同等の高速化比においてDistilBERTのような蒸留モデルを上回ること。

提案手法

• すべてのトークンが同時に終了できる文レベルの早期終了機構であるSENTEEを提案し、最大トークン不確実性がしきい値未満に下がった場合に全トークンが終了するようにする。
• 個々のトークンが自信がある段階で異なる層で終了できる、トークンレベルの早期終了機構であるTOKEEを導入し、計算効率を向上させる。
• 局所的文脈を考慮する窓ベースの不確実性基準を適用し、逐次的依存関係を捉えることで、早期終了の意思決定をより適切にする。
• トレーニング中に早期終了をシミュレートする自己サンプリング微調整ステージを実装し、不確実性に基づいて停止層をサンプリングすることで、トレーニングと推論のギャップを縮小する。
• 終了したトークンの表現を再計算せずに上位層に伝搬する「停止・コピー」機構を採用し、モデル構造を維持する。
• 早期終了機構をBERTの標準的なエンコーダ・デコーダフレームワークに統合し、標準的な序列ラベル付けヘッドとの互換性を保つ。

実験結果

リサーチクエスチョン

• RQ1 文レベルタスクに設計された早期終了機構を、予測が個々のトークンに分かれており、トークン間の依存関係が存在する序列ラベル付けタスクに効果的に適応できるか？
• RQ2 窓ベースの不確実性基準を用いたトークンレベルの早期終了は、長文シーケンスにおいて文レベルの早期終了よりも優れた計算効率を達成できるか？
• RQ3 推論時に停止および表現のコピーが行われることで生じるトレーニングと推論の乖離を、微調整段階で効果的に緩和できるか？
• RQ4 異なる高速化比において、TOKEEの性能とコストのトレードオフは、DistilBERTのような蒸留モデルと比較してどうなるか？
• RQ5 TOKEEによる細粒度で動的な計算コスト割り当ては、推論コストを顕著に削減しつつ、高い精度を維持できるか？

主な発見

• TOKEEは、序列ラベル付けタスクで最大75%の推論コスト削減を達成し、性能の低下は最小限に抑えられ、同等の高速化比においてDistilBERTを上回る性能を発揮する。
• SENTEEは最大66%の推論コスト削減を達成するが、全トークンが同時に終了する必要があるため、シーケンス長が延びるにつれて高速化比の向上が鈍化する。
• TOKEEは高加速比（例：4倍）でもSENTEEよりも高い性能を維持する。4倍加速時、SENTEEの平均有効深さはたった2.5層にまで低下するが、TOKEEはより多くの層を活用できる。
• 自己サンプリング微調整は、高加速比における性能低下を顕著に軽減し、ランダムサンプリングや微調整なしの手法を上回る性能を発揮する。
• TOKEEの層分布から、加速比が高くなるほどより多くのトークンが早期に終了することが明らかになり、一部のトークンは深い層を活用するため、細粒度なコスト割り当てが可能である。
• 窓ベースの不確実性基準は、局所的文脈を考慮することで、個々のトークンの信頼性に依存するのみの手法よりも、より頑健な早期終了意思決定を可能にする。