[Bug]: qKnowledgeGradient shape mismatch when used with SaasFullyBayesianSingleTaskGP

[ayesharoxx]

What happened?

Hi, I’m trying to use qKnowledgeGradient with a fully Bayesian SAAS GP (SaasFullyBayesianSingleTaskGP) in BoTorch. I'm doing so by writing a new class that inherits from both SaasFullyBayesianSingleTaskGP and FantasizeMixin. Then, I override the fantasize() method to define how fantasy data is generated for this model. I have used 256 samples, warmup of 512, thinning of 1, num_fantasies=2. However, on running the code, I keep getting a shape mismatch error even with raw_samples=1 and num_restarts=1. The error looks like this:
RuntimeError: shape '[2, 1, 16, 1]' is invalid for input of size 64

I don't understand why I keep getting this error and where it is coming from. Any guidance on what might be causing this and how to properly structure the fantasy model in this context would be greatly appreciated!

Thanks in advance.

Please provide a minimal, reproducible example of the unexpected behavior.

The error occurs whenever I combine a SaasFullyBayesianSingleTaskGP with qKnowledgeGradient. I fit the SAAS model via fit_fully_bayesian_model_nuts, and then construct qKnowledgeGradient(model=model, num_fantasies=2). Calling optimize_acqf with q=1, num_restarts=1 and raw_samples=1 leads to the runtime error above. Removing KG (e.g. using qExpectedImprovement) works without issue, so the problem seems isolated to qKnowledgeGradient and the fully Bayesian SAAS model.

Please paste any relevant traceback/logs produced by the example provided.

RuntimeError: shape '[2, 1, 16, 1]' is invalid for input of size 64

BoTorch Version

v0.14.0

Python Version

No response

Operating System

No response

(Optional) Describe any potential fixes you've considered to the issue outlined above.

No response

Pull Request

None

Code of Conduct

• I agree to follow BoTorch's Code of Conduct

[hvarfner]

Hi @ayesharoxx,
Thanks for reaching out. fantasize and the underlying condition_on_observations are not well-tested for all possible functionalities in the fully Bayesian case, such as with KG. Could you perhaps provide a snippet to reproduce the issues you are encountering?

Notably, KG has 64 fantasies by default, so are you sure the attribute is passed correctly?

[ayesharoxx]

Hi @hvarfner,
Thanks so much for your reply. Sure, below is a minimal version of my code to reproduce the issue.

--- Objective function-Branin2d---

import math
import numpy as np

lb = np.hstack((-5 * np.ones(50), 0 * np.ones(50))) # lower bounds for input domain
ub = np.hstack((10 * np.ones(50), 15 * np.ones(50))) # upper bounds for input domain

def branin100(x):
assert (x <= ub).all() and (x >= lb).all()
x1, x2 = x[19], x[64] # Only dimensions 19 and 64 affect the value of the function
t1 = x2 - 5.1 / (4 * math.pi ** 2) * x1 ** 2 + 5 / math.pi * x1 - 6
t2 = 10 * (1 - 1 / (8 * math.pi)) * np.cos(x1)
return t1 ** 2 + t2 + 10

--- SAAS GP model with fantasize override ---

class SaasFullyBayesianSingleTaskGPWithFantasy(
SaasFullyBayesianSingleTaskGP,
FantasizeMixin
):
def fantasize(
self,
X: torch.Tensor,
sampler: Optional[MCSampler] = None,
num_fantasies: int = 2,
**kwargs,
) -> Model:
if sampler is None:
sampler = SobolQMCNormalSampler(
sample_shape=torch.Size([num_fantasies]),
collapse_batch_dims=True,
)
X = torch.as_tensor(
X,
dtype=self.train_inputs[0].dtype,
device=self.train_inputs[0].device,
)
return FantasizeMixin.fantasize(
self,
X,
sampler=sampler,
**kwargs,
)

--- Run SAASBO with KG minimal code ---

def run_saasbo_botorch():
torch.manual_seed(0)
device = "cpu"
dtype = torch.double
d = 100
lb = np.zeros(50)
ub = np.ones(50)
bounds = torch.stack([torch.zeros(d), torch.ones(d)]).to(dtype=dtype)

