Unable to get good regret of suggested point result for qKnowledgeGradient

[yyongjiann]

Hi all,

I am currently trying to replicate the experimental finding w.r.t. to Botorch Paper, specifically the figure 9 in the paper.

I am only looking at the comparison between Botorch EI (orange curve) and OKG (red curve) for the Branin testing function (d=2). I have been struggling for the past few weeks trying to replicate the finding in the paper as the regret value for my implementation for the OKG method is not converging downwards sufficiently. I ran a comparison test between this implementation of OKG and qEI (logarithmic variant) for 20 steps over 30 trials on an ascending sequence of seed value from 0 to 29 and obtained the following result.

Noise of N~(0,0.5^2)

Noise of N~(0,0.25^2)

I understand that the paper used 100 trials and 40 steps (interpreted from the figure) and that may explain the difference in result, but I was not expecting the regret value result of OKG to still remain so high as the paper results shows a fairly quick drop in regret value for OKG after a few batches of seen examples. I am not exactly sure whether this is due to a difference in hyperparameter used or implementation error/conceptual misunderstanding but I would appreciate if anyone has any advice on this matter. Thank you!

Here is a snippet for the OKG code

def OKG_run(self):
    sampler = SobolQMCNormalSampler(torch.Size([128]), seed=self.seed)
    regrets = []
    self.numOfObs = []
    for step in range(self.n_steps):
        model = self.get_gp_model()
        self.numOfObs.append(self.Y_data.size(dim=0))

        # Optimize the PosteriorMean to find the current best value
        argmax_pmean, max_pmean = optimize_acqf(
            acq_function=PosteriorMean(model),
            bounds=self.bounds,
            q=1,
            num_restarts=20,
            raw_samples=2048,
        )

        qKG = qKnowledgeGradient(
            model=model,
            num_fantasies=128,
            sampler=sampler,
            current_value=max_pmean,
        )

        candidates, acq_value = optimize_acqf(
            acq_function=qKG,
            bounds=self.bounds,
            q=4,
            num_restarts=10,
            raw_samples=512,
        )

        # Evaluate the objective function
        new_y = self.evaluate(candidates)
        max_y = new_y.max().item()
        noisy_y = self.evaluate_noisy(candidates)

        # Update the data
        self.X_data = torch.cat([self.X_data, candidates])
        self.Y_data = torch.cat([self.Y_data, noisy_y])

        # Compute simple regret
        simple_regret = -max_y - self.minima
        regrets.append(simple_regret)

    self.okg_regrets.append(regrets)

`
Here is a snippet of my code for the other relevant functions.

#inside __init__
self.bounds = torch.stack(
        [
    torch.zeros(self.dim, dtype=dtype, device=device),
    torch.ones(self.dim, dtype=dtype, device=device),
        ]
 )

def unnormalise(self, x):
    lb, ub = Branin(negate=True).to(device=device).bounds  # from botorch.test_functions
    lb = lb.to(x.device)
    ub = ub.to(x.device)
    return lb + (ub - lb) * x

# noiseless eval
def evaluate(self, x):
    return self.branin(self.unnormalise(x)).unsqueeze(-1)

# noisy eval
def evaluate_noisy(self, x):
    noiseless = self.evaluate(x)
    noise = torch.randn_like(noiseless) * 0.5  # Add noise with variance 0.25
    return noiseless + noise

def initialise(self):
    n=2*self.dim+2
    self.X_data = SobolEngine(dimension=self.dim, scramble=True, seed=self.seed+42).draw(n).to(device=device, dtype=dtype)
    self.Y_data = self.evaluate_noisy(self.X_data) 

def get_gp_model(self):
    if self.X_data is None:
        self.initialise()

    train_Yvar = torch.full_like(self.Y_data, 0.0625)
    self.gp_model = SingleTaskGP(
        train_X=self.X_data,
        train_Y=self.Y_data,
        train_Yvar=train_Yvar)

    mll = ExactMarginalLogLikelihood(self.gp_model.likelihood, self.gp_model)
    fit_gpytorch_mll(mll)
    return self.gp_model

`

[Balandat]

Two things:

1. Mismatch between actual and reported noise variance

noise = torch.randn_like(noiseless) * 0.5 # Add noise with variance 0.25
train_Yvar = torch.full_like(self.Y_data, 0.0625)

