Improve all-future first-onset training

losses.py

@@ -37,6 +37,54 @@ def _zero_loss_like(logits: torch.Tensor) -> torch.Tensor:
     return logits.sum() * 0.0
 
 
+def _all_future_at_risk_mask(
+    vocab_size: int,
+    ignored_idx: Iterable[int],
+    logits: torch.Tensor,
+    history: torch.Tensor | None,
+) -> torch.Tensor:
+    """Return outcomes that have not occurred by the query time."""
+    batch_size = logits.shape[0]
+    at_risk = _valid_vocab_mask(
+        vocab_size,
+        ignored_idx,
+        logits.device,
+    ).unsqueeze(0).expand(batch_size, -1).clone()
+    if history is None or history.numel() == 0:
+        return at_risk
+    if history.dim() != 2 or history.shape[0] != batch_size:
+        raise ValueError(
+            "history must be (B, L). "
+            f"Got logits={tuple(logits.shape)}, history={tuple(history.shape)}"
+        )
+
+    history = history.to(device=logits.device, dtype=torch.long)
+    history_valid = (history >= 0) & (history < vocab_size)
+    for idx in ignored_idx:
+        history_valid &= history != int(idx)
+    safe_history = history.clamp(min=0, max=vocab_size - 1)
+    prevalent = torch.zeros(
+        (batch_size, vocab_size),
+        dtype=torch.long,
+        device=logits.device,
+    )
+    prevalent.scatter_add_(1, safe_history, history_valid.long())
+    return at_risk & ~prevalent.bool()
+
+
+def _all_future_target_mask(
+    targets: torch.Tensor,
+    at_risk: torch.Tensor,
+    ignored_idx: Iterable[int],
+) -> tuple[torch.Tensor, torch.Tensor]:
+    vocab_size = at_risk.shape[1]
+    in_vocab = (targets >= 0) & (targets < vocab_size)
+    for idx in ignored_idx:
+        in_vocab &= targets != int(idx)
+    safe_targets = targets.clamp(min=0, max=vocab_size - 1)
+    return in_vocab & at_risk.gather(1, safe_targets), safe_targets
+
+
 class Delphi2MLoss(nn.Module):
     """Next-token plus exponential time-to-next-token supervision."""
 
@@ -146,7 +194,7 @@ class Delphi2MLoss(nn.Module):
 
 
 class ExponentialLoss(nn.Module):
