Update modeling_geolip_vit.py

modeling_geolip_vit.py

@@ -862,31 +862,33 @@ class DualStreamViT(nn.Module):
 
 class MasteryQueue:
     """
-    Cross-batch embedding cache for progressive hard contrastive learning.
-
-    Stage 1: Inactive. Standard InfoNCE handles in-batch discrimination.
-    Stage 2: Activates when nce_acc=1.0 for `patience` consecutive batches.
-             Caches embeddings + labels from recent batches.
-             compute_loss uses the queue to find:
-               - hard negatives: closest different-class embedding
-               - hard positives: furthest same-class embedding
-             Margin loss forces the model to separate these.
-
-    This creates the class-level asymmetry that moves CV toward the
-    natural 0.20-0.23 band. Dense same-class regions + sparse boundaries
-    = volume variation = higher CV.
+    Cross-batch embedding cache with adaptive queue sizing.
+
+    Activation: when nce_acc >= 0.99 for `patience` consecutive batches.
+    Progressive margin: ramps from margin_start → margin_end over margin_warmup.
+
+    Adaptive queue sizing (call update_size each epoch):
+      - Monitors train_acc - val_acc gap (overfitting indicator)
+      - Gap growing → increase queue (more diverse negatives = regularization)
+      - Gap shrinking → decrease queue (tighter contrastive signal)
+      - Cooldown prevents oscillation: no resize for `resize_cooldown` epochs
+        after each change.
     """
-    def __init__(self, dim, max_size=4096, patience=50, device='cuda',
-                 margin_start=0.1, margin_end=0.3, margin_warmup=5000):
+    def __init__(self, dim, min_size=1024, max_size=8192, initial_size=4096,
+                 patience=50, device='cuda',
+                 margin_start=0.1, margin_end=0.3, margin_warmup=5000,
+                 resize_step=1024, resize_cooldown=5, overfit_threshold=3.0):
         self.dim = dim
+        self.min_size = min_size
         self.max_size = max_size
+        self._current_max = initial_size
         self.patience = patience
         self.device = device
         self.active = False
 
         # Queue storage
-        self._embs = None   # (Q, dim)
-        self._labels = None  # (Q,)
+        self._embs = None
+        self._labels = None
 
         # Activation tracking
         self._perfect_count = 0
@@ -896,11 +898,19 @@ class MasteryQueue:
         # Progressive margin
         self._margin_start = margin_start
         self._margin_end = margin_end
-        self._margin_warmup = margin_warmup  # batches after activation to reach max
-        self._mastery_steps = 0  # batches since activation
+        self._margin_warmup = margin_warmup
+        self._mastery_steps = 0
+
+        # Adaptive sizing
+        self._resize_step = resize_step
+        self._resize_cooldown = resize_cooldown
+        self._overfit_threshold = overfit_threshold
+        self._epochs_since_resize = resize_cooldown  # allow first resize
+        self._prev_gap = None
+        self._resize_history = []
 
     def check_activation(self, nce_acc):
-        """Call each batch. Activates when nce_acc=1.0 for patience steps."""
+        """Call each batch. Activates when nce_acc >= 0.99 for patience steps."""
         self._total_batches += 1
         if nce_acc >= 0.99:
             self._perfect_count += 1
@@ -912,21 +922,64 @@ class MasteryQueue:
             self._activated_at = self._total_batches
             print(f"\n  ★ MASTERY ACTIVATED at batch {self._total_batches} "
                   f"(nce_acc=1.0 for {self.patience} consecutive) "
-                  f"[InfoNCE stays ON, margin {self._margin_start}→{self._margin_end}]")
+                  f"[InfoNCE stays ON, margin {self._margin_start}→{self._margin_end}]"
+                  f" queue={self._current_max}")
 
         if self.active:
             self._mastery_steps += 1
 
+    def update_size(self, train_acc, val_acc, epoch):
+        """
+        Call once per epoch. Adjusts queue size based on overfit gap.
+
+        Gap = train_acc - val_acc.
+        Gap growing → queue grows (more negatives = regularization)
+        Gap shrinking → queue shrinks (tighter signal)
+        Cooldown prevents oscillation.
+        """
+        if not self.active:
+            return
+
+        self._epochs_since_resize += 1
+        gap = train_acc - val_acc
+
+        if self._prev_gap is not None and self._epochs_since_resize >= self._resize_cooldown:
+            delta = gap - self._prev_gap
+            old_size = self._current_max
+
+            if delta > self._overfit_threshold:
+                # Overfitting increasing → grow queue for regularization
+                self._current_max = min(
+                    self._current_max + self._resize_step, self.max_size)
+            elif delta < -self._overfit_threshold:
+                # Overfitting decreasing → shrink queue for sharper signal
+                self._current_max = max(
+                    self._current_max - self._resize_step, self.min_size)
+
+            if self._current_max != old_size:
+                direction = "↑" if self._current_max > old_size else "↓"
+                print(f"  ⚙ Queue {direction} {old_size}→{self._current_max} "
+                      f"(gap {self._prev_gap:.1f}→{gap:.1f}, Δ={delta:+.1f})")
+                self._epochs_since_resize = 0
+                self._resize_history.append(
+                    (epoch, old_size, self._current_max, gap))
+
+                # Trim queue if it shrunk
+                if self._embs is not None and self._embs.shape[0] > self._current_max:
+                    self._embs = self._embs[-self._current_max:]
+                    self._labels = self._labels[-self._current_max:]
+
+        self._prev_gap = gap
+
     @property
     def current_margin(self):
-        """Progressive margin: linearly ramps from start to end over warmup steps."""
         if not self.active:
             return self._margin_start
         t = min(self._mastery_steps / max(self._margin_warmup, 1), 1.0)
         return self._margin_start + t * (self._margin_end - self._margin_start)
 
     def push(self, emb, labels):
-        """Add batch to queue. FIFO eviction."""
+        """Add batch to queue. FIFO eviction at current_max."""
         emb = emb.detach().to(self.device)
         labels = labels.detach().to(self.device)
 
@@ -934,11 +987,10 @@ class MasteryQueue:
             self._embs = emb
             self._labels = labels
         else:
-            self._embs = torch.cat([self._embs, emb], 0)[-self.max_size:]
-            self._labels = torch.cat([self._labels, labels], 0)[-self.max_size:]
+            self._embs = torch.cat([self._embs, emb], 0)[-self._current_max:]
+            self._labels = torch.cat([self._labels, labels], 0)[-self._current_max:]
 
     def get(self):
-        """Return current queue contents."""
         if self._embs is None:
             return None, None
         return self._embs, self._labels
@@ -955,6 +1007,8 @@ class MasteryQueue:
             'activated_at': self._activated_at,
             'mastery_steps': self._mastery_steps,
             'current_margin': self.current_margin,
+            'current_max': self._current_max,
+            'resize_history': self._resize_history,
         }