caption-test/src/losses/multi_infonce.py

from __future__ import annotations

"""InfoNCE loss for cross-view geo-localization with optional text fusion.

Single symmetric InfoNCE between query (drone+text fused) and gallery (satellite).
Asymmetric weighting: query->gallery weighted higher (real use-case direction).

Supports both learnable temperature (CLIP-style logit_scale) and fixed/scheduled.
"""

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 | torch.Tensor,
    label_smoothing: float,
    weight_a2b: float = 0.5,
    weight_b2a: float = 0.5,
    queue_negatives: torch.Tensor | None = None,
) -> torch.Tensor:
    """Weighted symmetric InfoNCE with optional hard negative queue.

    Args:
        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.
    """
    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]
        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]
        targets_b = torch.arange(batch_size, device=emb_a.device)
        loss_b2a = F.cross_entropy(logits_b2a, targets_b, label_smoothing=label_smoothing)
    else:
        logits = emb_a_f @ emb_b_f.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."""
    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


@gin.configurable
class InfoNCELoss(nn.Module):
    """Symmetric InfoNCE with learnable or scheduled temperature.

    Args:
        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.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
        self.temperature_final = temperature_final
        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,
        embeddings: dict[str, torch.Tensor],
        epoch: int,
        total_epochs: int,
        queue_negatives: torch.Tensor | None = None,
    ) -> dict[str, torch.Tensor]:
        """Compute InfoNCE loss with optional hard negative queue.

        Args:
            embeddings: Dict with 'query' and 'gallery' [B, D] L2-normalized.
            epoch: Current epoch (0-indexed).
            total_epochs: Total epochs for temperature schedule.
            queue_negatives: Extra gallery negatives [Q, D] from memory bank.

        Returns:
            Dict with 'total', 'temperature', 'gate_q', 'gate_g'.
        """
        if self.learnable_temperature:
            # Clamp logit_scale in logit space first to prevent exp() overflow in fp16.
            # tau_min=0.01 -> max logit_scale=ln(1/0.01)=4.6
            # tau_max=0.5  -> min logit_scale=ln(1/0.5)=0.69
            clamped = self.logit_scale.float().clamp(
                min=math.log(1.0 / self.tau_max),
                max=math.log(1.0 / self.tau_min),
            )
            logit_scale = clamped.exp()
            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"],
            emb_b=embeddings["gallery"],
            temperature=tau,
            label_smoothing=self.label_smoothing,
            weight_a2b=self.weight_q2g,
            weight_b2a=self.weight_g2q,
            queue_negatives=queue_negatives,
        )

        gate_q = embeddings.get("gate_q", embeddings.get("gate", 1.0))
        gate_g = embeddings.get("gate_g", 1.0)

        if isinstance(tau, float):
            tau_out = torch.tensor(tau, device=loss.device)
        else:
            tau_out = tau.detach().clone()

        return {
            "total": loss,
            "temperature": tau_out,
            "gate_q": torch.tensor(gate_q, device=loss.device),
            "gate_g": torch.tensor(gate_g, device=loss.device),
        }