BanditSpec: バンディットアルゴリズムによる適応的投機的デコード

BanditSpec: Adaptive Speculative Decoding via Bandit Algorithms

May 21, 2025

著者: Yunlong Hou, Fengzhuo Zhang, Cunxiao Du, Xuan Zhang, Jiachun Pan, Tianyu Pang, Chao Du, Vincent Y. F. Tan, Zhuoran Yang

cs.AI

要旨

推測的デコーディングは、大規模言語モデル（LLMs）の推論を加速しつつ、その優れたテキスト生成性能を維持するための一般的な手法として登場している。従来の手法では、プレフィックストークンに関係なく固定された推測的デコーディング設定を採用するか、またはドラフトモデルをオフラインまたはオンラインでトレーニングしてコンテキストに適合させていた。本論文では、テキストが生成される際に推測的デコーディングのハイパーパラメータ設定を適応的に選択するための、トレーニング不要のオンライン学習フレームワークを提案する。まず、このハイパーパラメータ選択問題を多腕バンディット問題として定式化し、一般的な推測的デコーディングフレームワークであるBanditSpecを提供する。さらに、2つのバンディットベースのハイパーパラメータ選択アルゴリズム、UCBSpecとEXP3Specを設計し、新たな量である停止時間リグレットの観点から分析する。このリグレットを確率的および敵対的報酬設定の下で上限付ける。情報理論的な不可能性の結果を導出することで、UCBSpecのリグレット性能が普遍的な定数まで最適であることが示される。最後に、LLaMA3とQwen2を用いた広範な実証実験により、提案アルゴリズムが既存の手法と比較して有効であり、多様な入力プロンプトを伴うシミュレートされた実生活のLLMサービスシナリオにおいて、スループットがオラクル最適ハイパーパラメータに近いことが実証された。

English

Speculative decoding has emerged as a popular method to accelerate the inference of Large Language Models (LLMs) while retaining their superior text generation performance. Previous methods either adopt a fixed speculative decoding configuration regardless of the prefix tokens, or train draft models in an offline or online manner to align them with the context. This paper proposes a training-free online learning framework to adaptively choose the configuration of the hyperparameters for speculative decoding as text is being generated. We first formulate this hyperparameter selection problem as a Multi-Armed Bandit problem and provide a general speculative decoding framework BanditSpec. Furthermore, two bandit-based hyperparameter selection algorithms, UCBSpec and EXP3Spec, are designed and analyzed in terms of a novel quantity, the stopping time regret. We upper bound this regret under both stochastic and adversarial reward settings. By deriving an information-theoretic impossibility result, it is shown that the regret performance of UCBSpec is optimal up to universal constants. Finally, extensive empirical experiments with LLaMA3 and Qwen2 demonstrate that our algorithms are effective compared to existing methods, and the throughput is close to the oracle best hyperparameter in simulated real-life LLM serving scenarios with diverse input prompts.