v2 modeling fix: match released checkpoint (encoder.layers.*, no untrained geo path); ModelOutput; opt-in diagnostics. Verified 0 missing/0 unexpected keys, bit-reproduces faithful load.

modeling_geolip_vit.py

@@ -3,46 +3,44 @@
 #
 # Usage:
 #   from transformers import AutoModel
-#   model = AutoModel.from_pretrained("AbstractPhil/geolip-vit-base-x3",
-#                                      trust_remote_code=True)
+#   model = AutoModel.from_pretrained("AbstractPhil/geolip-vit-large-x3",
+#                                      trust_remote_code=True).eval()
 #
 #   from torchvision import transforms
 #   transform = transforms.Compose([
-#       transforms.Resize((224, 224)),
+#       transforms.Resize((224, 224)),          # squash resize (no crop)
 #       transforms.ToTensor(),
 #       transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
 #   ])
 #   pixel_values = transform(image).unsqueeze(0)
 #   outputs = model(pixel_values)
 #
-#   # 128-d embedding on hypersphere (L2-normalized)
-#   embedding = outputs.embedding              # (B, 128)
+#   outputs.embedding       # (B, 128)  L2-normalized consensus-space vector
+#   outputs.logits          # (B, 80)   multi-label COCO logits (soup head)
+#   outputs.triangulation   # (B, 256)  1 - cos to constellation anchors
+#   outputs.nearest         # (B,)      nearest anchor index
+#   outputs = model(px, compute_diagnostics=True)   # adds .diagnostics dict
 #
-#   # Multi-label classification logits (80 COCO classes)
-#   logits = outputs.logits                    # (B, 80) — if soup_enabled
-#
-#   # Triangulation distances to 256 constellation anchors
-#   triangulation = outputs.triangulation      # (B, 256)
-#
-#   # Nearest anchor index per sample
-#   nearest = outputs.nearest                  # (B,)
-#
-#   # Geometric diagnostics
-#   diagnostics = outputs.diagnostics          # dict
+# REVISION 2026-07-27 (v2): the original remote-code file did not match the
+# released checkpoint — the transformer was saved under `encoder.layers.*`
+# (nn.TransformerEncoder) but the class declared per-layer `layers.*` plus a
+# geometric-injection path (`geo_pool_proj`/`geo_tri_proj`) that was never
+# trained, so AutoModel silently returned a mostly-random model. This
+# revision matches the checkpoint exactly (verified: zero missing / zero
+# unexpected keys; COCO val anchor coverage and mAP reproduce the training
+# run's TensorBoard values). The output object is now a transformers
+# ModelOutput, and diagnostics are opt-in (they force a GPU sync).
 # ============================================================================
 
-import math
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from transformers import PretrainedConfig, PreTrainedModel
-from dataclasses import dataclass, field
-from typing import Optional, Dict, Any
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
 
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.utils import ModelOutput
 
-# ══════════════════════════════════════════════════════════════════
-# CONFIG
-# ══════════════════════════════════════════════════════════════════
 
 class GeoLIPViTConfig(PretrainedConfig):
     model_type = "geolip_vit"
@@ -81,25 +79,21 @@ class GeoLIPViTConfig(PretrainedConfig):
         self.hidden_dropout_prob = hidden_dropout_prob
         self.soup_enabled = soup_enabled
         self.consensus_cv = consensus_cv
-        self.experts = experts or ["clip_l14_openai", "dinov2_b14", "siglip_b16_384"]
-
+        self.experts = experts or ["clip_l14_openai", "dinov2_b14",
+                                   "siglip_b16_384"]
 
-# ══════════════════════════════════════════════════════════════════
-# OUTPUT
-# ══════════════════════════════════════════════════════════════════
 
 @dataclass
-class GeoLIPViTOutput:
+class GeoLIPViTOutput(ModelOutput):
     """
-    Output fields:
-        embedding:      (B, output_dim)   L2-normalized on hypersphere
-        logits:         (B, n_classes)    multi-label classification (if soup_enabled)
-        triangulation:  (B, n_anchors)    distances to constellation anchors
-        nearest:        (B,)              nearest anchor index
-        patch_tokens:   (B, n_patches, hidden_size)  pre-pooling patch representations
-        diagnostics:    dict              geometric metrics
+    embedding:      (B, output_dim)  L2-normalized on the consensus sphere
+    logits:         (B, n_classes)   multi-label logits (if soup_enabled)
+    triangulation:  (B, n_anchors)   1 - cos to constellation anchors
+    nearest:        (B,)             nearest anchor index
+    patch_tokens:   (B, n_patches, hidden_size)  if output_patch_tokens
+    diagnostics:    dict             only if compute_diagnostics=True
     """
-    embedding: torch.Tensor = None
+    embedding: Optional[torch.Tensor] = None
     logits: Optional[torch.Tensor] = None
     triangulation: Optional[torch.Tensor] = None
     nearest: Optional[torch.Tensor] = None
@@ -107,15 +101,12 @@ class GeoLIPViTOutput:
     diagnostics: Optional[Dict[str, Any]] = None
 
 
