[論文レビュー] Online Speculative Decoding

Xiaoxuan Liu, Lanxiang Hu+5 more|arXiv (Cornell University)|Oct 11, 2023

Topic ModelingComputer Science被引用数 3

ひとこと要約

本論文は、LLMサービングクラスタ内の空き計算リソースを活用して、リアルタイムのユーザークエリデータに基づき継続的に微調整を行うオンライン推論予測（speculative decoding）の手法を提案する。オンライン知識蒸留を適用することで、ドキュメントモデルはクエリの分布にリアルタイムで適応し、トークン受容率を10–65%向上させ、推論遅延を1.22倍から3.06倍まで低減する。

ABSTRACT

Speculative decoding is a pivotal technique to accelerate the inference of large language models (LLMs) by employing a smaller draft model to predict the target model's outputs. However, its efficacy can be limited due to the low predictive accuracy of the draft model, particularly when faced with diverse text inputs and a significant capability gap between the draft and target models. We introduce online speculative decoding to address this challenge. The main idea is to continuously update the (multiple) draft model(s) on observed user query data. Adapting to query distribution mitigates the shifts between the training distribution of the draft model and the query distribution, enabling the draft model to more accurately predict the target model's outputs. We develop a prototype of online speculative decoding based on knowledge distillation and evaluate it using both synthetic and real query data. The results show a substantial increase in the token acceptance rate by 0.1 to 0.65, bringing 1.42x to 2.17x latency reduction. Our code is available at https://github.com/LiuXiaoxuanPKU/OSD.

研究の動機と目的

• 分布シフトや能力格差のため、推論予測におけるドキュメントモデルの推論精度が低い問題に対処する。
• LLMサービングクラスタ内の未利用計算リソース（余剰FLOPs）を活用し、ドキュメントモデルの継続的再訓練を実現する。
• 推論コストを増加させることなく、進化するクエリ分布にリアルタイムでドキュメントモデルを適応可能にする。
• ユーザークエリデータに基づくオンライン知識蒸留を通じて、オンラインLLMサービングにおけるトークン受容率の向上と遅延低減を実現する。
• 限定的なデータでも、クエリ固有のデータに基づくオンライン再訓練が、静的でオフラインのファインチューニングを上回ることを示す。

提案手法

• LLMサービングクラスタ内の空き計算容量（余剰FLOPs）を活用し、ドキュメントモデルのオンライン再訓練を実施する。
• 一般化知識蒸留（GKD）を適用し、ドキュメントモデルとターゲットモデルの出力差異に基づいてドキュメントモデルを更新する。
• 誤った予測のバッファを維持し、余剰FLOPsが利用可能である場合にのみ勾配更新を実行することで、オーバーヘッドを最小限に抑える。
• 特定の言語やトピッククラスタに特化した複数のドキュメントモデルを訓練し、分布固有の精度を向上させる。
• オンライン蒸留中に、ドキュメントモデルとターゲットモデルの分布間の距離メトリクスを損失関数として計算する。
• 機会的トレーニングにより、推論遅延への影響を最小限に抑えながら、継続的なモデル適応を可能にする。

Figure 1: Online speculative decoding overview. For each prompt, the draft model suggests multiple tokens in a single step. The target model then verifies these tokens, accepting some and rejecting others. If the student proposes incorrect tokens, both the draft and target distributions are stored i

実験結果

リサーチクエスチョン

• RQ1 リアルユーザークエリデータを用いたドキュメントモデルのオンライン再訓練は、推論予測における推論精度を向上させることができるか？
• RQ2 クエリデータのわずかな割合しか利用できない状況で、オンライン推論予測はオフラインファインチューニングと比べてどのように性能を発揮するか？
• RQ3 LLMサービングクラスタ内の余剰FLOPsをドキュメントモデルのトレーニングに再利用できる範囲はどの程度で、推論パフォーマンスに影響を与えないか？
• RQ4 クエリ分布の変化にリアルタイムで適応することは、継続的なトークン受容率の向上をもたらすか？
• RQ5 特定の言語やトピックに特化した複数のドキュメントモデルは、1つの汎用ドキュメントモデルを上回る性能を発揮できるか？

主な発見

• オンライン推論予測は、多様なデータセットでトークン受容率を0.1から0.65（10–65%）まで向上させ、推論効率を顕著に改善した。
• 推論遅延は1.22倍から3.06倍まで低減され、特に実世界のLMSYS-chatデータで最も顕著な改善が得られた。
• 実際のLLMチャット会話において、OSDは1言語あたり2,000件未満のデータポイントでも、受容率を0.1から0.2まで向上させた。
• トピック固有の会話では、全トピックで受容率が0.6を超えた。特にソーシャルおよびコンピュータサイエンスの議論では、0.8に近い値に達した。
• SELECT、FROM、WHERE、+ などの高頻度トークンは、オンライン蒸留後、精度と再現率の両方で顕著な向上を示した。
• OSDは、クエリデータの70–100%をファインチューニングに使用したモデルと同等またはそれ以上の性能を発揮したが、使用データ量はごくわずかであった。

Figure 3: Online acceptance rate ( $\alpha$ ) across different datasets. The x-axis represents the number of records that OSD has processed. Alpha is averaged over the most recent 50 records.