Add exposure cache and keep absolute time only

train_next_step.py

@@ -47,6 +47,11 @@ MODEL_INPUT_KEYS = (
     "padding_mask",
 )
 
+EXPOSURE_INPUT_KEYS = (
+    "exposure_daily",
+    "exposure_monthly",
+)
+
 
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
@@ -68,9 +73,16 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--n_embd", type=int, default=120)
     parser.add_argument("--n_head", type=int, default=10)
     parser.add_argument("--n_hist_layer", type=int, default=12)
-    parser.add_argument("--time_mode", type=str, default="relative",
-                        choices=["relative", "absolute"])
     parser.add_argument("--dropout", type=float, default=0.0)
+    parser.add_argument("--exposure_cache_dir", type=str, default=None)
+    parser.add_argument("--mask_onset_exposure", action="store_true")
+    parser.add_argument("--exposure_d_model", type=int, default=None)
+    parser.add_argument("--exposure_n_layers", type=int, default=2)
+    parser.add_argument("--exposure_top_k", type=int, default=3)
+    parser.add_argument("--exposure_n_convnext_blocks", type=int, default=2)
+    parser.add_argument("--exposure_conv_kernel_size", type=int, default=7)
+    parser.add_argument("--exposure_mlp_ratio", type=float, default=4.0)
+    parser.add_argument("--no_exposure_gate", action="store_true")
 
     parser.add_argument("--target_mode", type=str, default="uts",
                         choices=["delphi2m", "uts"])
@@ -137,9 +149,16 @@ def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
         n_head=args.n_head,
         n_hist_layer=args.n_hist_layer,
         target_mode="next_token",
-        time_mode=args.time_mode,
         dist_mode="exponential",
         dropout=args.dropout,
+        use_exposure_encoder=args.exposure_cache_dir is not None,
+        exposure_d_model=args.exposure_d_model,
+        exposure_n_layers=args.exposure_n_layers,
+        exposure_top_k=args.exposure_top_k,
+        exposure_n_convnext_blocks=args.exposure_n_convnext_blocks,
+        exposure_conv_kernel_size=args.exposure_conv_kernel_size,
+        exposure_mlp_ratio=args.exposure_mlp_ratio,
+        exposure_use_gate=not args.no_exposure_gate,
     )
 
 
@@ -201,18 +220,24 @@ def compute_next_step_loss(
     device: torch.device,
 ) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]:
     batch_cpu = batch
+    input_keys = list(MODEL_INPUT_KEYS)
+    input_keys.extend(key for key in EXPOSURE_INPUT_KEYS if key in batch_cpu)
     batch = move_batch_to_device(
-        {key: batch_cpu[key] for key in MODEL_INPUT_KEYS},
+        {key: batch_cpu[key] for key in input_keys},
         device,
     )
-    model_out = model(
-        event_seq=batch["event_seq"],
-        time_seq=batch["time_seq"],
-        sex=batch["sex"],
-        padding_mask=batch["padding_mask"],
-        target_mode="next_token",
-        return_output=True,
-    )
+    model_kwargs = {
+        "event_seq": batch["event_seq"],
+        "time_seq": batch["time_seq"],
+        "sex": batch["sex"],
+        "padding_mask": batch["padding_mask"],
+        "target_mode": "next_token",
+        "return_output": True,
+    }
+    if "exposure_daily" in batch:
+        model_kwargs["exposure_daily"] = batch["exposure_daily"]
+        model_kwargs["exposure_monthly"] = batch["exposure_monthly"]
+    model_out = model(**model_kwargs)
     if not isinstance(model_out, DeepHealthOutput):
         raise TypeError("DeepHealth return_output=True must return DeepHealthOutput")
     targets = build_augmented_next_step_targets(
@@ -329,6 +354,16 @@ def build_metadata(
         "dataset_metadata": {
             "vocab_size": int(dataset.vocab_size),
         },
+        "use_exposure_encoder": args.exposure_cache_dir is not None,
+        "exposure_cache_dir": args.exposure_cache_dir,
+        "mask_onset_exposure": bool(args.mask_onset_exposure),
+        "exposure_d_model": args.exposure_d_model,
+        "exposure_n_layers": int(args.exposure_n_layers),
+        "exposure_top_k": int(args.exposure_top_k),
+        "exposure_n_convnext_blocks": int(args.exposure_n_convnext_blocks),
+        "exposure_conv_kernel_size": int(args.exposure_conv_kernel_size),
+        "exposure_mlp_ratio": float(args.exposure_mlp_ratio),
+        "exposure_use_gate": not bool(args.no_exposure_gate),
         "split_sizes": {
             "train": int(len(train_subset)),
             "val": int(len(val_subset)),
@@ -347,8 +382,10 @@ def main() -> None:
 
     run_dir, run_name = create_unique_run_dir(
         lambda timestamp: (
-            f"{args.time_mode}_exponential_next_token_{args.target_mode}_"
-            f"gap_{args.no_event_interval_years:g}y_{timestamp}"
+            f"absolute_exponential_next_token_{args.target_mode}_"
+            f"gap_{args.no_event_interval_years:g}y_"
+            f"{'exposure' if args.exposure_cache_dir else 'noexposure'}_"
+            f"{timestamp}"
         )
     )
     logger = setup_logging(run_dir)
@@ -356,12 +393,15 @@ def main() -> None:
     logger.info(f"Starting next-step training run: {run_name}")
     logger.info(f"Device: {device}")
     logger.info(f"readout={args.readout_name}, target_mode={args.target_mode}")
+    logger.info(f"exposure_cache_dir={args.exposure_cache_dir}")
 
     dataset = HealthDataset(
         data_prefix=args.data_prefix,
         labels_file=args.labels_file,
         no_event_interval_years=args.no_event_interval_years,
         include_no_event_in_uts_target=args.include_no_event_in_uts_target,
+        exposure_cache_dir=args.exposure_cache_dir,
+        mask_onset_exposure=args.mask_onset_exposure,
     )
     if args.train_eid_file and args.val_eid_file and args.test_eid_file:
         train_subset, val_subset, test_subset = split_dataset_by_eid_files(