diff --git a/README.md b/README.md index c564d2b..3ab581b 100644 --- a/README.md +++ b/README.md @@ -109,6 +109,15 @@ torchrun --standalone --nproc_per_node=4 train_next_step.py \ --exposure_cache_dir ukb_exposure_cache \ --batch_size 128 ``` + +To train the exposure-sequence ablation, which removes disease-token +embeddings while retaining sex and age encodings, add: + +```bash +python train_next_step.py \ + --exposure_cache_dir ukb_exposure_cache \ + --input_ablation exposure_only +``` Training-channel statistics are cached at `/train_channel_stats.npz`; use `--recompute_channel_stats` only when a forced refresh is needed. diff --git a/evaluate_auc.py b/evaluate_auc.py index 926cb39..9a02e16 100644 --- a/evaluate_auc.py +++ b/evaluate_auc.py @@ -326,6 +326,7 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data use_exposure_embeddings=bool( cfg_get(args, cfg, "use_exposure_embeddings", False) ), + input_ablation=str(cfg_get(args, cfg, "input_ablation", "none")), ) diff --git a/models.py b/models.py index 4780ab3..a090d02 100644 --- a/models.py +++ b/models.py @@ -30,6 +30,7 @@ class DeepHealth(nn.Module): dist_mode: str = "exponential", # "exponential", "weibull" or "mixed" dropout: float = 0.0, use_exposure_embeddings: bool = False, + input_ablation: str = "none", ): super().__init__() if target_mode not in ["next_token", "all_future"]: @@ -38,12 +39,19 @@ class DeepHealth(nn.Module): if dist_mode not in ["exponential", "weibull", "mixed"]: raise ValueError( "dist_mode must be either 'exponential', 'weibull' or 'mixed'") + if input_ablation not in {"none", "exposure_only"}: + raise ValueError("input_ablation must be 'none' or 'exposure_only'") + if input_ablation == "exposure_only" and not use_exposure_embeddings: + raise ValueError( + "input_ablation='exposure_only' requires exposure embeddings" + ) self.token_embedding = nn.Embedding(vocab_size, n_embd, padding_idx=0) self.gender_embedding = nn.Embedding( 2, n_embd) # Assuming binary gender self.target_mode = target_mode self.dist_mode = dist_mode self.use_exposure_embeddings = bool(use_exposure_embeddings) + self.input_ablation = input_ablation self.n_embd = n_embd self.vocab_size = vocab_size nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.02) @@ -61,6 +69,10 @@ class DeepHealth(nn.Module): if isinstance(module, nn.Linear): nn.init.normal_(module.weight, mean=0.0, std=0.02) nn.init.zeros_(module.bias) + if self.input_ablation == "exposure_only": + # The additive gate is bypassed in this ablation. Excluding it + # from optimization also avoids an unused parameter under DDP. + self.exposure_gate.requires_grad_(False) else: self.exposure_adapter = None self.register_parameter("exposure_gate", None) @@ -153,7 +165,10 @@ class DeepHealth(nn.Module): if self.exposure_adapter is None or self.exposure_gate is None: raise RuntimeError("Exposure adapter is not initialized") exposure = self.exposure_adapter(exposure) - h_disease = h_disease + torch.sigmoid(self.exposure_gate) * exposure + if self.input_ablation == "exposure_only": + h_disease = exposure + else: + h_disease = h_disease + torch.sigmoid(self.exposure_gate) * exposure elif exposure_embedding is not None: raise ValueError( "exposure_embedding provided but use_exposure_embeddings=False" diff --git a/train_next_step.py b/train_next_step.py index 68d00a0..f8fa6cb 100644 --- a/train_next_step.py +++ b/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}, "