Transformers documentation

Gemma 4 Unified Assistant

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

このモデルは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?"
)
</hfoption> <hfoption id="AutoModel">
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))
</hfoption>

Gemma4UnifiedAssistantConfig

class transformers.Gemma4UnifiedAssistantConfig

< >

( 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.config

Gemma4UnifiedAssistantForCausalLM

クラス transformers.Gemma4UnifiedAssistantForCausalLM

< >

( config: Gemma4UnifiedAssistantConfig )

パラメータ

  • config (Gemma4UnifiedAssistantConfig) — モデルのパラメータをすべて含むモデル設定クラスです。設定ファイルで初期化しても、モデルに関連する重みは読み込まれず、設定のみが読み込まれます。モデルの重みをロードするには、from_pretrained() メソッドを参照してください。

Gemma 4 におけるマルチトークン予測ベースの支援デコーディング用のモデル。

このモデルは PreTrainedModel を継承しています。すべてのモデルに実装されている汎用的なメソッド(ダウンロードや保存、入力埋め込みのリサイズ、ヘッドのプルーニングなど)については、スーパークラスのドキュメントを参照してください。

このモデルは、PyTorchのtorch.nn.Moduleのサブクラスでもあります。 通常のPyTorch Moduleとして使用し、一般的な使用方法や動作に関する詳細については、PyTorchのドキュメントを参照してください。

create_attention_masks

< >

( inputs_embedsattention_maskshared_kv_states )

アシスト対象モデル用のアテンションマスクを準備します。この場合、shared_kv_statesが過去キャッシュとして機能します。

因果性を考慮するために、双方向マスクを使用します

  • q_len == 1 というエッジケースについては、いずれにせよフルアテンションとして動作するため、違いはありません。
  • SWA は、 q_idx=1 かつ、 kv>=1
    • kv軸を反転させることで、未来の視点から過去の視点へと切り替えます
    • 位置不変パディングを考慮するために、初期作成の前にベースアテンションマスクを反転させます

forward

< >

( 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()を参照してください。

    入力IDとは何ですか?

  • 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].

    ポジションIDとは?

  • 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?"
Update on GitHub