支援デコーディング

Assisted decodingは、メインモデルがトークンを確定する前に、ヘルパーが候補トークンを提案できるようにすることで、テキスト生成を高速化します。メインモデルは、1回のフォワードパスで候補トークンを検証します。ヘルパーは高速かつ低コストであり、メインモデルによるより高コストなフォワードパスを数十回分置き換えることができます。

このガイドでは、Transformersにおける支援デコーディング（assisted decoding）の手法について説明します。

投機的デコーディング

Speculative decodingでは、より小さなアシスタントモデルが候補トークンをドラフトします。メインモデルはこれらトークンを1回のパスで検証します。検証されたトークンは最終出力に含まれ、拒否されたトークンは標準的なサンプリングをトリガーします。メインモデルが高コストなフォワードパスをより少ない回数で実行できるため、生成が高速化されます。

この手法は、アシスタントモデルがメインモデルよりもはるかに小さく、同じトークナイザーを使用する場合に最も効果的です。Speculative decodingは貪欲探索とサンプリングをサポートしますが、バッチ入力には対応していません。

assistant_model を generate() に渡します。トークンの検証に失敗した場合に再サンプリングを行うには、do_sample=True を設定します。

貪欲探索

サンプリング

プロンプト検索デコーディング

Prompt lookup decodingにはアシスタントモデルは不要です。プロンプト内の重複するn-gramを検出することで、候補トークンを提案します。一致が見つからない場合は、通常の自己回帰デコーディングにフォールバックします。要約や翻訳などの入力に依存するタスクでは、候補トークンがソーステキストの局所的なパターンを反映することが多いため、この手法が適しています。

prompt_lookup_num_tokens を generate() に渡します。これにより、アルゴリズムが繰り返しパターンを検出した際に、プロンプトの前方から何トークン分のコピーを試みるかが設定されます。

貪欲デコーディング

サンプリング

セルフ推測デコーディング

Self-speculative decodingでは、モデルの中間層をアシスタントとして使用し、候補トークンを提案します。提案が一致した場合、モデルは早期に終了し、残りの層がトークンを検証または修正します。

すべてが1つのモデルであるため、重みとキャッシュが共有され、追加のメモリオーバーヘッドなしに速度が向上します。この手法は、中間層からの早期終了ロジツをサポートするように訓練されたモデルでのみ機能します。

assistant_early_exit を generate() に渡して、終了レイヤーを設定します。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("facebook/layerskip-llama3.2-1B")
model = AutoModelForCausalLM.from_pretrained("facebook/layerskip-llama3.2-1B", dtype="auto")
inputs = tokenizer("Hugging Face is an open-source company", return_tensors="pt")

outputs = model.generate(**inputs, assistant_early_exit=4, do_sample=False, max_new_tokens=20)
tokenizer.batch_decode(outputs, skip_special_tokens=True)

マルチトークン予測（MTP）レイヤー

Multi-token prediction (MTP) レイヤーは、メインモデル自身のチェックポイントから由来する追加の予測レイヤーを使用して候補トークンを生成します。MTP レイヤーは、メインモデルの埋め込みと出力ヘッドを再利用します。各 MTP レイヤーは 1 つのトークンを生成するため、ステップごとの候補数はチェックポイントの MTP レイヤー数と一致します。メインモデルは、1 回のフォワードパスでこれらの生成結果を検証します。

MTPは、DeepSeek-V3やGLM-4.5のようにMTPレイヤーで学習されたチェックポイントでのみ動作します。これらのチェックポイントにはMTPの重みが含まれており、config.json内の設定ではnum_nextn_predict_layersとしてnum_mtp_layersが設定されています。

use_mtp=True を generate() に渡します。この引数は、モデルに MTP レイヤーがない場合にエラーを発生させます。投機的デコーディングと同様に、MTP は貪欲探索とサンプリングをサポートしますが、バッチ入力はサポートされていません。

ここでは、MTPとは、DeepSeek-V3やGLM-4.5のように、ターゲットのチェックポイントに含まれる予測レイヤーを指します。Gemma 4は、MTPデコーディングのために別のアシスタントチェックポイントを使用します。assistant_modelとして渡してください（Multi-token prediction drafterを参照）。

text-generation パイプラインと generate() は、どちらも use_mtp=True を受け付けます。

from transformers import pipeline

pipe = pipeline(
    "text-generation",
    model="zai-org/GLM-4.5-Air",
    device_map="auto",
)

outputs = pipe(
    "Hugging Face is an open-source company",
    use_mtp=True,
    do_sample=False,
    max_new_tokens=50,
    add_special_tokens=True,
    return_full_text=True,
)

