Train IC objectives on complete dates without NaN gradients

CHANGELOG.md

@@ -2,6 +2,27 @@
 
 ## Unreleased
 
+### Correlation-loss gradients
+
+- The trainer automatically uses one complete observed date per batch for
+  `ICLoss`/`RankICLoss`, shuffling dates rather than stocks across dates. Daily
+  losses are averaged equally over dates. `batch_mode="date"` supports custom
+  daily objectives; ordinary regression retains flat sample batching.
+- Time validation now holds out complete latest dates, with no date shared
+  between training and validation. `val_split` is a date fraction in this mode.
+  Explicit random daily splits also preserve whole dates. Flat training no longer
+  drops a dataset smaller than `batch_size`. These changes can alter historical
+  model diagnostics; no archived performance result is revalidated here.
+
+- `ICLoss` and `RankICLoss` return zero loss and zero gradients for constant
+  prediction/label batches and singleton batches. Previously the reported loss
+  could be zero while backward produced NaNs from `sqrt(0)`, poisoning optimizer
+  state. Undefined batches contribute no learning signal; nonconstant batches
+  retain the previous formula and gradients. Inputs must be aligned and finite.
+- Empty/broadcast shapes are rejected. Sign-aware `gamma` and soft-rank
+  `temperature` must be finite and positive. This is training correctness,
+  not predictive-quality or investment evidence.
+
 ### Full-pipeline validation: execution coverage
 
 - Reports now show joint factor-mask coverage, per-split training/prediction

docs/backtest_assumptions.md

@@ -50,13 +50,26 @@ A publishable result should document the train, validation, and test split, plus
 
 The CLI pipeline splits by date. `train_frac` (default `0.6`) sets the first
 out-of-sample date index; the model trains only on earlier dates, its
-validation set is the most recent slice of those training samples
+validation set is the most recent slice of complete training dates
 (`val_split_mode: time`), and the backtest scores only the out-of-sample window.
 `summary.json` records both windows under `evaluation`. The one-day label of the
 last training date uses the close of the first out-of-sample date, which is
 already known when the first out-of-sample signal is formed, so no purge gap is
 needed for the default one-day horizon; longer horizons need one.
 
+`Trainer` holds out whole latest observed dates for `val_split_mode="time"`;
+`val_split` therefore denotes a fraction of dates, not individual stock-days.
+The default `batch_mode="auto"` uses full daily cross-sections for `ICLoss` and
+`RankICLoss`, with one optimizer step per date and equally weighted daily losses.
+`batch_size` is ignored for daily batching: memory depends on the largest daily
+cross-section, and soft RankIC still requires quadratic pairwise rank memory.
+Use `batch_mode="date"` for a custom daily objective. Ordinary regression uses
+flat sample batches and sample-weighted epoch loss; it retains the final partial
+batch. Explicit `val_split_mode="random"` splits whole dates in daily mode and
+individual samples in flat mode. Random splits are reproducibility options,
+not chronological out-of-sample evidence. Constant/singleton IC batches contribute
+zero loss and zero gradients; they are retained, not counted as informative IC.
+
 This is a single split, not a walk-forward evaluation. Demo metrics come from
 synthetic data and are expected to hover around zero or below after costs.
 `scripts/public_data_validation.py` runs rolling walk-forward splits.

src/mlquant/models/losses.py

