Could I expect div-free and curl-free kernel functions to be add?

[team0ne1]

I am a beginner in gpytorch, and now need to implement a div-free RBF kernel in my project.

Ref paper: Learning divergence-free and curl-free vector field - Macedo and Castro.
Inspired by discussions #1504, I wrote a kernel as below:

class DivFreeKernel(gpytorch.kernels.Kernel):

    def __init__(self, num_tasks=2, length_prior=None, sigma_prior=None, length_constraint=None, sigma_constraint=None , **kwargs):

        super(CurlFreeKernel, self).__init__(**kwargs)

        # register the raw parameter
        self.register_parameter(
            name='raw_length', parameter=torch.nn.Parameter(torch.ones(*self.batch_shape, 1, 1))
        )
        self.register_parameter(
            name='raw_sigma', parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, 1, 1))
        )

        # set the parameter constraint to be positive, when nothing is specified
        if length_constraint is None:
            length_constraint = Positive()

        # register the constraint
        self.register_constraint("raw_length", length_constraint)

        # set the parameter prior, see
        # https://docs.gpytorch.ai/en/latest/module.html#gpytorch.Module.register_prior
        if length_prior is not None:
            self.register_prior(
                "length_prior",
                length_prior,
                lambda m: m.length,
                lambda m, v : m._set_length(v),
            )

        if sigma_constraint is None:
            sigma_constraint = Positive()

        # register the constraint
        self.register_constraint("raw_sigma", sigma_constraint)

        self.num_tasks = num_tasks

    # now set up the 'actual' paramter
    @property
    def length(self):
        # print('length:')
        # when accessing the parameter, apply the constraint transform
        return self.raw_length_constraint.transform(self.raw_length)

    @length.setter
    def length(self, value):
        return self._set_length(value)

    def _set_length(self, value):
        if not torch.is_tensor(value):
            value = torch.as_tensor(value).to(self.raw_length)
        # when setting the paramater, transform the actual value to a raw one by applying the inverse transform
        self.initialize(raw_length=self.raw_length_constraint.inverse_transform(value))
    
    @property
    def sigma(self):
        # when accessing the parameter, apply the constraint transform
        return self.raw_sigma_constraint.transform(self.raw_sigma)

    @sigma.setter
    def sigma(self, value):
        return self._set_sigma(value)

    def _set_sigma(self, value):
        if not torch.is_tensor(value):
            value = torch.as_tensor(value).to(self.raw_sigma)
        # when setting the paramater, transform the actual value to a raw one by applying the inverse transform
        self.initialize(raw_sigma=self.raw_sigma_constraint.inverse_transform(value))

    def forward(self, x1, x2, **params):
        
        n = x1.size(-2)
        m = x2.size(-2)
        d = self.num_tasks
        
        # Diff tensor: ... x n x m x d
        diff = x1.unsqueeze(-2) - x2.unsqueeze(-3)
        diff = diff.unsqueeze(-1).div(self.length)

        # Exp term
        norm = self.covar_dist(x1, x2, square_dist=True)  # ... x n x m
        exp_part = (-0.5 * norm / self.length.pow(2)).exp()
        exp_part = exp_part[..., :, None, :, None]  # ... x n x 1 x m x 1

        # Outer product term
        outer = diff * diff.mT                                       # ... x n x m x d x d
        outer = outer.transpose(-2, -3).contiguous()  # ... x n x d x m x d

        # Combine terms
        eye = torch.eye(d, dtype=x1.dtype, device=x1.device)  # ... x d x d
        eye = eye[..., None, :, None, :]  # ... x 1 x d x 1 x d
        kernel = (outer + ((d - 1) -  exp_part_n) * eye)
        kernel = kernel * exp_part
        kernel = torch.flatten(kernel, start_dim=-2, end_dim=-1)
        kernel = torch.flatten(kernel, start_dim=-3, end_dim=-2)
        kernel = kernel * self.sigma / self.length.pow(2)
        # print('kernel', kernel, kernel.size())
        return kernel

    def num_outputs_per_input(self, x1, x2):
        """
        Given `n` data points `x1` and `m` datapoints `x2`, this multitask
        kernel returns an `(n*num_tasks) x (m*num_tasks)` covariance matrix.
        """
        return x1.size(-1)

create a multitask model

class MultitaskGPModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super(MultitaskGPModel, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.MultitaskMean(
            gpytorch.means.ConstantMean(), num_tasks=2
        )  
        self.covar_module = DivFreeKernel(
            num_tasks=2,            
        )

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        
        return gpytorch.distributions.MultitaskMultivariateNormal(mean_x, covar_x)

likelihood = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=2, rank=0)
model = MultitaskGPModel(train_X, train_y, likelihood)

