Batch Acquisition Function Evaluations and Decouple Optimizer Updates for Faster Bayesian Optimization

AI summary

Overview

Research area: Bayesian optimization (BO), specifically the acquisition-function optimization step and its computational cost.

Technical level: Intermediate. The paper assumes familiarity with Bayesian optimization, quasi-Newton (QN) methods, inverse-Hessian approximations, and coroutines, but its argument is built around a single intuitive idea that a non-specialist can follow.

Scope: The paper diagnoses why the widely used practice of batching multi-start quasi-Newton optimization in Bayesian optimization slows convergence, and proposes a decoupled alternative that keeps batched evaluations.

What This Paper Is About

Bayesian optimization needs to maximize an acquisition function, and because that function is non-convex, implementations such as BoTorch and Optuna run multi-start optimization (MSO) with quasi-Newton methods. BoTorch's standard trick is to optimize the sum of the acquisition values from all restarts at once, which lets PyTorch batch the acquisition function evaluations and thereby speeds up MSO. This paper shows empirically that the summed formulation introduces approximation errors in the off-diagonal blocks of the inverse Hessian, which slows the QN solver's convergence and can cancel out or reverse the batching speedup; the authors then propose decoupling the QN updates per restart while keeping the batched evaluations.

Key Contributions

1. Identifies and characterizes a pitfall of coupled updates with batched evaluations (C-BE): approximation errors in the cross derivatives, or off-diagonal blocks, of the inverse Hessian, which the authors call "off-diagonal artifacts," distort search directions and slow convergence.
2. Proposes D-BE (Decoupling QN updates per restart while Batching acquisition function Evaluations), realized through a novel combination of a coroutine and batching, which sidesteps off-diagonal artifacts without requiring a new solver and without degrading solution quality.
3. Demonstrates experimentally that D-BE is faster than both sequential optimization (Seq. Opt.) and C-BE, and notes that the paper's approach has been successfully merged into GPSampler in Optuna, significantly accelerating its runtime.

Main Findings

• Off-diagonal artifacts are structural, not a memory artifact. The true Hessian of the summed acquisition function over B restarts is block-diagonal with zero cross-partial derivatives between different restarts. Structure-oblivious QN updates in the resulting BD-dimensional space maintain a dense inverse-Hessian approximation, injecting non-zero values into off-diagonal blocks. Contour plots on the Rosenbrock function (B=3, D=5, x in [0,3]^D, L-BFGS-B with memory m=10) show near-zero off-diagonal blocks under Seq. Opt. but dense off-diagonal blocks under C-BE.
• The artifacts translate into much slower convergence. On Rosenbrock (D=5, x in [0,3]^D, L-BFGS-B, m=10), the median reaches an objective value of 10^-12 in approximately 30 iterations for Seq. Opt. (B=1), whereas C-BE with B=2 needs about 50 iterations and B in {5,10} needs more than 120 iterations. MSO from 10 initial points would therefore require about 300 (=30 x 10) point evaluations for B=1 versus more than 1200 (=120 x 10) for B=10.
• The effect is not specific to limited memory. Appendix B repeats the analysis with full BFGS (unlimited memory) and finds C-BE still injects off-diagonal mass and still requires substantially more iterations as B grows (Figures 3–5), confirming the artifacts come from coupling QN updates for a separable problem rather than from limiting memory.
• On BO benchmarks, D-BE matches sequential iteration counts. On the Rastrigin function (300 trials, L-BFGS-B with m=10, B=10 restarts, D = 5, 10, 20, 40, termination at 200 iterations or ||∇α(x)||_∞ ≤ 10^-2), C-BE's L-BFGS-B iteration counts inflated by roughly 3.2x at D=20 (68.8 vs. 21.5) and 3.3x at D=40 (87.0 vs. 26.5) relative to Seq. Opt., whereas D-BE matched Seq. Opt.'s iteration counts while exploiting batched gradient evaluations.
• Median wall-clock speedups. The abstract reports up to 1.5x and 1.1x wall-clock speedups over Seq. Opt. and C-BE, respectively. The full Appendix C table reports D-BE attaining the shortest median runtime in nearly all objective-dimension pairs (14/16), up to 1.76x faster than Seq. Opt. (5D Attractive Sector, 127.8 s vs. 72.5 s) and 1.29x faster than C-BE (40D Attractive Sector, 139.5 s vs. 108.1 s).
• Higher-dimensional degradation on additional objectives. Across Sphere, Attractive Sector (AS), Step Ellipsoidal (SE), and Rastrigin at D=40, C-BE's Best Value was worse than D-BE's by +25.4% on AS, +16.5% on SE, +31.3% on Rastrigin, and +2.8% on Sphere (computed as (C-BE − D-BE)/D-BE). At D=40, C-BE required about 3.3x to 29x more L-BFGS-B iterations than D-BE, with the largest gap on Attractive Sector (102.2 vs. 3.5).
• Quality is preserved at the reported iteration cap. In Table 1 all methods achieved comparable median final objective values under the shared iteration cap, but C-BE paid for this either in extra iterations or, in principle, in degraded solution quality when the iteration count is capped.
• C-BE's occasional wins are explained by termination behavior. C-BE sometimes outperformed Seq. Opt. for D in {10, 20} on Rastrigin, which the authors attribute to differences in convergence behavior: the C-BE acquisition-function optimization is mostly terminated by the gradient tolerance rather than the iteration count, so its results are not premature and can differ substantially from Seq. Opt.'s.
• Prior art covers only part of the problem. BFGS and L-BFGS have been extended to handle this structure (Griewank and Toint; Bigeon et al.), but no principled way exists for L-BFGS-B, and to the authors' knowledge there is no practical block-structure-aware, bound-constrained QN algorithm for this setting. BoTorch v0.15.0 changed from C-BE to D-BE independently of this work, though it does not use a coroutine.

