Revert "Adding 2D pooling for image embeddings"

This reverts commit 65350cf.

Author: MayankChaturvedi

src/transformers/models/gemma3/configuration_gemma3.py

@@ -256,7 +256,6 @@ def __init__(
         layer_norm_eps: float = 0.000001,
         vision_use_head: bool = False,
         torch_dtype: str = "bfloat16",
-        pooled_seq_len: int = 256,
         **kwargs,
     ):
         super().__init__(
@@ -274,7 +273,6 @@ def __init__(
             **kwargs,
         )
 
-        self.pooled_seq_len = pooled_seq_len
         self.vision_use_head = vision_use_head
 
 

src/transformers/models/gemma3/modeling_gemma3.py

@@ -25,7 +25,6 @@
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
 from ...activations import ACT2FN
 from ...cache_utils import Cache, HybridCache, StaticCache
@@ -45,7 +44,7 @@
 from ...utils.deprecation import deprecate_kwarg
 from ..gemma import GemmaPreTrainedModel
 from ..siglip import SiglipVisionModel
-from .configuration_gemma3 import Gemma3Config, Gemma3RotaryEmbeddingConfig, Gemma3TextConfig, Gemma3VisionConfig
+from .configuration_gemma3 import Gemma3Config, Gemma3RotaryEmbeddingConfig, Gemma3TextConfig
 
 
 logger = logging.get_logger(__name__)
@@ -72,28 +71,6 @@ def extra_repr(self):
         return f"{tuple(self.weight.shape)}, eps={self.eps}"
 
 
-class Gemma3VisionAvgPool2D(nn.Module):
-    def __init__(self, config: Gemma3VisionConfig):
-        super().__init__()
-        self.config = config
-
-    def forward(self, x):
-        """
-        Applies average pooling on (B, width, width)
-        to make it (B, final_width, final_width).
-        """
-        batch_size, seq_len, channels = x.shape
-        width = int(seq_len**0.5)
-        if width * width != seq_len:
-            raise ValueError(f"Sequence length {seq_len} is not a perfect square. Cannot reshape to a square image.")
-        final_width = int(self.config.pooled_seq_len**0.5)
-        kernel_size = width // final_width
-        x = x.transpose(1, 2).reshape(batch_size, channels, width, width)
-        x = F.avg_pool2d(x, kernel_size=kernel_size, stride=kernel_size)
-        x = x.flatten(2).transpose(1, 2)
-        return x
-
-
 class Gemma3MultimodalInputProjection(nn.Module):
     def __init__(self, vision_dim: int, text_dim: int):
         super().__init__()
@@ -1035,6 +1012,7 @@ class Gemma3ForConditionalGeneration(PreTrainedModel, GenerationMixin):
 
     def __init__(self, config: Gemma3Config):
         super().__init__(config)
+
         self.config = config
         text_config = self.config.text_config
         vision_config = self.config.vision_config
@@ -1050,7 +1028,10 @@ def __init__(self, config: Gemma3Config):
             vision_dim=vision_config.hidden_size, text_dim=text_config.hidden_size
         )
         self.mm_soft_emb_norm = Gemma3RMSNorm(vision_config.hidden_size, eps=vision_config.layer_norm_eps)
-        self.avg_pool = Gemma3VisionAvgPool2D(config.vision_config)
+
+        patches_per_image = vision_config.image_size // vision_config.patch_size
+        avg_pool_k = patches_per_image**2 // text_config.mm_tokens_per_image
+        self.avg_pool = nn.AvgPool1d(kernel_size=avg_pool_k, stride=avg_pool_k)
         self.vocab_size = text_config.vocab_size
         self.pad_token_id = pad_token_id if (pad_token_id := text_config.pad_token_id) is not None else -1
         self.post_init()
