📊 statistics

Taking the GP Out of the Loop

This paper introduces Epistemic Nearest Neighbors (ENN), a lightweight surrogate model that replaces Gaussian processes in Bayesian optimization to achieve linear $\mathcal{O}(N)$ scaling for both fitting and acquisition, thereby reducing proposal time by one to two orders of magnitude in scenarios with up to 50,000 observations.

Original authors: Mehul Bafna, Siddhant anand Jadhav, David Sweet

Published 2026-05-06

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Mehul Bafna, Siddhant anand Jadhav, David Sweet

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to find the best recipe for a cake, but you don't have a cookbook. You have to bake cakes, taste them, and learn from the results. This is what Bayesian Optimization (BO) does: it's a smart search algorithm that tries to find the "best" solution (like the perfect cake) by testing different options and learning from the results.

The Problem: The "Slow Chef"

In the past, these "baking" experiments were very expensive and slow (like waiting days for a chemical reaction). So, the algorithm was designed to be very careful, only baking a few cakes. It used a complex mathematical tool called a Gaussian Process (GP) to guess what the next cake should taste like based on the ones it already baked.

However, the GP tool is like a slow, overthinking chef. Every time you add a new cake to the history, the chef has to re-calculate the entire history from scratch to update their guess.

If you have 10 cakes, it's fast.
If you have 1,000 cakes, it takes a long time.
If you have 50,000 cakes (which is easy to do with modern fast computers), the chef gets so bogged down in calculations that the whole process grinds to a halt.

The paper calls this situation BOMO (Bayesian Optimization with Many Observations). The old method just couldn't keep up with the speed of modern simulations.

The Solution: The "Neighborhood Watch"

The authors, Mehul Bafna, Siddhant Anand Jadhav, and David Sweet, propose a new method called TuRBO-ENN. Instead of using the slow, overthinking GP chef, they use a much simpler, faster approach called Epistemic Nearest Neighbors (ENN).

Think of ENN as a neighborhood watch or a local gossip network:

How it works: When you ask, "What will this new cake taste like?", ENN doesn't re-read the entire history of 50,000 cakes. Instead, it just looks at the K closest neighbors (the most similar cakes you've already baked).
The Math: It takes the average taste of those few neighbors and guesses the new taste. It also estimates how unsure it is (uncertainty) based on how far away the new cake is from those neighbors.
The Speed: Because it only looks at a small handful of neighbors, it doesn't matter if you have 10 cakes or 50,000 cakes. The time it takes to make a guess stays roughly the same. It scales linearly, meaning it's incredibly fast even with massive amounts of data.

The Two Modes

The paper shows that this new method works in two different scenarios:

Noisy World (Real Life): Sometimes, baking the same recipe twice gives slightly different results (maybe the oven temperature fluctuated). In this case, TuRBO-ENN still does a quick calculation to figure out how much "noise" is in the data, then picks the next best cake to bake using a simple formula called UCB (which balances trying new things vs. sticking to what works).
Perfect World (Simulations): Sometimes, the computer simulation is perfect and gives the exact same result every time. In this case, TuRBO-ENN is even faster. It skips the "guessing the noise" step entirely and just sorts the options to find the best ones, saving even more time.

The Results

The authors tested this new method against the old "slow chef" (TuRBO with GP) on various problems, from landing a virtual spaceship to optimizing robot movements.

Speed: TuRBO-ENN was 10 to 100 times faster at making decisions than the old method.
Quality: Despite being much faster and simpler, it found solutions just as good as the slow, complex method.
Scale: It successfully handled up to 50,000 observations without slowing down, whereas the old method would have crashed or taken hours just to make a single decision.

The Bottom Line

The paper essentially says: "We found a way to take the heavy, slow math out of the loop." By replacing a complex, global model with a simple, local "look at your neighbors" approach, we can optimize things much faster without losing quality. They even made this tool available as a free software package for others to use.

Technical Summary: Taking the GP Out of the Loop

Problem Statement

Bayesian Optimization (BO) is traditionally employed for black-box problems where function evaluations are expensive, resulting in few observations ( $N$ ). However, a growing class of problems, termed Bayesian Optimization with Many Observations (BOMO), involves fast, parallelizable evaluations (e.g., large-scale simulations) where $N$ can reach tens of thousands.

In the BOMO regime, standard BO methods face a scalability bottleneck. Traditional approaches rely on Gaussian Process (GP) surrogates. While modern implementations have reduced the acquisition time complexity from $O(N^3)$ to roughly $O(N^2)$ , the requirement to fit GP hyperparameters at every iteration remains a significant computational burden. This fitting process, often involving the reconstruction of the kernel matrix, scales poorly with $N$ , impeding the application of BO to problems with massive observation counts.

Methodology

The authors propose TuRBO-ENN, a method that replaces the GP surrogate in the Trust Region Bayesian Optimization (TuRBO) framework with a lightweight alternative called Epistemic Nearest Neighbors (ENN).

1. Epistemic Nearest Neighbors (ENN)

