Add exposure-only input ablation

train_next_step.py

@@ -220,6 +220,15 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--n_head", type=int, default=10)
     parser.add_argument("--n_hist_layer", type=int, default=12)
     parser.add_argument("--dropout", type=float, default=0.0)
+    parser.add_argument(
+        "--input_ablation",
+        choices=["none", "exposure_only"],
+        default="none",
+        help=(
+            "Input ablation. 'exposure_only' removes disease-token embeddings "
+            "while retaining exposure, sex, and age embeddings."
+        ),
+    )
     parser.add_argument("--exposure_cache_dir", type=str, default=None)
     parser.add_argument(
         "--exposure_embeddings_file",
@@ -308,6 +317,10 @@ def parse_args() -> argparse.Namespace:
         raise ValueError(
             "--exposure_cache_dir is required with --exposure_embeddings_file"
         )
+    if args.input_ablation == "exposure_only" and not args.exposure_cache_dir:
+        raise ValueError(
+            "--input_ablation exposure_only requires --exposure_cache_dir"
+        )
     if args.target_mode == "uts":
         args.readout_name = args.readout_name or "same_time_group_end"
         args.include_no_event_in_uts_target = True
@@ -403,11 +416,16 @@ def distributed_run_dir(
 ) -> tuple[Path, str]:
     payload: list[str | None] = [None, None]
     if rank == 0:
+        input_label = (
+            args.input_ablation
+            if args.input_ablation != "none"
+            else ("exposure" if args.exposure_cache_dir else "noexposure")
+        )
         run_dir, run_name = create_unique_run_dir(
             lambda 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"{input_label}_"
                 f"{timestamp}"
             )
         )
@@ -436,6 +454,7 @@ def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
         dist_mode="exponential",
         dropout=args.dropout,
         use_exposure_embeddings=args.exposure_embeddings_file is not None,
+        input_ablation=args.input_ablation,
     )
 
 
@@ -619,6 +638,7 @@ def build_metadata(
             "vocab_size": int(dataset.vocab_size),
         },
         "use_exposure_embeddings": args.exposure_embeddings_file is not None,
+        "input_ablation": args.input_ablation,
         "exposure_cache_dir": args.exposure_cache_dir,
         "exposure_embeddings_file": args.exposure_embeddings_file,
         "num_workers": int(args.num_workers),
@@ -655,6 +675,7 @@ def main() -> None:
     logger.info(f"readout={args.readout_name}, target_mode={args.target_mode}")
     logger.info(f"exposure_cache_dir={args.exposure_cache_dir}")
     logger.info(f"exposure_embeddings_file={args.exposure_embeddings_file}")
+    logger.info(f"input_ablation={args.input_ablation}")
     logger.info(
         "DataLoader IO: "
         f"num_workers={args.num_workers}, "