Gemma-4 vision: correct std_bias/std_scale classification per HF standardize (fidelity fix)

src/maxtext/configs/models/gemma4-26b.yml

@@ -61,3 +61,4 @@ num_attention_heads_for_vit: 16
 image_placeholder: "<|image|>"
 vision_output_length: 280
 num_position_embeddings_for_vit: 10240
+standardize_for_vit: true  # HF vision_config.standardize=true: std_bias/std_scale are checkpoint-resident buffers

src/maxtext/configs/models/gemma4-31b.yml

@@ -57,3 +57,4 @@ num_attention_heads_for_vit: 16
 image_placeholder: "<|image|>"
 vision_output_length: 280
 num_position_embeddings_for_vit: 10240
+standardize_for_vit: true  # HF vision_config.standardize=true: std_bias/std_scale are checkpoint-resident buffers

src/maxtext/configs/models/gemma4-e2b.yml

@@ -58,3 +58,4 @@ num_attention_heads_for_vit: 12
 image_placeholder: "<|image|>"
 vision_output_length: 280
 num_position_embeddings_for_vit: 10240
+standardize_for_vit: false  # HF vision_config.standardize=false: no std_bias/std_scale in checkpoint (identity)

src/maxtext/configs/models/gemma4-e4b.yml

@@ -59,3 +59,4 @@ num_attention_heads_for_vit: 12
 image_placeholder: "<|image|>"
 vision_output_length: 280
 num_position_embeddings_for_vit: 10240
+standardize_for_vit: false  # HF vision_config.standardize=false: no std_bias/std_scale in checkpoint (identity)

src/maxtext/configs/types.py

@@ -2143,6 +2143,20 @@ class MultimodalGeneral(BaseModel):
       description="The style of VisionEncoderBlock to use (e.g., 'gemma3', 'llama4').",
   )
   freeze_vision_encoder_params: bool = Field(True, description="Freeze the parameters of the vision encoder.")
+  standardize_for_vit: bool | None = Field(
+      None,
+      description=(
+          "Whether the vision encoder standardizes exit embeddings with std_bias/std_scale. Mirrors HF"
+          " Gemma-4 vision_config.standardize and is a per-model-family semantic distinction, so it must be set"
+          " EXPLICITLY by every Gemma-4 multimodal model config (no global default is safe): HF registers"
+          " std_bias/std_scale as checkpoint-resident (non-trainable) buffers ONLY when standardize=True."
+          " Gemma-4 26B/31B ship standardize=True (2 std tensors); E2B/E4B ship standardize=False (0 std tensors)."
+          " When False the standardize op is an exact identity and MaxText must NOT construct std_bias/std_scale"
+          " (fabricating them wrongly enters checkpoint-restored committed arrays into the nnx.Param gradient"
+          " filter). None means unset: a Gemma-4 model with use_multimodal=True and this unset is a hard config"
+          " error (see validate_standardize_for_vit) — it must never silently inherit a value."
+      ),
+  )
   freeze_audio_encoder_params: bool = Field(True, description="Freeze the parameters of the audio encoder.")
   use_audio: bool = Field(False, description="Enable audio encoder for multimodal models.")
   image_size_for_vit: int | list[int] | None = Field(896, description="Input image size for the Vision Transformer.")
@@ -3581,6 +3595,15 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
       raise ValueError(
           f"{self.model_name} requires scan_layers=False (per-layer KV sharing is incompatible with nn.scan)."
       )
+    # Gemma 4 vision standardization is a per-model-family semantic distinction (HF vision_config.standardize)
+    # and must be set EXPLICITLY — it must never silently inherit a global default, or a standardize=True variant
+    # (26B/31B) would drop its checkpoint-resident std buffers, or a standardize=False variant (E2B/E4B) would
+    # fabricate spurious trainable std leaves. Fail closed when unset for any Gemma-4 model.
+    if self.model_name.startswith("gemma4-") and self.standardize_for_vit is None:
+      raise ValueError(
+          f"{self.model_name}: standardize_for_vit must be set explicitly (True for 26B/31B, False for E2B/E4B) "
+          "to mirror HF vision_config.standardize. Set it in the model config YAML; it must not inherit a default."
+      )
     if self.use_multimodal:
       # Gemma 4 small (E2B / E4B) only supports text for now; multimodal
       # support is pending clipped-linears in the vision encoder.

