Active learning on the Hartmann test function using BoTorch's acquisition functions

[takafusui]

Issue description

I am working on active learning and have tested the PosteriorStandardDeviation (PSTD) and the qNegIntegratedPosteriorVariance (NIPV) acquisition function on the Hartmann 6 test function.
Active learning involves some randomness, therefore, I did 100 Monte-Carlo iterations.

As you see in the figure below, the leave-one-out error (errLOO) is lower with PSTD than with NIPV criterion. Moreover, the variance with PSTD is smaller than that of NIPV. I also did the same exercise with the space-filling Latin-Hypercube sampling (LHS), which shows a noisy result.

I believe that the NIPV criterion is one benchmark in the active learning literature. Thus, I want to understand possible reasons why in my implementation, the PTSD criteria works better than 'NIPV'. At least 'PTSD' shows more robust convergence. Maybe my hyper-parameter setting for optimize_acqf is not good or maybe NIPV cannot correctly integrate variance. I use the draw_sobol_samples helper function to generate points to be integrated.

Thank you very much for your insights on this.

This discussion should be closely related to #1366 and #2060.

Code example

import numpy as np
from smt.sampling_methods import Random, LHS
import torch
import gpytorch
from botorch.models import SingleTaskGP
from botorch.models.transforms import Normalize, Standardize
from botorch.fit import fit_gpytorch_mll
from botorch.optim import optimize_acqf
from botorch.acquisition.analytic import PosteriorStandardDeviation
from botorch.acquisition import qNegIntegratedPosteriorVariance
from botorch.utils.sampling import draw_sobol_samples
from botorch.test_functions.synthetic import Hartmann
# Leave-one-out error
from sklearn.model_selection import LeaveOneOut
import matplotlib.pyplot as plt
from matplotlib import rc
import pandas as pd

torch_dtype = torch.double
# --------------------------------------------------------------------------- #
# Plot setting
# --------------------------------------------------------------------------- #
# Use TeX font
rc('font', **{'family': 'sans-serif', 'serif': ['Helvetica']})
# rc('text', usetex=True)

# Figure size
fsize = (9, 6)

# Font size
plt.rcParams["font.size"] = 14
plt.rcParams["axes.labelsize"] = 14
plt.rcParams["axes.titlesize"] = 14
plt.rcParams["legend.title_fontsize"] = 14

# --------------------------------------------------------------------------- #
# Helper functions
# --------------------------------------------------------------------------- #
def initialize_GP(train_X, train_y, train_X_bounds):
	"""Initialize a GP model.
	trian_X.shape: [N, d]
	train_y.shape: [N, 1]

	Return the fitted GP model.
	"""
	# Get the dimensions of train_X and train_y to standardize inputs
	train_X_dim = train_X.shape[-1]
	train_y_dim = train_y.shape[-1]

	gp = SingleTaskGP(
		train_X, train_y,
		input_transform=Normalize(d=train_X_dim, bounds=train_X_bounds),
		outcome_transform=Standardize(m=train_y_dim)
	)
	mll = gpytorch.mlls.ExactMarginalLogLikelihood(gp.likelihood, gp)

	return mll, gp

def leave_one_out_GP(train_X, train_y, train_X_bounds):
	"""Compute the leave-one-out error.

	train_X.shape = [n, N_inputs]
	train_y.shape = [n, 1]
	errLOO <= 0.01 is a sufficient accuracy in practice.
	"""
	loo = LeaveOneOut()  # Leave-one-out error estimator instance
	err = torch.empty(loo.get_n_splits(train_X), dtype=torch_dtype)

	# Min-Max scale train_y in [0, 1]
	# Min-Max scaler prepserves the original shape of the distribution of train_y
	# min_max_scaler = preprocessing.MinMaxScaler()
	# train_y_minmax = torch.tensor(min_max_scaler.fit_transform(train_y))

	for train_idx, test_idx in loo.split(train_X):
		train_loo_X, test_loo_X = train_X[train_idx], train_X[test_idx]
		# train_loo_y, test_loo_y \
		#     = train_y_minmax[train_idx], train_y_minmax[test_idx]
		train_loo_y, test_loo_y \
			= train_y[train_idx], train_y[test_idx]

		# Train the GP model based on train_loo_X and train_loo_y
		mll_loo, gp_loo = initialize_GP(
			train_loo_X, train_loo_y, train_X_bounds)
		fit_gpytorch_mll(mll_loo)  # Using BoTorch's routine

		# Make prediction
		with torch.no_grad(), gpytorch.settings.fast_pred_var():
			# Posterior prediction
			pred_test_loo_X = gp_loo.posterior(test_loo_X)
			# Compute an squared approximation error
			err[test_idx] = (test_loo_y - pred_test_loo_X.mean)**2

	# Leave-one-out error
	errLOO = 1 / loo.get_n_splits(train_X) * torch.sum(err)

	return errLOO.reshape(1).numpy()