-    """Query-conditioned all-future-event exponential point-process loss."""
+    """First-onset all-future exponential survival likelihood."""
 
     def __init__(
         self,
@@ -162,24 +210,61 @@ class ExponentialLoss(nn.Module):
         logits: torch.Tensor,
         targets: torch.Tensor,
         exposure: torch.Tensor,
+        dt: torch.Tensor,
+        history: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        _, vocab_size = logits.shape
+        batch_size, vocab_size = logits.shape
+        if targets.dim() != 2 or targets.shape[0] != batch_size:
+            raise ValueError(
+                "targets must be (B, M). "
+                f"Got logits={tuple(logits.shape)}, targets={tuple(targets.shape)}"
+            )
+        if exposure.shape != (batch_size,):
+            raise ValueError(
+                f"exposure must be ({batch_size},), got {tuple(exposure.shape)}"
+            )
+        if dt.shape != targets.shape:
+            raise ValueError(
+                f"dt must match targets, got dt={tuple(dt.shape)}, "
+                f"targets={tuple(targets.shape)}"
+            )
+
         rate = F.softplus(logits) + self.eps
-        valid_vocab = _valid_vocab_mask(vocab_size, self.ignored_idx, logits.device)
+        at_risk = _all_future_at_risk_mask(
+            vocab_size,
+            self.ignored_idx,
+            logits,
+            history,
+        )
+        target_valid, safe_targets = _all_future_target_mask(
+            targets,
+            at_risk,
+            self.ignored_idx,
+        )
 
-        penalty = exposure.to(rate.dtype) * rate[:, valid_vocab].sum(dim=-1)
-        target_valid = torch.ones_like(targets, dtype=torch.bool, device=logits.device)
-        for idx in self.ignored_idx:
-            target_valid &= targets != idx
+        censor_time = exposure.to(rate.dtype).clamp_min(self.eps)
+        cumulative_at_censor = rate * censor_time.unsqueeze(1)
+        penalty = (cumulative_at_censor * at_risk.to(rate.dtype)).sum(dim=-1)
+
+        # A first-onset outcome leaves the risk set at its event time, not at
+        # the common end of follow-up.
+        if targets.numel() > 0:
+            event_time = dt.to(rate.dtype).clamp_min(self.eps)
+            event_time = torch.minimum(event_time, censor_time.unsqueeze(1))
+            event_cumulative = rate.gather(1, safe_targets) * event_time
+            censor_cumulative = cumulative_at_censor.gather(1, safe_targets)
+            penalty = penalty + (
+                (event_cumulative - censor_cumulative)
+                * target_valid.to(rate.dtype)
+            ).sum(dim=-1)
 
-        safe_targets = targets.clamp(min=0, max=vocab_size - 1)
         observed = rate.log().gather(1, safe_targets)
         observed = (observed * target_valid.to(rate.dtype)).sum(dim=-1)
         return (-observed + penalty).mean()
 
 
 class WeibullLoss(nn.Module):
-    """Query-conditioned all-future-event Weibull point-process loss."""
+    """First-onset all-future Weibull survival likelihood."""
 
     def __init__(
         self,
@@ -197,8 +282,9 @@ class WeibullLoss(nn.Module):
         targets: torch.Tensor,
         dt: torch.Tensor,
         exposure: torch.Tensor,
+        history: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        _, vocab_size = logits.shape
+        batch_size, vocab_size = logits.shape
         if weibull_rho is None:
             raise ValueError("weibull_rho is required for WeibullLoss")
         if weibull_rho.shape != logits.shape:
@@ -206,23 +292,50 @@ class WeibullLoss(nn.Module):
                 "weibull_rho must have the same shape as logits. "
                 f"Got logits={tuple(logits.shape)}, weibull_rho={tuple(weibull_rho.shape)}"
             )
+        if targets.dim() != 2 or targets.shape[0] != batch_size:
+            raise ValueError(
+                "targets must be (B, M). "
+                f"Got logits={tuple(logits.shape)}, targets={tuple(targets.shape)}"
+            )
+        if dt.shape != targets.shape:
+            raise ValueError(
+                f"dt must match targets, got dt={tuple(dt.shape)}, "
+                f"targets={tuple(targets.shape)}"
+            )
+        if exposure.shape != (batch_size,):
+            raise ValueError(
+                f"exposure must be ({batch_size},), got {tuple(exposure.shape)}"
+            )
 
         dtype = logits.dtype
         rate = F.softplus(logits) + self.eps
         rho = weibull_rho.to(device=logits.device, dtype=dtype).clamp_min(self.eps)
-        valid_vocab = _valid_vocab_mask(vocab_size, self.ignored_idx, logits.device)
+        at_risk = _all_future_at_risk_mask(
+            vocab_size,
+            self.ignored_idx,
+            logits,
+            history,
+        )
+        target_valid, safe_targets = _all_future_target_mask(
+            targets,
+            at_risk,
+            self.ignored_idx,
+        )
 
-        t_exp = exposure.to(dtype).clamp_min(self.eps).unsqueeze(1)
-        penalty = (rate * torch.pow(t_exp, rho))[:, valid_vocab].sum(dim=-1)
+        censor_time = exposure.to(dtype).clamp_min(self.eps)
+        cumulative_at_censor = rate * torch.pow(censor_time.unsqueeze(1), rho)
+        penalty = (cumulative_at_censor * at_risk.to(dtype)).sum(dim=-1)
 
-        target_valid = torch.ones_like(targets, dtype=torch.bool, device=logits.device)
-        for idx in self.ignored_idx:
-            target_valid &= targets != idx
-
-        safe_targets = targets.clamp(min=0, max=vocab_size - 1)
         target_rate = rate.gather(1, safe_targets)
         target_rho = rho.gather(1, safe_targets)
         target_dt = dt.to(dtype).clamp_min(self.eps)
+        target_dt = torch.minimum(target_dt, censor_time.unsqueeze(1))
+        event_cumulative = target_rate * torch.pow(target_dt, target_rho)
+        censor_cumulative = cumulative_at_censor.gather(1, safe_targets)
+        penalty = penalty + (
+            (event_cumulative - censor_cumulative) * target_valid.to(dtype)
+        ).sum(dim=-1)
+
         log_intensity = (
             target_rate.log()
             + target_rho.log()