Bound RankIC activations with checkpointed column blocks

CHANGELOG.md

@@ -2,6 +2,18 @@
 
 ## Unreleased
 
+### Chunked soft RankIC
+
+- `RankICLoss` computes soft-rank columns in blocks above `chunk_size` (default
+  256) and uses non-reentrant activation checkpointing in training. Pairwise
+  temporary matrices grow as N × chunk_size rather than N², and backward
+  recomputes blocks instead of retaining every sigmoid activation. The pairwise
+  arithmetic remains quadratic and recomputation can cost more runtime.
+- `chunk_size=None` retains the dense reference. The same soft-rank objective
+  supports both-input and higher-order gradients; reduction rounding can differ
+  across blocks/devices. Inputs must be aligned one-dimensional cross-sections.
+  This is memory scaling support, not predictive-quality or speedup evidence.
+
 ### Purged date validation
 
 - `TrainConfig.purge_dates` removes complete observed signal dates immediately

docs/backtest_assumptions.md

@@ -76,7 +76,12 @@ sample/date counts and numeric date-ID ranges so the chosen split can be inspect
 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.
+cross-section. Soft RankIC uses `chunk_size=256` by default above that many
+stocks, with N × chunk_size pairwise matrices and activation checkpointing during
+training. `RankICLoss(chunk_size=None)` retains the dense reference. Pairwise work
+remains quadratic and backward recomputation adds runtime; reduction rounding
+can differ across blocks/devices. This does not bound model activations or total
+training memory, and no GPU performance estimate is provided here.
 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

src/mlquant/models/losses.py

@@ -24,6 +24,7 @@
 
 import torch
 import torch.nn as nn
+from torch.utils.checkpoint import checkpoint
 
 
 class AdjMSELoss(nn.Module):
@@ -88,18 +89,43 @@ class RankICLoss(nn.Module):
     Spearman = Pearson on ranks. We use the *soft-rank* trick (a
     differentiable approximation that sorts the inputs through a
     softmax-temperature lens) so the loss is gradient-friendly.
