LMStudioに実装された「投機的デコーディング」を試す（LLM）

Aru

大規模言語モデル（LLM）を手軽に実行できるツール「LMStudio（0.3.10）」にSpeculative Decoding(投機的デコーディング)が実装されました。この技術に興味があったので少し調べて見ました。この記事では、Speculative Decodingについての解説と、LMStudioでの実際に動かした結果について解説します。

Contents

投機的デコーディング
LMStudioで投機的デコーディングを実行
- 投機的デコーディングを設定する
- 実行してみる
まとめ

投機的デコーディング

投機的デコーディングとは？

投機的デコーディング（Speculative Decoding）とは、小さなモデル（ドラフトモデル）を使って予測を行い、その結果を大きなモデル（ターゲットモデル）で検証することで、大規模言語モデル（LLM）の推論速度を高速化させる手法です。

上記の説明ではピンとこないと思いますが、わかりやすくいうと「高速な小さいモデルで予測を行って、その結果を重いモデルで検証することで計算負荷を下げる」という方法です。

動作の概要

投機的でコーディングでは、「ドラフトモデル」と「ターゲットモデル」の２つのモデルを組み合わせて２段階で処理を行います。

ドラフトモデル
軽量なモデルです。次に生成されるトークンを推論します。こちらでn個のトークンをまとめて生成します
ターゲットモデル
ドラフトモデルで生成されたn個のトークンに対して評価します。ドラフトモデルの出力が正しい場合は、それを採用します。間違っていた場合は修正します。

ポイントは、ドラフトモデルでn個のトークンを生成し、n個のトークンに対してターゲットモデルでまとめて評価するという点です。このようにすることで、重いターゲットモデルによる推論の回数を減らすことが可能となります。

LLMでは、入力された文章に続く単語（トークン）を予測することにより応答を生成しています（詳しくは以下の記事を参照してください）。

どの程度高速化する？

仮に、ドラフトモデルの速度を1、ターゲットモデルの速度を3（3倍遅い）、n=5として考えてみます。

ドラフトモデルを使う場合：　$5 \times 1 + 3 = 8$

ターゲットモデルだけで推論する場合：　$5 \times 3 = 15$

上記から分かるように$\frac{15}{8} = 1.875$と、約1.9倍高速になることがわかります。ただ、全ての推論がOKとなるわけではないので、推論がOKとなる確率p=0.3で考えます。

すると式は以下のようになります。

ドラフトモデルを使う場合：
$$
\begin{eqnarray}
p*(5 \times 1 + 3) + (1-p)(5 \times 3)
&=& 0.3*8 + 0.7*15\\
&=& 12.9
\end{eqnarray}
$$

この場合は、約1.16倍になります。このように、ドラフトモデルの予測がある程度正しければ速度が向上することになります。

ドラフトモデルの予測が全然正しくない場合は、逆に速度が落ちます。また、ドラフトモデルとターゲットモデルの処理速度にそれほど差がない場合も効果が小さいです

より詳しい解説

予測のより詳しい説明です。

ドラフトモデルがN個のトークンを生成（Drafting）

例えば、T_1, T_2, ..., T_N というトークンを予測。

ターゲットモデルが各トークンの確率を計算（Verification）

ターゲットモデルが P(T_1 | context), P(T_2 | context, T_1), …, P(T_N | context, T_1, ..., T_{N-1}) を計算する。

閾値またはランキングで評価

ドラフトモデルの予測 Pd(T_i | context, ...) とターゲットモデルの確率 P(T_i | context, ...) を比較し、以下のいずれかの基準で判断する：
- 確率差のしきい値を設定（例：| P(T_i) - Pd(T_i) | < ε なら合格）
- 確率比を設定（例：P(T_i) / Pd(T_i) > 0.9 なら合格）
- トップk候補内にあるか確認（本番モデルのトップkの出力内に含まれているか）

不合格なら修正

もし T_i の確率が本番モデルと大きくズレていた場合、そのトークン以降を修正し、通常のデコーディング（例えばGreedy, Beam Searchなど）を実行する。

基本的には上記のようなアルゴリズムで処理されると思います。

具体例として”Hello, my name is”という入力に対して予測する例を考えます。確率差が0.03以下を採用とした場合でドラフトモデルが”Hello, my name is John”と予測したとします。この場合は、ターゲットとの確率差が0.02なので採用されます。一方、”Hello, my name is Banana”と予測した場合は、確率差が閾値を超えるため不採用にし、ターゲットモデルで推論します。