Add a self-supervised window-summary head and the window-statistic probes

.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|scripts/cohort/cohort_check_(mimic|eicu)\.py'
+      # 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.50.2

_typos.toml

@@ -43,4 +43,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_mimic.py", "scripts/cohort/cohort_check_eicu.py"]
+# 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",
+]

odyssey/inference/run_inference.py

@@ -410,6 +410,18 @@
     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"),

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,12 @@ 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`."""
+    summary_target_weights: torch.Tensor | None = None
+    """Optional ``(K,)`` per-target weights inside the summary loss."""
 
 
 class ForwardWithFeatures(NamedTuple):
@@ -421,6 +429,23 @@ def _streaming_event_loss(
             event_heads.edges,
         )
 
+    def _streaming_summary_loss(
+        self,
+        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,
+            target_weights=target_weights,
+        )
+
     def _streaming_time_loss(
         self,
         time_head: TimeToEventHead | None,
@@ -486,6 +511,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 +551,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 +612,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 +623,18 @@ 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,
+            target_weights=objective.summary_target_weights,
+        )
         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 +643,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 +673,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 +743,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 +863,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 +918,18 @@ 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,
+            target_weights=objective.summary_target_weights,
+        )
         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 +951,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 +1033,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 +1101,18 @@ 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,
+            target_weights=objective.summary_target_weights,
+        )
         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 +1123,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 +1193,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

odyssey/models/summary_head.py

@@ -0,0 +1,74 @@
+"""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,
+    target_weights: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Huber loss averaged over the ``True`` entries of ``mask``.
+
+    ``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
+    per_entry = F.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()
+
+
+__all__ = ["SummaryHead", "masked_huber_loss"]