I try to train the model with a 2d-vector-field dataset, it seems that the DivFreeKernel work properly. However, an error occurred during I increasing the number of training_iterations:

Traceback (most recent call last):
  File "H:\Code\vector-field\gp_vector_field\src\main.py", line 125, in <module>
    loss = -mll(output, train_y)
  File "E:\anaconda3\envs\pytorch-gpu\lib\site-packages\gpytorch\module.py", line 30, in __call__
    outputs = self.forward(*inputs, **kwargs)
  File "E:\anaconda3\envs\pytorch-gpu\lib\site-packages\gpytorch\mlls\exact_marginal_log_likelihood.py", line 62, in forward
    res = output.log_prob(target)
  File "E:\anaconda3\envs\pytorch-gpu\lib\site-packages\gpytorch\distributions\multitask_multivariate_normal.py", line 211, in log_prob
    return super().log_prob(value.view(*value.shape[:-2], -1))
  File "E:\anaconda3\envs\pytorch-gpu\lib\site-packages\gpytorch\distributions\multivariate_normal.py", line 170, in log_prob
    inv_quad, logdet = covar.inv_quad_logdet(inv_quad_rhs=diff.unsqueeze(-1), logdet=True)
  File "E:\anaconda3\envs\pytorch-gpu\lib\site-packages\gpytorch\lazy\lazy_tensor.py", line 1291, in inv_quad_logdet
    cholesky = CholLazyTensor(TriangularLazyTensor(self.cholesky()))
  File "E:\anaconda3\envs\pytorch-gpu\lib\site-packages\gpytorch\lazy\lazy_tensor.py", line 1004, in cholesky    
    chol = self._cholesky(upper=False)
  File "E:\anaconda3\envs\pytorch-gpu\lib\site-packages\gpytorch\utils\memoize.py", line 59, in g
    return _add_to_cache(self, cache_name, method(self, *args, **kwargs), *args, kwargs_pkl=kwargs_pkl)
  File "E:\anaconda3\envs\pytorch-gpu\lib\site-packages\gpytorch\lazy\lazy_tensor.py", line 435, in _cholesky    
    cholesky = psd_safe_cholesky(evaluated_mat, upper=upper).contiguous()
  File "E:\anaconda3\envs\pytorch-gpu\lib\site-packages\gpytorch\utils\cholesky.py", line 68, in psd_safe_cholesky
    L = _psd_safe_cholesky(A, out=out, jitter=jitter, max_tries=max_tries)
  File "E:\anaconda3\envs\pytorch-gpu\lib\site-packages\gpytorch\utils\cholesky.py", line 49, in _psd_safe_cholesky
    raise NotPSDError(f"Matrix not positive definite after repeatedly adding jitter up to {jitter_new:.1e}.")    
gpytorch.utils.errors.NotPSDError: Matrix not positive definite after repeatedly adding jitter up to 1.0e-04.

According to issue #1479, I add a line with gpytorch.settings.cholesky_jitter(1e-1): and it doesn't show errors anymore.

for i in range(training_iterations):
    optimizer.zero_grad()
    output = model(train_X)
    with gpytorch.settings.cholesky_jitter(1e-1):
        loss = -mll(output, train_y)
    loss.backward()
    print('Iter %d/%d - Loss: %.3f' % (i + 1, training_iterations, loss.item()))
    optimizer.step()

Q1: Is this a good way to build a div-free RBF multitask kernel?

Expect to expand the kernel

According to this paper Kernels for Vector-Valued Functions: A Review, we can implement a Hessian matrix to a scalar kernel function to get a vector-value kernel with zero-divergence constraint.

Q2: Is it possible to create a div-free multitask kernel that converts scalar kernel to vector-value/multitask kernel?
look like:

self.covar_module = DivFreeKernel(
    gpytorch.kernels.RBFKernel()
    num_tasks=2,            
)

[gpleiss]

According to issue #1479, I add a line with gpytorch.settings.cholesky_jitter(1e-1): and it doesn't show errors anymore.

This makes me believe that there is a bug in your code. The "jitter" is supposed to be for kernels that are positive definite in exact arithmetic, but not positive definite in floating point arithmetic due to round-off errors. A jitter of 0.1 makes me think that your kernel is somehow not positive definite, which is likely due to a bug in your code.

I haven't had the time to go through your code to debug it, but what I would recommend is following the custom kernels tutorial. Try implementing your kernel without any hyperparameters first (i.e. set \sigma = 1, or to some fixed constant) and make sure that you aren't getting any NotPSDErrors. Then you can try adding all of the hyperparameter/prior/constraint complications and making sure it is still PSD.