Complex SBI Model Training

In this example we will demonstrate how to customise the training of your SBI model. In the Basic SBI Training tutorial we trained a default model, but here we will explore some of the cutomisation options available to you.

[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"
)

/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

We’ll use the default feature array settings for this example, but you can refer to the Feature Arrays tutorial to learn how to create your own custom feature arrays.

[2]:

fitter.create_feature_array();

2025-11-13 20:08:17,352 | synference | INFO     | ---------------------------------------------
2025-11-13 20:08:17,353 | synference | INFO     | Features: 8 features over 100 samples
2025-11-13 20:08:17,354 | synference | INFO     | ---------------------------------------------
2025-11-13 20:08:17,354 | synference | INFO     | Feature: Min - Max
2025-11-13 20:08:17,355 | synference | INFO     | ---------------------------------------------
2025-11-13 20:08:17,357 | synference | INFO     | JWST/NIRCam.F070W: 7.131974 - 42.758 AB
2025-11-13 20:08:17,358 | synference | INFO     | JWST/NIRCam.F090W: 7.108530 - 39.933 AB
2025-11-13 20:08:17,358 | synference | INFO     | JWST/NIRCam.F115W: 7.012560 - 38.354 AB
2025-11-13 20:08:17,360 | synference | INFO     | JWST/NIRCam.F150W: 6.969396 - 36.997 AB
2025-11-13 20:08:17,360 | synference | INFO     | JWST/NIRCam.F200W: 7.133157 - 35.470 AB
2025-11-13 20:08:17,361 | synference | INFO     | JWST/NIRCam.F277W: 7.670149 - 33.243 AB
2025-11-13 20:08:17,362 | synference | INFO     | JWST/NIRCam.F356W: 8.072730 - 32.490 AB
2025-11-13 20:08:17,362 | synference | INFO     | JWST/NIRCam.F444W: 8.353975 - 31.965 AB
2025-11-13 20:08:17,363 | synference | INFO     | ---------------------------------------------

We will break down the options into a few categories:

  1. Model Architecture Options: These options allow you to customize the architecture of the neural network used in the SBI model.

  • backend: This lets you choose between the SBI backends supported by LtU-ILI. The default is sbi, but you can also choose lampe if you prefer that framework.

  • engine: This option allows you to select the inference engine used for training. The default is NPE, but you can also choose NLE or NRE depending on your needs. Sequentual methods are also supported (SNPE, SNLE, SNRE).

  • model_type: This option lets you specify the type of model to use. The default is mdn (Mixture Density Network), but you can also choose maf (Masked Autoregressive Flow) or nsf (Neural Spline Flow) for more complex distributions. See the documentation for sbi or lampe for the full list of supported model types.

  • hidden_features: This option allows you to set the number of hidden features (neurons) in each layer of the neural network. The default is 50, but you can increase this for more complex models.

  • num_components: This option lets you specify the number of mixture components in the Mixture Density Network. The default is 4, but you can increase this for more complex distributions.

  • num_transforms: This option allows you to set the number of transforms in normalizing flow models (like MAF or NSF). The default is 4, but you can increase this for more complex distributions.

  • embedding_net: This option lets you provide a custom embedding network for processing the observations. By default no embedding network is used, but you can create your own using PyTorch and pass it here.

  • n_nets: This option allows you to specify the number of networks to use in an ensemble. The default is 1, but you can increase this to train an ensemble of models which may improve performance. Model parameters such as num_components and hidden_features can be specified as lists to vary them across the ensemble members.

  1. Training Options: These options allow you to customize the training process of the SBI model.

  • train_test_fraction: This option lets you specify the fraction of data to use for training versus testing. The default is 0.8 (80% training, 20% testing), but you can adjust this based on your dataset size.

  • training_batch_size: This option allows you to set the batch size for training. The default is 64, but you can vary this based on your hardware capabilities.

  • learning_rate: This option lets you specify the learning rate for the optimizer. The default is 1e-4, but you can adjust this based on your training dynamics.

  • override_prior_ranges: This option allows you to override the prior ranges defined in the model grid. You can provide a dictionary with parameter names as keys and tuples of (min, max) as values to set new prior ranges. The default parameter ranges are the minimum and maximum values defined in the model grid.

  • use_existing_indices: This option lets you choose whether to use existing training and testing indices if they are already saved on the object. The default is True, which will reuse the existing indices to ensure consistency across runs.

  • validation_fraction: This option allows you to specify the fraction of training data to use for validation during training. The default is 0.1 (10% of training data), but you can adjust this based on your needs.

  • stop_after_epochs: The number of epochs without improvement on the validation set after which training will stop early. The default is 15, but you can adjust this based on your training dynamics.

  • train_indices/test_indices: These options allow you to provide custom indices for training and testing data. By default, random indices are generated based on the train_test_fraction.

  1. Miscellaneous Options: These options provide additional customization for the SBI model training.

  • load_existing_model: This option allows you to load an existing trained model instead of training a new one if the named model file is found. The default is True.

  • evaluate_model: This option lets you choose whether to evaluate the model after training. The default is True, which will compute metrics on the test set.

  • save_model: This option allows you to save the trained model to disk. The default is True.

  • name_append: This option lets you append a custom string to the model name for easier identification. The default is a timestamp.

  • set_self: This option will set the trained model and various attributes to the current SBI object. The default is True.

  • plot: This option allows you to generate plots of the training results. The default is True.

  • verbose: This option lets you control the verbosity of the training output. The default is True, which will print detailed information during training.

  • save_method: This option allows you to specify the method for saving the model. The default is joblib, but you can choose other methods supported by the backend.

  • num_posterior_draws_per_sample: This option lets you specify the number of posterior draws to make per sample during evaluation. The default is 1000.

  • random_seed: This option allows you to set a random seed for reproducibility. The default is None, which means no specific seed is set.

  • out_dir: This option lets you specify the output directory where models and results will be saved. The default is the models/ folder.

  1. Online Training Options: These options allow you to customize the online training process of the SBI model.

  • learning_type: This option lets you specify the type of learning to perform. For online training, set this to online. The default is amortized.

  • num_online_rounds: This option lets you specify the number of online training rounds to perform. The default is 5, but you can adjust this based on your needs.

  • num_simulations: The number of simulations to run in each online training round. The default is 1000, but you can adjust this based on your computational resources.

  • online_training_xobs: As online training is not amortized, you need to provide the specific observation(s) for which you want to perform inference during online training. This option allows you to specify those observations.

  • initial_training_from_library: This option lets you specify whether to perform an initial training phase using a grid of simulations. The default is False. -simulator_: This option allows you to provide a custom simulator function for generating simulations during online training. By default, the simulator is attempted to be retrieved from the model grid.

Here’s an example where we use the lampe backend to train an ensemble of 2 models using different architectures (NSF and MAF) with varying hidden features and number of transforms. We also adjust the training parameters such as train-test split, batch size, and learning rate.

fitter.run_single_sbi(
    backend="lampe",
    engine="NPE",
    model_type=["nsf", "maf"],
    hidden_features=[12, 20],
    num_transforms=[9, 18],
    n_nets=2,
    train_test_fraction=0.9,
    training_batch_size=32,
    learning_rate=5e-2,
    override_prior_ranges={"log_mass": [7.75, 12.25]},
);