fix: zero turnover limit, CAE/SAE minibatch edge cases, attention with no neighbors or long lag

scripts/ci/hardware_qualification.py

@@ -11,6 +11,7 @@
 import tempfile
 import time
 from collections.abc import Callable
+from dataclasses import replace
 from pathlib import Path
 from typing import Any
 
@@ -299,11 +300,14 @@ def predict(value: LSTMPortfolioModel) -> np.ndarray:
 
 
 def _deep(batch: PortfolioSequenceBatch, device: str) -> FitOnce:
+    adjacency = np.zeros((batch.features.shape[2], batch.features.shape[2]), dtype=bool)
+    adjacency[0, :] = True
+    batch = replace(batch, adjacency_mask=adjacency)
     config = DeepPortfolioConfig(
         d_model=4,
         n_heads=1,
         cross_attention_heads=1,
-        macro_gnn_heads=1,
+        macro_gnn_heads=2,
         max_iters=1,
         eval_every=1,
         checkpoint_every=1,

src/ml4t/models/asset_prediction/sae.py

@@ -327,14 +327,27 @@ def _iterate_batches(
     seed: int,
 ):
     n_obs = int(features.shape[0])
+    if n_obs < 2:
+        raise ValueError(
+            "SAE training requires at least 2 valid observations; "
+            f"got {n_obs}. BatchNorm cannot update on a single row."
+        )
     if batch_size is None or batch_size <= 0 or batch_size >= n_obs:
         yield features, targets
         return
+    # BatchNorm1d rejects a one-row training batch. Keep the configured size
+    # when it is already safe, and fold a trailing singleton into the previous
+    # batch so every observation is still used.
+    step = max(int(batch_size), 2)
     generator = torch.Generator(device=features.device)
     generator.manual_seed(seed)
     order = torch.randperm(n_obs, generator=generator, device=features.device)
-    for start in range(0, n_obs, batch_size):
-        idx = order[start : start + batch_size]
+    starts = list(range(0, n_obs, step))
+    if len(starts) > 1 and n_obs - starts[-1] == 1:
+        starts.pop()
+    for index, start in enumerate(starts):
+        end = starts[index + 1] if index + 1 < len(starts) else n_obs
+        idx = order[start:end]
         yield features[idx], targets[idx]
 
 

src/ml4t/models/latent_factors/cae.py

@@ -146,6 +146,7 @@ def fit(
         loss_sums: dict[int, float] = dict.fromkeys(checkpoint_epochs, 0.0)
         val_best_losses: list[float] = []
         skipped_updates = 0
+        applied_updates = 0
 
         for ensemble_idx in range(self.config.n_ensemble):
             seed = self.config.seed + ensemble_idx
@@ -192,6 +193,7 @@ def fit(
 
                     epoch_loss += float(loss.item())
                     n_batches += 1
+                    applied_updates += 1
 
                 mean_loss = epoch_loss / max(n_batches, 1)
                 if epoch in checkpoint_epochs:
@@ -218,6 +220,13 @@ def fit(
                 self._checkpoint_states[0].append(best_state)
                 val_best_losses.append(best_val_loss)
 
+        if self.config.hidden_units and applied_updates == 0:
+            raise ValueError(
+                "CAE training updated no parameters because every minibatch had one row. "
+                "Beta BatchNorm cannot train on a single observation; provide at least two "
+                "valid rows and set batch_size to at least 2."
+            )
+
         history: list[dict[str, float | str]] = []
         for epoch in checkpoint_epochs:
             if epoch in self._checkpoint_states:

src/ml4t/models/portfolio/components.py

@@ -179,18 +179,24 @@ def __init__(self, *, d_model: int, n_heads: int, dropout: float, lag: int) -> N
     def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
         batch_size, n_periods, n_assets, d_model = x.shape
         if self.lag > 0:
-            pad = torch.zeros(
-                (batch_size, self.lag, n_assets, d_model),
-                device=x.device,
-                dtype=x.dtype,
-            )
-            kv = torch.cat([pad, x[:, : n_periods - self.lag]], dim=1)
-            pad_mask = torch.zeros(
-                (batch_size, self.lag, n_assets),
-                device=mask.device,
-                dtype=mask.dtype,
-            )
-            kv_mask = torch.cat([pad_mask, mask[:, : n_periods - self.lag]], dim=1)
+            history = n_periods - self.lag
+            if history <= 0:
+                # A negative slice would wrap to the end of the window.
+                kv = torch.zeros_like(x)
+                kv_mask = torch.zeros_like(mask)
+            else:
+                pad = torch.zeros(
+                    (batch_size, self.lag, n_assets, d_model),
+                    device=x.device,
+                    dtype=x.dtype,
+                )
+                kv = torch.cat([pad, x[:, :history]], dim=1)
+                pad_mask = torch.zeros(
+                    (batch_size, self.lag, n_assets),
+                    device=mask.device,
+                    dtype=mask.dtype,
+                )
+                kv_mask = torch.cat([pad_mask, mask[:, :history]], dim=1)
         else:
             kv = x
             kv_mask = mask
@@ -239,16 +245,21 @@ def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
         query = x.reshape(batch_size * n_periods, n_assets, d_model)
         key_padding_mask = mask.reshape(batch_size * n_periods, n_assets) < 0.5
         adjacency_mask = self.get_buffer("_adjacency_mask_buffer")
-        combined = adjacency_mask.unsqueeze(0) | key_padding_mask.unsqueeze(1)
-        all_blocked = combined.all(dim=-1)
-        if all_blocked.any():
-            key_padding_mask = key_padding_mask & ~all_blocked
-        attended, _ = self.attention(
-            query,
-            query,
-            query,
-            attn_mask=adjacency_mask,
-            key_padding_mask=key_padding_mask,
-        )
+        # True blocks attention. A query with no legal key makes softmax NaN,
+        # so fall back to self when that asset is present and otherwise allow
+        # a finite dummy attention that the caller mask zeros out.
+        blocked = adjacency_mask.unsqueeze(0) | key_padding_mask.unsqueeze(1)
+        fully_blocked = blocked.all(dim=-1)
+        if bool(fully_blocked.any()):
+            eye = torch.eye(n_assets, device=blocked.device, dtype=torch.bool)
+            allow_self = (
+                fully_blocked.unsqueeze(-1) & eye.unsqueeze(0) & (~key_padding_mask).unsqueeze(1)
+            )
+            blocked = blocked & ~allow_self
+            still_blocked = blocked.all(dim=-1)
+            if bool(still_blocked.any()):
+                blocked = blocked & ~still_blocked.unsqueeze(-1)
+        blocked = blocked.repeat_interleave(self.attention.num_heads, dim=0)
+        attended, _ = self.attention(query, query, query, attn_mask=blocked)
         attended = attended.reshape(batch_size, n_periods, n_assets, d_model)
         return self.layer_norm(x + self.dropout(attended))

src/ml4t/models/portfolio/postprocessors.py

@@ -217,10 +217,11 @@ def apply_turnover_limit(
 ) -> np.ndarray:
     """Scale cross-sectional weight changes to satisfy an L1 turnover cap."""
 
+    if not np.isfinite(turnover_limit) or turnover_limit < 0.0:
+        raise ValueError(f"turnover_limit must be finite and non-negative; got {turnover_limit}")
+
     constrained = np.asarray(weights, dtype=np.float64).copy()
     batch_size, n_periods, _ = constrained.shape
-    if turnover_limit <= 0:
-        return np.zeros_like(constrained)
 
     previous = (
         np.asarray(previous_weights, dtype=np.float64).copy()

tests/test_cae.py

@@ -189,3 +189,33 @@ def test_cae_pipeline_forecasts_returns_through_checkpointed_two_step_path() ->
     assert prediction.factor_forecast.factor_premia.shape == (3, 1)
     assert prediction.asset_forecast.expected_returns.shape == (3, n_assets)
     assert np.isfinite(prediction.asset_forecast.expected_returns).all()
+
+
+@pytest.mark.parametrize(
+    "batch",
+    [
+        CrossSectionBatch(
+            characteristics=np.ones((1, 1, 2), dtype=np.float64),
+            returns=np.ones((1, 1), dtype=np.float64),
+        ),
+        CrossSectionBatch(
+            characteristics=np.ones((2, 3, 2), dtype=np.float64),
+            returns=np.ones((2, 3), dtype=np.float64),
+        ),
+    ],
+)
+def test_cae_rejects_fits_that_never_update_parameters(batch: CrossSectionBatch) -> None:
+    model = CAEModel(
+        CAEConfig(
+            n_factors=1,
+            hidden_units=(4,),
+            n_epochs=2,
+            checkpoint_interval=2,
+            batch_size=1,
+        )
+    )
+
+    with pytest.raises(ValueError, match="updated no parameters"):
+        model.fit(batch)
+
+    assert not model.is_fitted

tests/test_financial_correctness.py

@@ -16,8 +16,12 @@
 from ml4t.models.latent_factors.ipca import _normalize_theta_y
 from ml4t.models.latent_factors.rp_pca import _risk_premium_matrix
 from ml4t.models.portfolio import runtime as portfolio_runtime
+from ml4t.models.portfolio.components import CrossSectionalAttention, MacroGraphAttention
 from ml4t.models.portfolio.losses import compute_net_portfolio_returns, softmin_sharpe
-from ml4t.models.portfolio.postprocessors import normalize_cross_sectional_weights
+from ml4t.models.portfolio.postprocessors import (
+    apply_turnover_limit,
+    normalize_cross_sectional_weights,
+)
 
 
 def test_rp_pca_default_matches_lettau_pelger_signal_matrix() -> None:
@@ -326,6 +330,91 @@ def fake_validation(
     )
 
 
+def test_cross_attention_lag_longer_than_the_window_does_not_wrap() -> None:
+    torch.manual_seed(0)
+    window = torch.randn(1, 3, 2, 4)
+    mask = torch.ones(1, 3, 2)
+    matched = CrossSectionalAttention(d_model=4, n_heads=1, dropout=0.0, lag=3)
+    longer = CrossSectionalAttention(d_model=4, n_heads=1, dropout=0.0, lag=8)
+    longer.load_state_dict(matched.state_dict())
+    matched.eval()
+    longer.eval()
+
+    matched_output = matched(window, mask)
+    longer_output = longer(window, mask)
+
+    assert matched_output.shape == window.shape
+    assert torch.isfinite(longer_output).all()
+    torch.testing.assert_close(longer_output, matched_output)
+
+
+def test_graph_attention_stays_finite_when_a_node_has_no_neighbors() -> None:
+    torch.manual_seed(0)
+    layer = MacroGraphAttention(
+        d_model=4,
+        n_heads=1,
+        dropout=0.0,
+        adjacency_mask=torch.ones(2, 2, dtype=torch.bool),
+    )
+    layer.eval()
+    features = torch.randn(1, 2, 2, 4, requires_grad=True)
+    mask = torch.ones(1, 2, 2)
+
+    output = layer(features, mask)
+    output.sum().backward()
+
+    assert torch.isfinite(output).all()
+    assert features.grad is not None
+    assert torch.isfinite(features.grad).all()
+
+
+def test_graph_attention_ignores_a_masked_only_neighbor() -> None:
+    layer = MacroGraphAttention(
+        d_model=4,
+        n_heads=1,
+        dropout=0.0,
+        adjacency_mask=torch.tensor([[True, False], [True, True]]),
+    )
+    layer.eval()
+    base = torch.randn(1, 1, 2, 4)
+    shifted = base.clone()
+    shifted[..., 1, :] += 10.0
+    mask = torch.tensor([[[1.0, 0.0]]])
+
+    base_output = layer(base, mask)
+    shifted_output = layer(shifted, mask)
+
+    assert torch.isfinite(base_output).all()
+    torch.testing.assert_close(base_output[:, :, 0, :], shifted_output[:, :, 0, :])
+
+
+def test_zero_turnover_limit_holds_the_previous_book() -> None:
+    weights = np.array([[[0.4, -0.4], [0.2, -0.1]]], dtype=np.float64)
+    previous = np.array([[0.5, -0.5]], dtype=np.float64)
+
+    actual = apply_turnover_limit(
+        weights,
+        previous_weights=previous,
+        mask=np.ones_like(weights, dtype=bool),
+        turnover_limit=0.0,
+    )
+
+    np.testing.assert_allclose(actual[0, 0], previous[0])
+    np.testing.assert_allclose(actual[0, 1], previous[0])
+
+
+def test_negative_turnover_limit_is_rejected() -> None:
+    weights = np.ones((1, 1, 2), dtype=np.float64)
+
+    with pytest.raises(ValueError, match="turnover_limit"):
+        apply_turnover_limit(
+            weights,
+            previous_weights=np.zeros((1, 2)),
+            mask=np.ones_like(weights, dtype=bool),
+            turnover_limit=-0.1,
+        )
+
+
 def test_portfolio_constraints_are_satisfied_jointly() -> None:
     raw = np.array([[[10.0, 1.0, -1.0, -2.0]]])
     mask = np.ones_like(raw, dtype=bool)

tests/test_portfolio.py

@@ -316,3 +316,42 @@ def test_deep_portfolio_reuses_fitted_topology_and_rejects_changes() -> None:
     )
     with pytest.raises(ValueError, match="adjacency_mask"):
         model.predict(changed)
+
+
+@pytest.mark.parametrize("heads", [1, 2])
+@pytest.mark.parametrize("isolated_asset", [False, True])
+def test_deep_portfolio_graph_masks_train_with_multiple_heads(
+    heads: int, isolated_asset: bool
+) -> None:
+    rng = np.random.default_rng(71)
+    adjacency = np.zeros((2, 2), dtype=bool)
+    if isolated_asset:
+        adjacency[0, :] = True
+    train = PortfolioSequenceBatch(
+        features=rng.normal(size=(2, 3, 2, 2)),
+        returns=0.01 * rng.normal(size=(2, 3, 2)),
+        adjacency_mask=adjacency,
+        asset_ids=("A", "B"),
+    )
+    model = DeepPortfolioModel(
+        DeepPortfolioConfig(
+            d_model=4,
+            n_heads=1,
+            cross_attention_heads=1,
+            macro_gnn_heads=heads,
+            dropout=0.0,
+            max_iters=1,
+            eval_every=1,
+            checkpoint_every=1,
+            default_checkpoint=1,
+            batch_size=2,
+            device="cpu",
+        )
+    )
+
+    fit = model.fit(train)
+    prediction = model.predict(train)
+
+    assert fit.converged
+    assert prediction.weights.shape == (2, 3, 2)
+    assert np.isfinite(prediction.weights).all()