class BayesianActiveLearning:
	"""Bayesian active learning to improve an overall GP model accuracy."""

	def __init__(
			self,
			acquisition,  # Acquisition function to be optimized
			N_restarts,  # Number of initial conditions (ICs)
			raw_samples,
			N_MC_samples,
			q_batch):
		self.acquisition = acquisition
		self.N_restarts = N_restarts
		self.raw_samples = raw_samples
		self.N_MC_samples = N_MC_samples
		self.q_batch = q_batch

	def forward(self, train_X, train_y, train_X_bounds):
		"""One step forward active learning"""
		# Instantiate a GP model
		mll, gp = initialize_GP(train_X, train_y, train_X_bounds)
		# Fit a GP model
		fit_gpytorch_mll(mll)

		if self.acquisition == 'qNIPV':
			# The points to use for MC-integrating the posterior variance
			# n: The number of (q-batch) samples. As a best practice, it
			# should be powers of 2
			# Reshape to have [N, d]
			qmc_samples = draw_sobol_samples(
				bounds=train_X_bounds, n=self.N_MC_samples, q=self.q_batch,
				batch_shape=None, seed=None
			).squeeze(-2)

			# Batch Integrated Negative Posterior Variance
			# mc_points are used for MC-integrating the posterior variance.
			acqf = qNegIntegratedPosteriorVariance(
					model=gp,
					mc_points=qmc_samples
					)
		elif self.acquisition == 'PSTD':
			# Single-outcome posterior standard deviation for active learning
			acqf = PosteriorStandardDeviation(gp)

		# Optimize the acquisition function
		train_X_add, acq_value = optimize_acqf(
			acqf, bounds=train_X_bounds, q=self.q_batch,
			num_restarts=self.N_restarts, raw_samples=self.raw_samples,
			# options={"method": "L-BFGS-B"}
			# options={"batch_limit": 10, "init_batch_limit": 512}
			)
		return train_X_add

# --------------------------------------------------------------------------- #
# Training data
# --------------------------------------------------------------------------- #
N_train_X_init = 100  # Initial experimental design size
# Six-dimensional test function
col_header = [r'$x_1$', r'$x_2$', r'$x_3$', r'$x_4$', r'$x_5$', r'$x_6$']
N_inputs = len(col_header)
# Evaluated on [0, 1]^N_inputs
Xlimits = np.tile([0., 1.], (N_inputs, 1))
# Hartmann test function
hartmann6 = Hartmann(dim=N_inputs)
train_X_bounds = torch.from_numpy(
	np.tile(Xlimits[0].reshape(2, 1), (1, N_inputs)))

# Latin-Hypercube sampling and the initial experimental design
sampler = LHS(xlimits=Xlimits)
train_X_init = torch.from_numpy(sampler(N_train_X_init))
train_y_init = hartmann6.evaluate_true(train_X_init)[:, None]