Seems like you're telling the model the observation noise variance is 0.0625 rather than 0.25? Presumably will cause the model to behave very poorly b/c it thinks that the observed values are much less noisy than they actually are and as a result you'll likely end up with a poor fit with very short lengthscales. Since KG performance depends much more on the surrogate model quality than EI performance, this could explain at least part of the behavior you're seeing.

2. Regret computation

You're computing the "suggested best point" for KG as the best function value of the q=4 newly suggested points at that iteration. However, what is shown in the plot in the paper is the inference regret, i.e. the regret w.r.t. to a single suggested point (the maximizer of the posterior mean), which generally is never actually evaluated during the optimization (i.e. is "out-of-sample"). In contrast, the regret w.r.t. to qEI is computed based on the maximizer of the posterior mean of all previously evaluated points.

From the paper:

We evaluate based on the true noiseless function value at the “suggested point” (i.e., the point to be chosen if BO were to end at this batch). OKG, MOE KG, and NEI use “out-of-sample” suggestions (introduced as χn in Section D.4), while the others use
“in-sample” suggestions

[yyongjiann]

Hi Max, thank you for the comments! I have adjusted for the mismatch in variance but I am not quite sure if my adjusted code for the implementation of the regret computation accurately reflects your comments.

 def qEI_run(self):
    sampler = SobolQMCNormalSampler(torch.Size([128]), seed=self.seed)
    regrets = []
    
    for step in range(self.n_steps):
        gp = self.get_gp_model()
        
        # Find the best observed value (in-sample suggestion)
        best_f = self.Y_data_noiseless.max().item()  
        
        qEI = qLogExpectedImprovement(model=gp, best_f=best_f, sampler=sampler)
        
        candidates, acq_value = optimize_acqf(
            acq_function=qEI,
            bounds=self.bounds,
            q=4,
            num_restarts=10,
            raw_samples=512,
        )

        new_y = self.evaluate(candidates)
        noisy_y = self.evaluate_noisy(candidates)

        self.X_data = torch.cat([self.X_data, candidates])
        self.Y_data = torch.cat([self.Y_data, noisy_y])
        self.Y_data_noiseless = torch.cat([self.Y_data_noiseless, new_y])

        # Compute regret based on the best in-sample point
        regret = - best_f - self.minima 
        regrets.append(regret)

        print(f"qEI Step {step}/{self.n_steps}:")
        print(f"Best value found: {-best_f}")
        print(f"Best location found: {self.X_data[self.Y_data_noiseless.argmax()]}\n")   

    self.qei_regrets.append(regrets)

def get_maximum_posterior_mean(self, model):
    posterior_mean_acq = PosteriorMean(model=model)
    max_point, max_value = optimize_acqf(
        acq_function=posterior_mean_acq,
        bounds=self.bounds,
        q=1,
        num_restarts=20,
        raw_samples=2024,
    )
    return max_point, max_value.item()

def OKG_run(self):
    sampler = SobolQMCNormalSampler(torch.Size([128]), seed=self.seed)
    regrets = []
    self.numOfObs = []
    
    for step in range(self.n_steps):
        model = self.get_gp_model()
        self.numOfObs.append(self.Y_data.size(dim=0))

        # Find the posterior mean of evaluated points (out-of-sample suggestion)
        out_of_sample_point, out_of_sample_value = self.get_maximum_posterior_mean(model)

        qKG = qKnowledgeGradient(
            model=model,
            num_fantasies=128,
            sampler=sampler,
            current_value=out_of_sample_value,
        )

        candidates, acq_value = optimize_acqf(
            acq_function=qKG,
            bounds=self.bounds,
            q=4,
            num_restarts=10,
            raw_samples=512,
        )

        noisy_y = self.evaluate_noisy(candidates)

        self.X_data = torch.cat([self.X_data, candidates])
        self.Y_data = torch.cat([self.Y_data, noisy_y])
    
        # Compute regret based on the out-of-sample suggestion
        true_out_of_sample_value = self.evaluate(out_of_sample_point).item()
        regret = -true_out_of_sample_value - self.minima 
        regrets.append(regret)

        print(f"OKG Step {step}/{self.n_steps}:")
        print(f"Best value found: {true_out_of_sample_value}")
        print(f"Best location found: {out_of_sample_point}\n")

    self.okg_regrets.append(regrets)

I plotted the regret curve for 30 trials and obtained the following result, which is still less than ideal:

[yyongjiann]

I realised that in my code I mislabeled the colours for the curves and now the graph makes more sense now. Thank you for the help!