Asymmetric encoders + MONA all 24 blocks + 1024-dim + hard negatives

Architecture changes:
- Asymmetric DINOv3: WEB (drone) + SAT (satellite) with separate MONA
- MONA on all 24 blocks per encoder (was last 12)
- Remove projection, native 1024-dim retrieval space (was 512)
- Total: 748M params, 17.6M trainable (2.35%)

Hard negative memory bank:
- MoCo-style FIFO queue of 4096 detached gallery embeddings
- Each batch: B in-batch + Q queue negatives in InfoNCE
- Queue updated after each forward pass

Training config:
- batch_size=8, grad_accum=8, effective_batch=64
- eval_every=1 (eval + train recall every epoch)
- Max bs=24 with grad checkpointing on RTX 4090

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

conf/gtauav_balanced.gin

@@ -1,6 +1,6 @@
-# GTA-UAV Balanced: GatedFusion with L1/L2/L3 captions on both branches.
+# GTA-UAV Balanced: Asymmetric DINOv3 (WEB+SAT) with L1/L2/L3 captions.
 # query = sigma(alpha) * drone + (1-sigma(alpha)) * text -> InfoNCE vs gallery
-# 10 epochs, DINOv3 + DGTRS-CLIP, MONA + LoRA adapters.
+# 10 epochs, MONA all 24 blocks, 1024-dim retrieval, hard negative bank.
 #
 # NOTE: TrainConfigGTAUAV is registered by train_gtauav.py before gin parsing.
 # InfoNCELoss is registered via import below.
@@ -15,9 +15,9 @@ TrainConfigGTAUAV.learning_rate = 1e-4
 TrainConfigGTAUAV.text_lr_factor = 0.1
 TrainConfigGTAUAV.weight_decay = 1e-4
 TrainConfigGTAUAV.grad_clip = 1.0
-TrainConfigGTAUAV.grad_accum_steps = 1
+TrainConfigGTAUAV.grad_accum_steps = 8
 TrainConfigGTAUAV.use_amp = True
-TrainConfigGTAUAV.eval_every = 2
+TrainConfigGTAUAV.eval_every = 1
 TrainConfigGTAUAV.warmup_epochs = 2
 TrainConfigGTAUAV.seed = 42
 TrainConfigGTAUAV.device = "cuda"
@@ -25,7 +25,7 @@ TrainConfigGTAUAV.device = "cuda"
 # ---- Model ----
 TrainConfigGTAUAV.init_gate = 0.7
 TrainConfigGTAUAV.baseline_mode = False
-TrainConfigGTAUAV.shared_encoder = True
+TrainConfigGTAUAV.shared_encoder = False
 TrainConfigGTAUAV.gradient_checkpointing = True
 
 # ---- Loss ----
@@ -34,6 +34,7 @@ TrainConfigGTAUAV.label_smoothing = 0.1
 TrainConfigGTAUAV.weight_q2g = 0.6
 TrainConfigGTAUAV.weight_g2q = 0.4
 TrainConfigGTAUAV.learnable_temperature = True
+TrainConfigGTAUAV.neg_bank_size = 4096
 
 # ---- Output ----
 TrainConfigGTAUAV.output_dir = "out/gtauav/with_text"

