Initial upload: Gemma 4 vision encoder (569.6M, 27-layer ViT with 2D RoPE)

README.md

@@ -16,7 +16,18 @@ pipeline_tag: image-feature-extraction
 base_model: google/gemma-4-31B-it
 model-index:
 - name: gemma4-vision-encoder
-  results: []
+  results:
+  - task:
+      type: image-classification
+      name: CIFAR-10 (10-class)
+    dataset:
+      name: CIFAR-10
+      type: cifar10
+      split: test
+    metrics:
+    - type: accuracy
+      value: 94.0
+      name: Linear Probe Accuracy
 ---
 
 # Gemma 4 Vision Encoder (27-layer ViT with 2D RoPE)
@@ -75,10 +86,11 @@ Unlike the audio encoder (which is identical across E2B and E4B), the vision enc
 
 ```python
 import torch
-from transformers import Gemma4VisionModel, Gemma4VisionConfig
+from transformers import Gemma4VisionModel, Gemma4VisionConfig, AutoProcessor
 from safetensors.torch import load_file
+from PIL import Image
 
-# Load vision encoder from this repo
+# Load vision encoder
 cfg = Gemma4VisionConfig.from_pretrained("rnagabh/gemma4-vision-encoder")
 vision_model = Gemma4VisionModel(cfg)
 state_dict = load_file("path/to/model.safetensors")  # or download from repo
@@ -86,42 +98,47 @@ vision_model.load_state_dict(state_dict, strict=True)
 vision_model = vision_model.to(dtype=torch.bfloat16, device="cuda")
 vision_model.eval()
 
-# Prepare image: patchify and create position IDs
-# Image must have sides divisible by patch_size (16) AND
-# num_patches must be divisible by pooling_kernel^2 (9)
-# Good sizes: 864 (54 patches/side), 768 (48), 576 (36)
-P = 16
-img_size = 864
-patches_per_side = img_size // P  # 54
-
-# Patchify: (B, C, H, W) → (B, num_patches, C*P*P)
-img = torch.randn(1, 3, img_size, img_size, dtype=torch.bfloat16, device="cuda")
-patches = img.unfold(2, P, P).unfold(3, P, P)
-patches = patches.contiguous().view(1, 3, -1, P, P)
-patches = patches.permute(0, 2, 1, 3, 4)
-patches = patches.reshape(1, -1, 3 * P * P)  # (1, 2916, 768)
-
-# Position IDs: (batch, num_patches, 2) as (x, y) coordinates
-ys, xs = torch.meshgrid(
-    torch.arange(patches_per_side),
-    torch.arange(patches_per_side),
-    indexing="ij",
-)
-position_ids = torch.stack([xs.flatten(), ys.flatten()], dim=-1)
-position_ids = position_ids.unsqueeze(0).to(device="cuda")  # (1, 2916, 2)
+# Use the parent model's image processor for correct preprocessing
+processor = AutoProcessor.from_pretrained("google/gemma-4-31B-it")
+image_processor = processor.image_processor
+
+# Process an image
+img = Image.open("your_image.jpg")
+processed = image_processor(images=[img], return_tensors="pt")
+
+pixel_values = processed["pixel_values"].to(dtype=torch.bfloat16, device="cuda")
+position_ids = processed["image_position_ids"].to(device="cuda")
+tokens_per_image = processed["num_soft_tokens_per_image"]  # for splitting batch output
 
 with torch.no_grad():
-    output = vision_model(pixel_values=patches, pixel_position_ids=position_ids)
-    embeddings = output.last_hidden_state  # (324, 1152) — pooled tokens
+    output = vision_model(pixel_values=pixel_values, pixel_position_ids=position_ids)
+    embeddings = output.last_hidden_state  # (num_tokens, 1152)
+    
+    # Mean-pool for a single image vector
+    image_embedding = embeddings.float().mean(dim=0)  # (1152,)
 ```
 
-> **Image size constraints:** The number of patches must be divisible by the pooling kernel² (9).
-> This means each image dimension divided by patch_size (16) must be divisible by 3.
-> Valid image sizes include: 576, 768, 864, 960, 1152, etc.
+> **Important:** Always use `Gemma4ImageProcessor` from the parent model for preprocessing.
+> It handles resizing, patchification, position ID generation, and pixel normalization.
+> Manual patchification without this processor will produce degraded results.
+
+## Benchmark Results (frozen 1152-dim embeddings, linear probe)
+
+### CIFAR-10 Classification
+
+| Metric | Value |
+|---|---|
+| Linear probe accuracy | **94.0%** |
+| Random baseline | 10.0% |
+| Improvement over chance | **9.4×** |
+| Dataset | CIFAR-10 test set (1000 samples, 100 per class) |
+| Probe | Logistic regression on L2-normalized mean-pooled embeddings |
+
+Strong performance across all classes: airplane (0.98 F1), ship (0.98 F1), truck (0.97 F1), automobile (0.97 F1). Weakest class is cat (0.86 F1) — a fine-grained category that is inherently harder.
 
-> **Output shape:** The batch dimension is collapsed — the pooler strips padding and returns
-> a flat `(num_valid_tokens, hidden_dim)` tensor. For a single 864×864 image, you get
-> `(324, 1152)` — 324 pooled visual tokens at 1152 dimensions.
+> **Important:** Use the parent model's image processor (`Gemma4ImageProcessor` from `google/gemma-4-31B-it`)
+> for correct preprocessing. Manual patchification without proper resizing and position ID generation
+> will produce significantly degraded results.
 
 ## Files in This Repo