from __future__ import annotations from typing import Iterable import torch import torch.nn as nn import torch.nn.functional as F PAD_IDX = 0 RESERVED_IDX = 1 NO_EVENT_IDX = 2 def _make_ignore_mask( vocab_size: int, ignored_idx: Iterable[int], device: torch.device, ) -> torch.Tensor: ignore_mask = torch.zeros(vocab_size, dtype=torch.bool, device=device) for idx in ignored_idx: idx = int(idx) if 0 <= idx < vocab_size: ignore_mask[idx] = True return ignore_mask def _valid_vocab_mask( vocab_size: int, ignored_idx: Iterable[int], device: torch.device, ) -> torch.Tensor: return ~_make_ignore_mask(vocab_size, ignored_idx, device) 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.""" def __init__( self, t_min: float = 1.0 / 365.25, ignored_tokens: Iterable[int] | None = None, exclude_ignored_from_intensity: bool = True, max_exp_input: float = 60.0, ce_weight: float = 1.0, time_weight: float = 1.0, ): super().__init__() self.t_min = float(t_min) self.ignored_tokens = ( [PAD_IDX, RESERVED_IDX] if ignored_tokens is None else [int(x) for x in ignored_tokens] ) self.exclude_ignored_from_intensity = bool(exclude_ignored_from_intensity) self.max_exp_input = float(max_exp_input) self.ce_weight = float(ce_weight) self.time_weight = float(time_weight) def forward( self, logits: torch.Tensor, target_events: torch.Tensor, target_times: torch.Tensor, current_times: torch.Tensor, padding_mask: torch.Tensor, return_components: bool = False, ) -> torch.Tensor | tuple[torch.Tensor, dict[str, torch.Tensor]]: if logits.dim() != 3: raise ValueError(f"logits must be (B, L, K), got {tuple(logits.shape)}") bsz, seq_len, vocab_size = logits.shape expected = (bsz, seq_len) if target_events.shape != expected: raise ValueError(f"target_events must be {expected}, got {tuple(target_events.shape)}") if target_times.shape != expected: raise ValueError(f"target_times must be {expected}, got {tuple(target_times.shape)}") if current_times.shape != expected: raise ValueError(f"current_times must be {expected}, got {tuple(current_times.shape)}") if padding_mask.shape != expected: raise ValueError(f"padding_mask must be {expected}, got {tuple(padding_mask.shape)}") valid_mask = padding_mask.bool() for idx in self.ignored_tokens: valid_mask = valid_mask & (target_events != int(idx)) valid_mask = valid_mask & (target_events > PAD_IDX) if not valid_mask.any(): total_loss = _zero_loss_like(logits) if return_components: return total_loss, { "ce": total_loss.detach(), "time": total_loss.detach(), "total": total_loss.detach(), } return total_loss logits_valid = logits[valid_mask] target_events_valid = target_events[valid_mask] target_times_valid = target_times[valid_mask] current_times_valid = current_times[valid_mask] logits_safe = torch.nan_to_num( logits_valid, nan=0.0, posinf=self.max_exp_input, neginf=-self.max_exp_input, ) loss_ce = F.cross_entropy( logits_safe, target_events_valid, reduction="mean", ) logits_for_lse = logits_safe if self.exclude_ignored_from_intensity: ignore_mask = _make_ignore_mask(vocab_size, self.ignored_tokens, logits.device) logits_for_lse = logits_safe.masked_fill(ignore_mask.unsqueeze(0), float("-inf")) log_lambda_total = torch.logsumexp(logits_for_lse, dim=-1) log_lambda_total = -torch.log(torch.exp(-log_lambda_total) + self.t_min) dt = torch.clamp(target_times_valid - current_times_valid, min=self.t_min) log_dt_inv = -torch.log(dt + self.t_min) loss_dt = -( log_lambda_total - torch.exp( torch.clamp(log_lambda_total - log_dt_inv, max=self.max_exp_input) ) ) loss_dt = loss_dt.mean() total_loss = self.ce_weight * loss_ce + self.time_weight * loss_dt if return_components: return total_loss, { "ce": loss_ce.detach(), "time": loss_dt.detach(), "total": total_loss.detach(), } return total_loss class ExponentialLoss(nn.Module): """First-onset all-future exponential survival likelihood.""" def __init__( self, ignored_idx: Iterable[int] = (PAD_IDX, RESERVED_IDX), eps: float = 1e-8, ): super().__init__() self.ignored_idx = tuple(int(i) for i in ignored_idx) self.eps = eps def forward( self, logits: torch.Tensor, targets: torch.Tensor, exposure: torch.Tensor, dt: torch.Tensor, history: torch.Tensor | None = None, ) -> torch.Tensor: 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 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, ) 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) 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): """First-onset all-future Weibull survival likelihood.""" def __init__( self, ignored_idx: Iterable[int] = (PAD_IDX, RESERVED_IDX), eps: float = 1e-8, ): super().__init__() self.ignored_idx = tuple(int(i) for i in ignored_idx) self.eps = eps def forward( self, logits: torch.Tensor, weibull_rho: torch.Tensor, targets: torch.Tensor, dt: torch.Tensor, exposure: torch.Tensor, history: torch.Tensor | None = None, ) -> torch.Tensor: 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: raise ValueError( "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) 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, ) 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_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() + (target_rho - 1.0) * target_dt.log() ) observed = (log_intensity * target_valid.to(dtype)).sum(dim=-1) return (-observed + penalty).mean() def build_loss(name: str, **kwargs) -> nn.Module: name = name.lower() if name == "delphi2m": return Delphi2MLoss(**kwargs) if name == "exponential": return ExponentialLoss(**kwargs) if name == "weibull": return WeibullLoss(**kwargs) raise ValueError( f"Unknown loss {name!r}. Available: delphi2m, exponential, weibull." )