@@ -1095,7 +1076,12 @@ def get_image_features(self, pixel_values: torch.Tensor) -> torch.Tensor:
             image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`).
         """
         vision_outputs = self.vision_model(pixel_values=pixel_values).last_hidden_state
-        pooled_vision_outputs = self.avg_pool(vision_outputs)
+        b, n, l = vision_outputs.shape
+        reshaped_vision_outputs = vision_outputs.permute(0, 2, 1)
+        reshaped_vision_outputs = reshaped_vision_outputs.contiguous()
+        reshaped_vision_outputs = reshaped_vision_outputs.view(b, l, n)
+        pooled_vision_outputs = self.avg_pool(reshaped_vision_outputs)
+        pooled_vision_outputs = pooled_vision_outputs.permute(0, 2, 1)
         image_features = self.encode_vision(pooled_vision_outputs)
         return image_features
 

src/transformers/models/gemma3/modular_gemma3.py

@@ -26,7 +26,6 @@
 import torch
 import torch.nn as nn
 import torch.utils.checkpoint
-import torch.nn.functional as F
 
 from ...activations import ACT2FN
 from ...cache_utils import Cache, HybridCache, StaticCache
@@ -333,7 +332,6 @@ def __init__(
         layer_norm_eps: float = 0.000001,
         vision_use_head: bool = False,
         torch_dtype: str = "bfloat16",
-        pooled_seq_len: int = 256,
         **kwargs,
     ):
         super().__init__(
@@ -351,7 +349,6 @@ def __init__(
             **kwargs,
         )
 
-        self.pooled_seq_len = pooled_seq_len
         self.vision_use_head = vision_use_head
 
 
@@ -713,28 +710,6 @@ def model_input_names(self):
 class Gemma3RMSNorm(GemmaRMSNorm):
     pass
 
-class Gemma3VisionAvgPool2D(nn.Module):
-  def __init__(self, config: Gemma3VisionConfig):
-    super().__init__()
-    self.config = config
-
-  def forward(self, x):
-      """
-        Applies average pooling on (B, width, width) 
-        to make it (B, final_width, final_width).
-      """
-      batch_size, seq_len, channels = x.shape
-      width = int(seq_len**0.5)
-      if width * width != seq_len:
-          raise ValueError(
-              f"Sequence length {seq_len} is not a perfect square. Cannot reshape to a square image."
-          )
-      final_width = int(self.config.pooled_seq_len**0.5)
-      kernel_size = width//final_width
-      x = x.transpose(1, 2).reshape(batch_size, channels, width, width)
-      x = F.avg_pool2d(x, kernel_size=kernel_size, stride=kernel_size)
-      x = x.flatten(2).transpose(1, 2)
-      return x
 
 class Gemma3MultimodalInputProjection(nn.Module):
 
@@ -1715,6 +1690,7 @@ class Gemma3ForConditionalGeneration(PreTrainedModel, GenerationMixin):
 
     def __init__(self, config: Gemma3Config):
         super().__init__(config)
+
         self.config = config
         text_config = self.config.text_config
         vision_config = self.config.vision_config
@@ -1732,7 +1708,10 @@ def __init__(self, config: Gemma3Config):
         self.mm_soft_emb_norm = Gemma3RMSNorm(
             vision_config.hidden_size, eps=vision_config.layer_norm_eps
         )
-        self.avg_pool = Gemma3VisionAvgPool2D(config.vision_config)
+
+        patches_per_image = vision_config.image_size // vision_config.patch_size
+        avg_pool_k = patches_per_image ** 2 // text_config.mm_tokens_per_image
+        self.avg_pool = nn.AvgPool1d(kernel_size=avg_pool_k, stride=avg_pool_k)
         self.vocab_size = text_config.vocab_size
         self.pad_token_id = (
             pad_token_id
@@ -1781,7 +1760,12 @@ def get_image_features(self, pixel_values: torch.Tensor) -> torch.Tensor:
             image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`).
         """
         vision_outputs = self.vision_model(pixel_values=pixel_values).last_hidden_state
-        pooled_vision_outputs = self.avg_pool(vision_outputs)
+        b, n, l = vision_outputs.shape
+        reshaped_vision_outputs = vision_outputs.permute(0, 2, 1)
+        reshaped_vision_outputs = reshaped_vision_outputs.contiguous()
+        reshaped_vision_outputs = reshaped_vision_outputs.view(b, l, n)
+        pooled_vision_outputs = self.avg_pool(reshaped_vision_outputs)
+        pooled_vision_outputs = pooled_vision_outputs.permute(0, 2, 1)
         image_features = self.encode_vision(pooled_vision_outputs)
         return image_features