from __future__ import annotations from typing import Iterable import torch import torch.nn as nn import torch.nn.functional as F PAD_IDX = 0 CHECKUP_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 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, CHECKUP_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 UniqueTimeSetExponentialLoss(nn.Module): """Next distinct timestamp event-set supervision with sum reduction.""" def __init__( self, ignored_idx: Iterable[int] = (PAD_IDX, CHECKUP_IDX), t_min: float = 1.0 / 365.25, max_exp_input: float = 60.0, exclude_ignored_from_intensity: bool = True, ): super().__init__() self.ignored_idx = [int(x) for x in ignored_idx] self.t_min = float(t_min) self.max_exp_input = float(max_exp_input) self.exclude_ignored_from_intensity = bool(exclude_ignored_from_intensity) def forward( self, logits: torch.Tensor, target_multi_hot: torch.Tensor, target_dt_unique: torch.Tensor, readout_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 if target_multi_hot.shape != (bsz, seq_len, vocab_size): raise ValueError( "target_multi_hot must match logits shape, " f"got {tuple(target_multi_hot.shape)} vs {tuple(logits.shape)}" ) if target_dt_unique.shape != (bsz, seq_len): raise ValueError( f"target_dt_unique must be {(bsz, seq_len)}, got {tuple(target_dt_unique.shape)}" ) if readout_mask.shape != (bsz, seq_len): raise ValueError(f"readout_mask must be {(bsz, seq_len)}, got {tuple(readout_mask.shape)}") ignore_mask = _make_ignore_mask(vocab_size, self.ignored_idx, logits.device) num_targets = target_multi_hot[:, :, ~ignore_mask].sum(dim=-1) valid_mask = readout_mask.bool() & (num_targets > 0) if not valid_mask.any(): total_loss = _zero_loss_like(logits) if return_components: return total_loss, { "observed": total_loss.detach(), "penalty": total_loss.detach(), "total": total_loss.detach(), } return total_loss logits_safe = torch.nan_to_num( logits[valid_mask], nan=0.0, posinf=self.max_exp_input, neginf=-self.max_exp_input, ) target_valid = target_multi_hot[valid_mask].to(logits_safe.dtype) target_valid[:, ignore_mask] = 0.0 observed_term = (logits_safe * target_valid).sum(dim=-1) penalty_scale = target_valid.sum(dim=-1) logits_for_lse = logits_safe if self.exclude_ignored_from_intensity: logits_for_lse = logits_safe.masked_fill(ignore_mask.unsqueeze(0), float("-inf")) dt_clamped = torch.clamp(target_dt_unique[valid_mask], min=self.t_min) log_lambda_total = torch.logsumexp(logits_for_lse, dim=-1) log_penalty = log_lambda_total + dt_clamped.log() penalty = torch.exp(torch.clamp(log_penalty, max=self.max_exp_input)) observed_loss = -observed_term penalty_loss = penalty_scale * penalty total_loss = (observed_loss + penalty_loss).mean() if return_components: return total_loss, { "observed": observed_loss.mean().detach(), "penalty": penalty_loss.mean().detach(), "total": total_loss.detach(), } return total_loss class ExponentialLoss(nn.Module): """Query-conditioned all-future-event exponential point-process loss.""" def __init__( self, ignored_idx: Iterable[int] = (PAD_IDX, CHECKUP_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, ) -> torch.Tensor: _, vocab_size = logits.shape rate = F.softplus(logits) + self.eps valid_vocab = _valid_vocab_mask(vocab_size, self.ignored_idx, logits.device) 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 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.""" def __init__( self, ignored_idx: Iterable[int] = (PAD_IDX, CHECKUP_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, ) -> torch.Tensor: _, 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)}" ) 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) t_exp = exposure.to(dtype).clamp_min(self.eps).unsqueeze(1) penalty = (rate * torch.pow(t_exp, rho))[:, 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 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) 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() class MixedLoss(nn.Module): """Exponential diseases plus one Weibull death endpoint.""" def __init__( self, death_idx: int, ignored_idx: Iterable[int] = (PAD_IDX, CHECKUP_IDX), eps: float = 1e-8, ): super().__init__() self.death_idx = int(death_idx) self.ignored_idx = tuple(int(i) for i in ignored_idx) self.eps = eps def forward( self, logits: torch.Tensor, death_rho: torch.Tensor, targets: torch.Tensor, dt: torch.Tensor, exposure: torch.Tensor, ) -> torch.Tensor: _, vocab_size = logits.shape dtype = logits.dtype rate = F.softplus(logits) + self.eps if death_rho.dim() == 2: death_rho = death_rho.squeeze(-1) death_rho = death_rho.to(device=logits.device, dtype=dtype).clamp_min(self.eps) valid_vocab = _valid_vocab_mask(vocab_size, self.ignored_idx, logits.device) valid_disease_vocab = valid_vocab.clone() valid_disease_vocab[self.death_idx] = False t_exp = exposure.to(dtype).clamp_min(self.eps) disease_penalty = t_exp * rate[:, valid_disease_vocab].sum(dim=-1) death_rate = rate[:, self.death_idx] death_penalty = death_rate * torch.pow(t_exp, death_rho) penalty = disease_penalty + death_penalty target_valid = torch.ones_like(targets, dtype=torch.bool, device=logits.device) for idx in self.ignored_idx: target_valid &= targets != idx disease_event_mask = target_valid & (targets != self.death_idx) safe_targets = targets.clamp(min=0, max=vocab_size - 1) disease_log_rate = rate.log().gather(1, safe_targets) observed_disease = (disease_log_rate * disease_event_mask.to(dtype)).sum(dim=-1) death_event_mask = target_valid & (targets == self.death_idx) death_observed = death_event_mask.any(dim=1) death_dt = (dt.to(dtype).clamp_min(self.eps) * death_event_mask.to(dtype)).sum(dim=1) death_log_intensity = ( death_rate.log() + death_rho.log() + (death_rho - 1.0) * death_dt.clamp_min(self.eps).log() ) observed_death = death_log_intensity * death_observed.to(dtype) return (-observed_disease - observed_death + penalty).mean() def build_loss(name: str, **kwargs) -> nn.Module: name = name.lower() if name in {"delphi2m", "d2m", "next_token"}: return Delphi2MLoss(**kwargs) if name in {"uts", "unique_time_set", "unique_time_exponential"}: return UniqueTimeSetExponentialLoss(**kwargs) if name in {"exponential", "query_exponential"}: return ExponentialLoss(**kwargs) if name in {"weibull", "query_weibull"}: return WeibullLoss(**kwargs) if name in {"mixed", "query_mixed"}: return MixedLoss(**kwargs) raise ValueError( f"Unknown loss {name!r}. Available: delphi2m, uts, exponential, weibull, mixed." )