Refactor loss computation and model input handling for improved clarity and efficiency

losses.py

@@ -99,29 +99,33 @@ class Delphi2MLoss(nn.Module):
                 }
             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,
+            logits_valid,
             nan=0.0,
             posinf=self.max_exp_input,
             neginf=-self.max_exp_input,
         )
 
         loss_ce = F.cross_entropy(
-            logits_safe.reshape(-1, vocab_size)[valid_mask.reshape(-1)],
-            target_events.reshape(-1)[valid_mask.reshape(-1)],
+            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.clone()
-            logits_for_lse[:, :, ignore_mask] = float("-inf")
+            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 - current_times, min=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
@@ -129,7 +133,7 @@ class Delphi2MLoss(nn.Module):
                 torch.clamp(log_lambda_total - log_dt_inv, max=self.max_exp_input)
             )
         )
-        loss_dt = loss_dt[valid_mask].mean()
+        loss_dt = loss_dt.mean()
 
         total_loss = self.ce_weight * loss_ce + self.time_weight * loss_dt
         if return_components:
@@ -181,17 +185,9 @@ class UniqueTimeSetExponentialLoss(nn.Module):
         if readout_mask.shape != (bsz, seq_len):
             raise ValueError(f"readout_mask must be {(bsz, seq_len)}, got {tuple(readout_mask.shape)}")
 
-        logits_safe = torch.nan_to_num(
-            logits,
-            nan=0.0,
-            posinf=self.max_exp_input,
-            neginf=-self.max_exp_input,
-        )
         ignore_mask = _make_ignore_mask(vocab_size, self.ignored_idx, logits.device)
 
-        target_valid = target_multi_hot.float().clone()
-        target_valid[:, :, ignore_mask] = 0.0
-        num_targets = target_valid.sum(dim=-1)
+        num_targets = target_multi_hot[:, :, ~ignore_mask].sum(dim=-1)
         valid_mask = readout_mask.bool() & (num_targets > 0)
 
         if not valid_mask.any():
@@ -204,34 +200,35 @@ class UniqueTimeSetExponentialLoss(nn.Module):
                 }
             return total_loss
 
-        log_lambda_targets = logits_safe * target_valid
-        log_lambda_targets = torch.where(
-            target_valid.bool(),
-            log_lambda_targets,
-            torch.zeros_like(log_lambda_targets),
+        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 = log_lambda_targets.sum(dim=-1)
-        penalty_scale = num_targets
+        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.clone()
-            logits_for_lse[:, :, ignore_mask] = float("-inf")
+            logits_for_lse = logits_safe.masked_fill(ignore_mask.unsqueeze(0), float("-inf"))
 
-        dt_clamped = torch.clamp(target_dt_unique, min=self.t_min)
+        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)[valid_mask].mean()
+        total_loss = (observed_loss + penalty_loss).mean()
 
         if return_components:
             return total_loss, {
-                "observed": observed_loss[valid_mask].mean().detach(),
-                "penalty": penalty_loss[valid_mask].mean().detach(),
+                "observed": observed_loss.mean().detach(),
+                "penalty": penalty_loss.mean().detach(),
                 "total": total_loss.detach(),
             }
         return total_loss

models.py

@@ -260,12 +260,27 @@ class DeepHealth(nn.Module):
         other_time: torch.Tensor,
         other_mask: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        batch_size, _n_other, n_embd = h_other.shape
