Use precomputed exposure embeddings

README.md

@@ -68,9 +68,8 @@ python train_exposure_autoencoder.py \
   --val_eid_file ukb_val_eid.csv
 ```
 
-The best checkpoint contains both `model_state_dict`, an `encoder_state_dict`
-compatible with the default gated `TimesNetExposureEncoder`, and the channel
-normalization statistics needed when the encoder is attached to DeepHealth.
+The best checkpoint contains the encoder and normalization statistics needed
+to generate fixed exposure embeddings.
 Multi-GPU pretraining follows the main trainer interface: add
 `--data_parallel --gpu_ids 0,1,2,3`.
 For efficient multi-GPU training, launch one process per GPU with DDP:
@@ -84,13 +83,21 @@ torchrun --standalone --nproc_per_node=4 train_exposure_autoencoder.py \
 In DDP mode, `--batch_size` is the global batch size and must be divisible by
 the number of processes.
 
-The end-to-end next-step trainer supports the same DDP launch pattern:
+Encode every cached exposure window once:
+
+```bash
+python encode_exposure_cache.py \
+  --exposure_cache_dir ukb_exposure_cache \
+  --checkpoint runs/exposure_ae_RUN/best.pt
+```
+
+The next-step trainer reads `exposure_embeddings.npy` directly and does not
+run TimesNet:
 
 ```bash
 torchrun --standalone --nproc_per_node=4 train_next_step.py \
   --exposure_cache_dir ukb_exposure_cache \
