[Paper Review] High-dimensional Bayesian Optimization of Hyperparameters for an Attention-based Network to Predict Materials Property: a Case Study on CrabNet using Ax and SAASBO

Sterling G. Baird, Marianne Liu+1 more|arXiv (Cornell University)|Mar 19, 2022

Machine Learning in Materials Science4 citations

TL;DR

This study applies high-dimensional Bayesian optimization using SAASBO and GPEI within the Ax platform to tune 23 hyperparameters of CrabNet, an attention-based materials property predictor. SAASBO achieves a new state-of-the-art on the Matbench experimental band gap regression task, reducing mean absolute error by 4.5% compared to the incumbent model.

ABSTRACT

Expensive-to-train deep learning models can benefit from an optimization of the hyperparameters that determine the model architecture. We optimize 23 hyperparameters of a materials informatics model, Compositionally-Restricted Attention-Based Network (CrabNet), over 100 adaptive design iterations using two models within the Adaptive Experimentation (Ax) Platform. This includes a recently developed Bayesian optimization (BO) algorithm, sparse axis-aligned subspaces Bayesian optimization (SAASBO), which has shown exciting performance on high-dimensional optimization tasks. Using SAASBO to optimize CrabNet hyperparameters, we demonstrate a new state-of-the-art on the experimental band gap regression task within the materials informatics benchmarking platform, Matbench (~4.5% decrease in mean absolute error (MAE) relative to incumbent). Characteristics of the adaptive design scheme as well as feature importances are described for each of the Ax models. SAASBO has great potential to both improve existing surrogate models, as shown in this work, and in future work, to efficiently discover new, high-performing materials in high-dimensional materials science search spaces.

Motivation & Objective

• To investigate whether hyperparameter optimization can further improve a model already at state-of-the-art performance.
• To evaluate the effectiveness of high-dimensional Bayesian optimization, particularly SAASBO, in tuning deep learning models for materials property prediction.
• To assess whether optimized hyperparameters generalize across different materials properties or data modalities.
• To establish a benchmark for hyperparameter optimization in materials informatics using a computationally expensive model.

Proposed method

• Utilizes the Ax Bayesian optimization platform with two models: Gaussian Process Expected Improvement (GPEI) and SAASBO.
• Applies SAASBO, a high-dimensional BO algorithm that uses sparse axis-aligned subspaces and strong priors on inverse lengthscales to prevent overfitting.
• Employs a nested 5-fold cross-validation scheme on the Matbench matbench_expt_gap dataset with 4,604 composition-based data points.
• Optimizes 23 hyperparameters of CrabNet, including learning rate, number of layers, dropout rate, and number of training epochs.
• Uses expected improvement as the acquisition function in GPEI and leverages automatic relevance determination for feature importance estimation.
• Conducts 100 adaptive design iterations, starting with Sobol sequences for initial exploration before BO-driven optimization.

Experimental results

Research questions

• RQ1 To what extent can hyperparameter optimization improve a model already at state-of-the-art performance?
• RQ2 How does SAASBO perform in high-dimensional hyperparameter spaces for materials informatics models compared to classical BO?
• RQ3 Which hyperparameters are most influential in improving CrabNet's predictive accuracy on the experimental band gap task?
• RQ4 Can the optimized hyperparameters generalize to other materials properties or data types beyond experimental band gap?

Key findings

• SAASBO achieved a 4.5% relative reduction in mean absolute error (MAE) on the Matbench experimental band gap regression task, setting a new state-of-the-art.
• The optimized number of training epochs reached the maximum allowed value (80) in all five cross-validation folds, indicating that longer training is beneficial and not limited by early stopping.
• SAASBO demonstrated a strong preference for exploration over exploitation in later iterations, suggesting robustness to overfitting in high-dimensional spaces.
• Feature importance analysis revealed that the number of epochs was the most influential hyperparameter, followed by learning rate and dropout rate.
• The adaptive design scheme with Sobol initialization followed by BO iterations effectively explored the high-dimensional hyperparameter space.
• The results validate SAASBO as a powerful tool for efficient hyperparameter optimization in materials informatics, with potential for broader application in materials discovery.