Methodology in Plain English

The authors start by comparing two ways of running multi-start optimization. The first, Seq. Opt., runs B completely independent optimizations one after another, each with its own optimizer state; it is accurate but pays for every acquisition-function evaluation separately. The second, C-BE, stacks all B points into one big matrix and optimizes the sum of their acquisition values with a single QN solver. Because the sum is additively separable, the gradients per point are identical to the independent case, so PyTorch can evaluate all B points in one batched call. The catch is that a QN solver working in BD dimensions builds one shared inverse-Hessian approximation for the whole stack, and its off-diagonal blocks — which should be zero because restarts do not interact — fill up with spurious values that corrupt the search directions.

To fix this, the authors propose D-BE: keep batching the acquisition function and gradient evaluations, but give each restart its own independent QN state. Naively this means running B solvers sequentially, which loses the throughput. Their solution is a coroutine: a batch evaluator performs one batched call, dispatches each restart's value and gradient to its own worker, and each worker updates its own QN state before yielding back. This avoids writing a new solver — important because libraries such as SciPy expose L-BFGS-B but not per-iteration hooks, and PyTorch's incompatibility with multi-processing makes ordinary parallelism awkward. The design also lets D-BE track a set of still-active restarts and prune converged ones, progressively shrinking the batch. The key claim is that under identical initialization and termination policies, each restart in D-BE theoretically reproduces the per-restart trajectory of Seq. Opt.; the authors note that automatic differentiation can yield slightly different gradients between a batched and a single evaluation due to floating-point nondeterminism, so trajectories may not be bit-identical.

The cost argument for why this is worthwhile: a single expected-improvement evaluation using a Gaussian process regressor with n training points costs O(n² + nD), so B points cost O(B(n² + nD)), while an L-BFGS-B update with memory m for B independent optimizations costs only O(BmD). Since n is much larger than m in practice (typically after tens to hundreds of trials, with m in [5, 20]), it pays to batch only the evaluations.

The empirical evaluation uses the COCO (BBOB) suite functions available in OptunaHub, starting with Rastrigin. Each BO iteration fits a Gaussian process regressor with a Matérn ν = 5/2 kernel to the observed parameter–objective pairs and selects the next point by maximizing log expected improvement (LogEI) via MSO with L-BFGS-B. Benchmarks run 300 trials with B = 10 restarts at D = 5, 10, 20, and 40, reporting medians over 20 independent runs with different random seeds. Full results for Sphere, Attractive Sector, and Step Ellipsoidal appear in an appendix table.

