Diversity-Aware Bayesian Optimization for Training Dataset Generation using Expensive Simulators

Large, representative training datasets are essential for developing state-of-the-art machine learning models. The quality and composition of training data can have a substantial impact on model performance, while generating and curating these datasets can be computationally expensive. We present a novel approach for strategically generating training data with expensive physics simulators using Bayesian optimization, targeting regions of the input space that are more likely to yield measurable phenomena (e.g. exceeding a given signal to noise threshold) while maintaining broad coverage of the plausible observable outcomes. Specifically, we use a Gaussian process to predict the posterior distribution of measurable properties given a set input parameters using existing simulations. These properties are then used to estimate the probability that an event will be detected, and this probability is optimized using expected improvement for composite functions (EI-CF) using an active search utility. We further impose a diversity constraint on the measurable properties using the ROBOT algorithm to prevent the optimizer from concentrating primarily on higher signal-to-noise ratio phenomena that are associated with a higher active search utility. We find that our EI-CF+ROBOT method performs better than EI+ROBOT and random sampling for both measurability and diversity metrics.