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>

CLAUDE.md

@@ -101,6 +101,7 @@ Eval: Resize(256) + CenterCrop(256) + ImageNet normalization.
 | `src/models/asymmetric_encoder.py` | DINOv3ViT + TextFusionMLP + AsymmetricEncoder + GatedFusion + encode_query/encode_gallery (per-sample caption mask) |
 | `src/datasets/gtauav_dataset.py` | GTA-UAV-LR loader + L1/L2/L3 captions + GTAUAVSatGallery/GTAUAVDroneQuery (full retrieval eval) |
 | `src/datasets/mutually_exclusive_sampler.py` | BatchSampler: drone'ы в батче не делят sat_candidates (no false negatives) |
+| `src/datasets/dynamic_similarity_sampler.py` | DSS: embedding-kNN + mutex — батчи из визуально похожих drone'ов (hard negatives) |
 | `src/losses/multi_infonce.py` | **Primary:** SymmetricInfoNCE + MoCo queue, learnable τ clamp [0.01, 0.1], weights q2g=0.6 g2q=0.4, `hard_mining_k` для top-K hardest negatives |
 | `src/losses/weighted_infonce.py` | Alternative: per-sample adaptive label smoothing (активируется `loss_type="weighted"`) |
 | `src/losses/hard_negatives.py` | NegativeMemoryBank (MoCo-style FIFO queue 4096 × 1024) |
@@ -213,7 +214,7 @@ Meta-файл `meta/seg_filter.json`: исключение изображени
 - **Scheduler:** linear warmup (2 epochs) + cosine annealing (per-step)
 - **Loss:** SymmetricInfoNCE (q2g=0.6, g2q=0.4) с learnable τ (init=0.07, clamp [0.01, 0.1])
 - **Hard mining:** top-K=512 hardest negatives per query из MoCo queue (размер 4096); `hard_mining_k=0` отключает
-- **Batch sampler:** MutuallyExclusiveSampler — batches disjoint по sat_candidates (на bs=8 сохраняет 100% entries)
+- **Batch sampler:** `sampler_type="dss"` (default) — DynamicSimilaritySampler с re-embedding каждую эпоху: пакует визуально похожих drone'ов в один батч (+hardness) с mutex-constraint (no false negatives). Первая эпоха warmup mutex-only. Средний in-batch cosine sim ~0.71 vs 0.26 у mutex. `sampler_type="mutex"` — без DSS
 - **Eval:** full satellite gallery (~2684 unique tiles для test_20) с multi-match R@K (учитывает все positive/semi-positive)
 - **Augmentations:**
   - Drone: RandomResizedCrop(0.7-1.0), HFlip, Rotation(15°), ColorJitter, Grayscale(5%), GaussianBlur

conf/gtauav_balanced.gin

@@ -35,7 +35,10 @@ TrainConfigGTAUAV.weight_g2q = 0.4
 TrainConfigGTAUAV.neg_bank_size = 4096
 
 # ---- Sampling ----
-TrainConfigGTAUAV.use_mutex_sampler = True
+TrainConfigGTAUAV.sampler_type = "dss"  # "dss" or "mutex"
+TrainConfigGTAUAV.dss_warmup_epochs = 1  # first N epochs use mutex-only (untrained embeds not useful)
+TrainConfigGTAUAV.dss_reembed_every = 1
+TrainConfigGTAUAV.use_mutex_sampler = True  # legacy flag, kept True unless disabling both samplers
 
 # ---- Output ----
 TrainConfigGTAUAV.output_dir = "out/gtauav/with_text"

scripts/smoke_dss.py