-        group_counts = [
-            int(torch.unique(other_time[b, other_mask[b]], sorted=True).numel())
-            for b in range(batch_size)
-        ]
-        max_groups = max(group_counts, default=0)
+        batch_size, n_other, n_embd = h_other.shape
+        if n_other == 0:
+            empty_h = h_other.new_zeros(batch_size, 0, n_embd)
+            empty_t = other_time.new_zeros(batch_size, 0)
+            empty_m = torch.zeros(batch_size, 0, dtype=torch.bool, device=h_other.device)
+            return empty_h, empty_t, empty_m
+
+        masked_time = other_time.masked_fill(~other_mask, float("inf"))
+        _sorted_time_with_pad, order = masked_time.sort(dim=1)
+        sorted_time = other_time.gather(1, order)
+        sorted_mask = other_mask.gather(1, order)
+        sorted_h = h_other.gather(1, order.unsqueeze(-1).expand(-1, -1, n_embd))
+
+        group_start = torch.zeros_like(sorted_mask)
+        group_start[:, 0] = sorted_mask[:, 0]
+        group_start[:, 1:] = sorted_mask[:, 1:] & (
+            sorted_time[:, 1:] != sorted_time[:, :-1]
+        )
+        group_id = group_start.long().cumsum(dim=1) - 1
+        max_groups = int(group_start.sum(dim=1).max().item())
+
         pooled_h = h_other.new_zeros(batch_size, max_groups, n_embd)
         pooled_time = other_time.new_zeros(batch_size, max_groups)
         pooled_mask = torch.zeros(
@@ -277,35 +292,29 @@ class DeepHealth(nn.Module):
         if max_groups == 0:
             return pooled_h, pooled_time, pooled_mask
 
-        for b in range(batch_size):
-            valid_time = other_time[b, other_mask[b]]
-            if valid_time.numel() == 0:
-                continue
-            valid_h = h_other[b, other_mask[b]]
-            unique_time, inverse = torch.unique(
-                valid_time,
-                sorted=True,
-                return_inverse=True,
+        safe_group_id = group_id.clamp_min(0)
+        pooled_h.scatter_add_(
+            1,
+            safe_group_id.unsqueeze(-1).expand_as(sorted_h),
+            sorted_h * sorted_mask.unsqueeze(-1).to(sorted_h.dtype),
+        )
+        if self.extra_pool_reduce == "mean":
+            counts = h_other.new_zeros(batch_size, max_groups, 1)
+            counts.scatter_add_(
+                1,
+                safe_group_id.unsqueeze(-1),
+                sorted_mask.unsqueeze(-1).to(h_other.dtype),
             )
-            n_groups = unique_time.numel()
-            group_h = valid_h.new_zeros(n_groups, n_embd)
-            group_h.scatter_add_(
-                0,
-                inverse[:, None].expand(-1, n_embd),
-                valid_h,
-            )
-            if self.extra_pool_reduce == "mean":
-                counts = valid_h.new_zeros(n_groups, 1)
-                counts.scatter_add_(
-                    0,
-                    inverse[:, None],
-                    torch.ones_like(valid_h[:, :1]),
-                )
-                group_h = group_h / counts.clamp_min(1.0)
+            pooled_h = pooled_h / counts.clamp_min(1.0)
 
-            pooled_h[b, :n_groups] = group_h
-            pooled_time[b, :n_groups] = unique_time
-            pooled_mask[b, :n_groups] = True
+        pooled_time.scatter_add_(
+            1,
+            safe_group_id,
+            sorted_time * group_start.to(sorted_time.dtype),
+        )
+        group_count = group_start.sum(dim=1)
+        arange_groups = torch.arange(max_groups, device=h_other.device)
+        pooled_mask = arange_groups.unsqueeze(0) < group_count.unsqueeze(1)
         return pooled_h, pooled_time, pooled_mask
 
     def _forward_shared(

train_all_future.py

@@ -44,6 +44,19 @@ from train_util import (
 )
 
 
+MODEL_INPUT_KEYS = (
+    "event_seq",
+    "time_seq",
+    "sex",
+    "padding_mask",
+    "t_query",
+    "other_type",
+    "other_value",
+    "other_value_kind",
+    "other_time",
+)
+
+
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
         description="Train DeepHealth with all-future supervision")
@@ -87,6 +100,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--min_lr_ratio", type=float, default=0.1)
     parser.add_argument("--num_workers", type=int, default=4)
     parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--progress_interval", type=int, default=20)
 
     args = parser.parse_args()
     if args.min_history_events < 1:
@@ -169,7 +183,14 @@ def compute_all_future_loss(
     batch: Dict[str, torch.Tensor],
     device: torch.device,
 ) -> torch.Tensor:
-    batch = move_batch_to_device(batch, device)
+    required_keys = set(MODEL_INPUT_KEYS)
+    required_keys.update(("future_targets", "exposure"))
+    if args.dist_mode in {"weibull", "mixed"}:
+        required_keys.add("future_dt")
+    batch = move_batch_to_device(
+        {key: batch[key] for key in required_keys},
+        device,
+    )
 
     hidden = model(
         event_seq=batch["event_seq"],
@@ -224,10 +245,11 @@ def run_epoch(
     is_train: bool,
 ) -> float:
     model.train(is_train)
-    total = 0.0
+    total = torch.zeros((), device=device)
     n_batches = 0
     skipped = 0
     desc = "train" if is_train else "val"
+    progress_interval = max(1, int(args.progress_interval))
 
     progress = tqdm(loader, desc=desc, leave=False, dynamic_ncols=True)
     for batch_idx, batch in enumerate(progress):
@@ -242,10 +264,15 @@ def run_epoch(
                     clip_grad_norm_(model.parameters(), args.grad_clip)
                 optimizer.step()
 
-            total += float(loss.detach().cpu())
+            total = total + loss.detach()
             n_batches += 1
-            avg = total / max(1, n_batches)
-            progress.set_postfix(loss=f"{float(loss.detach().cpu()):.4f}", avg=f"{avg:.4f}", skipped=skipped)
+            if (batch_idx + 1) % progress_interval == 0:
+                avg = total / max(1, n_batches)
+                progress.set_postfix(
+                    loss=f"{float(loss.detach().cpu()):.4f}",
+                    avg=f"{float(avg.detach().cpu()):.4f}",
+                    skipped=skipped,
+                )
         except RuntimeError as exc:
             if "Loss is not finite" not in str(exc):
                 raise
@@ -254,7 +281,7 @@ def run_epoch(
 
     if skipped:
         logger.info(f"Skipped {skipped} batches due to non-finite loss")
-    return total / max(1, n_batches) if n_batches else float("inf")
+    return float((total / max(1, n_batches)).detach().cpu()) if n_batches else float("inf")
 
 
 def build_metadata(

train_next_step.py

@@ -42,6 +42,18 @@ from train_util import (
 )
 
 
+MODEL_INPUT_KEYS = (
+    "event_seq",
+    "time_seq",
+    "sex",
+    "padding_mask",
+    "other_type",
+    "other_value",
+    "other_value_kind",
+    "other_time",
+)
+
+
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
         description="Train DeepHealth with next-token/point supervision")
@@ -94,6 +106,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--min_lr_ratio", type=float, default=0.1)
     parser.add_argument("--num_workers", type=int, default=4)
     parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--progress_interval", type=int, default=20)
 
     args = parser.parse_args()
     use_eid_split = all(
@@ -181,92 +194,111 @@ def build_next_step_loss(args: argparse.Namespace):
 
 
 def build_augmented_next_step_targets(
-    batch: Dict[str, torch.Tensor],
+    batch_cpu: Dict[str, torch.Tensor],
     model_out: DeepHealthOutput,
+    include_uts_targets: bool,
 ) -> Dict[str, torch.Tensor]:
     hidden_len = model_out.hidden.size(1)
     event_len = int(model_out.event_len)
     extra_len = hidden_len - event_len
-    if extra_len <= 0:
-        return {
-            "target_event_seq": batch["target_event_seq"],
-            "target_time_seq": batch["target_time_seq"],
-            "readout_mask": batch["readout_mask"],
-            "target_dt_unique": batch["target_dt_unique"],
-            "target_multi_hot": batch["target_multi_hot"],
-        }
-
     device = model_out.hidden.device
-    bsz, _seq_len, vocab_size = batch["target_multi_hot"].shape
-    extra_mask = model_out.padding_mask[:, event_len:]
-    extra_time = model_out.time_seq[:, event_len:]
+    non_blocking = device.type == "cuda"
+    if extra_len <= 0:
+        targets = {
+            "target_event_seq": batch_cpu["target_event_seq"].to(device, non_blocking=non_blocking),
+            "target_time_seq": batch_cpu["target_time_seq"].to(device, non_blocking=non_blocking),
+            "readout_mask": batch_cpu["readout_mask"].to(device, non_blocking=non_blocking),
+        }
+        if include_uts_targets:
+            targets["target_dt_unique"] = batch_cpu["target_dt_unique"].to(
+                device, non_blocking=non_blocking
+            )
+            targets["target_multi_hot"] = batch_cpu["target_multi_hot"].to(
+                device, non_blocking=non_blocking
+            )
+        return targets
+
+    bsz = batch_cpu["target_event_seq"].size(0)
+    vocab_size = (
+        batch_cpu["target_multi_hot"].size(2)
+        if include_uts_targets
+        else None
+    )
+    other_valid = batch_cpu["other_type"] > 0
+    extra_time = batch_cpu["other_time"].new_zeros(bsz, extra_len)
+    extra_mask = torch.zeros(bsz, extra_len, dtype=torch.bool)
+    for b in range(bsz):
+        unique_time = torch.unique(batch_cpu["other_time"][b, other_valid[b]], sorted=True)
+        n_time = min(int(unique_time.numel()), extra_len)
+        if n_time > 0:
+            extra_time[b, :n_time] = unique_time[:n_time]
+            extra_mask[b, :n_time] = True
 
     target_event_seq = torch.cat(
         [
-            batch["target_event_seq"],
+            batch_cpu["target_event_seq"],
             torch.full(
                 (bsz, extra_len),
                 PAD_IDX,
-                dtype=batch["target_event_seq"].dtype,
-                device=device,
+                dtype=batch_cpu["target_event_seq"].dtype,
             ),
         ],
         dim=1,
     )
     target_time_seq = torch.cat(
         [
-            batch["target_time_seq"],
+            batch_cpu["target_time_seq"],
             torch.zeros(
                 bsz,
                 extra_len,
-                dtype=batch["target_time_seq"].dtype,
-                device=device,
-            ),
-        ],
-        dim=1,
-    )
-    readout_mask = torch.cat([batch["readout_mask"], extra_mask], dim=1)
-    target_dt_unique = torch.cat(
-        [
-            batch["target_dt_unique"],
-            torch.zeros(
-                bsz,
-                extra_len,
-                dtype=batch["target_dt_unique"].dtype,
-                device=device,
-            ),
-        ],
-        dim=1,
-    )
-    target_multi_hot = torch.cat(
-        [
-            batch["target_multi_hot"],
-            torch.zeros(
-                bsz,
-                extra_len,
-                vocab_size,
-                dtype=batch["target_multi_hot"].dtype,
-                device=device,
+                dtype=batch_cpu["target_time_seq"].dtype,
             ),
         ],
         dim=1,
     )
+    readout_mask = torch.cat([batch_cpu["readout_mask"], extra_mask], dim=1)
+    target_dt_unique = None
+    target_multi_hot = None
+    if include_uts_targets:
+        target_dt_unique = torch.cat(
+            [
+                batch_cpu["target_dt_unique"],
+                torch.zeros(
+                    bsz,
+                    extra_len,
+                    dtype=batch_cpu["target_dt_unique"].dtype,
+                ),
+            ],
+            dim=1,
+        )
+        target_multi_hot = torch.cat(
+            [
+                batch_cpu["target_multi_hot"],
+                torch.zeros(
+                    bsz,
+                    extra_len,
+                    vocab_size,
+                    dtype=batch_cpu["target_multi_hot"].dtype,
+                ),
+            ],
+            dim=1,
+        )
 
     for b in range(bsz):
-        valid_event = batch["padding_mask"][b].bool()
+        valid_event = batch_cpu["padding_mask"][b].bool()
         if not valid_event.any():
             continue
         n_event = int(valid_event.sum().item())
         events = torch.cat(
             [
-                batch["event_seq"][b, :n_event],
-                batch["target_event_seq"][b, n_event - 1:n_event],
+                batch_cpu["event_seq"][b, :n_event],
+                batch_cpu["target_event_seq"][b, n_event - 1:n_event],
             ]
         )
         times = torch.cat(
             [
-                batch["time_seq"][b, :n_event],
-                batch["target_time_seq"][b, n_event - 1:n_event],
+                batch_cpu["time_seq"][b, :n_event],
+                batch_cpu["target_time_seq"][b, n_event - 1:n_event],
             ]
         )
         valid_full = events > PAD_IDX
@@ -291,6 +323,9 @@ def build_augmented_next_step_targets(
             target_event_seq[b, pos] = next_event
             target_time_seq[b, pos] = next_time
 
+            if not include_uts_targets:
+                continue
+
             same_next_time = times == next_time
             next_events = events[same_next_time]
             valid_next_events = next_events[
@@ -302,13 +337,15 @@ def build_augmented_next_step_targets(
             target_multi_hot[b, pos, valid_next_events] = True
             target_dt_unique[b, pos] = next_time - t
 
-    return {
-        "target_event_seq": target_event_seq,
-        "target_time_seq": target_time_seq,
-        "readout_mask": readout_mask,
-        "target_dt_unique": target_dt_unique,
-        "target_multi_hot": target_multi_hot,
+    targets = {
+        "target_event_seq": target_event_seq.to(device, non_blocking=non_blocking),
+        "target_time_seq": target_time_seq.to(device, non_blocking=non_blocking),
+        "readout_mask": readout_mask.to(device, non_blocking=non_blocking),
     }
+    if include_uts_targets:
+        targets["target_dt_unique"] = target_dt_unique.to(device, non_blocking=non_blocking)
+        targets["target_multi_hot"] = target_multi_hot.to(device, non_blocking=non_blocking)
+    return targets
 
 
 def compute_next_step_loss(
@@ -319,7 +356,11 @@ def compute_next_step_loss(
     batch: Dict[str, torch.Tensor],
     device: torch.device,
 ) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]:
-    batch = move_batch_to_device(batch, device)
+    batch_cpu = batch
+    batch = move_batch_to_device(
+        {key: batch_cpu[key] for key in MODEL_INPUT_KEYS},
+        device,
+    )
     model_out = model(
         event_seq=batch["event_seq"],
         time_seq=batch["time_seq"],
@@ -334,7 +375,11 @@ def compute_next_step_loss(
     )
     if not isinstance(model_out, DeepHealthOutput):
         raise TypeError("DeepHealth return_output=True must return DeepHealthOutput")
-    targets = build_augmented_next_step_targets(batch=batch, model_out=model_out)
+    targets = build_augmented_next_step_targets(
+        batch_cpu=batch_cpu,
+        model_out=model_out,
+        include_uts_targets=args.target_mode == "uts",
+    )
     readout_out = readout(
         hidden=model_out.hidden,
         time_seq=model_out.time_seq,
@@ -380,11 +425,12 @@ def run_epoch(
 ) -> float:
     model.train(is_train)
     readout.train(is_train)
-    total = 0.0
+    total = torch.zeros((), device=device)
     n_batches = 0
     skipped = 0
-    parts_sum: Dict[str, float] = {}
+    parts_sum: Dict[str, torch.Tensor] = {}
     desc = "train" if is_train else "val"
+    progress_interval = max(1, int(args.progress_interval))
 
     progress = tqdm(loader, desc=desc, leave=False, dynamic_ncols=True)
     for batch_idx, batch in enumerate(progress):
@@ -399,15 +445,20 @@ def run_epoch(
                     clip_grad_norm_(model.parameters(), args.grad_clip)
                 optimizer.step()
 
-            total += float(loss.detach().cpu())
+            total = total + loss.detach()
             n_batches += 1
             for name, value in parts.items():
-                parts_sum[name] = parts_sum.get(name, 0.0) + float(value.detach().cpu())
-            avg = total / max(1, n_batches)
-            postfix = {"loss": f"{float(loss.detach().cpu()):.4f}", "avg": f"{avg:.4f}", "skipped": skipped}
-            for name, value in parts_sum.items():
-                postfix[name] = f"{value / max(1, n_batches):.4f}"
-            progress.set_postfix(postfix)
+                parts_sum[name] = parts_sum.get(name, torch.zeros((), device=device)) + value.detach()
+            if (batch_idx + 1) % progress_interval == 0:
+                avg = total / max(1, n_batches)
+                postfix = {
+                    "loss": f"{float(loss.detach().cpu()):.4f}",
+                    "avg": f"{float(avg.detach().cpu()):.4f}",
+                    "skipped": skipped,
+                }
+                for name, value in parts_sum.items():
+                    postfix[name] = f"{float((value / max(1, n_batches)).detach().cpu()):.4f}"
+                progress.set_postfix(postfix)
         except RuntimeError as exc:
             if "Loss is not finite" not in str(exc):
                 raise
@@ -416,7 +467,7 @@ def run_epoch(
 
     if skipped:
         logger.info(f"Skipped {skipped} batches due to non-finite loss")
-    return total / max(1, n_batches) if n_batches else float("inf")
+    return float((total / max(1, n_batches)).detach().cpu()) if n_batches else float("inf")
 
 
 def build_metadata(