# --------------------------------------------------------------------------- #
# Botorch and optimize_acqf parameters
# --------------------------------------------------------------------------- #
# Number of initial conditions from which we optimize an acquisition function
# The larger NUM_RESTARTS, the more memory we need
N_restarts = 5  # Default number in Ax is 20
# The larger RAW_SAMPLES, the better initial condition we have
raw_samples = 1024  # Default number in Ax
# Larger N would lead to more precise integration of the posterior variance,
# but since GP posterior variance is a relatively smooth function, would not
# expect pushing N larger to make a big difference.
N_MC_samples = 256  # Should be a power of 2
q_batch = 1  # when q > 1 in optimize_acqf, q =! 1

# Error analysis setting
n_mc_iter = 100  # Number of Monte-Carlo iteration
verbose = 20
N_train_X_end = 200
N_train_X_list = np.arange(N_train_X_init, N_train_X_end + 1, verbose)

# --------------------------------------------------------------------------- #
# Using LHS
# --------------------------------------------------------------------------- #
errLOO = np.empty((n_mc_iter, len(N_train_X_list)))
N_train_X_LOO = []

print(r"Space filling design using LHS")

for jdx in range(n_mc_iter):
	print(r"Monte-Carlo iteration {} / {}".format(jdx + 1, n_mc_iter))

	for idx, N_train_X in enumerate(N_train_X_list):
		if idx == 0:
			# Use the same initial experimental design
			train_X = train_X_init
			train_y = train_y_init
		else:
			train_X = torch.from_numpy(sampler(N_train_X))
			train_y = hartmann6.evaluate_true(train_X)[:, None]
		# Compute a LOO error
		err = leave_one_out_GP(train_X, train_y, train_X_bounds)
		errLOO[jdx, idx] = err[0]

		if jdx == 0:
			# As N_train_X_LOO is the same for all jdx
			N_train_X_LOO.append(train_X.shape[0])

LHS_result = pd.DataFrame(np.vstack([N_train_X_LOO, errLOO]))
LHS_result.to_csv('hartmann6_LHS.csv', index=False, header=False)

