401 lines
14 KiB
Python
401 lines
14 KiB
Python
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."
|
|
)
|