🤖 AI Summary
This work proposes an end-to-end gradient-driven Bayesian optimization framework to address the high computational cost associated with posterior sampling and acquisition function optimization in traditional Bayesian neural network–based approaches. By introducing variational mutual information estimation into Bayesian optimization for the first time and integrating it within an actor-critic architecture, the method jointly optimizes input exploration and information gain assessment. This design eliminates the inner-loop acquisition function optimization, yielding a fully differentiable and computationally efficient optimization pipeline. Empirical evaluations demonstrate that the proposed approach achieves performance comparable to or better than existing baselines across multiple high-dimensional synthetic and real-world tasks, while reducing computational overhead by up to two orders of magnitude.
📝 Abstract
Many real-world tasks require optimizing expensive black-box functions accessible only through noisy evaluations, a setting commonly addressed with Bayesian optimization (BO). While Bayesian neural networks (BNNs) have recently emerged as scalable alternatives to Gaussian Processes (GPs), traditional BNN-BO frameworks remain burdened by expensive posterior sampling and acquisition function optimization. In this work, we propose {VBO-MI} (Variational Bayesian Optimization with Mutual Information), a fully gradient-based BO framework that leverages recent advances in variational mutual information estimation. To enable end-to-end gradient flow, we employ an actor-critic architecture consisting of an {action-net} to navigate the input space and a {variational critic} to estimate information gain. This formulation effectively eliminates the traditional inner-loop acquisition optimization bottleneck, achieving up to a {$10^2 \times$ reduction in FLOPs} compared to BNN-BO baselines. We evaluate our method on a diverse suite of benchmarks, including high-dimensional synthetic functions and complex real-world tasks such as PDE optimization, the Lunar Lander control problem, and categorical Pest Control. Our experiments demonstrate that VBO-MI consistently provides the same or superior optimization performance and computational scalability over the baselines.