@@ -0,0 +1,104 @@
+"""Smoke test for DynamicSimilaritySampler: re-embed + batch composition.
+
+Loads checkpoint, embeds a subset of train drones through model.encode_query,
+feeds embeddings to the sampler, and verifies:
+  - Sampler produces batches
+  - Mutex constraint preserved
+  - Batches differ from plain mutex (sorted by similarity to seed)
+"""
+import time
+import torch
+from torch.utils.data import DataLoader
+
+from src.datasets.dynamic_similarity_sampler import DynamicSimilaritySampler
+from src.datasets.gtauav_dataset import (
+    GTAUAVDataset,
+    GTAUAVDroneQuery,
+    collate_drone_query,
+)
+from src.models.asymmetric_encoder import AsymmetricEncoder, get_dino_transform
+
+CKPT = "out/gtauav/with_text/ckpt_epoch005.pt"
+N_SUBSET = 1024  # small subset for smoke speed
+
+
+def main() -> None:
+    model, _ = AsymmetricEncoder.load_checkpoint(CKPT, device="cuda")
+    model.eval()
+
+    tf = get_dino_transform(image_size=256)
+    ds = GTAUAVDataset(
+        pair_json="meta/train_80.json",
+        filter_meta="meta/seg_filter.json",
+        image_transform=tf,
+    )
+    # Subset the entries for smoke speed
+    ds.entries = ds.entries[:N_SUBSET]
+    print(f"Using {len(ds)} train entries")
+
+    query_ds = GTAUAVDroneQuery(ds)
+    loader = DataLoader(
+        query_ds, batch_size=32, shuffle=False, num_workers=2,
+        collate_fn=collate_drone_query, pin_memory=True,
+    )
+
+    # Embed
+    t0 = time.time()
+    embs = []
+    with torch.no_grad():
+        for batch in loader:
+            drone_img = batch["drone_img"].to("cuda", non_blocking=True)
+            q = model.encode_query(
+                drone_img,
+                batch["caption_l1"], batch["caption_l2"], batch["caption_l3"],
+            )
+            embs.append(q.cpu())
+    emb = torch.cat(embs, dim=0)
+    print(f"Embedded {emb.shape[0]} queries in {time.time()-t0:.1f}s, dim={emb.shape[1]}")
+
+    # Build sampler
+    sat_cand = [e["sat_candidates"] for e in ds.entries]
+    sampler = DynamicSimilaritySampler(sat_cand, batch_size=8, seed=42)
+    sampler.set_epoch(0)
+
+    # Mutex fallback
+    mutex_batches = list(iter(sampler))
+    print(f"Mutex-only: {len(mutex_batches)} batches")
+
+    # DSS mode
+    sampler.update_embeddings(emb)
+    t0 = time.time()
+    dss_batches = list(iter(sampler))
+    print(f"DSS: {len(dss_batches)} batches in {time.time()-t0:.2f}s")
+
+    # Verify mutex invariant in DSS output
+    n_ok = 0
+    for b in dss_batches[:20]:
+        all_sats = set()
+        overlap = False
+        for idx in b:
+            s = set(ds.entries[idx]["sat_candidates"])
+            if s & all_sats:
+                overlap = True; break
+            all_sats |= s
+        if not overlap: n_ok += 1
+    print(f"Mutex preserved: {n_ok}/20 DSS batches clean")
+
+    # Similarity check: within a DSS batch, drones should be more similar
+    # to each other than drones in a random mutex batch.
+    def batch_mean_pairwise_sim(batch):
+        e = emb[batch]
+        e = torch.nn.functional.normalize(e, dim=-1)
+        sim = e @ e.t()
+        # Exclude diagonal
+        mask = ~torch.eye(len(batch), dtype=torch.bool)
+        return sim[mask].mean().item()
+
+    mutex_sim = sum(batch_mean_pairwise_sim(b) for b in mutex_batches[:20]) / 20
+    dss_sim = sum(batch_mean_pairwise_sim(b) for b in dss_batches[:20]) / 20
+    print(f"Mean in-batch cosine sim — mutex: {mutex_sim:.4f}  DSS: {dss_sim:.4f}")
+    print(f"  → DSS batches {'MORE' if dss_sim > mutex_sim else 'LESS'} visually similar (expected: MORE)")
+
+
+if __name__ == "__main__":
+    main()

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

src/training/train_gtauav.py

@@ -38,6 +38,7 @@ from src.datasets.gtauav_dataset import (
     collate_gtauav_batch,
     collate_sat_gallery,
 )
+from src.datasets.dynamic_similarity_sampler import DynamicSimilaritySampler
 from src.datasets.mutually_exclusive_sampler import MutuallyExclusiveSampler
 from src.losses.multi_infonce import InfoNCELoss
 from src.losses.weighted_infonce import WeightedInfoNCELoss
@@ -114,7 +115,11 @@ class TrainConfigGTAUAV:
     neg_bank_size: int = 4096  # hard negative memory bank size (0 = disabled)
 
     # Sampling.
-    use_mutex_sampler: bool = True  # Mutually exclusive batches (no false negatives).
+    sampler_type: str = "mutex"  # "mutex" (no false negatives) or "dss" (DSS + mutex)
+    dss_reembed_every: int = 1  # Re-embed train queries every N epochs for DSS.
+    dss_warmup_epochs: int = 1  # Use mutex-only for the first N epochs (fresh model embeddings aren't useful)
+    # Legacy alias kept for backward compatibility.
+    use_mutex_sampler: bool = True
 
     # Tracking & diagnostics.
     use_wandb: bool = False