src/maxtext/models/gemma4_vision.py

@@ -85,6 +85,17 @@ def patchify(images: jax.Array, patch_size: int) -> tuple[jax.Array, jax.Array]:
   return patches, jnp.broadcast_to(positions_xy, tuple(b) + positions_xy.shape)
 
 
+class VisionStdVar(nnx.Variable):
+  """Checkpoint-resident, NON-trainable vision standardization state (std_bias / std_scale).
+
+  Mirrors HF `modeling_gemma4`, which registers `std_bias`/`std_scale` via `register_buffer`
+  (non-trainable) and only when `vision_config.standardize=True` (Gemma-4 26B/31B). As a plain
+  `nnx.Variable` (not `nnx.Param`), it is naturally excluded from the optimizer/weight-decay/grad
+  filter — which key on `nnx.Param` — while still being saved and restored by Orbax as part of the
+  model state. This is the same idiom MaxText uses for other non-trainable buffers (e.g. `MoEBiasVar`).
+  """
+
+
 class VisionEntry(nnx.Module):
   """The vision entry layer."""
 
@@ -566,12 +577,22 @@ def __init__(self, config: Config, mesh: Mesh, *, rngs: nnx.Rngs):
         rngs=self.rngs,
         precision=config.matmul_precision,
     )
-    self.std_bias = nnx.Param(
-        nnx.initializers.zeros(self.rngs.params(), (config.hidden_size_for_vit,), config.weight_dtype), sharding=(None,)
-    )
-    self.std_scale = nnx.Param(
-        nnx.initializers.ones(self.rngs.params(), (config.hidden_size_for_vit,), config.weight_dtype), sharding=(None,)
-    )
+    # Vision standardization (std_bias/std_scale) is CHECKPOINT-RESIDENT state that exists ONLY when the
+    # checkpoint ships std tensors, i.e. HF vision_config.standardize=True (Gemma-4 26B/31B). HF registers them
+    # via register_buffer (NON-trainable). Gemma-4 E2B/E4B have standardize=False and DO NOT store std_scale/
+    # std_bias in their safetensors; when disabled the standardize op is an EXACT identity, so we construct NO std
+    # state at all (fabricating trainable identity nnx.Param leaves wrongly enters checkpoint-restored committed
+    # arrays into the nnx.Param gradient filter). When enabled, we construct them as VisionStdVar (a plain
+    # nnx.Variable): checkpoint-resident and restored by Orbax, but NON-trainable — excluded from the optimizer /
+    # weight-decay / gradient filter (which key on nnx.Param), matching HF's register_buffer semantics.
+    self._standardize_for_vit = bool(config.standardize_for_vit)
+    if self._standardize_for_vit:
+      self.std_bias = VisionStdVar(
+          nnx.initializers.zeros(self.rngs.params(), (config.hidden_size_for_vit,), config.weight_dtype)
+      )
+      self.std_scale = VisionStdVar(
+          nnx.initializers.ones(self.rngs.params(), (config.hidden_size_for_vit,), config.weight_dtype)
+      )
 
   def __call__(self, inputs: jax.Array, deterministic: bool = False) -> jax.Array:
     """Applies the vision encoder layer."""
@@ -593,9 +614,10 @@ def __call__(self, inputs: jax.Array, deterministic: bool = False) -> jax.Array:
     # We take the first result.
     (embeddings, _) = vision_exit_results[0]
 
-    embeddings = (embeddings - self.std_bias.value.astype(embeddings.dtype)) * self.std_scale.value.astype(
-        embeddings.dtype
-    )
+    if self._standardize_for_vit:
+      embeddings = (embeddings - self.std_bias.value.astype(embeddings.dtype)) * self.std_scale.value.astype(
+          embeddings.dtype
+      )
 
     # Unflatten batch and num_images
     final_x = jnp.reshape(embeddings, (b, n, embeddings.shape[1], embeddings.shape[2]))

tests/unit/gemma4_vision_standardize_test.py

