[Feature Request] Non-default categorical kernel in mixedGP

[jduerholt]

🚀 Feature Request

I looked into the definition of MixedSingleTaskGP, so far one is only able to specify a cont_kernel_factory to use a non-default kernel for the continuous domains. From my point of view it would make sense to add als cat_kernel_factory to specify non-default categorical kernels.

If I understood it correctly, for the categorical kernel, ARD is used always. Having the possibility to specify a cat_kernel_factory one would by able to specify also a categorical kernel without ARD, or a totally different one.

A totally different categorical kernel could be for example a molecular kernel from the gauche paper https://realworldml.github.io/files/cr/paper8.pdf. This would allow for a mixed optimization over a categorical molecular species plus some continuous parameters by just using the already implemented MixedSingleTaskGP`.

What do you think?

Pitch

Describe the solution you'd like

Add a cat_kernel_factory argument to MixedSingleTaskGP to be able to use a non-default categorical kernel.

Are you willing to open a pull request? (See CONTRIBUTING)

Yes.

[jduerholt]

@Balandat @saitcakmak any thoughts on this? ;)

[Balandat]

Yeah I think this is sensible. The constructor of MixedSingleTaskGP is kind of a beast already so we should think about if/how we can simplify that a bit. Maybe factoring out the kernel assembly logic into a "private" helper function would get us at least part of the way there.

cc @Ryan-Rhys re GAUCHE kernels :)

[jduerholt]

The constructor of MixedSingleTaskGP is definitely a beast, I also run into other problems there today https://github.com/pytorch/botorch/pull/1568

So we first brainstorm how to refactor it before I implement the categorical_kernel_factory?

[Balandat]

So we first brainstorm how to refactor it before I implement the categorical_kernel_factory?

I don't think we need to do that - I think it's fine to implement this but I would recommend to compartmentalize the code and stick a bunch of the logic from https://github.com/pytorch/botorch/blob/main/botorch/models/gp_regression_mixed.py#L116-L182 and the new logic into one or more helper functions to keep things structured and readable.

There are some interesting deeper questions about input transforms that your other PR raises, but I think these are somewhat orthogonal to the proposal here.

[jduerholt]

Fine for me! I try to put a PR together for this over the course of the next week!

[Ryan-Rhys]

Happy to review/help if needed!