# Gemma 4 助手

## 概述

Gemma 4 Assistant 是一个小型的纯文本模型，它使用多 Token 预测（MTP）方法及其关联的候选生成器，为 Gemma 4 模型启用推测解码。我们提供了 Gemma 4 E2B、E4B、31B 和 26B-A4B（MoE）模型 IT 变体的预训练模型。

在架构上，Gemma 4 Assistant 与其他 Gemma 4 模型共享相同的 [Gemma4TextModel](/docs/transformers/main/zh/model_doc/gemma4#transformers.Gemma4TextModel) 主干，但在几个关键方面有所不同：

*   **整个模型使用 KV 共享**。这项技术最初随 [Gemma 3n](./gemma3n) 引入，允许模型复用由助手所支持的目标模型填充的 KV 缓存，从而使助手可以完全跳过 pre-fille 阶段，并显著降低前向传播过程中的注意力计算量。
*   **`position_ids` 的值是恒定的**。由于 KV 缓存是共享的，且助手没有更新缓存的方式，因此助手会从相同的位置 ID 预测所有 token。
*   **输入是嵌入和隐藏状态的拼接**。为了适配静态 KV 缓存和 `position_ids`，模型将目标模型最近一个 token 的 `embedding` 和 `hidden_states` 的拼接作为输入，并通过 `nn.Linear` 变换将其投影到辅助模型的空间中。在辅助解码循环中，“最近一个 token”的定义会不断变化。对于预填充后起草的第一个 token，最近一个 token 是提示词中的最后一个 token。对于后续的起草步骤，最近一个 token 是辅助模型生成的最后一个 token（在同一轮起草内）或目标模型接受的最后一个 token（在不同轮起草之间）。
*   **交叉注意力用于充分利用目标模型的上下文**。交叉注意力允许助手生成的查询状态关注来自目标模型的共享 KV 缓存值，从而使助手在每轮草稿生成中更准确地预测更多的草稿 token。

您可以在 [Gemma 4](https://huggingface.co/collections/google/gemma-4) 发布中找到所有原始的 Gemma 4 Assistant 检查点。

## 使用示例

以下示例演示了如何使用 [Pipeline](/docs/transformers/main/zh/main_classes/pipelines#transformers.Pipeline) 或 [AutoModel](/docs/transformers/main/zh/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]]`, *可选*) : 文本骨干网络的配置对象或字典。

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/zh/main_classes/configuration#transformers.PreTrainedConfig)，可用于控制模型输出。有关更多信息，请参阅 [PreTrainedConfig](/docs/transformers/main/zh/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/zh/model_doc/gemma4_assistant#transformers.Gemma4AssistantConfig))：模型配置类，包含模型的所有参数。使用配置文件进行初始化时，不会加载与模型关联的权重，仅加载配置。请查看 [from_pretrained()](/docs/transformers/main/zh/main_classes/model#transformers.PreTrainedModel.from_pretrained) 方法以加载模型权重。

用于 Gemma 4 的多 token 预测辅助解码模型。

该模型继承自 [PreTrainedModel](/docs/transformers/main/zh/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`，形状为 `(batch_size, sequence_length)`，*可选*)：输入序列中各个 token 在词汇表中的索引。默认情况下，填充（padding）会被忽略。可以使用 [AutoTokenizer](/docs/transformers/main/zh/model_doc/auto#transformers.AutoTokenizer) 获取这些索引。有关详细信息，请参阅 [PreTrainedTokenizer.encode()](/docs/transformers/main/zh/internal/tokenization_utils#transformers.PreTrainedTokenizerBase.encode) 和 [PreTrainedTokenizer.__call__()](/docs/transformers/main/zh/internal/tokenization_utils#transformers.PreTrainedTokenizerBase.__call__)。[什么是输入 ID？](../glossary#input-ids)

inputs_embeds (`torch.Tensor`，形状为 `(batch_size, sequence_length, hidden_size)`，*可选*)：可选地，你可以选择直接传入嵌入表示，而不是传入 `input_ids`。如果你希望比模型内部的嵌入查找矩阵更精细地控制如何将 `input_ids` 索引转换为对应的向量，这将非常有用。

position_ids (`torch.LongTensor`，形状为 `(batch_size, sequence_length)`，*可选*)：输入序列中每个 token 在位置嵌入中的位置索引。取值范围为 `[0, config.n_positions - 1]`。[什么是位置 ID？](../glossary#position-ids)

attention_mask (`dict[str, torch.Tensor]`，形状为 `(batch_size, sequence_length)`，*可选*)：用于避免对填充（padding）标记索引执行注意力机制的掩码。掩码值在 `[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/zh/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?"
```