-# ══════════════════════════════════════════════════════════════════
-# GEOMETRIC COMPONENTS
-# ══════════════════════════════════════════════════════════════════
-
 class Constellation(nn.Module):
     def __init__(self, n_anchors, d):
         super().__init__()
         self.n_anchors = n_anchors
-        self.anchors = nn.Parameter(F.normalize(torch.randn(n_anchors, d), dim=-1))
+        self.anchors = nn.Parameter(
+            F.normalize(torch.randn(n_anchors, d), dim=-1))
 
     def triangulate(self, emb):
         a = F.normalize(self.anchors, dim=-1)
@@ -130,53 +121,38 @@ class Patchwork(nn.Module):
         self.n_anchors = n_anchors
         asgn = torch.arange(n_anchors) % n_comp
         self.register_buffer("asgn", asgn)
-        # Compute input sizes from ints, not tensors (meta-tensor safe)
         anchors_per_comp = n_anchors // n_comp
         remainder = n_anchors % n_comp
         self.comps = nn.ModuleList([nn.Sequential(
-            nn.Linear(anchors_per_comp + (1 if k < remainder else 0), d_comp * 2),
+            nn.Linear(anchors_per_comp + (1 if k < remainder else 0),
+                      d_comp * 2),
             nn.GELU(),
             nn.Linear(d_comp * 2, d_comp), nn.LayerNorm(d_comp))
             for k in range(n_comp)])
 
     def forward(self, tri):
         return torch.cat([self.comps[k](tri[:, self.asgn == k])
-                         for k in range(self.n_comp)], -1)
-
+                          for k in range(self.n_comp)], -1)
 
-# ══════════════════════════════════════════════════════════════════
-# MODEL
-# ══════════════════════════════════════════════════════════════════
 
 class GeoLIPViTModel(PreTrainedModel):
     """
-    From-scratch Vision Transformer producing L2-normalized embeddings
-    on a 128-d hypersphere, geometrically anchored by a constellation
-    of 256 reference points trained via 3-expert consensus distillation.
-
-    The encoder is trained from Xavier initialization against consensus
-    targets from CLIP ViT-L/14, DINOv2 ViT-B/14, and SigLIP ViT-B/16.
-
-    Optional soup pipeline (constellation + patchwork + classifier)
-    provides multi-label COCO classification from the embedding.
-
-    Output fields:
-        embedding:      (B, 128)    L2-normalized, consensus-aligned
-        logits:         (B, 80)     multi-label COCO logits (if soup_enabled)
-        triangulation:  (B, 256)    distances to constellation anchors
-        nearest:        (B,)        nearest anchor index
-        patch_tokens:   (B, 196, 384)  pre-pooling patch representations
-        diagnostics:    dict        geometric metrics
+    From-scratch Vision Transformer producing L2-normalized embeddings on
+    a 128-d hypersphere — the 3-expert consensus space distilled from
+    CLIP ViT-L/14 + DINOv2 ViT-B/14 + SigLIP ViT-B/16-384 (GPA-aligned,
+    PCA-128, whitened-Procrustes calibrated) on COCO-2017.
+
+    The soup pipeline (constellation of 256 anchors + 8-compartment
+    patchwork + classifier) is the frozen calibrated teacher head and
+    yields 80-class multi-label COCO logits from the embedding.
     """
     config_class = GeoLIPViTConfig
 
     def __init__(self, config):
         super().__init__(config)
         self.config = config
-
         n_patches = (config.image_size // config.patch_size) ** 2
 
-        # ── Encoder ──
         self.patch_embed = nn.Conv2d(
             3, config.hidden_size,
             kernel_size=config.patch_size, stride=config.patch_size)
@@ -186,23 +162,20 @@ class GeoLIPViTModel(PreTrainedModel):
         self.embed_norm = nn.LayerNorm(config.hidden_size)
         self.embed_drop = nn.Dropout(config.hidden_dropout_prob)
 
-        # Individual layers for geometric injection between each
-        self.layers = nn.ModuleList([
-            nn.TransformerEncoderLayer(
-                d_model=config.hidden_size,
-                nhead=config.num_attention_heads,
-                dim_feedforward=config.intermediate_size,
-                dropout=config.hidden_dropout_prob,
-                activation="gelu",
-                batch_first=True,
-                norm_first=True)
-            for _ in range(config.num_hidden_layers)])
-
-        # Geometric injection: pool → anchor_dim → triangulate → hidden_size
-        self.geo_pool_proj = nn.Linear(config.hidden_size, config.output_dim)
-        self.geo_tri_proj = nn.Sequential(
-            nn.Linear(config.n_anchors, config.hidden_size), nn.GELU(),
-            nn.LayerNorm(config.hidden_size))
+        # Matches the released checkpoint: nn.TransformerEncoder saved as
+        # `encoder.layers.*`. (The v1 remote-code per-layer ModuleList +
+        # geometric-injection path never existed in the weights.)
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=config.hidden_size,
+            nhead=config.num_attention_heads,
+            dim_feedforward=config.intermediate_size,
+            dropout=config.hidden_dropout_prob,
+            activation="gelu",
+            batch_first=True,
+            norm_first=True)
+        self.encoder = nn.TransformerEncoder(
+            encoder_layer, num_layers=config.num_hidden_layers,
+            enable_nested_tensor=False)
 
         self.output_proj = nn.Sequential(
             nn.Linear(config.hidden_size, config.hidden_size),
@@ -211,9 +184,9 @@ class GeoLIPViTModel(PreTrainedModel):
             nn.Linear(config.hidden_size, config.output_dim),
         )
 
