We use Ax Bayesian optimization to adjust hyperparameters for (my fork of) CrabNet for the experimental band gap matbench task (matbench_expt_gap). We chose this task because it is a composition-only dataset (CrabNet is a composition-only model) and because CrabNet is currently (2021-01-08) listed at the top of this leaderboard (with MODNet just marginally worse). In other words, when a model whose defaults already produce state-of-the-art property prediction performance, to what extent can it benefit from hyperparameter optimization (i.e. tuning parameters such as Neural Network dimensions, learning rates, etc.)?
As of 2022-04-05, Ax/SAASBO CrabNet v1.2.7 holds the current per-task leaderboard on the matbench_expt_gap task (c3b910e4f0). For additional details, please consider reading the CMS article or the preprint. For more information on CrabNet's architecture, see the original CrabNet paper published in Nature Partner Journals: Computational Materials.
This case study can serve as an illustrative example of hyperparameter optimization using Bayesian adaptive design and could certainly be adapted to other models (e.g. Neural Networks), especially expensive-to-train models that have not undergone much by way of parameter tuning. See hyperparameterization.ipynb for a more in-depth walkthrough of the process and results.
- figures
- Ax experiment JSON files
- requirements.txt
- hyperparameterization.py (
.pyscript adapted into the Jupyter notebook mentioned above)
Eventually, I plan to incorporate this into (my fork of) CrabNet, but for now this can serve as an illustrative example of hyperparameter optimization using Bayesian adaptive design and could certainly be adapted to other models (e.g. Neural Networks), especially expensive-to-train models that have not undergone much by way of parameter tuning.
If you find this useful, please consider citing:
Baird, S. G.; Liu, M.; Sparks, T. D. High-Dimensional Bayesian Optimization of 23 Hyperparameters over 100 Iterations for an Attention-Based Network to Predict Materials Property: A Case Study on CrabNet Using Ax Platform and SAASBO. Computational Materials Science 2022, 211, 111505. https://doi.org/10.1016/j.commatsci.2022.111505.
In addition to the above manuscript citation, if you use this code, please also cite the following for all versions (alternatively, a specific version):
Sterling Baird, sgbaird-alt, & mliu7051. (2022). sparks-baird/crabnet-hyperparameter. Zenodo. https://doi.org/10.5281/zenodo.6355044
"Optimizing Training Trajectories in Variational Autoencoders via Latent Bayesian Optimization Approach" (https://arxiv.org/abs/2207.00128). It would be interesting to see whether Bayesian optimization across a latent space of hyperparameters or SAASBO would be more efficient for various problems. Of course, SAASBO is a type of BO, too, so technically SAASBO could be used to explore the hyperparameter latent space as well.