Add DynamicSimilaritySampler — embedding-kNN batches with mutex constraint

Batches assembled from visually-similar drone queries pressure the model to
learn finer discriminative features. Random mutex batches average ~0.26
pairwise cosine similarity in query embedding space; DSS batches average
~0.71 — confirming the lookalikes grouping works as intended.

Algorithm per batch:
  1. Pick a random seed drone from the remaining pool.
  2. Rank the entire remaining pool by cosine similarity to the seed.
  3. Walk the ranking in descending order; add items whose sat_candidates
     don't collide with the batch's already-claimed set.
  4. Drop the seed if no valid batch can be assembled (rare mutex deadlock).

Inherits MutuallyExclusiveSampler semantics — no false negatives. Degrades
gracefully to mutex-only when no embeddings are set (warmup epochs, or if
`sampler_type="mutex"` is chosen).

Integration in `train_gtauav.py`:
  - New `_embed_drone_queries` helper: model.encode_query forwarded over
    GTAUAVDroneQuery, returns [N, D] CPU tensor. ~13s per 1024 queries on
    a 4090 → ~5 min for the full 25K train set.
  - Epoch loop re-embeds every `dss_reembed_every` epochs after a `dss_warmup_epochs`
    warmup (first epochs use mutex-only since untrained embeddings aren't
    informative for kNN).
  - Config: `sampler_type` ∈ {"mutex", "dss"}. Default flipped to "dss".

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

src/datasets/dynamic_similarity_sampler.py

