Add per-query hard negative mining to InfoNCELoss

New `hard_mining_k` parameter on InfoNCELoss. When >0 and queue is non-empty,
each query row keeps only its K highest-similarity queue entries (via
`torch.topk`) as negatives, instead of the full queue. Fully vectorized —
no Python loop, no extra forward pass.

Rationale: the memory bank grows to 4096 detached gallery embeddings, but
most are easy negatives that contribute almost nothing to the gradient.
Hard mining focuses compute on the small subset that actually shapes the
decision boundary. +2-3% R@1 in similar contrastive setups.

Edge cases:
- K=0: mining disabled, full queue used (original behavior).
- K >= queue size: falls back to full queue (e.g. warmup when queue is small).
- Queue empty: in-batch only, no changes.

Default in `gtauav_balanced.gin`: K=512 (1/8 of queue). Smoke-train updated
to exercise the full memory-bank + mining path.

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

src/losses/multi_infonce.py

@@ -24,6 +24,7 @@ def _symmetric_info_nce(
     weight_a2b: float = 0.5,
     weight_b2a: float = 0.5,
     queue_negatives: torch.Tensor | None = None,
+    hard_mining_k: int = 0,
 ) -> torch.Tensor:
     """Weighted symmetric InfoNCE with optional hard negative queue.
 
@@ -31,20 +32,31 @@ def _symmetric_info_nce(
         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.
+        hard_mining_k: If > 0 and queue is non-empty, use only the top-K
+            hardest (highest-similarity) queue entries per query instead
+            of the full queue. Per-query selection — each row gets its
+            own K negatives gathered via `topk`.
     """
     batch_size = emb_a.size(0)
     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]
+        queue_f = queue_negatives.float()
+        sim_inbatch = emb_a_f @ emb_b_f.t() / temperature  # [B, B]
+        sim_queue = emb_a_f @ queue_f.t() / temperature    # [B, Q]
+
+        if hard_mining_k > 0 and hard_mining_k < queue_f.shape[0]:
+            # Per-row top-K — each query gets its own hardest negatives.
+            sim_queue, _ = sim_queue.topk(k=hard_mining_k, dim=1)  # [B, K]
+
+        # a→b: [B, B + (Q or K)]. Positive at column `i` for row `i`.
+        logits_a2b = torch.cat([sim_inbatch, sim_queue], dim=1)
         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]
+        # b→a: gallery sees B in-batch queries (queue is gallery-side, irrelevant here).
+        logits_b2a = sim_inbatch.t()  # [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:
@@ -83,6 +95,9 @@ class InfoNCELoss(nn.Module):
             (CLIP-style logit_scale). If False, uses cosine schedule.
         tau_min: Minimum clamp for learnable temperature.
         tau_max: Maximum clamp for learnable temperature.
+        hard_mining_k: If > 0, mine top-K hardest negatives per query from
+            the memory bank queue instead of using the full queue. 0 disables
+            mining (queue used whole). Typical values: 256-1024 for queue=4096.
     """
 
     def __init__(
@@ -95,6 +110,7 @@ class InfoNCELoss(nn.Module):
         learnable_temperature: bool = True,
         tau_min: float = 0.01,
         tau_max: float = 0.1,
+        hard_mining_k: int = 0,
     ) -> None:
         super().__init__()
         self.temperature_init = temperature_init
@@ -105,6 +121,7 @@ class InfoNCELoss(nn.Module):
         self.learnable_temperature = learnable_temperature
         self.tau_min = tau_min
         self.tau_max = tau_max
+        self.hard_mining_k = hard_mining_k
 
         if learnable_temperature:
             # Store as log(1/tau) like CLIP's logit_scale.
@@ -167,6 +184,7 @@ class InfoNCELoss(nn.Module):
             weight_a2b=self.weight_q2g,
             weight_b2a=self.weight_g2q,
             queue_negatives=queue_negatives,
+            hard_mining_k=self.hard_mining_k,
         )
 
         gate_q = embeddings.get("gate_q", embeddings.get("gate", 1.0))