Clover: 逐次知識を用いた回帰型軽量推測デコーディング

Clover: Regressive Lightweight Speculative Decoding with Sequential Knowledge

May 1, 2024

著者: Bin Xiao, Chunan Shi, Xiaonan Nie, Fan Yang, Xiangwei Deng, Lei Su, Weipeng Chen, Bin Cui

cs.AI

要旨

大規模言語モデル（LLM）は、自己回帰型デコードの要求と現代のGPUの設計とのミスマッチにより、効率性が低いという課題を抱えています。具体的には、数十億から数兆のパラメータをGPUキャッシュにロードするために限られたメモリ帯域幅を使用する必要がある一方で、実際に計算されるのは少量のトークンのみです。その結果、GPUは計算ではなくメモリ転送に大部分の時間を費やしています。最近では、並列デコードと呼ばれる一種の推測的デコードアルゴリズムが注目を集めており、生成効率の大幅な向上が実証されています。この手法は、大規模モデルに追加のデコードヘッドを導入し、複数の後続トークンを同時に予測し、これらの候補続きを単一のデコードステップで検証することを可能にします。しかし、このアプローチは事前学習で使用される次のトークン予測の訓練目的から逸脱しており、候補トークンのヒット率が低くなるという問題があります。本論文では、並列デコードプロセスに逐次的な知識を統合する新しい推測的デコードアルゴリズム「Clover」を提案します。この改良により、推測器のヒット率が向上し、全体の効率が向上します。Cloverは、事前に推測されたトークンからの逐次的な知識を回帰接続（Regressive Connection）を通じて伝達し、Attention Decoderを使用してこれらの推測トークンを統合します。さらに、Cloverは、次のトークン予測ではなく推測的生成の目的に合わせて隠れ状態を修正する拡張ブロック（Augmenting Block）を組み込んでいます。実験結果では、CloverがベースラインをBaichuan-Smallで最大91%、Baichuan-Largeで最大146%上回り、以前の最高性能を誇る手法であるMedusaをBaichuan-Smallで最大37%、Baichuan-Largeで最大57%上回る性能を示しました。

English

Large language models (LLMs) suffer from low efficiency as the mismatch between the requirement of auto-regressive decoding and the design of most contemporary GPUs. Specifically, billions to trillions of parameters must be loaded to the GPU cache through its limited memory bandwidth for computation, but only a small batch of tokens is actually computed. Consequently, the GPU spends most of its time on memory transfer instead of computation. Recently, parallel decoding, a type of speculative decoding algorithms, is becoming more popular and has demonstrated impressive efficiency improvement in generation. It introduces extra decoding heads to large models, enabling them to predict multiple subsequent tokens simultaneously and verify these candidate continuations in a single decoding step. However, this approach deviates from the training objective of next token prediction used during pre-training, resulting in a low hit rate for candidate tokens. In this paper, we propose a new speculative decoding algorithm, Clover, which integrates sequential knowledge into the parallel decoding process. This enhancement improves the hit rate of speculators and thus boosts the overall efficiency. Clover transmits the sequential knowledge from pre-speculated tokens via the Regressive Connection, then employs an Attention Decoder to integrate these speculated tokens. Additionally, Clover incorporates an Augmenting Block that modifies the hidden states to better align with the purpose of speculative generation rather than next token prediction. The experiment results demonstrate that Clover outperforms the baseline by up to 91% on Baichuan-Small and 146% on Baichuan-Large, respectively, and exceeds the performance of the previously top-performing method, Medusa, by up to 37% on Baichuan-Small and 57% on Baichuan-Large, respectively.