Sequential experimental design using BoTorch's ecosystem

[takafusui]

I want to improve the global accuracy of the GP model, and I would like to know how I apply BoTorch's routine to sequentially augment the experimental design point.
I know BoTorch efficiently addresses the exploitation-exploration tradeoff, but I am rather interested in how I can efficiently reduce the exploration part.
I appreciate any ideas and comments and thank you in advance.

[saitcakmak]

Sounds like you're interested in an active-learning type approach, where the goal is to learn the underlying function as well as possible rather than trying to optimize it. We do not have many acquisition functions defined for this setting but qNegIntegratedPosteriorVariance might be interesting to you.

If you're still interested in optimizing the function but want to give an additional emphasis on exploring the whole space (which should improve global model accuracy), artificially inflating the predictive variance would likely help the acquisition functions focus a bit more on this. You could use UCB with a large beta, or augment the posterior variance of any other acquisition function using a custom PosteriorTransform to achieve this in a bit of a hacky way. Passing something like this to the acquisition function could work (modified from ExpectationPosteriorTransform - untested).

class MultiplyVariancePosteriorTransform(PosteriorTransform):
    def __init__(self, variance_multiplier: int) -> None:
        super().__init__()
        self.variance_multiplier = variance_multiplier

    def evaluate(self, Y: Tensor) -> Tensor:
        return Y

    def forward(self, posterior: GPyTorchPosterior) -> GPyTorchPosterior:
        r"""Augment the posterior by scaling up the covariance matrix."""
        org_mvn = posterior.distribution
        if getattr(org_mvn, "_interleaved", False):
            raise UnsupportedError

        org_event_shape = org_mvn.event_shape
        m = 1 if len(org_event_shape) == 1 else org_event_shape[-1]

        new_cov = org_mvn.covariance_matrix * self.variance_multiplier
        if m == 1:
            new_mvn = MultivariateNormal(
                org_mvn.loc, to_linear_operator(new_cov)
            )
        else:
            # Using MTMVN since we pass a single loc and covar for all `m` outputs.
            new_mvn = MultitaskMultivariateNormal(
                org_mvn.loc, to_linear_operator(new_cov), interleaved=False
            )
        return GPyTorchPosterior(distribution=new_mvn)

[takafusui]

Hi @saitcakmak,

thank you very much for your comments. I, at first, go with the qNegIntegratedPosteriorVariance acquisition function.

Firstly, I define the SobolQMCNormalSampler sampler for the subsequent MC integration in qNegIntegratedPosteriorVariance:
qmc_sampler = SobolQMCNormalSampler(sample_shape=1024, seed=123)

Then prepared MC sample points to integrate posterior variance:

with torch.no_grad(), gpytorch.settings.fast_pred_var():
    # As we standardized inputs, we must use gp.posterior to get
    # rescaled outputs.
    pred_test_X = gp.posterior(test_X)
    qmc_samples = qmc_sampler(pred_test_X)

Finally, define the qNegIntegratedPosteriorVariance instance and optimize it by using optimize_acqf

qNIPV = qNegIntegratedPosteriorVariance(
    model=gp, mc_points=qmc_samples)
train_X_add, acq_value = optimize_acqf(
    qNIPV, bounds=bounds, q=1, num_restarts=10, raw_samples=2048
)

But when executing optimize_acqf, it says
RuntimeError: shape '[1, -1, 9]' is invalid for input of size 1024000

I guess 1024000 comes from the sample_shape=1024 x len(test_X)=1000. The last 9 indicates that the GP model has indeed 9 dimensions.
Please also note that the trained GP model is defined using SingleTaskGP, something like

gp = SingleTaskGP(
    train_X, train_y,
    input_transform=InputStandardize(d=train_X_dim),
    outcome_transform=Standardize(m=train_y_dim)
)

Do you notice where the error is coming from?
Does my implementation make sense to you?
I cannot find any tutorial on how to use qNegIntegratedPosteriorVariance...

