Improve all-future first-onset training

dataset.py

@@ -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.
+    Every split uses the same patient-equal query distribution: choose one
-    Valid/test use random-but-fixed query points. For each patient with N real
+    eligible inter-event interval uniformly, then choose a time uniformly in
-    disease events, N - 2 query points are sampled from the eligible observed
+    that interval. Train resamples on every ``__getitem__`` call; valid/test
-    time range, with at least one future event after every query.
+    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(
+    def _eligible_query_intervals(self, patient: Dict) -> List[Tuple[float, float]]:
-        self,
-        patient: Dict,
-        rng: np.random.RandomState,
-    ) -> List[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))
+        unique_times = np.unique(times[real_event_mask])
-        n_real_events = int(real_times.size)
+        intervals: List[Tuple[float, float]] = []
-        n_queries = max(0, n_real_events - 2)
+        for j in range(1, len(unique_times)):
-        if n_queries == 0:
+            left = float(unique_times[j - 1])
-            return []
+            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)
+    def sample_query(
-        min_future = int(self.min_future_events)
+        self,
-        if n_real_events < min_hist + min_future:
+        patient: Dict,
-            return []
+        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])
+        interval_idx = int(rng.randint(0, len(intervals)))
-        right_event_time = float(real_times[n_real_events - min_future])
+        left, right_event_time = intervals[interval_idx]
-        right = np.nextafter(np.float32(right_event_time), np.float32(-np.inf))
+        right = float(
-        if not np.isfinite(left) or not np.isfinite(right) or float(right) <= left:
+            np.nextafter(np.float32(right_event_time), np.float32(-np.inf))
-            return []
+        )
+        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] = []
+    def _sample_fixed_validation_queries(
-        max_attempts = max(100, n_queries * 50)
+        self,
-        for _ in range(max_attempts):
+        patient: Dict,
-            if len(queries) >= n_queries:
+        rng: np.random.RandomState,
-                break
+    ) -> List[float]:
-            t_query = float(rng.uniform(left, float(right)))
+        return [self.sample_query(patient, rng)]
-            if self._is_valid_query(patient, t_query):
-                queries.append(t_query)
-
-        return queries
 
     def _sample_train_query(self, patient: Dict) -> float:
-        unique_times = np.unique(patient["times"])
+        return self.sample_query(patient, np.random)
-        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)
 
     def _build_item(self, patient: Dict, t_query: float) -> Dict:
         times = patient["times"]

eval_data.py

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

evaluate_calibration.py

@@ -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())
 

losses.py

@@ -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)
+        censor_time = exposure.to(rate.dtype).clamp_min(self.eps)
-        target_valid = torch.ones_like(targets, dtype=torch.bool, device=logits.device)
+        cumulative_at_censor = rate * censor_time.unsqueeze(1)
-        for idx in self.ignored_idx:
+        penalty = (cumulative_at_censor * at_risk.to(rate.dtype)).sum(dim=-1)
-            target_valid &= targets != idx
+
+        # 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)
+        censor_time = exposure.to(dtype).clamp_min(self.eps)
-        penalty = (rate * torch.pow(t_exp, rho))[:, valid_vocab].sum(dim=-1)
+        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()

models.py

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

run_all_future_first_onset_experiments_linux.sh

@@ -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

tests/test_all_future_likelihood.py

@@ -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()

train_all_future.py

@@ -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
+All splits use the same patient/interval/time-uniform query distribution.
-times, then split by patient.
+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"))
+    required_keys.update(("future_targets", "future_dt", "exposure"))
-    if args.dist_mode == "weibull":
-        required_keys.add("future_dt")
     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(

train_util.py

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