[Feature Request] MultiTaskGP with multiple task features

[AdrianSosic]

🚀 Feature Request

Currently, the MultiTaskGP expects exactly one task dimension, specified via the task_feature argument. It would be great to extend the model such that it can handle several task dimensions.

Pitch

Describe the solution you'd like
I see no reason why the internal logic could not be straightforwardly generalized to handle multiple task dimensions, since the corresponding feature enters multiplicatively via an IndexKernel. That means, additional task dimensions could be treated in the exact same way, as additional factors to the kernel product.

Describe alternatives you've considered
Allowing to specify multiple task dimensions would enable learning models over structured tasks spaces – as opposed to the alternative where several such dimensions would be flattened into a single one before entering the MultiTaskGP.

[saitcakmak]

Hi @AdrianSosic. Your suggestion makes sense to me. It is a straightforward extension to support more broad use cases. I'd be curious to see how this works in practice.

Implementation-wise, we just need to make sure to keep the model constructor backwards compatible while adding support for this. If you want to put up a PR for this, I'd be happy to review it.

[AdrianSosic]

Hi @saitcakmak, thanks for the feedback. I can try to free some time for this in the next few weeks 👍🏼

Regarding backwards compatibility, perhaps it's worth discussing what the implications would be upfront. If you want full compatibility without introducing a new class, I think the we need to talk about two separate parts:

1. For the constructor, applying a 1:1 extension would mean that the task-related arguments need to be turned into corresponding unions, that is:

   • task_feature: int becomes task_feature: Union[int, Sequence[int]]
   • task_covar_prior: Optional[Prior] = None becomes task_covar_prior: Union[Prior, Sequence[Optional[Prior]], None] = None
   • and so on

   As you can see, this has the downsides that the types get rather complex + the parameter names are a bit misleading (e.g. singular task_feature for several ints).

2. In addition, full compatibility would also mean that public attributes remain consistent with the previous version. That is, something like self.num_tasks needs to carry a single integer in one case but a list of integers in the second case. Similar for other attributes.

If that is what you want, I'm happy to set it up this way. Let me know what you think.

[saitcakmak]

Thanks for the detailed thoughts on this. BoTorch is technically still at 0.x.x development stage, and does not have a defined public API, so we don't need to have 100% backwards compatibility. Though, we still try our best to avoid BC breaking changes.

task_feature: int becomes task_feature: Union[int, Sequence[int]]

I think this one is ok. We can convert it to List[int] internally while accepting both formats. I suspect most common use cases will have a single task feature.

task_covar_prior: Optional[Prior] = None becomes task_covar_prior: Union[Prior, Sequence[Optional[Prior]], None] = None

This one I do not like. I think task_covar_prior was a convenient input when this model was first implemented, but I find it restrictive. I'd advocate for (soft) deprecating task_covar_prior input here and instead introducing a task_covar_module: Module input in its place. If the user would like to customize the task_covar_module, they can provide a fully defined module rather than being restricted to a product of IndexKernels.

cc @Balandat, @sdaulton on potential changes to MTGP