Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

VBO-MI: A Fully Gradient-Based Bayesian Optimization Framework Using Variational Mutual Information Estimation

Farhad Mirkarimi

TL;DR

VBO-MI tackles expensive black-box optimization by removing GP/posterior sampling bottlenecks and replacing them with a fully gradient-based, variational mutual information framework. It employs an actor-critic architecture where an action-net generates candidate inputs and a neural MI estimator guides exploration via a DV-based MI bound, enabling end-to-end training and a significant reduction in FLOPs compared to BNNS-based surrogates. The method demonstrates competitive or superior optimization performance on high-dimensional synthetic benchmarks and complex real-world tasks (e.g., PDE optimization, Lunar Lander, Pest Control) while offering robustness to hyperparameter choices. This approach broadens scalable BO by combining flexible variational uncertainty modeling with principled information-theoretic exploration in a differentiable, gradient-friendly pipeline.

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.

VBO-MI: A Fully Gradient-Based Bayesian Optimization Framework Using Variational Mutual Information Estimation

TL;DR

VBO-MI tackles expensive black-box optimization by removing GP/posterior sampling bottlenecks and replacing them with a fully gradient-based, variational mutual information framework. It employs an actor-critic architecture where an action-net generates candidate inputs and a neural MI estimator guides exploration via a DV-based MI bound, enabling end-to-end training and a significant reduction in FLOPs compared to BNNS-based surrogates. The method demonstrates competitive or superior optimization performance on high-dimensional synthetic benchmarks and complex real-world tasks (e.g., PDE optimization, Lunar Lander, Pest Control) while offering robustness to hyperparameter choices. This approach broadens scalable BO by combining flexible variational uncertainty modeling with principled information-theoretic exploration in a differentiable, gradient-friendly pipeline.

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 { 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.
Paper Structure (24 sections, 2 theorems, 50 equations, 8 figures, 6 tables, 1 algorithm)

This paper contains 24 sections, 2 theorems, 50 equations, 8 figures, 6 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Notation.
  3. Related work
  4. Scalable Bayesian optimization
  5. Variational and Sampling-Based BNNs.
  6. Information-Theoretic acquisition function desgin
  7. Problem Setting and Background
  8. Bayesian Optimization.
  9. Gaussian BO
  10. Exploration-Exploitation trade-off
  11. Proposed Bayesian Optimization Framework with Variational MI Estimation
  12. Variational Mutual Information
  13. Exploration and Mutual Information Estimation
  14. Loss Function and Algorithm
  15. Numerical evaluation
  16. ...and 9 more sections

Key Result

Theorem A.1

Generalized AEP(Generalized Asymptotic Equipartition Property for Markov Approximation). AEPG Let $\mathcal{H}$ be a standard Borel space, and consider the infinite product space $\mathcal{H}_0^\infty = \prod_{t=0}^\infty \mathcal{H}$, equipped with its canonical Borel $\sigma$-algebra $(\Omega, \ma

Figures (8)

  • Figure 1: NN architecture for evaluating MI from samples of two distributions. Using the network output and Eq. \ref{['miestimate']}, an estimate of Mutual Information is computed.
  • Figure 2: An illustration of the neural architecture of the proposed algorithm. Each dashed line shows the flow of gradients for training the weights of the action-net (main) and helper networks. Concat and shuffle ($\pi$) blocks are used to produce the joint and product of marginal samples to feed the helper network.
  • Figure 3: Comparison of the proposed method with various baselines from li2024aNEURIPS2024_da30215e over different synthetic benchmark test functions.
  • Figure 4: Comparison of the proposed method with various baselines from li2024aNEURIPS2024_da30215e over different real-world tasks.
  • Figure 5: Impact of altering exploration and exploitation terms in Eq. \ref{['loss-f']} by substituting a GP posterior component and comparing performance with VBO.
  • ...and 3 more figures

Theorems & Definitions (2)

  • Theorem A.1
  • Theorem A.2