Default to MultiTaskGP + LearnedFeatureImputation for heterogeneous TL (#5193)

Summary:

Switches the default heterogeneous transfer learning model from a specialized per-task kernel model to a standard multi-task GP with learned feature imputation. The previous default model class is marked as deprecated.

Reviewed By: saitcakmak

Differential Revision: D102197137

Author: Carl Hvarfner

ax/generators/torch/botorch_modular/input_constructors/input_transforms.py

@@ -402,10 +402,14 @@ def _input_transform_argparse_learned_feature_imputation(
             torch.ones(d, dtype=dtype, device=torch_device),
         ]
     )
+    # The target task is at position 0 (target_dataset is prepended above), so
+    # at posterior time — when X arrives without a task column — LFI applies
+    # the target task's imputation pattern.
     kwargs: dict[str, Any] = {
         "feature_indices": feature_indices,
         "d": d,
         "task_feature_index": task_feature_index,
+        "target_task": 0,
         "bounds": bounds,
         "device": torch_device,
         "dtype": dtype,

ax/generators/torch/botorch_modular/utils.py

@@ -319,16 +319,15 @@ def choose_model_class(
         )
 
     # Check for heterogeneous multi-task datasets. If a model class was
-    # explicitly specified, respect it; otherwise default to HeterogeneousMTGP.
+    # explicitly specified, respect it; otherwise default to MultiTaskGP
+    # (LearnedFeatureImputation handles missing features).
     if (
         search_space_digest.task_features
         and isinstance(dataset, MultiTaskDataset)
         and dataset.has_heterogeneous_features
     ):
         model_class = (
-            specified_model_class
-            if specified_model_class is not None
-            else HeterogeneousMTGP
+            specified_model_class if specified_model_class is not None else MultiTaskGP
         )
         logger.debug(f"Chose BoTorch model class: {model_class}.")
         return model_class

ax/generators/torch/tests/test_utils.py

@@ -186,9 +186,9 @@ def test_choose_model_class_heterogeneous_task_features(self) -> None:
         mt_dataset = self._get_heterogeneous_mt_dataset()
         ssd = dataclasses.replace(self.search_space_digest, task_features=[-1])
 
-        # Default: HeterogeneousMTGP.
+        # Default: MultiTaskGP (LearnedFeatureImputation handles missing features).
         self.assertEqual(
-            HeterogeneousMTGP,
+            MultiTaskGP,
             choose_model_class(dataset=mt_dataset, search_space_digest=ssd),
         )
 
@@ -233,19 +233,23 @@ def test_copy_model_config_heterogeneous_mtgp(self) -> None:
         mt_dataset = self._get_heterogeneous_mt_dataset()
         ssd = dataclasses.replace(self.search_space_digest, task_features=[-1])
 
-        # Default (no model class specified) -> HeterogeneousMTGP.
-        # LFI is NOT injected; input_transform_classes stays DEFAULT.
+        # Default (no model class specified) -> MultiTaskGP.
+        # LFI is injected for MultiTaskGP with heterogeneous data.
         updated_config = copy_model_config_with_default_values(
             model_config=ModelConfig(),
             dataset=mt_dataset,
             search_space_digest=ssd,
         )
-        self.assertEqual(updated_config.botorch_model_class, HeterogeneousMTGP)
-        self.assertEqual(updated_config.input_transform_classes, [Normalize])
+        self.assertEqual(updated_config.botorch_model_class, MultiTaskGP)
         self.assertEqual(
-            none_throws(updated_config.input_transform_options),
-            {"Normalize": {}},
+            updated_config.input_transform_classes,
+            [Normalize, LearnedFeatureImputation],
         )
+        # LFI is present in transform classes but absent from options; its
+        # argparse computes kwargs from the dataset at construction time.
+        ito = none_throws(updated_config.input_transform_options)
+        self.assertEqual(ito, {"Normalize": {}})
+        self.assertNotIn("LearnedFeatureImputation", ito)
 
         # Explicit HeterogeneousMTGP behaves the same.
         updated_config = copy_model_config_with_default_values(