@@ -20,6 +20,8 @@
 """
 from __future__ import annotations
 
+import math
+
 import torch
 import torch.nn as nn
 
@@ -40,8 +42,8 @@ class AdjMSELoss(nn.Module):
 
     def __init__(self, gamma: float = 0.1) -> None:
         super().__init__()
-        if gamma <= 0:
-            raise ValueError("gamma must be positive")
+        if not math.isfinite(gamma) or gamma <= 0:
+            raise ValueError("gamma must be finite and positive")
         self.gamma = gamma
 
     def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
@@ -54,14 +56,27 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
 
 
 def _pearson(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    if x.shape != y.shape or not x.numel():
+        raise ValueError("correlation inputs must have the same nonempty shape")
     xm = x - x.mean()
     ym = y - y.mean()
-    denom = (xm.pow(2).sum().sqrt() * ym.pow(2).sum().sqrt()).clamp_min(1e-12)
-    return (xm * ym).sum() / denom
+    xss, yss = xm.pow(2).sum(), ym.pow(2).sum()
+    defined = (xss > 0) & (yss > 0)
+    # sqrt'(0) is infinite: clamping the product after sqrt cannot protect
+    # backward. Replace zero-variance operands before taking either root.
+    sx = torch.where(defined, xss, torch.ones_like(xss)).sqrt()
+    sy = torch.where(defined, yss, torch.ones_like(yss)).sqrt()
+    denom = (sx * sy).clamp_min(1e-12)
+    numerator = (xm * ym).sum()
+    return torch.where(defined, numerator / denom, numerator * 0.0)
 
 
 class ICLoss(nn.Module):
-    """Negative Pearson IC across a batch (e.g. one cross-section)."""
+    """Negative Pearson IC across a batch (e.g. one cross-section).
+
+    A constant prediction or label batch has undefined correlation and returns
+    zero loss/gradients. Pass aligned, finite tensors from one cross-section.
+    """
 
     def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         return -_pearson(pred, target)
@@ -77,6 +92,8 @@ class RankICLoss(nn.Module):
 
     def __init__(self, temperature: float = 1.0) -> None:
         super().__init__()
+        if not math.isfinite(temperature) or temperature <= 0:
+            raise ValueError("temperature must be finite and positive")
         self.temperature = temperature
 
     def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:

src/mlquant/training/__init__.py

@@ -2,5 +2,6 @@
 from .dataset import FactorDataset
 from .augment import gbm_augment
 from .trainer import Trainer, TrainConfig
+from .samplers import CrossSectionBatchSampler
 
-__all__ = ["FactorDataset", "gbm_augment", "Trainer", "TrainConfig"]
+__all__ = ["FactorDataset", "gbm_augment", "Trainer", "TrainConfig", "CrossSectionBatchSampler"]

src/mlquant/training/samplers.py

@@ -0,0 +1,36 @@
+"""Batch complete daily cross-sections without mixing their correlation targets."""
+from __future__ import annotations
+
+import torch
+from torch.utils.data import Sampler
+
+
+def date_boundaries(dates: torch.Tensor) -> list[int]:
+    """Start/end offsets for finite, date-ordered samples, including the sentinel."""
+    if (dates.ndim != 1 or not torch.isfinite(dates).all()
+            or (dates[1:] < dates[:-1]).any()):
+        raise ValueError("sample dates must be a finite nondecreasing vector")
+    if not dates.numel():
+        return [0]
+    starts = (dates[1:] != dates[:-1]).nonzero().flatten().add(1).tolist()
+    return [0, *starts, dates.numel()]
+
+
+class CrossSectionBatchSampler(Sampler[list[int]]):
+    """One full observed date per batch; optionally shuffle the order of dates.
+
+    Missing stocks stay missing. No padding, cross-date mixing, within-day
+    splitting or dropped small batches. Memory for boundaries grows with dates.
+    """
+
+    def __init__(self, dates: torch.Tensor, *, shuffle: bool = False) -> None:
+        self.boundaries = date_boundaries(dates)
+        self.shuffle = shuffle
+
+    def __len__(self) -> int:
+        return len(self.boundaries) - 1
+
+    def __iter__(self):
+        order = torch.randperm(len(self)).tolist() if self.shuffle else range(len(self))
+        for day in order:
+            yield list(range(self.boundaries[day], self.boundaries[day + 1]))

src/mlquant/training/trainer.py

@@ -8,14 +8,17 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
+import math
 from pathlib import Path
 from typing import Callable, Iterable
 
 import torch
 import torch.nn as nn
 from torch.utils.data import DataLoader, Subset
 
+from ..models.losses import ICLoss, RankICLoss
 from .dataset import FactorDataset
+from .samplers import CrossSectionBatchSampler, date_boundaries
 
 
 @dataclass
@@ -29,11 +32,14 @@ class TrainConfig:
     log_every:   int   = 50
     save_dir:    str   = "checkpoints"
     val_split:   float = 0.1
-    # "time" holds out the most recent samples (FactorDataset is ordered by
-    # date), so validation never sees days that precede training days.
-    # "random" restores the old shuffled split for backward compatibility.
+    # "time" holds out the most recent complete dates. "random" splits dates
+    # for daily batches, or individual samples for legacy regression batches.
     val_split_mode: str = "time"
     extra:       dict  = field(default_factory=dict)
+    # "auto" uses complete daily batches for IC/RankIC, flat batches otherwise.
+    # "date" also supports custom cross-sectional objectives. In date mode,
+    # batch_size is ignored and losses are averaged equally over observed dates.
+    batch_mode: str = "auto"
 
 
 class Trainer:
@@ -56,23 +62,53 @@ def __init__(
     # Public API
     # ------------------------------------------------------------------
     def fit(self, dataset: FactorDataset) -> dict:
-        n_val = max(1, int(len(dataset) * self.cfg.val_split))
-        n_train = len(dataset) - n_val
-        if self.cfg.val_split_mode == "time":
-            train_ds = Subset(dataset, range(n_train))
-            val_ds = Subset(dataset, range(n_train, len(dataset)))
-        elif self.cfg.val_split_mode == "random":
+        if not math.isfinite(self.cfg.val_split) or not 0 < self.cfg.val_split < 1:
+            raise ValueError("val_split must be finite and between zero and one")
+        if self.cfg.epochs < 1 or self.cfg.batch_size < 1:
+            raise ValueError("epochs and batch_size must be positive")
+        if self.cfg.val_split_mode not in ("time", "random"):
+            raise ValueError("val_split_mode must be 'time' or 'random'")
+        if self.cfg.batch_mode not in ("auto", "samples", "date"):
+            raise ValueError("batch_mode must be 'auto', 'samples' or 'date'")
+        correlation = isinstance(self.loss_fn, (ICLoss, RankICLoss))
+        if correlation and self.cfg.batch_mode == "samples":
+            raise ValueError("IC/RankIC require complete date batches")
+        daily = self.cfg.batch_mode == "date" or (self.cfg.batch_mode == "auto" and correlation)
+        if len(dataset) < 2 or dataset.dates.shape != (len(dataset),):
+            raise ValueError("training requires at least two samples with aligned dates")
+        if self.cfg.val_split_mode == "time" or daily:
+            boundaries = date_boundaries(dataset.dates)
+            n_dates = len(boundaries) - 1
+            if n_dates < 2:
+                raise ValueError("date splitting requires at least two observed dates")
+            n_val_dates = max(1, int(n_dates * self.cfg.val_split))
+            if self.cfg.val_split_mode == "time":
+                split = boundaries[n_dates - n_val_dates]
+                train_ds = Subset(dataset, range(split))
+                val_ds = Subset(dataset, range(split, len(dataset)))
+            else:
+                validation_days = set(torch.randperm(
+                    n_dates, generator=torch.Generator().manual_seed(0)
+                )[:n_val_dates].tolist())
+                train_indices, val_indices = [], []
+                for day in range(n_dates):
+                    target = val_indices if day in validation_days else train_indices
+                    target.extend(range(boundaries[day], boundaries[day + 1]))
+                train_ds, val_ds = Subset(dataset, train_indices), Subset(dataset, val_indices)
+        else:
+            n_val = max(1, int(len(dataset) * self.cfg.val_split))
             train_ds, val_ds = torch.utils.data.random_split(
-                dataset, [n_train, n_val],
+                dataset, [len(dataset) - n_val, n_val],
                 generator=torch.Generator().manual_seed(0),
             )
+        if daily:
+            train_loader = DataLoader(train_ds, batch_sampler=CrossSectionBatchSampler(
+                dataset.dates[train_ds.indices], shuffle=True))
+            val_loader = DataLoader(val_ds, batch_sampler=CrossSectionBatchSampler(
+                dataset.dates[val_ds.indices]))
         else:
-            raise ValueError(
-                f"val_split_mode must be 'time' or 'random', got {self.cfg.val_split_mode!r}"
-            )
-        train_loader = DataLoader(train_ds, batch_size=self.cfg.batch_size,
-                                  shuffle=True, drop_last=True)
-        val_loader   = DataLoader(val_ds,   batch_size=self.cfg.batch_size)
+            train_loader = DataLoader(train_ds, batch_size=self.cfg.batch_size, shuffle=True)
+            val_loader = DataLoader(val_ds, batch_size=self.cfg.batch_size)
 
         history: dict = {"train_loss": [], "val_loss": []}
         save_dir = Path(self.cfg.save_dir)
@@ -103,18 +139,21 @@ def predict(self, features: torch.Tensor) -> torch.Tensor:
     def _run_epoch(self, loader: Iterable, *, train: bool) -> float:
         self.model.train(train)
         total, n = 0.0, 0
+        daily = isinstance(getattr(loader, "batch_sampler", None), CrossSectionBatchSampler)
         for x, y in loader:
             x = x.to(self.cfg.device)
             y = y.to(self.cfg.device)
             if train:
                 self.optim.zero_grad(set_to_none=True)
-            preds = self.model(x)
-            loss = self.loss_fn(preds, y)
+            with torch.set_grad_enabled(train):
+                preds = self.model(x)
+                loss = self.loss_fn(preds, y)
             if train:
                 loss.backward()
                 if self.cfg.grad_clip:
                     nn.utils.clip_grad_norm_(self.model.parameters(), self.cfg.grad_clip)
                 self.optim.step()
-            total += loss.item() * x.shape[0]
-            n     += x.shape[0]
+            weight = 1 if daily else x.shape[0]
+            total += loss.item() * weight
+            n += weight
         return total / max(1, n)