Judge Decoding: Faster Speculative Sampling Requires Going Beyond Model Alignment

要約

大規模な言語モデル（LLMS）のパフォーマンスは、その基礎となるサイズに密接にリンクされており、増え続けるネットワークにつながり、したがって推論が遅くなります。
投機的デコードは、自己回帰生成を加速する手法として提案されており、高速ドラフトモデルを活用して候補トークンを提案し、ターゲットモデルの下での可能性に基づいて並行して検証されます。
このアプローチでは、ターゲット出力を再現することが保証されていますが、実質的なペナルティが発生します。客観的に有効な継続を表していても、多くの高品質のドラフトトークンが拒否されます。
実際、GPT-4Oなどの強力なドラフトモデルや、人間のテキストでさえ、標準検証スキームの下で高い受け入れ率を達成できないことを示しています。
これにより、ドラフトとターゲットの整合だけに依存する場合に早期拒絶が圧倒的に可能になるため、現在の投機的解読方法のスピードアップの可能性が大幅に制限されます。
したがって、次の質問をします。検証を適応させて、正しいが非整合した返信を認識することができますか？
この目的のために、LLMが汎用性の高い方法で回答を評価できることを実証したLLM-A-A-Judge Frameworkからインスピレーションを引き出しました。
埋め込みの上にコンパクトモジュールをトレーニングして現在の継続の「判断」を生成することにより、ターゲットモデルで同じ機能を引き出すようにデータセットを慎重に設計します。
私たちは、8B/405B-JudgeがLlama-405Bで9倍のスピードアップを達成しながら、幅広いベンチマークで品質を維持しながら、Llama-3.1ファミリーで戦略を紹介しています。
これらの利点は、最適化された推論フレームワークでも存在します。この方法では、2および8 H100でそれぞれ8B/70B-Judgeで最大141トークン、8B/405Bで129トークン/sに達します。

要約(オリジナル)

The performance of large language models (LLMs) is closely linked to their underlying size, leading to ever-growing networks and hence slower inference. Speculative decoding has been proposed as a technique to accelerate autoregressive generation, leveraging a fast draft model to propose candidate tokens, which are then verified in parallel based on their likelihood under the target model. While this approach guarantees to reproduce the target output, it incurs a substantial penalty: many high-quality draft tokens are rejected, even when they represent objectively valid continuations. Indeed, we show that even powerful draft models such as GPT-4o, as well as human text cannot achieve high acceptance rates under the standard verification scheme. This severely limits the speedup potential of current speculative decoding methods, as an early rejection becomes overwhelmingly likely when solely relying on alignment of draft and target. We thus ask the following question: Can we adapt verification to recognize correct, but non-aligned replies? To this end, we draw inspiration from the LLM-as-a-judge framework, which demonstrated that LLMs are able to rate answers in a versatile way. We carefully design a dataset to elicit the same capability in the target model by training a compact module on top of the embeddings to produce “judgements’ of the current continuation. We showcase our strategy on the Llama-3.1 family, where our 8b/405B-Judge achieves a speedup of 9x over Llama-405B, while maintaining its quality on a large range of benchmarks. These benefits remain present even in optimized inference frameworks, where our method reaches up to 141 tokens/s for 8B/70B-Judge and 129 tokens/s for 8B/405B on 2 and 8 H100s respectively.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google