Update AutoML to check for imbalanced binary or category output features

ludwig/automl/automl.py

@@ -25,6 +25,7 @@
     _ray_init,
     get_available_resources,
     get_model_type,
+    has_imbalanced_output,
     set_output_feature_metric,
 )
 from ludwig.constants import (
@@ -147,8 +148,10 @@ def create_auto_config(
     # Return
     :return: (dict) selected model configuration
     """
-    default_configs = _create_default_config(dataset, target, time_limit_s, random_seed)
-    model_config, model_category, row_count = _model_select(dataset, default_configs, user_config, use_reference_config)
+    default_configs, features_metadata = _create_default_config(dataset, target, time_limit_s, random_seed)
+    model_config, model_category, row_count = _model_select(
+        dataset, default_configs, features_metadata, user_config, use_reference_config
+    )
     if tune_for_memory:
         if ray.is_initialized():
             resources = get_available_resources()  # check if cluster has GPUS
@@ -218,6 +221,7 @@ def train_with_config(
 def _model_select(
     dataset: Union[str, pd.DataFrame, dd.core.DataFrame, DatasetInfo],
     default_configs,
+    features_metadata,
     user_config,
     use_reference_config: bool,
 ):
@@ -271,6 +275,10 @@ def _model_select(
                 if param in user_config[config_section]:
                     del base_config["hyperopt"]["parameters"][hyperopt_params]
 
+    # check if any binary or category output feature has highly imbalanced minority vs majority values
+    # note: check is done after any relevant user_config has been applied
+    has_imbalanced_output(base_config, features_metadata)
+
     # if single output feature, set relevant metric and goal if not already set
     base_config = set_output_feature_metric(base_config)
 

ludwig/automl/base_config.py

@@ -131,7 +131,7 @@ def _create_default_config(
     if not isinstance(dataset, DatasetInfo):
         dataset_info = get_dataset_info(dataset)
 
-    input_and_output_feature_config = get_features_config(
+    input_and_output_feature_config, features_metadata = get_features_config(
         dataset_info.fields, dataset_info.row_count, resources, target_name
     )
     # create set of all feature types appearing in the dataset
@@ -165,7 +165,7 @@ def _create_default_config(
         combiner_config = load_yaml(default_config)
         model_configs[COMBINER][combiner_type] = combiner_config
 
-    return model_configs
+    return model_configs, features_metadata
 
 
 # Read in the score and configuration of a reference model trained by Ludwig for each dataset in a list.
@@ -194,7 +194,9 @@ def get_dataset_info_from_source(source: DataSource) -> DatasetInfo:
     fields = []
     for field in source.columns:
         dtype = source.get_dtype(field)
-        num_distinct_values, distinct_values = source.get_distinct_values(field, MAX_DISTINCT_VALUES_TO_RETURN)
+        num_distinct_values, distinct_values, distinct_values_balance = source.get_distinct_values(
+            field, MAX_DISTINCT_VALUES_TO_RETURN
+        )
         nonnull_values = source.get_nonnull_values(field)
         image_values = source.get_image_values(field)
         audio_values = source.get_audio_values(field)
@@ -207,6 +209,7 @@ def get_dataset_info_from_source(source: DataSource) -> DatasetInfo:
                 dtype=dtype,
                 distinct_values=distinct_values,
                 num_distinct_values=num_distinct_values,
+                distinct_values_balance=distinct_values_balance,
                 nonnull_values=nonnull_values,
                 image_values=image_values,
                 audio_values=audio_values,
@@ -241,7 +244,7 @@ def get_features_config(
     targets = set(targets)
 
     metadata = get_field_metadata(fields, row_count, resources, targets)
-    return get_config_from_metadata(metadata, targets)
+    return get_config_from_metadata(metadata, targets), metadata
 
 
 def get_config_from_metadata(metadata: List[FieldMetadata], targets: Set[str] = None) -> dict:
@@ -297,6 +300,7 @@ def get_field_metadata(
                 excluded=should_exclude(idx, field, dtype, row_count, targets),
                 mode=infer_mode(field, targets),
                 missing_values=missing_value_percent,
+                imbalance_ratio=field.distinct_values_balance,
             )
         )
 

ludwig/automl/data_source.py

@@ -51,7 +51,11 @@ def get_dtype(self, column: str) -> str:
     def get_distinct_values(self, column, max_values_to_return: int) -> Tuple[int, List[str]]:
         unique_values = self.df[column].dropna().unique()
         num_unique_values = len(unique_values)
-        return num_unique_values, unique_values[:max_values_to_return]
+        unique_values_counts = self.df[column].value_counts()
+        unique_majority_values = unique_values_counts[unique_values_counts.idxmax()]
+        unique_minority_values = unique_values_counts[unique_values_counts.idxmin()]
+        unique_values_balance = unique_minority_values / unique_majority_values
+        return num_unique_values, unique_values[:max_values_to_return], unique_values_balance
 
     def get_nonnull_values(self, column: str) -> int:
         return len(self.df[column].notnull())

ludwig/automl/utils.py

@@ -8,7 +8,20 @@
 from numpy import nan_to_num
 from pandas import Series
 
-from ludwig.constants import COMBINER, CONFIG, HYPEROPT, NAME, NUMBER, PARAMETERS, SEARCH_ALG, TRAINER, TYPE
+from ludwig.constants import (
+    BINARY,
+    CATEGORY,
+    COMBINER,
+    CONFIG,
+    HYPEROPT,
+    IMBALANCE_DETECTION_RATIO,
+    NAME,
+    NUMBER,
+    PARAMETERS,
+    SEARCH_ALG,
+    TRAINER,
+    TYPE,
+)
 from ludwig.features.feature_registries import output_type_registry
 from ludwig.modules.metric_registry import metric_registry
 from ludwig.utils.defaults import default_combiner_type
@@ -29,6 +42,7 @@ class FieldInfo:
     dtype: str
     key: str = None
     distinct_values: List = None
+    distinct_values_balance: float = 1.0
     num_distinct_values: int = 0
     nonnull_values: int = 0
     image_values: int = 0
@@ -52,6 +66,7 @@ class FieldMetadata:
     excluded: bool
     mode: str
     missing_values: float
+    imbalance_ratio: float
 
 
 def avg_num_tokens(field: Series) -> int:
@@ -178,3 +193,20 @@ def set_output_feature_metric(base_config):
         base_config[HYPEROPT]["metric"] = output_metric
         base_config[HYPEROPT]["goal"] = output_goal
     return base_config
+
+
+def has_imbalanced_output(base_config, features_metadata) -> bool:
+    """Check binary and category output feature(s) for imbalance, i.e., low minority/majority instance count
+    ratio."""
+    imbalanced_output = False
+    for output_feature in base_config["output_features"]:
+        if output_feature[TYPE] == BINARY or output_feature[TYPE] == CATEGORY:
+            for feature_metadata in features_metadata:
+                if output_feature[NAME] == feature_metadata.name:
+                    if feature_metadata.imbalance_ratio < IMBALANCE_DETECTION_RATIO:
+                        logging.info(
+                            f"Imbalance in {output_feature[NAME]}: minority/majority={feature_metadata.imbalance_ratio}"
+                        )
+                        imbalanced_output = True
+                    break
+    return imbalanced_output

ludwig/constants.py

@@ -151,6 +151,7 @@
 COMBINER = "combiner"
 
 BALANCE_PERCENTAGE_TOLERANCE = 0.03
+IMBALANCE_DETECTION_RATIO = 0.05
 
 TABULAR = "tabular"
 AUTOML_DEFAULT_TABULAR_MODEL = "tabnet"