Curvature-aware Expected Free Energy for Bayesian Optimization

An acquisition function that optimizes a black-box function and learns it at the same time.

A. Anil Meera and W. Kouw, "Curvature-aware Expected Free Energy as an Acquisition Function for Bayesian Optimization", IEEE Control Systems Letters, 2026. ieeexplore.ieee.org/document/11589289

Fifty Bayesian optimization steps under seven acquisition functions, sharing one objective, surrogate and initial design. EFE is the only one that ends with both an accurate surrogate and the optimum found.

Method

Standard acquisition functions optimize at the expense of learning. Expected Free Energy scores a candidate query by a single quantity, to be minimised:

$$ G(x) = \frac{\big(\mu(x) - y^\star\big)^2 + \sigma^2(x)}{2,\tau^2(x)} - \frac{1}{2}\ln\left(1 + \frac{\sigma^2(x)}{\sigma_n^2}\right) $$

for a Gaussian process with posterior mean $\mu$ and standard deviation $\sigma$. The first term is pragmatic: it pulls the search towards a preferred outcome $y^\star$, held as a goal prior of variance $\tau^2$. The second is epistemic: it rewards queries that reduce uncertainty about the function.

The paper's contribution is a curvature-aware goal prior,

$$ \frac{1}{\tau_i^2(x)} = \big|\mu''(x)\big| + \frac{1}{\sigma^2(x)}, \qquad \tau^2(x) = \tau_{\min}^2 + \left(\tau_{\max}^2 - \tau_{\min}^2\right)\frac{\tau_i^2(x)}{\max_x \tau_i^2(x)} $$

which narrows where the surrogate is sharply curved and already well determined, and widens over flat, uncertain regions. Exploration is balanced against exploitation without a hand-tuned schedule.

Usage

pip install numpy scipy scikit-learn matplotlib pillow
python EFE_aq_BO.py          # compare all acquisitions and plot the results
python EFE_aq_BO.py --gif    # ... and write the animation above

Python 3.10 or newer, CPU only, under a minute per run. All of the code is in EFE_aq_BO.py; see --help for the options.

Citation

@article{anilmeera2026efe,
  author  = {Anil Meera, Ajith and Kouw, Wouter},
  title   = {Curvature-aware Expected Free Energy as an Acquisition Function
             for Bayesian Optimization},
  journal = {IEEE Control Systems Letters},
  year    = {2026},
  url     = {https://ieeexplore.ieee.org/document/11589289}
}