From a39a25a0530513646d982cbf30da95eeb7a8354c Mon Sep 17 00:00:00 2001 From: Amrit Krishnan Date: Sun, 13 Sep 2026 22:10:59 -0400 Subject: [PATCH 1/4] Add a self-supervised window-summary head and the window-statistic probes Frozen probes on full_run_v10 (scripts/probe_summary_signal.py, new) show the backbone's state holds current levels of the GBM's signal panel but not changes from the visit's first value (creatinine delta R^2 0.00) and not fine window counts, under a linear and a tree probe alike. Those are the contents of the GBM's margin on AKI and the counting events. This adds an auxiliary, training-only objective that asks the state to report those statistics: - odyssey/training/summary_targets.py: the target panel (per-signal min/max/mean over 6 h and 24 h, change from visit-first, per drug-class and per-family counts over 6 h and 24 h), computed by the GBM's own feature code at 4-hourly landmark rows; standardization (counts through log1p, winsorized at 0.5/99.5 percentiles); per-visit tables; and the per-chunk lookup that places each target on the last token of its bundle. Targets are computed from the input and never fed in. - odyssey/models/summary_head.py: a linear (or MLP) head and a masked Huber loss that is zero with a live graph where no target applies. - Both sequence models take summary_targets/summary_head_hidden, add the weighted loss to the streaming and steering objectives, and report a summary_loss component. - TrainingConfig gains summary_targets_dir, summary_weight, summary_head_hidden and summary_num_targets (recorded at training time so a checkpoint rebuilds the same head); the training loop, validation and load_run are wired accordingly. None (the default) reproduces every existing run. - scripts/build_summary_targets.py precomputes the targets per shard in parallel and fits the stats. Tests cover the target names, landmark rows, target values (baseline change, counts, NaN for unmeasured), stats, JSON round trip, chunk placement on bundle ends, padding, the head, the masked loss, and the model wiring. 1927 tests pass; ruff, format and mypy clean. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01StsjDmueuDqVNNuVoCEF1P --- odyssey/inference/run_inference.py | 12 + odyssey/models/sequence_model.py | 52 +++ odyssey/models/summary_head.py | 68 +++ odyssey/training/summary_targets.py | 395 ++++++++++++++++++ odyssey/training/train.py | 104 ++++- scripts/build_summary_targets.py | 100 +++++ scripts/probe_summary_signal.py | 389 +++++++++++++++++ .../odyssey/models/test_forecast_objective.py | 8 +- .../odyssey/models/test_streaming_training.py | 1 + tests/odyssey/models/test_summary_head.py | 131 ++++++ .../odyssey/training/test_summary_targets.py | 234 +++++++++++ 11 files changed, 1491 insertions(+), 3 deletions(-) create mode 100644 odyssey/models/summary_head.py create mode 100644 odyssey/training/summary_targets.py create mode 100644 scripts/build_summary_targets.py create mode 100644 scripts/probe_summary_signal.py create mode 100644 tests/odyssey/models/test_summary_head.py create mode 100644 tests/odyssey/training/test_summary_targets.py diff --git a/odyssey/inference/run_inference.py b/odyssey/inference/run_inference.py index ed69d62e..f8c02fb8 100644 --- a/odyssey/inference/run_inference.py +++ b/odyssey/inference/run_inference.py @@ -410,6 +410,18 @@ def load_run( config.event_head_hidden = ( int(first_layer.shape[0]) if first_layer is not None else 0 ) + # Summary head (training-only): rebuild it at the checkpoint's own width + # so the weights load; it is never read at inference. + summary_out = state.get("summary_head.proj.weight") + summary_hidden = state.get("summary_head.proj.0.weight") + if summary_out is None and summary_hidden is not None: + summary_out = state.get("summary_head.proj.2.weight") + config.summary_num_targets = ( + int(summary_out.shape[0]) if summary_out is not None else 0 + ) + config.summary_head_hidden = ( + int(summary_hidden.shape[0]) if summary_hidden is not None else 0 + ) concepts = concepts_for_source( getattr(config, "source", "mimic_iv"), diff --git a/odyssey/models/sequence_model.py b/odyssey/models/sequence_model.py index 30a063ee..83587ff4 100644 --- a/odyssey/models/sequence_model.py +++ b/odyssey/models/sequence_model.py @@ -57,6 +57,7 @@ fold_in_bottleneck_losses, ) from odyssey.models.injection import stream_injection +from odyssey.models.summary_head import SummaryHead, masked_huber_loss from odyssey.models.time_to_event import ( EventHazardHeads, TimeToEventHead, @@ -74,6 +75,7 @@ if TYPE_CHECKING: # the targets live in training; avoid a runtime import cycle from odyssey.training.event_targets import EventHazardTargets + from odyssey.training.summary_targets import SummaryTargets @dataclass @@ -122,6 +124,10 @@ class ForecastObjective: time_weight: float = 0.0 event_hazard_weight: float = 0.0 value_head_weight: float = 0.0 + summary_weight: float = 0.0 + """Weight of the self-supervised window-summary loss + (:mod:`odyssey.models.summary_head`), if the model has a summary head; + the targets come per chunk from :mod:`odyssey.training.summary_targets`.""" class ForwardWithFeatures(NamedTuple): @@ -421,6 +427,19 @@ def _streaming_event_loss( event_heads.edges, ) + def _streaming_summary_loss( + self, + summary_head: SummaryHead | None, + features: torch.Tensor, + summary_targets: Optional["SummaryTargets"], + ) -> torch.Tensor: + """Masked Huber loss of the window-summary head (zero-graph if absent).""" + if summary_head is None or summary_targets is None: + return features.sum() * 0.0 + return masked_huber_loss( + summary_head(features), summary_targets.values, summary_targets.mask + ) + def _streaming_time_loss( self, time_head: TimeToEventHead | None, @@ -486,6 +505,8 @@ def __init__( value_head: bool = False, value_head_hidden: int = 0, source: str = "mimic_iv", + summary_targets: int = 0, + summary_head_hidden: int = 0, ) -> None: """Initialize the baseline sequence model. @@ -524,6 +545,11 @@ def __init__( if value_head else None ) + self.summary_head: SummaryHead | None = ( + SummaryHead(head_in, summary_targets, hidden_size=summary_head_hidden) + if summary_targets > 0 + else None + ) def forward_features( self, @@ -580,6 +606,7 @@ def compute_streaming_loss( *, objective: ForecastObjective | None = None, event_targets: Optional["EventHazardTargets"] = None, + summary_targets: Optional["SummaryTargets"] = None, ) -> tuple[torch.Tensor, dict[str, torch.Tensor], TimeAwareState]: """Compute the forecasting loss over one packed, chunked training step.""" objective = objective or ForecastObjective() @@ -590,11 +617,15 @@ def compute_streaming_loss( time_loss, _ = self._streaming_time_loss(self.time_head, hidden, chunk) event_loss = self._streaming_event_loss(self.event_heads, hidden, event_targets) value_loss, _ = self._streaming_value_loss(self.value_head, hidden, chunk) + summary_loss = self._streaming_summary_loss( + self.summary_head, hidden, summary_targets + ) total = ( task_loss + objective.time_weight * time_loss + objective.event_hazard_weight * event_loss + objective.value_head_weight * value_loss + + objective.summary_weight * summary_loss ) return ( total, @@ -603,6 +634,7 @@ def compute_streaming_loss( "time_loss": time_loss.detach(), "event_loss": event_loss.detach(), "value_loss": value_loss.detach(), + "summary_loss": summary_loss.detach(), }, new_state, ) @@ -632,6 +664,8 @@ def __init__( value_head: bool = False, value_head_hidden: int = 0, source: str = "mimic_iv", + summary_targets: int = 0, + summary_head_hidden: int = 0, ) -> None: """Initialize the concept-bottleneck sequence model. @@ -700,6 +734,11 @@ def __init__( if value_head else None ) + self.summary_head: SummaryHead | None = ( + SummaryHead(head_in, summary_targets, hidden_size=summary_head_hidden) + if summary_targets > 0 + else None + ) def forward( self, @@ -815,6 +854,7 @@ def compute_steering_loss( lifted_ids: torch.Tensor, objective: ForecastObjective | None = None, event_targets: Optional["EventHazardTargets"] = None, + summary_targets: Optional["SummaryTargets"] = None, respond_weight: float = 1.0, express_weight: float = 1.0, forecast_at_injected: bool = True, @@ -869,11 +909,15 @@ def compute_steering_loss( self.event_heads, head_feats, event_targets ) value_loss, _ = self._streaming_value_loss(self.value_head, head_feats, scored) + summary_loss = self._streaming_summary_loss( + self.summary_head, head_feats, summary_targets + ) forecast_loss = ( next_token_loss + objective.time_weight * time_loss + objective.event_hazard_weight * event_loss + objective.value_head_weight * value_loss + + objective.summary_weight * summary_loss ) at = injected & chunk.real_mask if bool(at.any()): @@ -895,6 +939,7 @@ def compute_steering_loss( "time_loss": time_loss.detach(), "event_loss": event_loss.detach(), "value_loss": value_loss.detach(), + "summary_loss": summary_loss.detach(), "respond_loss": respond.detach(), "express_loss": express.detach(), "n_injected": at.sum().detach(), @@ -976,6 +1021,7 @@ def compute_streaming_loss( intervention: BottleneckIntervention | None = None, objective: ForecastObjective | None = None, event_targets: Optional["EventHazardTargets"] = None, + summary_targets: Optional["SummaryTargets"] = None, teacher_alpha_known: float = 0.0, teacher_alpha_unknown: float = 0.0, ) -> tuple[torch.Tensor, dict[str, torch.Tensor], TimeAwareState]: @@ -1043,11 +1089,15 @@ def compute_streaming_loss( self.event_heads, head_feats, event_targets ) value_loss, _ = self._streaming_value_loss(self.value_head, head_feats, chunk) + summary_loss = self._streaming_summary_loss( + self.summary_head, head_feats, summary_targets + ) forecast_loss = ( next_token_loss + objective.time_weight * time_loss + objective.event_hazard_weight * event_loss + objective.value_head_weight * value_loss + + objective.summary_weight * summary_loss ) pool_mask = chunk.patient_end if supervision == "stay" else chunk.visit_end @@ -1058,6 +1108,7 @@ def compute_streaming_loss( "time_loss": time_loss.detach(), "event_loss": event_loss.detach(), "value_loss": value_loss.detach(), + "summary_loss": summary_loss.detach(), "concept_loss": zero, "orthogonality_loss": zero, "observability_loss": zero, @@ -1127,4 +1178,5 @@ def compute_streaming_loss( components["time_loss"] = time_loss.detach() components["event_loss"] = event_loss.detach() components["value_loss"] = value_loss.detach() + components["summary_loss"] = summary_loss.detach() return total, components, new_state diff --git a/odyssey/models/summary_head.py b/odyssey/models/summary_head.py new file mode 100644 index 00000000..58d88ff7 --- /dev/null +++ b/odyssey/models/summary_head.py @@ -0,0 +1,68 @@ +"""Summary head: report window statistics of the chart from the state. + +An auxiliary, training-only readout. At landmark positions it predicts the +standardized window-summary panel of +:mod:`odyssey.training.summary_targets` (per-signal window min/max/mean, +change from the visit's first value, per-family occurrence counts) from +the same features the hazard heads read. The loss is a masked Huber loss: +robust to the residual heavy tails of clinical values, and zero (with a +live graph) where no target applies, so it can be summed unconditionally. +The head is never used at inference; it exists to shape the state. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F # noqa: N812 +from torch import nn + + +class SummaryHead(nn.Module): + """Features -> one standardized value per summary target.""" + + def __init__( + self, in_features: int, num_targets: int, hidden_size: int = 0 + ) -> None: + """Initialize a linear readout, or a GELU MLP when ``hidden_size`` > 0.""" + super().__init__() + self.num_targets = int(num_targets) + self.hidden_size = int(hidden_size) + self.proj: nn.Module = ( + nn.Sequential( + nn.Linear(in_features, self.hidden_size), + nn.GELU(), + nn.Linear(self.hidden_size, self.num_targets), + ) + if self.hidden_size > 0 + else nn.Linear(in_features, self.num_targets) + ) + + def forward(self, features: torch.Tensor) -> torch.Tensor: + """Return ``(..., num_targets)`` predictions.""" + out: torch.Tensor = self.proj(features) + return out + + +def masked_huber_loss( + prediction: torch.Tensor, + target: torch.Tensor, + mask: torch.Tensor, + *, + delta: float = 1.0, +) -> torch.Tensor: + """Huber loss averaged over the ``True`` entries of ``mask``. + + ``prediction``/``target``/``mask`` share the shape ``(..., K)``. Returns + ``0 * prediction.sum()`` when the mask is empty, so the result always + carries a graph and can be added to other losses unconditionally. + """ + if not bool(mask.any()): + return prediction.sum() * 0.0 + per_entry = F.huber_loss( + prediction.float(), target.float(), reduction="none", delta=delta + ) + weights = mask.to(per_entry.dtype) + return (per_entry * weights).sum() / weights.sum() + + +__all__ = ["SummaryHead", "masked_huber_loss"] diff --git a/odyssey/training/summary_targets.py b/odyssey/training/summary_targets.py new file mode 100644 index 00000000..cd63e962 --- /dev/null +++ b/odyssey/training/summary_targets.py @@ -0,0 +1,395 @@ +"""Self-supervised window-summary targets for the summary head. + +The tuned GBM the hazard heads are compared with is hand-fed windowed +statistics of the chart: per-signal minimum, maximum and mean over 6 h and +24 h, the change from the visit's first value, and per-family occurrence +counts. Frozen probes (2026-09-13, ``scripts/probe_summary_signal.py``) +showed the backbone's state holds current levels but not changes from +baseline and not fine window counts, linearly or otherwise. This module +turns those statistics into an auxiliary training target: at 4-hourly +landmark positions the model must *report* them from its own state. The +values are computed from the input stream by the GBM's own feature code, +never fed in as inputs, and the head that predicts them is discarded at +inference. The state is shaped to hold what the statistics need; nothing +is handed to it. + +Pipeline: :func:`compute_summary_targets` runs once per shard (offline, +``scripts/build_summary_targets.py``) and writes one parquet per shard; +:func:`fit_summary_stats` standardizes columns over the training set +(counts through ``log1p``, everything winsorized at the 0.5/99.5 +percentiles so sentinel values cannot dominate); :class:`SummaryTargetTables` +loads the parquets and :func:`summary_targets_for_chunk` looks up, for +every streaming chunk, which positions carry a target -- the last token of +the bundle at each landmark row -- and returns a ``(lanes, T, K)`` tensor +with a mask. Times are hours on the sequence origin, the same clock as +chunk time stamps and the alert harness' landmark rows. +""" + +from __future__ import annotations + +import json +from collections.abc import Iterable, Sequence +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import polars as pl +import torch + +from odyssey.data.alert_events import origin_hours +from odyssey.data.signal_panel import SIGNAL_PANEL +from odyssey.data.streaming import StreamingChunk +from odyssey.inference.baseline_features import ( + DRUG_CLASSES, + FAMILY_LABELS, + StrongFeatureBuilder, + feature_names, +) + + +SIGNAL_STATS: tuple[str, ...] = ( + "min_6h", + "max_6h", + "min_24h", + "max_24h", + "mean_24h", + "delta_visit_first", +) +"""Window statistics asked for per panel signal (the GBM's ``summary_stats`` +group without ``last``/``hours_since_last``, which recency probes already +showed the state holds, and without ``delta_prev``/``ratio_visit_min``, +which are functions of the others).""" + +COUNT_STATS: tuple[str, ...] = ("n_6h", "n_24h") +"""Occurrence counts asked for per drug class and per code family (the +GBM's ``counts_occurrence`` group at its two window lengths).""" + +DEFAULT_LANDMARK_HOURS = 4.0 +WINSOR_PERCENTILES = (0.5, 99.5) +KEY_COLUMNS = ("subject_id", "visit_id", "time_hours") + + +def summary_target_names() -> list[str]: + """Column names of the target panel, in the order the head predicts them.""" + names = [f"{label}.{stat}" for label, _ in SIGNAL_PANEL for stat in SIGNAL_STATS] + names += [ + f"drug.{label}.{stat}" for label, _ in DRUG_CLASSES for stat in COUNT_STATS + ] + names += [ + f"family.{label}.{stat}" for label in FAMILY_LABELS for stat in COUNT_STATS + ] + return names + + +def count_target_mask(names: Sequence[str] | None = None) -> np.ndarray: + """Boolean mask over the panel: which targets are occurrence counts.""" + names = list(names) if names is not None else summary_target_names() + return np.array([n.rsplit(".", 1)[-1] in COUNT_STATS for n in names]) + + +def landmark_rows( + events: pl.DataFrame, landmark_hours: float = DEFAULT_LANDMARK_HOURS +) -> tuple[list[int], list[int], list[float]]: + """``(subject_ids, visit_ids, times)`` of the landmark rows of every visit. + + Every ``landmark_hours`` from a visit's first event to its last, the + row is the visit's last event time at or before that instant, so a row + time always coincides with a real token's time stamp (hours since the + subject's first event, like :class:`~odyssey.data.sequences.PatientSequence`). + """ + origins = origin_hours(events) + timed = ( + events.filter(pl.col("time").is_not_null() & pl.col("hadm_id").is_not_null()) + .join(origins, on="subject_id", how="left") + .with_columns( + ((pl.col("time") - pl.col("_origin")).dt.total_seconds() / 3600.0).alias( + "_hours" + ) + ) + .group_by("subject_id", "hadm_id") + .agg(pl.col("_hours").unique().sort().alias("_hours")) + ) + sids: list[int] = [] + vids: list[int] = [] + times: list[float] = [] + for sid, vid, hours in zip( + timed["subject_id"].to_list(), + timed["hadm_id"].to_list(), + timed["_hours"].to_list(), + ): + arr = np.asarray(hours, dtype=np.float64) + grid = np.arange(arr[0], arr[-1] + 1e-9, landmark_hours) + idx = np.searchsorted(arr, grid, side="right") - 1 + picked = np.unique(arr[idx]) + sids.extend([int(sid)] * len(picked)) + vids.extend([int(vid)] * len(picked)) + times.extend(float(t) for t in picked) + return sids, vids, times + + +def compute_summary_targets( + events: pl.DataFrame, + *, + source: str = "mimic_iv", + landmark_hours: float = DEFAULT_LANDMARK_HOURS, +) -> pl.DataFrame: + """Raw (unstandardized) targets at every landmark row of ``events``. + + Columns: ``subject_id``, ``visit_id``, ``time_hours`` and one column per + :func:`summary_target_names` entry; NaN where a window is empty. + """ + names = summary_target_names() + all_names = feature_names() + idx = [all_names.index(n) for n in names] + sids, vids, times = landmark_rows(events, landmark_hours) + if not sids: + return pl.DataFrame( + schema={ + "subject_id": pl.Int64, + "visit_id": pl.Int64, + "time_hours": pl.Float64, + **dict.fromkeys(names, pl.Float32), + } + ) + builder = StrongFeatureBuilder(events, source=source) + values = builder.features(sids, vids, times)[:, idx].astype(np.float32) + frame = pl.DataFrame( + { + "subject_id": pl.Series(sids, dtype=pl.Int64), + "visit_id": pl.Series(vids, dtype=pl.Int64), + "time_hours": pl.Series(times, dtype=pl.Float64), + } + ) + return frame.with_columns( + [pl.Series(n, values[:, i], dtype=pl.Float32) for i, n in enumerate(names)] + ) + + +@dataclass(frozen=True) +class SummaryTargetStats: + """Per-column standardization fitted on the training targets.""" + + names: list[str] + lo: np.ndarray + """Winsorizing floor per column (after ``log1p`` for counts).""" + hi: np.ndarray + mean: np.ndarray + std: np.ndarray + + def transform(self, raw: np.ndarray) -> np.ndarray: + """Standardize raw targets ``(n, K)``; NaN stays NaN.""" + x = np.asarray(raw, dtype=np.float64).copy() + counts = count_target_mask(self.names) + x[:, counts] = np.log1p(np.clip(x[:, counts], 0.0, None)) + x = np.clip(x, self.lo, self.hi) + out: np.ndarray = ((x - self.mean) / self.std).astype(np.float32) + return out + + def save(self, path: str | Path) -> None: + """Write the stats as JSON.""" + Path(path).write_text( + json.dumps( + { + "names": self.names, + "lo": self.lo.tolist(), + "hi": self.hi.tolist(), + "mean": self.mean.tolist(), + "std": self.std.tolist(), + }, + indent=1, + ) + ) + + @classmethod + def load(cls, path: str | Path) -> SummaryTargetStats: + """Read stats written by :meth:`save`.""" + d = json.loads(Path(path).read_text()) + return cls( + names=list(d["names"]), + lo=np.asarray(d["lo"], dtype=np.float64), + hi=np.asarray(d["hi"], dtype=np.float64), + mean=np.asarray(d["mean"], dtype=np.float64), + std=np.asarray(d["std"], dtype=np.float64), + ) + + +def fit_summary_stats(frames: Iterable[pl.DataFrame]) -> SummaryTargetStats: + """Fit :class:`SummaryTargetStats` over the raw target frames of a split. + + Counts go through ``log1p``; every column is winsorized at + :data:`WINSOR_PERCENTILES` of its finite values before mean/std, so a + handful of sentinel readings cannot set the scale. A column with no + finite value, or no spread, gets mean 0 and std 1. + """ + names = summary_target_names() + counts = count_target_mask(names) + parts = [f.select(names).to_numpy().astype(np.float64) for f in frames] + x = np.concatenate(parts, axis=0) if parts else np.zeros((0, len(names))) + x[:, counts] = np.log1p(np.clip(x[:, counts], 0.0, None)) + lo = np.zeros(len(names)) + hi = np.ones(len(names)) + mean = np.zeros(len(names)) + std = np.ones(len(names)) + for j in range(len(names)): + col = x[:, j] + col = col[np.isfinite(col)] + if col.size == 0: + continue + lo[j], hi[j] = np.percentile(col, WINSOR_PERCENTILES) + clipped = np.clip(col, lo[j], hi[j]) + mean[j] = clipped.mean() + s = clipped.std() + std[j] = s if s > 0 else 1.0 + return SummaryTargetStats(names=names, lo=lo, hi=hi, mean=mean, std=std) + + +class SummaryTargetTables: + """Standardized targets per ``(subject, visit)``, queried per chunk. + + Each key maps to ``(times, values)``: sorted landmark row times and the + matching ``(rows, K)`` standardized targets, stored as float16 to keep + a whole split (millions of rows) in memory. + """ + + def __init__(self, stats: SummaryTargetStats) -> None: + self.stats = stats + self.num_targets = len(stats.names) + self._rows: dict[tuple[int, int], tuple[np.ndarray, np.ndarray]] = {} + + def __len__(self) -> int: + """Return the number of visits with targets.""" + return len(self._rows) + + def add_frame(self, frame: pl.DataFrame) -> None: + """Ingest one raw target frame from :func:`compute_summary_targets`.""" + if frame.height == 0: + return + frame = frame.sort("subject_id", "visit_id", "time_hours") + values = self.stats.transform(frame.select(self.stats.names).to_numpy()) + keys = frame.select("subject_id", "visit_id").to_numpy() + times = frame["time_hours"].to_numpy().astype(np.float64) + change = np.flatnonzero(np.any(keys[1:] != keys[:-1], axis=1)) + 1 + starts = np.concatenate([[0], change]) + ends = np.concatenate([change, [len(keys)]]) + for a, b in zip(starts, ends): + key = (int(keys[a, 0]), int(keys[a, 1])) + self._rows[key] = (times[a:b], values[a:b].astype(np.float16)) + + @classmethod + def from_dir( + cls, directory: str | Path, stats: SummaryTargetStats + ) -> SummaryTargetTables: + """Load every ``*.parquet`` under ``directory``.""" + tables = cls(stats) + for path in sorted(Path(directory).glob("*.parquet")): + tables.add_frame(pl.read_parquet(path)) + return tables + + def lookup( + self, subject_id: int, visit_id: int + ) -> tuple[np.ndarray, np.ndarray] | None: + """Return ``(times, values)`` for one visit, or ``None``.""" + return self._rows.get((subject_id, visit_id)) + + +@dataclass +class SummaryTargets: + """``(lanes, T, K)`` standardized targets and where they apply.""" + + values: torch.Tensor + mask: torch.Tensor + + +def _bundle_end_mask( + time_stamps: np.ndarray, subject_ids: np.ndarray, real: np.ndarray +) -> np.ndarray: + """Mark the last real token of each same-time bundle, per lane.""" + lanes, chunk_len = time_stamps.shape + end = np.zeros_like(real) + if chunk_len == 0: + return end + same_next = np.zeros((lanes, chunk_len), dtype=bool) + same_next[:, :-1] = ( + (time_stamps[:, 1:] == time_stamps[:, :-1]) + & (subject_ids[:, 1:] == subject_ids[:, :-1]) + & real[:, 1:] + ) + return real & ~same_next + + +def summary_targets_for_chunk( + chunk: StreamingChunk, tables: SummaryTargetTables +) -> SummaryTargets | None: + """Targets for one chunk, or ``None`` when no position carries one. + + A position carries a target when its ``(subject, visit, time)`` is a + landmark row of the tables and it is the last real token of its + bundle, so the state has read the whole bundle the row summarizes. + """ + sids = chunk.subject_ids.detach().cpu().numpy() + vids = chunk.visit_ids.detach().cpu().numpy() + times = chunk.batch.aux.time_stamps.detach().cpu().numpy().astype(np.float64) + real = chunk.real_mask.detach().cpu().numpy().astype(bool) + lanes, chunk_len = sids.shape + k = tables.num_targets + values = np.zeros((lanes, chunk_len, k), dtype=np.float32) + mask = np.zeros((lanes, chunk_len), dtype=bool) + ends = _bundle_end_mask(times, sids, real) + keys = np.stack([sids.reshape(-1), vids.reshape(-1)], axis=1) + unique_keys, inverse = np.unique(keys, axis=0, return_inverse=True) + inverse = inverse.reshape(-1) + for i, (s, v) in enumerate(unique_keys.tolist()): + rows = tables.lookup(int(s), int(v)) + if rows is None: + continue + row_times, row_values = rows + flat = np.flatnonzero((inverse == i) & ends.reshape(-1)) + if flat.size == 0: + continue + t = times.reshape(-1)[flat] + pos = np.searchsorted(row_times, t) + pos = np.clip(pos, 0, len(row_times) - 1) + hit = np.isclose(row_times[pos], t, rtol=0.0, atol=1e-6) + if not hit.any(): + continue + lane_idx, pos_idx = np.divmod(flat[hit], chunk_len) + values[lane_idx, pos_idx] = row_values[pos[hit]].astype(np.float32) + mask[lane_idx, pos_idx] = True + if not mask.any(): + return None + device = chunk.batch.concept_ids.device + finite = np.isfinite(values) + full_mask = finite & mask[:, :, None] + values = np.where(full_mask, values, 0.0).astype(np.float32) + return SummaryTargets( + values=torch.from_numpy(values).to(device), + mask=torch.from_numpy(full_mask).to(device), + ) + + +def load_summary_tables( + directory: str | Path, stats_path: str | Path | None = None +) -> SummaryTargetTables: + """Load a split's tables; the stats default to ``/../stats.json``.""" + directory = Path(directory) + stats = SummaryTargetStats.load( + stats_path if stats_path is not None else directory.parent / "stats.json" + ) + return SummaryTargetTables.from_dir(directory, stats) + + +__all__ = [ + "COUNT_STATS", + "DEFAULT_LANDMARK_HOURS", + "SIGNAL_STATS", + "SummaryTargetStats", + "SummaryTargetTables", + "SummaryTargets", + "compute_summary_targets", + "count_target_mask", + "fit_summary_stats", + "landmark_rows", + "load_summary_tables", + "summary_target_names", + "summary_targets_for_chunk", +] diff --git a/odyssey/training/train.py b/odyssey/training/train.py index e3744fce..ed7cbbf8 100644 --- a/odyssey/training/train.py +++ b/odyssey/training/train.py @@ -113,6 +113,12 @@ SteeringSchedule, choose_injection, ) +from odyssey.training.summary_targets import ( + SummaryTargetTables, + load_summary_tables, + summary_target_names, + summary_targets_for_chunk, +) from odyssey.utils.env_fingerprint import write_run_provenance @@ -454,6 +460,31 @@ class TrainingConfig: value_head -- lets an A/B separate "better input encoding" from "better output objective".""" + summary_targets_dir: str | None = None + """Directory of precomputed self-supervised window-summary targets + (scripts/build_summary_targets.py: ``/train/*.parquet``, + ``/tuning/*.parquet``, ``/stats.json``). When set, a summary + head (odyssey.models.summary_head) is trained to report, at 4-hourly + landmark positions, the GBM's window statistics of the chart so far + (per-signal window min/max/mean, change from the visit's first value, + per-family occurrence counts) from the same features the hazard heads + read. The statistics are computed from the input and never fed in; + the head is unused at inference. Frozen probes on 2026-09-13 showed the + state lacks exactly these (change from baseline R^2 0.00), which is + the content of the GBM's margin on AKI and the counting events. None + (the default) reproduces every existing run.""" + + summary_weight: float = 0.5 + """Weight of the summary head's masked Huber loss.""" + + summary_head_hidden: int = 0 + """Hidden width of the summary head; 0 = linear, which asks for the + same linear readability the probes measure.""" + + summary_num_targets: int = 0 + """Width of the summary head, recorded at training time so a checkpoint + rebuilds the same head even if the target panel definition changes.""" + randint_prob: float = 0.25 """Intervention-aware training (CEM's RandInt): at every training position, each observed concept's mixing probability is replaced by @@ -691,6 +722,8 @@ def build_model( value_head=bool(getattr(config, "value_head", False)), value_head_hidden=int(getattr(config, "value_head_hidden", 0) or 0), source=getattr(config, "source", "mimic_iv"), + summary_targets=int(getattr(config, "summary_num_targets", 0) or 0), + summary_head_hidden=int(getattr(config, "summary_head_hidden", 0) or 0), ) return ConceptBottleneckSequenceModel( backbone=backbone, @@ -710,6 +743,8 @@ def build_model( value_head=bool(getattr(config, "value_head", False)), value_head_hidden=int(getattr(config, "value_head_hidden", 0) or 0), source=getattr(config, "source", "mimic_iv"), + summary_targets=int(getattr(config, "summary_num_targets", 0) or 0), + summary_head_hidden=int(getattr(config, "summary_head_hidden", 0) or 0), ) @@ -978,6 +1013,11 @@ def build_objective( value_head_weight=( config.value_head_weight if getattr(config, "value_head", False) else 0.0 ), + summary_weight=( + config.summary_weight + if getattr(config, "summary_targets_dir", None) + else 0.0 + ), ) @@ -992,6 +1032,7 @@ def evaluate_streaming( supervision: ConceptSupervision = "stay", objective: ForecastObjective | None = None, event_tables: EventTimeTables | None = None, + summary_tables: SummaryTargetTables | None = None, ) -> dict[str, float]: """Average loss components over one (partial), gradient-free sampler pass.""" model.eval() @@ -1009,9 +1050,18 @@ def evaluate_streaming( if event_tables is not None else None ) + summary_targets = ( + summary_targets_for_chunk(chunk, summary_tables) + if summary_tables is not None + else None + ) if isinstance(model, BaselineSequenceModel): _, components, state = model.compute_streaming_loss( - chunk, state=state, objective=objective, event_targets=event_targets + chunk, + state=state, + objective=objective, + event_targets=event_targets, + summary_targets=summary_targets, ) else: _, components, state = model.compute_streaming_loss( @@ -1022,6 +1072,7 @@ def evaluate_streaming( supervision=supervision, objective=objective, event_targets=event_targets, + summary_targets=summary_targets, ) state = _detach_state(state) for key, value in components.items(): @@ -1042,6 +1093,8 @@ def _combined_val_loss( time_weight: float = 0.0, event_hazard_weight: float = 0.0, value_head_weight: float = 0.0, + *, + summary_weight: float = 0.0, ) -> float: """Compute the same task + weighted-auxiliary combination the training loss uses. @@ -1056,6 +1109,7 @@ def _combined_val_loss( + time_weight * components.get("time_loss", 0.0) + event_hazard_weight * components.get("event_loss", 0.0) + value_head_weight * components.get("value_loss", 0.0) + + summary_weight * components.get("summary_loss", 0.0) + weights.concept * components.get("concept_loss", 0.0) + weights.orthogonality * components.get("orthogonality_loss", 0.0) + weights.observability * components.get("observability_loss", 0.0) @@ -1078,6 +1132,9 @@ def train(config: TrainingConfig) -> Path: # noqa: PLR0912, PLR0915 config.event_hazards = any(k.startswith("event_heads.") for k in resume_keys) config.value_head = any(k.startswith("value_head.") for k in resume_keys) del resume_keys + config.summary_num_targets = ( + len(summary_target_names()) if config.summary_targets_dir else 0 + ) (output_dir / "config.json").write_text(json.dumps(asdict(config), indent=2)) _activate_run_sidecars(config) @@ -1224,6 +1281,7 @@ def make_train_patients(epoch: int) -> Iterator[PatientSequence]: shuffle_seed=config.seed + epoch, ) + train_summary_tables, tuning_summary_tables = _summary_tables(config) corpus = PreparedCorpus( vocab=vocab, concepts=concepts, @@ -1237,6 +1295,8 @@ def make_train_patients(epoch: int) -> Iterator[PatientSequence]: tuning_event_tables=tuning_event_tables, tuning_events_binned=tuning_events_binned, make_train_patients=make_train_patients, + train_summary_tables=train_summary_tables, + tuning_summary_tables=tuning_summary_tables, ) return _run_training(config, output_dir, device, corpus) @@ -1258,6 +1318,28 @@ class PreparedCorpus: tuning_events_binned: pl.DataFrame make_train_patients: Callable[[int], Iterator[PatientSequence]] """epoch -> the training patient stream for that epoch.""" + train_summary_tables: SummaryTargetTables | None = None + tuning_summary_tables: SummaryTargetTables | None = None + + +def _summary_tables( + config: TrainingConfig, +) -> tuple[SummaryTargetTables | None, SummaryTargetTables | None]: + """Load the precomputed window-summary targets for both splits, if configured.""" + directory = getattr(config, "summary_targets_dir", None) + if not directory: + return None, None + root = Path(directory) + logger.info("[data] loading window-summary targets from %s", root) + train = load_summary_tables(root / "train", root / "stats.json") + tuning = load_summary_tables(root / "tuning", root / "stats.json") + logger.info( + "[data] summary targets: %d train visits, %d tuning visits, %d targets", + len(train), + len(tuning), + train.num_targets, + ) + return train, tuning def _train_streaming(config: TrainingConfig, output_dir: Path, device: str) -> Path: @@ -1360,6 +1442,7 @@ def make_train_patients(epoch: int) -> Iterator[PatientSequence]: shuffle_seed=config.seed + epoch, ) + train_summary_tables, tuning_summary_tables = _summary_tables(config) corpus = PreparedCorpus( vocab=vocab, concepts=concepts, @@ -1373,6 +1456,8 @@ def make_train_patients(epoch: int) -> Iterator[PatientSequence]: tuning_event_tables=tuning_event_tables, tuning_events_binned=tuning_events_binned, make_train_patients=make_train_patients, + train_summary_tables=train_summary_tables, + tuning_summary_tables=tuning_summary_tables, ) return _run_training(config, output_dir, device, corpus) @@ -1420,6 +1505,8 @@ def _run_training( # noqa: PLR0912, PLR0915 tuning_labels, tuning_masks = corpus.tuning_labels, corpus.tuning_masks train_event_tables = corpus.train_event_tables tuning_event_tables = corpus.tuning_event_tables + train_summary_tables = corpus.train_summary_tables + tuning_summary_tables = corpus.tuning_summary_tables tuning_events_binned = corpus.tuning_events_binned model = build_model(config, vocab_size=len(vocab), num_concepts=len(concepts)).to( @@ -1616,9 +1703,18 @@ def make_tuning_sampler() -> StreamingSampler: if train_event_tables is not None else None ) + summary_targets = ( + summary_targets_for_chunk(chunk, train_summary_tables) + if train_summary_tables is not None + else None + ) if isinstance(model, BaselineSequenceModel): total, components, state = model.compute_streaming_loss( - chunk, state=state, objective=objective, event_targets=event_targets + chunk, + state=state, + objective=objective, + event_targets=event_targets, + summary_targets=summary_targets, ) elif ( steering is not None @@ -1639,6 +1735,7 @@ def make_tuning_sampler() -> StreamingSampler: lifted_ids=steering.lifted[injection.concept_index], objective=objective, event_targets=event_targets, + summary_targets=summary_targets, respond_weight=config.respond_weight, express_weight=config.express_weight, forecast_at_injected=config.steering_forecast_at_injected, @@ -1664,6 +1761,7 @@ def make_tuning_sampler() -> StreamingSampler: intervention=intervention, objective=objective, event_targets=event_targets, + summary_targets=summary_targets, teacher_alpha_known=annealed_alpha( global_step, config.teacher_anneal_steps, @@ -1709,6 +1807,7 @@ def make_tuning_sampler() -> StreamingSampler: supervision=config.concept_supervision, # type: ignore[arg-type] objective=objective, event_tables=tuning_event_tables, + summary_tables=tuning_summary_tables, ) fields = { "step": global_step, @@ -1727,6 +1826,7 @@ def make_tuning_sampler() -> StreamingSampler: objective.time_weight, objective.event_hazard_weight, objective.value_head_weight, + summary_weight=objective.summary_weight, ) if val_loss < best_val_loss: best_val_loss = val_loss diff --git a/scripts/build_summary_targets.py b/scripts/build_summary_targets.py new file mode 100644 index 00000000..0103e23a --- /dev/null +++ b/scripts/build_summary_targets.py @@ -0,0 +1,100 @@ +"""Precompute the self-supervised window-summary targets for a data root. + +Writes ``/train/.parquet`` and ``/tuning/.parquet`` +(one raw target frame per shard, see +:func:`odyssey.training.summary_targets.compute_summary_targets`) and +``/stats.json`` (standardization fitted on the train frames). Point a +training run at ```` with ``summary_targets_dir``. + +Shards are independent, so they are processed in parallel. The prepare +step (code normalization, history recap) is the training run's own, taken +from a run config or the defaults, so the target rows' time stamps match +the tokens the model streams. + + uv run python scripts/build_summary_targets.py \\ + --data-root ~/data/mimiciv_3.1_v1/data \\ + --out ~/data/mimiciv_3.1_v1/summary_targets \\ + --config ~/runs/full_run_v10/config.json --workers 12 +""" + +from __future__ import annotations + +import argparse +import json +import logging +from concurrent.futures import ProcessPoolExecutor +from pathlib import Path + +import polars as pl + +from odyssey.training.data import load_meds_shard +from odyssey.training.shard_stream import make_preparer, shard_paths +from odyssey.training.summary_targets import ( + DEFAULT_LANDMARK_HOURS, + compute_summary_targets, + fit_summary_stats, +) + + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") +logger = logging.getLogger("build_summary_targets") + + +def _one(args: tuple[str, str, str, bool, bool, float]) -> str: + path, out, source, normalize, recap, landmark_hours = args + out_path = Path(out) / (Path(path).stem + ".parquet") + if out_path.exists(): + return f"skip {out_path.name}" + prepare = make_preparer( + source=source, normalize_medications=normalize, history_recap=recap + ) + events = prepare(load_meds_shard(Path(path))) + frame = compute_summary_targets( + events, source=source, landmark_hours=landmark_hours + ) + frame.write_parquet(out_path) + return f"{out_path.name}: {frame.height} rows" + + +def main() -> None: + """Build targets for the train and tuning splits and fit the stats.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--data-root", required=True) + parser.add_argument("--out", required=True) + parser.add_argument( + "--config", default=None, help="a run's config.json for source/prepare flags" + ) + parser.add_argument("--max-train-shards", type=int, default=None) + parser.add_argument("--max-tuning-shards", type=int, default=None) + parser.add_argument("--landmark-hours", type=float, default=DEFAULT_LANDMARK_HOURS) + parser.add_argument("--workers", type=int, default=4) + args = parser.parse_args() + + cfg = json.loads(Path(args.config).read_text()) if args.config else {} + source = str(cfg.get("source", "mimic_iv")) + normalize = bool(cfg.get("normalize_medications", True)) + recap = bool(cfg.get("history_recap", False)) + out = Path(args.out) + for split, cap in ( + ("train", args.max_train_shards), + ("tuning", args.max_tuning_shards), + ): + (out / split).mkdir(parents=True, exist_ok=True) + paths = shard_paths(Path(args.data_root) / split, cap) + logger.info("%s: %d shards -> %s", split, len(paths), out / split) + jobs = [ + (str(p), str(out / split), source, normalize, recap, args.landmark_hours) + for p in paths + ] + with ProcessPoolExecutor(max_workers=args.workers) as pool: + for msg in pool.map(_one, jobs): + logger.info(" %s", msg) + logger.info("fitting standardization on the train frames") + frames = (pl.read_parquet(p) for p in sorted((out / "train").glob("*.parquet"))) + stats = fit_summary_stats(frames) + stats.save(out / "stats.json") + logger.info("wrote %s (%d targets)", out / "stats.json", len(stats.names)) + + +if __name__ == "__main__": + main() diff --git a/scripts/probe_summary_signal.py b/scripts/probe_summary_signal.py new file mode 100644 index 00000000..867848a1 --- /dev/null +++ b/scripts/probe_summary_signal.py @@ -0,0 +1,389 @@ +"""Frozen pre/post-bottleneck probes for the GBM's window-statistic features. + +Companion to scripts/probe_counting_signal.py. That script asks whether the +backbone's own hidden state already encodes the GBM's occurrence COUNTS +(the `counts_occurrence` group). This one asks the same question for the +`summary_stats` group -- per-signal window minimum/maximum/mean over 6 h and +24 h, the change from the previous value, the change from the visit's first +value and the ratio to the visit's minimum -- which the 2026-08-24 +feature-group ablation found leads the GBM's margin on acute kidney injury +(and is literally what KDIGO defines AKI on). Same frozen Ridge probe, same +row selection (4-hourly landmarks), same GBM feature code for the targets. + +Two differences from the counting script, both forced by the targets: + +- A window statistic is NaN when the signal was never measured in that + window, so each target is probed on its own finite rows and the row + count is reported beside its R^2. Multi-output fitting would discard + almost every row. +- Ratios and deltas have heavy tails, so Spearman rank correlation is + reported next to R^2 as a scale-free companion. + +`--feature-group both` also runs the counting group on the same extracted +embeddings, so one invocation gives a run both baselines at once. + +`--probe hgb` fits a HistGradientBoosting regressor beside the Ridge probe +(same rows, same targets), the nonlinear check the counting work used on +2026-08-29: a target the linear probe cannot read but the tree probe can is +"encoded, but not linearly". `--targets change` restricts to the change +statistics (delta_prev, delta_visit_first, ratio_visit_min) plus the level +anchors of a few key signals. Targets are winsorized at the train 0.5/99.5 +percentiles for BOTH probes when `--winsorize` is set, so artifact values +(sentinel MAPs) cannot dominate R^2. + +Not wired into any CI/registry path. Run directly: + + uv run python scripts/probe_summary_signal.py \ + --run-dir ~/runs/full_run_v10 \ + --train-shard-dir ~/data/mimiciv_3.1_v1/data/train \ + --held-out-shard-dir ~/data/mimiciv_3.1_v1/data/held_out \ + --max-train-shards 5 --max-held-out-shards 4 --feature-group both +""" + +from __future__ import annotations + +import argparse +import logging +from collections import defaultdict + +import numpy as np +import polars as pl +import torch +from scipy.stats import spearmanr +from sklearn.ensemble import HistGradientBoostingRegressor +from sklearn.linear_model import Ridge +from sklearn.metrics import r2_score +from sklearn.preprocessing import StandardScaler + +from odyssey.data.alert_events import alert_events_for +from odyssey.data.sidecars import activate_sidecars +from odyssey.data.value_binning import add_value_tokens +from odyssey.inference.alerts import _load_prepared_raw, _visit_starts +from odyssey.inference.baseline_features import ( + CONTEXT_FEATURES, + StrongFeatureBuilder, + feature_names, +) +from odyssey.inference.embedding_probe import collect_embeddings +from odyssey.inference.run_inference import load_run +from odyssey.models.sequence_model import ConceptBottleneckSequenceModel + + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") +logger = logging.getLogger("probe_summary_signal") + +SUMMARY_STATS: tuple[str, ...] = ( + "last", + "mean_24h", + "min_24h", + "max_24h", + "min_6h", + "max_6h", + "delta_prev", + "delta_visit_first", + "ratio_visit_min", +) +"""`last` is included as an anchor: the most recent value is the easiest +statistic to read off a sequence model, so its R^2 bounds the others.""" + +MIN_TRAIN_ROWS = 500 +MIN_HELD_ROWS = 200 + +CHANGE_STATS: tuple[str, ...] = ("delta_prev", "delta_visit_first", "ratio_visit_min") +ANCHOR_SIGNALS: tuple[str, ...] = ( + "creatinine", + "bun", + "lactate", + "map_noninvasive", + "sbp_noninvasive", + "heart_rate", + "platelets", + "hemoglobin", + "urine_output", +) +ANCHOR_STATS: tuple[str, ...] = ("last", "min_6h", "mean_24h") + + +def _is_change_target(name: str) -> bool: + """Change statistics for every signal, plus level anchors for key signals.""" + signal, stat = name.rsplit(".", 1) + return stat in CHANGE_STATS or (signal in ANCHOR_SIGNALS and stat in ANCHOR_STATS) + + +def _counting_columns(names: list[str]) -> list[int]: + """Occurrence-count columns, exactly as scripts/probe_counting_signal.py.""" + keep_suffixes = (".n_6h", ".n_24h", ".n_visit", ".ever_visit") + idx = [ + i + for i, n in enumerate(names) + if (n.startswith("drug.") or n.startswith("family.")) + and n.endswith(keep_suffixes) + ] + idx += [names.index("n_prior_visits"), names.index("n_events_visit")] + return sorted(idx) + + +def _summary_columns(names: list[str]) -> list[int]: + """Per-signal window-statistic columns (the GBM's `summary_stats` group).""" + context = set(CONTEXT_FEATURES) + return [ + i + for i, n in enumerate(names) + if n not in context + and not n.startswith(("drug.", "family.")) + and n.rsplit(".", 1)[-1] in SUMMARY_STATS + ] + + +def probe_one( + train_x: np.ndarray, + train_y: np.ndarray, + test_x: np.ndarray, + test_y: np.ndarray, + *, + alpha: float, + kind: str = "ridge", +) -> tuple[float, float]: + """Fit one frozen probe (Ridge or HistGradientBoosting); return (R^2, rho).""" + x_scaler = StandardScaler().fit(train_x) + y_mean, y_std = float(train_y.mean()), float(train_y.std() or 1.0) + reg: Ridge | HistGradientBoostingRegressor + if kind == "hgb": + reg = HistGradientBoostingRegressor( + max_iter=300, + learning_rate=0.1, + max_leaf_nodes=31, + early_stopping=True, + random_state=0, + ) + else: + reg = Ridge(alpha=alpha) + reg.fit(x_scaler.transform(train_x), (train_y - y_mean) / y_std) + pred = reg.predict(x_scaler.transform(test_x)) + r2 = float(r2_score((test_y - y_mean) / y_std, pred)) + rho = float(spearmanr(test_y, pred).correlation) + return r2, rho + + +def probe_group( # noqa: PLR0913 + label: str, + idx: list[int], + names: list[str], + *, + train_y_all: np.ndarray, + held_y_all: np.ndarray, + train_pre: np.ndarray, + train_post: np.ndarray, + held_pre: np.ndarray, + held_post: np.ndarray, + alpha: float, + kinds: tuple[str, ...] = ("ridge",), + winsorize: bool = False, + train_subsample: int = 0, + seed: int = 0, +) -> list[dict[str, object]]: + """Probe every column of a group on its own finite rows, per probe kind.""" + rng = np.random.default_rng(seed) + rows: list[dict[str, object]] = [] + for i in idx: + ty, hy = train_y_all[:, i], held_y_all[:, i] + m_tr, m_he = np.isfinite(ty), np.isfinite(hy) + if ( + m_tr.sum() < MIN_TRAIN_ROWS + or m_he.sum() < MIN_HELD_ROWS + or np.nanstd(hy) == 0 + ): + logger.info( + " skip %-40s (train %d, held %d rows)", + names[i], + m_tr.sum(), + m_he.sum(), + ) + continue + tr = np.flatnonzero(m_tr) + if train_subsample and len(tr) > train_subsample: + tr = rng.choice(tr, size=train_subsample, replace=False) + ty_use, hy_use = ty[tr], hy[m_he] + if winsorize: + lo, hi = np.percentile(ty_use, [0.5, 99.5]) + ty_use, hy_use = np.clip(ty_use, lo, hi), np.clip(hy_use, lo, hi) + row: dict[str, object] = {"feature": names[i], "n_held": int(m_he.sum())} + for kind in kinds: + pre = probe_one( + train_pre[tr], ty_use, held_pre[m_he], hy_use, alpha=alpha, kind=kind + ) + post = probe_one( + train_post[tr], ty_use, held_post[m_he], hy_use, alpha=alpha, kind=kind + ) + row[f"{kind}_pre_r2"], row[f"{kind}_pre_rho"] = pre + row[f"{kind}_post_r2"], row[f"{kind}_post_rho"] = post + rows.append(row) + logger.info( + " %-40s " + + " ".join( + f"{k}: pre={row[f'{k}_pre_r2']:.3f} post={row[f'{k}_post_r2']:.3f} rho={row[f'{k}_pre_rho']:.2f}" + for k in kinds + ), + names[i], + ) + logger.info("[%s] %d targets probed", label, len(rows)) + return rows + + +def report(label: str, rows: list[dict[str, object]], kinds: tuple[str, ...]) -> None: + """Log per-statistic medians for every probe kind, then print the CSV block.""" + if not rows: + logger.info("[%s] nothing to report", label) + return + cols = [ + f"{k}_{side}_{m}" + for k in kinds + for side in ("pre", "post") + for m in ("r2", "rho") + ] + by_stat: dict[str, list[dict[str, object]]] = defaultdict(list) + for r in rows: + by_stat[str(r["feature"]).rsplit(".", 1)[-1]].append(r) + logger.info("[%s] MEDIAN by statistic; columns: %s", label, " ".join(cols)) + for stat, rs in sorted(by_stat.items()): + meds = [float(np.median([float(r[c]) for r in rs])) for c in cols] # type: ignore[arg-type] + logger.info( + " %-20s %s (n=%d)", stat, " ".join(f"{m:6.3f}" for m in meds), len(rs) + ) + print(f"\n# group={label}") + print("feature,n_held," + ",".join(cols)) + for r in rows: + print(f"{r['feature']},{r['n_held']}," + ",".join(str(r[c]) for c in cols)) + + +def main() -> None: # noqa: PLR0915 + """Extract pre/post-bottleneck embeddings and report window-stat recovery.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--run-dir", required=True) + parser.add_argument("--train-shard-dir", required=True) + parser.add_argument("--held-out-shard-dir", required=True) + parser.add_argument("--max-train-shards", type=int, default=5) + parser.add_argument("--max-held-out-shards", type=int, default=4) + parser.add_argument("--landmark-hours", type=float, default=4.0) + parser.add_argument("--num-lanes", type=int, default=64) + parser.add_argument("--chunk-size", type=int, default=512) + parser.add_argument("--ridge-alpha", type=float, default=1.0) + parser.add_argument( + "--feature-group", choices=("summary", "counts", "both"), default="summary" + ) + parser.add_argument("--probe", choices=("ridge", "hgb", "both"), default="ridge") + parser.add_argument( + "--targets", + choices=("all", "change"), + default="all", + help="'change' keeps delta_prev/delta_visit_first/ratio_visit_min plus level anchors", + ) + parser.add_argument( + "--winsorize", + action="store_true", + help="clip targets at train 0.5/99.5 percentiles", + ) + parser.add_argument( + "--train-subsample", + type=int, + default=0, + help="rows per target for fitting (0 = all)", + ) + args = parser.parse_args() + kinds: tuple[str, ...] = ("ridge", "hgb") if args.probe == "both" else (args.probe,) + + device = "cuda" if torch.cuda.is_available() else "cpu" + model, vocab, binner, config = load_run(args.run_dir, device=device) + if not isinstance(model, ConceptBottleneckSequenceModel): + raise ValueError(f"{args.run_dir} is not a concept-bottleneck run") + source = getattr(config, "source", "mimic_iv") + task_set = getattr(config, "task_set", "v1") + landmark_alerts = [a for a in alert_events_for(task_set) if not a.next_visit] + + def load_split( + shard_dir: str, max_shards: int + ) -> tuple[pl.DataFrame, dict[tuple[int, int], float], StrongFeatureBuilder]: + activate_sidecars(shard_dir) + raw = _load_prepared_raw(shard_dir, max_shards, config, source) + visit_start = _visit_starts(raw) + binned = add_value_tokens(raw, binner, source=source) + builder = StrongFeatureBuilder(raw, source=source) + del raw + return binned, visit_start, builder + + logger.info( + "loading %d train shard(s) from %s", args.max_train_shards, args.train_shard_dir + ) + train_binned, train_visit_start, train_builder = load_split( + args.train_shard_dir, args.max_train_shards + ) + logger.info( + "loading %d held-out shard(s) from %s", + args.max_held_out_shards, + args.held_out_shard_dir, + ) + held_binned, held_visit_start, held_builder = load_split( + args.held_out_shard_dir, args.max_held_out_shards + ) + + def embed( + binned: pl.DataFrame, visit_start: dict[tuple[int, int], float] + ) -> tuple[list[tuple[int, int, float]], np.ndarray, np.ndarray]: + keys, pre, post, _, _, _ = collect_embeddings( + model, + binned, + vocab, + landmark_alerts=landmark_alerts, + visit_end_alerts=[], + visit_start=visit_start, + landmark_hours=args.landmark_hours, + num_lanes=args.num_lanes, + chunk_size=args.chunk_size, + device=device, + ) + return keys, pre, post + + logger.info("extracting train embeddings") + train_keys, train_pre, train_post = embed(train_binned, train_visit_start) + logger.info("train: %d landmark rows", len(train_keys)) + logger.info("extracting held-out embeddings") + held_keys, held_pre, held_post = embed(held_binned, held_visit_start) + logger.info("held-out: %d landmark rows", len(held_keys)) + + names = feature_names() + train_y = train_builder.features( + [k[0] for k in train_keys], + [k[1] for k in train_keys], + [k[2] for k in train_keys], + ) + held_y = held_builder.features( + [k[0] for k in held_keys], [k[1] for k in held_keys], [k[2] for k in held_keys] + ) + + groups = {"summary": _summary_columns, "counts": _counting_columns} + chosen = list(groups) if args.feature_group == "both" else [args.feature_group] + for label in chosen: + idx = groups[label](names) + if args.targets == "change" and label == "summary": + idx = [i for i in idx if _is_change_target(names[i])] + logger.info("[%s] %d candidate targets", label, len(idx)) + rows = probe_group( + label, + idx, + names, + train_y_all=train_y, + held_y_all=held_y, + train_pre=train_pre, + train_post=train_post, + held_pre=held_pre, + held_post=held_post, + alpha=args.ridge_alpha, + kinds=kinds, + winsorize=args.winsorize, + train_subsample=args.train_subsample, + ) + report(label, rows, kinds) + + +if __name__ == "__main__": + main() diff --git a/tests/odyssey/models/test_forecast_objective.py b/tests/odyssey/models/test_forecast_objective.py index ce5c60a1..b921f99f 100644 --- a/tests/odyssey/models/test_forecast_objective.py +++ b/tests/odyssey/models/test_forecast_objective.py @@ -219,7 +219,13 @@ def test_baseline_model_shares_objective_and_time_head() -> None: total, comp, _ = model.compute_streaming_loss( chunk, objective=ForecastObjective(bundle_invariant=True, time_weight=0.5) ) - assert set(comp) == {"task_loss", "time_loss", "event_loss", "value_loss"} + assert set(comp) == { + "task_loss", + "time_loss", + "event_loss", + "value_loss", + "summary_loss", + } total.backward() diff --git a/tests/odyssey/models/test_streaming_training.py b/tests/odyssey/models/test_streaming_training.py index 07759b70..03758588 100644 --- a/tests/odyssey/models/test_streaming_training.py +++ b/tests/odyssey/models/test_streaming_training.py @@ -102,6 +102,7 @@ def test_streaming_loss_is_finite_when_a_patient_ends_in_chunk() -> None: "time_loss", "event_loss", "value_loss", + "summary_loss", "concept_loss", "orthogonality_loss", "observability_loss", diff --git a/tests/odyssey/models/test_summary_head.py b/tests/odyssey/models/test_summary_head.py new file mode 100644 index 00000000..d721279a --- /dev/null +++ b/tests/odyssey/models/test_summary_head.py @@ -0,0 +1,131 @@ +"""Tests for the summary head, its masked loss, and its wiring into the models.""" + +import torch + +from odyssey.data.sequences import PatientSequence +from odyssey.data.streaming import PackedLaneSampler +from odyssey.models.backbones.tiny_gru import TinyGRUBackbone +from odyssey.models.sequence_model import ( + BaselineSequenceModel, + ConceptBottleneckSequenceModel, + ForecastObjective, +) +from odyssey.models.summary_head import SummaryHead, masked_huber_loss +from odyssey.training.summary_targets import SummaryTargets + + +VOCAB = 40 +HIDDEN = 8 +K = 5 + + +def _seq(sid: int, n: int) -> PatientSequence: + return PatientSequence( + subject_id=sid, + concept_ids=[1 + (i % (VOCAB - 1)) for i in range(n)], + type_ids=[1] * n, + time_stamps=[float(i) for i in range(n)], + ages=[40.0] * n, + visit_orders=[0] * n, + visit_segments=[0] * n, + ) + + +def _backbone() -> TinyGRUBackbone: + return TinyGRUBackbone( + vocab_size=VOCAB, hidden_size=HIDDEN, num_layers=1, padding_idx=0 + ) + + +def test_masked_huber_loss_averages_over_the_mask_only() -> None: + pred = torch.zeros(1, 3, 2) + target = torch.tensor([[[1.0, 100.0], [0.5, 0.0], [0.0, 0.0]]]) + mask = torch.tensor([[[True, False], [True, False], [False, False]]]) + loss = masked_huber_loss(pred, target, mask) + # huber(1.0) = 0.5, huber(0.5) = 0.125 -> mean 0.3125; the 100 is masked + assert loss.item() == 0.3125 + + +def test_masked_huber_loss_is_zero_with_a_graph_when_mask_is_empty() -> None: + pred = torch.zeros(1, 2, 2, requires_grad=True) + loss = masked_huber_loss( + pred, torch.ones(1, 2, 2), torch.zeros(1, 2, 2, dtype=torch.bool) + ) + assert loss.item() == 0.0 + loss.backward() + assert pred.grad is not None + + +def test_summary_head_shapes_linear_and_mlp() -> None: + x = torch.randn(2, 3, HIDDEN) + assert SummaryHead(HIDDEN, K)(x).shape == (2, 3, K) + mlp = SummaryHead(HIDDEN, K, hidden_size=6) + assert mlp(x).shape == (2, 3, K) + assert mlp.hidden_size == 6 + + +def test_models_without_the_head_report_a_zero_summary_loss() -> None: + model = ConceptBottleneckSequenceModel( + backbone=_backbone(), + vocab_size=VOCAB, + num_concepts=2, + embedding_dim=4, + padding_idx=0, + ) + assert model.summary_head is None + chunk = PackedLaneSampler( + iter([_seq(1, 6)]), num_lanes=1, chunk_size=6 + ).next_chunk() + total, components, _ = model.compute_streaming_loss(chunk, {1: torch.zeros(2)}) + assert components["summary_loss"].item() == 0.0 + assert torch.isfinite(total) + + +def test_summary_loss_trains_the_state_and_is_weighted() -> None: + torch.manual_seed(0) + model = ConceptBottleneckSequenceModel( + backbone=_backbone(), + vocab_size=VOCAB, + num_concepts=2, + embedding_dim=4, + padding_idx=0, + summary_targets=K, + ) + assert model.summary_head is not None + model.eval() # dropout would otherwise make the two passes differ + chunk = PackedLaneSampler( + iter([_seq(1, 6)]), num_lanes=1, chunk_size=6 + ).next_chunk() + values = torch.randn(1, 6, K) + mask = torch.zeros(1, 6, K, dtype=torch.bool) + mask[0, 3] = True + targets = SummaryTargets(values=values, mask=mask) + zero_w = ForecastObjective(summary_weight=0.0) + one_w = ForecastObjective(summary_weight=1.0) + total0, comp0, _ = model.compute_streaming_loss( + chunk, {1: torch.zeros(2)}, objective=zero_w, summary_targets=targets + ) + total1, comp1, _ = model.compute_streaming_loss( + chunk, {1: torch.zeros(2)}, objective=one_w, summary_targets=targets + ) + assert comp0["summary_loss"].item() > 0.0 + assert total1.item() > total0.item() + assert abs((total1 - total0).item() - comp1["summary_loss"].item()) < 1e-5 + total1.backward() + assert model.backbone.embeddings.embeddings.word_embeddings.weight.grad is not None + + +def test_baseline_model_accepts_summary_targets_too() -> None: + model = BaselineSequenceModel( + backbone=_backbone(), vocab_size=VOCAB, padding_idx=0, summary_targets=K + ) + chunk = PackedLaneSampler( + iter([_seq(1, 4)]), num_lanes=1, chunk_size=4 + ).next_chunk() + targets = SummaryTargets( + values=torch.zeros(1, 4, K), mask=torch.ones(1, 4, K, dtype=torch.bool) + ) + _, components, _ = model.compute_streaming_loss( + chunk, objective=ForecastObjective(summary_weight=1.0), summary_targets=targets + ) + assert "summary_loss" in components diff --git a/tests/odyssey/training/test_summary_targets.py b/tests/odyssey/training/test_summary_targets.py new file mode 100644 index 00000000..0824d07b --- /dev/null +++ b/tests/odyssey/training/test_summary_targets.py @@ -0,0 +1,234 @@ +"""Tests for the self-supervised window-summary targets and their chunk lookup.""" + +from datetime import datetime, timedelta + +import numpy as np +import polars as pl +import pytest +import torch + +from odyssey.data.streaming import StreamingChunk +from odyssey.data.types import AuxiliaryInputs, ClinicalSequenceBatch +from odyssey.training.summary_targets import ( + COUNT_STATS, + SIGNAL_STATS, + SummaryTargetStats, + SummaryTargetTables, + compute_summary_targets, + count_target_mask, + fit_summary_stats, + landmark_rows, + load_summary_tables, + summary_target_names, + summary_targets_for_chunk, +) + + +T0 = datetime(2024, 1, 1) + + +def _frame(rows): + return pl.DataFrame( + rows, + schema={ + "subject_id": pl.Int64, + "code": pl.Utf8, + "time": pl.Datetime, + "numeric_value": pl.Float32, + "hadm_id": pl.Int64, + }, + orient="row", + ) + + +def _creat(h: float, v: float, sid: int = 1, hadm: int = 10): + return (sid, "LAB//RESULT//50912//mg/dL::HIGH", T0 + timedelta(hours=h), v, hadm) + + +def _events() -> pl.DataFrame: + return _frame( + [ + (1, "MEDS_BIRTH", T0 - timedelta(days=365.25 * 60), None, None), + (1, "GENDER//F", None, None, None), + _creat(0.0, 1.0), + _creat(5.0, 1.2), + _creat(9.0, 2.2), + (1, "MEDICATION//norepinephrine", T0 + timedelta(hours=8.5), None, 10), + (1, "MEDICATION//norepinephrine", T0 + timedelta(hours=9.0), None, 10), + ] + ) + + +def test_target_names_cover_signals_and_counts() -> None: + names = summary_target_names() + assert len(names) == len(set(names)) + assert "creatinine.delta_visit_first" in names + assert "drug.vasopressor.n_6h" in names + assert "family.lab.n_24h" in names + assert all(n.rsplit(".", 1)[-1] in SIGNAL_STATS + COUNT_STATS for n in names) + counts = count_target_mask(names) + assert counts.sum() == sum(n.rsplit(".", 1)[-1] in COUNT_STATS for n in names) + + +def test_landmark_rows_are_every_4h_at_the_last_event_time() -> None: + sids, vids, times = landmark_rows(_events(), landmark_hours=4.0) + # visit spans 0..9 h: landmarks at 0, 4, 8 -> last event at or before: 0, 0, 5 + assert sids == [1, 1] # the 0 h and 4 h landmarks both map to the 0 h event + assert vids == [10, 10] + assert times == [0.0, 5.0] + + +def test_compute_summary_targets_reports_baseline_change_and_counts() -> None: + frame = compute_summary_targets(_events(), landmark_hours=4.0) + assert frame.columns[:3] == ["subject_id", "visit_id", "time_hours"] + at5 = frame.filter(pl.col("time_hours") == 5.0) + assert at5.height == 1 + assert at5["creatinine.delta_visit_first"][0] == pytest.approx(0.2) + assert at5["creatinine.max_24h"][0] == pytest.approx(1.2) + assert at5["creatinine.min_24h"][0] == pytest.approx(1.0) + assert at5["drug.vasopressor.n_6h"][0] == 0.0 + assert at5["family.lab.n_24h"][0] == 2.0 + at0 = frame.filter(pl.col("time_hours") == 0.0) + assert at0["creatinine.delta_visit_first"][0] == pytest.approx(0.0) + # a signal never measured is NaN, never a fake zero + assert np.isnan(at0["lactate.min_6h"][0]) + + +def test_compute_summary_targets_on_empty_frame_has_the_full_schema() -> None: + frame = compute_summary_targets(_events().head(0)) + assert frame.height == 0 + assert set(summary_target_names()) <= set(frame.columns) + + +def test_stats_standardize_counts_through_log1p_and_winsorize() -> None: + names = summary_target_names() + j_count = names.index("drug.vasopressor.n_6h") + j_delta = names.index("creatinine.delta_visit_first") + n = 1000 + raw = np.full((n, len(names)), np.nan) + rng = np.random.default_rng(0) + raw[:, j_count] = rng.poisson(2.0, n) + raw[:, j_delta] = rng.normal(0.0, 1.0, n) + raw[0, j_delta] = 1e6 # a sentinel that must not set the scale + frame = pl.DataFrame( + { + "subject_id": [1] * n, + "visit_id": [1] * n, + "time_hours": np.arange(n, dtype=float), + } + ) + frame = frame.with_columns( + [pl.Series(name, raw[:, i], dtype=pl.Float32) for i, name in enumerate(names)] + ) + stats = fit_summary_stats([frame]) + assert stats.std[j_delta] < 2.0 + z = stats.transform(raw) + assert abs(float(np.nanmean(z[:, j_delta]))) < 0.2 + assert float(z[0, j_delta]) < 5.0 # clipped, not 1e6 standard deviations + assert np.isnan(z[:, names.index("lactate.min_6h")]).all() + # counts: log1p(0) -> the smallest standardized value + zero = stats.transform( + np.where(np.arange(len(names)) == j_count, 0.0, np.nan)[None, :] + ) + assert float(zero[0, j_count]) < 0.0 + + +def test_stats_round_trip_through_json(tmp_path) -> None: + stats = fit_summary_stats([compute_summary_targets(_events())]) + stats.save(tmp_path / "stats.json") + back = SummaryTargetStats.load(tmp_path / "stats.json") + assert back.names == stats.names + np.testing.assert_allclose(back.mean, stats.mean) + np.testing.assert_allclose(back.std, stats.std) + + +def _chunk(subject_ids, visit_ids, times, real=None) -> StreamingChunk: + lanes, length = np.asarray(times).shape + real_mask = ( + torch.ones(lanes, length, dtype=torch.bool) + if real is None + else torch.tensor(real) + ) + return StreamingChunk( + batch=ClinicalSequenceBatch( + concept_ids=torch.ones(lanes, length, dtype=torch.long), + aux=AuxiliaryInputs( + type_ids=torch.ones(lanes, length, dtype=torch.long), + time_stamps=torch.tensor(times, dtype=torch.float32), + ages=torch.full((lanes, length), 40.0), + visit_orders=torch.zeros(lanes, length, dtype=torch.long), + visit_segments=torch.zeros(lanes, length, dtype=torch.long), + ), + ), + targets=torch.ones(lanes, length, dtype=torch.long), + reset_mask=torch.zeros(lanes, length, dtype=torch.bool), + real_mask=real_mask, + subject_ids=torch.tensor(subject_ids), + patient_end=torch.zeros(lanes, length, dtype=torch.bool), + visit_ids=torch.tensor(visit_ids), + visit_end=torch.zeros(lanes, length, dtype=torch.bool), + ) + + +def _tables() -> SummaryTargetTables: + frame = compute_summary_targets(_events()) + stats = fit_summary_stats([frame]) + tables = SummaryTargetTables(stats) + tables.add_frame(frame) + return tables + + +def test_chunk_targets_land_on_bundle_ends_of_landmark_rows() -> None: + tables = _tables() + assert len(tables) == 1 + # lane 0: patient 1, visit 10; the 5.0 h bundle has two tokens (a + # panel), the target must sit on the LAST of them. lane 1: unknown patient. + chunk = _chunk( + subject_ids=[[1, 1, 1, 1], [7, 7, 7, 7]], + visit_ids=[[10, 10, 10, 10], [1, 1, 1, 1]], + times=[[0.0, 5.0, 5.0, 9.0], [0.0, 4.0, 8.0, 12.0]], + ) + out = summary_targets_for_chunk(chunk, tables) + assert out is not None + hit = out.mask.any(dim=-1) + assert hit.tolist() == [[True, False, True, False], [False, False, False, False]] + k = tables.num_targets + assert out.values.shape == (2, 4, k) + j = tables.stats.names.index("creatinine.delta_visit_first") + # the 5 h row's standardized delta is above the 0 h row's (0.2 vs 0.0) + assert out.values[0, 2, j] > out.values[0, 0, j] + # NaN targets are masked out and zero-filled, never passed as NaN + assert torch.isfinite(out.values).all() + assert not out.mask[0, 2, tables.stats.names.index("lactate.min_6h")] + + +def test_chunk_targets_ignore_padding_and_return_none_when_nothing_matches() -> None: + tables = _tables() + chunk = _chunk( + subject_ids=[[1, 1]], + visit_ids=[[10, 10]], + times=[[5.0, 5.0]], + real=[[True, False]], + ) + out = summary_targets_for_chunk(chunk, tables) + assert out is not None + assert out.mask.any(dim=-1).tolist() == [[True, False]] + assert ( + summary_targets_for_chunk( + _chunk(subject_ids=[[1]], visit_ids=[[10]], times=[[3.0]]), tables + ) + is None + ) + + +def test_tables_load_from_a_directory_of_parquets(tmp_path) -> None: + frame = compute_summary_targets(_events()) + (tmp_path / "train").mkdir() + frame.write_parquet(tmp_path / "train" / "shard_0.parquet") + fit_summary_stats([frame]).save(tmp_path / "stats.json") + tables = load_summary_tables(tmp_path / "train") + assert len(tables) == 1 + times, values = tables.lookup(1, 10) + assert times.tolist() == [0.0, 5.0] + assert values.shape == (2, tables.num_targets) + assert tables.lookup(1, 11) is None From a214830ae66603cbc845f833f81762cc5e524de8 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 14 Sep 2026 02:31:15 +0000 Subject: [PATCH 2/4] [pre-commit.ci] Add auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- scripts/gemini/out/evals/gemini_full_14m_v1_eval_forecast.json | 2 +- scripts/gemini/out/evals/gemini_full_DEC_v12_alerts.json | 2 +- .../gemini/out/evals/gemini_full_DEC_v12_allshards_alerts.json | 2 +- scripts/gemini/out/evals/gemini_full_DEC_v12_concept_atlas.json | 2 +- scripts/gemini/out/evals/gemini_full_DEC_v12_eval_forecast.json | 2 +- scripts/gemini/out/evals/gemini_full_DEC_v12_steering_full.json | 2 +- scripts/gemini/out/evals/gemini_full_v10_15c_alerts.json | 2 +- scripts/gemini/out/evals/gemini_full_v10_15c_alerts_cis.json | 2 +- .../gemini/out/evals/gemini_full_v10_15c_counterfactual.json | 2 +- scripts/gemini/out/evals/gemini_full_v10_15c_eval_forecast.json | 2 +- .../gemini/out/evals/gemini_full_v10_15c_intervention_cis.json | 2 +- scripts/gemini/out/evals/gemini_smoke_2_eval_forecast.json | 2 +- 12 files changed, 12 insertions(+), 12 deletions(-) diff --git a/scripts/gemini/out/evals/gemini_full_14m_v1_eval_forecast.json b/scripts/gemini/out/evals/gemini_full_14m_v1_eval_forecast.json index 6a6b2a3d..a6adb1bc 100644 --- a/scripts/gemini/out/evals/gemini_full_14m_v1_eval_forecast.json +++ b/scripts/gemini/out/evals/gemini_full_14m_v1_eval_forecast.json @@ -129,4 +129,4 @@ "n_positive_gaps": 3639349 }, "tail_slice": null -} \ No newline at end of file +} diff --git a/scripts/gemini/out/evals/gemini_full_DEC_v12_alerts.json b/scripts/gemini/out/evals/gemini_full_DEC_v12_alerts.json index 84211bd8..d110a460 100644 --- a/scripts/gemini/out/evals/gemini_full_DEC_v12_alerts.json +++ b/scripts/gemini/out/evals/gemini_full_DEC_v12_alerts.json @@ -1913,4 +1913,4 @@ }, "landmark_protocol_version": 4 } -] \ No newline at end of file +] diff --git a/scripts/gemini/out/evals/gemini_full_DEC_v12_allshards_alerts.json b/scripts/gemini/out/evals/gemini_full_DEC_v12_allshards_alerts.json index efcf285e..be67255b 100644 --- a/scripts/gemini/out/evals/gemini_full_DEC_v12_allshards_alerts.json +++ b/scripts/gemini/out/evals/gemini_full_DEC_v12_allshards_alerts.json @@ -1913,4 +1913,4 @@ }, "landmark_protocol_version": 4 } -] \ No newline at end of file +] diff --git a/scripts/gemini/out/evals/gemini_full_DEC_v12_concept_atlas.json b/scripts/gemini/out/evals/gemini_full_DEC_v12_concept_atlas.json index 66eedf2b..96da3caa 100644 --- a/scripts/gemini/out/evals/gemini_full_DEC_v12_concept_atlas.json +++ b/scripts/gemini/out/evals/gemini_full_DEC_v12_concept_atlas.json @@ -13011,4 +13011,4 @@ "mean_activation": 0.4445817383871991 } ] -} \ No newline at end of file +} diff --git a/scripts/gemini/out/evals/gemini_full_DEC_v12_eval_forecast.json b/scripts/gemini/out/evals/gemini_full_DEC_v12_eval_forecast.json index f8b0699c..2a5636d7 100644 --- a/scripts/gemini/out/evals/gemini_full_DEC_v12_eval_forecast.json +++ b/scripts/gemini/out/evals/gemini_full_DEC_v12_eval_forecast.json @@ -543,4 +543,4 @@ }, "value_metrics": null, "tail_slice": null -} \ No newline at end of file +} diff --git a/scripts/gemini/out/evals/gemini_full_DEC_v12_steering_full.json b/scripts/gemini/out/evals/gemini_full_DEC_v12_steering_full.json index 41e55272..634f6db7 100644 --- a/scripts/gemini/out/evals/gemini_full_DEC_v12_steering_full.json +++ b/scripts/gemini/out/evals/gemini_full_DEC_v12_steering_full.json @@ -14428,4 +14428,4 @@ "sign_agreement": 1.0 } ] -} \ No newline at end of file +} diff --git a/scripts/gemini/out/evals/gemini_full_v10_15c_alerts.json b/scripts/gemini/out/evals/gemini_full_v10_15c_alerts.json index 7b04a88c..4c6dc770 100644 --- a/scripts/gemini/out/evals/gemini_full_v10_15c_alerts.json +++ b/scripts/gemini/out/evals/gemini_full_v10_15c_alerts.json @@ -1913,4 +1913,4 @@ }, "landmark_protocol_version": 4 } -] \ No newline at end of file +] diff --git a/scripts/gemini/out/evals/gemini_full_v10_15c_alerts_cis.json b/scripts/gemini/out/evals/gemini_full_v10_15c_alerts_cis.json index e559497c..80b3f7df 100644 --- a/scripts/gemini/out/evals/gemini_full_v10_15c_alerts_cis.json +++ b/scripts/gemini/out/evals/gemini_full_v10_15c_alerts_cis.json @@ -704,4 +704,4 @@ } } } -} \ No newline at end of file +} diff --git a/scripts/gemini/out/evals/gemini_full_v10_15c_counterfactual.json b/scripts/gemini/out/evals/gemini_full_v10_15c_counterfactual.json index 39b70901..0d486be5 100644 --- a/scripts/gemini/out/evals/gemini_full_v10_15c_counterfactual.json +++ b/scripts/gemini/out/evals/gemini_full_v10_15c_counterfactual.json @@ -219,4 +219,4 @@ } } } -} \ No newline at end of file +} diff --git a/scripts/gemini/out/evals/gemini_full_v10_15c_eval_forecast.json b/scripts/gemini/out/evals/gemini_full_v10_15c_eval_forecast.json index 2fb9ffb2..caf52962 100644 --- a/scripts/gemini/out/evals/gemini_full_v10_15c_eval_forecast.json +++ b/scripts/gemini/out/evals/gemini_full_v10_15c_eval_forecast.json @@ -383,4 +383,4 @@ }, "value_metrics": null, "tail_slice": null -} \ No newline at end of file +} diff --git a/scripts/gemini/out/evals/gemini_full_v10_15c_intervention_cis.json b/scripts/gemini/out/evals/gemini_full_v10_15c_intervention_cis.json index 7debd16e..033ff473 100644 --- a/scripts/gemini/out/evals/gemini_full_v10_15c_intervention_cis.json +++ b/scripts/gemini/out/evals/gemini_full_v10_15c_intervention_cis.json @@ -111,4 +111,4 @@ "n_boot": 2000 } } -} \ No newline at end of file +} diff --git a/scripts/gemini/out/evals/gemini_smoke_2_eval_forecast.json b/scripts/gemini/out/evals/gemini_smoke_2_eval_forecast.json index 6eba1956..d9622d0c 100644 --- a/scripts/gemini/out/evals/gemini_smoke_2_eval_forecast.json +++ b/scripts/gemini/out/evals/gemini_smoke_2_eval_forecast.json @@ -129,4 +129,4 @@ "n_positive_gaps": 3639349 }, "tail_slice": null -} \ No newline at end of file +} From e773363b90d7a22fe39838a7f4a532aa0877b777 Mon Sep 17 00:00:00 2001 From: Amrit Krishnan Date: Sun, 13 Sep 2026 22:35:17 -0400 Subject: [PATCH 3/4] Fix CI: docstring hook, and keep hooks off the verbatim GEMINI exports - check-docstring-first read the module-level attribute docstrings in the new target module and probe script as second module docstrings; they are comments now. - pre-commit.ci's autofix had rewritten 12 scripts/gemini/out/ JSONs (trailing newlines) on this branch, the same way it did on PR #250. This carries PR #250's hook config over: scripts/gemini/out/ is excluded from every hook, and mypy/typos skip the two cohort producers committed verbatim, which is also what has been failing main's code check. The 12 exports are byte-identical to main again. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01StsjDmueuDqVNNuVoCEF1P --- .pre-commit-config.yaml | 12 ++++++++---- _typos.toml | 8 +++++++- odyssey/training/summary_targets.py | 12 ++++++------ .../out/evals/gemini_full_14m_v1_eval_forecast.json | 2 +- .../gemini/out/evals/gemini_full_DEC_v12_alerts.json | 2 +- .../evals/gemini_full_DEC_v12_allshards_alerts.json | 2 +- .../out/evals/gemini_full_DEC_v12_concept_atlas.json | 2 +- .../out/evals/gemini_full_DEC_v12_eval_forecast.json | 2 +- .../out/evals/gemini_full_DEC_v12_steering_full.json | 2 +- .../gemini/out/evals/gemini_full_v10_15c_alerts.json | 2 +- .../out/evals/gemini_full_v10_15c_alerts_cis.json | 2 +- .../evals/gemini_full_v10_15c_counterfactual.json | 2 +- .../out/evals/gemini_full_v10_15c_eval_forecast.json | 2 +- .../evals/gemini_full_v10_15c_intervention_cis.json | 2 +- .../out/evals/gemini_smoke_2_eval_forecast.json | 2 +- scripts/probe_summary_signal.py | 4 ++-- 16 files changed, 35 insertions(+), 25 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index ef7b5d02..757cd061 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,5 +1,6 @@ -# Third-party LaTeX template files under paper/ are kept verbatim. -exclude: ^paper/(sn-jnl\.cls|sn-nature\.bst|template-.*|main\.pdf)$ +# Kept verbatim, never rewritten by hooks: third-party LaTeX template files +# under paper/, and scripts/gemini/out/ (GEMINI's own exported reports). +exclude: ^(paper/(sn-jnl\.cls|sn-nature\.bst|template-.*|main\.pdf)|scripts/gemini/out/.*)$ repos: - repo: https://github.com/pre-commit/pre-commit-hooks rev: v6.0.0 # Use the ref you want to point at @@ -42,7 +43,10 @@ repos: entry: python3 -m mypy --config-file pyproject.toml language: system types: [python] - exclude: "tests" + # tests, plus the cohort producers committed verbatim as they ran + # (scripts/cohort/README.md); pyproject's mypy exclude does not apply + # to files pre-commit passes by name. + exclude: ^(tests/|scripts/cohort/cohort_check_(mimic|eicu)\.py$) - repo: https://github.com/crate-ci/typos rev: v1 @@ -56,7 +60,7 @@ repos: # etc.), not ours to rename; see _typos.toml for the same exclusion # (this one is what actually stops the hook, since pre-commit's own # file filtering runs before typos ever sees its own config) - exclude: ^(docs/experiments\.md|odyssey/data/resources/.*\.csv|scripts/gemini/out/.*)$ + exclude: ^(docs/experiments\.md|odyssey/data/resources/.*\.csv|scripts/gemini/out/.*|scripts/cohort/cohort_check_(mimic|eicu)\.py)$ - repo: https://github.com/nbQA-dev/nbQA rev: 1.9.1 diff --git a/_typos.toml b/_typos.toml index 7aa8f26a..b83c2498 100644 --- a/_typos.toml +++ b/_typos.toml @@ -37,4 +37,10 @@ get_rolling_window_indicies = "get_rolling_window_indicies" # ours to rename and will keep showing up as new false positives every time # a fresh report is committed, so the whole directory is excluded rather # than growing an extend-words entry per column. -extend-exclude = ["scripts/gemini/out/"] +# scripts/cohort/cohort_check_*.py are committed verbatim as they ran +# (scripts/cohort/README.md); their `evn` (eICU visit number) stays. +extend-exclude = [ + "scripts/gemini/out/", + "scripts/cohort/cohort_check_mimic.py", + "scripts/cohort/cohort_check_eicu.py", +] diff --git a/odyssey/training/summary_targets.py b/odyssey/training/summary_targets.py index cd63e962..06bf16bb 100644 --- a/odyssey/training/summary_targets.py +++ b/odyssey/training/summary_targets.py @@ -47,6 +47,10 @@ ) +# Window statistics asked for per panel signal: the GBM's ``summary_stats`` +# group without ``last``/``hours_since_last`` (recency probes already showed +# the state holds those) and without ``delta_prev``/``ratio_visit_min`` +# (functions of the others). SIGNAL_STATS: tuple[str, ...] = ( "min_6h", "max_6h", @@ -55,14 +59,10 @@ "mean_24h", "delta_visit_first", ) -"""Window statistics asked for per panel signal (the GBM's ``summary_stats`` -group without ``last``/``hours_since_last``, which recency probes already -showed the state holds, and without ``delta_prev``/``ratio_visit_min``, -which are functions of the others).""" +# Occurrence counts asked for per drug class and per code family: the GBM's +# ``counts_occurrence`` group at its two window lengths. COUNT_STATS: tuple[str, ...] = ("n_6h", "n_24h") -"""Occurrence counts asked for per drug class and per code family (the -GBM's ``counts_occurrence`` group at its two window lengths).""" DEFAULT_LANDMARK_HOURS = 4.0 WINSOR_PERCENTILES = (0.5, 99.5) diff --git a/scripts/gemini/out/evals/gemini_full_14m_v1_eval_forecast.json b/scripts/gemini/out/evals/gemini_full_14m_v1_eval_forecast.json index a6adb1bc..6a6b2a3d 100644 --- a/scripts/gemini/out/evals/gemini_full_14m_v1_eval_forecast.json +++ b/scripts/gemini/out/evals/gemini_full_14m_v1_eval_forecast.json @@ -129,4 +129,4 @@ "n_positive_gaps": 3639349 }, "tail_slice": null -} +} \ No newline at end of file diff --git a/scripts/gemini/out/evals/gemini_full_DEC_v12_alerts.json b/scripts/gemini/out/evals/gemini_full_DEC_v12_alerts.json index d110a460..84211bd8 100644 --- a/scripts/gemini/out/evals/gemini_full_DEC_v12_alerts.json +++ b/scripts/gemini/out/evals/gemini_full_DEC_v12_alerts.json @@ -1913,4 +1913,4 @@ }, "landmark_protocol_version": 4 } -] +] \ No newline at end of file diff --git a/scripts/gemini/out/evals/gemini_full_DEC_v12_allshards_alerts.json b/scripts/gemini/out/evals/gemini_full_DEC_v12_allshards_alerts.json index be67255b..efcf285e 100644 --- a/scripts/gemini/out/evals/gemini_full_DEC_v12_allshards_alerts.json +++ b/scripts/gemini/out/evals/gemini_full_DEC_v12_allshards_alerts.json @@ -1913,4 +1913,4 @@ }, "landmark_protocol_version": 4 } -] +] \ No newline at end of file diff --git a/scripts/gemini/out/evals/gemini_full_DEC_v12_concept_atlas.json b/scripts/gemini/out/evals/gemini_full_DEC_v12_concept_atlas.json index 96da3caa..66eedf2b 100644 --- a/scripts/gemini/out/evals/gemini_full_DEC_v12_concept_atlas.json +++ b/scripts/gemini/out/evals/gemini_full_DEC_v12_concept_atlas.json @@ -13011,4 +13011,4 @@ "mean_activation": 0.4445817383871991 } ] -} +} \ No newline at end of file diff --git a/scripts/gemini/out/evals/gemini_full_DEC_v12_eval_forecast.json b/scripts/gemini/out/evals/gemini_full_DEC_v12_eval_forecast.json index 2a5636d7..f8b0699c 100644 --- a/scripts/gemini/out/evals/gemini_full_DEC_v12_eval_forecast.json +++ b/scripts/gemini/out/evals/gemini_full_DEC_v12_eval_forecast.json @@ -543,4 +543,4 @@ }, "value_metrics": null, "tail_slice": null -} +} \ No newline at end of file diff --git a/scripts/gemini/out/evals/gemini_full_DEC_v12_steering_full.json b/scripts/gemini/out/evals/gemini_full_DEC_v12_steering_full.json index 634f6db7..41e55272 100644 --- a/scripts/gemini/out/evals/gemini_full_DEC_v12_steering_full.json +++ b/scripts/gemini/out/evals/gemini_full_DEC_v12_steering_full.json @@ -14428,4 +14428,4 @@ "sign_agreement": 1.0 } ] -} +} \ No newline at end of file diff --git a/scripts/gemini/out/evals/gemini_full_v10_15c_alerts.json b/scripts/gemini/out/evals/gemini_full_v10_15c_alerts.json index 4c6dc770..7b04a88c 100644 --- a/scripts/gemini/out/evals/gemini_full_v10_15c_alerts.json +++ b/scripts/gemini/out/evals/gemini_full_v10_15c_alerts.json @@ -1913,4 +1913,4 @@ }, "landmark_protocol_version": 4 } -] +] \ No newline at end of file diff --git a/scripts/gemini/out/evals/gemini_full_v10_15c_alerts_cis.json b/scripts/gemini/out/evals/gemini_full_v10_15c_alerts_cis.json index 80b3f7df..e559497c 100644 --- a/scripts/gemini/out/evals/gemini_full_v10_15c_alerts_cis.json +++ b/scripts/gemini/out/evals/gemini_full_v10_15c_alerts_cis.json @@ -704,4 +704,4 @@ } } } -} +} \ No newline at end of file diff --git a/scripts/gemini/out/evals/gemini_full_v10_15c_counterfactual.json b/scripts/gemini/out/evals/gemini_full_v10_15c_counterfactual.json index 0d486be5..39b70901 100644 --- a/scripts/gemini/out/evals/gemini_full_v10_15c_counterfactual.json +++ b/scripts/gemini/out/evals/gemini_full_v10_15c_counterfactual.json @@ -219,4 +219,4 @@ } } } -} +} \ No newline at end of file diff --git a/scripts/gemini/out/evals/gemini_full_v10_15c_eval_forecast.json b/scripts/gemini/out/evals/gemini_full_v10_15c_eval_forecast.json index caf52962..2fb9ffb2 100644 --- a/scripts/gemini/out/evals/gemini_full_v10_15c_eval_forecast.json +++ b/scripts/gemini/out/evals/gemini_full_v10_15c_eval_forecast.json @@ -383,4 +383,4 @@ }, "value_metrics": null, "tail_slice": null -} +} \ No newline at end of file diff --git a/scripts/gemini/out/evals/gemini_full_v10_15c_intervention_cis.json b/scripts/gemini/out/evals/gemini_full_v10_15c_intervention_cis.json index 033ff473..7debd16e 100644 --- a/scripts/gemini/out/evals/gemini_full_v10_15c_intervention_cis.json +++ b/scripts/gemini/out/evals/gemini_full_v10_15c_intervention_cis.json @@ -111,4 +111,4 @@ "n_boot": 2000 } } -} +} \ No newline at end of file diff --git a/scripts/gemini/out/evals/gemini_smoke_2_eval_forecast.json b/scripts/gemini/out/evals/gemini_smoke_2_eval_forecast.json index d9622d0c..6eba1956 100644 --- a/scripts/gemini/out/evals/gemini_smoke_2_eval_forecast.json +++ b/scripts/gemini/out/evals/gemini_smoke_2_eval_forecast.json @@ -129,4 +129,4 @@ "n_positive_gaps": 3639349 }, "tail_slice": null -} +} \ No newline at end of file diff --git a/scripts/probe_summary_signal.py b/scripts/probe_summary_signal.py index 867848a1..76e06cc1 100644 --- a/scripts/probe_summary_signal.py +++ b/scripts/probe_summary_signal.py @@ -83,8 +83,8 @@ "delta_visit_first", "ratio_visit_min", ) -"""`last` is included as an anchor: the most recent value is the easiest -statistic to read off a sequence model, so its R^2 bounds the others.""" +# `last` is included as an anchor: the most recent value is the easiest +# statistic to read off a sequence model, so its R^2 bounds the others. MIN_TRAIN_ROWS = 500 MIN_HELD_ROWS = 200 From 3474f2c71f0d839807047f278aa7c4b8be88fa66 Mon Sep 17 00:00:00 2001 From: Amrit Krishnan Date: Mon, 14 Sep 2026 15:36:19 -0400 Subject: [PATCH 4/4] Weight the change-from-baseline and count targets inside the summary loss The first arm (weight 0.5, every target equal) tripled how readable the window levels are from the state but left the two targets the GBM's margin sits on almost untouched: creatinine change from admission went from R^2 0.00 to 0.23 after the bottleneck and stayed at ~0 before it, and the vasopressor 6 h count did not move. Levels are 241 of the 328 targets and won the average. `summary_change_weight` multiplies the delta_visit_first and count targets inside the masked Huber average (masked_huber_loss gains an optional per-target weight vector; ForecastObjective carries it). 1.0 reproduces the first arm exactly. Run A uses 4.0 with the overall summary weight raised to 2.0. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01StsjDmueuDqVNNuVoCEF1P --- odyssey/models/sequence_model.py | 23 +++++++++++++--- odyssey/models/summary_head.py | 12 ++++++--- odyssey/training/summary_targets.py | 18 +++++++++++++ odyssey/training/train.py | 18 +++++++++++++ tests/odyssey/models/test_summary_head.py | 32 ++++++++++++++++++++++- 5 files changed, 95 insertions(+), 8 deletions(-) diff --git a/odyssey/models/sequence_model.py b/odyssey/models/sequence_model.py index 83587ff4..87f648b0 100644 --- a/odyssey/models/sequence_model.py +++ b/odyssey/models/sequence_model.py @@ -128,6 +128,8 @@ class ForecastObjective: """Weight of the self-supervised window-summary loss (:mod:`odyssey.models.summary_head`), if the model has a summary head; the targets come per chunk from :mod:`odyssey.training.summary_targets`.""" + summary_target_weights: torch.Tensor | None = None + """Optional ``(K,)`` per-target weights inside the summary loss.""" class ForwardWithFeatures(NamedTuple): @@ -432,12 +434,16 @@ def _streaming_summary_loss( summary_head: SummaryHead | None, features: torch.Tensor, summary_targets: Optional["SummaryTargets"], + target_weights: torch.Tensor | None = None, ) -> torch.Tensor: """Masked Huber loss of the window-summary head (zero-graph if absent).""" if summary_head is None or summary_targets is None: return features.sum() * 0.0 return masked_huber_loss( - summary_head(features), summary_targets.values, summary_targets.mask + summary_head(features), + summary_targets.values, + summary_targets.mask, + target_weights=target_weights, ) def _streaming_time_loss( @@ -618,7 +624,10 @@ def compute_streaming_loss( event_loss = self._streaming_event_loss(self.event_heads, hidden, event_targets) value_loss, _ = self._streaming_value_loss(self.value_head, hidden, chunk) summary_loss = self._streaming_summary_loss( - self.summary_head, hidden, summary_targets + self.summary_head, + hidden, + summary_targets, + target_weights=objective.summary_target_weights, ) total = ( task_loss @@ -910,7 +919,10 @@ def compute_steering_loss( ) value_loss, _ = self._streaming_value_loss(self.value_head, head_feats, scored) summary_loss = self._streaming_summary_loss( - self.summary_head, head_feats, summary_targets + self.summary_head, + head_feats, + summary_targets, + target_weights=objective.summary_target_weights, ) forecast_loss = ( next_token_loss @@ -1090,7 +1102,10 @@ def compute_streaming_loss( ) value_loss, _ = self._streaming_value_loss(self.value_head, head_feats, chunk) summary_loss = self._streaming_summary_loss( - self.summary_head, head_feats, summary_targets + self.summary_head, + head_feats, + summary_targets, + target_weights=objective.summary_target_weights, ) forecast_loss = ( next_token_loss diff --git a/odyssey/models/summary_head.py b/odyssey/models/summary_head.py index 58d88ff7..bb327e0c 100644 --- a/odyssey/models/summary_head.py +++ b/odyssey/models/summary_head.py @@ -49,12 +49,16 @@ def masked_huber_loss( mask: torch.Tensor, *, delta: float = 1.0, + target_weights: torch.Tensor | None = None, ) -> torch.Tensor: """Huber loss averaged over the ``True`` entries of ``mask``. - ``prediction``/``target``/``mask`` share the shape ``(..., K)``. Returns - ``0 * prediction.sum()`` when the mask is empty, so the result always - carries a graph and can be added to other losses unconditionally. + ``prediction``/``target``/``mask`` share the shape ``(..., K)``. + ``target_weights`` (``(K,)``, optional) reweights the per-target terms + inside the average, so targets the state finds hardest (changes from + baseline, counts) can carry more of the gradient than the levels it + learns anyway. Returns ``0 * prediction.sum()`` when the mask is empty, + so the result always carries a graph and can be added unconditionally. """ if not bool(mask.any()): return prediction.sum() * 0.0 @@ -62,6 +66,8 @@ def masked_huber_loss( prediction.float(), target.float(), reduction="none", delta=delta ) weights = mask.to(per_entry.dtype) + if target_weights is not None: + weights = weights * target_weights.to(per_entry.dtype).to(weights.device) return (per_entry * weights).sum() / weights.sum() diff --git a/odyssey/training/summary_targets.py b/odyssey/training/summary_targets.py index 06bf16bb..ca026ac7 100644 --- a/odyssey/training/summary_targets.py +++ b/odyssey/training/summary_targets.py @@ -81,6 +81,23 @@ def summary_target_names() -> list[str]: return names +def summary_target_weights(change_weight: float = 1.0) -> torch.Tensor: + """Return ``(K,)`` per-target loss weights. + + ``change_weight`` on the ``delta_visit_first`` and count targets, 1 on + the window-level targets. + """ + names = summary_target_names() + hard = [ + n.rsplit(".", 1)[-1] == "delta_visit_first" + or n.rsplit(".", 1)[-1] in COUNT_STATS + for n in names + ] + return torch.where( + torch.tensor(hard), torch.tensor(float(change_weight)), torch.tensor(1.0) + ) + + def count_target_mask(names: Sequence[str] | None = None) -> np.ndarray: """Boolean mask over the panel: which targets are occurrence counts.""" names = list(names) if names is not None else summary_target_names() @@ -391,5 +408,6 @@ def load_summary_tables( "landmark_rows", "load_summary_tables", "summary_target_names", + "summary_target_weights", "summary_targets_for_chunk", ] diff --git a/odyssey/training/train.py b/odyssey/training/train.py index ed7cbbf8..1ee6a1db 100644 --- a/odyssey/training/train.py +++ b/odyssey/training/train.py @@ -117,6 +117,7 @@ SummaryTargetTables, load_summary_tables, summary_target_names, + summary_target_weights, summary_targets_for_chunk, ) from odyssey.utils.env_fingerprint import write_run_provenance @@ -485,6 +486,16 @@ class TrainingConfig: """Width of the summary head, recorded at training time so a checkpoint rebuilds the same head even if the target panel definition changes.""" + summary_change_weight: float = 1.0 + """Multiplier, inside the summary loss, on the change-from-baseline + (``delta_visit_first``) and occurrence-count targets relative to the + window-level targets. The first arm (weight 0.5, all targets equal) + tripled the readability of window levels but left changes and counts + almost where they were (creatinine change from admission R^2 0.00 -> + 0.23 post-bottleneck, ~0 pre; vasopressor 6 h count unchanged); levels + are 241 of the 328 targets and won the average. > 1 shifts the gradient + to the targets the state does not hold.""" + randint_prob: float = 0.25 """Intervention-aware training (CEM's RandInt): at every training position, each observed concept's mixing probability is replaced by @@ -1018,6 +1029,13 @@ def build_objective( if getattr(config, "summary_targets_dir", None) else 0.0 ), + summary_target_weights=( + summary_target_weights( + float(getattr(config, "summary_change_weight", 1.0)) + ).to(device) + if getattr(config, "summary_targets_dir", None) + else None + ), ) diff --git a/tests/odyssey/models/test_summary_head.py b/tests/odyssey/models/test_summary_head.py index d721279a..40a6ab56 100644 --- a/tests/odyssey/models/test_summary_head.py +++ b/tests/odyssey/models/test_summary_head.py @@ -11,7 +11,11 @@ ForecastObjective, ) from odyssey.models.summary_head import SummaryHead, masked_huber_loss -from odyssey.training.summary_targets import SummaryTargets +from odyssey.training.summary_targets import ( + SummaryTargets, + summary_target_names, + summary_target_weights, +) VOCAB = 40 @@ -129,3 +133,29 @@ def test_baseline_model_accepts_summary_targets_too() -> None: chunk, objective=ForecastObjective(summary_weight=1.0), summary_targets=targets ) assert "summary_loss" in components + + +def test_masked_huber_loss_target_weights_reweight_inside_the_average() -> None: + pred = torch.zeros(1, 1, 2) + target = torch.tensor([[[1.0, 0.5]]]) # huber 0.5 and 0.125 + mask = torch.ones(1, 1, 2, dtype=torch.bool) + plain = masked_huber_loss(pred, target, mask) + assert plain.item() == (0.5 + 0.125) / 2 + weighted = masked_huber_loss( + pred, target, mask, target_weights=torch.tensor([3.0, 1.0]) + ) + assert abs(weighted.item() - (3 * 0.5 + 0.125) / 4) < 1e-6 + # a weight of one everywhere is the plain average + same = masked_huber_loss(pred, target, mask, target_weights=torch.ones(2)) + assert same.item() == plain.item() + + +def test_summary_target_weights_mark_changes_and_counts() -> None: + names = summary_target_names() + w = summary_target_weights(4.0) + assert w.shape == (len(names),) + assert w[names.index("creatinine.delta_visit_first")].item() == 4.0 + assert w[names.index("drug.vasopressor.n_6h")].item() == 4.0 + assert w[names.index("family.lab.n_24h")].item() == 4.0 + assert w[names.index("creatinine.min_6h")].item() == 1.0 + assert summary_target_weights(1.0).eq(1.0).all()