# Gemma 4 アシスタント

## 概要

Gemma 4 Assistant は、Multi-Token Prediction (MTP) 手法および関連する候補生成器を用いた推測的デコーディングを Gemma 4 モデルで有効にする、小型のテキスト専用モデルです。Gemma 4 E2B、E4B、31B および 26B-A4B (MoE) モデルの IT バリアント用に事前学習済みモデルが提供されています。

アーキテクチャの観点では、Gemma 4 Assistant は他の Gemma 4 モデルと同様に [Gemma4TextModel](/docs/transformers/main/ja/model_doc/gemma4#transformers.Gemma4TextModel) バックボーンを共有していますが、いくつかの重要な点で異なります。

*   **モデル全体で KV シェアリングを使用します**。この技術は [Gemma 3n](./gemma3n) で初めて導入されたもので、アシスタントがサポートするターゲットモデルによって生成された KV キャッシュを再利用できるようにします。これにより、アシスタントは pre-fille フェーズを完全にスキップでき、フォワードパス中のアテンション計算を大幅に削減できます。
*   **`position_ids` の値は一定です**。KV キャッシュは共有されており、アシスタントがキャッシュを更新する手段がないため、アシスタントは同じ position ID からすべてのトークンを予測します。
*   **入力には、埋め込みと隠れ状態の連結が含まれます**。静的な KV キャッシュおよび`position_ids`への対応のため、モデルは、ターゲットモデルによって最後に確認されたトークンの`embedding`と`hidden_states`の連結を入力として受け取り、`nn.Linear`変換によってアシスタントモデルの空間へ射影します。アシストデコーディングループの進行に伴い、「最後に確認されたトークン」の定義は変化します。プリフィルの後に起草される最初のトークンについては、最後に確認されたトークンはプロンプトの最後のトークンになります。その後の起草ステップでは、最後に確認されたトークンは、アシスタントによって生成された最後のトークン（起草ラウンド内）またはターゲットモデルによって承認された最後のトークン（起草ラウンド間）になります。
*   **Cross-attentionは、ターゲットモデルのコンテキストを最大限に活用するために使用されます**。Cross-attentionにより、アシスタントが生成したクエリ状態が、ターゲットモデルからの共有KVキャッシュの値にアテンションを向けることができ、その結果、アシスタントはドラフティングラウンドごとに、より多くのドラフトトークンを正確に予測することが可能になります。

[Gemma 4](https://huggingface.co/collections/google/gemma-4)リリースの下に、Gemma 4 Assistantのすべてのオリジナルチェックポイントを見つけることができます。

## 使用例

以下の例では、[Pipeline](/docs/transformers/main/ja/main_classes/pipelines#transformers.Pipeline)または[AutoModel](/docs/transformers/main/ja/model_doc/auto#transformers.AutoModel)クラスを使用して、画像に基づいてテキストを生成する方法を示しています。

```py
import torch
from transformers import pipeline

pipeline = pipeline(
    task="image-text-to-text",
    model="google/gemma-4-E2B-it",
    assistant_model="google/gemma-4-E2B-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?"
)
```

```py
import torch
from transformers import AutoProcessor, AutoModelForImageTextToText

model = AutoModelForImageTextToText.from_pretrained(
    "google/gemma-4-E2B-it",
    dtype=torch.bfloat16,
    device_map="auto",
)
assistant_model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-4-E2B-it-assistant",
    dtype=torch.bfloat16,
    device_map="auto",
)

processor = AutoProcessor.from_pretrained(
    "google/gemma-4-E2B-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))
```

## Gemma4AssistantConfig[[transformers.Gemma4AssistantConfig]]

#### transformers.Gemma4AssistantConfig[[transformers.Gemma4AssistantConfig]]

```python
transformers.Gemma4AssistantConfig(transformers_version: str | None = None, architectures: list[str] | None = None, output_hidden_states: bool | None = False, return_dict: bool | None = True, dtype: str | torch.dtype | None = None, chunk_size_feed_forward: int = 0, is_encoder_decoder: bool = False, id2label: dict[int, str] | dict[str, str] | None = None, label2id: dict[str, int] | dict[str, str] | None = None, problem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] | None = None, text_config: transformers.configuration_utils.PreTrainedConfig | dict[str, typing.Any] | None = None, backbone_hidden_size: int = 1536, use_ordered_embeddings: bool = False, num_centroids: int = 2048, centroid_intermediate_top_k: int = 32, tie_word_embeddings: bool = True)
```

[ソース](https://github.com/huggingface/transformers/blob/main/src/transformers/models/gemma4_assistant/configuration_gemma4_assistant.py#L29)

**パラメータ:**

text_config (`Union[~configuration_utils.PreTrainedConfig, dict[str, Any]]`、*optional*) : テキストバックボーンの config オブジェクトまたは辞書。

backbone_hidden_size (`int`、デフォルトは 1536) : このアシスタントが学習されたターゲットモデルの隠れ層サイズです。

use_ordered_embeddings (`bool`、デフォルトは False) : True の場合、アシスタントモデルの性能を最適化するために並べ替えられた埋め込みテーブルを使用します。このテーブルは、メインモデルと整合させるために再並べ替えが必要です。

num_centroids (`int`、デフォルトは 2048) : セントロイドの総数。

centroid_intermediate_top_k (`int`、デフォルトは 32) : アクティブなセントロイドの数。

tie_word_embeddings (`bool`、*任意*、デフォルトは`True`) : モデルの`tied_weights_keys`マッピングに従って重み埋め込みを共有するかどうか。

これは、Gemma4 AssistantModel の設定を保存するための設定クラスです。指定された引数に基づいて Gemma4 Assistant モデルをインスタンス化し、モデルのアーキテクチャを定義するために使用されます。デフォルト値で設定をインスタンス化すると、[google/gemma-4-e2b-it](https://huggingface.co/google/gemma-4-e2b-it) と同様の設定が得られます。

設定オブジェクトは [PreTrainedConfig](/docs/transformers/main/ja/main_classes/configuration#transformers.PreTrainedConfig) を継承し、モデルの出力を制御するために使用できます。詳細については、[PreTrainedConfig](/docs/transformers/main/ja/main_classes/configuration#transformers.PreTrainedConfig) のドキュメントを参照してください。

例:

```python
>>> from transformers import (
>>>     Gemma4AssistantConfig,
>>>     Gemma4AssistantForCausalLM,
>>>     Gemma4TextConfig,
>>> )

>>> # Initializing a Gemma 4 Text config similar to google/gemma-4-e2b-it-assistant.
>>> text_config = Gemma4TextConfig(...)

>>> # Initializing a Gemma 4 Assistant config similar to google/gemma-4-e2b-it-assistant.
>>> configuration = Gemma4AssistantConfig(text_config)

>>> # Initializing a model from the google/gemma-4-e2b-it-assistant configuration.
>>> model = Gemma4AssistantForCausalLM(configuration)

>>> # Accessing the model configuration
>>> configuration = model.config
```

## Gemma4AssistantForCausalLM[[transformers.Gemma4AssistantForCausalLM]]

#### transformers.Gemma4AssistantForCausalLM[[transformers.Gemma4AssistantForCausalLM]]

```python
transformers.Gemma4AssistantForCausalLM(config: Gemma4AssistantConfig)
```

[ソース](https://github.com/huggingface/transformers/blob/main/src/transformers/models/gemma4_assistant/modeling_gemma4_assistant.py#L109)

**パラメータ:**

config ([Gemma4AssistantConfig](/docs/transformers/main/ja/model_doc/gemma4_assistant#transformers.Gemma4AssistantConfig)) : モデルのパラメータをすべて含むモデル設定クラスです。設定ファイルで初期化しても、モデルに関連する重みは読み込まれず、設定のみが読み込まれます。モデルの重みをロードするには、[from_pretrained()](/docs/transformers/main/ja/main_classes/model#transformers.PreTrainedModel.from_pretrained) メソッドを参照してください。

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

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

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

#### create_attention_masks[[transformers.Gemma4AssistantForCausalLM.create_attention_masks]]

```python
create_attention_masks(inputs_embeds, attention_mask, shared_kv_states)
```

[ソース](https://github.com/huggingface/transformers/blob/main/src/transformers/models/gemma4_assistant/modeling_gemma4_assistant.py#L204)

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

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

#### forward[[transformers.Gemma4AssistantForCausalLM.forward]]

```python
forward(input_ids: typing.Optional[torch.Tensor] = None, inputs_embeds: typing.Optional[torch.Tensor] = None, position_ids: typing.Optional[torch.LongTensor] = None, attention_mask: dict[str, torch.Tensor] | None = None, shared_kv_states: dict[str, tuple[torch.Tensor, torch.Tensor]] | None = None, use_cache: bool | None = None, **kwargs: *<class 'transformers.utils.generic.TransformersKwargs'>)
```

[ソース](https://github.com/huggingface/transformers/blob/main/src/transformers/models/gemma4_assistant/modeling_gemma4_assistant.py#L133)

**パラメータ:**

input_ids (`torch.Tensor`、shape `(batch_size, sequence_length)`、*任意*) : 語彙における入力シーケンスのトークンインデックスです。デフォルトでは、パディングは無視されます。インデックスは [AutoTokenizer](/docs/transformers/main/ja/model_doc/auto#transformers.AutoTokenizer) を使用して取得できます。詳細については、[PreTrainedTokenizer.encode()](/docs/transformers/main/ja/internal/tokenization_utils#transformers.PreTrainedTokenizerBase.encode) および [PreTrainedTokenizer.__call__()](/docs/transformers/main/ja/internal/tokenization_utils#transformers.PreTrainedTokenizerBase.__call__) を参照してください。 [入力IDとは何ですか？](../glossary#input-ids)

inputs_embeds (`torch.Tensor`、shape `(batch_size, sequence_length, hidden_size)`、*任意*) : `input_ids` を渡す代わりに、埋め込み表現を直接渡すこともできます。これは、モデル内部の埋め込み参照行列よりも、`input_ids` のインデックスを対応するベクトルに変換する方法をより細かく制御したい場合に有用です。

position_ids (`torch.LongTensor`、shape `(batch_size, sequence_length)`、*任意*) : 各入力シーケンスのトークンが位置埋め込み内で占める位置のインデックスです。`[0, config.n_positions - 1]` の範囲から選択されます。 [位置 ID とは何か？](../glossary#position-ids)

attention_mask (`dict[str, torch.Tensor]`、形状 `(batch_size, sequence_length)`、*任意*) : パディングトークンのインデックスに対するアテンションの計算を回避するためのマスク。`[0, 1]` から選択されるマスクの値:  - **マスクされない**トークンは 1、**マスクされる**トークンは 0。  [アテンションマスクとは?](../glossary#attention-mask)

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`を参照）。

**戻り値:**

`tuple[torch.Tensor, torch.Tensor]`

[Gemma4AssistantForCausalLM](/docs/transformers/main/ja/model_doc/gemma4_assistant#transformers.Gemma4AssistantForCausalLM)のforwardメソッドは、`__call__`特殊メソッドをオーバーライドします。

この関数内でフォワードパスの手順を定義する必要がありますが、その後、この関数ではなく`Module`インスタンスを呼び出すべきです。前者は前処理と後処理の実行を適切に処理するのに対し、後者はそれらを無視してしまいます。

例:

```python
>>> from transformers import AutoTokenizer, Gemma4AssistantForCausalLM, Gemma4ForCausalLM

>>> model = Gemma4ForCausalLM.from_pretrained("google/gemma-4-e2b-it")
>>> assistant_model = Gemma4AssistantForCausalLM.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?"
```