Gemma4 E-series staged export should PAD-mask image-token rows for per-layer embedding inputs

[kylejfrost]

Summary

Gemma4 E-series multimodal staged exports can lose image conditioning because the transformer-stage graph uses the raw image placeholder token id when computing per-layer embedding (PLE) identity inputs.

The stage also receives already-spliced hidden_states, so the image features are present at the graph boundary, but the PLE side channel still encodes token identity from image_token_id (258880) instead of the PAD token id used by the Hugging Face path for multimodal/image rows. The wrong PLE identity is injected into every decoder layer and can dominate or cancel the visual conditioning.

In our fork, decoupling the image-position projection mask from the PLE token identity fixed the issue:

• image/PAD rows still use the provided hidden_states for the per-layer projection input,
• PLE ids for those same rows are replaced with pad_token_id,
• non-image text rows continue to use their original token ids.

We are filing this as an issue, not a PR, because the repository notes that PRs are closed.

Affected Path

Observed in Gemma4 E-series staged transformer shards with hidden_size_per_layer_input > 0, specifically around:

• models/macos/gemma4.py, Gemma4TransformerLayerStage.project_per_layer_inputs
• the stage contract where transformer stages take both hidden_states and input_ids

The pre-fix behavior was effectively:

inputs_embeds = self.embed_tokens(input_ids)
image_mask = (input_ids == image_token_id).unsqueeze(-1)
inputs_embeds = torch.where(image_mask, hidden_states, inputs_embeds)
per_layer_projection = ...
per_layer_inputs = self.get_per_layer_inputs(input_ids)
return (per_layer_projection + per_layer_inputs) * self.per_layer_input_scale

That preserves the spliced image rows for per_layer_projection, but get_per_layer_inputs(input_ids) still sees 258880 on every image row.

The fix we validated is:

inputs_embeds = self.embed_tokens(input_ids)
image_token_id = int(getattr(self.config, "image_token_id", 258880))
pad_token_id = int(getattr(self.config, "pad_token_id", 0))
image_or_pad_mask = (input_ids == image_token_id) | (input_ids == pad_token_id)
inputs_embeds = torch.where(image_or_pad_mask.unsqueeze(-1), hidden_states, inputs_embeds)

per_layer_projection = ...
per_layer_input_ids = torch.where(
    image_or_pad_mask,
    torch.full_like(input_ids, pad_token_id),
    input_ids,
)
per_layer_inputs = self.get_per_layer_inputs(per_layer_input_ids)
return (per_layer_projection + per_layer_inputs) * self.per_layer_input_scale

Reproduction

Model/export shape:

• model: google/gemma-4-E4B-it-qat-q4_0-unquantized
• staged split: 0:21,21:42
• native context: 131072
• multimodal fixture: deterministic RGB image with red square, blue circle, yellow triangle, white diagonal line
• image token span: 266 rows, image_token_id=258880, pad_token_id=0
• fp32 vision tower output was validated against HF before LM execution (0 non-finite rows, cosine 1.0)

Symptoms before the exporter fix:

• caix/Core AI generated text like "No shapes are visible" or image-independent descriptions.
• Frozen HF soft tokens did not help, so preprocessing, vision tower, splice magnitude, and runtime insertion were ruled out.
• mm-on and mm-off teacher-forced agreement were identical in one QAT run: 43/64 == 43/64.
• Runtime-only PAD-masking of input_ids improved the PLE identity but disabled the graph's image-token projection mask, so it was not sufficient.

PLE instrumentation:

• Raw image-id PLE path:
  • HF-vs-CoreAI layer-0 image-row PLE cosine: 0.34706932306289673
  • all-layer raw PLE delta norm mean: about 102.109
• PAD-masked PLE path:
  • HF-vs-CoreAI layer-0 image-row PLE cosine: 1.0
  • all-layer PLE delta norm: 0.0

After the exporter fix:

• Split-cache QAT bundle first token: 818, matching HF.
• Teacher-forced image gate: mm_on_hf_top1=58/64, mm_off_hf_top1=43/64, first on/off divergence at step 4.
• Live HTTP response to the fixture image: "Red square, Blue circle, Yellow triangle".
• The same runtime also preserved existing non-mm staged behavior and the Gemma 12B unified multimodal route.

Why this looks like an exporter issue

The Hugging Face path appears to use PAD identity for multimodal/image rows in the PLE input while separately incorporating the projected image features. The current staged graph ties both decisions to the same raw input_ids tensor:

• image-token id is needed to decide where hidden_states should replace embed_tokens(input_ids),
• PAD id is needed for get_per_layer_inputs(...) on those image rows.

Using one tensor for both roles makes the exported stage unable to match HF for E-series multimodal rows unless the graph itself decouples the masks.

The 12B gemma4_unified path did not expose this because its transformer stages do not use this E-series input_ids / PLE side channel.

Ask

Is this a known/planned gap in the Gemma4 E-series macOS staged export path?

Would Apple consider decoupling the image-position projection mask from the PLE identity ids in Gemma4TransformerLayerStage.project_per_layer_inputs, so image/PAD multimodal rows use spliced hidden_states for projection while using pad_token_id for get_per_layer_inputs?

We can provide the minimal fixture, logs, and the validated fork patch if useful.