@@ -0,0 +1,195 @@
+from __future__ import annotations
+
+"""Dynamic Similarity Sampler — hard batch construction via embedding kNN.
+
+Builds on MutuallyExclusiveSampler: keeps the no-false-negatives guarantee
+(disjoint `sat_candidates` within a batch), but additionally pressures the
+model by packing visually-similar drone queries together. The intuition is
+that in-batch negatives which are easy to distinguish contribute little
+gradient — if the batch contains lookalikes, the model has to learn finer
+features.
+
+Workflow:
+  1. Training loop calls `update_embeddings(query_embeddings)` at the start
+     of each epoch (or every N epochs) with current query embeddings.
+  2. `_generate_batches` picks a random seed drone, ranks the rest by cosine
+     similarity to the seed, then greedily packs highest-similarity drones
+     whose `sat_candidates` don't conflict.
+  3. If no embeddings are set yet, falls back to mutex-only behavior (first
+     epoch, or warmup before the model has useful representations).
+"""
+
+import logging
+import random
+from typing import Iterator, Sequence
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data.sampler import Sampler
+
+LOGGER = logging.getLogger("caption_test.dss_sampler")
+
+
+class DynamicSimilaritySampler(Sampler[list[int]]):
+    """Batch sampler: visually-similar drones packed per batch, mutex-preserved.
+
+    Args:
+        sat_candidates_per_item: Valid satellite names per dataset index
+            (positive + semi-positive).
+        batch_size: Target batch size. Partial trailing batches are dropped.
+        shuffle: Randomize seed ordering each epoch.
+        seed: Base RNG seed (effective seed is `seed + epoch`).
+        allow_partial: Yield trailing partial batch if non-empty.
+    """
+
+    def __init__(
+        self,
+        sat_candidates_per_item: Sequence[Sequence[str]],
+        batch_size: int,
+        shuffle: bool = True,
+        seed: int = 0,
+        allow_partial: bool = False,
+    ) -> None:
+        if batch_size <= 0:
+            raise ValueError(f"batch_size must be positive, got {batch_size}")
+        self._item_sats: list[frozenset[str]] = [
+            frozenset(s) for s in sat_candidates_per_item
+        ]
+        self.batch_size = batch_size
+        self.shuffle = shuffle
+        self.seed = seed
+        self.allow_partial = allow_partial
+        self.epoch = 0
+        self._embeddings: torch.Tensor | None = None
+        self._cached_len: int | None = None
+
+    def set_epoch(self, epoch: int) -> None:
+        self.epoch = epoch
+        self._cached_len = None
+
+    def update_embeddings(self, embeddings: torch.Tensor) -> None:
+        """Store query embeddings [N, D] for the next epoch's sampling.
+
+        Call from the training loop before starting each epoch where you want
+        DSS to be active. Stored on CPU as float32 and L2-normalized.
+        """
+        if embeddings.size(0) != len(self._item_sats):
+            raise ValueError(
+                f"embeddings have {embeddings.size(0)} rows but sampler tracks "
+                f"{len(self._item_sats)} items",
+            )
+        self._embeddings = F.normalize(embeddings.detach().cpu().float(), dim=-1)
+        self._cached_len = None
+        LOGGER.info(
+            "DSS embeddings updated: %d × %d", *self._embeddings.shape,
+        )
+
+    def clear_embeddings(self) -> None:
+        """Revert to mutex-only sampling for the next epoch."""
+        self._embeddings = None
+        self._cached_len = None
+
+    def _generate_batches_mutex_only(self) -> list[list[int]]:
+        """Fallback: greedy mutex packing without similarity ranking."""
+        rng = random.Random(self.seed + self.epoch) if self.shuffle else None
+        remaining = list(range(len(self._item_sats)))
+        if rng is not None:
+            rng.shuffle(remaining)
+
+        batches: list[list[int]] = []
+        while remaining:
+            batch: list[int] = []
+            claimed: set[str] = set()
+            next_remaining: list[int] = []
+            for idx in remaining:
+                sats = self._item_sats[idx]
+                if len(batch) < self.batch_size and not (sats & claimed):
+                    batch.append(idx)
+                    claimed |= sats
+                else:
+                    next_remaining.append(idx)
+            if len(batch) == self.batch_size:
+                batches.append(batch)
+            elif self.allow_partial and batch:
+                batches.append(batch)
+                break
+            else:
+                break
+            remaining = next_remaining
+            if rng is not None:
+                rng.shuffle(remaining)
+        return batches
+
+    def _generate_batches_dss(self) -> list[list[int]]:
+        """Similarity-guided batches using stored embeddings."""
+        assert self._embeddings is not None
+        rng = random.Random(self.seed + self.epoch) if self.shuffle else None
+        emb = self._embeddings  # already L2-normalized, on CPU
+        n = emb.size(0)
+
+        remaining = set(range(n))
+        batches: list[list[int]] = []
+
+        while len(remaining) >= self.batch_size:
+            seed_idx = (
+                rng.choice(list(remaining))
+                if rng is not None
+                else next(iter(remaining))
+            )
+
+            # Cosine similarity from seed to all items (fp32 matmul, ~25K ops).
+            sims = emb @ emb[seed_idx]  # [N]
+            order = sims.argsort(descending=True).tolist()
+
+            batch: list[int] = []
+            claimed: set[str] = set()
+            for idx in order:
+                if idx not in remaining:
+                    continue
+                sats = self._item_sats[idx]
+                if sats & claimed:
+                    continue
+                batch.append(idx)
+                claimed |= sats
+                if len(batch) == self.batch_size:
+                    break
+
+            if len(batch) < self.batch_size:
+                # Couldn't fill from this seed's neighborhood (heavy mutex
+                # conflict). Drop the seed and retry with remaining pool.
+                remaining.discard(seed_idx)
+                continue
+
+            batches.append(batch)
+            for i in batch:
+                remaining.discard(i)
+
+        if self.allow_partial and remaining:
+            # Mutex-pack whatever's left.
+            claimed = set()
+            tail: list[int] = []
+            for idx in remaining:
+                sats = self._item_sats[idx]
+                if not (sats & claimed):
+                    tail.append(idx)
+                    claimed |= sats
+            if tail:
+                batches.append(tail)
+
+        return batches
+
+    def _generate_batches(self) -> list[list[int]]:
+        if self._embeddings is None:
+            return self._generate_batches_mutex_only()
+        return self._generate_batches_dss()
+
+    def __iter__(self) -> Iterator[list[int]]:
+        batches = self._generate_batches()
+        self._cached_len = len(batches)
+        for batch in batches:
+            yield batch
+
+    def __len__(self) -> int:
+        if self._cached_len is None:
+            self._cached_len = len(self._generate_batches())
+        return self._cached_len