-  --batch_size 128 \
-  --d_model 64
+  --batch_size 128
 ```
 Training-channel statistics are cached at
 `<exposure_cache_dir>/train_channel_stats.npz`; use

dataset.py

@@ -50,7 +50,11 @@ def _load_readonly_npy(path: Path) -> np.ndarray:
 class ExposureCache:
     """Eid-sequence-aligned exposure windows from prepare_exposure_cache.py."""
 
-    def __init__(self, cache_dir: str | Path):
+    def __init__(
+        self,
+        cache_dir: str | Path,
+        embeddings_file: str | Path | None = None,
+    ):
         cache_dir = Path(cache_dir)
         self.cache_dir = cache_dir
         manifest_path = cache_dir / "exposure_manifest.json"
@@ -101,6 +105,11 @@ class ExposureCache:
         self.eid_start = _load_readonly_npy(eid_start_path)
         self.daily = _load_readonly_npy(daily_path)
         self.monthly = _load_readonly_npy(monthly_path)
+        self.embeddings = (
+            _load_readonly_npy(Path(embeddings_file))
+            if embeddings_file is not None
+            else None
+        )
         quality_path = cache_dir / "exposure_quality.npy"
         self.quality = _load_readonly_npy(quality_path) if quality_path.is_file() else None
 
@@ -132,6 +141,16 @@ class ExposureCache:
             raise ValueError("exposure_eid_start.npy must have len(eid_index) + 1")
         if len(self.eid_start) and int(self.eid_start[-1]) != n_rows:
             raise ValueError("Last exposure eid offset must equal exposure row count")
+        if self.embeddings is not None:
+            if self.embeddings.ndim != 2:
+                raise ValueError(
+                    "Exposure embeddings must have shape (N, D), got "
+                    f"{self.embeddings.shape}"
+                )
+            if max_window_index >= len(self.embeddings):
+                raise ValueError(
+                    "Exposure row index points past exposure embeddings"
+                )
 
         self._eid_to_pos = {
             int(eid): idx
@@ -203,6 +222,20 @@ class ExposureCache:
     def monthly_windows(self, indices: np.ndarray) -> np.ndarray:
         return self._windows("monthly", indices)
 
+    def embedding_windows(self, indices: np.ndarray) -> np.ndarray:
+        if self.embeddings is None:
+            raise RuntimeError("Exposure embeddings are not enabled")
+        indices = np.asarray(indices, dtype=np.int64)
+        out = np.zeros(
+            (len(indices), self.embeddings.shape[1]), dtype=np.float32
+        )
+        valid_pos = np.nonzero(indices >= 0)[0]
+        if len(valid_pos):
+            out[valid_pos] = np.asarray(
+                self.embeddings[indices[valid_pos]], dtype=np.float32
+            )
+        return out
+
     def _windows(
         self,
         kind: Literal["daily", "monthly"],
@@ -288,18 +321,21 @@ class _ExpoBaseDataset(Dataset):
         no_event_interval_years: float = 5.0,
         include_no_event_in_uts_target: bool = False,
         exposure_cache_dir: str | Path | None = None,
-        mask_onset_exposure: bool = False,
+        exposure_embeddings_file: str | Path | None = None,
     ) -> None:
         self.data_prefix = data_prefix
         self.labels_file = labels_file
         self.no_event_interval_years = float(no_event_interval_years)
         self.include_no_event_in_uts_target = bool(include_no_event_in_uts_target)
         self.exposure_cache = (
-            ExposureCache(exposure_cache_dir)
+            ExposureCache(exposure_cache_dir, exposure_embeddings_file)
             if exposure_cache_dir is not None
             else None
         )
-        self.mask_onset_exposure = bool(mask_onset_exposure)
+        if exposure_embeddings_file is not None and exposure_cache_dir is None:
+            raise ValueError(
+                "exposure_cache_dir is required with exposure_embeddings_file"
+            )
 
         self.label_code_to_id, self.label_id_to_code = load_label_vocab(
             labels_file,
@@ -427,24 +463,14 @@ class _ExpoBaseDataset(Dataset):
             age_days=input_times_days,
         )
 
-    def _load_exposure_windows(self, exposure_index: np.ndarray) -> tuple[torch.Tensor, torch.Tensor]:
+    def _load_exposure_embeddings(
+        self, exposure_index: np.ndarray
+    ) -> torch.Tensor:
         if self.exposure_cache is None:
             raise RuntimeError("Exposure cache is not enabled")
-
-        daily = self.exposure_cache.daily_windows(exposure_index).astype(
-            np.float32,
-            copy=False,
-        )
-        monthly = self.exposure_cache.monthly_windows(exposure_index).astype(
-            np.float32,
-            copy=False,
-        )
-
-        if self.mask_onset_exposure:
-            daily[:, 0, :] = np.nan
-            monthly[:, 0, :] = np.nan
-
-        return torch.from_numpy(daily).float(), torch.from_numpy(monthly).float()
+        return torch.from_numpy(
+            self.exposure_cache.embedding_windows(exposure_index)
+        ).float()
 
 class NextStepHealthDataset(_ExpoBaseDataset):
     """
@@ -465,7 +491,7 @@ class NextStepHealthDataset(_ExpoBaseDataset):
         no_event_interval_years: float = 5.0,
         include_no_event_in_uts_target: bool = False,
         exposure_cache_dir: str | Path | None = None,
-        mask_onset_exposure: bool = False,
+        exposure_embeddings_file: str | Path | None = None,
     ) -> None:
         super().__init__(
             data_prefix=data_prefix,
@@ -473,7 +499,7 @@ class NextStepHealthDataset(_ExpoBaseDataset):
             no_event_interval_years=no_event_interval_years,
             include_no_event_in_uts_target=include_no_event_in_uts_target,
             exposure_cache_dir=exposure_cache_dir,
-            mask_onset_exposure=mask_onset_exposure,
+            exposure_embeddings_file=exposure_embeddings_file,
         )
 
         self.samples: List[Dict] = []
@@ -531,9 +557,9 @@ class NextStepHealthDataset(_ExpoBaseDataset):
             "target_multi_hot": torch.from_numpy(s["target_multi_hot"]).bool(),
         }
         if "exposure_index" in s:
-            daily, monthly = self._load_exposure_windows(s["exposure_index"])
-            out["exposure_daily"] = daily
-            out["exposure_monthly"] = monthly
+            out["exposure_embedding"] = self._load_exposure_embeddings(
+                s["exposure_index"]
+            )
         return out
 
 
@@ -561,7 +587,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
         min_future_events: int = 1,
         validation_query_seed: int = 42,
         exposure_cache_dir: str | Path | None = None,
-        mask_onset_exposure: bool = False,
+        exposure_embeddings_file: str | Path | None = None,
     ) -> None:
         if split not in {"train", "valid", "test"}:
             raise ValueError(f"split must be train/valid/test, got {split!r}")
@@ -572,7 +598,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
             no_event_interval_years=no_event_interval_years,
             include_no_event_in_uts_target=include_no_event_in_uts_target,
             exposure_cache_dir=exposure_cache_dir,
-            mask_onset_exposure=mask_onset_exposure,
+            exposure_embeddings_file=exposure_embeddings_file,
         )
 
         self.split = split
@@ -719,9 +745,9 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
                 input_times_days=times_days[hist],
             )
             if exposure_index is not None:
-                daily, monthly = self._load_exposure_windows(exposure_index)
-                out["exposure_daily"] = daily
-                out["exposure_monthly"] = monthly
+                out["exposure_embedding"] = self._load_exposure_embeddings(
+                    exposure_index
+                )
         return out
 
     def __len__(self) -> int:
@@ -745,19 +771,20 @@ def _collate_common_static(batch: List[Dict]) -> Dict:
     }
 
 
-def _pad_exposure(batch: List[Dict], key: str, shape: tuple[int, int]) -> torch.Tensor:
+def _pad_exposure_embedding(batch: List[Dict]) -> torch.Tensor:
     max_len = max(int(s["event_seq"].numel()) for s in batch)
-    out = torch.full(
-        (len(batch), max_len, shape[0], shape[1]),
-        float("nan"),
-        dtype=torch.float32,
+    embedding_dim = next(
+        int(s["exposure_embedding"].size(1))
+        for s in batch
+        if "exposure_embedding" in s
+    )
+    out = torch.zeros(
+        len(batch), max_len, embedding_dim, dtype=torch.float32
     )
     for idx, sample in enumerate(batch):
-        value = sample.get(key)
-        if value is None:
-            continue
-        seq_len = int(value.size(0))
-        out[idx, :seq_len] = value
+        value = sample.get("exposure_embedding")
+        if value is not None:
+            out[idx, :value.size(0)] = value
     return out
 
 
@@ -809,9 +836,8 @@ def next_step_collate_fn(batch: List[Dict]) -> Dict:
         "target_multi_hot": target_multi_hot,
     }
     out.update(_collate_common_static(batch))
-    if any("exposure_daily" in s for s in batch):
-        out["exposure_daily"] = _pad_exposure(batch, "exposure_daily", DAILY_EXPOSURE_SHAPE)
-        out["exposure_monthly"] = _pad_exposure(batch, "exposure_monthly", MONTHLY_EXPOSURE_SHAPE)
+    if any("exposure_embedding" in s for s in batch):
+        out["exposure_embedding"] = _pad_exposure_embedding(batch)
     return out
 
 
@@ -847,9 +873,8 @@ def all_future_collate_fn(batch: List[Dict]) -> Dict:
         "exposure": torch.stack([s["exposure"] for s in batch]),
     }
     out.update(_collate_common_static(batch))
-    if any("exposure_daily" in s for s in batch):
-        out["exposure_daily"] = _pad_exposure(batch, "exposure_daily", DAILY_EXPOSURE_SHAPE)
-        out["exposure_monthly"] = _pad_exposure(batch, "exposure_monthly", MONTHLY_EXPOSURE_SHAPE)
+    if any("exposure_embedding" in s for s in batch):
+        out["exposure_embedding"] = _pad_exposure_embedding(batch)
     return out
 
 

encode_exposure_cache.py

@@ -0,0 +1,147 @@
+"""Encode cached exposure windows once for embedding-only model training."""
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch.utils.data import DataLoader, Dataset
+from tqdm import tqdm
+
+from backbones import TimesNetExposureEncoder
+from dataset import ExposureCache
+from train_util import configure_torch_for_training, resolve_device
+
+
+class ExposureEncodingDataset(Dataset):
+    def __init__(self, cache: ExposureCache):
+        self.cache = cache
+
+    def __len__(self) -> int:
+        return len(self.cache.daily)
+
+    def __getitem__(self, index: int) -> tuple[int, torch.Tensor, torch.Tensor]:
+        return (
+            index,
+            torch.from_numpy(
+                np.array(self.cache.daily[index], dtype=np.float32, copy=True)
+            ),
+            torch.from_numpy(
+                np.array(self.cache.monthly[index], dtype=np.float32, copy=True)
+            ),
+        )
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Precompute exposure embeddings from an autoencoder checkpoint"
+    )
+    parser.add_argument("--exposure_cache_dir", required=True)
+    parser.add_argument("--checkpoint", required=True)
+    parser.add_argument("--output_file", default=None)
+    parser.add_argument("--batch_size", type=int, default=64)
+    parser.add_argument("--num_workers", type=int, default=4)
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument(
+        "--output_dtype", choices=["float16", "float32"], default="float16"
+    )
+    parser.add_argument("--overwrite", action="store_true")
+    args = parser.parse_args()
+    if args.batch_size <= 0:
+        parser.error("--batch_size must be positive")
+    return args
+
+
+def main() -> None:
+    args = parse_args()
+    device = resolve_device(args.device)
+    configure_torch_for_training(device)
+    cache_dir = Path(args.exposure_cache_dir)
+    output_path = (
+        Path(args.output_file)
+        if args.output_file
+        else cache_dir / "exposure_embeddings.npy"
+    )
+    if output_path.exists() and not args.overwrite:
+        raise FileExistsError(f"{output_path} exists; pass --overwrite")
+
+    checkpoint_data = torch.load(args.checkpoint, map_location="cpu")
+    model_cfg = checkpoint_data["model_config"]
+    normalization = checkpoint_data["normalization"]
+    encoder = TimesNetExposureEncoder(
+        n_embd=int(model_cfg["n_embd"]),
+        d_model=model_cfg["d_model"],
+        n_layers=int(model_cfg["n_layers"]),
+        top_k=int(model_cfg["top_k"]),
+        n_backbone_blocks=int(model_cfg["n_backbone_blocks"]),
+        backbone_kernel_size=int(model_cfg["backbone_kernel_size"]),
+        backbone_expansion=float(model_cfg["backbone_expansion"]),
+        dropout=float(model_cfg["dropout"]),
+        use_gate=True,
+    )
+    encoder.load_state_dict(checkpoint_data["encoder_state_dict"], strict=True)
+    encoder.to(device).eval()
+
+    cache = ExposureCache(cache_dir)
+    loader_kwargs = {
+        "batch_size": args.batch_size,
+        "shuffle": False,
+        "num_workers": args.num_workers,
+        "pin_memory": device.type == "cuda",
+        "persistent_workers": args.num_workers > 0,
+    }
+    loader = DataLoader(ExposureEncodingDataset(cache), **loader_kwargs)
+    output_dtype = np.float16 if args.output_dtype == "float16" else np.float32
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    output = np.lib.format.open_memmap(
+        output_path,
+        mode="w+",
+        dtype=output_dtype,
+        shape=(len(cache.daily), int(model_cfg["n_embd"])),
+    )
+    stats = {
+        key: torch.as_tensor(normalization[key], device=device).view(1, 1, -1)
+        for key in ("daily_mean", "daily_std", "monthly_mean", "monthly_std")
+    }
+
+    with torch.inference_mode():
+        for indices, daily, monthly in tqdm(loader, desc="Encoding exposure"):
+            daily = daily.to(device, non_blocking=True)
+            monthly = monthly.to(device, non_blocking=True)
+            daily_mask = torch.isfinite(daily)
+            monthly_mask = torch.isfinite(monthly)
+            has_observation = (
+                daily_mask.flatten(1).any(dim=1)
+                | monthly_mask.flatten(1).any(dim=1)
+            )
+            daily = (
+                torch.nan_to_num(daily) - stats["daily_mean"]
+            ) / stats["daily_std"]
+            monthly = (
+                torch.nan_to_num(monthly) - stats["monthly_mean"]
+            ) / stats["monthly_std"]
+            daily = daily * daily_mask
+            monthly = monthly * monthly_mask
+            encoded = encoder(daily, monthly, daily_mask, monthly_mask)
+            encoded = encoded * has_observation.unsqueeze(1)
+            output[indices.numpy()] = encoded.float().cpu().numpy().astype(
+                output_dtype, copy=False
+            )
+    output.flush()
+    metadata = {
+        "checkpoint": str(Path(args.checkpoint).resolve()),
+        "output_file": str(output_path.resolve()),
+        "rows": len(cache.daily),
+        "embedding_dim": int(model_cfg["n_embd"]),
+        "dtype": args.output_dtype,
+    }
+    output_path.with_suffix(".json").write_text(
+        json.dumps(metadata, indent=2), encoding="utf-8"
+    )
+    print(f"Wrote {len(cache.daily):,} embeddings to {output_path}")
+
+
+if __name__ == "__main__":
+    main()

eval_data.py

@@ -8,8 +8,6 @@ import torch
 from torch.nn.utils.rnn import pad_sequence
 
 from dataset import (
-    DAILY_EXPOSURE_SHAPE,
-    MONTHLY_EXPOSURE_SHAPE,
     AllFutureHealthDataset,
     HealthDataset,
 )
@@ -32,7 +30,7 @@ class AllFutureSequenceEvalDataset:
         min_history_events: int = 1,
         min_future_events: int = 1,
         exposure_cache_dir: str | Path | None = None,
-        mask_onset_exposure: bool = False,
+        exposure_embeddings_file: str | Path | None = None,
     ) -> None:
         base = AllFutureHealthDataset(
             data_prefix=data_prefix,
@@ -41,7 +39,7 @@ class AllFutureSequenceEvalDataset:
             min_history_events=min_history_events,
             min_future_events=min_future_events,
             exposure_cache_dir=exposure_cache_dir,
-            mask_onset_exposure=mask_onset_exposure,
+            exposure_embeddings_file=exposure_embeddings_file,
         )
 
         self.base = base
@@ -91,9 +89,9 @@ class AllFutureSequenceEvalDataset:
             "sex": torch.tensor(s["sex"], dtype=torch.long),
         }
         if "exposure_index" in s:
-            daily, monthly = self.base._load_exposure_windows(s["exposure_index"])
-            out["exposure_daily"] = daily
-            out["exposure_monthly"] = monthly
+            out["exposure_embedding"] = self.base._load_exposure_embeddings(
+                s["exposure_index"]
+            )
         return out
 
 
@@ -107,7 +105,7 @@ def load_sequence_eval_dataset(
     min_history_events: int,
     min_future_events: int,
     exposure_cache_dir: str | Path | None = None,
-    mask_onset_exposure: bool = False,
+    exposure_embeddings_file: str | Path | None = None,
 ):
     mode = str(model_target_mode).lower()
     if mode == "next_token":
@@ -117,7 +115,7 @@ def load_sequence_eval_dataset(
             no_event_interval_years=no_event_interval_years,
             include_no_event_in_uts_target=include_no_event_in_uts_target,
             exposure_cache_dir=exposure_cache_dir,
-            mask_onset_exposure=mask_onset_exposure,
+            exposure_embeddings_file=exposure_embeddings_file,
         )
     if mode == "all_future":
         return AllFutureSequenceEvalDataset(
@@ -126,7 +124,7 @@ def load_sequence_eval_dataset(
             min_history_events=min_history_events,
             min_future_events=min_future_events,
             exposure_cache_dir=exposure_cache_dir,
-            mask_onset_exposure=mask_onset_exposure,
+            exposure_embeddings_file=exposure_embeddings_file,
         )
     raise ValueError(f"Unknown model_target_mode: {model_target_mode!r}")
 
@@ -156,34 +154,17 @@ def sequence_eval_collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str,
         "readout_mask": readout_mask,
         "sex": torch.stack([s["sex"] for s in batch]),
     }
-    if any("exposure_daily" in s for s in batch):
-        out["exposure_daily"] = _pad_eval_exposure(
-            batch,
-            "exposure_daily",
-            DAILY_EXPOSURE_SHAPE,
+    if any("exposure_embedding" in s for s in batch):
+        embedding_dim = next(
+            int(s["exposure_embedding"].size(1))
+            for s in batch if "exposure_embedding" in s
         )
-        out["exposure_monthly"] = _pad_eval_exposure(
-            batch,
-            "exposure_monthly",
-            MONTHLY_EXPOSURE_SHAPE,
-        )
-    return out
-
-
-def _pad_eval_exposure(
-    batch: List[Dict[str, torch.Tensor]],
-    key: str,
-    shape: tuple[int, int],
-) -> torch.Tensor:
-    max_len = max(int(s["event_seq"].numel()) for s in batch)
-    out = torch.full(
-        (len(batch), max_len, shape[0], shape[1]),
-        float("nan"),
-        dtype=torch.float32,
+        encoded = torch.zeros(
+            len(batch), event_seq.size(1), embedding_dim, dtype=torch.float32
         )
         for idx, sample in enumerate(batch):
-        value = sample.get(key)
-        if value is None:
-            continue
-        out[idx, : int(value.size(0))] = value
+            value = sample.get("exposure_embedding")
+            if value is not None:
+                encoded[idx, :value.size(0)] = value
+        out["exposure_embedding"] = encoded
     return out

evaluate_auc.py

@@ -323,14 +323,9 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data
         target_mode=model_target_mode,
         dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")),
         dropout=float(cfg_get(args, cfg, "dropout", 0.0)),
-        use_exposure_encoder=bool(cfg_get(args, cfg, "use_exposure_encoder", False)),
-        exposure_d_model=cfg_get(args, cfg, "d_model", 64),
-        exposure_n_layers=int(cfg_get(args, cfg, "exposure_n_layers", 2)),
-        exposure_top_k=int(cfg_get(args, cfg, "exposure_top_k", 2)),
-        exposure_n_backbone_blocks=int(cfg_get(args, cfg, "exposure_n_backbone_blocks", 1)),
-        exposure_backbone_kernel_size=int(cfg_get(args, cfg, "exposure_backbone_kernel_size", 5)),
-        exposure_backbone_expansion=float(cfg_get(args, cfg, "exposure_backbone_expansion", 2.0)),
-        exposure_use_gate=bool(cfg_get(args, cfg, "exposure_use_gate", True)),
+        use_exposure_embeddings=bool(
+            cfg_get(args, cfg, "use_exposure_embeddings", False)
+        ),
     )
 
 
@@ -531,14 +526,9 @@ def infer_readout_hidden(
                         sex=batch_dev["sex"][active],
                         padding_mask=padding_mask[active],
                         t_query=time_seq[active, pos],
-                        exposure_daily=(
-                            batch_dev["exposure_daily"][active]
-                            if "exposure_daily" in batch_dev
-                            else None
-                        ),
-                        exposure_monthly=(
-                            batch_dev["exposure_monthly"][active]
-                            if "exposure_monthly" in batch_dev
+                        exposure_embedding=(
+                            batch_dev["exposure_embedding"][active]
+                            if "exposure_embedding" in batch_dev
                             else None
                         ),
                         target_mode="all_future",
@@ -551,8 +541,7 @@ def infer_readout_hidden(
                     time_seq=time_seq,
                     sex=batch_dev["sex"],
                     padding_mask=padding_mask,
-                    exposure_daily=batch_dev.get("exposure_daily"),
-                    exposure_monthly=batch_dev.get("exposure_monthly"),
+                    exposure_embedding=batch_dev.get("exposure_embedding"),
                     target_mode="next_token",
                 )
                 ro = readout(
@@ -1337,7 +1326,7 @@ def main() -> None:
         min_history_events=int(cfg.get("all_future_min_history_events", 1)),
         min_future_events=int(cfg.get("all_future_min_future_events", 1)),
         exposure_cache_dir=cfg.get("exposure_cache_dir", None),
-        mask_onset_exposure=bool(cfg.get("mask_onset_exposure", False)),
+        exposure_embeddings_file=cfg.get("exposure_embeddings_file", None),
     )
     validate_dataset_metadata(dataset, cfg)
 

models.py

@@ -7,7 +7,6 @@ import torch.nn.functional as F
 from backbones import (
     AgeSinusoidalEncoding,
     GPTBlock,
-    TimesNetExposureEncoder,
 )
 from targets import PAD_IDX
 
@@ -30,16 +29,7 @@ class DeepHealth(nn.Module):
         target_mode: str = "next_token",  # "next_token" or "all_future"
         dist_mode: str = "exponential",  # "exponential", "weibull" or "mixed"
         dropout: float = 0.0,
-        use_exposure_encoder: bool = False,
-        exposure_daily_input_dim: int = 4,
-        exposure_monthly_input_dim: int = 2,
-        exposure_d_model: int | None = 64,
-        exposure_n_layers: int = 2,
-        exposure_top_k: int = 2,
-        exposure_n_backbone_blocks: int = 1,
-        exposure_backbone_kernel_size: int = 5,
-        exposure_backbone_expansion: float = 2.0,
-        exposure_use_gate: bool = True,
+        use_exposure_embeddings: bool = False,
     ):
         super().__init__()
         if target_mode not in ["next_token", "all_future"]:
@@ -53,26 +43,9 @@ class DeepHealth(nn.Module):
             2, n_embd)  # Assuming binary gender
         self.target_mode = target_mode
         self.dist_mode = dist_mode
-        self.use_exposure_encoder = use_exposure_encoder
+        self.use_exposure_embeddings = bool(use_exposure_embeddings)
         self.n_embd = n_embd
         self.vocab_size = vocab_size
-        self.exposure_encoder = (
-            TimesNetExposureEncoder(
-                n_embd=n_embd,
-                daily_input_dim=exposure_daily_input_dim,
-                monthly_input_dim=exposure_monthly_input_dim,
-                d_model=exposure_d_model,
-                n_layers=exposure_n_layers,
-                top_k=exposure_top_k,
-                n_backbone_blocks=exposure_n_backbone_blocks,
-                backbone_kernel_size=exposure_backbone_kernel_size,
-                backbone_expansion=exposure_backbone_expansion,
-                dropout=dropout,
-                use_gate=exposure_use_gate,
-            )
-            if use_exposure_encoder
-            else None
-        )
         nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.02)
         nn.init.zeros_(self.token_embedding.weight[0])
         nn.init.normal_(self.gender_embedding.weight, mean=0.0, std=0.02)
@@ -121,95 +94,6 @@ class DeepHealth(nn.Module):
             dtype=dtype,
         ).masked_fill(~valid, -1e4)[:, None, :, :]
 
-    def _encode_event_exposure(
-        self,
-        exposure_daily: torch.Tensor | None,
-        exposure_monthly: torch.Tensor | None,
-        exposure_daily_mask: torch.Tensor | None,
-        exposure_monthly_mask: torch.Tensor | None,
-        event_shape: tuple[int, int],
-    ) -> torch.Tensor | None:
-        if self.exposure_encoder is None:
-            if exposure_daily is not None or exposure_monthly is not None:
-                raise ValueError(
-                    "Exposure tensors were provided but use_exposure_encoder=False"
-                )
-            return None
-        if exposure_daily is None or exposure_monthly is None:
-            raise ValueError(
-                "exposure_daily and exposure_monthly are required when "
-                "use_exposure_encoder=True"
-            )
-
-        batch_size, event_len = event_shape
-        if exposure_daily.shape[:2] != event_shape:
-            raise ValueError(
-                "exposure_daily must have shape (B, L, T_daily, C_daily), got "
-                f"{tuple(exposure_daily.shape)} for event shape {event_shape}"
-            )
-        if exposure_monthly.shape[:2] != event_shape:
-            raise ValueError(
-                "exposure_monthly must have shape (B, L, T_monthly, C_monthly), got "
-                f"{tuple(exposure_monthly.shape)} for event shape {event_shape}"
-            )
-        if exposure_daily.dim() != 4 or exposure_monthly.dim() != 4:
-            raise ValueError(
-                "exposure_daily and exposure_monthly must both be 4D tensors"
-            )
-
-        def flatten_mask(mask: torch.Tensor | None, ref: torch.Tensor) -> torch.Tensor | None:
-            if mask is None:
-                return None
-            if mask.shape[:2] != event_shape:
-                raise ValueError(
-                    "exposure mask must start with shape (B, L), got "
-                    f"{tuple(mask.shape)}"
-                )
-            if mask.dim() not in {3, 4}:
-                raise ValueError(
-                    "exposure mask must have shape (B, L, T) or (B, L, T, C), "
-                    f"got {tuple(mask.shape)}"
-                )
-            if mask.shape[2] != ref.shape[2]:
-                raise ValueError(
-                    "exposure mask time dimension does not match exposure tensor"
-                )
-            if mask.dim() == 4 and mask.shape[3] != ref.shape[3]:
-                raise ValueError(
-                    "exposure mask channel dimension does not match exposure tensor"
-                )
-            return mask.reshape(batch_size * event_len, *mask.shape[2:])
-
-        daily = exposure_daily.reshape(
-            batch_size * event_len,
-            exposure_daily.size(2),
-            exposure_daily.size(3),
-        )
-        monthly = exposure_monthly.reshape(
-            batch_size * event_len,
-            exposure_monthly.size(2),
-            exposure_monthly.size(3),
-        )
-        daily_mask = flatten_mask(exposure_daily_mask, exposure_daily)
-        monthly_mask = flatten_mask(exposure_monthly_mask, exposure_monthly)
-
-        exposure_device = exposure_daily.device
-        exposure_dtype = self.token_embedding.weight.dtype
-        daily = daily.to(device=exposure_device, dtype=exposure_dtype)
-        monthly = monthly.to(device=exposure_device, dtype=exposure_dtype)
-        if daily_mask is not None:
-            daily_mask = daily_mask.to(device=exposure_device)
-        if monthly_mask is not None:
-            monthly_mask = monthly_mask.to(device=exposure_device)
-
-        exposure_emb = self.exposure_encoder(
-            daily=daily,
-            monthly=monthly,
-            daily_mask=daily_mask,
-            monthly_mask=monthly_mask,
-        )
-        return exposure_emb.reshape(batch_size, event_len, self.n_embd)
-
     def _forward_shared(
         self,
         event_seq: torch.LongTensor,
@@ -218,10 +102,7 @@ class DeepHealth(nn.Module):
         mode: str,
         padding_mask: torch.Tensor | None = None,
         t_query: torch.FloatTensor | None = None,
-        exposure_daily: torch.Tensor | None = None,
-        exposure_monthly: torch.Tensor | None = None,
-        exposure_daily_mask: torch.Tensor | None = None,
-        exposure_monthly_mask: torch.Tensor | None = None,
+        exposure_embedding: torch.Tensor | None = None,
         return_output: bool = False,
         **unused_kwargs,
     ) -> torch.Tensor | DeepHealthOutput:
@@ -240,16 +121,24 @@ class DeepHealth(nn.Module):
 
         event_len = event_seq.size(1)
         h_disease = self.token_embedding(event_seq)
-        h_exposure = self._encode_event_exposure(
-            exposure_daily=exposure_daily,
-            exposure_monthly=exposure_monthly,
-            exposure_daily_mask=exposure_daily_mask,
-            exposure_monthly_mask=exposure_monthly_mask,
-            event_shape=(event_seq.size(0), event_len),
+        if self.use_exposure_embeddings:
+            if exposure_embedding is None:
+                raise ValueError(
+                    "exposure_embedding is required when "
+                    "use_exposure_embeddings=True"
+                )
+            if exposure_embedding.shape != h_disease.shape:
+                raise ValueError(
+                    "exposure_embedding must have shape "
+                    f"{tuple(h_disease.shape)}, got {tuple(exposure_embedding.shape)}"
+                )
+            h_disease = h_disease + exposure_embedding.to(
+                device=h_disease.device, dtype=h_disease.dtype
+            )
+        elif exposure_embedding is not None:
+            raise ValueError(
+                "exposure_embedding provided but use_exposure_embeddings=False"
             )
-        if h_exposure is not None:
-            h_exposure = h_exposure.to(device=event_seq.device, dtype=h_disease.dtype)
-            h_disease = h_disease + h_exposure
         t_disease = time_seq
         h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
 

train_next_step.py

@@ -58,8 +58,7 @@ MODEL_INPUT_KEYS = (
 )
 
 EXPOSURE_INPUT_KEYS = (
-    "exposure_daily",
-    "exposure_monthly",
+    "exposure_embedding",
 )
 
 
@@ -174,8 +173,7 @@ class NextStepTrainingModel(nn.Module):
         sex: torch.Tensor,
         padding_mask: torch.Tensor,
         readout_mask: torch.Tensor | None = None,
-        exposure_daily: torch.Tensor | None = None,
-        exposure_monthly: torch.Tensor | None = None,
+        exposure_embedding: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         hidden = self.model(
             event_seq=event_seq,
@@ -183,8 +181,7 @@ class NextStepTrainingModel(nn.Module):
             sex=sex,
             padding_mask=padding_mask,
             target_mode="next_token",
-            exposure_daily=exposure_daily,
-            exposure_monthly=exposure_monthly,
+            exposure_embedding=exposure_embedding,
         )
         if not isinstance(hidden, torch.Tensor):
             raise TypeError("DeepHealth forward must return a hidden-state tensor")
@@ -224,19 +221,15 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--n_hist_layer", type=int, default=12)
     parser.add_argument("--dropout", type=float, default=0.0)
     parser.add_argument("--exposure_cache_dir", type=str, default=None)
-    parser.add_argument("--mask_onset_exposure", action="store_true")
     parser.add_argument(
-        "--d_model",
-        type=int,
-        default=64,
-        help="Internal TimesNet channel dimension for exposure encoding.",
+        "--exposure_embeddings_file",
+        type=str,
+        default=None,
+        help=(
+            "Precomputed exposure embeddings. Defaults to "
+            "<exposure_cache_dir>/exposure_embeddings.npy."
+        ),
     )
-    parser.add_argument("--exposure_n_layers", type=int, default=2)
-    parser.add_argument("--exposure_top_k", type=int, default=2)
-    parser.add_argument("--exposure_n_backbone_blocks", type=int, default=1)
-    parser.add_argument("--exposure_backbone_kernel_size", type=int, default=5)
-    parser.add_argument("--exposure_backbone_expansion", type=float, default=2.0)
-    parser.add_argument("--no_exposure_gate", action="store_true")
     parser.add_argument("--target_mode", type=str, default="uts",
                         choices=["delphi2m", "uts"])
     parser.add_argument("--readout_name", type=str, default=None,
@@ -311,8 +304,10 @@ def parse_args() -> argparse.Namespace:
         raise ValueError("prefetch_factor must be positive when num_workers > 0")
     if args.exposure_locality_buffer_size < 0:
         raise ValueError("exposure_locality_buffer_size must be non-negative")
-    if args.d_model <= 0:
-        raise ValueError("--d_model must be positive")
+    if args.exposure_embeddings_file and not args.exposure_cache_dir:
+        raise ValueError(
+            "--exposure_cache_dir is required with --exposure_embeddings_file"
+        )
     if args.target_mode == "uts":
         args.readout_name = args.readout_name or "same_time_group_end"
         args.include_no_event_in_uts_target = True
@@ -440,14 +435,7 @@ def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
         target_mode="next_token",
         dist_mode="exponential",
         dropout=args.dropout,
-        use_exposure_encoder=args.exposure_cache_dir is not None,
-        exposure_d_model=args.d_model,
-        exposure_n_layers=args.exposure_n_layers,
-        exposure_top_k=args.exposure_top_k,
-        exposure_n_backbone_blocks=args.exposure_n_backbone_blocks,
-        exposure_backbone_kernel_size=args.exposure_backbone_kernel_size,
-        exposure_backbone_expansion=args.exposure_backbone_expansion,
-        exposure_use_gate=not args.no_exposure_gate,
+        use_exposure_embeddings=args.exposure_embeddings_file is not None,
     )
 
 
@@ -499,9 +487,8 @@ def compute_next_step_loss(
         "padding_mask": batch["padding_mask"],
         "readout_mask": batch["readout_mask"],
     }
-    if "exposure_daily" in batch:
-        model_kwargs["exposure_daily"] = batch["exposure_daily"]
-        model_kwargs["exposure_monthly"] = batch["exposure_monthly"]
+    if "exposure_embedding" in batch:
+        model_kwargs["exposure_embedding"] = batch["exposure_embedding"]
     logits, current_times, output_readout_mask = model(**model_kwargs)
     non_blocking = device.type == "cuda"
     targets = {
@@ -631,16 +618,9 @@ def build_metadata(
         "dataset_metadata": {
             "vocab_size": int(dataset.vocab_size),
         },
-        "use_exposure_encoder": args.exposure_cache_dir is not None,
+        "use_exposure_embeddings": args.exposure_embeddings_file is not None,
         "exposure_cache_dir": args.exposure_cache_dir,
-        "mask_onset_exposure": bool(args.mask_onset_exposure),
-        "d_model": int(args.d_model),
-        "exposure_n_layers": int(args.exposure_n_layers),
-        "exposure_top_k": int(args.exposure_top_k),
-        "exposure_n_backbone_blocks": int(args.exposure_n_backbone_blocks),
-        "exposure_backbone_kernel_size": int(args.exposure_backbone_kernel_size),
-        "exposure_backbone_expansion": float(args.exposure_backbone_expansion),
-        "exposure_use_gate": not bool(args.no_exposure_gate),
+        "exposure_embeddings_file": args.exposure_embeddings_file,
         "num_workers": int(args.num_workers),
         "prefetch_factor": int(args.prefetch_factor),
         "exposure_locality_buffer_size": int(args.exposure_locality_buffer_size),
@@ -659,6 +639,10 @@ def build_metadata(
 
 def main() -> None:
     args = parse_args()
+    if args.exposure_cache_dir and args.exposure_embeddings_file is None:
+        args.exposure_embeddings_file = str(
+            Path(args.exposure_cache_dir) / "exposure_embeddings.npy"
+        )
     device, rank, local_rank, world_size = init_distributed(args)
     set_seed(args.seed + rank)
     configure_torch_for_training(device)
@@ -670,6 +654,7 @@ def main() -> None:
     logger.info(f"Device: {device}")
     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(
         "DataLoader IO: "
         f"num_workers={args.num_workers}, "
@@ -683,7 +668,14 @@ def main() -> None:
         no_event_interval_years=args.no_event_interval_years,
         include_no_event_in_uts_target=args.include_no_event_in_uts_target,
         exposure_cache_dir=args.exposure_cache_dir,
-        mask_onset_exposure=args.mask_onset_exposure,
+        exposure_embeddings_file=args.exposure_embeddings_file,
+    )
+    if dataset.exposure_cache is not None:
+        embedding_dim = int(dataset.exposure_cache.embeddings.shape[1])
+        if embedding_dim != args.n_embd:
+            raise ValueError(
+                f"Exposure embedding dim {embedding_dim} must equal "
+                f"--n_embd={args.n_embd}"
             )
     if args.train_eid_file and args.val_eid_file and args.test_eid_file:
         train_subset, val_subset, test_subset = split_dataset_by_eid_files(