Two-step hyperparameter optimization method: Accelerating hyperparameter search by using a fraction of a training dataset
Sungduk Yu, Mike Pritchard, Po-Lun Ma, Balwinder Singh, Sam Silva
TL;DR
The paper addresses the high computational burden of hyperparameter optimization in data-intensive earth-system ML by proposing a universal two-step HPO framework that first evaluates configurations on a small subset of data and then retrains the top candidates on the full dataset. The method, demonstrated on neural network emulators for aerosol activation, shows that effective hyperparameters can be identified using as little as a few thousand samples, with substantial speedups (up to ~135×) while preserving or improving final performance after retraining. A detailed case study, parameter-setting guidance, and a distributed GPU-accelerated implementation illustrate how the approach can scale across data volumes and be integrated into climate modeling workflows. The work highlights the practical value of core-set ideas for efficient HPO and provides actionable guidance for practitioners seeking high-performance models with reduced compute costs in hybrid ML-physics climate simulations.
Abstract
Hyperparameter optimization (HPO) is an important step in machine learning (ML) model development, but common practices are archaic -- primarily relying on manual or grid searches. This is partly because adopting advanced HPO algorithms introduces added complexity to the workflow, leading to longer computation times. This poses a notable challenge to ML applications, as suboptimal hyperparameter selections curtail the potential of ML model performance, ultimately obstructing the full exploitation of ML techniques. In this article, we present a two-step HPO method as a strategic solution to curbing computational demands and wait times, gleaned from practical experiences in applied ML parameterization work. The initial phase involves a preliminary evaluation of hyperparameters on a small subset of the training dataset, followed by a re-evaluation of the top-performing candidate models post-retraining with the entire training dataset. This two-step HPO method is universally applicable across HPO search algorithms, and we argue it has attractive efficiency gains. As a case study, we present our recent application of the two-step HPO method to the development of neural network emulators for aerosol activation. Although our primary use case is a data-rich limit with many millions of samples, we also find that using up to 0.0025% of the data (a few thousand samples) in the initial step is sufficient to find optimal hyperparameter configurations from much more extensive sampling, achieving up to 135-times speedup. The benefits of this method materialize through an assessment of hyperparameters and model performance, revealing the minimal model complexity required to achieve the best performance. The assortment of top-performing models harvested from the HPO process allows us to choose a high-performing model with a low inference cost for efficient use in global climate models (GCMs).