Improving Hyperparameter Optimization with Checkpointed Model Weights

TL;DR

This paper introduces Forecasting Model Search (FMS), a novel hyperparameter optimization method that leverages logged checkpointed model weights and a deep kernel surrogate to improve efficiency in neural network training.

Contribution

FMS is the first HPO approach to incorporate logged checkpointed weights using a permutation-invariant graph metanetwork for data-efficient hyperparameter tuning.

Findings

FMS outperforms traditional HPO methods in efficiency.

The method effectively utilizes checkpointed weights for better hyperparameter predictions.

Open-source implementation promotes reproducibility and further research.

Abstract

When training deep learning models, the performance depends largely on the selected hyperparameters. However, hyperparameter optimization (HPO) is often one of the most expensive parts of model design. Classical HPO methods treat this as a black-box optimization problem. However, gray-box HPO methods, which incorporate more information about the setup, have emerged as a promising direction for more efficient optimization. For example, using intermediate loss evaluations to terminate bad selections. In this work, we propose an HPO method for neural networks using logged checkpoints of the trained weights to guide future hyperparameter selections. Our method, Forecasting Model Search (FMS), embeds weights into a Gaussian process deep kernel surrogate model, using a permutation-invariant graph metanetwork to be data-efficient with the logged network weights. To facilitate reproducibility and further research, we open-source our code at https://github.com/NVlabs/forecasting-model-search.