Arc Kernel for Conditional Spaces

🚀 Feature Request

Introduction

I am working on Bayesian optimization with gaussian processes for hyperparameter tuning of ML models. I am indirect user of your library (through Ax and Botorch. Many ML models have conditional configuration spaces. I understand that gaussian processes do not support directly conditional spaces. However, this article Raiders of the Lost Architecture:Kernels for Bayesian Optimization in Conditional Parameter Spaces defines a kernel which would be useful in such spaces.

Motivation

Would be useful to implement this kernel for use in upper libraries. The usage would allow the optimization of conditional configuration spaces by simply implementing this kernel (if I have understood correctly the article).

In the case where conditionality is related to hierarchy, this would solve issues such as this one. Also would be a nice try for a flaw consistently assigned to GPs.

Pitch

In a rough and simplistic way, the kernel consists in a cylindrical embedding in a standard RBF kernel. And then a product of the kernels for each dimension to build the final kernel. I leave the images of the descriptions found in the Github of the project.

I have made myself an implementation:

class ArcKernel(Kernel):

    def __init__(self, num_parameters, **kwargs):
        super(ArcKernel, self).__init__(has_lengthscale=True, **kwargs)

        self.register_parameter(
            name="raw_angle",
            parameter=torch.nn.Parameter(torch.randn(1, 1, self.ard_num_dims))
            )
        angle_constraint = Interval(0, 1)
        self.register_constraint("raw_angle", angle_constraint)

        self.register_parameter(
            name="raw_radius",
            parameter=torch.nn.Parameter(torch.randn(1, 1, self.ard_num_dims))
            )
        radius_constraint = Positive()
        self.register_constraint("raw_radius", radius_constraint)

        self.num_parameters = num_parameters

    def embedding(self, x):
        x_ = x.div(self.lengthscale)
        x_s = self.raw_radius*torch.sin(pi*self.raw_angle*x_)
        x_c = self.raw_radius*torch.sin(pi*self.raw_angle*x_)
        x_ = torch.cat((x_s, x_c), dim=-1).squeeze(0)
        return x_

    def forward(self, x1, x2, diag=False, **params):

        x1_, x2_ = self.embedding(x1), self.embedding(x2)
        return self.covar_dist(x1_, x2_, square_dist=True, diag=diag,
                               dist_postprocess_func=postprocess_rbf,
                               postprocess=True, **params)

And for calling it, it would be a composition of StructuredProductKernel, ScaleKernel, RBFKernel and the newly build ArcKernel.

I can make my self a pull request if the request is accepted. Also, I have not the enough knowledge to ensure that this kernel should work exactly as the explained in the paper. I have made my implementation but have not tested it in any experiment. if desired, I could do also the test bu would need orientation.

Thank you very much in advance.

P.S: you can find info summarized in the repository of the original authors, under the folder latex. I do attach also images from that source.

[Balandat]

Happy to help w/ debugging one the botorch end. Once your PR here is up, can you share the full code so I can reproduce?

[Balandat]

This is great, @BCJuan. Please let us know over at botorch if you need any help with hooking this up with the acquisition functions.