Initial commit: caption quality test on UAV-VisLoc

Self-contained experimental track validating generated text captions
via retrieval R@1 lift on UAV-VisLoc.

Architecture: GeoRSCLIP ViT-B/32 dual encoder, 512-dim shared space.
Loss: 4-term InfoNCE (img-img + sat-cap + drone-cap + cap-cap)
      with cosine temperature decay, PALW-like curriculum.
Metric: delta R@1 (with text - without text) >= +3% => PASS.

Gin-configured (balanced / baseline_no_text / text_heavy variants).
Follows NADEZHDA code style.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

src/losses/__init__.py

@@ -0,0 +1,8 @@
+"""Loss functions for caption quality test."""
+from src.losses.multi_infonce import (
+    MultiTermInfoNCE,
+    cosine_temperature,
+    curriculum_lambdas,
+)
+
+__all__ = ["MultiTermInfoNCE", "cosine_temperature", "curriculum_lambdas"]

src/losses/multi_infonce.py

@@ -0,0 +1,262 @@
+from __future__ import annotations
+
+"""Multi-term InfoNCE loss for caption quality validation.
+
+Four InfoNCE terms over projected embeddings:
+    L = lambda_ii * L_img_img
+      + lambda_sc * L_sat_cap
+      + lambda_dc * L_drone_cap
+      + lambda_cc * L_cap_cap
+where L_img_img is the classical symmetric CVGL contrastive loss
+with asymmetric weights (0.6 drone->sat + 0.4 sat->drone).
+"""
+
+import math
+
+import gin
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def _symmetric_info_nce(
+    emb_a: torch.Tensor,
+    emb_b: torch.Tensor,
+    temperature: float,
+    label_smoothing: float,
+    weight_a2b: float = 0.5,
+    weight_b2a: float = 0.5,
+) -> torch.Tensor:
+    """Compute weighted symmetric InfoNCE between two L2-normalized embeddings.
+
+    Args:
+        emb_a: First embedding set [B, D].
+        emb_b: Second embedding set [B, D]. Positive pairs are on the diagonal.
+        temperature: Softmax temperature (smaller = sharper distribution).
+        label_smoothing: Cross-entropy label smoothing epsilon.
+        weight_a2b: Weight for A-query direction.
+        weight_b2a: Weight for B-query direction.
+
+    Returns:
+        Scalar weighted loss.
+    """
+    batch_size = emb_a.size(0)
+    logits = emb_a @ emb_b.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
+
+
+def cosine_temperature(
+    epoch: int,
+    total_epochs: int,
+    tau_init: float = 0.1,
+    tau_final: float = 0.01,
+) -> float:
+    """Cosine-decay schedule for InfoNCE temperature.
+
+    Args:
+        epoch: Current training epoch (0-indexed).
+        total_epochs: Total number of epochs.
+        tau_init: Initial temperature.
+        tau_final: Final temperature.
+
+    Returns:
+        Temperature value for this epoch.
+    """
+    total_epochs = max(total_epochs, 1)
+    progress = min(max(epoch / total_epochs, 0.0), 1.0)
+    cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
+    return tau_final + (tau_init - tau_final) * cosine
+
+
+def curriculum_lambdas(
+    epoch: int,
+    warmup_epochs: int = 3,
+    text_ramp_epochs: int = 10,
+    lambda_ii: float = 1.0,
+    lambda_sc_max: float = 0.3,
+    lambda_dc_max: float = 0.3,
+    lambda_cc_max: float = 0.1,
+) -> dict[str, float]:
+    """Compute per-epoch loss weights under the curriculum schedule.
+
+    - Epochs 0..warmup_epochs: image-image only.
+    - Epochs warmup..text_ramp_epochs: linearly ramp sat-cap and drone-cap.
+    - Epochs >= text_ramp_epochs: full loss including caption-caption term.
+
+    Args:
+        epoch: Current epoch (0-indexed).
+        warmup_epochs: Number of warmup epochs (no text losses).
+        text_ramp_epochs: Epoch when text losses reach max.
+        lambda_ii: Constant weight for image-image loss.
+        lambda_sc_max: Max weight for satellite-caption loss.
+        lambda_dc_max: Max weight for drone-caption loss.
+        lambda_cc_max: Max weight for caption-caption loss.
+
+    Returns:
+        Dict with keys 'img_img', 'sat_cap', 'drone_cap', 'cap_cap'.
+    """
+    if epoch < warmup_epochs:
+        ramp = 0.0
+    elif epoch >= text_ramp_epochs:
+        ramp = 1.0
+    else:
+        denom = max(text_ramp_epochs - warmup_epochs, 1)
+        ramp = (epoch - warmup_epochs) / denom
+
+    return {
+        "img_img": lambda_ii,
+        "sat_cap": lambda_sc_max * ramp,
+        "drone_cap": lambda_dc_max * ramp,
+        "cap_cap": lambda_cc_max * ramp,
+    }
+
+
+@gin.configurable
+class MultiTermInfoNCE(nn.Module):
+    """Multi-term InfoNCE loss with curriculum and cosine temperature.
+
+    Produces total loss and per-component diagnostics. All inputs must be
+    L2-normalized embeddings of the same dimension.
+
+    Args:
+        temperature_init: Initial temperature (epoch 0).
+        temperature_final: Final temperature after cosine decay.
+        label_smoothing: Cross-entropy label smoothing epsilon.
+        asym_drone_to_sat: Weight for drone->sat InfoNCE direction.
+        asym_sat_to_drone: Weight for sat->drone InfoNCE direction.
+        warmup_epochs: Epochs with image-image loss only.
+        text_ramp_epochs: Epoch at which text losses reach max.
+        lambda_ii: Constant weight for image-image loss.
+        lambda_sc_max: Max weight for sat-caption loss.
+        lambda_dc_max: Max weight for drone-caption loss.
+        lambda_cc_max: Max weight for caption-caption loss.
+    """
+
+    def __init__(
+        self,
+        temperature_init: float = 0.1,
+        temperature_final: float = 0.01,
+        label_smoothing: float = 0.1,
+        asym_drone_to_sat: float = 0.6,
+        asym_sat_to_drone: float = 0.4,
+        warmup_epochs: int = 3,
+        text_ramp_epochs: int = 10,
+        lambda_ii: float = 1.0,
+        lambda_sc_max: float = 0.3,
+        lambda_dc_max: float = 0.3,
+        lambda_cc_max: float = 0.1,
+    ) -> None:
+        super().__init__()
+        self.temperature_init = temperature_init
+        self.temperature_final = temperature_final
+        self.label_smoothing = label_smoothing
+        self.asym_drone_to_sat = asym_drone_to_sat
+        self.asym_sat_to_drone = asym_sat_to_drone
+        self.warmup_epochs = warmup_epochs
+        self.text_ramp_epochs = text_ramp_epochs
+        self.lambda_ii = lambda_ii
+        self.lambda_sc_max = lambda_sc_max
+        self.lambda_dc_max = lambda_dc_max
+        self.lambda_cc_max = lambda_cc_max
+
+    def forward(
+        self,
+        embeddings: dict[str, torch.Tensor],
+        epoch: int,
+        total_epochs: int,
+    ) -> dict[str, torch.Tensor]:
+        """Compute multi-term loss.
+
+        Args:
+            embeddings: Dict with keys 'drone', 'sat', and optionally
+                'cap_drone', 'cap_sat'. Each [B, D] L2-normalized.
+            epoch: Current epoch (0-indexed).
+            total_epochs: Total epochs for temperature schedule.
+
+        Returns:
+            Dict with scalar tensors: 'total', 'img_img', 'sat_cap',
+            'drone_cap', 'cap_cap', plus 'temperature' and 'lambdas'.
+        """
+        tau = cosine_temperature(
+            epoch=epoch,
+            total_epochs=total_epochs,
+            tau_init=self.temperature_init,
+            tau_final=self.temperature_final,
+        )
+        lambdas = curriculum_lambdas(
+            epoch=epoch,
+            warmup_epochs=self.warmup_epochs,
+            text_ramp_epochs=self.text_ramp_epochs,
+            lambda_ii=self.lambda_ii,
+            lambda_sc_max=self.lambda_sc_max,
+            lambda_dc_max=self.lambda_dc_max,
+            lambda_cc_max=self.lambda_cc_max,
+        )
+
+        drone = embeddings["drone"]
+        sat = embeddings["sat"]
+
+        # Image-image symmetric InfoNCE with asymmetric weights.
+        loss_ii = _symmetric_info_nce(
+            emb_a=drone,
+            emb_b=sat,
+            temperature=tau,
+            label_smoothing=self.label_smoothing,
+            weight_a2b=self.asym_drone_to_sat,
+            weight_b2a=self.asym_sat_to_drone,
+        )
+
+        loss_sc = torch.zeros_like(loss_ii)
+        loss_dc = torch.zeros_like(loss_ii)
+        loss_cc = torch.zeros_like(loss_ii)
+
+        if "cap_sat" in embeddings and lambdas["sat_cap"] > 0:
+            loss_sc = _symmetric_info_nce(
+                emb_a=sat,
+                emb_b=embeddings["cap_sat"],
+                temperature=tau,
+                label_smoothing=self.label_smoothing,
+            )
+        if "cap_drone" in embeddings and lambdas["drone_cap"] > 0:
+            loss_dc = _symmetric_info_nce(
+                emb_a=drone,
+                emb_b=embeddings["cap_drone"],
+                temperature=tau,
+                label_smoothing=self.label_smoothing,
+            )
+        if (
+            "cap_drone" in embeddings
+            and "cap_sat" in embeddings
+            and lambdas["cap_cap"] > 0
+        ):
+            loss_cc = _symmetric_info_nce(
+                emb_a=embeddings["cap_drone"],
+                emb_b=embeddings["cap_sat"],
+                temperature=tau,
+                label_smoothing=self.label_smoothing,
+            )
+
+        total = (
+            lambdas["img_img"] * loss_ii
+            + lambdas["sat_cap"] * loss_sc
+            + lambdas["drone_cap"] * loss_dc
+            + lambdas["cap_cap"] * loss_cc
+        )
+
+        return {
+            "total": total,
+            "img_img": loss_ii.detach(),
+            "sat_cap": loss_sc.detach(),
+            "drone_cap": loss_dc.detach(),
+            "cap_cap": loss_cc.detach(),
+            "temperature": torch.tensor(tau, device=total.device),
+            "lambda_ii": torch.tensor(lambdas["img_img"], device=total.device),
+            "lambda_sc": torch.tensor(lambdas["sat_cap"], device=total.device),
+            "lambda_dc": torch.tensor(lambdas["drone_cap"], device=total.device),
+            "lambda_cc": torch.tensor(lambdas["cap_cap"], device=total.device),
+        }