src/losses/hard_negatives.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+"""Hard negative memory bank for contrastive learning.
+
+MoCo-style FIFO queue of recent gallery embeddings. Each batch gets
+B in-batch negatives + Q queue negatives, significantly increasing
+the effective number of negatives without extra VRAM for forward pass.
+
+Usage:
+    bank = NegativeMemoryBank(size=4096, dim=1024)
+    # In training loop:
+    sim = bank.compute_similarity(query, gallery)  # [B, B + Q]
+    bank.enqueue(gallery.detach())
+"""
+
+import torch
+import torch.nn as nn
+
+
+class NegativeMemoryBank(nn.Module):
+    """FIFO queue of detached gallery embeddings for hard negatives.
+
+    Args:
+        size: Queue capacity (number of stored embeddings).
+        dim: Embedding dimension.
+    """
+
+    def __init__(self, size: int = 4096, dim: int = 1024) -> None:
+        super().__init__()
+        self.size = size
+        self.dim = dim
+        # Queue stored as buffer (not a parameter, moves with .to(device)).
+        self.register_buffer("queue", torch.randn(size, dim))
+        self.queue = nn.functional.normalize(self.queue, dim=-1)
+        self.register_buffer("ptr", torch.zeros(1, dtype=torch.long))
+        self.register_buffer("full", torch.zeros(1, dtype=torch.bool))
+
+    @torch.no_grad()
+    def enqueue(self, embeddings: torch.Tensor) -> None:
+        """Add embeddings to the queue (FIFO). Oldest are overwritten."""
+        batch_size = embeddings.shape[0]
+        ptr = int(self.ptr.item())
+
+        if ptr + batch_size <= self.size:
+            self.queue[ptr:ptr + batch_size] = embeddings.detach()
+        else:
+            # Wrap around.
+            overflow = (ptr + batch_size) - self.size
+            self.queue[ptr:] = embeddings[:batch_size - overflow].detach()
+            self.queue[:overflow] = embeddings[batch_size - overflow:].detach()
+
+        new_ptr = (ptr + batch_size) % self.size
+        self.ptr[0] = new_ptr
+        if not self.full.item() and (new_ptr < ptr or new_ptr == 0):
+            self.full[0] = True
+
+    def get_queue(self) -> torch.Tensor:
+        """Return valid queue entries [Q, dim]."""
+        if self.full.item():
+            return self.queue
+        ptr = int(self.ptr.item())
+        if ptr == 0:
+            return self.queue[:0]  # empty
+        return self.queue[:ptr]
+
+    @property
+    def current_size(self) -> int:
+        if self.full.item():
+            return self.size
+        return int(self.ptr.item())

src/losses/multi_infonce.py

@@ -23,14 +23,36 @@ def _symmetric_info_nce(
     label_smoothing: float,
     weight_a2b: float = 0.5,
     weight_b2a: float = 0.5,
+    queue_negatives: torch.Tensor | None = None,
 ) -> torch.Tensor:
-    """Weighted symmetric InfoNCE. Positives on the diagonal."""
+    """Weighted symmetric InfoNCE with optional hard negative queue.
+
+    Args:
+        emb_a: Query embeddings [B, D].
+        emb_b: Gallery embeddings [B, D]. Positives on diagonal.
+        queue_negatives: Extra gallery negatives [Q, D] from memory bank.
+    """
     batch_size = emb_a.size(0)
-    # Compute logits in fp32 to avoid overflow with small temperature.
+    emb_a_f = emb_a.float()
-    logits = emb_a.float() @ emb_b.float().t() / temperature
+    emb_b_f = emb_b.float()
+
+    if queue_negatives is not None and queue_negatives.shape[0] > 0:
+        # a→b: query sees B in-batch + Q queue negatives.
+        all_b = torch.cat([emb_b_f, queue_negatives.float()], dim=0)  # [B+Q, D]
+        logits_a2b = emb_a_f @ all_b.t() / temperature  # [B, B+Q]
+        targets_a = torch.arange(batch_size, device=emb_a.device)
+        loss_a2b = F.cross_entropy(logits_a2b, targets_a, label_smoothing=label_smoothing)
+
+        # b→a: gallery sees B in-batch queries (no queue for this direction).
+        logits_b2a = emb_b_f @ emb_a_f.t() / temperature  # [B, B]
+        targets_b = torch.arange(batch_size, device=emb_a.device)
+        loss_b2a = F.cross_entropy(logits_b2a, targets_b, label_smoothing=label_smoothing)
+    else:
+        logits = emb_a_f @ emb_b_f.t() / temperature
         targets = torch.arange(batch_size, device=emb_a.device)
         loss_a2b = F.cross_entropy(logits, targets, label_smoothing=label_smoothing)
         loss_b2a = F.cross_entropy(logits.t(), targets, label_smoothing=label_smoothing)
+
     return weight_a2b * loss_a2b + weight_b2a * loss_b2a
 
 
@@ -106,17 +128,18 @@ class InfoNCELoss(nn.Module):
         embeddings: dict[str, torch.Tensor],
         epoch: int,
         total_epochs: int,
+        queue_negatives: torch.Tensor | None = None,
     ) -> dict[str, torch.Tensor]:
-        """Compute InfoNCE loss.
+        """Compute InfoNCE loss with optional hard negative queue.
 
         Args:
-            embeddings: Dict with 'query' and 'gallery' [B, D] L2-normalized,
+            embeddings: Dict with 'query' and 'gallery' [B, D] L2-normalized.
-                        plus 'gate' (float) from fusion module.
             epoch: Current epoch (0-indexed).
             total_epochs: Total epochs for temperature schedule.
+            queue_negatives: Extra gallery negatives [Q, D] from memory bank.
 
         Returns:
-            Dict with 'total', 'temperature', 'gate'.
+            Dict with 'total', 'temperature', 'gate_q', 'gate_g'.
         """
         if self.learnable_temperature:
             # Clamp logit_scale in logit space first to prevent exp() overflow in fp16.
@@ -143,6 +166,7 @@ class InfoNCELoss(nn.Module):
             label_smoothing=self.label_smoothing,
             weight_a2b=self.weight_q2g,
             weight_b2a=self.weight_g2q,
+            queue_negatives=queue_negatives,
         )
 
         gate_q = embeddings.get("gate_q", embeddings.get("gate", 1.0))

src/models/asymmetric_encoder.py

@@ -297,14 +297,14 @@ class AsymmetricEncoder(nn.Module):
         lrsclip_path: str = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt",
         init_gate: float = 0.7,
         baseline_mode: bool = False,
-        shared_encoder: bool = True,
+        shared_encoder: bool = False,
-        embed_dim: int = 512,
         mona_bottleneck: int = 64,
+        mona_last_n_blocks: int = 24,
         lora_rank: int = 4,
         device: str = "cuda",
     ) -> None:
         super().__init__()
-        self.embed_dim = embed_dim
+        self.embed_dim = self.DINO_DIM  # native 1024, no projection
         self.baseline_mode = baseline_mode
         self.shared_encoder = shared_encoder
         self.device = device
@@ -313,20 +313,17 @@ class AsymmetricEncoder(nn.Module):
         if shared_encoder:
             self.image_encoder = DINOv3ViT.from_pretrained(dino_web_path)
             self._freeze(self.image_encoder)
-            inject_mona_into_dinov3(self.image_encoder, bottleneck=mona_bottleneck)
+            inject_mona_into_dinov3(self.image_encoder, bottleneck=mona_bottleneck, last_n_blocks=mona_last_n_blocks)
             LOGGER.info("Shared encoder mode: single DINOv3 WEB for drone + satellite")
         else:
             self.drone_encoder = DINOv3ViT.from_pretrained(dino_web_path)
             self.sat_encoder = DINOv3ViT.from_pretrained(dino_sat_path)
             self._freeze(self.drone_encoder)
             self._freeze(self.sat_encoder)
-            inject_mona_into_dinov3(self.drone_encoder, bottleneck=mona_bottleneck)
+            inject_mona_into_dinov3(self.drone_encoder, bottleneck=mona_bottleneck, last_n_blocks=mona_last_n_blocks)
-            inject_mona_into_dinov3(self.sat_encoder, bottleneck=mona_bottleneck)
+            inject_mona_into_dinov3(self.sat_encoder, bottleneck=mona_bottleneck, last_n_blocks=mona_last_n_blocks)
             LOGGER.info("Asymmetric encoder mode: DINOv3 WEB (drone) + DINOv3 SAT (satellite)")
 
-        # Projection: DINOv3 1024-dim -> embed_dim (512).
-        self.image_projection = nn.Linear(self.DINO_DIM, embed_dim)
-
         # Text encoder — official DGTRS architecture (frozen + LoRA).
         if not baseline_mode:
             self.text_encoder = load_dgtrs_text_encoder(lrsclip_path)
@@ -335,11 +332,11 @@ class AsymmetricEncoder(nn.Module):
         else:
             self.text_encoder = None
 
-        # Shared text fusion MLP: 3×768 -> embed_dim (512).
+        # Shared text fusion MLP: 3×768 -> 1024 (native DINOv3 dim).
         if not baseline_mode:
             self.text_fusion = TextFusionMLP(
                 text_dim=self.TEXT_DIM,
-                out_dim=embed_dim,
+                out_dim=self.DINO_DIM,
             )
 
         # Separate gated fusion for query and gallery branches.
@@ -353,20 +350,16 @@ class AsymmetricEncoder(nn.Module):
         module.eval()
 
     def encode_drone(self, images: torch.Tensor) -> torch.Tensor:
-        """Encode drone images with MONA adapters + projection. Returns [B, embed_dim]."""
+        """Encode drone images with MONA adapters. Returns [B, 1024]."""
         if self.shared_encoder:
-            x = self.image_encoder(images)
+            return self.image_encoder(images)
-        else:
+        return self.drone_encoder(images)
-            x = self.drone_encoder(images)
-        return self.image_projection(x)
 
     def encode_satellite(self, images: torch.Tensor) -> torch.Tensor:
-        """Encode satellite images with MONA adapters + projection. Returns [B, embed_dim]."""
+        """Encode satellite images with MONA adapters. Returns [B, 1024]."""
         if self.shared_encoder:
-            x = self.image_encoder(images)
+            return self.image_encoder(images)
-        else:
+        return self.sat_encoder(images)
-            x = self.sat_encoder(images)
-        return self.image_projection(x)
 
     def encode_text_levels(
         self,
@@ -459,7 +452,6 @@ class AsymmetricEncoder(nn.Module):
             "model_state": self.state_dict(),
             "baseline_mode": self.baseline_mode,
             "shared_encoder": self.shared_encoder,
-            "embed_dim": self.embed_dim,
             **extra,
         }
         tmp = path.with_suffix(path.suffix + ".tmp")
@@ -494,8 +486,7 @@ class AsymmetricEncoder(nn.Module):
             dino_sat_path=dino_sat_path,
             lrsclip_path=lrsclip_path,
             baseline_mode=ckpt.get("baseline_mode", False),
-            shared_encoder=ckpt.get("shared_encoder", True),
+            shared_encoder=ckpt.get("shared_encoder", False),
-            embed_dim=ckpt.get("embed_dim", 512),
             device=device,
         )
         model.load_state_dict(ckpt["model_state"], strict=False)

src/training/train_gtauav.py

@@ -31,6 +31,7 @@ from tqdm import tqdm
 
 from src.datasets.gtauav_dataset import GTAUAVDataset, collate_gtauav_batch
 from src.losses.multi_infonce import InfoNCELoss
+from src.losses.hard_negatives import NegativeMemoryBank
 from src.training.plot_metrics import generate_plots
 from src.training.trackers import ExperimentTracker
 from src.training.grad_monitor import compute_gradient_norms, log_gradient_summary
@@ -73,7 +74,7 @@ class TrainConfigGTAUAV:
     lrsclip_path: str = _LRSCLIP
     init_gate: float = 0.7
     baseline_mode: bool = False
-    shared_encoder: bool = True  # single DINOv3 WEB for both branches (saves ~4-5 GB)
+    shared_encoder: bool = False  # asymmetric: WEB (drone) + SAT (satellite)
     gradient_checkpointing: bool = True  # trade compute for VRAM (allows larger batch)
 
     # Training.
@@ -99,6 +100,7 @@ class TrainConfigGTAUAV:
     weight_q2g: float = 0.6
     weight_g2q: float = 0.4
     learnable_temperature: bool = True
+    neg_bank_size: int = 4096  # hard negative memory bank size (0 = disabled)
 
     # Tracking & diagnostics.
     use_wandb: bool = False
@@ -405,6 +407,7 @@ def train(cfg: TrainConfigGTAUAV) -> None:
             shared_encoder=cfg.shared_encoder,
             device=cfg.device,
         ).to(cfg.device)
+        LOGGER.info("embed_dim=%d", model.embed_dim)
 
     # --- Gradient checkpointing (trade compute for VRAM) ---
     if cfg.gradient_checkpointing:
@@ -446,6 +449,12 @@ def train(cfg: TrainConfigGTAUAV) -> None:
         cfg.tau_init,
     )
 
+    # Hard negative memory bank.
+    neg_bank = None
+    if cfg.neg_bank_size > 0:
+        neg_bank = NegativeMemoryBank(size=cfg.neg_bank_size, dim=model.embed_dim).to(cfg.device)
+        LOGGER.info("Negative memory bank: size=%d, dim=%d", cfg.neg_bank_size, model.embed_dim)
+
     # Data — separate transforms for train (augmented) and eval (clean).
     drone_train_tf = get_drone_train_transform(image_size=256)
     sat_train_tf = get_satellite_train_transform(image_size=256)
@@ -599,13 +608,19 @@ def train(cfg: TrainConfigGTAUAV) -> None:
                         sat_caption_l2=batch["sat_caption_l2"],
                         sat_caption_l3=batch["sat_caption_l3"],
                     )
-            # Loss in fp32 (learnable temperature gradient overflows in fp16).
+            # Loss in fp32 with optional hard negative queue.
+            queue_neg = neg_bank.get_queue() if neg_bank is not None else None
             loss_dict = loss_fn(
                 embeddings=embeddings,
                 epoch=epoch,
                 total_epochs=cfg.epochs,
+                queue_negatives=queue_neg,
             )
 
+            # Enqueue current gallery embeddings (detached).
+            if neg_bank is not None:
+                neg_bank.enqueue(embeddings["gallery"].detach())
+
             # Scale loss by accumulation steps so gradients average correctly.
             raw_loss = float(loss_dict["total"].item())  # save before backward
             total_loss = loss_dict["total"] / accum