# --------------------------------------------------------------------------- #
# Bayesian active learning
# --------------------------------------------------------------------------- #
def MCBaisianActiveLearning(
		criterion, n_mc_iter, n_max_iter, train_X_init, train_y_init, verbose):
	"""Monte-Carlo Bayesian active learning.
	criterion: Acquisition function
	n_mc_iter: Number of Monte-Carlo iteration
	n_max_iter: Number of max Bayesian active learning iteration
	train_X_init: Initial training input
	train_y_init: Initial training output
	verbose: Print verbose
	"""
	# Instantiate active learning object
	activelearning = BayesianActiveLearning(
		criterion, N_restarts, raw_samples, N_MC_samples, q_batch)

	# Initialize output arrays
	errLOO = np.empty((n_mc_iter, n_max_iter//verbose + 1))
	N_train_X_LOO = []


	print(r"Bayesian active learning using the {} criterion".format(criterion))

	for jdx in range(n_mc_iter):
		print(r"Monte-Carlo iteration {} / {}".format(jdx + 1, n_mc_iter))
		# Initialization
		train_X, train_y = train_X_init, train_y_init

		for idx in range(n_max_iter):

			if idx == 0:
				err = leave_one_out_GP(
					train_X, train_y, train_X_bounds)
				errLOO[jdx, idx] = err[0]
				if jdx == 0:
					N_train_X_LOO.append(train_X.shape[0])

			train_X_add = activelearning.forward(train_X, train_y, train_X_bounds)
			train_y_add = hartmann6.evaluate_true(train_X_add)[:, None]

			# Add the candidate point
			train_X = torch.cat((train_X, train_X_add), axis=0)
			train_y = torch.cat((train_y, train_y_add), axis=0)


			if (idx + 1) % verbose == 0:
				print(r"Iteration: {} / {}".format(idx+ 1 , n_max_iter))
				err = leave_one_out_GP(train_X, train_y, train_X_bounds)
				errLOO[jdx, (idx + 1) // verbose] = err[0]
				if jdx == 0:
					N_train_X_LOO.append(train_X.shape[0])
	return N_train_X_LOO, errLOO

# --------------------------------------------------------------------------- #
# Using the PSTD acquisition function
# --------------------------------------------------------------------------- #
n_max_iter = N_train_X_list[-1] - N_train_X_list[0]
N_train_X_LOO, errLOO = MCBaisianActiveLearning(
	'PSTD', n_mc_iter, n_max_iter, train_X_init, train_y_init, verbose)

PSTD_result = pd.DataFrame(np.vstack([N_train_X_LOO, errLOO]))
PSTD_result.to_csv('hartmann6_PSTD.csv', index=False, header=False)

# --------------------------------------------------------------------------- #
# Using the qNIPV acquisition function
# --------------------------------------------------------------------------- #
N_train_X_LOO, errLOO = MCBaisianActiveLearning(
	'qNIPV', n_mc_iter, n_max_iter, train_X_init, train_y_init, verbose)

NIPV_result = pd.DataFrame(np.vstack([N_train_X_LOO, errLOO]))
NIPV_result.to_csv('hartmann6_NIPV.csv', index=False, header=False)

# import ipdb; ipdb.set_trace()
# --------------------------------------------------------------------------- #
# Plot the LOO error
# --------------------------------------------------------------------------- #
# Read the LOO error history
LHS_read = pd.read_csv('hartmann6_LHS.csv', header=None)
PSTD_read = pd.read_csv('hartmann6_PSTD.csv', header=None)
NIPV_read = pd.read_csv('hartmann6_NIPV.csv', header=None)

LHS_mu, LHS_std = LHS_read.loc[1:].mean(), LHS_read.loc[1:].std()
PSTD_mu, PSTD_std = PSTD_read.loc[1:].mean(), PSTD_read.loc[1:].std()
NIPV_mu, NIPV_std = NIPV_read.loc[1:].mean(), NIPV_read.loc[1:].std()

# Plot
fill_alpha = 0.3
fig, ax = plt.subplots(figsize=fsize)
ax.plot(LHS_read.loc[0], LHS_mu, label='LHS')
ax.fill_between(LHS_read.loc[0], LHS_mu - LHS_std, LHS_mu + LHS_std, alpha=fill_alpha)
ax.plot(PSTD_read.loc[0], PSTD_mu, label='PSTD')
ax.fill_between(PSTD_read.loc[0], PSTD_mu - PSTD_std, PSTD_mu + PSTD_std, alpha=fill_alpha)
ax.plot(NIPV_read.loc[0], NIPV_mu, label='NIPV')
ax.fill_between(NIPV_read.loc[0], NIPV_mu - NIPV_std, NIPV_mu + NIPV_std, alpha=fill_alpha)
ax.set_xlabel(r'Number of experimental design')
ax.set_ylabel(r'$err_{\mathrm{LOO}}$')
ax.set_yscale('log')
ax.legend()
plt.savefig('hartmann6_activelearning.png', bbox_inches="tight")
plt.close()

System Info

Please provide information about your setup, including

• BoTorch Version 0.10.0
• GPyTorch Version 1.11
• PyTorch Version 2.2.2
• Computer OS Ubuntu 22.04

[Balandat]

I can't say I have a great understanding of the nuances of the active learning acquisition functions and why you see this behavior (or whether that is surprising), but let me share some thoughts:

• The metric you're using for evaluation here is the MSE, while qNegIntegratedPosteriorVariance targets, as the name suggests, the integrated posterior variance. It's not particularly surprising that the this may not lead to the best MSE model. It would be interesting to also compute the integrated posterior variance, w.r.t. which one would indeed expect better performance from using qNegIntegratedPosteriorVariance. More generally, it would be good to also include metrics that include a measure of the quality of the model's uncertainty quantification, e.g. the NLL of the observations under the model posterior.
• Does this replicate on other test functions? It's not uncommon to see one acquisition strategy show better performance one one problem and worse performance on another based on the nature of the underlying function. I would at least try a couple of other test functions to see whether this behavior is systematic or not.
• It looks like you're starting off at 100 initial points, which is quite a lot for a GP model on Hartmann already. You'll probably have a much easier time distinguishing performance if you reduce this substantially, since right now the difference in model performance will be quite small from the get-go.