SingleTaskGP - what does subset_output do?

[bchen0]

I am trying to understand what subset_output does in the SingleTaskGP class.

In models.gp_regression.py, lines 205-216, there is this code snippet. I've tried looking at the subset_output method in the parent class BatchedMultiOutputGPyTorchModel but I am not sure what it is doing.

If I want to use a non-default covar_module, self._subset_batch_dict is not set. Should I be manually setting self._subset_batch_dict too if I pass in a covar_module?

if covar_module is None:
  covar_module = get_covar_module_with_dim_scaled_prior(
      ard_num_dims=transformed_X.shape[-1],
      batch_shape=self._aug_batch_shape,
  )
  # Used for subsetting along the output dimension. See Model.subset_output.
  self._subset_batch_dict = {
      "mean_module.raw_constant": -1,
      "covar_module.raw_lengthscale": -3,
  }
  if train_Yvar is None:
      self._subset_batch_dict["likelihood.noise_covar.raw_noise"] = -2

Thanks!

[saitcakmak]

Hi @bchen0. subset_output is used to extract a GP that models a subset of the outputs from a GP that models more outputs. If I have a SingleTaskGP with 3 outputs, I could use it to extract a SingleTaskGP with 2 outputs etc.

As you noticed, it only works right if self._subset_batch_dict is compatible with the submodules of the model (it needs to point to the output dimensions in the model state dict). In general, I wouldn't recommend using subset_output with custom covar modules. It is often better to simply train the model with the outputs we need, rather than training the model with more outputs and subsetting it later.

[bchen0]

Hi @saitcakmak, thanks for the quick response. The rationale behind subset_output makes sense.

Could you clarify a little more about this statement though: "it only works right if self._subset_batch_dict is compatible with the submodules of the model (it needs to point to the output dimensions in the model state dict)".

For instance, mean_module.raw_constant has the index -1 - does this mean the model would index along the last index to extract a SingleTaskGP with fewer outputs? But if so, why is it parameter dependent (e.g., why -3 for covar_module.raw_lengthscale)