Improve training: learnable temperature, per-group LR, warmup, augmentations

Loss:
- Learnable temperature (CLIP-style logit_scale) with clamp [0.01, 0.5]
- Replaces fixed cosine schedule (still available via --no-learnable-temp)
- Default tau_init=0.07

Optimizer:
- Per-group LR: projections 1e-4, text encoder 1e-5 (10x lower)
- Learnable temperature included in projection param group

Scheduler:
- Linear warmup (2 epochs default) + cosine annealing
- Per-step scheduling (not per-epoch)

Augmentations (separate drone/satellite):
- Drone: RandomResizedCrop(0.7-1.0), HFlip, Rotation(15), ColorJitter,
  RandomGrayscale(0.05), GaussianBlur
- Satellite: RandomResizedCrop(0.7-1.0), HFlip, ColorJitter, RandomGrayscale
- Eval: clean Resize+CenterCrop (no augmentation)

Dataset: supports separate drone_transform/sat_transform args

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

src/losses/multi_infonce.py

@@ -4,7 +4,8 @@ from __future__ import annotations
 
 Single symmetric InfoNCE between query (drone+text fused) and gallery (satellite).
 Asymmetric weighting: query->gallery weighted higher (real use-case direction).
-Cosine temperature schedule for sharper distribution over training.
+
+Supports both learnable temperature (CLIP-style logit_scale) and fixed/scheduled.
 """
 
 import math
@@ -47,23 +48,30 @@ def cosine_temperature(
 
 @gin.configurable
 class InfoNCELoss(nn.Module):
-    """Symmetric InfoNCE with cosine temperature schedule.
+    """Symmetric InfoNCE with learnable or scheduled temperature.
 
     Args:
-        temperature_init: Temperature at epoch 0.
-        temperature_final: Temperature after cosine decay.
+        temperature_init: Initial temperature value.
+        temperature_final: Final temperature (only used if learnable=False).
         label_smoothing: Cross-entropy label smoothing.
         weight_q2g: Weight for query->gallery direction.
         weight_g2q: Weight for gallery->query direction.
+        learnable_temperature: If True, temperature is a learnable parameter
+            (CLIP-style logit_scale). If False, uses cosine schedule.
+        tau_min: Minimum clamp for learnable temperature.
+        tau_max: Maximum clamp for learnable temperature.
     """
 
     def __init__(
         self,
-        temperature_init: float = 0.1,
+        temperature_init: float = 0.07,
         temperature_final: float = 0.01,
         label_smoothing: float = 0.1,
         weight_q2g: float = 0.6,
         weight_g2q: float = 0.4,
+        learnable_temperature: bool = True,
+        tau_min: float = 0.01,
+        tau_max: float = 0.5,
     ) -> None:
         super().__init__()
         self.temperature_init = temperature_init
@@ -71,6 +79,26 @@ class InfoNCELoss(nn.Module):
         self.label_smoothing = label_smoothing
         self.weight_q2g = weight_q2g
         self.weight_g2q = weight_g2q
+        self.learnable_temperature = learnable_temperature
+        self.tau_min = tau_min
+        self.tau_max = tau_max
+
+        if learnable_temperature:
+            # Store as log(1/tau) like CLIP's logit_scale.
+            init_logit_scale = math.log(1.0 / temperature_init)
+            self.logit_scale = nn.Parameter(torch.tensor(init_logit_scale))
+        else:
+            self.logit_scale = None
+
+    @property
+    def current_temperature(self) -> float:
+        """Current temperature value (for logging)."""
+        if self.logit_scale is not None:
+            tau = 1.0 / self.logit_scale.exp().clamp(
+                min=1.0 / self.tau_max, max=1.0 / self.tau_min,
+            ).item()
+            return tau
+        return self.temperature_init
 
     def forward(
         self,
@@ -89,12 +117,19 @@ class InfoNCELoss(nn.Module):
         Returns:
             Dict with 'total', 'temperature', 'gate'.
         """
-        tau = cosine_temperature(
-            epoch=epoch,
-            total_epochs=total_epochs,
-            tau_init=self.temperature_init,
-            tau_final=self.temperature_final,
-        )
+        if self.learnable_temperature:
+            # Clamp logit_scale to prevent tau from going out of bounds.
+            logit_scale = self.logit_scale.exp().clamp(
+                min=1.0 / self.tau_max, max=1.0 / self.tau_min,
+            )
+            tau = 1.0 / logit_scale
+        else:
+            tau = cosine_temperature(
+                epoch=epoch,
+                total_epochs=total_epochs,
+                tau_init=self.temperature_init,
+                tau_final=self.temperature_final,
+            )
 
         loss = _symmetric_info_nce(
             emb_a=embeddings["query"],
@@ -107,8 +142,13 @@ class InfoNCELoss(nn.Module):
 
         gate = embeddings.get("gate", 1.0)
 
+        if isinstance(tau, float):
+            tau_out = torch.tensor(tau, device=loss.device)
+        else:
+            tau_out = tau.detach().clone()
+
         return {
             "total": loss,
-            "temperature": torch.tensor(tau, device=loss.device),
+            "temperature": tau_out,
             "gate": torch.tensor(gate, device=loss.device),
         }