[論文レビュー] Magic Pyramid: Accelerating Inference with Early Exiting and Token Pruning

Xuanli He, Iman Keivanloo+5 more|arXiv (Cornell University)|Oct 30, 2021

Topic Modeling参考文献 22被引用数 5

ひとこと要約

Magic Pyramid (MP) は、トークンプルーニングとイ早退出を包括的に組み合わせることで、幅方向および深さ方向の計算を削減し、BERTの推論を高速化する。最大11.95倍の高速化を達成し、精度低下は0.5%未満に抑えられ、最先端手法に比べ最大2.13倍の高速化を達成しながら、さまざまな入力長に対して安定した性能を維持する。

ABSTRACT

Pre-training and then fine-tuning large language models is commonly used to achieve state-of-the-art performance in natural language processing (NLP) tasks. However, most pre-trained models suffer from low inference speed. Deploying such large models to applications with latency constraints is challenging. In this work, we focus on accelerating the inference via conditional computations. To achieve this, we propose a novel idea, Magic Pyramid (MP), to reduce both width-wise and depth-wise computation via token pruning and early exiting for Transformer-based models, particularly BERT. The former manages to save the computation via removing non-salient tokens, while the latter can fulfill the computation reduction by terminating the inference early before reaching the final layer, if the exiting condition is met. Our empirical studies demonstrate that compared to previous state of arts, MP is not only able to achieve a speed-adjustable inference but also to surpass token pruning and early exiting by reducing up to 70% giga floating point operations (GFLOPs) with less than 0.5% accuracy drop. Token pruning and early exiting express distinctive preferences to sequences with different lengths. However, MP is capable of achieving an average of 8.06x speedup on two popular text classification tasks, regardless of the sizes of the inputs.

研究の動機と目的

• BERT などの大規模事前学習モデルにおける高い推論遅延の課題に取り組むこと、特にリアルタイムでのプロダクション環境における制約を想定する。
• 既存手法の限界を克服すること——トークンプルーニング（長文に対して有効）とイ早退出（短文に対して有効）は、それぞれのシーケンス長の極端な場面で性能が著しく低下する。
• 幅方向（プルーニング）と深さ方向（イ早退出）の計算削減を統合的に組み合わせた包括的なフレームワークを構築し、入力長に応じて一貫性があり、スケーラブルな推論を実現すること。
• 特に低遅延・高スループットなデプロイ環境において、顕著な計算削減を実現しながらも、高い精度を維持すること。

提案手法

• MP は、トークンプルーニングとイ早退出を統合した階層的推論パイプラインとして実装され、プルーニングで幅方向の計算を削減し、イ早退出で深さ方向の計算を削減する。
• 各層で注意機構の不確実性に基づき、非重要とされるトークンを識別・削除するための学習可能なしきい値機構を採用する。
• 各トランスフォーマーブロックにサブ分類器を接続し、信頼度が動的しきい値 τ を超えると、その時点で予測を早期に終了できる。τ はトレーニング中に最適化される。
• プルーニングとイ早退出の最適化において、精度と速度のトレードオフをバランスさせるために、温度スケーリングと損失重み付け（λ）を採用する。
• トレーニングプロセスでは、精度、GFLOPs、イ早退出確率を同時に最適化するマルチタスク損失関数を用い、知識蒸留と不確実性正則化を組み合わせる。
• フレームワークは、BERTベースのモデルにエンドツーエンドで適用され、入力長に応じて動的にトークンをプルーニングし、条件を満たすとイ早退出を行う。

実験結果

リサーチクエスチョン

• RQ1 トークンプルーニングとイ早退出を効果的に統合することで、さまざまな入力長にわたって一貫した推論高速化を達成できるか？
• RQ2 幅方向（プルーニング）と深さ方向（イ早退出）の計算削減の相乗効果が、個別手法に比べ優れた高速化を実現するか？
• RQ3 統合手法は、顕著な計算削減を実現しながらも、高い精度を維持できるか。特に低遅延のプロダクション環境において顕著か？
• RQ4 提案手法の性能は、さまざまなシーケンス長や自然言語処理タスクでどのように変動するか？

主な発見

• MP は、AG News および Yelp データセットで最大11.95倍の高速化を達成し、すべてのシーケンス長グループで FastBERT（イ早退出）や LTP（トークンプルーニング）を大きく上回る。
• 平均して、2つのテキスト分類タスクにおいて、入力長にかかわらず 8.06倍の高速化を達成し、多様な入力に対して一貫した性能を示す。
• BERT に対して最大70%の GFLOPs 削減を実現し、精度低下は0.5%未満。LTP や FastBERT を上回る、精度-効率のトレードオフ性能を達成する。
• 長文シーケンス（70トークン以上）では、Yelp で 8.25倍、AG News で 11.95倍の高速化を達成。FastBERT の 6.18倍および 8.84倍を上回る。
• τ = 0.8 の条件下でも、AG News で 11.95倍、Yelp で 10.10倍の高速化を達成し、積極的なイ早退出しきい値でも高いスケーラビリティを示す。
• AG News では BERT と同等またはわずかに高い精度（94.3%）を維持しながら、GFLOPs を 1.8（4.95倍の高速化）に削減。FastBERT の 2.3 GFLOPs（3.97倍の高速化）を上回る。