-        # ── Soup Pipeline (optional) ──
         if getattr(config, "soup_enabled", False):
-            self.constellation = Constellation(config.n_anchors, config.output_dim)
+            self.constellation = Constellation(config.n_anchors,
+                                               config.output_dim)
             self.patchwork = Patchwork(
                 config.n_anchors, config.n_comp, config.d_comp)
             pw_dim = config.n_comp * config.d_comp
@@ -241,68 +214,38 @@ class GeoLIPViTModel(PreTrainedModel):
             nn.init.ones_(module.weight)
             nn.init.zeros_(module.bias)
 
-    def forward(self, pixel_values, output_patch_tokens=False, **kwargs):
+    def forward(self, pixel_values, output_patch_tokens=False,
+                compute_diagnostics=False, **kwargs):
         B = pixel_values.shape[0]
-
-        # ── Encode ──
-        x = self.patch_embed(pixel_values)
-        x = x.flatten(2).transpose(1, 2)
-
-        cls = self.cls_token.expand(B, -1, -1)
-        x = torch.cat([cls, x], dim=1)
+        x = self.patch_embed(pixel_values).flatten(2).transpose(1, 2)
+        x = torch.cat([self.cls_token.expand(B, -1, -1), x], dim=1)
         x = x + self.pos_embed
         x = self.embed_drop(self.embed_norm(x))
+        x = self.encoder(x)
 
-        # ── Transformer with geometric injection ──
-        # Get anchors for triangulation (from constellation if available)
-        if self.constellation is not None:
-            anchors_n = F.normalize(self.constellation.anchors.detach(), dim=-1)
-        else:
-            anchors_n = None
-
-        for layer in self.layers:
-            if anchors_n is not None:
-                # Pool → project → triangulate → geo token
-                pooled = x[:, 1:, :].mean(dim=1)
-                geo_128 = F.normalize(self.geo_pool_proj(pooled), dim=-1)
-                tri_dists = 1.0 - geo_128 @ anchors_n.T
-                geo_token = self.geo_tri_proj(tri_dists).unsqueeze(1)
-                x_with_geo = torch.cat([geo_token, x], dim=1)
-                x_with_geo = layer(x_with_geo)
-                x = x_with_geo[:, 1:, :]
-            else:
-                x = layer(x)
-
-        # ── Pool + Project ──
         patch_tokens = x[:, 1:, :]
-        pooled = patch_tokens.mean(dim=1)
-        embedding = F.normalize(self.output_proj(pooled), dim=-1)
-
-        # ── Soup Pipeline ──
-        logits = None
-        triangulation = None
-        nearest = None
-        diagnostics = {}
+        embedding = F.normalize(self.output_proj(patch_tokens.mean(dim=1)),
+                                dim=-1)
 
+        logits = triangulation = nearest = None
+        diagnostics = None
         if self.constellation is not None:
-            tri, near = self.constellation.triangulate(embedding)
-            triangulation = tri
-            nearest = near
-
+            triangulation, nearest = self.constellation.triangulate(embedding)
             if self.patchwork is not None and self.classifier is not None:
-                pw = self.patchwork(tri)
+                pw = self.patchwork(triangulation)
                 logits = self.classifier(torch.cat([pw, embedding], -1))
-
-            # Geometric diagnostics
-            with torch.no_grad():
-                anchors_n = F.normalize(self.constellation.anchors, dim=-1)
-                cos_to_anchors = embedding @ anchors_n.T
-                diagnostics = {
-                    "nearest_cos": cos_to_anchors.max(dim=-1).values.mean().item(),
-                    "mean_anchor_cos": cos_to_anchors.mean().item(),
-                    "n_active_anchors": near.unique().numel(),
-                    "embedding_norm": embedding.norm(dim=-1).mean().item(),
-                }
+            if compute_diagnostics:
+                with torch.no_grad():
+                    a_n = F.normalize(self.constellation.anchors, dim=-1)
+                    cos = embedding @ a_n.T
+                    diagnostics = {
+                        "nearest_cos":
+                            cos.max(dim=-1).values.mean().item(),
+                        "mean_anchor_cos": cos.mean().item(),
+                        "n_active_anchors": nearest.unique().numel(),
+                        "embedding_norm":
+                            embedding.norm(dim=-1).mean().item(),
+                    }
 
         return GeoLIPViTOutput(
             embedding=embedding,
@@ -311,4 +254,4 @@ class GeoLIPViTModel(PreTrainedModel):
             nearest=nearest,
             patch_tokens=patch_tokens if output_patch_tokens else None,
             diagnostics=diagnostics,
-        )
+        )