ApproximateGP vs ExactGP Performance Discrepancy in Bayesian Optimization Contexts

[AndrewFalkowski]

I've encountered unexpected performance differences between ApproximateGP and ExactGP models in Bayesian optimization contexts. Despite using identical priors and initial training points, the ApproximateGP model consistently underperforms compared to ExactGP when using fixed inducing point locations at the data.

I'm wondering if I've made some error in my model construction or training loop that might explain this discrepancy or whether there is reason to expect different performances from these models. If there are any best practices for training the approximate models I would greatly appreciate it. If there are other bits of code that I could provide to help identify the problem, please let me know. Below is my implementation for review:

Model classes:

class VarGP(ApproximateGP):
    def __init__(self, inducing_points):
        variational_distribution = gpytorch.variational.NaturalVariationalDistribution(
            inducing_points.size(0)
        )
        variational_strategy = VariationalStrategy(
            self,
            inducing_points,
            variational_distribution,
            learn_inducing_locations=False,  # Fixed inducing points
        )
        super(VarGP, self).__init__(variational_strategy)

        self.mean_module = ConstantMean()

        input_dim = inducing_points.size(1)
        lengthscale_prior = LogNormalPrior(
            loc=math.sqrt(2) + math.log(input_dim) * 0.5, scale=math.sqrt(3)
        )

        base_kernel = RBFKernel(
            ard_num_dims=input_dim,
            lengthscale_prior=lengthscale_prior,
            lengthscale_constraint=GreaterThan(
                2.5e-2, transform=None, initial_value=lengthscale_prior.mode
            ),
        )

        outputscale_prior = LogNormalPrior(loc=0.0, scale=1.0)
        self.covar_module = ScaleKernel(
            base_kernel,
            outputscale_prior=outputscale_prior,
            outputscale_constraint=GreaterThan(
                1e-4, transform=None, initial_value=outputscale_prior.mode
            ),
        )

        noise_prior = LogNormalPrior(loc=-4.0, scale=1.0)
        self.likelihood = gpytorch.likelihoods.GaussianLikelihood(
            noise_prior=noise_prior,
            noise_constraint=GreaterThan(
                1e-4,
                transform=None,
                initial_value=noise_prior.mode,
            ),
        )

    @property
    def num_outputs(self) -> int:
        return 1

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

    def posterior(self, X, observation_noise=True, num_samples=2048, **kwargs):
        self.eval()
        with torch.no_grad(), gpytorch.settings.num_likelihood_samples(num_samples):
            posterior = self(X)
            if observation_noise:
                posterior = self.likelihood(posterior)
        return posterior


class ExactGP(gpytorch.models.ExactGP):
    def __init__(self, X_train, y_train, input_transform=None):
        noise_prior = gpytorch.priors.LogNormalPrior(loc=-4.0, scale=1.0)
        likelihood = GaussianLikelihood(
            noise_prior=noise_prior,
            noise_constraint=gpytorch.constraints.GreaterThan(
                1e-4,
                transform=None,
                initial_value=noise_prior.mode,
            ),
        )

        super(ExactGP, self).__init__(X_train, y_train, likelihood)

        self.input_transform = input_transform
        self.mean_module = gpytorch.means.ConstantMean()

        input_dim = X_train.shape[1]

        lengthscale_prior = gpytorch.priors.LogNormalPrior(
            loc=math.sqrt(2) + math.log(input_dim) * 0.5, scale=math.sqrt(3)
        )

        base_kernel = gpytorch.kernels.RBFKernel(
            ard_num_dims=input_dim,
            lengthscale_prior=lengthscale_prior,
            lengthscale_constraint=GreaterThan(
                2.5e-2, transform=None, initial_value=lengthscale_prior.mode
            ),
        )

        outputscale_prior = gpytorch.priors.LogNormalPrior(loc=0.0, scale=1.0)
        self.covar_module = gpytorch.kernels.ScaleKernel(
            base_kernel,
            outputscale_prior=outputscale_prior,
            outputscale_constraint=gpytorch.constraints.GreaterThan(
                1e-4, transform=None, initial_value=outputscale_prior.mode
            ),
        )

        self.num_outputs = 1

    def transform_inputs(self, X=None, input_transform=None):
        if X is None:
            raise ValueError("X must be provided")

        transform = (
            input_transform if input_transform is not None else self.input_transform
        )
        if transform is not None:
            return transform(X)
        return X

    def forward(self, x):
        """Compute the prior distribution at input locations x."""
        if self.training:
            x = self.transform_inputs(x)
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

    def posterior(self, X, observation_noise=True, num_samples=2048, **kwargs):
        """Compute the posterior distribution."""
        self.eval()
        with torch.no_grad(), gpytorch.settings.fast_pred_var():
            if observation_noise:
                posterior = self.likelihood(self(X))
            else:
                posterior = self(X)
            return posterior

