From 33faa0f184bffaf86deb1f8aa44a3a9514dc4134 Mon Sep 17 00:00:00 2001 From: ChrisW09 <50968720+ChrisW09@users.noreply.github.com> Date: Mon, 27 Jul 2026 21:25:39 +0200 Subject: [PATCH 1/2] fix(hpo): sign-safe penalties, head_layer_size_length handling, honor random_state - The crashed-trial penalty best_val_loss*100 rewarded crashes whenever the best loss was negative (routine for LSS NLL), so gp_minimize could return a configuration that failed to fit as the best hyperparameters. It is now best + 100*abs(best) + 1. - The pruning threshold best*1.5 fell *below* a negative best loss, pruning every subsequent trial at prune_epoch; now best + 0.5*abs(best). - The final best-params loop lacked the head_layer_size_length branch the trial loop has, and 'head_layer_size_length'.startswith('head_layer_size_') is True: the optimized length was rounded to 0, prepended as a layer size, and the list was never truncated to that length. - gp_minimize now uses the estimator's random_state when set instead of a hardcoded 42. - on_validation_epoch_end no longer records the pre-training sanity-check loss, which shifted epoch_val_loss_at() -- the pruning baseline -- by one epoch. Fixes #420 Co-Authored-By: Claude Fable 5 --- deeptab/models/_mixins/hpo.py | 23 ++++++++--- deeptab/training/lightning_module.py | 4 ++ tests/test_hpo_pruning.py | 58 ++++++++++++++++++++++++++++ 3 files changed, 80 insertions(+), 5 deletions(-) create mode 100644 tests/test_hpo_pruning.py diff --git a/deeptab/models/_mixins/hpo.py b/deeptab/models/_mixins/hpo.py index a55a57e0..e18de2de 100644 --- a/deeptab/models/_mixins/hpo.py +++ b/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 diff --git a/deeptab/training/lightning_module.py b/deeptab/training/lightning_module.py index 9619add7..730de6af 100644 --- a/deeptab/training/lightning_module.py +++ b/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() diff --git a/tests/test_hpo_pruning.py b/tests/test_hpo_pruning.py new file mode 100644 index 00000000..c0f3e3e1 --- /dev/null +++ b/tests/test_hpo_pruning.py @@ -0,0 +1,58 @@ +"""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. +""" + +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(model._task_model.val_losses) == 2 From b8f4fdec8458d2977ee14be48b78066b0554890a Mon Sep 17 00:00:00 2001 From: ChrisW09 <50968720+ChrisW09@users.noreply.github.com> Date: Mon, 27 Jul 2026 21:57:04 +0200 Subject: [PATCH 2/2] test: satisfy pyright on the HPO pruning regression test model._task_model is typed ITaskModel | None, which has no .val_losses. Co-Authored-By: Claude Fable 5 --- tests/test_hpo_pruning.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_hpo_pruning.py b/tests/test_hpo_pruning.py index c0f3e3e1..231b7a4d 100644 --- a/tests/test_hpo_pruning.py +++ b/tests/test_hpo_pruning.py @@ -5,6 +5,8 @@ whose val_loss is a mean NLL. """ +from typing import Any, cast + import numpy as np import pandas as pd @@ -55,4 +57,4 @@ def test_val_losses_excludes_sanity_check(self): y = rng.randn(60) model = MLPRegressor() model.fit(X, y, max_epochs=2, batch_size=16, accelerator="cpu") - assert len(model._task_model.val_losses) == 2 + assert len(cast(Any, model._task_model).val_losses) == 2