Towards Scalable Bayesian Optimization via Gradient-Informed Bayesian Neural Networks

2025-04-14

AI Summary

To address the slow convergence and poor scalability of Bayesian optimization (BO) for high-dimensional black-box function optimization, this paper proposes a Gradient-Enhanced Bayesian Neural Network (GBNN) as a surrogate model. We introduce gradient observations explicitly into the Bayesian neural network (BNN) training framework for the first time, designing a gradient-aware variational loss function to overcome the fundamental limitation that conventional BNNs cannot incorporate derivative information in BO. Our method integrates automatic differentiation with gradient-enhanced variational inference, preserving principled uncertainty quantification while significantly improving predictive accuracy. Experimental results on standard benchmarks demonstrate that GBNN substantially enhances surrogate model fidelity; in high-dimensional settings, it reduces BO convergence steps by 30–50%, validating its superior scalability and optimization efficiency.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Optimization
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Bayesian optimization (BO) is a widely used method for data-driven optimization that generally relies on zeroth-order data of objective function to construct probabilistic surrogate models. These surrogates guide the exploration-exploitation process toward finding global optimum. While Gaussian processes (GPs) are commonly employed as surrogates of the unknown objective function, recent studies have highlighted the potential of Bayesian neural networks (BNNs) as scalable and flexible alternatives. Moreover, incorporating gradient observations into GPs, when available, has been shown to improve BO performance. However, the use of gradients within BNN surrogates remains unexplored. By leveraging automatic differentiation, gradient information can be seamlessly integrated into BNN training, resulting in more informative surrogates for BO. We propose a gradient-informed loss function for BNN training, effectively augmenting function observations with local gradient information. The effectiveness of this approach is demonstrated on well-known benchmarks in terms of improved BNN predictions and faster BO convergence as the number of decision variables increases.

Problem

• Scalable Bayesian optimization using gradient-informed BNNs
• Improving surrogate models with gradient information in BNNs
• Enhancing BO convergence speed with gradient-augmented BNN training

Innovation

• Gradient-informed Bayesian Neural Networks for BO
• Automatic differentiation integrates gradient data
• Enhanced BNN training with gradient loss function