[output["generated_text"] for output in outputs]

代わりに、次のように推論することもできます。

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("zai-org/GLM-4.5-Air")
model = AutoModelForCausalLM.from_pretrained("zai-org/GLM-4.5-Air", device_map="auto")
inputs = tokenizer("Hugging Face is an open-source company", return_tensors="pt").to(model.device)

outputs = model.generate(**inputs, use_mtp=True, do_sample=False, max_new_tokens=50)
tokenizer.batch_decode(outputs, skip_special_tokens=True)

マルチトークン予測ドラフター

MTPドラフターは、別のチェックポイントから読み込まれる軽量なアシスタントモデルです。ターゲットモデルが検証するための候補トークンを提案します。ドラフターは、ターゲットモデルの入力埋め込みを使用し、そのKVキャッシュにアテンションを適用します。ターゲットモデルの隠れ状態から開始し、自己回帰的にトークンをドラフトします。次の各トークンについて、投影された隠れ状態をドラフターにフィードバックします。ターゲットモデルは、speculative decodingと同様に、1回のフォワードパスでドラフトされたトークンを検証します。

ターゲットのチェックポイントの一部であり、その出力ヘッドを共有するMTPレイヤーとは異なり、MTPドラフターは独立したアシスタントモデルです。ターゲットの入力埋め込みとKVキャッシュを使用してトークンを提案します。

from transformers import AutoProcessor, AutoModelForCausalLM

processor = AutoProcessor.from_pretrained("google/gemma-4-E4B-it")
target_model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-4-E4B-it",
    dtype="auto",
    device_map="auto",

)

assistant_model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-4-E4B-it-assistant",
    dtype="auto",
    device_map="auto",
)

messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "Write a short joke about saving RAM."},
]

inputs = processor.apply_chat_template(
    messages,
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
    add_generation_prompt=True,
    enable_thinking=False
).to(target_model.device)
input_len = inputs["input_ids"].shape[-1]

outputs = target_model.generate(
    **inputs,
    assistant_model=assistant_model,
    max_new_tokens=256,
)
response = processor.decode(outputs[0][input_len:], skip_special_tokens=False)

processor.parse_response(response)

汎用的な支援デコーディング

Universal assisted decoding (UAD)は、メインモデルとアシスタントモデルのトークナイザーが異なる場合でも、speculative decodingを可能にします。これにより、任意の小型アシスタントモデルをメインモデルと組み合わせることができます。候補トークンは再エンコードされ、アルゴリズムが最長共通部分列を計算することで、継続部分が整列した状態に保たれます。

UADを有効にするには、tokenizer、assistant_tokenizer、assistant_modelをgenerate()に渡してください。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

assistant_tokenizer = AutoTokenizer.from_pretrained("double7/vicuna-68m")
tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b")
model = AutoModelForCausalLM.from_pretrained("google/gemma-2-9b", dtype="auto")
assistant_model = AutoModelForCausalLM.from_pretrained("double7/vicuna-68m", dtype="auto")
inputs = tokenizer("Hugging Face is an open-source company", return_tensors="pt")

outputs = model.generate(**inputs, assistant_model=assistant_model, tokenizer=tokenizer, assistant_tokenizer=assistant_tokenizer)
tokenizer.batch_decode(outputs, skip_special_tokens=True)
'Hugging Face is an open-source company that is dedicated to creating a better world through technology.'

静的アンサンブル検証

標準的な推測デコーディングはロスレスです。出力分布がターゲットモデルと完全に一致することを保証します。静的アンサンブル検証では、ターゲット分布とドラフト分布の混合分布に対して検証を行うことでこの制約を緩和し、出力分布が正確なターゲット分布ではなく混合分布になるというコストを払って、受入率を向上させます（推論が高速化されます）。推奨される初期値は0.7です。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM-1.7B")
model = AutoModelForCausalLM.from_pretrained("HuggingFaceTB/SmolLM-1.7B", dtype="auto")
assistant_model = AutoModelForCausalLM.from_pretrained("HuggingFaceTB/SmolLM-135M", dtype="auto")
inputs = tokenizer("Hugging Face is an open-source company", return_tensors="pt")

outputs = model.generate(**inputs, assistant_model=assistant_model, assistant_ensemble_weight=0.7, do_sample=True)
tokenizer.batch_decode(outputs, skip_special_tokens=True)

これにはアシスタントモデルからの候補ロジットが必要であり、プロンプト検索デコーディングではサポートされていません。

リソース

テキスト生成のレイテンシとアシスト生成に関する詳細な背景については、Assisted Generation: a new direction toward low-latency text generationブログ記事をご覧ください。

Update on GitHub