Implementation:Pyro ppl Pyro GP BayesOpt

Property	Value
Implementation Type	Pattern Doc
Source File	examples/contrib/oed/gp_bayes_opt.py
Module	pyro.contrib.oed
Pyro Features	pyro.contrib.gp.models.GPRegression, pyro.optim.multi.MultiOptimizer, TraceEnum_ELBO, Thompson sampling, LBFGS optimization
Pattern	Bayesian optimization using GP surrogate model with Thompson sampling acquisition

Overview

This file implements a GP-based Bayesian Optimizer (GPBayesOptimizer) that uses a Gaussian Process as an emulator (surrogate model) for an unknown objective function. It extends pyro.optim.multi.MultiOptimizer and follows a standard Bayesian optimization loop:

1. Acquire new points to evaluate using Thompson sampling (default) or a custom acquisition function
2. Evaluate the true objective function at those points
3. Update the GP posterior with the new observations

The Thompson sampling acquisition works by:

• Drawing a sample function from the GP posterior
• Optimizing this sample function using LBFGS with multiple random restarts
• Returning the optimal point as the next query location

The GP posterior is updated after each batch of acquisitions using Adam optimization of the GP hyperparameters.

Code Reference

class GPBayesOptimizer(pyro.optim.multi.MultiOptimizer):
    def __init__(self, constraints, gpmodel, num_acquisitions, acquisition_func=None):
        if acquisition_func is None:
            acquisition_func = self.acquire_thompson
        self.constraints = constraints
        self.gpmodel = gpmodel
        self.num_acquisitions = num_acquisitions
        self.acquisition_func = acquisition_func

    def update_posterior(self, X, y):
        X = torch.cat([self.gpmodel.X, X])
        y = torch.cat([self.gpmodel.y, y])
        self.gpmodel.set_data(X, y)
        optimizer = torch.optim.Adam(self.gpmodel.parameters(), lr=0.001)
        gp.util.train(self.gpmodel, optimizer,
            loss_fn=TraceEnum_ELBO(strict_enumeration_warning=False).differentiable_loss,
            retain_graph=True)

    def acquire_thompson(self, num_acquisitions=1, **opt_params):
        X = torch.empty(num_acquisitions, ...)
        for i in range(num_acquisitions):
            sampler = self.gpmodel.iter_sample(noiseless=False)
            x, _ = self.opt_differentiable(sampler, **opt_params)
            X[i, ...] = x
        return X

    def get_step(self, loss, params, verbose=False):
        X = self.acquisition_func(num_acquisitions=self.num_acquisitions)
        y = loss(X)
        self.update_posterior(X, y)
        return self.opt_differentiable(lambda x: self.gpmodel(x)[0])

I/O Contract

Parameter	Type	Description
constraints	torch.constraint	Domain constraints for the optimization
gpmodel	gp.models.GPRegression	GP regression model (kernel, etc.)
num_acquisitions	int	Number of points to acquire per step
acquisition_func	callable	Custom acquisition function (default: Thompson sampling)

Methods:

• get_step(loss, params): Main optimization step -- acquires points, evaluates loss, updates GP, returns current best
• update_posterior(X, y): Updates GP with new observations
• acquire_thompson(num_acquisitions): Thompson sampling acquisition
• find_a_candidate(differentiable, x_init): Single LBFGS step from initial point
• opt_differentiable(differentiable, num_candidates): Multi-start optimization

Usage Examples

import pyro.contrib.gp as gp
from gp_bayes_opt import GPBayesOptimizer

# Set up GP model
kernel = gp.kernels.RBF(input_dim=1)
gpmodel = gp.models.GPRegression(X_init, y_init, kernel)

# Create optimizer
constraints = torch.distributions.constraints.interval(0, 10)
optimizer = GPBayesOptimizer(constraints, gpmodel, num_acquisitions=1)

# Optimization loop
for step in range(50):
    x_best, y_best = optimizer.get_step(objective_fn, params=None)

Related Pages

• Pyro_ppl_Pyro_SV_DKL - Deep kernel learning with GPs
• Pyro_ppl_Pyro_GP_TimeSeries - GP-based time series models