ARD Periodic Kernel Implementation and Related Questions

[m-julian]

Hello!

I tried implementing a periodic kernel that supports ARD, where every input dimension can have its own period length and lengthscale. Everything seems to be working fine and the tests produce the same results (see below). I was wondering if something like this would be useful to have in gpytorch, as the current PeriodicKernel does not fully support multi-dimensional input (does not support different lengthscale for each dimension). I do have a couple of questions below, and would appreciate any comments/feedback on the implementation. Let me know if I should make a separate post for the questions because this post did get quite long.

Questions:

1. Would you recommend having ARD kernels implemented, which support different parameter values for each active dimensions or would it be better to just multiply (or add) 1D-kernels together, where each 1D-kernel has a different active dimension? I would be able to multiply many PeriodicKernel instances and change the active dimension for each one, but then it becomes a bit tedious to set parameters for separate dimensions (I didn't really find an easy way to set hyperparameter values for a more complex kernel composition).
2. For the batch dimension test, I am currently getting a RuntimeError: The expected shape of the kernel was torch.Size([3, 2, 5, 5]), but got torch.Size([3, 5, 5]). This is likely a bug in GPyTorch. in gpytorch/lazy/lazy_evaluated_kernel_tensor.py. Obviously caused by my implementation, but I can remove this by using gpytroch.settings.debug(False). Could I change my implementation to go around this without having to turn off debugging?
3. In the ARD implemntation, I temporarily have to store a n_dim x n_data_x1, n_data_x2 tensor because I cannot divide by the lengthscales before calculating the distance. Could KeOps be used here to not have to store the whole tensor at once?

ARD-Periodic Kernel:

import math
import torch
from gpytorch.constraints import Positive
from gpytorch.kernels import Kernel

class ARDPeriodicKernel(Kernel):

    has_lengthscale = True

    def __init__(self, period_length_prior=None, period_length_constraint=None, **kwargs):
        super().__init__(**kwargs)
        if period_length_constraint is None:
            period_length_constraint = Positive()

        period_length_num_dims = 1 if self.ard_num_dims is None else self.ard_num_dims
        self.register_parameter(
            name="raw_period_length", parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, 1, period_length_num_dims))
        )

        if period_length_prior is not None:
            self.register_prior(
                "period_length_prior",
                period_length_prior,
                lambda m: m.period_length,
                lambda m, v: m._set_period_length(v),
            )

        self.register_constraint("raw_period_length", period_length_constraint)

    @property
    def period_length(self):
        return self.raw_period_length_constraint.transform(self.raw_period_length)

    @period_length.setter
    def period_length(self, value):
        self._set_period_length(value)

    def _set_period_length(self, value):
        if not torch.is_tensor(value):
            value = torch.as_tensor(value).to(self.raw_period_length)
        self.initialize(raw_period_length=self.raw_period_length_constraint.inverse_transform(value))

    def forward(self, x1, x2, **kwargs):

        # active dimensions are handled by __call__ in Kernel class
        # divide by period length and multiply by pi beforehand
        x1_ = x1.div(self.period_length).mul(math.pi)
        x2_ = x2.div(self.period_length).mul(math.pi)

        if kwargs.get("last_dim_is_batch", False):
            x1_ = x1_.transpose(-1, -2).unsqueeze(-1)
            x2_ = x2_.transpose(-1, -2).unsqueeze(-2)
            # shape is n_batch, n_dim, n_data_x1, n_data_x2
            diff = x1_ - x2_
            # shape is n_batch, n_dim, 1, 1
            lengthscales = self.lengthscale.transpose(-2,-1).unsqueeze(-1)
        else:
            # The shape is n_dim, n_data_x1, n_data_x2
            diff = x1_.T.unsqueeze(-1) - x2_.T.unsqueeze(-2)
            # shape is n_dim, 1, 1
            lengthscales = self.lengthscale.reshape(-1, 1, 1)

        # get a n_batch, n_data_x1, n_data_x2 or n_data_x1, n_data_x2 if not batch
        res = torch.sin(diff).pow(2).div(lengthscales).sum(dim=-3).mul(-2.0).exp_()

        # TODO: add diagonal
        return res

Here are the tests I have ran so far to check if the covariance matrices match.

from gpytorch.kernels import PeriodicKernel
from ard_periodic_kernel import ARDPeriodicKernel
import torch
import numpy as np
import gpytorch

def euclidean_distance(x1: torch.Tensor, x2: torch.Tensor):
    """ Calculates Euclidean distance between two tensors."""
    
    if x1.ndim == 1:
        x1 = x1.reshape(-1,1)
        x2 = x2.reshape(-1,1)

    return torch.cdist(x1, x2)


def periodic_kernel(euc_dist: torch.Tensor, period: float, lengthscale: float):
    """ Calculates covariance matrix using periodic kernel, given euclidean distance matrix,
    period length, and lengthscale. Lengthscale is already squared."""

    # lengthscale is already squared as in gpytorch
    return torch.exp(-2 * torch.pow(torch.sin(np.pi*torch.abs(euc_dist) / period), 2) / lengthscale)

