Model Optimization

Synference supports Optuna for hyperparameter optimization. Optuna is an automatic hyperparameter optimization software framework, particularly designed for machine learning.

Optuna works through the creation of a ‘Study’, which is an optimization task to be executed. Each Study consists of multiple ‘Trials’, where each Trial represents a single execution of the objective function with a specific set of hyperparameters.

Synference supports several different objective functions for optimization, including:

  • log_probability: Maximizes the log probability of the model given the data.

First let’s create a fitter.

[1]:

from synference import SBI_Fitter, test_data_dir

fitter = SBI_Fitter.init_from_hdf5(
    model_name="test", hdf5_path=f"{test_data_dir}/example_model_library.hdf5"
)
fitter.create_feature_array(verbose=False);

/opt/hostedtoolcache/Python/3.10.19/x64/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

Optimizing a model is primarily done through the optimize_sbi method of the SBI_Fitter class.

You must choose a study name, and optionally a storage location for the study. If no storage location is provided, an in-memory storage will be used, which means that the study will not be saved after the program ends.

Optuna works best with a database backend for storage, such as SQLite or PostgreSQL. This allows you to save and resume studies, and also to share studies between different processes or machines. This can be set with the persistent_storage parameter.

The model is setup through the configuration of two dictionaries, suggested_hyperparameters and fixed_hyperparameters. The suggested_hyperparameters dictionary contains hyperparameters that will be optimized by Optuna, while the fixed_hyperparameters dictionary contains hyperparameters that will remain constant during the optimization process.

The keys of these dictionaries correspond to the names of the hyperparameters in the model, and the values in the suggested_hyperparameters dictionary are either two component lists/tuples specifying the minimum and maximum values for the hyperparameter, or a list of categorical values to choose from. The values in the fixed_hyperparameters dictionary are simply the fixed values for those hyperparameters.

We can also set the number of trials to run for the optimization process using the n_trials parameter. Note that if you are parallelizing the optimization process across multiple workers, each worker will run n_trials trials, so the total number of trials will be n_trials multiplied by the number of workers.

Optimization Metrics

We must also choose our optimization metric(s). Multi-objective optimization is supported by providing a list of metrics. The available metrics are:

  • “log_prob”: Maximizes the log probability of the model given the data.

  • “log_prob-pit”: Maximizes the log probability adjusted by the Probability Integral Transform (PIT) to encourage calibrated posteriors.

  • “loss”: Minimizes the negative log probability of the model given the data.

  • Custom callable: You can also provide a custom callable function that takes in the posterior, test data, and test labels, and returns a scalar score.

Note that you should set the ‘direction’ parameter to either ‘minimize’ or ‘maximize’ depending on whether you want to minimize or maximize the chosen metric(s).

fitter.optimize_sbi(
    study_name="my_study",
    suggested_hyperparameters={
        "learning_rate": (1e-5, 1e-2),  # Optimize learning rate between 1e-5 and 1e-2
        "num_transforms": [5, 80],     # Optimize number of transforms between 5 and 80
        "hidden_features": (16, 128),   # Optimize hidden features between 16 and 128
    },
    fixed_hyperparameters={
        "training_batch_size": 64,                # Fixed batch size
        "model_type": "maf",                       # Fixed model type
    },
    n_trials=50,                        # Number of trials to run
    score_metrics="log_prob",  # Single-objective optimization
    directions='maximize',
    persistent_storage=True
)