Changes to the default Multi-Task parametrization

[hvarfner]

We are introducing a new multitask parametrization in BoTorch. This update aims to improve the modeling of correlations between tasks and address some apparent shortcomings of the existing parametrization. For the MultiTaskGP models, we make the following changes to the defaults:

• ConstantMean --> MultiTaskMean
• IndexKernel --> PositiveIndexKernel

However, these can easily be swapped out for the older versions by providing a custom mean_module or task_covar_module to the model.

Motivation
The previous multitask parametrization in BoTorch had limitations in capturing the correlations between tasks, often leading to suboptimal performance. In general, especially in the very low data regime, it is extremely hard to infer cross-task correlations. By reducing flexibility of the model and being optimistic about the correlations between tasks being positive we can improve modeling and optimization performance. With very little data and unconstrained task correlations, both perfect negative and positive correlation are plausible. The MLL objective is rewarded simply for correlating the data while minimizing the residuals, making negative task correlations a fairly common occurrence.

Implementation Details
MultiTaskMean infers one mean function per task. For any type of constant offset between tasks, or with the use of StratifiedStandardize, this parametrization is paramount. In low-data regimes, StratifiedStandardize's mean parameters will likely differ even if the tasks in the MTGP are identical, making MultiTaskMean particularly important when this transform is used.
PositiveIndexKernel: Inferring task correlation is difficult. When data is scarce and non-identical across tasks (e.g. 3 random points per task on 2 Branin tasks), task correlation inference can be very spotty, to the point where -1 and +1 is inferred with almost equal frequency in extreme cases (see image below).

In addition to more accurate task correlation inference (assuming task correlation is positive), the change to PositiveIndexKernel is made to align better with people's intuition - that source tasks are indeed positively correlated with the target task.

Results
In the figure below, we show the inferred task correlation of GPs with different parametrizations, ranging from the current default MultiTaskGP (red), the addition of either component in isolation (yellow, green) to the new default (blue).

For all three functions (Branin (2D), Forrester (1D), Hartmann (3D), we use the same function as the source as for the target, so we try to infer the correlation between Branin and itself. For the target task, we use three data points, and two source tasks (total of 16 data points) with three different degrees of data imbalance (1:1 means 8 points from each, 3:1 means 12 points from task A, 4 from task B) and 7:1 means 14 from task A and 2 from task B). All data is drawn uniformly at random. We evaluate the rank correlation on a held-out set on the target task, i.e. the ability of the model to predict the rankings of unseen data points. Both the old and the partial new parametrizations perform rather poorly on this task, whereas the new one successfully infers the rankings of unseen data about as well as if the source and target task were all slotted into one shared SingleTaskGP (black dashed line). This is (arguably) the performance that is expected in this setting, as source data and target data are perfectly correlated and on the same scale.

Potential future changes:
The SaasFullyBayesianMultiTaskGP will likely undergo similar changes. Moreover, the transforms applied to multi-task modeling (specifically StratifiedStandardize) are currently under evaluation.

Potential downsides - when you’d not want to use this
The change from IndexKernel from PositiveIndexKernel is made with optimization in mind. Specifically, this means typical BO scenarios where data is scarce and tasks are assumed to be correlated, but there isn’t enough evidence to robustly estimate the true correlation. In these settings, constraining correlations to be positive can improve modeling stability and optimization performance.
However, PositiveIndexKernel may not be ideal for regression tasks. If you have access to large amounts of data and your goal is to determine whether tasks are correlated (positively or negatively), the standard IndexKernel is preferable.