Thank you very much in advance.

[takafusui]

I believe my last post is not good enough to explain what I want to do here, so I prepared a simple code snippet here.
The snippet should be self-explained, but I would like to add some comments below.

import numpy as np
import pandas as pd
from smt.sampling_methods import LHS

# PyTorch
import torch
# GPyTorch
import gpytorch
# BoTorch
from botorch.models import SingleTaskGP
from botorch.fit import fit_gpytorch_model
from botorch.models.transforms.outcome import Standardize
from botorch.models.transforms.input import InputStandardize
from botorch.acquisition import qNegIntegratedPosteriorVariance
from botorch.sampling.normal import SobolQMCNormalSampler
from botorch.optim import optimize_acqf

# Leave-one-out error
from sklearn.model_selection import LeaveOneOut

# Fix random seeds
np.random.seed(123)
torch.manual_seed(123)


# --------------------------------------------------------------------------- #
# Two-dimensional Gaussian peak function
# --------------------------------------------------------------------------- #
def peak_func(x):
    """Two-dimensional Gaussian peak function.

    x.shape: (N, 2)
    """
    ai = np.array([5, 5])
    ui = np.array([0.8, 0.8])

    _peak_func = np.exp(-np.sum(ai**2 * (x - ui)**2, axis=1, keepdims=True))

    return _peak_func


# --------------------------------------------------------------------------- #
# Test dataset
# --------------------------------------------------------------------------- #
# Define the function on [0, 1]^2
xlimits = np.array([[0, 1], [0, 1]])
# Latin-Hypercube sampling
sampler_LHS = LHS(xlimits=xlimits, random_state=123)
# Test samples
N_test_X = 1000
test_X = sampler_LHS(N_test_X)
test_y = peak_func(test_X)


