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

Improve all-future first-onset training

Browse Source
This commit is contained in:
Jiarui Li
2026-08-21 13:48:59 +08:00
parent 9ccc6b56ec
commit d728d8585c
9 changed files with 1002 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files

105
dataset.py
View File

@@ -17,9 +17,6 @@ from targets import (
build_next_token_targets,
)
ONE_DAY_YEARS = 1.0 / DAYS_PER_YEAR
DISEASE_HISTORY_MODE_TIMED = "timed"
DISEASE_HISTORY_MODE_ORDERED = "ordered"
DISEASE_HISTORY_MODE_SET = "set"
@@ -518,10 +515,10 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
Dataset with unified other-info tokens and DeepHealthV2-style all-future
targets.
Train samples one query time per patient at each __getitem__ call.
Valid/test use random-but-fixed query points. For each patient with N real
disease events, N - 2 query points are sampled from the eligible observed
time range, with at least one future event after every query.
Every split uses the same patient-equal query distribution: choose one
eligible inter-event interval uniformly, then choose a time uniformly in
that interval. Train resamples on every ``__getitem__`` call; valid/test
keep one deterministic draw per patient.
"""
CACHE_VERSION = 5
@@ -591,6 +588,11 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
**features,
}
query_intervals = self._eligible_query_intervals(patient)
if not query_intervals:
continue
patient["query_intervals"] = query_intervals
pidx = len(self.patients)
self.patients.append(patient)
@@ -615,7 +617,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
labels,
np.array([PAD_IDX, RESERVED_IDX, NO_EVENT_IDX], dtype=np.int64),
)
n_hist = int((times <= t_query).sum())
n_hist = int(((times <= t_query) & real_event_mask).sum())
n_future = int(((times > t_query) & real_event_mask).sum())
return (
n_hist >= self.min_history_events
@@ -623,62 +625,59 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
and patient["t_obs"] > t_query
)
def _sample_fixed_validation_queries(
self,
patient: Dict,
rng: np.random.RandomState,
) -> List[float]:
def _eligible_query_intervals(self, patient: Dict) -> List[Tuple[float, float]]:
times = np.asarray(patient["times"], dtype=np.float32)
labels = np.asarray(patient["labels"], dtype=np.int64)
real_event_mask = ~np.isin(
labels,
np.array([PAD_IDX, RESERVED_IDX, NO_EVENT_IDX], dtype=np.int64),
)
real_times = np.sort(times[real_event_mask].astype(np.float32, copy=False))
n_real_events = int(real_times.size)
n_queries = max(0, n_real_events - 2)
if n_queries == 0:
return []
unique_times = np.unique(times[real_event_mask])
intervals: List[Tuple[float, float]] = []
for j in range(1, len(unique_times)):
left = float(unique_times[j - 1])
right = float(unique_times[j])
probe = float(
np.nextafter(np.float32(right), np.float32(-np.inf))
)
if np.isfinite(left) and np.isfinite(probe) and probe >= left:
if self._is_valid_query(patient, probe):
intervals.append((left, right))
return intervals
min_hist = int(self.min_history_events)
min_future = int(self.min_future_events)
if n_real_events < min_hist + min_future:
return []
def sample_query(
self,
patient: Dict,
rng,
) -> float:
intervals = patient.get("query_intervals")
if intervals is None:
intervals = self._eligible_query_intervals(patient)
if not intervals:
raise RuntimeError("Patient has no eligible all-future query interval.")
left = float(real_times[min_hist - 1])
right_event_time = float(real_times[n_real_events - min_future])
right = np.nextafter(np.float32(right_event_time), np.float32(-np.inf))
if not np.isfinite(left) or not np.isfinite(right) or float(right) <= left:
return []
interval_idx = int(rng.randint(0, len(intervals)))
left, right_event_time = intervals[interval_idx]
right = float(
np.nextafter(np.float32(right_event_time), np.float32(-np.inf))
)
if right <= left:
t_query = float(left)
else:
t_query = float(rng.uniform(left, right))
if not self._is_valid_query(patient, t_query):
raise RuntimeError("Sampled an invalid all-future query time.")
return t_query
queries: List[float] = []
max_attempts = max(100, n_queries * 50)
for _ in range(max_attempts):
if len(queries) >= n_queries:
break
t_query = float(rng.uniform(left, float(right)))
if self._is_valid_query(patient, t_query):
queries.append(t_query)
return queries
def _sample_fixed_validation_queries(
self,
patient: Dict,
rng: np.random.RandomState,
) -> List[float]:
return [self.sample_query(patient, rng)]
def _sample_train_query(self, patient: Dict) -> float:
unique_times = np.unique(patient["times"])
if len(unique_times) < 2:
raise RuntimeError("Training patient has fewer than two unique times.")
j = np.random.randint(1, len(unique_times))
left = float(unique_times[j - 1])
right = float(unique_times[j])
if right - left <= ONE_DAY_YEARS:
t_query = right - ONE_DAY_YEARS
else:
t_query = np.random.uniform(left, right - ONE_DAY_YEARS)
if not self._is_valid_query(patient, t_query):
t_query = right - 1e-6
return float(t_query)
return self.sample_query(patient, np.random)
def _build_item(self, patient: Dict, t_query: float) -> Dict:
times = patient["times"]

7
eval_data.py
View File

@@ -66,7 +66,6 @@ def validate_training_mode_config(cfg: Dict[str, Any]) -> None:
"model_target_mode must be next_token or all_future, got "
f"{model_target_mode!r}"
)
disease_history_mode = normalize_disease_history_mode(
cfg.get("disease_history_mode", DISEASE_HISTORY_MODE_TIMED)
)
@@ -183,6 +182,11 @@ def build_model_from_dataset(
"model_target_mode must be next_token or all_future, got "
f"{model_target_mode!r}"
)
risk_head_bias = bool(cfg_get(args, cfg, "risk_head_bias", False))
if state_dict is not None:
# The checkpoint schema is authoritative. This keeps all older
# bias-free checkpoints loadable while restoring the new baseline bias.
risk_head_bias = "risk_head.bias" in state_dict
model_architecture = resolve_model_architecture(cfg, state_dict)
continuous_value_center = None
continuous_value_scale = None
@@ -231,6 +235,7 @@ def build_model_from_dataset(
dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")),
dropout=float(cfg_get(args, cfg, "dropout", 0.0)),
model_architecture=model_architecture,
risk_head_bias=risk_head_bias,
)

8
evaluate_calibration.py
View File

@@ -66,7 +66,7 @@ from evaluate_auc_v2 import (
)
from losses import build_loss
from model_architectures import resolve_model_architecture
from targets import PAD_IDX, RESERVED_IDX
from targets import NO_EVENT_IDX, PAD_IDX, RESERVED_IDX
from train_util import load_eid_file
@@ -1483,7 +1483,7 @@ def build_calibration_summary(metrics: pd.DataFrame) -> pd.DataFrame:
def _build_point_process_criterion(
dist_mode: str,
) -> Any:
ignored = {PAD_IDX, RESERVED_IDX}
ignored = {PAD_IDX, RESERVED_IDX, NO_EVENT_IDX}
if dist_mode == "exponential":
return build_loss("exponential", ignored_idx=ignored)
if dist_mode == "weibull":
@@ -1545,6 +1545,8 @@ def evaluate_point_process_nll(
logits=logits,
targets=batch_device["future_targets"],
exposure=batch_device["exposure"],
dt=batch_device["future_dt"],
history=batch_device["event_seq"],
)
elif dist_mode == "weibull":
loss = criterion(
@@ -1553,6 +1555,7 @@ def evaluate_point_process_nll(
targets=batch_device["future_targets"],
dt=batch_device["future_dt"],
exposure=batch_device["exposure"],
history=batch_device["event_seq"],
)
else:
raise ValueError(f"Unsupported dist_mode: {dist_mode!r}")
@@ -1564,6 +1567,7 @@ def evaluate_point_process_nll(
query_count += batch_size
valid_targets = batch["future_targets"] > PAD_IDX
valid_targets &= batch["future_targets"] != RESERVED_IDX
valid_targets &= batch["future_targets"] != NO_EVENT_IDX
future_event_count += int(valid_targets.sum().item())
exposure_sum += float(batch["exposure"].sum().item())

149
losses.py
View File

@@ -37,6 +37,54 @@ def _zero_loss_like(logits: torch.Tensor) -> torch.Tensor:
return logits.sum() * 0.0
def _all_future_at_risk_mask(
vocab_size: int,
ignored_idx: Iterable[int],
logits: torch.Tensor,
history: torch.Tensor | None,
) -> torch.Tensor:
"""Return outcomes that have not occurred by the query time."""
batch_size = logits.shape[0]
at_risk = _valid_vocab_mask(
vocab_size,
ignored_idx,
logits.device,
).unsqueeze(0).expand(batch_size, -1).clone()
if history is None or history.numel() == 0:
return at_risk
if history.dim() != 2 or history.shape[0] != batch_size:
raise ValueError(
"history must be (B, L). "
f"Got logits={tuple(logits.shape)}, history={tuple(history.shape)}"
)
history = history.to(device=logits.device, dtype=torch.long)
history_valid = (history >= 0) & (history < vocab_size)
for idx in ignored_idx:
history_valid &= history != int(idx)
safe_history = history.clamp(min=0, max=vocab_size - 1)
prevalent = torch.zeros(
(batch_size, vocab_size),
dtype=torch.long,
device=logits.device,
)
prevalent.scatter_add_(1, safe_history, history_valid.long())
return at_risk & ~prevalent.bool()
def _all_future_target_mask(
targets: torch.Tensor,
at_risk: torch.Tensor,
ignored_idx: Iterable[int],
) -> tuple[torch.Tensor, torch.Tensor]:
vocab_size = at_risk.shape[1]
in_vocab = (targets >= 0) & (targets < vocab_size)
for idx in ignored_idx:
in_vocab &= targets != int(idx)
safe_targets = targets.clamp(min=0, max=vocab_size - 1)
return in_vocab & at_risk.gather(1, safe_targets), safe_targets
class Delphi2MLoss(nn.Module):
"""Next-token plus exponential time-to-next-token supervision."""
@@ -146,7 +194,7 @@ class Delphi2MLoss(nn.Module):
class ExponentialLoss(nn.Module):
"""Query-conditioned all-future-event exponential point-process loss."""
"""First-onset all-future exponential survival likelihood."""
def __init__(
self,
@@ -162,24 +210,61 @@ class ExponentialLoss(nn.Module):
logits: torch.Tensor,
targets: torch.Tensor,
exposure: torch.Tensor,
dt: torch.Tensor,
history: torch.Tensor | None = None,
) -> torch.Tensor:
_, vocab_size = logits.shape
batch_size, vocab_size = logits.shape
if targets.dim() != 2 or targets.shape[0] != batch_size:
raise ValueError(
"targets must be (B, M). "
f"Got logits={tuple(logits.shape)}, targets={tuple(targets.shape)}"
)
if exposure.shape != (batch_size,):
raise ValueError(
f"exposure must be ({batch_size},), got {tuple(exposure.shape)}"
)
if dt.shape != targets.shape:
raise ValueError(
f"dt must match targets, got dt={tuple(dt.shape)}, "
f"targets={tuple(targets.shape)}"
)
rate = F.softplus(logits) + self.eps
valid_vocab = _valid_vocab_mask(vocab_size, self.ignored_idx, logits.device)
at_risk = _all_future_at_risk_mask(
vocab_size,
self.ignored_idx,
logits,
history,
)
target_valid, safe_targets = _all_future_target_mask(
targets,
at_risk,
self.ignored_idx,
)
penalty = exposure.to(rate.dtype) * rate[:, valid_vocab].sum(dim=-1)
target_valid = torch.ones_like(targets, dtype=torch.bool, device=logits.device)
for idx in self.ignored_idx:
target_valid &= targets != idx
censor_time = exposure.to(rate.dtype).clamp_min(self.eps)
cumulative_at_censor = rate * censor_time.unsqueeze(1)
penalty = (cumulative_at_censor * at_risk.to(rate.dtype)).sum(dim=-1)
# A first-onset outcome leaves the risk set at its event time, not at
# the common end of follow-up.
if targets.numel() > 0:
event_time = dt.to(rate.dtype).clamp_min(self.eps)
event_time = torch.minimum(event_time, censor_time.unsqueeze(1))
event_cumulative = rate.gather(1, safe_targets) * event_time
censor_cumulative = cumulative_at_censor.gather(1, safe_targets)
penalty = penalty + (
(event_cumulative - censor_cumulative)
* target_valid.to(rate.dtype)
).sum(dim=-1)
safe_targets = targets.clamp(min=0, max=vocab_size - 1)
observed = rate.log().gather(1, safe_targets)
observed = (observed * target_valid.to(rate.dtype)).sum(dim=-1)
return (-observed + penalty).mean()
class WeibullLoss(nn.Module):
"""Query-conditioned all-future-event Weibull point-process loss."""
"""First-onset all-future Weibull survival likelihood."""
def __init__(
self,
@@ -197,8 +282,9 @@ class WeibullLoss(nn.Module):
targets: torch.Tensor,
dt: torch.Tensor,
exposure: torch.Tensor,
history: torch.Tensor | None = None,
) -> torch.Tensor:
_, vocab_size = logits.shape
batch_size, vocab_size = logits.shape
if weibull_rho is None:
raise ValueError("weibull_rho is required for WeibullLoss")
if weibull_rho.shape != logits.shape:
@@ -206,23 +292,50 @@ class WeibullLoss(nn.Module):
"weibull_rho must have the same shape as logits. "
f"Got logits={tuple(logits.shape)}, weibull_rho={tuple(weibull_rho.shape)}"
)
if targets.dim() != 2 or targets.shape[0] != batch_size:
raise ValueError(
"targets must be (B, M). "
f"Got logits={tuple(logits.shape)}, targets={tuple(targets.shape)}"
)
if dt.shape != targets.shape:
raise ValueError(
f"dt must match targets, got dt={tuple(dt.shape)}, "
f"targets={tuple(targets.shape)}"
)
if exposure.shape != (batch_size,):
raise ValueError(
f"exposure must be ({batch_size},), got {tuple(exposure.shape)}"
)
dtype = logits.dtype
rate = F.softplus(logits) + self.eps
rho = weibull_rho.to(device=logits.device, dtype=dtype).clamp_min(self.eps)
valid_vocab = _valid_vocab_mask(vocab_size, self.ignored_idx, logits.device)
at_risk = _all_future_at_risk_mask(
vocab_size,
self.ignored_idx,
logits,
history,
)
target_valid, safe_targets = _all_future_target_mask(
targets,
at_risk,
self.ignored_idx,
)
t_exp = exposure.to(dtype).clamp_min(self.eps).unsqueeze(1)
penalty = (rate * torch.pow(t_exp, rho))[:, valid_vocab].sum(dim=-1)
censor_time = exposure.to(dtype).clamp_min(self.eps)
cumulative_at_censor = rate * torch.pow(censor_time.unsqueeze(1), rho)
penalty = (cumulative_at_censor * at_risk.to(dtype)).sum(dim=-1)
target_valid = torch.ones_like(targets, dtype=torch.bool, device=logits.device)
for idx in self.ignored_idx:
target_valid &= targets != idx
safe_targets = targets.clamp(min=0, max=vocab_size - 1)
target_rate = rate.gather(1, safe_targets)
target_rho = rho.gather(1, safe_targets)
target_dt = dt.to(dtype).clamp_min(self.eps)
target_dt = torch.minimum(target_dt, censor_time.unsqueeze(1))
event_cumulative = target_rate * torch.pow(target_dt, target_rho)
censor_cumulative = cumulative_at_censor.gather(1, safe_targets)
penalty = penalty + (
(event_cumulative - censor_cumulative) * target_valid.to(dtype)
).sum(dim=-1)
log_intensity = (
target_rate.log()
+ target_rho.log()

30
models.py
View File

@@ -221,6 +221,8 @@ class DeepHealth(nn.Module):
extra_pool_reduce: str = "mean",
dropout: float = 0.0,
model_architecture: str | None = None,
risk_head_bias: bool = False,
risk_head_bias_init: torch.Tensor | list[float] | None = None,
):
super().__init__()
if target_mode not in ["next_token", "all_future"]:
@@ -302,7 +304,33 @@ class DeepHealth(nn.Module):
self.rbf = GaussianRBFTimeBasis(n_bases=16, max_time_diff=40.0)
self.final_ln = nn.LayerNorm(n_embd)
self.risk_head = nn.Linear(n_embd, vocab_size, bias=False)
self.risk_head = nn.Linear(
n_embd,
vocab_size,
bias=bool(risk_head_bias),
)
if risk_head_bias_init is not None:
if self.risk_head.bias is None:
raise ValueError(
"risk_head_bias_init requires risk_head_bias=True"
)
initial_bias = torch.as_tensor(
risk_head_bias_init,
dtype=self.risk_head.bias.dtype,
device=self.risk_head.bias.device,
)
if initial_bias.shape != (vocab_size,):
raise ValueError(
"risk_head_bias_init must have shape "
f"({vocab_size},), got {tuple(initial_bias.shape)}"
)
if not torch.isfinite(initial_bias).all():
raise ValueError("risk_head_bias_init must contain only finite values")
with torch.no_grad():
# Start exactly at the fitted marginal baseline; covariate and
# history effects are learned away from zero during training.
self.risk_head.weight.zero_()
self.risk_head.bias.copy_(initial_bias)
if target_mode == "next_token":
self.risk_head.weight = self.token_embedding.weight
self.query_token = nn.Parameter(torch.zeros(n_embd))

390
run_all_future_first_onset_experiments_linux.sh Normal file
View File

@@ -0,0 +1,390 @@
#!/usr/bin/env bash
#
# Run the matched experiments required after the all-future first-onset update.
#
# Shared by every experiment:
# - corrected first-onset likelihood and outcome-specific risk exposure;
# - patient/interval/time-uniform query sampling in train/valid/test;
# - timed disease history and smoking/alcohol/BMI extra information.
#
# Per seed, the script trains:
# 1. TrajMixer + relative Weibull + fitted risk baseline (primary model);
# 2. the primary model without the fitted risk baseline (baseline ablation);
# 3. TrajMixer + relative exponential (time-distribution control);
# 4. FFN + relative Weibull (architecture control).
#
# After training, AUC and calibration/point-process-NLL evaluation run
# automatically unless --train-only is supplied.
set -uo pipefail
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"
GPU_CSV=""
SEED_CSV="42,43,44"
NUM_WORKERS=4
PYTHON_BIN="${PYTHON_BIN:-python}"
CAMPAIGN_NAME="all_future_first_onset_v2_multiseed"
DRY_RUN=0
TRAIN_ONLY=0
BATCH_SIZE=256
EXTRA_INFO_TYPES_FILE="$SCRIPT_DIR/extra_info_types_smoking_alcohol_bmi.txt"
TRAIN_EID_FILE="$SCRIPT_DIR/ukb_train_eid.csv"
VAL_EID_FILE="$SCRIPT_DIR/ukb_val_eid.csv"
TEST_EID_FILE="$SCRIPT_DIR/ukb_test_eid.csv"
usage() {
cat <<'EOF'
Usage:
bash run_all_future_first_onset_experiments_linux.sh --gpus LIST [options]
Required:
--gpus LIST Comma-separated GPU ids, for example 0 or 0,1,2,3.
Options:
--seeds LIST Comma-separated seeds (default: 42,43,44).
--num-workers N DataLoader workers per active GPU (default: 4).
--python PATH Python executable (default: $PYTHON_BIN or python).
--campaign NAME Output campaign directory name.
--train-only Skip AUC and calibration/NLL evaluation.
--dry-run Print commands without running them.
-h, --help Show this help message.
Fixed settings:
batch_size 256
extra information smoking + alcohol + BMI
patient split fixed train/validation/test EID files
disease history timed
experiments per seed 4
Outputs:
runs/<campaign>/seed_<seed>/<architecture>/...
batch_logs/<campaign>/seed_<seed>/<experiment>.log
EOF
}
while (($# > 0)); do
case "$1" in
--gpus)
[[ $# -ge 2 ]] || {
echo "ERROR: --gpus requires a value." >&2
exit 2
}
GPU_CSV="$2"
shift 2
;;
--seeds)
[[ $# -ge 2 ]] || {
echo "ERROR: --seeds requires a value." >&2
exit 2
}
SEED_CSV="$2"
shift 2
;;
--num-workers)
[[ $# -ge 2 ]] || {
echo "ERROR: --num-workers requires a value." >&2
exit 2
}
NUM_WORKERS="$2"
shift 2
;;
--python)
[[ $# -ge 2 ]] || {
echo "ERROR: --python requires a value." >&2
exit 2
}
PYTHON_BIN="$2"
shift 2
;;
--campaign)
[[ $# -ge 2 ]] || {
echo "ERROR: --campaign requires a value." >&2
exit 2
}
CAMPAIGN_NAME="$2"
shift 2
;;
--train-only)
TRAIN_ONLY=1
shift
;;
--dry-run)
DRY_RUN=1
shift
;;
-h|--help)
usage
exit 0
;;
*)
echo "ERROR: unknown argument: $1" >&2
usage >&2
exit 2
;;
esac
done
[[ -n "$GPU_CSV" ]] || {
echo "ERROR: --gpus is required." >&2
usage >&2
exit 2
}
[[ -n "$SEED_CSV" ]] || {
echo "ERROR: --seeds must not be empty." >&2
exit 2
}
[[ "$NUM_WORKERS" =~ ^[0-9]+$ ]] || {
echo "ERROR: --num-workers must be a non-negative integer." >&2
exit 2
}
[[ "$CAMPAIGN_NAME" =~ ^[A-Za-z0-9._-]+$ ]] || {
echo "ERROR: --campaign may contain only letters, numbers, ., _, and -." >&2
exit 2
}
required_files=(
"$SCRIPT_DIR/train_all_future.py"
"$SCRIPT_DIR/evaluate_all_runs_linux.sh"
"$SCRIPT_DIR/evaluate_calibration_all_runs_linux.sh"
"$EXTRA_INFO_TYPES_FILE"
"$TRAIN_EID_FILE"
"$VAL_EID_FILE"
"$TEST_EID_FILE"
)
for required_file in "${required_files[@]}"; do
[[ -f "$required_file" ]] || {
echo "ERROR: missing required file: $required_file" >&2
exit 2
}
done
command -v "$PYTHON_BIN" >/dev/null 2>&1 || {
echo "ERROR: Python executable not found: $PYTHON_BIN" >&2
exit 2
}
IFS=',' read -r -a GPU_IDS <<< "$GPU_CSV"
declare -A SEEN_GPUS=()
for gpu in "${GPU_IDS[@]}"; do
[[ -n "$gpu" && "$gpu" =~ ^[A-Za-z0-9._:-]+$ ]] || {
echo "ERROR: invalid GPU id: $gpu" >&2
exit 2
}
[[ -z "${SEEN_GPUS[$gpu]+x}" ]] || {
echo "ERROR: duplicate GPU id: $gpu" >&2
exit 2
}
SEEN_GPUS["$gpu"]=1
done
IFS=',' read -r -a SEEDS <<< "$SEED_CSV"
declare -A SEEN_SEEDS=()
for seed in "${SEEDS[@]}"; do
[[ "$seed" =~ ^[0-9]+$ ]] || {
echo "ERROR: invalid seed: $seed" >&2
exit 2
}
[[ -z "${SEEN_SEEDS[$seed]+x}" ]] || {
echo "ERROR: duplicate seed: $seed" >&2
exit 2
}
SEEN_SEEDS["$seed"]=1
done
RUNS_ROOT="$SCRIPT_DIR/runs/$CAMPAIGN_NAME"
LOG_ROOT="$SCRIPT_DIR/batch_logs/$CAMPAIGN_NAME"
if ((DRY_RUN == 0)); then
mkdir -p "$RUNS_ROOT" "$LOG_ROOT"
fi
declare -a JOB_NAMES=()
declare -a JOB_SEEDS=()
declare -a JOB_ARCHITECTURES=()
declare -a JOB_DIST_MODES=()
declare -a JOB_BASELINES=()
add_job() {
JOB_NAMES+=("$1")
JOB_SEEDS+=("$2")
JOB_ARCHITECTURES+=("$3")
JOB_DIST_MODES+=("$4")
JOB_BASELINES+=("$5")
}
for seed in "${SEEDS[@]}"; do
add_job \
"traj_mixer_relative_weibull_first_onset" \
"$seed" \
"traj_mixer_v5" \
"weibull" \
"on"
add_job \
"traj_mixer_relative_weibull_no_rate_baseline" \
"$seed" \
"traj_mixer_v5" \
"weibull" \
"off"
add_job \
"traj_mixer_relative_exponential_first_onset" \
"$seed" \
"traj_mixer_v5" \
"exponential" \
"on"
add_job \
"ffn_relative_weibull_first_onset" \
"$seed" \
"transformer_ffn_v1" \
"weibull" \
"on"
done
print_command() {
printf '%q ' "$@"
printf '\n'
}
run_job() {
local job_index="$1"
local gpu="$2"
local job_name="${JOB_NAMES[$job_index]}"
local seed="${JOB_SEEDS[$job_index]}"
local architecture="${JOB_ARCHITECTURES[$job_index]}"
local dist_mode="${JOB_DIST_MODES[$job_index]}"
local baseline="${JOB_BASELINES[$job_index]}"
local seed_runs_root="$RUNS_ROOT/seed_$seed"
local seed_log_root="$LOG_ROOT/seed_$seed"
local log_file="$seed_log_root/$job_name.log"
local -a command=(
"$PYTHON_BIN"
-u
"$SCRIPT_DIR/train_all_future.py"
--runs_root "$seed_runs_root"
--seed "$seed"
--batch_size "$BATCH_SIZE"
--num_workers "$NUM_WORKERS"
--device cuda
--model_architecture "$architecture"
--time_mode relative
--dist_mode "$dist_mode"
--disease_history_mode timed
--extra_info_types_file "$EXTRA_INFO_TYPES_FILE"
--train_eid_file "$TRAIN_EID_FILE"
--val_eid_file "$VAL_EID_FILE"
--test_eid_file "$TEST_EID_FILE"
)
if [[ "$baseline" == "on" ]]; then
command+=(--risk_head_bias)
else
command+=(--no-risk_head_bias)
fi
if ((DRY_RUN == 0)); then
mkdir -p "$seed_runs_root" "$seed_log_root"
fi
echo "[$(date '+%F %T')] START seed=$seed job=$job_name gpu=$gpu"
echo " log=$log_file"
if ((DRY_RUN)); then
printf ' CUDA_VISIBLE_DEVICES=%q ' "$gpu"
print_command "${command[@]}"
return 0
fi
if CUDA_VISIBLE_DEVICES="$gpu" PYTHONUNBUFFERED=1 \
"${command[@]}" >"$log_file" 2>&1; then
echo "[$(date '+%F %T')] DONE seed=$seed job=$job_name gpu=$gpu"
return 0
else
local exit_code=$?
echo "[$(date '+%F %T')] FAIL seed=$seed job=$job_name gpu=$gpu exit=$exit_code" >&2
echo " See: $log_file" >&2
return "$exit_code"
fi
}
worker() {
local slot="$1"
local gpu="${GPU_IDS[$slot]}"
local job_index
local failed=0
for ((job_index = slot; job_index < ${#JOB_NAMES[@]}; job_index += ${#GPU_IDS[@]})); do
run_job "$job_index" "$gpu" || failed=1
done
return "$failed"
}
echo "Campaign: $CAMPAIGN_NAME"
echo "Seeds: ${SEEDS[*]}"
echo "GPUs: ${GPU_IDS[*]}"
echo "Experiments per seed: 4"
echo "Total training tasks: ${#JOB_NAMES[@]}"
echo "Runs root: $RUNS_ROOT"
echo "Log root: $LOG_ROOT"
echo
declare -a WORKER_PIDS=()
for ((slot = 0; slot < ${#GPU_IDS[@]}; slot++)); do
worker "$slot" &
WORKER_PIDS+=("$!")
done
overall_status=0
for pid in "${WORKER_PIDS[@]}"; do
wait "$pid" || overall_status=1
done
if ((overall_status != 0)); then
echo "One or more training tasks failed. Inspect: $LOG_ROOT" >&2
exit 1
fi
if ((TRAIN_ONLY == 0)); then
evaluation_status=0
auc_command=(
bash
"$SCRIPT_DIR/evaluate_all_runs_linux.sh"
--gpus "$GPU_CSV"
--runs-root "$RUNS_ROOT"
--log-root "$LOG_ROOT/evaluate_auc"
--python "$PYTHON_BIN"
--num-workers "$NUM_WORKERS"
--num-workers-auc "$NUM_WORKERS"
)
calibration_command=(
bash
"$SCRIPT_DIR/evaluate_calibration_all_runs_linux.sh"
--gpus "$GPU_CSV"
--runs-root "$RUNS_ROOT"
--log-root "$LOG_ROOT/evaluate_calibration"
--python "$PYTHON_BIN"
--num-workers "$NUM_WORKERS"
--num-workers-calibration "$NUM_WORKERS"
)
echo ">> AUC evaluation"
print_command "${auc_command[@]}"
if ((DRY_RUN == 0)); then
"${auc_command[@]}" || evaluation_status=1
fi
echo ">> Calibration and point-process NLL evaluation"
print_command "${calibration_command[@]}"
if ((DRY_RUN == 0)); then
"${calibration_command[@]}" || evaluation_status=1
fi
if ((evaluation_status != 0)); then
echo "One or more evaluation workflows failed. Inspect: $LOG_ROOT" >&2
exit 1
fi
fi
if ((DRY_RUN)); then
echo "Dry run completed successfully."
else
echo "All required all-future experiments completed successfully."
fi

173
tests/test_all_future_likelihood.py Normal file
View File

@@ -0,0 +1,173 @@
import unittest
import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.data import Subset
from dataset import AllFutureHealthDataset
from losses import ExponentialLoss, WeibullLoss
from models import DeepHealth
from train_util import fit_all_future_baseline
def _inverse_softplus(value: torch.Tensor) -> torch.Tensor:
return torch.log(torch.expm1(value))
class AllFutureLikelihoodTests(unittest.TestCase):
def test_exponential_uses_event_specific_first_onset_exposure(self):
desired_rate = torch.tensor(
[[0.01, 0.01, 0.01, 0.20, 0.30, 0.40]],
dtype=torch.float64,
)
logits = _inverse_softplus(desired_rate)
criterion = ExponentialLoss(ignored_idx={0, 1, 2}, eps=1e-12)
loss = criterion(
logits=logits,
targets=torch.tensor([[4, 0]]),
exposure=torch.tensor([5.0], dtype=torch.float64),
dt=torch.tensor([[2.0, 0.0]], dtype=torch.float64),
history=torch.tensor([[3, 0]]),
)
rate = F.softplus(logits) + criterion.eps
expected = -(rate[0, 4].log()) + rate[0, 4] * 2.0 + rate[0, 5] * 5.0
torch.testing.assert_close(loss, expected)
def test_weibull_uses_event_time_and_excludes_prevalent_outcome(self):
desired_rate = torch.tensor(
[[0.01, 0.01, 0.01, 0.20, 0.30, 0.40]],
dtype=torch.float64,
)
logits = _inverse_softplus(desired_rate)
rho = torch.tensor(
[[1.0, 1.0, 1.0, 1.2, 1.5, 0.8]],
dtype=torch.float64,
)
criterion = WeibullLoss(ignored_idx={0, 1, 2}, eps=1e-12)
loss = criterion(
logits=logits,
weibull_rho=rho,
targets=torch.tensor([[4, 0]]),
dt=torch.tensor([[2.0, 0.0]], dtype=torch.float64),
exposure=torch.tensor([5.0], dtype=torch.float64),
history=torch.tensor([[3, 0]]),
)
rate = F.softplus(logits) + criterion.eps
log_hazard = (
rate[0, 4].log()
+ rho[0, 4].log()
+ (rho[0, 4] - 1.0) * torch.tensor(2.0).log()
)
expected = (
-log_hazard
+ rate[0, 4] * torch.pow(torch.tensor(2.0), rho[0, 4])
+ rate[0, 5] * torch.pow(torch.tensor(5.0), rho[0, 5])
)
torch.testing.assert_close(loss, expected)
class AllFutureQueryDistributionTests(unittest.TestCase):
@staticmethod
def _dataset() -> AllFutureHealthDataset:
dataset = AllFutureHealthDataset.__new__(AllFutureHealthDataset)
dataset.min_history_events = 1
dataset.min_future_events = 1
return dataset
@staticmethod
def _patient():
return {
"times": np.asarray([1.0, 2.0, 3.0, 4.0], dtype=np.float32),
"labels": np.asarray([3, 4, 5, 6], dtype=np.int64),
"t_obs": 4.0,
}
def test_train_validation_and_test_share_one_query_sampler(self):
dataset = self._dataset()
patient = self._patient()
patient["query_intervals"] = dataset._eligible_query_intervals(patient)
self.assertEqual(len(patient["query_intervals"]), 3)
rng = np.random.RandomState(123)
fixed = dataset._sample_fixed_validation_queries(patient, rng)
self.assertEqual(len(fixed), 1)
self.assertTrue(dataset._is_valid_query(patient, fixed[0]))
seen_intervals = set()
rng = np.random.RandomState(456)
for _ in range(200):
query = dataset.sample_query(patient, rng)
self.assertTrue(dataset._is_valid_query(patient, query))
seen_intervals.add(int(np.floor(query)))
self.assertEqual(seen_intervals, {1, 2, 3})
class AllFutureBaselineTests(unittest.TestCase):
def test_training_baseline_matches_first_onset_exposure(self):
dataset = AllFutureHealthDataset.__new__(AllFutureHealthDataset)
dataset.vocab_size = 7
dataset.min_history_events = 1
dataset.min_future_events = 2
patient = {
"times": np.asarray([1.0, 3.0, 5.0], dtype=np.float32),
"labels": np.asarray([3, 4, 5], dtype=np.int64),
"t_obs": 5.0,
"query_intervals": [(1.0, 3.0)],
}
dataset.patients = [patient]
subset = Subset(dataset, [0])
stats = fit_all_future_baseline(dataset, subset, seed=17)
query = dataset.sample_query(patient, np.random.RandomState(17))
self.assertEqual(stats.event_count[3], 0)
self.assertEqual(stats.at_risk_exposure[3], 0.0)
self.assertEqual(stats.event_count[4], 1)
self.assertEqual(stats.event_count[5], 1)
self.assertAlmostEqual(stats.at_risk_exposure[4], 3.0 - query, places=5)
self.assertAlmostEqual(stats.at_risk_exposure[5], 5.0 - query, places=5)
self.assertAlmostEqual(
float(F.softplus(torch.tensor(stats.bias[4]))),
float(stats.rate[4]),
places=6,
)
def test_model_starts_exactly_at_fitted_output_baseline(self):
baseline_rate = torch.tensor([0.0, 0.0, 0.0, 0.02, 0.05, 0.10])
baseline_bias = torch.zeros_like(baseline_rate)
baseline_bias[3:] = _inverse_softplus(baseline_rate[3:])
model = DeepHealth(
vocab_size=6,
n_embd=4,
n_head=1,
n_layer=1,
n_types=1,
n_cont_types=0,
n_categories=1,
cont_type_ids=[],
target_mode="all_future",
time_mode="absolute",
dist_mode="weibull",
model_architecture="transformer_ffn_v1",
risk_head_bias=True,
risk_head_bias_init=baseline_bias,
)
torch.testing.assert_close(model.risk_head.weight, torch.zeros_like(model.risk_head.weight))
output_rate = F.softplus(model.risk_head(torch.randn(3, 4)))
torch.testing.assert_close(output_rate[:, 3:], baseline_rate[None, 3:].expand(3, -1))
torch.testing.assert_close(
F.softplus(model.rho_head.bias),
torch.ones_like(model.rho_head.bias),
atol=2e-5,
rtol=0.0,
)
if __name__ == "__main__":
unittest.main()

95
train_all_future.py
View File

@@ -5,8 +5,8 @@ Training samples are patient-level. For each patient and each __getitem__ call,
AllFutureHealthDataset randomly samples a query time t_query, uses events at or
before t_query as history, and uses events after t_query as the future target set.
Validation/test samples are deterministic query points built from future event
times, then split by patient.
All splits use the same patient/interval/time-uniform query distribution.
Validation/test keep one deterministic query draw per patient.
"""
from __future__ import annotations
@@ -37,13 +37,15 @@ from model_architectures import (
SUPPORTED_MODEL_ARCHITECTURES,
)
from models import DeepHealth
from targets import PAD_IDX, RESERVED_IDX
from targets import NO_EVENT_IDX, PAD_IDX, RESERVED_IDX
from train_util import (
AllFutureBaselineStats,
ContinuousRobustScalerStats,
configure_torch_for_training,
create_unique_run_dir,
format_extra_info_types,
fit_continuous_robust_scaler,
fit_all_future_baseline,
get_lr,
get_model_parameter_counts,
load_extra_info_types_file,
@@ -113,6 +115,15 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--dist_mode", type=str, default="exponential",
choices=["exponential", "weibull"])
parser.add_argument("--dropout", type=float, default=0.0)
parser.add_argument(
"--risk_head_bias",
action=argparse.BooleanOptionalAction,
default=True,
help=(
"Initialize a learnable all-future output bias from training-only "
"marginal first-onset rates"
),
)
parser.add_argument(
"--model_architecture",
type=str,
@@ -179,6 +190,7 @@ def build_model(
args: argparse.Namespace,
dataset: AllFutureHealthDataset,
scaler_stats: ContinuousRobustScalerStats,
baseline_stats: AllFutureBaselineStats | None,
) -> DeepHealth:
if tuple(int(x) for x in dataset.cont_type_ids) != scaler_stats.cont_type_ids:
raise ValueError(
@@ -186,6 +198,8 @@ def build_model(
)
center = scaler_stats.center if dataset.n_cont_types > 0 else None
scale = scaler_stats.scale if dataset.n_cont_types > 0 else None
if args.risk_head_bias and baseline_stats is None:
raise ValueError("baseline_stats is required when risk_head_bias is enabled")
return DeepHealth(
vocab_size=dataset.vocab_size,
n_embd=args.n_embd,
@@ -204,11 +218,17 @@ def build_model(
dist_mode=args.dist_mode,
dropout=args.dropout,
model_architecture=args.model_architecture,
risk_head_bias=args.risk_head_bias,
risk_head_bias_init=(
baseline_stats.bias
if args.risk_head_bias and baseline_stats is not None
else None
),
)
def build_criterion(args: argparse.Namespace):
ignored_idx = {PAD_IDX, RESERVED_IDX}
ignored_idx = {PAD_IDX, RESERVED_IDX, NO_EVENT_IDX}
if args.dist_mode == "exponential":
return build_loss("exponential", ignored_idx=ignored_idx)
if args.dist_mode == "weibull":
@@ -224,9 +244,7 @@ def compute_all_future_loss(
device: torch.device,
) -> torch.Tensor:
required_keys = set(MODEL_INPUT_KEYS)
required_keys.update(("future_targets", "exposure"))
if args.dist_mode == "weibull":
required_keys.add("future_dt")
required_keys.update(("future_targets", "future_dt", "exposure"))
batch = move_batch_to_device(
{key: batch[key] for key in required_keys},
device,
@@ -250,6 +268,8 @@ def compute_all_future_loss(
logits=logits,
targets=batch["future_targets"],
exposure=batch["exposure"],
dt=batch["future_dt"],
history=batch["event_seq"],
)
elif args.dist_mode == "weibull":
loss = criterion(
@@ -258,6 +278,7 @@ def compute_all_future_loss(
targets=batch["future_targets"],
dt=batch["future_dt"],
exposure=batch["exposure"],
history=batch["event_seq"],
)
else:
raise ValueError(f"Unknown dist_mode: {args.dist_mode}")
@@ -325,6 +346,7 @@ def build_metadata(
val_subset,
test_subset,
scaler_stats: ContinuousRobustScalerStats,
baseline_stats: AllFutureBaselineStats | None,
) -> Dict[str, Any]:
scaler_metadata = scaler_stats.as_metadata()
return {
@@ -342,6 +364,17 @@ def build_metadata(
"all_future_min_history_events": int(args.min_history_events),
"all_future_min_future_events": int(args.min_future_events),
"all_future_validation_query_seed": int(args.validation_query_seed),
"all_future_query_distribution": "patient_interval_time_uniform",
"all_future_queries_per_validation_patient": 1,
"all_future_likelihood": "first_onset_survival_v2",
"all_future_prevalent_outcomes_excluded": True,
"risk_head_bias": bool(args.risk_head_bias),
"risk_head_bias_weight_decay": 0.0 if args.risk_head_bias else None,
"risk_head_baseline": (
baseline_stats.as_metadata()
if args.risk_head_bias and baseline_stats is not None
else {"method": "none"}
),
"extra_info_types_file": (
Path(args.extra_info_types_file).name
if args.extra_info_types_file is not None
@@ -476,6 +509,27 @@ def main() -> None:
f"max_per_feature={int(scaler_stats.observation_count.max()):,}"
)
baseline_stats = None
if args.risk_head_bias:
logger.info(
"Fitting all-future output baseline on one seeded query per "
f"training patient: patients={len(train_subset):,}"
)
baseline_stats = fit_all_future_baseline(
train_dataset,
train_subset,
seed=args.seed,
)
baseline_summary = baseline_stats.as_metadata()
rate_summary = baseline_summary["rate_per_year"]
logger.info(
"All-future baseline rates/year: "
f"min={rate_summary['min']:.8g}, "
f"median={rate_summary['median']:.8g}, "
f"max={rate_summary['max']:.8g}, "
f"first_onsets={baseline_summary['observed_first_onsets']:,}"
)
train_loader = DataLoader(
train_subset,
batch_size=args.batch_size,
@@ -507,15 +561,37 @@ def main() -> None:
prefetch_factor=2 if args.num_workers > 0 else None,
)
model = build_model(args, train_dataset, scaler_stats=scaler_stats).to(device)
model = build_model(
args,
train_dataset,
scaler_stats=scaler_stats,
baseline_stats=baseline_stats,
).to(device)
parameter_counts = get_model_parameter_counts(model)
logger.info(
"Model parameters: "
f"total={parameter_counts['model_parameter_count']:,}, "
f"trainable={parameter_counts['trainable_parameter_count']:,}"
)
if model.risk_head.bias is None:
optimizer_parameters = model.parameters()
else:
optimizer_parameters = [
{
"params": [
parameter
for parameter in model.parameters()
if parameter is not model.risk_head.bias
],
"weight_decay": args.weight_decay,
},
{
"params": [model.risk_head.bias],
"weight_decay": 0.0,
},
]
optimizer = AdamW(
model.parameters(),
optimizer_parameters,
lr=args.base_lr,
betas=tuple(args.betas),
weight_decay=args.weight_decay,
@@ -531,6 +607,7 @@ def main() -> None:
val_subset,
test_subset,
scaler_stats,
baseline_stats,
)
train_metadata.update(parameter_counts)
save_config(

129
train_util.py
View File

@@ -18,6 +18,7 @@ from torch.utils.data import Subset
from dataset import AllFutureHealthDataset, HealthDataset
from models import DeepHealth
from targets import NO_EVENT_IDX, PAD_IDX, RESERVED_IDX
@dataclass(frozen=True)
@@ -42,6 +43,134 @@ class ContinuousRobustScalerStats:
}
@dataclass(frozen=True)
class AllFutureBaselineStats:
"""Training-only marginal first-onset rates for output initialization."""
event_count: np.ndarray
at_risk_exposure: np.ndarray
rate: np.ndarray
bias: np.ndarray
query_count: int
rate_floor: float
rate_ceiling: float
def as_metadata(self) -> Dict[str, Any]:
positive = self.rate[self.at_risk_exposure > 0]
if positive.size:
rate_summary = {
"min": float(positive.min()),
"median": float(np.median(positive)),
"max": float(positive.max()),
}
else:
rate_summary = {"min": 0.0, "median": 0.0, "max": 0.0}
return {
"method": "training_marginal_first_onset_rate",
"fitted_on": "one_seeded_query_per_training_patient",
"query_distribution": "patient_interval_time_uniform",
"query_count": int(self.query_count),
"observed_first_onsets": int(self.event_count.sum()),
"rate_floor": float(self.rate_floor),
"rate_ceiling": float(self.rate_ceiling),
"rate_per_year": rate_summary,
"checkpoint_parameter": "risk_head.bias",
}
def fit_all_future_baseline(
dataset: AllFutureHealthDataset,
subset: Subset,
*,
seed: int,
ignored_idx: Iterable[int] = (PAD_IDX, RESERVED_IDX, NO_EVENT_IDX),
rate_floor: float = 1e-6,
rate_ceiling: float = 10.0,
) -> AllFutureBaselineStats:
"""Fit marginal per-outcome rates under the training query distribution."""
if subset.dataset is not dataset:
raise ValueError("subset must reference the all-future training dataset")
if rate_floor <= 0 or rate_ceiling <= rate_floor:
raise ValueError("Require 0 < rate_floor < rate_ceiling")
subset_indices = np.asarray(subset.indices, dtype=np.int64)
if subset_indices.ndim != 1 or subset_indices.size == 0:
raise ValueError("training subset must contain at least one patient")
vocab_size = int(dataset.vocab_size)
ignored = {
int(idx)
for idx in ignored_idx
if 0 <= int(idx) < vocab_size
}
event_count = np.zeros(vocab_size, dtype=np.int64)
exposure_adjustment = np.zeros(vocab_size, dtype=np.float64)
common_exposure = 0.0
rng = np.random.RandomState(int(seed))
for patient_index in subset_indices.tolist():
patient = dataset.patients[int(patient_index)]
t_query = dataset.sample_query(patient, rng)
times = np.asarray(patient["times"], dtype=np.float64)
labels = np.asarray(patient["labels"], dtype=np.int64)
censor_time = max(float(patient["t_obs"]) - t_query, 0.0)
common_exposure += censor_time
prevalent = {
int(label)
for label in labels[times <= t_query].tolist()
if 0 <= int(label) < vocab_size and int(label) not in ignored
}
for label in prevalent:
exposure_adjustment[label] -= censor_time
# Keep the earliest future occurrence defensively; prepared disease
# events are already deduplicated to their first occurrence.
first_future_dt: Dict[int, float] = {}
for label, event_time in zip(labels[times > t_query], times[times > t_query]):
label = int(label)
if label in ignored or label in prevalent or not (0 <= label < vocab_size):
continue
event_dt = max(float(event_time) - t_query, 0.0)
previous = first_future_dt.get(label)
if previous is None or event_dt < previous:
first_future_dt[label] = event_dt
for label, event_dt in first_future_dt.items():
event_count[label] += 1
exposure_adjustment[label] += event_dt - censor_time
at_risk_exposure = common_exposure + exposure_adjustment
at_risk_exposure = np.maximum(at_risk_exposure, 0.0)
rate = np.full(vocab_size, float(rate_floor), dtype=np.float64)
has_exposure = at_risk_exposure > 0
rate[has_exposure] = (
event_count[has_exposure].astype(np.float64)
/ at_risk_exposure[has_exposure]
)
rate = np.clip(rate, float(rate_floor), float(rate_ceiling))
for idx in ignored:
at_risk_exposure[idx] = 0.0
rate[idx] = 0.0
bias = np.zeros(vocab_size, dtype=np.float64)
valid = np.ones(vocab_size, dtype=bool)
if ignored:
valid[np.asarray(sorted(ignored), dtype=np.int64)] = False
bias[valid] = np.log(np.expm1(rate[valid]))
if not np.isfinite(bias).all():
raise RuntimeError("All-future baseline initialization is non-finite")
return AllFutureBaselineStats(
event_count=event_count,
at_risk_exposure=at_risk_exposure.astype(np.float32),
rate=rate.astype(np.float32),
bias=bias.astype(np.float32),
query_count=int(subset_indices.size),
rate_floor=float(rate_floor),
rate_ceiling=float(rate_ceiling),
)
def fit_continuous_robust_scaler(
dataset: HealthDataset | AllFutureHealthDataset,
subset: Subset,