Why This Matters

Impact on research. The paper reframes a batching trick that is standard practice in BoTorch as a convergence hazard, and it does so with a diagnosis rather than just a benchmark: the off-diagonal artifacts are a concrete, measurable property of the inverse-Hessian approximation. It also shows that a practical fix does not require a new bound-constrained QN algorithm, which is notable given that the structurally aware extensions that exist for BFGS and L-BFGS do not carry over to L-BFGS-B. For anyone benchmarking acquisition-function optimizers, the distinction between coupled and decoupled updates becomes a variable that must be controlled.

Real-world applications (drawn from the paper's own list of expensive-objective domains):

• Materials discovery, where each candidate evaluation is an expensive experiment or simulation.
• Hyperparameter optimization of machine learning models, the setting in which BoTorch and Optuna are most heavily used.
• Drug discovery, where candidate screening is costly and trial-and-error iterations must be minimized.

Industry relevance. The approach is already merged into GPSampler in Optuna, one of the most widely used open-source hyperparameter optimization frameworks, so the speedup is available to practitioners rather than remaining a research prototype. BoTorch v0.15.0 also moved from C-BE to D-BE independently of this work, which suggests the coupled formulation was a recognized bottleneck even before this paper formalized why. The code for the experiments is released at https://github.com/Kaichi-Irie/faster-batched-acqf-opt-experiments/tree/submission.

Future Directions

• A bound-constrained, structure-aware QN solver. The paper states that no practical block-structure-aware, bound-constrained QN algorithm exists that removes off-diagonal artifacts, and that extending BFGS and L-BFGS work in this direction to L-BFGS-B is challenging. A solver that exploits block-diagonality while respecting box constraints would remove the need for the coroutine workaround entirely.
• Quantifying the floating-point nondeterminism. D-BE is claimed to be theoretically identical to sequential MSO, but the authors note that automatic differentiation can produce slightly different gradients between batched and single evaluations due to modulo floating-point nondeterminism. How much this perturbs trajectories in practice is left open.
• Coroutine versus native D-BE in BoTorch v0.15.0. Since BoTorch v0.15.0 adopted D-BE without a coroutine, a direct comparison of the coroutine-based implementation against that native implementation is an obvious next experiment.
• Beyond the four benchmark functions tested. Results are reported for Rastrigin, Sphere, Attractive Sector, and Step Ellipsoidal from the COCO (BBOB) suite at D = 5, 10, 20, and 40; whether the same pattern holds for real BO workloads with noisy, heterogeneous objectives is not reported.

Target Audience

Researchers and engineers working on Bayesian optimization, acquisition-function optimization, or hyperparameter optimization tooling — particularly those using or contributing to BoTorch, Optuna, or similar libraries, and anyone who has noticed that batched multi-start optimization does not deliver the expected speedup in high dimensions. Readers who need a practical, low-implementation-cost optimization for their BO stack benefit most; readers looking for new convergence theory for quasi-Newton methods will find the paper's argument empirical, with the theoretical guarantee limited to the claim that D-BE reproduces sequential MSO trajectories.

Authors’ abstract

Bayesian optimization (BO) efficiently finds high-performing parameters by maximizing an acquisition function, which models the promise of parameters. A major computational bottleneck arises in acquisition function optimization, where multi-start optimization (MSO) with quasi-Newton (QN) methods is required due to the non-convexity of the acquisition function. BoTorch, a widely used BO library, currently optimizes the summed acquisition function over multiple points, leading to the speedup of MSO owing to PyTorch batching. Nevertheless, this paper empirically demonstrates the suboptimality of this approach in terms of off-diagonal approximation errors in the inverse Hessian of a QN method, slowing down its convergence. To address this problem, we propose to decouple QN updates using a coroutine while batching the acquisition function calls. Our approach not only yields the theoretically identical convergence to the sequential MSO but also drastically reduces the wall-clock time compared to the previous approaches. Our approach is available in GPSampler in Optuna, effectively reducing its computational overhead.