[Question] Extract one GP from Hadamard Multitask

Hi team 👋

Context

I’m using a Multitask Hadamard GP for Regression as exactly as described in the tutorial Hadamard Multitask GP Regression with 2 tasks.

After training, the tutorial proceeds to evaluate both tasks separately with

with torch.no_grad(), gpytorch.settings.fast_pred_var():
      observed_pred_y1 = likelihood(model(test_x, test_i_task1))
      observed_pred_y2 = likelihood(model(test_x, test_i_task2))

Problem

The variables observed_pred_y1 and observed_pred_y2 are of the type gpytorch.distributions.MultivariateNormal. I would like to extract from model one GP from one of the tasks, say task 1, called model1 that is of type gpytorch.models.ExactGP instead of a Multivariate distribution. This might sound weird but after training, I need to pass a GP to other custom software and a distribution won’t work. It is necessary to use Multitask because both tasks are related and need to be fitted together.

(failed) Attempts

I’ve tried to create a class similar to the classic ExactGPModel like here but initializing the modules according to the tasK 2 of the Multitask model

class ExactGPModel_Task1(gpytorch.models.ExactGP):
    def __init__(self, x_train, y_train, likelihood, model):
        super().__init__(x_train, y_train, likelihood)
        self.mean_module = model.mean_module 
        self.covar_module = model.covar_module
        
    def forward(self, x, model):
        mean_x = self.mean_module(x)
        this_test_i_task1 = torch.full((x.shape[0],1), dtype=torch.long, fill_value=0)
        pred_y1 = model(x, this_test_i_task1)
        covar_x = pred_y1._covar
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
 

model1 = ExactGPModel_Task1(train_x1, train_y1, likelihood, model)

and when evaluating

model1.eval()
likelihood.eval()
model1(test_x, model)

(EDIT) I run into AttributeError: 'MultitaskGPModel' object has no attribute 'ndimension' at gpytorch/module.py in __getattr__(self, name) starting on this line.

Here is the full stack

--------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
~/Documents/PROJECTS/AL_Exotics_notebooks/.env/lib/python3.8/site-packages/gpytorch/module.py in __getattr__(self, name)
    432             try:
--> 433                 return super().__getattribute__(name)
    434             except AttributeError:

AttributeError: 'MultitaskGPModel' object has no attribute 'ndimension'

During handling of the above exception, another exception occurred:

AttributeError                            Traceback (most recent call last)
/tmp/ipykernel_3770492/3093790522.py in <module>
      2 likelihood.eval()
      3 
----> 4 model1(test_x, model)

~/Documents/PROJECTS/AL_Exotics_notebooks/.env/lib/python3.8/site-packages/gpytorch/models/exact_gp.py in __call__(self, *args, **kwargs)
    243     def __call__(self, *args, **kwargs):
    244         train_inputs = list(self.train_inputs) if self.train_inputs is not None else []
--> 245         inputs = [i.unsqueeze(-1) if i.ndimension() == 1 else i for i in args]
    246 
    247         # Training mode: optimizing

~/Documents/PROJECTS/AL_Exotics_notebooks/.env/lib/python3.8/site-packages/gpytorch/models/exact_gp.py in <listcomp>(.0)
    243     def __call__(self, *args, **kwargs):
    244         train_inputs = list(self.train_inputs) if self.train_inputs is not None else []
--> 245         inputs = [i.unsqueeze(-1) if i.ndimension() == 1 else i for i in args]
    246 
    247         # Training mode: optimizing

~/Documents/PROJECTS/AL_Exotics_notebooks/.env/lib/python3.8/site-packages/gpytorch/module.py in __getattr__(self, name)
    433                 return super().__getattribute__(name)
    434             except AttributeError:
--> 435                 raise e
    436 
    437 

~/Documents/PROJECTS/AL_Exotics_notebooks/.env/lib/python3.8/site-packages/gpytorch/module.py in __getattr__(self, name)
    428     def __getattr__(self, name):
    429         try:
--> 430             return super().__getattr__(name)
    431         except AttributeError as e:
    432             try:

~/Documents/PROJECTS/AL_Exotics_notebooks/.env/lib/python3.8/site-packages/torch/nn/modules/module.py in __getattr__(self, name)
   1128             if name in modules:
   1129                 return modules[name]
-> 1130         raise AttributeError("'{}' object has no attribute '{}'".format(
   1131             type(self).__name__, name))
   1132 

AttributeError: 'MultitaskGPModel' object has no attribute 'ndimension'

Do you have any advice?

Thank you so much! 😃

[irinaespejo]

Hi @wjmaddox that worked, thank you! I had to pass the multitask model via kwargs so it wasn’t a major refactor.

I guess I’d be curious if there are some more general patterns in terms of extracting individual tasks that emerge in different setups that we could generalize and provide as utilities as part of botorch.

Personally, I am going to use extracting a task from multitask model quite a lot. I guess one pattern, our pattern, is to use a multitask GP to incorporate correlations of weaker but faster experiments into a high fidelity experiment. And then, use that task to do optimal experiment design. I’d be down to help with this if you decide to support this feature.

Thanks again!

[irinaespejo]

Hi @wjmaddox, thanks! That would almost work but downstream there are calls like gp.train_inputs, gp.train_targets, gp.parameters or gp.get_fantasy_model. It is possible to refactor the downstream code so that it works with a Multivariate distribution but I’d rather leave it as last option and try first to accommodate the requirement of input a GP.