bayesflow-hpo

Automated hyperparameter optimization for BayesFlow 2.x neural posterior estimation (NPE) models, powered by Optuna.

Tunes inference networks, summary networks, and training hyperparameters — so you can focus on the model, not the plumbing.

Installation

pip install bayesflow-hpo            # core
pip install bayesflow-hpo[dashboard] # + Optuna dashboard
pip install bayesflow-hpo[sklearn]   # + C2ST metrics

Requires Python >= 3.11 and a Keras 3 backend (PyTorch recommended: export KERAS_BACKEND=torch).

How it works

                              ┌──────────────────────────┐
                              │   CompositeSearchSpace   │
                              │  ┌────────────────────┐  │
 Your simulator + adapter ──► │  │ Inference networks │  │
                              │  │ Summary networks   │  │ ──► Optuna Study
 Validation conditions ─────► │  │ Training params    │  │      (n_trials)
                              │  └────────────────────┘  │
                              └──────────────────────────┘
                                          │
                                 Each trial:
                                 1. Sample hyperparameters
                                 2. Budget check (reject if too large)
                                 3. Build approximator
                                 4. Compile with the mode's LR schedule
                                 5. Train
                                 6. Validate on fixed dataset
                                 7. Report metrics to Optuna

Features

• 10 network search spaces — CouplingFlow, FlowMatching, DiffusionModel, ConsistencyModel, StableConsistencyModel, DeepSet, SetTransformer, FusionTransformer, TimeSeriesNetwork, TimeSeriesTransformer
• FlowMatching solver controls — tune or fix fm_integrate_method / fm_integrate_steps plus TimeMLP kwargs (fm_merge, fm_norm, fm_residual, fm_spectral_normalization) with BayesFlow-default constants when untuned
• Network selection — let Optuna choose the best architecture via NetworkSelectionSpace / SummarySelectionSpace
• Sampler presets — TPE, GP, BoTorch, NSGA-II/III, Auto, Random with auto-wired budget constraints
• QMC warm-up — Sobol sequences for better space-filling startup coverage
• Multi-objective pruning — dominance, MO-SHA, primary-median strategies
• Metric constraints — hard post-validation rejection, soft feasibility guidance
• Coherent training modes — equal-budget cosine runs, plus open-ended inverse-sqrt runs with validation early stopping
• Derived hyperparameters — express exact simulation-budget relationships without redundant sampler dimensions
• Custom hooks — replace build, train, or validate steps while reusing the full trial lifecycle
• Pre-flight validation — check_pipeline() catches interface errors before GPU hours are wasted
• Multi-objective — single-metric, mean-aggregated, or full Pareto-front optimization (2-3 objectives)
• Rich visualization — plot_study() adaptive dashboard with pairwise Pareto projections, per-objective history, and parameter importance; plus standalone plot_pareto_3d(), plot_parallel_coordinates(), and more
• 19 built-in validation metrics — calibration, accuracy, SBC diagnostics, C2ST, plus a registry for custom metrics
• Budget constraints — reject infeasible architectures before training (max_param_count, max_memory_mb="auto")
• Custom network registration — plug in your own inference/summary networks
• Study management — resume, warm-start, save/load workflows, analyze results
• Lexicographic-Pareto selection — select_best_trial() with satisficing thresholds

Examples

See the examples/ directory for complete walkthroughs:

# Speed-sensitive FlowMatching search space
import bayesflow_hpo as hpo

inference_space = hpo.FlowMatchingSpace.fast()
search_space = hpo.CompositeSearchSpace(inference_space=inference_space)

Contributing

Each checkout needs its own virtual environment — including each git worktree, since an editable install resolves to the source tree it was installed from.

git clone git@github.com:matthiaskloft/bayesflow-hpo.git
cd bayesflow-hpo
python -m venv .venv
PY=.venv/bin/python                             # .venv/Scripts/python on Windows
$PY -m pip install torch --index-url https://download.pytorch.org/whl/cpu
$PY -m pip install -e ".[dev]"                  # [dev] pulls no backend of its own

export KERAS_BACKEND=torch
$PY -m pytest tests/ -v
$PY -m ruff check src/ tests/

License

MIT