Training loop for the approximate model. The ExactGP is trained with fit_gpytorch_mll(mll)

def train_natural_variational_model(model, X, y, epochs=500, lr=0.05, verbose=False, y_standardize=True):

    if y_standardize:
        y = TorchStandardScaler().fit_transform(y)

    model.train()
    model.likelihood.train()

    variational_ngd_optimizer = gpytorch.optim.NGD(
        model.variational_parameters(), num_data=y.size(0), lr=lr
    )

    hyperparameter_optimizer = torch.optim.Adam(
        [
            {"params": model.hyperparameters()},
        ],
        lr=0.01,
    )

    mll = gpytorch.mlls.VariationalELBO(model.likelihood, model, num_data=y.size(0))

    loss_history = []
    best_loss = float("inf")
    best_state = None
    best_epoch = 0

    for i in range(epochs):
        variational_ngd_optimizer.zero_grad()
        hyperparameter_optimizer.zero_grad()
        output = model(X)
        loss = -mll(output, y)
        current_loss = loss.item()
        loss_history.append(current_loss)

        # Update best model if we found a better loss
        if current_loss < best_loss:
            best_loss = current_loss
            best_state = {
                "model": model.state_dict(),
                "var_opt": variational_ngd_optimizer.state_dict(),
                "hyper_opt": hyperparameter_optimizer.state_dict(),
            }
            best_epoch = i

        loss.backward()
        variational_ngd_optimizer.step()
        hyperparameter_optimizer.step()

    # Load the best model state
    if best_state is not None:
        model.load_state_dict(best_state["model"])
        variational_ngd_optimizer.load_state_dict(best_state["var_opt"])
        hyperparameter_optimizer.load_state_dict(best_state["hyper_opt"])

    return None

[Balandat]

How many inducing points are you using and how did you select them?

I'm wondering if I've made some error in my model construction or training loop that might explain this discrepancy or whether there is reason to expect different performances from these models.

At the end of the day the ApproximateGP is, well, approximate. So I wouldn't necessarily expect it to perform as well as the exact model (if the exact model is indeed a good modeling choice). It'll matter a lot how you select the inducing points; as you increase their number and density across the space you should get closer to the Exact model performance. That said, even if you do that you are still optimizing a variational objective, and this comes with some known issues (e.g. underestimating predictive variances). I'd recommend checking out the Parametric Gaussian Process Regression (PPGPR) model from https://proceedings.mlr.press/v119/jankowiak20a/jankowiak20a.pdf that may help with some of these issues.

As to the training of the approximate model: I don't see anything wrong at first glance, but one would have to take a closer look at the actual results (e.g. learning curve) to understand what one might want to tweak.

[AndrewFalkowski]

Thanks for the quick reply. In the interest of trying to create a comparable benchmark, I set the number of inducing points to the number of observed data points at a given optimization iteration. For the same purpose, I set them to be fixed. Are there any heuristics to choosing the number of inducing points relative to data size?

Ultimately, I am looking to benchmark the effects of choosing different likelihood distributions on optimization performance. Perhaps this is a sufficiently apples to oranges comparison that I am better off comparing my alt-likelihood GPs to the standard variational GP rather than to an exact GP?

[jacobrgardner]

Couple add-ons here since I use variational GPs almost exclusively in my BO problems:

• Try switching gpytorch.mlls.VariationalELBO out for gpytorch.mlls.PredictiveLogLikelihood as Max suggests. This usually works either a bit better or the same.
• Using fixed inducing points is almost always worse in my experience over initializing with data and then learning them.
• In my experience for reasons I don't yet understand, the dimension scaled lengthscale priors interact very badly with variational GPs, at least on my own optimization problems. I get worse performance quite often (despite reproducing better performance with Exact GPs).
• At the end of the day, as Max suggests, these are different models so you're going to see at least different performance.

[AndrewFalkowski]

Hi @jacobrgardner, thanks for the reply! In your optimization problems, how do you typically go about selecting the number of inducing locations for a given optimization iteration? A simple M = n where n is the number of data points or is there a better heuristic?

Good to know on the length-scale prior issues! I was hoping to keep things comparable in the interest of benchmarking, but I certainly don't want to artificially damped the variational model performance. Any advice on choosing a more appropriate length-scale prior for the variational models?