[Paper Review] Scalable and Flexible Deep Bayesian Optimization with Auxiliary Information for Scientific Problems.

Samuel Kim, Peter Y. Lu+4 more|arXiv (Cornell University)|Apr 23, 2021

Machine Learning in Materials Science7 references4 citations

TL;DR

This paper proposes using Bayesian Neural Networks (BNNs) as scalable, flexible surrogate models in Bayesian optimization (BO) for scientific problems with structured data and auxiliary information. By leveraging BNNs' ability to model complex patterns and incorporate known symmetries or intermediate data, the method achieves superior sampling efficiency and lower computational cost than Gaussian Processes (GPs) on physics and chemistry optimization tasks, including photonic crystal topology and molecular property optimization.

ABSTRACT

Bayesian optimization (BO) is a popular paradigm for global optimization of expensive black-box functions, but there are many domains where the function is not completely black-box. The data may have some known structure, e.g. symmetries, and the data generation process can yield useful intermediate or auxiliary information in addition to the value of the optimization objective. However, surrogate models traditionally employed in BO, such as Gaussian Processes (GPs), scale poorly with dataset size and struggle to incorporate known structure or auxiliary information. Instead, we propose performing BO on complex, structured problems by using Bayesian Neural Networks (BNNs), a class of scalable surrogate models that have the representation power and flexibility to handle structured data and exploit auxiliary information. We demonstrate BO on a number of realistic problems in physics and chemistry, including topology optimization of photonic crystal materials using convolutional neural networks, and chemical property optimization of molecules using graph neural networks. On these complex tasks, we show that BNNs often outperform GPs as surrogate models for BO in terms of both sampling efficiency and computational cost.

Motivation & Objective

• To address the limitations of Gaussian Processes (GPs) in Bayesian optimization (BO) for large-scale, structured scientific problems.
• To enable effective use of known structural properties (e.g. symmetries) and auxiliary information in the optimization process.
• To develop a scalable and flexible surrogate modeling approach that maintains uncertainty quantification while handling complex data like graphs and images.
• To improve sampling efficiency and reduce computational cost in BO for real-world scientific applications.
• To demonstrate the method’s effectiveness on challenging problems in physics and chemistry, such as photonic crystal design and molecular property optimization.

Proposed method

• Using Bayesian Neural Networks (BNNs) as surrogate models in Bayesian optimization, replacing traditional Gaussian Processes (GPs).
• Leveraging the representation power of BNNs to model complex, structured data such as convolutional grids (e.g. photonic crystals) and graph-structured data (e.g. molecules).
• Incorporating auxiliary information—such as intermediate outputs or structural symmetries—directly into the BNN architecture during training.
• Applying active learning acquisition functions (e.g. Expected Improvement) with uncertainty estimates from the BNN to guide the optimization process.
• Training BNNs with variational inference to maintain uncertainty quantification while enabling scalability to larger datasets.
• Using structured neural network architectures (e.g. CNNs for spatial data, GNNs for molecular graphs) to encode domain-specific inductive biases.

Experimental results

Research questions

• RQ1 Can Bayesian Neural Networks outperform Gaussian Processes as surrogate models in Bayesian optimization for scientific problems?
• RQ2 How effectively can BNNs incorporate known structural properties and auxiliary information in optimization tasks?
• RQ3 Does the use of BNNs lead to improved sampling efficiency and lower computational cost compared to GPs in real-world scientific optimization?
• RQ4 To what extent can BNNs scale to large datasets while preserving uncertainty estimates in Bayesian optimization?
• RQ5 Can BNN-based BO achieve better performance on complex scientific problems such as photonic crystal topology and molecular property optimization?

Key findings

• BNNs outperform GPs in sampling efficiency on complex scientific optimization tasks, reducing the number of function evaluations required to find optimal solutions.
• The use of BNNs enables effective integration of auxiliary information and structural priors, such as symmetries, into the optimization process.
• BNNs scale more effectively than GPs to larger datasets, maintaining computational feasibility for real-world scientific applications.
• On photonic crystal topology optimization using convolutional neural networks, BNN-based BO achieved faster convergence and better performance than GP-based BO.
• In molecular property optimization using graph neural networks, BNNs demonstrated superior sample efficiency and lower computational cost compared to GPs.
• The method maintains uncertainty quantification through variational inference, enabling reliable exploration-exploitation trade-offs in BO.