Transformers documentation
Gemma 4 Unified Assistant
このモデルは2026年6月3日にHugging Face Transformersに貢献されました。
Gemma 4 Unified Assistant
概要
Gemma 4 Unified Assistant は、Multi-Token Prediction (MTP) 手法および関連する候補生成器を用いて Gemma 4 Unified モデルの投機的デコーディングを可能にする、小型のテキスト専用モデルです。Gemma 4 12B モデルの IT バリアント向けに事前学習済みモデルが提供されています。
詳細については、Gemma4 Assistantを参照してください。アーキテクチャおよび概念の面では、ベースモデルとの共通点と相違点は同じです。
- モデル全体で KV シェアリングを使用します。この技術は Gemma 3n で初めて導入されたもので、アシスタントがサポートするターゲットモデルによって生成された KV キャッシュを再利用できるようにします。これにより、アシスタントは pre-fille フェーズを完全にスキップでき、フォワードパス中のアテンション計算を大幅に削減できます。
position_idsの値は一定です。KV キャッシュは共有されており、アシスタントがキャッシュを更新する手段がないため、アシスタントは同じ position ID からすべてのトークンを予測します。- 入力には、埋め込みと隠れ状態の連結が含まれます。静的な KV キャッシュおよび
position_idsへの対応のため、モデルは、ターゲットモデルによって最後に確認されたトークンのembeddingとhidden_statesの連結を入力として受け取り、nn.Linear変換によってアシスタントモデルの空間に射影します。最後に確認されたトークンの定義は、アシストデコーディングループの進行に伴って変化します。プリフィル後の最初のドラフトトークンでは、最後に確認されたトークンはプロンプトの最後のトークンになります。その後のドラフトステップでは、最後に確認されたトークンは、アシスタントによって生成された最後のトークン（ドラフトラウンド内）またはターゲットモデルによって承認された最後のトークン（ドラフトラウンド間）になります。 - Cross-attentionは、ターゲットモデルのコンテキストを最大限に活用するために使用されます。Cross-attentionにより、アシスタントが生成したクエリ状態が、ターゲットモデルからの共有KVキャッシュの値にアテンションを向けることができ、その結果、アシスタントはドラフティングラウンドごとに、より多くのドラフトトークンを正確に予測できるようになります。
使用例
以下の例では、PipelineまたはAutoModelクラスを使用して、画像に基づいてテキストを生成する方法を示しています。
<hfoptions id="usage"> <hfoption id="Pipeline">import torch
from transformers import pipeline
pipeline = pipeline(
task="image-text-to-text",
model="google/gemma-4-12B-it",
assistant_model="google/gemma-4-12B-it-assistant",
)
pipeline(
images="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg",
text="<|image|>\n\nWhat is shown in this image?"
)import torch
from transformers import AutoProcessor, AutoModelForImageTextToText
model = AutoModelForImageTextToText.from_pretrained(
"google/gemma-4-12B-it",
dtype=torch.bfloat16,
device_map="auto",
)
assistant_model = AutoModelForCausalLM.from_pretrained(
"google/gemma-4-12B-it-assistant",
dtype=torch.bfloat16,
device_map="auto",
)
processor = AutoProcessor.from_pretrained(
"google/gemma-4-12B-it",
padding_side="left"
)
messages = [
{
"role": "user", "content": [
{"type": "image", "url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"},
{"type": "text", "text": "What is shown in this image?"},
]
},
]
inputs = processor.apply_chat_template(
messages,
tokenize=True,
return_dict=True,
return_tensors="pt",
add_generation_prompt=True,
).to(model.device)
input_len = inputs["input_ids"].shape[-1]
output = model.generate(**inputs, max_new_tokens=50, assistant_model=assistant_model)
print(processor.decode(output[0][input_len:], skip_special_tokens=True))Gemma4UnifiedAssistantConfig
class transformers.Gemma4UnifiedAssistantConfig
< source >( transformers_version: str | None = Nonearchitectures: list[str] | None = Noneoutput_hidden_states: bool | None = Falsereturn_dict: bool | None = Truedtype: str | torch.dtype | None = Nonechunk_size_feed_forward: int = 0is_encoder_decoder: bool = Falseid2label: dict[int, str] | dict[str, str] | None = Nonelabel2id: dict[str, int] | dict[str, str] | None = Noneproblem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] | None = Nonetext_config: transformers.configuration_utils.PreTrainedConfig | dict[str, typing.Any] | None = Nonebackbone_hidden_size: int = 3840use_ordered_embeddings: bool = Falsenum_centroids: int = 2048centroid_intermediate_top_k: int = 32tie_word_embeddings: bool = True )
パラメータ
- text_config (
Union[~configuration_utils.PreTrainedConfig, dict[str, Any]]、省略可) — テキストバックボーンの設定オブジェクトまたは辞書です。 - backbone_hidden_size (
int、デフォルト値は 3840) — このアシスタントの学習に使用された、対象の統合モデルの隠れ層サイズです。 - use_ordered_embeddings (
bool、デフォルトはFalse) —Trueの場合、最適に順序付けられた埋め込みテーブルを使用します。この場合、再順序付けが必要です。 - num_centroids (
int、デフォルト値は 2048) — centroid ベースの masked embedder における centroid の総数。 - centroid_intermediate_top_k (
int、デフォルト値は 32)— 推論時に選択されるアクティブなセントロイドの数。 - tie_word_embeddings (
bool、任意、デフォルトはTrue) — モデルのtied_weights_keysマッピングに従って重み埋め込みを共有するかどうかを指定します。
これは、Gemma4 Unified AssistantModel の設定を保存するための設定クラスです。指定された引数に基づいて Gemma4 Unified Assistant モデルをインスタンス化し、モデルのアーキテクチャを定義するために使用されます。デフォルト値で設定をインスタンス化すると、google/gemma-4-12b-it と同様の設定が得られます。
設定オブジェクトはPreTrainedConfigを継承し、モデルの出力を制御するために使用できます。詳細については、PreTrainedConfigのドキュメントを参照してください。
例:
>>> from transformers import (
... Gemma4UnifiedAssistantConfig,
... Gemma4UnifiedAssistantForCausalLM,
... Gemma4UnifiedTextConfig,
... )
>>> # Initializing a Gemma 4 Unified Text config for the assistant.
>>> text_config = Gemma4UnifiedTextConfig(...)
>>> # Initializing a Gemma 4 Unified Assistant config.
>>> configuration = Gemma4UnifiedAssistantConfig(text_config=text_config)
>>> # Initializing a model from the configuration.
>>> model = Gemma4UnifiedAssistantForCausalLM(configuration)
>>> # Accessing the model configuration
>>> configuration = model.configGemma4UnifiedAssistantForCausalLM
クラス transformers.Gemma4UnifiedAssistantForCausalLM
< source >( config: Gemma4UnifiedAssistantConfig )
パラメータ
- config (Gemma4UnifiedAssistantConfig) — モデルのパラメータをすべて含むモデル設定クラスです。設定ファイルで初期化しても、モデルに関連する重みは読み込まれず、設定のみが読み込まれます。モデルの重みをロードするには、from_pretrained() メソッドを参照してください。
Gemma 4 におけるマルチトークン予測ベースの支援デコーディング用のモデル。
このモデルは PreTrainedModel を継承しています。すべてのモデルに実装されている汎用的なメソッド（ダウンロードや保存、入力埋め込みのリサイズ、ヘッドのプルーニングなど）については、スーパークラスのドキュメントを参照してください。
このモデルは、PyTorchのtorch.nn.Moduleのサブクラスでもあります。 通常のPyTorch Moduleとして使用し、一般的な使用方法や動作に関する詳細については、PyTorchのドキュメントを参照してください。
アシスト対象モデル用のアテンションマスクを準備します。この場合、shared_kv_statesが過去キャッシュとして機能します。
因果性を考慮するために、双方向マスクを使用します
q_len == 1というエッジケースについては、いずれにせよフルアテンションとして動作するため、違いはありません。- SWA は、
q_idx=1かつ、kv>=1- kv軸を反転させることで、未来の視点から過去の視点へと切り替えます
- 位置不変パディングを考慮するために、初期作成の前にベースアテンションマスクを反転させます
forward
< source >( input_ids: typing.Optional[torch.Tensor] = Noneinputs_embeds: typing.Optional[torch.Tensor] = Noneposition_ids: typing.Optional[torch.LongTensor] = Noneattention_mask: dict[str, torch.Tensor] | None = Noneshared_kv_states: dict[str, tuple[torch.Tensor, torch.Tensor]] | None = Noneuse_cache: bool | None = None**kwargs: *<class 'transformers.utils.generic.TransformersKwargs'> )
パラメータ
- input_ids (
torch.Tensorの形状(batch_size, sequence_length), 、任意) — 語彙における入力シーケンスのトークンインデックスです。デフォルトでは、パディングは無視されます。インデックスはAutoTokenizerを使用して取得できます。詳細については、PreTrainedTokenizer.encode()およびPreTrainedTokenizer.call()を参照してください。
- inputs_embeds (
torch.Tensor、形状(batch_size, sequence_length, hidden_size)、任意) —input_idsを渡す代わりに、埋め込み表現を直接渡すこともできます。これは、モデル内部の埋め込み参照行列よりも、input_idsのインデックスを対応するベクトルに変換する方法をより細かく制御したい場合に有用です。 - position_ids (
torch.LongTensorの形状(batch_size, sequence_length), 、任意) — 各入力シーケンスのトークンが位置埋め込み内で占める位置のインデックス。範囲[0, config.n_positions - 1]. - attention_mask (
dict[str, torch.Tensor]形状(batch_size, sequence_length), 、任意) — パディングトークンのインデックスに対するアテンションを回避するためのマスク。[0, 1]:- マスクされていないトークンに対して1
- マスクされたトークンに対しては 0 です。
- shared_kv_states (
dict[str, torch.Tensor、形状(batch_size, 1, q_len, kv_len)、任意) — このモデルの各layer_typeの最終層で計算された KV 値を含む辞書。 - use_cache (
bool、任意) —Trueに設定すると、past_key_valuesのキーバリュー状態が返され、デコーディングの高速化に利用できます（past_key_valuesを参照）。
Gemma4UnifiedAssistantForCausalLM の forward メソッドは、__call__ という特殊メソッドをオーバーライドします。
この関数内でフォワードパスの手順を定義する必要がありますが、その後、この関数ではなく
Moduleインスタンスを呼び出すべきです。前者は前処理と後処理の実行を適切に処理するのに対し、後者はそれらを黙って無視するためです。
例:
>>> from transformers import AutoTokenizer, Gemma4UnifiedAssistantForCausalLM, Gemma4UnifiedForConditionalGeneration
>>> model = Gemma4UnifiedForConditionalGeneration.from_pretrained("google/gemma-4-e2b-it")
>>> assistant_model = Gemma4UnifiedAssistantForCausalLM.from_pretrained("google/gemma-4-e2b-it-assistant")
>>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-4-e2b-it")
>>> prompt = "What is your favorite condiment?"
>>> inputs = tokenizer(prompt, return_tensors="pt")
>>> # Generate
>>> generate_ids = model.generate(inputs.input_ids, assistant_model=assistant_model, max_length=30)
>>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True)[0]
"What is your favorite condiment?"