# 1D INPUTS test
x = torch.rand(5)
x_prime = torch.rand(5)
lengthscale = torch.tensor(2.0)
period_length = torch.tensor(1.0)
distance = euclidean_distance(x, x_prime)
cov_matrix = periodic_kernel(distance, period_length, lengthscale)
gpytorch_periodic_kernel = PeriodicKernel()
gpytorch_periodic_kernel.lengthscale = lengthscale
gpytorch_periodic_kernel.period_length = period_length
gpytorch_cov_matrix = gpytorch_periodic_kernel(x, x_prime).evaluate()
ard_per_1d = ARDPeriodicKernel()
ard_per_1d.lengthscale = lengthscale
ard_per_1d.period_length = period_length
ard_cov_1d = ard_per_1d(x, x_prime).evaluate()

print(f"1D inputs have same covariance matrices (own and gpytorch): {torch.allclose(cov_matrix, gpytorch_cov_matrix)}")
print(f"1D inputs have same covariance matrices (gpytorch and ard-periodic): {torch.allclose(gpytorch_cov_matrix, ard_cov_1d)}")

# 2D INPUTS test
# separate lengthscale and period length for each dimension
x_2d = torch.rand(5,2)
x_prime_2d = torch.rand(5,2)
lengthscale_2d = torch.tensor([2.0, 1.0])
period_length_2d = torch.tensor([0.5, 1.0])
euclidean_distance_1st_dim = euclidean_distance(x_2d[:, 0], x_prime_2d[:, 0])
euclidean_distance_2nd_dim = euclidean_distance(x_2d[:, 1], x_prime_2d[:, 1])
cov_matrix_2d = periodic_kernel(euclidean_distance_1st_dim, period_length_2d[0], lengthscale_2d[0]) \
    * periodic_kernel(euclidean_distance_2nd_dim, period_length_2d[1], lengthscale_2d[1])
gpytorch_periodic_kernel_2d_1 = PeriodicKernel()
gpytorch_periodic_kernel_2d_2 = PeriodicKernel()
gpytorch_periodic_kernel_2d_1.lengthscale = lengthscale_2d[0]
gpytorch_periodic_kernel_2d_2.lengthscale = lengthscale_2d[1]
gpytorch_periodic_kernel_2d_1.period_length = period_length_2d[0]
gpytorch_periodic_kernel_2d_2.period_length = period_length_2d[1]

# multiply 1D gpytorch covariance matrices for each dimension
gpytorch_cov_matrix_2d = gpytorch_periodic_kernel_2d_1(x_2d[:, 0], x_prime_2d[:, 0]).evaluate() \
    * gpytorch_periodic_kernel_2d_2(x_2d[:, 1], x_prime_2d[:, 1]).evaluate()

ard_periodic_kernel = ARDPeriodicKernel(ard_num_dims=2, active_dims=[0,1])
ard_periodic_kernel.lengthscale = lengthscale_2d
ard_periodic_kernel.period_length = period_length_2d
ard_periodic_covariance_2d = ard_periodic_kernel(x_2d, x_prime_2d).evaluate()

print(f"2D inputs have same covariance matrices (own and gpytorch 1D Periodic): {torch.allclose(cov_matrix_2d, gpytorch_cov_matrix_2d)}")
print(f"2D inputs have same covariance matrices (gpytorch 1D Periodic and ARD-Periodic): {torch.allclose(gpytorch_cov_matrix_2d, ard_periodic_covariance_2d)}")

# Batch dimensions test
# have to set this, otherwise getting
# RuntimeError: The expected shape of the kernel was torch.Size([3, 2, 5, 5]), but got torch.Size([3, 5, 5]). This is likely a bug in GPyTorch.
with gpytorch.settings.debug(False):

    # (n_batch, n_train, n_feats)
    batch_x = torch.rand(3,5,2)
    batch_x_prime = torch.rand(3,5,2)
    ard_periodic_kernel_batch = ARDPeriodicKernel(batch_shape=torch.Size([3]), ard_num_dims=2, active_dims=[0,1])
    batch_cov_matrix = ard_periodic_kernel_batch(batch_x, batch_x_prime, last_dim_is_batch=True).evaluate()

    gpytorch_periodic_1d = PeriodicKernel()
    covar_1st_batch = gpytorch_periodic_1d(batch_x[0], batch_x_prime[0]).evaluate()
    covar_2nd_batch = gpytorch_periodic_1d(batch_x[1], batch_x_prime[1]).evaluate()
    covar_3rd_batch = gpytorch_periodic_1d(batch_x[2], batch_x_prime[2]).evaluate()

print(f"First batch equal to gpytorch periodic kernel: {torch.allclose(batch_cov_matrix[0], covar_1st_batch)}")
print(f"Second batch equal to gpytorch periodic kernel: {torch.allclose(batch_cov_matrix[1], covar_2nd_batch)}")
print(f"Third batch equal to gpytorch periodic kernel: {torch.allclose(batch_cov_matrix[2], covar_3rd_batch)}")

Any feedback is much appreciated.

[gpleiss]

@m-julian as of the 1.7.0 release, the gpytorch periodic kernel supports ARD. See #1919.