# Generate initial data
# Initial Sobol samples
sobol = SobolEngine(dim, scramble=True, seed=seed)
X = sobol.draw(num_init_evals).to(dtype=torch_dtype, device=torch_device)
Y = torch.tensor([f(lb + (ub - lb) * x.cpu().numpy()) for x in X], dtype=torch_dtype, device=torch_device).unsqueeze(-1)


# Normalize Y
train_Y = (Y - Y.mean()) / Y.std()

# Fit SAAS GP
model = SaasFullyBayesianSingleTaskGPWithFantasy(X, Y)
fit_fully_bayesian_model_nuts(model, warmup_steps=512, num_samples=256, thinning=16)

# Define posterior transform to average MCMC samples
weights = torch.ones(2, dtype=dtype) / 2
post_tf = ScalarizedPosteriorTransform(weights=weights)

# Define KG
qkg = qKnowledgeGradient(
    model=model,
    num_fantasies=2,
    current_value=train_Y.min(),
    posterior_transform=post_tf,
)

# Optimize acquisition
candidate, _ = optimize_acqf(
    acq_function=qkg,
    bounds=bounds,
    q=1,
    raw_samples=1,
    num_restarts=1,
)

run_saasbo_botorch()

[hvarfner]

Thanks. On my end, I ran
class FantasizeSaasFullyBayesianSingleTaskGP( SaasFullyBayesianSingleTaskGP, FantasizeMixin): pass
and set up KG the same way.

It seems like the issue is in AbstractFullyBayesianSingleTaskGP.condition_on_observations, which doesn't handle the three batch dimensions (fantasy_batch x raw_samples_batch x acq_batch) correctly. Having looked into this a while back, I don't think this is a quick fix but that is where I would start, since fantasize() calls condition_on_observations() on a num_fantasies batch of samples.

If I were you, I would override condition_on_observations() and check what is going on in there.

[ayesharoxx]

Thank you for the suggestion ! I overrode condition_on_observations() as you advised (code below), but I’m now running into the following error : -

Output shape not equal to that of weights. Output shape is 1 and weights are torch.Size([64])

Here’s my current implementation of condition_on_observations() :-

def condition_on_observations(self, X: torch.Tensor, Y: torch.Tensor, **kwargs):
model_batch_ndim = len(self.batch_shape)

    if X.ndim == 2 and Y.ndim == 2:
        X = X.repeat(self.batch_shape + (1, 1)).contiguous()
        Y = Y.repeat(self.batch_shape + (1, 1)).contiguous()
        return super().condition_on_observations(X, Y, **kwargs)

    
    start_idx = Y.ndim - (2 + model_batch_ndim)
    model_batch_indices = list(range(start_idx, start_idx + model_batch_ndim))
    extra_indices = list(range(0, start_idx))
    remaining_indices = list(range(start_idx + model_batch_ndim, Y.ndim - 2))
    permute_order = model_batch_indices + extra_indices + remaining_indices + [Y.ndim - 2, Y.ndim - 1]
    Y_perm = Y.permute(*permute_order).contiguous()
    

    
    if X.shape[:model_batch_ndim] != self.batch_shape:
        X = X.expand(self.batch_shape + X.shape[-2:])  

   
    extra_dims = len(extra_indices) + len(remaining_indices)
    for _ in range(extra_dims):
        X = X.unsqueeze(model_batch_ndim)

    
    expand_shape = list(X.shape)
    for i in range(extra_dims):
        expand_shape[model_batch_ndim + i] = Y_perm.shape[model_batch_ndim + i]
    X_expanded = X.expand(*expand_shape).contiguous()

    
    flat_size = int(torch.tensor(Y_perm.shape[model_batch_ndim:-1]).prod())
    X_flat = X_expanded.reshape(*X_expanded.shape[:model_batch_ndim], flat_size, X_expanded.shape[-1]).clone()
    Y_flat = Y_perm.reshape(*Y_perm.shape[:model_batch_ndim], flat_size, Y_perm.shape[-1]).clone()

    return super().condition_on_observations(X_flat, Y_flat, **kwargs)