+
+    Above ``chunk_size`` stocks, compute columns in blocks and checkpoint their
+    activations. Pairwise work remains quadratic; temporary matrices scale as
+    N * chunk_size. ``chunk_size=None`` restores the dense reference path.
     """
 
-    def __init__(self, temperature: float = 1.0) -> None:
+    def __init__(self, temperature: float = 1.0, *, chunk_size: int | None = 256) -> None:
         super().__init__()
         if not math.isfinite(temperature) or temperature <= 0:
             raise ValueError("temperature must be finite and positive")
+        if chunk_size is not None and (type(chunk_size) is not int or chunk_size < 1):
+            raise ValueError("chunk_size must be a positive integer or None")
         self.temperature = temperature
+        self.chunk_size = chunk_size
 
     def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        if pred.shape != target.shape or not pred.numel():
+            raise ValueError("correlation inputs must have the same nonempty shape")
+        if pred.ndim != 1:
+            raise ValueError("RankIC inputs must be one-dimensional cross-sections")
         return -_pearson(self._soft_rank(pred), self._soft_rank(target))
 
+    @staticmethod
+    def _rank_columns(full, columns, temperature):
+        return torch.sigmoid((columns.unsqueeze(0) - full.unsqueeze(1)) / temperature).sum(0)
+
     def _soft_rank(self, x: torch.Tensor) -> torch.Tensor:
-        # Differentiable rank: softmax over pairwise differences.
-        diffs = (x.unsqueeze(0) - x.unsqueeze(1)) / self.temperature
-        return torch.sigmoid(diffs).sum(dim=0)
+        if self.chunk_size is None or x.numel() <= self.chunk_size:
+            return self._rank_columns(x, x, self.temperature)
+        ranks = []
+        for columns in x.split(self.chunk_size):
+            if torch.is_grad_enabled() and x.requires_grad:
+                # Non-reentrant checkpointing preserves autograd.grad and higher
+                # derivatives while recomputing one block during backward.
+                rank = checkpoint(self._rank_columns, x, columns, self.temperature,
+                                  use_reentrant=False, preserve_rng_state=False)
+            else:
+                rank = self._rank_columns(x, columns, self.temperature)
+            ranks.append(rank)
+        return torch.cat(ranks)

tests/test_chunked_rankic.py

@@ -0,0 +1,89 @@
+"""Chunked soft rank keeps its objective/derivatives with bounded activations."""
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from mlquant.models.losses import RankICLoss
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
+@pytest.mark.parametrize("temperature", [.05, 1., 7.])
+@pytest.mark.parametrize("chunk", [1, 4, 9])
+def test_chunked_loss_and_both_gradients_match_dense_with_ties_and_strides(dtype, temperature, chunk):
+    values = torch.tensor([.1, 9., -.7, 9., .1, 9., .9, 9., 1.2, 9., -.4, 9., .5, 9., -.2, 9., 2.1, 9.], dtype=dtype)
+    pred = values[::2].requires_grad_()
+    target = torch.tensor([.6, -.1, .6, .3, -.7, .2, .9, -.4, .8], dtype=dtype, requires_grad=True)
+    assert not pred.is_contiguous()
+    dense = RankICLoss(temperature, chunk_size=None)(pred, target)
+    actual = RankICLoss(temperature, chunk_size=chunk)(pred, target)
+    expected_grads = torch.autograd.grad(dense, (pred, target), retain_graph=True)
+    actual_grads = torch.autograd.grad(actual, (pred, target))
+    tolerance = 2e-6 if dtype == torch.float32 else 1e-12
+    torch.testing.assert_close(actual, dense, atol=tolerance, rtol=tolerance)
+    for actual_grad, expected_grad in zip(actual_grads, expected_grads):
+        torch.testing.assert_close(actual_grad, expected_grad, atol=tolerance, rtol=tolerance)
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
+@pytest.mark.parametrize("constant", ["prediction", "target", "both"])
+def test_constant_chunked_batches_keep_zero_gradients(dtype, constant):
+    pred = torch.ones(17, dtype=dtype) if constant != "target" else torch.arange(17, dtype=dtype)
+    target = torch.ones(17, dtype=dtype) if constant != "prediction" else torch.arange(17, dtype=dtype)
+    pred.requires_grad_()
+    target.requires_grad_()
+    loss = RankICLoss(chunk_size=4)(pred, target)
+    loss.backward()
+    assert loss.item() == 0
+    assert pred.grad.eq(0).all() and target.grad.eq(0).all()
+
+
+def test_checkpointed_rankic_supports_first_and_second_derivatives():
+    pred = torch.tensor([.4, -.2, .7, 1.1], dtype=torch.float64, requires_grad=True)
+    target = torch.tensor([-.1, .3, .6, .8], dtype=torch.float64, requires_grad=True)
+    loss = RankICLoss(.8, chunk_size=2)
+    assert torch.autograd.gradcheck(loss, (pred, target))
+    assert torch.autograd.gradgradcheck(loss, (pred, target))
+
+
+def test_large_cross_section_never_builds_a_dense_pairwise_matrix_or_saves_quadratic_activations():
+    n, chunk = 513, 32
+    pred = torch.linspace(-2, 2, n, requires_grad=True)
+    target = torch.sin(torch.arange(n).float())
+    saved, shapes = {}, []
+    real_sigmoid = torch.sigmoid
+
+    def keep(tensor):
+        storage = tensor.untyped_storage()
+        saved[storage.data_ptr()] = storage.nbytes()
+        return tensor
+
+    def observe(values):
+        shapes.append(tuple(values.shape))
+        return real_sigmoid(values)
+
+    with patch("torch.sigmoid", side_effect=observe):
+        with torch.autograd.graph.saved_tensors_hooks(keep, lambda tensor: tensor):
+            loss = RankICLoss(chunk_size=chunk)(pred, target)
+        retained_bytes = sum(saved.values())
+        loss.backward()
+    assert all(rows == n and columns <= chunk for rows, columns in shapes)
+    assert len(shapes) > 2 * ((n + chunk - 1) // chunk)  # Backward recomputes blocks.
+    assert retained_bytes < 64 * n * pred.element_size()
+    assert pred.grad.isfinite().all()
+    saved.clear()
+    with torch.autograd.graph.saved_tensors_hooks(keep, lambda tensor: tensor):
+        dense = RankICLoss(chunk_size=None)(pred, target)
+    assert sum(saved.values()) >= n * n * pred.element_size()
+    torch.testing.assert_close(loss, dense, atol=2e-6, rtol=2e-6)
+
+
+@pytest.mark.parametrize("bad", [0, -1, True, 1.5])
+def test_invalid_chunk_size_is_rejected(bad):
+    with pytest.raises(ValueError, match="positive integer or None"):
+        RankICLoss(chunk_size=bad)
+
+
+def test_rankic_rejects_multidimensional_cross_sections():
+    with pytest.raises(ValueError, match="one-dimensional"):
+        RankICLoss()(torch.ones(2, 3), torch.ones(2, 3))