@@ -186,6 +191,46 @@ def _cosine_warmup_schedule(
     return lr_lambda
 
 
+@torch.no_grad()
+def _embed_drone_queries(
+    model: AsymmetricEncoder,
+    train_ds: GTAUAVDataset,
+    device: str,
+    batch_size: int,
+    num_workers: int,
+) -> torch.Tensor:
+    """Forward all drone queries and return [N, D] embeddings on CPU.
+
+    Used by DynamicSimilaritySampler to rank drones by visual similarity.
+    Runs with model.eval() but restores original train state afterwards.
+    """
+    was_training = model.training
+    model.eval()
+
+    query_ds = GTAUAVDroneQuery(train_ds)
+    loader = DataLoader(
+        query_ds,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        collate_fn=collate_drone_query,
+        pin_memory=True,
+    )
+
+    embs: list[torch.Tensor] = []
+    for batch in tqdm(loader, desc="  dss-embed-queries", unit="batch", leave=False):
+        drone_img = batch["drone_img"].to(device, non_blocking=True)
+        q = model.encode_query(
+            drone_img,
+            batch["caption_l1"], batch["caption_l2"], batch["caption_l3"],
+        )
+        embs.append(q.cpu())
+
+    if was_training:
+        model.train()
+    return torch.cat(embs, dim=0)
+
+
 @torch.no_grad()
 def _evaluate(
     model: AsymmetricEncoder,
@@ -570,25 +615,37 @@ def train(cfg: TrainConfigGTAUAV) -> None:
         filter_meta=cfg.filter_meta,
     )
 
-    if cfg.use_mutex_sampler:
-        mutex_sampler = MutuallyExclusiveSampler(
-            [entry["sat_candidates"] for entry in train_ds.entries],
-            batch_size=cfg.batch_size,
-            shuffle=True,
-            seed=cfg.seed,
+    sat_cand_list = [entry["sat_candidates"] for entry in train_ds.entries]
+    # Backward compat: `use_mutex_sampler=False` overrides to plain shuffle.
+    effective_sampler_type = cfg.sampler_type if cfg.use_mutex_sampler else "none"
+
+    if effective_sampler_type == "dss":
+        batch_sampler = DynamicSimilaritySampler(
+            sat_cand_list, batch_size=cfg.batch_size, shuffle=True, seed=cfg.seed,
         )
         LOGGER.info(
-            "Sampler: MutuallyExclusive — no false negatives within a batch",
+            "Sampler: DynamicSimilarity — embedding-ranked batches with mutex constraint "
+            "(warmup=%d epochs mutex-only, re-embed every %d epochs)",
+            cfg.dss_warmup_epochs, cfg.dss_reembed_every,
         )
+    elif effective_sampler_type == "mutex":
+        batch_sampler = MutuallyExclusiveSampler(
+            sat_cand_list, batch_size=cfg.batch_size, shuffle=True, seed=cfg.seed,
+        )
+        LOGGER.info("Sampler: MutuallyExclusive — no false negatives within a batch")
+    else:
+        batch_sampler = None
+        LOGGER.info("Sampler: default shuffle (no mutex / no DSS)")
+
+    if batch_sampler is not None:
         train_loader = DataLoader(
             train_ds,
-            batch_sampler=mutex_sampler,
+            batch_sampler=batch_sampler,
             num_workers=cfg.num_workers,
             collate_fn=collate_gtauav_batch,
             pin_memory=True,
         )
     else:
-        mutex_sampler = None
         train_loader = DataLoader(
             train_ds,
             batch_size=cfg.batch_size,
@@ -687,8 +744,25 @@ def train(cfg: TrainConfigGTAUAV) -> None:
 
     for epoch in range(start_epoch, cfg.epochs):
         model.train()
-        if mutex_sampler is not None:
-            mutex_sampler.set_epoch(epoch)
+        if batch_sampler is not None:
+            batch_sampler.set_epoch(epoch)
+
+        # DSS re-embedding: refresh query embeddings before the epoch starts.
+        if (
+            isinstance(batch_sampler, DynamicSimilaritySampler)
+            and epoch >= cfg.dss_warmup_epochs
+            and (epoch - cfg.dss_warmup_epochs) % cfg.dss_reembed_every == 0
+        ):
+            LOGGER.info("DSS: re-embedding %d train queries (epoch=%d)", len(train_ds), epoch)
+            t_embed = time.time()
+            query_embs = _embed_drone_queries(
+                model, train_ds, cfg.device,
+                batch_size=cfg.batch_size * cfg.grad_accum_steps,
+                num_workers=cfg.num_workers,
+            )
+            batch_sampler.update_embeddings(query_embs)
+            LOGGER.info("DSS: re-embed took %.1fs", time.time() - t_embed)
+
         epoch_start = time.time()
         agg: dict[str, float] = {}
         n_batches = 0