ENN is a K-nearest-neighbor (KNN) based surrogate designed to estimate both function values and uncertainty (epistemic and aleatoric) with linear scaling.

Surrogate Construction: For a query point $x$ , ENN identifies the $K$ nearest neighbors in the observation set $D$ . It treats each neighbor as an independent estimate of the function value.
Uncertainty Quantification:
- Aleatoric Variance ( $\sigma_a^2$ ): Modeled as the sum of a base noise level ( $s_0^2$ ) and the specific observation's noise scale ( $s_m^2$ ).
- Epistemic Variance ( $\sigma_e^2$ ): Modeled as proportional to the squared Euclidean distance between the query point and the neighbor ( $c_e d^2(x, x_m)$ ).
Aggregation: The final mean estimate $\mu(x)$ and variance estimates are computed using a precision-weighted average (inverse variance weighting) of the $K$ neighbors. This avoids the $O(N^2)$ pairwise covariance calculations required by GPs.
Complexity: Finding $K$ nearest neighbors using a max-heap or Faiss library takes $O(N \ln K)$ per query. The fitting of hyperparameters ( $s_0, c_e$ ) is performed via Type-II Maximum Likelihood Estimation (MLE) on a Leave-One-Out Cross-Validation (LOOCV) pseudolikelihood, approximated using a fixed-size subsample to maintain $O(N)$ scaling.

2. TuRBO-ENN Integration

The method integrates ENN into the TuRBO algorithm, which utilizes a trust region to focus search efforts.

Noisy Objectives: For problems with noise, the method fits ENN hyperparameters and uses the Upper Confidence Bound (UCB) acquisition function: $UCB(x) = \mu(x) + \sigma(x)$ .
Noise-Free Objectives: For deterministic problems, the authors propose a fitting-free variant. Since hyperparameter fitting is unnecessary when noise is known to be zero, the method selects candidates via Non-Dominated Sorting (NDS) over the pair $(\mu(x), \sigma(x))$ . This eliminates the fitting step entirely, further reducing proposal time.
Pseudo-Bayesian Optimization (PBO): The authors verify that TuRBO-ENN satisfies the conditions for Pseudo-Bayesian Optimization (Chen & Lam, 2025), ensuring theoretical convergence guarantees.

Key Contributions

ENN Surrogate: Introduction of a lightweight KNN surrogate that estimates function values and uncertainty with $O(N)$ fitting and querying time, contrasting with the $O(N^2)$ or $O(N^3)$ scaling of GPs.
TuRBO-ENN Algorithm: Development of a BO method that replaces TuRBO's GP and Thompson Sampling with ENN and UCB (or NDS for noise-free cases).
Empirical Validation: Demonstration that TuRBO-ENN reduces proposal time (fitting + acquisition) by one to two orders of magnitude compared to TuRBO-GP on problems with up to 50,000 observations, without sacrificing solution quality.
Theoretical Guarantee: Verification that TuRBO-ENN qualifies as a PBO method, inheriting convergence guarantees.
Open Source: Release of the ennbo Python package.

Experimental Results

The authors benchmarked TuRBO-ENN against state-of-the-art methods (TuRBO-GP, Optuna, SMAC, CMA-ES, DNGO, Vecchia GP) on five simulation tasks (LunarLander, Hopper, BipedalWalker, Push) and a hyperparameter optimization task (LASSO-DNA) with dimensions ranging from 12 to 180.

Scaling: While GP-based methods (e.g., turbo-one, ucb) exhibited approximately $O(N^2)$ scaling in proposal time, TuRBO-ENN scaled linearly ( $O(N)$ ).
Speedup: TuRBO-ENN achieved speedups ranging from 8.5x to 151x over TuRBO-GP, with larger gains observed at higher observation counts ( $N$ ) and in frozen-noise (noise-free) settings where hyperparameter fitting was omitted.
Solution Quality: TuRBO-ENN matched the optimization performance (solution quality) of TuRBO-GP across all tested tasks.
Fitting Tasks: In pure regression tasks (synthetic data), ENN was orders of magnitude faster to fit than Exact GPs, SVGP, and neural network surrogates, while maintaining competitive test-set log-likelihoods.

Significance and Claims

The paper argues that as simulation capabilities grow, the bottleneck in Bayesian Optimization shifts from function evaluation cost to the computational cost of the optimization algorithm itself (proposal time). By removing the Gaussian Process from the loop, TuRBO-ENN enables BO to scale to regimes with tens of thousands of observations where traditional methods become computationally prohibitive.

The authors modestly claim that their method provides a practical, scalable alternative for BOMO problems, offering a "fitting-free" option for deterministic simulations and a linear-scaling option for noisy ones, all while maintaining the theoretical convergence properties of Bayesian optimization. They note that while they achieve convergence guarantees via PBO axioms, they do not provide specific regret bounds. Future work directions suggested include adapting ENN to approximate nearest neighbor algorithms for $O(\ln N)$ scaling and exploring dimension-specific weighting in the distance metric.