feat: Optimise sigmoid log-likelihood transformation from data

Author: saattrupdan

config/config.yaml

@@ -14,7 +14,7 @@ focus: null
 seed: 4242
 
 # Using a subset of the questions
-subset_csv: null
+subset_csv: data/processed/optimisation-davies-bouldin-penalty10/davies-bouldin-penalty10-eufocus-1000it.csv
 top_num_questions_in_subset: null
 
 plotting:

src/european_values/generative_training.py

@@ -1,13 +1,18 @@
 """Training generative on the dataset."""
 
 import logging
+from functools import partial
 from pathlib import Path
 
 import joblib
+import numpy as np
 import pandas as pd
+import scipy.optimize as opt
+from scipy.special import expit as sigmoid
+from scipy.special import logit as inverse_sigmoid
 from sklearn.model_selection import GridSearchCV
 from sklearn.neighbors import KernelDensity
-from sklearn.pipeline import Pipeline
+from sklearn.pipeline import Pipeline, check_is_fitted
 from sklearn.preprocessing import MinMaxScaler
 
 logger = logging.getLogger(__name__)
@@ -37,70 +42,171 @@ def train_generative_model(
     # Split data by country
     logger.info("Splitting data into train/test sets...")
     train_dfs: list[pd.DataFrame] = []
+    val_dfs: list[pd.DataFrame] = []
     test_dfs: list[pd.DataFrame] = []
     for country in eu_df["country_code"].unique():
         country_data = eu_df.query("country_code == @country").sample(
             frac=1, random_state=seed
         )
         n_test = min(test_samples_per_country, len(country_data) // 5)
         test_dfs.append(country_data.iloc[:n_test])
-        train_dfs.append(country_data.iloc[n_test:])
+        val_dfs.append(country_data.iloc[n_test : 2 * n_test])
+        train_dfs.append(country_data.iloc[2 * n_test :])
 
     # Set up the data as NumPy arrays
     train_matrix = scaler.transform(pd.concat(train_dfs)[question_columns].values)
+    val_matrix = scaler.transform(pd.concat(val_dfs)[question_columns].values)
     test_matrix = scaler.transform(pd.concat(test_dfs)[question_columns].values)
     logger.info(
-        f"There are {len(train_matrix):,} training samples and {len(test_matrix):,} "
-        "test samples."
+        f"There are {len(train_matrix):,} training samples, "
+        f"{len(val_matrix):,} validation samples, "
+        f"and {len(test_matrix):,} test samples."
     )
 
     # Initialise the model
     grid = GridSearchCV(
         estimator=KernelDensity(),
         param_grid=dict(
-            bandwidth=[0.1, 0.2, 0.3, 0.4, 0.5, 1.0, "scott", "silverman"],
-            leaf_size=[10, 20, 30, 40, 50],
+            bandwidth=[0.1, 0.2, 0.3, 0.4, 0.5], leaf_size=[10, 20, 30, 40, 50]
         ),
         n_jobs=-1,
     )
 
     # Fit the model
     logger.info("Training the model on the training data...")
     grid.fit(train_matrix)
-    model = grid.best_estimator_
+    model: KernelDensity = grid.best_estimator_
     logger.info(f"Best model found with the parameters {grid.best_params_}.")
 
+    # Set the `transform` method of the model to the score_samples method, as this will
+    # allow us to use the scaler, model and scorer in the same pipeline
+    model.transform = model.score_samples.__get__(model)
+
+    # logger.info("Computing the log-likelihoods for the training data...")
+    train_log_likelihoods = model.transform(train_matrix)
+
+    logger.info("Computing the log-likelihoods for the validation data...")
+    val_log_likelihoods = model.transform(val_matrix)
+
+    logger.info("Computing the log-likelihoods for the test data...")
+    test_log_likelihoods = model.transform(test_matrix)
+
+    # Fit the log-likelihood transform
+    logger.info("Fitting the sigmoid transform on the validation data...")
+    scorer = SigmoidTransformer().fit(val_log_likelihoods)
+
     # Evaluate the model
-    logger.info("Evaluating the model on the training and test data...")
-    train_log_likelihoods = model.score_samples(train_matrix)
+    logger.info("Evaluating the model on the training, validation and test data...")
     logger.info(
         f"Log-likelihoods for train:\n"
         f"\t- Mean: {train_log_likelihoods.mean():.4f}\n"
         f"\t- Std: {train_log_likelihoods.std():.4f}\n"
         f"\t- Min: {train_log_likelihoods.min():.4f}\n"
         f"\t- 10% quantile: {pd.Series(train_log_likelihoods).quantile(q=0.1):.4f}\n"
         f"\t- 90% quantile: {pd.Series(train_log_likelihoods).quantile(q=0.9):.4f}\n"
-        f"\t- Max: {train_log_likelihoods.max():.4f}"
+        f"\t- Max: {train_log_likelihoods.max():.4f}\n"
+        f"Mean score for train: {scorer.transform(train_log_likelihoods).mean():.0%}"
+    )
+    logger.info(
+        f"Log-likelihoods for validation:\n"
+        f"\t- Mean: {val_log_likelihoods.mean():.4f}\n"
+        f"\t- Std: {val_log_likelihoods.std():.4f}\n"
+        f"\t- Min: {val_log_likelihoods.min():.4f}\n"
+        f"\t- 10% quantile: {pd.Series(val_log_likelihoods).quantile(q=0.1):.4f}\n"
+        f"\t- 90% quantile: {pd.Series(val_log_likelihoods).quantile(q=0.9):.4f}\n"
+        f"\t- Max: {val_log_likelihoods.max():.4f}\n"
+        f"Mean score for validation: {scorer.transform(val_log_likelihoods).mean():.0%}"
     )
-    test_log_likelihoods = model.score_samples(test_matrix)
     logger.info(
         f"Log-likelihoods for test:\n"
         f"\t- Mean: {test_log_likelihoods.mean():.4f}\n"
         f"\t- Std: {test_log_likelihoods.std():.4f}\n"
         f"\t- Min: {test_log_likelihoods.min():.4f}\n"
         f"\t- 10% quantile: {pd.Series(test_log_likelihoods).quantile(q=0.1):.4f}\n"
         f"\t- 90% quantile: {pd.Series(test_log_likelihoods).quantile(q=0.9):.4f}\n"
-        f"\t- Max: {test_log_likelihoods.max():.4f}"
+        f"\t- Max: {test_log_likelihoods.max():.4f}\n"
+        f"Mean score for test: {scorer.transform(test_log_likelihoods).mean():.0%}"
     )
 
     # Train final model on all data
     logger.info("Training final model on entire EU dataset...")
     full_matrix = scaler.transform(eu_df[question_columns].values)
     model.fit(full_matrix)
-    pipeline = Pipeline([("scaler", scaler), ("model", model)])
+    pipeline = Pipeline([("scaler", scaler), ("model", model), ("scorer", scorer)])
 
     # Save the complete pipeline
     model_path = Path("models", "model.pkl")
     model_path.parent.mkdir(exist_ok=True)
     joblib.dump(pipeline, model_path)
     logger.info(f"Pipeline saved to {model_path.resolve()}")
+
+
+class SigmoidTransformer:
+    """Transformer to apply a sigmoid function to log-likelihoods."""
+
+    def fit(self, X: np.ndarray) -> "SigmoidTransformer":
+        """Fit the transformer to the data.
+
+        Args:
+            X:
+                The input array of log-likelihoods.
+
+        Returns:
+            The fitted transformer.
+        """
+        # We choose the alpha parameter such that the range is shrunk down to a length
+        # of 10, as that gives a smooth sigmoid curve that is not too flat
+        lower, upper = np.quantile(X, q=[0.05, 0.95])
+        self.alpha_ = 10 / (upper.item() - lower.item())
+
+        # Optimise the center of the sigmoid function to fit the target value
+        result: opt.OptimizeResult = opt.minimize(
+            fun=partial(self._loss, array=X, target=0.99, alpha=self.alpha_),
+            x0=np.array([0.0]),
+        )
+        self.center_ = result.x[0]
+        logger.info(
+            f"Fitted sigmoid transformer with alpha={self.alpha_:.2f} and "
+            f"center={self.center_:.2f}."
+        )
+        return self
+
+    def transform(self, X: np.ndarray) -> np.ndarray:
+        """Transform the input data using the fitted sigmoid function.
+
+        Args:
+            X:
+                The input array of log-likelihoods.
+
+        Returns:
+            The transformed values between 0 and 1.
+        """
+        check_is_fitted(estimator=self, attributes=["alpha_", "center_"])
+        return sigmoid(self.alpha_ * (X - self.center_))
+
+    @staticmethod
+    def _loss(
+        center: np.ndarray, array: np.ndarray, target: float, alpha: float
+    ) -> float:
+        """Calculate the loss for the sigmoid transformation.
+
+        The loss aims to get the sigmoid values of the array as close to a given target
+        value as possible.
+
+        Args:
+            center:
+                The center of the sigmoid curve.
+            array:
+                The input array of log-likelihoods.
+            target:
+                The target value for the sigmoid transformation.
+            alpha:
+                The steepness of the sigmoid curve.
+
+        Returns:
+            The l2 loss between the transformed values and the target sigmoid values.
+        """
+        target = inverse_sigmoid(target)
+        errors = (alpha * (array - center) - target) ** 2
+        l2_loss = np.mean(errors).item()
+        return l2_loss

src/scripts/evaluate_llm_benchmark.py

@@ -7,28 +7,13 @@
 import numpy as np
 import pandas as pd
 from omegaconf import DictConfig
-from sklearn.preprocessing import FunctionTransformer
 
 from european_values.data_loading import load_evs_trend_data, load_evs_wvs_data
 from european_values.data_processing import process_data
 
 logger = logging.getLogger("evaluate_llm")
 
 
-def sigmoid_transform(log_likelihoods, alpha=0.05, center=-50.0):
-    """Apply sigmoid transformation to log-likelihood values.
-
-    Args:
-        log_likelihoods: Array of log-likelihood values
-        alpha: Scaling parameter for sigmoid steepness (default 0.05)
-        center: Center point of the sigmoid (default -50.0)
-
-    Returns:
-        Transformed values between 0 and 1
-    """
-    return 1 / (1 + np.exp(-alpha * (log_likelihoods - center)))
-
-
 @hydra.main(config_path="../../config", config_name="config", version_base=None)
 def main(config: DictConfig) -> None:
     """Main evaluation function."""
@@ -63,7 +48,10 @@ def main(config: DictConfig) -> None:
             if col.startswith("question_") and col not in question_subset
         ]
         df.drop(columns=question_cols_to_remove, inplace=True)
-        logger.info(f"Using {len(question_subset)} questions from subset")
+        logger.info(
+            f"Using {len(question_subset)} questions from the subset "
+            f"{config.subset_csv!r}."
+        )
 
     # Process data without normalization (let pipeline handle it)
     logger.info("Processing the data WITHOUT normalization...")
@@ -76,25 +64,15 @@ def main(config: DictConfig) -> None:
     for country_group in df.country_group.unique():
         group_df = df.query("country_group == @country_group")
         responses = group_df[question_cols].values
-        log_likelihoods = pipeline.score_samples(responses)
-
-        # Apply sigmoid transformation using FunctionTransformer
-        # Ensures EU countries (around -31 mean) stay above 99%
-        sigmoid_transformer = FunctionTransformer(
-            func=sigmoid_transform,
-            validate=False
-        )
-        normalised_scores = sigmoid_transformer.transform(log_likelihoods.reshape(-1, 1)).flatten()
-
+        scores = pipeline.transform(responses)
         logger.info(
-            f"Log-likelihoods for {country_group}:\n"
-            f"\t- Mean: {log_likelihoods.mean():.2f}\n"
-            f"\t- Std: {log_likelihoods.std():.2f}\n"
-            f"\t- Min: {log_likelihoods.min():.2f}\n"
-            f"\t- 10% quantile: {np.quantile(log_likelihoods, q=0.1):.2f}\n"
-            f"\t- 90% quantile: {np.quantile(log_likelihoods, q=0.9):.2f}\n"
-            f"\t- Max: {log_likelihoods.max():.2f}\n"
-            f"\t- Mean normalised score: {normalised_scores.mean():.2%} "
+            f"Scores for {country_group}:\n"
+            f"\t- Mean: {scores.mean():.0%}\n"
+            f"\t- Std: {scores.std():.0%}\n"
+            f"\t- Min: {scores.min():.0%}\n"
+            f"\t- 10% quantile: {np.quantile(scores, q=0.1):.0%}\n"
+            f"\t- 90% quantile: {np.quantile(scores, q=0.9):.0%}\n"
+            f"\t- Max: {scores.max():.0%}\n"
         )
 
 

src/scripts/train_generative_model.py

@@ -47,7 +47,10 @@ def main(config: DictConfig) -> None:
             if col.startswith("question_") and col not in question_subset
         ]
         df.drop(columns=question_cols_to_remove, inplace=True)
-        logger.info(f"Using {len(question_subset)} questions from subset")
+        logger.info(
+            f"Using {len(question_subset)} questions from the subset "
+            f"{config.subset_csv!r}."
+        )
 
     # Process data but SKIP normalization (let pipeline handle it)
     logger.info("Processing the data WITHOUT normalization...")