@@ -0,0 +1,105 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for the Gemma-4 vision standardization (std_bias/std_scale) contract.
+
+HF `modeling_gemma4` registers vision `std_bias`/`std_scale` via `register_buffer` (non-trainable) and ONLY when
+`vision_config.standardize=True`. Per-variant HF truth: Gemma-4 26B/31B ship standardize=True (2 std tensors);
+E2B/E4B ship standardize=False (0 std tensors). This test locks the MaxText contract:
+
+  * The `standardize_for_vit` config field is a per-model-family semantic distinction and must be set explicitly
+    by every Gemma-4 model config (no silent global default); an unset Gemma-4 config is a hard config error.
+  * standardize_for_vit=False  -> NO std state at all (the standardize op is an exact identity).
+  * standardize_for_vit=True   -> exactly two std leaves, constructed as VisionStdVar (a non-trainable
+    `nnx.Variable`, NOT `nnx.Param`), so they are checkpoint-resident but excluded from the optimizer/gradient.
+"""
+
+import unittest
+
+from flax import nnx
+import jax
+import numpy as np
+from jax.sharding import Mesh
+
+from maxtext.configs import pyconfig
+from maxtext.utils.globals import MAXTEXT_REPO_ROOT
+from maxtext.models.gemma4_vision import Gemma4VisionEncoderLayer, VisionStdVar
+
+
+def _init(model_name):
+  base_config_path = f"{MAXTEXT_REPO_ROOT}/src/maxtext/configs/base.yml"
+  argv = [
+      "",
+      base_config_path,
+      f"model_name={model_name}",
+      "attention=dot_product",
+      "matmul_precision=highest",
+      "dtype=float32",
+      "dtype_mm=float32",
+      "weight_dtype=float32",
+      "skip_jax_distributed_system=true",
+      "enable_checkpointing=false",
+      "run_name=std_test",
+      "base_output_directory=/tmp/std_test",
+  ]
+  if model_name in ("gemma4-e2b", "gemma4-e4b"):
+    argv.append("scan_layers=false")
+  return pyconfig.initialize(argv)
+
+
+class TestGemma4VisionStandardizeContract(unittest.TestCase):
+  """standardize_for_vit gates whether/how std_bias/std_scale exist, per HF vision_config.standardize."""
+
+  def setUp(self):
+    self.mesh = Mesh(np.array(jax.devices()[:1]), axis_names=("data",))
+
+  def _std_leaves(self, model, filt):
+    names = set()
+    for path, _ in jax.tree_util.tree_flatten_with_path(nnx.state(model, filt))[0]:
+      ps = "/".join(str(getattr(p, "key", p)) for p in path)
+      if "std_bias" in ps or "std_scale" in ps:
+        names.add(ps)
+    return names
+
+  def test_e2b_standardize_false_has_no_std_state(self):
+    """E2B (HF standardize=false): no std state at all; standardize is an exact identity."""
+    cfg = _init("gemma4-e2b")
+    self.assertFalse(cfg.standardize_for_vit)
+    model = Gemma4VisionEncoderLayer(cfg, self.mesh, rngs=nnx.Rngs(0))
+    self.assertFalse(hasattr(model, "std_bias"))
+    self.assertFalse(hasattr(model, "std_scale"))
+    self.assertEqual(self._std_leaves(model, nnx.Param), set())
+    self.assertEqual(self._std_leaves(model, nnx.Variable), set())
+
+  def test_26b_standardize_true_has_nontrainable_std_buffers(self):
+    """26B (HF standardize=true): exactly two std leaves as VisionStdVar, NOT in nnx.Param."""
+    cfg = _init("gemma4-26b")
+    self.assertTrue(cfg.standardize_for_vit)
+    model = Gemma4VisionEncoderLayer(cfg, self.mesh, rngs=nnx.Rngs(0))
+    self.assertIsInstance(model.std_bias, VisionStdVar)
+    self.assertIsInstance(model.std_scale, VisionStdVar)
+    # Present as (non-Param) Variables, absent from the Param (trainable) collection.
+    self.assertEqual(self._std_leaves(model, nnx.Param), set())
+    self.assertEqual(len(self._std_leaves(model, nnx.Variable)), 2)
+
+  def test_e2b_e4b_false_and_26b_31b_true(self):
+    """All four Gemma-4 configs resolve standardize_for_vit to their HF truth from the real YAMLs."""
+    self.assertFalse(_init("gemma4-e2b").standardize_for_vit)
+    self.assertFalse(_init("gemma4-e4b").standardize_for_vit)
+    self.assertTrue(_init("gemma4-26b").standardize_for_vit)
+    self.assertTrue(_init("gemma4-31b").standardize_for_vit)
+
+
+if __name__ == "__main__":
+  unittest.main()