# --------------------------------------------------------------------------- #
# Train a GP model
# --------------------------------------------------------------------------- #
def train_GP(train_X, train_y):
    """Train a GP model using GPyTorch and BoTorch.

    Return the posterior.
    """
    # 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]

    # Train the GP model with the standardized inputs
    gp = SingleTaskGP(
        train_X, train_y,
        input_transform=InputStandardize(d=train_X_dim),
        outcome_transform=Standardize(m=train_y_dim)
    )
    mll = gpytorch.mlls.ExactMarginalLogLikelihood(gp.likelihood, gp)
    fit_gpytorch_model(mll)  # Using BoTorch's routine

    # Get into the evaluation mode
    gp.eval()

    with torch.no_grad(), gpytorch.settings.fast_pred_var():
        # As we standardized inputs, we must use gp.posterior to get
        # rescaled outputs.
        pred_test_X = gp.posterior(torch.from_numpy(test_X))

    return pred_test_X


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

    train_X.shape: (N, d)
    train_y.shape: (N, 1)

    return the leave-ont-out error
    """
    loo = LeaveOneOut()  # Leave-one-out error estimator instance

    # Get the dimensions of train_X and train_y
    train_X_dim = train_X.shape[1]
    train_y_dim = train_y.shape[1]

    err = torch.empty(loo.get_n_splits(train_X), dtype=torch.float64)

    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[train_idx], train_y[test_idx]
        # Train the GP model based on train_loo_X and train_loo_y
        gp_loo = SingleTaskGP(
            train_loo_X, train_loo_y,
            input_transform=InputStandardize(d=train_X_dim),
            outcome_transform=Standardize(m=train_y_dim)
        )
        mll_loo = gpytorch.mlls.ExactMarginalLogLikelihood(
            gp_loo.likelihood, gp_loo)
        fit_gpytorch_model(mll_loo)  # Using BoTorch's routine to fit GP

        # Get into evaluation (predictive posterior) mode
        gp_loo.eval()

        # Make prediction
        with torch.no_grad(), gpytorch.settings.fast_pred_var():
            pred_test_loo_X = gp_loo.posterior(test_loo_X)

        with torch.no_grad():
            # 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


# --------------------------------------------------------------------------- #
# Active learning
# --------------------------------------------------------------------------- #
N_AL_iter_list = np.arange(0, 51, 5)
N_train_X_init = 100  # Initialize a pilot GP model
# Initial training set sampling from LHS
train_X = sampler_LHS(N_train_X_init)
train_y = peak_func(train_X)
# Convert to torch.tensor
train_X = torch.from_numpy(train_X)
train_y = torch.from_numpy(train_y)
# 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]

# Test dataset to evaluate a GP model
test_X = torch.from_numpy(test_X)
# Bounds on train_X
bounds = torch.tensor(
    [[0.] * train_X.shape[1], [1.] * train_X.shape[1]], dtype=torch.float64)

# Store errLOO errors
errLOO_AL = np.empty(len(N_AL_iter_list))

for AL_iter_idx, N_AL_iter in enumerate(N_AL_iter_list):
    print(r"Actively learning {} new training data...".format(N_AL_iter))

    for iter in range(N_AL_iter):
        # Instantiate a GP model
        gp = SingleTaskGP(
            train_X, train_y,
            input_transform=InputStandardize(d=train_X_dim),
            outcome_transform=Standardize(m=train_y_dim)
        )
        mll = gpytorch.mlls.ExactMarginalLogLikelihood(gp.likelihood, gp)
        fit_gpytorch_model(mll)  # Using BoTorch's routine to fit the GP model

        # Quasi-MC base sampler using Sobol sequences
        qmc_sampler = SobolQMCNormalSampler(sample_shape=1, seed=123)

        with torch.no_grad(), gpytorch.settings.fast_pred_var():
            # As we standardized inputs, we must use gp.posterior to get
            # rescaled outputs.
            pred_test_X = gp.posterior(test_X)

        # The points to use for MC-integrating the posterior variance
        qmc_samples = qmc_sampler(pred_test_X)
        # Batch Integrated Negative Posterior Variance for Active Learning
        qNIPV = qNegIntegratedPosteriorVariance(model=gp, mc_points=qmc_samples)
        # Optimize the acquisition function
        train_X_add, acq_value = optimize_acqf(
            qNIPV, bounds=bounds, q=1, num_restarts=5, raw_samples=128
        )

        # Compute the new label
        train_y_add = peak_func(train_X_add.numpy())

        # Concatenate the training data and the labels to be added
        train_X = torch.cat((train_X, train_X_add), axis=0)
        train_y = torch.cat((train_y, torch.from_numpy(train_y_add)), axis=0)

    # Compute LOO error when we add N_AL_iter number of samples
    errLOO_AL[AL_iter_idx] = compute_errLOO(train_X, train_y)

# --------------------------------------------------------------------------- #
# Write-out errLOO to csv
# --------------------------------------------------------------------------- #
errLOO_AL_df = pd.DataFrame(errLOO_AL)
errLOO_AL_df.columns = ['errLOO']
errLOO_AL_df.index = N_train_X_init + N_AL_iter_list
errLOO_AL_df.to_csv('err_LOO_AL.csv')

In the above example, I aim to interpolate a 2D function peak_func using a small number of training datasets as my actual function is expensive to be evaluated.
I used the qNegIntegratedPosteriorVariance method to increase the global model accuracy and sequentially refine the model by adding a new training dataset by optimizing the acquisition function.
I also used the leave-one-out error as the number of training samples is limited.
The initial training dataset is sampled from LHS to instantiate a pilot GP model.

The main active learning loop is started from l. 160 (for AL_iter_idx, N_AL_iter in enumerate(N_AL_iter_list):).
I defined quasi-MC based sampler qmc_sampler, and integrated variance on these points when defining qNIPV.

My first question is whether my implementation is consistent with BoTorch's design idea. I am unsure as I cannot find a tutorial on how to use qNegIntegratedPosteriorVariance...

I also want to clarify: what is sample_shape in SobolQMCNormalSampler?
Thank you in advance for your comments.