[論文レビュー] TR-BERT: Dynamic Token Reduction for Accelerating BERT Inference

Deming Ye, Yankai Lin+2 more|arXiv (Cornell University)|May 25, 2021

Topic Modeling参考文献 48被引用数 7

ひとこと要約

TR-BERTは、強化学習を用いて各層でのタスク関連性に基づき、どのトークンを保持するかを学習することで、BERTの推論を高速化する動的トークン削減手法を提案する。最小限の精度低下で2–5倍の高速化を達成し、特に自己注意機構の計算量がシーケンス長に二次的に増加する長文処理タスクにおいて、選択的トークン伝搬により計算複雑度を低減する点で顕著な効果を示す。

ABSTRACT

Existing pre-trained language models (PLMs) are often computationally expensive in inference, making them impractical in various resource-limited real-world applications. To address this issue, we propose a dynamic token reduction approach to accelerate PLMs' inference, named TR-BERT, which could flexibly adapt the layer number of each token in inference to avoid redundant calculation. Specially, TR-BERT formulates the token reduction process as a multi-step token selection problem and automatically learns the selection strategy via reinforcement learning. The experimental results on several downstream NLP tasks show that TR-BERT is able to speed up BERT by 2-5 times to satisfy various performance demands. Moreover, TR-BERT can also achieve better performance with less computation in a suite of long-text tasks since its token-level layer number adaption greatly accelerates the self-attention operation in PLMs. The source code and experiment details of this paper can be obtained from https://github.com/thunlp/TR-BERT.

研究の動機と目的

• リソース制限のある環境における事前学習言語モデル（PLM）の高い推論遅延を解消すること。
• 固定層のプルーニングに起因する制限を克服し、複雑なタスクで性能が著しく低下するのを防ぐこと。
• 各トークンごとに適応的計算を可能にするために、深層層を通過させるべきトークンを動的に選択すること。
• シーケンス長に二次的に増加する自己注意機構の計算量が大きな課題となる長文処理シナリオでの効率性を向上させること。
• 実世界の展開に適した、オンデマンドで速度と精度のトレードオフを柔軟に制御できる仕組みを提供すること。

提案手法

• TR-BERTは、各トークンが各層で選択されるかスキップされるかを判断する多段階選択問題としてトークン削減を定式化する。
• タスク固有の重要度に基づき、動的トークン選択戦略を学習するため、強化学習フレームワークにポリシーネットワークを採用する。
• 推論の信頼度と選択トークン数に対するペナルティを組み合わせた報酬をRLエージェントが受信し、速度と精度のバランスを取る。
• 重要度が低いと判断されたトークンは早期にプルーニングされ、その後の計算が削減され、無駄な自己注意演算が回避される。
• 本手法はトークンレベルで動作し、層レベルのプルーニングとは異なり、各入力トークンごとの計算を細かく制御可能である。
• フレームワークはプラグアンドプレイであり、ベースとなるBERTモデルの再訓練が不要で、既存のPLMとも互換性がある。

実験結果

リサーチクエスチョン

• RQ1 動的でトークンレベルのプルーニングは、モデル性能を損なわせることなく推論速度を向上させることができるか？
• RQ2 タスク関連性に基づく動的トークン選択は、固定層または層単位のプルーニングと比較してどのように優れているか？
• RQ3 人為的アノテーションなしで、強化学習が効果的にトークンの重要度信号を学習できるか？
• RQ4 長文NLPタスクにおいて、トークンレベルのプルーニングはどの程度効率性を向上させるか？
• RQ5 本手法は多様な下流タスクにおいて、柔軟な速度-精度トレードオフを達成できるか？

主な発見

• TR-BERTは、目的の精度レベルに応じて11のNLPベンチマークで2–5倍の高速化を達成する。
• 質問応答タスクにおいて、特に固定プルーニングが失敗するような複雑な推論シナリオでは、ベースライン手法を上回る性能を示す。
• 長文処理タスク（例：1,024トークン入力）では、ヴァニラBERTと比較して、はるかに高い性能と著しい推論時間短縮を達成する。
• プルーニングされたシーケンスに対しても予測の信頼度が高く維持されており、意味的に重要なトークンの適切な選択が行われていることが示唆される。
• 人為的アノテーションなしで、強化学習ベースのトークン選択ポリシーは、質問語や感情を示すキーワードなどのタスク関連トークンを効果的に保持するよう学習する。
• 本手法は、テキスト分類、自然言語推論、質問応答を含む多様なNLPタスクにおいて、強固な性能を示す。