Possible addition of diffusion kernel over discrete/categorical input spaces

[aryandeshwal]

Hi, I am Aryan. I am a PhD student at Washington State University. I like GPyTorch a lot and regularly use it in my own research on combinatorial Bayesian optimization. Thanks for the great library! I think adding diffusion kernel for discrete/categorical input spaces will be a nice addition to the library. It is a very useful (extension of RBF to discrete spaces) kernel for Bayesian optimization over discrete/combinatorial spaces (potentially good for BoTorch).

I am providing my implementation of the diffusion kernel in GPyTorch format below, hoping it will be useful. It is a batch-compatible implementation for the ARD version of the kernel that supports arbitrary number of categories in each dimension.

from typing import Optional
import torch
from gpytorch.constraints import Interval, Positive
from gpytorch.priors import Prior
from gpytorch.kernels import Kernel

class DiffusionKernel(Kernel):
    r"""
        Computes diffusion kernel over discrete spaces with arbitrary number of categories. 
        Input type: n dimensional discrete input with c_i possible categories/choices for each dimension i 
        As an example, binary {0,1} combinatorial space corresponds to c_i = 2 for each dimension i
        References:
        - https://www.ml.cmu.edu/research/dap-papers/kondor-diffusion-kernels.pdf (Section 4.4)
        - https://arxiv.org/abs/1902.00448
        - https://arxiv.org/abs/2012.07762
        
        Args:
        :attr:`categories`(tensor, list):
            array with number of possible categories in each dimension            
    """
    has_lengthscale = True
    def __init__(self, categories, **kwargs):
        if categories is None:
            raise RunTimeError("Can't create a diffusion kernel without number of categories. Please define them!")
        super().__init__(**kwargs)
        self.cats = categories

    def forward(self, x1, x2, diag: Optional[bool] = False, last_dim_is_batch: Optional[bool] = False, **params):
        if last_dim_is_batch:
            x1 = x1.transpose(-1, -2).unsqueeze(-1)
            x2 = x2.transpose(-1, -2).unsqueeze(-1)        

        if diag:
            res = 1.
            for i in range(x1.shape[-1]):
                res *= ((1 - torch.exp(-self.lengthscale[..., i] * self.cats[i]))/(1 + (self.cats[i] - 1) * torch.exp(-self.lengthscale[..., i]*self.cats[i]))).unsqueeze(-1) ** ((x1[..., i] != x2[..., i])[:, 0, ...])
            return res

        res = 1.
        for i in range(x1.shape[-1]): 
            res *= ((1 - torch.exp(-self.lengthscale[..., i] * self.cats[i]))/(1 + (self.cats[i] - 1) * torch.exp(-self.lengthscale[..., i]*self.cats[i]))).unsqueeze(-1) ** ((x1[..., i].unsqueeze(-2)[..., None] != x2[..., i].unsqueeze(-2))[0, ...])
        return res

Thanks and Happy new year!

[Balandat]

@dme65

[gpleiss]

This seems like a good kernel to have! I'd be open to a PR. Make sure to add unit tests (see https://github.com/cornellius-gp/gpytorch/blob/master/test/kernels/test_rq_kernel.py for an example).

[dme65]

Agreed that this would be a great kernel to have!

It looks like this implementation works directly with the graph Laplacian which has cubic complexity w.r.t. the number of vertices in the graph. It would be great to extend this (maybe at a later point) to work with the graph Cartesian product and exploit the resulting Kronecker structure, just like for COMBO here. This is particularly interesting for Bayesian optimization with categorical variables. Additionally, using a horseshoe prior on the inverse lengthscale makes a lot of sense to support as well.

[martinitus]

Hi Folks,
I played around with that kernel a little, ultimately failing using it with mixed continuous/categorical features in an ApproximateGP / CholeskyVariationalDistribution / VariationalStrategy setup, where the inducing points only "live" in the continuous feature subspace.

I am thinking of it as "the categorical features only contribute to the distance between data points".

I didn't find any active_dim kwargs for the variational distribution or strategy that would allow me to configure them for the respective subspace and I start wondering whether my plans make any sense at all or would be possible with gpytorch without massive API changes on the variational classes.

Any hints about where to look or read are highly appreciated ;-)

Regards,
Martin

[gpleiss]

@martinitus sorry for the slow reply!

Is it necessary that this kernel works for variational models? You are right - our inducing point methods are not currently well-suited to handle categorical inputs.

One option would be to used fixed inducing points (set learn_inducing_points=False in VariationalStrategy). This is sub-optimal, but would work for now.

[martinitus]

Hi there - thanks for the reply 👍🏼 and no worries, I got lots of other things to work on already :-)

For now, I am using onehot encoding for the categoricals, which obviously isn't optimal, but also seems to work reasonably well for out problem.

The idea with learn_inducing_points=False is actually great. If i find time I will definitively try that - especially because the continuous feature space is rather low dimensional, and i could essentially span a uniform grid over and using this grid as inducing points.

[cottrell]

Did this get anywhere? Trying to understand more generally best practices for mixed features (float and categorical/int). I haven't found examples yet.

And I mean in situations with high-arity categoricals where naively one-hotting is not a good approach.