[SYSTEMDS-3947] Introduce robustness check in hyperparameter tuning

src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py

@@ -40,6 +40,10 @@
 from systemds.scuro.utils.checkpointing import CheckpointManager
 
 
+def _wandb_safe_tag(s: str, max_len: int = 64) -> str:
+    return s if len(s) <= max_len else s[: max_len - 3] + "..."
+
+
 def _get_params_for_node(node_id, params):
     return {
         k.split("-")[-1]: v for k, v in params.items() if k.startswith(node_id + "-")
@@ -53,6 +57,7 @@ def _param_values_to_spec(
         return {"name": full_name, "type": "categorical", "domain": list(param_values)}
     if isinstance(param_values, tuple) and len(param_values) == 2:
         lo, hi = param_values
+        lo, hi = min(lo, hi), max(lo, hi)
         if isinstance(lo, int) and isinstance(hi, int):
             return {"name": full_name, "type": "integer", "domain": (lo, hi)}
         return {"name": full_name, "type": "real", "domain": (float(lo), float(hi))}
@@ -456,16 +461,25 @@ def visit_node(node_id):
         modalities_override = (
             self._get_cached_modalities_for_task(task, modality_ids) if mm_opt else None
         )
+
+        baseline_params, baseline_raw_scores = self.evaluate_dag_config(
+            dag,
+            {},
+            node_order,
+            modality_ids,
+            task,
+            modalities_override=modalities_override,
+        )
+        baseline = (
+            baseline_params,
+            [
+                self._score_value(baseline_raw_scores[0]),
+                self._score_value_list(baseline_raw_scores[1]),
+                self._score_value(baseline_raw_scores[2]),
+            ],
+        )
+
         if not hyperparams:
-            # TODO: extract the information from the unimodal optimization results
-            baseline = self.evaluate_dag_config(
-                dag,
-                {},
-                node_order,
-                modality_ids,
-                task,
-                modalities_override=modalities_override,
-            )
             all_results = [baseline]
         else:
             param_specs = self._build_param_specs(hyperparams)
@@ -482,6 +496,7 @@ def visit_node(node_id):
                 initial_config=None,
                 rep_name=rep_name,
             )
+            all_results.append(baseline)
 
         if not all_results:
             return None
@@ -491,14 +506,37 @@ def get_score(result):
             if isinstance(score, PerformanceMeasure):
                 return score.average_scores[self.scoring_metric]
             elif isinstance(score, list):
-                return score[1]
+                score = score[1]
+
+            if isinstance(score, list):
+                score = np.mean(score)
             return score
 
-        if self.maximize_metric:
-            best_params, best_score = max(all_results, key=get_score)
-        else:
-            best_params, best_score = min(all_results, key=get_score)
+        best_params, best_score = all_results[0][0], get_score(all_results[0])
+        for params, score in all_results[1:]:
+            candidate_score = get_score((params, score))
+            if self._is_better(candidate_score, best_score):
+                best_params, best_score = params, candidate_score
+
+        if hyperparams and best_params != baseline_params:
+            baseline_folds = self._score_value_list(baseline_raw_scores[1])
+            candidate_folds = next(s for p, s in all_results if p == best_params)[1]
 
+            if baseline_folds is not None and candidate_folds is not None:
+                accept, candidate_range, baseline_range = self._should_accept_optimized(
+                    baseline_folds, candidate_folds
+                )
+                if not accept:
+                    self.logger.info(
+                        f"{rep_name}: optimized config too variable across folds "
+                        f"(range={candidate_range:.4f} vs baseline={baseline_range:.4f}) "
+                        "— keeping baseline"
+                    )
+                    best_params, best_score = baseline_params, get_score(baseline)
+            else:
+                self.logger.warning(
+                    f"{rep_name}: fold-level scores unavailable, skipping variance gate"
+                )
         tuning_time = time.time() - start_time
 
         best_result = HyperparamResult(
@@ -572,6 +610,23 @@ def _score_value(self, score: Any) -> float:
             return score.average_scores.get(self.scoring_metric, np.nan)
         return score
 
+    def _score_value_list(self, score: Any) -> List[float]:
+        if isinstance(score, PerformanceMeasure):
+            return score.scores.get(self.scoring_metric, [])
+        return [score]
+
+    def _should_accept_optimized(
+        self,
+        baseline_folds: List[float],
+        candidate_folds: List[float],
+        range_threshold: float = 0.03,
+    ) -> Tuple[bool, float, float]:
+        baseline_range = max(baseline_folds) - min(baseline_folds)
+        candidate_range = max(candidate_folds) - min(candidate_folds)
+        if candidate_range > baseline_range * (1 + range_threshold):
+            return False, candidate_range, baseline_range
+        return True, candidate_range, baseline_range
+
     def _is_better(self, candidate_score: float, best_score: float) -> bool:
         if np.isnan(candidate_score):
             return False
@@ -789,7 +844,10 @@ def _search_best_configs(
                     "project": self.wandb_project,
                     "entity": self.wandb_entity,
                     "group": self.wandb_group or task.model.name,
-                    "tags": self.wandb_tags + [rep_name, task.model.name],
+                    "tags": [
+                        _wandb_safe_tag(t)
+                        for t in (self.wandb_tags + [rep_name, task.model.name])
+                    ],
                     "name": f"{task.model.name}-{rep_name}-{int(time.time())}",
                     "config": {
                         "task": task.model.name,
@@ -835,10 +893,18 @@ def objective(trial: optuna.Trial) -> float:
 
             seen[self._config_key(params)] = (
                 params,
-                [train_score, val_score, test_score],
+                [train_score, self._score_value_list(scores[1]), test_score],
+            )
+
+            trial_results.append(
+                (params, [train_score, self._score_value_list(scores[1]), test_score])
             )
+            val_folds = self._score_value_list(scores[1])
+            if val_folds is not None and len(val_folds) > 1:
+                lam = 0.5
+                robust_val = np.mean(val_folds) - lam * np.std(val_folds, ddof=1)
+                return robust_val
 
-            trial_results.append((params, [train_score, val_score, test_score]))
             return val_score
 
         callbacks = [c for c in [wandb_cb] if c is not None]

src/main/python/systemds/scuro/representations/color_histogram.py

@@ -49,6 +49,11 @@ def __init__(
         super().__init__(
             "ColorHistogram", ModalityType.EMBEDDING, self._get_parameters()
         )
+        if params is not None:
+            color_space = params.get("color_space", color_space)
+            bins = params.get("bins", bins)
+            normalize = params.get("normalize", normalize)
+            aggregation = params.get("aggregation", aggregation)
         self.color_space = color_space
         self.bins = bins
         self.normalize = normalize

src/main/python/systemds/scuro/representations/lstm.py

@@ -51,19 +51,27 @@ def __init__(
             "depth": [1, 2, 3],
             "dropout_rate": [0.1, 0.2, 0.3, 0.4, 0.5],
             "learning_rate": [0.001, 0.0001, 0.01, 0.1],
-            "epochs": [10, 2050, 100, 200],
+            "epochs": [10, 20, 50, 100, 200],
             "batch_size": [8, 16, 32, 64, 128],
         }
 
         super().__init__("LSTM", parameters)
 
+        if params is not None:
+            width = params.get("width", width)
+            depth = params.get("depth", depth)
+            dropout_rate = params.get("dropout_rate", dropout_rate)
+            learning_rate = params.get("learning_rate", learning_rate)
+            epochs = params.get("epochs", epochs)
+            batch_size = params.get("batch_size", batch_size)
+
         self.width = int(width)
         self.depth = int(depth)
         self.dropout_rate = float(dropout_rate)
         self.learning_rate = float(learning_rate)
         self.epochs = int(epochs)
         self.batch_size = int(batch_size)
-
+        self.device = get_device()
         self.needs_training = True
         self.needs_alignment = True
         self.model = None
@@ -180,7 +188,6 @@ def execute(self, modalities: List[Modality], labels: np.ndarray = None):
         self.input_dim = X.shape[2]
 
         self.model = self._build_model(self.input_dim, self.num_classes)
-        self.device = get_device_for_model(self.model, memory_factor=1.5)
         self.model = self.model.to(self.device)
 
         if self.is_multilabel:
@@ -245,7 +252,6 @@ def apply_representation(self, modalities: List[Modality]) -> np.ndarray:
 
         X = self._prepare_data(modalities)
 
-        self.device = get_device_for_model(self.model, memory_factor=1.5)
         self.model = self.model.to(self.device)
 
         X_tensor = torch.FloatTensor(X)

src/main/python/systemds/scuro/representations/mlp_averaging.py

@@ -54,6 +54,9 @@ def __init__(self, output_dim=512, batch_size=32, params=None):
             "batch_size": [8, 16, 32, 64, 128],
         }
         super().__init__("MLPAveraging", parameters)
+        if params is not None:
+            output_dim = params.get("output_dim", output_dim)
+            batch_size = params.get("batch_size", batch_size)
         self.output_dim = output_dim
         self.batch_size = batch_size
         self.device = None

src/main/python/systemds/scuro/representations/resnet.py

@@ -56,6 +56,7 @@ def __init__(
         if params is not None:
             self.batch_size = int(params.get("batch_size", batch_size))
             self.layer_name = params.get("layer_name", layer_name)
+            model_name = params.get("model_name", model_name)
         else:
             self.batch_size = batch_size
             self.layer_name = layer_name

src/main/python/systemds/scuro/representations/swin_video_transformer.py

@@ -59,6 +59,8 @@ def __init__(self, layer_name="avgpool", params=None):
         }
         self.data_type = torch.float32
         super().__init__("SwinVideoTransformer", ModalityType.EMBEDDING, parameters)
+        if params is not None:
+            layer_name = params.get("layer_name", layer_name)
         self.layer_name = layer_name
         self.model = swin3d_t(weights=models.video.Swin3D_T_Weights.KINETICS400_V1)
         self.device = get_device_for_model(self.model, memory_factor=1.5)

src/main/python/systemds/scuro/representations/text_context_with_indices.py

@@ -296,6 +296,11 @@ def __init__(self, max_words=55, overlap=0.5, stride=None, params=None):
             "stride": [10, 15, 20, 30],
         }
         super().__init__("OverlappingSplit", parameters)
+        if params is not None:
+            max_words = params.get("max_words", max_words)
+            overlap = params.get("overlap", overlap)
+            overlap_words = int(max_words * overlap)
+            stride = params.get("stride", max_words - overlap_words)
         self.max_words = max_words
         self.overlap = overlap
         self.stride = stride

src/main/python/systemds/scuro/representations/tfidf.py

@@ -36,6 +36,8 @@ class TfIdf(UnimodalRepresentation):
     def __init__(self, min_df=2, output_file=None, params=None):
         parameters = {"min_df": [min_df, 4, 8]}
         super().__init__("TF-IDF", ModalityType.EMBEDDING, parameters)
+        if params is not None:
+            min_df = params.get("min_df", min_df)
         self.min_df = int(min_df)
         self.output_file = output_file
         self.data_type = np.float32

src/main/python/systemds/scuro/representations/timeseries_representations.py

@@ -199,6 +199,8 @@ def compute_feature(self, signal, axis=-1):
 class ACF(TimeSeriesRepresentation):
     def __init__(self, k=1, params=None):
         super().__init__("ACF", {"k": [1, 2, 5, 10, 20, 25, 50, 100, 200, 500]})
+        if params is not None:
+            k = params.get("k", k)
         self.k = k
 
     def compute_feature(self, signal, axis=-1):
@@ -246,6 +248,8 @@ def compute_feature(self, signal, axis=-1):
 class SpectralCentroid(TimeSeriesRepresentation):
     def __init__(self, fs=1.0, params=None):
         super().__init__("SpectralCentroid", parameters={"fs": [0.5, 1.0, 2.0]})
+        if params is not None:
+            fs = params.get("fs", fs)
         self.fs = fs
 
     def compute_feature(self, signal, axis=-1):
@@ -271,6 +275,10 @@ def __init__(self, fs=1.0, f1=0.0, f2=0.5, params=None):
             "BandpowerFFT",
             parameters={"fs": [0.5, 1.0], "f1": [0.0, 1.0], "f2": [0.5, 1.0]},
         )
+        if params is not None:
+            fs = params.get("fs", fs)
+            f1 = params.get("f1", f1)
+            f2 = params.get("f2", f2)
         self.fs = fs
         self.f1 = f1
         self.f2 = f2

src/main/python/systemds/scuro/representations/vgg.py

@@ -56,6 +56,8 @@ def __init__(
         self.model = self.model.to(self.device)
         parameters = self._get_parameters()
         super().__init__("VGG19", ModalityType.EMBEDDING, parameters)
+        if params is not None:
+            layer = params.get("layer_name", layer)
         self.output_file = output_file
         self.layer_name = layer
         self.model.eval()

src/main/python/systemds/scuro/representations/word2vec.py

@@ -49,6 +49,9 @@ def __init__(self, vector_size=150, min_count=1, output_file=None, params=None):
             "min_count": [1, 2, 4, 8],
         }
         super().__init__("Word2Vec", ModalityType.EMBEDDING, parameters)
+        if params is not None:
+            vector_size = params.get("vector_size", vector_size)
+            min_count = params.get("min_count", min_count)
         self.vector_size = vector_size
         self.min_count = min_count
         self.output_file = output_file

src/main/python/systemds/scuro/representations/x3d.py

@@ -50,6 +50,9 @@ def __init__(
         self, layer="classifier.1", model_name="s3d", output_file=None, params=None
     ):
         self.data_type = torch.float32
+        if params is not None:
+            model_name = params.get("model_name", model_name)
+            layer = params.get("layer_name", layer)
         self.model_name = model_name
         parameters = self._get_parameters()
         super().__init__("X3D", ModalityType.EMBEDDING, parameters)
@@ -127,7 +130,7 @@ def model_name(self, model_name):
 
     def _get_parameters(self, high_level=True):
         parameters = {"model_name": [], "layer_name": []}
-        for m in ["c3d", "s3d"]:
+        for m in ["r3d", "s3d"]:
             parameters["model_name"].append(m)
 
         # TODO: add embedding dimensions for each layer