fix(hpo): sign-safe penalties, head_layer_size_length handling, honor random_state

deeptab/models/_mixins/hpo.py

@@ -162,10 +162,13 @@ def _objective(hyperparams):
                 build_kwargs["embeddings_val"] = embeddings_val
             self.build_model(X, y, **build_kwargs)
 
+            # "50% worse than the best" — computed sign-safely, since val_loss
+            # can be negative (LSS NLL) and a plain *1.5 would then be *below*
+            # the best loss, pruning every trial.
             if prune_by_epoch:
-                early_pruning_threshold = best_epoch_val_loss * 1.5
+                early_pruning_threshold = best_epoch_val_loss + 0.5 * abs(best_epoch_val_loss)
             else:
-                early_pruning_threshold = best_val_loss * 1.5  # type: ignore[operator]
+                early_pruning_threshold = best_val_loss + 0.5 * abs(best_val_loss)  # type: ignore[operator]
 
             self._task_model.early_pruning_threshold = early_pruning_threshold  # type: ignore
             self._task_model.pruning_epoch = prune_epoch  # type: ignore
@@ -190,16 +193,24 @@ def _objective(hyperparams):
 
             except Exception as e:
                 print(f"Error encountered during fit with hyperparameters {hyperparams}: {e}")
-                return best_val_loss * 100  # type: ignore[operator]
+                # Penalty must be much *worse* (higher) than the best loss even
+                # when the best loss is negative; best*100 would reward crashes.
+                return best_val_loss + 100.0 * abs(best_val_loss) + 1.0  # type: ignore[operator]
 
-        result = gp_minimize(_objective, param_space, n_calls=time, random_state=42)
+        _hpo_seed = getattr(self, "random_state", None)
+        result = gp_minimize(_objective, param_space, n_calls=time, random_state=42 if _hpo_seed is None else _hpo_seed)
 
         best_hparams = result.x  # type: ignore
         head_layer_sizes = [] if "head_layer_sizes" in self.config.__dataclass_fields__ else None
         layer_sizes = [] if "layer_sizes" in self.config.__dataclass_fields__ else None
 
+        best_head_layer_size_length = None
         for key, param_value in zip(param_names, best_hparams, strict=False):
-            if key.startswith("head_layer_size_") and head_layer_sizes is not None:
+            if key == "head_layer_size_length":
+                # Must be checked before the startswith() branch below, which
+                # would round the length to 0 and prepend it as a layer size.
+                best_head_layer_size_length = param_value
+            elif key.startswith("head_layer_size_") and head_layer_sizes is not None:
                 head_layer_sizes.append(round_to_nearest_16(param_value))
             elif key.startswith("layer_size_") and layer_sizes is not None:
                 layer_sizes.append(round_to_nearest_16(param_value))
@@ -209,6 +220,8 @@ def _objective(hyperparams):
                 setattr(self.config, key, param_value)
 
         if head_layer_sizes is not None and head_layer_sizes:
+            if best_head_layer_size_length is not None:
+                head_layer_sizes = head_layer_sizes[:best_head_layer_size_length]
             self.config.head_layer_sizes = head_layer_sizes
         if layer_sizes is not None and layer_sizes:
             self.config.layer_sizes = layer_sizes

deeptab/training/lightning_module.py

@@ -577,6 +577,10 @@ def on_validation_epoch_end(self):
         loss exceeds the `early_pruning_threshold`, the training is stopped early by setting
         `self.trainer.should_stop` to True.
         """
+        # The sanity-check validation runs before training starts; recording it
+        # would shift every epoch_val_loss_at() lookup off by one.
+        if self.trainer.sanity_checking:
+            return
         val_loss = self.trainer.callback_metrics.get("val_loss")
         if val_loss is not None:
             val_loss_value = val_loss.item()

tests/test_hpo_pruning.py

@@ -0,0 +1,60 @@
+"""Regression tests for the HPO driver's pruning and best-params handling.
+
+The objective and pruning arithmetic used ``best * 1.5`` / ``best * 100``,
+which invert when the validation loss is negative -- routine for LSS models
+whose val_loss is a mean NLL.
+"""
+
+from typing import Any, cast
+
+import numpy as np
+import pandas as pd
+
+from deeptab.models import MLPRegressor
+from deeptab.models._mixins.hpo import round_to_nearest_16
+
+
+def _pruning_threshold(best):
+    """Mirror of the threshold expression in _FitMixin.optimize_hparams."""
+    return best + 0.5 * abs(best)
+
+
+def _failure_penalty(best):
+    """Mirror of the failed-trial penalty in _FitMixin.optimize_hparams."""
+    return best + 100.0 * abs(best) + 1.0
+
+
+class TestSignSafeObjective:
+    def test_pruning_threshold_is_worse_than_best(self):
+        for best in (-12.5, -1.0, -0.001, 0.0, 0.5, 42.0):
+            assert _pruning_threshold(best) >= best
+
+    def test_failure_penalty_is_worse_than_best(self):
+        """A crashing trial must never look better than the incumbent."""
+        for best in (-12.5, -1.0, -0.001, 0.0, 0.5, 42.0):
+            assert _failure_penalty(best) > best
+
+
+class TestHeadLayerSizeLength:
+    def test_length_key_matches_the_size_prefix(self):
+        """Why the final apply loop needs an explicit length branch."""
+        assert "head_layer_size_length".startswith("head_layer_size_")
+
+    def test_length_value_would_round_to_zero(self):
+        """Rounding the length as if it were a layer size yields a 0-width layer."""
+        assert all(round_to_nearest_16(n) == 0 for n in range(1, 6))
+
+
+class TestValLossTracking:
+    def test_val_losses_excludes_sanity_check(self):
+        """The sanity-check validation must not enter val_losses.
+
+        Recording it shifts epoch_val_loss_at(e) to the loss of epoch e-1,
+        which skews the HPO pruning baseline.
+        """
+        rng = np.random.RandomState(0)
+        X = pd.DataFrame({"a": rng.randn(60), "b": rng.randn(60)})
+        y = rng.randn(60)
+        model = MLPRegressor()
+        model.fit(X, y, max_epochs=2, batch_size=16, accelerator="cpu")
+        assert len(cast(Any, model._task_model).val_losses) == 2