[論文レビュー] Speculative Decoding with Big Little Decoder

Sehoon Kim, Karttikeya Mangalam+5 more|arXiv (Cornell University)|Feb 15, 2023

Topic Modeling被引用数 8

ひとこと要約

BiLDは小型の自己回帰デコーダと大規模な非自己回帰デコーダを組み合わせ、フォールバックとロールバック方針でテキスト生成を加速する。品質の損失を最小限に抑える。

ABSTRACT

The recent emergence of Large Language Models based on the Transformer architecture has enabled dramatic advancements in the field of Natural Language Processing. However, these models have long inference latency, which limits their deployment and makes them prohibitively expensive for various real-time applications. The inference latency is further exacerbated by autoregressive generative tasks, as models need to run iteratively to generate tokens sequentially without leveraging token-level parallelization. To address this, we propose Big Little Decoder (BiLD), a framework that can improve inference efficiency and latency for a wide range of text generation applications. The BiLD framework contains two models with different sizes that collaboratively generate text. The small model runs autoregressively to generate text with a low inference cost, and the large model is only invoked occasionally to refine the small model's inaccurate predictions in a non-autoregressive manner. To coordinate the small and large models, BiLD introduces two simple yet effective policies: (1) the fallback policy that determines when to hand control over to the large model; and (2) the rollback policy that determines when the large model needs to correct the small model's inaccurate predictions. To evaluate our framework across different tasks and models, we apply BiLD to various text generation scenarios encompassing machine translation on IWSLT 2017 De-En and WMT 2014 De-En, and summarization on XSUM and CNN/DailyMail. On an NVIDIA T4 GPU, our framework achieves a speedup of up to 2.12x speedup with minimal generation quality degradation. Furthermore, our framework is fully plug-and-play and can be applied without any modifications in the training process or model architecture. Our code is open-sourced

研究の動機と目的

• autoregressive テキスト生成の推論遅延の低減を、品質の損失を最小限に抑えつつ動機づける。
• ベースモデルを再学習せずに、サイズの異なる2つのデコーダを調整するプラグアンドプレー型のフレームワークを導入する。
• 大規模モデルが介入すべき場合と予測を巻き戻すべき場合を決定する、単純なフォールバックおよびロールバックの方針を提案する。
• オープンソースの実装を用いた機械翻訳と要約のベンチマークで適用性を示す。

提案手法

• 小型モデルが自己回帰的にトークンを生成し、 大型モデルが非自己回帰的に予測を refine する二モデルデコード。
• フォールバック方針: 小型モデルの信頼度（max p_S）が閾値 alpha_FB を下回ると大型モデル推論を誘発。
• ロールバック方針: 大型モデルの分布が閾値 alpha_RB を超えて分岐する場合、初期の小型モデル予測を上書きし、後続トークンを巻き戻す。
• 予測の整合性を高めるため、キャリブレーションデータを用いて小型モデルと大型モデルの出力を整合させ、不要なロールバックを減らす。
• アルゴリズム1はエンドツーエンドの BiLD プロセスと、各デコードステップごとの方針チェックの詳細を示す。

実験結果

リサーチクエスチョン

• RQ1 小型の自己回帰モデルが大部分の生成を担い、時折の大型モデルの refined により遅延を削減できるか。
• RQ2 単純なフォールバックとロールバック方針は、品質の損失を許容しつつタスク全体で意味のあるスピードアップをもたらすか。
• RQ3 独立して訓練された小型・大型モデル間の整合性は BiLD の性能を改善するか。
• RQ4 BiLDを早期終了戦略と結びつけ、特化した学習パイプラインと競合できるか。
• RQ5 BiLDは全自己回帰デコーダと比較して翻訳と要約のベンチマークでどのような性能を示すか。

主な発見

Task (Model)	BLEU (IWSLT)	Speedup (IWSLT)	BLEU (WMT)	Speedup (WMT)	ROUGE-L (XSUM)	Speedup (XSUM)	ROUGE-L (CNN/DM)	Speedup (CNN/DM)
Vanilla Inference (large)	40.32	-	31.38	-	35.08	-	41.54	-
BiLD (Unaligned)	40.33	1.43x	31.28	1.34x	35.12	1.48x	41.44	1.71x
BiLD (Unaligned) Degraded	39.44	1.58x	30.47	1.43x	34.02	1.72x	40.57	2.05x
BiLD (Aligned)	40.24	1.62x	31.26	1.47x	35.05	1.50x	41.52	1.85x
BiLD (Aligned) Degraded	39.13	1.78x	30.33	1.70x	33.95	1.80x	40.96	2.12x

• BiLDは特定のタスクでエンドツーエンドの速度を最大2.12x向上させ、BLEU/ROUGE-Lの約1ポイントの劣化を示す。
• アンアライン BiLD はベンチマーク全体で平均1.50xの速度向上を達成し、一部タスクで品質の劣化なし; アライン BiLD は平均1.61xで最大1.85xの速度向上。
• 予測の整合性を高めることで、未整合 BiLD よりも速度向上と品質が改善された。
• アブレーションにより、ロールバックとフォールバックの各方針が品質と遅延の優位性を維持するために重要であることが示された。
• BiLDは早期終了シナリオへ拡張可能で、MTベンチマークで最小のBLEU損失で最大1.74xの速度向上を達成。