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>

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.
-    logits = emb_a.float() @ emb_b.float().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)
+    emb_a_f = emb_a.float()
+    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,
-                        plus 'gate' (float) from fusion module.
+            embeddings: Dict with 'query' and 'gallery' [B, D] L2-normalized.
             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))