Resolving Must provide inverse transform to be able to sample from prior

[Hrovatin]

I tried to use IndexKernel similar as in MultiTaskGP (here and for more details on prior setting here).
I used IndexKernel as follows:

sd_prior = GammaPrior(1.0, 0.15)
sd_prior._event_shape = torch.Size([1]) # If I put here n_tasks as in MultiTaskGP the shapes mismatch
task_covar_prior = LKJCovariancePrior(n=n_tasks, eta=eta, sd_prior=sd_prior)
task_covar_module = gpytorch.kernels.IndexKernel(
    num_tasks=n_tasks,
    active_dims=ask_idx,
    rank=n_tasks,  
    prior=task_covar_prior,
)

Which is passed as product kernel to SingleTaskGP resulting in model:

SingleTaskGP(
  (likelihood): GaussianLikelihood(
    (noise_covar): HomoskedasticNoise(
      (noise_prior): GammaPrior()
      (raw_noise_constraint): GreaterThan(1.000E-04)
    )
  )
  (mean_module): ConstantMean()
  (covar_module): ProductKernel(
    (kernels): ModuleList(
      (0): ScaleKernel(
        (base_kernel): MaternKernel(
          (lengthscale_prior): GammaPrior()
          (raw_lengthscale_constraint): Positive()
        )
        (outputscale_prior): GammaPrior()
        (raw_outputscale_constraint): Positive()
      )
      (1): IndexKernel(
        (IndexKernelPrior): LKJCovariancePrior()
        (raw_var_constraint): Positive()
      )
    )
  )
  (outcome_transform): Standardize()
  (input_transform): Normalize()
)

Then I fit the model with botorch.fit.fit_gpytorch_mll

Depending on eta the sklearn fitting fails with

OptimizationWarning: `scipy_minimize` terminated with status 3, displaying original message from `scipy.optimize.minimize`: ABNORMAL_TERMINATION_IN_LNSRCH

which I will inspect myself. However, when botorch then attempts to sample new initial hyperparameter values in sample_all_priors I get the error:

Resolving Must provide inverse transform to be able to sample from prior

Any ideas what I must change to make this sampling work?

I already tried updating to botorch 0.11.3 as there was report of similar error which should have been resolved in 0.11.3
gpytorch==1.12
botorch==0.11.3

[saitcakmak]

Hi @Hrovatin. I believe this is an issue with the LKJ prior you're using. See #1860

[Hrovatin]

I resolved it as follows:

Define closure that is able to use samples from LKJCovariancePrior to set parameters of IndexKernel. I did not find an exact way to solve the equation for decomposing sampled cov matrix into kernel parameters, so I used gradient-based solver.

def set_index_kernel(index_kernel,cov):

    # Find kernel parameters from cov
    
    var_transform=torch.nn.functional.softplus
    cov=torch.detach(cov)
    def loss_fn(x, y):
        cov_pred= x @ x.transpose(-1,-2) + torch.diag_embed(var_transform(y))
        return torch.square(cov - cov_pred).sum()
    
    # Initial guess for the variables as tensors
    x = torch.randn_like(index_kernel.covar_factor,requires_grad=True)
    y =torch.randn_like(index_kernel.var,requires_grad=True)
    
    # Define an optimizer
    optimizer = torch.optim.Adam([x, y], lr=0.01)
    
    # Iteratively minimize the norm of the equations
    for _ in range(10000):
        optimizer.zero_grad()
        loss = loss_fn(x,y)
        loss.backward()
        optimizer.step()

    # Set parameters
    index_kernel.var = var_transform(torch.detach(y)) # Must be positive
    index_kernel.covar_factor.data=torch.detach(x)

Set prior with this closure on IndexKernel

sd_prior = GammaPrior(1.0, index_prior_rate)
sd_prior._event_shape = torch.Size([1])
task_covar_prior = LKJCovariancePrior(n=n_tasks, eta=index_prior_eta, sd_prior=sd_prior)
task_covar_prior._event_shape = torch.Size([1])
task_covar_prior._batch_shape = torch.Size([1])
task_covar_module = gpytorch.kernels.IndexKernel(
    num_tasks=n_tasks,
    active_dims=task_idx,
    rank=n_tasks,  
)
task_covar_module.register_prior(
    name='cov_prior',
    prior=task_covar_prior,
    param_or_closure=lambda m: m._eval_covar_matrix(),
    setting_closure= set_index_kernel,
)
   

I also needed to increase n fitting attempts as more re-initializations of priors were needed for successful fit.

        botorch.fit.fit_gpytorch_mll(
            mll,
            max_attempts=50,
            optimizer=optimizer
        )