Permutationally Invariant Matern Kernel

[jennevv]

I've been working on implementing a permutationally invariant Matern kernel in GPytorch for an ExactGP model. My approach involves modifying the Matern kernel from the GPytorch source and incorporating permutational symmetry operations as described in the following paper:

Bartók and Csányi - 2015 - Gaussian approximation potentials A brief tutoria.pdf

Kernel Details:

Input: Tensor with each row containing xyz coordinates of all atoms.
Transformation: Converts xyz coordinates to internuclear distances, scaled by the exponential of the negative inverse distance.

I want to perform the transformation of the xyz coordinates to the distances inside the kernel class so that I have easy access to gradients and hessians w.r.t. the xyz inputs during simulations.

Here is the implementation of the kernel function:

class PermutationallyInvariantMatern(gpytorch.kernels.Kernel):
    has_lengthscale = True

    def __init__(self, n_atoms, idx_equiv_atoms, ard=False, nu=2.5, **kwargs):
        if nu not in {0.5, 1.5, 2.5}:
            raise RuntimeError("nu expected to be 0.5, 1.5, or 2.5")
        super(PermutationallyInvariantMatern, self).__init__(**kwargs)
        self.nu = nu
        self.idx_equiv_atoms = idx_equiv_atoms

        dims_cart = torch.arange(0, n_atoms * 3).reshape(n_atoms, 3)
        self.dims = dims_cart

        self.permutations = self.generate_permutations()

        if ard is True:
            self.ard_num_dims = n_atoms * (n_atoms - 1) // 2

    def generate_permutations(self):
        perms = torch.tensor([list(p) for p in itertools.permutations(self.idx_equiv_atoms)])
        return perms

    def forward(self, x1, x2, diag=False, **params):
        k_sum = 0
        init_perm = self.permutations[0]

        for p in self.permutations:
            x1_perm = x1.clone()
            x2_perm = x2.clone()
            x1_perm[:, self.dims[init_perm, :].flatten()] = x1[:, self.dims[p, :].flatten()]
            x2_perm[:, self.dims[init_perm, :].flatten()] = x2[:, self.dims[p, :].flatten()]

            # Transform xyz coordinates to internuclear distances
            x1_interdist = xyz_to_expinvinterdist_torch(x1)
            x2_perm_interdist = xyz_to_expinvinterdist_torch(x2_perm)

            mean = x1_interdist.mean(dim=-2, keepdim=True)

            x1_interdist_ = (x1_interdist - mean) / (self.lengthscale)
            x2_perm_interdist_ = (x2_perm_interdist - mean) / (self.lengthscale)

            distance = self.covar_dist(x1_interdist_, x2_perm_interdist_, diag=diag, **params)

            exp_component = torch.exp(-math.sqrt(self.nu * 2) * distance)

            if self.nu == 0.5:
                constant_component = 1
            elif self.nu == 1.5:
                constant_component = (math.sqrt(3) * distance).add(1)
            elif self.nu == 2.5:
                constant_component = (math.sqrt(5) * distance).add(1).add(5.0 / 3.0 * distance**2)

            k_sum += constant_component * exp_component

        return k_sum

I face two issues:

1. Normalizing the kernel as suggested in the paper above results in a kernel matrix of all ones.
2. When training a model with ARD using the non-normalized kernel, I encounter a NotPSDError after about 15 iterations, indicating the matrix is not positive definite. This issue does not occur with a regular Matern kernel and an already transformed dataset.

Is there an issue with the kernel class that makes it not possible to use ARD? Would the normalization (if implemented correctly) be able to help?