JiaruiLi/DeepHealthExpo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add exposure-only input ablation

Browse Source
This commit is contained in:
Jiarui Li
2026-07-12 16:00:34 +08:00
parent 0339f1a278
commit 14608ed04d
4 changed files with 48 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
README.md
View File

@@ -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
`<exposure_cache_dir>/train_channel_stats.npz`; use
`--recompute_channel_stats` only when a forced refresh is needed.

1
evaluate_auc.py
View File

@@ -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")),
)

15
models.py
View File

@@ -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,6 +165,9 @@ 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)
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(

23
train_next_step.py
View File

@@ -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}, "